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
19
        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
6
            .await
187
4
            .unwrap();
188
51
        for (id, category) in &operation.initial_data.categories {
189
51
            let mut args = PgArguments::default();
190
51
            args.reserve(2, 0);
191
51
            args.add(*id as i32);
192
51
            args.add(&category.name);
193
77
            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
338
        for (&id, product) in &operation.initial_data.products {
205
338
            let mut args = PgArguments::default();
206
338
            args.reserve(2, 0);
207
338
            args.add(id as i32);
208
338
            args.add(&product.name);
209
619
            tx.execute(insert_product.query_with(args)).await.unwrap();
210
1119
            for &category_id in &product.category_ids {
211
781
                let mut args = PgArguments::default();
212
781
                args.reserve(2, 0);
213
781
                args.add(id as i32);
214
781
                args.add(category_id as i32);
215
781
                tx.execute(insert_product_category.query_with(args))
216
1449
                    .await
217
781
                    .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
8
            .await
224
4
            .unwrap();
225
797
        for (id, customer) in &operation.initial_data.customers {
226
797
            let mut args = PgArguments::default();
227
797
            args.reserve(9, 0);
228
797
            args.add(*id as i32);
229
797
            args.add(&customer.name);
230
797
            args.add(&customer.email);
231
797
            args.add(&customer.address);
232
797
            args.add(&customer.city);
233
797
            args.add(&customer.region);
234
797
            args.add(&customer.country);
235
797
            args.add(&customer.postal_code);
236
797
            args.add(&customer.phone);
237
1575
            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
8
            .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
4
            .await
249
4
            .unwrap();
250
845
        for (&id, order) in &operation.initial_data.orders {
251
845
            let mut args = PgArguments::default();
252
845
            args.reserve(2, 0);
253
845
            args.add(id as i32);
254
845
            args.add(order.customer_id as i32);
255
1611
            tx.execute(insert_order.query_with(args)).await.unwrap();
256
6896
            for &product_id in &order.product_ids {
257
6051
                let mut args = PgArguments::default();
258
6051
                args.reserve(2, 0);
259
6051
                args.add(id as i32);
260
6051
                args.add(product_id as i32);
261
6051
                tx.execute(insert_order_product.query_with(args))
262
11047
                    .await
263
6051
                    .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
8
            .await
270
4
            .unwrap();
271
334
        for review in &operation.initial_data.reviews {
272
334
            let mut args = PgArguments::default();
273
334
            args.reserve(4, 0);
274
334
            args.add(review.product_id as i32);
275
334
            args.add(review.customer_id as i32);
276
334
            args.add(review.rating as i32);
277
334
            args.add(&review.review);
278
614
            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
462
    async fn operate(
301
462
        &mut self,
302
462
        _operation: &CreateCart,
303
462
        _results: &[OperationResult<u32>],
304
462
        measurements: &Measurements,
305
462
    ) -> OperationResult<u32> {
306
462
        let measurement = measurements.begin("postgresql", Metric::CreateCart);
307
704
        let mut conn = self.sqlite.acquire().await.unwrap();
308
462
        let mut tx = conn.begin().await.unwrap();
309
462
        let statement = tx
310
462
            .prepare("insert into commerce_bench.carts (customer_id) values (null) returning id")
311
38
            .await
312
462
            .unwrap();
313

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

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

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

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

            
354
1865
        tx.execute(statement.query_with(args)).await.unwrap();
355
2346
        tx.commit().await.unwrap();
356
1184
        measurement.finish();
357
1184

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

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

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

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

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

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

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

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

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

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

            
552
132
        measurement.finish();
553
132

            
554
132
        OperationResult::Ok
555
264
    }
556
}
557

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

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

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

            
602
106
        OperationResult::Ok
603
212
    }
604
}