fix alter table call in scripting service KIKIMR-19889

2 files changed, 83 insertions, 3 deletions

diff --git a/ydb/core/kqp/host/kqp_gateway_proxy.cpp b/ydb/core/kqp/host/kqp_gateway_proxy.cpp
index 6c65fa03ff..97a2f8ddca 100644
--- a/ydb/core/kqp/host/kqp_gateway_proxy.cpp
+++ b/ydb/core/kqp/host/kqp_gateway_proxy.cpp
@@ -658,7 +658,11 @@ public:
 
         if (!IsPrepare()) {
             SessionCtx->Query().PrepareOnly = false;
-            if (SessionCtx->Query().PreparingQuery) {
+            if (!SessionCtx->Query().PreparingQuery) {
+                SessionCtx->Query().PreparingQuery = std::make_unique<NKikimrKqp::TPreparedQuery>();
+            }
+
+            if (SessionCtx->Query().PreparingQuery->MutablePhysicalQuery()->GetTransactions().size() > 0) {
                 auto code = Ydb::StatusIds::BAD_REQUEST;
                 auto error = TStringBuilder() << "multiple transactions are not supported for alter table operation.";
                 IKqpGateway::TGenericResult errResult;
@@ -667,8 +671,6 @@ public:
                 tablePromise.SetValue(errResult);
                 return tablePromise.GetFuture();
             }
-
-            SessionCtx->Query().PreparingQuery = std::make_unique<NKikimrKqp::TPreparedQuery>();
         }
 
         auto prepareFuture = PrepareAlterTable(cluster, std::move(req), requestType, flags);
diff --git a/ydb/core/kqp/ut/yql/kqp_scripting_ut.cpp b/ydb/core/kqp/ut/yql/kqp_scripting_ut.cpp
index 1301e84fbb..10d3f56e5b 100644
--- a/ydb/core/kqp/ut/yql/kqp_scripting_ut.cpp
+++ b/ydb/core/kqp/ut/yql/kqp_scripting_ut.cpp
@@ -46,6 +46,84 @@ Y_UNIT_TEST_SUITE(KqpScripting) {
         UNIT_ASSERT_VALUES_EQUAL(rs0.ColumnParser(0).GetUint64(), 2u);
     }
 
+    Y_UNIT_TEST(ScriptingCreateAndAlterTableTest) {
+        TKikimrRunner kikimr;
+        TScriptingClient client(kikimr.GetDriver());
+
+        auto result = client.ExecuteYqlScript(R"(
+            CREATE TABLE `/Root/ScriptingCreateAndAlterTableTest` (
+                Key Uint64,
+                Value String,
+                PRIMARY KEY (Key)
+            );
+            COMMIT;
+
+            REPLACE INTO `/Root/ScriptingCreateAndAlterTableTest` (Key, Value) VALUES
+                (1, "One"),
+                (2, "Two");
+        )").GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 0);
+
+        result = client.ExecuteYqlScript(R"(
+            SELECT COUNT(*) FROM `/Root/ScriptingCreateAndAlterTableTest`;
+        )").GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+        UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
+        UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets()[0].RowsCount(), 1);
+        TResultSetParser rs0(result.GetResultSets()[0]);
+        UNIT_ASSERT(rs0.TryNextRow());
+        UNIT_ASSERT_VALUES_EQUAL(rs0.ColumnParser(0).GetUint64(), 2u);
+
+        result = client.ExecuteYqlScript(R"(
+            ALTER TABLE `/Root/ScriptingCreateAndAlterTableTest` SET (AUTO_PARTITIONING_BY_SIZE = ENABLED);
+        )").GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+    
+        result = client.ExecuteYqlScript(R"(
+            ALTER TABLE `/Root/ScriptingCreateAndAlterTableTest` SET (AUTO_PARTITIONING_BY_SIZE = ENABLED);
+            COMMIT;
+            ALTER TABLE `/Root/ScriptingCreateAndAlterTableTest` SET (AUTO_PARTITIONING_PARTITION_SIZE_MB = 500);
+            COMMIT;
+            ALTER TABLE `/Root/ScriptingCreateAndAlterTableTest` SET (AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = 4);
+        )").GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+    
+        result = client.ExecuteYqlScript(R"(
+            ALTER TABLE `/Root/ScriptingCreateAndAlterTableTest` SET (AUTO_PARTITIONING_BY_SIZE = ENABLED);
+            COMMIT;
+            CREATE TABLE `/Root/ScriptingCreateAndAlterTableTest2` (
+                Key Uint64,
+                Value String,
+                PRIMARY KEY (Key)
+            );
+            COMMIT;
+        )").GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        result = client.ExecuteYqlScript(R"(
+            SELECT COUNT(*) FROM `/Root/ScriptingCreateAndAlterTableTest2`;
+        )").GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        result = client.ExecuteYqlScript(R"(
+            CREATE TABLE `/Root/ScriptingCreateAndAlterTableTest3` (
+                Key Uint64,
+                Value String,
+                PRIMARY KEY (Key)
+            );
+            COMMIT;
+            ALTER TABLE `/Root/ScriptingCreateAndAlterTableTest3` SET (AUTO_PARTITIONING_BY_SIZE = ENABLED);
+            COMMIT;
+        )").GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        result = client.ExecuteYqlScript(R"(
+            SELECT COUNT(*) FROM `/Root/ScriptingCreateAndAlterTableTest3`;
+        )").GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+    }
+
     Y_UNIT_TEST(UnsafeTimestampCast) {
         auto setting = NKikimrKqp::TKqpSetting();
         setting.SetName("_KqpYqlSyntaxVersion");