From 43c73fe9c1ea50a3d428e1601cbf0e63b319cb24 Mon Sep 17 00:00:00 2001
From: Nikita Saveliev <savnik@ydb.tech>
Date: Tue, 22 Oct 2024 09:46:12 +0000
Subject: [PATCH] Allow short auto partitioning strategies yql

---
 ydb/core/kqp/provider/yql_kikimr_gateway.cpp | 18 ++++++++--
 ydb/services/persqueue_v1/topic_yql_ut.cpp   | 37 +++++++++++---------
 2 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/ydb/core/kqp/provider/yql_kikimr_gateway.cpp b/ydb/core/kqp/provider/yql_kikimr_gateway.cpp
index f43d42a98298..cae2db91ba9a 100644
--- a/ydb/core/kqp/provider/yql_kikimr_gateway.cpp
+++ b/ydb/core/kqp/provider/yql_kikimr_gateway.cpp
@@ -198,7 +198,7 @@ bool SetColumnType(const TTypeAnnotationNode* typeNode, bool notNull, Ydb::Type&
         const TStringBuf typeName = dataTypeNode->GetName();
         NUdf::EDataSlot dataSlot = NUdf::GetDataSlot(typeName);
         if (dataSlot == NUdf::EDataSlot::Decimal) {
-            auto dataExprTypeNode = typeNode->Cast<TDataExprParamsType>();  
+            auto dataExprTypeNode = typeNode->Cast<TDataExprParamsType>();
             ui32 precision = FromString(dataExprTypeNode->GetParamOne());
             ui32 scale = FromString(dataExprTypeNode->GetParamTwo());
             if (!NKikimr::NScheme::TDecimalType::Validate(precision, scale, error)) {
@@ -348,8 +348,20 @@ static std::shared_ptr<THashMap<TString, Ydb::Topic::Codec>> GetCodecsMapping()
 
 static std::shared_ptr<THashMap<TString, Ydb::Topic::AutoPartitioningStrategy>> GetAutoPartitioningStrategiesMapping() {
     static std::shared_ptr<THashMap<TString, Ydb::Topic::AutoPartitioningStrategy>> strategiesMapping;
-    if (strategiesMapping == nullptr) {
-        strategiesMapping = MakeEnumMapping<Ydb::Topic::AutoPartitioningStrategy>(Ydb::Topic::AutoPartitioningStrategy_descriptor(), "auto_partitioning_strategy_");
+    if (!strategiesMapping) {
+        strategiesMapping = MakeEnumMapping<Ydb::Topic::AutoPartitioningStrategy>(
+            Ydb::Topic::AutoPartitioningStrategy_descriptor(), "auto_partitioning_strategy_");
+
+        const TString prefix = "scale_";
+        for (const auto& [key, value] : *strategiesMapping) {
+            if (key.StartsWith(prefix)) {
+                TString newKey = key;
+                newKey.erase(0, prefix.length());
+
+                Y_ABORT_UNLESS(strategiesMapping->find(newKey) == strategiesMapping->end());
+                (*strategiesMapping)[newKey] = value;
+            }
+        }
     }
     return strategiesMapping;
 }
diff --git a/ydb/services/persqueue_v1/topic_yql_ut.cpp b/ydb/services/persqueue_v1/topic_yql_ut.cpp
index bda80c105b43..92534059a64a 100644
--- a/ydb/services/persqueue_v1/topic_yql_ut.cpp
+++ b/ydb/services/persqueue_v1/topic_yql_ut.cpp
@@ -23,31 +23,34 @@ Y_UNIT_TEST_SUITE(TTopicYqlTest) {
         UNIT_ASSERT_VALUES_EQUAL(after + 1, before);
     }
 
-    Y_UNIT_TEST(CreateAndAlterTopicYqlBackCompatibility) {
+    Y_UNIT_TEST(CreateTopicYqlBackCompatibility) {
         NKikimrConfig::TFeatureFlags ff;
         ff.SetEnableTopicSplitMerge(true);
         auto settings = NKikimr::NPersQueueTests::PQSettings();
         settings.SetFeatureFlags(ff);
 
         NPersQueue::TTestServer server(settings);
+        {
+            const char *query = R"__(
+                CREATE TOPIC `/Root/PQ/rt3.dc1--legacy--topic1` (
+                    CONSUMER c1
+                ) WITH (min_active_partitions = 2,
+                        partition_count_limit = 5,
+                        auto_partitioning_strategy = 'scale_up'
+                );
+            )__";
 
-        const char *query = R"__(
-            CREATE TOPIC `/Root/PQ/rt3.dc1--legacy--topic1` (
-                CONSUMER c1
-            ) WITH (min_active_partitions = 2,
-                    partition_count_limit = 5,
-                    auto_partitioning_strategy = 'scale_up'
-            );
-        )__";
+            server.AnnoyingClient->RunYqlSchemeQuery(query);
+            auto pqGroup = server.AnnoyingClient->Ls("/Root/PQ/rt3.dc1--legacy--topic1")->Record.GetPathDescription()
+                                                                                                .GetPersQueueGroup();
+            const auto& describeAfterCreate = pqGroup.GetPQTabletConfig();
+            Cerr <<"=== PATH DESCRIPTION: \n" << pqGroup.DebugString();
 
-        server.AnnoyingClient->RunYqlSchemeQuery(query);
-        auto pqGroup = server.AnnoyingClient->Ls("/Root/PQ/rt3.dc1--legacy--topic1")->Record.GetPathDescription()
-                                                                                            .GetPersQueueGroup();
-        const auto& describeAfterCreate = pqGroup.GetPQTabletConfig();
-        Cerr <<"=== PATH DESCRIPTION: \n" << pqGroup.DebugString();
+            UNIT_ASSERT_VALUES_EQUAL(describeAfterCreate.GetPartitionStrategy().GetMinPartitionCount(), 2);
+            UNIT_ASSERT_VALUES_EQUAL(describeAfterCreate.GetPartitionStrategy().GetMaxPartitionCount(), 5);
+            UNIT_ASSERT_VALUES_EQUAL(static_cast<int>(describeAfterCreate.GetPartitionStrategy().GetPartitionStrategyType()), static_cast<int>(::NKikimrPQ::TPQTabletConfig_TPartitionStrategyType::TPQTabletConfig_TPartitionStrategyType_CAN_SPLIT));
+        }
 
-        UNIT_ASSERT_VALUES_EQUAL(describeAfterCreate.GetPartitionStrategy().GetMinPartitionCount(), 2);
-        UNIT_ASSERT_VALUES_EQUAL(describeAfterCreate.GetPartitionStrategy().GetMaxPartitionCount(), 5);
     }
 
     Y_UNIT_TEST(CreateAndAlterTopicYql) {
@@ -77,7 +80,7 @@ Y_UNIT_TEST_SUITE(TTopicYqlTest) {
                     supported_codecs = 'RAW, GZIP',
                     partition_write_speed_bytes_per_second = 9000,
                     partition_write_burst_bytes = 100500,
-                    auto_partitioning_strategy = 'scale_up'
+                    auto_partitioning_strategy = 'up'
             );
         )__";