customize reading from bindings in kqp

10 files changed, 155 insertions, 1 deletions

diff --git a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
index a79f5a096b2..c7cc1b608db 100644
--- a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
+++ b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
@@ -32,6 +32,23 @@ using namespace NThreading;
 using namespace NYql;
 using namespace NYql::NDq;
 
+namespace {
+NSQLTranslation::EBindingsMode RemapBindingsMode(NKikimrConfig::TTableServiceConfig::EBindingsMode mode) {
+    switch (mode) {
+    case NKikimrConfig::TTableServiceConfig::BM_ENABLED:
+        return NSQLTranslation::EBindingsMode::ENABLED;
+    case NKikimrConfig::TTableServiceConfig::BM_DISABLED:
+        return NSQLTranslation::EBindingsMode::DISABLED;
+    case NKikimrConfig::TTableServiceConfig::BM_DROP_WITH_WARNING:
+        return NSQLTranslation::EBindingsMode::DROP_WITH_WARNING;
+    case NKikimrConfig::TTableServiceConfig::BM_DROP:
+        return NSQLTranslation::EBindingsMode::DROP;
+    default:
+        return NSQLTranslation::EBindingsMode::ENABLED;
+    }
+}
+}
+
 class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
 public:
     using TBase = TActorBootstrapped<TKqpCompileActor>;
@@ -388,6 +405,7 @@ void ApplyServiceConfig(TKikimrConfiguration& kqpConfig, const TTableServiceConf
     kqpConfig.EnablePredicateExtractForScanQuery = serviceConfig.GetEnablePredicateExtractForScanQueries();
     kqpConfig.EnableSequentialReads = serviceConfig.GetEnableSequentialReads();
     kqpConfig.EnableKqpImmediateEffects = serviceConfig.GetEnableKqpImmediateEffects();
+    kqpConfig.BindingsMode = RemapBindingsMode(serviceConfig.GetBindingsMode());
 }
 
 IActor* CreateKqpCompileActor(const TActorId& owner, const TKqpSettings::TConstPtr& kqpSettings,
diff --git a/ydb/core/kqp/host/kqp_host.cpp b/ydb/core/kqp/host/kqp_host.cpp
index 566fa324311..18ddfed7391 100644
--- a/ydb/core/kqp/host/kqp_host.cpp
+++ b/ydb/core/kqp/host/kqp_host.cpp
@@ -1079,7 +1079,11 @@ private:
                 settings.V0Behavior = NSQLTranslation::EV0Behavior::Silent;
             }
 
-            settings.DynamicClusterProvider = SessionCtx->Config().FeatureFlags.GetEnableExternalDataSources() ? NYql::KikimrProviderName : TString{};
+            if (SessionCtx->Config().FeatureFlags.GetEnableExternalDataSources()) {
+                settings.DynamicClusterProvider = NYql::KikimrProviderName;
+                settings.BindingsMode = SessionCtx->Config().BindingsMode;
+            }
+
             settings.InferSyntaxVersion = true;
             settings.V0ForceDisable = false;
             settings.WarnOnV0 = false;
diff --git a/ydb/core/kqp/provider/CMakeLists.darwin-x86_64.txt b/ydb/core/kqp/provider/CMakeLists.darwin-x86_64.txt
index d78b2eaf38f..4b843449f59 100644
--- a/ydb/core/kqp/provider/CMakeLists.darwin-x86_64.txt
+++ b/ydb/core/kqp/provider/CMakeLists.darwin-x86_64.txt
@@ -46,6 +46,7 @@ target_link_libraries(core-kqp-provider PUBLIC
   providers-dq-expr_nodes
   providers-result-expr_nodes
   providers-result-provider
+  yql-sql-settings
   tools-enum_parser-enum_serialization_runtime
 )
 target_sources(core-kqp-provider PRIVATE
diff --git a/ydb/core/kqp/provider/CMakeLists.linux-aarch64.txt b/ydb/core/kqp/provider/CMakeLists.linux-aarch64.txt
index d8cfb2499f9..46d277a850a 100644
--- a/ydb/core/kqp/provider/CMakeLists.linux-aarch64.txt
+++ b/ydb/core/kqp/provider/CMakeLists.linux-aarch64.txt
@@ -47,6 +47,7 @@ target_link_libraries(core-kqp-provider PUBLIC
   providers-dq-expr_nodes
   providers-result-expr_nodes
   providers-result-provider
+  yql-sql-settings
   tools-enum_parser-enum_serialization_runtime
 )
 target_sources(core-kqp-provider PRIVATE
diff --git a/ydb/core/kqp/provider/CMakeLists.linux-x86_64.txt b/ydb/core/kqp/provider/CMakeLists.linux-x86_64.txt
index d8cfb2499f9..46d277a850a 100644
--- a/ydb/core/kqp/provider/CMakeLists.linux-x86_64.txt
+++ b/ydb/core/kqp/provider/CMakeLists.linux-x86_64.txt
@@ -47,6 +47,7 @@ target_link_libraries(core-kqp-provider PUBLIC
   providers-dq-expr_nodes
   providers-result-expr_nodes
   providers-result-provider
+  yql-sql-settings
   tools-enum_parser-enum_serialization_runtime
 )
 target_sources(core-kqp-provider PRIVATE
diff --git a/ydb/core/kqp/provider/CMakeLists.windows-x86_64.txt b/ydb/core/kqp/provider/CMakeLists.windows-x86_64.txt
index d78b2eaf38f..4b843449f59 100644
--- a/ydb/core/kqp/provider/CMakeLists.windows-x86_64.txt
+++ b/ydb/core/kqp/provider/CMakeLists.windows-x86_64.txt
@@ -46,6 +46,7 @@ target_link_libraries(core-kqp-provider PUBLIC
   providers-dq-expr_nodes
   providers-result-expr_nodes
   providers-result-provider
+  yql-sql-settings
   tools-enum_parser-enum_serialization_runtime
 )
 target_sources(core-kqp-provider PRIVATE
diff --git a/ydb/core/kqp/provider/ya.make b/ydb/core/kqp/provider/ya.make
index c7807553adb..c25d4d5aa8e 100644
--- a/ydb/core/kqp/provider/ya.make
+++ b/ydb/core/kqp/provider/ya.make
@@ -45,6 +45,7 @@ PEERDIR(
     ydb/library/yql/providers/dq/expr_nodes
     ydb/library/yql/providers/result/expr_nodes
     ydb/library/yql/providers/result/provider
+    ydb/library/yql/sql/settings
 )
 
 YQL_LAST_ABI_VERSION()
diff --git a/ydb/core/kqp/provider/yql_kikimr_settings.h b/ydb/core/kqp/provider/yql_kikimr_settings.h
index cf62a2740e4..c48673f4a13 100644
--- a/ydb/core/kqp/provider/yql_kikimr_settings.h
+++ b/ydb/core/kqp/provider/yql_kikimr_settings.h
@@ -3,6 +3,7 @@
 #include <ydb/library/yql/dq/common/dq_common.h>
 #include <ydb/library/yql/providers/common/config/yql_dispatch.h>
 #include <ydb/library/yql/providers/common/config/yql_setting.h>
+#include <ydb/library/yql/sql/settings/translation_settings.h>
 #include <ydb/core/protos/config.pb.h>
 
 namespace NYql {
@@ -143,6 +144,7 @@ struct TKikimrConfiguration : public TKikimrSettings, public NCommon::TSettingDi
     bool EnablePredicateExtractForDataQuery = false;
     bool EnableKqpImmediateEffects = false;
     bool EnableSequentialReads = false;
+    NSQLTranslation::EBindingsMode BindingsMode = NSQLTranslation::EBindingsMode::ENABLED;
 };
 
 }
diff --git a/ydb/core/kqp/ut/federated_query/kqp_federated_query_ut.cpp b/ydb/core/kqp/ut/federated_query/kqp_federated_query_ut.cpp
index 4fe2c48bf21..02596b6e123 100644
--- a/ydb/core/kqp/ut/federated_query/kqp_federated_query_ut.cpp
+++ b/ydb/core/kqp/ut/federated_query/kqp_federated_query_ut.cpp
@@ -537,6 +537,122 @@ Y_UNIT_TEST_SUITE(KqpFederatedQuery) {
         UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(0).GetUtf8(), "2");
         UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(1).GetUtf8(), "hello world");
     }
+
+    std::pair<NYdb::NQuery::TScriptExecutionOperation, TFetchScriptResultsResult> ExecuteScriptOverBinding(NKikimrConfig::TTableServiceConfig::EBindingsMode mode) {
+        using namespace fmt::literals;
+        const TString externalDataSourceName = "/Root/external_data_source";
+        const TString externalTableName = "/Root/test_binding_resolve";
+        const TString bucket = "test_bucket1";
+        const TString object = "test_object";
+
+        CreateBucketWithObject(bucket, object, TEST_CONTENT);
+
+        NKikimrConfig::TAppConfig appConfig;
+        appConfig.MutableTableServiceConfig()->SetBindingsMode(mode);
+
+        NKikimrConfig::TFeatureFlags featureFlags;
+        featureFlags.SetEnableExternalDataSources(true);
+        featureFlags.SetEnableScriptExecutionOperations(true);
+
+        auto settings = TKikimrSettings()
+            .SetAppConfig(appConfig)
+            .SetFeatureFlags(featureFlags);
+
+        TKikimrRunner kikimr(settings);
+
+        auto tc = kikimr.GetTableClient();
+        auto session = tc.CreateSession().GetValueSync().GetSession();
+        const TString query = fmt::format(R"(
+            CREATE EXTERNAL DATA SOURCE `{external_source}` WITH (
+                SOURCE_TYPE="ObjectStorage",
+                LOCATION="{location}",
+                AUTH_METHOD="NONE"
+            );
+            CREATE EXTERNAL TABLE `{external_table}` (
+                key Utf8 NOT NULL,
+                value Utf8 NOT NULL
+            ) WITH (
+                DATA_SOURCE="{external_source}",
+                LOCATION="{object}",
+                FORMAT="json_each_row"
+            );)",
+            "external_source"_a = externalDataSourceName,
+            "external_table"_a = externalTableName,
+            "location"_a = GetEnv("S3_ENDPOINT") + "/" + bucket + "/",
+            "object"_a = object
+            );
+        auto result = session.ExecuteSchemeQuery(query).GetValueSync();
+        UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+
+        const TString sql = fmt::format(R"(
+                SELECT * FROM bindings.`{external_table}`
+            )", "external_table"_a=externalTableName);
+
+        auto db = kikimr.GetQueryClient();
+        auto scriptExecutionOperation = db.ExecuteScript(sql).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(scriptExecutionOperation.Status().GetStatus(), EStatus::SUCCESS, scriptExecutionOperation.Status().GetIssues().ToString());
+        UNIT_ASSERT(scriptExecutionOperation.Metadata().ExecutionId);
+
+        NYdb::NQuery::TScriptExecutionOperation readyOp = WaitScriptExecutionOperation(scriptExecutionOperation.Id(), kikimr.GetDriver());
+        TFetchScriptResultsResult results(TStatus(EStatus::SUCCESS, {}));
+        if (readyOp.Metadata().ExecStatus == EExecStatus::Completed) {
+            results = db.FetchScriptResults(scriptExecutionOperation.Metadata().ExecutionId, 0).ExtractValueSync();
+            UNIT_ASSERT_C(results.IsSuccess(), results.GetIssues().ToString());
+        }
+        return {readyOp, results};
+    }
+
+    Y_UNIT_TEST(ExecuteScriptWithDifferentBindingsMode) {
+        {
+            auto [readyOp, results] = ExecuteScriptOverBinding(NKikimrConfig::TTableServiceConfig::BM_DROP);
+            UNIT_ASSERT_EQUAL(readyOp.Metadata().ExecStatus, EExecStatus::Completed);
+            UNIT_ASSERT_VALUES_EQUAL(readyOp.Status().GetIssues().ToString(), "");
+
+            TResultSetParser resultSet(results.ExtractResultSet());
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnsCount(), 2);
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.RowsCount(), 2);
+
+            UNIT_ASSERT(resultSet.TryNextRow());
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(0).GetUtf8(), "1");
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(1).GetUtf8(), "trololo");
+
+            UNIT_ASSERT(resultSet.TryNextRow());
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(0).GetUtf8(), "2");
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(1).GetUtf8(), "hello world");
+
+        }
+
+        {
+            auto [readyOp, results] = ExecuteScriptOverBinding(NKikimrConfig::TTableServiceConfig::BM_DROP_WITH_WARNING);
+            UNIT_ASSERT_EQUAL(readyOp.Metadata().ExecStatus, EExecStatus::Completed);
+            UNIT_ASSERT_VALUES_EQUAL(readyOp.Status().GetIssues().ToString(), "<main>:2:31: Warning: Please remove 'bindings.' from your query, the support for this syntax will be dropped soon, code: 4538\n");
+
+            TResultSetParser resultSet(results.ExtractResultSet());
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnsCount(), 2);
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.RowsCount(), 2);
+
+            UNIT_ASSERT(resultSet.TryNextRow());
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(0).GetUtf8(), "1");
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(1).GetUtf8(), "trololo");
+
+            UNIT_ASSERT(resultSet.TryNextRow());
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(0).GetUtf8(), "2");
+            UNIT_ASSERT_VALUES_EQUAL(resultSet.ColumnParser(1).GetUtf8(), "hello world");
+        }
+
+        {
+            auto [readyOp, results] = ExecuteScriptOverBinding(NKikimrConfig::TTableServiceConfig::BM_ENABLED);
+            UNIT_ASSERT_EQUAL(readyOp.Metadata().ExecStatus, EExecStatus::Failed);
+            UNIT_ASSERT_VALUES_EQUAL(readyOp.Status().GetIssues().ToString(), "<main>:2:40: Error: Table binding `/Root/test_binding_resolve` is not defined\n");
+        }
+
+        {
+            auto [readyOp, results] = ExecuteScriptOverBinding(NKikimrConfig::TTableServiceConfig::BM_DISABLED);
+            UNIT_ASSERT_EQUAL(readyOp.Metadata().ExecStatus, EExecStatus::Failed);
+            UNIT_ASSERT_VALUES_EQUAL(readyOp.Status().GetIssues().ToString(), "<main>:2:31: Error: Please remove 'bindings.' from your query, the support for this syntax has ended, code: 4601\n");
+        }
+    }
+
 }
 
 } // namespace NKqp
diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto
index c17c13715cd..f2df45d4ec1 100644
--- a/ydb/core/protos/config.proto
+++ b/ydb/core/protos/config.proto
@@ -1296,6 +1296,15 @@ message TTableServiceConfig {
     optional bool EnablePredicateExtractForDataQueries = 37 [default = true];
     optional bool EnableSequentialReads = 38 [default = false];
     optional bool EnableKqpImmediateEffects = 39 [default = true];
+
+    enum EBindingsMode {
+        BM_ENABLED = 0;
+        BM_DISABLED = 1;
+        BM_DROP_WITH_WARNING = 2;
+        BM_DROP = 3;
+    }
+
+    optional EBindingsMode BindingsMode = 40 [default = BM_ENABLED];
 };
 
 // Config describes immediate controls and allows