disable not null data columns

fix(kqp): disable not null data columns

7 files changed, 167 insertions, 53 deletions

diff --git a/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp b/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp
index 7a03f49a796..be7f3c9f5b6 100644
--- a/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp
+++ b/ydb/core/kqp/ut/kqp_not_null_columns_ut.cpp
@@ -7,6 +7,36 @@ using namespace NYdb;
 using namespace NYdb::NTable;
 
 Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
+    Y_UNIT_TEST_NEW_ENGINE(CreateTableWithDisabledNotNullDataColumns) {
+        TKikimrRunner kikimr;
+        auto client = kikimr.GetTableClient();
+        auto session = client.CreateSession().GetValueSync().GetSession();
+
+        {
+            const auto query = Q_(R"(
+                CREATE TABLE `/Root/TestCreateTable` (
+                    Key Uint64 NOT NULL,
+                    Value String NOT NULL,
+                    PRIMARY KEY (Key))
+            )");
+
+            auto result = session.ExecuteSchemeQuery(query).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString());
+        }
+
+        {
+            const auto query = Q_(R"(
+                CREATE TABLE `/Root/TestCreateTable` (
+                    Key Uint64 NOT NULL,
+                    Value String,
+                    PRIMARY KEY (Key))
+            )");
+
+            auto result = session.ExecuteSchemeQuery(query).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+    }
+
     Y_UNIT_TEST_NEW_ENGINE(InsertNotNullPk) {
         TKikimrRunner kikimr;
         auto client = kikimr.GetTableClient();
@@ -170,7 +200,10 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST_NEW_ENGINE(SelectNotNullColumns) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -218,7 +251,10 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST_NEW_ENGINE(InsertNotNull) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -256,7 +292,10 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST_NEW_ENGINE(UpsertNotNull) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -294,7 +333,10 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST_NEW_ENGINE(ReplaceNotNull) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -332,7 +374,10 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST_NEW_ENGINE(UpdateNotNull) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -381,7 +426,10 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST_NEW_ENGINE(UpdateOnNotNull) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -443,19 +491,15 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
         {
             const auto query = Q_("ALTER TABLE `/Root/TestAddNotNullColumn` ADD COLUMN Value2 String NOT NULL");
             auto result = session.ExecuteSchemeQuery(query).ExtractValueSync();
-            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
-        }
-
-        {  /* set NULL to not null column */
-            const auto query = Q_("UPSERT INTO `/Root/TestAddNotNullColumn` (Key, Value1, Value2) VALUES (1, 'Value1', NULL)");
-            auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx()).ExtractValueSync();
-            UNIT_ASSERT(!result.IsSuccess());
-            UNIT_ASSERT_C(HasIssue(result.GetIssues(), NYql::TIssuesIds::KIKIMR_BAD_COLUMN_TYPE), result.GetIssues().ToString());
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString());
         }
     }
 
     Y_UNIT_TEST_NEW_ENGINE(AlterDropNotNullColumn) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -480,7 +524,10 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST_NEW_ENGINE(FailedMultiEffects) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
@@ -516,13 +563,52 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
         }
     }
 
-    Y_UNIT_TEST_NEW_ENGINE(SecondaryKeyWithNotNullColumn) {
+    Y_UNIT_TEST_NEW_ENGINE(CreateIndexedTableWithDisabledNotNullDataColumns) {
         TKikimrRunner kikimr;
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
         {
             const auto query = Q1_(R"(
+                CREATE TABLE `/Root/TestCreateIndexedTable` (
+                    Key Uint64 NOT NULL,
+                    SecondaryKey Uint64,
+                    Value String NOT NULL,
+                    PRIMARY KEY (Key),
+                    INDEX Index GLOBAL ON (SecondaryKey)
+                    COVER (Value))
+            )");
+
+            auto result = session.ExecuteSchemeQuery(query).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString());
+        }
+
+        {
+            const auto query = Q1_(R"(
+                CREATE TABLE `/Root/TestCreateIndexedTable` (
+                    Key Uint64 NOT NULL,
+                    SecondaryKey Uint64,
+                    Value String,
+                    PRIMARY KEY (Key),
+                    INDEX Index GLOBAL ON (SecondaryKey)
+                    COVER (Value))
+            )");
+
+            auto result = session.ExecuteSchemeQuery(query).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        }
+    }
+
+    Y_UNIT_TEST_NEW_ENGINE(SecondaryKeyWithNotNullColumn) {
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
+        auto client = kikimr.GetTableClient();
+        auto session = client.CreateSession().GetValueSync().GetSession();
+
+        {
+            const auto query = Q1_(R"(
                 CREATE TABLE `/Root/TestNotNullSecondaryKey` (
                     Key1 Uint64 NOT NULL,
                     Key2 Uint64 NOT NULL,
@@ -616,7 +702,10 @@ Y_UNIT_TEST_SUITE(KqpNotNullColumns) {
     }
 
     Y_UNIT_TEST_NEW_ENGINE(SecondaryIndexWithNotNullDataColumn) {
-        TKikimrRunner kikimr;
+        auto settings = TKikimrSettings()
+            .SetEnableNotNullDataColumns(true);
+
+        TKikimrRunner kikimr(settings);
         auto client = kikimr.GetTableClient();
         auto session = client.CreateSession().GetValueSync().GetSession();
 
diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto
index 1ffcbe49d40..45dfa91e214 100644
--- a/ydb/core/protos/config.proto
+++ b/ydb/core/protos/config.proto
@@ -714,6 +714,7 @@ message TFeatureFlags {
     // enable http handle for self termination
     optional bool EnableFailureInjectionTermination = 71 [default = false];
     optional bool EnableChunkLocking = 72 [default = false];
+    optional bool EnableNotNullDataColumns = 73 [default = false];
 }
 
 
diff --git a/ydb/core/testlib/basics/feature_flags.h b/ydb/core/testlib/basics/feature_flags.h
index a2bc92ed6de..c12bb5657da 100644
--- a/ydb/core/testlib/basics/feature_flags.h
+++ b/ydb/core/testlib/basics/feature_flags.h
@@ -38,6 +38,7 @@ public:
     FEATURE_FLAG_SETTER(EnableKqpScanQueryStreamLookup)
     FEATURE_FLAG_SETTER(EnableMoveIndex)
     FEATURE_FLAG_SETTER(EnablePredicateExtractForDataQueries)
+    FEATURE_FLAG_SETTER(EnableNotNullDataColumns)
 
     TDerived& SetEnableMvcc(std::optional<bool> value) {
         if (value) {
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_alter_table.cpp b/ydb/core/tx/schemeshard/schemeshard__operation_alter_table.cpp
index 5378350ecf3..30ecac43b2d 100644
--- a/ydb/core/tx/schemeshard/schemeshard__operation_alter_table.cpp
+++ b/ydb/core/tx/schemeshard/schemeshard__operation_alter_table.cpp
@@ -87,6 +87,11 @@ TTableInfo::TAlterDataPtr ParseParams(const TPath& path, TTableInfo::TPtr table,
 
     // Ignore column ids if they were passed by user!
     for (auto& col : *copyAlter.MutableColumns()) {
+        if (col.GetNotNull()) {
+            errStr = Sprintf("Not null columns is not supported for alter command");
+            status = NKikimrScheme::StatusInvalidParameter;
+            return nullptr;
+        }
         col.ClearId();
     }
 
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_create_indexed_table.cpp b/ydb/core/tx/schemeshard/schemeshard__operation_create_indexed_table.cpp
index b115d0f22e5..a789376f988 100644
--- a/ydb/core/tx/schemeshard/schemeshard__operation_create_indexed_table.cpp
+++ b/ydb/core/tx/schemeshard/schemeshard__operation_create_indexed_table.cpp
@@ -216,9 +216,18 @@ TVector<ISubOperationBase::TPtr> CreateIndexedTable(TOperationId nextId, const T
     }
 
     for (auto& column : baseTableDescription.GetColumns()) {
-        if (column.GetNotNull() && !AppData()->FeatureFlags.GetEnableNotNullColumns()) {
-            TString msg = TStringBuilder() << "It is not allowed to create not null column";
-            return {CreateReject(nextId, NKikimrScheme::EStatus::StatusPreconditionFailed, msg)};
+        if (column.GetNotNull()) {
+            bool isPrimaryKey =  keys.contains(column.GetName());
+
+            if (isPrimaryKey && !AppData()->FeatureFlags.GetEnableNotNullColumns()) {
+                TString msg = TStringBuilder() << "It is not allowed to create not null pk: " << column.GetName();
+                return {CreateReject(nextId, NKikimrScheme::EStatus::StatusPreconditionFailed, msg)};
+            }
+
+            if (!isPrimaryKey && !AppData()->FeatureFlags.GetEnableNotNullDataColumns()) {
+                TString msg = TStringBuilder() << "It is not allowed to create not null data column: " << column.GetName();
+                return {CreateReject(nextId, NKikimrScheme::EStatus::StatusPreconditionFailed, msg)};
+            }
         }
 
         if (column.HasDefaultFromSequence()) {
diff --git a/ydb/core/tx/schemeshard/schemeshard__operation_create_table.cpp b/ydb/core/tx/schemeshard/schemeshard__operation_create_table.cpp
index 07825405ee4..7dd80607321 100644
--- a/ydb/core/tx/schemeshard/schemeshard__operation_create_table.cpp
+++ b/ydb/core/tx/schemeshard/schemeshard__operation_create_table.cpp
@@ -19,6 +19,28 @@ void PrepareScheme(NKikimrSchemeOp::TTableDescription& schema) {
     schema.ClearKeyColumnIds();
 }
 
+bool CheckColumnTypesConstraints(NKikimrSchemeOp::TTableDescription& desc, TString& errMsg) {
+    THashSet<TString> keyColumns(desc.GetKeyColumnNames().begin(), desc.GetKeyColumnNames().end());
+
+    for (const auto& column : desc.GetColumns()) {
+        if (column.GetNotNull()) {
+            bool isPrimaryKey = keyColumns.contains(column.GetName());
+
+            if (isPrimaryKey && !AppData()->FeatureFlags.GetEnableNotNullColumns()) {
+                errMsg = TStringBuilder() << "It is not allowed to create not null pk: " << column.GetName();
+                return false;
+            }
+
+            if (!isPrimaryKey && !AppData()->FeatureFlags.GetEnableNotNullDataColumns()) {
+                errMsg = TStringBuilder() << "It is not allowed to create not null data column: " << column.GetName();
+                return false;
+            }
+        }
+    }
+
+    return true;
+}
+
 bool InitPartitioning(const NKikimrSchemeOp::TTableDescription& op,
                       const NScheme::TTypeRegistry* typeRegistry,
                       const TVector<ui32>& keyColIds,
@@ -542,6 +564,11 @@ public:
 
         TString errStr;
 
+        if (!CheckColumnTypesConstraints(schema, errStr)) {
+            result->SetError(NKikimrScheme::StatusPreconditionFailed, errStr);
+            return result;
+        }
+
         NKikimrSchemeOp::TPartitionConfig compilationPartitionConfig;
         if (!TPartitionConfigMerger::ApplyChanges(compilationPartitionConfig, TPartitionConfigMerger::DefaultConfig(AppData()), schema.GetPartitionConfig(), AppData(), errStr)
             || !TPartitionConfigMerger::VerifyCreateParams(compilationPartitionConfig, AppData(), IsShadowDataAllowed(), errStr)) {
diff --git a/ydb/services/ydb/ydb_bulk_upsert_ut.cpp b/ydb/services/ydb/ydb_bulk_upsert_ut.cpp
index 7baa87871dc..1b7655e1580 100644
--- a/ydb/services/ydb/ydb_bulk_upsert_ut.cpp
+++ b/ydb/services/ydb/ydb_bulk_upsert_ut.cpp
@@ -230,7 +230,7 @@ Y_UNIT_TEST_SUITE(YdbTableBulkUpsert) {
 
         TString tableName = "/Root/TestNotNullColumns";
 
-        {
+        {  /* create table with disabled not null data column */
             auto tableBuilder = client.GetTableBuilder();
             tableBuilder
                     .AddNonNullableColumn("Key", EPrimitiveType::Uint64)
@@ -238,6 +238,19 @@ Y_UNIT_TEST_SUITE(YdbTableBulkUpsert) {
                 .SetPrimaryKeyColumns({"Key"});
             auto result = session.CreateTable(tableName, tableBuilder.Build()).ExtractValueSync();
 
+            Cerr << result.GetIssues().ToString() << Endl;
+            UNIT_ASSERT_EQUAL(result.IsTransportError(), false);
+            UNIT_ASSERT_EQUAL(result.GetStatus(), EStatus::PRECONDITION_FAILED);
+        }
+
+        {  /* create table with not null primary key column */
+            auto tableBuilder = client.GetTableBuilder();
+            tableBuilder
+                    .AddNonNullableColumn("Key", EPrimitiveType::Uint64)
+                    .AddNullableColumn("Value", EPrimitiveType::Uint64)
+                .SetPrimaryKeyColumns({"Key"});
+            auto result = session.CreateTable(tableName, tableBuilder.Build()).ExtractValueSync();
+
             Cerr << result.GetIssues().ToString();
             UNIT_ASSERT_EQUAL(result.IsTransportError(), false);
             UNIT_ASSERT_EQUAL(result.GetStatus(), EStatus::SUCCESS);
@@ -274,21 +287,6 @@ Y_UNIT_TEST_SUITE(YdbTableBulkUpsert) {
             UNIT_ASSERT_EQUAL(res.GetStatus(), EStatus::SCHEME_ERROR);
         }
 
-        {  /* missing not null value column */
-            TValueBuilder rows;
-            rows.BeginList();
-            rows.AddListItem()
-                .BeginStruct()
-                    .AddMember("Key").Uint64(2)
-                .EndStruct();
-            rows.EndList();
-
-            auto res = client.BulkUpsert(tableName, rows.Build()).GetValueSync();
-            Cerr << res.GetIssues().ToString();
-            UNIT_ASSERT_STRING_CONTAINS(res.GetIssues().ToString(), "Missing not null columns: Value");
-            UNIT_ASSERT_EQUAL(res.GetStatus(), EStatus::SCHEME_ERROR);
-        }
-
         {  /* set null to not null primary key column */
             TValueBuilder rows;
             rows.BeginList();
@@ -305,22 +303,6 @@ Y_UNIT_TEST_SUITE(YdbTableBulkUpsert) {
             UNIT_ASSERT_EQUAL(res.GetStatus(), EStatus::BAD_REQUEST);
         }
 
-        {  /* set null to not null value column */
-            TValueBuilder rows;
-            rows.BeginList();
-            rows.AddListItem()
-                .BeginStruct()
-                    .AddMember("Key").Uint64(2)
-                    .AddMember("Value").EmptyOptional(EPrimitiveType::Uint64)
-                .EndStruct();
-            rows.EndList();
-
-            auto res = client.BulkUpsert(tableName, rows.Build()).GetValueSync();
-            Cerr << res.GetIssues().ToString();
-            UNIT_ASSERT_STRING_CONTAINS(res.GetIssues().ToString(), "Received NULL value for not null column");
-            UNIT_ASSERT_EQUAL(res.GetStatus(), EStatus::BAD_REQUEST);
-        }
-
         {
             auto result = session.DropTable(tableName).ExtractValueSync();