Browse Source

Intermediate changes

robot-piglet 1 year ago
parent
commit
38ae703deb

+ 4 - 0
yt/yt/client/api/delegating_client.h

@@ -750,6 +750,10 @@ public:
         const TAlterQueryOptions& options),
         (queryId, options))
 
+    DELEGATE_METHOD(TFuture<TGetQueryTrackerInfoResult>, GetQueryTrackerInfo, (
+        const TGetQueryTrackerInfoOptions& options),
+        (options))
+
     // Bundle Controller
     DELEGATE_METHOD(TFuture<NBundleControllerClient::TBundleConfigDescriptorPtr>, GetBundleConfig, (
         const TString& bundleName,

+ 16 - 0
yt/yt/client/api/query_tracker_client.h

@@ -140,6 +140,20 @@ struct TAlterQueryOptions
     std::optional<TString> AccessControlObject;
 };
 
+struct TGetQueryTrackerInfoOptions
+    : public TTimeoutOptions
+    , public TQueryTrackerOptions
+{
+    NYTree::TAttributeFilter Attributes;
+};
+
+struct TGetQueryTrackerInfoResult
+{
+    TString ClusterName;
+    NYson::TYsonString SupportedFeatures;
+    std::vector<TString> AccessControlObjects;
+};
+
 ////////////////////////////////////////////////////////////////////////////////
 
 struct IQueryTrackerClient
@@ -174,6 +188,8 @@ struct IQueryTrackerClient
     virtual TFuture<void> AlterQuery(
         NQueryTrackerClient::TQueryId queryId,
         const TAlterQueryOptions& options = {}) = 0;
+
+    virtual TFuture<TGetQueryTrackerInfoResult> GetQueryTrackerInfo(const TGetQueryTrackerInfoOptions& options = {}) = 0;
 };
 
 ////////////////////////////////////////////////////////////////////////////////

+ 6 - 0
yt/yt/client/api/rpc_proxy/client_impl.cpp

@@ -2067,6 +2067,12 @@ TFuture<void> TClient::AlterQuery(
     ThrowUnimplemented("AlterQuery");
 }
 
+TFuture<TGetQueryTrackerInfoResult> TClient::GetQueryTrackerInfo(
+    const TGetQueryTrackerInfoOptions& /*options*/)
+{
+    ThrowUnimplemented("GetQueryTrackerInfo");
+}
+
 TFuture<NBundleControllerClient::TBundleConfigDescriptorPtr> TClient::GetBundleConfig(
     const TString& /*bundleName*/,
     const NBundleControllerClient::TGetBundleConfigOptions& /*options*/)

+ 3 - 0
yt/yt/client/api/rpc_proxy/client_impl.h

@@ -459,6 +459,9 @@ public:
         NQueryTrackerClient::TQueryId queryId,
         const TAlterQueryOptions& options = {}) override;
 
+    TFuture<TGetQueryTrackerInfoResult> GetQueryTrackerInfo(
+        const TGetQueryTrackerInfoOptions& options = {}) override;
+
     // Authentication
 
     virtual TFuture<void> SetUserPassword(

+ 1 - 0
yt/yt/client/driver/driver.cpp

@@ -358,6 +358,7 @@ public:
         REGISTER    (TGetQueryResultCommand,               "get_query_result",                Null,       Structured, false, false, ApiVersion4);
         REGISTER    (TReadQueryResultCommand,              "read_query_result",               Null,       Tabular,    false, true, ApiVersion4);
         REGISTER    (TAlterQueryCommand,                   "alter_query",                     Null,       Tabular,    false, false, ApiVersion4);
+        REGISTER    (TGetQueryTrackerInfoCommand,          "get_query_tracker_info",          Null,       Structured, false, false, ApiVersion4);
 
         REGISTER_ALL(TGetBundleConfigCommand,              "get_bundle_config",               Null,       Structured, false,  false);
         REGISTER_ALL(TSetBundleConfigCommand,              "set_bundle_config",               Structured, Null,       false,  false);

+ 39 - 0
yt/yt/client/driver/query_commands.cpp

@@ -330,6 +330,13 @@ void TAlterQueryCommand::Register(TRegistrar registrar)
             return command->Options.AccessControlObject;
         })
         .Optional(/*init*/ false);
+
+    registrar.ParameterWithUniversalAccessor<TString>(
+        "stage",
+        [] (TThis* command) -> auto& {
+            return command->Options.QueryTrackerStage;
+        })
+        .Default("production");
 }
 
 void TAlterQueryCommand::DoExecute(ICommandContextPtr context)
@@ -341,4 +348,36 @@ void TAlterQueryCommand::DoExecute(ICommandContextPtr context)
 
 //////////////////////////////////////////////////////////////////////////////
 
+void TGetQueryTrackerInfoCommand::Register(TRegistrar registrar)
+{
+    registrar.ParameterWithUniversalAccessor<TString>(
+        "stage",
+        [] (TThis* command) -> auto& {
+            return command->Options.QueryTrackerStage;
+        })
+        .Default("production");
+
+    registrar.ParameterWithUniversalAccessor<TAttributeFilter>(
+        "attributes",
+        [] (TThis* command) -> auto& {
+            return command->Options.Attributes;
+        })
+        .Optional(/*init*/ false);
+}
+
+void TGetQueryTrackerInfoCommand::DoExecute(ICommandContextPtr context)
+{
+    auto result = WaitFor(context->GetClient()->GetQueryTrackerInfo(Options))
+        .ValueOrThrow();
+
+    context->ProduceOutputValue(BuildYsonStringFluently()
+        .BeginMap()
+            .Item("cluster_name").Value(result.ClusterName)
+            .Item("supported_features").Value(result.SupportedFeatures)
+            .Item("access_control_objects").Value(result.AccessControlObjects)
+        .EndMap());
+}
+
+//////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT::NDriver

+ 14 - 0
yt/yt/client/driver/query_commands.h

@@ -123,4 +123,18 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////
 
+class TGetQueryTrackerInfoCommand
+    : public TTypedCommand<NApi::TGetQueryTrackerInfoOptions>
+{
+public:
+    REGISTER_YSON_STRUCT_LITE(TGetQueryTrackerInfoCommand);
+
+    static void Register(TRegistrar registrar);
+
+private:
+    void DoExecute(ICommandContextPtr context) override;
+};
+
+////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT::NDriver

+ 1 - 0
yt/yt/client/federated/client.cpp

@@ -421,6 +421,7 @@ public:
     UNIMPLEMENTED_METHOD(TFuture<TQuery>, GetQuery, (NQueryTrackerClient::TQueryId, const TGetQueryOptions&));
     UNIMPLEMENTED_METHOD(TFuture<TListQueriesResult>, ListQueries, (const TListQueriesOptions&));
     UNIMPLEMENTED_METHOD(TFuture<void>, AlterQuery, (NQueryTrackerClient::TQueryId, const TAlterQueryOptions&));
+    UNIMPLEMENTED_METHOD(TFuture<TGetQueryTrackerInfoResult>, GetQueryTrackerInfo, (const TGetQueryTrackerInfoOptions&));
     UNIMPLEMENTED_METHOD(TFuture<NBundleControllerClient::TBundleConfigDescriptorPtr>, GetBundleConfig, (const TString&, const NBundleControllerClient::TGetBundleConfigOptions&));
     UNIMPLEMENTED_METHOD(TFuture<void>, SetBundleConfig, (const TString&, const NBundleControllerClient::TBundleTargetConfigPtr&, const NBundleControllerClient::TSetBundleConfigOptions&));
     UNIMPLEMENTED_METHOD(TFuture<ITableReaderPtr>, CreateTableReader, (const NYPath::TRichYPath&, const TTableReaderOptions&));

+ 1 - 0
yt/yt/client/hedging/hedging.cpp

@@ -211,6 +211,7 @@ public:
     UNSUPPORTED_METHOD(TFuture<TQuery>, GetQuery, (NQueryTrackerClient::TQueryId, const TGetQueryOptions&));
     UNSUPPORTED_METHOD(TFuture<TListQueriesResult>, ListQueries, (const TListQueriesOptions&));
     UNSUPPORTED_METHOD(TFuture<void>, AlterQuery, (NQueryTrackerClient::TQueryId, const TAlterQueryOptions&));
+    UNSUPPORTED_METHOD(TFuture<TGetQueryTrackerInfoResult>, GetQueryTrackerInfo, (const TGetQueryTrackerInfoOptions&));
     UNSUPPORTED_METHOD(TFuture<NBundleControllerClient::TBundleConfigDescriptorPtr>, GetBundleConfig, (const TString&, const NBundleControllerClient::TGetBundleConfigOptions&));
     UNSUPPORTED_METHOD(TFuture<TGetPipelineSpecResult>, GetPipelineSpec, (const NYPath::TYPath&, const TGetPipelineSpecOptions&));
     UNSUPPORTED_METHOD(TFuture<TSetPipelineSpecResult>, SetPipelineSpec, (const NYPath::TYPath&, const NYson::TYsonString&, const TSetPipelineSpecOptions&));

+ 4 - 0
yt/yt/client/unittests/mock/client.h

@@ -739,6 +739,10 @@ public:
         NQueryTrackerClient::TQueryId queryId, const TAlterQueryOptions& options),
         (override));
 
+    MOCK_METHOD(TFuture<TGetQueryTrackerInfoResult>, GetQueryTrackerInfo, (
+        const TGetQueryTrackerInfoOptions& options),
+        (override));
+
     MOCK_METHOD(TFuture<NBundleControllerClient::TBundleConfigDescriptorPtr>, GetBundleConfig, (
         const TString& bundleName,
         const NBundleControllerClient::TGetBundleConfigOptions& options), (override));