KIKIMR-18545: add optional<bool> not_null into describe API

7 files changed, 99 insertions, 32 deletions

diff --git a/ydb/core/kqp/ut/opt/kqp_not_null_ut.cpp b/ydb/core/kqp/ut/opt/kqp_not_null_ut.cpp
index 6ddd08b7719..ca29e5a5fba 100644
--- a/ydb/core/kqp/ut/opt/kqp_not_null_ut.cpp
+++ b/ydb/core/kqp/ut/opt/kqp_not_null_ut.cpp
@@ -1,6 +1,7 @@
 #include <ydb/core/kqp/ut/common/kqp_ut_common.h>
 #include <ydb/core/tx/tx_proxy/proxy.h>
 #include <ydb/core/tx/schemeshard/schemeshard.h>
+#include <ydb/public/sdk/cpp/client/ydb_proto/accessor.h>
 
 namespace NKikimr {
 namespace NKqp {
@@ -1480,7 +1481,7 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST(Describe) {
-        TKikimrRunner kikimr;
+        TKikimrRunner kikimr(NKqp::TKikimrSettings().SetWithSampleTables(false));
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -1488,7 +1489,9 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
             CREATE TABLE `/Root/DescribeTest` (
                 Key1 Uint64 NOT NULL,
                 Key2 Uint64,
-                Value String,
+                Value1 String,
+                Value2 PgInt2,
+                Value3 PgInt2 NOT NULL,
                 PRIMARY KEY (Key1, Key2)
             );
         )")).ExtractValueSync();
@@ -1500,12 +1503,60 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
         const THashMap<std::string_view, std::string_view> columnTypes = {
             {"Key1", "Uint64"},
             {"Key2", "Uint64?"},
-            {"Value", "String?"}
+            {"Value1", "String?"},
+            {"Value2", "Pg('pgint2','',21,0,0)"},
+            {"Value3", "Pg('pgint2','',21,0,0)"}
+        };
+
+        const THashMap<std::string_view, bool> columnNullability = {
+            {"Key1", true},
+            {"Key2", false},
+            {"Value1", false},
+            {"Value2", false},
+            {"Value3", true}
         };
 
         const auto& columns = describeTableResult.GetTableDescription().GetTableColumns();
         for (const auto& column : columns) {
-            UNIT_ASSERT_VALUES_EQUAL(column.Type.ToString(), columnTypes.at(column.Name));
+            UNIT_ASSERT_VALUES_EQUAL_C(column.Type.ToString(), columnTypes.at(column.Name), column.Name);
+            UNIT_ASSERT_VALUES_EQUAL_C(column.NotNull.value(), columnNullability.at(column.Name), column.Name);
+        }
+    }
+
+    Y_UNIT_TEST(CreateTableWithNotNullColumns) {
+        TKikimrRunner kikimr(NKqp::TKikimrSettings().SetWithSampleTables(false));
+        auto client = kikimr.GetTableClient();
+        auto session = client.CreateSession().GetValueSync().GetSession();
+        {
+            TTableBuilder builder;
+            builder.AddNullableColumn("1", EPrimitiveType::Uint64);
+            builder.AddNonNullableColumn("2", EPrimitiveType::Uint64);
+            builder.AddNullableColumn("3", TPgType("pgint2"));
+            builder.AddNonNullableColumn("4", TPgType("pgint2"));
+            builder.SetPrimaryKeyColumn("1");
+            auto result = session.CreateTable("/Root/NotNullCheck", builder.Build()).GetValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+        auto describeTableResult = session.DescribeTable("/Root/NotNullCheck").GetValueSync();
+        UNIT_ASSERT_C(describeTableResult.IsSuccess(), describeTableResult.GetIssues().ToString());
+        const THashMap<std::string_view, bool> columnNullability = {
+            {"1", false},
+            {"2", true},
+            {"3", false},
+            {"4", true},
+        };
+
+        const auto& columns = describeTableResult.GetTableDescription().GetTableColumns();
+        for (const auto& column : columns) {
+            UNIT_ASSERT_VALUES_EQUAL_C(column.NotNull.value(), columnNullability.at(column.Name), column.Name);
+        }
+
+        {
+            auto proto = NYdb::TProtoAccessor::GetProto(describeTableResult.GetTableDescription());
+            proto.mutable_columns()->begin()->set_not_null(true);
+            auto result = session.CreateTable("/Root/NotNullCheck2", TTableDescription(std::move(proto), {})).GetValueSync();
+            UNIT_ASSERT_C(!result.GetIssues().Empty(), "ok with faulty protobuf");
+            UNIT_ASSERT(result.GetIssues().ToString().Contains("Error: Not consistent column type and not_null option for column: 1"));
         }
     }
 
diff --git a/ydb/core/tx/schemeshard/ut_export.cpp b/ydb/core/tx/schemeshard/ut_export.cpp
index cc078fc6b1f..a86c54227b0 100644
--- a/ydb/core/tx/schemeshard/ut_export.cpp
+++ b/ydb/core/tx/schemeshard/ut_export.cpp
@@ -345,6 +345,7 @@ Y_UNIT_TEST_SUITE(TExportToS3Tests) {
       }
     }
   }
+  not_null: false
 }
 columns {
   name: "value"
@@ -355,6 +356,7 @@ columns {
       }
     }
   }
+  not_null: false
 }
 primary_key: "key"
 storage_settings {
@@ -1044,7 +1046,7 @@ partitioning_settings {
             });
             runtime.DispatchEvents(opts);
         }
-        
+
         THolder<IEventHandle> proposeTxResult;
         runtime.SetObserverFunc([&](TTestActorRuntimeBase&, TAutoPtr<IEventHandle>& ev) {
             if (ev->GetTypeRewrite() == TEvDataShard::EvProposeTransactionResult) {
diff --git a/ydb/core/ydb_convert/table_description.cpp b/ydb/core/ydb_convert/table_description.cpp
index 65fb0f675fe..6beb381edd1 100644
--- a/ydb/core/ydb_convert/table_description.cpp
+++ b/ydb/core/ydb_convert/table_description.cpp
@@ -59,7 +59,6 @@ static Ydb::Type* AddColumn(Ydb::Table::ColumnMeta* newColumn, const TColumn& co
         } else {
             columnType = newColumn->mutable_type()->mutable_optional_type()->mutable_item();
         }
-
         Y_ENSURE(columnType);
         if (protoType == NYql::NProto::TypeIds::Decimal) {
             auto typeParams = columnType->mutable_decimal_type();
@@ -70,6 +69,8 @@ static Ydb::Type* AddColumn(Ydb::Table::ColumnMeta* newColumn, const TColumn& co
             NMiniKQL::ExportPrimitiveTypeToProto(protoType, *columnType);
         }
     }
+    newColumn->set_not_null(column.GetNotNull());
+
     return columnType;
 }
 
@@ -120,10 +121,6 @@ void FillColumnDescriptionImpl(TYdbProto& out,
 
     for (const auto& column : in.GetColumns()) {
         auto newColumn = out.add_columns();
-        Y_ENSURE(
-            column.GetTypeId() != NScheme::NTypeIds::Pg || !column.GetNotNull(),
-            "It is not allowed to create NOT NULL column with pg type"
-        );
         Ydb::Type* columnType = AddColumn(newColumn, column);
 
         if (columnIdToKeyPos.count(column.GetId())) {
@@ -259,16 +256,23 @@ bool FillColumnDescription(NKikimrSchemeOp::TTableDescription& out,
     for (const auto& column : in) {
         NKikimrSchemeOp:: TColumnDescription* cd = out.AddColumns();
         cd->SetName(column.name());
-        if (!column.type().has_optional_type()) {
-            if (!AppData()->FeatureFlags.GetEnableNotNullColumns()) {
-                status = Ydb::StatusIds::UNSUPPORTED;
-                error = "Not null columns feature is not supported yet";
-                return false;
-            }
-
+        bool notOptional = !column.type().has_optional_type();
+        if (!column.has_not_null()) {
             if (!column.type().has_pg_type()) {
-                cd->SetNotNull(true);
+                cd->SetNotNull(notOptional);
+            }
+        } else {
+            if (!column.type().has_pg_type() && notOptional != column.not_null()) {
+                status = Ydb::StatusIds::BAD_REQUEST;
+                error = "Not consistent column type and not_null option for column: " + column.name();
+                return false;
             }
+            cd->SetNotNull(column.not_null());
+        }
+        if (cd->GetNotNull() && !AppData()->FeatureFlags.GetEnableNotNullColumns()) {
+            status = Ydb::StatusIds::UNSUPPORTED;
+            error = "Not null columns feature is not supported yet";
+            return false;
         }
 
         NScheme::TTypeInfo typeInfo;
diff --git a/ydb/public/api/protos/ydb_table.proto b/ydb/public/api/protos/ydb_table.proto
index 82074236f17..6ac181353a5 100644
--- a/ydb/public/api/protos/ydb_table.proto
+++ b/ydb/public/api/protos/ydb_table.proto
@@ -330,6 +330,8 @@ message ColumnMeta {
     Type type = 2;
     // Column family name of the column
     string family = 3;
+    // Column nullability
+    optional bool not_null = 4;
 }
 
 message DateTypeColumnModeSettings {
@@ -1134,4 +1136,3 @@ message ExecuteScanQueryPartialResult {
     reserved 5; // query_plan
     Ydb.TableStats.QueryStats query_stats = 6;
 }
-
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.cpp b/ydb/public/sdk/cpp/client/ydb_table/table.cpp
index 10209fce5b9..de3e4ebed32 100644
--- a/ydb/public/sdk/cpp/client/ydb_table/table.cpp
+++ b/ydb/public/sdk/cpp/client/ydb_table/table.cpp
@@ -262,7 +262,7 @@ class TTableDescription::TImpl {
 
         // columns
         for (const auto& col : proto.columns()) {
-            Columns_.emplace_back(col.name(), col.type(), col.family());
+            Columns_.emplace_back(col.name(), col.type(), col.family(), col.not_null());
         }
 
         // indexes
@@ -426,8 +426,8 @@ public:
         return Proto_;
     }
 
-    void AddColumn(const TString& name, const Ydb::Type& type, const TString& family) {
-        Columns_.emplace_back(name, type, family);
+    void AddColumn(const TString& name, const Ydb::Type& type, const TString& family, std::optional<bool> notNull) {
+        Columns_.emplace_back(name, type, family, notNull);
     }
 
     void SetPrimaryKeyColumns(const TVector<TString>& primaryKeyColumns) {
@@ -685,8 +685,8 @@ const TVector<TKeyRange>& TTableDescription::GetKeyRanges() const {
     return Impl_->GetKeyRanges();
 }
 
-void TTableDescription::AddColumn(const TString& name, const Ydb::Type& type, const TString& family) {
-    Impl_->AddColumn(name, type, family);
+void TTableDescription::AddColumn(const TString& name, const Ydb::Type& type, const TString& family, std::optional<bool> notNull) {
+    Impl_->AddColumn(name, type, family, notNull);
 }
 
 void TTableDescription::SetPrimaryKeyColumns(const TVector<TString>& primaryKeyColumns) {
@@ -835,6 +835,9 @@ void TTableDescription::SerializeTo(Ydb::Table::CreateTableRequest& request) con
         protoColumn.set_name(column.Name);
         protoColumn.mutable_type()->CopyFrom(TProtoAccessor::GetProto(column.Type));
         protoColumn.set_family(column.Family);
+        if (column.NotNull.has_value()) {
+            protoColumn.set_not_null(column.NotNull.value());
+        }
     }
 
     for (const auto& pk : Impl_->GetPrimaryKeyColumns()) {
@@ -1038,7 +1041,7 @@ TTableBuilder& TTableBuilder::AddNullableColumn(const TString& name, const EPrim
         .EndOptional()
         .Build();
 
-    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family);
+    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family, false);
     return *this;
 }
 
@@ -1048,7 +1051,7 @@ TTableBuilder& TTableBuilder::AddNullableColumn(const TString& name, const TDeci
             .Decimal(type)
         .EndOptional()
         .Build();
-    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family);
+    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family, false);
     return *this;
 }
 
@@ -1057,7 +1060,7 @@ TTableBuilder& TTableBuilder::AddNullableColumn(const TString& name, const TPgTy
         .Pg(type)
         .Build();
 
-    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family);
+    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family, false);
     return *this;
 }
 
@@ -1066,7 +1069,7 @@ TTableBuilder& TTableBuilder::AddNonNullableColumn(const TString& name, const EP
         .Primitive(type)
         .Build();
 
-    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family);
+    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family, true);
     return *this;
 }
 
@@ -1075,7 +1078,7 @@ TTableBuilder& TTableBuilder::AddNonNullableColumn(const TString& name, const TD
         .Decimal(type)
         .Build();
 
-    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family);
+    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family, true);
     return *this;
 }
 
@@ -1084,7 +1087,7 @@ TTableBuilder& TTableBuilder::AddNonNullableColumn(const TString& name, const TP
         .Pg(type)
         .Build();
 
-    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family);
+    TableDescription_.AddColumn(name, TProtoAccessor::GetProto(columnType), family, true);
     return *this;
 }
 
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.h b/ydb/public/sdk/cpp/client/ydb_table/table.h
index 02016e4c500..38765041e6d 100644
--- a/ydb/public/sdk/cpp/client/ydb_table/table.h
+++ b/ydb/public/sdk/cpp/client/ydb_table/table.h
@@ -90,13 +90,15 @@ struct TTableColumn {
     TString Name;
     TType Type;
     TString Family;
+    std::optional<bool> NotNull;
 
     TTableColumn() = default;
 
-    TTableColumn(TString name, TType type, TString family = TString())
+    TTableColumn(TString name, TType type, TString family = TString(), std::optional<bool> notNull = std::nullopt)
         : Name(std::move(name))
         , Type(std::move(type))
         , Family(std::move(family))
+        , NotNull(std::move(notNull))
     { }
 
     // Conversion from TColumn for API compatibility
@@ -521,7 +523,7 @@ private:
     TTableDescription();
     explicit TTableDescription(const Ydb::Table::CreateTableRequest& request);
 
-    void AddColumn(const TString& name, const Ydb::Type& type, const TString& family);
+    void AddColumn(const TString& name, const Ydb::Type& type, const TString& family, std::optional<bool> notNull);
     void SetPrimaryKeyColumns(const TVector<TString>& primaryKeyColumns);
 
     // common
diff --git a/ydb/services/ydb/ydb_ut.cpp b/ydb/services/ydb/ydb_ut.cpp
index 526d8a58379..8668b0f8b28 100644
--- a/ydb/services/ydb/ydb_ut.cpp
+++ b/ydb/services/ydb/ydb_ut.cpp
@@ -1262,6 +1262,7 @@ columns {
       }
     }
   }
+  not_null: false
 }
 columns {
   name: "Value"
@@ -1272,6 +1273,7 @@ columns {
       }
     }
   }
+  not_null: false
 }
 primary_key: "Key"
 partitioning_settings {
@@ -1599,6 +1601,7 @@ columns {
       }
     }
   }
+  not_null: false
 }
 columns {
   name: "IValue"
@@ -1609,6 +1612,7 @@ columns {
       }
     }
   }
+  not_null: false
 }
 primary_key: "Key"
 indexes {