YT-20504: Add path/query tags to trace contexts of client rpc requests ModifyRows/SelectRows

4 files changed, 31 insertions, 1 deletions

diff --git a/yt/yt/client/api/rpc_proxy/client_base.cpp b/yt/yt/client/api/rpc_proxy/client_base.cpp
index 53e542c77f..6b5dbced6a 100644
--- a/yt/yt/client/api/rpc_proxy/client_base.cpp
+++ b/yt/yt/client/api/rpc_proxy/client_base.cpp
@@ -52,6 +52,19 @@ using NYT::FromProto;
 
 ////////////////////////////////////////////////////////////////////////////////
 
+constexpr i64 MaxTracingTagLength = 1000;
+static const TString DisabledSelectQueryTracingTag = "Tag is disabled, look for enable_select_query_tracing_tag parameter";
+
+TString SanitizeTracingTag(const TString& originalTag)
+{
+    if (originalTag.size() <= MaxTracingTagLength) {
+        return originalTag;
+    }
+    return Format("%v ... TRUNCATED", originalTag.substr(0, MaxTracingTagLength));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 IConnectionPtr TClientBase::GetConnection()
 {
     return GetRpcProxyConnection();
@@ -887,7 +900,7 @@ TFuture<std::vector<TUnversionedLookupRowsResult>> TClientBase::MultiLookup(
         for (const auto& subrequest : subrequests) {
             paths.emplace_back(subrequest.Path);
         }
-        req->TracingTags().emplace_back("yt.table_paths", NYson::ConvertToYsonString(paths).ToString());
+        req->TracingTags().emplace_back("yt.table_paths", SanitizeTracingTag(NYson::ConvertToYsonString(paths).ToString()));
     }
 
     req->set_replica_consistency(static_cast<NProto::EReplicaConsistency>(options.ReplicaConsistency));
@@ -951,6 +964,14 @@ TFuture<TSelectRowsResult> TClientBase::SelectRows(
 
     const auto& config = GetRpcProxyConnection()->GetConfig();
 
+    if (NTracing::IsCurrentTraceContextRecorded()) {
+        if (config->EnableSelectQueryTracingTag) {
+            req->TracingTags().emplace_back("yt.query", SanitizeTracingTag(query));
+        } else {
+            req->TracingTags().emplace_back("yt.query", DisabledSelectQueryTracingTag);
+        }
+    }
+
     FillRequestBySelectRowsOptionsBase(options, config->UdfRegistryPath, req);
     // TODO(ifsmirnov): retention timestamp in explain_query.
     req->set_retention_timestamp(options.RetentionTimestamp);
diff --git a/yt/yt/client/api/rpc_proxy/config.cpp b/yt/yt/client/api/rpc_proxy/config.cpp
index 1fbd8eb1f6..808d4f6262 100644
--- a/yt/yt/client/api/rpc_proxy/config.cpp
+++ b/yt/yt/client/api/rpc_proxy/config.cpp
@@ -112,6 +112,9 @@ void TConnectionConfig::Register(TRegistrar registrar)
     registrar.Parameter("udf_registry_path", &TThis::UdfRegistryPath)
         .Optional();
 
+    registrar.Parameter("enable_select_query_tracing_tag", &TThis::EnableSelectQueryTracingTag)
+        .Default(false);
+
     registrar.Postprocessor([] (TThis* config) {
         if (!config->ProxyEndpoints && !config->ClusterUrl && !config->ProxyAddresses && !config->ProxyUnixDomainSocket) {
             THROW_ERROR_EXCEPTION("Either \"endpoints\" or \"cluster_url\" or \"proxy_addresses\" or \"proxy_unix_domain_socket\" must be specified");
diff --git a/yt/yt/client/api/rpc_proxy/config.h b/yt/yt/client/api/rpc_proxy/config.h
index fd92333c39..3d3a41540f 100644
--- a/yt/yt/client/api/rpc_proxy/config.h
+++ b/yt/yt/client/api/rpc_proxy/config.h
@@ -74,6 +74,9 @@ public:
     //! Path in Cypress with UDFs.
     std::optional<NYPath::TYPath> UdfRegistryPath;
 
+    //! If |true| select query will be added to tracing tags of SelectRows span.
+    bool EnableSelectQueryTracingTag;
+
     REGISTER_YSON_STRUCT(TConnectionConfig);
 
     static void Register(TRegistrar registrar);
diff --git a/yt/yt/client/api/rpc_proxy/transaction_impl.cpp b/yt/yt/client/api/rpc_proxy/transaction_impl.cpp
index 4ab161b34b..6777b7109f 100644
--- a/yt/yt/client/api/rpc_proxy/transaction_impl.cpp
+++ b/yt/yt/client/api/rpc_proxy/transaction_impl.cpp
@@ -385,6 +385,9 @@ void TTransaction::ModifyRows(
     req->set_sequence_number_source_id(SequenceNumberSourceId_);
     ToProto(req->mutable_transaction_id(), GetId());
     req->set_path(path);
+    if (NTracing::IsCurrentTraceContextRecorded()) {
+        req->TracingTags().emplace_back("yt.table_path", path);
+    }
     req->set_require_sync_replica(options.RequireSyncReplica);
     ToProto(req->mutable_upstream_replica_id(), options.UpstreamReplicaId);
     req->set_allow_missing_key_columns(options.AllowMissingKeyColumns);