feat(opentenbase): add PostgreSQL-compatible database support

crates/dbx-core/src/db/mod.rs

@@ -25,6 +25,7 @@ pub mod mysql;
 pub mod mysql_compatible;
 pub mod ob_oracle;
 pub mod oceanbase_mysql;
+pub mod opentenbase;
 pub mod postgres;
 pub mod proxy_tunnel;
 pub mod questdb;

crates/dbx-core/src/db/opentenbase.rs

@@ -0,0 +1,151 @@
+use deadpool_postgres::Pool;
+
+use crate::db;
+use crate::models::connection::ConnectionConfig;
+
+pub const DRIVER_PROFILE: &str = "opentenbase";
+
+const TABLE_DISTRIBUTION_SQL: &str = "SELECT c.pclocatortype::text, \
+            COALESCE(( \
+                SELECT string_agg(a.attname, E'\\n' ORDER BY distributed_column.ordinality) \
+                FROM unnest(c.discolnums::smallint[]) WITH ORDINALITY \
+                    AS distributed_column(attnum, ordinality) \
+                JOIN pg_catalog.pg_attribute a \
+                  ON a.attrelid = c.pcrelid AND a.attnum = distributed_column.attnum \
+                WHERE NOT a.attisdropped \
+            ), '')::text \
+     FROM pg_catalog.pgxc_class c \
+     JOIN pg_catalog.pg_class relation ON relation.oid = c.pcrelid \
+     JOIN pg_catalog.pg_namespace namespace ON namespace.oid = relation.relnamespace \
+     WHERE namespace.nspname = $1 AND relation.relname = $2 \
+     LIMIT 1";
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum DistributionPolicy {
+    Hash(Vec<String>),
+    Modulo(Vec<String>),
+    Shard(Vec<String>),
+    Replication,
+    RoundRobin,
+}
+
+pub fn is_config(config: &ConnectionConfig) -> bool {
+    matches!(config.driver_profile.as_deref(), Some(DRIVER_PROFILE))
+}
+
+pub async fn table_distribution(pool: &Pool, schema: &str, table: &str) -> Result<Option<DistributionPolicy>, String> {
+    let client = db::postgres::checkout_postgres_client(pool, None, super::connection_timeout()).await?;
+    let Some(row) =
+        client.query_opt(TABLE_DISTRIBUTION_SQL, &[&schema, &table]).await.map_err(|error| error.to_string())?
+    else {
+        return Ok(None);
+    };
+    let locator_type = row.try_get::<_, String>(0).map_err(|error| error.to_string())?;
+    let columns = row
+        .try_get::<_, String>(1)
+        .map_err(|error| error.to_string())?
+        .lines()
+        .map(str::trim)
+        .filter(|column| !column.is_empty())
+        .map(str::to_string)
+        .collect::<Vec<_>>();
+    parse_distribution_policy(&locator_type, columns)
+}
+
+pub fn append_distribution_clause(ddl: &str, distribution: &DistributionPolicy) -> Result<String, String> {
+    let clause = render_distribution_clause(distribution)?;
+    let insertion = ddl
+        .find(";\n")
+        .or_else(|| ddl.find(';'))
+        .ok_or_else(|| "OpenTenBase DDL has no CREATE TABLE terminator".to_string())?;
+    let mut output = String::with_capacity(ddl.len() + clause.len() + 2);
+    output.push_str(&ddl[..insertion]);
+    output.push('\n');
+    output.push_str(&clause);
+    output.push_str(&ddl[insertion..]);
+    Ok(output)
+}
+
+fn parse_distribution_policy(locator_type: &str, columns: Vec<String>) -> Result<Option<DistributionPolicy>, String> {
+    let column_distribution = |kind| {
+        if columns.is_empty() {
+            Err(format!("OpenTenBase {kind} distribution has no distribution columns"))
+        } else {
+            Ok(columns)
+        }
+    };
+    match locator_type {
+        "H" => column_distribution("HASH").map(|columns| Some(DistributionPolicy::Hash(columns))),
+        "M" => column_distribution("MODULO").map(|columns| Some(DistributionPolicy::Modulo(columns))),
+        "S" => column_distribution("SHARD").map(|columns| Some(DistributionPolicy::Shard(columns))),
+        "R" => Ok(Some(DistributionPolicy::Replication)),
+        "N" => Ok(Some(DistributionPolicy::RoundRobin)),
+        _ => Ok(None),
+    }
+}
+
+fn render_distribution_clause(distribution: &DistributionPolicy) -> Result<String, String> {
+    let render_columns = |kind: &str, columns: &[String]| {
+        if columns.is_empty() {
+            return Err(format!("OpenTenBase {kind} distribution has no distribution columns"));
+        }
+        Ok(format!(
+            "DISTRIBUTE BY {kind} ({})",
+            columns.iter().map(|column| db::postgres::pg_quote_ident(column)).collect::<Vec<_>>().join(", ")
+        ))
+    };
+    match distribution {
+        DistributionPolicy::Hash(columns) => render_columns("HASH", columns),
+        DistributionPolicy::Modulo(columns) => render_columns("MODULO", columns),
+        DistributionPolicy::Shard(columns) => render_columns("SHARD", columns),
+        DistributionPolicy::Replication => Ok("DISTRIBUTE BY REPLICATION".to_string()),
+        DistributionPolicy::RoundRobin => Ok("DISTRIBUTE BY ROUNDROBIN".to_string()),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn renders_hash_distribution_before_create_table_terminator() {
+        let ddl = "CREATE TABLE \"public\".\"orders\" (\n  \"tenant_id\" integer\n);\n";
+        let rendered = append_distribution_clause(
+            ddl,
+            &DistributionPolicy::Hash(vec!["tenant_id".to_string(), "order id".to_string()]),
+        )
+        .unwrap();
+
+        assert_eq!(
+            rendered,
+            "CREATE TABLE \"public\".\"orders\" (\n  \"tenant_id\" integer\n)\nDISTRIBUTE BY HASH (\"tenant_id\", \"order id\");\n"
+        );
+    }
+
+    #[test]
+    fn renders_replication_and_round_robin_distribution() {
+        assert_eq!(render_distribution_clause(&DistributionPolicy::Replication).unwrap(), "DISTRIBUTE BY REPLICATION");
+        assert_eq!(render_distribution_clause(&DistributionPolicy::RoundRobin).unwrap(), "DISTRIBUTE BY ROUNDROBIN");
+    }
+
+    #[test]
+    fn maps_current_opentenbase_locator_types() {
+        assert_eq!(
+            parse_distribution_policy("S", vec!["tenant_id".to_string()]).unwrap(),
+            Some(DistributionPolicy::Shard(vec!["tenant_id".to_string()]))
+        );
+        assert_eq!(
+            parse_distribution_policy("M", vec!["tenant_id".to_string()]).unwrap(),
+            Some(DistributionPolicy::Modulo(vec!["tenant_id".to_string()]))
+        );
+        assert_eq!(parse_distribution_policy("D", Vec::new()).unwrap(), None);
+    }
+
+    #[test]
+    fn distribution_query_matches_opentenbase_catalog_contract() {
+        assert!(TABLE_DISTRIBUTION_SQL.contains("pg_catalog.pgxc_class"));
+        assert!(TABLE_DISTRIBUTION_SQL.contains("c.discolnums::smallint[]"));
+        assert!(TABLE_DISTRIBUTION_SQL.contains("WITH ORDINALITY"));
+        assert!(TABLE_DISTRIBUTION_SQL.contains("NOT a.attisdropped"));
+    }
+}

crates/dbx-core/src/db/postgres.rs

@@ -2554,14 +2554,29 @@ pub async fn completion_assistant_search(
     }
 
     if candidates.len() < limit && kinds.iter().any(CompletionAssistantObjectKind::is_routine_like) {
-        let prokinds = postgres_completion_prokinds(&kinds);
-        let rows = postgres_query_cached(
-            &client,
-            postgres_completion_routines_sql(),
-            &[&routine_schema, &pattern, &prokinds, &((limit - candidates.len()) as i64)],
-        )
-        .await
-        .map_err(|e| e.to_string())?;
+        let has_proc_prokind = postgres_proc_has_prokind(&client).await?;
+        let rows = if has_proc_prokind {
+            let prokinds = postgres_completion_prokinds(&kinds);
+            postgres_query_cached(
+                &client,
+                postgres_completion_routines_sql(true),
+                &[&routine_schema, &pattern, &prokinds, &((limit - candidates.len()) as i64)],
+            )
+            .await
+            .map_err(|e| e.to_string())?
+        } else if kinds.iter().any(|kind| {
+            matches!(kind, CompletionAssistantObjectKind::Function | CompletionAssistantObjectKind::Routine)
+        }) {
+            postgres_query_cached(
+                &client,
+                postgres_completion_routines_sql(false),
+                &[&routine_schema, &pattern, &((limit - candidates.len()) as i64)],
+            )
+            .await
+            .map_err(|e| e.to_string())?
+        } else {
+            Vec::new()
+        };
         for row in rows {
             let routine_type: String = pg_row_try_string(&row, 2);
             candidates.push(CompletionAssistantCandidate {
@@ -2665,15 +2680,26 @@ fn postgres_completion_tables_sql() -> &'static str {
      ORDER BY c.relname LIMIT $4"
 }
 
-fn postgres_completion_routines_sql() -> &'static str {
-    "SELECT p.proname, n.nspname, CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \
+fn postgres_completion_routines_sql(has_proc_prokind: bool) -> &'static str {
+    if has_proc_prokind {
+        return "SELECT p.proname, n.nspname, CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \
             obj_description(p.oid) AS routine_comment, COALESCE(pg_get_function_result(p.oid), '') AS data_type, \
             pg_get_function_identity_arguments(p.oid) AS signature \
      FROM pg_catalog.pg_proc p \
      JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \
      WHERE n.nspname = $1 AND p.prokind::text = ANY($3::text[]) \
        AND ($2 = '%%' OR p.proname ILIKE $2 ESCAPE '~') \
-     ORDER BY p.proname LIMIT $4"
+     ORDER BY p.proname LIMIT $4";
+    }
+
+    "SELECT p.proname, n.nspname, 'FUNCTION'::text, \
+            obj_description(p.oid) AS routine_comment, COALESCE(pg_get_function_result(p.oid), '') AS data_type, \
+            pg_get_function_identity_arguments(p.oid) AS signature \
+     FROM pg_catalog.pg_proc p \
+     JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \
+     WHERE n.nspname = $1 \
+       AND ($2 = '%%' OR p.proname ILIKE $2 ESCAPE '~') \
+     ORDER BY p.proname LIMIT $3"
 }
 
 fn postgres_completion_sequences_sql() -> &'static str {
@@ -9456,10 +9482,15 @@ mod tests {
         assert!(postgres_completion_tables_sql().contains("pg_catalog.pg_table_is_visible(c.oid)"));
         assert!(postgres_completion_tables_sql().contains("c.relkind::text = ANY($3::text[])"));
         assert!(postgres_completion_tables_sql().contains("ORDER BY c.relname LIMIT $4"));
-        assert!(postgres_completion_routines_sql().contains("p.proname ILIKE $2 ESCAPE '~'"));
-        assert!(postgres_completion_routines_sql().contains("p.prokind::text = ANY($3::text[])"));
-        assert!(postgres_completion_routines_sql().contains("pg_get_function_identity_arguments(p.oid) AS signature"));
-        assert!(postgres_completion_routines_sql().contains("ORDER BY p.proname LIMIT $4"));
+        assert!(postgres_completion_routines_sql(true).contains("p.proname ILIKE $2 ESCAPE '~'"));
+        assert!(postgres_completion_routines_sql(true).contains("p.prokind::text = ANY($3::text[])"));
+        assert!(
+            postgres_completion_routines_sql(true).contains("pg_get_function_identity_arguments(p.oid) AS signature")
+        );
+        assert!(postgres_completion_routines_sql(true).contains("ORDER BY p.proname LIMIT $4"));
+        assert!(!postgres_completion_routines_sql(false).contains("p.prokind"));
+        assert!(postgres_completion_routines_sql(false).contains("'FUNCTION'::text"));
+        assert!(postgres_completion_routines_sql(false).contains("ORDER BY p.proname LIMIT $3"));
         assert!(postgres_completion_sequences_sql().contains("c.relkind = 'S'"));
         assert!(postgres_completion_sequences_sql().contains("has_schema_privilege(n.oid, 'USAGE')"));
         assert!(postgres_completion_sequences_sql().contains("pg_catalog.pg_table_is_visible(c.oid)"));

crates/dbx-core/src/schema.rs

@@ -6616,6 +6616,9 @@ async fn get_table_ddl_once(
         PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_cloudberry_config) => {
             cloudberry_ddl(p, schema, table).await
         }
+        PoolKind::Postgres(p) if db_config.as_ref().is_some_and(db::opentenbase::is_config) => {
+            opentenbase_ddl(p, schema, table).await
+        }
         PoolKind::Postgres(p)
             if include_postgres_access && db_config.as_ref().is_some_and(is_native_postgres_config) =>
         {
@@ -9451,6 +9454,34 @@ pub async fn cloudberry_ddl(pool: &deadpool_postgres::Pool, schema: &str, table:
     }
 }
 
+pub async fn opentenbase_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result<String, String> {
+    let ddl = pg_ddl(pool, schema, table).await?;
+    match db::opentenbase::table_distribution(pool, schema, table).await {
+        Ok(Some(distribution)) => match db::opentenbase::append_distribution_clause(&ddl, &distribution) {
+            Ok(ddl) => Ok(ddl),
+            Err(error) => {
+                log::warn!(
+                    "[schema][opentenbase:table-ddl-distribution-render-fallback] schema={} table={} error={}",
+                    schema,
+                    table,
+                    error
+                );
+                Ok(ddl)
+            }
+        },
+        Ok(None) => Ok(ddl),
+        Err(error) => {
+            log::warn!(
+                "[schema][opentenbase:table-ddl-distribution-query-fallback] schema={} table={} error={}",
+                schema,
+                table,
+                error
+            );
+            Ok(ddl)
+        }
+    }
+}
+
 pub fn render_postgres_table_ddl(
     schema: &str,
     table: &str,

crates/dbx-core/src/schema/providers/native.rs

@@ -260,6 +260,9 @@ pub(in crate::schema) async fn table_ddl(
         PoolKind::Postgres(p) if config.is_some_and(is_cloudberry_config) => {
             super::super::cloudberry_ddl(p, schema, table).await
         }
+        PoolKind::Postgres(p) if config.is_some_and(db::opentenbase::is_config) => {
+            super::super::opentenbase_ddl(p, schema, table).await
+        }
         PoolKind::Postgres(p) => super::super::pg_ddl(p, schema, table).await,
         PoolKind::Sqlite(p) => super::super::sqlite_ddl(p, schema, table).await,
         PoolKind::Rqlite(client) => db::rqlite_driver::table_ddl(client, table).await,