YQ-4317 PQ provider to ydb

ydb/core/external_sources/external_data_source.cpp

@@ -2,6 +2,7 @@
 #include "validation_functions.h"
 
 #include <ydb/core/protos/flat_scheme_op.pb.h>
+#include <ydb/library/yql/providers/common/db_id_async_resolver/database_type.h>
 
 namespace NKikimr::NExternalSource {
 
@@ -62,6 +63,9 @@ struct TExternalDataSource : public IExternalSource {
             throw TExternalSourceException() << proto.GetSourceType() << " source must provide service_name";
         }
 
+        if (proto.GetSourceType() == ToString(NExternalSource::YdbTopicsType)) {
+            throw TExternalSourceException() << "External source with type " << proto.GetSourceType() << " is not allowed, use " << ToString(NYql::EDatabaseType::Ydb)  << " source type to read from topics ";
+        }
         ValidateHostname(HostnamePatterns, proto.GetLocation());
     }
 

ydb/core/external_sources/external_source.h

@@ -10,6 +10,8 @@
 
 namespace NKikimr::NExternalSource {
 
+constexpr TStringBuf YdbTopicsType = "YdbTopics";
+
 struct TExternalSourceException: public yexception {
 };
 

ydb/core/external_sources/external_source_factory.cpp

@@ -115,7 +115,7 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector<TStri
         },
         {
             ToString(NYql::EDatabaseType::Ydb),
-            CreateExternalDataSource(TString{NYql::GenericProviderName}, {"BASIC", "SERVICE_ACCOUNT"}, {"database_name", "use_tls", "database_id"}, hostnamePatternsRegEx)
+            CreateExternalDataSource(TString{NYql::GenericProviderName}, {"NONE", "BASIC", "SERVICE_ACCOUNT", "TOKEN"}, {"database_name", "use_tls", "database_id"}, hostnamePatternsRegEx)
         },
         {
             ToString(NYql::EDatabaseType::YT),
@@ -160,6 +160,10 @@ IExternalSourceFactory::TPtr CreateExternalSourceFactory(const std::vector<TStri
         {
             ToString(NYql::EDatabaseType::OpenSearch),
             CreateExternalDataSource(TString{NYql::GenericProviderName}, {"BASIC"}, {"database_name", "use_tls"}, hostnamePatternsRegEx)
+        },
+        {
+            ToString(YdbTopicsType),
+            CreateExternalDataSource(TString{NYql::PqProviderName}, {"NONE", "BASIC", "TOKEN"}, {"database_name", "use_tls"}, hostnamePatternsRegEx)
         }
     },
     allExternalDataSourcesAreAvailable,

ydb/core/kqp/compile_service/kqp_compile_actor.cpp

@@ -309,7 +309,7 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
         counters->TxProxyMon = new NTxProxy::TTxProxyMon(AppData(ctx)->Counters);
         std::shared_ptr<NYql::IKikimrGateway::IKqpTableMetadataLoader> loader =
             std::make_shared<TKqpTableMetadataLoader>(
-                QueryId.Cluster, TlsActivationContext->ActorSystem(), Config, true, TempTablesState);
+                QueryId.Cluster, TlsActivationContext->ActorSystem(), Config, true, TempTablesState, FederatedQuerySetup);
         Gateway = CreateKikimrIcGateway(QueryId.Cluster, QueryId.Settings.QueryType, QueryId.Database, QueryId.DatabaseId, std::move(loader),
             ctx.ActorSystem(), ctx.SelfID.NodeId(), counters, QueryServiceConfig);
         Gateway->SetToken(QueryId.Cluster, UserToken);

ydb/core/kqp/compute_actor/kqp_compute_actor.cpp

@@ -14,7 +14,8 @@
 #include <ydb/library/formats/arrow/protos/ssa.pb.h>
 #include <ydb/library/yql/dq/proto/dq_tasks.pb.h>
 #include <ydb/library/yql/providers/solomon/actors/dq_solomon_read_actor.h>
-
+#include <ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.h>
+#include <ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h>
 
 namespace NKikimr {
 namespace NMiniKQL {
@@ -88,6 +89,8 @@ NYql::NDq::IDqAsyncIoFactory::TPtr CreateKqpAsyncIoFactory(
         }
 
         NYql::NDq::RegisterDQSolomonReadActorFactory(*factory, federatedQuerySetup->CredentialsFactory);
+        NYql::NDq::RegisterDqPqReadActorFactory(*factory, *federatedQuerySetup->Driver, federatedQuerySetup->CredentialsFactory, federatedQuerySetup->PqGateway, nullptr);
+        NYql::NDq::RegisterDqPqWriteActorFactory(*factory, *federatedQuerySetup->Driver, federatedQuerySetup->CredentialsFactory, federatedQuerySetup->PqGateway, nullptr);
     }
 
     return factory;

ydb/core/kqp/federated_query/kqp_federated_query_helpers.cpp

@@ -2,6 +2,7 @@
 
 #include <ydb/library/actors/http/http_proxy.h>
 #include <ydb/library/yql/providers/common/db_id_async_resolver/database_type.h>
+#include <ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h>
 
 #include <ydb/core/base/counters.h>
 #include <ydb/core/base/feature_flags.h>
@@ -13,7 +14,6 @@
 #include <ydb/core/fq/libs/db_id_async_resolver_impl/http_proxy.h>
 #include <ydb/core/fq/libs/db_id_async_resolver_impl/mdb_endpoint_generator.h>
 #include <ydb/library/actors/http/http_proxy.h>
-
 #include <yql/essentials/public/issue/yql_issue_utils.h>
 
 #include <yt/yql/providers/yt/comp_nodes/dq/dq_yt_factory.h>
@@ -71,6 +71,16 @@ namespace NKikimr::NKqp {
         return NYql::IHTTPGateway::Make(&httpGatewayConfig, httpGatewayGroup);
     }
 
+    NYql::IPqGateway::TPtr MakePqGateway(const std::shared_ptr<NYdb::TDriver>& driver, const NYql::TPqGatewayConfig& pqGatewayConfig) {
+        NYql::TPqGatewayServices pqServices(
+            *driver,
+            nullptr,
+            nullptr,
+            std::make_shared<NYql::TPqGatewayConfig>(pqGatewayConfig),
+            nullptr);
+        return CreatePqNativeGateway(pqServices);
+    }
+
     NYql::THttpGatewayConfig DefaultHttpGatewayConfig() {
         NYql::THttpGatewayConfig config;
         config.SetMaxInFlightCount(2000);
@@ -119,6 +129,14 @@ namespace NKikimr::NKqp {
         YtGateway = MakeYtGateway(appData->FunctionRegistry, queryServiceConfig);
         DqTaskTransformFactory = NYql::CreateYtDqTaskTransformFactory(true);
 
+        ActorSystemPtr = std::make_shared<NKikimr::TDeferredActorLogBackend::TAtomicActorSystemPtr>(nullptr);
+        NYdb::TDriverConfig cfg;
+        cfg.SetLog(std::make_unique<NKikimr::TDeferredActorLogBackend>(ActorSystemPtr, NKikimrServices::EServiceKikimr::YDB_SDK));
+        Driver = std::make_shared<NYdb::TDriver>(cfg);
+
+        PqGatewayConfig = NYql::TPqGatewayConfig{};
+        PqGateway = MakePqGateway(Driver, NYql::TPqGatewayConfig{});
+
         // Initialize Token Accessor
         if (appConfig.GetAuthConfig().HasTokenAccessorConfig()) {
             const auto& tokenAccessorConfig = appConfig.GetAuthConfig().GetTokenAccessorConfig();
@@ -175,7 +193,11 @@ namespace NKikimr::NKqp {
             SolomonGateway,
             nullptr,
             S3ReadActorFactoryConfig,
-            DqTaskTransformFactory};
+            DqTaskTransformFactory,
+            PqGatewayConfig,
+            PqGateway,
+            ActorSystemPtr,
+            Driver};
 
         // Init DatabaseAsyncResolver only if all requirements are met
         if (DatabaseResolverActorId && MdbEndpointGenerator &&
@@ -270,4 +292,37 @@ namespace NKikimr::NKqp {
         return issues;
     }
 
+    NThreading::TFuture<TGetSchemeEntryResult> GetSchemeEntryType(
+        const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup,
+        const TString& endpoint,
+        const TString& database,
+        bool useTls,
+        const TString& structuredTokenJson,
+        const TString& path) {
+
+        if (!federatedQuerySetup || !federatedQuerySetup->Driver) {
+            return NThreading::MakeFuture<TGetSchemeEntryResult>(Nothing()); 
+        }
+        std::shared_ptr<NYdb::ICredentialsProviderFactory> credentialsProviderFactory = NYql::CreateCredentialsProviderFactoryForStructuredToken(nullptr, structuredTokenJson, false);
+        auto driver = federatedQuerySetup->Driver;
+
+        NYdb::TCommonClientSettings opts;
+        opts
+            .DiscoveryEndpoint(endpoint)
+            .Database(database)
+            .SslCredentials(NYdb::TSslCredentials(useTls))
+            .CredentialsProviderFactory(credentialsProviderFactory);
+        auto schemeClient = NYdb::NScheme::TSchemeClient(*driver, opts);
+        return schemeClient.DescribePath(path)
+            .Apply([actorSystem = NActors::TActivationContext::ActorSystem()](const NThreading::TFuture<NYdb::NScheme::TDescribePathResult>& result) {
+                auto describePathResult = result.GetValue();
+                if (!describePathResult.IsSuccess()) {
+                    LOG_WARN_S(*actorSystem, NKikimrServices::KQP_GATEWAY, "DescribePath failed, " << describePathResult.GetIssues().ToString());
+                    return NThreading::MakeFuture<TGetSchemeEntryResult>(Nothing()); 
+                }
+                NYdb::NScheme::TSchemeEntry entry = describePathResult.GetEntry();
+                return NThreading::MakeFuture<TGetSchemeEntryResult>(entry.Type);
+            });
+    };
+
 }  // namespace NKikimr::NKqp

ydb/core/kqp/federated_query/kqp_federated_query_helpers.h

@@ -6,6 +6,9 @@
 #include <ydb/library/yql/providers/common/db_id_async_resolver/mdb_endpoint_generator.h>
 #include <ydb/library/yql/providers/common/http_gateway/yql_http_gateway.h>
 #include <ydb/library/yql/providers/common/token_accessor/client/factory.h>
+
+#include <ydb/library/yql/providers/pq/gateway/native/yql_pq_gateway.h>
+
 #include <ydb/library/yql/providers/generic/connector/libcpp/client.h>
 #include <ydb/library/yql/providers/s3/actors_factory/yql_s3_actors_factory.h>
 #include <ydb/library/yql/providers/solomon/gateway/yql_solomon_gateway.h>
@@ -15,6 +18,7 @@
 #include <yql/essentials/public/issue/yql_issue_message.h>
 
 #include <yt/yql/providers/yt/provider/yql_yt_gateway.h>
+#include <ydb/library/logger/actor.h>
 
 namespace NKikimrConfig {
     class TQueryServiceConfig;
@@ -28,6 +32,9 @@ namespace NKikimr::NKqp {
 
     NYql::IHTTPGateway::TPtr MakeHttpGateway(const NYql::THttpGatewayConfig& httpGatewayConfig, NMonitoring::TDynamicCounterPtr countersRoot);
 
+    NYql::IPqGateway::TPtr MakePqGateway(const std::shared_ptr<NYdb::TDriver>& driver, const NYql::TPqGatewayConfig& pqGatewayConfig);
+
+
     struct TKqpFederatedQuerySetup {
         NYql::IHTTPGateway::TPtr HttpGateway;
         NYql::NConnector::IClient::TPtr ConnectorClient;
@@ -42,6 +49,10 @@ namespace NKikimr::NKqp {
         NMiniKQL::TComputationNodeFactory ComputationFactory;
         NYql::NDq::TS3ReadActorFactoryConfig S3ReadActorFactoryConfig;
         NYql::TTaskTransformFactory DqTaskTransformFactory;
+        NYql::TPqGatewayConfig PqGatewayConfig;
+        NYql::IPqGateway::TPtr PqGateway;
+        NKikimr::TDeferredActorLogBackend::TSharedAtomicActorSystemPtr ActorSystemPtr;
+        std::shared_ptr<NYdb::TDriver> Driver;
     };
 
     struct IKqpFederatedQuerySetupFactory {
@@ -81,6 +92,10 @@ namespace NKikimr::NKqp {
         NYql::IMdbEndpointGenerator::TPtr MdbEndpointGenerator;
         NYql::NDq::TS3ReadActorFactoryConfig S3ReadActorFactoryConfig;
         NYql::TTaskTransformFactory DqTaskTransformFactory;
+        NYql::TPqGatewayConfig PqGatewayConfig;
+        NYql::IPqGateway::TPtr PqGateway;
+        NKikimr::TDeferredActorLogBackend::TSharedAtomicActorSystemPtr ActorSystemPtr;
+        std::shared_ptr<NYdb::TDriver> Driver;
     };
 
     struct TKqpFederatedQuerySetupFactoryMock: public IKqpFederatedQuerySetupFactory {
@@ -99,7 +114,11 @@ namespace NKikimr::NKqp {
             const NYql::ISolomonGateway::TPtr& solomonGateway,
             NMiniKQL::TComputationNodeFactory computationFactory,
             const NYql::NDq::TS3ReadActorFactoryConfig& s3ReadActorFactoryConfig,
-            NYql::TTaskTransformFactory dqTaskTransformFactory)
+            NYql::TTaskTransformFactory dqTaskTransformFactory,
+            const NYql::TPqGatewayConfig& pqGatewayConfig,
+            NYql::IPqGateway::TPtr pqGateway,
+            NKikimr::TDeferredActorLogBackend::TSharedAtomicActorSystemPtr actorSystemPtr,
+            std::shared_ptr<NYdb::TDriver> driver)
             : HttpGateway(httpGateway)
             , ConnectorClient(connectorClient)
             , CredentialsFactory(credentialsFactory)
@@ -113,6 +132,10 @@ namespace NKikimr::NKqp {
             , ComputationFactory(computationFactory)
             , S3ReadActorFactoryConfig(s3ReadActorFactoryConfig)
             , DqTaskTransformFactory(dqTaskTransformFactory)
+            , PqGatewayConfig(pqGatewayConfig)
+            , PqGateway(pqGateway)
+            , ActorSystemPtr(actorSystemPtr)
+            , Driver(driver)
         {
         }
 
@@ -122,7 +145,7 @@ namespace NKikimr::NKqp {
                 DatabaseAsyncResolver, S3GatewayConfig, GenericGatewayConfig,
                 YtGatewayConfig, YtGateway, SolomonGatewayConfig,
                 SolomonGateway, ComputationFactory, S3ReadActorFactoryConfig,
-                DqTaskTransformFactory};
+                DqTaskTransformFactory, PqGatewayConfig, PqGateway, ActorSystemPtr, Driver};
         }
 
     private:
@@ -139,6 +162,10 @@ namespace NKikimr::NKqp {
         NMiniKQL::TComputationNodeFactory ComputationFactory;
         NYql::NDq::TS3ReadActorFactoryConfig S3ReadActorFactoryConfig;
         NYql::TTaskTransformFactory DqTaskTransformFactory;
+        NYql::TPqGatewayConfig PqGatewayConfig;
+        NYql::IPqGateway::TPtr PqGateway;
+        NKikimr::TDeferredActorLogBackend::TSharedAtomicActorSystemPtr ActorSystemPtr;
+        std::shared_ptr<NYdb::TDriver> Driver;
     };
 
     IKqpFederatedQuerySetupFactory::TPtr MakeKqpFederatedQuerySetupFactory(
@@ -162,4 +189,14 @@ namespace NKikimr::NKqp {
 
     NYql::TIssues ValidateResultSetColumns(const google::protobuf::RepeatedPtrField<Ydb::Column>& columns, ui32 maxNestingDepth = 90);
 
+    using TGetSchemeEntryResult = TMaybe<NYdb::NScheme::ESchemeEntryType>;
+
+    NThreading::TFuture<TGetSchemeEntryResult> GetSchemeEntryType(
+        const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup,
+        const TString& endpoint,
+        const TString& database,
+        bool useTls,
+        const TString& structuredTokenJson,
+        const TString& path);
+
 }  // namespace NKikimr::NKqp

ydb/core/kqp/federated_query/ya.make

@@ -10,9 +10,11 @@ PEERDIR(
     ydb/core/fq/libs/grpc
     ydb/core/fq/libs/db_id_async_resolver_impl
     ydb/library/db_pool/protos
+    ydb/library/logger
     ydb/library/yql/providers/common/http_gateway
     ydb/library/yql/providers/generic/connector/libcpp
     ydb/library/yql/providers/s3/actors_factory
+    ydb/library/yql/providers/pq/gateway/native
     ydb/library/yql/providers/solomon/gateway
     yql/essentials/core/dq_integration/transform
     yql/essentials/public/issue