feat(cubesql): Support pg_total_relation_size UDF

rust/cubesql/cubesql/src/compile/engine/udf.rs

@@ -2375,3 +2375,29 @@ pub fn create_has_schema_privilege_udf(state: Arc<SessionState>) -> ScalarUDF {
         &fun,
     )
 }
+
+pub fn create_pg_total_relation_size_udf() -> ScalarUDF {
+    let fun = make_scalar_function(move |args: &[ArrayRef]| {
+        assert!(args.len() == 1);
+
+        let relids = downcast_primitive_arg!(args[0], "relid", OidType);
+
+        // 8192 is the lowest size for a table that has at least one column
+        // TODO: check if the requested table actually exists
+        let result = relids
+            .iter()
+            .map(|relid| relid.map(|_| 8192))
+            .collect::<PrimitiveArray<Int64Type>>();
+
+        Ok(Arc::new(result))
+    });
+
+    let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Int64)));
+
+    ScalarUDF::new(
+        "pg_total_relation_size",
+        &Signature::exact(vec![DataType::UInt32], Volatility::Immutable),
+        &return_type,
+        &fun,
+    )
+}

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

@@ -58,11 +58,11 @@ use self::{
             create_pg_get_constraintdef_udf, create_pg_get_expr_udf, create_pg_get_userbyid_udf,
             create_pg_is_other_temp_schema, create_pg_my_temp_schema,
             create_pg_numeric_precision_udf, create_pg_numeric_scale_udf,
-            create_pg_table_is_visible_udf, create_pg_truetypid_udf, create_pg_truetypmod_udf,
-            create_pg_type_is_visible_udf, create_quarter_udf, create_second_udf,
-            create_session_user_udf, create_str_to_date_udf, create_time_format_udf,
-            create_timediff_udf, create_to_char_udf, create_ucase_udf, create_unnest_udtf,
-            create_user_udf, create_version_udf, create_year_udf,
+            create_pg_table_is_visible_udf, create_pg_total_relation_size_udf,
+            create_pg_truetypid_udf, create_pg_truetypmod_udf, create_pg_type_is_visible_udf,
+            create_quarter_udf, create_second_udf, create_session_user_udf, create_str_to_date_udf,
+            create_time_format_udf, create_timediff_udf, create_to_char_udf, create_ucase_udf,
+            create_unnest_udtf, create_user_udf, create_version_udf, create_year_udf,
         },
     },
     parser::parse_sql_to_statement,
@@ -2485,6 +2485,7 @@ WHERE `TABLE_SCHEMA` = '{}'",
         ctx.register_udf(create_pg_my_temp_schema());
         ctx.register_udf(create_pg_is_other_temp_schema());
         ctx.register_udf(create_has_schema_privilege_udf(self.state.clone()));
+        ctx.register_udf(create_pg_total_relation_size_udf());
 
         // udaf
         ctx.register_udaf(create_measure_udaf());
@@ -8053,6 +8054,27 @@ ORDER BY \"COUNT(count)\" DESC"
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_pg_total_relation_size() -> Result<(), CubeError> {
+        insta::assert_snapshot!(
+            "pg_total_relation_size",
+            execute_query(
+                "SELECT
+                    oid,
+                    relname,
+                    pg_total_relation_size(oid) relsize
+                FROM pg_class
+                ORDER BY oid ASC
+                "
+                .to_string(),
+                DatabaseProtocol::PostgreSQL
+            )
+            .await?
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_discard_postgres() -> Result<(), CubeError> {
         insta::assert_snapshot!(

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

@@ -0,0 +1,10 @@
+---
+source: cubesql/src/compile/mod.rs
+expression: "execute_query(\"SELECT\n                    oid,\n                    relname,\n                    pg_total_relation_size(oid) relsize\n                FROM pg_class\n                ORDER BY oid ASC\n                \".to_string(),\n            DatabaseProtocol::PostgreSQL).await?"
+---
++-------+---------------------------+---------+
+| oid   | relname                   | relsize |
++-------+---------------------------+---------+
+| 18000 | KibanaSampleDataEcommerce | 8192    |
+| 18013 | Logs                      | 8192    |
++-------+---------------------------+---------+