Skip to content

Commit 487593b

Browse files
YDB FQ: refactoring work with tokens in YQL Generic Provider (#26015)
Co-authored-by: Pisarenko Grigoriy <[email protected]>
1 parent 6e01802 commit 487593b

22 files changed

+353
-229
lines changed

ydb/core/kqp/ut/federated_query/common/common.cpp

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,4 +142,57 @@ namespace NKikimr::NKqp::NFederatedQueryTest {
142142
return kikimr;
143143
}
144144

145+
class TStaticCredentialsProvider: public NYdb::ICredentialsProvider {
146+
public:
147+
TStaticCredentialsProvider(const TString& yqlToken)
148+
: YqlToken_(yqlToken)
149+
{
150+
}
151+
152+
std::string GetAuthInfo() const override {
153+
return YqlToken_;
154+
}
155+
156+
bool IsValid() const override {
157+
return true;
158+
}
159+
160+
private:
161+
std::string YqlToken_;
162+
};
163+
164+
class TStaticCredentialsProviderFactory: public NYdb::ICredentialsProviderFactory {
165+
public:
166+
TStaticCredentialsProviderFactory(const TString& yqlToken)
167+
: YqlToken_(yqlToken)
168+
{
169+
}
170+
171+
std::shared_ptr<NYdb::ICredentialsProvider> CreateProvider() const override {
172+
return std::make_shared<TStaticCredentialsProvider>(YqlToken_);
173+
}
174+
175+
private:
176+
TString YqlToken_;
177+
};
178+
179+
class TStaticSecuredCredentialsFactory: public NYql::ISecuredServiceAccountCredentialsFactory {
180+
public:
181+
TStaticSecuredCredentialsFactory(const TString& yqlToken)
182+
: YqlToken_(yqlToken)
183+
{
184+
}
185+
186+
std::shared_ptr<NYdb::ICredentialsProviderFactory> Create(const TString&, const TString&) override {
187+
return std::make_shared<TStaticCredentialsProviderFactory>(YqlToken_);
188+
}
189+
190+
private:
191+
TString YqlToken_;
192+
};
193+
194+
std::shared_ptr<NYql::ISecuredServiceAccountCredentialsFactory> CreateCredentialsFactory(const TString& token) {
195+
return std::make_shared<TStaticSecuredCredentialsFactory>(token);
196+
}
197+
145198
} // namespace NKikimr::NKqp::NFederatedQueryTest

ydb/core/kqp/ut/federated_query/common/common.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
#include <ydb/core/kqp/ut/common/kqp_ut_common.h>
44
#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/query.h>
55
#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/operation/operation.h>
6+
#include <ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/defaults.h>
67
#include <ydb/library/yql/providers/s3/actors_factory/yql_s3_actors_factory.h>
78

89
namespace NKikimr::NKqp::NFederatedQueryTest {
@@ -34,4 +35,7 @@ namespace NKikimr::NKqp::NFederatedQueryTest {
3435
std::shared_ptr<NYql::NDq::IS3ActorsFactory> s3ActorsFactory = nullptr,
3536
const TKikimrRunnerOptions& options = {});
3637

38+
std::shared_ptr<NYql::ISecuredServiceAccountCredentialsFactory> CreateCredentialsFactory(
39+
const TString& token = NYql::NConnector::NTest::DEFAULT_PASSWORD);
40+
3741
} // namespace NKikimr::NKqp::NFederatedQueryTest

ydb/core/kqp/ut/federated_query/datastreams/datastreams_ut.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,9 @@ using namespace NYdb;
2121
using namespace NYdb::NQuery;
2222
using namespace NKikimr::NKqp::NFederatedQueryTest;
2323
using namespace fmt::literals;
24-
using namespace NTestUtils;
2524
using namespace NYql::NConnector::NTest;
2625
using namespace NYql::NConnector::NApi;
26+
using namespace NTestUtils;
2727

2828
namespace {
2929

@@ -92,6 +92,7 @@ class TStreamingTestFixture : public NUnitTest::TBaseFixture {
9292
queryServiceConfig.SetProgressStatsPeriodMs(1000);
9393

9494
Kikimr = MakeKikimrRunner(true, ConnectorClient, nullptr, AppConfig, NYql::NDq::CreateS3ActorsFactory(), {
95+
.CredentialsFactory = CreateCredentialsFactory(),
9596
.PqGateway = PqGateway,
9697
.CheckpointPeriod = CheckpointPeriod,
9798
});

ydb/core/kqp/ut/federated_query/generic_ut/iceberg_ut_data.cpp

Lines changed: 0 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -182,51 +182,6 @@ void TIcebergTestData::ExecuteCreateHadoopExternalDataSource(const std::shared_p
182182
ExecuteQuery(kikimr, CreateQuery(hadoopCatalog));
183183
}
184184

185-
class TStaticCredentialsProvider : public NYdb::ICredentialsProvider {
186-
public:
187-
TStaticCredentialsProvider(const TString& yqlToken)
188-
: YqlToken_(yqlToken)
189-
{}
190-
191-
std::string GetAuthInfo() const override {
192-
return YqlToken_;
193-
}
194-
195-
bool IsValid() const override {
196-
return true;
197-
}
198-
199-
private:
200-
std::string YqlToken_;
201-
};
202-
203-
class TStaticCredentialsProviderFactory : public NYdb::ICredentialsProviderFactory {
204-
public:
205-
TStaticCredentialsProviderFactory(const TString& yqlToken)
206-
: YqlToken_(yqlToken)
207-
{}
208-
209-
std::shared_ptr<NYdb::ICredentialsProvider> CreateProvider() const override {
210-
return std::make_shared<TStaticCredentialsProvider>(YqlToken_);
211-
}
212-
213-
private:
214-
TString YqlToken_;
215-
};
216-
217-
class TStaticSecuredCredentialsFactory : public NYql::ISecuredServiceAccountCredentialsFactory {
218-
public:
219-
TStaticSecuredCredentialsFactory(const TString& yqlToken)
220-
: YqlToken_(yqlToken)
221-
{}
222-
223-
std::shared_ptr<NYdb::ICredentialsProviderFactory> Create(const TString&, const TString&) override {
224-
return std::make_shared<TStaticCredentialsProviderFactory>(YqlToken_);
225-
}
226-
227-
private:
228-
TString YqlToken_;
229-
};
230185

231186
TIcebergTestData CreateIcebergBasic(const TString& dataSourceName, const TString& database, const TString& userName, const TString& password){
232187
return TIcebergTestData({TIcebergTestData::EAuthType::AuthBasic, userName, password}, dataSourceName, database, false);
@@ -240,8 +195,4 @@ TIcebergTestData CreateIcebergSa(const TString& dataSourceName, const TString& d
240195
return TIcebergTestData({TIcebergTestData::EAuthType::AuthSa,VALUE_IAM, token}, dataSourceName, database, false);
241196
}
242197

243-
std::shared_ptr<NYql::ISecuredServiceAccountCredentialsFactory> CreateCredentialProvider(const TString& token) {
244-
return std::make_shared<TStaticSecuredCredentialsFactory>(token);
245-
}
246-
247198
} // NTestUtils

ydb/core/kqp/ut/federated_query/generic_ut/iceberg_ut_data.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,6 @@ TIcebergTestData CreateIcebergSa(const TString& dataSourceName = NYql::NConnecto
6161
const TString& database = NYql::NConnector::NTest::DEFAULT_DATABASE,
6262
const TString& token = NYql::NConnector::NTest::DEFAULT_PASSWORD);
6363

64-
std::shared_ptr<NYql::ISecuredServiceAccountCredentialsFactory> CreateCredentialProvider(const TString& token = NYql::NConnector::NTest::DEFAULT_PASSWORD);
64+
std::shared_ptr<NYql::ISecuredServiceAccountCredentialsFactory> CreateCredentialsFactory(const TString& token = NYql::NConnector::NTest::DEFAULT_PASSWORD);
6565

6666
} // NTestUtils

ydb/core/kqp/ut/federated_query/generic_ut/kqp_generic_provider_ut.cpp

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ namespace NKikimr::NKqp {
198198
auto appConfig = CreateDefaultAppConfig();
199199
auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory();
200200
auto kikimr = MakeKikimrRunner(false, clientMock, databaseAsyncResolverMock, appConfig, s3ActorsFactory,
201-
{.CredentialsFactory = NTestUtils::CreateCredentialProvider()});
201+
{.CredentialsFactory = CreateCredentialsFactory()});
202202

203203
CreateExternalDataSource(providerType, kikimr);
204204

@@ -314,7 +314,7 @@ namespace NKikimr::NKqp {
314314
auto appConfig = CreateDefaultAppConfig();
315315
auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory();
316316
auto kikimr = MakeKikimrRunner(false, clientMock, databaseAsyncResolverMock, appConfig, s3ActorsFactory,
317-
{.CredentialsFactory = NTestUtils::CreateCredentialProvider()});
317+
{.CredentialsFactory = CreateCredentialsFactory()});
318318

319319
CreateExternalDataSource(providerType, kikimr);
320320

@@ -429,7 +429,7 @@ namespace NKikimr::NKqp {
429429
auto appConfig = CreateDefaultAppConfig();
430430
auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory();
431431
auto kikimr = MakeKikimrRunner(false, clientMock, databaseAsyncResolverMock, appConfig, s3ActorsFactory,
432-
{.CredentialsFactory = NTestUtils::CreateCredentialProvider()});
432+
{.CredentialsFactory = CreateCredentialsFactory()});
433433

434434
CreateExternalDataSource(providerType, kikimr);
435435

@@ -563,7 +563,7 @@ namespace NKikimr::NKqp {
563563
auto appConfig = CreateDefaultAppConfig();
564564
auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory();
565565
auto kikimr = MakeKikimrRunner(false, clientMock, databaseAsyncResolverMock, appConfig, s3ActorsFactory,
566-
{.CredentialsFactory = NTestUtils::CreateCredentialProvider()});
566+
{.CredentialsFactory = CreateCredentialsFactory()});
567567

568568
CreateExternalDataSource(providerType, kikimr);
569569

@@ -631,7 +631,7 @@ namespace NKikimr::NKqp {
631631
auto appConfig = CreateDefaultAppConfig();
632632
auto s3ActorsFactory = NYql::NDq::CreateS3ActorsFactory();
633633
auto kikimr = MakeKikimrRunner(false, clientMock, databaseAsyncResolverMock, appConfig, s3ActorsFactory,
634-
{.CredentialsFactory = NTestUtils::CreateCredentialProvider()});
634+
{.CredentialsFactory = CreateCredentialsFactory()});
635635

636636
// Create trash query
637637
NYdbGrpc::TGRpcClientLow clientLow;

ydb/library/yql/dq/actors/compute/dq_compute_actor_async_io.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,7 @@ struct IDqAsyncIoFactory : public TThrRefBase {
276276
const NKikimr::NMiniKQL::TStructType* PayloadType;
277277
const NKikimr::NMiniKQL::TTypeEnvironment& TypeEnv;
278278
const NKikimr::NMiniKQL::THolderFactory& HolderFactory;
279+
const THashMap<TString, TString>& SecureParams;
279280
size_t MaxKeysInRequest;
280281
};
281282

ydb/library/yql/dq/actors/input_transforms/dq_input_transform_lookup.cpp

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ class TInputTransformStreamLookupCommonBase
4747
const NMiniKQL::TStructType* lookupPayloadType,
4848
const NMiniKQL::TMultiType* outputRowType,
4949
TOutputRowColumnOrder&& outputRowColumnOrder,
50+
const THashMap<TString, TString>& secureParams,
5051
size_t maxDelayedRows,
5152
size_t cacheLimit,
5253
std::chrono::seconds cacheTtl
@@ -61,6 +62,7 @@ class TInputTransformStreamLookupCommonBase
6162
, TaskCounters(taskCounters)
6263
, Factory(factory)
6364
, Settings(std::move(settings))
65+
, SecureParams(secureParams)
6466
, LookupInputIndexes(std::move(lookupInputIndexes))
6567
, OtherInputIndexes(std::move(otherInputIndexes))
6668
, InputRowType(inputRowType)
@@ -157,6 +159,7 @@ class TInputTransformStreamLookupCommonBase
157159
.PayloadType = LookupPayloadType,
158160
.TypeEnv = TypeEnv,
159161
.HolderFactory = HolderFactory,
162+
.SecureParams = SecureParams,
160163
.MaxKeysInRequest = 1000 // TODO configure me
161164
};
162165
auto [lookupSource, lookupSourceActor] = Factory->CreateDqLookupSource(Settings.GetRightSource().GetProviderName(), std::move(lookupSourceArgs));
@@ -221,6 +224,7 @@ class TInputTransformStreamLookupCommonBase
221224
::NMonitoring::TDynamicCounterPtr TaskCounters;
222225
IDqAsyncIoFactory::TPtr Factory;
223226
NDqProto::TDqInputTransformLookupSettings Settings;
227+
const THashMap<TString, TString> SecureParams;
224228
protected:
225229
NActors::TActorId LookupSourceId;
226230
size_t MaxKeysInRequest;
@@ -1094,6 +1098,7 @@ std::pair<IDqComputeActorAsyncInput*, NActors::IActor*> CreateInputTransformStre
10941098
lookupPayloadType,
10951099
outputRowType,
10961100
std::move(outputColumnsOrder),
1101+
args.SecureParams,
10971102
settings.GetMaxDelayedRows(),
10981103
settings.GetCacheLimit(),
10991104
std::chrono::seconds(settings.GetCacheTtlSeconds())
@@ -1115,6 +1120,7 @@ std::pair<IDqComputeActorAsyncInput*, NActors::IActor*> CreateInputTransformStre
11151120
lookupPayloadType,
11161121
outputRowType,
11171122
std::move(outputColumnsOrder),
1123+
args.SecureParams,
11181124
settings.GetMaxDelayedRows(),
11191125
settings.GetCacheLimit(),
11201126
std::chrono::seconds(settings.GetCacheTtlSeconds())
@@ -1139,6 +1145,7 @@ std::pair<IDqComputeActorAsyncInput*, NActors::IActor*> CreateInputTransformStre
11391145
lookupPayloadType,
11401146
outputRowType,
11411147
std::move(outputColumnsOrder),
1148+
args.SecureParams,
11421149
settings.GetMaxDelayedRows(),
11431150
settings.GetCacheLimit(),
11441151
std::chrono::seconds(settings.GetCacheTtlSeconds())
@@ -1160,6 +1167,7 @@ std::pair<IDqComputeActorAsyncInput*, NActors::IActor*> CreateInputTransformStre
11601167
lookupPayloadType,
11611168
outputRowType,
11621169
std::move(outputColumnsOrder),
1170+
args.SecureParams,
11631171
settings.GetMaxDelayedRows(),
11641172
settings.GetCacheLimit(),
11651173
std::chrono::seconds(settings.GetCacheTtlSeconds())

ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,8 @@ Y_UNIT_TEST_SUITE(GenericProviderLookupActor) {
190190
outputypeBuilder.Build(),
191191
typeEnv,
192192
holderFactory,
193-
1'000'000);
193+
1'000'000,
194+
{});
194195
auto lookupActor = runtime.Register(actor);
195196

196197
auto request = std::make_shared<NYql::NDq::IDqAsyncLookupSource::TUnboxedValueMap>(3, keyTypeHelper->GetValueHash(), keyTypeHelper->GetValueEqual());
@@ -381,7 +382,8 @@ Y_UNIT_TEST_SUITE(GenericProviderLookupActor) {
381382
outputypeBuilder.Build(),
382383
typeEnv,
383384
holderFactory,
384-
1'000'000);
385+
1'000'000,
386+
{});
385387
auto lookupActor = runtime.Register(actor);
386388

387389
auto request = std::make_shared<NYql::NDq::IDqAsyncLookupSource::TUnboxedValueMap>(3, keyTypeHelper->GetValueHash(), keyTypeHelper->GetValueEqual());

ydb/library/yql/providers/generic/actors/ya.make

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
LIBRARY()
22

33
SRCS(
4-
yql_generic_read_actor.cpp
4+
yql_generic_credentials_provider.cpp
55
yql_generic_lookup_actor.cpp
66
yql_generic_provider_factories.cpp
7-
yql_generic_token_provider.cpp
7+
yql_generic_read_actor.cpp
88
)
99

1010
PEERDIR(

0 commit comments

Comments
 (0)