diff --git a/CMakeLists.txt b/CMakeLists.txt index db39d34fc..bf410e9e7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -218,7 +218,7 @@ message(STATUS "Generate coverage from unit tests: ${TEST_COVERAGE}") # includes include_directories(${src_dir}) include_directories(SYSTEM ${GTEST_INCLUDE_DIRS} ${GMOCK_INCLUDE_DIRS}) -include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/libs) +include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/libs) # cppitertools # needed to include configured files (plan_compiler_flags.hpp) set(generated_headers_dir ${CMAKE_BINARY_DIR}/generated_headers) include_directories(${generated_headers_dir}) diff --git a/src/durability/file_writer_buffer.hpp b/src/durability/file_writer_buffer.hpp index 852a65f43..fbee521e4 100644 --- a/src/durability/file_writer_buffer.hpp +++ b/src/durability/file_writer_buffer.hpp @@ -1,8 +1,8 @@ #pragma once #include -#include "utils/bswap.hpp" #include "hasher.hpp" +#include "utils/bswap.hpp" /** * Buffer that writes data to file and calculates hash of written data. @@ -19,6 +19,11 @@ class FileWriterBuffer { output_stream_.exceptions(std::ifstream::failbit | std::ifstream::badbit); } + /** + * Constructor which also takes a file path and opens it immediately. + */ + FileWriterBuffer(const std::string &path) : FileWriterBuffer() { Open(path); } + /** * Opens ofstream to file given in constructor. * @param file: @@ -72,8 +77,7 @@ class FileWriterBuffer { */ void WriteLong(uint64_t val) { uint64_t bval = bswap(val); - output_stream_.write(reinterpret_cast(&bval), - sizeof(bval)); + output_stream_.write(reinterpret_cast(&bval), sizeof(bval)); } /** diff --git a/src/utils/string.hpp b/src/utils/string.hpp index 31562ce54..876b2bde9 100644 --- a/src/utils/string.hpp +++ b/src/utils/string.hpp @@ -2,9 +2,8 @@ #include #include -#include - #include +#include #include #include #include @@ -91,9 +90,13 @@ inline std::string Replace(std::string src, const std::string &match, /** * Split string by delimeter and return vector of results. + * If the delimiter is not provided, a different splitting algorithm is used. + * Runs of consecutive whitespace are regarded as a single delimiter. + * Additionally, the result will not contain empty strings at the start of end + * as if the string was trimmed before splitting. */ inline std::vector Split(const std::string &src, - const std::string &delimiter = " ") { + const std::string &delimiter) { if (src.empty()) { return {}; } @@ -102,13 +105,31 @@ inline std::vector Split(const std::string &src, std::vector res; do { n = src.find(delimiter, index); - auto word = src.substr(index, n - index); - if (!word.empty()) res.push_back(word); + res.emplace_back(src.substr(index, n - index)); index = n + delimiter.size(); } while (n != std::string::npos); return res; } +/** + * Split string by whitespace and return vector of results. + */ +inline std::vector Split(const std::string &src) { + if (src.empty()) { + return {}; + } + std::regex not_whitespace("[^\\s]+"); + auto matches_begin = + std::sregex_iterator(src.begin(), src.end(), not_whitespace); + auto matches_end = std::sregex_iterator(); + std::vector res; + res.reserve(std::distance(matches_begin, matches_end)); + for (auto match = matches_begin; match != matches_end; ++match) { + res.emplace_back(match->str()); + } + return res; +} + /** * Parse double using classic locale, throws BasicException if it wasn't able to * parse whole string. @@ -132,4 +153,12 @@ inline bool EndsWith(const std::string &s, const std::string &suffix) { return s.size() >= suffix.size() && s.compare(s.size() - suffix.size(), std::string::npos, suffix) == 0; } + +/** + * Checks if the given string `s` starts with the given `prefix`. + */ +inline bool StartsWith(const std::string &s, const std::string &prefix) { + return s.size() >= prefix.size() && s.compare(0, prefix.size(), prefix) == 0; +} + } diff --git a/tests/public_benchmark/ldbc/build_dataset b/tests/public_benchmark/ldbc/build_dataset index edb72a7e4..5d7028d34 100755 --- a/tests/public_benchmark/ldbc/build_dataset +++ b/tests/public_benchmark/ldbc/build_dataset @@ -131,6 +131,6 @@ memgraph_snapshot_dir=${dataset_dir}/memgraph/default mkdir -p ${memgraph_snapshot_dir} cd ${memgraph_snapshot_dir} echo "Converting CSV dataset to '${memgraph_snapshot_dir}/snapshot'" -${base_dir}/tools/csv_to_snapshot -o snapshot ${csv_dataset} --csv-delimiter "|" --array-delimiter ";" +${base_dir}/tools/csv_to_snapshot --out snapshot ${csv_dataset} --csv-delimiter "|" --array-delimiter ";" echo "Done!" diff --git a/tools/.gitignore b/tools/.gitignore new file mode 100644 index 000000000..cc38133f6 --- /dev/null +++ b/tools/.gitignore @@ -0,0 +1,2 @@ +build +csv_to_snapshot diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt new file mode 100644 index 000000000..95b307551 --- /dev/null +++ b/tools/CMakeLists.txt @@ -0,0 +1,77 @@ +# MemGraph Tools CMake configuration + +cmake_minimum_required(VERSION 3.1) + +if (NOT UNIX) + message(FATAL, "Unsupported operating system.") +endif() + +# ccache setup +# ccache isn't enabled all the time because it makes some problem +# during the code coverage process +find_program(CCACHE_FOUND ccache) +option(USE_CCACHE "ccache:" ON) +message(STATUS "CCache: ${USE_CCACHE}") +if(CCACHE_FOUND AND USE_CCACHE) + set_property(GLOBAL PROPERTY RULE_LAUNCH_COMPILE ccache) + set_property(GLOBAL PROPERTY RULE_LAUNCH_LINK ccache) +endif(CCACHE_FOUND AND USE_CCACHE) + +# choose a compiler +# NOTE: must be choosen before use of project() or enable_language() +set(CMAKE_C_COMPILER "clang") +set(CMAKE_CXX_COMPILER "clang++") + +project("memgraph_tools") + +# setup CMake module path, defines path for include() and find_package() +# https://cmake.org/cmake/help/latest/variable/CMAKE_MODULE_PATH.html +set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${PROJECT_SOURCE_DIR}/../cmake) +# custom function definitions +include(functions) + +disallow_in_source_build() + +# threading +find_package(Threads REQUIRED) + +# optional readline +find_package(Readline REQUIRED) +if (READLINE_FOUND) + include_directories(SYSTEM ${READLINE_INCLUDE_DIR}) + add_definitions(-DHAS_READLINE) +endif() + +# c++14 +# TODO: set here 17 once it will be available in the cmake version (3.8) +set(cxx_standard 14) +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++1z -Wall -Wno-c++1z-extensions") +# Don't omit frame pointer in RelWithDebInfo, for additional callchain debug. +set(CMAKE_CXX_FLAGS_RELWITHDEBINFO + "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -fno-omit-frame-pointer") + +set(PREFERRED_DEBUGGER "gdb" CACHE STRING + "Tunes the debug output for your preferred debugger (gdb or lldb).") +if ("${PREFERRED_DEBUGGER}" STREQUAL "gdb" AND + "${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang|GNU") + set(CMAKE_CXX_FLAGS_DEBUG "-ggdb") +elseif ("${PREFERRED_DEBUGGER}" STREQUAL "lldb" AND + "${CMAKE_CXX_COMPILER_ID}" STREQUAL "Clang") + set(CMAKE_CXX_FLAGS_DEBUG "-glldb") +else() + message(WARNING "Unable to tune for PREFERRED_DEBUGGER: " + "'${PREFERRED_DEBUGGER}' with compiler: '${CMAKE_CXX_COMPILER_ID}'") + set(CMAKE_CXX_FLAGS_DEBUG "-g") +endif() + +# Setup external dependencies. Use EXCLUDE_FROM_ALL to prevent *installing* libs. +add_subdirectory(${PROJECT_SOURCE_DIR}/../libs libs EXCLUDE_FROM_ALL) +include_directories(${CMAKE_BINARY_DIR}/libs/gflags/include) +include_directories(${GLOG_INCLUDE_DIR}) +include_directories(${PROJECT_SOURCE_DIR}/../libs) # cppitertools + +# Include memgraph headers +set(memgraph_src_dir ${PROJECT_SOURCE_DIR}/../src) +include_directories(${memgraph_src_dir}) + +add_subdirectory(src) diff --git a/tools/csv_to_snapshot b/tools/csv_to_snapshot deleted file mode 100755 index 187f4ce89..000000000 --- a/tools/csv_to_snapshot +++ /dev/null @@ -1,351 +0,0 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- - -''' -Create a Memgraph recovery snapshot file from CSV. -''' - -import argparse -import csv -import itertools as it -import logging -import struct - -log = logging.getLogger(__name__) - -_CSV_TYPE_TO_PY_TYPE = { - 'int': int, - 'long': int, - 'float': float, - 'double': float, - 'boolean': bool, - 'byte': int, - 'short': int, - 'char': str, - 'string': str, - } - - -def csv_to_py_val(value, csv_type, array_delimiter): - if not value.strip(): - # Empty string signifies null or None in Python. - return None - if not csv_type.endswith('[]'): - return _CSV_TYPE_TO_PY_TYPE[csv_type](value) - # Otherwise we have an array type, so convert it to a list. - csv_type = csv_type[:-2] - py_type = _CSV_TYPE_TO_PY_TYPE[csv_type] - return [py_type(val.strip()) for val in value.split(array_delimiter)] - - -class NodeId: - def __init__(self, id_, id_space): - if not id_: - raise ValueError('ID must not be empty') - self.id = id_ - self.id_space = id_space - - def __eq__(self, other): - if not isinstance(other, NodeId): - return NotImplemented - return self.id == other.id and self.id_space == other.id_space - - def __hash__(self): - return hash((self.id, self.id_space)) - - def __str__(self): - if self.id_space is None: - return self.id - return '{}({})'.format(self.id, self.id_space) - - -class Hasher: - '''Implementation of memgraph/src/durability/hasher. - The API mimics hashlib, so that it will be easier to switch to something - more sane (e.g. sha256).''' - _PRIME = 3137 - - def __init__(self): - self._hash = 0 - - def update(self, data): - if not isinstance(data, bytes): - raise TypeError("Expected 'bytes', but got '{}'" - .format(type(data).__name__)) - for byte in data: - self._hash = self._hash * self._PRIME + byte + 1 - self._hash %= 2**64 # Make hash fit in uint64_t - - def digest(self): - '''Return the digest value as an int (which fits in uint64_t) and - *not* as bytes. (This is different from hashlib objects.)''' - return self._hash - - -class BoltEncoder: - # Type markers - _NULL_MARKER = b'\xC0' - _FLOAT64_MARKER = b'\xC1' - _FALSE_MARKER = b'\xC2' - _TRUE_MARKER = b'\xC3' - _INT64_MARKER = b'\xCB' - _STRING32_MARKER = b'\xD2' - _LIST32_MARKER = b'\xD6' - _MAP32_MARKER = b'\xDA' - _NODE_MARKER = b'\xB3\x4E' - _RELATIONSHIP_MARKER = b'\xB5\x52' - # Struct formats - _INT64_STRUCT = struct.Struct('>q') - _UINT32_STRUCT = struct.Struct('>I') - _UINT64_STRUCT = struct.Struct('>Q') - _FLOAT64_STRUCT = struct.Struct('>d') - - def __init__(self, file, hasher, skip_duplicate_nodes): - self._file = file - self._hasher = hasher - self._relationship_id = 0 - self._node_id = 0 - self._csv_to_mg_node_id = {} - self._skip_duplicate_nodes = skip_duplicate_nodes - - def write(self, value): - if value is None: - return self.write_null() - write = getattr(self, 'write_' + type(value).__name__) - write(value) - - def write_null(self): - self._write(self._NULL_MARKER) - - def write_bool(self, value): - if value: - self._write(self._TRUE_MARKER) - else: - self._write(self._FALSE_MARKER) - - def write_int(self, value): - self._write(self._INT64_MARKER) - self._write(self._INT64_STRUCT.pack(value)) - - def write_float(self, value): - self._write(self._FLOAT64_MARKER) - self._write(self._FLOAT64_STRUCT.pack(value)) - - def write_str(self, value): - self._write(self._STRING32_MARKER) - data = value.encode('utf-8') - self._write(self._UINT32_STRUCT.pack(len(data))) - self._write(data) - - def write_list(self, values): - self._write(self._LIST32_MARKER) - self._write(self._UINT32_STRUCT.pack(len(values))) - for value in values: - self.write(value) - - def write_dict(self, dict_value): - self._write(self._MAP32_MARKER) - self._write(self._UINT32_STRUCT.pack(len(dict_value))) - for key, value in dict_value.items(): - self.write_str(key) - self.write(value) - - def write_summary(self, node_count, relationship_count): - # It's a bit silly that the summary isn't considered for hashing - # (see: memgraph/src/durability/file_writer_buffer) - self._write(self._UINT64_STRUCT.pack(node_count), update_hash=False) - self._write(self._UINT64_STRUCT.pack(relationship_count), - update_hash=False) - self._write(self._UINT64_STRUCT.pack(self._hasher.digest()), - update_hash=False) - - def write_node(self, node_id, labels, properties): - id_ = None - try: - id_ = self._add_node_id(node_id) - except ValueError: - if self._skip_duplicate_nodes: - return - else: - raise - properties['id'] = node_id.id - self._write(self._NODE_MARKER) - self.write_int(id_) - self.write_list(labels) - self.write_dict(properties) - - def write_relationship(self, start_id, end_id, type_, properties): - self._write(self._RELATIONSHIP_MARKER) - self.write_int(self._relationship_id) - self.write_int(self._csv_to_mg_node_id[start_id]) - self.write_int(self._csv_to_mg_node_id[end_id]) - self._relationship_id += 1 - self.write_str(type_) - self.write_dict(properties) - - def _write(self, byte_data, update_hash=True): - log.debug("Writing bytes '0x{}'".format(byte_data.hex())) - if update_hash: - self._hasher.update(byte_data) - self._file.write(byte_data) - - def _add_node_id(self, node_id): - '''Add a new mapping from CSV node ID to Memgraph ID and - return the Memgraph ID.''' - if node_id in self._csv_to_mg_node_id: - raise ValueError("Node '{}' already exists".format(node_id)) - id_ = self._node_id - self._csv_to_mg_node_id[node_id] = id_ - self._node_id += 1 - return id_ - - -def parse_args(): - argp = argparse.ArgumentParser(description=__doc__) - argp.add_argument('-o', '--out', required=True, - help='Destination for the created snapshot file') - argp.add_argument('-n', '--nodes', action='append', required=True, - help='CSV file containing graph nodes (vertices)') - argp.add_argument('-r', '--relationships', default=[], action='append', - help='CSV file containing graph relationships (edges)') - argp.add_argument('--overwrite', action='store_true', default=False, - help='Overwrite the output file if it exists') - argp.add_argument('--log_level', default='WARNING', - choices=['INFO', 'WARNING', 'DEBUG'], - help='Log level, default is WARNING') - argp.add_argument('--array-delimiter', default=';', - help='Delimiter between elements of array values, ' - "default is ';'") - argp.add_argument('--csv-delimiter', default=',', - help='Delimiter between each field in the CSV, ' - "default is ','") - argp.add_argument('--skip-duplicate-nodes', action='store_true', default=False, - help='Skip duplicate nodes or raise an error (default)') - return argp.parse_args() - - -def get_field_name_and_type(field): - '''Return (field_name, field_type) from the field string. - If there is no type, field_type is returned as None.''' - field_name_and_type = field.split(':', maxsplit=1) - name = field_name_and_type[0] - if len(field_name_and_type) == 1: - return name, None - field_type = field_name_and_type[1].strip().lower() - return name, field_type - - -def get_id_space(field_type): - group_start = field_type.find('(') - if group_start == -1: - return None - return field_type[1 + field_type.find('('):-1] - - -def write_node_row(node_row, array_delimiter, encoder): - node_id = None - node_labels = [] - properties = {} - for field, value in node_row.items(): - value = value.strip() - name, field_type = get_field_name_and_type(field) - if field_type is not None and field_type.startswith('id'): - if node_id is not None: - raise ValueError('Only one node ID must be specified') - node_id = NodeId(value, get_id_space(field_type)) - elif field_type == 'label': - labels = map(str.strip, value.split(array_delimiter)) - node_labels.extend(label for label in labels if label) - elif field_type != 'ignore': - # Everything else is a property. - # Missing field_type defaults to string. - if not field_type: - field_type = 'string' - properties[name] = csv_to_py_val(value, field_type, array_delimiter) - if node_id is None: - raise ValueError('Node ID must be specified') - encoder.write_node(node_id, node_labels, properties) - - -def convert_nodes(node_filenames, csv_delimiter, array_delimiter, encoder): - node_count = 0 - for node_filename in node_filenames: - with open(node_filename, newline='', encoding='utf-8') as node_file: - nodes = csv.DictReader(node_file, delimiter=csv_delimiter) - for node in nodes: - write_node_row(node, array_delimiter, encoder) - node_count += 1 - return node_count - - -def write_relationship_row(relationship_row, array_delimiter, encoder): - start_id = None - end_id = None - relationship_type = None - properties = {} - for field, value in relationship_row.items(): - value = value.strip() - name, field_type = get_field_name_and_type(field) - if field_type is not None and field_type.startswith('start_id'): - if start_id is not None: - raise ValueError('Only one node ID must be specified') - start_id = NodeId(value, get_id_space(field_type)) - elif field_type is not None and field_type.startswith('end_id'): - if end_id is not None: - raise ValueError('Only one node ID must be specified') - end_id = NodeId(value, get_id_space(field_type)) - elif field_type == 'type': - if relationship_type is not None: - raise ValueError('Only one relationship TYPE must be specified') - relationship_type = value - elif field_type != 'ignore': - # Everything else is a property. - # Missing field_type defaults to string. - if not field_type: - field_type = 'string' - properties[name] = csv_to_py_val(value, field_type, array_delimiter) - if None in (start_id, end_id, relationship_type): - raise ValueError('Relationship TYPE, START_ID and END_ID must be set') - encoder.write_relationship(start_id, end_id, relationship_type, properties) - - -def convert_relationships(relationship_filenames, csv_delimiter, - array_delimiter, encoder): - relationship_count = 0 - for relationship_filename in relationship_filenames: - with open(relationship_filename, newline='', encoding='utf-8') as \ - relationship_file: - relationships = csv.DictReader(relationship_file, - delimiter=csv_delimiter) - for relationship in relationships: - write_relationship_row(relationship, array_delimiter, encoder) - relationship_count += 1 - return relationship_count - - -def main(): - args = parse_args() - logging.basicConfig(level=args.log_level) - all_input_names = ', '.join(it.chain(args.nodes, args.relationships)) - log.info("Converting {} to '{}'".format(all_input_names, args.out)) - with open(args.out, 'wb' if args.overwrite else 'xb') as dest_file: - hasher = Hasher() - encoder = BoltEncoder(dest_file, hasher, args.skip_duplicate_nodes) - # Snapshot file has the following contents in order: - # 1) list of label+property index - # 2) all nodes, sequantially, but not encoded as a list - # 3) all relationships, sequantially, but not encoded as a list - # 3) summary with node count, relationship count and hash digest - encoder.write_list([]) # Label + property indexes. - node_count = convert_nodes(args.nodes, args.csv_delimiter, - args.array_delimiter, encoder) - relationship_count = convert_relationships(args.relationships, - args.csv_delimiter, - args.array_delimiter, - encoder) - encoder.write_summary(node_count, relationship_count) - log.info("Created '{}'".format(args.out)) - - -if __name__ == '__main__': - main() diff --git a/tools/setup b/tools/setup new file mode 100755 index 000000000..9f04c900f --- /dev/null +++ b/tools/setup @@ -0,0 +1,16 @@ +#!/bin/bash -e + +# Builds the memgraph tools and installs them in this directory. + +script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" + +mkdir -p ${script_dir}/build +cd ${script_dir}/build + +# Setup cmake +cmake -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX=${script_dir} \ + ${script_dir} + +# Install the tools +make -j$(nproc) install diff --git a/tools/src/CMakeLists.txt b/tools/src/CMakeLists.txt new file mode 100644 index 000000000..404fff43a --- /dev/null +++ b/tools/src/CMakeLists.txt @@ -0,0 +1,20 @@ +add_executable(csv_to_snapshot + csv_to_snapshot/main.cpp + # This is just friggin terrible. csv_to_snapshot needs to depend almost on + # the whole memgraph, just to use TypedValue and BaseEncoder. + ${memgraph_src_dir}/data_structures/concurrent/skiplist_gc.cpp + ${memgraph_src_dir}/database/graph_db_accessor.cpp + ${memgraph_src_dir}/query/typed_value.cpp + ${memgraph_src_dir}/storage/edge_accessor.cpp + ${memgraph_src_dir}/storage/locking/record_lock.cpp + ${memgraph_src_dir}/storage/property_value.cpp + ${memgraph_src_dir}/storage/record_accessor.cpp + ${memgraph_src_dir}/storage/vertex_accessor.cpp + ${memgraph_src_dir}/transactions/transaction.cpp + ) + + +target_link_libraries(csv_to_snapshot stdc++fs Threads::Threads fmt gflags ${GLOG_LIBRARY}) +add_dependencies(csv_to_snapshot glog) +install(TARGETS csv_to_snapshot + RUNTIME DESTINATION .) diff --git a/tools/src/csv_to_snapshot/main.cpp b/tools/src/csv_to_snapshot/main.cpp new file mode 100644 index 000000000..41e696e2a --- /dev/null +++ b/tools/src/csv_to_snapshot/main.cpp @@ -0,0 +1,359 @@ +#include +#include +#include +#include +#include + +#include "cppitertools/chain.hpp" +#include "gflags/gflags.h" +#include "glog/logging.h" + +#include "communication/bolt/v1/encoder/base_encoder.hpp" +#include "durability/file_writer_buffer.hpp" +#include "utils/string.hpp" + +bool ValidateNotEmpty(const char *flagname, const std::string &value) { + if (utils::Trim(value).empty()) { + printf("The argument '%s' is required\n", flagname); + return false; + } + return true; +} + +DEFINE_string(out, "", "Destination for the created snapshot file"); +DEFINE_validator(out, &ValidateNotEmpty); +DEFINE_bool(overwrite, false, "Overwrite the output file if it exists"); +DEFINE_string(array_delimiter, ";", + "Delimiter between elements of array values, default is ';'"); +DEFINE_string(csv_delimiter, ",", + "Delimiter between each field in the CSV, default is ','"); +DEFINE_bool(skip_duplicate_nodes, false, + "Skip duplicate nodes or raise an error (default)"); +// Arguments `--nodes` and `--relationships` can be input multiple times and are +// handled with custom parsing. +DEFINE_string(nodes, "", "CSV file containing graph nodes (vertices)"); +DEFINE_validator(nodes, &ValidateNotEmpty); +DEFINE_string(relationships, "", + "CSV file containing graph relationships (edges)"); + +auto ParseRepeatedFlag(const std::string &flagname, int argc, char *argv[]) { + std::vector values; + for (int i = 1; i < argc; i += 2) { + std::string flag(argv[i]); + if ((flag == "--" + flagname || flag == "-" + flagname) && i + 1 < argc) + values.push_back(argv[i + 1]); + } + return values; +} + +// A field describing the CSV column. +struct Field { + // Name of the field. + std::string name; + // Type of the values under this field. + std::string type; +}; + +// A node ID from CSV format. +struct NodeId { + std::string id; + // Group/space of IDs. ID must be unique in a single group. + std::string id_space; +}; + +bool operator==(const NodeId &a, const NodeId &b) { + return a.id == b.id && a.id_space == b.id_space; +} + +auto &operator<<(std::ostream &stream, const NodeId &node_id) { + return stream << fmt::format("{}({})", node_id.id, node_id.id_space); +} + +namespace std { + +template <> +struct hash { + size_t operator()(const NodeId &node_id) const { + size_t id_hash = std::hash{}(node_id.id); + size_t id_space_hash = std::hash{}(node_id.id_space); + return id_hash ^ (id_space_hash << 1); + } +}; + +} // namespace std + +class MemgraphNodeIdMap { + public: + std::experimental::optional Get(const NodeId &node_id) const { + auto found_it = node_id_to_mg_.find(node_id); + if (found_it == node_id_to_mg_.end()) return std::experimental::nullopt; + return found_it->second; + } + + int64_t Insert(const NodeId &node_id) { + int64_t id = mg_id_++; + node_id_to_mg_[node_id] = id; + return id; + } + + private: + int64_t mg_id_ = 0; + std::unordered_map node_id_to_mg_; +}; + +std::vector ReadRow(std::istream &stream) { + std::vector row; + char quoting = 0; + std::vector column; + char c; + while (!stream.get(c).eof()) { + if (quoting) { + if (c == quoting) + quoting = 0; + else + column.push_back(c); + } else if (c == '"') { + // Hopefully, escaping isn't needed. + quoting = c; + } else if (c == FLAGS_csv_delimiter.front()) { + row.emplace_back(column.begin(), column.end()); + column.clear(); + } else if (c == '\n') { + row.emplace_back(column.begin(), column.end()); + return row; + } else { + column.push_back(c); + } + } + if (!column.empty()) row.emplace_back(column.begin(), column.end()); + return row; +} + +std::vector ReadHeader(std::istream &stream) { + auto row = ReadRow(stream); + std::vector fields; + fields.reserve(row.size()); + for (const auto &value : row) { + auto name_and_type = utils::Split(value, ":"); + CHECK(name_and_type.size() == 1U || name_and_type.size() == 2U) + << "Expected a name and optionally a type"; + auto name = name_and_type[0]; + // When type is missing, default is string. + std::string type("string"); + if (name_and_type.size() == 2U) + type = utils::ToLowerCase(utils::Trim(name_and_type[1])); + fields.push_back(Field{name, type}); + } + return fields; +} + +query::TypedValue StringToTypedValue(const std::string &str, + const std::string &type) { + // Empty string signifies Null. + if (str.empty()) return query::TypedValue::Null; + auto convert = [](const auto &str, const auto &type) -> query::TypedValue { + if (type == "int" || type == "long" || type == "byte" || type == "short") { + std::istringstream ss(str); + int64_t val; + ss >> val; + return val; + } else if (type == "float" || type == "double") { + return utils::ParseDouble(str); + } else if (type == "boolean") { + return utils::ToLowerCase(str) == "true" ? true : false; + } else if (type == "char" || type == "string") { + return str; + } + LOG(FATAL) << "Unexpected type: " << type; + return query::TypedValue::Null; + }; + // Type *not* ending with '[]', signifies regular value. + if (!utils::EndsWith(type, "[]")) return convert(str, type); + // Otherwise, we have an array type. + auto elem_type = type.substr(0, type.size() - 2); + auto elems = utils::Split(str, FLAGS_array_delimiter); + std::vector array; + array.reserve(elems.size()); + for (const auto &elem : elems) { + array.push_back(convert(utils::Trim(elem), elem_type)); + } + return array; +} + +std::string GetIdSpace(const std::string &type) { + auto start = type.find("("); + if (start == std::string::npos) return ""; + return type.substr(start + 1, type.size() - 1); +} + +void WriteNodeRow(const std::vector &fields, + const std::vector &row, + MemgraphNodeIdMap &node_id_map, + communication::bolt::BaseEncoder &encoder) { + std::experimental::optional id; + std::vector labels; + std::map properties; + for (int i = 0; i < row.size(); ++i) { + const auto &field = fields[i]; + auto value = utils::Trim(row[i]); + if (utils::StartsWith(field.type, "id")) { + CHECK(!id) << "Only one node ID must be specified"; + NodeId node_id{value, GetIdSpace(field.type)}; + if (node_id_map.Get(node_id)) { + if (FLAGS_skip_duplicate_nodes) { + LOG(WARNING) << fmt::format("Skipping duplicate node with id '{}'", + node_id); + return; + } else { + LOG(FATAL) << fmt::format("Node with id '{}' already exists", + node_id); + } + } + id = node_id_map.Insert(node_id); + properties["id"] = *id; + } else if (field.type == "label") { + for (const auto &label : utils::Split(value, FLAGS_array_delimiter)) { + labels.emplace_back(utils::Trim(label)); + } + } else if (field.type != "ignore") { + properties[field.name] = StringToTypedValue(value, field.type); + } + } + CHECK(id) << "Node ID must be specified"; + // write node + encoder.WriteRAW(underlying_cast(communication::bolt::Marker::TinyStruct) + + 3); + encoder.WriteRAW(underlying_cast(communication::bolt::Signature::Node)); + encoder.WriteInt(*id); + encoder.WriteList(labels); + encoder.WriteMap(properties); +} + +auto ConvertNodes(const std::string &nodes_path, MemgraphNodeIdMap &node_id_map, + communication::bolt::BaseEncoder &encoder) { + int64_t node_count = 0; + std::ifstream nodes_file(nodes_path); + CHECK(nodes_file) << fmt::format("Unable to open '{}'", nodes_path); + auto fields = ReadHeader(nodes_file); + auto row = ReadRow(nodes_file); + while (!row.empty()) { + CHECK_EQ(row.size(), fields.size()) + << "Expected as many values as there are header fields"; + WriteNodeRow(fields, row, node_id_map, encoder); + // Increase count and move to next row. + node_count += 1; + row = ReadRow(nodes_file); + } + return node_count; +} + +void WriteRelationshipsRow( + const std::vector &fields, const std::vector &row, + const MemgraphNodeIdMap &node_id_map, int64_t relationship_id, + communication::bolt::BaseEncoder &encoder) { + std::experimental::optional start_id; + std::experimental::optional end_id; + std::experimental::optional relationship_type; + std::map properties; + for (int i = 0; i < row.size(); ++i) { + const auto &field = fields[i]; + auto value = utils::Trim(row[i]); + if (utils::StartsWith(field.type, "start_id")) { + CHECK(!start_id) << "Only one node ID must be specified"; + NodeId node_id{value, GetIdSpace(field.type)}; + start_id = node_id_map.Get(node_id); + if (!start_id) + LOG(FATAL) << fmt::format("Node with id '{}' does not exist", node_id); + } else if (utils::StartsWith(field.type, "end_id")) { + CHECK(!end_id) << "Only one node ID must be specified"; + NodeId node_id{value, GetIdSpace(field.type)}; + end_id = node_id_map.Get(node_id); + if (!end_id) + LOG(FATAL) << fmt::format("Node with id '{}' does not exist", node_id); + } else if (field.type == "type") { + CHECK(!relationship_type) + << "Only one relationship TYPE must be specified"; + relationship_type = value; + } else if (field.type != "ignore") { + properties[field.name] = StringToTypedValue(value, field.type); + } + } + CHECK(start_id) << "START_ID must be set"; + CHECK(end_id) << "END_ID must be set"; + CHECK(relationship_type) << "Relationship TYPE must be set"; + // write relationship + encoder.WriteRAW(underlying_cast(communication::bolt::Marker::TinyStruct) + + 5); + encoder.WriteRAW( + underlying_cast(communication::bolt::Signature::Relationship)); + encoder.WriteInt(relationship_id); + encoder.WriteInt(*start_id); + encoder.WriteInt(*end_id); + encoder.WriteString(*relationship_type); + encoder.WriteMap(properties); +} + +auto ConvertRelationships( + const std::string &relationships_path, const MemgraphNodeIdMap &node_id_map, + communication::bolt::BaseEncoder &encoder) { + int64_t relationship_count = 0; + std::ifstream relationships_file(relationships_path); + CHECK(relationships_file) + << fmt::format("Unable to open '{}'", relationships_path); + auto fields = ReadHeader(relationships_file); + auto row = ReadRow(relationships_file); + while (!row.empty()) { + CHECK_EQ(row.size(), fields.size()) + << "Expected as many values as there are header fields"; + auto relationship_id = relationship_count; + WriteRelationshipsRow(fields, row, node_id_map, relationship_id, encoder); + // Increase count and move to next row. + relationship_count += 1; + row = ReadRow(relationships_file); + } + return relationship_count; +} + +void Convert(const std::vector &nodes, + const std::vector &relationships) { + FileWriterBuffer buffer(FLAGS_out); + communication::bolt::BaseEncoder encoder(buffer); + int64_t node_count = 0; + int64_t relationship_count = 0; + MemgraphNodeIdMap node_id_map; + // Snapshot file has the following contents in order: + // 1) list of label+property index + // 2) all nodes, sequentially, but not encoded as a list + // 3) all relationships, sequentially, but not encoded as a list + // 3) summary with node count, relationship count and hash digest + encoder.WriteList({}); // Label + property indexes. + for (const auto &nodes_file : nodes) { + node_count += ConvertNodes(nodes_file, node_id_map, encoder); + } + for (const auto &relationships_file : relationships) { + relationship_count += + ConvertRelationships(relationships_file, node_id_map, encoder); + } + buffer.WriteSummary(node_count, relationship_count); +} + +int main(int argc, char *argv[]) { + gflags::SetUsageMessage("Create a Memgraph recovery snapshot file from CSV."); + auto nodes = ParseRepeatedFlag("nodes", argc, argv); + auto relationships = ParseRepeatedFlag("relationships", argc, argv); + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + if (std::experimental::filesystem::exists(FLAGS_out) && !FLAGS_overwrite) { + LOG(FATAL) << fmt::format( + "File exists: '{}'. Pass --overwrite if you want to overwrite.", + FLAGS_out); + } + auto iter_all_inputs = iter::chain(nodes, relationships); + std::vector all_inputs(iter_all_inputs.begin(), + iter_all_inputs.end()); + LOG(INFO) << fmt::format("Converting {} to '{}'", + utils::Join(all_inputs, ", "), FLAGS_out); + Convert(nodes, relationships); + LOG(INFO) << fmt::format("Created '{}'", FLAGS_out); + return 0; +}