Added --ignore-unsupported option in pg-convert command

Added --ignore-unsupported option in pg-convert command

Pull Request resolved: https://github.com/ydb-platform/ydb/pull/397

5 files changed, 39 insertions, 5 deletions

diff --git a/ydb/public/lib/ydb_cli/commands/ydb_tools.cpp b/ydb/public/lib/ydb_cli/commands/ydb_tools.cpp
index 874a3f2363..a96b7935fa 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_tools.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_tools.cpp
@@ -392,6 +392,7 @@ void TCommandPgConvert::Config(TConfig& config) {
     config.SetFreeArgsNum(0);
 
     config.Opts->AddLongOption('i', "input", "Path to input SQL file. Read from stdin if not specified.").StoreResult(&Path);
+    config.Opts->AddLongOption("ignore-unsupported", "Comment unsupported statements in result dump file if specified.").StoreTrue(&IgnoreUnsupported);
 }
 
 void TCommandPgConvert::Parse(TConfig& config) {
@@ -400,7 +401,7 @@ void TCommandPgConvert::Parse(TConfig& config) {
 
 int TCommandPgConvert::Run(TConfig& config) {
     Y_UNUSED(config);
-    TPgDumpParser parser(Cout);
+    TPgDumpParser parser(Cout, IgnoreUnsupported);
     if (Path) {
         std::unique_ptr<TFileInput> fileInput = std::make_unique<TFileInput>(Path);
         parser.Prepare(*fileInput);
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_tools.h b/ydb/public/lib/ydb_cli/commands/ydb_tools.h
index 6b38794fc2..f4c22d894c 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_tools.h
+++ b/ydb/public/lib/ydb_cli/commands/ydb_tools.h
@@ -116,6 +116,7 @@ public:
 
 private:
     TString Path;
+    bool IgnoreUnsupported = false;
 };
 
 }
diff --git a/ydb/public/lib/ydb_cli/common/pg_dump_parser.cpp b/ydb/public/lib/ydb_cli/common/pg_dump_parser.cpp
index 8d802c428d..634d3dc372 100644
--- a/ydb/public/lib/ydb_cli/common/pg_dump_parser.cpp
+++ b/ydb/public/lib/ydb_cli/common/pg_dump_parser.cpp
@@ -1,5 +1,7 @@
 #include "pg_dump_parser.h"
 
+#include <ydb/public/lib/ydb_cli/common/common.h>
+
 namespace NYdb::NConsoleClient {
 
 size_t TFixedStringStream::DoRead(void* buf, size_t len) {
@@ -66,7 +68,7 @@ TPgDumpParser::TSQLCommandNode* TPgDumpParser::TSQLCommandNode::GetNextCommand(c
     return nullptr;
 }
 
-TPgDumpParser::TPgDumpParser(IOutputStream& out) : Out(out) {
+TPgDumpParser::TPgDumpParser(IOutputStream& out, bool ignoreUnsupported) : Out(out), IgnoreUnsupported(ignoreUnsupported) {
     auto saveTableName = [this] {
         FixPublicScheme();
         TableName = LastToken;
@@ -199,6 +201,10 @@ void TPgDumpParser::PgCatalogCheck() {
     if (IsSelect) {
         IsSelect = false;
         if (TableName.StartsWith("pg_catalog.set_config")) {
+            if (!IgnoreUnsupported) {
+                throw yexception() << "\"SELECT pg_catalog.set_config.*\" statement is not supported.\n" <<
+                    "Use \"--ignore-unsupported\" option if you want to ignore this statement.";
+            }
             TString tmpBuffer;
             while (!Buffer.empty()) {
                 auto token = ExtractToken(&tmpBuffer, [](char c){return !std::isspace(c);});
@@ -209,16 +215,32 @@ void TPgDumpParser::PgCatalogCheck() {
             }
             std::reverse(tmpBuffer.begin(), tmpBuffer.vend());
             Buffer += TStringBuilder() << "-- " << tmpBuffer;
+            Cerr << TStringBuilder() << "-- " << tmpBuffer << Endl;
+        }
+    }
+}
+
+bool TPgDumpParser::IsPrimaryKeyTokens(const TVector<TString>& tokens) {
+    const TVector<TString> desired = {"ALTER", "TABLE", "ONLY", "", "ADD", "CONSTRAINT", "", "PRIMARY", "KEY", "(", "", ")", ";"};
+    if (tokens.size() != desired.size()) {
+        return false;
+    }
+    for (size_t i = 0; i < desired.size(); ++i) {
+        if (!desired[i].empty() && desired[i] != tokens[i]) {
+            return false;
         }
     }
+    return true;
 }
 
 void TPgDumpParser::AlterTableCheck() {
     if (IsAlterTable) {
         IsAlterTable = false;
         TString tmpBuffer;
+        TVector <TString> tokens;
         while (!Buffer.empty()) {
             auto token = ExtractToken(&tmpBuffer, [](char c){return !std::isspace(c);});
+            tokens.push_back(token);
             if (token == "ALTER") {
                 break;
             }
@@ -228,8 +250,14 @@ void TPgDumpParser::AlterTableCheck() {
             }
             ExtractToken(&tmpBuffer, [](char c){return std::isspace(c);});
         }
+        std::reverse(tokens.begin(), tokens.end());
+        if (!IgnoreUnsupported && !IsPrimaryKeyTokens(tokens)) {
+            throw yexception() << "\"ALTER TABLE\" statement is not supported.\n" <<
+                "Use \"--ignore-unsupported\" option if you want to ignore this statement.";
+        }
         std::reverse(tmpBuffer.begin(), tmpBuffer.vend());
         Buffer += TStringBuilder() << "-- " << tmpBuffer;
+        Cerr << TStringBuilder() << "-- " << tmpBuffer << Endl;
     }
 }
 
diff --git a/ydb/public/lib/ydb_cli/common/pg_dump_parser.h b/ydb/public/lib/ydb_cli/common/pg_dump_parser.h
index b249a0dce8..39028ede23 100644
--- a/ydb/public/lib/ydb_cli/common/pg_dump_parser.h
+++ b/ydb/public/lib/ydb_cli/common/pg_dump_parser.h
@@ -57,12 +57,13 @@ class TPgDumpParser {
     };
 
 public:
-    explicit TPgDumpParser(IOutputStream& out);
+    explicit TPgDumpParser(IOutputStream& out, bool ignoreUnsupported);
 
     void Prepare(IInputStream& in);
     void WritePgDump(IInputStream& in);
 
 private:
+    static bool IsPrimaryKeyTokens(const TVector<TString>& tokens);
     static bool IsNewTokenSymbol(char c) {
         return c == '(' || c == ')' || c == ',' || c == ';' || c == '\'';
     }
@@ -88,9 +89,12 @@ private:
     bool IsPrimaryKey = false;
     bool IsCommented = false;
     bool NotFlush = false;
+
     size_t BracesCount = 0;
 
     IOutputStream& Out;
+    const bool IgnoreUnsupported;
+
     TSQLCommandNode RunRoot, PrepareRoot;
     TSQLCommandNode* CurrentNode = nullptr;
 };
diff --git a/ydb/public/lib/ydb_cli/common/pg_dump_parser_ut.cpp b/ydb/public/lib/ydb_cli/common/pg_dump_parser_ut.cpp
index 206c45f900..3b5d9bf233 100644
--- a/ydb/public/lib/ydb_cli/common/pg_dump_parser_ut.cpp
+++ b/ydb/public/lib/ydb_cli/common/pg_dump_parser_ut.cpp
@@ -8,7 +8,7 @@ using namespace NYdb::NConsoleClient;
 Y_UNIT_TEST_SUITE(PgDumpParserTests) {
     TString ParseDump(const TString& data) {
         TStringStream in, out;
-        TPgDumpParser parser(out);
+        TPgDumpParser parser(out, true);
 
         in << data;
         parser.Prepare(in);
@@ -21,7 +21,7 @@ Y_UNIT_TEST_SUITE(PgDumpParserTests) {
     TString ParseDumpFixedString(const TString& data) {
         TStringStream out;
         TFixedStringStream in(data);
-        TPgDumpParser parser(out);
+        TPgDumpParser parser(out, true);
 
         parser.Prepare(in);
         in.MovePointer();