lbug 0.21.2

#include "binder/binder.h"
#include "binder/bound_import_database.h"
#include "common/copier_config/csv_reader_config.h"
#include "common/exception/binder.h"
#include "common/file_system/virtual_file_system.h"
#include "common/string_utils.h"
#include "main/client_context.h"
#include "parser/copy.h"
#include "parser/ddl/create_table.h"
#include "parser/parser.h"
#include "parser/port_db.h"
#include <format>

using namespace lbug::common;
using namespace lbug::parser;

namespace lbug {
namespace binder {

static std::string getQueryFromFile(VirtualFileSystem* vfs, const std::string& boundFilePath,
    const std::string& fileName, main::ClientContext* context) {
    auto filePath = vfs->joinPath(boundFilePath, fileName);
    if (!vfs->fileOrPathExists(filePath, context)) {
        if (fileName == PortDBConstants::COPY_FILE_NAME) {
            return "";
        }
        if (fileName == PortDBConstants::INDEX_FILE_NAME) {
            return "";
        }
        throw BinderException(std::format("File {} does not exist.", filePath));
    }
    auto fileInfo = vfs->openFile(filePath, FileOpenFlags(FileFlags::READ_ONLY
#ifdef _WIN32
                                                          | FileFlags::BINARY
#endif
                                                ));
    auto fsize = fileInfo->getFileSize();
    auto buffer = std::make_unique<char[]>(fsize);
    fileInfo->readFile(buffer.get(), fsize);
    return std::string(buffer.get(), fsize);
}

static std::string getColumnNamesToCopy(const CopyFrom& copyFrom) {
    std::string columns = "";
    std::string delimiter = "";
    for (auto& column : copyFrom.getCopyColumnInfo().columnNames) {
        columns += delimiter;
        columns += common::StringUtils::quoteIdentifier(column);
        if (delimiter == "") {
            delimiter = ",";
        }
    }
    if (columns.empty()) {
        return columns;
    }
    return std::format("({})", columns);
}

static std::string getCopyFilePath(const std::string& boundFilePath, const std::string& filePath) {
    if (filePath[0] == '/' || (std::isalpha(filePath[0]) && filePath[1] == ':')) {
        // Note:
        // Unix absolute path starts with '/'
        // Windows absolute path starts with "[DiskID]://"
        // This code path is for backward compatibility, we used to export the absolute path for
        // csv files to copy.cypher files.
        return filePath;
    }

    auto path = boundFilePath + "/" + filePath;
#if defined(_WIN32)
    // TODO(Ziyi): This is a temporary workaround because our parser requires input cypher queries
    // to escape all special characters in string literal. E.g. The user input query is: 'IMPORT
    // DATABASE 'C:\\db\\uw''. The parser removes any escaped characters and this function accepts
    // the path parameter as 'C:\db\uw'. Then the ImportDatabase operator gives the file path to
    // antlr4 parser directly without escaping any special characters in the path, which causes a
    // parser exception. However, the parser exception is not thrown properly which leads to the
    // undefined behaviour.
    size_t pos = 0;
    while ((pos = path.find('\\', pos)) != std::string::npos) {
        path.replace(pos, 1, "\\\\");
        pos += 2;
    }
#endif
    return path;
}

// Removes the `WITH (storage = '<dir>', format = 'icebug-disk')` clauses that EXPORT DATABASE
// attaches to CREATE TABLE statements in schema.cypher, so that IMPORT DATABASE (without an
// explicit `storage_format = 'icebug-disk'` override) creates native tables instead of mounting
// the exported parquet files in place.
static std::string stripIcebugStorageClauses(const std::string& schema) {
    std::string result = schema;
    size_t searchFrom = 0;
    size_t pos;
    while ((pos = result.find(" WITH (", searchFrom)) != std::string::npos) {
        auto end = result.find(')', pos);
        if (end == std::string::npos) {
            break;
        }
        auto clause = result.substr(pos, end - pos);
        if (clause.find("format = 'icebug-disk'") != std::string::npos &&
            clause.find("storage = '") != std::string::npos) {
            result.erase(pos, end - pos + 1);
            searchFrom = pos;
        } else {
            searchFrom = end;
        }
    }
    return result;
}

// Generates COPY statements that ingest the exported parquet data files into the native
// tables created by the (stripped) schema.cypher statements. The exported layout is
// nodes_<table>.parquet for node tables and rels_<rel>.parquet (from/to primary-key
// columns + edge properties) for relationship tables.
static std::string getNativeCopyStatements(VirtualFileSystem* vfs, const std::string& boundFilePath,
    const std::string& schemaQuery, main::ClientContext* context) {
    std::string copies;
    if (schemaQuery.empty()) {
        return copies;
    }
    for (auto& parsedStatement : Parser::parseQuery(schemaQuery)) {
        if (parsedStatement->getStatementType() != StatementType::CREATE_TABLE) {
            continue;
        }
        auto& createTable = parsedStatement->constCast<CreateTable>();
        const auto& info = *createTable.getInfo();
        auto isRel = info.type == TableType::REL;
        auto dataFile = isRel ? std::format("rels_{}.parquet", info.tableName) :
                                std::format("nodes_{}.parquet", info.tableName);
        auto filePath = getCopyFilePath(boundFilePath, dataFile);
        if (!vfs->fileOrPathExists(vfs->joinPath(boundFilePath, dataFile), context)) {
            continue;
        }
        // Explicit column list for node tables (covers SERIAL pk values, which COPY FROM
        // would otherwise auto-generate). Rel tables rely on positional binding: the flat
        // file already stores from/to key columns followed by properties in schema order.
        std::string columns;
        if (!isRel) {
            std::vector<std::string> columnList;
            for (auto& property : info.propertyDefinitions) {
                columnList.push_back(StringUtils::quoteIdentifier(property.getName()));
            }
            if (!columnList.empty()) {
                columns = std::format(" ({})", StringUtils::join(columnList, ", "));
            }
        }
        copies += std::format("COPY {}{} FROM \"{}\";\n",
            StringUtils::quoteIdentifier(info.tableName), columns, filePath);
    }
    return copies;
}

std::unique_ptr<BoundStatement> Binder::bindImportDatabaseClause(const Statement& statement) {
    auto& importDB = statement.constCast<ImportDB>();
    auto fs = VirtualFileSystem::GetUnsafe(*clientContext);
    auto boundFilePath = fs->expandPath(clientContext, importDB.getFilePath());
    if (!fs->fileOrPathExists(boundFilePath, clientContext)) {
        throw BinderException(std::format("Directory {} does not exist.", boundFilePath));
    }
    // By default the exported icebug-disk data files are ingested into native tables.
    // An explicit `storage_format = 'icebug-disk'` override mounts the parquet files in
    // place (i.e. the database is opened directly against the exported directory).
    bool mountIcebug = false;
    auto parsingOptions = bindParsingOptions(importDB.getParsingOptionsRef());
    if (parsingOptions.contains(PortDBConstants::IMPORT_STORAGE_FORMAT_OPTION)) {
        auto& value = parsingOptions.at(PortDBConstants::IMPORT_STORAGE_FORMAT_OPTION);
        if (value.getDataType().getLogicalTypeID() != LogicalTypeID::STRING) {
            throw BinderException("The type of storage_format option must be a string.");
        }
        auto valueStr = StringUtils::getUpper(value.getValue<std::string>());
        if (valueStr != "ICEBUG-DISK") {
            throw BinderException(std::format(
                "Unsupported import database storage format '{}'. Valid options are: icebug-disk.",
                value.getValue<std::string>()));
        }
        mountIcebug = true;
        parsingOptions.erase(PortDBConstants::IMPORT_STORAGE_FORMAT_OPTION);
    }
    if (!parsingOptions.empty()) {
        throw BinderException(
            std::format("Unrecognized import database option: {}.", parsingOptions.begin()->first));
    }
    std::string finalQueryStatements;
    auto schemaQuery =
        getQueryFromFile(fs, boundFilePath, PortDBConstants::SCHEMA_FILE_NAME, clientContext);
    if (!mountIcebug) {
        schemaQuery = stripIcebugStorageClauses(schemaQuery);
        finalQueryStatements += schemaQuery;
        finalQueryStatements +=
            getNativeCopyStatements(fs, boundFilePath, schemaQuery, clientContext);
    } else {
        finalQueryStatements += schemaQuery;
    }
    // replace the path in copy from statements with the bound path
    auto copyQuery =
        getQueryFromFile(fs, boundFilePath, PortDBConstants::COPY_FILE_NAME, clientContext);
    if (!copyQuery.empty()) {
        auto parsedStatements = Parser::parseQuery(copyQuery);
        for (auto& parsedStatement : parsedStatements) {
            DASSERT(parsedStatement->getStatementType() == StatementType::COPY_FROM);
            auto& copyFromStatement = parsedStatement->constCast<CopyFrom>();
            DASSERT(copyFromStatement.getSource()->type == ScanSourceType::FILE);
            auto filePaths =
                copyFromStatement.getSource()->constPtrCast<FileScanSource>()->filePaths;
            DASSERT(filePaths.size() == 1);
            auto fileTypeInfo = bindFileTypeInfo(filePaths);
            std::string query;
            auto copyFilePath = getCopyFilePath(boundFilePath, filePaths[0]);
            auto columnNames = getColumnNamesToCopy(copyFromStatement);
            auto parsingOptions = bindParsingOptions(copyFromStatement.getParsingOptions());
            std::unordered_map<std::string, std::string> copyFromOptions;
            if (parsingOptions.contains(CopyConstants::FROM_OPTION_NAME)) {
                DASSERT(parsingOptions.contains(CopyConstants::TO_OPTION_NAME));
                copyFromOptions[CopyConstants::FROM_OPTION_NAME] = std::format("'{}'",
                    parsingOptions.at(CopyConstants::FROM_OPTION_NAME).getValue<std::string>());
                copyFromOptions[CopyConstants::TO_OPTION_NAME] = std::format("'{}'",
                    parsingOptions.at(CopyConstants::TO_OPTION_NAME).getValue<std::string>());
                parsingOptions.erase(CopyConstants::FROM_OPTION_NAME);
                parsingOptions.erase(CopyConstants::TO_OPTION_NAME);
            }
            if (fileTypeInfo.fileType == FileType::CSV) {
                auto csvConfig = CSVReaderConfig::construct(parsingOptions);
                csvConfig.option.autoDetection = false;
                auto optionsMap = csvConfig.option.toOptionsMap(csvConfig.parallel);
                if (!copyFromOptions.empty()) {
                    optionsMap.insert(copyFromOptions.begin(), copyFromOptions.end());
                }
                query = std::format("COPY {} {} FROM \"{}\" {};",
                    StringUtils::quoteIdentifier(copyFromStatement.getTableName()), columnNames,
                    copyFilePath, CSVOption::toCypher(optionsMap));
            } else {
                query = std::format("COPY {} {} FROM \"{}\" {};",
                    StringUtils::quoteIdentifier(copyFromStatement.getTableName()), columnNames,
                    copyFilePath, CSVOption::toCypher(copyFromOptions));
            }
            finalQueryStatements += query;
        }
    }
    return std::make_unique<BoundImportDatabase>(boundFilePath, finalQueryStatements,
        getQueryFromFile(fs, boundFilePath, PortDBConstants::INDEX_FILE_NAME, clientContext));
}

} // namespace binder
} // namespace lbug