Merge pull request #1728 from joto/add-rollback-to-copy-mgr

Add rollback function to COPY manager

Author: lonvia

src/db-copy-mgr.hpp

@@ -43,6 +43,13 @@ class db_copy_mgr_t
 
             m_current = std::make_unique<db_cmd_copy_delete_t<DELETER>>(table);
         }
+        m_committed = m_current->buffer.size();
+    }
+
+    void rollback_line()
+    {
+        assert(m_current);
+        m_current->buffer.resize(m_committed);
     }
 
     /**
@@ -356,6 +363,7 @@ class db_copy_mgr_t
 
     std::shared_ptr<db_copy_thread_t> m_processor;
     std::unique_ptr<db_cmd_copy_delete_t<DELETER>> m_current;
+    std::size_t m_committed = 0;
 };
 
 #endif // OSM2PGSQL_DB_COPY_MGR_HPP

tests/test-db-copy-mgr.cpp

@@ -196,4 +196,46 @@ TEST_CASE("copy_mgr_t")
             CHECK(res == v);
         }
     }
+
+    SECTION("Insert something and roll back")
+    {
+        auto const t = setup_table("t text");
+
+        mgr.new_line(t);
+        mgr.add_column(0);
+        mgr.add_column("foo");
+        mgr.rollback_line();
+        mgr.sync();
+
+        auto const conn = db.connect();
+        CHECK(conn.get_count("test_copy_mgr") == 0);
+    }
+
+    SECTION("Insert something, insert more, roll back, insert something else")
+    {
+        auto const t = setup_table("t text");
+
+        mgr.new_line(t);
+        mgr.add_column(0);
+        mgr.add_column("good");
+        mgr.finish_line();
+
+        mgr.new_line(t);
+        mgr.add_column(1);
+        mgr.add_column("bad");
+        mgr.rollback_line();
+
+        mgr.new_line(t);
+        mgr.add_column(2);
+        mgr.add_column("better");
+        mgr.finish_line();
+        mgr.sync();
+
+        auto const conn = db.connect();
+        auto const res = conn.query(PGRES_TUPLES_OK,
+                                    "SELECT t FROM test_copy_mgr ORDER BY id");
+        CHECK(res.num_tuples() == 2);
+        CHECK(res.get_value_as_string(0, 0) == "good");
+        CHECK(res.get_value_as_string(1, 0) == "better");
+    }
 }