initial implementation of buffer filling functions in db-copy

Author: lonvia

db-copy.cpp

@@ -23,16 +23,32 @@ db_copy_thread_t::db_copy_thread_t(std::string const &conninfo)
     });
 }
 
+db_copy_thread_t::~db_copy_thread_t() { finish(); }
+
 void db_copy_thread_t::add_buffer(std::unique_ptr<db_cmd_t> &&buffer)
 {
+    assert(m_worker.joinable()); // thread must not have been finished
     std::unique_lock<std::mutex> lock(m_queue_mutex);
     m_worker_queue.push_back(std::move(buffer));
+    m_queue_cond.notify_one();
+}
+
+void db_copy_thread_t::sync_and_wait()
+{
+    std::promise<void> barrier;
+    std::future<void> sync = barrier.get_future();
+    add_buffer(std::unique_ptr<db_cmd_t>(new db_cmd_sync_t(std::move(barrier))));
+    sync.wait();
 }
 
 void db_copy_thread_t::finish()
 {
-    add_buffer(std::unique_ptr<db_cmd_t>(new db_cmd_finish_t()));
-    m_worker.join();
+    if (m_worker.joinable()) {
+        finish_copy();
+
+        add_buffer(std::unique_ptr<db_cmd_t>(new db_cmd_finish_t()));
+        m_worker.join();
+    }
 }
 
 void db_copy_thread_t::worker_thread()
@@ -61,6 +77,7 @@ void db_copy_thread_t::worker_thread()
                 execute_sql(static_cast<db_cmd_sql_t *>(item.get())->buffer);
                 break;
             case db_cmd_t::Cmd_sync:
+                finish_copy();
                 static_cast<db_cmd_sync_t *>(item.get())->barrier.set_value();
                 break;
             case db_cmd_t::Cmd_finish:
@@ -69,8 +86,7 @@ void db_copy_thread_t::worker_thread()
         }
     }
 
-    if (m_inflight)
-        finish_copy();
+    finish_copy();
 
     disconnect();
 }
@@ -93,8 +109,7 @@ void db_copy_thread_t::connect()
 
 void db_copy_thread_t::execute_sql(std::string const &sql_cmd)
 {
-    if (m_inflight)
-        finish_copy();
+    finish_copy();
 
     pgsql_exec_simple(m_conn, PGRES_COMMAND_OK, sql_cmd.c_str());
 }
@@ -117,7 +132,8 @@ void db_copy_thread_t::write_to_db(db_cmd_copy_t *buffer)
     if (!buffer->deletables.empty())
         delete_rows(buffer);
 
-    start_copy(buffer->target);
+    if (!m_inflight)
+        start_copy(buffer->target);
 
     pgsql_CopyData(buffer->target->name.c_str(), m_conn, buffer->buffer);
 }
@@ -130,24 +146,21 @@ void db_copy_thread_t::delete_rows(db_cmd_copy_t *buffer)
     sql.reserve(buffer->target->name.size() + buffer->deletables.size() * 15 +
                 30);
     sql += buffer->target->name;
-    sql += "WHERE ";
+    sql += " WHERE ";
     sql += buffer->target->id;
     sql += " IN (";
     for (auto id : buffer->deletables) {
         sql += std::to_string(id);
         sql += ',';
     }
-    sql += ')';
+    sql[sql.size() - 1] = ')';
 
     pgsql_exec_simple(m_conn, PGRES_COMMAND_OK, sql);
 }
 
 void db_copy_thread_t::start_copy(std::shared_ptr<db_target_descr_t> const &target)
 {
-    if (m_inflight)
-        return;
-
-    assert(m_inflight.get() == target.get());
+    m_inflight = target;
 
     std::string copystr = "COPY ";
     copystr.reserve(target->name.size() + target->rows.size() + 14);
@@ -165,6 +178,9 @@ void db_copy_thread_t::start_copy(std::shared_ptr<db_target_descr_t> const &targ
 
 void db_copy_thread_t::finish_copy()
 {
+    if (!m_inflight)
+        return;
+
     if (PQputCopyEnd(m_conn, nullptr) != 1)
         throw std::runtime_error((fmt("stop COPY_END for %1% failed: %2%\n") %
                                   m_inflight->name %
@@ -215,10 +231,12 @@ void db_copy_mgr_t::exec_sql(std::string const &sql_cmd)
 
 void db_copy_mgr_t::sync()
 {
-    std::promise<void> barrier;
-    std::future<void> sync = barrier.get_future();
-    m_processor->add_buffer(std::unique_ptr<db_cmd_t>(new db_cmd_sync_t(std::move(barrier))));
-    sync.wait();
+    // finish any ongoing copy operations
+    if (m_current) {
+        m_processor->add_buffer(std::move(m_current));
+    }
+
+    m_processor->sync_and_wait();
 }
 
 void db_copy_mgr_t::finish_line()

db-copy.hpp

@@ -3,6 +3,7 @@
 
 #include <condition_variable>
 #include <deque>
+#include <future>
 #include <memory>
 #include <mutex>
 #include <string>
@@ -108,12 +109,18 @@ class db_copy_thread_t
 {
 public:
     db_copy_thread_t(std::string const &conninfo);
+    ~db_copy_thread_t();
 
     /**
      * Add another command for the worker.
      */
     void add_buffer(std::unique_ptr<db_cmd_t> &&buffer);
 
+    /**
+     * Send sync command and wait for the notification.
+     */
+    void sync_and_wait();
+
     /**
      * Finish the copy process.
      *
@@ -166,21 +173,82 @@ class db_copy_mgr_t
      */
     void finish_line();
 
-    /**
-     * Add a column entry of simple type.
-     */
+    template <typename T, typename ...ARGS>
+    void add_columns(T value, ARGS&&... args)
+    {
+        add_column(value);
+        add_columns(args...);
+    }
+
+    template <typename T>
+    void add_columns(T value)
+    {
+        add_column(value);
+    }
+
+    /// Add a column entry of simple type.
     template <typename T>
     void add_column(T value)
     {
-        auto s = std::to_string(value);
-        add_column<>(s);
+        add_value(value);
+        m_current->buffer += '\t';
+    }
+
+    /// Add an empty column.
+    void add_null_column() { m_current->buffer += "\\N\t"; }
+
+    /// Start an array column.
+    void new_array() { m_current->buffer += "{"; }
+
+    /// Add a single value to an array column
+    template <typename T>
+    void add_array_elem(T value)
+    {
+        add_value(value);
+        m_current->buffer += ',';
     }
 
-    template <>
-    void add_column(std::string const &s)
+    void add_array_elem(std::string const &s) { add_array_elem(s.c_str()); }
+
+    void add_array_elem(char const *s)
     {
         assert(m_current);
-        m_current->buffer += s;
+        m_current->buffer += '"';
+        for (char const *c = s; *c; ++c) {
+            switch (*c) {
+                case '"':
+                    m_current->buffer += "\\\\\"";
+                    break;
+                case '\\':
+                    m_current->buffer += "\\\\\\\\";
+                    break;
+                case '\n':
+                    m_current->buffer += "\\n";
+                    break;
+                case '\r':
+                    m_current->buffer += "\\r";
+                    break;
+                case '\t':
+                    m_current->buffer += "\\t";
+                    break;
+                default:
+                    m_current->buffer += *c;
+                    break;
+            }
+        }
+
+        m_current->buffer += "\",";
+    }
+
+    /// Finish an array column.
+    void finish_array()
+    {
+        auto idx = m_current->buffer.size() - 1;
+        if (m_current->buffer[idx] == '{')
+            m_current->buffer += '}';
+        else
+            m_current->buffer[idx] = '}';
+        m_current->buffer += '\t';
     }
 
     /**
@@ -207,6 +275,42 @@ class db_copy_mgr_t
     void sync();
 
 private:
+    template <typename T>
+    void add_value(T value)
+    {
+        assert(m_current);
+        m_current->buffer += std::to_string(value);
+    }
+
+    void add_value(std::string const &s) { add_value(s.c_str()); }
+
+    void add_value(char const *s)
+    {
+        assert(m_current);
+        for (char const *c = s; *c; ++c) {
+            switch (*c) {
+                case '"':
+                    m_current->buffer += "\\\"";
+                    break;
+                case '\\':
+                    m_current->buffer += "\\\\";
+                    break;
+                case '\n':
+                    m_current->buffer += "\\n";
+                    break;
+                case '\r':
+                    m_current->buffer += "\\r";
+                    break;
+                case '\t':
+                    m_current->buffer += "\\t";
+                    break;
+                default:
+                    m_current->buffer += *c;
+                    break;
+            }
+        }
+    }
+
     std::shared_ptr<db_copy_thread_t> m_processor;
     std::unique_ptr<db_cmd_copy_t> m_current;
     unsigned m_last_line;