pg_syntax create table in ydb

6 files changed, 125 insertions, 11 deletions

diff --git a/ydb/core/kqp/provider/yql_kikimr_datasink.cpp b/ydb/core/kqp/provider/yql_kikimr_datasink.cpp
index ebbe3af7915..17bb92e34b9 100644
--- a/ydb/core/kqp/provider/yql_kikimr_datasink.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_datasink.cpp
@@ -508,12 +508,17 @@ public:
                         settings.PartitionBy = Build<TCoAtomList>(ctx, node->Pos()).Done();
                     }
 
+                    if (!settings.NotNullColumns.IsValid()) {
+                        settings.NotNullColumns = Build<TCoAtomList>(ctx, node->Pos()).Done();
+                    }
+
                     return Build<TKiCreateTable>(ctx, node->Pos())
                         .World(node->Child(0))
                         .DataSink(node->Child(1))
                         .Table().Build(key.GetTablePath())
                         .Columns(settings.Columns.Cast())
                         .PrimaryKey(settings.PrimaryKey.Cast())
+                        .NotNullColumns(settings.NotNullColumns.Cast())
                         .Settings(settings.Other)
                         .Indexes(settings.Indexes.Cast())
                         .Changefeeds(settings.Changefeeds.Cast())
diff --git a/ydb/core/kqp/provider/yql_kikimr_expr_nodes.json b/ydb/core/kqp/provider/yql_kikimr_expr_nodes.json
index fb404d36e46..8ec4e9a4821 100644
--- a/ydb/core/kqp/provider/yql_kikimr_expr_nodes.json
+++ b/ydb/core/kqp/provider/yql_kikimr_expr_nodes.json
@@ -118,7 +118,8 @@
                 {"Index": 8, "Name": "ColumnFamilies", "Type": "TExprList"},
                 {"Index": 9, "Name": "TableSettings", "Type": "TCoNameValueTupleList"},
                 {"Index": 10, "Name": "Changefeeds", "Type": "TCoChangefeedList"},
-                {"Index": 11, "Name": "TableType", "Type": "TCoAtom"}
+                {"Index": 11, "Name": "TableType", "Type": "TCoAtom"},
+                {"Index": 12, "Name": "NotNullColumns", "Type": "TCoAtomList"}
             ]
         },
         {
diff --git a/ydb/core/kqp/provider/yql_kikimr_type_ann.cpp b/ydb/core/kqp/provider/yql_kikimr_type_ann.cpp
index 9131d8f089f..46f5a3b0ef4 100644
--- a/ydb/core/kqp/provider/yql_kikimr_type_ann.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_type_ann.cpp
@@ -559,6 +559,11 @@ virtual TStatus HandleCreateTable(TKiCreateTable create, TExprContext& ctx) over
             meta->TableSettings.PartitionBy.emplace_back(column.Value());
         }
 
+        THashSet<TString> notNullColumns;
+        for (const auto& column : create.NotNullColumns()) {
+            notNullColumns.emplace(column.Value());
+        }
+
         for (auto item : create.Columns()) {
             auto columnTuple = item.Cast<TExprList>();
             auto nameNode = columnTuple.Item(0).Cast<TCoAtom>();
@@ -569,10 +574,31 @@ virtual TStatus HandleCreateTable(TKiCreateTable create, TExprContext& ctx) over
             YQL_ENSURE(columnType && columnType->GetKind() == ETypeAnnotationKind::Type);
 
             auto type = columnType->Cast<TTypeExprType>()->GetType();
-            auto notNull = type->GetKind() != ETypeAnnotationKind::Optional;
-            auto actualType = notNull ? type : type->Cast<TOptionalExprType>()->GetItemType();
-            if (
-                actualType->GetKind() != ETypeAnnotationKind::Data
+
+            auto isOptional = type->GetKind() == ETypeAnnotationKind::Optional;
+            auto actualType = !isOptional ? type : type->Cast<TOptionalExprType>()->GetItemType();
+
+            bool notNull;
+            if (actualType->GetKind() == ETypeAnnotationKind::Pg) {
+                if (notNullColumns.contains(columnName)) {
+                    if (std::find(meta->KeyColumnNames.begin(), meta->KeyColumnNames.end(), columnName) == meta->KeyColumnNames.end()) {
+                        ctx.AddError(TIssue(ctx.GetPosition(create.NotNullColumns().Pos()), TStringBuilder()
+                            << "notnull option for pg column " << columnName << " is forbidden"));
+                        return TStatus::Error;
+                    } else {
+                        //TODO: KIKIMR-17471
+                        //Right now YDB ignores the constraint native Postgres enforces
+                        //on primary key values; it should be used very carefully.
+                        ctx.AddWarning(TIssue(ctx.GetPosition(create.NotNullColumns().Pos()), TStringBuilder()
+                            << "notnull option for primary key column " << columnName << " will be ignored"));
+                    }
+                }
+                //TODO: set notnull for pg types
+                notNull = false;
+            } else {
+                notNull = !isOptional;
+            }
+            if (actualType->GetKind() != ETypeAnnotationKind::Data
                 && actualType->GetKind() != ETypeAnnotationKind::Pg
             ) {
                 columnTypeError(typeNode.Pos(), columnName, "Only YQL data types and PG types are currently supported");
@@ -596,7 +622,6 @@ virtual TStatus HandleCreateTable(TKiCreateTable create, TExprContext& ctx) over
                     NKikimr::NScheme::NTypeIds::Pg,
                     NKikimr::NPg::TypeDescFromPgTypeId(pgTypeId)
                 );
-                YQL_ENSURE(!notNull, "notNull is forbidden for pg types");
             }
             columnMeta.NotNull = notNull;
 
diff --git a/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp b/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
index 9f2d030874a..56603beb680 100644
--- a/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
+++ b/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
@@ -1121,7 +1121,7 @@ Y_UNIT_TEST_SUITE(KqpPg) {
         }
     }
 
-    Y_UNIT_TEST(CreateTable) {
+    Y_UNIT_TEST(V1CreateTable) {
         TKikimrRunner kikimr(NKqp::TKikimrSettings().SetWithSampleTables(false));
 
         auto testSingleType = [&kikimr] (const TPgTypeTestSpec& spec, bool isArray) {
@@ -1183,6 +1183,67 @@ Y_UNIT_TEST_SUITE(KqpPg) {
         }
     }
 
+    Y_UNIT_TEST(PgCreateTable) {
+        TKikimrRunner kikimr(NKqp::TKikimrSettings().SetWithSampleTables(false));
+
+        auto testSingleType = [&kikimr] (const TPgTypeTestSpec& spec, bool isArray) {
+            NYdb::NScripting::TScriptingClient client(kikimr.GetDriver());
+
+            auto tableName = "/Root/Pg" + ToString(spec.TypeId) + (isArray ? "array" : "");
+            auto typeName = ((isArray) ? "_" : "") + NYql::NPg::LookupType(spec.TypeId).Name;
+            auto keyEntry = spec.IsKey ? ("key "+ typeName) : "key int2";
+            auto valueEntry = "value " + typeName;
+            auto req = Sprintf("\
+            --!syntax_pg\n\
+            CREATE TABLE \"%s\" (\n\
+                %s PRIMARY KEY,\n\
+                %s\n\
+            );", tableName.Data(), keyEntry.Data(), valueEntry.Data());
+            Cerr << req << Endl;
+            auto result = client.ExecuteYqlScript(req).GetValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            if (!isArray) {
+                ExecutePgInsert(kikimr, tableName, spec);
+                result = ExecutePgSelect(kikimr, tableName);
+                ValidatePgYqlResult(result, spec);
+            } else {
+                ExecutePgArrayInsert(kikimr, tableName, spec);
+                result = ExecutePgSelect(kikimr, tableName);
+                TResultSetParser parser(result.GetResultSetParser(0));
+                for (size_t i = 0; parser.TryNextRow(); ++i) {
+                    auto check = [&parser, &spec] (const TString& column, const TString& expected) {
+                        auto& c = parser.ColumnParser(column);
+                        UNIT_ASSERT_VALUES_EQUAL(expected, c.GetPg().Content_);
+                    };
+                    TString expected = spec.TextOut(i);
+                    check("value", expected);
+                }
+            }
+        };
+
+        auto testType = [&] (const TPgTypeTestSpec& spec) {
+            auto textOutArray = [&spec] (auto i) {
+                auto str = spec.TextOut(i);
+                return spec.ArrayPrint(str);
+            };
+
+            TPgTypeTestSpec arraySpec{spec.TypeId, false, spec.TextIn, textOutArray};
+
+            testSingleType(spec, false);
+            testSingleType(arraySpec, true);
+        };
+
+        for (const auto& spec : typeSpecs) {
+            Cerr << spec.TypeId << Endl;
+            if (spec.TypeId == CHAROID) {
+                continue;
+                // I cant come up with a query with explicit char conversion.
+                // ::char, ::character casts to pg_bpchar
+            }
+            testType(spec);
+        }
+    }
+
     Y_UNIT_TEST(CreateNotNullPgColumn) {
         TKikimrRunner kikimr(NKqp::TKikimrSettings().SetWithSampleTables(false));
 
@@ -1191,16 +1252,32 @@ Y_UNIT_TEST_SUITE(KqpPg) {
 
         NYdb::NScripting::TScriptingClient client(kikimr.GetDriver());
         auto req = TStringBuilder() << R"(
+        --!syntax_pg
+        CREATE TABLE "/Root/Pg" (
+        key int2 PRIMARY KEY,
+        value int2 NOT NULL
+        );)";
+        Cerr << req << Endl;
+        auto result = client.ExecuteYqlScript(req).GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
+        UNIT_ASSERT_NO_DIFF(result.GetIssues().ToString(), "<main>: Error: Type annotation, code: 1030\n"
+        "    <main>:1:1: Error: At function: KiCreateTable!\n"
+        "        <main>:1:1: Error: notnull option for primary key column key will be ignored\n"
+        "        <main>:1:1: Error: notnull option for pg column value is forbidden\n");
+
+        TString reqV1 = TStringBuilder() << R"(
         --!syntax_v1
         CREATE TABLE `/Root/Pg` (
         key pg_int2,
         value pg_int2 NOT NULL,
         PRIMARY KEY (key)
         );)";
-        Cerr << req << Endl;
-        auto result = client.ExecuteYqlScript(req).GetValueSync();
-        UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::INTERNAL_ERROR);
-        UNIT_ASSERT(result.GetIssues().begin()->GetMessage().EndsWith("notNull is forbidden for pg types"));
+        Cerr << reqV1 << Endl;
+        result = client.ExecuteYqlScript(reqV1).GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL(result.GetStatus(), EStatus::GENERIC_ERROR);
+                UNIT_ASSERT_NO_DIFF(result.GetIssues().ToString(), "<main>: Error: Type annotation, code: 1030\n"
+        "    <main>:6:22: Error: At function: KiCreateTable!\n"
+        "        <main>:6:22: Error: notnull option for pg column value is forbidden\n");
     }
 
     Y_UNIT_TEST(ValuesInsert) {
diff --git a/ydb/library/yql/providers/common/provider/yql_provider.cpp b/ydb/library/yql/providers/common/provider/yql_provider.cpp
index f6f9f5c6e22..809dfe35cbf 100644
--- a/ydb/library/yql/providers/common/provider/yql_provider.cpp
+++ b/ydb/library/yql/providers/common/provider/yql_provider.cpp
@@ -206,6 +206,7 @@ TWriteTableSettings ParseWriteTableSettings(TExprList node, TExprContext& ctx) {
     TMaybeNode<TCoAtom> mode;
     TMaybeNode<TExprList> columns;
     TMaybeNode<TCoAtomList> primaryKey;
+    TMaybeNode<TCoAtomList> notNullColumns;
     TMaybeNode<TCoAtomList> partitionBy;
     TMaybeNode<TCoNameValueTupleList> orderBy;
     TMaybeNode<TCoLambda> filter;
@@ -294,6 +295,9 @@ TWriteTableSettings ParseWriteTableSettings(TExprList node, TExprContext& ctx) {
             } else if (name == "tableType") {
                 YQL_ENSURE(tuple.Value().Maybe<TCoAtom>());
                 tableType = tuple.Value().Cast<TCoAtom>();
+            } else if (name == "notnull") {
+                YQL_ENSURE(tuple.Value().Maybe<TCoAtomList>());
+                notNullColumns = tuple.Value().Cast<TCoAtomList>();
             } else {
                 other.push_back(tuple);
             }
@@ -328,6 +332,7 @@ TWriteTableSettings ParseWriteTableSettings(TExprList node, TExprContext& ctx) {
     ret.Mode = mode;
     ret.Columns = columns;
     ret.PrimaryKey = primaryKey;
+    ret.NotNullColumns = notNullColumns;
     ret.PartitionBy = partitionBy;
     ret.OrderBy = orderBy;
     ret.Filter = filter;
diff --git a/ydb/library/yql/providers/common/provider/yql_provider.h b/ydb/library/yql/providers/common/provider/yql_provider.h
index 87242eab37f..b7541edaed0 100644
--- a/ydb/library/yql/providers/common/provider/yql_provider.h
+++ b/ydb/library/yql/providers/common/provider/yql_provider.h
@@ -36,6 +36,7 @@ struct TWriteTableSettings {
     NNodes::TMaybeNode<NNodes::TCoAtom> Mode;
     NNodes::TMaybeNode<NNodes::TExprList> Columns;
     NNodes::TMaybeNode<NNodes::TCoAtomList> PrimaryKey;
+    NNodes::TMaybeNode<NNodes::TCoAtomList> NotNullColumns;
     NNodes::TMaybeNode<NNodes::TCoAtomList> PartitionBy;
     NNodes::TMaybeNode<NNodes::TCoNameValueTupleList> OrderBy;
     NNodes::TMaybeNode<NNodes::TCoLambda> Filter;