diff --git a/.github/workflows/IntegrationTests.yml b/.github/workflows/IntegrationTests.yml index 7159e7e40..6da8afbe6 100644 --- a/.github/workflows/IntegrationTests.yml +++ b/.github/workflows/IntegrationTests.yml @@ -38,6 +38,119 @@ jobs: clang-format --dump-config make format-check + linux-yugabyte: + name: YugabyteDB Tests + needs: format-check + runs-on: ubuntu-latest + + env: + GEN: ninja + CC: 'ccache gcc' + CXX: 'ccache g++' + CCACHE_DIR: ${{ github.workspace }}/ccache + VCPKG_TARGET_TRIPLET: x64-linux-release + VCPKG_HOST_TRIPLET: x64-linux-release + VCPKG_TOOLCHAIN_PATH: ${{ github.workspace }}/vcpkg/scripts/buildsystems/vcpkg.cmake + + steps: + - name: Checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + submodules: 'true' + + - name: Install Dependencies + run: | + sudo apt-get update -y -q -o=Dpkg::Use-Pty=0 + sudo apt-get install -y -q -o=Dpkg::Use-Pty=0 \ + build-essential \ + ccache \ + cmake \ + ninja-build \ + postgresql-client + + - name: Start YugabyteDB + run: | + docker run -d --name yugabyte \ + -p 5433:5433 -p 7000:7000 -p 9000:9000 \ + yugabytedb/yugabyte:2025.1.4.0-b103 \ + bin/yugabyted start --daemon=false + + - name: Cache Key + id: cache_key + run: | + DUCKDB_VERSION=$(cd duckdb && git rev-parse --short HEAD) + EXT_VERSION=$(git rev-parse --short HEAD) + KEY="${{ runner.os }}-${{ runner.arch }}-${DUCKDB_VERSION}-${EXT_VERSION}-yugabyte" + echo "value=${KEY}" >> "${GITHUB_OUTPUT}" + + - name: Restore Cache + uses: actions/cache/restore@v5 + with: + path: ${{ github.workspace }}/ccache + key: ${{ steps.cache_key.outputs.value }} + + - name: Setup vcpkg + uses: lukka/run-vcpkg@v11.1 + with: + vcpkgGitCommitId: 84bab45d415d22042bd0b9081aea57f362da3f35 + + - name: Build extension + run: | + make release + + - name: Save Cache + uses: actions/cache/save@v5 + with: + path: ${{ github.workspace }}/ccache + key: ${{ steps.cache_key.outputs.value }} + + - name: Wait for YugabyteDB readiness + env: + PGHOST: localhost + PGPORT: '5433' + PGUSER: yugabyte + PGPASSWORD: yugabyte + run: | + echo "Waiting for YugabyteDB to accept connections..." + for i in $(seq 1 40); do + if psql -d yugabyte -c "SELECT count(*) FROM yb_servers()" 2>/dev/null; then + echo "YugabyteDB cluster is ready (attempt $i)" + psql -d yugabyte -c "SELECT host, port, node_type FROM yb_servers()" + exit 0 + fi + echo "Attempt $i/40 — waiting 10s..." + sleep 10 + done + echo "ERROR: YugabyteDB did not become ready after 400 seconds" + docker logs yugabyte 2>&1 | tail -100 + exit 1 + + - name: Setup YugabyteDB test data + env: + PGHOST: localhost + PGPORT: '5433' + PGUSER: yugabyte + PGPASSWORD: yugabyte + run: | + source ./create-yugabyte-tables.sh + + - name: Run YugabyteDB tests + env: + PGHOST: localhost + PGPORT: '5433' + PGUSER: yugabyte + PGPASSWORD: yugabyte + YUGABYTE_TEST_DATABASE_AVAILABLE: 1 + LOCAL_EXTENSION_REPO: 'build/release/repository' + run: | + make test + + - name: YugabyteDB logs on failure + if: failure() + run: | + docker logs yugabyte 2>&1 | tail -200 + linux-tests: name: Linux Tests needs: format-check @@ -262,7 +375,7 @@ jobs: threadsan: [0, 1] env: - CMAKE_BUILD_PARALLEL_LEVEL: 2 + CMAKE_BUILD_PARALLEL_LEVEL: 1 CC: 'ccache gcc' CXX: 'ccache g++' CCACHE_DIR: ${{ github.workspace }}/ccache @@ -367,11 +480,8 @@ jobs: - name: Dependencies shell: bash run: | - choco install \ - ccache \ - make \ - ninja \ - -y --force --no-progress + choco install ccache make -y --force --no-progress + choco install ninja -y --force --no-progress || pip install ninja - name: Build Environment shell: bash diff --git a/.gitignore b/.gitignore index 76c2cc0a0..d3ab301ff 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ cmake-build-debug .vscode .cache duckdb_unittest_tempdir +# CocoIndex Code (ccc) +/.cocoindex_code/ diff --git a/create-yugabyte-tables.sh b/create-yugabyte-tables.sh new file mode 100755 index 000000000..dc287e2fe --- /dev/null +++ b/create-yugabyte-tables.sh @@ -0,0 +1,198 @@ +#!/bin/bash +set -e +set -x + +# Create test tables on YugabyteDB for integration tests. +# Expects YSQL connection via PGHOST/PGPORT/PGUSER/PGPASSWORD env vars. +# +# YugabyteDB can be slow to accept connections after startup. +# Retry with backoff before giving up. + +MAX_RETRIES=30 +RETRY_DELAY=5 + +echo "Waiting for YugabyteDB to accept connections..." +for i in $(seq 1 $MAX_RETRIES); do + if psql -d yugabyte -c "SELECT 1" >/dev/null 2>&1; then + echo "YugabyteDB is ready (attempt $i)" + break + fi + if [ "$i" -eq "$MAX_RETRIES" ]; then + echo "ERROR: YugabyteDB did not become ready after $((MAX_RETRIES * RETRY_DELAY)) seconds" + exit 1 + fi + echo "Attempt $i/$MAX_RETRIES failed, retrying in ${RETRY_DELAY}s..." + sleep $RETRY_DELAY +done + +dropdb --if-exists postgresscanner || true +createdb postgresscanner + +# Hash-partitioned table (default for YugabyteDB) — 100k rows to exercise parallel scan +psql -d postgresscanner -c " +CREATE TABLE hash_test ( + id INTEGER PRIMARY KEY, + name TEXT, + value INTEGER +); +INSERT INTO hash_test SELECT g, 'row_' || g, g * 10 FROM generate_series(1, 100000) g; +ANALYZE hash_test; +" + +# Wide table with various types to stress the COPY path +psql -d postgresscanner -c " +CREATE TABLE wide_test ( + id INTEGER PRIMARY KEY, + col_text TEXT, + col_int BIGINT, + col_float DOUBLE PRECISION, + col_bool BOOLEAN, + col_ts TIMESTAMP, + col_date DATE +); +INSERT INTO wide_test +SELECT g, + 'text_' || g, + g * 100000::BIGINT, + g * 3.14159, + (g % 2 = 0), + '2024-01-01'::TIMESTAMP + (g || ' seconds')::INTERVAL, + '2024-01-01'::DATE + g +FROM generate_series(1, 50000) g; +ANALYZE wide_test; +" + +# Simple test table for attach/detach cycles +psql -d postgresscanner -c " +CREATE TABLE test (i INTEGER); +INSERT INTO test VALUES (1), (2), (3), (NULL); +" + +# Null test table +psql -d postgresscanner -c " +CREATE TABLE nulltest ( + c1 INTEGER, c2 INTEGER, c3 INTEGER, c4 INTEGER, c5 INTEGER, + c6 INTEGER, c7 INTEGER, c8 INTEGER, c9 INTEGER, c10 INTEGER +); +INSERT INTO nulltest VALUES (1, 2, 3, 4, 5, 6, 7, 8, 9, 10); +INSERT INTO nulltest VALUES (1, NULL, 3, 4, NULL, 6, 7, 8, NULL, 10); +INSERT INTO nulltest VALUES (NULL, NULL, 3, 4, 5, 6, 7, NULL, NULL, NULL); +INSERT INTO nulltest VALUES (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL); +" + +# Multi-column hash key table — tests compound partition key parallel scan +psql -d postgresscanner -c " +CREATE TABLE multi_hash ( + region TEXT, + id INTEGER, + data TEXT, + PRIMARY KEY (region, id) +); +INSERT INTO multi_hash SELECT 'region_' || (g % 5), g, 'data_' || g FROM generate_series(1, 10000) g; +ANALYZE multi_hash; +" + +# Range-partitioned tables (no hash key — yb_num_hash_key_columns = 0) +# These must fall back to single-threaded scan, NOT yb_hash_code() ranges. + +# Single-column range key +psql -d postgresscanner -c " +CREATE TABLE range_single ( + id INTEGER, + name TEXT, + value INTEGER, + PRIMARY KEY (id ASC) +); +INSERT INTO range_single SELECT g, 'range_' || g, g * 10 FROM generate_series(1, 10000) g; +ANALYZE range_single; +" + +# Compound range key (timeseries pattern) +psql -d postgresscanner -c " +CREATE TABLE range_ts ( + ts TIMESTAMP, + sensor_id INTEGER, + reading DOUBLE PRECISION, + PRIMARY KEY (ts ASC, sensor_id ASC) +); +INSERT INTO range_ts +SELECT '2024-01-01'::TIMESTAMP + (g || ' seconds')::INTERVAL, + g % 100, + random() * 1000 +FROM generate_series(1, 20000) g; +ANALYZE range_ts; +" + +# Range key with DESC ordering +psql -d postgresscanner -c " +CREATE TABLE range_desc ( + created_at TIMESTAMP, + id INTEGER, + payload TEXT, + PRIMARY KEY (created_at DESC, id DESC) +); +INSERT INTO range_desc +SELECT '2024-06-01'::TIMESTAMP - (g || ' seconds')::INTERVAL, + g, + 'payload_' || g +FROM generate_series(1, 15000) g; +ANALYZE range_desc; +" + +# Colocated database — all tables share a single tablet (no hash partitioning) +# This exercises the non-parallel scan fallback and tests that yb_table_properties +# returns 0 tablets for colocated tables. +# YugabyteDB colocated databases must be created with colocation=true at CREATE DATABASE time. +dropdb --if-exists postgresscanner_colocated || true +psql -d yugabyte -c "CREATE DATABASE postgresscanner_colocated WITH colocation = true" +psql -d postgresscanner_colocated -c " +CREATE TABLE coloc_test ( + id INTEGER PRIMARY KEY, + name TEXT, + value INTEGER +); +INSERT INTO coloc_test SELECT g, 'coloc_' || g, g * 10 FROM generate_series(1, 10000) g; +ANALYZE coloc_test; +" + +psql -d postgresscanner_colocated -c " +CREATE TABLE coloc_wide ( + id INTEGER PRIMARY KEY, + col_text TEXT, + col_int BIGINT, + col_float DOUBLE PRECISION, + col_bool BOOLEAN +); +INSERT INTO coloc_wide +SELECT g, 'text_' || g, g * 100000::BIGINT, g * 3.14, (g % 2 = 0) +FROM generate_series(1, 5000) g; +ANALYZE coloc_wide; +" + +# Non-colocated table in the same database for contrast +psql -d postgresscanner_colocated -c " +CREATE TABLE non_coloc_test ( + id INTEGER PRIMARY KEY, + name TEXT, + value INTEGER +) WITH (colocation = false); +INSERT INTO non_coloc_test SELECT g, 'nocoloc_' || g, g * 10 FROM generate_series(1, 10000) g; +ANALYZE non_coloc_test; +" + +echo "YugabyteDB test tables created successfully" +echo "" +echo " Database: postgresscanner (non-colocated, default)" +echo " hash_test: 100,000 rows (hash-partitioned, single key)" +echo " wide_test: 50,000 rows (hash-partitioned, multiple types)" +echo " multi_hash: 10,000 rows (hash-partitioned, compound key)" +echo " range_single: 10,000 rows (range ASC, single key)" +echo " range_ts: 20,000 rows (range ASC compound key, timeseries)" +echo " range_desc: 15,000 rows (range DESC compound key)" +echo " test: 4 rows (simple)" +echo " nulltest: 4 rows (null patterns)" +echo "" +echo " Database: postgresscanner_colocated" +echo " coloc_test: 10,000 rows (colocated, single tablet)" +echo " coloc_wide: 5,000 rows (colocated, multiple types)" +echo " non_coloc_test: 10,000 rows (non-colocated in colocated db)" diff --git a/src/include/postgres_connection.hpp b/src/include/postgres_connection.hpp index 8d5cc7b72..3511d3f72 100644 --- a/src/include/postgres_connection.hpp +++ b/src/include/postgres_connection.hpp @@ -31,6 +31,7 @@ struct OwnedPostgresConnection { PGconn *connection; mutex connection_lock; + PostgresInstanceType instance_type = PostgresInstanceType::POSTGRES; }; class PostgresConnection { @@ -63,6 +64,8 @@ class PostgresConnection { void BeginCopyTo(ClientContext &context, PostgresCopyState &state, PostgresCopyFormat format, const string &schema_name, const string &table_name, const vector &column_names); + void CommitAndRestartCopy(ClientContext &context, PostgresCopyState &state, PostgresCopyFormat format, + const string &schema_name, const string &table_name, const vector &column_names); void CopyData(data_ptr_t buffer, idx_t size); void CopyData(PostgresBinaryWriter &writer); void CopyData(PostgresTextWriter &writer); diff --git a/src/include/postgres_scanner.hpp b/src/include/postgres_scanner.hpp index 3bf4bd5ad..a459b6241 100644 --- a/src/include/postgres_scanner.hpp +++ b/src/include/postgres_scanner.hpp @@ -50,6 +50,10 @@ struct PostgresBindData : public FunctionData { bool use_text_protocol = false; idx_t max_threads = 1; + idx_t yb_num_tablets = 0; + idx_t yb_num_hash_key_columns = 0; + vector yb_hash_partition_columns; + public: void SetTablePages(idx_t approx_num_pages); diff --git a/src/include/postgres_version.hpp b/src/include/postgres_version.hpp index d26b0526a..b7ea9f96b 100644 --- a/src/include/postgres_version.hpp +++ b/src/include/postgres_version.hpp @@ -12,7 +12,7 @@ namespace duckdb { -enum class PostgresInstanceType { UNKNOWN, POSTGRES, AURORA, REDSHIFT }; +enum class PostgresInstanceType { UNKNOWN, POSTGRES, AURORA, REDSHIFT, YUGABYTE }; struct PostgresVersion { PostgresVersion() { diff --git a/src/include/storage/postgres_catalog.hpp b/src/include/storage/postgres_catalog.hpp index 74fa3577e..582c2d4cd 100644 --- a/src/include/storage/postgres_catalog.hpp +++ b/src/include/storage/postgres_catalog.hpp @@ -15,6 +15,7 @@ #include "postgres_connection.hpp" #include "storage/postgres_schema_set.hpp" #include "storage/postgres_connection_pool.hpp" +#include "yugabyte_topology.hpp" namespace duckdb { class PostgresCatalog; @@ -87,6 +88,10 @@ class PostgresCatalog : public Catalog { return connection_pool; } + const YugabyteTopology &GetYugabyteTopology() const { + return yb_topology; + } + void ClearCache(); //! Whether or not this catalog should search a specific type with the standard priority @@ -111,6 +116,7 @@ class PostgresCatalog : public Catalog { PostgresSchemaSet schemas; shared_ptr connection_pool; string default_schema; + YugabyteTopology yb_topology; }; } // namespace duckdb diff --git a/src/include/storage/postgres_connection_pool.hpp b/src/include/storage/postgres_connection_pool.hpp index d14176429..ddabb7607 100644 --- a/src/include/storage/postgres_connection_pool.hpp +++ b/src/include/storage/postgres_connection_pool.hpp @@ -47,6 +47,8 @@ class PostgresConnectionPool : public dbconnector::pool::ConnectionPool CreateConnectionToHost(const string &host, int32_t port); + protected: std::unique_ptr CreateNewConnection() override; bool CheckConnectionHealthy(PostgresConnection &conn) override; diff --git a/src/include/storage/postgres_table_entry.hpp b/src/include/storage/postgres_table_entry.hpp index b97677c39..5eb8910db 100644 --- a/src/include/storage/postgres_table_entry.hpp +++ b/src/include/storage/postgres_table_entry.hpp @@ -36,6 +36,9 @@ struct PostgresTableInfo { vector postgres_types; vector postgres_names; int64_t approx_num_pages = 0; + idx_t yb_num_tablets = 0; + idx_t yb_num_hash_key_columns = 0; + vector yb_hash_partition_columns; }; class PostgresTableEntry : public TableCatalogEntry { @@ -65,6 +68,9 @@ class PostgresTableEntry : public TableCatalogEntry { vector postgres_names; //! The approximate number of pages a table consumes in Postgres std::atomic approx_num_pages; + idx_t yb_num_tablets = 0; + idx_t yb_num_hash_key_columns = 0; + vector yb_hash_partition_columns; }; } // namespace duckdb diff --git a/src/include/yugabyte_topology.hpp b/src/include/yugabyte_topology.hpp new file mode 100644 index 000000000..c8aa16f38 --- /dev/null +++ b/src/include/yugabyte_topology.hpp @@ -0,0 +1,31 @@ +//===----------------------------------------------------------------------===// +// DuckDB +// +// yugabyte_topology.hpp +// +// +//===----------------------------------------------------------------------===// + +#pragma once + +#include "duckdb/common/common.hpp" +#include "duckdb/common/vector.hpp" + +namespace duckdb { + +struct YugabyteTserver { + string host; + int32_t port = 5433; + string cloud; + string region; + string zone; + string ip_address; + bool reachable = false; +}; + +struct YugabyteTopology { + vector tservers; + bool direct_connect_available = false; +}; + +} // namespace duckdb diff --git a/src/postgres_connection.cpp b/src/postgres_connection.cpp index b5a10689f..7df5257dc 100644 --- a/src/postgres_connection.cpp +++ b/src/postgres_connection.cpp @@ -172,6 +172,13 @@ PostgresVersion PostgresConnection::GetPostgresVersion(ClientContext &context) { if (StringUtil::Contains(pg_version_string, "Redshift")) { version.type_v = PostgresInstanceType::REDSHIFT; } + auto yb_pos = pg_version_string.find("-YB-"); + if (yb_pos != string::npos) { + version.type_v = PostgresInstanceType::YUGABYTE; + } + if (connection) { + connection->instance_type = version.type_v; + } return version; } @@ -212,7 +219,13 @@ void PostgresConnection::Reset(const std::string &health_check_query) { PGresult *res = PQexec(conn, "ROLLBACK"); PostgresResult res_holder(res); } - { + if (connection->instance_type == PostgresInstanceType::YUGABYTE) { + PGresult *res = PQexec(conn, "RESET ALL; DEALLOCATE ALL; CLOSE ALL; UNLISTEN *"); + PostgresResult res_holder(res); + if (PQresultStatus(res) == PGRES_COMMAND_OK) { + return; + } + } else { PGresult *res = PQexec(conn, "DISCARD ALL"); PostgresResult res_holder(res); if (PQresultStatus(res) == PGRES_COMMAND_OK) { diff --git a/src/postgres_copy_to.cpp b/src/postgres_copy_to.cpp index 72580a358..27f8d88a6 100644 --- a/src/postgres_copy_to.cpp +++ b/src/postgres_copy_to.cpp @@ -26,6 +26,12 @@ void PostgresCopyState::Initialize(ClientContext &context) { void PostgresConnection::BeginCopyTo(ClientContext &context, PostgresCopyState &state, PostgresCopyFormat format, const string &schema_name, const string &table_name, const vector &column_names) { + Value yb_disable_txn_writes; + if (context.TryGetCurrentSetting("pg_yb_disable_transactional_writes", yb_disable_txn_writes) && + !yb_disable_txn_writes.IsNull() && BooleanValue::Get(yb_disable_txn_writes)) { + Execute(context, "SET yb_disable_transactional_writes = true"); + } + string query = "COPY "; if (!schema_name.empty()) { query += KeywordHelper::WriteQuoted(schema_name, '"') + "."; @@ -87,6 +93,15 @@ void PostgresConnection::CopyData(PostgresTextWriter &writer) { CopyData(writer.stream.GetData(), writer.stream.GetPosition()); } +void PostgresConnection::CommitAndRestartCopy(ClientContext &context, PostgresCopyState &state, + PostgresCopyFormat format, const string &schema_name, + const string &table_name, const vector &column_names) { + FinishCopyTo(state); + Execute(context, "COMMIT"); + Execute(context, "BEGIN"); + BeginCopyTo(context, state, format, schema_name, table_name, column_names); +} + void PostgresConnection::FinishCopyTo(PostgresCopyState &state) { if (state.format == PostgresCopyFormat::BINARY) { // binary copy requires a footer diff --git a/src/postgres_extension.cpp b/src/postgres_extension.cpp index 208092ff7..82ae07d0e 100644 --- a/src/postgres_extension.cpp +++ b/src/postgres_extension.cpp @@ -117,6 +117,8 @@ unique_ptr CreatePostgresSecretFunction(ClientContext &context, Crea result->secret_map["port"] = named_param.second.ToString(); } else if (lower_name == "passfile") { result->secret_map["passfile"] = named_param.second.ToString(); + } else if (lower_name == "options") { + result->secret_map["options"] = named_param.second.ToString(); } else { throw InternalException("Unknown named parameter passed to CreatePostgresSecretFunction: " + lower_name); } @@ -135,6 +137,7 @@ void SetPostgresSecretParameters(CreateSecretFunction &function) { function.named_parameters["database"] = LogicalType::VARCHAR; // alias for dbname function.named_parameters["dbname"] = LogicalType::VARCHAR; function.named_parameters["passfile"] = LogicalType::VARCHAR; + function.named_parameters["options"] = LogicalType::VARCHAR; } void SetPostgresNullByteReplacement(ClientContext &context, SetScope scope, Value ¶meter) { @@ -302,6 +305,23 @@ static void LoadInternal(ExtensionLoader &loader) { LogicalType::VARCHAR, PostgresConnectionPool::DefaultHealthCheckQuery(), nullptr, SetScope::GLOBAL); + // YugabyteDB-specific options + config.AddExtensionOption( + "pg_yb_parallel_scan", + "Enable hash-code parallel scanning on YugabyteDB. Workers use separate REPEATABLE READ " + "transactions without a shared snapshot, so concurrent writes may cause inconsistent reads. " + "Safe for read-only or append-only tables.", + LogicalType::BOOLEAN, Value::BOOLEAN(false)); + config.AddExtensionOption("pg_yb_tserver_probe_timeout", + "Connect timeout in seconds for tserver reachability probes during ATTACH", + LogicalType::UINTEGER, Value::UINTEGER(2)); + config.AddExtensionOption("pg_yb_rows_per_transaction", + "Number of rows per transaction batch for COPY FROM on YugabyteDB (0 to disable)", + LogicalType::UBIGINT, Value::UBIGINT(10000)); + config.AddExtensionOption("pg_yb_disable_transactional_writes", + "Disable transactional writes for bulk COPY FROM on YugabyteDB (no rollback on failure)", + LogicalType::BOOLEAN, Value::BOOLEAN(false)); + OptimizerExtension postgres_optimizer; postgres_optimizer.optimize_function = PostgresOptimizer::Optimize; OptimizerExtension::Register(config, std::move(postgres_optimizer)); diff --git a/src/postgres_scanner.cpp b/src/postgres_scanner.cpp index 6679746cd..f293874ea 100644 --- a/src/postgres_scanner.cpp +++ b/src/postgres_scanner.cpp @@ -51,6 +51,9 @@ struct PostgresGlobalState : public GlobalTableFunctionState { bool used_main_thread = false; string snapshot; + idx_t yb_hash_idx = 0; + idx_t yb_num_tasks = 0; + PostgresConnection &GetConnection(); void SetConnection(PostgresConnection connection); void SetConnection(shared_ptr connection); @@ -75,6 +78,9 @@ static void PostgresGetSnapshot(ClientContext &context, PostgresVersion version, if (version.type_v == PostgresInstanceType::AURORA) { return; } + if (version.type_v == PostgresInstanceType::YUGABYTE) { + return; + } // SET TRANSACTION SNAPSHOT requires REPEATABLE READ or SERIALIZABLE auto pg_catalog = bind_data.GetCatalog(); if (pg_catalog && pg_catalog->isolation_level == PostgresIsolationLevel::READ_COMMITTED) { @@ -130,6 +136,9 @@ void PostgresScanFunction::PrepareBind(PostgresVersion version, ClientContext &c // see https://github.com/duckdb/postgres_scanner/issues/186 use_ctid_scan = false; } + if (version.type_v == PostgresInstanceType::YUGABYTE) { + use_ctid_scan = false; + } if (approx_num_pages < 0) { // negative relpages (e.g. partitioned tables) cannot use ctid scan use_ctid_scan = false; @@ -142,6 +151,18 @@ void PostgresScanFunction::PrepareBind(PostgresVersion version, ClientContext &c if (version.type_v == PostgresInstanceType::REDSHIFT) { bind_data.use_text_protocol = true; } + if (version.type_v == PostgresInstanceType::YUGABYTE && bind_data.yb_num_tablets > 0) { + bool yb_parallel = false; + Value yb_parallel_val; + if (context.TryGetCurrentSetting("pg_yb_parallel_scan", yb_parallel_val) && !yb_parallel_val.IsNull()) { + yb_parallel = BooleanValue::Get(yb_parallel_val); + } + if (!yb_parallel || !bind_data.read_only || bind_data.use_text_protocol) { + bind_data.max_threads = 1; + } else { + bind_data.max_threads = bind_data.yb_num_tablets; + } + } } PostgresBindData::PostgresBindData(ClientContext &context) { @@ -272,6 +293,16 @@ static void PostgresInitInternal(ClientContext &context, const PostgresBindData lstate.done = false; if (bind_data->pages_approx > 0) { filter = StringUtil::Format("WHERE ctid BETWEEN '(%d,0)'::tid AND '(%d,0)'::tid", task_min, task_max); + } else if (bind_data->version.type_v == PostgresInstanceType::YUGABYTE && bind_data->yb_num_hash_key_columns > 0 && + !bind_data->yb_hash_partition_columns.empty()) { + string hash_cols; + for (idx_t i = 0; i < bind_data->yb_hash_partition_columns.size(); i++) { + if (i > 0) { + hash_cols += ", "; + } + hash_cols += KeywordHelper::WriteQuoted(bind_data->yb_hash_partition_columns[i], '"'); + } + filter = StringUtil::Format("WHERE yb_hash_code(%s) BETWEEN %d AND %d", hash_cols, task_min, task_max); } if (!filter_string.empty()) { if (filter.empty()) { @@ -313,11 +344,17 @@ static unique_ptr GetLocalState(ClientContext &context, PostgresGlobalState &gstate); static void PostgresScanConnect(ClientContext &context, PostgresConnection &conn, const string &snapshot, - AccessMode access_mode, PostgresIsolationLevel isolation_level) { - conn.Execute(context, PostgresTransaction::GetBeginTransactionQuery(isolation_level, access_mode)); - if (!snapshot.empty()) { - D_ASSERT(isolation_level != PostgresIsolationLevel::READ_COMMITTED); - conn.Query(context, StringUtil::Format("SET TRANSACTION SNAPSHOT '%s'", snapshot)); + AccessMode access_mode, PostgresIsolationLevel isolation_level, + PostgresInstanceType instance_type = PostgresInstanceType::POSTGRES) { + if (instance_type == PostgresInstanceType::YUGABYTE) { + conn.Execute(context, PostgresTransaction::GetBeginTransactionQuery(PostgresIsolationLevel::REPEATABLE_READ, + AccessMode::READ_ONLY)); + } else { + conn.Execute(context, PostgresTransaction::GetBeginTransactionQuery(isolation_level, access_mode)); + if (!snapshot.empty()) { + D_ASSERT(isolation_level != PostgresIsolationLevel::READ_COMMITTED); + conn.Query(context, StringUtil::Format("SET TRANSACTION SNAPSHOT '%s'", snapshot)); + } } Value statement_timeout; if (context.TryGetCurrentSetting("pg_statement_timeout_millis", statement_timeout) && !statement_timeout.IsNull()) { @@ -375,6 +412,12 @@ static unique_ptr PostgresInitGlobalState(ClientContex // we create a transaction here, and get the snapshot id to enable transaction-safe parallelism PostgresGetSnapshot(context, bind_data.version, bind_data, *result); } + + if (bind_data.version.type_v == PostgresInstanceType::YUGABYTE && bind_data.yb_num_hash_key_columns > 0 && + bind_data.max_threads > 1) { + result->yb_num_tasks = bind_data.max_threads; + } + return std::move(result); } @@ -385,6 +428,21 @@ static bool PostgresParallelStateNext(ClientContext &context, const FunctionData lock_guard parallel_lock(gstate.lock); lstate.batch_idx = gstate.batch_idx++; + + if (bind_data->version.type_v == PostgresInstanceType::YUGABYTE && bind_data->yb_num_hash_key_columns > 0 && + gstate.yb_num_tasks > 0) { + if (gstate.yb_hash_idx >= gstate.yb_num_tasks) { + lstate.done = true; + return false; + } + idx_t range_size = 65536 / gstate.yb_num_tasks; + idx_t range_min = gstate.yb_hash_idx * range_size; + idx_t range_max = (gstate.yb_hash_idx == gstate.yb_num_tasks - 1) ? 65535 : range_min + range_size - 1; + gstate.yb_hash_idx++; + PostgresInitInternal(context, bind_data, lstate, range_min, range_max); + return true; + } + if (gstate.page_idx < bind_data->pages_approx) { auto page_max = gstate.page_idx + bind_data->pages_per_task; if (page_max >= bind_data->pages_approx || page_max > POSTGRES_TID_MAX) { @@ -420,15 +478,37 @@ bool PostgresGlobalState::TryOpenNewConnection(ClientContext &context, PostgresL } if (pg_catalog) { + if (pg_catalog->GetPostgresVersion().type_v == PostgresInstanceType::YUGABYTE && + pg_catalog->GetYugabyteTopology().direct_connect_available) { + auto &topology = pg_catalog->GetYugabyteTopology(); + lock_guard parallel_lock(lock); + idx_t ts_idx = batch_idx % topology.tservers.size(); + for (idx_t i = 0; i < topology.tservers.size(); i++) { + auto &ts = topology.tservers[(ts_idx + i) % topology.tservers.size()]; + if (ts.reachable) { + try { + auto conn = pg_catalog->GetConnectionPool().CreateConnectionToHost(ts.ip_address, ts.port); + lstate.connection = std::move(*conn); + PostgresScanConnect(context, lstate.connection, snapshot, pg_catalog->access_mode, + pg_catalog->isolation_level, bind_data.version.type_v); + return true; + } catch (...) { + continue; + } + } + } + } + if (!pg_catalog->GetConnectionPool().TryGetConnection(lstate.pool_connection)) { return false; } lstate.connection = PostgresConnection(lstate.pool_connection.GetConnection().GetConnection()); - PostgresScanConnect(context, lstate.connection, snapshot, pg_catalog->access_mode, pg_catalog->isolation_level); + PostgresScanConnect(context, lstate.connection, snapshot, pg_catalog->access_mode, pg_catalog->isolation_level, + bind_data.version.type_v); } else { lstate.connection = PostgresConnection::Open(bind_data.dsn, bind_data.attach_path); PostgresScanConnect(context, lstate.connection, snapshot, AccessMode::READ_ONLY, - PostgresIsolationLevel::REPEATABLE_READ); + PostgresIsolationLevel::REPEATABLE_READ, bind_data.version.type_v); } return true; } @@ -449,7 +529,7 @@ static unique_ptr GetLocalState(ClientContext &context, local_state->no_connection = true; return std::move(local_state); } - if (bind_data.pages_approx == 0 || bind_data.requires_materialization) { + if ((bind_data.pages_approx == 0 && gstate.yb_num_tasks == 0) || bind_data.requires_materialization) { PostgresInitInternal(context, &bind_data, *local_state, 0, POSTGRES_TID_MAX); lock_guard parallel_lock(gstate.lock); gstate.page_idx = POSTGRES_TID_MAX; @@ -555,6 +635,9 @@ double PostgresScanProgress(ClientContext &context, const FunctionData *bind_dat auto &gstate = global_state->Cast(); lock_guard parallel_lock(gstate.lock); + if (gstate.yb_num_tasks > 0) { + return MinValue(100, 100.0 * double(gstate.yb_hash_idx) / double(gstate.yb_num_tasks)); + } double progress = 100 * double(gstate.page_idx) / double(bind_data.pages_approx); return MinValue(100, progress); } diff --git a/src/storage/postgres_catalog.cpp b/src/storage/postgres_catalog.cpp index 868c3291f..96a0fa3e8 100644 --- a/src/storage/postgres_catalog.cpp +++ b/src/storage/postgres_catalog.cpp @@ -1,4 +1,6 @@ #include "storage/postgres_catalog.hpp" +#include "yugabyte_topology.hpp" +#include "duckdb/common/string_util.hpp" #include "storage/postgres_schema_entry.hpp" #include "storage/postgres_transaction.hpp" #include "postgres_connection.hpp" @@ -10,6 +12,45 @@ namespace duckdb { +static void DiscoverYugabyteTopology(ClientContext &context, PostgresConnection &conn, const string &connection_string, + YugabyteTopology &topology) { + auto result = + conn.TryQuery(context, "SELECT host, port, node_type, cloud, region, zone, public_ip FROM yb_servers()"); + if (!result) { + return; + } + auto rows = result->Count(); + for (idx_t r = 0; r < rows; r++) { + YugabyteTserver ts; + ts.host = result->GetString(r, 0); + ts.port = result->IsNull(r, 1) ? 5433 : static_cast(result->GetInt64(r, 1)); + ts.cloud = result->IsNull(r, 3) ? "" : result->GetString(r, 3); + ts.region = result->IsNull(r, 4) ? "" : result->GetString(r, 4); + ts.zone = result->IsNull(r, 5) ? "" : result->GetString(r, 5); + ts.ip_address = result->IsNull(r, 6) ? ts.host : result->GetString(r, 6); + topology.tservers.push_back(std::move(ts)); + } + + uint32_t probe_timeout = 2; + Value timeout_val; + if (context.TryGetCurrentSetting("pg_yb_tserver_probe_timeout", timeout_val) && !timeout_val.IsNull()) { + probe_timeout = UIntegerValue::Get(timeout_val); + } + + for (auto &ts : topology.tservers) { + string probe_dsn = connection_string + StringUtil::Format(" host='%s' port=%d connect_timeout=%d", + ts.ip_address, ts.port, probe_timeout); + PGconn *probe = PQconnectdb(probe_dsn.c_str()); + if (probe && PQstatus(probe) == CONNECTION_OK) { + ts.reachable = true; + topology.direct_connect_available = true; + } + if (probe) { + PQfinish(probe); + } + } +} + PostgresCatalog::PostgresCatalog(AttachedDatabase &db_p, string connection_string_p, string attach_path_p, AccessMode access_mode, string schema_to_load, PostgresIsolationLevel isolation_level, ClientContext &context) @@ -22,6 +63,10 @@ PostgresCatalog::PostgresCatalog(AttachedDatabase &db_p, string connection_strin auto connection = connection_pool->GetConnection(); this->version = connection.GetConnection().GetPostgresVersion(context); + + if (version.type_v == PostgresInstanceType::YUGABYTE) { + DiscoverYugabyteTopology(context, connection.GetConnection(), connection_string, yb_topology); + } } string EscapeConnectionString(const string &input) { @@ -89,6 +134,7 @@ string PostgresCatalog::GetConnectionString(ClientContext &context, const string new_connection_info += AddConnectionOption(kv_secret, "port"); new_connection_info += AddConnectionOption(kv_secret, "dbname"); new_connection_info += AddConnectionOption(kv_secret, "passfile"); + new_connection_info += AddConnectionOption(kv_secret, "options"); connection_string = new_connection_info + connection_string; } else if (explicit_secret) { diff --git a/src/storage/postgres_connection_pool.cpp b/src/storage/postgres_connection_pool.cpp index 098405b64..fddb7bc93 100644 --- a/src/storage/postgres_connection_pool.cpp +++ b/src/storage/postgres_connection_pool.cpp @@ -67,6 +67,13 @@ std::unique_ptr PostgresConnectionPool::CreateNewConnection( return make_uniq(std::move(conn)); } +std::unique_ptr PostgresConnectionPool::CreateConnectionToHost(const string &host, int32_t port) { + string tserver_dsn = postgres_catalog.connection_string; + tserver_dsn += StringUtil::Format(" host='%s' port=%d", host, port); + auto conn = PostgresConnection::Open(tserver_dsn, postgres_catalog.attach_path); + return make_uniq(std::move(conn)); +} + bool PostgresConnectionPool::CheckConnectionHealthy(PostgresConnection &conn) { if (!conn.IsOpen()) { return false; diff --git a/src/storage/postgres_insert.cpp b/src/storage/postgres_insert.cpp index 5310b1d53..4f694b6f5 100644 --- a/src/storage/postgres_insert.cpp +++ b/src/storage/postgres_insert.cpp @@ -32,13 +32,22 @@ PostgresInsert::PostgresInsert(PhysicalPlan &physical_plan, LogicalOperator &op, class PostgresInsertGlobalState : public GlobalSinkState { public: explicit PostgresInsertGlobalState(ClientContext &context, PostgresTableEntry &table, PostgresCopyFormat format) - : table(table), insert_count(0), format(format) { + : table(table), insert_count(0), rows_in_batch(0), yb_rows_per_transaction(0), format(format) { + auto &pg_catalog = table.catalog.Cast(); + if (pg_catalog.GetPostgresVersion().type_v == PostgresInstanceType::YUGABYTE) { + Value val; + if (context.TryGetCurrentSetting("pg_yb_rows_per_transaction", val) && !val.IsNull()) { + yb_rows_per_transaction = UBigIntValue::Get(val); + } + } } PostgresTableEntry &table; PostgresCopyState copy_state; DataChunk varchar_chunk; idx_t insert_count; + idx_t rows_in_batch; + idx_t yb_rows_per_transaction; PostgresCopyFormat format; vector insert_column_names; bool copy_is_active = false; @@ -121,8 +130,12 @@ SinkResultType PostgresInsert::Sink(ExecutionContext &context, DataChunk &chunk, } connection.CopyChunk(context.client, gstate.copy_state, chunk, gstate.varchar_chunk); gstate.insert_count += chunk.size(); - if (!keep_copy_alive) { - // if we are can't keep the copy alive we need to restart the copy during every sink + gstate.rows_in_batch += chunk.size(); + if (gstate.yb_rows_per_transaction > 0 && gstate.rows_in_batch >= gstate.yb_rows_per_transaction) { + connection.CommitAndRestartCopy(context.client, gstate.copy_state, gstate.format, gstate.table.schema.name, + gstate.table.name, gstate.insert_column_names); + gstate.rows_in_batch = 0; + } else if (!keep_copy_alive) { gstate.FinishCopyTo(connection); } return SinkResultType::NEED_MORE_INPUT; diff --git a/src/storage/postgres_table_entry.cpp b/src/storage/postgres_table_entry.cpp index a746a456b..e952b8680 100644 --- a/src/storage/postgres_table_entry.cpp +++ b/src/storage/postgres_table_entry.cpp @@ -23,7 +23,9 @@ PostgresTableEntry::PostgresTableEntry(Catalog &catalog, SchemaCatalogEntry &sch PostgresTableEntry::PostgresTableEntry(Catalog &catalog, SchemaCatalogEntry &schema, PostgresTableInfo &info) : TableCatalogEntry(catalog, schema, *info.create_info), postgres_types(std::move(info.postgres_types)), - postgres_names(std::move(info.postgres_names)) { + postgres_names(std::move(info.postgres_names)), yb_num_tablets(info.yb_num_tablets), + yb_num_hash_key_columns(info.yb_num_hash_key_columns), + yb_hash_partition_columns(std::move(info.yb_hash_partition_columns)) { D_ASSERT(postgres_types.size() == columns.LogicalColumnCount()); approx_num_pages.store(info.approx_num_pages, std::memory_order_release); } @@ -54,6 +56,9 @@ TableFunction PostgresTableEntry::GetScanFunction(ClientContext &context, unique result->names = postgres_names; result->postgres_types = postgres_types; result->read_only = transaction.IsReadOnly(); + result->yb_num_tablets = yb_num_tablets; + result->yb_num_hash_key_columns = yb_num_hash_key_columns; + result->yb_hash_partition_columns = yb_hash_partition_columns; PostgresScanFunction::PrepareBind(pg_catalog.GetPostgresVersion(), context, *result, approx_num_pages.load(std::memory_order_acquire)); diff --git a/src/storage/postgres_table_set.cpp b/src/storage/postgres_table_set.cpp index 8df2c7b82..7aec18a5c 100644 --- a/src/storage/postgres_table_set.cpp +++ b/src/storage/postgres_table_set.cpp @@ -12,6 +12,8 @@ #include "duckdb/parser/parser.hpp" #include "duckdb/common/string_util.hpp" #include "postgres_conversion.hpp" +#include "storage/postgres_catalog.hpp" +#include "postgres_version.hpp" namespace duckdb { @@ -118,7 +120,51 @@ void PostgresTableSet::AddColumnOrConstraint(optional_ptr t } } +static void LoadYugabyteTableProperties(PostgresTransaction &transaction, PostgresTableInfo &table_info, + const string &schema_name) { + string qualified = + KeywordHelper::WriteQuoted(schema_name, '"') + "." + KeywordHelper::WriteQuoted(table_info.GetTableName(), '"'); + string escaped = StringUtil::Replace(qualified, "'", "''"); + + string props_query = StringUtil::Format( + "SELECT num_tablets, num_hash_key_columns FROM yb_table_properties('%s'::regclass)", escaped); + try { + auto result = transaction.Query(props_query); + if (result && result->Count() > 0) { + table_info.yb_num_tablets = result->IsNull(0, 0) ? 0 : result->GetInt64(0, 0); + table_info.yb_num_hash_key_columns = result->IsNull(0, 1) ? 0 : result->GetInt64(0, 1); + } + } catch (...) { + return; + } + + if (table_info.yb_num_hash_key_columns > 0) { + string pk_query = + StringUtil::Format("SELECT a.attname " + "FROM pg_index i " + "JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) " + "WHERE i.indrelid = '%s'::regclass AND i.indisprimary " + "ORDER BY array_position(i.indkey, a.attnum) " + "LIMIT %d", + escaped, table_info.yb_num_hash_key_columns); + try { + auto result = transaction.Query(pk_query); + if (result) { + for (idx_t r = 0; r < result->Count(); r++) { + table_info.yb_hash_partition_columns.push_back(result->GetString(r, 0)); + } + } + } catch (...) { + table_info.yb_hash_partition_columns.clear(); + table_info.yb_num_hash_key_columns = 0; + } + } +} + void PostgresTableSet::CreateEntries(PostgresTransaction &transaction, PostgresResult &result, idx_t start, idx_t end) { + auto &pg_catalog = catalog.Cast(); + bool is_yugabyte = pg_catalog.GetPostgresVersion().type_v == PostgresInstanceType::YUGABYTE; + vector> tables; unique_ptr info; @@ -137,6 +183,9 @@ void PostgresTableSet::CreateEntries(PostgresTransaction &transaction, PostgresR tables.push_back(std::move(info)); } for (auto &tbl_info : tables) { + if (is_yugabyte) { + LoadYugabyteTableProperties(transaction, *tbl_info, schema.name); + } auto table_entry = make_shared_ptr(catalog, schema, *tbl_info); CreateEntry(transaction, std::move(table_entry)); } diff --git a/test/sql/storage/attach_secret.test b/test/sql/storage/attach_secret.test index 1c92fdf96..eb46c2aab 100644 --- a/test/sql/storage/attach_secret.test +++ b/test/sql/storage/attach_secret.test @@ -61,6 +61,30 @@ unknown_database statement ok ATTACH 'dbname=postgresscanner' AS secret_attach (TYPE POSTGRES, SECRET postgres_db) +statement ok +DETACH secret_attach + +# OPTIONS keyword carries libpq startup options through the secret +statement ok +CREATE OR REPLACE SECRET postgres_db ( + TYPE POSTGRES, + HOST '127.0.0.1', + DBNAME postgresscanner, + PASSWORD 'postgres', + OPTIONS '-c statement_timeout=12345' +); + +statement ok +ATTACH '' AS secret_attach (TYPE POSTGRES, SECRET postgres_db) + +query I +SELECT setting FROM postgres_query('secret_attach', 'SELECT setting FROM pg_settings WHERE name = ''statement_timeout''') +---- +12345 + +statement ok +DETACH secret_attach + statement error CREATE SECRET new_secret ( TYPE POSTGRES, diff --git a/test/sql/storage/attach_yugabyte_basic.test b/test/sql/storage/attach_yugabyte_basic.test new file mode 100644 index 000000000..6866394f4 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_basic.test @@ -0,0 +1,120 @@ +# name: test/sql/storage/attach_yugabyte_basic.test +# description: Basic YugabyteDB attach, scan, and detection tests +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +# Attach to YugabyteDB +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# Hash-partitioned table: 100k rows +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query II +SELECT min(id), max(id) FROM yb.hash_test +---- +1 100000 + +query III +SELECT id, name, value FROM yb.hash_test WHERE id = 42 +---- +42 row_42 420 + +# Aggregation over full table to verify no rows lost in scan +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +# Wide table with multiple types: 50k rows +query I +SELECT count(*) FROM yb.wide_test +---- +50000 + +query I +SELECT count(*) FROM yb.wide_test WHERE col_bool = true +---- +25000 + +query I +SELECT count(*) FROM yb.wide_test WHERE col_int > 2500000000 +---- +25000 + +# Range-partitioned tables (no hash key — must NOT use yb_hash_code scan) +query I +SELECT count(*) FROM yb.range_single +---- +10000 + +query I +SELECT count(*) FROM yb.range_ts +---- +20000 + +query I +SELECT count(DISTINCT sensor_id) FROM yb.range_ts +---- +100 + +query I +SELECT count(*) FROM yb.range_desc +---- +15000 + +# Scan with NULLs +query I +SELECT count(*) FROM yb.nulltest +---- +4 + +query I +SELECT count(*) FROM yb.nulltest WHERE c1 IS NULL +---- +2 + +# Simple table scan +query I +SELECT count(*) FROM yb.test +---- +4 + +query I +SELECT count(*) FROM yb.test WHERE i IS NOT NULL +---- +3 + +# Multi-column hash key table: 10k rows +query I +SELECT count(*) FROM yb.multi_hash +---- +10000 + +query I +SELECT count(DISTINCT region) FROM yb.multi_hash +---- +5 + +# pg_catalog scan — exercises system table scan path on YugabyteDB +query I +SELECT count(*) > 0 FROM yb.pg_catalog.pg_class +---- +true + +query I +SELECT count(*) > 0 FROM yb.pg_catalog.pg_tables +---- +true + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_catalog.test b/test/sql/storage/attach_yugabyte_catalog.test new file mode 100644 index 000000000..a66daf10b --- /dev/null +++ b/test/sql/storage/attach_yugabyte_catalog.test @@ -0,0 +1,198 @@ +# name: test/sql/storage/attach_yugabyte_catalog.test +# description: Test catalog operations, system tables, schema handling on YugabyteDB +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# === pg_catalog system table access === +query I +SELECT count(*) > 0 FROM yb.pg_catalog.pg_class +---- +true + +query I +SELECT count(*) > 0 FROM yb.pg_catalog.pg_tables +---- +true + +query I +SELECT count(*) > 0 FROM yb.pg_catalog.pg_namespace +---- +true + +query I +SELECT count(*) > 0 FROM yb.pg_catalog.pg_attribute +---- +true + +query I +SELECT count(*) > 0 FROM yb.pg_catalog.pg_type +---- +true + +# === information_schema access === +query I +SELECT count(*) > 0 FROM yb.information_schema.tables +---- +true + +query I +SELECT count(*) > 0 FROM yb.information_schema.columns +---- +true + +# === Verify our test tables appear in catalog === +query I +SELECT count(*) FROM yb.pg_catalog.pg_tables WHERE tablename = 'hash_test' +---- +1 + +query I +SELECT count(*) FROM yb.pg_catalog.pg_tables WHERE tablename = 'multi_hash' +---- +1 + +query I +SELECT count(*) FROM yb.pg_catalog.pg_tables WHERE tablename = 'range_single' +---- +1 + +# === Schema operations === +statement ok +DROP TABLE IF EXISTS yb.catalog_test + +statement ok +CREATE TABLE yb.catalog_test (id INTEGER PRIMARY KEY, name TEXT, value DOUBLE) + +# Table should appear in catalog +query I +SELECT count(*) FROM yb.pg_catalog.pg_tables WHERE tablename = 'catalog_test' +---- +1 + +# Column metadata should be accessible +query I +SELECT count(*) FROM yb.information_schema.columns WHERE table_name = 'catalog_test' +---- +3 + +query I +INSERT INTO yb.catalog_test SELECT g, 'cat_' || g, g * 2.5 FROM generate_series(1, 100) t(g) +---- +100 + +query I +SELECT count(*) FROM yb.catalog_test +---- +100 + +# === CREATE TABLE with various types === +statement ok +DROP TABLE IF EXISTS yb.types_test + +statement ok +CREATE TABLE yb.types_test ( + id INTEGER PRIMARY KEY, + col_text TEXT, + col_varchar VARCHAR, + col_bigint BIGINT, + col_smallint SMALLINT, + col_double DOUBLE, + col_float FLOAT, + col_bool BOOLEAN, + col_date DATE, + col_timestamp TIMESTAMP +) + +query I +INSERT INTO yb.types_test VALUES ( + 1, 'hello', 'world', 9223372036854775807, 32767, + 3.14159265358979, 2.71828, true, '2024-06-15', '2024-06-15 10:30:00' +) +---- +1 + +query IIIIIIIIII +SELECT * FROM yb.types_test WHERE id = 1 +---- +1 hello world 9223372036854775807 32767 3.14159265358979 2.71828 true 2024-06-15 2024-06-15 10:30:00 + +# === CREATE OR REPLACE TABLE === +statement ok +CREATE OR REPLACE TABLE yb.catalog_test (id INTEGER PRIMARY KEY, new_col TEXT) + +query I +SELECT count(*) FROM yb.information_schema.columns WHERE table_name = 'catalog_test' +---- +2 + +# === Read from one YB table, insert into another (cross-table COPY) === +statement ok +DROP TABLE IF EXISTS yb.cross_copy + +statement ok +CREATE TABLE yb.cross_copy (id INTEGER PRIMARY KEY, name TEXT, value INTEGER) + +query I +INSERT INTO yb.cross_copy SELECT id, name, value FROM yb.hash_test WHERE id <= 500 +---- +500 + +query I +SELECT count(*) FROM yb.cross_copy +---- +500 + +query I +SELECT sum(id) FROM yb.cross_copy +---- +125250 + +# === Multi-schema support === +statement ok +CREATE SCHEMA yb.test_schema + +statement ok +CREATE TABLE yb.test_schema.schema_table (id INTEGER PRIMARY KEY, data TEXT) + +query I +INSERT INTO yb.test_schema.schema_table SELECT g, 'schema_' || g FROM generate_series(1, 100) t(g) +---- +100 + +query I +SELECT count(*) FROM yb.test_schema.schema_table +---- +100 + +query I +SELECT sum(id) FROM yb.test_schema.schema_table +---- +5050 + +statement ok +DROP TABLE yb.test_schema.schema_table + +statement ok +DROP SCHEMA yb.test_schema + +# === Cleanup === +statement ok +DROP TABLE yb.catalog_test + +statement ok +DROP TABLE yb.types_test + +statement ok +DROP TABLE yb.cross_copy + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_colocated.test b/test/sql/storage/attach_yugabyte_colocated.test new file mode 100644 index 000000000..d71f81759 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_colocated.test @@ -0,0 +1,119 @@ +# name: test/sql/storage/attach_yugabyte_colocated.test +# description: Test YugabyteDB colocated vs non-colocated table behavior +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +# Attach to the colocated database +statement ok +ATTACH 'dbname=postgresscanner_colocated' AS ybc (TYPE POSTGRES) + +# === Colocated table: single tablet, no hash partitioning === +# yb_table_properties should return 0 hash key columns for colocated tables, +# forcing single-threaded scan (no yb_hash_code ranges) + +query I +SELECT count(*) FROM ybc.coloc_test +---- +10000 + +query II +SELECT min(id), max(id) FROM ybc.coloc_test +---- +1 10000 + +query I +SELECT sum(id) FROM ybc.coloc_test +---- +50005000 + +query III +SELECT id, name, value FROM ybc.coloc_test WHERE id = 500 +---- +500 coloc_500 5000 + +# === Colocated wide table: multiple types === +query I +SELECT count(*) FROM ybc.coloc_wide +---- +5000 + +query I +SELECT count(*) FROM ybc.coloc_wide WHERE col_bool = true +---- +2500 + +query I +SELECT sum(col_int) FROM ybc.coloc_wide +---- +1250250000000 + +# === Non-colocated table in the same colocated database === +# This table opts out of colocation, so it gets its own tablets +# and should use hash-code parallel scan + +query I +SELECT count(*) FROM ybc.non_coloc_test +---- +10000 + +query I +SELECT sum(id) FROM ybc.non_coloc_test +---- +50005000 + +query III +SELECT id, name, value FROM ybc.non_coloc_test WHERE id = 7777 +---- +7777 nocoloc_7777 77770 + +# === Cross-table join: colocated and non-colocated in same query === +query I +SELECT count(*) FROM ybc.coloc_test c JOIN ybc.non_coloc_test n ON c.id = n.id +---- +10000 + +# === Write to colocated table === +statement ok +DROP TABLE IF EXISTS ybc.coloc_write_test + +statement ok +CREATE TABLE ybc.coloc_write_test (id INTEGER PRIMARY KEY, data TEXT) + +query I +INSERT INTO ybc.coloc_write_test SELECT g, 'data_' || g FROM generate_series(1, 1000) t(g) +---- +1000 + +query I +SELECT count(*) FROM ybc.coloc_write_test +---- +1000 + +query I +SELECT sum(id) FROM ybc.coloc_write_test +---- +500500 + +statement ok +DROP TABLE ybc.coloc_write_test + +# === Detach and re-attach to verify pool reset works for colocated db === +statement ok +DETACH ybc + +statement ok +ATTACH 'dbname=postgresscanner_colocated' AS ybc (TYPE POSTGRES) + +query I +SELECT count(*) FROM ybc.coloc_test +---- +10000 + +statement ok +DETACH ybc diff --git a/test/sql/storage/attach_yugabyte_concurrent.test b/test/sql/storage/attach_yugabyte_concurrent.test new file mode 100644 index 000000000..dec141885 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_concurrent.test @@ -0,0 +1,80 @@ +# name: test/sql/storage/attach_yugabyte_concurrent.test +# description: Test concurrent queries on YugabyteDB — connection pool under contention +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# Concurrent reads — should not crash or return wrong results +concurrentloop i 0 10 + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +endloop + +# Concurrent reads across different tables +concurrentloop i 0 5 + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT count(*) FROM yb.wide_test +---- +50000 + +query I +SELECT count(*) FROM yb.multi_hash +---- +10000 + +endloop + +# Create tables sequentially then write concurrently +loop i 0 5 + +statement ok +CREATE OR REPLACE TABLE yb.concurrent_${i} (id INTEGER PRIMARY KEY, val INTEGER) + +endloop + +# Concurrent writes to separate tables +concurrentloop i 0 5 + +statement ok +INSERT INTO yb.concurrent_${i} SELECT g, g FROM generate_series(1, 100) t(g) + +endloop + +# Verify +loop i 0 5 + +query I +SELECT count(*) FROM yb.concurrent_${i} +---- +100 + +endloop + +# Cleanup +loop i 0 5 + +statement ok +DROP TABLE IF EXISTS yb.concurrent_${i} + +endloop + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_copy_batch.test b/test/sql/storage/attach_yugabyte_copy_batch.test new file mode 100644 index 000000000..c058310a0 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_copy_batch.test @@ -0,0 +1,141 @@ +# name: test/sql/storage/attach_yugabyte_copy_batch.test +# description: Test COPY batch splitting via pg_yb_rows_per_transaction on YugabyteDB +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# === Batch COPY: rows_per_transaction triggers CommitAndRestartCopy === +# Set a small batch size so 10k rows forces multiple commit cycles +statement ok +SET pg_yb_rows_per_transaction=1000 + +statement ok +DROP TABLE IF EXISTS yb.batch_test + +statement ok +CREATE TABLE yb.batch_test (id INTEGER PRIMARY KEY, data TEXT) + +# 10k rows with batch size 1000 = ~10 commit cycles +query I +INSERT INTO yb.batch_test SELECT g, 'batch_' || g FROM generate_series(1, 10000) t(g) +---- +10000 + +# Verify all rows landed despite mid-stream commits +query I +SELECT count(*) FROM yb.batch_test +---- +10000 + +query I +SELECT sum(id) FROM yb.batch_test +---- +50005000 + +query II +SELECT min(id), max(id) FROM yb.batch_test +---- +1 10000 + +# === Batch COPY with very small batch (stress the restart path) === +statement ok +SET pg_yb_rows_per_transaction=100 + +statement ok +DROP TABLE IF EXISTS yb.batch_small + +statement ok +CREATE TABLE yb.batch_small (id INTEGER PRIMARY KEY, val DOUBLE) + +query I +INSERT INTO yb.batch_small SELECT g, g * 1.5 FROM generate_series(1, 5000) t(g) +---- +5000 + +query I +SELECT count(*) FROM yb.batch_small +---- +5000 + +query I +SELECT sum(id) FROM yb.batch_small +---- +12502500 + +# === Batch COPY disabled (0 = single transaction) === +statement ok +SET pg_yb_rows_per_transaction=0 + +statement ok +DROP TABLE IF EXISTS yb.batch_disabled + +statement ok +CREATE TABLE yb.batch_disabled (id INTEGER PRIMARY KEY, name TEXT) + +query I +INSERT INTO yb.batch_disabled SELECT g, 'no_batch_' || g FROM generate_series(1, 3000) t(g) +---- +3000 + +query I +SELECT count(*) FROM yb.batch_disabled +---- +3000 + +# === Batch COPY combined with yb_disable_transactional_writes === +statement ok +SET pg_yb_rows_per_transaction=500 + +statement ok +SET pg_yb_disable_transactional_writes=true + +statement ok +DROP TABLE IF EXISTS yb.batch_notxn + +statement ok +CREATE TABLE yb.batch_notxn (id INTEGER PRIMARY KEY, data TEXT) + +query I +INSERT INTO yb.batch_notxn SELECT g, 'notxn_' || g FROM generate_series(1, 2000) t(g) +---- +2000 + +query I +SELECT count(*) FROM yb.batch_notxn +---- +2000 + +query I +SELECT sum(id) FROM yb.batch_notxn +---- +2001000 + +# === Cleanup === +statement ok +RESET pg_yb_rows_per_transaction + +statement ok +RESET pg_yb_disable_transactional_writes + +statement ok +DROP TABLE yb.batch_test + +statement ok +DROP TABLE yb.batch_small + +statement ok +DROP TABLE yb.batch_disabled + +statement ok +DROP TABLE yb.batch_notxn + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_edge_cases.test b/test/sql/storage/attach_yugabyte_edge_cases.test new file mode 100644 index 000000000..a97053643 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_edge_cases.test @@ -0,0 +1,193 @@ +# name: test/sql/storage/attach_yugabyte_edge_cases.test +# description: Edge cases and boundary conditions for YugabyteDB support +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# === Empty table scan (hash-partitioned) === +statement ok +CREATE TABLE yb.edge_empty_hash (id INTEGER PRIMARY KEY, data TEXT) + +query I +SELECT count(*) FROM yb.edge_empty_hash +---- +0 + +query II +SELECT * FROM yb.edge_empty_hash +---- + +statement ok +DROP TABLE yb.edge_empty_hash + +# === Single row table === +statement ok +CREATE TABLE yb.edge_single (id INTEGER PRIMARY KEY, val TEXT) + +query I +INSERT INTO yb.edge_single VALUES (1, 'only_row') +---- +1 + +query II +SELECT * FROM yb.edge_single +---- +1 only_row + +statement ok +DROP TABLE yb.edge_single + +# === Table with all NULL columns (except PK) === +statement ok +CREATE TABLE yb.edge_nulls (id INTEGER PRIMARY KEY, a TEXT, b INTEGER, c DOUBLE, d BOOLEAN) + +query I +INSERT INTO yb.edge_nulls VALUES (1, NULL, NULL, NULL, NULL), (2, NULL, NULL, NULL, NULL) +---- +2 + +query I +SELECT count(*) FROM yb.edge_nulls WHERE a IS NULL AND b IS NULL AND c IS NULL AND d IS NULL +---- +2 + +statement ok +DROP TABLE yb.edge_nulls + +# === Large text values === +statement ok +CREATE TABLE yb.edge_large_text (id INTEGER PRIMARY KEY, big_text TEXT) + +query I +INSERT INTO yb.edge_large_text VALUES (1, repeat('x', 100000)) +---- +1 + +query I +SELECT length(big_text) FROM yb.edge_large_text WHERE id = 1 +---- +100000 + +statement ok +DROP TABLE yb.edge_large_text + +# === Special characters in data === +statement ok +CREATE TABLE yb.edge_special (id INTEGER PRIMARY KEY, data TEXT) + +query I +INSERT INTO yb.edge_special VALUES + (1, 'single''quote'), + (2, 'back\slash'), + (3, 'new +line'), + (4, ''), + (5, 'tab here'), + (6, 'unicode: 日本語'), + (7, 'emoji: 🎉') +---- +7 + +query I +SELECT count(*) FROM yb.edge_special +---- +7 + +query I +SELECT data FROM yb.edge_special WHERE id = 1 +---- +single'quote + +query I +SELECT data FROM yb.edge_special WHERE id = 6 +---- +unicode: 日本語 + +statement ok +DROP TABLE yb.edge_special + +# === Boundary values for integer types === +statement ok +CREATE TABLE yb.edge_bounds (id INTEGER PRIMARY KEY, big BIGINT, small SMALLINT) + +query I +INSERT INTO yb.edge_bounds VALUES + (1, 9223372036854775807, 32767), + (2, -9223372036854775808, -32768), + (3, 0, 0) +---- +3 + +query III +SELECT * FROM yb.edge_bounds WHERE id = 1 +---- +1 9223372036854775807 32767 + +query III +SELECT * FROM yb.edge_bounds WHERE id = 2 +---- +2 -9223372036854775808 -32768 + +statement ok +DROP TABLE yb.edge_bounds + +# === Table with many columns === +statement ok +CREATE TABLE yb.edge_wide ( + id INTEGER PRIMARY KEY, + c1 TEXT, c2 TEXT, c3 TEXT, c4 TEXT, c5 TEXT, + c6 TEXT, c7 TEXT, c8 TEXT, c9 TEXT, c10 TEXT, + c11 INTEGER, c12 INTEGER, c13 INTEGER, c14 INTEGER, c15 INTEGER, + c16 DOUBLE, c17 DOUBLE, c18 DOUBLE, c19 DOUBLE, c20 DOUBLE +) + +query I +INSERT INTO yb.edge_wide VALUES ( + 1, 'a','b','c','d','e','f','g','h','i','j', + 1,2,3,4,5, 1.1,2.2,3.3,4.4,5.5 +) +---- +1 + +query I +SELECT count(*) FROM yb.edge_wide +---- +1 + +statement ok +DROP TABLE yb.edge_wide + +# === Concurrent attach to both colocated and non-colocated databases === +statement ok +ATTACH 'dbname=postgresscanner_colocated' AS ybc (TYPE POSTGRES) + +# Query both databases in same session +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT count(*) FROM ybc.coloc_test +---- +10000 + +# Cross-database join (DuckDB handles this by materializing both sides) +query I +SELECT count(*) FROM yb.hash_test h JOIN ybc.coloc_test c ON h.id = c.id +---- +10000 + +statement ok +DETACH ybc + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_errors.test b/test/sql/storage/attach_yugabyte_errors.test new file mode 100644 index 000000000..2daf061d9 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_errors.test @@ -0,0 +1,110 @@ +# name: test/sql/storage/attach_yugabyte_errors.test +# description: Test error handling and schema mismatch resilience on YugabyteDB +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# === Non-existent table — should error gracefully, not crash === +statement error +SELECT * FROM yb.this_table_does_not_exist +---- +this_table_does_not_exist + +# === Non-existent schema === +statement error +SELECT * FROM yb.fake_schema.fake_table +---- + +# === Create, query, drop, query again — stale catalog reference === +statement ok +CREATE TABLE yb.ephemeral_test (id INTEGER PRIMARY KEY, data TEXT) + +query I +INSERT INTO yb.ephemeral_test SELECT g, 'data_' || g FROM generate_series(1, 100) t(g) +---- +100 + +query I +SELECT count(*) FROM yb.ephemeral_test +---- +100 + +statement ok +DROP TABLE yb.ephemeral_test + +statement error +SELECT * FROM yb.ephemeral_test +---- + +# === Re-create with different schema — type change === +statement ok +CREATE TABLE yb.ephemeral_test (id TEXT PRIMARY KEY, value DOUBLE) + +query I +INSERT INTO yb.ephemeral_test VALUES ('key1', 3.14), ('key2', 2.71) +---- +2 + +query II +SELECT id, value FROM yb.ephemeral_test ORDER BY id +---- +key1 3.14 +key2 2.71 + +statement ok +DROP TABLE yb.ephemeral_test + +# === Invalid connection string === +statement error +ATTACH 'dbname=this_database_does_not_exist host=localhost port=5433' AS bad_yb (TYPE POSTGRES) +---- + +# === Empty table scans === +statement ok +CREATE TABLE yb.empty_hash (id INTEGER PRIMARY KEY, data TEXT) + +query I +SELECT count(*) FROM yb.empty_hash +---- +0 + +query II +SELECT * FROM yb.empty_hash +---- + +statement ok +DROP TABLE yb.empty_hash + +# === Table with no primary key (heap-like on YugabyteDB) === +statement ok +CREATE TABLE yb.no_pk (a INTEGER, b TEXT) + +query I +INSERT INTO yb.no_pk SELECT g, 'val_' || g FROM generate_series(1, 100) t(g) +---- +100 + +query I +SELECT count(*) FROM yb.no_pk +---- +100 + +statement ok +DROP TABLE yb.no_pk + +# === Verify the connection is still healthy after all error cases === +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_parallel.test_slow b/test/sql/storage/attach_yugabyte_parallel.test_slow new file mode 100644 index 000000000..79b239a29 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_parallel.test_slow @@ -0,0 +1,161 @@ +# name: test/sql/storage/attach_yugabyte_parallel.test_slow +# description: Test parallel hash-code scanning on YugabyteDB with substantial data volumes +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +# ATTACH first to ensure extension settings are registered +statement ok +ATTACH 'dbname=postgresscanner' AS yb_init (TYPE POSTGRES) + +statement ok +DETACH yb_init + +statement ok +SET threads=4 + +statement ok +SET pg_yb_parallel_scan=true + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# === 100k-row hash-partitioned table — exercises yb_hash_code() BETWEEN ranges === + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +# Order-independent aggregation verifies no rows lost or duplicated across hash ranges +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +query I +SELECT sum(value) FROM yb.hash_test +---- +50000500000 + +# Filter pushdown combined with hash-code parallel scan +query I +SELECT count(*) FROM yb.hash_test WHERE id > 50000 +---- +50000 + +query I +SELECT sum(id) FROM yb.hash_test WHERE id <= 1000 +---- +500500 + +query I +SELECT count(*) FROM yb.hash_test WHERE name LIKE 'row_99%' +---- +1111 + +# === Wide table — multiple column types through parallel COPY path === + +query I +SELECT count(*) FROM yb.wide_test +---- +50000 + +query I +SELECT count(*) FROM yb.wide_test WHERE col_bool = true +---- +25000 + +query I +SELECT count(*) FROM yb.wide_test WHERE col_int > 2500000000 +---- +25000 + +# === Multi-column hash key parallel scan — 10k rows, compound partition key === + +query I +SELECT count(*) FROM yb.multi_hash +---- +10000 + +query II +SELECT region, count(*) FROM yb.multi_hash GROUP BY region ORDER BY region +---- +region_0 2000 +region_1 2000 +region_2 2000 +region_3 2000 +region_4 2000 + +# Cross-table join that forces parallel scan on both sides +query I +SELECT count(*) FROM yb.hash_test h JOIN yb.multi_hash m ON h.id = m.id +---- +10000 + +# === Range-partitioned tables — must NOT use hash-code parallel scan === +# These have yb_num_hash_key_columns = 0, so they fall back to single-thread + +query I +SELECT count(*) FROM yb.range_single +---- +10000 + +query I +SELECT sum(id) FROM yb.range_single +---- +50005000 + +query I +SELECT count(*) FROM yb.range_ts WHERE sensor_id = 0 +---- +200 + +query I +SELECT count(*) FROM yb.range_desc +---- +15000 + +# === Vary thread counts to exercise different hash range splits === + +statement ok +DETACH yb + +statement ok +SET threads=8 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +statement ok +DETACH yb + +statement ok +SET threads=1 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# Single-thread should still work correctly +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_pool.test b/test/sql/storage/attach_yugabyte_pool.test new file mode 100644 index 000000000..e7cd2ea83 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_pool.test @@ -0,0 +1,98 @@ +# name: test/sql/storage/attach_yugabyte_pool.test +# description: Test connection pool behavior on YugabyteDB +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +# === Default pool settings === +statement ok +SET pg_connection_cache=true + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb + +# === Pool with limit=1 (forces serial access) === +statement ok +SET pg_connection_limit=1 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT count(*) FROM yb.wide_test +---- +50000 + +statement ok +DETACH yb + +# === Pool with high limit === +statement ok +SET pg_connection_limit=1000 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb + +# === Pool disabled === +statement ok +SET pg_connection_cache=false + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +# Multiple queries without pool — each gets fresh connection through Reset path +query I +SELECT count(*) FROM yb.wide_test +---- +50000 + +query I +SELECT count(*) FROM yb.multi_hash +---- +10000 + +query I +SELECT count(*) FROM yb.range_single +---- +10000 + +statement ok +DETACH yb + +# === Restore defaults === +statement ok +RESET pg_connection_cache + +statement ok +RESET pg_connection_limit diff --git a/test/sql/storage/attach_yugabyte_range.test b/test/sql/storage/attach_yugabyte_range.test new file mode 100644 index 000000000..3d712e5fa --- /dev/null +++ b/test/sql/storage/attach_yugabyte_range.test @@ -0,0 +1,167 @@ +# name: test/sql/storage/attach_yugabyte_range.test +# description: Test range-partitioned tables on YugabyteDB — must NOT use yb_hash_code parallel scan +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +statement ok +SET threads=4 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# === range_single: ASC range key === +# yb_num_hash_key_columns = 0, so no hash-code parallel scan. +# Must still return correct results via single-threaded scan. + +query I +SELECT count(*) FROM yb.range_single +---- +10000 + +query II +SELECT min(id), max(id) FROM yb.range_single +---- +1 10000 + +query I +SELECT sum(id) FROM yb.range_single +---- +50005000 + +query III +SELECT id, name, value FROM yb.range_single WHERE id = 5000 +---- +5000 range_5000 50000 + +# Filter pushdown on range key +query I +SELECT count(*) FROM yb.range_single WHERE id > 7500 +---- +2500 + +query I +SELECT count(*) FROM yb.range_single WHERE id BETWEEN 2500 AND 5000 +---- +2501 + +# === range_ts: compound ASC range key (timeseries pattern) === + +query I +SELECT count(*) FROM yb.range_ts +---- +20000 + +query I +SELECT count(DISTINCT sensor_id) FROM yb.range_ts +---- +100 + +query I +SELECT count(*) FROM yb.range_ts WHERE sensor_id = 0 +---- +200 + +# Range scan with timestamp filter (data is 1..20000 seconds past midnight) +query I +SELECT count(*) FROM yb.range_ts WHERE ts > '2024-01-01 02:46:40' +---- +10000 + +# Aggregation correctness +query I +SELECT sum(sensor_id) FROM yb.range_ts +---- +990000 + +# === range_desc: DESC ordering === +# Tests that DESC range keys don't confuse the scan path + +query I +SELECT count(*) FROM yb.range_desc +---- +15000 + +query II +SELECT min(id), max(id) FROM yb.range_desc +---- +1 15000 + +query I +SELECT sum(id) FROM yb.range_desc +---- +112507500 + +query I +SELECT count(*) FROM yb.range_desc WHERE id <= 5000 +---- +5000 + +# === Contrast: hash table in same session should still use parallel scan === +# This proves the scan path selection is per-table, not per-connection + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +# === Mix range and hash in same query === + +query I +SELECT count(*) +FROM yb.range_single r +JOIN yb.hash_test h ON r.id = h.id +---- +10000 + +query I +SELECT count(*) +FROM yb.range_single r +JOIN yb.multi_hash m ON r.id = m.id +---- +10000 + +# === Write to range table === +statement ok +DROP TABLE IF EXISTS yb.range_write_test + +statement ok +CREATE TABLE yb.range_write_test ( + ts TIMESTAMP, + id INTEGER, + data TEXT, + PRIMARY KEY (ts, id) +) + +query I +INSERT INTO yb.range_write_test +SELECT '2024-01-01'::TIMESTAMP + (g || ' seconds')::INTERVAL, g, 'data_' || g +FROM generate_series(1, 5000) t(g) +---- +5000 + +query I +SELECT count(*) FROM yb.range_write_test +---- +5000 + +query I +SELECT sum(id) FROM yb.range_write_test +---- +12502500 + +statement ok +DROP TABLE yb.range_write_test + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_reconnect.test b/test/sql/storage/attach_yugabyte_reconnect.test new file mode 100644 index 000000000..709e1f6d4 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_reconnect.test @@ -0,0 +1,116 @@ +# name: test/sql/storage/attach_yugabyte_reconnect.test +# description: Test attach/detach/re-attach cycles and connection pool resilience on YugabyteDB +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +# === Cycle 1: attach, query, detach === +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb + +# === Cycle 2: re-attach exercises connection pool Reset path === +# On YugabyteDB, DISCARD ALL would clobber session state — the replacement should preserve it +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +# Query a different table to exercise catalog reload after reconnect +query I +SELECT count(*) FROM yb.wide_test +---- +50000 + +statement ok +DETACH yb + +# === Cycle 3: verify session state survives another pool recycle === +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.test +---- +4 + +query I +SELECT count(*) FROM yb.multi_hash +---- +10000 + +statement ok +DETACH yb + +# === Cycle 4: rapid attach/detach without queries === +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +statement ok +DETACH yb + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +statement ok +DETACH yb + +# === Cycle 5: verify data after rapid cycling === +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +statement ok +DETACH yb + +# === Cycle 6: attach with connection pool settings === +statement ok +SET pg_connection_cache=true + +statement ok +SET pg_connection_limit=2 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb + +# === Cycle 7: re-attach after pool limit change === +statement ok +SET pg_connection_limit=1000 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_scan_modes.test b/test/sql/storage/attach_yugabyte_scan_modes.test new file mode 100644 index 000000000..31f71c786 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_scan_modes.test @@ -0,0 +1,247 @@ +# name: test/sql/storage/attach_yugabyte_scan_modes.test +# description: Test all scan mode combinations on YugabyteDB — parallel vs serial, hash vs range vs colocated +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +# === Mode 1: Default (pg_yb_parallel_scan=false) — all tables scan single-threaded === +statement ok +SET threads=4 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# Hash table with parallel disabled — must still return correct results +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +# Range table — always single-threaded +query I +SELECT count(*) FROM yb.range_single +---- +10000 + +# Multi-hash — single-threaded without parallel enabled +query I +SELECT count(*) FROM yb.multi_hash +---- +10000 + +statement ok +DETACH yb + +# === Mode 2: parallel enabled, threads=1 — degrades to single-thread gracefully === +statement ok +SET pg_yb_parallel_scan=true + +statement ok +SET threads=1 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +statement ok +DETACH yb + +# === Mode 3: parallel enabled, threads=2 === +statement ok +SET threads=2 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +# Range table still single-threaded even with parallel enabled +query I +SELECT count(*) FROM yb.range_single +---- +10000 + +statement ok +DETACH yb + +# === Mode 4: parallel enabled, high thread count === +statement ok +SET threads=16 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT sum(id) FROM yb.hash_test +---- +5000050000 + +# Compound hash key with many threads +query I +SELECT count(*) FROM yb.multi_hash +---- +10000 + +query II +SELECT region, count(*) FROM yb.multi_hash GROUP BY region ORDER BY region +---- +region_0 2000 +region_1 2000 +region_2 2000 +region_3 2000 +region_4 2000 + +statement ok +DETACH yb + +# === Mode 5: filter pushdown combined with hash-code scan === +statement ok +SET threads=4 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# Equality filter +query I +SELECT count(*) FROM yb.hash_test WHERE id = 42 +---- +1 + +# Range filter +query I +SELECT count(*) FROM yb.hash_test WHERE id BETWEEN 10000 AND 20000 +---- +10001 + +# String filter +query I +SELECT count(*) FROM yb.hash_test WHERE name LIKE 'row_1%' +---- +11112 + +# Compound filter +query I +SELECT count(*) FROM yb.hash_test WHERE id > 50000 AND name LIKE 'row_9%' +---- +10000 + +# NULL filter (no NULLs in hash_test) +query I +SELECT count(*) FROM yb.hash_test WHERE name IS NULL +---- +0 + +# Wide table filters with parallel +query I +SELECT count(*) FROM yb.wide_test WHERE col_bool = false AND col_int < 1000000 +---- +5 + +statement ok +DETACH yb + +# === Mode 6: parallel enabled but write transaction (INSERT...SELECT) === +# Source scan should be single-threaded because the transaction is read-write +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +statement ok +DROP TABLE IF EXISTS yb.scan_mode_write + +statement ok +CREATE TABLE yb.scan_mode_write (id INTEGER PRIMARY KEY, name TEXT, value INTEGER) + +query I +INSERT INTO yb.scan_mode_write SELECT id, name, value FROM yb.hash_test WHERE id <= 5000 +---- +5000 + +query I +SELECT count(*) FROM yb.scan_mode_write +---- +5000 + +query I +SELECT sum(id) FROM yb.scan_mode_write +---- +12502500 + +statement ok +DROP TABLE yb.scan_mode_write + +statement ok +DETACH yb + +# === Mode 7: parallel scan + colocated table (0 tablets → must fallback to serial) === +statement ok +ATTACH 'dbname=postgresscanner_colocated' AS ybc (TYPE POSTGRES) + +query I +SELECT count(*) FROM ybc.coloc_test +---- +10000 + +query I +SELECT sum(id) FROM ybc.coloc_test +---- +50005000 + +statement ok +DETACH ybc + +# === Mode 8: tserver probe timeout setting === +statement ok +SET pg_yb_tserver_probe_timeout=1 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb + +# === Cleanup === +statement ok +RESET pg_yb_parallel_scan + +statement ok +RESET pg_yb_tserver_probe_timeout + +statement ok +RESET threads diff --git a/test/sql/storage/attach_yugabyte_settings.test b/test/sql/storage/attach_yugabyte_settings.test new file mode 100644 index 000000000..49be2ed1a --- /dev/null +++ b/test/sql/storage/attach_yugabyte_settings.test @@ -0,0 +1,88 @@ +# name: test/sql/storage/attach_yugabyte_settings.test +# description: Test YugabyteDB-specific settings and their effect on connections +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +# First ATTACH to ensure extension is fully loaded and settings registered +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +statement ok +DETACH yb + +# Now verify YB-specific settings are registered and can be set +statement ok +SET pg_yb_parallel_scan=true + +statement ok +SET pg_yb_tserver_probe_timeout=5 + +statement ok +SET pg_yb_rows_per_transaction=5000 + +statement ok +SET pg_yb_disable_transactional_writes=true + +# Verify they can be read back +query I +SELECT current_setting('pg_yb_parallel_scan') +---- +true + +query I +SELECT current_setting('pg_yb_tserver_probe_timeout') +---- +5 + +query I +SELECT current_setting('pg_yb_rows_per_transaction') +---- +5000 + +query I +SELECT current_setting('pg_yb_disable_transactional_writes') +---- +true + +# Verify settings don't break attach/query flow +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb + +# Reset to defaults +statement ok +RESET pg_yb_parallel_scan + +statement ok +RESET pg_yb_tserver_probe_timeout + +statement ok +RESET pg_yb_rows_per_transaction + +statement ok +RESET pg_yb_disable_transactional_writes + +# Verify defaults restored +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_stress.test_slow b/test/sql/storage/attach_yugabyte_stress.test_slow new file mode 100644 index 000000000..5c9f52c5e --- /dev/null +++ b/test/sql/storage/attach_yugabyte_stress.test_slow @@ -0,0 +1,147 @@ +# name: test/sql/storage/attach_yugabyte_stress.test_slow +# description: Stress test — large scans, bulk writes, cross-table operations on YugabyteDB +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +# ATTACH first to ensure extension settings are registered +statement ok +ATTACH 'dbname=postgresscanner' AS yb_init (TYPE POSTGRES) + +statement ok +DETACH yb_init + +statement ok +SET pg_yb_parallel_scan=true + +statement ok +SET threads=4 + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# === Full table scan with aggregation across all pre-loaded tables === + +query I +SELECT count(*) FROM yb.hash_test +---- +100000 + +query I +SELECT count(*) FROM yb.wide_test +---- +50000 + +query I +SELECT count(*) FROM yb.multi_hash +---- +10000 + +query I +SELECT count(*) FROM yb.range_single +---- +10000 + +query I +SELECT count(*) FROM yb.range_ts +---- +20000 + +query I +SELECT count(*) FROM yb.range_desc +---- +15000 + +# === Total row count across all tables === +query I +SELECT (SELECT count(*) FROM yb.hash_test) + + (SELECT count(*) FROM yb.wide_test) + + (SELECT count(*) FROM yb.multi_hash) + + (SELECT count(*) FROM yb.range_single) + + (SELECT count(*) FROM yb.range_ts) + + (SELECT count(*) FROM yb.range_desc) +---- +205000 + +# === Large bulk write and read-back === +statement ok +DROP TABLE IF EXISTS yb.stress_write + +statement ok +CREATE TABLE yb.stress_write (id INTEGER PRIMARY KEY, payload TEXT, value BIGINT) + +# Write 50k rows — exercises multiple COPY batches with default pg_yb_rows_per_transaction=10000 +query I +INSERT INTO yb.stress_write SELECT g, 'payload_' || g, g * 100000::BIGINT FROM generate_series(1, 50000) t(g) +---- +50000 + +query I +SELECT count(*) FROM yb.stress_write +---- +50000 + +query I +SELECT sum(id) FROM yb.stress_write +---- +1250025000 + +query I +SELECT sum(value) FROM yb.stress_write +---- +125002500000000 + +# === Complex query with JOIN, GROUP BY, ORDER BY === +query I +SELECT count(*) FROM yb.hash_test h JOIN yb.stress_write s ON h.id = s.id +---- +50000 + +query II rowsort +SELECT h.name, s.payload +FROM yb.hash_test h JOIN yb.stress_write s ON h.id = s.id +WHERE h.id IN (1, 50000) +---- +row_1 payload_1 +row_50000 payload_50000 + +# === Subquery with hash-code parallel scan === +query I +SELECT count(*) FROM yb.hash_test WHERE id IN (SELECT id FROM yb.multi_hash) +---- +10000 + +# === UNION across hash and range tables === +query I +SELECT count(*) FROM ( + SELECT id FROM yb.hash_test WHERE id <= 1000 + UNION ALL + SELECT id FROM yb.range_single WHERE id <= 1000 +) +---- +2000 + +# === Window function over parallel scan === +query I +SELECT count(*) FROM ( + SELECT id, ROW_NUMBER() OVER (ORDER BY id) as rn + FROM yb.hash_test + WHERE id <= 100 +) +---- +100 + +# === Cleanup === +statement ok +DROP TABLE yb.stress_write + +statement ok +RESET pg_yb_parallel_scan + +statement ok +DETACH yb diff --git a/test/sql/storage/attach_yugabyte_write.test b/test/sql/storage/attach_yugabyte_write.test new file mode 100644 index 000000000..efd97fc59 --- /dev/null +++ b/test/sql/storage/attach_yugabyte_write.test @@ -0,0 +1,164 @@ +# name: test/sql/storage/attach_yugabyte_write.test +# description: Test INSERT/UPDATE/DELETE and bulk COPY write path on YugabyteDB +# group: [storage] + +require postgres_scanner + +require-env YUGABYTE_TEST_DATABASE_AVAILABLE + +statement ok +PRAGMA enable_verification + +statement ok +ATTACH 'dbname=postgresscanner' AS yb (TYPE POSTGRES) + +# === Bulk insert — exercises COPY write path === +statement ok +DROP TABLE IF EXISTS yb.write_test + +statement ok +CREATE TABLE yb.write_test (id INTEGER PRIMARY KEY, name TEXT, value DOUBLE) + +query I +INSERT INTO yb.write_test SELECT g, 'name_' || g, g * 1.5 FROM generate_series(1, 10000) t(g) +---- +10000 + +query I +SELECT count(*) FROM yb.write_test +---- +10000 + +query III +SELECT id, name, value FROM yb.write_test WHERE id = 5000 +---- +5000 name_5000 7500.0 + +# Checksum to catch data corruption in COPY path +query I +SELECT sum(id) FROM yb.write_test +---- +50005000 + +# === INSERT INTO from another YB table (cross-table COPY) === +statement ok +DROP TABLE IF EXISTS yb.write_copy + +statement ok +CREATE TABLE yb.write_copy (id INTEGER PRIMARY KEY, name TEXT, value DOUBLE) + +query I +INSERT INTO yb.write_copy SELECT * FROM yb.write_test WHERE id <= 1000 +---- +1000 + +query I +SELECT count(*) FROM yb.write_copy +---- +1000 + +query I +SELECT sum(id) FROM yb.write_copy +---- +500500 + +# NOTE: UPDATE/DELETE not tested here — the extension's UPDATE/DELETE path uses ctid +# which YugabyteDB does not support. That requires a separate fix to use primary key +# based deletes instead of ctid-based deletes. + +# === Multi-type table write === +statement ok +DROP TABLE IF EXISTS yb.write_types + +statement ok +CREATE TABLE yb.write_types ( + id INTEGER PRIMARY KEY, + col_text TEXT, + col_bigint BIGINT, + col_double DOUBLE, + col_bool BOOLEAN, + col_date DATE +) + +query I +INSERT INTO yb.write_types +SELECT g, 'text_' || g, g * 100000::BIGINT, g * 3.14, (g % 2 = 0), '2024-01-01'::DATE + INTERVAL (g) DAY +FROM generate_series(1, 5000) t(g) +---- +5000 + +query I +SELECT count(*) FROM yb.write_types +---- +5000 + +query I +SELECT count(*) FROM yb.write_types WHERE col_bool = true +---- +2500 + +# === yb_disable_transactional_writes COPY path === +# When enabled, BeginCopyTo sends SET yb_disable_transactional_writes = true +# This disables per-row transactions for bulk loads (faster but no atomicity) +statement ok +SET pg_yb_disable_transactional_writes=true + +statement ok +DROP TABLE IF EXISTS yb.write_notxn + +statement ok +CREATE TABLE yb.write_notxn (id INTEGER PRIMARY KEY, data TEXT) + +query I +INSERT INTO yb.write_notxn SELECT g, 'notxn_' || g FROM generate_series(1, 5000) t(g) +---- +5000 + +query I +SELECT count(*) FROM yb.write_notxn +---- +5000 + +query I +SELECT sum(id) FROM yb.write_notxn +---- +12502500 + +statement ok +RESET pg_yb_disable_transactional_writes + +# === Verify normal COPY still works after resetting the flag === +statement ok +DROP TABLE IF EXISTS yb.write_after_reset + +statement ok +CREATE TABLE yb.write_after_reset (id INTEGER PRIMARY KEY, val INTEGER) + +query I +INSERT INTO yb.write_after_reset SELECT g, g * 10 FROM generate_series(1, 1000) t(g) +---- +1000 + +query I +SELECT count(*) FROM yb.write_after_reset +---- +1000 + +# === Cleanup === +statement ok +DROP TABLE yb.write_test + +statement ok +DROP TABLE yb.write_copy + +statement ok +DROP TABLE yb.write_types + +statement ok +DROP TABLE yb.write_notxn + +statement ok +DROP TABLE yb.write_after_reset + +statement ok +DETACH yb