From 80589160b7402b8b289f7bf7764e0177186ad7d0 Mon Sep 17 00:00:00 2001 From: Fan Yang Date: Thu, 26 Dec 2024 12:23:09 +0000 Subject: [PATCH] test: migrate psql test cases to BATS --- .github/workflows/bats-tests.yml | 2 +- .github/workflows/psql.yml | 56 ------------- test/bats/postgres/copy_tests.bats | 123 +++++++++++++++++++++++++++++ 3 files changed, 124 insertions(+), 57 deletions(-) delete mode 100644 .github/workflows/psql.yml create mode 100644 test/bats/postgres/copy_tests.bats diff --git a/.github/workflows/bats-tests.yml b/.github/workflows/bats-tests.yml index 31aacb1c..2d064a34 100644 --- a/.github/workflows/bats-tests.yml +++ b/.github/workflows/bats-tests.yml @@ -32,7 +32,7 @@ jobs: run: | go get . - pip3 install "sqlglot[rs]" + pip3 install "sqlglot[rs]" pyarrow pandas curl -LJO https://github.com/duckdb/duckdb/releases/download/v1.1.3/duckdb_cli-linux-amd64.zip unzip duckdb_cli-linux-amd64.zip diff --git a/.github/workflows/psql.yml b/.github/workflows/psql.yml deleted file mode 100644 index 75c08a04..00000000 --- a/.github/workflows/psql.yml +++ /dev/null @@ -1,56 +0,0 @@ -name: psql Test - -on: - push: - branches: [ "main" ] - pull_request: - branches: [ "main" ] - -jobs: - - build: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Set up Go - uses: actions/setup-go@v5 - with: - go-version: '1.23' - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.10' - - - name: Install system packages - uses: awalsh128/cache-apt-pkgs-action@latest - with: - packages: postgresql-client - version: 1.0 - - - name: Install dependencies - run: | - go get . - - pip3 install "sqlglot[rs]" pyarrow pandas - - curl -LJO https://github.com/duckdb/duckdb/releases/latest/download/duckdb_cli-linux-amd64.zip - unzip duckdb_cli-linux-amd64.zip - chmod +x duckdb - sudo mv duckdb /usr/local/bin - - - name: Build - run: go build -v - - - name: Start MyDuck Server - run: | - ./myduckserver & - sleep 5 - - - name: Run the SQL scripts - run: | - # for each SQL script in the `pgtest/psql` directory (recursively) - for f in pgtest/psql/**/*.sql; do - psql -h 127.0.0.1 -U postgres -d postgres -f $f - done diff --git a/test/bats/postgres/copy_tests.bats b/test/bats/postgres/copy_tests.bats new file mode 100644 index 00000000..d892533c --- /dev/null +++ b/test/bats/postgres/copy_tests.bats @@ -0,0 +1,123 @@ +#!/usr/bin/env bats + +load helper + +setup() { + psql_exec "CREATE SCHEMA IF NOT EXISTS test_copy;" + psql_exec "USE test_copy;" + psql_exec "CREATE TABLE t (a int, b text, c float);" + psql_exec "INSERT INTO t VALUES (1, 'one', 1.1), (2, 'two', 2.2), (3, 'three', 3.3);" +} + +teardown() { + psql_exec "DROP SCHEMA IF EXISTS test_copy CASCADE;" + rm -f test_*.{csv,parquet,arrow,db} 2>/dev/null +} + +@test "copy with csv format" { + # Test copy to CSV file + tmpfile=$(mktemp) + run psql_exec "\copy t TO '${tmpfile}' (FORMAT CSV, HEADER false);" + [ "$status" -eq 0 ] + run cat "${tmpfile}" + [ "$status" -eq 0 ] + [[ "${lines[0]}" == "1,one,1.1" ]] + rm "${tmpfile}" + + # Test copy from CSV with headers + run psql_exec_stdin <<-EOF + CREATE TABLE csv_test (a int, b text); + \copy csv_test FROM 'pgtest/testdata/basic.csv' (FORMAT CSV, HEADER); + \copy csv_test FROM 'pgtest/testdata/basic.csv' WITH DELIMITER ',' CSV HEADER; + SELECT COUNT(*) FROM csv_test; +EOF + [ "$status" -eq 0 ] + [[ "${output}" != "0" ]] + + # Test various CSV output formats + run psql_exec_stdin <<-EOF + \copy t TO STDOUT; + \copy t (a, b) TO STDOUT (FORMAT CSV); + \copy t TO STDOUT (FORMAT CSV, HEADER false, DELIMITER '|'); + \copy (SELECT a * a, b, c + a FROM t) TO STDOUT (FORMAT CSV, HEADER false, DELIMITER '|'); +EOF + [ "$status" -eq 0 ] + [ "${#lines[@]}" -ge 12 ] + [[ "${lines[0]}" == "1,one,1.1" ]] + [[ "${lines[3]}" == "1,one" ]] + [[ "${lines[6]}" == "1|one|1.1" ]] + [[ "${lines[9]}" == "1|one|2.1" ]] +} + +@test "copy with parquet format" { + # Test basic COPY TO PARQUET + tmpfile=$(mktemp).parquet + run psql_exec "\copy t TO '${tmpfile}' (FORMAT PARQUET);" + [ "$status" -eq 0 ] + run duckdb -c "SELECT COUNT(*) FROM '${tmpfile}'" + [ "$status" -eq 0 ] + [[ "${output}" == "3" ]] + rm "${tmpfile}" + + # Test with column selection + outfile="test_cols.parquet" + run psql_exec "\copy t (a, b) TO '${outfile}' (FORMAT PARQUET);" + [ "$status" -eq 0 ] + run duckdb -c "SELECT COUNT(*) FROM '${outfile}'" + [ "$status" -eq 0 ] + [[ "${output}" == "3" ]] + + # Test with transformed data + outfile="test_transform.parquet" + run psql_exec "(SELECT a * a, b, c + a FROM t) TO '${outfile}' (FORMAT PARQUET);" + [ "$status" -eq 0 ] + run duckdb -c "SELECT COUNT(*) FROM '${outfile}'" + [ "$status" -eq 0 ] + [[ "${output}" == "3" ]] +} + +@test "copy with arrow format" { + # Test basic COPY TO ARROW + outfile="test_out.arrow" + run psql_exec "\copy t TO '${outfile}' (FORMAT ARROW);" + [ "$status" -eq 0 ] + run python3 -c "import pyarrow as pa; reader = pa.ipc.open_stream('${outfile}'); print(len(reader.read_all()))" + [ "$status" -eq 0 ] + [[ "${output}" == "3" ]] + + # Test with column selection + outfile="test_cols.arrow" + run psql_exec "\copy t (a, b) TO '${outfile}' (FORMAT ARROW);" + [ "$status" -eq 0 ] + run python3 -c "import pyarrow as pa; reader = pa.ipc.open_stream('${outfile}'); print(len(reader.read_all()))" + [ "$status" -eq 0 ] + [[ "${output}" == "3" ]] + + # Test with transformed data + outfile="test_transform.arrow" + run psql_exec "\copy (SELECT a * a, b, c + a FROM t) TO '${outfile}' (FORMAT ARROW);" + [ "$status" -eq 0 ] + run python3 -c "import pyarrow as pa; reader = pa.ipc.open_stream('${outfile}'); print(len(reader.read_all()))" + [ "$status" -eq 0 ] + [[ "${output}" == "3" ]] + + # Test COPY FROM ARROW + run psql_exec_stdin <<-EOF + CREATE TABLE arrow_test (a int, b text, c float); + \copy arrow_test FROM '${outfile}' (FORMAT ARROW); + SELECT COUNT(*) FROM arrow_test; +EOF + [ "$status" -eq 0 ] + [[ "${output}" == "3" ]] +} + +@test "copy from database" { + run psql_exec_stdin <<-EOF + CREATE TABLE db_test (a int, b text); + INSERT INTO db_test VALUES (1, 'a'), (2, 'b'), (3, 'c'); + ATTACH 'test_copy.db' AS tmp; + COPY FROM DATABASE mysql TO tmp; + DETACH tmp; +EOF + [ "$status" -eq 0 ] +}