fix(cubesql): Support Sigma Computing table schema sync

rust/cubesql/cubesql/src/compile/mod.rs

@@ -10831,6 +10831,50 @@ ORDER BY \"COUNT(count)\" DESC"
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_sigma_computing_attnames() -> Result<(), CubeError> {
+        insta::assert_snapshot!(
+            "sigma_computing_attnames",
+            execute_query(
+                "
+                with
+                    nsp as (
+                        select oid as relnamespace
+                        from pg_catalog.pg_namespace
+                        where nspname = 'public'
+                    ),
+                    tbl as (
+                        select
+                            nsp.relnamespace as connamespace,
+                            tbl.oid as conrelid
+                        from pg_catalog.pg_class tbl
+                        inner join nsp using (relnamespace)
+                        where relname = 'emptytbl'
+                    ),
+                    con as (
+                        select
+                            conrelid,
+                            conkey
+                        from pg_catalog.pg_constraint
+                        inner join tbl using (connamespace, conrelid)
+                        where contype = 'p'
+                    )
+                select attname
+                from pg_catalog.pg_attribute att
+                inner join con on
+                    conrelid = attrelid
+                    and attnum = any(con.conkey)
+                order by attnum
+                "
+                .to_string(),
+                DatabaseProtocol::PostgreSQL
+            )
+            .await?
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_google_sheets_pg_database_query() -> Result<(), CubeError> {
         insta::assert_snapshot!(

rust/cubesql/cubesql/src/compile/parser.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 
+use regex::Regex;
 use sqlparser::{
     ast::Statement,
     dialect::{Dialect, PostgreSqlDialect},
@@ -39,7 +40,7 @@ impl Dialect for MySqlDialectWithBackTicks {
 }
 
 lazy_static! {
-    static ref SIGMA_WORKAROUND: regex::Regex = regex::Regex::new(r#"(?s)^\s*with\s+nsp\sas\s\(.*nspname\s=\s(?P<nspname>'[^']+'|\$\d+).*\),\s+tbl\sas\s\(.*relname\s=\s(?P<relname>'[^']+'|\$\d+).*\).*$"#).unwrap();
+    static ref SIGMA_WORKAROUND: Regex = Regex::new(r#"(?s)^\s*with\s+nsp\sas\s\(.*nspname\s=\s.*\),\s+tbl\sas\s\(.*relname\s=\s.*\).*select\s+attname.*from\spg_attribute.*$"#).unwrap();
 }
 
 pub fn parse_sql_to_statements(
@@ -184,33 +185,17 @@ pub fn parse_sql_to_statements(
     );
 
     // Sigma Computing WITH query workaround
-    let query = match SIGMA_WORKAROUND.captures(&query) {
-        Some(c) => {
-            let nspname = c.name("nspname").unwrap().as_str();
-            let relname = c.name("relname").unwrap().as_str();
-            format!(
-                "
-                select
-                    attname,
-                    typname,
-                    description
-                from pg_attribute a
-                join pg_type on atttypid = pg_type.oid
-                left join pg_description on
-                    attrelid = objoid and
-                    attnum = objsubid
-                join pg_catalog.pg_namespace nsp ON nspname = {}
-                join pg_catalog.pg_class tbl ON relname = {} and relnamespace = nsp.oid
-                where
-                    attnum > 0 and
-                    attrelid = tbl.oid
-                order by attnum
-                ;
-                ",
-                nspname, relname
-            )
-        }
-        None => query,
+    let query = if SIGMA_WORKAROUND.is_match(&query) {
+        let relnamespace_re = Regex::new(r#"(?s)from\spg_catalog\.pg_class\s+where\s+relname\s=\s(?P<relname>'(?:[^']|'')+'|\$\d+)\s+and\s+relnamespace\s=\s\(select\soid\sfrom\snsp\)"#).unwrap();
+        let relnamespace_replaced = relnamespace_re.replace(
+            &query,
+            "from pg_catalog.pg_class join nsp on relnamespace = nsp.oid where relname = $relname",
+        );
+        let attrelid_re = Regex::new(r#"(?s)left\sjoin\spg_description\son\s+attrelid\s=\sobjoid\sand\s+attnum\s=\sobjsubid\s+where\s+attnum\s>\s0\s+and\s+attrelid\s=\s\(select\soid\sfrom\stbl\)"#).unwrap();
+        let attrelid_replaced = attrelid_re.replace(&relnamespace_replaced, "left join pg_description on attrelid = objoid and attnum = objsubid join tbl on attrelid = tbl.oid where attnum > 0");
+        attrelid_replaced.to_string()
+    } else {
+        query
     };
 
     // Metabase

rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__sigma_computing_attnames.snap

@@ -0,0 +1,8 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_query(\"\n                with\n                    nsp as (\n                        select oid as relnamespace\n                        from pg_catalog.pg_namespace\n                        where nspname = 'public'\n                    ),\n                    tbl as (\n                        select\n                            nsp.relnamespace as connamespace,\n                            tbl.oid as conrelid\n                        from pg_catalog.pg_class tbl\n                        inner join nsp using (relnamespace)\n                        where relname = 'emptytbl'\n                    ),\n                    con as (\n                        select\n                            conrelid,\n                            conkey\n                        from pg_catalog.pg_constraint\n                        inner join tbl using (connamespace, conrelid)\n                        where contype = 'p'\n                    )\n                select attname\n                from pg_catalog.pg_attribute att\n                inner join con on\n                    conrelid = attrelid\n                    and attnum = any(con.conkey)\n                order by attnum\n                \".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
+---
++---------+
+| attname |
++---------+
++---------+