feat: support postgres wire protocol (#62)

* feat: added basic support for postgres

* fix: use proxy to chose pg or mysql
---------

Co-authored-by: tanruixiang <tanruixiang0104@gmail.com>
Co-authored-by: jiacai2050 <dev@liujiacai.net>

.github/workflows/ci.yml

@@ -78,6 +78,24 @@ jobs:
   sqlness-cli:
     runs-on: ubuntu-latest
     timeout-minutes: 60
+    # Service containers to run with `container-job`
+    services:
+      # Label used to access the service container
+      postgres:
+        # Docker Hub image
+        image: postgres
+        # Provide the password for postgres
+        env:
+          POSTGRES_PASSWORD: postgres
+        # Set health checks to wait until postgres has started
+        options: >-
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+        ports:
+          # Maps tcp port 5432 on service container to the host
+          - 5432:5432
     strategy:
       matrix:
         rust: [stable]

Cargo.toml

@@ -1,5 +1,6 @@
 [workspace]
 members = ["sqlness", "sqlness-cli"]
+resolver = "2"
 
 [workspace.package]
 version = "0.5.0"

Makefile

@@ -22,7 +22,8 @@ clippy:
 	cd $(DIR); cargo clippy --all-targets --all-features --workspace -- -D warnings
 
 cli-test:
-	cd $(DIR)/sqlness-cli;  cargo run -- -c tests -i 127.0.0.1 -p 3306 -u root -P 1a2b3c -d public
+	cd $(DIR)/sqlness-cli;  cargo run -- -t mysql -c tests/mysql -i 127.0.0.1 -p 3306 -u root -P 1a2b3c -d public
+	cd $(DIR)/sqlness-cli;  cargo run -- -t postgresql -c tests/postgresql -i 127.0.0.1 -p 5432 -u postgres -P postgres -d postgres
 
 example: good-example bad-example
 

sqlness-cli/Cargo.toml

@@ -13,4 +13,4 @@ readme = { workspace = true }
 async-trait = "0.1.64"
 clap = { version = "4.1.8", features = ["derive"] }
 futures = "0.3.26"
-sqlness = { path = "../sqlness", version = "0.5", features = ["mysql"] }
+sqlness = { path = "../sqlness", version = "0.5", features = ["mysql", "postgres"] }

sqlness-cli/src/main.rs

@@ -1,13 +1,14 @@
 // Copyright 2023 CeresDB Project Authors. Licensed under Apache-2.0.
 
-use std::path::Path;
+use std::{fmt::Display, path::Path};
 
 use async_trait::async_trait;
 use clap::Parser;
 use futures::executor::block_on;
 use sqlness::{
-    database_impl::mysql::MysqlDatabase, ConfigBuilder, DatabaseConfig, DatabaseConfigBuilder,
-    EnvController, Runner,
+    database_impl::{mysql::MysqlDatabase, postgresql::PostgresqlDatabase},
+    ConfigBuilder, Database, DatabaseConfig, DatabaseConfigBuilder, EnvController, QueryContext,
+    Runner,
 };
 
 #[derive(Parser, Debug)]
@@ -39,27 +40,59 @@ struct Args {
     db: Option<String>,
 
     /// Which DBMS to test against
-    #[clap(short, long)]
+    #[clap(short('t'), long("type"))]
     #[arg(value_enum, default_value_t)]
-    r#type: DBType,
+    db_type: DBType,
 }
 
-#[derive(clap::ValueEnum, Clone, Debug, Default)]
+#[derive(clap::ValueEnum, Clone, Debug, Default, Copy)]
 enum DBType {
     #[default]
     Mysql,
+    Postgresql,
+}
+
+struct DBProxy {
+    database: Box<dyn Database + Sync + Send>,
+}
+
+#[async_trait]
+impl Database for DBProxy {
+    async fn query(&self, context: QueryContext, query: String) -> Box<dyn Display> {
+        self.database.query(context, query).await
+    }
+}
+
+impl DBProxy {
+    pub fn new(db_config: DatabaseConfig, db_type: DBType) -> Self {
+        let database: Box<dyn Database + Sync + Send> = match db_type {
+            DBType::Mysql => Box::new(MysqlDatabase::try_new(db_config).expect("build mysql db")),
+            DBType::Postgresql => {
+                Box::new(PostgresqlDatabase::try_new(&db_config).expect("build postgresql db"))
+            }
+        };
+
+        DBProxy { database }
+    }
 }
 
 struct CliController {
     db_config: DatabaseConfig,
+    db_type: DBType,
+}
+
+impl CliController {
+    fn new(db_config: DatabaseConfig, db_type: DBType) -> Self {
+        Self { db_config, db_type }
+    }
 }
 
 #[async_trait]
 impl EnvController for CliController {
-    type DB = MysqlDatabase;
+    type DB = DBProxy;
 
     async fn start(&self, _env: &str, _config: Option<&Path>) -> Self::DB {
-        MysqlDatabase::try_new(self.db_config.clone()).expect("build db")
+        DBProxy::new(self.db_config.clone(), self.db_type)
     }
 
     async fn stop(&self, _env: &str, _db: Self::DB) {}
@@ -77,15 +110,14 @@ fn main() {
         .build()
         .expect("build db config");
 
-    let ctrl = CliController { db_config };
     let config = ConfigBuilder::default()
         .case_dir(args.case_dir)
         .build()
         .expect("build config");
 
     block_on(async {
+        let ctrl = CliController::new(db_config, args.db_type);
         let runner = Runner::new(config, ctrl);
-
         runner.run().await.expect("run testcase")
     });
 

sqlness-cli/tests/postgresql/local/input.result

@@ -0,0 +1,38 @@
+DROP TABLE if exists categories;
+
+(Empty response)
+
+CREATE TABLE categories (
+  category_id SERIAL NOT NULL PRIMARY KEY,
+  category_name VARCHAR(255),
+  description VARCHAR(255)
+);
+
+(Empty response)
+
+INSERT INTO categories (category_name, description)
+VALUES
+  ('Beverages', 'Soft drinks, coffees, teas, beers, and ales'),
+  ('Condiments', 'Sweet and savory sauces, relishes, spreads, and seasonings'),
+  ('Confections', 'Desserts, candies, and sweet breads'),
+  ('Dairy Products', 'Cheeses'),
+  ('Grains/Cereals', 'Breads, crackers, pasta, and cereal'),
+  ('Meat/Poultry', 'Prepared meats'),
+  ('Produce', 'Dried fruit and bean curd'),
+  ('Seafood', 'Seaweed and fish');
+
+(Empty response)
+
+select * from categories;
+
+category_id,category_name,description,
+Row { columns: [Column { name: "category_id", type: Int4 }, Column { name: "category_name", type: Varchar }, Column { name: "description", type: Varchar }] }
+Row { columns: [Column { name: "category_id", type: Int4 }, Column { name: "category_name", type: Varchar }, Column { name: "description", type: Varchar }] }
+Row { columns: [Column { name: "category_id", type: Int4 }, Column { name: "category_name", type: Varchar }, Column { name: "description", type: Varchar }] }
+Row { columns: [Column { name: "category_id", type: Int4 }, Column { name: "category_name", type: Varchar }, Column { name: "description", type: Varchar }] }
+Row { columns: [Column { name: "category_id", type: Int4 }, Column { name: "category_name", type: Varchar }, Column { name: "description", type: Varchar }] }
+Row { columns: [Column { name: "category_id", type: Int4 }, Column { name: "category_name", type: Varchar }, Column { name: "description", type: Varchar }] }
+Row { columns: [Column { name: "category_id", type: Int4 }, Column { name: "category_name", type: Varchar }, Column { name: "description", type: Varchar }] }
+Row { columns: [Column { name: "category_id", type: Int4 }, Column { name: "category_name", type: Varchar }, Column { name: "description", type: Varchar }] }
+
+

sqlness-cli/tests/postgresql/local/input.sql

@@ -0,0 +1,22 @@
+
+DROP TABLE if exists categories;
+
+CREATE TABLE categories (
+  category_id SERIAL NOT NULL PRIMARY KEY,
+  category_name VARCHAR(255),
+  description VARCHAR(255)
+);
+
+INSERT INTO categories (category_name, description)
+VALUES
+  ('Beverages', 'Soft drinks, coffees, teas, beers, and ales'),
+  ('Condiments', 'Sweet and savory sauces, relishes, spreads, and seasonings'),
+  ('Confections', 'Desserts, candies, and sweet breads'),
+  ('Dairy Products', 'Cheeses'),
+  ('Grains/Cereals', 'Breads, crackers, pasta, and cereal'),
+  ('Meat/Poultry', 'Prepared meats'),
+  ('Produce', 'Dried fruit and bean curd'),
+  ('Seafood', 'Seaweed and fish');
+
+
+select * from categories;

sqlness/Cargo.toml

@@ -13,6 +13,7 @@ readme = { workspace = true }
 async-trait = "0.1"
 derive_builder = "0.11"
 mysql = { version = "23.0.1", optional = true }
+postgres = { version = "0.19.7", optional = true }
 prettydiff = { version = "0.6.2", default_features = false }
 regex = "1.7.1"
 thiserror = "1.0"

sqlness/src/database_impl/mod.rs

@@ -2,3 +2,5 @@
 
 #[cfg(feature = "mysql")]
 pub mod mysql;
+#[cfg(feature = "postgres")]
+pub mod postgresql;

sqlness/src/database_impl/postgresql.rs

@@ -0,0 +1,92 @@
+// Copyright 2022 CeresDB Project Authors. Licensed under Apache-2.0.
+
+use async_trait::async_trait;
+use postgres::{Client, Config, NoTls, Row};
+use std::{
+    fmt::Display,
+    sync::{Arc, Mutex},
+};
+
+use crate::{Database, DatabaseConfig, QueryContext};
+
+pub struct PostgresqlDatabase {
+    client: Arc<Mutex<Client>>,
+}
+
+impl PostgresqlDatabase {
+    pub fn try_new(config: &DatabaseConfig) -> Result<Self, postgres::Error> {
+        let mut postgres_config = Config::new();
+        postgres_config
+            .port(config.tcp_port)
+            .host(&config.ip_or_host);
+
+        if let Some(user) = &config.user {
+            postgres_config.user(user);
+        }
+        if let Some(password) = &config.pass {
+            postgres_config.password(password);
+        }
+        if let Some(dbname) = &config.db_name {
+            postgres_config.dbname(dbname);
+        }
+        let client = postgres_config.connect(NoTls)?;
+        Ok(PostgresqlDatabase {
+            client: Arc::new(Mutex::new(client)),
+        })
+    }
+
+    pub fn execute(query: &str, client: Arc<Mutex<Client>>) -> Box<dyn Display> {
+        let mut client = match client.lock() {
+            Ok(client) => client,
+            Err(err) => {
+                return Box::new(format!("Failed to get connection, encountered: {:?}", err))
+            }
+        };
+
+        let result = match client.query(query, &[]) {
+            Ok(rows) => {
+                format!("{}", PostgresqlFormatter { rows })
+            }
+            Err(err) => format!("Failed to execute query, encountered: {:?}", err),
+        };
+
+        Box::new(result)
+    }
+}
+
+struct PostgresqlFormatter {
+    pub rows: Vec<Row>,
+}
+
+impl Display for PostgresqlFormatter {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.rows.is_empty() {
+            return f.write_fmt(format_args!("(Empty response)"));
+        }
+
+        let top = &self.rows[0];
+        let columns = top
+            .columns()
+            .iter()
+            .map(|column| column.name())
+            .collect::<Vec<_>>();
+        for col in &columns {
+            f.write_fmt(format_args!("{},", col))?;
+        }
+
+        f.write_str("\n")?;
+
+        for row in &self.rows {
+            f.write_fmt(format_args!("{:?}\n", row))?;
+        }
+
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl Database for PostgresqlDatabase {
+    async fn query(&self, _: QueryContext, query: String) -> Box<dyn Display> {
+        Self::execute(&query, Arc::clone(&self.client))
+    }
+}

sqlness/src/interceptor/sort_result.rs

@@ -63,8 +63,8 @@ impl Interceptor for SortResultInterceptor {
 
         let new_lines = head
             .into_iter()
-            .chain(lines.into_iter())
-            .chain(tail.into_iter())
+            .chain(lines)
+            .chain(tail)
             .collect::<Vec<_>>();
         *result = new_lines.join("\n");
     }