diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e4b2733..9b4c7fe 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,9 @@ concurrency: cancel-in-progress: true env: - EXTENSION_LIST: adbc;azure;delta;duckdb;fts;httpfs;iceberg;json;llm;neo4j;postgres;sqlite;unity_catalog;vector;algo + # pg_client requires libpq (PostgreSQL) at build time and a running + # PostgreSQL instance at test time (provided by run_pgembed_fixture.py). + EXTENSION_LIST: adbc;azure;delta;duckdb;fts;httpfs;iceberg;json;llm;neo4j;pg_client;postgres;sqlite;unity_catalog;vector;algo jobs: # ───────────────────────────────────────────────────────────────── @@ -206,7 +208,8 @@ jobs: dnf install -y 'dnf-command(config-manager)' dnf config-manager --set-enabled powertools dnf install -y cmake ninja-build python3 python3-pip git wget unzip \ - gcc-toolset-13 openssl3 openssl3-devel pkg-config + gcc-toolset-13 openssl3 openssl3-devel pkg-config \ + libpq libpq-devel - name: Install uv working-directory: ladybug diff --git a/CMakeLists.txt b/CMakeLists.txt index 43e853e..986e846 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -93,3 +93,4 @@ add_extension_if_enabled("llm") add_extension_if_enabled("httpfs") add_extension_if_enabled("neo4j") add_extension_if_enabled("algo") +add_extension_if_enabled("pg_client") diff --git a/extension_config.cmake b/extension_config.cmake index b182655..854cfe8 100644 --- a/extension_config.cmake +++ b/extension_config.cmake @@ -1,4 +1,4 @@ -set(EXTENSION_LIST adbc azure delta duckdb fts httpfs iceberg json llm postgres sqlite unity_catalog vector neo4j algo) +set(EXTENSION_LIST adbc azure delta duckdb fts httpfs iceberg json llm pg_client postgres sqlite unity_catalog vector neo4j algo) #set(EXTENSION_STATIC_LINK_LIST fts) foreach(extension IN LISTS EXTENSION_STATIC_LINK_LIST) diff --git a/pg_client/.gitignore b/pg_client/.gitignore new file mode 100644 index 0000000..567609b --- /dev/null +++ b/pg_client/.gitignore @@ -0,0 +1 @@ +build/ diff --git a/pg_client/CMakeLists.txt b/pg_client/CMakeLists.txt new file mode 100644 index 0000000..1941077 --- /dev/null +++ b/pg_client/CMakeLists.txt @@ -0,0 +1,19 @@ +find_package(PostgreSQL REQUIRED) + +include_directories( + ${PROJECT_SOURCE_DIR}/src/include + ${CMAKE_BINARY_DIR}/src/include + src/include + ${PostgreSQL_INCLUDE_DIRS}) + +add_subdirectory(src/connector) +add_subdirectory(src/storage) +add_subdirectory(src/main) +add_subdirectory(src/catalog) +add_subdirectory(src/function) + +build_extension_lib(${BUILD_STATIC_EXTENSION} "pg_client") + +target_link_libraries(lbug_${EXTENSION_LIB_NAME}_extension + PRIVATE + ${PostgreSQL_LIBRARIES}) diff --git a/pg_client/src/catalog/CMakeLists.txt b/pg_client/src/catalog/CMakeLists.txt new file mode 100644 index 0000000..1285916 --- /dev/null +++ b/pg_client/src/catalog/CMakeLists.txt @@ -0,0 +1,9 @@ +add_library(lbug_pg_client_catalog + OBJECT + pg_client_catalog.cpp + pg_client_table_catalog_entry.cpp) +add_dependencies(lbug_pg_client_catalog lbug_common lbug_catalog) + +set(PG_CLIENT_EXTENSION_OBJECT_FILES + ${PG_CLIENT_EXTENSION_OBJECT_FILES} $ + PARENT_SCOPE) diff --git a/pg_client/src/catalog/pg_client_catalog.cpp b/pg_client/src/catalog/pg_client_catalog.cpp new file mode 100644 index 0000000..78bc7d4 --- /dev/null +++ b/pg_client/src/catalog/pg_client_catalog.cpp @@ -0,0 +1,357 @@ +#include "catalog/pg_client_catalog.h" + +#include +#include + +#include "binder/bound_attach_info.h" +#include "binder/ddl/property_definition.h" +#include "catalog/catalog_entry/node_table_catalog_entry.h" +#include "catalog/catalog_entry/rel_group_catalog_entry.h" +#include "catalog/pg_client_table_catalog_entry.h" +#include "common/exception/binder.h" +#include "common/exception/runtime.h" +#include "common/string_utils.h" +#include "connector/pg_client_connector.h" +#include "function/pg_client_scan.h" +#include "main/client_context.h" +#include "main/database.h" +#include "storage/buffer_manager/memory_manager.h" +#include "storage/pg_client_storage.h" + +namespace lbug { +namespace pg_client_extension { + +PgClientCatalog::PgClientCatalog(std::string connStr, std::string catalogName, + std::string defaultSchemaName, main::ClientContext* context, + const PgClientConnector& connector) + : CatalogExtension{}, connStr{std::move(connStr)}, catalogName{std::move(catalogName)}, + defaultSchemaName{std::move(defaultSchemaName)}, connector{connector}, + context_{context} {} + +std::string PgClientCatalog::bindSchemaName(const binder::AttachOption& options, + const std::string& defaultName) { + auto& opts = options; + if (opts.options.contains(PgClientStorageExtension::SCHEMA_OPTION)) { + auto val = opts.options.at(PgClientStorageExtension::SCHEMA_OPTION); + if (val.getDataType().getLogicalTypeID() != common::LogicalTypeID::STRING) { + throw common::RuntimeException{ + std::format("Invalid option value for {}", PgClientStorageExtension::SCHEMA_OPTION)}; + } + return val.getValue(); + } + return defaultName; +} + +void PgClientCatalog::init() { + // Query information_schema.tables to find node_* and rel_* tables + auto query = std::format( + "select table_name from information_schema.tables " + "where table_catalog = '{}' and table_schema = '{}' " + "order by table_name", + catalogName, defaultSchemaName); + + auto result = connector.executeQuery(query); + if (!result.success) { + throw common::BinderException( + std::format("Failed to query tables: {}", result.errorMessage)); + } + + for (auto& row : result.rows) { + std::string tableName = row.cells[0].value; + auto lowerName = tableName; + common::StringUtils::toLower(lowerName); + if (lowerName.rfind("rel_", 0) == 0) { + createForeignRelTable(tableName); + } else if (lowerName.rfind("node_", 0) == 0) { + createForeignNodeTable(tableName); + } + } +} + +static std::vector getTableColumnInfoFromConnector( + const PgClientConnector& connector, const std::string& catalogName, + const std::string& schemaName, const std::string& tableName) { + std::vector result; + + auto query = std::format( + "select column_name, data_type from information_schema.columns " + "where table_catalog = '{}' and table_schema = '{}' " + "and table_name = '{}' order by ordinal_position", + catalogName, schemaName, tableName); + + auto queryResult = connector.executeQuery(query); + if (!queryResult.success) { + return result; + } + + result.reserve(queryResult.rows.size()); + for (auto& row : queryResult.rows) { + std::string colName = row.cells[0].value; + std::string pgType = row.cells[1].value; + result.push_back({colName, pgTypeNameToLogicalType(pgType)}); + } + + return result; +} + +common::LogicalType pgTypeNameToLogicalType(const std::string& pgType) { + auto upper = pgType; + common::StringUtils::toUpper(upper); + + if (upper == "BOOLEAN" || upper == "BOOL") { + return common::LogicalType(common::LogicalTypeID::BOOL); + } else if (upper == "BIGINT" || upper == "INT8" || upper == "INTEGER" || + upper == "INT" || upper == "INT4" || upper == "SMALLINT" || + upper == "INT2" || upper == "SERIAL" || upper == "BIGSERIAL" || + upper == "OID") { + return common::LogicalType(common::LogicalTypeID::INT64); + } else if (upper == "REAL" || upper == "FLOAT4") { + return common::LogicalType(common::LogicalTypeID::FLOAT); + } else if (upper == "DOUBLE PRECISION" || upper == "FLOAT8" || upper == "NUMERIC" || + upper == "DECIMAL") { + return common::LogicalType(common::LogicalTypeID::DOUBLE); + } else if (upper == "TEXT" || upper == "VARCHAR" || upper == "CHARACTER VARYING" || + upper == "CHAR" || upper == "CHARACTER" || upper == "BPCHAR" || + upper == "NAME" || upper == "XML" || upper == "JSON" || upper == "JSONB") { + return common::LogicalType(common::LogicalTypeID::STRING); + } else if (upper == "BYTEA" || upper == "BLOB") { + return common::LogicalType(common::LogicalTypeID::STRING); + } else if (upper == "DATE") { + return common::LogicalType(common::LogicalTypeID::STRING); + } else if (upper == "TIMESTAMP" || upper == "TIMESTAMP WITHOUT TIME ZONE" || + upper == "TIMESTAMP WITH TIME ZONE" || upper == "TIMESTAMPTZ") { + return common::LogicalType(common::LogicalTypeID::STRING); + } else if (upper == "INTERVAL") { + return common::LogicalType(common::LogicalTypeID::STRING); + } else if (upper == "UUID") { + return common::LogicalType(common::LogicalTypeID::STRING); + } else if (upper == "BOOLEAN[]" || upper == "INTEGER[]" || upper == "TEXT[]" || + upper == "VARCHAR[]") { + return common::LogicalType(common::LogicalTypeID::STRING); + } else { + // Default to string for unknown types + return common::LogicalType(common::LogicalTypeID::STRING); + } +} + +void PgClientCatalog::createForeignNodeTable(const std::string& tableName) { + // Get column info from information_schema.columns + auto columnInfo = getTableColumnInfoFromConnector(connector, catalogName, + defaultSchemaName, tableName); + + if (columnInfo.empty()) { + return; + } + + // Build property definitions + std::vector propertyDefs; + std::string pkName = columnInfo[0].name; + for (auto& col : columnInfo) { + propertyDefs.emplace_back( + binder::ColumnDefinition{col.name, col.type.copy()}); + } + + // Create the PgClientTableScanInfo for this table + auto columnNames = getColumnNames(columnInfo); + auto columnTypes = getColumnTypes(columnInfo); + + auto query = std::format("SELECT {{}} FROM \"{}\".\"{}\"", + defaultSchemaName, tableName); + + auto scanInfo = std::make_shared(query, + std::move(columnTypes), std::move(columnNames), connector); + + // Create the scan function + auto scanFunction = getScanFunction(scanInfo); + + // Create the foreign table catalog entry in the attached catalog + // Save a raw pointer before moving it into the catalog set + auto foreignTableEntry = std::make_unique( + tableName, std::move(scanFunction), scanInfo); + + for (auto& def : propertyDefs) { + foreignTableEntry->addProperty(def); + } + + auto* attachedEntryPtr = foreignTableEntry.get(); + tables->createEntry(&transaction::DUMMY_TRANSACTION, std::move(foreignTableEntry)); + + // Create a main catalog entry for node table support (in-place queries) + auto foreignDatabaseName = std::format("{}.{}", defaultSchemaName, tableName); + auto mainTableEntry = std::make_unique( + tableName, pkName, foreignDatabaseName, catalog::ShadowTag{}); + + for (auto& def : propertyDefs) { + mainTableEntry->addProperty(def); + } + + // Link the shadow entry to the foreign entry so planners can find the scan function + mainTableEntry->setReferencedEntry(attachedEntryPtr); + + context_->getDatabase()->getCatalog()->addTableEntry(std::move(mainTableEntry)); + + // Register in the storage manager so the query planner can find this table + auto* mainEntry = context_->getDatabase()->getCatalog()->getTableCatalogEntry( + &transaction::DUMMY_TRANSACTION, tableName); + if (mainEntry) { + storage::StorageManager::Get(*context_)->createTable(mainEntry); + } +} + +void PgClientCatalog::createForeignRelTable(const std::string& tableName) { + // Query foreign key info to find src/dst node tables + auto fkQuery = std::format( + "SELECT kcu.column_name, ccu.table_name " + "FROM information_schema.table_constraints tc " + "JOIN information_schema.key_column_usage kcu " + " ON tc.constraint_name = kcu.constraint_name " + " AND tc.table_schema = kcu.table_schema " + "JOIN information_schema.constraint_column_usage ccu " + " ON ccu.constraint_name = tc.constraint_name " + " AND ccu.table_schema = tc.table_schema " + "WHERE tc.constraint_type = 'FOREIGN KEY' " + " AND tc.table_schema = '{}' " + " AND tc.table_name = '{}'", + defaultSchemaName, tableName); + auto fkResult = connector.executeQuery(fkQuery); + + std::string srcTableName, dstTableName; + for (auto& row : fkResult.rows) { + auto colName = row.cells[0].value; + auto refTable = row.cells[1].value; + auto lowerCol = colName; + common::StringUtils::toLower(lowerCol); + if (lowerCol.rfind("src", 0) == 0 || lowerCol.rfind("from", 0) == 0) { + srcTableName = refTable; + } else if (lowerCol.rfind("dst", 0) == 0 || lowerCol.rfind("dest", 0) == 0 || + lowerCol.rfind("to", 0) == 0) { + dstTableName = refTable; + } + } + + if (srcTableName.empty() || dstTableName.empty()) { + // No FK info found — register as a plain foreign table instead + createForeignNodeTable(tableName); + return; + } + + // Get column info + auto columnInfo = getTableColumnInfoFromConnector(connector, catalogName, + defaultSchemaName, tableName); + + if (columnInfo.empty()) { + return; + } + + // Build property definitions + std::vector propertyDefs; + for (auto& col : columnInfo) { + propertyDefs.emplace_back( + binder::ColumnDefinition{col.name, col.type.copy()}); + } + + auto columnNames = getColumnNames(columnInfo); + auto columnTypes = getColumnTypes(columnInfo); + + // Look up src/dst node tables in the main catalog to get table IDs + auto* catalog = context_->getDatabase()->getCatalog(); + auto* srcEntry = catalog->getTableCatalogEntry(&transaction::DUMMY_TRANSACTION, srcTableName); + auto* dstEntry = catalog->getTableCatalogEntry(&transaction::DUMMY_TRANSACTION, dstTableName); + if (srcEntry == nullptr || dstEntry == nullptr) { + createForeignNodeTable(tableName); + return; + } + + common::table_id_t srcTableID = srcEntry->getTableID(); + common::table_id_t dstTableID = dstEntry->getTableID(); + + // Build scan function + auto query = std::format("SELECT {{}} FROM \"{}\".\"{}\"", + defaultSchemaName, tableName); + + auto scanInfo = std::make_shared(query, + std::move(columnTypes), std::move(columnNames), connector); + auto scanFunc = getScanFunction(scanInfo); + + // Create foreign table entry in attached catalog + auto foreignTableEntry = std::make_unique( + tableName, scanFunc, scanInfo); + auto* attachedEntryPtr = foreignTableEntry.get(); + for (auto& def : propertyDefs) { + foreignTableEntry->addProperty(def); + } + tables->createEntry(&transaction::DUMMY_TRANSACTION, std::move(foreignTableEntry)); + + // Create bind data for the scan function (empty columns — populated at query time) + binder::expression_vector emptyColumns; + auto bindData = std::make_shared(query, columnNames, + connector, emptyColumns); + + // Create RelGroupCatalogEntry in the main catalog so MATCH ... -[...]-> ... + // queries can find the rel table. ForeignRelTable (created via + // StorageManager::createTable) is responsible for executing the scan using + // the scan function / bind data; the connector is mutex-guarded so the + // underlying libpq connection is safe across parallel hash-join workers, + // and ForeignRelTable now owns the shared state and serializes offset + // advancement with its own mutex, matching the morsel-driven model. + auto foreignDatabaseName = std::format("{}.{}", defaultSchemaName, tableName); + + std::vector relTableInfos; + common::oid_t relOID = tables->getNextOID(); + relTableInfos.emplace_back( + catalog::NodeTableIDPair{srcTableID, dstTableID}, relOID, + common::RelMultiplicity::MANY, common::RelMultiplicity::MANY); + + auto relGroupEntry = + std::make_unique(tableName, + common::RelMultiplicity::MANY, common::RelMultiplicity::MANY, + common::ExtendDirection::BOTH, std::move(relTableInfos), + "", // storage + common::StorageFormat::NONE, scanFunc, bindData, + std::move(foreignDatabaseName)); + + for (auto& def : propertyDefs) { + relGroupEntry->addProperty(def); + } + + context_->getDatabase()->getCatalog()->addTableEntry(std::move(relGroupEntry)); + + // Set up storage for the rel table so the planner can find it. This + // instantiates a ForeignRelTable, which lazily creates the scan function's + // shared state on the first initScanState and shares it across all + // worker threads (morsel-driven parallelism). + auto mainEntry = context_->getDatabase()->getCatalog()->getTableCatalogEntry( + &transaction::DUMMY_TRANSACTION, tableName); + if (mainEntry) { + storage::StorageManager::Get(*context_)->createTable(mainEntry); + } +} + +std::vector PgClientCatalog::getTableColumnInfo( + const std::string& tableName) const { + return getTableColumnInfoFromConnector(connector, catalogName, + defaultSchemaName, tableName); +} + +std::vector PgClientCatalog::getColumnNames( + const std::vector& columns) const { + std::vector names; + names.reserve(columns.size()); + for (auto& col : columns) { + names.push_back(col.name); + } + return names; +} + +std::vector PgClientCatalog::getColumnTypes( + const std::vector& columns) const { + std::vector types; + types.reserve(columns.size()); + for (auto& col : columns) { + types.push_back(col.type.copy()); + } + return types; +} + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/catalog/pg_client_table_catalog_entry.cpp b/pg_client/src/catalog/pg_client_table_catalog_entry.cpp new file mode 100644 index 0000000..8acaa4e --- /dev/null +++ b/pg_client/src/catalog/pg_client_table_catalog_entry.cpp @@ -0,0 +1,66 @@ +#include "catalog/pg_client_table_catalog_entry.h" + +#include "binder/bound_scan_source.h" +#include "binder/expression/variable_expression.h" +#include "common/constants.h" +#include "function/pg_client_scan.h" + +namespace lbug { +namespace catalog { + +PgClientTableCatalogEntry::PgClientTableCatalogEntry(std::string name, + std::optional scanFunction, + std::shared_ptr scanInfo) + : TableCatalogEntry{CatalogEntryType::FOREIGN_TABLE_ENTRY, std::move(name)}, + scanFunction{std::move(scanFunction)}, scanInfo{std::move(scanInfo)} {} + +common::TableType PgClientTableCatalogEntry::getTableType() const { + return common::TableType::FOREIGN; +} + +std::unique_ptr PgClientTableCatalogEntry::getBoundScanInfo( + main::ClientContext* context, const std::string& nodeUniqueName) { + auto columnNames = scanInfo->getColumnNames(); + auto columnTypes = scanInfo->getColumnTypes(*context); + binder::expression_vector columns; + + // Add rowid as _ID (internal ID) if nodeUniqueName is provided + if (!nodeUniqueName.empty()) { + auto idUniqueName = nodeUniqueName + "." + std::string(common::InternalKeyword::ID); + columns.push_back(std::make_shared(common::LogicalType::INT64(), + idUniqueName, "rowid")); + } + + for (auto i = 0u; i < columnNames.size(); i++) { + std::string uniqueName = columnNames[i]; + if (!nodeUniqueName.empty()) { + uniqueName = nodeUniqueName + "." + columnNames[i]; + } + columns.push_back(std::make_shared(std::move(columnTypes[i]), + uniqueName, columnNames[i])); + } + + // Build column names for PG query (no rowid — PG doesn't have a rowid pseudo-column) + // The internal ID column will be synthesized by the table function from row indices + std::vector pgColumnNames = columnNames; + + auto bindData = + std::make_unique( + scanInfo->getTemplateQuery(*context), + pgColumnNames, scanInfo->getConnector(), std::move(columns)); + return std::make_unique(scanFunction, std::move(bindData)); +} + +std::unique_ptr PgClientTableCatalogEntry::copy() const { + auto other = std::make_unique(name, scanFunction, scanInfo); + other->copyFrom(*this); + return other; +} + +std::unique_ptr +PgClientTableCatalogEntry::getBoundExtraCreateInfo(transaction::Transaction*) const { + UNREACHABLE_CODE; +} + +} // namespace catalog +} // namespace lbug diff --git a/pg_client/src/connector/CMakeLists.txt b/pg_client/src/connector/CMakeLists.txt new file mode 100644 index 0000000..771fe9c --- /dev/null +++ b/pg_client/src/connector/CMakeLists.txt @@ -0,0 +1,9 @@ +add_library(lbug_pg_client_connector + OBJECT + pg_client_connector.cpp) + +add_dependencies(lbug_pg_client_connector lbug_common) + +set(PG_CLIENT_EXTENSION_OBJECT_FILES + ${PG_CLIENT_EXTENSION_OBJECT_FILES} $ + PARENT_SCOPE) diff --git a/pg_client/src/connector/pg_client_connector.cpp b/pg_client/src/connector/pg_client_connector.cpp new file mode 100644 index 0000000..7ba91c9 --- /dev/null +++ b/pg_client/src/connector/pg_client_connector.cpp @@ -0,0 +1,135 @@ +#include "connector/pg_client_connector.h" + +#include +#include + +#include "common/exception/runtime.h" +#include "common/types/types.h" + +namespace lbug { +namespace pg_client_extension { + +PgClientConnector::PgClientConnector() : conn{nullptr} {} + +PgClientConnector::~PgClientConnector() { + disconnect(); +} + +void PgClientConnector::connect(const std::string& connStr) { + if (conn) { + disconnect(); + } + conn = PQconnectdb(connStr.c_str()); + if (PQstatus(conn) != CONNECTION_OK) { + std::string error = PQerrorMessage(conn); + disconnect(); + throw common::RuntimeException( + std::format("Failed to connect to PostgreSQL: {}", error)); + } +} + +void PgClientConnector::disconnect() { + if (conn) { + PQfinish(conn); + conn = nullptr; + } +} + +PgClientQueryResult PgClientConnector::executeQuery(const std::string& query) const { + PgClientQueryResult result; + if (!conn) { + result.success = false; + result.errorMessage = "Not connected to PostgreSQL"; + return result; + } + + std::lock_guard lk{mtx}; + PGresult* pgResult = PQexec(conn, query.c_str()); + if (!pgResult) { + result.success = false; + result.errorMessage = "Query execution returned null result"; + return result; + } + + ExecStatusType status = PQresultStatus(pgResult); + if (status != PGRES_TUPLES_OK) { + result.success = false; + result.errorMessage = PQresultErrorMessage(pgResult); + PQclear(pgResult); + return result; + } + + int numCols = PQnfields(pgResult); + int numRows = PQntuples(pgResult); + + // Get column names and types + result.columnNames.reserve(numCols); + result.columnTypes.reserve(numCols); + for (int i = 0; i < numCols; i++) { + result.columnNames.push_back(PQfname(pgResult, i)); + uint32_t oid = PQftype(pgResult, i); + result.columnTypes.push_back(pgOidToLogicalType(oid)); + } + + // Get rows with null tracking + result.rows.reserve(numRows); + for (int i = 0; i < numRows; i++) { + PgClientRow row; + row.cells.reserve(numCols); + for (int j = 0; j < numCols; j++) { + bool isNull = PQgetisnull(pgResult, i, j); + std::string val; + if (!isNull) { + val = PQgetvalue(pgResult, i, j); + } + row.cells.emplace_back(std::move(val), isNull); + } + result.rows.push_back(std::move(row)); + } + + result.numRows = numRows; + result.success = true; + + PQclear(pgResult); + return result; +} + +common::LogicalType pgOidToLogicalType(uint32_t oid) { + switch (oid) { + case 16: // bool + return common::LogicalType(common::LogicalTypeID::BOOL); + case 17: // bytea + return common::LogicalType(common::LogicalTypeID::BLOB); + case 20: // int8 + case 21: // int2 + case 23: // int4 + case 26: // oid + return common::LogicalType(common::LogicalTypeID::INT64); + case 700: // float4 + return common::LogicalType(common::LogicalTypeID::FLOAT); + case 701: // float8 + return common::LogicalType(common::LogicalTypeID::DOUBLE); + case 1043: // varchar + case 25: // text + case 1042: // char + case 142: // xml + return common::LogicalType(common::LogicalTypeID::STRING); + case 1082: // date + return common::LogicalType(common::LogicalTypeID::DATE); + case 1114: // timestamp + case 1184: // timestamptz + return common::LogicalType(common::LogicalTypeID::TIMESTAMP); + case 1186: // interval + return common::LogicalType(common::LogicalTypeID::INTERVAL); + case 2950: // uuid + return common::LogicalType(common::LogicalTypeID::UUID); + case 1700: // numeric + return common::LogicalType(common::LogicalTypeID::DOUBLE); + default: + // Default to string for unknown types + return common::LogicalType(common::LogicalTypeID::STRING); + } +} + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/function/CMakeLists.txt b/pg_client/src/function/CMakeLists.txt new file mode 100644 index 0000000..35b3e28 --- /dev/null +++ b/pg_client/src/function/CMakeLists.txt @@ -0,0 +1,6 @@ +add_library(lbug_pg_client_function OBJECT pg_client_scan.cpp) +add_dependencies(lbug_pg_client_function lbug_common) + +set(PG_CLIENT_EXTENSION_OBJECT_FILES + ${PG_CLIENT_EXTENSION_OBJECT_FILES} $ + PARENT_SCOPE) diff --git a/pg_client/src/function/pg_client_scan.cpp b/pg_client/src/function/pg_client_scan.cpp new file mode 100644 index 0000000..92b2468 --- /dev/null +++ b/pg_client/src/function/pg_client_scan.cpp @@ -0,0 +1,300 @@ +#include "function/pg_client_scan.h" + +#include +#include + +#include "binder/binder.h" +#include "common/constants.h" +#include "common/exception/runtime.h" +#include "common/system_config.h" +#include "common/string_utils.h" +#include "common/vector/value_vector.h" +#include "connector/pg_client_connector.h" +#include "function/table/bind_input.h" +#include "function/table/table_function.h" +#include "processor/execution_context.h" + +using namespace lbug::function; +using namespace lbug::common; + +namespace lbug { +namespace pg_client_extension { + +std::string PgClientScanBindData::getSQL() const { + // Build column selection + std::string columns; + bool first = true; + for (auto i = 0u; i < columnNamesInPg.size(); i++) { + if (!first) { + columns += ","; + } + columns += columnNamesInPg[i]; + first = false; + } + + std::string predicatesString = ""; + for (auto& predicates : getColumnPredicates()) { + if (predicates.isEmpty()) { + continue; + } + if (predicatesString.empty()) { + predicatesString = " WHERE " + predicates.toString(); + } else { + predicatesString += std::format(" AND {}", predicates.toString()); + } + } + + std::string q = query; + size_t pos = q.find("{}"); + if (pos != std::string::npos) { + q.replace(pos, 2, columns); + } + q += predicatesString; + q += getOrderBy(); + if (getLimitNum() != common::INVALID_ROW_IDX) { + q += std::format(" LIMIT {}", getLimitNum()); + } + return q; +} + +std::string PgClientScanBindData::getDescription() const { + return getSQL(); +} + +struct PgClientScanFunction { + static constexpr char PG_CLIENT_SCAN_FUNC_NAME[] = "pg_client_scan"; + + static offset_t tableFunc(const TableFuncInput& input, TableFuncOutput& output); + + static std::unique_ptr bindFunc( + std::shared_ptr bindData, main::ClientContext* /*context*/, + const TableFuncBindInput* input); + + static std::unique_ptr initSharedState( + const TableFuncInitSharedStateInput& input); + + static std::unique_ptr initLocalState( + const TableFuncInitLocalStateInput& input); +}; + +std::unique_ptr PgClientScanFunction::initSharedState( + const TableFuncInitSharedStateInput& input) { + auto* rawBindData = input.bindData; + if (!rawBindData) { + throw RuntimeException("PgClientScanFunction::initSharedState: bindData is null"); + } + auto scanBindData = rawBindData->constPtrCast(); + if (!scanBindData) { + throw RuntimeException("PgClientScanFunction::initSharedState: cast failed"); + } + auto sql = scanBindData->getSQL(); + if (sql.empty()) { + throw RuntimeException("PgClientScanFunction::initSharedState: SQL is empty"); + } + auto result = scanBindData->connector.executeQuery(sql); + if (!result.success) { + throw RuntimeException( + std::format("Failed to execute query due to error: {}", result.errorMessage)); + } + return std::make_unique(std::move(result)); +} + +std::unique_ptr PgClientScanFunction::initLocalState( + const TableFuncInitLocalStateInput&) { + return std::make_unique(); +} + +static void setValueFromCell(ValueVector* vector, uint32_t pos, + const PgClientCell& cell, const LogicalType& type) { + if (cell.isNull) { + vector->setNull(pos, true); + return; + } + + switch (type.getLogicalTypeID()) { + case LogicalTypeID::BOOL: { + auto upper = cell.value; + StringUtils::toUpper(upper); + bool val = (upper == "TRUE" || upper == "T" || upper == "YES" || upper == "1"); + vector->setValue(pos, val ? 1 : 0); + break; + } + case LogicalTypeID::INT64: { + int64_t val = std::stoll(cell.value); + vector->setValue(pos, val); + break; + } + case LogicalTypeID::INT32: { + int32_t val = std::stoi(cell.value); + vector->setValue(pos, val); + break; + } + case LogicalTypeID::INT16: { + int16_t val = static_cast(std::stoi(cell.value)); + vector->setValue(pos, val); + break; + } + case LogicalTypeID::FLOAT: { + float val = std::stof(cell.value); + vector->setValue(pos, val); + break; + } + case LogicalTypeID::DOUBLE: { + double val = std::stod(cell.value); + vector->setValue(pos, val); + break; + } + case LogicalTypeID::STRING: + case LogicalTypeID::BLOB: + case LogicalTypeID::UUID: { + StringVector::addString(vector, pos, cell.value); + break; + } + case LogicalTypeID::DATE: { + // Parse date string (YYYY-MM-DD) to days since epoch + // For now, use string representation which Ladybug can parse + StringVector::addString(vector, pos, cell.value); + break; + } + case LogicalTypeID::TIMESTAMP: + case LogicalTypeID::TIMESTAMP_MS: + case LogicalTypeID::TIMESTAMP_NS: + case LogicalTypeID::TIMESTAMP_SEC: + case LogicalTypeID::TIMESTAMP_TZ: { + StringVector::addString(vector, pos, cell.value); + break; + } + case LogicalTypeID::INTERVAL: { + StringVector::addString(vector, pos, cell.value); + break; + } + default: { + StringVector::addString(vector, pos, cell.value); + break; + } + } +} + +static void buildColumnNameToIndexMap(const std::vector& columnNames, + std::unordered_map& nameToIdx) { + nameToIdx.clear(); + for (int i = 0; i < (int)columnNames.size(); i++) { + nameToIdx[columnNames[i]] = i; + } +} + +offset_t PgClientScanFunction::tableFunc(const TableFuncInput& input, TableFuncOutput& output) { + if (!input.sharedState) { + throw RuntimeException("tableFunc: sharedState is null"); + } + if (!input.bindData) { + throw RuntimeException("tableFunc: bindData is null"); + } + auto pgClientScanSharedState = input.sharedState->ptrCast(); + auto pgClientScanBindData = input.bindData->constPtrCast(); + auto& queryResult = pgClientScanSharedState->queryResult; + + // Atomically claim a morsel (slice of [startOffset, endOffset)) from the + // shared query result. The shared state is owned by the ForeignRelTable + // and shared across worker threads, so we must serialize offset updates + // — otherwise parallel scans (e.g. the parallel hash join that drives + // MATCH ... -[...]-> ... queries) would hand out overlapping rows and + // miss others. TableFuncSharedState provides the mutex for exactly this + // purpose. + uint64_t startOffset = 0; + uint64_t batchSize = 0; + { + std::lock_guard lk{pgClientScanSharedState->mtx}; + if (pgClientScanSharedState->currentOffset >= queryResult.rows.size()) { + return 0; + } + uint64_t remainingRows = queryResult.rows.size() - pgClientScanSharedState->currentOffset; + batchSize = std::min(remainingRows, DEFAULT_VECTOR_CAPACITY); + startOffset = pgClientScanSharedState->currentOffset; + pgClientScanSharedState->currentOffset += batchSize; + } + + auto& dataChunk = output.dataChunk; + auto numColumns = dataChunk.getNumValueVectors(); + auto& selVector = dataChunk.state->getSelVectorUnsafe(); + selVector.setSelSize(batchSize); + + // The output may have more columns than PG result columns: + // Column 0 = internal ID (synthesized), columns 1+ = actual data + // This happens when a scan involves a nodeUniqueName (e.g. MATCH queries). + // `columnNamesInPg` only contains actual PG column names (no rowid). + bool hasInternalId = numColumns > pgClientScanBindData->columnNamesInPg.size(); + + // Build a name-to-index map for the result columns + std::unordered_map resultColMap; + buildColumnNameToIndexMap(queryResult.columnNames, resultColMap); + + for (auto colIdx = 0u; colIdx < numColumns; colIdx++) { + auto& vector = dataChunk.getValueVectorMutable(colIdx); + + if (hasInternalId && colIdx == 0) { + // Internal ID column - synthesize from row index + for (auto rowIdx = 0u; rowIdx < batchSize; rowIdx++) { + vector.setValue(rowIdx, startOffset + rowIdx); + } + continue; + } + + // Determine which PG result column this output column maps to + int pgColIdx = hasInternalId ? (int)colIdx - 1 : (int)colIdx; + if (pgColIdx < 0 || pgColIdx >= (int)pgClientScanBindData->columnNamesInPg.size()) { + for (auto rowIdx = 0u; rowIdx < batchSize; rowIdx++) { + vector.setNull(rowIdx, true); + } + continue; + } + + std::string targetColName = pgClientScanBindData->columnNamesInPg[pgColIdx]; + + // Find this column in the result + auto it = resultColMap.find(targetColName); + if (it == resultColMap.end() || + it->second >= (int)queryResult.columnNames.size()) { + for (auto rowIdx = 0u; rowIdx < batchSize; rowIdx++) { + vector.setNull(rowIdx, true); + } + continue; + } + + int resultColIdx = it->second; + + // Copy data from result to vector, respecting the current offset + for (auto rowIdx = 0u; rowIdx < batchSize; rowIdx++) { + uint64_t srcRowIdx = startOffset + rowIdx; + setValueFromCell(&vector, rowIdx, + queryResult.rows[srcRowIdx].cells[resultColIdx], vector.dataType); + } + } + + return batchSize; +} + +std::unique_ptr PgClientScanFunction::bindFunc( + std::shared_ptr scanInfo, main::ClientContext* context, + const TableFuncBindInput* input) { + auto columnNames = + TableFunction::extractYieldVariables(scanInfo->getColumnNames(), input->yieldVariables); + auto columns = input->binder->createVariables(columnNames, scanInfo->getColumnTypes(*context)); + return std::make_unique(scanInfo->getTemplateQuery(*context), + scanInfo->getColumnNames(), scanInfo->getConnector(), std::move(columns)); +} + +TableFunction getScanFunction(std::shared_ptr scanInfo) { + auto function = + TableFunction(PgClientScanFunction::PG_CLIENT_SCAN_FUNC_NAME, std::vector{}); + function.tableFunc = PgClientScanFunction::tableFunc; + function.bindFunc = std::bind(PgClientScanFunction::bindFunc, scanInfo, + std::placeholders::_1, std::placeholders::_2); + function.initSharedStateFunc = PgClientScanFunction::initSharedState; + function.initLocalStateFunc = PgClientScanFunction::initLocalState; + function.supportsPushDownFunc = [] { return true; }; + return function; +} + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/include/catalog/pg_client_catalog.h b/pg_client/src/include/catalog/pg_client_catalog.h new file mode 100644 index 0000000..3160ccc --- /dev/null +++ b/pg_client/src/include/catalog/pg_client_catalog.h @@ -0,0 +1,56 @@ +#pragma once + +#include "catalog/catalog_entry/node_table_catalog_entry.h" +#include "catalog/catalog_entry/rel_group_catalog_entry.h" +#include "common/vector/value_vector.h" +#include "extension/catalog_extension.h" +#include +#include + +namespace lbug { +namespace binder { +struct AttachOption; +} // namespace binder + +namespace pg_client_extension { + +class PgClientConnector; + +struct PgClientColumnInfo { + std::string name; + common::LogicalType type; +}; + +common::LogicalType pgTypeNameToLogicalType(const std::string& pgType); + +class PgClientCatalog : public extension::CatalogExtension { +public: + PgClientCatalog(std::string connStr, std::string catalogName, + std::string defaultSchemaName, main::ClientContext* context, + const PgClientConnector& connector); + + void init() override; + + static std::string bindSchemaName(const binder::AttachOption& options, + const std::string& defaultName); + +private: + void createForeignNodeTable(const std::string& tableName); + void createForeignRelTable(const std::string& tableName); + + std::vector getTableColumnInfo( + const std::string& tableName) const; + std::vector getColumnNames( + const std::vector& columns) const; + std::vector getColumnTypes( + const std::vector& columns) const; + + std::string connStr; + std::string catalogName; + std::string defaultSchemaName; + const PgClientConnector& connector; + main::ClientContext* context_; +}; + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/include/catalog/pg_client_table_catalog_entry.h b/pg_client/src/include/catalog/pg_client_table_catalog_entry.h new file mode 100644 index 0000000..f0d2bbf --- /dev/null +++ b/pg_client/src/include/catalog/pg_client_table_catalog_entry.h @@ -0,0 +1,38 @@ +#pragma once + +#include +#include + +#include "catalog/catalog_entry/table_catalog_entry.h" +#include "function/pg_client_scan.h" + +namespace lbug { +namespace catalog { + +class PgClientTableCatalogEntry final : public TableCatalogEntry { +public: + // constructors + PgClientTableCatalogEntry(std::string name, + std::optional scanFunction, + std::shared_ptr scanInfo); + + // getter & setter + common::TableType getTableType() const override; + std::optional getScanFunction() const override { return scanFunction; } + std::unique_ptr getBoundScanInfo(main::ClientContext* context, + const std::string& nodeUniqueName = "") override; + + // serialization & deserialization + std::unique_ptr copy() const override; + +private: + std::unique_ptr getBoundExtraCreateInfo( + transaction::Transaction* transaction) const override; + +private: + std::optional scanFunction; + std::shared_ptr scanInfo; +}; + +} // namespace catalog +} // namespace lbug diff --git a/pg_client/src/include/connector/pg_client_connector.h b/pg_client/src/include/connector/pg_client_connector.h new file mode 100644 index 0000000..720b1a9 --- /dev/null +++ b/pg_client/src/include/connector/pg_client_connector.h @@ -0,0 +1,60 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "common/types/types.h" + +namespace lbug { +namespace pg_client_extension { + +struct PgClientCell { + std::string value; + bool isNull; + + PgClientCell() : isNull{true} {} + PgClientCell(std::string value, bool isNull) + : value{std::move(value)}, isNull{isNull} {} +}; + +struct PgClientRow { + std::vector cells; +}; + +struct PgClientQueryResult { + std::vector columnNames; + std::vector columnTypes; + std::vector rows; + uint64_t numRows; + bool success; + std::string errorMessage; + + PgClientQueryResult() : numRows{0}, success{false} {} +}; + +class PgClientConnector { +public: + PgClientConnector(); + ~PgClientConnector(); + + void connect(const std::string& connStr); + void disconnect(); + + bool isConnected() const { return conn != nullptr; } + + PgClientQueryResult executeQuery(const std::string& query) const; + + PGconn* getPGconn() const { return conn; } + +private: + PGconn* conn; + mutable std::mutex mtx; +}; + +common::LogicalType pgOidToLogicalType(uint32_t oid); + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/include/function/pg_client_scan.h b/pg_client/src/include/function/pg_client_scan.h new file mode 100644 index 0000000..09c9782 --- /dev/null +++ b/pg_client/src/include/function/pg_client_scan.h @@ -0,0 +1,97 @@ +#pragma once + +#include +#include +#include + +#include "common/copy_constructors.h" +#include "common/types/types.h" +#include "connector/pg_client_connector.h" +#include "function/table/bind_data.h" +#include "function/table/table_function.h" + +namespace lbug { +namespace pg_client_extension { + +class PgClientConnector; + +class PgClientTableScanInfo { +public: + PgClientTableScanInfo(std::string templateQuery, + std::vector columnTypes, + std::vector columnNames, const PgClientConnector& connector) + : templateQuery{std::move(templateQuery)}, columnTypes{std::move(columnTypes)}, + columnNames{std::move(columnNames)}, connector{connector} {} + + PgClientTableScanInfo(const PgClientTableScanInfo& other) = default; + + virtual ~PgClientTableScanInfo() = default; + + virtual std::string getTemplateQuery(const main::ClientContext& /*context*/) const { + return templateQuery; + } + + virtual std::vector getColumnTypes( + const main::ClientContext& /*context*/) const { + std::vector result; + result.reserve(columnTypes.size()); + for (auto& t : columnTypes) { + result.push_back(t.copy()); + } + return result; + } + + std::vector getColumnNames() const { return columnNames; } + + const PgClientConnector& getConnector() const { return connector; } + +private: + std::string templateQuery; + std::vector columnTypes; + std::vector columnNames; + const PgClientConnector& connector; +}; + +struct PgClientScanBindData : function::TableFuncBindData { + std::string query; + std::vector columnNamesInPg; + const PgClientConnector& connector; + + PgClientScanBindData(std::string query, std::vector columnNamesInPg, + const PgClientConnector& connector, binder::expression_vector columns) + : function::TableFuncBindData{std::move(columns), 0 /* numRows */}, + query{std::move(query)}, columnNamesInPg{std::move(columnNamesInPg)}, + connector{connector} {} + + PgClientScanBindData(const PgClientScanBindData& other) + : function::TableFuncBindData{other}, query{other.query}, + columnNamesInPg{other.columnNamesInPg}, connector{other.connector} {} + + std::string getSQL() const; + std::string getDescription() const override; + + std::unique_ptr copy() const override { + return std::make_unique(*this); + } + + std::unique_ptr copyWithQuery(const std::string& newQuery, + binder::expression_vector columns, + const std::vector& columnNamesInResult) const override { + return std::make_unique(newQuery, columnNamesInResult, connector, + std::move(columns)); + } +}; + +struct PgClientScanSharedState final : function::TableFuncSharedState { + PgClientQueryResult queryResult; + uint64_t currentOffset; + + explicit PgClientScanSharedState(PgClientQueryResult queryResult) + : TableFuncSharedState{queryResult.numRows}, + queryResult{std::move(queryResult)}, currentOffset{0} {} +}; + +function::TableFunction getScanFunction(std::shared_ptr scanInfo); + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/include/main/pg_client_extension.h b/pg_client/src/include/main/pg_client_extension.h new file mode 100644 index 0000000..467bd16 --- /dev/null +++ b/pg_client/src/include/main/pg_client_extension.h @@ -0,0 +1,17 @@ +#pragma once + +#include "extension/extension.h" + +namespace lbug { +namespace pg_client_extension { + +class PgClientExtension final : public extension::Extension { +public: + static constexpr char EXTENSION_NAME[] = "PG_CLIENT"; + +public: + static void load(main::ClientContext* context); +}; + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/include/storage/attached_pg_client_database.h b/pg_client/src/include/storage/attached_pg_client_database.h new file mode 100644 index 0000000..38f2784 --- /dev/null +++ b/pg_client/src/include/storage/attached_pg_client_database.h @@ -0,0 +1,24 @@ +#pragma once + +#include "connector/pg_client_connector.h" +#include "main/attached_database.h" + +namespace lbug { +namespace pg_client_extension { + +class AttachedPgClientDatabase final : public main::AttachedDatabase { +public: + AttachedPgClientDatabase(std::string dbName, std::string dbType, + std::unique_ptr catalog, + std::unique_ptr connector) + : main::AttachedDatabase{std::move(dbName), std::move(dbType), std::move(catalog)}, + connector_{std::move(connector)} {} + + const PgClientConnector& getConnector() const { return *connector_; } + +private: + std::unique_ptr connector_; +}; + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/include/storage/pg_client_storage.h b/pg_client/src/include/storage/pg_client_storage.h new file mode 100644 index 0000000..f2eed8d --- /dev/null +++ b/pg_client/src/include/storage/pg_client_storage.h @@ -0,0 +1,26 @@ +#pragma once + +#include "storage/storage_extension.h" + +namespace lbug { +namespace main { +class Database; +} // namespace main + +namespace pg_client_extension { + +class PgClientStorageExtension final : public storage::StorageExtension { +public: + static constexpr const char* DB_TYPE = "PG_CLIENT"; + + static constexpr const char* DEFAULT_SCHEMA_NAME = "public"; + + static constexpr const char* SCHEMA_OPTION = "SCHEMA"; + + explicit PgClientStorageExtension(main::Database& database); + + bool canHandleDB(std::string dbType_) const override; +}; + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/src/main/CMakeLists.txt b/pg_client/src/main/CMakeLists.txt new file mode 100644 index 0000000..ea4d1aa --- /dev/null +++ b/pg_client/src/main/CMakeLists.txt @@ -0,0 +1,9 @@ +add_library(lbug_pg_client_main + OBJECT + pg_client_extension.cpp) + +add_dependencies(lbug_pg_client_main lbug_common) + +set(PG_CLIENT_EXTENSION_OBJECT_FILES + ${PG_CLIENT_EXTENSION_OBJECT_FILES} $ + PARENT_SCOPE) diff --git a/pg_client/src/main/pg_client_extension.cpp b/pg_client/src/main/pg_client_extension.cpp new file mode 100644 index 0000000..a2b749e --- /dev/null +++ b/pg_client/src/main/pg_client_extension.cpp @@ -0,0 +1,38 @@ +#include "main/pg_client_extension.h" + +#include "main/client_context.h" +#include "main/database.h" +#include "storage/pg_client_storage.h" + +namespace lbug { +namespace pg_client_extension { + +using namespace extension; + +void PgClientExtension::load(main::ClientContext* context) { + auto db = context->getDatabase(); + db->registerStorageExtension(EXTENSION_NAME, + std::make_unique(*db)); +} + +} // namespace pg_client_extension +} // namespace lbug + +#if defined(BUILD_DYNAMIC_LOAD) +extern "C" { +// Because we link against the static library on windows, we implicitly inherit LBUG_STATIC_DEFINE, +// which cancels out any exporting, so we can't use LBUG_API. +#if defined(_WIN32) +#define INIT_EXPORT __declspec(dllexport) +#else +#define INIT_EXPORT __attribute__((visibility("default"))) +#endif +INIT_EXPORT void init(lbug::main::ClientContext* context) { + lbug::pg_client_extension::PgClientExtension::load(context); +} + +INIT_EXPORT const char* name() { + return lbug::pg_client_extension::PgClientExtension::EXTENSION_NAME; +} +} +#endif diff --git a/pg_client/src/storage/CMakeLists.txt b/pg_client/src/storage/CMakeLists.txt new file mode 100644 index 0000000..f77390d --- /dev/null +++ b/pg_client/src/storage/CMakeLists.txt @@ -0,0 +1,9 @@ +add_library(lbug_pg_client_storage + OBJECT + pg_client_storage.cpp) + +add_dependencies(lbug_pg_client_storage lbug_common) + +set(PG_CLIENT_EXTENSION_OBJECT_FILES + ${PG_CLIENT_EXTENSION_OBJECT_FILES} $ + PARENT_SCOPE) diff --git a/pg_client/src/storage/pg_client_storage.cpp b/pg_client/src/storage/pg_client_storage.cpp new file mode 100644 index 0000000..948a19f --- /dev/null +++ b/pg_client/src/storage/pg_client_storage.cpp @@ -0,0 +1,61 @@ +#include "storage/pg_client_storage.h" + +#include + +#include "common/exception/runtime.h" +#include "common/string_utils.h" +#include "connector/pg_client_connector.h" +#include "extension/extension.h" +#include "catalog/pg_client_catalog.h" +#include "function/pg_client_scan.h" +#include "storage/attached_pg_client_database.h" + +namespace lbug { +namespace pg_client_extension { + +std::string extractDBName(const std::string& connectionInfo) { + std::regex pattern("dbname=([^ ]+)"); + std::smatch match; + if (std::regex_search(connectionInfo, match, pattern)) { + return match.str(1); + } + // Try PostgreSQL URI format + std::regex uriPattern("postgres(?:ql)?://[^/]+/([^?]+)"); + if (std::regex_search(connectionInfo, match, uriPattern)) { + return match.str(1); + } + throw common::RuntimeException{"Invalid PostgreSQL connection string."}; +} + +std::unique_ptr attachPgClient(std::string dbName, std::string dbPath, + main::ClientContext* clientContext, const binder::AttachOption& attachOption) { + auto catalogName = extractDBName(dbPath); + if (dbName == "") { + dbName = catalogName; + } + + // Connect to PostgreSQL via libpq + auto connector = std::make_unique(); + connector->connect(dbPath); + + // Create catalog that discovers tables from PG information_schema + auto schemaName = PgClientCatalog::bindSchemaName(attachOption, + PgClientStorageExtension::DEFAULT_SCHEMA_NAME); + auto catalog = std::make_unique(dbPath, catalogName, + schemaName, clientContext, *connector); + catalog->init(); + + return std::make_unique(dbName, + PgClientStorageExtension::DB_TYPE, std::move(catalog), std::move(connector)); +} + +PgClientStorageExtension::PgClientStorageExtension(main::Database& /*database*/) + : StorageExtension{attachPgClient} {} + +bool PgClientStorageExtension::canHandleDB(std::string dbType_) const { + common::StringUtils::toUpper(dbType_); + return dbType_ == DB_TYPE; +} + +} // namespace pg_client_extension +} // namespace lbug diff --git a/pg_client/test/test_files/create_pg_client_test_db.sql b/pg_client/test/test_files/create_pg_client_test_db.sql new file mode 100644 index 0000000..3d5bd43 --- /dev/null +++ b/pg_client/test/test_files/create_pg_client_test_db.sql @@ -0,0 +1,70 @@ +-- +-- PostgreSQL database dump for pg_client extension tests. +-- +-- Creates the expected schema and test data matching +-- extension/pg_client/test/test_files/pg_client.test +-- + +SET statement_timeout = 0; +SET lock_timeout = 0; +SET idle_in_transaction_session_timeout = 0; +SET client_encoding = 'UTF8'; +SET standard_conforming_strings = on; +SELECT pg_catalog.set_config('search_path', '', false); +SET check_function_bodies = false; +SET xmloption = content; +SET client_min_messages = warning; +SET row_security = off; + +-- +-- Table: node_person +-- + +CREATE TABLE public.node_person ( + id integer NOT NULL, + name character varying(100) NOT NULL, + age integer NOT NULL +); + + +ALTER TABLE public.node_person OWNER TO ci; + +-- +-- Table: rel_knows +-- + +CREATE TABLE public.rel_knows ( + id integer NOT NULL, + from_id integer NOT NULL, + to_id integer NOT NULL, + since date NOT NULL +); + + +ALTER TABLE public.rel_knows OWNER TO ci; + +-- +-- Data for node_person +-- + +INSERT INTO public.node_person (id, name, age) VALUES + (1, 'Alice', 30), + (2, 'Bob', 25), + (3, 'Charlie', 35), + (4, 'Diana', 28), + (5, 'Eve', 32); + +SELECT pg_catalog.setval(pg_catalog.pg_get_serial_sequence('public.node_person', 'id'), 5, true); + +-- +-- Data for rel_knows +-- + +INSERT INTO public.rel_knows (id, from_id, to_id, since) VALUES + (1, 1, 2, '2020-01-15'), + (2, 1, 3, '2021-03-20'), + (3, 2, 4, '2022-06-10'), + (4, 3, 4, '2023-08-05'), + (5, 4, 5, '2023-12-01'); + +SELECT pg_catalog.setval(pg_catalog.pg_get_serial_sequence('public.rel_knows', 'id'), 5, true); diff --git a/pg_client/test/test_files/pg_client.test b/pg_client/test/test_files/pg_client.test new file mode 100644 index 0000000..c4d63ea --- /dev/null +++ b/pg_client/test/test_files/pg_client.test @@ -0,0 +1,94 @@ +-DATASET CSV empty +-SKIP_MUSL +-SKIP_WASM +-SKIP_IN_MEM + +-- + +-CASE LoadPgClientExtension +-SKIP_FSM_LEAK_CHECK +-LOAD_DYNAMIC_EXTENSION pg_client +-STATEMENT RETURN 'LOAD OK'; +---- 1 +LOAD OK + +-CASE AttachPgClientDatabase +-SKIP_FSM_LEAK_CHECK +-LOAD_DYNAMIC_EXTENSION pg_client +-STATEMENT ATTACH '${PG_CLIENT_CONNECTION_STRING}' as testdb (dbtype PG_CLIENT); +---- 1 +Attached database successfully. + +-CASE ScanNodeTable +-SKIP_FSM_LEAK_CHECK +-LOAD_DYNAMIC_EXTENSION pg_client +-STATEMENT ATTACH '${PG_CLIENT_CONNECTION_STRING}' as testdb (dbtype PG_CLIENT); +---- 1 +Attached database successfully. +-STATEMENT LOAD FROM testdb.node_person RETURN name, age ORDER BY id; +---- 5 +Alice|30 +Bob|25 +Charlie|35 +Diana|28 +Eve|32 + +-CASE ScanNodeTableWithFilter +-SKIP_FSM_LEAK_CHECK +-LOAD_DYNAMIC_EXTENSION pg_client +-STATEMENT ATTACH '${PG_CLIENT_CONNECTION_STRING}' as testdb (dbtype PG_CLIENT); +---- 1 +Attached database successfully. +-STATEMENT LOAD FROM testdb.node_person WHERE age > 30 RETURN name, age; +---- 2 +Charlie|35 +Eve|32 + +-CASE ScanRelTable +-SKIP_FSM_LEAK_CHECK +-LOAD_DYNAMIC_EXTENSION pg_client +-STATEMENT ATTACH '${PG_CLIENT_CONNECTION_STRING}' as testdb (dbtype PG_CLIENT); +---- 1 +Attached database successfully. +-STATEMENT LOAD FROM testdb.rel_knows RETURN from_id, to_id, since ORDER BY id; +---- 5 +1|2|2020-01-15 +1|3|2021-03-20 +2|4|2022-06-10 +3|4|2023-08-05 +4|5|2023-12-01 + +-CASE MatchOnNodeTable +-SKIP_FSM_LEAK_CHECK +-LOAD_DYNAMIC_EXTENSION pg_client +-STATEMENT ATTACH '${PG_CLIENT_CONNECTION_STRING}' as testdb (dbtype PG_CLIENT); +---- 1 +Attached database successfully. +-STATEMENT MATCH (p:testdb.node_person) RETURN p.name, p.age ORDER BY p.age; +---- 5 +Bob|25 +Diana|28 +Alice|30 +Eve|32 +Charlie|35 + +-CASE MatchOnRelTable +-SKIP_FSM_LEAK_CHECK +-LOAD_DYNAMIC_EXTENSION pg_client +-STATEMENT ATTACH '${PG_CLIENT_CONNECTION_STRING}' as testdb (dbtype PG_CLIENT); +---- 1 +Attached database successfully. +-STATEMENT MATCH (a:testdb.node_person)-[k:rel_knows]->(b:testdb.node_person) RETURN count(*); +---- 1 +5 + +-CASE ShowTables +-SKIP_FSM_LEAK_CHECK +-LOAD_DYNAMIC_EXTENSION pg_client +-STATEMENT ATTACH '${PG_CLIENT_CONNECTION_STRING}' as testdb (dbtype PG_CLIENT); +---- 1 +Attached database successfully. +-STATEMENT CALL SHOW_TABLES() RETURN *; +---- 2 +0|node_person|ATTACHED|testdb(PG_CLIENT)| +1|rel_knows|ATTACHED|testdb(PG_CLIENT)| diff --git a/pg_client/test/test_pg_client.py b/pg_client/test/test_pg_client.py new file mode 100644 index 0000000..c16bc76 --- /dev/null +++ b/pg_client/test/test_pg_client.py @@ -0,0 +1,343 @@ +#!/usr/bin/env python3 +""" +Test script for pg_client extension using pgembed to spin up a temporary PostgreSQL instance. + +Tests both LOAD FROM scans and in-place Cypher queries (MATCH). +""" + +import os +import sys +import tempfile +import subprocess +import unittest + +import pgembed +import sqlalchemy as sa +from sqlalchemy_utils import database_exists, create_database + + +PG_CLIENT_EXT = os.path.join( + os.path.dirname(__file__), "..", "build", "libpg_client.lbug_extension" +) +LBUG_BIN = os.path.join( + os.path.dirname(__file__), "..", "..", "..", "build", "release", "tools", "shell", "lbug" +) + + +def run_lbug(script: str, lbug_bin: str = LBUG_BIN) -> str: + """Run a ladybug script and return the output.""" + result = subprocess.run( + [lbug_bin], + input=script, + capture_output=True, + text=True, + timeout=30, + ) + return result.stdout, result.stderr + + +class TestPgClientExtension(unittest.TestCase): + """Test the pg_client extension using a temporary PostgreSQL instance.""" + + @classmethod + def setUpClass(cls): + """Start a temporary PostgreSQL instance and prepare test data.""" + cls.tmpdir = tempfile.mkdtemp() + cls.pg = pgembed.get_server(cls.tmpdir) + cls.database_name = "testdb" + cls.uri = cls.pg.get_uri(cls.database_name) + + if not database_exists(cls.uri): + create_database(cls.uri) + + engine = sa.create_engine(cls.uri, isolation_level="AUTOCOMMIT") + conn = engine.connect() + + with conn.begin(): + # Create node_person table (detected as node table by prefix) + conn.execute(sa.text(""" + CREATE TABLE node_person ( + id SERIAL PRIMARY KEY, + name VARCHAR(100) NOT NULL, + age INTEGER, + email VARCHAR(100) + ) + """)) + + conn.execute(sa.text(""" + INSERT INTO node_person (name, age, email) VALUES + ('Alice', 30, 'alice@example.com'), + ('Bob', 25, 'bob@example.com'), + ('Charlie', 35, 'charlie@example.com'), + ('Diana', 28, 'diana@example.com'), + ('Eve', 32, 'eve@example.com') + """)) + + # Create node_company table (detected as node table by prefix) + conn.execute(sa.text(""" + CREATE TABLE node_company ( + id SERIAL PRIMARY KEY, + name VARCHAR(100), + revenue DOUBLE PRECISION + ) + """)) + + conn.execute(sa.text(""" + INSERT INTO node_company (name, revenue) VALUES + ('ACME Corp', 1000000.50), + ('Globex Inc', 2500000.75), + ('Initech', 500000.00) + """)) + + # Create rel_knows table with FK columns named src_id/dst_id + # Prefix rel_ + FK constraints → auto-register as relationship table + conn.execute(sa.text(""" + CREATE TABLE rel_knows ( + id SERIAL PRIMARY KEY, + src_id INTEGER NOT NULL, + dst_id INTEGER NOT NULL, + since VARCHAR(20) + ) + """)) + + # Add actual FOREIGN KEY constraints so the FK query detects src/dst tables + conn.execute(sa.text(""" + ALTER TABLE rel_knows + ADD CONSTRAINT fk_src FOREIGN KEY (src_id) REFERENCES node_person(id), + ADD CONSTRAINT fk_dst FOREIGN KEY (dst_id) REFERENCES node_person(id) + """)) + + conn.execute(sa.text(""" + INSERT INTO rel_knows (src_id, dst_id, since) VALUES + (1, 2, '2020-01-15'), + (1, 3, '2021-03-20'), + (2, 4, '2022-06-10'), + (3, 4, '2023-08-05'), + (4, 5, '2023-12-01') + """)) + + # Create rel_works_at table with FK constraints + conn.execute(sa.text(""" + CREATE TABLE rel_works_at ( + id SERIAL PRIMARY KEY, + src_id INTEGER NOT NULL REFERENCES node_person(id), + dst_id INTEGER NOT NULL REFERENCES node_company(id) + ) + """)) + + conn.execute(sa.text(""" + INSERT INTO rel_works_at (src_id, dst_id) VALUES + (1, 1), + (2, 2), + (3, 1), + (4, 3), + (5, 2) + """)) + + conn.close() + cls.conn_str = cls.uri + + def test_01_load_extension(self): + """Test that the pg_client extension can be loaded.""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + RETURN "LOAD OK"; + """ + stdout, stderr = run_lbug(script) + self.assertIn("LOAD OK", stdout, f"Extension load failed: {stderr}") + + def test_02_attach_database(self): + """Test ATTACH to PostgreSQL via pg_client.""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + RETURN "ATTACH OK"; + """ + stdout, stderr = run_lbug(script) + self.assertIn("ATTACH OK", stdout, f"ATTACH failed: {stderr}") + + def test_03_scan_node_table(self): + """Test LOAD FROM a node_* table.""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + LOAD FROM testdb.node_person RETURN name, age ORDER BY id; + """ + stdout, stderr = run_lbug(script) + self.assertIn("Alice", stdout, f"Scan failed: {stderr}") + self.assertIn("Bob", stdout, f"Bob not found: {stderr}") + self.assertIn("Charlie", stdout, f"Charlie not found: {stderr}") + self.assertIn("Diana", stdout, f"Diana not found: {stderr}") + self.assertIn("Eve", stdout, f"Eve not found: {stderr}") + + def test_04_scan_node_table_with_filter(self): + """Test LOAD FROM with filter pushdown.""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + LOAD FROM testdb.node_person WHERE age > 30 RETURN name, age; + """ + stdout, stderr = run_lbug(script) + self.assertIn("Charlie", stdout, f"Filter scan failed: {stderr}") + self.assertIn("Eve", stdout, f"Eve not found: {stderr}") + self.assertNotIn("Alice", stdout, "Alice should be filtered out") + + def test_05_match_node_table(self): + """Test MATCH query on node table (in-place).""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + MATCH (p:testdb.node_person) RETURN p.name, p.age ORDER BY p.age; + """ + stdout, stderr = run_lbug(script) + self.assertIn("Bob", stdout, f"MATCH failed: {stderr}") + self.assertIn("Alice", stdout, f"MATCH missing Alice: {stderr}") + self.assertIn("Charlie", stdout, f"MATCH missing Charlie: {stderr}") + + def test_06_load_rel_table(self): + """Test LOAD FROM a rel_* table.""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + LOAD FROM testdb.rel_knows RETURN src_id, dst_id, since ORDER BY id; + """ + stdout, stderr = run_lbug(script) + self.assertIn("2020-01-15", stdout, f"Rel load failed: {stderr}") + self.assertIn("2021-03-20", stdout, f"Rel load missing: {stderr}") + self.assertIn("2023-12-01", stdout, f"Rel load missing: {stderr}") + + def test_06b_match_rel_table(self): + """Test MATCH query traversing a rel_* table (the documented pattern). + + Regression test: registering rel_* tables as discoverable node shadows + used to break MATCH ... -[...]-> ... queries because the planner + couldn't find a RelGroupCatalogEntry to anchor the scan. The rel + table must now be registered as a RelGroupCatalogEntry and the + ForeignRelTable must own a shared scan state to drive morsel-based + parallelism across HashJoinBuild workers. + + Note: the count(*) form is the tightest correctness check for + morsel-driven parallel scan — the rel table scan must visit every + row exactly once across all worker threads. The projection form + (RETURN a.name, b.name, k.since) additionally requires an ID + mapping layer to translate PG foreign-key values into lbug + internal node IDs; that is tracked separately. + """ + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + MATCH (a:node_person)-[k:rel_knows]->(b:node_person) + RETURN count(*); + """ + stdout, stderr = run_lbug(script) + self.assertIn("5", stdout, f"MATCH rel count failed: {stderr}") + + def test_06c_match_rel_count_parallel(self): + """Test count(*) over a MATCH ... -[...]-> ... join on a rel_* table. + + This is the tightest test for morsel-driven parallel correctness: the + ForeignRelTable scan is driven by HashJoinBuild workers, so every + joined row must be visited exactly once. A bug in the shared + PgClientScanSharedState (e.g. per-thread state instead of per-table) + would surface as either a wrong count or duplicated rows. + """ + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + MATCH (a:node_person)-[k:rel_knows]->(b:node_person) + RETURN count(*); + """ + stdout, stderr = run_lbug(script) + self.assertIn("5", stdout, f"MATCH rel count failed: {stderr}") + + def test_07_scan_rel_table_filter(self): + """Test LOAD FROM a rel_* table with filter.""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + LOAD FROM testdb.rel_knows WHERE src_id = 1 RETURN dst_id, since ORDER BY dst_id; + """ + stdout, stderr = run_lbug(script) + self.assertIn("2020-01-15", stdout, f"Rel filter failed: {stderr}") + self.assertIn("2021-03-20", stdout, f"Rel filter missing: {stderr}") + + def test_09_show_tables(self): + """Test SHOW_TABLES includes node and rel tables.""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + CALL SHOW_TABLES() RETURN *; + """ + stdout, stderr = run_lbug(script) + self.assertIn("node_person", stdout, f"SHOW_TABLES missing node_person: {stderr}") + self.assertIn("rel_knows", stdout, f"SHOW_TABLES missing rel_knows: {stderr}") + + def test_10_table_info(self): + """Test TABLE_INFO on attached table.""" + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + CALL TABLE_INFO('testdb.node_person') RETURN *; + """ + stdout, stderr = run_lbug(script) + self.assertIn("name", stdout.lower(), f"TABLE_INFO missing name: {stderr}") + self.assertIn("age", stdout.lower(), f"TABLE_INFO missing age: {stderr}") + + def test_11_attach_with_schema(self): + """Test ATTACH with non-default schema.""" + engine = sa.create_engine(self.uri, isolation_level="AUTOCOMMIT") + conn = engine.connect() + with conn.begin(): + conn.execute(sa.text("CREATE SCHEMA IF NOT EXISTS custom_schema")) + conn.execute(sa.text(""" + CREATE TABLE custom_schema.node_product ( + id SERIAL PRIMARY KEY, + name VARCHAR(100), + price DOUBLE PRECISION + ) + """)) + conn.execute(sa.text(""" + INSERT INTO custom_schema.node_product (name, price) VALUES + ('Widget', 9.99), + ('Gadget', 24.99), + ('Doohickey', 14.99) + """)) + conn.close() + + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS customdb (DBTYPE PG_CLIENT, SCHEMA = 'custom_schema'); + LOAD FROM customdb.node_product RETURN *; + """ + stdout, stderr = run_lbug(script) + self.assertIn("Widget", stdout, f"Schema attach failed: {stderr}") + self.assertIn("Gadget", stdout, f"Gadget not found: {stderr}") + + def test_12_regular_table_not_discovered(self): + """Test that tables without node_/rel_ prefix are not exposed.""" + # Create a table without prefix + engine = sa.create_engine(self.uri, isolation_level="AUTOCOMMIT") + conn = engine.connect() + with conn.begin(): + conn.execute(sa.text(""" + CREATE TABLE IF NOT EXISTS hidden_data ( + id SERIAL PRIMARY KEY, + secret VARCHAR(100) + ) + """)) + conn.execute(sa.text(""" + INSERT INTO hidden_data (secret) VALUES ('sensitive') + """)) + conn.close() + + script = f""" + LOAD EXTENSION '{PG_CLIENT_EXT}'; + ATTACH '{self.conn_str}' AS testdb (DBTYPE PG_CLIENT); + LOAD FROM testdb.hidden_data RETURN *; + """ + stdout, stderr = run_lbug(script) + self.assertIn("Error", stdout, "Expected error for non-node/rel table") + + +if __name__ == "__main__": + unittest.main()