Print location in error messages (#10948)

Author: CyberROFL

ydb/public/lib/ydb_cli/dump/restore_compat.cpp

@@ -46,23 +46,23 @@ class TDataAccumulator: public NPrivate::IDataAccumulator, protected TQueryBuild
         Begin();
     }
 
-    bool Fits(const TString& line) const override {
+    EStatus Check(const NPrivate::TLine& line) const override {
         if (RowsPerRequest > 0 && (Rows + 1) > RowsPerRequest) {
-            return Rows == 0;
+            return Rows == 0 ? OK : FULL;
         }
 
         if (BytesPerRequest > 0 && (Bytes + RecordSize(line)) > BytesPerRequest) {
-            return Bytes == 0;
+            return Bytes == 0 ? OK : FULL;
         }
 
         if (RequestUnitsPerRequest > 0 && CalcRequestUnits(RequestUnitsX2 + RecordRequestUnitsX2(line)) > RequestUnitsPerRequest) {
-            return RequestUnitsX2 == 0;
+            return RequestUnitsX2 == 0 ? OK : FULL;
         }
 
-        return true;
+        return OK;
     }
 
-    void Feed(TString&& line) override {
+    void Feed(NPrivate::TLine&& line) override {
         AddLine(line);
 
         ++Rows;
@@ -90,7 +90,7 @@ class TDataAccumulator: public NPrivate::IDataAccumulator, protected TQueryBuild
         return false;
     }
 
-    TString GetData(bool) override {
+    NPrivate::TBatch GetData(bool) override {
         Rows = 0;
         Bytes = 0;
         RequestUnitsX2 = 0;
@@ -133,7 +133,7 @@ class TDataWriter: public NPrivate::IDataWriter {
         Uploader = MakeHolder<TUploader>(opts, TableClient, Accumulator->GetQueryString());
     }
 
-    bool Push(TString&&) override {
+    bool Push(NPrivate::TBatch&&) override {
         bool ok;
 
         if (UseBulkUpsert) {

ydb/public/lib/ydb_cli/dump/restore_impl.cpp

@@ -15,7 +15,6 @@
 #include <util/generic/maybe.h>
 #include <util/generic/vector.h>
 #include <util/stream/file.h>
-#include <util/string/builder.h>
 #include <util/string/join.h>
 
 namespace NYdb {
@@ -83,6 +82,61 @@ bool IsOperationStarted(TStatus operationStatus) {
 
 } // anonymous
 
+namespace NPrivate {
+
+TLocation::TLocation(TStringBuf file, ui64 lineNo)
+    : File(file)
+    , LineNo(lineNo)
+{
+}
+
+void TLocation::Out(IOutputStream& out) const {
+    out << File << ":" << LineNo;
+}
+
+TLine::TLine(TString&& data, TStringBuf file, ui64 lineNo)
+    : Data(std::move(data))
+    , Location(file, lineNo)
+{
+}
+
+TLine::TLine(TString&& data, const TLocation& location)
+    : Data(std::move(data))
+    , Location(location)
+{
+}
+
+void TBatch::Add(const TLine& line) {
+    Data << line.GetData() << "\n";
+    Locations.push_back(line.GetLocation());
+}
+
+TString TBatch::GetLocation() const {
+    THashMap<TStringBuf, std::pair<ui64, ui64>> locations;
+    for (const auto& location : Locations) {
+        auto it = locations.find(location.File);
+        if (it == locations.end()) {
+            it = locations.emplace(location.File, std::make_pair(Max<ui64>(), Min<ui64>())).first;
+        }
+        it->second.first = Min(location.LineNo, it->second.first);
+        it->second.second = Max(location.LineNo, it->second.second);
+    }
+
+    TStringBuilder result;
+    bool comma = false;
+    for (const auto& [file, range] : locations) {
+        if (comma) {
+            result << ", ";
+        }
+        result << file << ":" << range.first << "-" << range.second;
+        comma = true;
+    }
+
+    return result;
+}
+
+} // NPrivate
+
 TRestoreClient::TRestoreClient(const TDriver& driver, const std::shared_ptr<TLog>& log)
     : ImportClient(driver)
     , OperationClient(driver)
@@ -371,25 +425,35 @@ TRestoreResult TRestoreClient::RestoreData(const TFsPath& fsPath, const TString&
                 return Result<TRestoreResult>(dbPath, std::move(descResult));
             }
 
-            accumulator.Reset(CreateImportDataAccumulator(desc, *actualDesc, settings));
+            accumulator.Reset(CreateImportDataAccumulator(desc, *actualDesc, settings, Log));
             writer.Reset(CreateImportDataWriter(dbPath, desc, ImportClient, TableClient, accumulator.Get(), settings, Log));
 
             break;
         }
     }
 
     TWriterWaiter waiter(*writer);
+
     ui32 dataFileId = 0;
     TFsPath dataFile = fsPath.Child(DataFileName(dataFileId));
+    TVector<TString> dataFileNames;
 
     while (dataFile.Exists()) {
         LOG_D("Read data from " << dataFile.GetPath().Quote());
 
+        dataFileNames.push_back(dataFile);
         TFileInput input(dataFile, settings.FileBufferSize_);
         TString line;
+        ui64 lineNo = 0;
 
         while (input.ReadLine(line)) {
-            while (!accumulator->Fits(line)) {
+            auto l = NPrivate::TLine(std::move(line), dataFileNames.back(), ++lineNo);
+            for (auto status = accumulator->Check(l); status != NPrivate::IDataAccumulator::OK; status = accumulator->Check(l)) {
+                if (status == NPrivate::IDataAccumulator::ERROR) {
+                    return Result<TRestoreResult>(dbPath, EStatus::GENERIC_ERROR,
+                        TStringBuilder() << "Invalid data: " << l.GetLocation());
+                }
+
                 if (!accumulator->Ready(true)) {
                     LOG_E("Error reading data from " << dataFile.GetPath().Quote());
                     return Result<TRestoreResult>(dbPath, EStatus::INTERNAL_ERROR, "Data is not ready");
@@ -401,7 +465,7 @@ TRestoreResult TRestoreClient::RestoreData(const TFsPath& fsPath, const TString&
                 }
             }
 
-            accumulator->Feed(std::move(line));
+            accumulator->Feed(std::move(l));
             if (accumulator->Ready()) {
                 if (!writer->Push(accumulator->GetData())) {
                     LOG_E("Error writing data to " << dbPath.Quote() << ", file: " << dataFile.GetPath().Quote());
@@ -518,3 +582,7 @@ TRestoreResult TRestoreClient::RestoreEmptyDir(const TFsPath& fsPath, const TStr
 
 } // NDump
 } // NYdb
+
+Y_DECLARE_OUT_SPEC(, NYdb::NDump::NPrivate::TLocation, o, x) {
+    return x.Out(o);
+}

ydb/public/lib/ydb_cli/dump/restore_impl.h

@@ -9,6 +9,8 @@
 
 #include <util/folder/path.h>
 #include <util/generic/hash_set.h>
+#include <util/stream/fwd.h>
+#include <util/string/builder.h>
 
 namespace NYdb {
 namespace NDump {
@@ -18,19 +20,91 @@ extern const char DOC_API_REQUEST_TYPE[22];
 
 namespace NPrivate {
 
+class TBatch;
+
+class TLocation {
+    friend class TBatch;
+    TStringBuf File;
+    ui64 LineNo;
+
+public:
+    explicit TLocation(TStringBuf file, ui64 lineNo);
+    void Out(IOutputStream& out) const;
+};
+
+class TLine {
+    TString Data;
+    TLocation Location;
+
+public:
+    explicit TLine(TString&& data, TStringBuf file, ui64 lineNo);
+    explicit TLine(TString&& data, const TLocation& location);
+
+    inline const TString& GetData() const {
+        return Data;
+    }
+
+    inline const TLocation& GetLocation() const {
+        return Location;
+    }
+
+    inline operator TString() const {
+        return Data;
+    }
+
+    inline operator TStringBuf() const {
+        return Data;
+    }
+
+    inline auto size() const {
+        return Data.size();
+    }
+};
+
+class TBatch {
+    TStringBuilder Data;
+    TVector<TLocation> Locations;
+
+public:
+    void Add(const TLine& line);
+    TString GetLocation() const;
+
+    inline const TString& GetData() const {
+        return Data;
+    }
+
+    inline const TLocation& GetLocation(ui32 index) const {
+        return Locations.at(index);
+    }
+
+    inline operator TString() const {
+        return Data;
+    }
+
+    inline auto size() const {
+        return Data.size();
+    }
+};
+
 class IDataAccumulator {
 public:
+    enum EStatus {
+        OK,
+        FULL,
+        ERROR,
+    };
+
     virtual ~IDataAccumulator() = default;
-    virtual bool Fits(const TString& line) const = 0;
-    virtual void Feed(TString&& line) = 0;
+    virtual EStatus Check(const TLine& line) const = 0;
+    virtual void Feed(TLine&& line) = 0;
     virtual bool Ready(bool force = false) const = 0;
-    virtual TString GetData(bool force = false) = 0;
+    virtual TBatch GetData(bool force = false) = 0;
 };
 
 class IDataWriter {
 public:
     virtual ~IDataWriter() = default;
-    virtual bool Push(TString&& data) = 0;
+    virtual bool Push(TBatch&& data) = 0;
     virtual void Wait() = 0;
 };