Added PostgreSQL support to harness client.

Reviewers: mislav.bradac

Reviewed By: mislav.bradac

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D753

tests/macro_benchmark/CMakeLists.txt

@@ -7,6 +7,9 @@ get_filename_component(test_type ${CMAKE_CURRENT_SOURCE_DIR} NAME)
 file(GLOB_RECURSE test_type_cpps *.cpp)
 message(STATUS "Available ${test_type} cpp files are: ${test_type_cpps}")
 
+# postgres directory
+set(postgres_dir ${libs_dir}/postgresql)
+
 # for each cpp file build binary and register test
 foreach(test_cpp ${test_type_cpps})
 
@@ -28,4 +31,8 @@ foreach(test_cpp ${test_type_cpps})
     # link libraries
     target_link_libraries(${target_name} memgraph_lib)
 
+    # libpq includes
+    target_link_libraries(${target_name} "${postgres_dir}/lib/libpq.so")
+    target_include_directories(${target_name} PUBLIC "${postgres_dir}/include")
+
 endforeach()

tests/macro_benchmark/harness/clients/bolt_client.hpp

@@ -0,0 +1,45 @@
+#include <fstream>
+#include <vector>
+
+#include <glog/logging.h>
+
+#include "communication/bolt/client.hpp"
+#include "communication/bolt/v1/decoder/decoded_value.hpp"
+#include "io/network/network_endpoint.hpp"
+#include "io/network/socket.hpp"
+
+using SocketT = io::network::Socket;
+using EndpointT = io::network::NetworkEndpoint;
+using ClientT = communication::bolt::Client<SocketT>;
+using QueryDataT = communication::bolt::QueryData;
+using communication::bolt::DecodedValue;
+
+class BoltClient {
+ public:
+  BoltClient(std::string &address, std::string &port, std::string &username,
+             std::string &password, std::string database = "") {
+    SocketT socket;
+    EndpointT endpoint;
+
+    try {
+      endpoint = EndpointT(address, port);
+    } catch (const io::network::NetworkEndpointException &e) {
+      LOG(FATAL) << "Invalid address or port: " << address << ":" << port;
+    }
+    if (!socket.Connect(endpoint)) {
+      LOG(FATAL) << "Could not connect to: " << address << ":" << port;
+    }
+
+    client_ = std::make_unique<ClientT>(std::move(socket), username, password);
+  }
+
+  QueryDataT Execute(const std::string &query,
+                     const std::map<std::string, DecodedValue> &parameters) {
+    return client_->Execute(query, parameters);
+  }
+
+  void Close() { client_->Close(); }
+
+ private:
+  std::unique_ptr<ClientT> client_;
+};

tests/macro_benchmark/harness/clients/common.hpp

@@ -47,17 +47,16 @@ void PrintJsonDecodedValue(std::ostream &os,
   }
 }
 
-template <typename SocketT>
+template <typename ClientT, typename ExceptionT>
 communication::bolt::QueryData ExecuteNTimesTillSuccess(
-    communication::bolt::Client<SocketT> &client, const std::string &query,
-    int times) {
+    ClientT &client, const std::string &query, int times) {
   for (int i = 0; i < times; ++i) {
     try {
       auto ret = client.Execute(query, {});
       return ret;
-    } catch (const communication::bolt::ClientQueryException &e) {
+    } catch (const ExceptionT &e) {
     }
   }
-  throw communication::bolt::ClientQueryException();
+  throw ExceptionT();
 }
 }

tests/macro_benchmark/harness/clients/harness_client.cpp

@@ -4,26 +4,27 @@
 #include <gflags/gflags.h>
 #include <glog/logging.h>
 
-#include "common.hpp"
-#include "communication/bolt/client.hpp"
 #include "communication/bolt/v1/decoder/decoded_value.hpp"
-#include "io/network/network_endpoint.hpp"
-#include "io/network/socket.hpp"
 #include "threading/sync/spinlock.hpp"
 #include "utils/algorithm.hpp"
 #include "utils/timer.hpp"
 
-using SocketT = io::network::Socket;
-using EndpointT = io::network::NetworkEndpoint;
-using ClientT = communication::bolt::Client<SocketT>;
-using communication::bolt::DecodedValue;
+#include "bolt_client.hpp"
+#include "common.hpp"
+#include "postgres_client.hpp"
+
+DEFINE_string(protocol, "bolt", "Protocol to use (available: bolt, postgres)");
+DEFINE_int32(num_workers, 1, "Number of workers");
+DEFINE_string(input, "", "Input file");
+DEFINE_string(output, "", "Output file");
 
 DEFINE_string(address, "127.0.0.1", "Server address");
-DEFINE_string(port, "7687", "Server port");
-DEFINE_int32(num_workers, 1, "Number of workers");
-DEFINE_string(output, "", "Output file");
+DEFINE_string(port, "", "Server port");
 DEFINE_string(username, "", "Username for the database");
 DEFINE_string(password, "", "Password for the database");
+DEFINE_string(database, "", "Database for the database");
+
+using communication::bolt::DecodedValue;
 
 const int MAX_RETRIES = 1000;
 
@@ -46,10 +47,11 @@ void PrintSummary(
   os << "}";
 }
 
-int main(int argc, char **argv) {
-  gflags::ParseCommandLineFlags(&argc, &argv, true);
-  google::InitGoogleLogging(argv[0]);
-
+template <typename ClientT, typename ExceptionT>
+void ExecuteQueries(std::istream &istream, int num_workers,
+                    std::ostream &ostream, std::string &address,
+                    std::string &port, std::string &username,
+                    std::string &password, std::string &database) {
   std::string query;
   std::vector<std::thread> threads;
 
@@ -58,30 +60,16 @@ int main(int argc, char **argv) {
   std::vector<std::string> queries;
   std::vector<std::map<std::string, DecodedValue>> metadata;
 
-  while (std::getline(std::cin, query)) {
+  while (std::getline(istream, query)) {
     queries.push_back(query);
   }
   metadata.resize(queries.size());
 
   utils::Timer timer;
 
-  for (int i = 0; i < FLAGS_num_workers; ++i) {
+  for (int i = 0; i < num_workers; ++i) {
     threads.push_back(std::thread([&]() {
-      SocketT socket;
-      EndpointT endpoint;
-
-      try {
-        endpoint = EndpointT(FLAGS_address, FLAGS_port);
-      } catch (const io::network::NetworkEndpointException &e) {
-        LOG(FATAL) << "Invalid address or port: " << FLAGS_address << ":"
-                   << FLAGS_port;
-      }
-      if (!socket.Connect(endpoint)) {
-        LOG(FATAL) << "Could not connect to: " << FLAGS_address << ":"
-                   << FLAGS_port;
-      }
-
-      ClientT client(std::move(socket), FLAGS_username, FLAGS_password);
+      ClientT client(address, port, username, password, database);
 
       std::string str;
       while (true) {
@@ -95,9 +83,10 @@ int main(int argc, char **argv) {
           str = queries[pos];
         }
         try {
-          metadata[pos] =
-              ExecuteNTimesTillSuccess(client, str, MAX_RETRIES).metadata;
-        } catch (const communication::bolt::ClientQueryException &e) {
+          metadata[pos] = ExecuteNTimesTillSuccess<ClientT, ExceptionT>(
+                              client, str, MAX_RETRIES)
+                              .metadata;
+        } catch (const ExceptionT &e) {
           LOG(FATAL) << "Could not execute query '" << str << "' "
                      << MAX_RETRIES << "times";
         }
@@ -106,19 +95,57 @@ int main(int argc, char **argv) {
     }));
   }
 
-  for (int i = 0; i < FLAGS_num_workers; ++i) {
+  for (int i = 0; i < num_workers; ++i) {
     threads[i].join();
   }
 
   auto elapsed = timer.Elapsed();
   double duration = elapsed.count();
 
+  PrintSummary(ostream, duration, metadata);
+}
+
+using BoltClientT = BoltClient;
+using BoltExceptionT = communication::bolt::ClientQueryException;
+
+using PostgresClientT = postgres::Client;
+using PostgresExceptionT = postgres::ClientQueryException;
+
+int main(int argc, char **argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+  google::InitGoogleLogging(argv[0]);
+
+  std::ifstream ifile;
+  std::istream *istream{&std::cin};
+
+  std::ofstream ofile;
+  std::ostream *ostream{&std::cout};
+
+  if (FLAGS_input != "") {
+    ifile.open(FLAGS_input);
+    istream = &ifile;
+  }
+
   if (FLAGS_output != "") {
-    std::ofstream ofile;
     ofile.open(FLAGS_output);
-    PrintSummary(ofile, duration, metadata);
-  } else {
-    PrintSummary(std::cout, duration, metadata);
+    ostream = &ofile;
+  }
+
+  std::string port = FLAGS_port;
+  if (FLAGS_protocol == "bolt") {
+    if (port == "") port = "7687";
+    ExecuteQueries<BoltClientT, BoltExceptionT>(
+        *istream, FLAGS_num_workers, *ostream, FLAGS_address, port,
+        FLAGS_username, FLAGS_password, FLAGS_database);
+  } else if (FLAGS_protocol == "postgres") {
+    permanent_assert(FLAGS_username != "",
+                     "Username can't be empty for postgres!");
+    permanent_assert(FLAGS_database != "",
+                     "Database can't be empty for postgres!");
+    if (port == "") port = "5432";
+    ExecuteQueries<PostgresClientT, PostgresExceptionT>(
+        *istream, FLAGS_num_workers, *ostream, FLAGS_address, port,
+        FLAGS_username, FLAGS_password, FLAGS_database);
   }
 
   return 0;

tests/macro_benchmark/harness/clients/postgres_client.hpp

@@ -0,0 +1,105 @@
+#pragma once
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include <fmt/format.h>
+#include <glog/logging.h>
+
+#include <libpq-fe.h>
+
+#include "communication/bolt/client.hpp"
+#include "communication/bolt/v1/decoder/decoded_value.hpp"
+#include "utils/exceptions.hpp"
+
+using communication::bolt::QueryData;
+
+namespace postgres {
+
+class ClientException : public utils::BasicException {
+  using utils::BasicException::BasicException;
+};
+
+class ClientQueryException : public ClientException {
+ public:
+  using ClientException::ClientException;
+  ClientQueryException() : ClientException("Couldn't execute query!") {}
+};
+
+class Client {
+ public:
+  Client(std::string &host, std::string &port, std::string &username,
+         std::string &password, std::string database = "") {
+    // https://www.postgresql.org/docs/9.4/static/libpq-connect.html#LIBPQ-PARAMKEYWORDS
+    std::string pass = "";
+    if (password != "") {
+      pass = "password=" + password;
+    }
+    std::string conninfo =
+        fmt::format("host={} port={} user={} {} dbname={} sslmode=disable",
+                    host, port, username, pass, database);
+
+    // Make a connection to the database.
+    connection_ = PQconnectdb(conninfo.c_str());
+
+    // Check to see that the backend connection was successfully made
+    if (PQstatus(connection_) != CONNECTION_OK) {
+      throw ClientException(PQerrorMessage(connection_));
+    }
+  }
+
+  QueryData Execute(const std::string &query,
+                    const std::map<std::string, std::string> &parameters) {
+    QueryData ret;
+
+    DLOG(INFO) << "Sending run message with statement: '" << query << "'";
+
+    result_ = PQexec(connection_, query.c_str());
+    if (PQresultStatus(result_) == PGRES_TUPLES_OK) {
+      // get fields
+      int num_fields = PQnfields(result_);
+      for (int i = 0; i < num_fields; ++i) {
+        ret.fields.push_back(PQfname(result_, i));
+      }
+
+      // get records
+      int num_records = PQntuples(result_);
+      ret.records.resize(num_records);
+      for (int i = 0; i < num_records; ++i) {
+        for (int j = 0; j < num_fields; ++j) {
+          ret.records[i].push_back(std::string(PQgetvalue(result_, i, j)));
+        }
+      }
+
+      // get metadata
+      ret.metadata.insert({"status", std::string(PQcmdStatus(result_))});
+      ret.metadata.insert({"rows_affected", std::string(PQcmdTuples(result_))});
+    } else if (PQresultStatus(result_) != PGRES_COMMAND_OK) {
+      throw ClientQueryException(PQerrorMessage(connection_));
+    }
+
+    PQclear(result_);
+    result_ = nullptr;
+
+    return ret;
+  }
+
+  void Close() {
+    if (result_ != nullptr) {
+      PQclear(result_);
+      result_ = nullptr;
+    }
+    if (connection_ != nullptr) {
+      PQfinish(connection_);
+      connection_ = nullptr;
+    }
+  }
+
+  ~Client() { Close(); }
+
+ private:
+  PGconn *connection_{nullptr};
+  PGresult *result_{nullptr};
+};
+}

tools/apollo/generate

@@ -29,6 +29,7 @@ BASE_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", ".."))
 WORKSPACE_DIR = os.path.normpath(os.path.join(BASE_DIR, ".."))
 BASE_DIR_NAME = os.path.basename(BASE_DIR)
 BUILD_DIR = os.path.join(BASE_DIR, "build")
+LIBS_DIR = os.path.join(BASE_DIR, "libs")
 TESTS_DIR = os.path.join(BUILD_DIR, "tests")
 OUTPUT_DIR = os.path.join(BUILD_DIR, "apollo")
 
@@ -195,11 +196,12 @@ MACRO_BENCHMARK_ARGS = (
     "--groups aggregation 1000_create unwind_create dense_expand "
     "--no-strict")
 macro_bench_path = os.path.join(BASE_DIR, "tests", "macro_benchmark")
-harness_client_binary = os.path.join(BUILD_RELEASE_DIR, "tests",
-        "macro_benchmark", "harness_client")
+harness_client_binaries = os.path.join(BUILD_RELEASE_DIR, "tests",
+        "macro_benchmark")
+postgresql_lib_dir = os.path.join(LIBS_DIR, "postgresql", "lib")
 infile = create_archive("macro_benchmark", [binary_release_path,
         binary_release_link_path, macro_bench_path, config_path,
-        harness_client_binary], cwd = WORKSPACE_DIR)
+        harness_client_binaries, postgresql_lib_dir], cwd = WORKSPACE_DIR)
 supervisor = "./memgraph/tests/macro_benchmark/harness/harness.py"
 outfile_paths = "\./memgraph/tests/macro_benchmark/harness/\.harness_summary"
 RUNS.append(generate_run("macro_benchmark", supervisor = supervisor,
@@ -210,16 +212,19 @@ RUNS.append(generate_run("macro_benchmark", supervisor = supervisor,
 if mode == "diff":
     PARENT_DIR = os.path.join(WORKSPACE_DIR, "parent")
     BUILD_PARENT_DIR = os.path.join(PARENT_DIR, "build")
+    LIBS_PARENT_DIR = os.path.join(PARENT_DIR, "libs")
     binary_parent_name = find_memgraph_binary(BUILD_PARENT_DIR)
     binary_parent_path = os.path.join(BUILD_PARENT_DIR, binary_parent_name)
     binary_parent_link_path = os.path.join(BUILD_PARENT_DIR, "memgraph")
     parent_config_path = os.path.join(PARENT_DIR, "config")
     parent_macro_bench_path = os.path.join(PARENT_DIR, "tests", "macro_benchmark")
-    parent_harness_client_binary = os.path.join(BUILD_PARENT_DIR, "tests",
-            "macro_benchmark", "harness_client")
+    parent_harness_client_binaries = os.path.join(BUILD_PARENT_DIR, "tests",
+            "macro_benchmark")
+    parent_postgresql_lib_dir = os.path.join(LIBS_PARENT_DIR, "postgresql", "lib")
     infile = create_archive("macro_benchmark_parent", [binary_parent_path,
             binary_parent_link_path, parent_macro_bench_path, parent_config_path,
-            parent_harness_client_binary], cwd = WORKSPACE_DIR)
+            parent_harness_client_binaries, parent_postgresql_lib_dir],
+            cwd = WORKSPACE_DIR)
     supervisor = "./parent/tests/macro_benchmark/harness/harness.py"
     args = MACRO_BENCHMARK_ARGS + " --RunnerBin " + binary_parent_path
     outfile_paths = "\./parent/tests/macro_benchmark/harness/\.harness_summary"