feat(postgres): add extension management

crates/dbx-core/src/database_export.rs

@@ -1300,8 +1300,13 @@ pub async fn export_database_sql_core(
             PostgresExtensionMembers::default()
         };
     let postgres_extensions = if request.include_structure && matches!(db_type, DatabaseType::Postgres) {
-        match crate::schema::list_extensions_core(state, &request.connection_id, &request.database, &request.schema)
-            .await
+        match crate::schema::list_extensions_core(
+            state,
+            &request.connection_id,
+            &request.database,
+            Some(&request.schema),
+        )
+        .await
         {
             Ok(extensions) => extensions
                 .into_iter()

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

@@ -3329,19 +3329,32 @@ pub async fn list_rules(pool: &Pool, schema: &str) -> Result<Vec<RuleInfo>, Stri
         .collect())
 }
 
-pub async fn list_extensions(pool: &Pool, schema: &str) -> Result<Vec<ExtensionInfo>, String> {
+pub async fn list_extensions(pool: &Pool, schema: Option<&str>) -> Result<Vec<ExtensionInfo>, String> {
     let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?;
-    let rows = postgres_query_cached(
-        &client,
-        "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \
-         FROM pg_catalog.pg_extension e \
-         JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace \
-         LEFT JOIN pg_catalog.pg_description d ON d.objoid = e.oid AND d.classoid = 'pg_extension'::regclass \
-         WHERE n.nspname = $1 \
-         ORDER BY e.extname",
-        &[&schema],
-    )
-    .await
+    let rows = if let Some(schema) = schema.filter(|value| !value.is_empty()) {
+        postgres_query_cached(
+            &client,
+            "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \
+             FROM pg_catalog.pg_extension e \
+             JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace \
+             LEFT JOIN pg_catalog.pg_description d ON d.objoid = e.oid AND d.classoid = 'pg_extension'::regclass \
+             WHERE n.nspname = $1 \
+             ORDER BY e.extname",
+            &[&schema],
+        )
+        .await
+    } else {
+        postgres_query_cached(
+            &client,
+            "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \
+             FROM pg_catalog.pg_extension e \
+             JOIN pg_catalog.pg_namespace n ON n.oid = e.extnamespace \
+             LEFT JOIN pg_catalog.pg_description d ON d.objoid = e.oid AND d.classoid = 'pg_extension'::regclass \
+             ORDER BY n.nspname, e.extname",
+            &[],
+        )
+        .await
+    }
     .map_err(|e| e.to_string())?;
 
     Ok(rows
@@ -3350,7 +3363,7 @@ pub async fn list_extensions(pool: &Pool, schema: &str) -> Result<Vec<ExtensionI
             name: pg_row_try_string(row, 0),
             version: pg_row_try_string(row, 1),
             comment: row.try_get::<_, Option<String>>(2).ok().flatten().filter(|s| !s.is_empty()),
-            schema: Some(schema.to_string()),
+            schema: row.try_get::<_, Option<String>>(3).ok().flatten().filter(|s| !s.is_empty()),
         })
         .collect())
 }

crates/dbx-core/src/schema.rs

@@ -3217,6 +3217,57 @@ mod tests {
         assert!(!is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Mysql)));
     }
 
+    #[test]
+    fn kingbase_extension_sql_uses_sys_catalog_and_escapes_schema() {
+        let sql = super::kingbase_list_extensions_sql(Some("app's"), super::KingbaseExtensionCatalog::Sys);
+
+        assert!(sql.contains("FROM sys_catalog.sys_extension e"));
+        assert!(sql.contains("JOIN sys_catalog.sys_namespace n"));
+        assert!(sql.contains("LEFT JOIN sys_catalog.sys_description d"));
+        assert!(sql.contains("WHERE n.nspname = 'app''s'"));
+    }
+
+    #[test]
+    fn kingbase_available_extension_sql_supports_pg_catalog_fallback() {
+        let sql = super::kingbase_list_available_extensions_sql(super::KingbaseExtensionCatalog::Pg);
+
+        assert!(sql.contains("FROM pg_catalog.pg_available_extensions"));
+        assert!(sql.contains("WHERE installed_version IS NULL"));
+    }
+
+    #[test]
+    fn extension_infos_from_query_result_maps_installed_extensions() {
+        let result = db::QueryResult {
+            columns: vec![
+                "extname".to_string(),
+                "extversion".to_string(),
+                "description".to_string(),
+                "nspname".to_string(),
+            ],
+            column_types: Vec::new(),
+            column_sortables: Vec::new(),
+            rows: vec![vec![
+                serde_json::json!("kdb_utils"),
+                serde_json::json!("1.0"),
+                serde_json::json!("KingBase utilities"),
+                serde_json::json!("public"),
+            ]],
+            affected_rows: 0,
+            execution_time_ms: 0,
+            truncated: false,
+            session_id: None,
+            has_more: false,
+        };
+
+        let extensions = super::extension_infos_from_query_result(result, true);
+
+        assert_eq!(extensions.len(), 1);
+        assert_eq!(extensions[0].name, "kdb_utils");
+        assert_eq!(extensions[0].version, "1.0");
+        assert_eq!(extensions[0].comment.as_deref(), Some("KingBase utilities"));
+        assert_eq!(extensions[0].schema.as_deref(), Some("public"));
+    }
+
     #[test]
     fn agent_metadata_timeout_defaults_to_sixty_seconds_and_honors_longer_config() {
         assert_eq!(super::agent_metadata_timeout(None), Some(std::time::Duration::from_secs(60)));
@@ -4827,10 +4878,25 @@ pub async fn list_extensions_core(
     state: &AppState,
     connection_id: &str,
     database: &str,
-    schema: &str,
+    schema: Option<&str>,
 ) -> Result<Vec<db::ExtensionInfo>, String> {
     retry_metadata_connection(state, connection_id, Some(database), || async {
         let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
+        let db_config = connection_config(state, connection_id).await;
+        if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) {
+            let connections = state.connections.read().await;
+            if let Some(client) = extract_pool!(&connections, &pool_key, Agent) {
+                drop(connections);
+                return kingbase_agent_list_extensions(
+                    client,
+                    database,
+                    schema,
+                    agent_metadata_timeout(db_config.as_ref()),
+                )
+                .await;
+            }
+        }
+
         let connections = state.connections.read().await;
         let pool = connections.get(&pool_key).ok_or("Pool not found")?;
 
@@ -4849,6 +4915,20 @@ pub async fn list_available_extensions_core(
 ) -> Result<Vec<db::ExtensionInfo>, String> {
     retry_metadata_connection(state, connection_id, Some(database), || async {
         let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
+        let db_config = connection_config(state, connection_id).await;
+        if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) {
+            let connections = state.connections.read().await;
+            if let Some(client) = extract_pool!(&connections, &pool_key, Agent) {
+                drop(connections);
+                return kingbase_agent_list_available_extensions(
+                    client,
+                    database,
+                    agent_metadata_timeout(db_config.as_ref()),
+                )
+                .await;
+            }
+        }
+
         let connections = state.connections.read().await;
         let pool = connections.get(&pool_key).ok_or("Pool not found")?;
 
@@ -4860,6 +4940,140 @@ pub async fn list_available_extensions_core(
     .await
 }
 
+#[derive(Clone, Copy)]
+enum KingbaseExtensionCatalog {
+    Sys,
+    Pg,
+}
+
+impl KingbaseExtensionCatalog {
+    fn catalog_name(self) -> &'static str {
+        match self {
+            Self::Sys => "sys_catalog",
+            Self::Pg => "pg_catalog",
+        }
+    }
+
+    fn prefix(self) -> &'static str {
+        match self {
+            Self::Sys => "sys",
+            Self::Pg => "pg",
+        }
+    }
+}
+
+fn kingbase_list_extensions_sql(schema: Option<&str>, catalog: KingbaseExtensionCatalog) -> String {
+    let catalog_name = catalog.catalog_name();
+    let prefix = catalog.prefix();
+    let schema_filter = schema
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+        .map(|schema| format!("WHERE n.nspname = {}", sql_string(schema)))
+        .unwrap_or_default();
+    let order_by = if schema_filter.is_empty() { "n.nspname, e.extname" } else { "e.extname" };
+    format!(
+        "SELECT e.extname, COALESCE(e.extversion, '') AS extversion, d.description, n.nspname \
+         FROM {catalog_name}.{prefix}_extension e \
+         JOIN {catalog_name}.{prefix}_namespace n ON n.oid = e.extnamespace \
+         LEFT JOIN {catalog_name}.{prefix}_description d ON d.objoid = e.oid AND d.objsubid = 0 \
+         {schema_filter} \
+         ORDER BY {order_by}"
+    )
+}
+
+fn kingbase_list_available_extensions_sql(catalog: KingbaseExtensionCatalog) -> String {
+    let catalog_name = catalog.catalog_name();
+    let prefix = catalog.prefix();
+    format!(
+        "SELECT name, default_version, comment \
+         FROM {catalog_name}.{prefix}_available_extensions \
+         WHERE installed_version IS NULL \
+         ORDER BY name"
+    )
+}
+
+async fn kingbase_agent_query_result(
+    client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
+    database: &str,
+    sql: &str,
+    max_rows: usize,
+    timeout_duration: Option<Duration>,
+) -> Result<db::QueryResult, String> {
+    let params = agent_execute_query_params(
+        sql,
+        if database.is_empty() { None } else { Some(database) },
+        None,
+        QueryExecutionOptions { max_rows: Some(max_rows), ..Default::default() },
+    );
+    let mut client = client.lock().await;
+    client.execute_query_with_timeout(params, timeout_duration).await
+}
+
+async fn kingbase_agent_query_result_with_catalog_fallback(
+    client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
+    database: &str,
+    sys_sql: String,
+    pg_sql: String,
+    max_rows: usize,
+    timeout_duration: Option<Duration>,
+) -> Result<db::QueryResult, String> {
+    match kingbase_agent_query_result(client.clone(), database, &sys_sql, max_rows, timeout_duration).await {
+        Ok(result) => Ok(result),
+        Err(sys_error) => kingbase_agent_query_result(client, database, &pg_sql, max_rows, timeout_duration)
+            .await
+            .map_err(|pg_error| format!("{sys_error}; pg_catalog fallback failed: {pg_error}")),
+    }
+}
+
+async fn kingbase_agent_list_extensions(
+    client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
+    database: &str,
+    schema: Option<&str>,
+    timeout_duration: Option<Duration>,
+) -> Result<Vec<db::ExtensionInfo>, String> {
+    let result = kingbase_agent_query_result_with_catalog_fallback(
+        client,
+        database,
+        kingbase_list_extensions_sql(schema, KingbaseExtensionCatalog::Sys),
+        kingbase_list_extensions_sql(schema, KingbaseExtensionCatalog::Pg),
+        10_000,
+        timeout_duration,
+    )
+    .await?;
+    Ok(extension_infos_from_query_result(result, true))
+}
+
+async fn kingbase_agent_list_available_extensions(
+    client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
+    database: &str,
+    timeout_duration: Option<Duration>,
+) -> Result<Vec<db::ExtensionInfo>, String> {
+    let result = kingbase_agent_query_result_with_catalog_fallback(
+        client,
+        database,
+        kingbase_list_available_extensions_sql(KingbaseExtensionCatalog::Sys),
+        kingbase_list_available_extensions_sql(KingbaseExtensionCatalog::Pg),
+        10_000,
+        timeout_duration,
+    )
+    .await?;
+    Ok(extension_infos_from_query_result(result, false))
+}
+
+fn extension_infos_from_query_result(result: db::QueryResult, include_schema: bool) -> Vec<db::ExtensionInfo> {
+    result
+        .rows
+        .into_iter()
+        .filter_map(|row| {
+            let name = query_result_cell_string(&row, 0)?;
+            let version = query_result_cell_string(&row, 1).unwrap_or_default();
+            let comment = query_result_cell_string(&row, 2).filter(|value| !value.trim().is_empty());
+            let schema = include_schema.then(|| query_result_cell_string(&row, 3)).flatten();
+            Some(db::ExtensionInfo { name, version, comment, schema })
+        })
+        .collect()
+}
+
 pub async fn list_owners_core(
     state: &AppState,
     connection_id: &str,