Support CREATE TABLE in Postgres syntax

3 files changed, 377 insertions, 4 deletions

diff --git a/ydb/library/yql/parser/pg_catalog/catalog.h b/ydb/library/yql/parser/pg_catalog/catalog.h
index 12cc8812f02..634be813599 100644
--- a/ydb/library/yql/parser/pg_catalog/catalog.h
+++ b/ydb/library/yql/parser/pg_catalog/catalog.h
@@ -6,6 +6,14 @@
 
 namespace NYql::NPg {
 
+// copied from pg_class.h
+enum class ERelPersistence : char
+{
+    Permanent = 'p',
+    Unlogged = 'u',
+    Temp = 't',
+};
+
 enum class EOperKind {
     Binary,
     LeftUnary,
diff --git a/ydb/library/yql/sql/pg/pg_sql.cpp b/ydb/library/yql/sql/pg/pg_sql.cpp
index 635951202fb..368c9ddb8de 100644
--- a/ydb/library/yql/sql/pg/pg_sql.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql.cpp
@@ -2,6 +2,7 @@
 #include <ydb/library/yql/sql/settings/partitioning.h>
 #include <ydb/library/yql/parser/pg_wrapper/interface/parser.h>
 #include <ydb/library/yql/parser/pg_wrapper/parser.h>
+#include <ydb/library/yql/parser/pg_catalog/catalog.h>
 #include <ydb/library/yql/providers/common/provider/yql_provider_names.h>
 #include <ydb/library/yql/core/yql_callable_names.h>
 #include <ydb/library/yql/parser/pg_catalog/catalog.h>
@@ -267,6 +268,8 @@ public:
             return ParseInsertStmt(CAST_NODE(InsertStmt, node)) != nullptr;
         case T_ViewStmt:
             return ParseViewStmt(CAST_NODE(ViewStmt, node)) != nullptr;
+        case T_CreateStmt:
+            return ParseCreateStmt(CAST_NODE(CreateStmt, node)) != nullptr;
         case T_DropStmt:
             return ParseDropStmt(CAST_NODE(DropStmt, node)) != nullptr;
         case T_VariableSetStmt:
@@ -1020,6 +1023,262 @@ public:
         return Statements.back();
     }
 
+#pragma region CreateTable
+private:
+    struct TCreateTableCtx {
+        std::vector<TAstNode*> Columns;
+        std::unordered_set<TString> ColumnsSet;
+        std::vector<TAstNode*> PrimaryKey;
+        std::vector<TAstNode*> NotNullColumns;
+        std::unordered_set<TString> NotNullColSet;
+    };
+
+    bool CheckConstraintSupported(const Constraint* pk) {
+        bool isSupported = true;
+
+        if (pk->deferrable) {
+            AddError("DEFERRABLE constraints not supported");
+            isSupported = false;
+        }
+
+        if (pk->initdeferred) {
+            AddError("INITIALLY DEFERRED constraints not supported");
+            isSupported = false;
+        }
+
+        if (0 < ListLength(pk->including)) {
+            AddError("INCLUDING columns not supported");
+            isSupported = false;
+        }
+
+        if (0 < ListLength(pk->options)) {
+            AddError("WITH options not supported");
+            isSupported = false;
+        }
+
+        if (pk->indexname) {
+            AddError("INDEX name not supported");
+            isSupported = false;
+        }
+
+        if (pk->indexspace) {
+            AddError("USING INDEX TABLESPACE not supported");
+            isSupported = false;
+        }
+
+        return isSupported;
+    }
+
+    bool FillPrimaryKeyColumns(TCreateTableCtx& ctx, const Constraint* pk) {
+        if (!CheckConstraintSupported(pk))
+            return false;
+
+        for (auto i = 0; i < ListLength(pk->keys); ++i) {
+            auto node = ListNodeNth(pk->keys, i);
+
+            AddNonNullColumn(ctx, StrVal(node));
+            ctx.PrimaryKey.push_back(QA(StrVal(node)));
+        }
+
+        Y_ENSURE(0 < ctx.PrimaryKey.size());
+
+        return true;
+    }
+
+    bool AddNonNullColumn(TCreateTableCtx& ctx, const char* colName) {
+        auto [it, inserted] = ctx.NotNullColSet.insert(colName);
+        if (inserted)
+            ctx.NotNullColumns.push_back(QA(colName));
+
+        return inserted;
+    }
+
+    bool AddColumn(TCreateTableCtx& ctx, const ColumnDef* node) {
+        auto success = true;
+
+        if (node->constraints) {
+            for (ui32 i = 0; i < ListLength(node->constraints); ++i) {
+                auto constraintNode =
+                        CAST_NODE(Constraint, ListNodeNth(node->constraints, i));
+
+                switch (constraintNode->contype) {
+                    case CONSTR_NOTNULL:
+                        AddNonNullColumn(ctx, node->colname);
+                        break;
+
+                    case CONSTR_PRIMARY: {
+                        if (!ctx.PrimaryKey.empty()) {
+                            AddError("Only a single PK is allowed per table");
+                            success = false;
+                            break;
+                        }
+                        AddNonNullColumn(ctx, node->colname);
+                        ctx.PrimaryKey.push_back(QA(node->colname));
+                    } break;
+
+                    default:
+                        AddError("column constraint not supported");
+                        success = false;
+                }
+            }
+        }
+        auto [it, inserted] = ctx.ColumnsSet.insert(node->colname);
+        if (!inserted) {
+            AddError("duplicated column names found");
+            success = false;
+        }
+
+        if (!success)
+            return success;
+
+        // for now we pass just the last part of the type name
+        auto colType = StrVal( ListNodeNth(node->typeName->names,
+                                           ListLength(node->typeName->names) - 1));
+
+        ctx.Columns.push_back(
+                QL(QA(node->colname), L(A("PgType"), QA(colType)))
+                );
+
+        return success;
+    }
+
+    bool AddConstraint(TCreateTableCtx& ctx, const Constraint* node) {
+        auto success = true;
+
+        switch (node->contype) {
+            case CONSTR_PRIMARY: {
+                if (!ctx.PrimaryKey.empty()) {
+                    AddError("Only a single PK is allowed per table");
+                    success = false;
+                    break;
+                }
+                success &= FillPrimaryKeyColumns(ctx, node);
+            } break;
+
+            // TODO: support table-level not null constraints like:
+            // CHECK (col1 is not null [OR col2 is not null])
+
+            default:
+                AddError("table constraint not supported");
+                success = false;
+        }
+        return success;
+    }
+
+    TAstNode* BuildCreateTableOptions(TCreateTableCtx& ctx) {
+        return QL(
+                QL(QA("mode"), QA("create")),
+                QL(QA("columns"), QVL(ctx.Columns.data(), ctx.Columns.size())),
+                QL(QA("primarykey"), QVL(ctx.PrimaryKey.data(), ctx.PrimaryKey.size())),
+                QL(QA("notnull"), QVL(ctx.NotNullColumns.data(), ctx.NotNullColumns.size())));
+    }
+
+public:
+    [[nodiscard]]
+    TAstNode* ParseCreateStmt(const CreateStmt* value) {
+        auto success = true;
+
+        // See also transformCreateStmt() in parse_utilcmd.c
+        if (0 < ListLength(value->inhRelations)) {
+            AddError("table inheritance not supported");
+            success = false;
+        }
+
+        if (value->partspec) {
+            AddError("PARTITION BY clause not supported");
+            success = false;
+        }
+
+        if (value->partbound) {
+            AddError("FOR VALUES clause not supported");
+            success = false;
+        }
+
+        // if we ever support typed tables, check transformOfType() in parse_utilcmd.c
+        if (value->ofTypename) {
+            AddError("typed tables not supported");
+            success = false;
+        }
+
+        if (0 < ListLength(value->options)) {
+            AddError("table options not supported");
+            success = false;
+        }
+
+        if (value->oncommit != ONCOMMIT_NOOP && value->oncommit != ONCOMMIT_PRESERVE_ROWS) {
+            AddError("ON COMMIT actions not supported");
+            success = false;
+        }
+
+        if (value->tablespacename) {
+            AddError("TABLESPACE not supported");
+            success = false;
+        }
+
+        if (value->accessMethod) {
+            AddError("USING not supported");
+            success = false;
+        }
+
+        if (value->if_not_exists) {
+            AddError("IF NOT EXISTS not supported");
+            success = false;
+        }
+
+        { auto relPersistence = static_cast<NPg::ERelPersistence>(value->relation->relpersistence);
+        if (relPersistence != NPg::ERelPersistence::Permanent) {
+            switch (relPersistence) {
+                case NPg::ERelPersistence::Temp:
+                    AddError("CREATE TEMP TABLE not supported");
+                    break;
+
+                case NPg::ERelPersistence::Unlogged:
+                    AddError("UNLOGGED tables not supported");
+                    break;
+
+                default:
+                    Y_UNREACHABLE();
+            }
+            success = false;
+        }}
+
+        auto [sink, key] = ParseWriteRangeVar(value->relation, true);
+
+        if (!sink || !key)
+            success = false;
+
+        TCreateTableCtx ctx;
+
+        for (ui32 i = 0; i < ListLength(value->tableElts); ++i) {
+            auto rawNode = ListNodeNth(value->tableElts, i);
+
+            switch (NodeTag(rawNode)) {
+                case T_ColumnDef:
+                    success &= AddColumn(ctx, CAST_NODE(ColumnDef, rawNode));
+                    break;
+
+                case T_Constraint:
+                    success &= AddConstraint(ctx, CAST_NODE(Constraint, rawNode));
+                    break;
+
+                default:
+                    NodeNotImplemented(value, rawNode);
+                    success = false;
+            }
+        }
+
+        if (!success)
+            return nullptr;
+
+        Statements.push_back(
+                L(A("let"), A("world"),
+                  L(A("Write!"), A("world"), sink, key, L(A("Void")),
+                    BuildCreateTableOptions(ctx))));
+
+        return Statements.back();
+    }
+#pragma endregion CreateTable
+
     [[nodiscard]]
     TAstNode* ParseDropStmt(const DropStmt* value) {
         if (value->removeType != OBJECT_VIEW) {
@@ -1269,7 +1528,7 @@ public:
         return true;
     }
 
-    TWriteRangeDesc ParseWriteRangeVar(const RangeVar* value) {
+    TWriteRangeDesc ParseWriteRangeVar(const RangeVar* value, bool isScheme = false) {
         AT_LOCATION(value);
         if (StrLength(value->catalogname) > 0) {
             AddError("catalogname is not supported");
@@ -1294,7 +1553,7 @@ public:
         }
 
         auto sink = L(A("DataSink"), QAX(*p), QAX(schemaname));
-        auto key = L(A("Key"), QL(QA("table"), L(A("String"), QAX(TablePathPrefix + value->relname))));
+        auto key = L(A("Key"), QL(QA(isScheme ? "tablescheme" : "table"), L(A("String"), QAX(TablePathPrefix + value->relname))));
         return { sink, key };
     }
 
diff --git a/ydb/library/yql/sql/pg/pg_sql_ut.cpp b/ydb/library/yql/sql/pg/pg_sql_ut.cpp
index eb29a07cc85..728b11b65ca 100644
--- a/ydb/library/yql/sql/pg/pg_sql_ut.cpp
+++ b/ydb/library/yql/sql/pg/pg_sql_ut.cpp
@@ -30,13 +30,14 @@ NYql::TAstParseResult SqlToYqlWithMode(const TString& query, NSQLTranslation::ES
     settings.ClusterMapping[cluster] = service;
     settings.ClusterMapping["hahn"] = NYql::YtProviderName;
     settings.ClusterMapping["mon"] = NYql::SolomonProviderName;
+    settings.ClusterMapping[""] = NYql::KikimrProviderName;
     settings.MaxErrors = maxErrors;
     settings.Mode = mode;
     settings.Arena = &arena;
     settings.AnsiLexer = ansiLexer;
     settings.SyntaxVersion = 1;
-    auto q = TStringBuilder() << "--!syntax_pg\n" << query;
-    auto res = SqlToYql(q, settings);
+    settings.PgParser = true;
+    auto res = SqlToYql(query, settings);
     if (debug == EDebugOutput::ToCerr) {
         Err2Str(res, debug);
     }
@@ -70,4 +71,109 @@ Y_UNIT_TEST_SUITE(PgSqlParsingOnly) {
         const auto expectedAst = NYql::ParseAst(program);
         UNIT_ASSERT_STRINGS_EQUAL(res.Root->ToString(), expectedAst.Root->ToString());
     }
+
+    Y_UNIT_TEST(CreateTableStmt_Basic) {
+        auto res = PgSqlToYql("CREATE TABLE t (a int, b text)");
+        UNIT_ASSERT(res.Root);
+
+        TString program = R"(
+        (
+            (let world (Configure! world (DataSource 'config) 'OrderedColumns))
+            (let world (Write! world (DataSink '"kikimr" '"") (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('a (PgType 'int4)) '('b (PgType 'text)))) '('primarykey '()) '('notnull '()))))
+            (let world (CommitAll! world))
+            (return world)
+        )
+        )";
+        const auto expectedAst = NYql::ParseAst(program);
+        UNIT_ASSERT_STRINGS_EQUAL(res.Root->ToString(), expectedAst.Root->ToString());
+    }
+
+    Y_UNIT_TEST(CreateTableStmt_NotNull) {
+        auto res = PgSqlToYql("CREATE TABLE t (a int NOT NULL, b text)");
+        UNIT_ASSERT(res.Root);
+
+        TString program = R"(
+        (
+            (let world (Configure! world (DataSource 'config) 'OrderedColumns))
+            (let world (Write! world (DataSink '"kikimr" '"") (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('a (PgType 'int4)) '('b (PgType 'text)))) '('primarykey '()) '('notnull '('a)))))
+            (let world (CommitAll! world))
+            (return world)
+        )
+        )";
+        const auto expectedAst = NYql::ParseAst(program);
+        UNIT_ASSERT_STRINGS_EQUAL(res.Root->ToString(), expectedAst.Root->ToString());
+    }
+
+    Y_UNIT_TEST(CreateTableStmt_JustPK) {
+        auto res = PgSqlToYql("CREATE TABLE t (a int PRIMARY KEY, b text)");
+        UNIT_ASSERT(res.Root);
+
+        TString program = R"(
+        (
+            (let world (Configure! world (DataSource 'config) 'OrderedColumns))
+            (let world (Write! world (DataSink '"kikimr" '"") (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('a (PgType 'int4)) '('b (PgType 'text)))) '('primarykey '('a)) '('notnull '('a)))))
+            (let world (CommitAll! world))
+            (return world)
+        )
+        )";
+        const auto expectedAst = NYql::ParseAst(program);
+        UNIT_ASSERT_STRINGS_EQUAL(res.Root->ToString(), expectedAst.Root->ToString());
+    }
+
+    Y_UNIT_TEST(CreateTableStmt_PKAndNotNull) {
+        auto res = PgSqlToYql("CREATE TABLE t (a int PRIMARY KEY NOT NULL, b text)");
+        UNIT_ASSERT(res.Root);
+
+        TString program = R"(
+        (
+            (let world (Configure! world (DataSource 'config) 'OrderedColumns))
+            (let world (Write! world (DataSink '"kikimr" '"") (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('a (PgType 'int4)) '('b (PgType 'text)))) '('primarykey '('a)) '('notnull '('a)))))
+            (let world (CommitAll! world))
+            (return world)
+        )
+        )";
+        const auto expectedAst = NYql::ParseAst(program);
+        UNIT_ASSERT_STRINGS_EQUAL(res.Root->ToString(), expectedAst.Root->ToString());
+    }
+
+    Y_UNIT_TEST(CreateTableStmt_PKAndOtherNotNull) {
+        auto res = PgSqlToYql("CREATE TABLE t (a int PRIMARY KEY, b text NOT NULL)");
+        UNIT_ASSERT(res.Root);
+
+        TString program = R"(
+        (
+            (let world (Configure! world (DataSource 'config) 'OrderedColumns))
+            (let world (Write! world (DataSink '"kikimr" '"") (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('a (PgType 'int4)) '('b (PgType 'text)))) '('primarykey '('a)) '('notnull '('a 'b)))))
+            (let world (CommitAll! world))
+            (return world)
+        )
+        )";
+        const auto expectedAst = NYql::ParseAst(program);
+        UNIT_ASSERT_STRINGS_EQUAL(res.Root->ToString(), expectedAst.Root->ToString());
+    }
+
+    Y_UNIT_TEST(CreateTableStmt_TableLevelPK) {
+        auto res = PgSqlToYql("CREATE TABLE t (a int, b text NOT NULL, PRIMARY KEY (a, b))");
+        UNIT_ASSERT(res.Root);
+
+        TString program = R"(
+        (
+            (let world (Configure! world (DataSource 'config) 'OrderedColumns))
+            (let world (Write! world (DataSink '"kikimr" '"") (Key '('tablescheme (String '"t"))) (Void) '('('mode 'create) '('columns '('('a (PgType 'int4)) '('b (PgType 'text)))) '('primarykey '('a 'b)) '('notnull '('b 'a)))))
+            (let world (CommitAll! world))
+            (return world)
+        )
+        )";
+        const auto expectedAst = NYql::ParseAst(program);
+        UNIT_ASSERT_STRINGS_EQUAL(res.Root->ToString(), expectedAst.Root->ToString());
+    }
+
+    Y_UNIT_TEST(CreateTableStmt_RepeatingColumnNames) {
+        auto res = PgSqlToYql("CREATE TABLE t (a int, a text)");
+        UNIT_ASSERT(!res.Root);
+        UNIT_ASSERT_EQUAL(res.Issues.Size(), 1);
+
+        auto issue = *(res.Issues.begin());
+        UNIT_ASSERT(issue.GetMessage().find("duplicate") != TString::npos);
+    }
 }