risingwave_e2e_extended_mode_test/
test.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::time::Duration;
16
17use anyhow::anyhow;
18use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
19use pg_interval::Interval;
20use rust_decimal::Decimal;
21use tokio::select;
22use tokio_postgres::types::Type;
23use tokio_postgres::{Client, NoTls};
24
25pub struct TestSuite {
26    config: String,
27}
28
29macro_rules! test_eq {
30    ($left:expr, $right:expr $(,)?) => {
31        match (&$left, &$right) {
32            (left_val, right_val) => {
33                if !(*left_val == *right_val) {
34                    return Err(anyhow!(
35                        "{}:{} assertion failed: `(left == right)` \
36                                (left: `{:?}`, right: `{:?}`)",
37                        file!(),
38                        line!(),
39                        left_val,
40                        right_val
41                    ));
42                }
43            }
44        }
45    };
46}
47
48impl TestSuite {
49    pub fn new(
50        db_name: String,
51        user_name: String,
52        server_host: String,
53        server_port: u16,
54        password: String,
55    ) -> Self {
56        let config = if !password.is_empty() {
57            format!(
58                "dbname={} user={} host={} port={} password={}",
59                db_name, user_name, server_host, server_port, password
60            )
61        } else {
62            format!(
63                "dbname={} user={} host={} port={}",
64                db_name, user_name, server_host, server_port
65            )
66        };
67        Self { config }
68    }
69
70    fn init_logger() {
71        let _ = tracing_subscriber::fmt()
72            .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
73            .with_ansi(false)
74            .try_init();
75    }
76
77    pub async fn test(&self) -> anyhow::Result<()> {
78        Self::init_logger();
79        self.binary_param_and_result().await?;
80        self.dql_dml_with_param().await?;
81        self.max_row().await?;
82        self.multiple_on_going_portal().await?;
83        self.create_with_parameter().await?;
84        self.simple_cancel(false).await?;
85        self.simple_cancel(true).await?;
86        self.complex_cancel(false).await?;
87        self.complex_cancel(true).await?;
88        self.subscription_fetch_cancel(false).await?;
89        self.subscription_fetch_cancel(true).await?;
90        self.subquery_with_param().await?;
91        self.create_mview_with_parameter().await?;
92        Ok(())
93    }
94
95    async fn create_client(&self, is_distributed: bool) -> anyhow::Result<Client> {
96        let (client, connection) = tokio_postgres::connect(&self.config, NoTls).await?;
97
98        // The connection object performs the actual communication with the database,
99        // so spawn it off to run on its own.
100        tokio::spawn(async move {
101            if let Err(e) = connection.await {
102                eprintln!("connection error: {}", e);
103            }
104        });
105
106        if is_distributed {
107            client.execute("set query_mode = distributed", &[]).await?;
108        } else {
109            client.execute("set query_mode = local", &[]).await?;
110        }
111
112        Ok(client)
113    }
114
115    pub async fn binary_param_and_result(&self) -> anyhow::Result<()> {
116        let client = self.create_client(false).await?;
117
118        for row in client.query("select $1::SMALLINT;", &[&1024_i16]).await? {
119            let data: i16 = row.try_get(0)?;
120            test_eq!(data, 1024);
121        }
122
123        for row in client.query("select $1::INT;", &[&144232_i32]).await? {
124            let data: i32 = row.try_get(0)?;
125            test_eq!(data, 144232);
126        }
127
128        for row in client.query("select $1::BIGINT;", &[&99999999_i64]).await? {
129            let data: i64 = row.try_get(0)?;
130            test_eq!(data, 99999999);
131        }
132
133        for row in client
134            .query(
135                "select $1::DECIMAL;",
136                &[&Decimal::try_from(2.33454_f32).ok()],
137            )
138            .await?
139        {
140            let data: Decimal = row.try_get(0)?;
141            test_eq!(data, Decimal::try_from(2.33454_f32).unwrap());
142        }
143
144        for row in client.query("select $1::BOOL;", &[&true]).await? {
145            let data: bool = row.try_get(0)?;
146            assert!(data);
147        }
148
149        for row in client.query("select $1::REAL;", &[&1.234234_f32]).await? {
150            let data: f32 = row.try_get(0)?;
151            test_eq!(data, 1.234234);
152        }
153
154        for row in client
155            .query("select $1::DOUBLE PRECISION;", &[&234234.23490238483_f64])
156            .await?
157        {
158            let data: f64 = row.try_get(0)?;
159            test_eq!(data, 234234.23490238483);
160        }
161
162        for row in client
163            .query(
164                "select $1::date;",
165                &[&NaiveDate::from_ymd_opt(2022, 1, 1).unwrap()],
166            )
167            .await?
168        {
169            let data: NaiveDate = row.try_get(0)?;
170            test_eq!(data, NaiveDate::from_ymd_opt(2022, 1, 1).unwrap());
171        }
172
173        for row in client
174            .query(
175                "select $1::time",
176                &[&NaiveTime::from_hms_opt(10, 0, 0).unwrap()],
177            )
178            .await?
179        {
180            let data: NaiveTime = row.try_get(0)?;
181            test_eq!(data, NaiveTime::from_hms_opt(10, 0, 0).unwrap());
182        }
183
184        for row in client
185            .query(
186                "select $1::timestamp",
187                &[&NaiveDate::from_ymd_opt(2022, 1, 1)
188                    .unwrap()
189                    .and_hms_opt(10, 0, 0)
190                    .unwrap()],
191            )
192            .await?
193        {
194            let data: NaiveDateTime = row.try_get(0)?;
195            test_eq!(
196                data,
197                NaiveDate::from_ymd_opt(2022, 1, 1)
198                    .unwrap()
199                    .and_hms_opt(10, 0, 0)
200                    .unwrap()
201            );
202        }
203
204        let timestamptz = DateTime::<Utc>::from_naive_utc_and_offset(
205            NaiveDate::from_ymd_opt(2022, 1, 1)
206                .unwrap()
207                .and_hms_opt(10, 0, 0)
208                .unwrap(),
209            Utc,
210        );
211        for row in client
212            .query("select $1::timestamptz", &[&timestamptz])
213            .await?
214        {
215            let data: DateTime<Utc> = row.try_get(0)?;
216            test_eq!(data, timestamptz);
217        }
218
219        for row in client
220            .query("select $1::interval", &[&Interval::new(1, 1, 24000000)])
221            .await?
222        {
223            let data: Interval = row.try_get(0)?;
224            test_eq!(data, Interval::new(1, 1, 24000000));
225        }
226
227        Ok(())
228    }
229
230    async fn dql_dml_with_param(&self) -> anyhow::Result<()> {
231        let client = self.create_client(false).await?;
232
233        client.query("create table t(id int)", &[]).await?;
234
235        let insert_statement = client
236            .prepare_typed("insert INTO t (id) VALUES ($1)", &[])
237            .await?;
238
239        for i in 0..20 {
240            client.execute(&insert_statement, &[&i]).await?;
241        }
242        client.execute("flush", &[]).await?;
243
244        let update_statement = client
245            .prepare_typed(
246                "update t set id = $1 where id < $2",
247                &[Type::INT4, Type::INT4],
248            )
249            .await?;
250        let query_statement = client
251            .prepare_typed(
252                "select * FROM t where id < $1 order by id ASC",
253                &[Type::INT4],
254            )
255            .await?;
256        let delete_statement = client
257            .prepare_typed("delete FROM t where id < $1", &[Type::INT4])
258            .await?;
259
260        let mut i = 0;
261        for row in client.query(&query_statement, &[&10_i32]).await? {
262            let id: i32 = row.try_get(0)?;
263            test_eq!(id, i);
264            i += 1;
265        }
266        test_eq!(i, 10);
267
268        client
269            .execute(&update_statement, &[&100_i32, &10_i32])
270            .await?;
271        client.execute("flush", &[]).await?;
272
273        let mut i = 0;
274        for _ in client.query(&query_statement, &[&10_i32]).await? {
275            i += 1;
276        }
277        test_eq!(i, 0);
278
279        client.execute(&delete_statement, &[&20_i32]).await?;
280        client.execute("flush", &[]).await?;
281
282        let mut i = 0;
283        for row in client.query(&query_statement, &[&101_i32]).await? {
284            let id: i32 = row.try_get(0)?;
285            test_eq!(id, 100);
286            i += 1;
287        }
288        test_eq!(i, 10);
289
290        client.execute("drop table t", &[]).await?;
291
292        Ok(())
293    }
294
295    async fn max_row(&self) -> anyhow::Result<()> {
296        let mut client = self.create_client(false).await?;
297
298        client.query("create table t(id int)", &[]).await?;
299
300        let insert_statement = client
301            .prepare_typed("insert INTO t (id) VALUES ($1)", &[])
302            .await?;
303
304        for i in 0..10 {
305            client.execute(&insert_statement, &[&i]).await?;
306        }
307        client.execute("flush", &[]).await?;
308
309        let transaction = client.transaction().await?;
310        let statement = transaction
311            .prepare_typed("SELECT * FROM t order by id", &[])
312            .await?;
313        let portal = transaction.bind(&statement, &[]).await?;
314
315        for t in 0..5 {
316            let rows = transaction.query_portal(&portal, 1).await?;
317            test_eq!(rows.len(), 1);
318            let row = rows.first().unwrap();
319            let id: i32 = row.get(0);
320            test_eq!(id, t);
321        }
322
323        let mut i = 5;
324        for row in transaction.query_portal(&portal, 3).await? {
325            let id: i32 = row.get(0);
326            test_eq!(id, i);
327            i += 1;
328        }
329        test_eq!(i, 8);
330
331        for row in transaction.query_portal(&portal, 5).await? {
332            let id: i32 = row.get(0);
333            test_eq!(id, i);
334            i += 1;
335        }
336        test_eq!(i, 10);
337
338        transaction.rollback().await?;
339
340        client.execute("drop table t", &[]).await?;
341
342        Ok(())
343    }
344
345    async fn multiple_on_going_portal(&self) -> anyhow::Result<()> {
346        let mut client = self.create_client(false).await?;
347
348        let transaction = client.transaction().await?;
349        let statement = transaction
350            .prepare_typed("SELECT generate_series(1,5,1)", &[])
351            .await?;
352        let portal_1 = transaction.bind(&statement, &[]).await?;
353        let portal_2 = transaction.bind(&statement, &[]).await?;
354
355        let rows = transaction.query_portal(&portal_1, 1).await?;
356        test_eq!(rows.len(), 1);
357        test_eq!(rows.first().unwrap().get::<usize, i32>(0), 1);
358
359        let rows = transaction.query_portal(&portal_2, 1).await?;
360        test_eq!(rows.len(), 1);
361        test_eq!(rows.first().unwrap().get::<usize, i32>(0), 1);
362
363        let rows = transaction.query_portal(&portal_2, 3).await?;
364        test_eq!(rows.len(), 3);
365        test_eq!(rows.first().unwrap().get::<usize, i32>(0), 2);
366        test_eq!(rows.get(1).unwrap().get::<usize, i32>(0), 3);
367        test_eq!(rows.get(2).unwrap().get::<usize, i32>(0), 4);
368
369        let rows = transaction.query_portal(&portal_1, 1).await?;
370        test_eq!(rows.len(), 1);
371        test_eq!(rows.first().unwrap().get::<usize, i32>(0), 2);
372
373        Ok(())
374    }
375
376    // Can't support these sql
377    async fn create_with_parameter(&self) -> anyhow::Result<()> {
378        let client = self.create_client(false).await?;
379
380        test_eq!(
381            client
382                .query("create table t as select $1", &[])
383                .await
384                .is_err(),
385            true
386        );
387        test_eq!(
388            client
389                .query("create view v as select $1", &[])
390                .await
391                .is_err(),
392            true
393        );
394
395        Ok(())
396    }
397
398    async fn create_mview_with_parameter(&self) -> anyhow::Result<()> {
399        let client = self.create_client(false).await?;
400
401        let statement = client
402            .prepare_typed(
403                "create materialized view mv as select $1 as x",
404                &[Type::INT4],
405            )
406            .await?;
407
408        client.execute(&statement, &[&42_i32]).await?;
409
410        let rows = client.query("select * from mv", &[]).await?;
411        test_eq!(rows.len(), 1);
412        test_eq!(rows.first().unwrap().get::<usize, i32>(0), 42);
413
414        // Test renaming mv because it relies on parsing and rewrite the `create MV` query
415        client
416            .execute("alter materialized view mv rename to mv2", &[])
417            .await?;
418
419        let rows = client.query("select * from mv2", &[]).await?;
420        test_eq!(rows.len(), 1);
421        test_eq!(rows.first().unwrap().get::<usize, i32>(0), 42);
422
423        client.execute("drop materialized view mv2", &[]).await?;
424
425        Ok(())
426    }
427
428    async fn simple_cancel(&self, is_distributed: bool) -> anyhow::Result<()> {
429        let client = self.create_client(is_distributed).await?;
430        client.execute("create table t(id int)", &[]).await?;
431
432        let insert_statement = client
433            .prepare_typed("insert INTO t (id) VALUES ($1)", &[])
434            .await?;
435
436        for i in 0..1000 {
437            client.execute(&insert_statement, &[&i]).await?;
438        }
439
440        client.execute("flush", &[]).await?;
441
442        let cancel_token = client.cancel_token();
443
444        let query_handle = tokio::spawn(async move {
445            client.query("select * from t", &[]).await.unwrap();
446        });
447
448        select! {
449            _ = query_handle => {
450                tracing::error!("Failed to cancel query")
451            },
452            _ = cancel_token.cancel_query(NoTls) => {
453                tracing::trace!("Cancel query successfully")
454            },
455        }
456
457        let new_client = self.create_client(is_distributed).await?;
458
459        let rows = new_client
460            .query("select * from t order by id limit 10", &[])
461            .await?;
462
463        test_eq!(rows.len(), 10);
464        for (expect_id, row) in rows.iter().enumerate() {
465            let id: i32 = row.get(0);
466            test_eq!(id, expect_id as i32);
467        }
468
469        new_client.execute("drop table t", &[]).await?;
470
471        Ok(())
472    }
473
474    async fn complex_cancel(&self, is_distributed: bool) -> anyhow::Result<()> {
475        let client = self.create_client(is_distributed).await?;
476
477        client
478            .execute("create table t1(name varchar, id int)", &[])
479            .await?;
480        client
481            .execute("create table t2(name varchar, id int)", &[])
482            .await?;
483        client
484            .execute("create table t3(name varchar, id int)", &[])
485            .await?;
486
487        let insert_statement = client
488            .prepare_typed("insert INTO t1 (name, id) VALUES ($1, $2)", &[])
489            .await?;
490        let insert_statement2 = client
491            .prepare_typed("insert INTO t2 (name, id) VALUES ($1, $2)", &[])
492            .await?;
493        let insert_statement3 = client
494            .prepare_typed("insert INTO t3 (name, id) VALUES ($1, $2)", &[])
495            .await?;
496        for i in 0..1000 {
497            client
498                .execute(&insert_statement, &[&i.to_string(), &i])
499                .await?;
500            client
501                .execute(&insert_statement2, &[&i.to_string(), &i])
502                .await?;
503            client
504                .execute(&insert_statement3, &[&i.to_string(), &i])
505                .await?;
506        }
507
508        client.execute("flush", &[]).await?;
509
510        client.execute("set query_mode=local", &[]).await?;
511
512        let cancel_token = client.cancel_token();
513
514        let query_sql = "SELECT t1.name, t2.id, t3.name
515        FROM t1
516        INNER JOIN (
517          SELECT id, name
518          FROM t2
519          WHERE id IN (
520            SELECT id
521            FROM t1
522            WHERE name LIKE '%1%'
523          )
524        ) AS t2 ON t1.id = t2.id
525        LEFT JOIN t3 ON t2.name = t3.name
526        WHERE t3.id IN (
527          SELECT MAX(id)
528          FROM t3
529          GROUP BY name
530        )
531        ORDER BY t1.name ASC, t3.id DESC
532        ";
533
534        let query_handle = tokio::spawn(async move {
535            let result = client.query(query_sql, &[]).await;
536            match result {
537                Ok(_) => {
538                    tracing::error!("Query should be canceled");
539                }
540                Err(e) => {
541                    tracing::error!("Query failed with error: {:?}", e);
542                }
543            };
544        });
545
546        select! {
547            _ = query_handle => {
548                tracing::error!("Failed to cancel query")
549            },
550            _ = cancel_token.cancel_query(NoTls) => {
551                tracing::info!("Cancel query successfully")
552            },
553        }
554
555        let new_client = self.create_client(is_distributed).await?;
556
557        let rows = new_client
558            .query(&format!("{} LIMIT 10", query_sql), &[])
559            .await?;
560        let expect_ans = [
561            (1, 1, 1),
562            (10, 10, 10),
563            (100, 100, 100),
564            (101, 101, 101),
565            (102, 102, 102),
566            (103, 103, 103),
567            (104, 104, 104),
568            (105, 105, 105),
569            (106, 106, 106),
570            (107, 107, 107),
571        ];
572        for (i, row) in rows.iter().enumerate() {
573            test_eq!(
574                row.get::<_, String>(0).parse::<i32>().unwrap(),
575                expect_ans[i].0
576            );
577            test_eq!(row.get::<_, i32>(1), expect_ans[i].1);
578            test_eq!(
579                row.get::<_, String>(2).parse::<i32>().unwrap(),
580                expect_ans[i].2
581            );
582        }
583
584        new_client.execute("drop table t1", &[]).await?;
585        new_client.execute("drop table t2", &[]).await?;
586        new_client.execute("drop table t3", &[]).await?;
587        Ok(())
588    }
589
590    async fn subscription_fetch_cancel(&self, is_distributed: bool) -> anyhow::Result<()> {
591        let client = self.create_client(is_distributed).await?;
592        let suffix = if is_distributed { "dist" } else { "local" };
593        let table_name = format!("sub_cancel_t_{suffix}");
594        let subscription_name = format!("sub_cancel_{suffix}");
595
596        client
597            .execute(&format!("create table {table_name}(v int)"), &[])
598            .await?;
599        client
600            .execute(
601                &format!("create subscription {subscription_name} from {table_name} with(retention = '1D')"),
602                &[],
603            )
604            .await?;
605        client
606            .execute(
607                &format!("declare cur subscription cursor for {subscription_name} since now()"),
608                &[],
609            )
610            .await?;
611
612        let cancel_token = client.cancel_token();
613        let fetch_handle = tokio::spawn(async move {
614            client
615                .query("fetch 1 from cur with (timeout = '60s')", &[])
616                .await
617        });
618
619        tokio::time::sleep(Duration::from_secs(1)).await;
620        cancel_token.cancel_query(NoTls).await?;
621
622        let result = tokio::time::timeout(Duration::from_secs(10), fetch_handle).await??;
623        if result.is_ok() {
624            return Err(anyhow!(
625                "subscription cursor fetch should be cancelled by CancelRequest"
626            ));
627        }
628
629        let cleanup_client = self.create_client(is_distributed).await?;
630        cleanup_client
631            .execute(&format!("drop subscription {subscription_name}"), &[])
632            .await?;
633        cleanup_client
634            .execute(&format!("drop table {table_name}"), &[])
635            .await?;
636        Ok(())
637    }
638
639    async fn subquery_with_param(&self) -> anyhow::Result<()> {
640        let client = self.create_client(false).await?;
641
642        let res = client
643            .query("select (select $1::SMALLINT)", &[&1024_i16])
644            .await
645            .unwrap();
646
647        assert_eq!(res[0].get::<usize, i16>(0), 1024_i16);
648
649        Ok(())
650    }
651}
risingwave_e2e_extended_mode_test/test.rs

risingwave_e2e_extended_mode_test/
test.rs
