1
use bonsaidb::core::actionable::async_trait;
2
use futures::StreamExt;
3
use sqlx::postgres::PgArguments;
4
use sqlx::{Arguments, Connection, Executor, PgPool, Row, Statement};
5

            
6
use crate::execute::{Backend, BackendOperator, Measurements, Metric, Operator};
7
use crate::model::Product;
8
use crate::plan::{
9
    AddProductToCart, Checkout, CreateCart, FindProduct, Load, LookupProduct, OperationResult,
10
    ReviewProduct,
11
};
12

            
13
pub struct Postgres {
14
    pool: PgPool,
15
}
16

            
17
#[async_trait]
18
impl Backend for Postgres {
19
    type Config = String;
20
    type Operator = PostgresOperator;
21

            
22
    fn label(&self) -> &'static str {
23
        "postgresql"
24
    }
25

            
26
4
    async fn new(url: Self::Config) -> Self {
27
20
        let pool = PgPool::connect(&url).await.unwrap();
28

            
29
4
        let mut conn = pool.acquire().await.unwrap();
30
4
        conn.execute(r#"DROP SCHEMA IF EXISTS commerce_bench CASCADE"#)
31
7
            .await
32
4
            .unwrap();
33
4
        conn.execute(r#"CREATE SCHEMA commerce_bench"#)
34
8
            .await
35
4
            .unwrap();
36
4
        conn.execute("SET search_path='commerce_bench';")
37
8
            .await
38
4
            .unwrap();
39
4
        conn.execute(
40
4
            r#"CREATE TABLE customers (
41
4
            id SERIAL PRIMARY KEY,
42
4
            name TEXT,
43
4
            email TEXT,
44
4
            address TEXT,
45
4
            city TEXT,
46
4
            region TEXT,
47
4
            country TEXT,
48
4
            postal_code TEXT,
49
4
            phone TEXT
50
4
        )"#,
51
4
        )
52
8
        .await
53
4
        .unwrap();
54
4
        conn.execute(
55
4
            r#"CREATE TABLE products (
56
4
                    id SERIAL PRIMARY KEY,
57
4
                    name TEXT
58
4
                )"#,
59
4
        )
60
8
        .await
61
4
        .unwrap();
62
4
        conn.execute(r#"CREATE INDEX products_by_name ON products(name)"#)
63
8
            .await
64
4
            .unwrap();
65
4
        conn.execute(
66
4
            r#"CREATE TABLE product_reviews (
67
4
                product_id INTEGER NOT NULL,-- REFERENCES products(id),
68
4
                customer_id INTEGER NOT NULL,-- REFERENCES customers(id),
69
4
                rating INTEGER NOT NULL,
70
4
                review TEXT
71
4
            )"#,
72
4
        )
73
8
        .await
74
4
        .unwrap();
75
4
        conn.execute(r#"CREATE INDEX product_reviews_by_product ON product_reviews(product_id)"#)
76
8
            .await
77
4
            .unwrap();
78
4
        conn.execute(r#"CREATE UNIQUE INDEX product_reviews_by_customer ON product_reviews(customer_id, product_id)"#)
79
8
            .await
80
4
            .unwrap();
81
4
        conn.execute(
82
4
            r#"CREATE MATERIALIZED VIEW
83
4
                    product_ratings
84
4
                AS
85
4
                    SELECT
86
4
                        product_id,
87
4
                        sum(rating)::int as total_rating,
88
4
                        count(rating)::int as ratings
89
4
                    FROM
90
4
                        product_reviews
91
4
                    GROUP BY product_id
92
4
            "#,
93
4
        )
94
8
        .await
95
4
        .unwrap();
96
4
        conn.execute(
97
4
            r#"CREATE TABLE categories (
98
4
                    id SERIAL PRIMARY KEY,
99
4
                    name TEXT
100
4
                )"#,
101
4
        )
102
8
        .await
103
4
        .unwrap();
104
4
        conn.execute(
105
4
            r#"CREATE TABLE product_categories (
106
4
                    product_id INTEGER,-- REFERENCES products(id),
107
4
                    category_id INTEGER-- REFERENCES categories(id)
108
4
                )"#,
109
4
        )
110
8
        .await
111
4
        .unwrap();
112
4
        conn.execute(
113
4
            r#"CREATE INDEX product_categories_by_product ON product_categories(product_id)"#,
114
4
        )
115
8
        .await
116
4
        .unwrap();
117
4
        conn.execute(
118
4
            r#"CREATE INDEX product_categories_by_category ON product_categories(category_id)"#,
119
4
        )
120
8
        .await
121
4
        .unwrap();
122
4
        conn.execute(
123
4
            r#"CREATE TABLE orders (
124
4
                    id SERIAL PRIMARY KEY,
125
4
                    customer_id INTEGER -- REFERENCES customers(id)
126
4
                )"#,
127
4
        )
128
8
        .await
129
4
        .unwrap();
130
4
        conn.execute(
131
4
            r#"CREATE TABLE order_products (
132
4
                    order_id INTEGER NOT NULL,-- REFERENCES orders(id),
133
4
                    product_id INTEGER NOT NULL -- REFERENCES products(id)
134
4
                )"#,
135
4
        )
136
8
        .await
137
4
        .unwrap();
138
4
        conn.execute(
139
4
            r#"CREATE TABLE carts (
140
4
            id SERIAL PRIMARY KEY,
141
4
            customer_id INTEGER
142
4
        )"#,
143
4
        )
144
8
        .await
145
4
        .unwrap();
146
4
        conn.execute(
147
4
            r#"CREATE TABLE cart_products (
148
4
                cart_id INTEGER,
149
4
                product_id INTEGER
150
4
            )"#,
151
4
        )
152
8
        .await
153
4
        .unwrap();
154
4

            
155
4
        Self { pool }
156
8
    }
157

            
158
12
    async fn new_operator_async(&self) -> Self::Operator {
159
12
        PostgresOperator {
160
12
            sqlite: self.pool.clone(),
161
12
        }
162
12
    }
163
}
164

            
165
pub struct PostgresOperator {
166
    sqlite: PgPool,
167
}
168

            
169
impl BackendOperator for PostgresOperator {
170
    type Id = u32;
171
}
172

            
173
#[async_trait]
174
impl Operator<Load, u32> for PostgresOperator {
175
4
    async fn operate(
176
4
        &mut self,
177
4
        operation: &Load,
178
4
        _results: &[OperationResult<u32>],
179
4
        measurements: &Measurements,
180
4
    ) -> OperationResult<u32> {
181
4
        let measurement = measurements.begin("postgresql", Metric::Load);
182
11
        let mut conn = self.sqlite.acquire().await.unwrap();
183
4
        let mut tx = conn.begin().await.unwrap();
184
4
        let insert_category = tx
185
4
            .prepare("INSERT INTO commerce_bench.categories (id, name) VALUES ($1, $2)")
186
8
            .await
187
4
            .unwrap();
188
60
        for (id, category) in &operation.initial_data.categories {
189
60
            let mut args = PgArguments::default();
190
60
            args.reserve(2, 0);
191
60
            args.add(*id as i32);
192
60
            args.add(&category.name);
193
97
            tx.execute(insert_category.query_with(args)).await.unwrap();
194
        }
195

            
196
4
        let insert_product = tx
197
4
            .prepare("INSERT INTO commerce_bench.products (id, name) VALUES ($1, $2)")
198
6
            .await
199
4
            .unwrap();
200
4
        let insert_product_category = tx
201
4
            .prepare("INSERT INTO commerce_bench.product_categories (product_id, category_id) VALUES ($1, $2)")
202
3
            .await
203
4
            .unwrap();
204
370
        for (&id, product) in &operation.initial_data.products {
205
370
            let mut args = PgArguments::default();
206
370
            args.reserve(2, 0);
207
370
            args.add(id as i32);
208
370
            args.add(&product.name);
209
467
            tx.execute(insert_product.query_with(args)).await.unwrap();
210
1483
            for &category_id in &product.category_ids {
211
1113
                let mut args = PgArguments::default();
212
1113
                args.reserve(2, 0);
213
1113
                args.add(id as i32);
214
1113
                args.add(category_id as i32);
215
1113
                tx.execute(insert_product_category.query_with(args))
216
1732
                    .await
217
1113
                    .unwrap();
218
            }
219
        }
220

            
221
4
        let insert_customer = tx
222
4
            .prepare("INSERT INTO commerce_bench.customers (id, name, email, address, city, region, country, postal_code, phone) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)")
223
6
            .await
224
4
            .unwrap();
225
557
        for (id, customer) in &operation.initial_data.customers {
226
557
            let mut args = PgArguments::default();
227
557
            args.reserve(9, 0);
228
557
            args.add(*id as i32);
229
557
            args.add(&customer.name);
230
557
            args.add(&customer.email);
231
557
            args.add(&customer.address);
232
557
            args.add(&customer.city);
233
557
            args.add(&customer.region);
234
557
            args.add(&customer.country);
235
557
            args.add(&customer.postal_code);
236
557
            args.add(&customer.phone);
237
911
            tx.execute(insert_customer.query_with(args)).await.unwrap();
238
        }
239

            
240
4
        let insert_order = tx
241
4
            .prepare("INSERT INTO commerce_bench.orders (id, customer_id) VALUES ($1, $2)")
242
6
            .await
243
4
            .unwrap();
244
4
        let insert_order_product = tx
245
4
            .prepare(
246
4
                "INSERT INTO commerce_bench.order_products (order_id, product_id) VALUES ($1, $2)",
247
4
            )
248
3
            .await
249
4
            .unwrap();
250
1263
        for (&id, order) in &operation.initial_data.orders {
251
1263
            let mut args = PgArguments::default();
252
1263
            args.reserve(2, 0);
253
1263
            args.add(id as i32);
254
1263
            args.add(order.customer_id as i32);
255
2098
            tx.execute(insert_order.query_with(args)).await.unwrap();
256
12942
            for &product_id in &order.product_ids {
257
11679
                let mut args = PgArguments::default();
258
11679
                args.reserve(2, 0);
259
11679
                args.add(id as i32);
260
11679
                args.add(product_id as i32);
261
11679
                tx.execute(insert_order_product.query_with(args))
262
19047
                    .await
263
11679
                    .unwrap();
264
            }
265
        }
266

            
267
4
        let insert_review = tx
268
4
            .prepare("INSERT INTO commerce_bench.product_reviews (product_id, customer_id, rating, review) VALUES ($1, $2, $3, $4)")
269
6
            .await
270
4
            .unwrap();
271
446
        for review in &operation.initial_data.reviews {
272
446
            let mut args = PgArguments::default();
273
446
            args.reserve(4, 0);
274
446
            args.add(review.product_id as i32);
275
446
            args.add(review.customer_id as i32);
276
446
            args.add(review.rating as i32);
277
446
            args.add(&review.review);
278
778
            tx.execute(insert_review.query_with(args)).await.unwrap();
279
        }
280
4
        tx.execute(
281
4
            "SELECT setval('commerce_bench.orders_id_seq', COALESCE((SELECT MAX(id)+1 FROM commerce_bench.orders), 1), false)",
282
4
        )
283
6
        .await
284
4
        .unwrap();
285
4

            
286
8
        tx.commit().await.unwrap();
287
4
        // Make sure all ratings show up in the view.
288
4
        conn.execute("REFRESH MATERIALIZED VIEW commerce_bench.product_ratings")
289
8
            .await
290
4
            .unwrap();
291
4
        // This makes a significant difference.
292
8
        conn.execute("ANALYZE").await.unwrap();
293
4
        measurement.finish();
294
4

            
295
4
        OperationResult::Ok
296
8
    }
297
}
298
#[async_trait]
299
impl Operator<CreateCart, u32> for PostgresOperator {
300
461
    async fn operate(
301
461
        &mut self,
302
461
        _operation: &CreateCart,
303
461
        _results: &[OperationResult<u32>],
304
461
        measurements: &Measurements,
305
461
    ) -> OperationResult<u32> {
306
461
        let measurement = measurements.begin("postgresql", Metric::CreateCart);
307
695
        let mut conn = self.sqlite.acquire().await.unwrap();
308
461
        let mut tx = conn.begin().await.unwrap();
309
461
        let statement = tx
310
461
            .prepare("insert into commerce_bench.carts (customer_id) values (null) returning id")
311
36
            .await
312
461
            .unwrap();
313

            
314
686
        let result = tx.fetch_one(statement.query()).await.unwrap();
315
910
        tx.commit().await.unwrap();
316
461
        let id: i32 = result.get(0);
317
461
        measurement.finish();
318
461

            
319
461
        OperationResult::Cart { id: id as u32 }
320
922
    }
321
}
322
#[async_trait]
323
impl Operator<AddProductToCart, u32> for PostgresOperator {
324
1188
    async fn operate(
325
1188
        &mut self,
326
1188
        operation: &AddProductToCart,
327
1188
        results: &[OperationResult<u32>],
328
1188
        measurements: &Measurements,
329
1188
    ) -> OperationResult<u32> {
330
1188
        let cart = match &results[operation.cart.0] {
331
1188
            OperationResult::Cart { id } => *id,
332
            _ => unreachable!("Invalid operation result"),
333
        };
334
1188
        let product = match &results[operation.product.0] {
335
1188
            OperationResult::Product { id, .. } => *id,
336
            _ => unreachable!("Invalid operation result"),
337
        };
338

            
339
1188
        let measurement = measurements.begin("postgresql", Metric::AddProductToCart);
340
1822
        let mut conn = self.sqlite.acquire().await.unwrap();
341
1188
        let mut tx = conn.begin().await.unwrap();
342
1188
        let statement = tx
343
1188
            .prepare(
344
1188
                "insert into commerce_bench.cart_products (cart_id, product_id) values ($1, $2)",
345
1188
            )
346
40
            .await
347
1188
            .unwrap();
348
1188

            
349
1188
        let mut args = PgArguments::default();
350
1188
        args.reserve(2, 0);
351
1188
        args.add(cart as i32);
352
1188
        args.add(product as i32);
353
1188

            
354
1737
        tx.execute(statement.query_with(args)).await.unwrap();
355
2318
        tx.commit().await.unwrap();
356
1188
        measurement.finish();
357
1188

            
358
1188
        OperationResult::CartProduct { id: product }
359
2376
    }
360
}
361
#[async_trait]
362
impl Operator<FindProduct, u32> for PostgresOperator {
363
2363
    async fn operate(
364
2363
        &mut self,
365
2363
        operation: &FindProduct,
366
2363
        _results: &[OperationResult<u32>],
367
2363
        measurements: &Measurements,
368
2363
    ) -> OperationResult<u32> {
369
2363
        let measurement = measurements.begin("postgresql", Metric::FindProduct);
370
3758
        let mut conn = self.sqlite.acquire().await.unwrap();
371
2363
        let statement = conn
372
2363
            .prepare(
373
2363
                r#"
374
2363
                SELECT
375
2363
                    id,
376
2363
                    name,
377
2363
                    category_id,
378
2363
                    commerce_bench.product_ratings.total_rating as "total_rating: Option<i32>",
379
2363
                    commerce_bench.product_ratings.ratings as "ratings: Option<i32>"
380
2363
                FROM
381
2363
                    commerce_bench.products
382
2363
                LEFT OUTER JOIN commerce_bench.product_categories ON
383
2363
                    commerce_bench.product_categories.product_id = id
384
2363
                LEFT OUTER JOIN commerce_bench.product_ratings ON
385
2363
                    commerce_bench.product_ratings.product_id = id
386
2363
                WHERE name = $1
387
2363
                GROUP BY id, name, category_id, commerce_bench.product_ratings.total_rating, commerce_bench.product_ratings.ratings
388
2363
            "#,
389
2363
            )
390
27
            .await
391
2363
            .unwrap();
392
2363

            
393
2363
        let mut args = PgArguments::default();
394
2363
        args.reserve(1, 0);
395
2363
        args.add(&operation.name);
396
2363

            
397
2363
        let mut results = conn.fetch(statement.query_with(args));
398
2363
        let mut id: Option<i32> = None;
399
2363
        let mut name = None;
400
2363
        let mut category_ids = Vec::new();
401
2363
        let mut total_rating: Option<i32> = None;
402
2363
        let mut rating_count: Option<i32> = None;
403
15401
        while let Some(row) = results.next().await {
404
13038
            let row = row.unwrap();
405
13038
            id = Some(row.get(0));
406
13038
            name = Some(row.get(1));
407
13038
            total_rating = row.get(2);
408
13038
            rating_count = row.get(3);
409
13038
            if let Some(category_id) = row.get::<Option<i32>, _>(2) {
410
12760
                category_ids.push(category_id as u32);
411
12760
            }
412
        }
413
2363
        let rating_count = rating_count.unwrap_or_default();
414
2363
        let total_rating = total_rating.unwrap_or_default();
415
2363
        measurement.finish();
416
2363
        OperationResult::Product {
417
2363
            id: id.unwrap() as u32,
418
2363
            product: Product {
419
2363
                name: name.unwrap(),
420
2363
                category_ids,
421
2363
            },
422
2363
            rating: if rating_count > 0 {
423
1791
                Some(total_rating as f32 / rating_count as f32)
424
            } else {
425
572
                None
426
            },
427
        }
428
4726
    }
429
}
430
#[async_trait]
431
impl Operator<LookupProduct, u32> for PostgresOperator {
432
2300
    async fn operate(
433
2300
        &mut self,
434
2300
        operation: &LookupProduct,
435
2300
        _results: &[OperationResult<u32>],
436
2300
        measurements: &Measurements,
437
2300
    ) -> OperationResult<u32> {
438
2300
        let measurement = measurements.begin("postgresql", Metric::LookupProduct);
439
3656
        let mut conn = self.sqlite.acquire().await.unwrap();
440
2300
        let statement = conn
441
2300
            .prepare(
442
2300
                r#"
443
2300
                    SELECT
444
2300
                        id,
445
2300
                        name,
446
2300
                        category_id,
447
2300
                        commerce_bench.product_ratings.total_rating as "total_rating: Option<i32>",
448
2300
                        commerce_bench.product_ratings.ratings as "ratings: Option<i32>"
449
2300
                    FROM
450
2300
                        commerce_bench.products
451
2300
                    LEFT OUTER JOIN commerce_bench.product_categories ON
452
2300
                        commerce_bench.product_categories.product_id = id
453
2300
                    LEFT OUTER JOIN commerce_bench.product_ratings ON
454
2300
                        commerce_bench.product_ratings.product_id = id
455
2300
                    WHERE id = $1
456
2300
                    GROUP BY id, name, category_id, commerce_bench.product_ratings.total_rating, commerce_bench.product_ratings.ratings
457
2300
                "#,
458
2300
            )
459
22
            .await
460
2300
            .unwrap();
461
2300

            
462
2300
        let mut args = PgArguments::default();
463
2300
        args.reserve(1, 0);
464
2300
        args.add(operation.id as i32);
465
2300

            
466
2300
        let mut results = conn.fetch(statement.query_with(args));
467
2300
        let mut id: Option<i32> = None;
468
2300
        let mut name = None;
469
2300
        let mut category_ids = Vec::new();
470
2300
        let mut total_rating: Option<i32> = None;
471
2300
        let mut rating_count: Option<i32> = None;
472
14871
        while let Some(row) = results.next().await {
473
12571
            let row = row.unwrap();
474
12571
            id = Some(row.get(0));
475
12571
            name = Some(row.get(1));
476
12571
            total_rating = row.get(2);
477
12571
            rating_count = row.get(3);
478
12571
            if let Some(category_id) = row.get::<Option<i32>, _>(2) {
479
12306
                category_ids.push(category_id as u32);
480
12306
            }
481
        }
482
2300
        let rating_count = rating_count.unwrap_or_default();
483
2300
        let total_rating = total_rating.unwrap_or_default();
484
2300

            
485
2300
        measurement.finish();
486
2300
        OperationResult::Product {
487
2300
            id: id.unwrap() as u32,
488
2300
            product: Product {
489
2300
                name: name.unwrap(),
490
2300
                category_ids,
491
2300
            },
492
2300
            rating: if rating_count > 0 {
493
1799
                Some(total_rating as f32 / rating_count as f32)
494
            } else {
495
501
                None
496
            },
497
        }
498
4600
    }
499
}
500

            
501
#[async_trait]
502
impl Operator<Checkout, u32> for PostgresOperator {
503
121
    async fn operate(
504
121
        &mut self,
505
121
        operation: &Checkout,
506
121
        results: &[OperationResult<u32>],
507
121
        measurements: &Measurements,
508
121
    ) -> OperationResult<u32> {
509
121
        let cart = match &results[operation.cart.0] {
510
121
            OperationResult::Cart { id } => *id as i32,
511
            _ => unreachable!("Invalid operation result"),
512
        };
513

            
514
121
        let measurement = measurements.begin("postgresql", Metric::Checkout);
515
177
        let mut conn = self.sqlite.acquire().await.unwrap();
516
121
        let mut tx = conn.begin().await.unwrap();
517
        // Create a new order
518
121
        let statement = tx
519
121
            .prepare(r#"INSERT INTO commerce_bench.orders (customer_id) VALUES ($1) RETURNING ID"#)
520
30
            .await
521
121
            .unwrap();
522
121
        let mut args = PgArguments::default();
523
121
        args.reserve(1, 0);
524
121
        args.add(operation.customer_id as i32);
525
154
        let result = tx.fetch_one(statement.query_with(args)).await.unwrap();
526
121
        let order_id: i32 = result.get(0);
527

            
528
121
        let statement = tx
529
121
            .prepare(r#"
530
121
                WITH products_in_cart AS (
531
121
                    DELETE FROM commerce_bench.cart_products WHERE cart_id = $1 RETURNING $2::int as order_id, product_id
532
121
                )
533
121
                INSERT INTO commerce_bench.order_products (order_id, product_id) SELECT * from products_in_cart;"#)
534
30
            .await
535
121
            .unwrap();
536
121
        let mut args = PgArguments::default();
537
121
        args.reserve(2, 0);
538
121
        args.add(cart);
539
121
        args.add(order_id);
540
161
        tx.execute(statement.query_with(args)).await.unwrap();
541

            
542
121
        let statement = tx
543
121
            .prepare(r#"DELETE FROM commerce_bench.carts WHERE id = $1"#)
544
38
            .await
545
121
            .unwrap();
546
121
        let mut args = PgArguments::default();
547
121
        args.reserve(1, 0);
548
121
        args.add(cart);
549
157
        tx.execute(statement.query_with(args)).await.unwrap();
550
231
        tx.commit().await.unwrap();
551
121

            
552
121
        measurement.finish();
553
121

            
554
121
        OperationResult::Ok
555
242
    }
556
}
557

            
558
#[async_trait]
559
impl Operator<ReviewProduct, u32> for PostgresOperator {
560
99
    async fn operate(
561
99
        &mut self,
562
99
        operation: &ReviewProduct,
563
99
        results: &[OperationResult<u32>],
564
99
        measurements: &Measurements,
565
99
    ) -> OperationResult<u32> {
566
99
        let product = match &results[operation.product_id.0] {
567
            OperationResult::Product { id, .. } => *id,
568
99
            OperationResult::CartProduct { id, .. } => *id,
569
            _ => unreachable!("Invalid operation result"),
570
        };
571
99
        let measurement = measurements.begin("postgresql", Metric::RateProduct);
572
151
        let mut conn = self.sqlite.acquire().await.unwrap();
573
99
        let mut tx = conn.begin().await.unwrap();
574
99
        let statement = tx
575
99
            .prepare(
576
99
                r#"INSERT INTO commerce_bench.product_reviews (
577
99
                        product_id,
578
99
                        customer_id,
579
99
                        rating,
580
99
                        review)
581
99
                    VALUES ($1, $2, $3, $4)
582
99
                    ON CONFLICT (customer_id, product_id) DO UPDATE SET rating = $3, review = $4"#,
583
99
            )
584
32
            .await
585
99
            .unwrap();
586
99

            
587
99
        let mut args = PgArguments::default();
588
99
        args.reserve(4, 0);
589
99
        args.add(product as i32);
590
99
        args.add(operation.customer_id as i32);
591
99
        args.add(operation.rating as i32);
592
99
        args.add(&operation.review);
593
99

            
594
141
        tx.execute(statement.query_with(args)).await.unwrap();
595
196
        tx.commit().await.unwrap();
596
99
        // Make this rating show up
597
99
        conn.execute("REFRESH MATERIALIZED VIEW commerce_bench.product_ratings")
598
198
            .await
599
99
            .unwrap();
600
99
        measurement.finish();
601
99

            
602
99
        OperationResult::Ok
603
198
    }
604
}