add support for tags

3 files changed, 15 insertions, 3 deletions

diff --git a/ydb/core/local_pgwire/local_pgwire_connection.cpp b/ydb/core/local_pgwire/local_pgwire_connection.cpp
index 3b43f753bd..b634b22bad 100644
--- a/ydb/core/local_pgwire/local_pgwire_connection.cpp
+++ b/ydb/core/local_pgwire/local_pgwire_connection.cpp
@@ -185,6 +185,11 @@ public:
         auto rpcFuture = NRpcService::DoLocalRpc<TRpcEv>(std::move(request), database, token, actorSystem);
         rpcFuture.Subscribe([actorSystem, ev](NThreading::TFuture<Ydb::Scripting::ExecuteYqlResponse> future) {
             auto response = std::make_unique<NPG::TEvPGEvents::TEvQueryResponse>();
+            // HACK
+            if (ev->Get()->Message->GetQuery().starts_with("BEGIN")) {
+                response->Tag = "BEGIN";
+            }
+            // HACK
             try {
                 Ydb::Scripting::ExecuteYqlResponse yqlResponse(future.ExtractValueSync());
                 if (yqlResponse.has_operation()) {
@@ -216,6 +221,10 @@ public:
                                             }
                                         }
                                     }
+
+                                    // HACK
+                                    response->Tag = TStringBuilder() << "SELECT " << response->DataRows.size();
+                                    // HACK
                                 }
                             }
                         }
diff --git a/ydb/core/pgproxy/pg_connection.cpp b/ydb/core/pgproxy/pg_connection.cpp
index 3c9a1adaec..aabceee6e9 100644
--- a/ydb/core/pgproxy/pg_connection.cpp
+++ b/ydb/core/pgproxy/pg_connection.cpp
@@ -491,7 +491,7 @@ protected:
     void HandleConnected(TEvPGEvents::TEvQueryResponse::TPtr& ev) {
         if (IsEventExpected(ev)) {
             if (ev->Get()->ErrorFields.empty()) {
-                TString tag = "OK";
+                TString tag = ev->Get()->Tag ? ev->Get()->Tag : "OK";
                 { // rowDescription
                     TPGStreamOutput<TPGRowDescription> rowDescription;
                     rowDescription << uint16_t(ev->Get()->DataFields.size()); // number of fields
@@ -540,7 +540,7 @@ protected:
 
     void HandleConnected(TEvPGEvents::TEvDescribeResponse::TPtr& ev) {
         if (IsEventExpected(ev)) {
-            TString tag = "OK";
+            TString tag = ev->Get()->Tag ? ev->Get()->Tag : "OK";
             { // rowDescription
                 TPGStreamOutput<TPGRowDescription> rowDescription;
                 rowDescription << uint16_t(ev->Get()->DataFields.size()); // number of fields
@@ -568,7 +568,7 @@ protected:
     void HandleConnected(TEvPGEvents::TEvExecuteResponse::TPtr& ev) {
         if (IsEventExpected(ev)) {
             if (ev->Get()->ErrorFields.empty()) {
-                TString tag = "OK";
+                TString tag = ev->Get()->Tag ? ev->Get()->Tag : "OK";
                 { // dataFields
                     for (const auto& row : ev->Get()->DataRows) {
                         TPGStreamOutput<TPGDataRow> dataRow;
diff --git a/ydb/core/pgproxy/pg_proxy_events.h b/ydb/core/pgproxy/pg_proxy_events.h
index 63d85f4052..e312760e24 100644
--- a/ydb/core/pgproxy/pg_proxy_events.h
+++ b/ydb/core/pgproxy/pg_proxy_events.h
@@ -99,6 +99,7 @@ struct TEvPGEvents {
         std::vector<TRowDescriptionField> DataFields;
         std::vector<TDataRow> DataRows;
         std::vector<std::pair<char, TString>> ErrorFields;
+        TString Tag;
     };
 
         /*
@@ -181,6 +182,7 @@ struct TEvPGEvents {
     struct TEvDescribeResponse : NActors::TEventLocal<TEvDescribeResponse, EvDescribeResponse> {
         std::vector<TRowDescriptionField> DataFields;
         std::vector<std::pair<char, TString>> ErrorFields;
+        TString Tag;
     };
 
     struct TEvExecute : NActors::TEventLocal<TEvExecute, EvExecute> {
@@ -194,6 +196,7 @@ struct TEvPGEvents {
     struct TEvExecuteResponse : NActors::TEventLocal<TEvExecuteResponse, EvExecuteResponse> {
         std::vector<TDataRow> DataRows;
         std::vector<std::pair<char, TString>> ErrorFields;
+        TString Tag;
     };
 };