Moved commit "persqueue: add user-agent counters" from ydb repo

Author: qyryq

src/client/persqueue_public/ut/ut_utils/data_plane_helpers.cpp

@@ -1,4 +1,6 @@
 #include "data_plane_helpers.h"
+#include <ydb-cpp-sdk/client/resources/ydb_resources.h>
+#include <ydb-cpp-sdk/client/topic/client.h>
 
 namespace NKikimr::NPersQueueTests {
 
@@ -51,7 +53,8 @@ namespace NKikimr::NPersQueueTests {
         std::optional<ui32> partitionGroup,
         std::optional<TString> codec,
         std::optional<bool> reconnectOnFailure,
-        std::unordered_map<std::string, std::string> sessionMeta
+        std::unordered_map<std::string, std::string> sessionMeta,
+        const TString& userAgent
     ) {
         auto settings = TWriteSessionSettings().Path(topic).MessageGroupId(sourceId);
         if (partitionGroup) settings.PartitionGroupId(*partitionGroup);
@@ -66,6 +69,9 @@ namespace NKikimr::NPersQueueTests {
         }
         settings.MaxMemoryUsage(1024*1024*1024*1024ll);
         settings.Meta_.Fields = sessionMeta;
+        if (!userAgent.empty()) {
+            settings.Header({{NYdb::YDB_APPLICATION_NAME, userAgent}});
+        }
         return CreateSimpleWriter(driver, settings);
     }
 
@@ -79,6 +85,21 @@ namespace NKikimr::NPersQueueTests {
         return TPersQueueClient(driver, clientSettings).CreateReadSession(TReadSessionSettings(settings).DisableClusterDiscovery(true));
     }
 
+    std::shared_ptr<NYdb::NTopic::IReadSession> CreateReader(
+        NYdb::TDriver& driver,
+        const NYdb::NTopic::TReadSessionSettings& settings,
+        std::shared_ptr<NYdb::ICredentialsProviderFactory> creds,
+        const TString& userAgent
+    ) {
+        NYdb::NTopic::TTopicClientSettings clientSettings;
+        if (creds) clientSettings.CredentialsProviderFactory(creds);
+        auto readerSettings = settings;
+        if (!userAgent.empty()) {
+            readerSettings.Header({{NYdb::YDB_APPLICATION_NAME, userAgent}});
+        }
+        return NYdb::NTopic::TTopicClient(driver, clientSettings).CreateReadSession(readerSettings);
+    }
+
     TMaybe<TReadSessionEvent::TDataReceivedEvent> GetNextMessageSkipAssignment(std::shared_ptr<IReadSession>& reader, TDuration timeout) {
         while (true) {
             auto future = reader->WaitEvent();
@@ -99,4 +120,24 @@ namespace NKikimr::NPersQueueTests {
         }
         return {};
     }
+
+    TMaybe<NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent> GetNextMessageSkipAssignment(std::shared_ptr<NYdb::NTopic::IReadSession>& reader, TDuration timeout) {
+        while (true) {
+            auto future = reader->WaitEvent();
+            future.Wait(timeout);
+            std::optional<NYdb::NTopic::TReadSessionEvent::TEvent> event = reader->GetEvent(false, 1);
+            if (!event)
+                return {};
+            if (auto e = std::get_if<NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent>(&*event)) {
+                return *e;
+            } else if (auto* e = std::get_if<NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent>(&*event)) {
+                e->Confirm();
+            } else if (auto* e = std::get_if<NYdb::NTopic::TReadSessionEvent::TStopPartitionSessionEvent>(&*event)) {
+                e->Confirm();
+            } else if (std::get_if<NYdb::NTopic::TSessionClosedEvent>(&*event)) {
+                return {};
+            }
+        }
+        return {};
+    }
 }

src/client/persqueue_public/ut/ut_utils/data_plane_helpers.h

@@ -3,6 +3,7 @@
 #include <ydb-cpp-sdk/client/driver/driver.h>
 #include <ydb-cpp-sdk/client/table/table.h>
 #include <src/client/persqueue_public/persqueue.h>
+#include <ydb-cpp-sdk/client/topic/client.h>
 
 #include <util/generic/hash.h>
 
@@ -36,16 +37,24 @@ namespace NKikimr::NPersQueueTests {
         std::optional<ui32> partitionGroup = {},
         std::optional<TString> codec = {},
         std::optional<bool> reconnectOnFailure = {},
-        std::unordered_map<std::string, std::string> sessionMeta = {}
+        std::unordered_map<std::string, std::string> sessionMeta = {},
+        const TString& userAgent = {}
     );
 
     std::shared_ptr<NYdb::NPersQueue::IReadSession> CreateReader(
         NYdb::TDriver& driver,
         const NYdb::NPersQueue::TReadSessionSettings& settings,
         std::shared_ptr<NYdb::ICredentialsProviderFactory> creds = nullptr
+        );
 
+    std::shared_ptr<NYdb::NTopic::IReadSession> CreateReader(
+        NYdb::TDriver& driver,
+        const NYdb::NTopic::TReadSessionSettings& settings,
+        std::shared_ptr<NYdb::ICredentialsProviderFactory> creds = nullptr,
+        const TString& userAgent = ""
     );
 
     TMaybe<NYdb::NPersQueue::TReadSessionEvent::TDataReceivedEvent> GetNextMessageSkipAssignment(std::shared_ptr<NYdb::NPersQueue::IReadSession>& reader, TDuration timeout = TDuration::Max());
+    TMaybe<NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent> GetNextMessageSkipAssignment(std::shared_ptr<NYdb::NTopic::IReadSession>& reader, TDuration timeout = TDuration::Max());
 
 }