YQ-2715 add ydb data source

Woodey777 · Woodey777 · commit 2c06274d8733 · 2024-01-12T13:00:07.000Z
diff --git a/ydb/core/external_sources/external_source_factory.cpp b/ydb/core/external_sources/external_source_factory.cpp
@@ -46,6 +46,10 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector<TStri
         {
             ToString(NYql::EDatabaseType::PostgreSQL),
             CreateExternalDataSource(TString{NYql::GenericProviderName}, {"MDB_BASIC", "BASIC"}, {"database_name", "protocol", "mdb_cluster_id", "use_tls", "schema"}, hostnamePatternsRegEx)
+        },
+        {
+            ToString(NYql::EDatabaseType::Ydb),
+            CreateExternalDataSource(TString{NYql::GenericProviderName}, {"BASIC", "SERVICE_ACCOUNT"}, {"database_name", "use_tls", "database_id"}, hostnamePatternsRegEx)
         }
     });
 }
diff --git a/ydb/core/kqp/ut/federated_query/generic_ut/kqp_generic_provider_ut.cpp b/ydb/core/kqp/ut/federated_query/generic_ut/kqp_generic_provider_ut.cpp
@@ -31,6 +31,7 @@ namespace NKikimr::NKqp {
     enum class EProviderType {
         PostgreSQL,
         ClickHouse,
+        Ydb,
     };
 
     NApi::TDataSourceInstance MakeDataSourceInstance(EProviderType providerType) {
@@ -39,6 +40,8 @@ namespace NKikimr::NKqp {
                 return TConnectorClientMock::TPostgreSQLDataSourceInstanceBuilder<>().GetResult();
             case EProviderType::ClickHouse:
                 return TConnectorClientMock::TClickHouseDataSourceInstanceBuilder<>().GetResult();
+            case EProviderType::Ydb:
+                return TConnectorClientMock::TYdbDataSourceInstanceBuilder<>().GetResult();
         }
     }
 
@@ -48,6 +51,8 @@ namespace NKikimr::NKqp {
                 return CreatePostgreSQLExternalDataSource(kikimr);
             case EProviderType::ClickHouse:
                 return CreateClickHouseExternalDataSource(kikimr);
+            case EProviderType::Ydb:
+                return CreateYdbExternalDataSource(kikimr);
         }
     }
 
@@ -165,6 +170,10 @@ namespace NKikimr::NKqp {
             TestSelectAllFields(EProviderType::ClickHouse);
         }
 
+        Y_UNIT_TEST(YdbManaged) {
+            TestSelectAllFields(EProviderType::Ydb);
+        }
+
         void TestSelectConstant(EProviderType providerType) {
             // prepare mock
             auto clientMock = std::make_shared<TConnectorClientMock>();
@@ -257,6 +266,10 @@ namespace NKikimr::NKqp {
             TestSelectConstant(EProviderType::ClickHouse);
         }
 
+        Y_UNIT_TEST(YdbManagedSelectConstant) {
+            TestSelectConstant(EProviderType::Ydb);
+        }
+
         void TestSelectCount(EProviderType providerType) {
             // prepare mock
             auto clientMock = std::make_shared<TConnectorClientMock>();
@@ -345,6 +358,10 @@ namespace NKikimr::NKqp {
             TestSelectCount(EProviderType::ClickHouse);
         }
 
+        Y_UNIT_TEST(YdbSelectCount) {
+            TestSelectCount(EProviderType::Ydb);
+        }
+
         void TestFilterPushdown(EProviderType providerType) {
             // prepare mock
             auto clientMock = std::make_shared<TConnectorClientMock>();
@@ -450,5 +467,9 @@ namespace NKikimr::NKqp {
         Y_UNIT_TEST(ClickHouseFilterPushdown) {
             TestFilterPushdown(EProviderType::ClickHouse);
         }
+
+        Y_UNIT_TEST(YdbFilterPushdown) {
+            TestFilterPushdown(EProviderType::Ydb);
+        }
     }
 }
diff --git a/ydb/library/yql/providers/common/db_id_async_resolver/db_async_resolver.h b/ydb/library/yql/providers/common/db_id_async_resolver/db_async_resolver.h
@@ -22,6 +22,8 @@ inline EDatabaseType DatabaseTypeFromDataSourceKind(NConnector::NApi::EDataSourc
             return EDatabaseType::PostgreSQL;
         case NConnector::NApi::EDataSourceKind::CLICKHOUSE:
             return EDatabaseType::ClickHouse;
+        case NConnector::NApi::EDataSourceKind::YDB:
+            return EDatabaseType::Ydb;
         default:
             ythrow yexception() << "Unknown data source kind: " << NConnector::NApi::EDataSourceKind_Name(dataSourceKind);
     }
@@ -33,6 +35,8 @@ inline NConnector::NApi::EDataSourceKind DatabaseTypeToDataSourceKind(EDatabaseT
             return  NConnector::NApi::EDataSourceKind::POSTGRESQL;
         case EDatabaseType::ClickHouse:
             return  NConnector::NApi::EDataSourceKind::CLICKHOUSE;
+        case EDatabaseType::Ydb:
+            return  NConnector::NApi::EDataSourceKind::YDB;
         default:
             ythrow yexception() << "Unknown database type: " << ToString(databaseType);
     }
diff --git a/ydb/library/yql/providers/generic/actors/yql_generic_source_factory.cpp b/ydb/library/yql/providers/generic/actors/yql_generic_source_factory.cpp
@@ -16,7 +16,7 @@ namespace NYql::NDq {
                                           args.SecureParams, args.TaskParams, args.ComputeActorId, credentialsFactory, args.HolderFactory);
         };
 
-        for (auto& sourceName : {"ClickHouseGeneric", "PostgreSqlGeneric"}) {
+        for (auto& sourceName : {"ClickHouseGeneric", "PostgreSqlGeneric", "YdbGeneric"}) {
             factory.RegisterSource<Generic::TSource>(sourceName, genericFactory);
         }
     }
diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.cpp b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.cpp
@@ -107,6 +107,43 @@ namespace NYql::NConnector::NTest {
         UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
     }
 
+    void CreateYdbExternalDataSource(
+        const std::shared_ptr<NKikimr::NKqp::TKikimrRunner>& kikimr,
+        const TString& dataSourceName,
+        const TString& login,
+        const TString& password,
+        const TString& endpoint,
+        bool useTls,
+        const TString& databaseName)
+    {
+        auto tc = kikimr->GetTableClient();
+        auto session = tc.CreateSession().GetValueSync().GetSession();
+        const TString query = fmt::format(
+            R"(
+            CREATE OBJECT sa_signature (TYPE SECRET) WITH (value=sa_signature);
+            CREATE OBJECT {data_source_name}_password (TYPE SECRET) WITH (value={password});
+
+            CREATE EXTERNAL DATA SOURCE {data_source_name} WITH (
+                SOURCE_TYPE="{source_type}",
+                LOCATION="{endpoint}",
+                AUTH_METHOD="BASIC",
+                LOGIN="{login}",
+                PASSWORD_SECRET_NAME="{data_source_name}_password",
+                USE_TLS="{use_tls}",
+                DATABASE_NAME="{database}"
+            );
+        )",
+            "data_source_name"_a = dataSourceName,
+            "login"_a = login,
+            "password"_a = password,
+            "use_tls"_a = useTls ? "TRUE" : "FALSE",
+            "source_type"_a = ToString(NYql::EDatabaseType::Ydb),
+            "endpoint"_a = endpoint,
+            "database"_a = databaseName);
+        auto result = session.ExecuteSchemeQuery(query).GetValueSync();
+        UNIT_ASSERT_C(result.GetStatus() == NYdb::EStatus::SUCCESS, result.GetIssues().ToString());
+    }
+
     std::shared_ptr<arrow::RecordBatch> MakeEmptyRecordBatch(size_t rowsCount) {
         return arrow::RecordBatch::Make(
             std::make_shared<arrow::Schema>(arrow::FieldVector()),
diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/connector_client_mock.h
@@ -54,6 +54,11 @@ namespace NYql::NConnector::NTest {
         return TClickHouseDataSourceInstanceBuilder<TBuilder>(                      \
             this->Result_->mutable_data_source_instance(),                          \
             static_cast<TBuilder*>(this));                                          \
+    }                                                                               \
+    TYdbDataSourceInstanceBuilder<TBuilder> YdbDataSourceInstance() {               \
+        return TYdbDataSourceInstanceBuilder<TBuilder>(                             \
+            this->Result_->mutable_data_source_instance(),                          \
+            static_cast<TBuilder*>(this));                                          \
     }
 
     MATCHER_P(ProtobufRequestMatcher, expected, "request does not match") {
@@ -200,6 +205,15 @@ namespace NYql::NConnector::NTest {
         const TString& serviceAccountIdSignature = DEFAULT_CH_SERVICE_ACCOUNT_ID_SIGNATURE,
         const TString& databaseName = DEFAULT_DATABASE);
 
+    void CreateYdbExternalDataSource(
+        const std::shared_ptr<NKikimr::NKqp::TKikimrRunner>& kikimr,
+        const TString& dataSourceName = DEFAULT_DATA_SOURCE_NAME,
+        const TString& login = DEFAULT_LOGIN,
+        const TString& password = DEFAULT_PASSWORD,
+        const TString& endpoint = DEFAULT_YDB_ENDPOINT,
+        bool useTls = DEFAULT_USE_TLS,
+        const TString& databaseName = DEFAULT_DATABASE);
+
     class TConnectorClientMock: public NYql::NConnector::IClient {
     public:
         MOCK_METHOD(TResult<NApi::TDescribeTableResponse>, DescribeTableImpl, (const NApi::TDescribeTableRequest& request));
@@ -277,6 +291,24 @@ namespace NYql::NConnector::NTest {
             }
         };
 
+        template <class TParent = void /* no parent by default */>
+        struct TYdbDataSourceInstanceBuilder: public TBaseDataSourceInstanceBuilder<TYdbDataSourceInstanceBuilder<TParent>, TParent> {
+            using TBase = TBaseDataSourceInstanceBuilder<TYdbDataSourceInstanceBuilder<TParent>, TParent>;
+
+            explicit TYdbDataSourceInstanceBuilder(NApi::TDataSourceInstance* result = nullptr, TParent* parent = nullptr)
+                : TBase(result, parent)
+            {
+                FillWithDefaults();
+            }
+
+            void FillWithDefaults() {
+                TBase::FillWithDefaults();
+                this->Host(DEFAULT_YDB_HOST);
+                this->Port(DEFAULT_YDB_PORT);
+                this->Kind(NApi::EDataSourceKind::YDB);
+            }
+        };
+
         template <class TParent = void /* no parent by default */>
         struct TDescribeTableResultBuilder: public TResponseBuilder<TParent, NApi::TDescribeTableResponse> {
             using TBuilder = TDescribeTableResultBuilder<TParent>;
diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.cpp b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.cpp
@@ -16,4 +16,8 @@ namespace NYql::NConnector::NTest {
     extern const TString DEFAULT_CH_CLUSTER_ID = "ch-managed";
     extern const TString DEFAULT_CH_SERVICE_ACCOUNT_ID = "sa";
     extern const TString DEFAULT_CH_SERVICE_ACCOUNT_ID_SIGNATURE = "sa_signature";
+
+    extern const TString DEFAULT_YDB_HOST = "localhost";
+    extern const TString DEFAULT_YDB_DATABASE = "local";
+    extern const TString DEFAULT_YDB_ENDPOINT = TStringBuilder() << DEFAULT_YDB_HOST << ':' << DEFAULT_YDB_PORT;
 } // namespace NYql::NConnector::NTest
diff --git a/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.h b/ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.h
@@ -25,4 +25,9 @@ namespace NYql::NConnector::NTest {
     constexpr NApi::EProtocol DEFAULT_CH_PROTOCOL = NApi::EProtocol::HTTP;
     extern const TString DEFAULT_CH_SERVICE_ACCOUNT_ID;
     extern const TString DEFAULT_CH_SERVICE_ACCOUNT_ID_SIGNATURE;
+
+    extern const TString DEFAULT_YDB_DATABASE;
+    extern const TString DEFAULT_YDB_HOST;
+    constexpr int DEFAULT_YDB_PORT = 2136;
+    extern const TString DEFAULT_YDB_ENDPOINT;
 } // namespace NYql::NConnector::NTest
diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_cluster_config.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_cluster_config.cpp
@@ -172,6 +172,11 @@ namespace NYql {
                        NYql::TGenericClusterConfig& clusterConfig) {
         using namespace NConnector::NApi;
 
+        if (clusterConfig.GetKind() == EDataSourceKind::YDB) {
+            clusterConfig.SetProtocol(EProtocol::NATIVE);
+            return;
+        }
+
         auto it = properties.find("protocol");
         if (it == properties.cend()) {
             ythrow yexception() << "missing 'PROTOCOL' value";
diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_dq_integration.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_dq_integration.cpp
@@ -162,6 +162,9 @@ namespace NYql {
                         case NYql::NConnector::NApi::POSTGRESQL:
                             sourceType = "PostgreSqlGeneric";
                             break;
+                        case NYql::NConnector::NApi::YDB:
+                            sourceType = "YdbGeneric";
+                            break;
                         default:
                             ythrow yexception() << "Data source kind is unknown or not specified";
                             break;
@@ -193,6 +196,9 @@ namespace NYql {
                         case NConnector::NApi::POSTGRESQL:
                             properties["SourceType"] = "PostgreSql";
                             break;
+                        case NConnector::NApi::YDB:
+                            properties["SourceType"] = "YDB";
+                            break;
                         case NConnector::NApi::DATA_SOURCE_KIND_UNSPECIFIED:
                             break;
                         default:
diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_load_meta.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_load_meta.cpp
@@ -269,6 +269,8 @@ namespace NYql {
             switch (dataSourceKind) {
                 case NYql::NConnector::NApi::CLICKHOUSE:
                     break;
+                case NYql::NConnector::NApi::YDB:
+                    break;
                 case NYql::NConnector::NApi::POSTGRESQL: {
                     // for backward compability set schema "public" by default
                     // TODO: simplify during https://st.yandex-team.ru/YQ-2494
@@ -324,7 +326,7 @@ namespace NYql {
                         dbNameTarget = "postgres";
                         break;
                     default:
-                        ythrow yexception() << "Unexpected data source kind: '"
+                        ythrow yexception() << "You must provide database name explicitly for data source kind: '"
                                             << NYql::NConnector::NApi::EDataSourceKind_Name(dataSourceKind) << "'";
                 }
             } // else take database name from table path

Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,10 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector<TStri`
`46`	`46`	`{`
`47`	`47`	`ToString(NYql::EDatabaseType::PostgreSQL),`
`48`	`48`	`CreateExternalDataSource(TString{NYql::GenericProviderName}, {"MDB_BASIC", "BASIC"}, {"database_name", "protocol", "mdb_cluster_id", "use_tls", "schema"}, hostnamePatternsRegEx)`
	`49`	`+ },`
	`50`	`+ {`
	`51`	`+ ToString(NYql::EDatabaseType::Ydb),`
	`52`	`+ CreateExternalDataSource(TString{NYql::GenericProviderName}, {"BASIC", "SERVICE_ACCOUNT"}, {"database_name", "use_tls", "database_id"}, hostnamePatternsRegEx)`
`49`	`53`	`}`
`50`	`54`	`});`
`51`	`55`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ namespace NYql::NDq {`
`16`	`16`	`args.SecureParams, args.TaskParams, args.ComputeActorId, credentialsFactory, args.HolderFactory);`
`17`	`17`	`};`
`18`	`18`
`19`		`- for (auto& sourceName : {"ClickHouseGeneric", "PostgreSqlGeneric"}) {`
	`19`	`+ for (auto& sourceName : {"ClickHouseGeneric", "PostgreSqlGeneric", "YdbGeneric"}) {`
`20`	`20`	`factory.RegisterSource<Generic::TSource>(sourceName, genericFactory);`
`21`	`21`	`}`
`22`	`22`	`}`