YQ-3697 Add partition count to dqrun

ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.cpp

@@ -54,7 +54,7 @@ TDummyPqGateway& TDummyPqGateway::AddDummyTopic(const TDummyTopic& topic) {
     with_lock (Mutex) {
         Y_ENSURE(topic.Cluster);
         Y_ENSURE(topic.Path);
-        const auto key = std::make_pair(topic.Cluster, topic.Path);
+        const auto key = std::make_pair(topic.Cluster, topic.TopicName);
         Y_ENSURE(Topics.emplace(key, topic).second, "Already inserted dummy topic {" << topic.Cluster << ", " << topic.Path << "}");
         return *this;
     }

ydb/library/yql/providers/pq/gateway/dummy/yql_pq_dummy_gateway.h

@@ -9,10 +9,11 @@
 namespace NYql {
 
 struct TDummyTopic {
-    TDummyTopic(const TString& cluster, const TString& path, const TMaybe<TString>& filePath = {})
+    TDummyTopic(const TString& cluster, const TString& topicName, const TMaybe<TString>& path = {}, size_t partitionCount = 1)
         : Cluster(cluster)
+        , TopicName(topicName)
         , Path(path)
-        , FilePath(filePath)
+        , PartitionsCount(partitionCount)
     {
     }
 
@@ -22,9 +23,9 @@ struct TDummyTopic {
     }
 
     TString Cluster;
-    TString Path;
-    TMaybe<TString> FilePath;
-    size_t PartitionsCount = 1;
+    TString TopicName;
+    TMaybe<TString> Path;
+    size_t PartitionsCount;
 };
 
 // Dummy Pq gateway for tests.

ydb/library/yql/providers/pq/gateway/dummy/yql_pq_file_topic_client.cpp

@@ -6,6 +6,7 @@
 #include <library/cpp/threading/blocking_queue/blocking_queue.h>
 #include <library/cpp/threading/future/async.h>
 
+#include <util/folder/path.h>
 #include <util/system/file.h>
 
 namespace NYql {
@@ -214,18 +215,26 @@ struct TDummyPartitionSession: public NYdb::NTopic::TPartitionSession {
 std::shared_ptr<NYdb::NTopic::IReadSession> TFileTopicClient::CreateReadSession(const NYdb::NTopic::TReadSessionSettings& settings) {
     Y_ENSURE(!settings.Topics_.empty());
     TString topicPath = settings.Topics_.front().Path_;
-
+    Y_ENSURE(settings.Topics_.front().PartitionIds_.size() == 1);
+    ui64 partitionId = settings.Topics_.front().PartitionIds_.front();
     auto topicsIt = Topics_.find(make_pair("pq", topicPath));
     Y_ENSURE(topicsIt != Topics_.end());
-    auto filePath = topicsIt->second.FilePath;
-    Y_ENSURE(filePath);
-
-    // TODO
-    ui64 sessionId = 0;
-    ui64 partitionId = 0;
+    Y_ENSURE(topicsIt->second.Path);
+
+    TString filePath;
+    if (TFsPath(*topicsIt->second.Path).IsDirectory()) {
+        filePath = TString(*topicsIt->second.Path) + "/" + ToString(partitionId);
+    } else {
+        if (TFsPath(*topicsIt->second.Path).Exists() && topicsIt->second.PartitionsCount == 1) {
+            filePath = *topicsIt->second.Path;
+        } else {
+            filePath = TString(*topicsIt->second.Path) + "_" + ToString(partitionId);
+        }
+    }
 
+    ui64 sessionId = 0;
     return std::make_shared<TFileTopicReadSession>(
-        TFile(*filePath, EOpenMode::TEnum::RdOnly),
+        TFile(filePath, EOpenMode::TEnum::RdOnly),
         MakeIntrusive<TDummyPartitionSession>(sessionId, topicPath, partitionId)
     );
 }

ydb/library/yql/tools/dqrun/dqrun.cpp

@@ -985,13 +985,16 @@ int RunMain(int argc, const char* argv[])
             auto fileGateway = MakeIntrusive<TDummyPqGateway>();
 
             for (auto& s : pqFileList) {
-                TStringBuf topicName, filePath;
-                TStringBuf(s).Split('@', topicName, filePath);
-                if (topicName.empty() || filePath.empty()) {
-                    Cerr << "Incorrect table mapping, expected form topic@file" << Endl;
+                TStringBuf topicName, others;
+                TStringBuf(s).Split('@', topicName, others);
+                TStringBuf path, partitionCountStr;
+                TStringBuf(others).Split(':', path, partitionCountStr);
+                if (topicName.empty() || path.empty()) {
+                    Cerr << "Incorrect table mapping, expected form topic@path[:partitions_count]" << Endl;
                     return 1;
                 }
-                fileGateway->AddDummyTopic(TDummyTopic("pq", TString(topicName), TString(filePath)));
+                size_t partitionCount = !partitionCountStr.empty() ? partitionCount = FromString<size_t>(partitionCountStr) : 1;
+                fileGateway->AddDummyTopic(TDummyTopic("pq", TString(topicName), TString(path), partitionCount));
             }
             pqGateway = std::move(fileGateway);
         } else {