From a3110b85acbd811e3a08eedffb7143d5e6251c08 Mon Sep 17 00:00:00 2001
From: hjiang <hjiang@anyscale.com>
Date: Wed, 4 Dec 2024 08:32:32 +0000
Subject: [PATCH] disallow out of range numerics type

---
 pg_duckdb                           |  1 +
 src/columnstore_handler.cpp         | 67 ++++++++++++++++++++++-------
 test/expected/unsupported/types.out |  4 ++
 test/sql/unsupported/types.sql      |  3 ++
 4 files changed, 60 insertions(+), 15 deletions(-)
 create mode 160000 pg_duckdb

diff --git a/pg_duckdb b/pg_duckdb
new file mode 160000
index 0000000..0b400af
--- /dev/null
+++ b/pg_duckdb
@@ -0,0 +1 @@
+Subproject commit 0b400af176f9e3e4fa54cf55fbfb22e93ae55543
diff --git a/src/columnstore_handler.cpp b/src/columnstore_handler.cpp
index 205bce6..60e1d8c 100644
--- a/src/columnstore_handler.cpp
+++ b/src/columnstore_handler.cpp
@@ -9,6 +9,38 @@ extern "C" {
 #include "utils/syscache.h"
 }
 
+namespace {
+
+// Get precision for postgres "numerics" type, copied from pg_duckdb.
+int numeric_typmod_precision(int32 typmod) {
+    return ((typmod - VARHDRSZ) >> 16) & 0xffff;
+}
+// Get scale for postgres "numerics" type, copied from pg_duckdb.
+int numeric_typmod_scale(int32 typmod) {
+    return (((typmod - VARHDRSZ) & 0x7ff) ^ 1024) - 1024;
+}
+
+// If column is of type "numeric", check whether it's acceptable for mooncake;
+// If not, exception is thrown via `elog`.
+void ValidateColumnNumericType(Form_pg_attribute &attribute) {
+    auto &type = attribute->atttypid;
+    if (type != NUMERICOID) {
+        return;
+    }
+    auto &typmod = attribute->atttypmod;
+    auto precision = numeric_typmod_precision(typmod);
+    auto scale = numeric_typmod_scale(typmod);
+    // duckdb's "numeric" type's max supported precision is 38.
+    if (typmod == -1 || precision < 0 || scale < 0 || precision > 38) {
+        elog(ERROR,
+             "Unsupported type when creating column: type precision %d with scale %d is not supported by duckdb, which "
+             "only allows maximum precision 38",
+             precision, scale);
+    }
+}
+
+} // namespace
+
 const TupleTableSlotOps *columnstore_slot_callbacks(Relation rel) {
     elog(ERROR, "columnstore_slot_callbacks not implemented");
 }
@@ -129,24 +161,29 @@ void columnstore_relation_set_new_filenode(Relation rel, const RelFileNode *newr
                                            TransactionId *freezeXid, MultiXactId *minmulti) {
 #endif
     HeapTuple tp = SearchSysCache1(RELOID, ObjectIdGetDatum(rel->rd_id));
-    if (!HeapTupleIsValid(tp)) {
-        TupleDesc desc = RelationGetDescr(rel);
-        for (int i = 0; i < desc->natts; i++) {
-            Form_pg_attribute attr = &desc->attrs[i];
-            auto duck_type = pgduckdb::ConvertPostgresToDuckColumnType(attr);
-            if (duck_type.id() == duckdb::LogicalTypeId::USER) {
-                elog(ERROR, "column \"%s\" has unsupported type", NameStr(attr->attname));
-            }
-            if (attr->attgenerated) {
-                elog(ERROR, "unsupported generated column \"%s\"", NameStr(attr->attname));
-            }
-        }
-
-        duckdb::Columnstore::CreateTable(rel->rd_id);
-    } else {
+    if (HeapTupleIsValid(tp)) {
         ReleaseSysCache(tp);
         duckdb::Columnstore::TruncateTable(rel->rd_id);
     }
+
+    TupleDesc desc = RelationGetDescr(rel);
+    for (int i = 0; i < desc->natts; i++) {
+        // Check whether column type is acceptable.
+        Form_pg_attribute attr = &desc->attrs[i];
+        auto duck_type = pgduckdb::ConvertPostgresToDuckColumnType(attr);
+        if (duck_type.id() == duckdb::LogicalTypeId::USER) {
+            const auto type_info = duck_type.ToString();
+            elog(ERROR, "column \"%s\" has unsupported user type: %s", NameStr(attr->attname), type_info.data());
+        }
+        if (attr->attgenerated) {
+            elog(ERROR, "unsupported generated column \"%s\"", NameStr(attr->attname));
+        }
+
+        // Check numeric types, which have different support for duckdb and postgres.
+        ValidateTableCreationColumnType(attr);
+    }
+
+    duckdb::Columnstore::CreateTable(rel->rd_id);
 }
 
 void columnstore_relation_nontransactional_truncate(Relation rel) {
diff --git a/test/expected/unsupported/types.out b/test/expected/unsupported/types.out
index 364b025..524cfe5 100644
--- a/test/expected/unsupported/types.out
+++ b/test/expected/unsupported/types.out
@@ -3,3 +3,7 @@ ERROR:  column "a" has unsupported type
 CREATE TYPE point AS (x int, y int);
 CREATE TABLE t (a point) USING columnstore;
 ERROR:  column "a" has unsupported type
+CREATE TABLE mooncake_numeric_tbl (val numeric(40, 5)) USING columnstore;
+ERROR:  Unsupported type when creating column: type precision 40 with scale 5 is not supported by duckdb, which only allows maximum precision 38
+CREATE TABLE mooncake_numeric_tbl (val numeric) USING columnstore;
+ERROR:  Unsupported type when creating column: type precision 65535 with scale -5 is not supported by duckdb, which only allows maximum precision 38
diff --git a/test/sql/unsupported/types.sql b/test/sql/unsupported/types.sql
index 93c53a9..ac6a46f 100644
--- a/test/sql/unsupported/types.sql
+++ b/test/sql/unsupported/types.sql
@@ -2,3 +2,6 @@ CREATE TABLE t (a jsonb) USING columnstore;
 
 CREATE TYPE point AS (x int, y int);
 CREATE TABLE t (a point) USING columnstore;
+
+CREATE TABLE mooncake_numeric_tbl (val numeric(40, 5)) USING columnstore;
+CREATE TABLE mooncake_numeric_tbl (val numeric) USING columnstore;