diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6456b6a..c0d8ab3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -36,6 +36,48 @@ jobs: - run: cargo clippy --all-targets -- -D warnings - run: cargo fmt --all -- --check + emulator: + name: BigQuery emulator integration + runs-on: ubuntu-latest + services: + bigquery: + image: ghcr.io/goccy/bigquery-emulator:latest + env: + BIGQUERY_EMULATOR_PROJECT: test + BIGQUERY_EMULATOR_DATASET: ds + ports: + - 9050:9050 + steps: + - uses: actions/checkout@v7 + - uses: dtolnay/rust-toolchain@stable + - name: Wait for the emulator + run: | + for attempt in $(seq 1 30); do + curl -sf http://localhost:9050/bigquery/v2/projects/test/datasets && exit 0 + sleep 2 + done + exit 1 + - run: cargo test --test emulator --test emulator_metadata + env: + BIGQUERY_EMULATOR_HOST: localhost:9050 + + explain: + name: EXPLAIN parser package + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: pnpm/action-setup@v4 + with: + version: 9 + - uses: actions/setup-node@v4 + with: + node-version: 20 + cache: pnpm + cache-dependency-path: explain/pnpm-lock.yaml + - run: pnpm --dir explain install --frozen-lockfile + - run: pnpm --dir explain typecheck + - run: pnpm --dir explain test + validate-manifest: name: Validate .tabularium manifest runs-on: ubuntu-latest diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index d968509..de6509f 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -18,14 +18,23 @@ jobs: with: node-version: 20 cache: pnpm - cache-dependency-path: ui/pnpm-lock.yaml + cache-dependency-path: | + ui/pnpm-lock.yaml + explain/pnpm-lock.yaml - run: pnpm --dir ui install --frozen-lockfile - run: pnpm --dir ui typecheck - run: pnpm --dir ui build + - run: pnpm --dir explain install --frozen-lockfile + - run: pnpm --dir explain typecheck + - run: pnpm --dir explain test - uses: actions/upload-artifact@v7 with: name: ui-dist path: ui/dist/ + - uses: actions/upload-artifact@v7 + with: + name: explain-dist + path: explain/dist/index.iife.js build: name: ${{ matrix.platform-label }} @@ -74,22 +83,29 @@ jobs: with: name: ui-dist path: ui/dist/ + - uses: actions/download-artifact@v8 + with: + name: explain-dist + path: explain/dist/ - name: Package on Unix if: runner.os != 'Windows' run: | - mkdir -p staging/ui/dist + mkdir -p staging/ui/dist staging/explain/dist cp target/${{ matrix.target }}/release/bigquery-plugin${{ matrix.binary-suffix }} staging/ cp .tabularium staging/ cp ui/dist/*.js staging/ui/dist/ + cp explain/dist/index.iife.js staging/explain/dist/ (cd staging && zip -r ../bigquery-plugin-${{ matrix.platform-label }}.zip .) - name: Package on Windows if: runner.os == 'Windows' shell: pwsh run: | New-Item -ItemType Directory -Force -Path "staging\ui\dist" | Out-Null + New-Item -ItemType Directory -Force -Path "staging\explain\dist" | Out-Null Copy-Item "target\${{ matrix.target }}\release\bigquery-plugin${{ matrix.binary-suffix }}" staging Copy-Item ".tabularium" staging Copy-Item "ui\dist\*.js" "staging\ui\dist" + Copy-Item "explain\dist\index.iife.js" "staging\explain\dist" Compress-Archive -Path "staging\*" -DestinationPath "bigquery-plugin-${{ matrix.platform-label }}.zip" - name: Smoke test on Unix if: runner.os != 'Windows' && !matrix.cross diff --git a/.gitignore b/.gitignore index 603a9e6..46c02ad 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ ui/dist/ .DS_Store .idea/ .vscode/ +explain/node_modules/ +explain/dist/ diff --git a/.tabularium b/.tabularium index 159649e..a1be161 100644 --- a/.tabularium +++ b/.tabularium @@ -1,13 +1,27 @@ { "$schema": "https://registry.tabularis.dev/manifest.schema.json?kind=driver", - "name": "bigquery", - "version": "0.1.0", - "description": "Google BigQuery driver for Tabularis with GoogleSQL, datasets, views, routines, editable results and cost controls.", + "id": "bigquery", + "name": "BigQuery", + "version": "0.2.0", + "description": "Google BigQuery driver for Tabularis with multi-project browsing, cancellable paged queries, sessions and transactions, query plans, GoogleSQL, materialized views, routine management, typed row editing and cost controls.", + "category": "database", + "tags": ["driver", "bigquery", "google-cloud", "gcp", "sql", "googlesql", "analytics", "data-warehouse", "database-driver"], "kind": "driver", "engine": "bigquery", + "explain_parsers": [ + { + "engine": "bigquery", + "format": "bigquery-job-json", + "label": "BigQuery job plan (JSON)", + "module": "explain/dist/index.iife.js" + } + ], "paradigms": ["relational", "columnar"], "license": "Apache-2.0", + "readme": "README.md", + "min_runtime_version": "0.29.0", "homepage": "https://github.com/TabularisDB/tabularis-bigquery-plugin", + "documentation_url": "https://github.com/TabularisDB/tabularis-bigquery-plugin#readme", "support": { "issues_url": "https://github.com/TabularisDB/tabularis-bigquery-plugin/issues" }, @@ -15,7 +29,7 @@ "default_username": "", "executable": "bigquery-plugin", "color": "#4285F4", - "icon": "cloud", + "icon": "https://raw.githubusercontent.com/TabularisDB/tabularis-bigquery-plugin/main/bigquery-icon.svg", "connection_fields": { "host": {"hidden": true}, "port": {"hidden": true}, @@ -23,7 +37,7 @@ "password": {"hidden": true}, "database": { "label": "GCP Project ID", - "placeholder": "billing-project" + "placeholder": "my-data-project" } }, "ui_extensions": [ @@ -39,14 +53,35 @@ "views": true, "materialized_views": true, "routines": true, - "routine_management": false, + "routine_management": true, "triggers": false, "file_based": false, "folder_based": false, "no_connection_required": false, "connection_string": false, "identifier_quote": "`", - "sql_dialect": "generic", + "sql_dialect": "mysql", + "sql_syntax": { + "extends": "generic", + "quotes": [ + {"open": "'", "kind": "string", "escape": "backslash"}, + {"open": "\"", "kind": "string", "escape": "backslash"}, + {"open": "`", "kind": "identifier", "escape": "backslash"} + ], + "triple_quoted_strings": true, + "string_prefixes": ["r", "b", "rb", "br"], + "hash_comments": true, + "nested_block_comments": false, + "line_comment_requires_space": false, + "script_blocks": true, + "dollar_quoting": false, + "string_escape": "backslash", + "like_escape": false, + "boolean_literals": "keywords", + "formatter": "bigquery", + "keywords": ["QUALIFY", "PIVOT", "UNPIVOT", "TABLESAMPLE SYSTEM", "UNNEST", "WITH OFFSET", "EXCEPT", "REPLACE", "WINDOW", "OVER", "PARTITION BY", "CLUSTER BY", "OPTIONS", "FOR SYSTEM_TIME AS OF", "IGNORE NULLS", "RESPECT NULLS", "UNION ALL", "UNION DISTINCT", "INTERSECT DISTINCT", "EXCEPT DISTINCT", "MERGE", "WHEN MATCHED", "WHEN NOT MATCHED", "WHEN NOT MATCHED BY SOURCE", "CREATE OR REPLACE TABLE", "CREATE TEMP TABLE", "CREATE TEMP FUNCTION", "CREATE OR REPLACE PROCEDURE", "CREATE SCHEMA", "RETURNS", "LANGUAGE js", "EXPORT DATA", "LOAD DATA", "ASSERT", "DECLARE", "BEGIN", "BEGIN TRANSACTION", "COMMIT TRANSACTION", "ROLLBACK TRANSACTION", "EXCEPTION WHEN ERROR THEN", "ELSEIF", "END IF", "LOOP", "END LOOP", "WHILE", "END WHILE", "REPEAT", "UNTIL", "END REPEAT", "FOR", "DO", "END FOR", "END CASE", "LEAVE", "BREAK", "CONTINUE", "ITERATE", "RETURN", "CALL", "EXECUTE IMMEDIATE", "RAISE USING MESSAGE", "INT64", "FLOAT64", "NUMERIC", "BIGNUMERIC", "BOOL", "STRING", "BYTES", "DATE", "DATETIME", "TIME", "TIMESTAMP", "INTERVAL", "GEOGRAPHY", "JSON", "STRUCT", "ARRAY", "RANGE"], + "functions": ["SAFE_CAST", "SAFE.", "SAFE_DIVIDE", "IF", "IFNULL", "COALESCE", "NULLIF", "COUNTIF", "ARRAY_AGG", "STRING_AGG", "LOGICAL_AND", "LOGICAL_OR", "ANY_VALUE", "APPROX_COUNT_DISTINCT", "APPROX_QUANTILES", "APPROX_TOP_COUNT", "APPROX_TOP_SUM", "ROW_NUMBER", "RANK", "DENSE_RANK", "LAG", "LEAD", "FIRST_VALUE", "LAST_VALUE", "NTILE", "PERCENTILE_CONT", "ARRAY_LENGTH", "ARRAY_CONCAT", "ARRAY_TO_STRING", "GENERATE_ARRAY", "GENERATE_DATE_ARRAY", "SPLIT", "FORMAT", "STARTS_WITH", "ENDS_WITH", "CONTAINS_SUBSTR", "REGEXP_CONTAINS", "REGEXP_EXTRACT", "REGEXP_EXTRACT_ALL", "REGEXP_REPLACE", "REGEXP_INSTR", "FARM_FINGERPRINT", "GENERATE_UUID", "JSON_VALUE", "JSON_QUERY", "JSON_VALUE_ARRAY", "JSON_QUERY_ARRAY", "JSON_EXTRACT", "JSON_EXTRACT_SCALAR", "PARSE_JSON", "TO_JSON", "TO_JSON_STRING", "CURRENT_DATE", "CURRENT_DATETIME", "CURRENT_TIMESTAMP", "EXTRACT", "DATE_TRUNC", "DATE_ADD", "DATE_SUB", "DATE_DIFF", "DATETIME_TRUNC", "DATETIME_DIFF", "TIMESTAMP_TRUNC", "TIMESTAMP_ADD", "TIMESTAMP_SUB", "TIMESTAMP_DIFF", "FORMAT_DATE", "FORMAT_TIMESTAMP", "PARSE_DATE", "PARSE_TIMESTAMP", "UNIX_SECONDS", "UNIX_MILLIS", "TIMESTAMP_SECONDS", "TIMESTAMP_MILLIS", "LAST_DAY", "ST_GEOGPOINT", "ST_GEOGFROMTEXT", "ST_ASTEXT", "ST_DISTANCE", "ST_DWITHIN", "ST_CONTAINS", "ST_INTERSECTS", "ST_AREA", "ST_LENGTH", "ST_BUFFER"] + }, "alter_primary_key": true, "alter_column": true, "create_foreign_keys": true, @@ -56,20 +91,33 @@ "readonly": false, "serial_type": "", "auto_increment_keyword": "", - "inline_pk": false + "inline_pk": false, + "column_comments": true, + "drop_table_sql": true }, "type_mappings": { "INTEGER": "INT64", + "BIGINT": "INT64", + "SERIAL": "INT64", "REAL": "FLOAT64", + "DOUBLE": "FLOAT64", + "DECIMAL": "NUMERIC", "BOOLEAN": "BOOL", + "TEXT": "STRING", + "VARCHAR": "STRING", + "UUID": "STRING", + "BLOB": "BYTES", + "JSON": "JSON", + "DATE": "DATE", + "TIME": "TIME", "DATETIME": "DATETIME", - "TEXT": "STRING" + "TIMESTAMP": "TIMESTAMP" }, "data_types": [ {"name": "INT64", "category": "numeric", "requires_length": false, "requires_precision": false}, {"name": "FLOAT64", "category": "numeric", "requires_length": false, "requires_precision": false}, - {"name": "NUMERIC", "category": "numeric", "requires_length": false, "requires_precision": true}, - {"name": "BIGNUMERIC", "category": "numeric", "requires_length": false, "requires_precision": true}, + {"name": "NUMERIC", "category": "numeric", "requires_length": false, "requires_precision": false}, + {"name": "BIGNUMERIC", "category": "numeric", "requires_length": false, "requires_precision": false}, {"name": "STRING", "category": "string", "requires_length": false, "requires_precision": false}, {"name": "BYTES", "category": "binary", "requires_length": false, "requires_precision": false}, {"name": "BOOL", "category": "other", "requires_length": false, "requires_precision": false}, diff --git a/CHANGELOG.md b/CHANGELOG.md index 508b84e..d6334f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,11 +2,64 @@ All notable changes to this project are documented in this file. -## Unreleased +## 0.2.0 + +Requires Tabularis 0.29.0 or later (`min_runtime_version`): it relies on plugin connection fields and secrets, the DDL connection context, the declarative `sql_syntax` capability and the bundled Visual EXPLAIN parser. + +### Query runtime + +- Handle JSON-RPC requests concurrently with a single ordered writer, cancel the running BigQuery job when Tabularis sends a `cancel` notification, and answer with an internal error instead of never replying when a request handler fails. +- Start jobs with a `requestId`, the `client: tabularis` label and a server-side `jobTimeoutMs`; cancel jobs that exceed the query timeout. The default query timeout is now 110 seconds, below the host's 120 second call timeout. +- Retry HTTP 429, 5xx and 403 `rateLimitExceeded`/`backendError` responses with jittered backoff and refresh the token once after HTTP 401. Judge job failure by `status.errorResult`. +- Fill result pages across `pageToken` responses, keep the rows returned by `jobs.query`, and read later pages from the finished job instead of running the SQL again. +- Return TIMESTAMP values with exact microseconds. +- Serve plain `SELECT * FROM table` table browsing through the free `tabledata.list` API. +- Support Tabularis editor sessions with BigQuery sessions, transaction tracking (including COMMIT and ROLLBACK inside scripting blocks) and `release_session`; run stateful batches in a short-lived session. +- Return the last result set and summed DML row count of multi-statement scripts. +- Return Visual EXPLAIN as a raw `bigquery-job-json` payload with a bundled parser: dry-run bytes and on-demand cost estimates, and executed stage plans for EXPLAIN ANALYZE. +- Map common BigQuery error reasons to messages with hints. +- Run metadata queries, DDL, row edits and routine statements through the same job path as editor queries, with typed named query parameters, so they get request IDs, retries, job timeouts, labels, error hints and cancellation. +- Clear the metadata cache after any editor statement other than a plain query, so DDL and DML run from a query tab show up at once. + +### Projects and metadata + +- Browse several projects from one connection: `get_databases` lists visible projects plus the `additional_projects` setting, and metadata follows the selected project. +- Run and bill every job in the optional `billing_project` (default: the GCP Project ID), while unqualified names and table browsing resolve in the selected project. +- Read routines through the free routines REST API, return host-compatible routine parameters with ordinal positions and a named return row, and add routine management: run, create from templates, edit with rebuilt `CREATE OR REPLACE` scripts (SQL, JavaScript, Python, remote, table and aggregate functions, procedures) and drop. +- Fetch each table resource once with bounded concurrency for schema snapshots, ER diagram batches and AI context, cache metadata briefly, and return the requested tables from `get_all_columns_batch` and `get_all_foreign_keys_batch`, skipping tables that cannot be read instead of failing the batch. +- Fill table descriptions, parameterized column types, numeric precision and scale, string lengths and composite foreign-key order. +- Add `get_table_ddl` and use `INFORMATION_SCHEMA.TABLES.ddl` for table, view and materialized view definitions, with a metadata-based fallback that keeps partitioning, clustering and options. +- Keep view options when saving a view and refuse to replace legacy SQL views. +- Refresh materialized views on demand with `BQ.REFRESH_MATERIALIZED_VIEW`. + +### DDL and row editing + +- Quote every part of generated identifiers and fully qualify objects as `` `project`.`dataset`.`object` ``. Table, column, foreign-key and routine template builders use the browsed project from the connection parameters the host sends with every DDL request, and fall back to dataset-qualified names on older hosts. +- Declare `drop_table_sql` and implement `get_drop_table_sql`, so the explorer drops tables, external tables, snapshots, views and materialized views with the matching statement. +- Declare `column_comments`: column descriptions from the column editor are written with `OPTIONS(description = ...)` on create and add, and changed or cleared with `ALTER COLUMN ... SET OPTIONS(description = ...)`. +- Drop views and materialized views with the matching statement, and reject row edits and column changes on snapshots and external tables. +- Generate only DDL that BigQuery accepts: no `NOT NULL` on added columns, no `SET NOT NULL`, type changes limited to supported widenings, column descriptions through `SET OPTIONS`, primary keys added and dropped as `NOT ENFORCED`, and optional partitioning, clustering and description on new tables. +- Edit rows with typed query parameters; support tables without a primary key by matching comparable columns inside a transaction that rolls back when more than one row matches; report real affected-row counts and explain streaming-buffer errors. +- Declare the GoogleSQL syntax (backtick identifiers, triple-quoted and raw strings, `#` comments, procedural `BEGIN … END` blocks, backslash literals, `TRUE`/`FALSE` boolean literals, keywords and functions) through `sql_syntax`, so Tabularis splits scripts, formats SQL and builds filters the GoogleSQL way. Older hosts fall back to the `mysql` dialect. +- Write valid GoogleSQL literals for strings with newlines and control characters, NUMERIC, BIGNUMERIC, GEOGRAPHY, INTERVAL, RANGE, JSON, BYTES, ARRAY and STRUCT values. + +### Connection and authentication - Hide unused common connection fields on compatible Tabularis hosts. - Store plugin-owned credentials in the OS keychain while preserving legacy Password-based connections. -- Author the connection UI in React and TypeScript and generate its release bundle with Vite. +- Author the connection UI in React and TypeScript and generate its release bundle with Vite, with advanced fields for the billing project, additional projects, API endpoint and impersonation. +- Support `impersonated_service_account` ADC files and an optional service account to impersonate, reject service account keys whose `token_uri` is not a googleapis.com endpoint, and tolerate a byte order mark in credentials. +- Add an API endpoint override for Private Service Connect and emulators, and a `none` authentication method for localhost endpoints. +- Fix the client cache key to cover the plugin-owned credential and every connection option, including the billing and data project. +- Use a single `datasets.list` call for `ping`. +- Fix BigQuery REST URLs that contained a double slash after the API version. + +### Packaging and tests + +- Fix the GCP Project ID placeholder, which suggested the billing project. +- Let `pnpm install --frozen-lockfile` run the esbuild install script with pnpm 10 and later. +- Complete the manifest type mappings, stop requiring precision for NUMERIC and BIGNUMERIC, add registry metadata and declare the Visual EXPLAIN parser. +- Add mocked HTTP, host-shape conformance and BigQuery emulator integration tests for the query runtime and for metadata, DDL and row editing. ## 0.1.0 diff --git a/Cargo.lock b/Cargo.lock index 239f304..c0730b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + [[package]] name = "android_system_properties" version = "0.1.6" @@ -11,6 +20,16 @@ dependencies = [ "libc", ] +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-trait" version = "0.1.92" @@ -42,8 +61,9 @@ checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" [[package]] name = "bigquery-plugin" -version = "0.1.0" +version = "0.2.0" dependencies = [ + "async-trait", "base64", "chrono", "gcp_auth", @@ -53,6 +73,8 @@ dependencies = [ "sha2", "tokio", "url", + "uuid", + "wiremock", ] [[package]] @@ -173,6 +195,24 @@ dependencies = [ "typenum", ] +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + [[package]] name = "digest" version = "0.10.7" @@ -221,6 +261,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.34" @@ -228,6 +283,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -236,6 +292,34 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "futures-sink" version = "0.3.34" @@ -254,8 +338,13 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", + "futures-sink", "futures-task", + "memchr", "pin-project-lite", "slab", ] @@ -349,6 +438,12 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "hermit-abi" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" + [[package]] name = "http" version = "1.5.0" @@ -388,6 +483,12 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + [[package]] name = "hyper" version = "1.11.0" @@ -402,6 +503,7 @@ dependencies = [ "http", "http-body", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -609,6 +711,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" + [[package]] name = "libc" version = "0.2.189" @@ -659,6 +767,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -818,6 +936,35 @@ dependencies = [ "rand_core", ] +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + [[package]] name = "reqwest" version = "0.12.28" @@ -1338,6 +1485,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "uuid" +version = "1.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cc1186384beb7dd8eedea376413fd654937285ea6c9cfbb928dc3043ea4b606" +dependencies = [ + "getrandom 0.4.3", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "version_check" version = "0.9.5" @@ -1584,6 +1742,29 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "wiremock" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" +dependencies = [ + "assert-json-diff", + "base64", + "deadpool", + "futures", + "http", + "http-body-util", + "hyper", + "hyper-util", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", + "url", +] + [[package]] name = "writeable" version = "0.6.3" diff --git a/Cargo.toml b/Cargo.toml index 532b786..a98c695 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "bigquery-plugin" -version = "0.1.0" +version = "0.2.0" edition = "2021" description = "Google BigQuery driver plugin for Tabularis" license = "Apache-2.0" @@ -20,6 +20,11 @@ serde_json = "1" sha2 = "0.10" tokio = { version = "1", features = ["fs", "io-std", "io-util", "macros", "rt-multi-thread", "sync", "time"] } url = "2" +uuid = { version = "1", features = ["v4"] } + +[dev-dependencies] +async-trait = "0.1" +wiremock = "0.6" [profile.release] lto = true diff --git a/README.md b/README.md index 0bb43f6..7e58082 100644 --- a/README.md +++ b/README.md @@ -6,20 +6,34 @@ A native Google BigQuery driver for [Tabularis](https://github.com/TabularisDB/tabularis). The plugin is a standalone Rust executable that uses the BigQuery REST API v2 and communicates with Tabularis through newline-delimited JSON-RPC 2.0 over standard input and output. +Requires Tabularis 0.29.0 or later. + ## Features -- Application Default Credentials, service account JSON, refreshable OAuth authorized user credentials and direct OAuth access token authentication -- Dataset, table, external table, snapshot, clone, view and materialized view browsing +- Application Default Credentials (including `impersonated_service_account` ADC files), service account JSON, refreshable OAuth authorized user credentials and direct OAuth access token authentication +- Optional service account impersonation for every authentication method +- Optional API endpoint override for Private Service Connect and local emulators +- Multi-project browsing: every project the credential can see, plus extra projects such as `bigquery-public-data`, with an optional separate billing project for jobs +- Dataset, table, external table, snapshot, clone, view and materialized view browsing with descriptions - Nested `STRUCT`, repeated `ARRAY`, `JSON`, `BYTES`, temporal, numeric and geography value conversion -- GoogleSQL query execution with result pagination, job polling, retry handling and DML affected-row counts -- Visual EXPLAIN cost summaries backed by BigQuery dry runs +- GoogleSQL editor support declared through `sql_syntax`: script splitting with `BEGIN … END` blocks, triple-quoted and raw strings, `#` comments, BigQuery formatting and autocomplete keywords +- GoogleSQL query execution with idempotent, labelled jobs, server-side job timeouts and cancellation +- Exact result pagination that reuses the finished job for later pages instead of running the SQL again +- Free table browsing through `tabledata.list` for plain `SELECT * FROM table` queries +- Editor sessions: temporary tables, variables and multi-statement transactions survive between runs +- Multi-statement scripts return the last result set and the total DML affected-row count +- Visual EXPLAIN with dry-run cost estimates and executed stage plans (`EXPLAIN ANALYZE`) +- Clear error messages with hints for common BigQuery failures - Per-connection location, timeout and maximum bytes billed controls -- Table columns, descriptions, primary keys and foreign keys from BigQuery metadata -- Routine discovery and routine definition inspection -- Inline insert, update and delete operations with type-aware GoogleSQL literals -- Table, column, view and unenforced key DDL generation +- Table columns, descriptions, parameterized types, primary keys and foreign keys from BigQuery metadata +- Table, view and materialized view DDL from `INFORMATION_SCHEMA`, with partitioning, clustering and options +- Routine management for SQL, JavaScript and Python functions, table functions, aggregate functions and procedures: browse, inspect parameters, run, create, edit and drop +- Manual materialized view refresh +- Inline insert, update and delete with typed query parameters, including tables without a primary key +- Table, column, view and unenforced key DDL generation that only emits changes BigQuery accepts, with column descriptions +- Dropping tables, external tables, snapshots, views and materialized views from the explorer with the matching `DROP` statement - BLOB preview and export for `BYTES` columns -- Batch metadata methods for ER diagrams and AI schema context +- Batch metadata methods for ER diagrams and AI schema context, fetched with bounded concurrency - Linux, macOS and Windows release packaging ## Connection setup @@ -28,15 +42,21 @@ On a compatible Tabularis host, the BigQuery plugin hides unused common fields a | Tabularis field | BigQuery meaning | | --- | --- | -| GCP Project ID | Project used for jobs and billing | +| GCP Project ID | Project that is browsed and, unless a billing project is set, runs jobs | | Credential | Service account JSON/path, authorized user JSON/path, or access token | The BigQuery section in the connection form adds: -- **Authentication**: ADC, service account JSON, OAuth authorized user JSON, or OAuth access token +- **Authentication**: ADC, service account JSON, OAuth authorized user JSON, OAuth access token, or no authentication (local emulator only) - **Location**: optional job location such as `US`, `EU`, or `europe-west1` -- **Query timeout**: 30 to 3600 seconds, default 300 +- **Query timeout**: 30 to 3600 seconds, default 110 - **Maximum bytes billed**: optional per-query cost guard +- **Advanced options**: + - **Billing project** (`billing_project`): optional project that runs and pays for every job. Set it when browsing several projects, so jobs never run in a project you can only read, such as `bigquery-public-data` + - **Additional projects** (`additional_projects`): optional comma-separated projects added to the project list, for public or shared projects the credential cannot list + - **API endpoint** and **Impersonate service account**, described below + +The query timeout is also sent to BigQuery as the job timeout, and the plugin cancels a job that exceeds it. Tabularis waits 120 seconds for a plugin call by default, so the default of 110 seconds lets the plugin report and cancel a slow query itself. To allow longer queries, raise both the connection's query timeout and the Tabularis plugin call timeout (`pluginCallTimeoutSeconds`, or `callTimeoutSeconds` for the `bigquery` plugin). When Tabularis gives up on a call it sends a cancel notification and the plugin cancels the running job. The Credential field is plugin-owned. Tabularis stores its value in the OS keychain and never writes it to `connections.json`. Existing connections that used the Password field remain supported and do not require migration. @@ -51,9 +71,19 @@ Leave Credential empty. The plugin tries the standard Google credential chain: 3. A service account attached to the Google Cloud runtime 4. The active `gcloud` CLI account +If the ADC file (`GOOGLE_APPLICATION_CREDENTIALS` or the gcloud default file) is an `impersonated_service_account` credential, such as one created by `gcloud auth application-default login --impersonate-service-account`, the plugin exchanges its source credential for a token of the target service account through the IAM Credentials API. + ### Service account -Select **Service account JSON** and put either the credential file path or the complete JSON document in Credential. +Select **Service account JSON** and put either the credential file path or the complete JSON document in Credential. The key's `token_uri` must be a `googleapis.com` HTTPS endpoint; keys pointing elsewhere are rejected. A leading byte order mark is tolerated. + +### Service account impersonation + +Enter a service account email under **Advanced options > Impersonate service account** to run every request as that account. The configured credential must hold the Service Account Token Creator role on it. + +### API endpoint and emulators + +**Advanced options > API endpoint** replaces `https://bigquery.googleapis.com`, for example with a Private Service Connect endpoint. The REST path `/bigquery/v2/` is appended when missing. HTTPS is required except for `localhost`, `127.0.0.1` and `[::1]`. The **No authentication** method is accepted only together with such a local endpoint, for example `http://localhost:9050` for a BigQuery emulator. ### OAuth access token @@ -65,7 +95,18 @@ Select **OAuth authorized user JSON** and put either an `authorized_user` creden ## Required Google Cloud permissions -The exact IAM role depends on the enabled operations. Read-only use generally needs BigQuery Metadata Viewer and BigQuery Job User permissions. Editing and DDL need the corresponding BigQuery Data Editor or Data Owner permissions. Grant the least privilege required by your workflow. +The exact IAM role depends on the enabled operations. Grant the least privilege required by your workflow. + +| Operation | Permissions | Typical role | +| --- | --- | --- | +| Run queries, DDL and DML | `bigquery.jobs.create` on the billing project | BigQuery Job User | +| Browse datasets, tables, views and routines | `bigquery.datasets.get`, `bigquery.tables.list`, `bigquery.tables.get`, `bigquery.routines.list`, `bigquery.routines.get` | BigQuery Metadata Viewer | +| Read table data and `INFORMATION_SCHEMA` DDL | `bigquery.tables.getData` | BigQuery Data Viewer | +| List projects for multi-project browsing | `resourcemanager.projects.get` on each listed project | Any role granting BigQuery access to the project | +| Edit rows, alter tables, create and drop views and routines | `bigquery.tables.update`, `bigquery.tables.updateData`, `bigquery.tables.create`, `bigquery.tables.delete`, `bigquery.routines.create`, `bigquery.routines.update`, `bigquery.routines.delete` | BigQuery Data Editor | +| Refresh materialized views | `bigquery.tables.updateData` on the view and read access to its base tables | BigQuery Data Editor | + +When projects cannot be listed, the project list falls back to the configured, billing and additional projects. ## Installation @@ -75,6 +116,9 @@ Download the archive for your platform from GitHub Releases, extract it, and cop bigquery/ ├── .tabularium ├── bigquery-plugin +├── explain/ +│ └── dist/ +│ └── index.iife.js └── ui/ └── dist/ └── connection-fields.js @@ -84,24 +128,36 @@ Use `bigquery-plugin.exe` on Windows. | OS | Plugin directory | | --- | --- | -| Linux | `~/.local/share/tabularis/plugins/bigquery/` | -| macOS | `~/Library/Application Support/com.debba.tabularis/plugins/bigquery/` | -| Windows | `%APPDATA%\debba\tabularis\data\plugins\bigquery\` | +| Linux | `~/.local/share/tabularis/plugins/drivers/bigquery/` | +| macOS | `~/Library/Application Support/tabularis/plugins/drivers/bigquery/` | +| Windows | `%APPDATA%\tabularis\plugins\drivers\bigquery\` | ## Build and test -Rust stable, Node.js 20 and pnpm 9 are required. The UI extension is authored in React and TypeScript under `ui/src`; Vite produces the IIFE bundle loaded by Tabularis. +Rust stable, Node.js 20 and pnpm 9 or later are required. Each package directory is its own pnpm workspace whose `pnpm-workspace.yaml` allows the esbuild install script, so frozen installs also work with pnpm 10 and later. The UI extension is authored in React and TypeScript under `ui/src`; Vite produces the IIFE bundle loaded by Tabularis. The Visual EXPLAIN parser lives under `explain/` and is built with tsup. ```bash pnpm --dir ui install --frozen-lockfile pnpm --dir ui typecheck pnpm --dir ui build +pnpm --dir explain install --frozen-lockfile +pnpm --dir explain typecheck +pnpm --dir explain test cargo build cargo test --all-targets cargo clippy --all-targets -- -D warnings cargo fmt --all -- --check ``` +Integration tests run against the BigQuery emulator when `BIGQUERY_EMULATOR_HOST` is set and are skipped otherwise. `tests/emulator.rs` covers the query runtime and `tests/emulator_metadata.rs` covers metadata, DDL and row editing; both expect the project `test` and the dataset `ds`: + +```bash +docker run -d -p 9050:9050 ghcr.io/goccy/bigquery-emulator:latest --project=test --dataset=ds +BIGQUERY_EMULATOR_HOST=localhost:9050 cargo test --all-targets +``` + +The emulator does not implement sessions, dry runs, DML row counts, script child jobs, `tabledata.list` paging, `@@row_count`, `INFORMATION_SCHEMA.TABLES.ddl`, informational keys or the routines API; the plugin falls back gracefully and the mocked tests in `tests/runtime_mock.rs` and the host-shape tests in `tests/conformance.rs` and `tests/conformance_metadata.rs` cover those paths. + Validate the release manifest with Tabularium: ```bash @@ -117,16 +173,32 @@ echo '{"jsonrpc":"2.0","method":"initialize","id":1}' \ | cargo run --quiet ``` -The plugin writes protocol responses only to stdout. Diagnostics are written to stderr. +The plugin writes protocol responses only to stdout. Diagnostics are written to stderr. Requests are handled concurrently, so responses can arrive in a different order than requests; each carries its request `id`. Notifications, such as `cancel`, are never answered. ## BigQuery-specific behavior -- Datasets are exposed as Tabularis schemas. -- The selected project is used as the BigQuery job project and billing project. -- BigQuery primary and foreign keys are informational and emitted as `NOT ENFORCED` in generated DDL. -- BigQuery materialized views refresh automatically, so manual refresh returns an unsupported-operation error. -- General-purpose indexes are not available in BigQuery. Search indexes and vector indexes require specialized statements and are not exposed through Tabularis' generic index UI. -- Visual EXPLAIN uses a BigQuery dry run. It reports estimated bytes processed without executing the query, rather than exposing an operator tree. +- Projects are exposed as Tabularis databases and datasets as schemas. Select several projects in the connection dialog to browse them side by side. +- Jobs always run in the billing project: the `billing_project` setting when present, otherwise the GCP Project ID. Unqualified table names in the editor resolve in the project being browsed (`defaultDataset.projectId`), and metadata, DDL and row edits target that project too. Generated SQL is fully qualified as `` `project`.`dataset`.`object` `` with every part quoted separately, so project IDs containing dashes work. +- Metadata comes from the free REST API (`tables.get`, `routines.get`). Resources are fetched once per object, with at most eight requests in flight for ER diagrams, schema snapshots and AI context, and cached for five seconds; any DDL sent through the plugin, and any statement other than a plain query run from the editor, clears the cache. Table descriptions are loaded for datasets with up to 100 tables or views. +- Jobs carry a `requestId`, so retried requests never start a second job, and the label `client: tabularis` for cost attribution. Requests are retried on HTTP 429 and 5xx and on 403 `rateLimitExceeded` or `backendError`, with exponential backoff and jitter. A rejected token is refreshed once. +- Pages after the first are read from the finished job with `getQueryResults` for 10 minutes; the SQL is not executed again. The first page always runs the statement. +- A query of the exact form `SELECT * FROM table [LIMIT n [OFFSET m]]` on a native table, which is what a Tabularis table tab sends, is served by `tabledata.list` without a query job and without scanning cost. Views, external tables, filters, ordering and any other SQL run as normal jobs. +- TIMESTAMP values are returned with exact microseconds. +- When Tabularis pins an editor tab to a session, the plugin creates a BigQuery session for it and runs every statement of that tab inside it. `BEGIN`, `COMMIT` and `ROLLBACK` are tracked for the transaction indicator. Closing the tab aborts the session. A batch without a session that contains `DECLARE`, `SET`, temporary tables or transaction statements runs in a short-lived session. BigQuery has no savepoints; savepoint statements are answered with an error without being sent. +- Visual EXPLAIN runs a dry run and shows bytes processed and an estimated on-demand cost (US$6.25 per TiB list price, clearly labelled as an estimate) plus the referenced tables. EXPLAIN ANALYZE executes the statement, including DML, and shows each query stage with slot time, wall time, records read and written, shuffle bytes and steps. The plan is parsed by the bundled `explain/dist/index.iife.js` parser. +- Table, view and materialized view definitions come from `INFORMATION_SCHEMA.TABLES.ddl`, a small billed query run in the dataset location, so options, partitioning and clustering are preserved. When that query is not available the DDL is rebuilt from table metadata. +- Saving a view keeps its description, friendly name and labels. Legacy SQL views can be read but not saved from the view editor. +- Routines are listed as functions and procedures. Table functions, aggregate functions, JavaScript and Python UDFs and remote functions are rebuilt as complete `CREATE OR REPLACE` scripts for editing. Running a function previews `SELECT fn(...)` or `SELECT * FROM table_fn(...)`; procedures use `CALL`, with `OUT` and `INOUT` arguments returned through script variables. +- Materialized views refresh automatically. **Refresh** runs `CALL BQ.REFRESH_MATERIALIZED_VIEW(...)` for an immediate refresh. +- BigQuery primary and foreign keys are informational and emitted as `NOT ENFORCED`. Foreign keys support no `ON DELETE` / `ON UPDATE` actions. +- Column changes follow BigQuery limits: new columns cannot be `NOT NULL`, existing columns can drop `NOT NULL` but never gain it, and `SET DATA TYPE` only widens types (`INT64` to `NUMERIC`, `BIGNUMERIC` or `FLOAT64`; `NUMERIC` to `BIGNUMERIC` or `FLOAT64`; larger `STRING`, `BYTES` and decimal parameters). Other changes are rejected with an explanation instead of failing at run time. +- Snapshots and external tables are read-only for row editing and column changes. +- Column descriptions entered in the column editor become `OPTIONS(description = ...)` on new tables and columns, and `ALTER COLUMN ... SET OPTIONS(description = ...)` on existing columns; clearing the description sets it to `NULL`. +- Dropping an object from the explorer looks up its type and runs `DROP TABLE`, `DROP EXTERNAL TABLE`, `DROP SNAPSHOT TABLE`, `DROP VIEW` or `DROP MATERIALIZED VIEW` on the fully qualified name. +- Row edits use typed query parameters. On tables without a primary key, rows are matched on every comparable column (floating point, JSON, geography, arrays, structs, ranges and bytes are skipped), and the change runs in a transaction that rolls back if more than one row would change. +- Rows still in the streaming buffer cannot be updated or deleted; the plugin reports this with a clear message. +- General-purpose indexes are not available in BigQuery. Search and vector indexes use `CREATE SEARCH INDEX` / `CREATE VECTOR INDEX` in a query tab; the index editor only drops the informational primary key. +- Metadata queries, DDL and row edits run through the same job path as editor queries: `requestId`, retries, `jobTimeoutMs`, labels, error hints and cancellation apply to every statement the plugin sends. ## License diff --git a/bigquery-icon.svg b/bigquery-icon.svg new file mode 100644 index 0000000..a04833e --- /dev/null +++ b/bigquery-icon.svg @@ -0,0 +1,23 @@ + + Google BigQuery + White magnifying glass over a bar chart on a blue hexagon + + + + + + + + + + + + + + + + + + + + diff --git a/explain/LICENSE b/explain/LICENSE new file mode 100644 index 0000000..57bc88a --- /dev/null +++ b/explain/LICENSE @@ -0,0 +1,202 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + diff --git a/explain/README.md b/explain/README.md new file mode 100644 index 0000000..1dcba85 --- /dev/null +++ b/explain/README.md @@ -0,0 +1,38 @@ +# `@tabularis/explain-bigquery` + +Google BigQuery job plan parser for Tabularis Visual EXPLAIN. It is maintained +and versioned with the +[Tabularis BigQuery plugin](https://github.com/TabularisDB/tabularis-bigquery-plugin), +which ships `dist/index.iife.js` in each release archive. + +The parser reads the `bigquery-job-json` payload produced by the plugin's +`explain_query`: + +- `mode: "estimate"`: a dry run. The plan shows bytes processed, an on-demand + cost estimate and the referenced tables. Nothing is executed. +- `mode: "analyze"`: an executed job. Each query stage becomes a node with slot + milliseconds as cost, stage wall time, records read and written, shuffle + bytes and the stage steps. + +`@tabularis/explain` 0.2.0 or newer is required for the parser registry. + +```ts +import "@tabularis/explain-bigquery"; +import { parseRawExplain } from "@tabularis/explain"; + +const plan = parseRawExplain({ + engine: "bigquery", + format: "bigquery-job-json", + payload, + original_query: "SELECT * FROM sales.orders", +}); +``` + +## Development + +```bash +pnpm install +pnpm typecheck +pnpm test +pnpm build +``` diff --git a/explain/package.json b/explain/package.json new file mode 100644 index 0000000..2237200 --- /dev/null +++ b/explain/package.json @@ -0,0 +1,59 @@ +{ + "name": "@tabularis/explain-bigquery", + "version": "0.2.0", + "description": "Google BigQuery job plan parser for Tabularis Visual EXPLAIN", + "license": "Apache-2.0", + "homepage": "https://github.com/TabularisDB/tabularis-bigquery-plugin/tree/main/explain", + "repository": { + "type": "git", + "url": "https://github.com/TabularisDB/tabularis-bigquery-plugin.git", + "directory": "explain" + }, + "bugs": { + "url": "https://github.com/TabularisDB/tabularis-bigquery-plugin/issues" + }, + "keywords": [ + "explain", + "bigquery", + "query-plan", + "tabularis" + ], + "type": "module", + "sideEffects": [ + "./dist/index.js" + ], + "exports": { + ".": { + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + } + }, + "files": [ + "dist", + "README.md", + "LICENSE" + ], + "scripts": { + "build": "tsup", + "clean": "rm -rf dist", + "typecheck": "tsc --noEmit", + "test": "pnpm run build && vitest run", + "prepublishOnly": "pnpm run clean && pnpm run typecheck && pnpm run test" + }, + "peerDependencies": { + "@tabularis/explain": "^0.2.0" + }, + "devDependencies": { + "@tabularis/explain": "file:tests/host", + "@types/node": "^26.4.0", + "tsup": "^8.3.5", + "typescript": "~5.9.3", + "vitest": "^3.2.4" + }, + "publishConfig": { + "access": "public" + }, + "engines": { + "node": ">=20" + } +} diff --git a/explain/pnpm-lock.yaml b/explain/pnpm-lock.yaml new file mode 100644 index 0000000..c1ddd27 --- /dev/null +++ b/explain/pnpm-lock.yaml @@ -0,0 +1,1591 @@ +lockfileVersion: '9.0' + +settings: + autoInstallPeers: true + excludeLinksFromLockfile: false + +importers: + + .: + devDependencies: + '@tabularis/explain': + specifier: file:tests/host + version: file:tests/host + '@types/node': + specifier: ^26.4.0 + version: 26.6.5 + tsup: + specifier: ^8.3.5 + version: 8.5.1(postcss@8.5.30)(typescript@5.9.3) + typescript: + specifier: ~5.9.3 + version: 5.9.3 + vitest: + specifier: ^3.2.4 + version: 3.2.7(@types/node@26.6.5) + +packages: + + '@esbuild/aix-ppc64@0.27.7': + resolution: {integrity: sha512-EKX3Qwmhz1eMdEJokhALr0YiD0lhQNwDqkPYyPhiSwKrh7/4KRjQc04sZ8db+5DVVnZ1LmbNDI1uAMPEUBnQPg==} + engines: {node: '>=18'} + cpu: [ppc64] + os: [aix] + + '@esbuild/aix-ppc64@0.28.2': + resolution: {integrity: sha512-XExcO+dvLKvVtNTibSTBej1NCAbaGhWn9Ww1ZPx80qsahhPFe/8jgWP0IchNe0F3HwkU7n8ejhH8bjonqht8mQ==} + engines: {node: '>=18'} + cpu: [ppc64] + os: [aix] + + '@esbuild/android-arm64@0.27.7': + resolution: {integrity: sha512-62dPZHpIXzvChfvfLJow3q5dDtiNMkwiRzPylSCfriLvZeq0a1bWChrGx/BbUbPwOrsWKMn8idSllklzBy+dgQ==} + engines: {node: '>=18'} + cpu: [arm64] + os: [android] + + '@esbuild/android-arm64@0.28.2': + resolution: {integrity: sha512-5YfKeeI8qWfBZIX+u2xZC3Zlb3Os/gLS2sbEKM+I4ZOcsWmHS2WLysCcQZDAFRslDUU5Oiq44gf6PYN1vGwG5A==} + engines: {node: '>=18'} + cpu: [arm64] + os: [android] + + '@esbuild/android-arm@0.27.7': + resolution: {integrity: sha512-jbPXvB4Yj2yBV7HUfE2KHe4GJX51QplCN1pGbYjvsyCZbQmies29EoJbkEc+vYuU5o45AfQn37vZlyXy4YJ8RQ==} + engines: {node: '>=18'} + cpu: [arm] + os: [android] + + '@esbuild/android-arm@0.28.2': + resolution: {integrity: sha512-kXXoiPVVGQcnIYGOeaovwOURpniDBpSq4A03qkQ+BMQqtGG6HYap3xne9C1O1yo4TR3qxlCX5IqqmX6fFo2Lqg==} + engines: {node: '>=18'} + cpu: [arm] + os: [android] + + '@esbuild/android-x64@0.27.7': + resolution: {integrity: sha512-x5VpMODneVDb70PYV2VQOmIUUiBtY3D3mPBG8NxVk5CogneYhkR7MmM3yR/uMdITLrC1ml/NV1rj4bMJuy9MCg==} + engines: {node: '>=18'} + cpu: [x64] + os: [android] + + '@esbuild/android-x64@0.28.2': + resolution: {integrity: sha512-O387ite7SzUyCcy3JQX4P4bLtEA7bLLkx+esve5JHnyYfNTxcVpXZo9jhdB0lTKN44gztELTdU7nS8Nr16Fs1Q==} + engines: {node: '>=18'} + cpu: [x64] + os: [android] + + '@esbuild/darwin-arm64@0.27.7': + resolution: {integrity: sha512-5lckdqeuBPlKUwvoCXIgI2D9/ABmPq3Rdp7IfL70393YgaASt7tbju3Ac+ePVi3KDH6N2RqePfHnXkaDtY9fkw==} + engines: {node: '>=18'} + cpu: [arm64] + os: [darwin] + + '@esbuild/darwin-arm64@0.28.2': + resolution: {integrity: sha512-n4KqkOQrraxHJcgjM1RvwbigfQKIKJVpM7xp+KsxiyUSrRdIXnt73VhrPAx0fV44hgfmIVKjxMN9J1t5jySVkw==} + engines: {node: '>=18'} + cpu: [arm64] + os: [darwin] + + '@esbuild/darwin-x64@0.27.7': + resolution: {integrity: sha512-rYnXrKcXuT7Z+WL5K980jVFdvVKhCHhUwid+dDYQpH+qu+TefcomiMAJpIiC2EM3Rjtq0sO3StMV/+3w3MyyqQ==} + engines: {node: '>=18'} + cpu: [x64] + os: [darwin] + + '@esbuild/darwin-x64@0.28.2': + resolution: {integrity: sha512-uq6suIWYP37qzGddBKPw5QEQPi6HiLGsO7UmkpfyaYNQ3D+rN6w6WfwH+nuqcGXWvawGwxOEroO4YGnFh95azw==} + engines: {node: '>=18'} + cpu: [x64] + os: [darwin] + + '@esbuild/freebsd-arm64@0.27.7': + resolution: {integrity: sha512-B48PqeCsEgOtzME2GbNM2roU29AMTuOIN91dsMO30t+Ydis3z/3Ngoj5hhnsOSSwNzS+6JppqWsuhTp6E82l2w==} + engines: {node: '>=18'} + cpu: [arm64] + os: [freebsd] + + '@esbuild/freebsd-arm64@0.28.2': + resolution: {integrity: sha512-n+I0BTSRIoy+d6RPKnEVwql5UwBJolytvY4mAOIEJorKlqgPII8ix6slVVrfZ5Tnj7glIZvloylbB/EJPMWEXw==} + engines: {node: '>=18'} + cpu: [arm64] + os: [freebsd] + + '@esbuild/freebsd-x64@0.27.7': + resolution: {integrity: sha512-jOBDK5XEjA4m5IJK3bpAQF9/Lelu/Z9ZcdhTRLf4cajlB+8VEhFFRjWgfy3M1O4rO2GQ/b2dLwCUGpiF/eATNQ==} + engines: {node: '>=18'} + cpu: [x64] + os: [freebsd] + + '@esbuild/freebsd-x64@0.28.2': + resolution: {integrity: sha512-78XJTJkvPs0kz2w61301PJjXl4g7q3JqiYMZ/M/yVI73EHBrCRTgkhu9oqG7vPqq+a/yadEW8aD+agKlk5xrmg==} + engines: {node: '>=18'} + cpu: [x64] + os: [freebsd] + + '@esbuild/linux-arm64@0.27.7': + resolution: {integrity: sha512-RZPHBoxXuNnPQO9rvjh5jdkRmVizktkT7TCDkDmQ0W2SwHInKCAV95GRuvdSvA7w4VMwfCjUiPwDi0ZO6Nfe9A==} + engines: {node: '>=18'} + cpu: [arm64] + os: [linux] + + '@esbuild/linux-arm64@0.28.2': + resolution: {integrity: sha512-pW4AC0P3it8c7do9MVM4p51FzHzdM/TZrerurgRcHJ2WTa1VQ1CIq18xncfpBJw4ojkiZZrKW2yIBWBP92j6Ug==} + engines: {node: '>=18'} + cpu: [arm64] + os: [linux] + + '@esbuild/linux-arm@0.27.7': + resolution: {integrity: sha512-RkT/YXYBTSULo3+af8Ib0ykH8u2MBh57o7q/DAs3lTJlyVQkgQvlrPTnjIzzRPQyavxtPtfg0EopvDyIt0j1rA==} + engines: {node: '>=18'} + cpu: [arm] + os: [linux] + + '@esbuild/linux-arm@0.28.2': + resolution: {integrity: sha512-XlDnu2q5yoqems+xay6wSAcg9DDD7K9RLKZEBOMZm3ckNpJBvOX20tSfby8KfrrhINDyv9V2YVZKY/SpoGJI8w==} + engines: {node: '>=18'} + cpu: [arm] + os: [linux] + + '@esbuild/linux-ia32@0.27.7': + resolution: {integrity: sha512-GA48aKNkyQDbd3KtkplYWT102C5sn/EZTY4XROkxONgruHPU72l+gW+FfF8tf2cFjeHaRbWpOYa/uRBz/Xq1Pg==} + engines: {node: '>=18'} + cpu: [ia32] + os: [linux] + + '@esbuild/linux-ia32@0.28.2': + resolution: {integrity: sha512-CYbnj78HsIeA+DhgUKgFCfvNsTHFhMMrinUrMZpDXJXKN8T3XViTZ/+wtHeVxEWY8ewSzTFN+nRmSwO2tZaLUQ==} + engines: {node: '>=18'} + cpu: [ia32] + os: [linux] + + '@esbuild/linux-loong64@0.27.7': + resolution: {integrity: sha512-a4POruNM2oWsD4WKvBSEKGIiWQF8fZOAsycHOt6JBpZ+JN2n2JH9WAv56SOyu9X5IqAjqSIPTaJkqN8F7XOQ5Q==} + engines: {node: '>=18'} + cpu: [loong64] + os: [linux] + + '@esbuild/linux-loong64@0.28.2': + resolution: {integrity: sha512-buwkd8nsph4R+ajRvw0qM5Hja/TXQow3ptzWO2EbG/cqcIkHloRrdlBtQlshyYGTNFvfkfJ5tpPLVkY4DtsPfQ==} + engines: {node: '>=18'} + cpu: [loong64] + os: [linux] + + '@esbuild/linux-mips64el@0.27.7': + resolution: {integrity: sha512-KabT5I6StirGfIz0FMgl1I+R1H73Gp0ofL9A3nG3i/cYFJzKHhouBV5VWK1CSgKvVaG4q1RNpCTR2LuTVB3fIw==} + engines: {node: '>=18'} + cpu: [mips64el] + os: [linux] + + '@esbuild/linux-mips64el@0.28.2': + resolution: {integrity: sha512-ZVykbDyk7519VwiNb9Lcj9m8XM6v5V9uKPvrEMkkEedVewf+0itkhahp4HDpgERXhwLRpWFypsGbG/J8s0QjJA==} + engines: {node: '>=18'} + cpu: [mips64el] + os: [linux] + + '@esbuild/linux-ppc64@0.27.7': + resolution: {integrity: sha512-gRsL4x6wsGHGRqhtI+ifpN/vpOFTQtnbsupUF5R5YTAg+y/lKelYR1hXbnBdzDjGbMYjVJLJTd2OFmMewAgwlQ==} + engines: {node: '>=18'} + cpu: [ppc64] + os: [linux] + + '@esbuild/linux-ppc64@0.28.2': + resolution: {integrity: sha512-CAXl+Dtd9UUuJd8pKKdwh6MLm3MUMiqMPmhZ3tTSXPqfyQ3vDl6R5hZdZ/kYojK4ofXtdfSv1tFq8XzWx3heNQ==} + engines: {node: '>=18'} + cpu: [ppc64] + os: [linux] + + '@esbuild/linux-riscv64@0.27.7': + resolution: {integrity: sha512-hL25LbxO1QOngGzu2U5xeXtxXcW+/GvMN3ejANqXkxZ/opySAZMrc+9LY/WyjAan41unrR3YrmtTsUpwT66InQ==} + engines: {node: '>=18'} + cpu: [riscv64] + os: [linux] + + '@esbuild/linux-riscv64@0.28.2': + resolution: {integrity: sha512-GeXCej4IQtU1B+QlDV8W/RRvbzI3O/Stss+/bCXv4lZls5WGRtu2a+3JkA3i4qIUlMXpcHebWpF8AkJhATowuA==} + engines: {node: '>=18'} + cpu: [riscv64] + os: [linux] + + '@esbuild/linux-s390x@0.27.7': + resolution: {integrity: sha512-2k8go8Ycu1Kb46vEelhu1vqEP+UeRVj2zY1pSuPdgvbd5ykAw82Lrro28vXUrRmzEsUV0NzCf54yARIK8r0fdw==} + engines: {node: '>=18'} + cpu: [s390x] + os: [linux] + + '@esbuild/linux-s390x@0.28.2': + resolution: {integrity: sha512-3H1weTYZPxt/WOhByszQZybS9w5lKzUn1FDMsgEChbHWQwHYQQRfBxgCcZvPhjHfKyJjIievvMmEUawJrdY9Dg==} + engines: {node: '>=18'} + cpu: [s390x] + os: [linux] + + '@esbuild/linux-x64@0.27.7': + resolution: {integrity: sha512-hzznmADPt+OmsYzw1EE33ccA+HPdIqiCRq7cQeL1Jlq2gb1+OyWBkMCrYGBJ+sxVzve2ZJEVeePbLM2iEIZSxA==} + engines: {node: '>=18'} + cpu: [x64] + os: [linux] + + '@esbuild/linux-x64@0.28.2': + resolution: {integrity: sha512-4xTZr1FUmSoQW4XIWmit3tzQrUTZM+N3P0XV8xROKYF50XfI7xeO90+1bZvNwxIufQ9hDQVRJH5YhgPVF8A/HQ==} + engines: {node: '>=18'} + cpu: [x64] + os: [linux] + + '@esbuild/netbsd-arm64@0.27.7': + resolution: {integrity: sha512-b6pqtrQdigZBwZxAn1UpazEisvwaIDvdbMbmrly7cDTMFnw/+3lVxxCTGOrkPVnsYIosJJXAsILG9XcQS+Yu6w==} + engines: {node: '>=18'} + cpu: [arm64] + os: [netbsd] + + '@esbuild/netbsd-arm64@0.28.2': + resolution: {integrity: sha512-sSATRjPeDBg3pdgHoQfoYBob11Kk1FGa9lui5RIHZCoCkJa9QKlvl3/vKz2usCmYYjs7ymJR/2Nnsqe+Hjt5nw==} + engines: {node: '>=18'} + cpu: [arm64] + os: [netbsd] + + '@esbuild/netbsd-x64@0.27.7': + resolution: {integrity: sha512-OfatkLojr6U+WN5EDYuoQhtM+1xco+/6FSzJJnuWiUw5eVcicbyK3dq5EeV/QHT1uy6GoDhGbFpprUiHUYggrw==} + engines: {node: '>=18'} + cpu: [x64] + os: [netbsd] + + '@esbuild/netbsd-x64@0.28.2': + resolution: {integrity: sha512-lqnzCV+mM0gIADaKihiCg6ifgfU2L3h5E33rNQBN1Y4MaVGnzryzmvvf7UHxprpQdE8hpqLolJ9Rl+SkIRDpyw==} + engines: {node: '>=18'} + cpu: [x64] + os: [netbsd] + + '@esbuild/openbsd-arm64@0.27.7': + resolution: {integrity: sha512-AFuojMQTxAz75Fo8idVcqoQWEHIXFRbOc1TrVcFSgCZtQfSdc1RXgB3tjOn/krRHENUB4j00bfGjyl2mJrU37A==} + engines: {node: '>=18'} + cpu: [arm64] + os: [openbsd] + + '@esbuild/openbsd-arm64@0.28.2': + resolution: {integrity: sha512-AL2qJILH7lNjrDmCQDvdxMfAUIv8KMNZOvrwAQ8i8//ntL9FflhOyMJ8OZSMBb8/AWXe3/5v5S20y3zCoZWKoQ==} + engines: {node: '>=18'} + cpu: [arm64] + os: [openbsd] + + '@esbuild/openbsd-x64@0.27.7': + resolution: {integrity: sha512-+A1NJmfM8WNDv5CLVQYJ5PshuRm/4cI6WMZRg1by1GwPIQPCTs1GLEUHwiiQGT5zDdyLiRM/l1G0Pv54gvtKIg==} + engines: {node: '>=18'} + cpu: [x64] + os: [openbsd] + + '@esbuild/openbsd-x64@0.28.2': + resolution: {integrity: sha512-QtiuPytchRyC4rwUKhexJdQKvDuZ6hWloi3igqPQNUJCS1/v9EiO3UTOXR6A3FoMo4fnAKbWJdqaIwhOzh8qEw==} + engines: {node: '>=18'} + cpu: [x64] + os: [openbsd] + + '@esbuild/openharmony-arm64@0.27.7': + resolution: {integrity: sha512-+KrvYb/C8zA9CU/g0sR6w2RBw7IGc5J2BPnc3dYc5VJxHCSF1yNMxTV5LQ7GuKteQXZtspjFbiuW5/dOj7H4Yw==} + engines: {node: '>=18'} + cpu: [arm64] + os: [openharmony] + + '@esbuild/openharmony-arm64@0.28.2': + resolution: {integrity: sha512-WkhYDmpTjLvGlScA1rwjRUmhl4k8oXR3cIbtqWmELgU/dFeHHlEllxDvdWcNJV9rbzCexB5vz8gtNewWLgCT7Q==} + engines: {node: '>=18'} + cpu: [arm64] + os: [openharmony] + + '@esbuild/sunos-x64@0.27.7': + resolution: {integrity: sha512-ikktIhFBzQNt/QDyOL580ti9+5mL/YZeUPKU2ivGtGjdTYoqz6jObj6nOMfhASpS4GU4Q/Clh1QtxWAvcYKamA==} + engines: {node: '>=18'} + cpu: [x64] + os: [sunos] + + '@esbuild/sunos-x64@0.28.2': + resolution: {integrity: sha512-GPMSkTOtMnv2U2F8gxe4Io6qmVs+YKyp832Etqqxr0hFngmXQ3rzwytelm3GIn7T4VviRUlf3sOgBOiTdvaf7g==} + engines: {node: '>=18'} + cpu: [x64] + os: [sunos] + + '@esbuild/win32-arm64@0.27.7': + resolution: {integrity: sha512-7yRhbHvPqSpRUV7Q20VuDwbjW5kIMwTHpptuUzV+AA46kiPze5Z7qgt6CLCK3pWFrHeNfDd1VKgyP4O+ng17CA==} + engines: {node: '>=18'} + cpu: [arm64] + os: [win32] + + '@esbuild/win32-arm64@0.28.2': + resolution: {integrity: sha512-PIhhEkE9uPBleRBrQEJpUn7MBnibZzbGzYWPmY3x+YoVg/95zbjB4CxPPOQ8l5tYYM4mMaCthF8/1DIfBQQyWQ==} + engines: {node: '>=18'} + cpu: [arm64] + os: [win32] + + '@esbuild/win32-ia32@0.27.7': + resolution: {integrity: sha512-SmwKXe6VHIyZYbBLJrhOoCJRB/Z1tckzmgTLfFYOfpMAx63BJEaL9ExI8x7v0oAO3Zh6D/Oi1gVxEYr5oUCFhw==} + engines: {node: '>=18'} + cpu: [ia32] + os: [win32] + + '@esbuild/win32-ia32@0.28.2': + resolution: {integrity: sha512-YmJbfTlvU7Sdn9BB+4PRES4oB6pxgS37MAONj+hBr/cpXS1aBPKXxNnDbu+QCWPj0o9dgyxeq79g6c5P8KeuYA==} + engines: {node: '>=18'} + cpu: [ia32] + os: [win32] + + '@esbuild/win32-x64@0.27.7': + resolution: {integrity: sha512-56hiAJPhwQ1R4i+21FVF7V8kSD5zZTdHcVuRFMW0hn753vVfQN8xlx4uOPT4xoGH0Z/oVATuR82AiqSTDIpaHg==} + engines: {node: '>=18'} + cpu: [x64] + os: [win32] + + '@esbuild/win32-x64@0.28.2': + resolution: {integrity: sha512-5ebpxr3nWMzrL/rnUI755Jkuee0bHL/Gq0WTF9lvcpv73wAp5eu8MfBUgWK9bhWvZjj7yX8etf/8tI8Ney695g==} + engines: {node: '>=18'} + cpu: [x64] + os: [win32] + + '@jridgewell/gen-mapping@0.3.13': + resolution: {integrity: sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==} + + '@jridgewell/resolve-uri@3.1.2': + resolution: {integrity: sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw==} + engines: {node: '>=6.0.0'} + + '@jridgewell/sourcemap-codec@1.6.0': + resolution: {integrity: sha512-T7jf+5zgsZHwNJ4lvQ7/aezbyk0nNX+zJVWpmHA7VYsEx7a7qr5Rg5IbtJFqkgze5Y2sruq1RUY8Q837Od7iFw==} + + '@jridgewell/trace-mapping@0.3.31': + resolution: {integrity: sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==} + + '@napi-rs/lzma-linux-x64-gnu@1.5.1': + resolution: {integrity: sha512-oTXEIha4SsuXdTA4Iyskj0kpdx2yVXdhd75c2v3xGrHFfVMsbhTPZU/nMPL4sWKo4pBHm3aucLaqGlF696dTyQ==} + engines: {node: ^22.20 || ^24.12 || >=25} + cpu: [x64] + os: [linux] + libc: [glibc] + + '@rollup/rollup-android-arm-eabi@4.64.7': + resolution: {integrity: sha512-GFF2tHk0NqYVObrK0+/ps/LEXsV/v2q6Lyiy7mZr+Ej24fpQScvqK2ziz9NHeEQYnKWKym38Ml7dZbPtC9tuAw==} + cpu: [arm] + os: [android] + + '@rollup/rollup-android-arm64@4.64.7': + resolution: {integrity: sha512-BOymcXJFYYXqluDidj6sLelNjXruC3u1BqVxLn2OPCXH9xaPSOTH+MS4uychsGI1NU7tJoiv9KvqijDHFe+hhQ==} + cpu: [arm64] + os: [android] + + '@rollup/rollup-darwin-arm64@4.64.7': + resolution: {integrity: sha512-jgVpGweIbls2H5jfBB70tnX52abMd5pJr3xmtoDILx6JdmGl+Un+tMb9R+MbHPVB46tTBpY8emtn7M71SnQhEQ==} + cpu: [arm64] + os: [darwin] + + '@rollup/rollup-darwin-x64@4.64.7': + resolution: {integrity: sha512-AVMpt94mrXH7WjyWKmd9nT4RC5yM0AN3Lt6A89xVIsuLDi3dDX6dxMytLbW36FlppY9JsoqUoS3LUUpUrYxiMQ==} + cpu: [x64] + os: [darwin] + + '@rollup/rollup-freebsd-arm64@4.64.7': + resolution: {integrity: sha512-Nje29c/MyajxRNFGU05oyeTDS6SWdMg/wWMSlZSBTBju5BKz2NVc0kFiFSQGzRELNiB+CBk/ZvAVsWsENgjGag==} + cpu: [arm64] + os: [freebsd] + + '@rollup/rollup-freebsd-x64@4.64.7': + resolution: {integrity: sha512-xxgN5/kMzCNz/NYJ4t49B0/WS6InNlWv6xvgwUQqNnWM9WJMMwMUGWrAxdNEgYWdHL0sMnuh8I87StGrIMTA+w==} + cpu: [x64] + os: [freebsd] + + '@rollup/rollup-linux-arm-gnueabihf@4.64.7': + resolution: {integrity: sha512-5U/9gVn/X657I8+NHoHO9Wnfyd0lEXwqbslobRoB9TC8SwYzOitcIr6TLE7531Q/NWs3L4mb+YOGzlf3vHIeYw==} + cpu: [arm] + os: [linux] + libc: [glibc] + + '@rollup/rollup-linux-arm-musleabihf@4.64.7': + resolution: {integrity: sha512-hcDS0a/Y8ggM3dcOfzH7UpSaogCViWTtNDSl0EXpWIj8g8/r86ITqptjil2hayTXS0Umzb+Yt3mjgpQVFVf2+A==} + cpu: [arm] + os: [linux] + libc: [musl] + + '@rollup/rollup-linux-arm64-gnu@4.64.7': + resolution: {integrity: sha512-awbAr4N7Qd/yjZAQBEEPzgcPveYG2Gi5L6gq1/Bp+3SMRjn5F68Wv1zKsntzvvoioBR6yKvhLtWqu11am1TsxA==} + cpu: [arm64] + os: [linux] + libc: [glibc] + + '@rollup/rollup-linux-arm64-musl@4.64.7': + resolution: {integrity: sha512-Ut7xQJfeLhzu2yN3rHzVvL9ccyPJvN4aCdP1V7hJ5lGxX/eCzgB9SiMk/Md3J+rWWZgvoYsDW4BUFggJr32pBA==} + cpu: [arm64] + os: [linux] + libc: [musl] + + '@rollup/rollup-linux-loong64-gnu@4.64.7': + resolution: {integrity: sha512-trqoni/vHkuM5oo5iQhoKPh6fjjD7E80pMgCXWgGxU8vSRPUlwvG8GeYvEEcB6eJAxNa99WH9iVHPHelgtKiHA==} + cpu: [loong64] + os: [linux] + libc: [glibc] + + '@rollup/rollup-linux-loong64-musl@4.64.7': + resolution: {integrity: sha512-gvlUy0XFIwtUWiQmhc+4GT2Knm3KosinpUVhr7CtTFA78gFDnE5NWwMPzHrJsYpKkh2LlY7yvBIboN1FAV6Xeg==} + cpu: [loong64] + os: [linux] + libc: [musl] + + '@rollup/rollup-linux-ppc64-gnu@4.64.7': + resolution: {integrity: sha512-8yS/kSXTJ/47Z62JW+C2AJ3vd/8zZw016vFjtNpZwmZlwJumLua4s1i78dkYuSSqQ+0aIMxJcBp/4uIwp2ehjA==} + cpu: [ppc64] + os: [linux] + libc: [glibc] + + '@rollup/rollup-linux-ppc64-musl@4.64.7': + resolution: {integrity: sha512-mxVSKvvty0hOP1AFo/l4Kpv3i3HK0zx4rYUrirZDToYlWhS3xzmIkq3vUJ7rtrM2cPyXveDyDT6M1FfXZ9hTRA==} + cpu: [ppc64] + os: [linux] + libc: [musl] + + '@rollup/rollup-linux-riscv64-gnu@4.64.7': + resolution: {integrity: sha512-lEse1zra1CagM+uHH2bl28GInDsZZ45NJzHei2OMaXQFzZWKBO1hFo07rS0xB6ayV4WU360UrSmPDSyCMpYsUQ==} + cpu: [riscv64] + os: [linux] + libc: [glibc] + + '@rollup/rollup-linux-riscv64-musl@4.64.7': + resolution: {integrity: sha512-d/cqCoiG9WMZ5VuJZ5Qdlv9g97QyoGwnmvHccvGhKSacTy7xFrAS+jMcJSBdMNOw1oacSkl6p7K6O/DRYTKL3g==} + cpu: [riscv64] + os: [linux] + libc: [musl] + + '@rollup/rollup-linux-s390x-gnu@4.64.7': + resolution: {integrity: sha512-kVEA//CeDOZZP1XCH2JBgqXb9+ALlwuA/4civaazMVdgzaSQC9WM38w65zG/LLk0lA7F9Od19CSBZW7J9mCbNg==} + cpu: [s390x] + os: [linux] + libc: [glibc] + + '@rollup/rollup-linux-x64-gnu@4.64.7': + resolution: {integrity: sha512-14oLfDsOSr0TSntT/34jRg/V8B2XW/Jomh9zX9CjIM9gwwNlI1OCymS/GVn0DV32h67eRf7U0GGem0rvfXgNwQ==} + cpu: [x64] + os: [linux] + libc: [glibc] + + '@rollup/rollup-linux-x64-musl@4.64.7': + resolution: {integrity: sha512-dFVtPfjZqAnUPpPY4Bfg+87yo1IBW9HiY8iCFtg49Pn46MfZTebSHZh8PL6SOsgyt9LQP8Nj/sE6jm3aWpySmQ==} + cpu: [x64] + os: [linux] + libc: [musl] + + '@rollup/rollup-openbsd-x64@4.64.7': + resolution: {integrity: sha512-Dm9CZ2SM3u2NgvBj4awW/gGqkOwzH8f6MSYIJE6+bKcKLI+T0+LhpD6vtetTfrIIx+rLXDf5RJqBP6v7V5O+wg==} + cpu: [x64] + os: [openbsd] + + '@rollup/rollup-openharmony-arm64@4.64.7': + resolution: {integrity: sha512-oaxsvS2mL6+78741CT2HCORBvy1USg3Wqpt5qvZwy3aMXRMl4zxyDo8maeFRWJkF3dp0Wbi71fKH5Gr9YI3vYA==} + cpu: [arm64] + os: [openharmony] + + '@rollup/rollup-win32-arm64-msvc@4.64.7': + resolution: {integrity: sha512-UKZz1Y4mxUqkWh60uH7B1aZn6UcAimPk4ipRktpb3wfZq7JrzZYyAXTwiafkb3T+kd5P3Sf+vzcRdsDks4WtjQ==} + cpu: [arm64] + os: [win32] + + '@rollup/rollup-win32-ia32-msvc@4.64.7': + resolution: {integrity: sha512-WWkSuC1J+mZJnsDh0mRaCAFSWfQdl6y3fDjZeWElI6zHo8CcOSG5lY3yurKJjEq+FrX0WD60sFrVVpbHqzMmrw==} + cpu: [ia32] + os: [win32] + + '@rollup/rollup-win32-x64-gnu@4.64.7': + resolution: {integrity: sha512-u9bp2yXCI3/BUjUxY6uWknAY32pfvvDCPysyUTySotx0dj3ydlFHrRQYfn/b/kZgnfjyEwySMtI7Ga054uiysQ==} + cpu: [x64] + os: [win32] + + '@rollup/rollup-win32-x64-msvc@4.64.7': + resolution: {integrity: sha512-EhX6N3XPMEnVU5cRfNgzPOKcSucsbu3nNNWuoAaEuKIqmddBKHN5b9d4rPe65T1I+FaOT176JKiZNsUBp3rqbA==} + cpu: [x64] + os: [win32] + + '@tabularis/explain@file:tests/host': + resolution: {directory: tests/host, type: directory} + + '@types/chai@5.2.3': + resolution: {integrity: sha512-Mw558oeA9fFbv65/y4mHtXDs9bPnFMZAL/jxdPFUpOHHIXX91mcgEHbS5Lahr+pwZFR8A7GQleRWeI6cGFC2UA==} + + '@types/deep-eql@4.0.2': + resolution: {integrity: sha512-c9h9dVVMigMPc4bwTvC5dxqtqJZwQPePsWjPlpSOnojbor6pGqdk541lfA7AqFQr5pB1BRdq0juY9db81BwyFw==} + + '@types/estree@1.0.9': + resolution: {integrity: sha512-GhdPgy1el4/ImP05X05Uw4cw2/M93BCUmnEvWZNStlCzEKME4Fkk+YpoA5OiHNQmoS7Cafb8Xa3Pya8m1Qrzeg==} + + '@types/node@26.6.5': + resolution: {integrity: sha512-uMzsv9stDdofIZBx4UGh4Z9ZgXnGC9seRDscilBcTHAeOFk67MZg/k9bWVfO+bk4qZmG+l52hRnLwUiH2LnJTg==} + + '@vitest/expect@3.2.7': + resolution: {integrity: sha512-E8eBXaKibuvH2pSZErOjdVb5vF4PbKYcrnluBTYxEk1l/VhhwZg1kZQsdtjq+CsF5CFydf2Rdkz7jDHKSisi3w==} + + '@vitest/mocker@3.2.7': + resolution: {integrity: sha512-Trr0hYO9CM3Wj6ksWHRhK9IZpIY6wTMO5u/MqXurMxT57sWBaOPEtP3Oq60ihZuh5JsiagKfz95OcxdEP6dBrA==} + peerDependencies: + msw: ^2.4.9 + vite: ^5.0.0 || ^6.0.0 || ^7.0.0-0 + peerDependenciesMeta: + msw: + optional: true + vite: + optional: true + + '@vitest/pretty-format@3.2.7': + resolution: {integrity: sha512-KUHlwqVu0sRlhCdyPdQ/wBoTfRahjUky1MubOmYw9fWfIZy1gNoHpuaaQBPAaMaVYdQYHJLurzj8ECCj5OwTqA==} + + '@vitest/runner@3.2.7': + resolution: {integrity: sha512-sB9y4ovltoQP+WaUPwmSxO9WIg9Ig694Di5PalVPsYHklAdE027mehpWF2SQSVq+k6sFgaivbTjTJwZLSHbedA==} + + '@vitest/snapshot@3.2.7': + resolution: {integrity: sha512-7C+MwShwtBSI5Buwoyg3s/iY1eHL9PKAf+O1wVh/TdnjXUtkoL/9YQtre90i4MtNXM6edP1wJ2zOBpfCyhIS7g==} + + '@vitest/spy@3.2.7': + resolution: {integrity: sha512-Q2eQGI6d2L/hBtZ0qNuKcAGid68XK6cv1xsoaIma6PaJhHPoqcEJhYpXZ/5myCMqkNgtP6UKuBhbc0nHKnrkuQ==} + + '@vitest/utils@3.2.7': + resolution: {integrity: sha512-x6BDOd7dyo3PFLY3I9/HJ25X/6OurhGXk2/B9gOZNPF7XDVjeBK4k01lQE5uvDpbuheErh91qYuE1E2OEjK3Rw==} + + acorn@8.19.0: + resolution: {integrity: sha512-oJlA3XiRm7Cyk6qFD2Jn8ak9B7jSy0qy00ADO3+8dpT0LSjFihQYv4C02LFCSYJV3Q37xYwwRy5m+IpIiUzqWw==} + engines: {node: '>=0.4.0'} + hasBin: true + + any-promise@1.3.0: + resolution: {integrity: sha512-7UvmKalWRt1wgjL1RrGxoSJW/0QZFIegpeGvZG9kjp8vrRu55XTHbwnqq2GpXm9uLbcuhxm3IqX9OB4MZR1b2A==} + + assertion-error@2.0.1: + resolution: {integrity: sha512-Izi8RQcffqCeNVgFigKli1ssklIbpHnCYc6AknXGYoB6grJqyeby7jv12JUQgmTAnIDnbck1uxksT4dzN3PWBA==} + engines: {node: '>=12'} + + bundle-require@5.1.0: + resolution: {integrity: sha512-3WrrOuZiyaaZPWiEt4G3+IffISVC9HYlWueJEBWED4ZH4aIAC2PnkdnuRrR94M+w6yGWn4AglWtJtBI8YqvgoA==} + engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0} + peerDependencies: + esbuild: '>=0.18' + + cac@6.7.14: + resolution: {integrity: sha512-b6Ilus+c3RrdDk+JhLKUAQfzzgLEPy6wcXqS7f/xe1EETvsDP6GORG7SFuOs6cID5YkqchW/LXZbX5bc8j7ZcQ==} + engines: {node: '>=8'} + + chai@5.3.3: + resolution: {integrity: sha512-4zNhdJD/iOjSH0A05ea+Ke6MU5mmpQcbQsSOkgdaUMJ9zTlDTD/GYlwohmIE2u0gaxHYiVHEn1Fw9mZ/ktJWgw==} + engines: {node: '>=18'} + + check-error@2.1.3: + resolution: {integrity: sha512-PAJdDJusoxnwm1VwW07VWwUN1sl7smmC3OKggvndJFadxxDRyFJBX/ggnu/KE4kQAB7a3Dp8f/YXC1FlUprWmA==} + engines: {node: '>= 16'} + + chokidar@4.0.3: + resolution: {integrity: sha512-Qgzu8kfBvo+cA4962jnP1KkS6Dop5NS6g7R5LFYJr4b8Ub94PPQXUksCw9PvXoeXPRRddRNC5C1JQUR2SMGtnA==} + engines: {node: '>= 14.16.0'} + + commander@4.1.1: + resolution: {integrity: sha512-NOKm8xhkzAjzFx8B2v5OAHT+u5pRQc2UCa2Vq9jYL/31o2wi9mxBA7LIFs3sV5VSC49z6pEhfbMULvShKj26WA==} + engines: {node: '>= 6'} + + confbox@0.1.8: + resolution: {integrity: sha512-RMtmw0iFkeR4YV+fUOSucriAQNb9g8zFR52MWCtl+cCZOFRNL6zeB395vPzFhEjjn4fMxXudmELnl/KF/WrK6w==} + + consola@3.4.2: + resolution: {integrity: sha512-5IKcdX0nnYavi6G7TtOhwkYzyjfJlatbjMjuLSfE2kYT5pMDOilZ4OvMhi637CcDICTmz3wARPoyhqyX1Y+XvA==} + engines: {node: ^14.18.0 || >=16.10.0} + + debug@4.4.3: + resolution: {integrity: sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==} + engines: {node: '>=6.0'} + peerDependencies: + supports-color: '*' + peerDependenciesMeta: + supports-color: + optional: true + + deep-eql@5.0.2: + resolution: {integrity: sha512-h5k/5U50IJJFpzfL6nO9jaaumfjO/f2NjK/oYB2Djzm4p9L+3T9qWpZqZ2hAbLPuuYq9wrU08WQyBTL5GbPk5Q==} + engines: {node: '>=6'} + + es-module-lexer@1.7.0: + resolution: {integrity: sha512-jEQoCwk8hyb2AZziIOLhDqpm5+2ww5uIE6lkO/6jcOCusfk6LhMHpXXfBLXTZ7Ydyt0j4VoUQv6uGNYbdW+kBA==} + + esbuild@0.27.7: + resolution: {integrity: sha512-IxpibTjyVnmrIQo5aqNpCgoACA/dTKLTlhMHihVHhdkxKyPO1uBBthumT0rdHmcsk9uMonIWS0m4FljWzILh3w==} + engines: {node: '>=18'} + hasBin: true + + esbuild@0.28.2: + resolution: {integrity: sha512-HKVLS8dvII+xoKW9kmqxbRKrnWEXfJJr/FZhhJmiqIB0e053QNYFqOBouTMO/k5sID4MvCiUCvv8b9M4h32wIA==} + engines: {node: '>=18'} + hasBin: true + + estree-walker@3.0.3: + resolution: {integrity: sha512-7RUKfXgSMMkzt6ZuXmqapOurLGPPfgj6l9uRZ7lRGolvk0y2yocc35LdcxKC5PQZdn2DMqioAQ2NoWcrTKmm6g==} + + expect-type@1.4.0: + resolution: {integrity: sha512-KfYbmpRm0VbLjEvVa9yGwCi9GI34xvi7A/HXYWQO65CSD2u3MczUJSuwXKFIxlGsgBQizV9q5J9NHj4VG0n+pA==} + engines: {node: '>=12.0.0'} + + fdir@6.5.0: + resolution: {integrity: sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==} + engines: {node: '>=12.0.0'} + peerDependencies: + picomatch: ^3 || ^4 + peerDependenciesMeta: + picomatch: + optional: true + + fix-dts-default-cjs-exports@1.0.1: + resolution: {integrity: sha512-pVIECanWFC61Hzl2+oOCtoJ3F17kglZC/6N94eRWycFgBH35hHx0Li604ZIzhseh97mf2p0cv7vVrOZGoqhlEg==} + + fsevents@2.3.3: + resolution: {integrity: sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==} + engines: {node: ^8.16.0 || ^10.6.0 || >=11.0.0} + os: [darwin] + + joycon@3.1.1: + resolution: {integrity: sha512-34wB/Y7MW7bzjKRjUKTa46I2Z7eV62Rkhva+KkopW7Qvv/OSWBqvkSY7vusOPrNuZcUG3tApvdVgNB8POj3SPw==} + engines: {node: '>=10'} + + js-tokens@9.0.1: + resolution: {integrity: sha512-mxa9E9ITFOt0ban3j6L5MpjwegGz6lBQmM1IJkWeBZGcMxto50+eWdjC/52xDbS2vy0k7vIMK0Fe2wfL9OQSpQ==} + + lilconfig@3.1.4: + resolution: {integrity: sha512-zJZyBmQhArDRefuWTypZxwDpFDIYmrj8ypSBvwGD5P13Y95s6DjzKgu4DVdGbeQa5PDEsAOCAPNk4X+QOR006w==} + engines: {node: '>=14'} + + lines-and-columns@1.2.4: + resolution: {integrity: sha512-7ylylesZQ/PV29jhEDl3Ufjo6ZX7gCqJr5F7PKrqc93v7fzSymt1BpwEU8nAUXs8qzzvqhbjhK5QZg6Mt/HkBg==} + + load-tsconfig@0.2.5: + resolution: {integrity: sha512-IXO6OCs9yg8tMKzfPZ1YmheJbZCiEsnBdcB03l0OcfK9prKnJb96siuHCr5Fl37/yo9DnKU+TLpxzTUspw9shg==} + engines: {node: ^12.20.0 || ^14.13.1 || >=16.0.0} + + loupe@3.2.1: + resolution: {integrity: sha512-CdzqowRJCeLU72bHvWqwRBBlLcMEtIvGrlvef74kMnV2AolS9Y8xUv1I0U/MNAWMhBlKIoyuEgoJ0t/bbwHbLQ==} + + magic-string@0.30.21: + resolution: {integrity: sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ==} + + mlly@1.8.2: + resolution: {integrity: sha512-d+ObxMQFmbt10sretNDytwt85VrbkhhUA/JBGm1MPaWJ65Cl4wOgLaB1NYvJSZ0Ef03MMEU/0xpPMXUIQ29UfA==} + + ms@2.1.3: + resolution: {integrity: sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==} + + mz@2.7.0: + resolution: {integrity: sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q==} + + nanoid@3.3.20: + resolution: {integrity: sha512-uKdg2G3GNCKQn9byYOpxbGqrT2fGO5KRt5J/8b3pok8rT6qxGWF6hxMyJiEYtAf+FVyYuD9hRaDqX5uPFYJ4ZQ==} + engines: {node: ^10 || ^12 || ^13.7 || ^14 || >=15.0.1} + hasBin: true + + object-assign@4.1.1: + resolution: {integrity: sha512-rJgTQnkUnH1sFw8yT6VSU3zD3sWmu6sZhIseY8VX+GRu3P6F7Fu+JNDoXfklElbLJSnc3FUQHVe4cU5hj+BcUg==} + engines: {node: '>=0.10.0'} + + pathe@2.0.3: + resolution: {integrity: sha512-WUjGcAqP1gQacoQe+OBJsFA7Ld4DyXuUIjZ5cc75cLHvJ7dtNsTugphxIADwspS+AraAUePCKrSVtPLFj/F88w==} + + pathval@2.0.1: + resolution: {integrity: sha512-//nshmD55c46FuFw26xV/xFAaB5HF9Xdap7HJBBnrKdAd6/GxDBaNA1870O79+9ueg61cZLSVc+OaFlfmObYVQ==} + engines: {node: '>= 14.16'} + + picocolors@1.1.1: + resolution: {integrity: sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==} + + picomatch@4.0.7: + resolution: {integrity: sha512-qcJu88Q2IWqJsDD529JKMdwGm/dvInW4HvQnRwiH9JtihJvzGOscDtHE3x1pBKeUOTysQ8kVmLnJ2kJu7yhcGA==} + engines: {node: '>=12'} + + pirates@4.0.7: + resolution: {integrity: sha512-TfySrs/5nm8fQJDcBDuUng3VOUKsd7S+zqvbOTiGXHfxX4wK31ard+hoNuvkicM/2YFzlpDgABOevKSsB4G/FA==} + engines: {node: '>= 6'} + + pkg-types@1.3.1: + resolution: {integrity: sha512-/Jm5M4RvtBFVkKWRu2BLUTNP8/M2a+UwuAX+ae4770q1qVGtfjG+WTCupoZixokjmHiry8uI+dlY8KXYV5HVVQ==} + + postcss-load-config@6.0.1: + resolution: {integrity: sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g==} + engines: {node: '>= 18'} + peerDependencies: + jiti: '>=1.21.0' + postcss: '>=8.0.9' + tsx: ^4.8.1 + yaml: ^2.4.2 + peerDependenciesMeta: + jiti: + optional: true + postcss: + optional: true + tsx: + optional: true + yaml: + optional: true + + postcss@8.5.30: + resolution: {integrity: sha512-zkd+e9S1AQHhmaBKye9zlFHWjOwVgqXl1q3ECYy8561K7qxR2cREf/STLskulBFd78k6AOFeA+ff1TtEzk9bAA==} + engines: {node: ^10 || ^12 || >=14} + + readdirp@4.1.2: + resolution: {integrity: sha512-GDhwkLfywWL2s6vEjyhri+eXmfH6j1L7JE27WhqLeYzoh/A3DBaYGEj2H/HFZCn/kMfim73FXxEJTw06WtxQwg==} + engines: {node: '>= 14.18.0'} + + resolve-from@5.0.0: + resolution: {integrity: sha512-qYg9KP24dD5qka9J47d0aVky0N+b4fTU89LN9iDnjB5waksiC49rvMB0PrUJQGoTmH50XPiqOvAjDfaijGxYZw==} + engines: {node: '>=8'} + + rollup@4.64.7: + resolution: {integrity: sha512-aCiJONoZp0lNGJ6sYqspLkh6UXaWEey/7Ko2TZyXQ4kgzAhDUxAgbt/oulXKmD+zFHXw1k+RF6ZZWOrjp7DZnQ==} + engines: {node: '>=18.0.0', npm: '>=8.0.0'} + hasBin: true + + siginfo@2.0.0: + resolution: {integrity: sha512-ybx0WO1/8bSBLEWXZvEd7gMW3Sn3JFlW3TvX1nREbDLRNQNaeNN8WK0meBwPdAaOI7TtRRRJn/Es1zhrrCHu7g==} + + source-map-js@1.2.2: + resolution: {integrity: sha512-KGj/8Y43x35aZVDtt+J4mK1hoLGHULMYfSkODJNQjNDC3oW1PqPoxMwo0pLUsWM/UEGzON/NxeHywEfNXNP3Vw==} + engines: {node: '>=0.10.0'} + + source-map@0.7.6: + resolution: {integrity: sha512-i5uvt8C3ikiWeNZSVZNWcfZPItFQOsYTUAOkcUPGd8DqDy1uOUikjt5dG+uRlwyvR108Fb9DOd4GvXfT0N2/uQ==} + engines: {node: '>= 12'} + + stackback@0.0.2: + resolution: {integrity: sha512-1XMJE5fQo1jGH6Y/7ebnwPOBEkIEnT4QF32d5R1+VXdXveM0IBMJt8zfaxX1P3QhVwrYe+576+jkANtSS2mBbw==} + + std-env@3.10.0: + resolution: {integrity: sha512-5GS12FdOZNliM5mAOxFRg7Ir0pWz8MdpYm6AY6VPkGpbA7ZzmbzNcBJQ0GPvvyWgcY7QAhCgf9Uy89I03faLkg==} + + strip-literal@3.1.0: + resolution: {integrity: sha512-8r3mkIM/2+PpjHoOtiAW8Rg3jJLHaV7xPwG+YRGrv6FP0wwk/toTpATxWYOW0BKdWwl82VT2tFYi5DlROa0Mxg==} + + sucrase@3.35.1: + resolution: {integrity: sha512-DhuTmvZWux4H1UOnWMB3sk0sbaCVOoQZjv8u1rDoTV0HTdGem9hkAZtl4JZy8P2z4Bg0nT+YMeOFyVr4zcG5Tw==} + engines: {node: '>=16 || 14 >=14.17'} + hasBin: true + + thenify-all@1.6.0: + resolution: {integrity: sha512-RNxQH/qI8/t3thXJDwcstUO4zeqo64+Uy/+sNVRBx4Xn2OX+OZ9oP+iJnNFqplFra2ZUVeKCSa2oVWi3T4uVmA==} + engines: {node: '>=0.8'} + + thenify@3.3.1: + resolution: {integrity: sha512-RVZSIV5IG10Hk3enotrhvz0T9em6cyHBLkH/YAZuKqd8hRkKhSfCGIcP2KUY0EPxndzANBmNllzWPwak+bheSw==} + + tinybench@2.9.0: + resolution: {integrity: sha512-0+DUvqWMValLmha6lr4kD8iAMK1HzV0/aKnCtWb9v9641TnP/MFb7Pc2bxoxQjTXAErryXVgUOfv2YqNllqGeg==} + + tinyexec@0.3.2: + resolution: {integrity: sha512-KQQR9yN7R5+OSwaK0XQoj22pwHoTlgYqmUscPYoknOoWCWfj/5/ABTMRi69FrKU5ffPVh5QcFikpWJI/P1ocHA==} + + tinyglobby@0.2.17: + resolution: {integrity: sha512-wXR/dYpcqKmfWpEdZjiKJOwCNFndD0DMnrW/cYjVGttEkBfVgcLFHoNrlj47mjOVic9yyNu65alsgF4NQyTa2g==} + engines: {node: '>=12.0.0'} + + tinypool@1.1.1: + resolution: {integrity: sha512-Zba82s87IFq9A9XmjiX5uZA/ARWDrB03OHlq+Vw1fSdt0I+4/Kutwy8BP4Y/y/aORMo61FQ0vIb5j44vSo5Pkg==} + engines: {node: ^18.0.0 || >=20.0.0} + + tinyrainbow@2.0.0: + resolution: {integrity: sha512-op4nsTR47R6p0vMUUoYl/a+ljLFVtlfaXkLQmqfLR1qHma1h/ysYk4hEXZ880bf2CYgTskvTa/e196Vd5dDQXw==} + engines: {node: '>=14.0.0'} + + tinyspy@4.0.6: + resolution: {integrity: sha512-u8KszXvGfU68hVcZpRHKG28T0krMuv2G5nDhiHaMLen/gIuFEgIJhaJuO69qjnXg5paSrbPMFfx3brNuN8eVSg==} + engines: {node: '>=14.0.0'} + + tree-kill@1.2.2: + resolution: {integrity: sha512-L0Orpi8qGpRG//Nd+H90vFB+3iHnue1zSSGmNOOCh1GLJ7rUKVwV2HvijphGQS2UmhUZewS9VgvxYIdgr+fG1A==} + hasBin: true + + ts-interface-checker@0.1.13: + resolution: {integrity: sha512-Y/arvbn+rrz3JCKl9C4kVNfTfSm2/mEp5FSz5EsZSANGPSlQrpRI5M4PKF+mJnE52jOO90PnPSc3Ur3bTQw0gA==} + + tsup@8.5.1: + resolution: {integrity: sha512-xtgkqwdhpKWr3tKPmCkvYmS9xnQK3m3XgxZHwSUjvfTjp7YfXe5tT3GgWi0F2N+ZSMsOeWeZFh7ZZFg5iPhing==} + engines: {node: '>=18'} + hasBin: true + peerDependencies: + '@microsoft/api-extractor': ^7.36.0 + '@swc/core': ^1 + postcss: ^8.4.12 + typescript: '>=4.5.0' + peerDependenciesMeta: + '@microsoft/api-extractor': + optional: true + '@swc/core': + optional: true + postcss: + optional: true + typescript: + optional: true + + typescript@5.9.3: + resolution: {integrity: sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==} + engines: {node: '>=14.17'} + hasBin: true + + ufo@1.6.4: + resolution: {integrity: sha512-JFNbkD1Svwe0KvGi8GOeLcP4kAWQ609twvCdcHxq1oSL8svv39ZuSvajcD8B+5D0eL4+s1Is2D/O6KN3qcTeRA==} + + undici-types@8.9.0: + resolution: {integrity: sha512-KTDyRTYX8sWmKXAikPHHSyc63CRPETMctyjKFupcC6OBLXT3xsN0e9aF7m+mIXutFWpUXuedtowG7iLOzp0kQg==} + + vite-node@3.2.4: + resolution: {integrity: sha512-EbKSKh+bh1E1IFxeO0pg1n4dvoOTt0UDiXMd/qn++r98+jPO1xtJilvXldeuQ8giIB5IkpjCgMleHMNEsGH6pg==} + engines: {node: ^18.0.0 || ^20.0.0 || >=22.0.0} + hasBin: true + + vite@7.3.7: + resolution: {integrity: sha512-AHEzTORMbi7/7aXp7w7KZ5+DCUdSeLsJz0grY8xHpc9DHdYb0vQu6DqJfm2tOw3gDExNgPgnJTXxM6syYFHaCA==} + engines: {node: ^20.19.0 || >=22.12.0} + hasBin: true + peerDependencies: + '@types/node': ^20.19.0 || >=22.12.0 + jiti: '>=1.21.0' + less: ^4.0.0 + lightningcss: ^1.21.0 + sass: ^1.70.0 + sass-embedded: ^1.70.0 + stylus: '>=0.54.8' + sugarss: ^5.0.0 + terser: ^5.16.0 + tsx: ^4.8.1 + yaml: ^2.4.2 + peerDependenciesMeta: + '@types/node': + optional: true + jiti: + optional: true + less: + optional: true + lightningcss: + optional: true + sass: + optional: true + sass-embedded: + optional: true + stylus: + optional: true + sugarss: + optional: true + terser: + optional: true + tsx: + optional: true + yaml: + optional: true + + vitest@3.2.7: + resolution: {integrity: sha512-KrxIJ62Fd89gfysR4WotlgZABiz2dqFPgqGzX7s+CwsqLFomRH7777ZcrOD6+WVAh7khPQP41A+BKbpcJFrdEg==} + engines: {node: ^18.0.0 || ^20.0.0 || >=22.0.0} + hasBin: true + peerDependencies: + '@edge-runtime/vm': '*' + '@types/debug': ^4.1.12 + '@types/node': ^18.0.0 || ^20.0.0 || >=22.0.0 + '@vitest/browser': 3.2.7 + '@vitest/ui': 3.2.7 + happy-dom: '*' + jsdom: '*' + peerDependenciesMeta: + '@edge-runtime/vm': + optional: true + '@types/debug': + optional: true + '@types/node': + optional: true + '@vitest/browser': + optional: true + '@vitest/ui': + optional: true + happy-dom: + optional: true + jsdom: + optional: true + + why-is-node-running@2.3.0: + resolution: {integrity: sha512-hUrmaWBdVDcxvYqnyh09zunKzROWjbZTiNy8dBEjkS7ehEDQibXJ7XvlmtbwuTclUiIyN+CyXQD4Vmko8fNm8w==} + engines: {node: '>=8'} + hasBin: true + +snapshots: + + '@esbuild/aix-ppc64@0.27.7': + optional: true + + '@esbuild/aix-ppc64@0.28.2': + optional: true + + '@esbuild/android-arm64@0.27.7': + optional: true + + '@esbuild/android-arm64@0.28.2': + optional: true + + '@esbuild/android-arm@0.27.7': + optional: true + + '@esbuild/android-arm@0.28.2': + optional: true + + '@esbuild/android-x64@0.27.7': + optional: true + + '@esbuild/android-x64@0.28.2': + optional: true + + '@esbuild/darwin-arm64@0.27.7': + optional: true + + '@esbuild/darwin-arm64@0.28.2': + optional: true + + '@esbuild/darwin-x64@0.27.7': + optional: true + + '@esbuild/darwin-x64@0.28.2': + optional: true + + '@esbuild/freebsd-arm64@0.27.7': + optional: true + + '@esbuild/freebsd-arm64@0.28.2': + optional: true + + '@esbuild/freebsd-x64@0.27.7': + optional: true + + '@esbuild/freebsd-x64@0.28.2': + optional: true + + '@esbuild/linux-arm64@0.27.7': + optional: true + + '@esbuild/linux-arm64@0.28.2': + optional: true + + '@esbuild/linux-arm@0.27.7': + optional: true + + '@esbuild/linux-arm@0.28.2': + optional: true + + '@esbuild/linux-ia32@0.27.7': + optional: true + + '@esbuild/linux-ia32@0.28.2': + optional: true + + '@esbuild/linux-loong64@0.27.7': + optional: true + + '@esbuild/linux-loong64@0.28.2': + optional: true + + '@esbuild/linux-mips64el@0.27.7': + optional: true + + '@esbuild/linux-mips64el@0.28.2': + optional: true + + '@esbuild/linux-ppc64@0.27.7': + optional: true + + '@esbuild/linux-ppc64@0.28.2': + optional: true + + '@esbuild/linux-riscv64@0.27.7': + optional: true + + '@esbuild/linux-riscv64@0.28.2': + optional: true + + '@esbuild/linux-s390x@0.27.7': + optional: true + + '@esbuild/linux-s390x@0.28.2': + optional: true + + '@esbuild/linux-x64@0.27.7': + optional: true + + '@esbuild/linux-x64@0.28.2': + optional: true + + '@esbuild/netbsd-arm64@0.27.7': + optional: true + + '@esbuild/netbsd-arm64@0.28.2': + optional: true + + '@esbuild/netbsd-x64@0.27.7': + optional: true + + '@esbuild/netbsd-x64@0.28.2': + optional: true + + '@esbuild/openbsd-arm64@0.27.7': + optional: true + + '@esbuild/openbsd-arm64@0.28.2': + optional: true + + '@esbuild/openbsd-x64@0.27.7': + optional: true + + '@esbuild/openbsd-x64@0.28.2': + optional: true + + '@esbuild/openharmony-arm64@0.27.7': + optional: true + + '@esbuild/openharmony-arm64@0.28.2': + optional: true + + '@esbuild/sunos-x64@0.27.7': + optional: true + + '@esbuild/sunos-x64@0.28.2': + optional: true + + '@esbuild/win32-arm64@0.27.7': + optional: true + + '@esbuild/win32-arm64@0.28.2': + optional: true + + '@esbuild/win32-ia32@0.27.7': + optional: true + + '@esbuild/win32-ia32@0.28.2': + optional: true + + '@esbuild/win32-x64@0.27.7': + optional: true + + '@esbuild/win32-x64@0.28.2': + optional: true + + '@jridgewell/gen-mapping@0.3.13': + dependencies: + '@jridgewell/sourcemap-codec': 1.6.0 + '@jridgewell/trace-mapping': 0.3.31 + + '@jridgewell/resolve-uri@3.1.2': {} + + '@jridgewell/sourcemap-codec@1.6.0': {} + + '@jridgewell/trace-mapping@0.3.31': + dependencies: + '@jridgewell/resolve-uri': 3.1.2 + '@jridgewell/sourcemap-codec': 1.6.0 + + '@napi-rs/lzma-linux-x64-gnu@1.5.1': + optional: true + + '@rollup/rollup-android-arm-eabi@4.64.7': + optional: true + + '@rollup/rollup-android-arm64@4.64.7': + optional: true + + '@rollup/rollup-darwin-arm64@4.64.7': + optional: true + + '@rollup/rollup-darwin-x64@4.64.7': + optional: true + + '@rollup/rollup-freebsd-arm64@4.64.7': + optional: true + + '@rollup/rollup-freebsd-x64@4.64.7': + optional: true + + '@rollup/rollup-linux-arm-gnueabihf@4.64.7': + optional: true + + '@rollup/rollup-linux-arm-musleabihf@4.64.7': + optional: true + + '@rollup/rollup-linux-arm64-gnu@4.64.7': + optional: true + + '@rollup/rollup-linux-arm64-musl@4.64.7': + optional: true + + '@rollup/rollup-linux-loong64-gnu@4.64.7': + optional: true + + '@rollup/rollup-linux-loong64-musl@4.64.7': + optional: true + + '@rollup/rollup-linux-ppc64-gnu@4.64.7': + optional: true + + '@rollup/rollup-linux-ppc64-musl@4.64.7': + optional: true + + '@rollup/rollup-linux-riscv64-gnu@4.64.7': + optional: true + + '@rollup/rollup-linux-riscv64-musl@4.64.7': + optional: true + + '@rollup/rollup-linux-s390x-gnu@4.64.7': + optional: true + + '@rollup/rollup-linux-x64-gnu@4.64.7': + optional: true + + '@rollup/rollup-linux-x64-musl@4.64.7': + optional: true + + '@rollup/rollup-openbsd-x64@4.64.7': + optional: true + + '@rollup/rollup-openharmony-arm64@4.64.7': + optional: true + + '@rollup/rollup-win32-arm64-msvc@4.64.7': + optional: true + + '@rollup/rollup-win32-ia32-msvc@4.64.7': + optional: true + + '@rollup/rollup-win32-x64-gnu@4.64.7': + optional: true + + '@rollup/rollup-win32-x64-msvc@4.64.7': + optional: true + + '@tabularis/explain@file:tests/host': {} + + '@types/chai@5.2.3': + dependencies: + '@types/deep-eql': 4.0.2 + assertion-error: 2.0.1 + + '@types/deep-eql@4.0.2': {} + + '@types/estree@1.0.9': {} + + '@types/node@26.6.5': + dependencies: + undici-types: 8.9.0 + + '@vitest/expect@3.2.7': + dependencies: + '@types/chai': 5.2.3 + '@vitest/spy': 3.2.7 + '@vitest/utils': 3.2.7 + chai: 5.3.3 + tinyrainbow: 2.0.0 + + '@vitest/mocker@3.2.7(vite@7.3.7(@types/node@26.6.5))': + dependencies: + '@vitest/spy': 3.2.7 + estree-walker: 3.0.3 + magic-string: 0.30.21 + optionalDependencies: + vite: 7.3.7(@types/node@26.6.5) + + '@vitest/pretty-format@3.2.7': + dependencies: + tinyrainbow: 2.0.0 + + '@vitest/runner@3.2.7': + dependencies: + '@vitest/utils': 3.2.7 + pathe: 2.0.3 + strip-literal: 3.1.0 + + '@vitest/snapshot@3.2.7': + dependencies: + '@vitest/pretty-format': 3.2.7 + magic-string: 0.30.21 + pathe: 2.0.3 + + '@vitest/spy@3.2.7': + dependencies: + tinyspy: 4.0.6 + + '@vitest/utils@3.2.7': + dependencies: + '@vitest/pretty-format': 3.2.7 + loupe: 3.2.1 + tinyrainbow: 2.0.0 + + acorn@8.19.0: {} + + any-promise@1.3.0: {} + + assertion-error@2.0.1: {} + + bundle-require@5.1.0(esbuild@0.27.7): + dependencies: + esbuild: 0.27.7 + load-tsconfig: 0.2.5 + + cac@6.7.14: {} + + chai@5.3.3: + dependencies: + assertion-error: 2.0.1 + check-error: 2.1.3 + deep-eql: 5.0.2 + loupe: 3.2.1 + pathval: 2.0.1 + + check-error@2.1.3: {} + + chokidar@4.0.3: + dependencies: + readdirp: 4.1.2 + + commander@4.1.1: {} + + confbox@0.1.8: {} + + consola@3.4.2: {} + + debug@4.4.3: + dependencies: + ms: 2.1.3 + + deep-eql@5.0.2: {} + + es-module-lexer@1.7.0: {} + + esbuild@0.27.7: + optionalDependencies: + '@esbuild/aix-ppc64': 0.27.7 + '@esbuild/android-arm': 0.27.7 + '@esbuild/android-arm64': 0.27.7 + '@esbuild/android-x64': 0.27.7 + '@esbuild/darwin-arm64': 0.27.7 + '@esbuild/darwin-x64': 0.27.7 + '@esbuild/freebsd-arm64': 0.27.7 + '@esbuild/freebsd-x64': 0.27.7 + '@esbuild/linux-arm': 0.27.7 + '@esbuild/linux-arm64': 0.27.7 + '@esbuild/linux-ia32': 0.27.7 + '@esbuild/linux-loong64': 0.27.7 + '@esbuild/linux-mips64el': 0.27.7 + '@esbuild/linux-ppc64': 0.27.7 + '@esbuild/linux-riscv64': 0.27.7 + '@esbuild/linux-s390x': 0.27.7 + '@esbuild/linux-x64': 0.27.7 + '@esbuild/netbsd-arm64': 0.27.7 + '@esbuild/netbsd-x64': 0.27.7 + '@esbuild/openbsd-arm64': 0.27.7 + '@esbuild/openbsd-x64': 0.27.7 + '@esbuild/openharmony-arm64': 0.27.7 + '@esbuild/sunos-x64': 0.27.7 + '@esbuild/win32-arm64': 0.27.7 + '@esbuild/win32-ia32': 0.27.7 + '@esbuild/win32-x64': 0.27.7 + + esbuild@0.28.2: + optionalDependencies: + '@esbuild/aix-ppc64': 0.28.2 + '@esbuild/android-arm': 0.28.2 + '@esbuild/android-arm64': 0.28.2 + '@esbuild/android-x64': 0.28.2 + '@esbuild/darwin-arm64': 0.28.2 + '@esbuild/darwin-x64': 0.28.2 + '@esbuild/freebsd-arm64': 0.28.2 + '@esbuild/freebsd-x64': 0.28.2 + '@esbuild/linux-arm': 0.28.2 + '@esbuild/linux-arm64': 0.28.2 + '@esbuild/linux-ia32': 0.28.2 + '@esbuild/linux-loong64': 0.28.2 + '@esbuild/linux-mips64el': 0.28.2 + '@esbuild/linux-ppc64': 0.28.2 + '@esbuild/linux-riscv64': 0.28.2 + '@esbuild/linux-s390x': 0.28.2 + '@esbuild/linux-x64': 0.28.2 + '@esbuild/netbsd-arm64': 0.28.2 + '@esbuild/netbsd-x64': 0.28.2 + '@esbuild/openbsd-arm64': 0.28.2 + '@esbuild/openbsd-x64': 0.28.2 + '@esbuild/openharmony-arm64': 0.28.2 + '@esbuild/sunos-x64': 0.28.2 + '@esbuild/win32-arm64': 0.28.2 + '@esbuild/win32-ia32': 0.28.2 + '@esbuild/win32-x64': 0.28.2 + + estree-walker@3.0.3: + dependencies: + '@types/estree': 1.0.9 + + expect-type@1.4.0: {} + + fdir@6.5.0(picomatch@4.0.7): + optionalDependencies: + picomatch: 4.0.7 + + fix-dts-default-cjs-exports@1.0.1: + dependencies: + magic-string: 0.30.21 + mlly: 1.8.2 + rollup: 4.64.7 + + fsevents@2.3.3: + optional: true + + joycon@3.1.1: {} + + js-tokens@9.0.1: {} + + lilconfig@3.1.4: {} + + lines-and-columns@1.2.4: {} + + load-tsconfig@0.2.5: {} + + loupe@3.2.1: {} + + magic-string@0.30.21: + dependencies: + '@jridgewell/sourcemap-codec': 1.6.0 + + mlly@1.8.2: + dependencies: + acorn: 8.19.0 + pathe: 2.0.3 + pkg-types: 1.3.1 + ufo: 1.6.4 + + ms@2.1.3: {} + + mz@2.7.0: + dependencies: + any-promise: 1.3.0 + object-assign: 4.1.1 + thenify-all: 1.6.0 + + nanoid@3.3.20: {} + + object-assign@4.1.1: {} + + pathe@2.0.3: {} + + pathval@2.0.1: {} + + picocolors@1.1.1: {} + + picomatch@4.0.7: {} + + pirates@4.0.7: {} + + pkg-types@1.3.1: + dependencies: + confbox: 0.1.8 + mlly: 1.8.2 + pathe: 2.0.3 + + postcss-load-config@6.0.1(postcss@8.5.30): + dependencies: + lilconfig: 3.1.4 + optionalDependencies: + postcss: 8.5.30 + + postcss@8.5.30: + dependencies: + nanoid: 3.3.20 + picocolors: 1.1.1 + source-map-js: 1.2.2 + + readdirp@4.1.2: {} + + resolve-from@5.0.0: {} + + rollup@4.64.7: + dependencies: + '@types/estree': 1.0.9 + optionalDependencies: + '@napi-rs/lzma-linux-x64-gnu': 1.5.1 + '@rollup/rollup-android-arm-eabi': 4.64.7 + '@rollup/rollup-android-arm64': 4.64.7 + '@rollup/rollup-darwin-arm64': 4.64.7 + '@rollup/rollup-darwin-x64': 4.64.7 + '@rollup/rollup-freebsd-arm64': 4.64.7 + '@rollup/rollup-freebsd-x64': 4.64.7 + '@rollup/rollup-linux-arm-gnueabihf': 4.64.7 + '@rollup/rollup-linux-arm-musleabihf': 4.64.7 + '@rollup/rollup-linux-arm64-gnu': 4.64.7 + '@rollup/rollup-linux-arm64-musl': 4.64.7 + '@rollup/rollup-linux-loong64-gnu': 4.64.7 + '@rollup/rollup-linux-loong64-musl': 4.64.7 + '@rollup/rollup-linux-ppc64-gnu': 4.64.7 + '@rollup/rollup-linux-ppc64-musl': 4.64.7 + '@rollup/rollup-linux-riscv64-gnu': 4.64.7 + '@rollup/rollup-linux-riscv64-musl': 4.64.7 + '@rollup/rollup-linux-s390x-gnu': 4.64.7 + '@rollup/rollup-linux-x64-gnu': 4.64.7 + '@rollup/rollup-linux-x64-musl': 4.64.7 + '@rollup/rollup-openbsd-x64': 4.64.7 + '@rollup/rollup-openharmony-arm64': 4.64.7 + '@rollup/rollup-win32-arm64-msvc': 4.64.7 + '@rollup/rollup-win32-ia32-msvc': 4.64.7 + '@rollup/rollup-win32-x64-gnu': 4.64.7 + '@rollup/rollup-win32-x64-msvc': 4.64.7 + fsevents: 2.3.3 + + siginfo@2.0.0: {} + + source-map-js@1.2.2: {} + + source-map@0.7.6: {} + + stackback@0.0.2: {} + + std-env@3.10.0: {} + + strip-literal@3.1.0: + dependencies: + js-tokens: 9.0.1 + + sucrase@3.35.1: + dependencies: + '@jridgewell/gen-mapping': 0.3.13 + commander: 4.1.1 + lines-and-columns: 1.2.4 + mz: 2.7.0 + pirates: 4.0.7 + tinyglobby: 0.2.17 + ts-interface-checker: 0.1.13 + + thenify-all@1.6.0: + dependencies: + thenify: 3.3.1 + + thenify@3.3.1: + dependencies: + any-promise: 1.3.0 + + tinybench@2.9.0: {} + + tinyexec@0.3.2: {} + + tinyglobby@0.2.17: + dependencies: + fdir: 6.5.0(picomatch@4.0.7) + picomatch: 4.0.7 + + tinypool@1.1.1: {} + + tinyrainbow@2.0.0: {} + + tinyspy@4.0.6: {} + + tree-kill@1.2.2: {} + + ts-interface-checker@0.1.13: {} + + tsup@8.5.1(postcss@8.5.30)(typescript@5.9.3): + dependencies: + bundle-require: 5.1.0(esbuild@0.27.7) + cac: 6.7.14 + chokidar: 4.0.3 + consola: 3.4.2 + debug: 4.4.3 + esbuild: 0.27.7 + fix-dts-default-cjs-exports: 1.0.1 + joycon: 3.1.1 + picocolors: 1.1.1 + postcss-load-config: 6.0.1(postcss@8.5.30) + resolve-from: 5.0.0 + rollup: 4.64.7 + source-map: 0.7.6 + sucrase: 3.35.1 + tinyexec: 0.3.2 + tinyglobby: 0.2.17 + tree-kill: 1.2.2 + optionalDependencies: + postcss: 8.5.30 + typescript: 5.9.3 + transitivePeerDependencies: + - jiti + - supports-color + - tsx + - yaml + + typescript@5.9.3: {} + + ufo@1.6.4: {} + + undici-types@8.9.0: {} + + vite-node@3.2.4(@types/node@26.6.5): + dependencies: + cac: 6.7.14 + debug: 4.4.3 + es-module-lexer: 1.7.0 + pathe: 2.0.3 + vite: 7.3.7(@types/node@26.6.5) + transitivePeerDependencies: + - '@types/node' + - jiti + - less + - lightningcss + - sass + - sass-embedded + - stylus + - sugarss + - supports-color + - terser + - tsx + - yaml + + vite@7.3.7(@types/node@26.6.5): + dependencies: + esbuild: 0.28.2 + fdir: 6.5.0(picomatch@4.0.7) + picomatch: 4.0.7 + postcss: 8.5.30 + rollup: 4.64.7 + tinyglobby: 0.2.17 + optionalDependencies: + '@types/node': 26.6.5 + fsevents: 2.3.3 + + vitest@3.2.7(@types/node@26.6.5): + dependencies: + '@types/chai': 5.2.3 + '@vitest/expect': 3.2.7 + '@vitest/mocker': 3.2.7(vite@7.3.7(@types/node@26.6.5)) + '@vitest/pretty-format': 3.2.7 + '@vitest/runner': 3.2.7 + '@vitest/snapshot': 3.2.7 + '@vitest/spy': 3.2.7 + '@vitest/utils': 3.2.7 + chai: 5.3.3 + debug: 4.4.3 + expect-type: 1.4.0 + magic-string: 0.30.21 + pathe: 2.0.3 + picomatch: 4.0.7 + std-env: 3.10.0 + tinybench: 2.9.0 + tinyexec: 0.3.2 + tinyglobby: 0.2.17 + tinypool: 1.1.1 + tinyrainbow: 2.0.0 + vite: 7.3.7(@types/node@26.6.5) + vite-node: 3.2.4(@types/node@26.6.5) + why-is-node-running: 2.3.0 + optionalDependencies: + '@types/node': 26.6.5 + transitivePeerDependencies: + - jiti + - less + - lightningcss + - msw + - sass + - sass-embedded + - stylus + - sugarss + - supports-color + - terser + - tsx + - yaml + + why-is-node-running@2.3.0: + dependencies: + siginfo: 2.0.0 + stackback: 0.0.2 diff --git a/explain/pnpm-workspace.yaml b/explain/pnpm-workspace.yaml new file mode 100644 index 0000000..3ee2585 --- /dev/null +++ b/explain/pnpm-workspace.yaml @@ -0,0 +1,8 @@ +# Own workspace root so pnpm 9 installs this package alone; the build +# allowlist lets pnpm 10 and later run the esbuild install script. +packages: + - . +allowBuilds: + esbuild: true +onlyBuiltDependencies: + - esbuild diff --git a/explain/src/iife.ts b/explain/src/iife.ts new file mode 100644 index 0000000..d775da5 --- /dev/null +++ b/explain/src/iife.ts @@ -0,0 +1,5 @@ +import { bigQueryExplainParser } from "./parser"; + +// The desktop loader registers this descriptor after matching it to the +// plugin manifest. This entry intentionally has no registration side effect. +export default bigQueryExplainParser; diff --git a/explain/src/index.ts b/explain/src/index.ts new file mode 100644 index 0000000..15c89cb --- /dev/null +++ b/explain/src/index.ts @@ -0,0 +1,9 @@ +import { registerExplainParser } from "@tabularis/explain"; + +import { bigQueryExplainParser } from "./parser"; + +registerExplainParser(bigQueryExplainParser); + +export { bigQueryExplainParser } from "./parser"; +export { parseBigQueryJob } from "./plan"; +export type { BigQueryJobPayload, BigQueryStage } from "./plan"; diff --git a/explain/src/parser.ts b/explain/src/parser.ts new file mode 100644 index 0000000..767957f --- /dev/null +++ b/explain/src/parser.ts @@ -0,0 +1,16 @@ +import type { RegisteredExplainParser } from "@tabularis/explain"; + +import { parseBigQueryJob } from "./plan"; + +/** BigQuery parser descriptor consumed by both package and plugin loaders. */ +export const bigQueryExplainParser: RegisteredExplainParser = { + engine: "bigquery", + format: "bigquery-job-json", + label: "BigQuery job plan (JSON)", + parse: parseBigQueryJob, + // Cheap check for the plugin payload; the parser still validates fully. + sniff: (payload) => { + const head = payload.slice(0, 4096); + return /^\s*\{/.test(head) && /"query_plan"\s*:/.test(payload) && /"mode"\s*:\s*"(estimate|analyze)"/.test(head); + }, +}; diff --git a/explain/src/plan.ts b/explain/src/plan.ts new file mode 100644 index 0000000..25de3f2 --- /dev/null +++ b/explain/src/plan.ts @@ -0,0 +1,327 @@ +import type { ExplainNode, ExplainPlan } from "@tabularis/explain"; + +/** One `statistics.query.queryPlan` entry. BigQuery encodes int64 as strings. */ +export interface BigQueryStage { + id?: string; + name?: string; + status?: string; + inputStages?: string[]; + startMs?: string; + endMs?: string; + slotMs?: string; + recordsRead?: string; + recordsWritten?: string; + parallelInputs?: string; + completedParallelInputs?: string; + shuffleOutputBytes?: string; + shuffleOutputBytesSpilled?: string; + waitMsAvg?: string; + waitMsMax?: string; + readMsAvg?: string; + readMsMax?: string; + computeMsAvg?: string; + computeMsMax?: string; + writeMsAvg?: string; + writeMsMax?: string; + steps?: { kind?: string; substeps?: string[] }[]; +} + +/** Payload written by the plugin's `explain_query` (`src/explain.rs`). */ +export interface BigQueryJobPayload { + version: number; + mode: "estimate" | "analyze"; + job_id?: string | null; + location?: string | null; + statement_type?: string | null; + total_bytes_processed?: number | null; + total_bytes_processed_accuracy?: string | null; + total_bytes_billed?: number | null; + estimated_cost_usd?: number | null; + cost_note?: string | null; + cache_hit?: boolean | null; + total_slot_ms?: number | null; + elapsed_ms?: number | null; + billing_tier?: number | null; + referenced_tables?: string[]; + query_plan?: BigQueryStage[]; +} + +function int(value: string | number | null | undefined): number | null { + if (value === null || value === undefined || value === "") return null; + const parsed = typeof value === "number" ? value : Number(value); + return Number.isFinite(parsed) ? parsed : null; +} + +function emptyNode(id: string, nodeType: string): ExplainNode { + return { + id, + node_type: nodeType, + relation: null, + startup_cost: null, + total_cost: null, + plan_rows: null, + actual_rows: null, + actual_time_ms: null, + actual_loops: null, + buffers_hit: null, + buffers_read: null, + filter: null, + index_condition: null, + join_type: null, + hash_condition: null, + extra: {}, + children: [], + }; +} + +function formatBytes(bytes: number): string { + const units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]; + let value = bytes; + let unit = 0; + while (value >= 1024 && unit < units.length - 1) { + value /= 1024; + unit += 1; + } + return unit === 0 ? `${bytes} B` : `${value.toFixed(2)} ${units[unit]}`; +} + +function formatUsd(value: number): string { + if (value === 0) return "$0.00"; + if (value < 0.01) return `$${value.toPrecision(2)}`; + return `$${value.toFixed(2)}`; +} + +function substeps(stage: BigQueryStage, kind: string): string[] { + return (stage.steps ?? []) + .filter((step) => step.kind === kind) + .flatMap((step) => step.substeps ?? []); +} + +/** Tables a stage reads, ignoring intermediate stage outputs. */ +function relationOf(stage: BigQueryStage): string | null { + const tables = substeps(stage, "READ") + .filter((substep) => substep.startsWith("FROM ")) + .map((substep) => substep.slice(5).trim()) + .filter((source) => !source.startsWith("__stage") && !source.startsWith("$")); + return tables.length > 0 ? tables.join(", ") : null; +} + +function filterOf(stage: BigQueryStage): string | null { + const filters = [...substeps(stage, "READ"), ...substeps(stage, "FILTER")] + .filter((substep) => substep.startsWith("WHERE ")) + .map((substep) => substep.slice(6).trim()); + return filters.length > 0 ? filters.join(" AND ") : null; +} + +function joinOf(stage: BigQueryStage): { type: string | null; condition: string | null } { + const join = substeps(stage, "JOIN")[0]; + if (join === undefined) return { type: null, condition: null }; + const on = join.indexOf(" ON "); + return on === -1 + ? { type: join, condition: null } + : { type: join.slice(0, on).trim(), condition: join.slice(on + 4).trim() }; +} + +function stageNode(stage: BigQueryStage, key: string): ExplainNode { + const node = emptyNode(`stage-${key}`, stage.name ?? `Stage ${key}`); + const start = int(stage.startMs); + const end = int(stage.endMs); + const join = joinOf(stage); + node.relation = relationOf(stage); + node.filter = filterOf(stage); + node.join_type = join.type; + node.hash_condition = join.condition; + node.actual_rows = int(stage.recordsWritten); + node.actual_time_ms = start !== null && end !== null && end >= start ? end - start : null; + node.total_cost = int(stage.slotMs); + const shuffle = int(stage.shuffleOutputBytes); + const spilled = int(stage.shuffleOutputBytesSpilled); + node.extra = { + status: stage.status ?? null, + slot_ms: int(stage.slotMs), + records_read: int(stage.recordsRead), + records_written: int(stage.recordsWritten), + parallel_inputs: int(stage.parallelInputs), + completed_parallel_inputs: int(stage.completedParallelInputs), + shuffle_output: shuffle === null ? null : formatBytes(shuffle), + shuffle_spilled: spilled === null ? null : formatBytes(spilled), + wait_ms_avg: int(stage.waitMsAvg), + wait_ms_max: int(stage.waitMsMax), + read_ms_avg: int(stage.readMsAvg), + read_ms_max: int(stage.readMsMax), + compute_ms_avg: int(stage.computeMsAvg), + compute_ms_max: int(stage.computeMsMax), + write_ms_avg: int(stage.writeMsAvg), + write_ms_max: int(stage.writeMsMax), + steps: (stage.steps ?? []).map( + (step) => `${step.kind ?? "STEP"}: ${(step.substeps ?? []).join("; ")}`, + ), + }; + return node; +} + +/** + * Turn per-node own values into the inclusive values Visual EXPLAIN expects + * (a node's cost and time cover its children, as in PostgreSQL plans). + */ +function accumulate(node: ExplainNode): void { + let childCost = 0; + let childTime = 0; + let anyCost = false; + let anyTime = false; + for (const child of node.children) { + accumulate(child); + if (child.total_cost !== null) { + childCost += child.total_cost; + anyCost = true; + } + if (child.actual_time_ms !== null) { + childTime += child.actual_time_ms; + anyTime = true; + } + } + if (anyCost) node.total_cost = (node.total_cost ?? 0) + childCost; + if (anyTime) node.actual_time_ms = (node.actual_time_ms ?? 0) + childTime; +} + +/** + * Build the stage tree. BigQuery stages form a DAG through `inputStages`; a + * stage feeding several consumers is attached to the first one and listed in + * the others' `extra.also_reads_stages`. + */ +function stageTree(stages: BigQueryStage[]): ExplainNode[] { + const byId = new Map(); + stages.forEach((stage, index) => byId.set(stage.id ?? String(index), stage)); + const parentOf = new Map(); + const shared = new Map(); + for (const [id, stage] of byId) { + for (const input of stage.inputStages ?? []) { + if (!byId.has(input) || input === id) continue; + if (parentOf.has(input)) { + shared.set(id, [...(shared.get(id) ?? []), input]); + } else { + parentOf.set(input, id); + } + } + } + + const visited = new Set(); + const build = (id: string): ExplainNode => { + visited.add(id); + const node = stageNode(byId.get(id)!, id); + const alsoReads = shared.get(id); + if (alsoReads) node.extra.also_reads_stages = alsoReads; + for (const input of byId.get(id)!.inputStages ?? []) { + if (parentOf.get(input) === id && !visited.has(input)) { + node.children.push(build(input)); + } + } + return node; + }; + + const roots = [...byId.keys()].filter((id) => !parentOf.has(id)); + const tree = roots.map(build); + // A malformed cycle leaves stages unvisited; keep them visible at the top. + for (const id of byId.keys()) { + if (!visited.has(id)) tree.push(build(id)); + } + return tree.reverse(); +} + +function parsePayload(payload: string): BigQueryJobPayload { + let value: unknown; + try { + value = JSON.parse(payload); + } catch (error) { + throw new Error(`Invalid BigQuery plan JSON: ${(error as Error).message}`); + } + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error("BigQuery plan payload must be a JSON object"); + } + const job = value as BigQueryJobPayload; + if (job.version !== 1) { + throw new Error(`Unsupported BigQuery plan payload version ${String(job.version)}`); + } + if (job.mode !== "estimate" && job.mode !== "analyze") { + throw new Error("BigQuery plan payload has an unknown mode"); + } + if (job.query_plan !== undefined && !Array.isArray(job.query_plan)) { + throw new Error("BigQuery plan payload query_plan must be an array"); + } + job.query_plan?.forEach(validateStage); + return job; +} + +function isObject(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function validateStage(stage: unknown, index: number): void { + const at = `BigQuery plan payload query_plan[${index}]`; + if (!isObject(stage)) { + throw new Error(`${at} must be an object`); + } + if (stage.inputStages !== undefined && !Array.isArray(stage.inputStages)) { + throw new Error(`${at}.inputStages must be an array`); + } + if (stage.steps === undefined) return; + if (!Array.isArray(stage.steps)) { + throw new Error(`${at}.steps must be an array`); + } + stage.steps.forEach((step: unknown, stepIndex: number) => { + if (!isObject(step)) { + throw new Error(`${at}.steps[${stepIndex}] must be an object`); + } + if (step.substeps !== undefined && !Array.isArray(step.substeps)) { + throw new Error(`${at}.steps[${stepIndex}].substeps must be an array`); + } + }); +} + +/** Parse the `bigquery-job-json` payload into a Visual EXPLAIN plan. */ +export function parseBigQueryJob(payload: string): ExplainPlan { + const job = parsePayload(payload); + const analyze = job.mode === "analyze"; + const root = emptyNode("query", analyze ? "BigQuery job" : "Dry run estimate"); + const processed = int(job.total_bytes_processed); + const billed = int(job.total_bytes_billed); + const cost = int(job.estimated_cost_usd); + root.extra = { + statement_type: job.statement_type ?? null, + bytes_processed: processed === null ? null : formatBytes(processed), + bytes_processed_exact: processed, + bytes_processed_accuracy: job.total_bytes_processed_accuracy ?? null, + bytes_billed: billed === null ? null : formatBytes(billed), + estimated_cost: cost === null ? null : `${formatUsd(cost)} (estimate)`, + cost_note: job.cost_note ?? null, + cache_hit: job.cache_hit ?? null, + slot_ms: int(job.total_slot_ms), + billing_tier: int(job.billing_tier), + job_id: job.job_id ?? null, + location: job.location ?? null, + }; + + if (analyze) { + root.children = stageTree(job.query_plan ?? []); + accumulate(root); + if (root.children.length === 0 && job.cache_hit) { + root.node_type = "BigQuery job (cached result)"; + } + } else { + root.children = (job.referenced_tables ?? []).map((table, index) => { + const node = emptyNode(`table-${index}`, "Referenced table"); + node.relation = table; + return node; + }); + } + + return { + root, + planning_time_ms: null, + execution_time_ms: analyze ? int(job.elapsed_ms) : null, + original_query: "", + driver: "bigquery", + has_analyze_data: analyze && (job.query_plan ?? []).length > 0, + raw_output: JSON.stringify(job, null, 2), + }; +} diff --git a/explain/tests/fixtures/estimate.json b/explain/tests/fixtures/estimate.json new file mode 100644 index 0000000..806a7a7 --- /dev/null +++ b/explain/tests/fixtures/estimate.json @@ -0,0 +1,19 @@ +{ + "version": 1, + "mode": "estimate", + "job_id": null, + "location": "EU", + "statement_type": "SELECT", + "total_bytes_processed": 1099511627776, + "total_bytes_processed_accuracy": "PRECISE", + "total_bytes_billed": null, + "estimated_cost_usd": 6.25, + "cost_note": "Estimate at the on-demand list price of $6.25 per TiB.", + "cache_hit": false, + "total_slot_ms": null, + "elapsed_ms": null, + "billing_tier": null, + "referenced_tables": ["p.sales.orders"], + "query_plan": [], + "timeline": [] +} diff --git a/explain/tests/fixtures/join-analyzed.json b/explain/tests/fixtures/join-analyzed.json new file mode 100644 index 0000000..462f9d7 --- /dev/null +++ b/explain/tests/fixtures/join-analyzed.json @@ -0,0 +1,56 @@ +{ + "version": 1, + "mode": "analyze", + "job_id": "job_join", + "location": "US", + "statement_type": "SELECT", + "total_bytes_processed": 52428800, + "total_bytes_billed": 52428800, + "estimated_cost_usd": 0.0002980232238769531, + "cost_note": "Estimate at the on-demand list price of $6.25 per TiB.", + "cache_hit": false, + "total_slot_ms": 1450, + "elapsed_ms": 812, + "billing_tier": 1, + "referenced_tables": ["p.sales.orders", "p.sales.customers"], + "query_plan": [ + { + "id": "0", "name": "S00: Input", "status": "COMPLETE", + "startMs": "1000", "endMs": "1300", "slotMs": "600", + "recordsRead": "100000", "recordsWritten": "2500", + "parallelInputs": "4", "completedParallelInputs": "4", + "shuffleOutputBytes": "40960", "shuffleOutputBytesSpilled": "0", + "waitMsAvg": "2", "readMsAvg": "40", "computeMsAvg": "80", "writeMsAvg": "5", + "steps": [ + {"kind": "READ", "substeps": ["$1:customer_id, $2:amount", "FROM p.sales.orders", "WHERE greater($2, 100)"]}, + {"kind": "WRITE", "substeps": ["$1, $2", "TO __stage00_output"]} + ] + }, + { + "id": "1", "name": "S01: Input", "status": "COMPLETE", + "startMs": "1000", "endMs": "1200", "slotMs": "300", + "recordsRead": "5000", "recordsWritten": "5000", + "steps": [ + {"kind": "READ", "substeps": ["$10:id, $11:name", "FROM p.sales.customers"]}, + {"kind": "WRITE", "substeps": ["$10, $11", "TO __stage01_output"]} + ] + }, + { + "id": "2", "name": "S02: Join+", "status": "COMPLETE", "inputStages": ["0", "1"], + "startMs": "1300", "endMs": "1700", "slotMs": "500", + "recordsRead": "7500", "recordsWritten": "2500", + "steps": [ + {"kind": "READ", "substeps": ["FROM __stage00_output"]}, + {"kind": "JOIN", "substeps": ["INNER HASH JOIN EACH WITH ALL ON $1 = $10"]}, + {"kind": "WRITE", "substeps": ["TO __stage02_output"]} + ] + }, + { + "id": "3", "name": "S03: Output", "status": "COMPLETE", "inputStages": ["2"], + "startMs": "1700", "endMs": "1800", "slotMs": "50", + "recordsRead": "2500", "recordsWritten": "2500", + "steps": [{"kind": "WRITE", "substeps": ["TO __stage03_output"]}] + } + ], + "timeline": [] +} diff --git a/explain/tests/host/index.d.ts b/explain/tests/host/index.d.ts new file mode 100644 index 0000000..047f79e --- /dev/null +++ b/explain/tests/host/index.d.ts @@ -0,0 +1,40 @@ +export interface ExplainNode { + id: string; + node_type: string; + relation: string | null; + startup_cost: number | null; + total_cost: number | null; + plan_rows: number | null; + actual_rows: number | null; + actual_time_ms: number | null; + actual_loops: number | null; + buffers_hit: number | null; + buffers_read: number | null; + filter: string | null; + index_condition: string | null; + join_type: string | null; + hash_condition: string | null; + extra: Record; + children: ExplainNode[]; +} + +export interface ExplainPlan { + root: ExplainNode; + planning_time_ms: number | null; + execution_time_ms: number | null; + original_query: string; + driver: string; + has_analyze_data: boolean; + raw_output: string | null; +} + +export interface RegisteredExplainParser { + readonly engine: string; + readonly format: string; + readonly label?: string; + parse(payload: string): ExplainPlan; + sniff?(payload: string): boolean; +} + +export const registrations: RegisteredExplainParser[]; +export function registerExplainParser(parser: RegisteredExplainParser): void; diff --git a/explain/tests/host/index.js b/explain/tests/host/index.js new file mode 100644 index 0000000..b2b4ce3 --- /dev/null +++ b/explain/tests/host/index.js @@ -0,0 +1,5 @@ +export const registrations = []; + +export function registerExplainParser(parser) { + registrations.push(parser); +} diff --git a/explain/tests/host/package.json b/explain/tests/host/package.json new file mode 100644 index 0000000..7f4b8d9 --- /dev/null +++ b/explain/tests/host/package.json @@ -0,0 +1,11 @@ +{ + "name": "@tabularis/explain", + "version": "0.2.0", + "type": "module", + "exports": { + ".": { + "types": "./index.d.ts", + "import": "./index.js" + } + } +} diff --git a/explain/tests/package.test.ts b/explain/tests/package.test.ts new file mode 100644 index 0000000..15a8ecf --- /dev/null +++ b/explain/tests/package.test.ts @@ -0,0 +1,44 @@ +import { readFile } from "node:fs/promises"; + +import { + registrations, + type RegisteredExplainParser, +} from "@tabularis/explain"; +import { describe, expect, it } from "vitest"; + +import { bigQueryExplainParser, parseBigQueryJob } from "../src/index"; + +const fixtureUrl = new URL("./fixtures/join-analyzed.json", import.meta.url); + +describe("package entry points", () => { + it("registers on ESM import and exports the direct parser API", async () => { + expect(registrations).toEqual([bigQueryExplainParser]); + expect(bigQueryExplainParser).toMatchObject({ + engine: "bigquery", + format: "bigquery-job-json", + label: "BigQuery job plan (JSON)", + parse: parseBigQueryJob, + }); + expect(bigQueryExplainParser.sniff?.(await readFile(fixtureUrl, "utf8"))).toBe(true); + expect(bigQueryExplainParser.sniff?.("")).toBe(false); + }); + + it("builds an isolated IIFE descriptor without self-registration", async () => { + const source = await readFile(new URL("../dist/index.iife.js", import.meta.url), "utf8"); + const registrationsBeforeEvaluation = registrations.length; + const evaluate = new Function( + "__TABULARIS_EXPLAIN__", + `${source}\nreturn typeof __tabularis_explain_parser__ !== "undefined" ? __tabularis_explain_parser__ : null;`, + ); + const raw = evaluate({}) as Record; + const descriptor = (raw.default ?? raw) as RegisteredExplainParser; + + expect(descriptor).toMatchObject({ + engine: "bigquery", + format: "bigquery-job-json", + label: "BigQuery job plan (JSON)", + }); + expect(descriptor.parse(await readFile(fixtureUrl, "utf8")).driver).toBe("bigquery"); + expect(registrations).toHaveLength(registrationsBeforeEvaluation); + }); +}); diff --git a/explain/tests/plan.test.ts b/explain/tests/plan.test.ts new file mode 100644 index 0000000..8591ab2 --- /dev/null +++ b/explain/tests/plan.test.ts @@ -0,0 +1,106 @@ +import { readFile } from "node:fs/promises"; + +import type { ExplainNode } from "@tabularis/explain"; +import { describe, expect, it } from "vitest"; + +import { parseBigQueryJob } from "../src/plan"; + +async function fixture(name: string): Promise { + return readFile(new URL(`./fixtures/${name}.json`, import.meta.url), "utf8"); +} + +function flatten(node: ExplainNode): ExplainNode[] { + return [node, ...node.children.flatMap(flatten)]; +} + +describe("parseBigQueryJob", () => { + it("builds the stage tree from inputStages with inclusive values", async () => { + const plan = parseBigQueryJob(await fixture("join-analyzed")); + + expect(plan.driver).toBe("bigquery"); + expect(plan.has_analyze_data).toBe(true); + expect(plan.execution_time_ms).toBe(812); + expect(plan.root.node_type).toBe("BigQuery job"); + expect(flatten(plan.root).map((node) => node.id)).toEqual([ + "query", + "stage-3", + "stage-2", + "stage-0", + "stage-1", + ]); + + const [output] = plan.root.children; + const join = output.children[0]; + expect(join.join_type).toBe("INNER HASH JOIN EACH WITH ALL"); + expect(join.hash_condition).toBe("$1 = $10"); + // Inclusive: own slot ms plus both inputs. + expect(join.total_cost).toBe(500 + 600 + 300); + expect(output.total_cost).toBe(1450); + expect(plan.root.total_cost).toBe(1450); + expect(join.actual_time_ms).toBe(400 + 300 + 200); + + const orders = join.children[0]; + expect(orders.relation).toBe("p.sales.orders"); + expect(orders.filter).toBe("greater($2, 100)"); + expect(orders.actual_rows).toBe(2500); + expect(orders.extra.records_read).toBe(100000); + expect(orders.extra.shuffle_output).toBe("40.00 KiB"); + }); + + it("labels the dry-run cost as an estimate", async () => { + const plan = parseBigQueryJob(await fixture("estimate")); + + expect(plan.has_analyze_data).toBe(false); + expect(plan.root.node_type).toBe("Dry run estimate"); + expect(plan.root.extra.bytes_processed).toBe("1.00 TiB"); + expect(plan.root.extra.estimated_cost).toBe("$6.25 (estimate)"); + expect(plan.root.children.map((node) => node.relation)).toEqual(["p.sales.orders"]); + }); + + it("marks cached results that have no stages", () => { + const plan = parseBigQueryJob( + JSON.stringify({ version: 1, mode: "analyze", cache_hit: true, query_plan: [] }), + ); + expect(plan.root.node_type).toBe("BigQuery job (cached result)"); + expect(plan.has_analyze_data).toBe(false); + }); + + it("keeps stages of a malformed cyclic plan", () => { + const plan = parseBigQueryJob( + JSON.stringify({ + version: 1, + mode: "analyze", + query_plan: [ + { id: "a", name: "A", inputStages: ["b"] }, + { id: "b", name: "B", inputStages: ["a"] }, + ], + }), + ); + expect(flatten(plan.root).map((node) => node.node_type).sort()).toEqual([ + "A", + "B", + "BigQuery job", + ]); + }); + + it("rejects payloads it does not understand", () => { + expect(() => parseBigQueryJob("not json")).toThrow(/Invalid BigQuery plan JSON/); + expect(() => parseBigQueryJob("[]")).toThrow(/JSON object/); + expect(() => parseBigQueryJob('{"version":2,"mode":"analyze"}')).toThrow(/version 2/); + expect(() => parseBigQueryJob('{"version":1,"mode":"other"}')).toThrow(/unknown mode/); + }); + + it("rejects malformed stages and steps with clear errors", () => { + const plan = (query_plan: unknown) => + parseBigQueryJob(JSON.stringify({ version: 1, mode: "analyze", query_plan })); + expect(() => plan([null])).toThrow(/query_plan\[0\] must be an object/); + expect(() => plan([{ id: "0" }, 7])).toThrow(/query_plan\[1\] must be an object/); + expect(() => plan([["stage"]])).toThrow(/query_plan\[0\] must be an object/); + expect(() => plan([{ id: "0", inputStages: "1" }])).toThrow(/inputStages must be an array/); + expect(() => plan([{ id: "0", steps: {} }])).toThrow(/steps must be an array/); + expect(() => plan([{ id: "0", steps: [null] }])).toThrow(/steps\[0\] must be an object/); + expect(() => plan([{ id: "0", steps: [{ kind: "READ", substeps: "x" }] }])).toThrow( + /substeps must be an array/, + ); + }); +}); diff --git a/explain/tsconfig.json b/explain/tsconfig.json new file mode 100644 index 0000000..e20274f --- /dev/null +++ b/explain/tsconfig.json @@ -0,0 +1,25 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "ESNext", + "moduleResolution": "bundler", + "lib": ["ES2022", "DOM"], + "strict": true, + "noImplicitAny": true, + "noUnusedLocals": true, + "noUnusedParameters": true, + "noFallthroughCasesInSwitch": true, + "declaration": true, + "declarationMap": false, + "sourceMap": false, + "esModuleInterop": true, + "isolatedModules": true, + "skipLibCheck": true, + "verbatimModuleSyntax": true, + "erasableSyntaxOnly": true, + "noEmit": true, + "types": ["vitest/globals"] + }, + "include": ["src/**/*", "tests/**/*.ts"], + "exclude": ["dist", "node_modules"] +} diff --git a/explain/tsup.config.ts b/explain/tsup.config.ts new file mode 100644 index 0000000..39d8d46 --- /dev/null +++ b/explain/tsup.config.ts @@ -0,0 +1,30 @@ +import { defineConfig } from "tsup"; + +const common = { + platform: "browser" as const, + target: "es2022" as const, + sourcemap: false, + minify: false, + splitting: false, + treeshake: true, + external: ["@tabularis/explain"], +}; + +export default defineConfig([ + { + ...common, + entry: { index: "src/index.ts" }, + format: ["esm"], + dts: true, + clean: true, + }, + { + ...common, + entry: { "index.iife": "src/iife.ts" }, + format: ["iife"], + globalName: "__tabularis_explain_parser__", + dts: false, + clean: false, + outExtension: () => ({ js: ".js" }), + }, +]); diff --git a/justfile b/justfile index 04a4069..f064a94 100644 --- a/justfile +++ b/justfile @@ -13,6 +13,9 @@ ui-install: ui-build: pnpm --dir ui build +explain-build: + pnpm --dir explain build + ui-typecheck: pnpm --dir ui typecheck @@ -30,15 +33,17 @@ check-ui: pnpm --dir ui build [linux] -dev-install: build ui-build - mkdir -p ~/.local/share/tabularis/plugins/bigquery/ui/dist - cp target/debug/bigquery-plugin ~/.local/share/tabularis/plugins/bigquery/ - cp .tabularium ~/.local/share/tabularis/plugins/bigquery/ - cp ui/dist/*.js ~/.local/share/tabularis/plugins/bigquery/ui/dist/ +dev-install: build ui-build explain-build + mkdir -p ~/.local/share/tabularis/plugins/drivers/bigquery/ui/dist ~/.local/share/tabularis/plugins/drivers/bigquery/explain/dist + cp target/debug/bigquery-plugin ~/.local/share/tabularis/plugins/drivers/bigquery/ + cp .tabularium ~/.local/share/tabularis/plugins/drivers/bigquery/ + cp ui/dist/*.js ~/.local/share/tabularis/plugins/drivers/bigquery/ui/dist/ + cp explain/dist/index.iife.js ~/.local/share/tabularis/plugins/drivers/bigquery/explain/dist/ [macos] -dev-install: build ui-build - mkdir -p "$HOME/Library/Application Support/com.debba.tabularis/plugins/bigquery/ui/dist" - cp target/debug/bigquery-plugin "$HOME/Library/Application Support/com.debba.tabularis/plugins/bigquery/" - cp .tabularium "$HOME/Library/Application Support/com.debba.tabularis/plugins/bigquery/" - cp ui/dist/*.js "$HOME/Library/Application Support/com.debba.tabularis/plugins/bigquery/ui/dist/" +dev-install: build ui-build explain-build + mkdir -p "$HOME/Library/Application Support/tabularis/plugins/drivers/bigquery/ui/dist" "$HOME/Library/Application Support/tabularis/plugins/drivers/bigquery/explain/dist" + cp target/debug/bigquery-plugin "$HOME/Library/Application Support/tabularis/plugins/drivers/bigquery/" + cp .tabularium "$HOME/Library/Application Support/tabularis/plugins/drivers/bigquery/" + cp ui/dist/*.js "$HOME/Library/Application Support/tabularis/plugins/drivers/bigquery/ui/dist/" + cp explain/dist/index.iife.js "$HOME/Library/Application Support/tabularis/plugins/drivers/bigquery/explain/dist/" diff --git a/src/auth.rs b/src/auth.rs index 57da9b1..59b4653 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1,23 +1,76 @@ -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant}; use gcp_auth::{CustomServiceAccount, TokenProvider}; +use reqwest::Url; use serde::Deserialize; +use serde_json::{json, Value}; use tokio::sync::Mutex; +use crate::endpoint::ApiEndpoint; use crate::models::ConnectionParams; const BIGQUERY_SCOPES: &[&str] = &[ "https://www.googleapis.com/auth/bigquery", "https://www.googleapis.com/auth/cloud-platform", ]; +const IAM_CREDENTIALS_BASE: &str = "https://iamcredentials.googleapis.com/v1/"; + +type ProviderFuture<'a> = + std::pin::Pin> + Send + 'a>>; + +type GoogleFuture = std::pin::Pin< + Box, String>> + Send>, +>; +/// Builds a fresh gcp_auth provider, so a rejected token is not served again +/// from the old provider's internal cache. +type GoogleBuild = Arc GoogleFuture + Send + Sync>; #[derive(Clone)] pub enum AuthProvider { - Google(Arc), + Google(Arc), AuthorizedUser(Arc), + Impersonated(Arc), AccessToken(String), + /// No Authorization header; only allowed for loopback emulator endpoints. + None, +} + +/// A gcp_auth provider (service account key, metadata server, gcloud CLI). +/// gcp_auth caches tokens internally without a way to drop them, so after a +/// 401 the provider is rebuilt instead. +pub struct GoogleProvider { + current: Mutex>, + build: GoogleBuild, +} + +impl GoogleProvider { + async fn new(build: GoogleBuild) -> Result { + let current = build().await?; + Ok(Self { + current: Mutex::new(current), + build, + }) + } + + async fn access_token(&self) -> Result { + let provider = self.current.lock().await.clone(); + provider + .token(BIGQUERY_SCOPES) + .await + .map(|token| token.as_str().to_string()) + .map_err(|error| format!("Unable to obtain a Google access token: {error}")) + } + + async fn rebuild(&self) { + match (self.build)().await { + Ok(provider) => *self.current.lock().await = provider, + Err(error) => { + eprintln!("bigquery-plugin: unable to refresh Google credentials: {error}") + } + } + } } pub struct AuthorizedUserProvider { @@ -26,10 +79,18 @@ pub struct AuthorizedUserProvider { cached: Mutex>, } +/// Exchanges a source credential for a short-lived token of another service +/// account through the IAM Credentials `generateAccessToken` method. +pub struct ImpersonatedProvider { + http: reqwest::Client, + source: AuthProvider, + url: Url, + delegates: Vec, + cached: Mutex>, +} + #[derive(Clone, Deserialize)] struct AuthorizedUserCredentials { - #[serde(rename = "type")] - credential_type: String, client_id: String, client_secret: String, refresh_token: String, @@ -40,74 +101,97 @@ struct CachedToken { expires_at: Instant, } +impl CachedToken { + fn fresh(&self) -> Option { + (self.expires_at > Instant::now() + Duration::from_secs(60)).then(|| self.value.clone()) + } +} + impl AuthProvider { - pub async fn from_params(params: &ConnectionParams) -> Result { - match params.auth_method() { - "adc" => gcp_auth::provider() - .await - .map(Self::Google) - .map_err(|error| format!("Application Default Credentials are unavailable: {error}")), + pub async fn from_params( + params: &ConnectionParams, + endpoint: &ApiEndpoint, + ) -> Result { + let base = match params.auth_method() { + "adc" => application_default().await?, "service_account" => { - let credential = params - .credential() - .ok_or_else(|| { - "Service account authentication requires a JSON key or file path" - .to_string() - })?; - let account = if credential.starts_with('{') { - CustomServiceAccount::from_json(credential) - } else { - CustomServiceAccount::from_file(Path::new(credential)) - } - .map_err(|error| format!("Invalid service account credential: {error}"))?; - Ok(Self::Google(Arc::new(account))) + let credential = params.credential().ok_or_else(|| { + "Service account authentication requires a JSON key or file path".to_string() + })?; + let text = load_json_credential(credential).await?; + from_credential_json(&text, &["service_account", "impersonated_service_account"]) + .await? } "authorized_user" => { - let credential = read_json_credential(params).await?; - let credentials: AuthorizedUserCredentials = serde_json::from_str(&credential) - .map_err(|error| format!("Invalid authorized user credential JSON: {error}"))?; - if credentials.credential_type != "authorized_user" { - return Err("OAuth credential JSON must have type 'authorized_user'".to_string()); + let credential = params.credential().ok_or_else(|| { + "Authentication requires credential JSON or a file path".to_string() + })?; + let text = load_json_credential(credential).await?; + from_credential_json(&text, &["authorized_user", "impersonated_service_account"]) + .await? + } + "access_token" => params + .credential() + .map(|value| Self::AccessToken(strip_bom(value).trim().to_string())) + .ok_or_else(|| "Access token authentication requires a token".to_string())?, + "none" => { + if !endpoint.local { + return Err( + "Authentication 'none' is only allowed with a localhost API endpoint, such as a BigQuery emulator" + .to_string(), + ); } - if credentials.client_id.trim().is_empty() - || credentials.client_secret.trim().is_empty() - || credentials.refresh_token.trim().is_empty() - { + Self::None + } + method => { + return Err(format!( + "Unsupported authentication method '{method}'. Use adc, service_account, authorized_user, access_token or none" + )) + } + }; + + match params.impersonate_service_account() { + Some(target) => { + if matches!(base, Self::None) { return Err( - "OAuth credential JSON must contain client_id, client_secret and refresh_token" + "Service account impersonation requires an authentication method other than 'none'" .to_string(), ); } - let http = reqwest::Client::builder() - .connect_timeout(Duration::from_secs(20)) - .timeout(Duration::from_secs(30)) - .build() - .map_err(|error| format!("Unable to initialize OAuth client: {error}"))?; - Ok(Self::AuthorizedUser(Arc::new(AuthorizedUserProvider { - http, - credentials, - cached: Mutex::new(None), - }))) + let url = impersonation_url_for(target)?; + Ok(Self::Impersonated(Arc::new(ImpersonatedProvider::new( + base, + url, + Vec::new(), + )?))) } - "access_token" => params - .credential() - .map(|value| Self::AccessToken(value.to_string())) - .ok_or_else(|| "Access token authentication requires a token".to_string()), - method => Err(format!( - "Unsupported authentication method '{method}'. Use adc, service_account, authorized_user, or access_token" - )), + None => Ok(base), } } - pub async fn access_token(&self) -> Result { + /// A bearer token, or `None` when the connection is unauthenticated. + pub async fn access_token(&self) -> Result, String> { match self { - Self::Google(provider) => provider - .token(BIGQUERY_SCOPES) - .await - .map(|token| token.as_str().to_string()) - .map_err(|error| format!("Unable to obtain a Google access token: {error}")), - Self::AuthorizedUser(provider) => provider.access_token().await, - Self::AccessToken(token) => Ok(token.clone()), + Self::Google(provider) => provider.access_token().await.map(Some), + Self::AuthorizedUser(provider) => provider.access_token().await.map(Some), + Self::Impersonated(provider) => provider.access_token().await.map(Some), + Self::AccessToken(token) => Ok(Some(token.clone())), + Self::None => Ok(None), + } + } + + /// Drop cached tokens after the API rejected one with HTTP 401, so the + /// next request fetches a new one. A fixed access token cannot be + /// refreshed. + pub async fn invalidate(&self) { + match self { + Self::Google(provider) => provider.rebuild().await, + Self::AuthorizedUser(provider) => *provider.cached.lock().await = None, + Self::Impersonated(provider) => { + *provider.cached.lock().await = None; + Box::pin(provider.source.invalidate()).await; + } + Self::AccessToken(_) | Self::None => {} } } } @@ -115,10 +199,8 @@ impl AuthProvider { impl AuthorizedUserProvider { async fn access_token(&self) -> Result { let mut cached = self.cached.lock().await; - if let Some(token) = cached.as_ref() { - if token.expires_at > Instant::now() + Duration::from_secs(30) { - return Ok(token.value.clone()); - } + if let Some(token) = cached.as_ref().and_then(CachedToken::fresh) { + return Ok(token); } let response = self @@ -134,7 +216,7 @@ impl AuthorizedUserProvider { .await .map_err(|error| format!("OAuth token refresh failed: {error}"))?; let status = response.status(); - let payload: serde_json::Value = response + let payload: Value = response .json() .await .map_err(|error| format!("Invalid OAuth token response: {error}"))?; @@ -142,18 +224,18 @@ impl AuthorizedUserProvider { let message = payload .get("error_description") .or_else(|| payload.get("error")) - .and_then(serde_json::Value::as_str) + .and_then(Value::as_str) .unwrap_or("unknown OAuth error"); return Err(format!("OAuth token refresh failed: {message}")); } let value = payload .get("access_token") - .and_then(serde_json::Value::as_str) + .and_then(Value::as_str) .ok_or_else(|| "OAuth token response did not contain access_token".to_string())? .to_string(); let expires_in = payload .get("expires_in") - .and_then(serde_json::Value::as_u64) + .and_then(Value::as_u64) .unwrap_or(3600); *cached = Some(CachedToken { value: value.clone(), @@ -163,15 +245,528 @@ impl AuthorizedUserProvider { } } -async fn read_json_credential(params: &ConnectionParams) -> Result { - let credential = params - .credential() - .ok_or_else(|| "Authentication requires credential JSON or a file path".to_string())?; - if credential.starts_with('{') { - Ok(credential.to_string()) +impl ImpersonatedProvider { + pub(crate) fn new( + source: AuthProvider, + url: Url, + delegates: Vec, + ) -> Result { + Ok(Self { + http: oauth_http()?, + source, + url, + delegates: delegates + .into_iter() + .map(|delegate| { + if delegate.starts_with("projects/") { + delegate + } else { + format!("projects/-/serviceAccounts/{delegate}") + } + }) + .collect(), + cached: Mutex::new(None), + }) + } + + async fn access_token(&self) -> Result { + let mut cached = self.cached.lock().await; + if let Some(token) = cached.as_ref().and_then(CachedToken::fresh) { + return Ok(token); + } + let source = Box::pin(self.source.access_token()).await?.ok_or_else(|| { + "Impersonation requires an authenticated source credential".to_string() + })?; + let response = self + .http + .post(self.url.clone()) + .bearer_auth(source) + .json(&json!({ + "scope": BIGQUERY_SCOPES, + "delegates": self.delegates, + "lifetime": "3600s" + })) + .send() + .await + .map_err(|error| format!("Service account impersonation failed: {error}"))?; + let status = response.status(); + let payload: Value = response + .json() + .await + .map_err(|error| format!("Invalid impersonation response: {error}"))?; + if !status.is_success() { + let message = payload + .pointer("/error/message") + .and_then(Value::as_str) + .unwrap_or("unknown IAM Credentials error"); + return Err(format!( + "Service account impersonation failed ({status}): {message}. The source account needs roles/iam.serviceAccountTokenCreator on the target" + )); + } + let value = payload + .get("accessToken") + .and_then(Value::as_str) + .ok_or_else(|| "Impersonation response did not contain accessToken".to_string())? + .to_string(); + let lifetime = payload + .get("expireTime") + .and_then(Value::as_str) + .and_then(|expiry| chrono::DateTime::parse_from_rfc3339(expiry).ok()) + .and_then(|expiry| { + (expiry.with_timezone(&chrono::Utc) - chrono::Utc::now()) + .to_std() + .ok() + }) + .unwrap_or(Duration::from_secs(3600)); + *cached = Some(CachedToken { + value: value.clone(), + expires_at: Instant::now() + lifetime, + }); + Ok(value) + } +} + +fn oauth_http() -> Result { + reqwest::Client::builder() + .connect_timeout(Duration::from_secs(20)) + .timeout(Duration::from_secs(30)) + .build() + .map_err(|error| format!("Unable to initialize OAuth client: {error}")) +} + +/// Resolve Application Default Credentials. Credential files are parsed here +/// so `impersonated_service_account` and `authorized_user` files referenced by +/// GOOGLE_APPLICATION_CREDENTIALS work; everything else uses gcp_auth's chain +/// (metadata server, gcloud CLI). +async fn application_default() -> Result { + if let Some(path) = adc_file_path() { + let text = tokio::fs::read_to_string(&path).await.map_err(|error| { + format!( + "Unable to read Application Default Credentials file '{}': {error}", + path.display() + ) + })?; + let text = strip_bom(&text); + let credential_type = serde_json::from_str::(text) + .ok() + .and_then(|value| value.get("type").and_then(Value::as_str).map(str::to_owned)); + if matches!( + credential_type.as_deref(), + Some("service_account" | "authorized_user" | "impersonated_service_account") + ) { + return from_credential_json( + text, + &[ + "service_account", + "authorized_user", + "impersonated_service_account", + ], + ) + .await; + } + } + let build: GoogleBuild = Arc::new(|| { + Box::pin(async { + gcp_auth::provider().await.map_err(|error| { + format!("Application Default Credentials are unavailable: {error}") + }) + }) + }); + GoogleProvider::new(build) + .await + .map(|provider| AuthProvider::Google(Arc::new(provider))) +} + +fn adc_file_path() -> Option { + if let Some(path) = std::env::var_os("GOOGLE_APPLICATION_CREDENTIALS") + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + { + return Some(path); + } + let directory = if cfg!(windows) { + PathBuf::from(std::env::var_os("APPDATA")?).join("gcloud") } else { - tokio::fs::read_to_string(Path::new(credential)) + PathBuf::from(std::env::var_os("HOME")?) + .join(".config") + .join("gcloud") + }; + let path = directory.join("application_default_credentials.json"); + path.is_file().then_some(path) +} + +/// Build a provider from a Google credential JSON document of an allowed type. +fn from_credential_json<'a>(text: &'a str, allowed: &'a [&'a str]) -> ProviderFuture<'a> { + Box::pin(async move { + let text = strip_bom(text); + let value: Value = serde_json::from_str(text) + .map_err(|error| format!("Invalid credential JSON: {error}"))?; + let credential_type = value + .get("type") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + if !allowed.contains(&credential_type.as_str()) { + return Err(format!( + "Credential JSON has type '{credential_type}', expected {}", + allowed.join(" or ") + )); + } + match credential_type.as_str() { + "service_account" => { + validate_token_uri(&value)?; + let key: Arc = Arc::from(text); + let build: GoogleBuild = Arc::new(move || { + let key = key.clone(); + Box::pin(async move { + CustomServiceAccount::from_json(&key) + .map(|account| Arc::new(account) as Arc) + .map_err(|error| format!("Invalid service account credential: {error}")) + }) + }); + GoogleProvider::new(build) + .await + .map(|provider| AuthProvider::Google(Arc::new(provider))) + } + "authorized_user" => { + let credentials: AuthorizedUserCredentials = serde_json::from_value(value) + .map_err(|error| format!("Invalid authorized user credential JSON: {error}"))?; + if credentials.client_id.trim().is_empty() + || credentials.client_secret.trim().is_empty() + || credentials.refresh_token.trim().is_empty() + { + return Err( + "OAuth credential JSON must contain client_id, client_secret and refresh_token" + .to_string(), + ); + } + Ok(AuthProvider::AuthorizedUser(Arc::new( + AuthorizedUserProvider { + http: oauth_http()?, + credentials, + cached: Mutex::new(None), + }, + ))) + } + _ => { + let url = value + .get("service_account_impersonation_url") + .and_then(Value::as_str) + .ok_or_else(|| { + "Impersonated credential JSON must contain service_account_impersonation_url" + .to_string() + }) + .and_then(validate_impersonation_url)?; + let source = value + .get("source_credentials") + .ok_or_else(|| { + "Impersonated credential JSON must contain source_credentials".to_string() + })? + .to_string(); + let source = + from_credential_json(&source, &["service_account", "authorized_user"]).await?; + let delegates = value + .get("delegates") + .and_then(Value::as_array) + .map(|delegates| { + delegates + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect() + }) + .unwrap_or_default(); + Ok(AuthProvider::Impersonated(Arc::new( + ImpersonatedProvider::new(source, url, delegates)?, + ))) + } + } + }) +} + +/// A service account key must only send its signed assertion to Google. +fn validate_token_uri(value: &Value) -> Result<(), String> { + let Some(token_uri) = value.get("token_uri").and_then(Value::as_str) else { + return Err("Service account JSON must contain token_uri".to_string()); + }; + let url = Url::parse(token_uri) + .map_err(|_| format!("Service account token_uri '{token_uri}' is not a valid URL"))?; + if url.scheme() == "https" && is_googleapis_host(&url) { + Ok(()) + } else { + Err(format!( + "Service account token_uri '{token_uri}' is not a googleapis.com HTTPS endpoint" + )) + } +} + +fn validate_impersonation_url(value: &str) -> Result { + let url = Url::parse(value) + .map_err(|_| format!("Invalid service_account_impersonation_url '{value}'"))?; + if url.scheme() == "https" + && is_googleapis_host(&url) + && url.path().ends_with(":generateAccessToken") + { + Ok(url) + } else { + Err(format!( + "service_account_impersonation_url '{value}' must be a googleapis.com generateAccessToken URL" + )) + } +} + +fn impersonation_url_for(target: &str) -> Result { + let valid = target.contains('@') + && target + .chars() + .all(|character| character.is_ascii_alphanumeric() || "@._-".contains(character)); + if !valid { + return Err(format!( + "Impersonated service account '{target}' must be a service account email" + )); + } + Url::parse(&format!( + "{IAM_CREDENTIALS_BASE}projects/-/serviceAccounts/{target}:generateAccessToken" + )) + .map_err(|error| format!("Invalid impersonation target: {error}")) +} + +fn is_googleapis_host(url: &Url) -> bool { + url.host_str() + .is_some_and(|host| host == "googleapis.com" || host.ends_with(".googleapis.com")) +} + +fn strip_bom(value: &str) -> &str { + value.strip_prefix('\u{feff}').unwrap_or(value) +} + +/// Inline JSON is used as-is; anything else is read as a file path. +async fn load_json_credential(credential: &str) -> Result { + let credential = strip_bom(credential).trim(); + if credential.starts_with('{') { + return Ok(credential.to_string()); + } + let text = tokio::fs::read_to_string(Path::new(credential)) + .await + .map_err(|error| format!("Unable to read credential file '{credential}': {error}"))?; + Ok(strip_bom(&text).to_string()) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + use serde_json::json; + use wiremock::matchers::{body_partial_json, header, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::{ + from_credential_json, impersonation_url_for, load_json_credential, validate_token_uri, + AuthProvider, GoogleBuild, GoogleProvider, ImpersonatedProvider, + }; + use crate::endpoint; + use crate::models::ConnectionParams; + + fn service_account(token_uri: &str) -> serde_json::Value { + json!({ + "type": "service_account", + "project_id": "p", + "private_key_id": "k", + "private_key": "-----BEGIN PRIVATE KEY-----\nMIIB\n-----END PRIVATE KEY-----\n", + "client_email": "sa@p.iam.gserviceaccount.com", + "client_id": "1", + "token_uri": token_uri + }) + } + + #[test] + fn accepts_only_googleapis_token_endpoints() { + assert!( + validate_token_uri(&service_account("https://oauth2.googleapis.com/token")).is_ok() + ); + assert!(validate_token_uri(&service_account("https://evil.example.com/token")).is_err()); + assert!( + validate_token_uri(&service_account("http://oauth2.googleapis.com/token")).is_err() + ); + assert!(validate_token_uri(&service_account("https://googleapis.com.evil.io/t")).is_err()); + assert!(validate_token_uri(&json!({"type": "service_account"})).is_err()); + } + + #[tokio::test] + async fn rejects_service_accounts_with_foreign_token_uri() { + let text = service_account("https://attacker.example/token").to_string(); + let error = from_credential_json(&text, &["service_account"]) + .await + .err() + .unwrap(); + assert!(error.contains("googleapis.com"), "{error}"); + } + + #[tokio::test] + async fn tolerates_a_byte_order_mark() { + let text = format!( + "\u{feff}{}", + json!({"type": "authorized_user", "client_id": "c", "client_secret": "s", "refresh_token": "r"}) + ); + assert!(load_json_credential(&text).await.unwrap().starts_with('{')); + let provider = from_credential_json(&text, &["authorized_user"]) + .await + .unwrap(); + assert!(matches!(provider, AuthProvider::AuthorizedUser(_))); + } + + #[tokio::test] + async fn parses_impersonated_adc_files() { + let text = json!({ + "type": "impersonated_service_account", + "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/target@p.iam.gserviceaccount.com:generateAccessToken", + "delegates": ["middle@p.iam.gserviceaccount.com"], + "source_credentials": {"type": "authorized_user", "client_id": "c", "client_secret": "s", "refresh_token": "r"} + }) + .to_string(); + let provider = from_credential_json(&text, &["impersonated_service_account"]) + .await + .unwrap(); + let AuthProvider::Impersonated(provider) = provider else { + panic!("expected an impersonated provider"); + }; + assert_eq!( + provider.delegates, + vec!["projects/-/serviceAccounts/middle@p.iam.gserviceaccount.com"] + ); + + let foreign = text.replace("iamcredentials.googleapis.com", "iam.example.com"); + assert!( + from_credential_json(&foreign, &["impersonated_service_account"]) + .await + .is_err() + ); + } + + #[tokio::test] + async fn impersonation_exchanges_the_source_token() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path( + "/v1/projects/-/serviceAccounts/t@p.iam.gserviceaccount.com:generateAccessToken", + )) + .and(header("authorization", "Bearer source-token")) + .and(body_partial_json(json!({"lifetime": "3600s"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "accessToken": "impersonated-token", + "expireTime": "2999-01-01T00:00:00Z" + }))) + .expect(1) + .mount(&server) + .await; + let url = format!( + "{}/v1/projects/-/serviceAccounts/t@p.iam.gserviceaccount.com:generateAccessToken", + server.uri() + ) + .parse() + .unwrap(); + let provider = ImpersonatedProvider::new( + AuthProvider::AccessToken("source-token".into()), + url, + Vec::new(), + ) + .unwrap(); + assert_eq!(provider.access_token().await.unwrap(), "impersonated-token"); + // The second call is served from the cache; `expect(1)` verifies it. + assert_eq!(provider.access_token().await.unwrap(), "impersonated-token"); + } + + struct FixedToken(&'static str); + + #[async_trait::async_trait] + impl gcp_auth::TokenProvider for FixedToken { + async fn token(&self, _scopes: &[&str]) -> Result, gcp_auth::Error> { + Ok(Arc::new( + serde_json::from_value(json!({"access_token": self.0, "expires_in": 3600})) + .unwrap(), + )) + } + + async fn project_id(&self) -> Result, gcp_auth::Error> { + Ok(Arc::from("p")) + } + } + + #[tokio::test] + async fn invalidating_google_credentials_rebuilds_the_provider() { + let builds = Arc::new(AtomicUsize::new(0)); + let counter = builds.clone(); + let build: GoogleBuild = Arc::new(move || { + let token = if counter.fetch_add(1, Ordering::SeqCst) == 0 { + "first" + } else { + "second" + }; + Box::pin( + async move { Ok(Arc::new(FixedToken(token)) as Arc) }, + ) + }); + let provider = AuthProvider::Google(Arc::new(GoogleProvider::new(build).await.unwrap())); + assert_eq!( + provider.access_token().await.unwrap().as_deref(), + Some("first") + ); + assert_eq!( + provider.access_token().await.unwrap().as_deref(), + Some("first") + ); + provider.invalidate().await; + assert_eq!( + provider.access_token().await.unwrap().as_deref(), + Some("second") + ); + assert_eq!(builds.load(Ordering::SeqCst), 2); + } + + #[test] + fn validates_impersonation_targets() { + assert!(impersonation_url_for("reader@p.iam.gserviceaccount.com").is_ok()); + assert!(impersonation_url_for("not-an-email").is_err()); + assert!(impersonation_url_for("a@b/../../x").is_err()); + } + + #[tokio::test] + async fn none_requires_a_local_endpoint() { + let params = ConnectionParams::from_value(&json!({ + "database": "p", + "extra": {"auth_method": "none"} + })); + let public = endpoint::resolve(None).unwrap(); + assert!(AuthProvider::from_params(¶ms, &public).await.is_err()); + let local = endpoint::resolve(Some("http://localhost:9050")).unwrap(); + let provider = AuthProvider::from_params(¶ms, &local).await.unwrap(); + assert_eq!(provider.access_token().await.unwrap(), None); + } + + #[tokio::test] + async fn impersonation_wraps_any_authenticated_mode() { + let params = ConnectionParams::from_value(&json!({ + "database": "p", + "extra": { + "auth_method": "access_token", + "credential": "token", + "impersonate_service_account": "reader@p.iam.gserviceaccount.com" + } + })); + let provider = AuthProvider::from_params(¶ms, &endpoint::resolve(None).unwrap()) .await - .map_err(|error| format!("Unable to read credential file '{credential}': {error}")) + .unwrap(); + assert!(matches!(provider, AuthProvider::Impersonated(_))); + + let params = ConnectionParams::from_value(&json!({ + "database": "p", + "extra": { + "auth_method": "none", + "impersonate_service_account": "reader@p.iam.gserviceaccount.com" + } + })); + let local = endpoint::resolve(Some("http://localhost:9050")).unwrap(); + assert!(AuthProvider::from_params(¶ms, &local).await.is_err()); } } diff --git a/src/bigquery.rs b/src/bigquery.rs index 0b4ab79..bd2dc94 100644 --- a/src/bigquery.rs +++ b/src/bigquery.rs @@ -5,7 +5,7 @@ use serde_json::{json, Map, Value}; const JS_MAX_SAFE_INTEGER: i64 = 9_007_199_254_740_991; -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] +#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)] #[serde(rename_all = "camelCase")] pub struct FieldSchema { pub name: String, @@ -19,6 +19,20 @@ pub struct FieldSchema { pub fields: Vec, #[serde(default)] pub default_value_expression: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub range_element_type: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_length: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub precision: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scale: Option, +} + +#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq)] +pub struct RangeElementType { + #[serde(rename = "type")] + pub element_type: String, } impl FieldSchema { @@ -54,6 +68,8 @@ pub struct QueryOutput { pub total_bytes_processed: Option, pub cache_hit: Option, pub job_id: Option, + /// How the rows were read, for example `job` or `tabledata.list`. + pub source: Option, } impl QueryOutput { @@ -72,7 +88,8 @@ impl QueryOutput { "metadata": { "total_bytes_processed": self.total_bytes_processed, "cache_hit": self.cache_hit, - "job_id": self.job_id + "job_id": self.job_id, + "source": self.source } }) } @@ -181,8 +198,7 @@ fn convert_scalar_or_record(value: &Value, field: &FieldSchema) -> Value { .unwrap_or_else(|| value.clone()), "TIMESTAMP" => value .as_str() - .and_then(|value| value.parse::().ok()) - .and_then(epoch_to_rfc3339) + .and_then(timestamp_to_rfc3339) .map(Value::String) .unwrap_or_else(|| value.clone()), "BYTES" => value @@ -199,10 +215,40 @@ fn convert_scalar_or_record(value: &Value, field: &FieldSchema) -> Value { } } -fn epoch_to_rfc3339(epoch: f64) -> Option { - let seconds = epoch.floor() as i64; - let nanos = ((epoch - seconds as f64) * 1_000_000_000.0).round() as u32; - DateTime::::from_timestamp(seconds, nanos) +/// Format a wire TIMESTAMP exactly. Results requested with +/// `formatOptions.useInt64Timestamp` carry integer microseconds; older +/// responses carry decimal or scientific epoch seconds. +pub fn timestamp_to_rfc3339(value: &str) -> Option { + let value = value.trim(); + let micros = if let Ok(micros) = value.parse::() { + micros + } else if value.contains(['e', 'E']) { + (value.parse::().ok()? * 1_000_000.0).round() as i64 + } else { + let (negative, digits) = match value.strip_prefix('-') { + Some(rest) => (true, rest), + None => (false, value), + }; + let (seconds, fraction) = digits.split_once('.').unwrap_or((digits, "")); + let mut fraction = fraction.to_string(); + if fraction.len() > 6 || !fraction.chars().all(|digit| digit.is_ascii_digit()) { + return None; + } + while fraction.len() < 6 { + fraction.push('0'); + } + let micros = seconds + .parse::() + .ok()? + .checked_mul(1_000_000)? + .checked_add(fraction.parse::().ok()?)?; + if negative { + -micros + } else { + micros + } + }; + DateTime::::from_timestamp_micros(micros) .map(|date| date.to_rfc3339_opts(SecondsFormat::Micros, true)) } @@ -231,7 +277,35 @@ pub fn api_error(response: &Value) -> Option { mod tests { use serde_json::json; - use super::{fields_from_response, rows_from_response, FieldSchema}; + use super::{fields_from_response, rows_from_response, timestamp_to_rfc3339, FieldSchema}; + + #[test] + fn formats_timestamps_with_exact_microseconds() { + assert_eq!( + timestamp_to_rfc3339("1700000000123457").unwrap(), + "2023-11-14T22:13:20.123457Z" + ); + assert_eq!( + timestamp_to_rfc3339("1700000000.123457").unwrap(), + "2023-11-14T22:13:20.123457Z" + ); + assert_eq!( + timestamp_to_rfc3339("1.700000000123457E9").unwrap(), + "2023-11-14T22:13:20.123457Z" + ); + assert_eq!( + timestamp_to_rfc3339("-1").unwrap(), + "1969-12-31T23:59:59.999999Z" + ); + assert_eq!( + timestamp_to_rfc3339("-0.5").unwrap(), + "1969-12-31T23:59:59.500000Z" + ); + assert_eq!( + timestamp_to_rfc3339("253402300799999999").unwrap(), + "9999-12-31T23:59:59.999999Z" + ); + } #[test] fn renders_nested_and_repeated_type_names() { @@ -246,9 +320,9 @@ mod tests { mode: None, description: None, fields: vec![], - default_value_expression: None, + ..Default::default() }], - default_value_expression: None, + ..Default::default() }; assert_eq!(field.display_type(), "ARRAY>"); } diff --git a/src/browse.rs b/src/browse.rs new file mode 100644 index 0000000..b718a80 --- /dev/null +++ b/src/browse.rs @@ -0,0 +1,210 @@ +//! Recognition of plain table-browsing queries. +//! +//! Tabularis opens a table with `SELECT * FROM `dataset`.`table`` and an +//! optional `LIMIT`. Such a query can be served by `tabledata.list`, which is +//! free and scans nothing. The match is deliberately exact: any WHERE, ORDER +//! BY, column list, wildcard table or decorator falls back to a query job. + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BrowseTarget { + pub project: Option, + pub dataset: String, + pub table: String, + pub limit: Option, + pub offset: u64, +} + +pub fn parse(sql: &str, default_dataset: Option<&str>) -> Option { + let mut rest = sql.trim(); + rest = rest.strip_suffix(';').unwrap_or(rest).trim_end(); + rest = keyword(rest, "SELECT")?; + rest = rest.strip_prefix('*')?.trim_start(); + rest = keyword(rest, "FROM")?; + let (parts, mut rest) = table_reference(rest)?; + rest = rest.trim_start(); + + let mut limit = None; + let mut offset = 0; + if !rest.is_empty() { + rest = keyword(rest, "LIMIT")?; + let (value, tail) = number(rest)?; + limit = Some(value); + rest = tail.trim_start(); + if !rest.is_empty() { + rest = keyword(rest, "OFFSET")?; + let (value, tail) = number(rest)?; + offset = value; + rest = tail.trim_start(); + } + } + if !rest.is_empty() { + return None; + } + if parts + .iter() + .any(|part| part.eq_ignore_ascii_case("INFORMATION_SCHEMA") || part.starts_with("region-")) + { + return None; + } + + let (project, dataset, table) = match parts.as_slice() { + [table] => ( + None, + default_dataset + .filter(|dataset| !dataset.is_empty())? + .to_string(), + table.clone(), + ), + [dataset, table] => (None, dataset.clone(), table.clone()), + [project, dataset, table] => (Some(project.clone()), dataset.clone(), table.clone()), + _ => return None, + }; + if !valid_dataset(&dataset) { + return None; + } + Some(BrowseTarget { + project, + dataset, + table, + limit, + offset, + }) +} + +/// Consume a case-insensitive keyword followed by whitespace. +fn keyword<'a>(input: &'a str, word: &str) -> Option<&'a str> { + let head = input.get(..word.len())?; + let tail = &input[word.len()..]; + (head.eq_ignore_ascii_case(word) && tail.starts_with(char::is_whitespace)) + .then(|| tail.trim_start()) +} + +fn number(input: &str) -> Option<(u64, &str)> { + let end = input + .find(|character: char| !character.is_ascii_digit()) + .unwrap_or(input.len()); + if end == 0 { + return None; + } + Some((input[..end].parse().ok()?, &input[end..])) +} + +/// Parse `a.b.c`, `` `a`.`b` `` or `` `a.b` `` into its path parts. +fn table_reference(input: &str) -> Option<(Vec, &str)> { + let mut parts = Vec::new(); + let mut rest = input; + loop { + if let Some(quoted) = rest.strip_prefix('`') { + let end = quoted.find('`')?; + let content = "ed[..end]; + if content.is_empty() + || content.contains(['\\', '*', '$', '@', '{', '}']) + || content.trim() != content + { + return None; + } + parts.extend(content.split('.').map(str::to_owned)); + rest = "ed[end + 1..]; + } else { + let end = rest + .find(|character: char| !(character.is_ascii_alphanumeric() || character == '_')) + .unwrap_or(rest.len()); + let name = &rest[..end]; + if name.is_empty() || name.starts_with(|character: char| character.is_ascii_digit()) { + return None; + } + parts.push(name.to_string()); + rest = &rest[end..]; + } + match rest.strip_prefix('.') { + Some(tail) => rest = tail, + None => break, + } + } + if parts.iter().any(String::is_empty) || parts.len() > 3 { + return None; + } + // Reject anything glued to the reference, such as an alias or a comma. + if !(rest.is_empty() || rest.starts_with(char::is_whitespace)) { + return None; + } + Some((parts, rest)) +} + +fn valid_dataset(dataset: &str) -> bool { + !dataset.is_empty() + && dataset + .chars() + .all(|character| character.is_ascii_alphanumeric() || character == '_') +} + +#[cfg(test)] +mod tests { + use super::{parse, BrowseTarget}; + + fn target(project: Option<&str>, dataset: &str, table: &str) -> BrowseTarget { + BrowseTarget { + project: project.map(str::to_owned), + dataset: dataset.into(), + table: table.into(), + limit: None, + offset: 0, + } + } + + #[test] + fn matches_host_table_open_queries() { + assert_eq!( + parse("SELECT * FROM `sales`.`orders`", None), + Some(target(None, "sales", "orders")) + ); + assert_eq!( + parse("select * from `my-project.sales.orders` LIMIT 100;", None), + Some(BrowseTarget { + limit: Some(100), + ..target(Some("my-project"), "sales", "orders") + }) + ); + assert_eq!( + parse("SELECT * FROM sales.orders LIMIT 10 OFFSET 20", None), + Some(BrowseTarget { + limit: Some(10), + offset: 20, + ..target(None, "sales", "orders") + }) + ); + assert_eq!( + parse("SELECT * FROM `orders`", Some("sales")), + Some(target(None, "sales", "orders")) + ); + assert_eq!( + parse("SELECT * FROM `sales`.`order items`", None), + Some(target(None, "sales", "order items")) + ); + } + + #[test] + fn rejects_everything_else() { + for sql in [ + "SELECT * FROM `orders`", + "SELECT * FROM sales.orders WHERE id = 1", + "SELECT * FROM sales.orders ORDER BY id LIMIT 10", + "SELECT id FROM sales.orders", + "SELECT * EXCEPT (id) FROM sales.orders", + "SELECT * FROM `sales.events_*`", + "SELECT * FROM `sales.events$20240101`", + "SELECT * FROM sales.INFORMATION_SCHEMA.TABLES", + "SELECT * FROM `region-us`.INFORMATION_SCHEMA.JOBS", + "SELECT * FROM sales.orders o", + "SELECT * FROM sales.orders, sales.items", + "SELECT * FROM (SELECT * FROM `sales`.`orders` LIMIT 10) AS limited_subset", + "SELECT * FROM sales.orders LIMIT @n", + "SELECT * FROM sales.orders; DELETE FROM sales.orders WHERE true", + "SELECT * FROM a.b.c.d", + "SELECT * FROM `sales`.`orders` FOR SYSTEM_TIME AS OF TIMESTAMP '2024-01-01'", + "SELECT * FROM sales.orders TABLESAMPLE SYSTEM (10 PERCENT)", + ] { + assert_eq!(parse(sql, None), None, "{sql}"); + } + } +} diff --git a/src/cancel.rs b/src/cancel.rs new file mode 100644 index 0000000..28e1bef --- /dev/null +++ b/src/cancel.rs @@ -0,0 +1,179 @@ +//! Cancellation of in-flight requests. +//! +//! Every JSON-RPC request runs inside [`scope`], which registers it under its +//! request id. Query code registers the BigQuery job it started with +//! [`register_job`]; a host `cancel` notification for that id then stops the +//! local wait and cancels the job on the server. + +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; + +use serde_json::Value; +use tokio::sync::Notify; + +pub type CancelFuture = Pin + Send>>; +pub type JobCanceller = Arc CancelFuture + Send + Sync>; + +#[derive(Default)] +pub struct RequestHandle { + cancelled: AtomicBool, + notify: Notify, + job: Mutex>, +} + +impl RequestHandle { + fn cancel(&self) { + if self.cancelled.swap(true, Ordering::SeqCst) { + return; + } + self.notify.notify_waiters(); + let canceller = self.job.lock().ok().and_then(|job| job.clone()); + if let Some(canceller) = canceller { + spawn_cancel(canceller); + } + } +} + +tokio::task_local! { + static CURRENT: Arc; +} + +static ACTIVE: OnceLock>>> = OnceLock::new(); + +fn active() -> &'static Mutex>> { + ACTIVE.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn key(id: &Value) -> String { + id.to_string() +} + +/// Run `future` as the request identified by `id`, so it can be cancelled. +pub async fn scope(id: Value, future: F) -> F::Output { + let handle = Arc::new(RequestHandle::default()); + let id = key(&id); + if let Ok(mut active) = active().lock() { + active.insert(id.clone(), handle.clone()); + } + let output = CURRENT.scope(handle.clone(), future).await; + if let Ok(mut active) = active().lock() { + if active + .get(&id) + .is_some_and(|current| Arc::ptr_eq(current, &handle)) + { + active.remove(&id); + } + } + output +} + +/// Handle a host `cancel` notification. Unknown ids are ignored. +pub fn cancel(id: &Value) -> bool { + let handle = active() + .lock() + .ok() + .and_then(|active| active.get(&key(id)).cloned()); + match handle { + Some(handle) => { + handle.cancel(); + true + } + None => false, + } +} + +/// Whether the current request has been cancelled. +pub fn is_cancelled() -> bool { + CURRENT + .try_with(|handle| handle.cancelled.load(Ordering::SeqCst)) + .unwrap_or(false) +} + +/// Resolves when the current request is cancelled; never outside a scope. +pub async fn cancelled() { + let Ok(handle) = CURRENT.try_with(Arc::clone) else { + return std::future::pending().await; + }; + loop { + let notified = handle.notify.notified(); + if handle.cancelled.load(Ordering::SeqCst) { + return; + } + notified.await; + } +} + +/// Attach the server-side job of the current request. When the request was +/// already cancelled the job is cancelled immediately. +pub fn register_job(canceller: JobCanceller) { + let _ = CURRENT.try_with(|handle| { + if let Ok(mut job) = handle.job.lock() { + *job = Some(canceller.clone()); + } + if handle.cancelled.load(Ordering::SeqCst) { + spawn_cancel(canceller); + } + }); +} + +/// Detach the job once it has finished, so a late cancel is a no-op. +pub fn clear_job() { + let _ = CURRENT.try_with(|handle| { + if let Ok(mut job) = handle.job.lock() { + *job = None; + } + }); +} + +fn spawn_cancel(canceller: JobCanceller) { + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(canceller()); + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use std::time::Duration; + + use serde_json::json; + + use super::{cancel, cancelled, is_cancelled, register_job, scope, JobCanceller}; + + #[tokio::test] + async fn cancel_wakes_the_request_and_cancels_its_job() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let canceller: JobCanceller = Arc::new(move || { + let counter = counter.clone(); + Box::pin(async move { + counter.fetch_add(1, Ordering::SeqCst); + }) + }); + let request = tokio::spawn(scope(json!(41), async move { + register_job(canceller); + tokio::select! { + _ = cancelled() => is_cancelled(), + _ = tokio::time::sleep(Duration::from_secs(5)) => false, + } + })); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(cancel(&json!(41))); + assert!(request.await.unwrap()); + tokio::time::sleep(Duration::from_millis(20)).await; + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(!cancel(&json!(41)), "finished requests are unregistered"); + } + + #[tokio::test] + async fn unknown_ids_and_unscoped_code_are_ignored() { + assert!(!cancel(&json!(999_999))); + assert!(!is_cancelled()); + let waited = tokio::time::timeout(Duration::from_millis(20), cancelled()).await; + assert!(waited.is_err()); + } +} diff --git a/src/catalog.rs b/src/catalog.rs new file mode 100644 index 0000000..c4c0033 --- /dev/null +++ b/src/catalog.rs @@ -0,0 +1,227 @@ +//! Short-lived metadata cache and bounded-concurrency catalog fetches. +//! +//! Opening a table makes Tabularis ask for columns, keys and indexes in quick +//! succession, and every one of them is answered from the same `tables.get` +//! resource. Only per-object resources are cached, for [`METADATA_TTL`]; +//! object lists are always fetched so created and dropped objects show up +//! at once. Every DDL statement issued by this plugin, and every DDL builder +//! call, clears the cache. + +use std::collections::HashMap; +use std::future::Future; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::{Duration, Instant}; + +use serde_json::Value; +use tokio::sync::Semaphore; +use tokio::task::JoinSet; + +use crate::client::BigQueryClient; + +pub const METADATA_TTL: Duration = Duration::from_secs(5); +/// Concurrent `tables.get` calls for bulk metadata. +pub const FETCH_CONCURRENCY: usize = 8; +const MAX_ENTRIES: usize = 4_096; + +type Entries = HashMap; + +static CACHE: OnceLock> = OnceLock::new(); + +fn entries() -> &'static Mutex { + CACHE.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn cache_key(client: &BigQueryClient, kind: &str, parts: &[&str]) -> String { + format!( + "{:p}\u{1f}{}\u{1f}{kind}\u{1f}{}", + client as *const BigQueryClient, + client.project_id, + parts.join("\u{1f}") + ) +} + +fn lookup(key: &str) -> Option { + let mut entries = entries().lock().ok()?; + match entries.get(key) { + Some((stored, value)) if stored.elapsed() < METADATA_TTL => Some(value.clone()), + Some(_) => { + entries.remove(key); + None + } + None => None, + } +} + +fn store(key: String, value: &Value) { + if let Ok(mut entries) = entries().lock() { + if entries.len() >= MAX_ENTRIES { + entries.retain(|_, (stored, _)| stored.elapsed() < METADATA_TTL); + if entries.len() >= MAX_ENTRIES { + entries.clear(); + } + } + entries.insert(key, (Instant::now(), value.clone())); + } +} + +/// Drops every cached resource. Called after any DDL statement. +pub fn invalidate() { + if let Ok(mut entries) = entries().lock() { + entries.clear(); + } +} + +async fn cached(key: String, load: F) -> Result +where + F: FnOnce() -> Fut, + Fut: Future>, +{ + if let Some(value) = lookup(&key) { + return Ok(value); + } + let value = load().await?; + store(key, &value); + Ok(value) +} + +pub async fn dataset( + client: &BigQueryClient, + project: &str, + dataset: &str, +) -> Result { + cached(cache_key(client, "dataset", &[project, dataset]), || { + client.get_dataset(project, dataset) + }) + .await +} + +/// Dataset location, used to run INFORMATION_SCHEMA jobs in the right region. +pub async fn dataset_location( + client: &BigQueryClient, + project: &str, + name: &str, +) -> Option { + dataset(client, project, name) + .await + .ok() + .and_then(|resource| { + resource + .get("location") + .and_then(Value::as_str) + .map(str::to_owned) + }) +} + +pub async fn tables( + client: &BigQueryClient, + project: &str, + dataset: &str, +) -> Result, String> { + client.list_tables_in(project, dataset).await +} + +pub async fn table( + client: &BigQueryClient, + project: &str, + dataset: &str, + table: &str, +) -> Result { + cached( + cache_key(client, "table", &[project, dataset, table]), + || client.get_table_in(project, dataset, table), + ) + .await +} + +pub async fn routines( + client: &BigQueryClient, + project: &str, + dataset: &str, +) -> Result, String> { + client.list_routines(project, dataset).await +} + +pub async fn routine( + client: &BigQueryClient, + project: &str, + dataset: &str, + routine: &str, +) -> Result { + cached( + cache_key(client, "routine", &[project, dataset, routine]), + || client.get_routine(project, dataset, routine), + ) + .await +} + +/// Fetches each table resource once with at most `concurrency` requests in +/// flight. Results keep the order of `names`. +pub async fn tables_batch( + client: Arc, + project: &str, + dataset: &str, + names: Vec, + concurrency: usize, +) -> Vec<(String, Result)> { + let semaphore = Arc::new(Semaphore::new(concurrency.max(1))); + let mut tasks = JoinSet::new(); + for (index, name) in names.iter().cloned().enumerate() { + let client = client.clone(); + let semaphore = semaphore.clone(); + let project = project.to_string(); + let dataset = dataset.to_string(); + tasks.spawn(async move { + let _permit = semaphore.acquire_owned().await; + let result = table(&client, &project, &dataset, &name).await; + (index, result) + }); + } + let mut results: Vec>> = names.iter().map(|_| None).collect(); + while let Some(joined) = tasks.join_next().await { + if let Ok((index, result)) = joined { + results[index] = Some(result); + } + } + names + .into_iter() + .zip(results) + .map(|(name, result)| { + let result = result + .unwrap_or_else(|| Err(format!("Metadata request for '{name}' was cancelled"))); + (name, result) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{cached, invalidate, lookup}; + + #[tokio::test] + async fn caches_until_invalidated() { + let key = "test\u{1f}cache".to_string(); + let first = cached(key.clone(), || async { Ok(json!(1)) }) + .await + .unwrap(); + let second = cached(key.clone(), || async { Ok(json!(2)) }) + .await + .unwrap(); + assert_eq!(first, json!(1)); + assert_eq!(second, json!(1)); + invalidate(); + assert!(lookup(&key).is_none()); + let third = cached(key, || async { Ok(json!(3)) }).await.unwrap(); + assert_eq!(third, json!(3)); + } + + #[tokio::test] + async fn failed_loads_are_not_cached() { + let key = "test\u{1f}failure".to_string(); + assert!(cached(key.clone(), || async { Err("boom".to_string()) }) + .await + .is_err()); + assert!(lookup(&key).is_none()); + } +} diff --git a/src/client.rs b/src/client.rs index 8ab1448..7fa7748 100644 --- a/src/client.rs +++ b/src/client.rs @@ -2,36 +2,47 @@ use std::collections::HashMap; use std::sync::{Arc, OnceLock}; use std::time::Duration; -use reqwest::{Method, StatusCode, Url}; +use reqwest::Url; use serde_json::{json, Value}; use sha2::{Digest, Sha256}; use tokio::sync::Mutex; -use tokio::time::{sleep, Instant}; use crate::auth::AuthProvider; -use crate::bigquery::{ - api_error, fields_from_response, rows_from_response, FieldSchema, QueryOutput, -}; +use crate::bigquery::{FieldSchema, QueryOutput}; use crate::models::ConnectionParams; -const API_BASE: &str = "https://bigquery.googleapis.com/bigquery/v2/"; +mod jobs; +mod tabledata; +mod transport; + +pub use jobs::{JobRef, QueryRequest}; + const DEFAULT_PAGE_SIZE: u32 = 10_000; static CLIENTS: OnceLock>>> = OnceLock::new(); +#[derive(Clone)] pub struct BigQueryClient { http: reqwest::Client, auth: AuthProvider, + /// Project that runs and bills query jobs (the billing project). pub project_id: String, + /// Project of unqualified names and of table browsing (the selected + /// database, falling back to the billing project). + data_project: String, location: Option, max_bytes_billed: Option, timeout: Duration, + api_base: Url, + /// Fingerprint of the connection settings, used to scope cached results. + key: String, } impl BigQueryClient { async fn new(params: &ConnectionParams) -> Result { let timeout = Duration::from_secs(params.timeout_seconds()); - let project_id = params.project_id()?.to_string(); + let project_id = params.billing_project_id()?.to_string(); + let data_project = params.data_project_id()?.to_string(); let max_bytes_billed = params .max_bytes_billed() .map(|value| { @@ -50,119 +61,185 @@ impl BigQueryClient { )) .build() .map_err(|error| format!("Unable to initialize the HTTP client: {error}"))?; + let endpoint = crate::endpoint::resolve(params.api_endpoint())?; Ok(Self { http, - auth: AuthProvider::from_params(params).await?, + auth: AuthProvider::from_params(params, &endpoint).await?, project_id, + data_project, location: params.location().map(str::to_owned), max_bytes_billed, timeout, + api_base: endpoint.base, + key: cache_key(params), }) } fn url(&self, segments: &[&str]) -> Result { - let mut url = Url::parse(API_BASE).map_err(|error| error.to_string())?; + let mut url = self.api_base.clone(); url.path_segments_mut() .map_err(|_| "Invalid BigQuery API base URL".to_string())? + .pop_if_empty() .extend(segments); Ok(url) } - async fn request_json( + pub async fn execute_query( &self, - method: Method, - url: Url, - query: &[(String, String)], - body: Option<&Value>, - ) -> Result { - let mut last_error = String::new(); - for attempt in 0..4u32 { - let token = self.auth.access_token().await?; - let mut request = self - .http - .request(method.clone(), url.clone()) - .bearer_auth(token) - .query(query); - if let Some(body) = body { - request = request.json(body); - } + sql: &str, + dataset: Option<&str>, + limit: Option, + page: u32, + ) -> Result { + self.run_query( + &QueryRequest { + sql, + dataset, + page_size: limit.unwrap_or(DEFAULT_PAGE_SIZE), + page: page.max(1), + fetch_rows: true, + ..QueryRequest::default() + }, + None, + ) + .await + } +} - match request.send().await { - Ok(response) => { - let status = response.status(); - let text = response.text().await.map_err(|error| { - format!("Unable to read the BigQuery response: {error}") - })?; - let payload: Value = - serde_json::from_str(&text).unwrap_or_else(|_| json!({"raw": text})); - if status.is_success() { - return Ok(payload); - } - let message = api_error(&payload).unwrap_or_else(|| format!("HTTP {status}")); - last_error = format!("BigQuery API request failed ({status}): {message}"); - if !is_retryable(status) || attempt == 3 { - return Err(last_error); - } - } - Err(error) => { - last_error = format!("BigQuery API request failed: {error}"); - if attempt == 3 { - return Err(last_error); - } - } - } - sleep(Duration::from_millis(250 * 2u64.pow(attempt))).await; +/// Every setting that changes how the client talks to BigQuery must be part +/// of the key, or an edited connection would keep using a stale client. +fn cache_key(params: &ConnectionParams) -> String { + let timeout = params.timeout_seconds().to_string(); + let mut hasher = Sha256::new(); + for part in [ + params.billing_project_id().unwrap_or_default(), + params.data_project_id().unwrap_or_default(), + params.auth_method(), + params.credential().unwrap_or_default(), + params.location().unwrap_or_default(), + params.max_bytes_billed().unwrap_or_default(), + timeout.as_str(), + params.api_endpoint().unwrap_or_default(), + params.impersonate_service_account().unwrap_or_default(), + ] { + hasher.update(part); + hasher.update([0]); + } + format!("{:x}", hasher.finalize()) +} + +pub async fn client_for(params: &ConnectionParams) -> Result, String> { + let fingerprint = cache_key(params); + let key = params + .connection_id + .as_deref() + .map(|connection_id| format!("{connection_id}:{fingerprint}")) + .unwrap_or(fingerprint); + let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); + if let Some(client) = cache.lock().await.get(&key).cloned() { + return Ok(client); + } + let client = Arc::new(BigQueryClient::new(params).await?); + cache.lock().await.insert(key, client.clone()); + Ok(client) +} + +pub async fn clear_cache() { + if let Some(cache) = CLIENTS.get() { + cache.lock().await.clear(); + } +} + +pub fn table_fields(table: &Value) -> Result, String> { + let fields = table + .pointer("/schema/fields") + .cloned() + .unwrap_or_else(|| json!([])); + serde_json::from_value(fields) + .map_err(|error| format!("Invalid BigQuery table schema: {error}")) +} + +// --------------------------------------------------------------------------- +// Catalog REST calls and parameterized statements used by the metadata, DDL, +// CRUD and routine handlers. Every call takes the target project explicitly so +// multi-project browsing never depends on the billing project. + +/// One GoogleSQL statement or script with optional named query parameters. +#[derive(Debug, Clone, Default)] +pub struct Statement { + pub sql: String, + /// Complete `QueryParameter` resources (`name`, `parameterType`, + /// `parameterValue`). Non-empty switches the job to named parameters. + pub parameters: Vec, + /// `(project, dataset)` used to resolve unqualified names. + pub default_dataset: Option<(String, String)>, + /// Overrides the connection location, e.g. with the dataset location. + pub location: Option, + /// Rows to return; 0 runs the statement for its side effects only. + pub max_results: u32, +} + +impl Statement { + pub fn new(sql: impl Into) -> Self { + Self { + sql: sql.into(), + ..Self::default() } - Err(last_error) } - pub async fn list_datasets(&self) -> Result, String> { - let url = self.url(&["projects", &self.project_id, "datasets"])?; - let mut page_token: Option = None; - let mut datasets = Vec::new(); - loop { - let mut query = vec![("maxResults".to_string(), "1000".to_string())]; - if let Some(token) = &page_token { - query.push(("pageToken".to_string(), token.clone())); - } - let response = self - .request_json(Method::GET, url.clone(), &query, None) - .await?; - if let Some(entries) = response.get("datasets").and_then(Value::as_array) { - datasets.extend(entries.iter().filter_map(|entry| { - entry - .pointer("/datasetReference/datasetId") - .and_then(Value::as_str) - .map(str::to_owned) - })); - } - page_token = response - .get("nextPageToken") - .and_then(Value::as_str) - .map(str::to_owned); - if page_token.is_none() { - break; - } + pub fn with_parameters(mut self, parameters: Vec) -> Self { + self.parameters = parameters; + self + } + + pub fn in_dataset(mut self, project: &str, dataset: &str) -> Self { + self.default_dataset = Some((project.to_string(), dataset.to_string())); + self + } + + pub fn at_location(mut self, location: Option) -> Self { + self.location = location; + self + } + + pub fn returning(mut self, max_results: u32) -> Self { + self.max_results = max_results; + self + } + + /// The job request for this statement, run through the same job path as + /// user queries (requestId, retries, jobTimeoutMs, labels, cancellation). + pub fn request(&self) -> QueryRequest<'_> { + let (project, dataset) = match &self.default_dataset { + Some((project, dataset)) => (Some(project.as_str()), Some(dataset.as_str())), + None => (None, None), + }; + QueryRequest { + sql: &self.sql, + dataset, + project, + location: self.location.as_deref(), + parameters: &self.parameters, + page_size: self.max_results, + page: 1, + fetch_rows: self.max_results > 0, } - datasets.sort(); - Ok(datasets) } +} - pub async fn list_tables(&self, dataset: &str) -> Result, String> { - let url = self.url(&["projects", &self.project_id, "datasets", dataset, "tables"])?; +impl BigQueryClient { + async fn list_paged(&self, segments: &[&str], key: &str) -> Result, String> { let mut page_token: Option = None; - let mut tables = Vec::new(); + let mut items = Vec::new(); loop { let mut query = vec![("maxResults".to_string(), "1000".to_string())]; if let Some(token) = &page_token { query.push(("pageToken".to_string(), token.clone())); } - let response = self - .request_json(Method::GET, url.clone(), &query, None) - .await?; - tables.extend( + let response = self.get(segments, &query).await.map_err(String::from)?; + items.extend( response - .get("tables") + .get(key) .and_then(Value::as_array) .cloned() .unwrap_or_default(), @@ -170,237 +247,156 @@ impl BigQueryClient { page_token = response .get("nextPageToken") .and_then(Value::as_str) + .filter(|token| !token.is_empty()) .map(str::to_owned); if page_token.is_none() { - break; + return Ok(items); } } - Ok(tables) } - pub async fn get_table(&self, dataset: &str, table: &str) -> Result { - let url = self.url(&[ - "projects", - &self.project_id, - "datasets", - dataset, - "tables", - table, - ])?; - self.request_json(Method::GET, url, &[], None).await + /// Projects visible to the credential (`GET projects`). + pub async fn list_projects(&self) -> Result, String> { + let projects = self.list_paged(&["projects"], "projects").await?; + Ok(projects + .iter() + .filter_map(|project| { + project + .pointer("/projectReference/projectId") + .and_then(Value::as_str) + .map(str::to_owned) + }) + .collect()) } - pub async fn execute_query( - &self, - sql: &str, - dataset: Option<&str>, - limit: Option, - page: u32, - ) -> Result { - self.run_query(sql, dataset, limit, page, false).await + pub async fn list_datasets_in(&self, project: &str) -> Result, String> { + let datasets = self + .list_paged(&["projects", project, "datasets"], "datasets") + .await?; + let mut names = datasets + .iter() + .filter_map(|dataset| { + dataset + .pointer("/datasetReference/datasetId") + .and_then(Value::as_str) + .map(str::to_owned) + }) + .collect::>(); + names.sort(); + Ok(names) } - pub async fn dry_run(&self, sql: &str, dataset: Option<&str>) -> Result { - let output = self.run_query(sql, dataset, Some(0), 1, true).await?; - Ok(json!({ - "kind": "bigquery_dry_run", - "total_bytes_processed": output.total_bytes_processed, - "job_id": output.job_id - })) + pub async fn get_dataset(&self, project: &str, dataset: &str) -> Result { + self.get(&["projects", project, "datasets", dataset], &[]) + .await + .map_err(String::from) } - async fn run_query( - &self, - sql: &str, - dataset: Option<&str>, - limit: Option, - page: u32, - dry_run: bool, - ) -> Result { - if sql.trim().is_empty() { - return Err("Query must not be empty".to_string()); - } - let url = self.url(&["projects", &self.project_id, "queries"])?; - let requested_size = limit.unwrap_or(DEFAULT_PAGE_SIZE); - let initial_timeout_ms = self.timeout.as_millis().min(120_000) as u64; - let mut body = json!({ - "query": sql, - "useLegacySql": false, - "timeoutMs": initial_timeout_ms, - "maxResults": if dry_run { 0 } else { requested_size.saturating_add(1) }, - "dryRun": dry_run - }); - if let Some(location) = &self.location { - body["location"] = Value::String(location.clone()); - } - if let Some(maximum) = &self.max_bytes_billed { - body["maximumBytesBilled"] = Value::String(maximum.clone()); - } - if let Some(dataset) = dataset.filter(|value| !value.is_empty()) { - body["defaultDataset"] = json!({"projectId": self.project_id, "datasetId": dataset}); - } + pub async fn list_tables_in(&self, project: &str, dataset: &str) -> Result, String> { + self.list_paged( + &["projects", project, "datasets", dataset, "tables"], + "tables", + ) + .await + } - let initial = self - .request_json(Method::POST, url, &[], Some(&body)) - .await?; - if let Some(error) = api_error(&initial) { - return Err(error); - } - let job_id = initial - .pointer("/jobReference/jobId") - .and_then(Value::as_str) - .map(str::to_owned); - if dry_run { - return Ok(output_from_response(initial, requested_size, job_id)); - } + pub async fn get_table_in( + &self, + project: &str, + dataset: &str, + table: &str, + ) -> Result { + self.get( + &["projects", project, "datasets", dataset, "tables", table], + &[], + ) + .await + .map_err(String::from) + } - let job_id = job_id.ok_or_else(|| "BigQuery did not return a query job ID".to_string())?; - let location = initial - .pointer("/jobReference/location") - .and_then(Value::as_str) - .map(str::to_owned) - .or_else(|| self.location.clone()); - let offset = u64::from(page.saturating_sub(1)) * u64::from(requested_size); - let final_response = self - .wait_for_results( - &job_id, - location.as_deref(), - offset, - requested_size.saturating_add(1), - ) - .await?; - let mut output = output_from_response(final_response, requested_size, Some(job_id)); - if output.rows.len() > requested_size as usize { - output.rows.truncate(requested_size as usize); - output.has_more = true; - } - if let Some(total) = output.total_rows { - output.has_more = offset.saturating_add(output.rows.len() as u64) < total; - } - Ok(output) + pub async fn list_routines(&self, project: &str, dataset: &str) -> Result, String> { + self.list_paged( + &["projects", project, "datasets", dataset, "routines"], + "routines", + ) + .await } - async fn wait_for_results( + pub async fn get_routine( &self, - job_id: &str, - location: Option<&str>, - start_index: u64, - max_results: u32, + project: &str, + dataset: &str, + routine: &str, ) -> Result { - let url = self.url(&["projects", &self.project_id, "queries", job_id])?; - let started = Instant::now(); - loop { - let mut query = vec![ - ("timeoutMs".to_string(), "10000".to_string()), - ("startIndex".to_string(), start_index.to_string()), - ("maxResults".to_string(), max_results.to_string()), - ]; - if let Some(location) = location { - query.push(("location".to_string(), location.to_string())); - } - let response = self - .request_json(Method::GET, url.clone(), &query, None) - .await?; - if let Some(error) = api_error(&response) { - return Err(error); - } - if response - .get("jobComplete") - .and_then(Value::as_bool) - .unwrap_or(false) - { - return Ok(response); - } - if started.elapsed() >= self.timeout { - return Err(format!( - "BigQuery query timed out after {} seconds", - self.timeout.as_secs() - )); - } - sleep(Duration::from_millis(250)).await; - } + self.get( + &[ + "projects", project, "datasets", dataset, "routines", routine, + ], + &[], + ) + .await + .map_err(String::from) } -} -fn output_from_response(response: Value, page_size: u32, job_id: Option) -> QueryOutput { - let fields = fields_from_response(&response).unwrap_or_default(); - let columns = fields.iter().map(|field| field.name.clone()).collect(); - let rows = rows_from_response(&response, &fields); - let affected_rows = response - .get("numDmlAffectedRows") - .and_then(Value::as_str) - .and_then(|value| value.parse().ok()) - .unwrap_or_default(); - let total_rows = response - .get("totalRows") - .and_then(Value::as_str) - .and_then(|value| value.parse().ok()); - let has_more = response.get("pageToken").is_some() || rows.len() > page_size as usize; - QueryOutput { - columns, - fields, - rows, - affected_rows, - total_rows, - has_more, - total_bytes_processed: response - .get("totalBytesProcessed") - .and_then(Value::as_str) - .map(str::to_owned), - cache_hit: response.get("cacheHit").and_then(Value::as_bool), - job_id, + /// Runs a statement in the billing project, with optional named + /// parameters, and waits for it to finish. Uses the query job runtime, + /// so it is retried, labelled, bounded by the query timeout and + /// cancellable like any user query. + pub async fn run_statement(&self, statement: &Statement) -> Result { + self.run_query(&statement.request(), None).await } } -fn is_retryable(status: StatusCode) -> bool { - status == StatusCode::TOO_MANY_REQUESTS - || status == StatusCode::INTERNAL_SERVER_ERROR - || status == StatusCode::BAD_GATEWAY - || status == StatusCode::SERVICE_UNAVAILABLE - || status == StatusCode::GATEWAY_TIMEOUT -} +#[cfg(test)] +mod tests { + use serde_json::json; -fn cache_key(params: &ConnectionParams) -> String { - let mut hasher = Sha256::new(); - hasher.update(params.project_id.as_deref().unwrap_or_default()); - hasher.update([0]); - hasher.update(params.auth_method()); - hasher.update([0]); - hasher.update(params.password.as_deref().unwrap_or_default()); - hasher.update([0]); - hasher.update(params.location().unwrap_or_default()); - hasher.update([0]); - hasher.update(params.max_bytes_billed().unwrap_or_default()); - format!("{:x}", hasher.finalize()) -} + use super::{cache_key, Statement}; + use crate::models::ConnectionParams; -pub async fn client_for(params: &ConnectionParams) -> Result, String> { - let fingerprint = cache_key(params); - let key = params - .connection_id - .as_deref() - .map(|connection_id| format!("{connection_id}:{fingerprint}")) - .unwrap_or(fingerprint); - let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); - if let Some(client) = cache.lock().await.get(&key).cloned() { - return Ok(client); + fn key(extra: serde_json::Value) -> String { + cache_key(&ConnectionParams::from_value( + &json!({"database": "p", "password": "legacy", "extra": extra}), + )) } - let client = Arc::new(BigQueryClient::new(params).await?); - cache.lock().await.insert(key, client.clone()); - Ok(client) -} -pub async fn clear_cache() { - if let Some(cache) = CLIENTS.get() { - cache.lock().await.clear(); + #[test] + fn cache_key_tracks_every_client_setting() { + let base = key(json!({"credential": "a"})); + assert_ne!(base, key(json!({"credential": "b"}))); + for (field, value) in [ + ("query_timeout_seconds", "600"), + ("location", "EU"), + ("max_bytes_billed", "1000"), + ("api_endpoint", "http://localhost:9050"), + ( + "impersonate_service_account", + "sa@p.iam.gserviceaccount.com", + ), + ("auth_method", "access_token"), + ("billing_project", "billing"), + ] { + let mut extra = json!({"credential": "a"}); + extra[field] = json!(value); + assert_ne!(base, key(extra), "{field}"); + } } -} -pub fn table_fields(table: &Value) -> Result, String> { - let fields = table - .pointer("/schema/fields") - .cloned() - .unwrap_or_else(|| json!([])); - serde_json::from_value(fields) - .map_err(|error| format!("Invalid BigQuery table schema: {error}")) + #[test] + fn statement_request_carries_dataset_location_and_parameters() { + let parameter = json!({"name": "id", "parameterType": {"type": "INT64"}, "parameterValue": {"value": "1"}}); + let statement = Statement::new("SELECT @id") + .with_parameters(vec![parameter.clone()]) + .in_dataset("data", "ds") + .at_location(Some("EU".into())) + .returning(5); + let request = statement.request(); + assert_eq!(request.project, Some("data")); + assert_eq!(request.dataset, Some("ds")); + assert_eq!(request.location, Some("EU")); + assert_eq!(request.parameters, &[parameter]); + assert_eq!(request.page_size, 5); + assert!(request.fetch_rows); + assert!(!Statement::new("DROP TABLE t").request().fetch_rows); + } } diff --git a/src/client/jobs.rs b/src/client/jobs.rs new file mode 100644 index 0000000..f687ae0 --- /dev/null +++ b/src/client/jobs.rs @@ -0,0 +1,848 @@ +//! Query job lifecycle: start, wait, cancel, page through results and reuse +//! a finished job for later pages. + +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; + +use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; +use tokio::time::{sleep_until, Instant}; + +use super::transport::CallError; +use super::BigQueryClient; +use crate::bigquery::{fields_from_response, rows_from_response, QueryOutput}; +use crate::cancel::{self, JobCanceller}; +use crate::errors::ApiError; +use crate::session::SessionState; +use crate::statements; + +/// How long one jobs.query / getQueryResults call waits server-side. +const WAIT_SLICE_MS: u64 = 10_000; +/// Finished results stay addressable for later pages for this long. +const RESULT_TTL: Duration = Duration::from_secs(10 * 60); +const RESULT_CACHE_CAPACITY: usize = 256; +/// Upper bound on jobs.list pages read for one script's child jobs. +const MAX_CHILD_JOB_PAGES: usize = 100; + +#[derive(Debug, Clone, Default)] +pub struct QueryRequest<'a> { + pub sql: &'a str, + pub dataset: Option<&'a str>, + /// Project of `dataset`; defaults to the connection's data project. + pub project: Option<&'a str>, + /// Overrides the connection location, e.g. with the dataset location. + pub location: Option<&'a str>, + /// Named `QueryParameter` resources; non-empty switches the job to + /// `parameterMode: NAMED`. + pub parameters: &'a [Value], + pub page_size: u32, + pub page: u32, + /// False when only the job statistics matter (EXPLAIN ANALYZE). + pub fetch_rows: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct JobRef { + pub project: String, + pub job_id: String, + pub location: Option, +} + +struct CachedResult { + job: JobRef, + affected_rows: u64, + stored: std::time::Instant, +} + +static RESULTS: OnceLock>> = OnceLock::new(); + +fn results() -> &'static Mutex> { + RESULTS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn lookup_result(key: &str) -> Option<(JobRef, u64)> { + let cache = results().lock().ok()?; + cache + .get(key) + .filter(|entry| entry.stored.elapsed() < RESULT_TTL) + .map(|entry| (entry.job.clone(), entry.affected_rows)) +} + +fn store_result(key: String, job: JobRef, affected_rows: u64) { + let Ok(mut cache) = results().lock() else { + return; + }; + cache.retain(|_, entry| entry.stored.elapsed() < RESULT_TTL); + if cache.len() >= RESULT_CACHE_CAPACITY { + if let Some(oldest) = cache + .iter() + .min_by_key(|(_, entry)| entry.stored) + .map(|(key, _)| key.clone()) + { + cache.remove(&oldest); + } + } + cache.insert( + key, + CachedResult { + job, + affected_rows, + stored: std::time::Instant::now(), + }, + ); +} + +/// Rows read for one page, before value conversion. +struct PageData { + schema: Value, + rows: Vec, + total_rows: Option, + page_token: Option, +} + +impl PageData { + fn into_output(self, offset: u64, page_size: u32) -> QueryOutput { + let response = json!({"schema": self.schema, "rows": self.rows}); + let fields = fields_from_response(&response).unwrap_or_default(); + let mut rows = rows_from_response(&response, &fields); + rows.truncate(page_size as usize); + let has_more = match self.total_rows { + Some(total) => offset.saturating_add(rows.len() as u64) < total, + None => self.page_token.is_some(), + }; + QueryOutput { + columns: fields.iter().map(|field| field.name.clone()).collect(), + fields, + rows, + affected_rows: 0, + total_rows: self.total_rows, + has_more, + total_bytes_processed: None, + cache_hit: None, + job_id: None, + source: None, + } + } +} + +fn string_u64(value: Option<&Value>) -> Option { + value.and_then(|value| match value { + Value::String(text) => text.parse().ok(), + other => other.as_u64(), + }) +} + +fn rows_of(response: &Value) -> Vec { + response + .get("rows") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default() +} + +fn page_token_of(response: &Value) -> Option { + response + .get("pageToken") + .and_then(Value::as_str) + .filter(|token| !token.is_empty()) + .map(str::to_owned) +} + +impl BigQueryClient { + // Project selection for jobs and unqualified names lives only in these two + // helpers. + + /// Project that runs and bills query jobs (`billing_project`, or the + /// connection's project). + fn job_project(&self) -> &str { + &self.project_id + } + + /// Project of unqualified table names (`defaultDataset.projectId`) and of + /// table browsing: the selected database, not the billing project. + pub(crate) fn data_project(&self) -> &str { + &self.data_project + } + + /// Run a statement and return one page of its result. + pub async fn run_query( + &self, + request: &QueryRequest<'_>, + session: Option<&mut SessionState>, + ) -> Result { + self.run_job(request, session) + .await + .map(|(output, _)| output) + } + + /// Execute a statement and return its finished job resource, for + /// EXPLAIN ANALYZE. + pub async fn run_for_statistics( + &self, + sql: &str, + dataset: Option<&str>, + session: Option<&mut SessionState>, + ) -> Result { + let request = QueryRequest { + sql, + dataset, + page_size: 0, + page: 1, + fetch_rows: false, + ..QueryRequest::default() + }; + let (_, job) = self.run_job(&request, session).await?; + let job = job.ok_or_else(|| "BigQuery did not return a query job".to_string())?; + self.get_job(&job).await.map_err(String::from) + } + + async fn run_job( + &self, + request: &QueryRequest<'_>, + mut session: Option<&mut SessionState>, + ) -> Result<(QueryOutput, Option), String> { + if request.sql.trim().is_empty() { + return Err("Query must not be empty".to_string()); + } + let page_size = request.page_size; + let offset = u64::from(request.page.saturating_sub(1)) * u64::from(page_size); + let session_key = session.as_ref().and_then(|state| state.bigquery_id.clone()); + // Only used for the lookup: `finish` re-keys after it captured the id + // of a session created by this statement, so later pages find it. + let cache_key = self.result_key(request, session_key.as_deref()); + + if request.page > 1 && request.fetch_rows { + if let Some((job, affected_rows)) = lookup_result(&cache_key) { + let page = self.fetch_page(&job, offset, page_size).await?; + let mut output = page.into_output(offset, page_size); + output.affected_rows = affected_rows; + output.job_id = Some(job.job_id.clone()); + output.source = Some("job".to_string()); + return Ok((output, Some(job))); + } + if !statements::is_read_only_query(request.sql) { + return Err( + "The result of this statement has expired. Run it again to fetch more rows." + .to_string(), + ); + } + } + + if session.is_none() && request.fetch_rows && request.parameters.is_empty() { + if let Some(mut target) = crate::browse::parse(request.sql, request.dataset) { + if target.project.is_none() { + target.project = request.project.map(str::to_owned); + } + if let Ok(Some(output)) = self.browse_table(&target, request.page, page_size).await + { + return Ok((output, None)); + } + } + } + + let script = statements::may_be_script(request.sql); + let rows_requested = request.fetch_rows && offset == 0 && !script; + let max_results = if rows_requested { page_size } else { 0 }; + let deadline = Instant::now() + self.timeout; + + let (job, initial) = self + .start_query(request, session.as_deref_mut(), max_results) + .await?; + let canceller = self.canceller(job.clone()); + cancel::register_job(canceller); + let outcome = async { + let done = self + .wait_for_completion(&job, initial.clone(), max_results, deadline) + .await?; + self.finish( + request, + &job, + &initial, + done, + session, + rows_requested, + offset, + ) + .await + } + .await; + cancel::clear_job(); + outcome.map(|output| (output, Some(job))) + } + + #[allow(clippy::too_many_arguments)] + async fn finish( + &self, + request: &QueryRequest<'_>, + job: &JobRef, + initial: &Value, + done: Value, + session: Option<&mut SessionState>, + rows_requested: bool, + offset: u64, + ) -> Result { + let mut resource: Option = None; + let has_errors = done + .get("errors") + .and_then(Value::as_array) + .is_some_and(|errors| !errors.is_empty()); + if has_errors { + let job_resource = self.get_job(job).await?; + if let Some(error) = job_resource.pointer("/status/errorResult") { + // A failed `createSession` job may still have created the + // session: remember it so releasing the tab aborts it. + let mut failed = Some(job_resource.clone()); + if let Some(state) = session { + self.capture_session(state, job, initial, &mut failed).await; + } + return Err(ApiError::from_error_proto(error).describe()); + } + resource = Some(job_resource); + } + + let mut session_key: Option = None; + if let Some(state) = session { + self.capture_session(state, job, initial, &mut resource) + .await; + session_key = state.bigquery_id.clone(); + } + let cache_key = self.result_key(request, session_key.as_deref()); + + let mut result_job = Some(job.clone()); + let mut affected_rows = string_u64(done.get("numDmlAffectedRows")).unwrap_or_default(); + let mut use_initial_rows = rows_requested; + if statements::may_be_script(request.sql) { + if resource.is_none() { + resource = self.get_job(job).await.ok(); + } + let statement_type = resource + .as_ref() + .and_then(|value| value.pointer("/statistics/query/statementType")) + .and_then(Value::as_str); + if statement_type == Some("SCRIPT") { + let children = self.list_child_jobs(job).await?; + let (child, affected) = script_result(&children, job); + result_job = child; + affected_rows = affected; + use_initial_rows = false; + } + } + + let mut output = match (&result_job, request.fetch_rows) { + (Some(result_job), true) => { + let page = if use_initial_rows && result_job == job { + self.complete_page(job, done.clone(), request.page_size) + .await? + } else { + self.fetch_page(result_job, offset, request.page_size) + .await? + }; + page.into_output(offset, request.page_size) + } + _ => PageData { + schema: Value::Null, + rows: Vec::new(), + total_rows: None, + page_token: None, + } + .into_output(offset, request.page_size), + }; + output.affected_rows = affected_rows; + output.total_bytes_processed = done + .get("totalBytesProcessed") + .and_then(Value::as_str) + .map(str::to_owned); + output.cache_hit = done.get("cacheHit").and_then(Value::as_bool); + output.job_id = Some(job.job_id.clone()); + output.source = Some("job".to_string()); + + if let (Some(result_job), true) = (result_job, request.fetch_rows) { + if !output.fields.is_empty() { + store_result(cache_key, result_job, affected_rows); + } + } + Ok(output) + } + + /// Record the BigQuery session a `createSession` job created, from the + /// query response or, failing that, the job resource (fetched on demand). + async fn capture_session( + &self, + state: &mut SessionState, + job: &JobRef, + initial: &Value, + resource: &mut Option, + ) { + if state.bigquery_id.is_some() { + return; + } + let mut id = initial + .pointer("/sessionInfo/sessionId") + .and_then(Value::as_str) + .map(str::to_owned); + if id.is_none() { + if resource.is_none() { + *resource = self.get_job(job).await.ok(); + } + id = resource + .as_ref() + .and_then(|value| value.pointer("/statistics/sessionInfo/sessionId")) + .and_then(Value::as_str) + .map(str::to_owned); + } + if id.is_some() { + state.bigquery_id = id; + state.location = job.location.clone(); + } + } + + async fn start_query( + &self, + request: &QueryRequest<'_>, + session: Option<&mut SessionState>, + max_results: u32, + ) -> Result<(JobRef, Value), String> { + let mut body = json!({ + "query": request.sql, + "useLegacySql": false, + "requestId": uuid::Uuid::new_v4().to_string(), + "timeoutMs": WAIT_SLICE_MS.min(self.timeout.as_millis() as u64), + "maxResults": max_results, + "jobTimeoutMs": self.timeout.as_millis().to_string(), + "labels": {"client": "tabularis"}, + "formatOptions": {"useInt64Timestamp": true} + }); + if !request.parameters.is_empty() { + body["parameterMode"] = json!("NAMED"); + body["queryParameters"] = Value::Array(request.parameters.to_vec()); + } + let mut location = request + .location + .map(str::to_owned) + .or_else(|| self.location.clone()); + if let Some(state) = session { + match &state.bigquery_id { + Some(id) => { + body["connectionProperties"] = json!([{"key": "session_id", "value": id}]); + if state.location.is_some() { + location = state.location.clone(); + } + } + None => body["createSession"] = Value::Bool(true), + } + } + if let Some(location) = &location { + body["location"] = Value::String(location.clone()); + } + if let Some(maximum) = &self.max_bytes_billed { + body["maximumBytesBilled"] = Value::String(maximum.clone()); + } + if let Some(dataset) = request.dataset.filter(|value| !value.is_empty()) { + let project = request.project.unwrap_or(self.data_project()); + body["defaultDataset"] = json!({"projectId": project, "datasetId": dataset}); + } + + let response = self + .post(&["projects", self.job_project(), "queries"], &[], &body) + .await + .map_err(String::from)?; + let job_id = response + .pointer("/jobReference/jobId") + .and_then(Value::as_str) + .ok_or_else(|| "BigQuery did not return a query job ID".to_string())?; + let job = JobRef { + project: response + .pointer("/jobReference/projectId") + .and_then(Value::as_str) + .unwrap_or(self.job_project()) + .to_string(), + job_id: job_id.to_string(), + location: response + .pointer("/jobReference/location") + .and_then(Value::as_str) + .map(str::to_owned) + .or(location), + }; + Ok((job, response)) + } + + async fn wait_for_completion( + &self, + job: &JobRef, + mut response: Value, + max_results: u32, + deadline: Instant, + ) -> Result { + loop { + if response + .get("jobComplete") + .and_then(Value::as_bool) + .unwrap_or(false) + { + return Ok(response); + } + if cancel::is_cancelled() { + return Err("Query cancelled. The BigQuery job was asked to stop.".to_string()); + } + tokio::select! { + biased; + _ = cancel::cancelled() => { + return Err("Query cancelled. The BigQuery job was asked to stop.".to_string()); + } + _ = sleep_until(deadline) => { + let _ = self.cancel_job(job).await; + return Err(format!( + "BigQuery query exceeded the {} second query timeout and was cancelled.\nHint: raise Query timeout in the connection settings; above 120 seconds also raise the Tabularis plugin call timeout.", + self.timeout.as_secs() + )); + } + next = self.get_query_results(job, Some(0), max_results, None) => { + response = next.map_err(String::from)?; + } + } + } + } + + /// Fill a page from a completed response, following `pageToken` when the + /// server returned fewer rows than requested (responses are size capped). + async fn complete_page( + &self, + job: &JobRef, + first: Value, + page_size: u32, + ) -> Result { + let mut rows = rows_of(&first); + let mut page_token = page_token_of(&first); + let total_rows = string_u64(first.get("totalRows")); + let schema = first.get("schema").cloned().unwrap_or(Value::Null); + let mut seen_tokens: HashSet = HashSet::new(); + while rows.len() < page_size as usize { + let Some(token) = page_token.clone() else { + break; + }; + // A repeated token would read the same rows twice. + if !seen_tokens.insert(token.clone()) { + return Err(format!( + "BigQuery returned the page token '{token}' twice while reading job {}", + job.job_id + )); + } + let remaining = page_size - rows.len() as u32; + let next = self + .get_query_results(job, None, remaining, Some(&token)) + .await + .map_err(String::from)?; + let batch = rows_of(&next); + page_token = page_token_of(&next); + if batch.is_empty() { + break; + } + rows.extend(batch); + } + Ok(PageData { + schema, + rows, + total_rows, + page_token, + }) + } + + /// Read rows `offset..offset + page_size` of a finished job. + async fn fetch_page( + &self, + job: &JobRef, + offset: u64, + page_size: u32, + ) -> Result { + let first = self + .get_query_results(job, Some(offset), page_size, None) + .await + .map_err(String::from)?; + if !first + .get("jobComplete") + .and_then(Value::as_bool) + .unwrap_or(true) + { + return Err("BigQuery job is still running".to_string()); + } + self.complete_page(job, first, page_size).await + } + + async fn get_query_results( + &self, + job: &JobRef, + start_index: Option, + max_results: u32, + page_token: Option<&str>, + ) -> Result { + let mut query = vec![ + ("timeoutMs".to_string(), WAIT_SLICE_MS.to_string()), + ("maxResults".to_string(), max_results.to_string()), + ( + "formatOptions.useInt64Timestamp".to_string(), + "true".to_string(), + ), + ]; + if let Some(token) = page_token { + query.push(("pageToken".to_string(), token.to_string())); + } else if let Some(start) = start_index { + query.push(("startIndex".to_string(), start.to_string())); + } + if let Some(location) = &job.location { + query.push(("location".to_string(), location.clone())); + } + self.get(&["projects", &job.project, "queries", &job.job_id], &query) + .await + } + + pub async fn get_job(&self, job: &JobRef) -> Result { + let mut query = Vec::new(); + if let Some(location) = &job.location { + query.push(("location".to_string(), location.clone())); + } + self.get(&["projects", &job.project, "jobs", &job.job_id], &query) + .await + } + + async fn list_child_jobs(&self, job: &JobRef) -> Result, String> { + let mut children = Vec::new(); + let mut page_token: Option = None; + let mut seen_tokens: HashSet = HashSet::new(); + for _ in 0..MAX_CHILD_JOB_PAGES { + let mut query = vec![ + ("parentJobId".to_string(), job.job_id.clone()), + ("maxResults".to_string(), "1000".to_string()), + ]; + if let Some(token) = &page_token { + query.push(("pageToken".to_string(), token.clone())); + } + let response = self + .get(&["projects", &job.project, "jobs"], &query) + .await + .map_err(String::from)?; + children.extend( + response + .get("jobs") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(), + ); + page_token = response + .get("nextPageToken") + .and_then(Value::as_str) + .filter(|token| !token.is_empty()) + .map(str::to_owned); + match &page_token { + None => return Ok(children), + Some(token) if !seen_tokens.insert(token.clone()) => { + return Err(format!( + "BigQuery returned the page token '{token}' twice while listing the child jobs of {}", + job.job_id + )); + } + Some(_) => {} + } + } + Err(format!( + "Script job {} has more than {} pages of child jobs", + job.job_id, MAX_CHILD_JOB_PAGES + )) + } + + /// Ask BigQuery to stop a job. Errors are ignored by callers: the job may + /// already have finished. + pub async fn cancel_job(&self, job: &JobRef) -> Result<(), String> { + let url = self.url(&["projects", &job.project, "jobs", &job.job_id, "cancel"])?; + let mut query = Vec::new(); + if let Some(location) = &job.location { + query.push(("location".to_string(), location.clone())); + } + self.call(reqwest::Method::POST, url, &query, None) + .await + .map(|_| ()) + .map_err(String::from) + } + + fn canceller(&self, job: JobRef) -> JobCanceller { + let client = self.clone(); + Arc::new(move || { + let client = client.clone(); + let job = job.clone(); + Box::pin(async move { + if let Err(error) = client.cancel_job(&job).await { + eprintln!( + "bigquery-plugin: unable to cancel job {}: {error}", + job.job_id + ); + } + }) + }) + } + + /// Abort a BigQuery session from outside it. + pub async fn abort_session( + &self, + session_id: &str, + location: Option<&str>, + ) -> Result<(), String> { + let escaped = session_id.replace('\\', "\\\\").replace('\'', "\\'"); + let mut body = json!({ + "query": format!("CALL BQ.ABORT_SESSION('{escaped}')"), + "useLegacySql": false, + "requestId": uuid::Uuid::new_v4().to_string(), + "timeoutMs": WAIT_SLICE_MS, + "labels": {"client": "tabularis"} + }); + if let Some(location) = location.or(self.location.as_deref()) { + body["location"] = Value::String(location.to_string()); + } + self.post(&["projects", self.job_project(), "queries"], &[], &body) + .await + .map(|_| ()) + .map_err(String::from) + } + + /// Dry-run a statement and return the job resource with its estimate. + pub async fn dry_run_job( + &self, + sql: &str, + dataset: Option<&str>, + session: Option<&SessionState>, + ) -> Result { + let mut query = json!({"query": sql, "useLegacySql": false}); + if let Some(dataset) = dataset.filter(|value| !value.is_empty()) { + query["defaultDataset"] = + json!({"projectId": self.data_project(), "datasetId": dataset}); + } + if let Some(maximum) = &self.max_bytes_billed { + query["maximumBytesBilled"] = Value::String(maximum.clone()); + } + let mut location = self.location.clone(); + if let Some(state) = session { + if let Some(id) = &state.bigquery_id { + query["connectionProperties"] = json!([{"key": "session_id", "value": id}]); + location = state.location.clone().or(location); + } + } + let mut reference = json!({"projectId": self.job_project()}); + if let Some(location) = location { + reference["location"] = Value::String(location); + } + let body = json!({ + "jobReference": reference, + "configuration": {"dryRun": true, "query": query, "labels": {"client": "tabularis"}} + }); + self.post(&["projects", self.job_project(), "jobs"], &[], &body) + .await + .map_err(String::from) + } + + fn result_key(&self, request: &QueryRequest<'_>, session: Option<&str>) -> String { + let mut hasher = Sha256::new(); + let parameters = Value::Array(request.parameters.to_vec()).to_string(); + for part in [ + self.key.as_str(), + request.sql, + request.dataset.unwrap_or_default(), + request.project.unwrap_or_default(), + request.location.unwrap_or_default(), + parameters.as_str(), + session.unwrap_or_default(), + ] { + hasher.update(part); + hasher.update([0]); + } + format!("{:x}", hasher.finalize()) + } +} + +/// Pick the child job whose rows a script returns and sum DML counts. +fn script_result(children: &[Value], parent: &JobRef) -> (Option, u64) { + let mut ordered: Vec<&Value> = children.iter().collect(); + ordered.sort_by_key(|child| { + ( + string_u64(child.pointer("/statistics/creationTime")).unwrap_or_default(), + string_u64(child.pointer("/statistics/startTime")).unwrap_or_default(), + ) + }); + let affected = ordered + .iter() + .filter_map(|child| string_u64(child.pointer("/statistics/query/numDmlAffectedRows"))) + .sum(); + let result = ordered + .iter() + .rev() + .find(|child| { + child + .pointer("/statistics/query/statementType") + .and_then(Value::as_str) + == Some("SELECT") + && child + .pointer("/statistics/scriptStatistics/evaluationKind") + .and_then(Value::as_str) + != Some("EXPRESSION") + && child.pointer("/status/errorResult").is_none() + }) + .and_then(|child| { + Some(JobRef { + project: child + .pointer("/jobReference/projectId") + .and_then(Value::as_str) + .unwrap_or(&parent.project) + .to_string(), + job_id: child.pointer("/jobReference/jobId")?.as_str()?.to_string(), + location: child + .pointer("/jobReference/location") + .and_then(Value::as_str) + .map(str::to_owned) + .or_else(|| parent.location.clone()), + }) + }); + (result, affected) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{script_result, JobRef}; + + #[test] + fn script_returns_last_select_statement_and_sums_dml() { + let parent = JobRef { + project: "p".into(), + job_id: "parent".into(), + location: Some("EU".into()), + }; + let child = + |id: &str, created: &str, kind: &str, statement: &str, affected: Option<&str>| { + let mut value = json!({ + "jobReference": {"projectId": "p", "jobId": id}, + "statistics": { + "creationTime": created, + "query": {"statementType": statement}, + "scriptStatistics": {"evaluationKind": kind} + } + }); + if let Some(affected) = affected { + value["statistics"]["query"]["numDmlAffectedRows"] = json!(affected); + } + value + }; + // jobs.list returns newest first. + let children = vec![ + child("declare", "5", "EXPRESSION", "SELECT", None), + child("update", "4", "STATEMENT", "UPDATE", Some("3")), + child("second", "3", "STATEMENT", "SELECT", None), + child("insert", "2", "STATEMENT", "INSERT", Some("2")), + child("first", "1", "STATEMENT", "SELECT", None), + ]; + let (job, affected) = script_result(&children, &parent); + let job = job.unwrap(); + assert_eq!(job.job_id, "second"); + assert_eq!(job.location.as_deref(), Some("EU")); + assert_eq!(affected, 5); + assert_eq!(script_result(&[], &parent), (None, 0)); + } +} diff --git a/src/client/tabledata.rs b/src/client/tabledata.rs new file mode 100644 index 0000000..fe2ea01 --- /dev/null +++ b/src/client/tabledata.rs @@ -0,0 +1,122 @@ +//! Free table browsing through `tabledata.list`. + +use serde_json::{json, Value}; + +use super::transport::CallError; +use super::BigQueryClient; +use crate::bigquery::{fields_from_response, rows_from_response, QueryOutput}; +use crate::browse::BrowseTarget; + +fn string_u64(value: Option<&Value>) -> Option { + value.and_then(|value| match value { + Value::String(text) => text.parse().ok(), + other => other.as_u64(), + }) +} + +impl BigQueryClient { + /// Serve a `SELECT * FROM table [LIMIT n [OFFSET m]]` page without a query + /// job. Returns `Ok(None)` when the table is not a native table or the + /// server's paging looks inconsistent, so the caller runs a normal job. + pub(crate) async fn browse_table( + &self, + target: &BrowseTarget, + page: u32, + page_size: u32, + ) -> Result, CallError> { + let project = target.project.as_deref().unwrap_or(self.data_project()); + let table = self + .get( + &[ + "projects", + project, + "datasets", + &target.dataset, + "tables", + &target.table, + ], + &[], + ) + .await?; + if table.get("type").and_then(Value::as_str) != Some("TABLE") { + return Ok(None); + } + let stored_rows = string_u64(table.get("numRows")).unwrap_or_default(); + let available = stored_rows.saturating_sub(target.offset); + let total = target.limit.map_or(available, |limit| limit.min(available)); + let page_offset = u64::from(page.saturating_sub(1)) * u64::from(page_size); + let wanted = total.saturating_sub(page_offset).min(u64::from(page_size)) as usize; + let start = target.offset + page_offset; + + let mut rows: Vec = Vec::new(); + let mut page_token: Option = None; + while rows.len() < wanted { + let remaining = wanted - rows.len(); + let mut query = vec![ + ("maxResults".to_string(), remaining.to_string()), + ( + "formatOptions.useInt64Timestamp".to_string(), + "true".to_string(), + ), + ]; + match &page_token { + Some(token) => query.push(("pageToken".to_string(), token.clone())), + None => query.push(("startIndex".to_string(), start.to_string())), + } + let response = self + .get( + &[ + "projects", + project, + "datasets", + &target.dataset, + "tables", + &target.table, + "data", + ], + &query, + ) + .await?; + let batch = response + .get("rows") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + // A server that ignores startIndex/maxResults (some emulators) + // returns more rows than asked for; do not trust its pages. + if batch.len() > remaining { + return Ok(None); + } + page_token = response + .get("pageToken") + .and_then(Value::as_str) + .filter(|token| !token.is_empty()) + .map(str::to_owned); + if batch.is_empty() { + break; + } + rows.extend(batch); + if page_token.is_none() { + break; + } + } + + let response = + json!({"schema": table.get("schema").cloned().unwrap_or(Value::Null), "rows": rows}); + let fields = fields_from_response(&response).map_err(CallError::Other)?; + let rows = rows_from_response(&response, &fields); + let has_more = page_offset.saturating_add(rows.len() as u64) < total; + Ok(Some(QueryOutput { + columns: fields.iter().map(|field| field.name.clone()).collect(), + fields, + rows, + affected_rows: 0, + total_rows: Some(total), + has_more, + total_bytes_processed: Some("0".to_string()), + cache_hit: None, + job_id: None, + source: Some("tabledata.list".to_string()), + })) + } +} diff --git a/src/client/transport.rs b/src/client/transport.rs new file mode 100644 index 0000000..8d5c826 --- /dev/null +++ b/src/client/transport.rs @@ -0,0 +1,157 @@ +use std::time::Duration; + +use reqwest::{Method, StatusCode, Url}; +use serde_json::{json, Value}; +use tokio::time::sleep; + +use super::BigQueryClient; +use crate::errors::ApiError; + +const MAX_ATTEMPTS: u32 = 5; +const BASE_BACKOFF_MS: u64 = 250; +const MAX_BACKOFF_MS: u64 = 8_000; + +/// Failure of one REST call: either a decoded API error or a local problem. +#[derive(Debug, Clone)] +pub enum CallError { + Api(ApiError), + Other(String), +} + +impl CallError { + pub fn describe(&self) -> String { + match self { + Self::Api(error) => error.describe(), + Self::Other(message) => message.clone(), + } + } + + pub fn status(&self) -> Option { + match self { + Self::Api(error) => error.status, + Self::Other(_) => None, + } + } +} + +impl From for String { + fn from(error: CallError) -> Self { + error.describe() + } +} + +impl BigQueryClient { + /// Send a REST call with bearer auth. Transient failures (429, 5xx and 403 + /// rate-limit or backend errors, transport errors) are retried with + /// exponential backoff and jitter; a 401 drops the cached token (or + /// rebuilds a gcp_auth provider) and retries once, which cannot help a + /// fixed access token. Every + /// POST this plugin sends is idempotent (jobs.query carries a requestId), + /// so retrying is safe. + pub(crate) async fn call( + &self, + method: Method, + url: Url, + query: &[(String, String)], + body: Option<&Value>, + ) -> Result { + let mut refreshed = false; + let mut attempt = 0; + loop { + let token = self.auth.access_token().await.map_err(CallError::Other)?; + let mut request = self.http.request(method.clone(), url.clone()).query(query); + if let Some(token) = token { + request = request.bearer_auth(token); + } + if let Some(body) = body { + request = request.json(body); + } + + let error = match request.send().await { + Ok(response) => { + let status = response.status(); + let text = response.text().await.map_err(|error| { + CallError::Other(format!("Unable to read the BigQuery response: {error}")) + })?; + if status.is_success() { + return serde_json::from_str(&text).map_err(|error| { + CallError::Other(format!("Invalid BigQuery response: {error}")) + }); + } + let payload: Value = + serde_json::from_str(&text).unwrap_or_else(|_| json!({"raw": text})); + let error = ApiError::from_response(status.as_u16(), &payload); + if status == StatusCode::UNAUTHORIZED && !refreshed { + refreshed = true; + self.auth.invalidate().await; + continue; + } + if !error.is_retryable() { + return Err(CallError::Api(error)); + } + CallError::Api(error) + } + Err(error) => CallError::Other(format!("BigQuery API request failed: {error}")), + }; + attempt += 1; + if attempt >= MAX_ATTEMPTS { + return Err(error); + } + sleep(backoff(attempt)).await; + } + } + + pub(crate) async fn get( + &self, + segments: &[&str], + query: &[(String, String)], + ) -> Result { + let url = self.url(segments).map_err(CallError::Other)?; + self.call(Method::GET, url, query, None).await + } + + pub(crate) async fn post( + &self, + segments: &[&str], + query: &[(String, String)], + body: &Value, + ) -> Result { + let url = self.url(segments).map_err(CallError::Other)?; + self.call(Method::POST, url, query, Some(body)).await + } + + /// Cheapest authenticated call: one dataset of the project. + pub async fn ping(&self) -> Result<(), String> { + self.get( + &["projects", self.data_project(), "datasets"], + &[("maxResults".to_string(), "1".to_string())], + ) + .await + .map(|_| ()) + .map_err(String::from) + } +} + +/// Exponential backoff with up to 50 % random jitter. +fn backoff(attempt: u32) -> Duration { + let base = (BASE_BACKOFF_MS << attempt.saturating_sub(1).min(6)).min(MAX_BACKOFF_MS); + let jitter = (uuid::Uuid::new_v4().as_u128() % 1_000) as u64 * base / 2_000; + Duration::from_millis(base + jitter) +} + +#[cfg(test)] +mod tests { + use super::backoff; + + #[test] + fn backoff_grows_and_stays_bounded() { + for attempt in 1..10 { + let delay = backoff(attempt).as_millis() as u64; + let base = (250u64 << (attempt - 1).min(6)).min(8_000); + assert!( + delay >= base && delay <= base + base / 2, + "{attempt}: {delay}" + ); + } + } +} diff --git a/src/endpoint.rs b/src/endpoint.rs new file mode 100644 index 0000000..415c2fa --- /dev/null +++ b/src/endpoint.rs @@ -0,0 +1,114 @@ +use reqwest::Url; + +pub const DEFAULT_API_BASE: &str = "https://bigquery.googleapis.com/bigquery/v2/"; + +/// The BigQuery REST base URL used for every API call of a connection. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ApiEndpoint { + pub base: Url, + /// True when the endpoint is a loopback address, such as a local emulator. + pub local: bool, +} + +impl ApiEndpoint { + pub fn default_endpoint() -> Self { + Self { + base: Url::parse(DEFAULT_API_BASE).expect("valid default BigQuery API base"), + local: false, + } + } +} + +/// Resolve the optional `api_endpoint` override. HTTPS is required except for +/// loopback hosts, so credentials are never sent in clear text over a network. +pub fn resolve(value: Option<&str>) -> Result { + let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else { + return Ok(ApiEndpoint::default_endpoint()); + }; + let mut url = + Url::parse(value).map_err(|error| format!("Invalid API endpoint '{value}': {error}"))?; + let local = is_loopback(&url); + match url.scheme() { + "https" => {} + "http" if local => {} + "http" => { + return Err( + "API endpoint must use https unless it points to localhost or 127.0.0.1" + .to_string(), + ) + } + scheme => return Err(format!("Unsupported API endpoint scheme '{scheme}'")), + } + if url.host_str().is_none_or(str::is_empty) { + return Err("API endpoint must include a host".to_string()); + } + if !url.username().is_empty() || url.password().is_some() { + return Err("API endpoint must not contain user information".to_string()); + } + if url.query().is_some() || url.fragment().is_some() { + return Err("API endpoint must not contain a query string or fragment".to_string()); + } + + let path = url.path().trim_end_matches('/').to_string(); + let path = if path.ends_with("/bigquery/v2") { + format!("{path}/") + } else { + format!("{path}/bigquery/v2/") + }; + url.set_path(&path); + Ok(ApiEndpoint { base: url, local }) +} + +fn is_loopback(url: &Url) -> bool { + matches!( + url.host_str(), + Some("localhost") | Some("127.0.0.1") | Some("[::1]") + ) +} + +#[cfg(test)] +mod tests { + use super::{resolve, DEFAULT_API_BASE}; + + #[test] + fn defaults_to_the_public_api() { + let endpoint = resolve(None).unwrap(); + assert_eq!(endpoint.base.as_str(), DEFAULT_API_BASE); + assert!(!endpoint.local); + assert_eq!(resolve(Some(" ")).unwrap().base.as_str(), DEFAULT_API_BASE); + } + + #[test] + fn appends_the_rest_path_once() { + assert_eq!( + resolve(Some("http://localhost:9050")) + .unwrap() + .base + .as_str(), + "http://localhost:9050/bigquery/v2/" + ); + assert_eq!( + resolve(Some("https://bigquery-psc.p.googleapis.com/bigquery/v2/")) + .unwrap() + .base + .as_str(), + "https://bigquery-psc.p.googleapis.com/bigquery/v2/" + ); + } + + #[test] + fn allows_plain_http_only_on_loopback() { + assert!(resolve(Some("http://127.0.0.1:9050")).unwrap().local); + assert!(resolve(Some("http://[::1]:9050")).unwrap().local); + assert!(resolve(Some("http://bigquery.internal:9050")).is_err()); + assert!(!resolve(Some("https://bigquery.internal")).unwrap().local); + } + + #[test] + fn rejects_unsafe_urls() { + assert!(resolve(Some("ftp://localhost")).is_err()); + assert!(resolve(Some("https://user:secret@example.com")).is_err()); + assert!(resolve(Some("https://example.com/?key=1")).is_err()); + assert!(resolve(Some("not a url")).is_err()); + } +} diff --git a/src/errors.rs b/src/errors.rs new file mode 100644 index 0000000..f28abc3 --- /dev/null +++ b/src/errors.rs @@ -0,0 +1,222 @@ +use serde_json::Value; + +/// A BigQuery error decoded from an HTTP error body or a job `errorResult`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ApiError { + pub status: Option, + pub reason: Option, + pub message: String, + pub location: Option, +} + +impl ApiError { + /// Decode the `{"error": {...}}` body of a failed REST call. + pub fn from_response(status: u16, payload: &Value) -> Self { + let error = payload.get("error").unwrap_or(payload); + let first = error + .get("errors") + .and_then(Value::as_array) + .and_then(|errors| errors.first()); + let message = error + .get("message") + .and_then(Value::as_str) + .or_else(|| { + first + .and_then(|first| first.get("message")) + .and_then(Value::as_str) + }) + .or_else(|| payload.get("raw").and_then(Value::as_str)) + .map(str::trim) + .filter(|message| !message.is_empty()) + .map(str::to_owned) + .unwrap_or_else(|| format!("BigQuery request failed with HTTP {status}")); + Self { + status: Some(status), + reason: first + .and_then(|first| first.get("reason")) + .or_else(|| error.get("status")) + .and_then(Value::as_str) + .filter(|reason| !reason.is_empty()) + .map(str::to_owned), + message, + location: first.and_then(location_of), + } + } + + /// Decode a job's `status.errorResult` (or one entry of `errors`). + pub fn from_error_proto(error: &Value) -> Self { + Self { + status: None, + reason: error + .get("reason") + .and_then(Value::as_str) + .filter(|reason| !reason.is_empty()) + .map(str::to_owned), + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or("BigQuery job failed") + .to_string(), + location: location_of(error), + } + } + + pub fn reason(&self) -> &str { + self.reason.as_deref().unwrap_or_default() + } + + /// Whether the same request may succeed when sent again after a delay. + pub fn is_retryable(&self) -> bool { + match self.status { + Some(429 | 500 | 502 | 503 | 504) => true, + Some(403) => matches!(self.reason(), "rateLimitExceeded" | "backendError"), + _ => matches!(self.reason(), "backendError" | "internalError"), + } + } + + /// A user-facing message with an actionable hint when one is known. + pub fn describe(&self) -> String { + let mut text = self.message.clone(); + if let Some(position) = self.position() { + text.push_str(&format!(" (line {}, column {})", position.0, position.1)); + } + let suffix = match (self.reason.as_deref(), self.status) { + (Some(reason), Some(status)) => format!(" [{reason}, HTTP {status}]"), + (Some(reason), None) => format!(" [{reason}]"), + (None, Some(status)) if !self.message.ends_with(&format!("HTTP {status}")) => { + format!(" [HTTP {status}]") + } + (None, _) => String::new(), + }; + text.push_str(&suffix); + if let Some(hint) = self.hint() { + text.push_str("\nHint: "); + text.push_str(hint); + } + text + } + + fn position(&self) -> Option<(u32, u32)> { + if self.reason() != "invalidQuery" || self.message.contains("[at ") { + return None; + } + let location = self.location.as_deref()?; + let (line, column) = location.split_once(':')?; + Some((line.trim().parse().ok()?, column.trim().parse().ok()?)) + } + + fn hint(&self) -> Option<&'static str> { + let message = self.message.to_ascii_lowercase(); + if message.contains("partition elimination") || message.contains("without a filter over") { + return Some("The table requires a partition filter. Add a WHERE condition on the partitioning column."); + } + if message.contains("streaming buffer") { + return Some("Rows still in the streaming buffer cannot be modified by DML. Wait until the buffer is flushed (up to 90 minutes) or restrict the statement to older rows."); + } + match self.reason() { + "accessDenied" => Some("Grant the account BigQuery Job User on the billing project and BigQuery Data Viewer (or Editor for writes) on the dataset."), + "notFound" => Some("Check the project, dataset and table names, and that the connection Location matches the dataset location."), + "quotaExceeded" => Some("A BigQuery quota was exceeded. Wait and retry, or review the project quotas in the Google Cloud console."), + "rateLimitExceeded" => Some("Too many requests in a short time. Retry in a few seconds."), + "bytesBilledLimitExceeded" => Some("The query would bill more than the connection's Maximum bytes billed. Raise the limit or reduce scanned data with partition filters and fewer columns."), + "responseTooLarge" => Some("The result is too large to return. Add LIMIT, select fewer columns, or write the result to a destination table."), + "billingTierLimitExceeded" => Some("The query needs more compute per byte than the on-demand tier allows. Simplify joins or use reservations."), + "invalidQuery" => Some("Fix the GoogleSQL statement at the reported position."), + "stopped" | "jobBackendError" => None, + "duplicate" => Some("The object already exists. Use CREATE OR REPLACE or IF NOT EXISTS."), + _ if self.status == Some(401) => Some("The credential was rejected or has expired. Refresh the credential in the connection settings."), + _ => None, + } + } +} + +fn location_of(error: &Value) -> Option { + error + .get("location") + .and_then(Value::as_str) + .filter(|location| !location.is_empty()) + .map(str::to_owned) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::ApiError; + + fn error(status: u16, reason: &str, message: &str) -> ApiError { + ApiError::from_response( + status, + &json!({"error": {"code": status, "message": message, "errors": [{"reason": reason, "message": message}]}}), + ) + } + + #[test] + fn maps_common_reasons_to_hints() { + for (reason, fragment) in [ + ("accessDenied", "BigQuery Job User"), + ("notFound", "Location"), + ("quotaExceeded", "quota"), + ("rateLimitExceeded", "Retry"), + ("bytesBilledLimitExceeded", "Maximum bytes billed"), + ("responseTooLarge", "LIMIT"), + ("billingTierLimitExceeded", "reservations"), + ] { + let text = error(400, reason, "failed").describe(); + assert!(text.contains(reason), "{text}"); + assert!(text.contains(fragment), "{reason}: {text}"); + } + } + + #[test] + fn recognizes_partition_and_streaming_buffer_errors_by_message() { + let text = error( + 400, + "invalidQuery", + "Cannot query over table 'd.t' without a filter over column(s) 'day' that can be used for partition elimination", + ) + .describe(); + assert!(text.contains("partition filter")); + let text = error( + 400, + "invalidQuery", + "UPDATE or DELETE statement over table d.t would affect rows in the streaming buffer, which is not supported", + ) + .describe(); + assert!(text.contains("streaming buffer cannot be modified")); + } + + #[test] + fn reports_invalid_query_positions() { + let job_error = ApiError::from_error_proto( + &json!({"reason": "invalidQuery", "location": "3:15", "message": "Unrecognized name: foo"}), + ); + assert!(job_error + .describe() + .starts_with("Unrecognized name: foo (line 3, column 15)")); + let inline = error( + 400, + "invalidQuery", + "Syntax error: Unexpected keyword [at 1:8]", + ); + assert!(!inline.describe().contains("(line")); + } + + #[test] + fn classifies_retryable_errors() { + assert!(error(429, "rateLimitExceeded", "x").is_retryable()); + assert!(error(503, "backendError", "x").is_retryable()); + assert!(error(403, "rateLimitExceeded", "x").is_retryable()); + assert!(error(403, "backendError", "x").is_retryable()); + assert!(!error(403, "accessDenied", "x").is_retryable()); + assert!(!error(403, "quotaExceeded", "x").is_retryable()); + assert!(!error(400, "invalidQuery", "x").is_retryable()); + } + + #[test] + fn falls_back_to_raw_bodies() { + let parsed = ApiError::from_response(502, &json!({"raw": "Bad gateway"})); + assert_eq!(parsed.message, "Bad gateway"); + assert!(parsed.is_retryable()); + } +} diff --git a/src/explain.rs b/src/explain.rs new file mode 100644 index 0000000..e77aea4 --- /dev/null +++ b/src/explain.rs @@ -0,0 +1,134 @@ +//! Raw EXPLAIN payloads for the `bigquery-job-json` parser bundle. +//! +//! The plugin hands the host a JSON document built from a dry-run or finished +//! job; `explain/src` turns it into a Visual EXPLAIN tree. + +use serde_json::{json, Value}; + +pub const ENGINE: &str = "bigquery"; +pub const FORMAT: &str = "bigquery-job-json"; +/// On-demand list price used for the estimate, in US dollars per TiB. +pub const ON_DEMAND_USD_PER_TIB: f64 = 6.25; +const TIB: f64 = 1_099_511_627_776.0; + +pub const COST_NOTE: &str = "Estimate at the on-demand list price of $6.25 per TiB. Actual cost depends on your pricing model (editions or reservations), free tier, regional pricing and the 10 MiB minimum per referenced table."; + +fn number(value: Option<&Value>) -> Option { + value.and_then(|value| match value { + Value::String(text) => text.parse().ok(), + other => other.as_u64(), + }) +} + +pub fn estimated_cost_usd(bytes: u64) -> f64 { + bytes as f64 / TIB * ON_DEMAND_USD_PER_TIB +} + +/// Build the payload from a job resource. `mode` is `estimate` for a dry run +/// and `analyze` for an executed job. +pub fn payload(mode: &str, job: &Value) -> Value { + let statistics = job.get("statistics").cloned().unwrap_or(Value::Null); + let query = statistics.get("query").cloned().unwrap_or(Value::Null); + let processed = number(query.get("totalBytesProcessed")) + .or_else(|| number(statistics.get("totalBytesProcessed"))); + let billed = number(query.get("totalBytesBilled")); + let cost_bytes = if mode == "analyze" { + billed.or(processed) + } else { + processed + }; + let elapsed_ms = match ( + number(statistics.get("startTime")), + number(statistics.get("endTime")), + ) { + (Some(start), Some(end)) if end >= start => Some(end - start), + _ => None, + }; + let referenced_tables: Vec = query + .get("referencedTables") + .and_then(Value::as_array) + .map(|tables| { + tables + .iter() + .filter_map(|table| { + Some(format!( + "{}.{}.{}", + table.get("projectId")?.as_str()?, + table.get("datasetId")?.as_str()?, + table.get("tableId")?.as_str()? + )) + }) + .collect() + }) + .unwrap_or_default(); + + json!({ + "version": 1, + "mode": mode, + "job_id": job.pointer("/jobReference/jobId"), + "location": job.pointer("/jobReference/location"), + "statement_type": query.get("statementType"), + "total_bytes_processed": processed, + "total_bytes_processed_accuracy": query.get("totalBytesProcessedAccuracy"), + "total_bytes_billed": billed, + "estimated_cost_usd": cost_bytes.map(estimated_cost_usd), + "cost_note": COST_NOTE, + "cache_hit": query.get("cacheHit"), + "total_slot_ms": number(query.get("totalSlotMs")).or_else(|| number(statistics.get("totalSlotMs"))), + "elapsed_ms": elapsed_ms, + "billing_tier": query.get("billingTier"), + "referenced_tables": referenced_tables, + "query_plan": query.get("queryPlan").cloned().unwrap_or_else(|| json!([])), + "timeline": query.get("timeline").cloned().unwrap_or_else(|| json!([])) + }) +} + +/// The raw `explain_query` result recognised by Tabularis 0.23.0 and later. +pub fn raw_output(payload: &Value, original_query: &str) -> Value { + json!({ + "engine": ENGINE, + "format": FORMAT, + "payload": payload.to_string(), + "original_query": original_query + }) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{estimated_cost_usd, payload}; + + #[test] + fn estimates_on_demand_cost_per_tib() { + assert!((estimated_cost_usd(1_099_511_627_776) - 6.25).abs() < 1e-9); + assert_eq!(estimated_cost_usd(0), 0.0); + } + + #[test] + fn summarizes_dry_runs_and_executed_jobs() { + let job = json!({ + "jobReference": {"jobId": "j", "location": "US"}, + "statistics": { + "startTime": "1000", "endTime": "1600", + "query": { + "statementType": "SELECT", + "totalBytesProcessed": "2199023255552", + "totalBytesBilled": "1099511627776", + "cacheHit": false, + "totalSlotMs": "42", + "referencedTables": [{"projectId": "p", "datasetId": "d", "tableId": "t"}], + "queryPlan": [{"id": "0", "name": "S00: Input"}] + } + } + }); + let estimate = payload("estimate", &job); + assert_eq!(estimate["estimated_cost_usd"], 12.5); + assert_eq!(estimate["referenced_tables"], json!(["p.d.t"])); + let analyzed = payload("analyze", &job); + assert_eq!(analyzed["estimated_cost_usd"], 6.25); + assert_eq!(analyzed["elapsed_ms"], 600); + assert_eq!(analyzed["total_slot_ms"], 42); + assert_eq!(analyzed["query_plan"][0]["name"], "S00: Input"); + } +} diff --git a/src/handlers/blob.rs b/src/handlers/blob.rs index 92fd146..e588ed8 100644 --- a/src/handlers/blob.rs +++ b/src/handlers/blob.rs @@ -1,24 +1,25 @@ use base64::Engine; use serde_json::{json, Value}; -use crate::client; -use crate::models::{inner_params, ConnectionParams}; +use crate::catalog; +use crate::client::{self, Statement}; +use crate::handlers::crud::{build_predicate, row_identity}; +use crate::handlers::metadata::Target; use crate::rpc::{error_response, ok_response}; -use crate::sql::{fully_qualified, quote_identifier, value_literal}; +use crate::schema_ddl; +use crate::sql::{quote_identifier, QueryParams}; pub async fn save_blob_to_file(id: Value, params: &Value) -> Value { + let Some(path) = params.get("file_path").and_then(Value::as_str) else { + return error_response(id, -32602, "Invalid params: missing 'file_path'"); + }; match fetch_blob(params).await { - Ok(bytes) => { - let Some(path) = params.get("file_path").and_then(Value::as_str) else { - return error_response(id, -32602, "Invalid params: missing 'file_path'"); - }; - match tokio::fs::write(path, bytes).await { - Ok(()) => ok_response(id, Value::Null), - Err(error) => { - error_response(id, -32603, &format!("Unable to write BLOB file: {error}")) - } + Ok(bytes) => match tokio::fs::write(path, bytes).await { + Ok(()) => ok_response(id, Value::Null), + Err(error) => { + error_response(id, -32603, &format!("Unable to write BLOB file: {error}")) } - } + }, Err(error) => error_response(id, -32603, &error), } } @@ -39,44 +40,42 @@ pub async fn fetch_blob_as_data_url(id: Value, params: &Value) -> Value { } } +/// `SELECT TO_BASE64(column)` for one row, identified like row edits. +pub fn blob_sql(path: &str, column: &str, predicate: &str) -> String { + format!( + "SELECT TO_BASE64({}) AS data FROM {path} WHERE {predicate} LIMIT 1", + quote_identifier(column) + ) +} + async fn fetch_blob(params: &Value) -> Result, String> { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required(params, "schema")?; let table = required(params, "table")?; let column = required(params, "col_name")?; - let pk_map = params - .get("pk_map") - .and_then(Value::as_object) - .filter(|map| !map.is_empty()) - .ok_or_else(|| "A non-empty pk_map is required".to_string())?; - let client = client::client_for(&connection).await?; - let resource = client.get_table(schema, table).await?; + let identity = row_identity(params)?; + let target = Target::from_params(params).await?; + let resource = catalog::table(&target.client, &target.project, schema, table).await?; let fields = client::table_fields(&resource)?; - let find_field = |name: &str| fields.iter().find(|field| field.name == name); - let predicates = pk_map - .iter() - .map(|(name, value)| { - let expression = if value.is_null() { - "IS NULL".to_string() - } else { - format!("= {}", value_literal(value, find_field(name))?) - }; - Ok(format!("{} {expression}", quote_identifier(name))) - }) - .collect::, String>>()?; - let sql = format!( - "SELECT TO_BASE64({}) AS data FROM {} WHERE {} LIMIT 1", - quote_identifier(column), - fully_qualified(&client.project_id, schema, table), - predicates.join(" AND ") - ); - let output = client.execute_query(&sql, Some(schema), Some(1), 1).await?; - let encoded = output + if !fields.iter().any(|field| field.name == column) { + return Err(format!("Column '{column}' does not exist in {table}")); + } + let keyless = schema_ddl::primary_key_columns(&resource).is_empty(); + let mut parameters = QueryParams::new(); + let predicate = build_predicate(&identity, &fields, keyless, &mut parameters)?; + let statement = Statement::new(blob_sql(&target.path(schema, table), column, &predicate)) + .with_parameters(parameters.into_vec()) + .in_dataset(&target.project, schema) + .at_location(target.location(schema).await) + .returning(1); + let output = target.client.run_statement(&statement).await?; + let row = output .rows .first() - .and_then(|row| row.first()) + .ok_or_else(|| "The row was not found; it may have been changed or deleted".to_string())?; + let encoded = row + .first() .and_then(Value::as_str) - .ok_or_else(|| "BLOB value was not found".to_string())?; + .ok_or_else(|| "The BYTES value is NULL".to_string())?; base64::engine::general_purpose::STANDARD .decode(encoded) .map_err(|error| format!("BigQuery returned invalid base64 data: {error}")) @@ -90,3 +89,16 @@ fn required<'a>(params: &'a Value, key: &str) -> Result<&'a str, String> { .filter(|value| !value.is_empty()) .ok_or_else(|| format!("Invalid params: missing '{key}'")) } + +#[cfg(test)] +mod tests { + use super::blob_sql; + + #[test] + fn selects_one_base64_value() { + assert_eq!( + blob_sql("`p-1`.`d`.`t`", "data`x", "`id` = @p0"), + "SELECT TO_BASE64(`data\\`x`) AS data FROM `p-1`.`d`.`t` WHERE `id` = @p0 LIMIT 1" + ); + } +} diff --git a/src/handlers/connection.rs b/src/handlers/connection.rs index 40455b1..5c3d4a9 100644 --- a/src/handlers/connection.rs +++ b/src/handlers/connection.rs @@ -21,7 +21,7 @@ async fn test(id: Value, params: &Value, lightweight: bool) -> Value { let result = async { let client = client::client_for(¶ms).await?; if lightweight { - client.list_datasets().await.map(|_| ()) + client.ping().await } else { client .execute_query("SELECT 1 AS connected", None, Some(1), 1) @@ -38,6 +38,11 @@ async fn test(id: Value, params: &Value, lightweight: bool) -> Value { } pub async fn shutdown(id: Value, _params: &Value) -> Value { + let _ = tokio::time::timeout( + std::time::Duration::from_secs(5), + crate::session::release_all(), + ) + .await; client::clear_cache().await; ok_response(id, Value::Null) } diff --git a/src/handlers/crud.rs b/src/handlers/crud.rs index 92f64af..90d342f 100644 --- a/src/handlers/crud.rs +++ b/src/handlers/crud.rs @@ -1,153 +1,297 @@ -use std::collections::HashMap; - -use serde_json::{json, Value}; +use serde_json::{json, Map, Value}; use crate::bigquery::FieldSchema; -use crate::client; -use crate::models::{inner_params, ConnectionParams}; -use crate::rpc::{error_response, ok_response}; -use crate::sql::{fully_qualified, quote_identifier, value_literal}; +use crate::catalog; +use crate::client::{self, Statement}; +use crate::handlers::metadata::{respond, Target}; +use crate::schema_ddl; +use crate::sql::{base_type, is_repeated, quote_identifier, QueryParams, USE_DEFAULT_SENTINEL}; + +/// A table resolved for row editing. +pub struct EditTarget { + pub target: Target, + pub dataset: String, + pub table: String, + pub fields: Vec, + /// No informational primary key: rows are matched on comparable columns. + pub keyless: bool, +} + +impl EditTarget { + pub async fn load(params: &Value) -> Result { + let dataset = required_string(params, "schema")?.to_string(); + let table = params + .get("table") + .or_else(|| params.get("table_name")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .ok_or_else(|| "Invalid params: missing 'table'".to_string())? + .to_string(); + let target = Target::from_params(params).await?; + let resource = catalog::table(&target.client, &target.project, &dataset, &table).await?; + schema_ddl::ensure_dml_allowed(&resource)?; + Ok(Self { + fields: client::table_fields(&resource)?, + keyless: schema_ddl::primary_key_columns(&resource).is_empty(), + target, + dataset, + table, + }) + } + + pub fn path(&self) -> String { + self.target.path(&self.dataset, &self.table) + } + + pub fn field(&self, name: &str) -> Result<&FieldSchema, String> { + self.fields + .iter() + .find(|field| field.name == name) + .ok_or_else(|| format!("Column '{name}' does not exist in {}", self.table)) + } + + async fn run( + &self, + sql: String, + parameters: QueryParams, + ) -> Result { + let statement = Statement::new(sql) + .with_parameters(parameters.into_vec()) + .in_dataset(&self.target.project, &self.dataset) + .at_location(self.target.location(&self.dataset).await) + .returning(1); + self.target + .client + .run_statement(&statement) + .await + .map_err(|error| friendly_dml_error(&error)) + } +} + +/// Columns that can identify a row with `=`. Floating point, JSON, +/// geography, arrays, structs, ranges and bytes are excluded. +pub fn is_comparable(field: &FieldSchema) -> bool { + !is_repeated(field) + && !matches!( + base_type(field).as_str(), + "FLOAT64" | "JSON" | "GEOGRAPHY" | "STRUCT" | "RANGE" | "BYTES" + ) +} + +/// WHERE clause from the row identity sent by the host. For keyless tables +/// the host sends every column; only comparable ones are used. +pub fn build_predicate( + identity: &Map, + fields: &[FieldSchema], + keyless: bool, + parameters: &mut QueryParams, +) -> Result { + let mut columns = identity.keys().collect::>(); + columns.sort(); + let mut predicates = Vec::new(); + for column in columns { + let field = fields + .iter() + .find(|field| &field.name == column) + .ok_or_else(|| format!("Column '{column}' does not exist"))?; + if keyless && !is_comparable(field) { + continue; + } + let value = &identity[column]; + if value.is_null() { + predicates.push(format!("{} IS NULL", quote_identifier(column))); + } else { + predicates.push(format!( + "{} = {}", + quote_identifier(column), + parameters.bind(value, Some(field))? + )); + } + } + if predicates.is_empty() { + return Err( + "This row cannot be identified: the table has no primary key and no comparable columns" + .to_string(), + ); + } + Ok(predicates.join(" AND ")) +} + +/// Wraps a single-row DML statement for keyless tables in a transaction +/// that rolls back when more than one row would change. +pub fn guarded_script(dml: &str) -> String { + format!( + "DECLARE tabularis_affected INT64 DEFAULT 0;\n\ + BEGIN TRANSACTION;\n\ + {dml};\n\ + SET tabularis_affected = @@row_count;\n\ + IF tabularis_affected > 1 THEN\n \ + ROLLBACK TRANSACTION;\n \ + RAISE USING MESSAGE = FORMAT('The row is not unique: %d identical rows match it and the change was rolled back. Add a primary key or edit the rows with SQL.', tabularis_affected);\n\ + END IF;\n\ + COMMIT TRANSACTION;\n\ + SELECT tabularis_affected AS affected_rows;" + ) +} + +pub fn friendly_dml_error(error: &str) -> String { + let lower = error.to_ascii_lowercase(); + if lower.contains("streaming buffer") { + return format!( + "BigQuery cannot UPDATE or DELETE rows that are still in the streaming buffer (rows streamed in the last few minutes, up to about 90). Retry later. Details: {error}" + ); + } + error.to_string() +} + +fn affected_rows(output: &crate::bigquery::QueryOutput, keyless: bool) -> u64 { + if keyless { + if let Some(value) = output.rows.first().and_then(|row| row.first()) { + return value + .as_u64() + .or_else(|| value.as_str().and_then(|value| value.parse().ok())) + .unwrap_or_default(); + } + } + output.affected_rows +} + +pub fn insert_sql( + path: &str, + data: &Map, + fields: &[FieldSchema], + parameters: &mut QueryParams, +) -> Result { + if data.is_empty() { + return Err("Cannot insert an empty row".to_string()); + } + let mut columns = Vec::new(); + let mut values = Vec::new(); + for (name, value) in data { + let field = fields + .iter() + .find(|field| &field.name == name) + .ok_or_else(|| format!("Column '{name}' does not exist"))?; + columns.push(quote_identifier(name)); + if value.as_str() == Some(USE_DEFAULT_SENTINEL) { + values.push("DEFAULT".to_string()); + } else { + values.push(parameters.bind(value, Some(field))?); + } + } + Ok(format!( + "INSERT INTO {path} ({}) VALUES ({})", + columns.join(", "), + values.join(", ") + )) +} pub async fn insert_record(id: Value, params: &Value) -> Value { respond(id, async { - let (connection, schema, table) = connection_target(params)?; let data = params .get("data") .and_then(Value::as_object) .ok_or_else(|| "Invalid params: missing or malformed 'data'".to_string())?; - if data.is_empty() { - return Err("Cannot insert an empty row".to_string()); - } - let client = client::client_for(&connection).await?; - let resource = client.get_table(schema, table).await?; - let fields = client::table_fields(&resource)?; - let by_name = field_map(&fields); - let columns = data - .keys() - .map(|name| quote_identifier(name)) - .collect::>(); - let values = data - .iter() - .map(|(name, value)| value_literal(value, by_name.get(name).copied())) - .collect::, _>>()?; - let sql = format!( - "INSERT INTO {} ({}) VALUES ({})", - fully_qualified(&client.project_id, schema, table), - columns.join(", "), - values.join(", ") - ); - let result = client.execute_query(&sql, Some(schema), Some(0), 1).await?; - Ok(json!(result.affected_rows.max(1))) + let edit = EditTarget::load(params).await?; + let mut parameters = QueryParams::new(); + let sql = insert_sql(&edit.path(), data, &edit.fields, &mut parameters)?; + let output = edit.run(sql, parameters).await?; + Ok(json!(output.affected_rows)) }) .await } +pub fn update_sql( + path: &str, + column: &str, + value: &Value, + identity: &Map, + fields: &[FieldSchema], + keyless: bool, + parameters: &mut QueryParams, +) -> Result { + let field = fields + .iter() + .find(|field| field.name == column) + .ok_or_else(|| format!("Column '{column}' does not exist"))?; + let assignment = if value.as_str() == Some(USE_DEFAULT_SENTINEL) { + "DEFAULT".to_string() + } else { + parameters.bind(value, Some(field))? + }; + let predicate = build_predicate(identity, fields, keyless, parameters)?; + Ok(format!( + "UPDATE {path} SET {} = {assignment} WHERE {predicate}", + quote_identifier(column) + )) +} + pub async fn update_record(id: Value, params: &Value) -> Value { respond(id, async { - let (connection, schema, table) = connection_target(params)?; let column = required_string(params, "col_name")?; let value = params.get("new_val").cloned().unwrap_or(Value::Null); - let predicates = primary_key_map(params)?; - let client = client::client_for(&connection).await?; - let resource = client.get_table(schema, table).await?; - let fields = client::table_fields(&resource)?; - let by_name = field_map(&fields); - let assignment = value_literal(&value, by_name.get(column).copied())?; - let where_clause = build_predicate(&predicates, &by_name)?; - let sql = format!( - "UPDATE {} SET {} = {} WHERE {}", - fully_qualified(&client.project_id, schema, table), - quote_identifier(column), - assignment, - where_clause - ); - let result = client.execute_query(&sql, Some(schema), Some(0), 1).await?; - Ok(json!(result.affected_rows)) + let identity = row_identity(params)?; + let edit = EditTarget::load(params).await?; + let mut parameters = QueryParams::new(); + let sql = update_sql( + &edit.path(), + column, + &value, + &identity, + &edit.fields, + edit.keyless, + &mut parameters, + )?; + let sql = if edit.keyless { + guarded_script(&sql) + } else { + sql + }; + let output = edit.run(sql, parameters).await?; + Ok(json!(affected_rows(&output, edit.keyless))) }) .await } pub async fn delete_record(id: Value, params: &Value) -> Value { respond(id, async { - let (connection, schema, table) = connection_target(params)?; - let predicates = primary_key_map(params)?; - let client = client::client_for(&connection).await?; - let resource = client.get_table(schema, table).await?; - let fields = client::table_fields(&resource)?; - let by_name = field_map(&fields); - let where_clause = build_predicate(&predicates, &by_name)?; - let sql = format!( - "DELETE FROM {} WHERE {}", - fully_qualified(&client.project_id, schema, table), - where_clause - ); - let result = client.execute_query(&sql, Some(schema), Some(0), 1).await?; - Ok(json!(result.affected_rows)) + let identity = row_identity(params)?; + let edit = EditTarget::load(params).await?; + let mut parameters = QueryParams::new(); + let predicate = build_predicate(&identity, &edit.fields, edit.keyless, &mut parameters)?; + let sql = format!("DELETE FROM {} WHERE {predicate}", edit.path()); + let sql = if edit.keyless { + guarded_script(&sql) + } else { + sql + }; + let output = edit.run(sql, parameters).await?; + Ok(json!(affected_rows(&output, edit.keyless))) }) .await } -fn connection_target(params: &Value) -> Result<(ConnectionParams, &str, &str), String> { - let connection = ConnectionParams::from_value(inner_params(params)); - let schema = required_string(params, "schema")?; - let table = params - .get("table") - .or_else(|| params.get("table_name")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .ok_or_else(|| "Invalid params: missing 'table'".to_string())?; - Ok((connection, schema, table)) -} - -fn primary_key_map(params: &Value) -> Result, String> { +/// `pk_map`, or the legacy `pk_col` / `pk_val` pair. +pub fn row_identity(params: &Value) -> Result, String> { if let Some(map) = params.get("pk_map").and_then(Value::as_object) { if !map.is_empty() { - return Ok(map - .iter() - .map(|(key, value)| (key.clone(), value.clone())) - .collect()); + return Ok(map.clone()); } } match ( params.get("pk_col").and_then(Value::as_str), params.get("pk_val"), ) { - (Some(column), Some(value)) => Ok(HashMap::from([(column.to_string(), value.clone())])), + (Some(column), Some(value)) => { + let mut map = Map::new(); + map.insert(column.to_string(), value.clone()); + Ok(map) + } _ => Err("A non-empty primary-key map is required for update and delete".to_string()), } } -fn field_map(fields: &[FieldSchema]) -> HashMap { - fields - .iter() - .map(|field| (field.name.clone(), field)) - .collect() -} - -fn build_predicate( - values: &HashMap, - fields: &HashMap, -) -> Result { - let mut keys = values.keys().collect::>(); - keys.sort(); - keys.into_iter() - .map(|column| { - let value = &values[column]; - if value.is_null() { - Ok(format!("{} IS NULL", quote_identifier(column))) - } else { - Ok(format!( - "{} = {}", - quote_identifier(column), - value_literal(value, fields.get(column).copied())? - )) - } - }) - .collect::, String>>() - .map(|parts| parts.join(" AND ")) -} - fn required_string<'a>(params: &'a Value, key: &str) -> Result<&'a str, String> { params .get(key) @@ -157,12 +301,124 @@ fn required_string<'a>(params: &'a Value, key: &str) -> Result<&'a str, String> .ok_or_else(|| format!("Invalid params: missing '{key}'")) } -async fn respond(id: Value, future: F) -> Value -where - F: std::future::Future>, -{ - match future.await { - Ok(value) => ok_response(id, value), - Err(error) => error_response(id, -32603, &error), +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn fields() -> Vec { + serde_json::from_value(json!([ + {"name": "id", "type": "INTEGER"}, + {"name": "name", "type": "STRING"}, + {"name": "score", "type": "FLOAT"}, + {"name": "payload", "type": "JSON"}, + {"name": "location", "type": "GEOGRAPHY"}, + {"name": "tags", "type": "STRING", "mode": "REPEATED"}, + {"name": "price", "type": "NUMERIC"}, + {"name": "data", "type": "BYTES"} + ])) + .unwrap() + } + + #[test] + fn inserts_with_typed_parameters_and_defaults() { + let data = json!({"name": "O'Reilly\n", "price": "1.25", "location": "POINT(1 2)", "id": "__USE_DEFAULT__"}); + let mut parameters = QueryParams::new(); + let sql = insert_sql( + "`p-1`.`d`.`t`", + data.as_object().unwrap(), + &fields(), + &mut parameters, + ) + .unwrap(); + assert_eq!( + sql, + "INSERT INTO `p-1`.`d`.`t` (`id`, `location`, `name`, `price`) VALUES (DEFAULT, ST_GEOGFROMTEXT(@p0), @p1, CAST(@p2 AS NUMERIC))" + ); + let parameters = parameters.into_vec(); + assert_eq!(parameters[1]["parameterValue"]["value"], "O'Reilly\n"); + assert_eq!(parameters[2]["parameterValue"]["value"], "1.25"); + let unknown = json!({"nope": 1}); + assert!(insert_sql( + "t", + unknown.as_object().unwrap(), + &fields(), + &mut QueryParams::new() + ) + .is_err()); + } + + #[test] + fn updates_by_primary_key() { + let identity = json!({"id": 7}).as_object().unwrap().clone(); + let mut parameters = QueryParams::new(); + let sql = update_sql( + "`p`.`d`.`t`", + "payload", + &json!({"a": [1, 2]}), + &identity, + &fields(), + false, + &mut parameters, + ) + .unwrap(); + assert_eq!( + sql, + "UPDATE `p`.`d`.`t` SET `payload` = PARSE_JSON(@p0) WHERE `id` = @p1" + ); + let parameters = parameters.into_vec(); + assert_eq!(parameters[0]["parameterValue"]["value"], "{\"a\":[1,2]}"); + assert_eq!(parameters[1]["parameterType"]["type"], "INT64"); + } + + #[test] + fn keyless_predicates_skip_non_comparable_columns() { + let identity = json!({ + "id": 1, "name": null, "score": 1.5, "payload": {"x": 1}, + "location": "POINT(0 0)", "tags": ["a"], "price": "2.50", "data": "AQID" + }); + let mut parameters = QueryParams::new(); + let predicate = build_predicate( + identity.as_object().unwrap(), + &fields(), + true, + &mut parameters, + ) + .unwrap(); + assert_eq!( + predicate, + "`id` = @p0 AND `name` IS NULL AND `price` = CAST(@p1 AS NUMERIC)" + ); + let only_floats = json!({"score": 1.0}); + assert!(build_predicate( + only_floats.as_object().unwrap(), + &fields(), + true, + &mut QueryParams::new() + ) + .is_err()); + let guarded = guarded_script("DELETE FROM `p`.`d`.`t` WHERE `id` = @p0"); + assert!(guarded.contains("BEGIN TRANSACTION;\nDELETE FROM `p`.`d`.`t` WHERE `id` = @p0;")); + assert!(guarded.contains("ROLLBACK TRANSACTION")); + assert!(guarded + .trim_end() + .ends_with("SELECT tabularis_affected AS affected_rows;")); + } + + #[test] + fn maps_streaming_buffer_errors() { + let error = "UPDATE or DELETE statement over table p.d.t would affect rows in the streaming buffer, which is not supported"; + assert!(friendly_dml_error(error).starts_with( + "BigQuery cannot UPDATE or DELETE rows that are still in the streaming buffer" + )); + assert_eq!(friendly_dml_error("other"), "other"); + } + + #[test] + fn reads_legacy_identity_shape() { + let identity = row_identity(&json!({"pk_col": "id", "pk_val": 3})).unwrap(); + assert_eq!(identity["id"], 3); + assert!(row_identity(&json!({"pk_map": {}})).is_err()); } } diff --git a/src/handlers/ddl.rs b/src/handlers/ddl.rs index 4e0262e..e62cc85 100644 --- a/src/handlers/ddl.rs +++ b/src/handlers/ddl.rs @@ -1,11 +1,71 @@ -use serde_json::Value; +use serde_json::{json, Value}; -use crate::client; -use crate::models::{inner_params, ColumnDefinition, ConnectionParams}; +use crate::catalog; +use crate::handlers::metadata::{respond, Target}; +use crate::models::{optional_data_project, ColumnDefinition}; use crate::rpc::{error_response, ok_response}; -use crate::sql::{qualified, quote_identifier}; +use crate::schema_ddl::{self, TableKind}; +use crate::sql::{canonical_type, fully_qualified, object_path, quote_identifier, string_literal}; + +/// Rejects column changes on snapshots, external tables and views when the +/// host sends connection parameters; without them the table is not looked up. +async fn ensure_editable(params: &Value, schema: &str, table: &str) -> Result<(), String> { + if params.get("params").is_none() { + return Ok(()); + } + // Best effort: a failed lookup leaves the decision to BigQuery. + let Ok(target) = Target::from_params(params).await else { + return Ok(()); + }; + match catalog::table(&target.client, &target.project, schema, table).await { + Ok(resource) => schema_ddl::ensure_schema_editable(&resource), + Err(_) => Ok(()), + } +} + +/// Optional table settings. Current hosts send only columns; these keys are +/// accepted for callers that provide them. +#[derive(Debug, Default, Clone, PartialEq)] +pub struct TableOptions { + pub partition_by: Option, + pub cluster_by: Vec, + pub description: Option, +} + +impl TableOptions { + fn from_params(params: &Value) -> Self { + let text = |key: &str| { + params + .get(key) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + }; + let cluster_by = match params.get("cluster_by") { + Some(Value::Array(values)) => values + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect(), + Some(Value::String(value)) => value + .split(',') + .map(str::trim) + .filter(|column| !column.is_empty()) + .map(str::to_owned) + .collect(), + _ => Vec::new(), + }; + Self { + partition_by: text("partition_by"), + cluster_by, + description: text("description").or_else(|| text("comment")), + } + } +} pub async fn get_create_table_sql(id: Value, params: &Value) -> Value { + catalog::invalidate(); let table = params .get("table_name") .and_then(Value::as_str) @@ -22,30 +82,33 @@ pub async fn get_create_table_sql(id: Value, params: &Value) -> Value { "Table name, dataset and at least one column are required", ); } - ok_response( - id, - Value::from(vec![create_table_sql(schema, table, &columns)]), - ) + let path = object_path(optional_data_project(params).as_deref(), schema, table); + match create_table_sql(&path, &columns, &TableOptions::from_params(params)) { + Ok(sql) => ok_response(id, Value::from(vec![sql])), + Err(error) => error_response(id, -32602, &error), + } } pub async fn get_add_column_sql(id: Value, params: &Value) -> Value { + catalog::invalidate(); let Some((schema, table)) = target(params) else { return error_response(id, -32602, "Table and dataset are required"); }; let Some(column) = parse_column(params.get("column")) else { return error_response(id, -32602, "A valid column definition is required"); }; - ok_response( - id, - Value::from(vec![format!( - "ALTER TABLE {} ADD COLUMN {}", - qualified(schema, table), - column_sql(&column) - )]), - ) + if let Err(error) = ensure_editable(params, schema, table).await { + return error_response(id, -32602, &error); + } + let path = object_path(optional_data_project(params).as_deref(), schema, table); + match add_column_sql(&path, &column) { + Ok(statements) => ok_response(id, Value::from(statements)), + Err(error) => error_response(id, -32602, &error), + } } pub async fn get_alter_column_sql(id: Value, params: &Value) -> Value { + catalog::invalidate(); let Some((schema, table)) = target(params) else { return error_response(id, -32602, "Table and dataset are required"); }; @@ -55,7 +118,11 @@ pub async fn get_alter_column_sql(id: Value, params: &Value) -> Value { let Some(new) = parse_column(params.get("new_column")) else { return error_response(id, -32602, "A valid new_column definition is required"); }; - match alter_column_sql(schema, table, &old, &new) { + if let Err(error) = ensure_editable(params, schema, table).await { + return error_response(id, -32602, &error); + } + let path = object_path(optional_data_project(params).as_deref(), schema, table); + match alter_column_sql(&path, &old, &new) { Ok(statements) => ok_response(id, Value::from(statements)), Err(error) => error_response(id, -32603, &error), } @@ -65,7 +132,7 @@ pub async fn get_create_index_sql(id: Value, _params: &Value) -> Value { error_response( id, -32601, - "General-purpose CREATE INDEX is not supported by BigQuery", + "BigQuery has no general-purpose indexes. Search and vector indexes use CREATE SEARCH INDEX / CREATE VECTOR INDEX in a query tab.", ) } @@ -77,6 +144,7 @@ pub async fn get_create_foreign_key_sql(id: Value, params: &Value) -> Value { params .get(key) .and_then(Value::as_str) + .map(str::trim) .filter(|value| !value.is_empty()) }; let (Some(name), Some(column), Some(ref_table), Some(ref_column)) = ( @@ -91,54 +159,178 @@ pub async fn get_create_foreign_key_sql(id: Value, params: &Value) -> Value { "Foreign-key name, columns and referenced table are required", ); }; - let sql = format!( - "ALTER TABLE {} ADD CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {} ({}) NOT ENFORCED", - qualified(schema, table), - quote_identifier(name), - quote_identifier(column), - qualified(schema, ref_table), - quote_identifier(ref_column) - ); - ok_response(id, Value::from(vec![sql])) + let actions = [ + params.get("on_delete").and_then(Value::as_str), + params.get("on_update").and_then(Value::as_str), + ]; + let project = optional_data_project(params); + match foreign_key_sql( + &object_path(project.as_deref(), schema, table), + &object_path(project.as_deref(), schema, ref_table), + name, + column, + ref_column, + &actions, + ) { + Ok(sql) => ok_response(id, Value::from(vec![sql])), + Err(error) => error_response(id, -32602, &error), + } } -pub async fn drop_index(id: Value, _params: &Value) -> Value { - error_response( - id, - -32601, - "General-purpose indexes are not supported by BigQuery", - ) +pub async fn drop_index(id: Value, params: &Value) -> Value { + respond(id, async { + let (schema, table) = + target(params).ok_or_else(|| "Table and dataset are required".to_string())?; + let index = params + .get("index_name") + .and_then(Value::as_str) + .unwrap_or_default(); + if !index.eq_ignore_ascii_case("PRIMARY") { + return Err( + "BigQuery has no general-purpose indexes. Drop search or vector indexes with DROP SEARCH INDEX / DROP VECTOR INDEX in a query tab." + .to_string(), + ); + } + let target_handle = Target::from_params(params).await?; + let sql = format!( + "ALTER TABLE {} DROP PRIMARY KEY IF EXISTS", + target_handle.path(schema, table) + ); + target_handle.run_ddl(&sql, schema).await?; + Ok(Value::Null) + }) + .await } pub async fn drop_foreign_key(id: Value, params: &Value) -> Value { - let connection = ConnectionParams::from_value(inner_params(params)); - let Some((schema, table)) = target(params) else { - return error_response(id, -32602, "Table and dataset are required"); - }; - let Some(name) = params.get("fk_name").and_then(Value::as_str) else { - return error_response(id, -32602, "Foreign-key name is required"); - }; - let sql = format!( - "ALTER TABLE {} DROP CONSTRAINT {}", - qualified(schema, table), - quote_identifier(name) - ); - let result = async { - let client = client::client_for(&connection).await?; - client - .execute_query(&sql, Some(schema), Some(0), 1) + respond(id, async { + let (schema, table) = + target(params).ok_or_else(|| "Table and dataset are required".to_string())?; + let name = params + .get("fk_name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .ok_or_else(|| "Foreign-key name is required".to_string())?; + let target_handle = Target::from_params(params).await?; + let sql = format!( + "ALTER TABLE {} DROP CONSTRAINT IF EXISTS {}", + target_handle.path(schema, table), + quote_identifier(name) + ); + target_handle.run_ddl(&sql, schema).await?; + Ok(Value::Null) + }) + .await +} + +/// Statement for the explorer's drop table action. BigQuery needs the object +/// kind in the keyword (`DROP SNAPSHOT TABLE`, `DROP EXTERNAL TABLE`, +/// `DROP MATERIALIZED VIEW`, ...) and a project-qualified name, which the +/// host's generic `DROP TABLE` lacks. Returns `null` (host fallback) when the +/// request has no connection parameters or no dataset. +pub async fn get_drop_table_sql(id: Value, params: &Value) -> Value { + respond(id, async { + let Some((schema, table)) = target(params) else { + return Ok(Value::Null); + }; + let Some(connection) = params.get("params") else { + return Ok(Value::Null); + }; + // `database` repeats the browsed project; it only matters when the + // connection parameters carry none. + let mut request = json!({ "params": connection }); + let browsed = params + .get("database") + .and_then(Value::as_str) + .map(str::trim) + .filter(|database| !database.is_empty()); + let has_database = connection + .get("database") + .is_some_and(|database| match database { + Value::String(value) => !value.trim().is_empty(), + Value::Array(values) => !values.is_empty(), + _ => false, + }); + if let (Some(database), false) = (browsed, has_database) { + request["params"]["database"] = json!(database); + } + // Without a usable client (missing or invalid credentials) the kind + // cannot be looked up: drop it as a table and let BigQuery report the + // real error rather than blocking the drop in the explorer. + let target_handle = match Target::from_params(&request).await { + Ok(target_handle) => target_handle, + Err(_) => { + return Ok( + optional_data_project(&request).map_or(Value::Null, |project| { + json!(schema_ddl::drop_sql( + TableKind::Table, + &fully_qualified(&project, schema, table) + )) + }), + ); + } + }; + // Best effort: an unknown object is dropped as a table and BigQuery + // reports the real error. + let kind = catalog::table(&target_handle.client, &target_handle.project, schema, table) .await - .map(|_| ()) + .map(|resource| TableKind::from_resource(&resource)) + .unwrap_or(TableKind::Table); + Ok(json!(schema_ddl::drop_sql( + kind, + &target_handle.path(schema, table) + ))) + }) + .await +} + +fn split_columns(columns: &str) -> Vec { + columns + .split(',') + .map(str::trim) + .filter(|column| !column.is_empty()) + .map(quote_identifier) + .collect() +} + +pub fn foreign_key_sql( + table: &str, + referenced: &str, + name: &str, + columns: &str, + referenced_columns: &str, + actions: &[Option<&str>], +) -> Result { + for action in actions.iter().flatten() { + let action = action.trim().to_ascii_uppercase(); + if !action.is_empty() && action != "NO ACTION" && action != "RESTRICT" { + return Err(format!( + "BigQuery foreign keys are informational (NOT ENFORCED) and do not support ON DELETE/ON UPDATE {action}" + )); + } } - .await; - match result { - Ok(()) => ok_response(id, Value::Null), - Err(error) => error_response(id, -32603, &error), + let columns = split_columns(columns); + let referenced_columns = split_columns(referenced_columns); + if columns.len() != referenced_columns.len() { + return Err("Foreign-key and referenced column counts must match".to_string()); } + Ok(format!( + "ALTER TABLE {table} ADD CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {referenced}({}) NOT ENFORCED", + quote_identifier(name), + columns.join(", "), + referenced_columns.join(", ") + )) } -fn create_table_sql(schema: &str, table: &str, columns: &[ColumnDefinition]) -> String { - let mut definitions = columns.iter().map(column_sql).collect::>(); +pub fn create_table_sql( + path: &str, + columns: &[ColumnDefinition], + options: &TableOptions, +) -> Result { + let mut definitions = columns + .iter() + .map(|column| column_sql(column, true)) + .collect::, _>>()?; let primary_keys = columns .iter() .filter(|column| column.is_pk) @@ -150,76 +342,228 @@ fn create_table_sql(schema: &str, table: &str, columns: &[ColumnDefinition]) -> primary_keys.join(", ") )); } - format!( - "CREATE TABLE {} (\n {}\n)", - qualified(schema, table), - definitions.join(",\n ") - ) + let mut sql = format!("CREATE TABLE {path} (\n {}\n)", definitions.join(",\n ")); + if let Some(partition) = &options.partition_by { + sql.push_str(&format!("\nPARTITION BY {partition}")); + } + if !options.cluster_by.is_empty() { + sql.push_str(&format!( + "\nCLUSTER BY {}", + options + .cluster_by + .iter() + .map(|column| quote_identifier(column)) + .collect::>() + .join(", ") + )); + } + if let Some(description) = &options.description { + sql.push_str(&format!( + "\nOPTIONS(description = {})", + string_literal(description) + )); + } + Ok(sql) } -fn column_sql(column: &ColumnDefinition) -> String { - let mut output = format!( - "{} {}", - quote_identifier(&column.name), - column.data_type.trim() - ); +fn column_sql(column: &ColumnDefinition, allow_required: bool) -> Result { + if column.is_auto_increment { + return Err(format!( + "BigQuery has no auto-increment columns ('{}'). Use a DEFAULT such as GENERATE_UUID() instead.", + column.name + )); + } + let data_type = column.data_type.trim(); + if data_type.is_empty() { + return Err(format!("Column '{}' needs a data type", column.name)); + } + let mut output = format!("{} {data_type}", quote_identifier(&column.name)); + // GoogleSQL column_schema: name type [NOT NULL] [DEFAULT expr] [OPTIONS(...)] if !column.is_nullable { + if !allow_required { + return Err(format!( + "BigQuery cannot add a NOT NULL (REQUIRED) column to an existing table. Add '{}' as nullable instead.", + column.name + )); + } output.push_str(" NOT NULL"); } if let Some(default) = column .default_value .as_deref() - .filter(|value| !value.trim().is_empty()) + .map(str::trim) + .filter(|value| !value.is_empty()) { output.push_str(" DEFAULT "); - output.push_str(default.trim()); + output.push_str(default); + } + if let Some(description) = column.comment.as_deref().filter(|value| !value.is_empty()) { + output.push_str(&format!( + " OPTIONS(description = {})", + string_literal(description) + )); } - output + Ok(output) } -fn alter_column_sql( - schema: &str, - table: &str, +pub fn add_column_sql(path: &str, column: &ColumnDefinition) -> Result, String> { + let mut statements = vec![format!( + "ALTER TABLE {path} ADD COLUMN {}", + column_sql(column, false)? + )]; + if column.is_pk { + statements.push(format!( + "ALTER TABLE {path} ADD PRIMARY KEY ({}) NOT ENFORCED", + quote_identifier(&column.name) + )); + } + Ok(statements) +} + +/// A type name split into its canonical base and its parameters, e.g. +/// `numeric(10, 2)` → `("NUMERIC", [10, 2])`. +fn parse_type(data_type: &str) -> (String, Vec) { + let trimmed = data_type.trim(); + match trimmed.find('(') { + Some(open) if trimmed.ends_with(')') => { + let parameters = trimmed[open + 1..trimmed.len() - 1] + .split(',') + .filter_map(|value| value.trim().parse::().ok()) + .collect(); + (canonical_type(&trimmed[..open]), parameters) + } + _ => (canonical_type(trimmed), Vec::new()), + } +} + +/// Precision and scale of a decimal type; unparameterized types use the +/// BigQuery maximums. +fn decimal_bounds(base: &str, parameters: &[u32]) -> (u32, u32) { + let (default_precision, default_scale) = if base == "BIGNUMERIC" { + (76, 38) + } else { + (38, 9) + }; + match parameters { + [precision, scale] => (*precision, *scale), + [precision] => (*precision, 0), + _ => (default_precision, default_scale), + } +} + +/// Column type changes BigQuery accepts in `ALTER COLUMN SET DATA TYPE`. +pub fn check_type_change(old: &str, new: &str) -> Result<(), String> { + let (old_base, old_parameters) = parse_type(old); + let (new_base, new_parameters) = parse_type(new); + let allowed = match (old_base.as_str(), new_base.as_str()) { + ("INT64", "NUMERIC" | "BIGNUMERIC") => { + let (precision, scale) = decimal_bounds(&new_base, &new_parameters); + precision.saturating_sub(scale) >= 19 + } + ("INT64", "FLOAT64") | ("NUMERIC", "FLOAT64") | ("BIGNUMERIC", "FLOAT64") => true, + ("NUMERIC", "NUMERIC" | "BIGNUMERIC") | ("BIGNUMERIC", "BIGNUMERIC") => { + let (old_precision, old_scale) = decimal_bounds(&old_base, &old_parameters); + let (new_precision, new_scale) = decimal_bounds(&new_base, &new_parameters); + new_scale >= old_scale + && new_precision.saturating_sub(new_scale) + >= old_precision.saturating_sub(old_scale) + } + ("STRING", "STRING") | ("BYTES", "BYTES") => { + match (old_parameters.first(), new_parameters.first()) { + (_, None) => true, + (None, Some(_)) => false, + (Some(old_length), Some(new_length)) => new_length >= old_length, + } + } + _ => false, + }; + if allowed { + Ok(()) + } else { + Err(format!( + "BigQuery cannot change a column from {} to {}. Supported changes: INT64 to NUMERIC, BIGNUMERIC or FLOAT64; NUMERIC to BIGNUMERIC or FLOAT64; and widening STRING, BYTES, NUMERIC or BIGNUMERIC parameters. Rewrite the table with CREATE OR REPLACE TABLE ... AS SELECT for other conversions.", + old.trim(), + new.trim() + )) + } +} + +pub fn alter_column_sql( + path: &str, old: &ColumnDefinition, new: &ColumnDefinition, ) -> Result, String> { - let target = qualified(schema, table); + if new.is_auto_increment && !old.is_auto_increment { + return Err("BigQuery has no auto-increment columns".to_string()); + } let mut statements = Vec::new(); if old.name != new.name { statements.push(format!( - "ALTER TABLE {target} RENAME COLUMN {} TO {}", + "ALTER TABLE {path} RENAME COLUMN {} TO {}", quote_identifier(&old.name), quote_identifier(&new.name) )); } let column = quote_identifier(&new.name); - if !old.data_type.eq_ignore_ascii_case(&new.data_type) { + let same_type = old + .data_type + .split_whitespace() + .collect::() + .eq_ignore_ascii_case(&new.data_type.split_whitespace().collect::()) + || parse_type(&old.data_type) == parse_type(&new.data_type); + if !same_type { + check_type_change(&old.data_type, &new.data_type)?; statements.push(format!( - "ALTER TABLE {target} ALTER COLUMN {column} SET DATA TYPE {}", + "ALTER TABLE {path} ALTER COLUMN {column} SET DATA TYPE {}", new.data_type.trim() )); } if old.is_nullable != new.is_nullable { - let action = if new.is_nullable { - "DROP NOT NULL" - } else { - "SET NOT NULL" - }; + if !new.is_nullable { + return Err(format!( + "BigQuery cannot make an existing column NOT NULL ('{}'). Only relaxing REQUIRED to NULLABLE is supported.", + new.name + )); + } statements.push(format!( - "ALTER TABLE {target} ALTER COLUMN {column} {action}" + "ALTER TABLE {path} ALTER COLUMN {column} DROP NOT NULL" )); } - if old.default_value != new.default_value { - let action = new - .default_value + let normalize = |value: &Option| { + value .as_deref() - .filter(|value| !value.trim().is_empty()) - .map(|value| format!("SET DEFAULT {}", value.trim())) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + }; + if normalize(&old.default_value) != normalize(&new.default_value) { + let action = normalize(&new.default_value) + .map(|value| format!("SET DEFAULT {value}")) .unwrap_or_else(|| "DROP DEFAULT".to_string()); + statements.push(format!("ALTER TABLE {path} ALTER COLUMN {column} {action}")); + } + // The host omits `comment` when the description is empty, so a missing + // new comment next to an old one clears it. + let description = |value: &Option| value.clone().filter(|text| !text.is_empty()); + let new_description = description(&new.comment); + if description(&old.comment) != new_description { + let value = new_description + .as_deref() + .map(string_literal) + .unwrap_or_else(|| "NULL".to_string()); statements.push(format!( - "ALTER TABLE {target} ALTER COLUMN {column} {action}" + "ALTER TABLE {path} ALTER COLUMN {column} SET OPTIONS(description = {value})" )); } + if old.is_pk != new.is_pk { + if new.is_pk { + statements.push(format!( + "ALTER TABLE {path} ADD PRIMARY KEY ({column}) NOT ENFORCED" + )); + } else { + statements.push(format!("ALTER TABLE {path} DROP PRIMARY KEY IF EXISTS")); + } + } if statements.is_empty() { Err("No column changes were detected".to_string()) } else { @@ -230,7 +574,11 @@ fn alter_column_sql( fn target(params: &Value) -> Option<(&str, &str)> { Some(( params.get("schema")?.as_str()?.trim(), - params.get("table")?.as_str()?.trim(), + params + .get("table") + .or_else(|| params.get("table_name"))? + .as_str()? + .trim(), )) .filter(|(schema, table)| !schema.is_empty() && !table.is_empty()) } @@ -250,8 +598,9 @@ fn parse_column(value: Option<&Value>) -> Option { #[cfg(test)] mod tests { - use super::{alter_column_sql, create_table_sql}; - use crate::models::ColumnDefinition; + use serde_json::json; + + use super::*; fn column(name: &str, data_type: &str, is_pk: bool) -> ColumnDefinition { ColumnDefinition { @@ -261,21 +610,199 @@ mod tests { is_pk, is_auto_increment: false, default_value: None, + comment: None, } } + const PATH: &str = "`my-proj`.`d`.`t`"; + + #[test] + fn creates_tables_with_keys_partitioning_and_options() { + let mut payload = column("payload", "JSON", false); + payload.comment = Some("Raw 'event'\nbody".into()); + let mut created = column("created_at", "TIMESTAMP", false); + created.default_value = Some("CURRENT_TIMESTAMP()".into()); + let options = TableOptions::from_params(&json!({ + "partition_by": "DATE(created_at)", + "cluster_by": "id, payload", + "description": "Events" + })); + let sql = create_table_sql( + PATH, + &[column("id", "INT64", true), created, payload], + &options, + ) + .unwrap(); + assert_eq!( + sql, + "CREATE TABLE `my-proj`.`d`.`t` (\n \ + `id` INT64 NOT NULL,\n \ + `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP(),\n \ + `payload` JSON OPTIONS(description = 'Raw \\'event\\'\\nbody'),\n \ + PRIMARY KEY (`id`) NOT ENFORCED\n)\n\ + PARTITION BY DATE(created_at)\n\ + CLUSTER BY `id`, `payload`\n\ + OPTIONS(description = 'Events')" + ); + let mut serial = column("id", "INT64", true); + serial.is_auto_increment = true; + assert!(create_table_sql(PATH, &[serial], &TableOptions::default()).is_err()); + } + #[test] - fn creates_not_enforced_primary_key() { - let sql = create_table_sql("analytics", "events", &[column("id", "INT64", true)]); - assert!(sql.contains("PRIMARY KEY (`id`) NOT ENFORCED")); + fn required_defaulted_column_puts_not_null_before_default() { + let mut status = column("status", "STRING", false); + status.is_nullable = false; + status.default_value = Some("'new'".into()); + status.comment = Some("State".into()); + let sql = create_table_sql(PATH, &[status], &TableOptions::default()).unwrap(); + assert_eq!( + sql, + "CREATE TABLE `my-proj`.`d`.`t` (\n \ + `status` STRING NOT NULL DEFAULT 'new' OPTIONS(description = 'State')\n)" + ); } #[test] - fn builds_supported_column_alterations() { + fn add_column_refuses_required_and_adds_primary_key() { + let required = column("id", "INT64", true); + assert!(add_column_sql(PATH, &required) + .unwrap_err() + .contains("NOT NULL")); + let mut key = column("id", "INT64", false); + key.is_pk = true; + assert_eq!( + add_column_sql(PATH, &key).unwrap(), + vec![ + "ALTER TABLE `my-proj`.`d`.`t` ADD COLUMN `id` INT64".to_string(), + "ALTER TABLE `my-proj`.`d`.`t` ADD PRIMARY KEY (`id`) NOT ENFORCED".to_string(), + ] + ); + } + + #[test] + fn alters_only_supported_column_changes() { let old = column("value", "INT64", false); let mut new = column("amount", "NUMERIC", false); - new.is_nullable = false; - let sql = alter_column_sql("d", "t", &old, &new).unwrap(); - assert_eq!(sql.len(), 3); + new.comment = Some("Net".into()); + assert_eq!( + alter_column_sql(PATH, &old, &new).unwrap(), + vec![ + "ALTER TABLE `my-proj`.`d`.`t` RENAME COLUMN `value` TO `amount`", + "ALTER TABLE `my-proj`.`d`.`t` ALTER COLUMN `amount` SET DATA TYPE NUMERIC", + "ALTER TABLE `my-proj`.`d`.`t` ALTER COLUMN `amount` SET OPTIONS(description = 'Net')", + ] + ); + let mut required = column("value", "INT64", false); + required.is_nullable = false; + assert_eq!( + alter_column_sql(PATH, &required, &column("value", "INT64", false)).unwrap(), + vec!["ALTER TABLE `my-proj`.`d`.`t` ALTER COLUMN `value` DROP NOT NULL"] + ); + assert!( + alter_column_sql(PATH, &column("value", "INT64", false), &required) + .unwrap_err() + .contains("NOT NULL") + ); + assert!( + alter_column_sql(PATH, &old, &column("value", "STRING", false)) + .unwrap_err() + .contains("cannot change") + ); + assert!(alter_column_sql(PATH, &old, &old.clone()).is_err()); + let mut keyed = old.clone(); + keyed.is_pk = true; + assert_eq!( + alter_column_sql(PATH, &old, &keyed).unwrap(), + vec!["ALTER TABLE `my-proj`.`d`.`t` ADD PRIMARY KEY (`value`) NOT ENFORCED"] + ); + assert_eq!( + alter_column_sql(PATH, &keyed, &old).unwrap(), + vec!["ALTER TABLE `my-proj`.`d`.`t` DROP PRIMARY KEY IF EXISTS"] + ); + let mut defaulted = old.clone(); + defaulted.default_value = Some("0".into()); + assert_eq!( + alter_column_sql(PATH, &old, &defaulted).unwrap(), + vec!["ALTER TABLE `my-proj`.`d`.`t` ALTER COLUMN `value` SET DEFAULT 0"] + ); + } + + #[test] + fn validates_type_widenings() { + assert!(check_type_change("INT64", "FLOAT64").is_ok()); + assert!(check_type_change("INTEGER", "BIGNUMERIC").is_ok()); + assert!(check_type_change("NUMERIC", "BIGNUMERIC").is_ok()); + assert!(check_type_change("NUMERIC(10, 2)", "NUMERIC(12, 4)").is_ok()); + assert!(check_type_change("NUMERIC(10, 2)", "NUMERIC(10, 4)").is_err()); + assert!(check_type_change("STRING(10)", "STRING(20)").is_ok()); + assert!(check_type_change("STRING(10)", "STRING").is_ok()); + assert!(check_type_change("STRING", "STRING(10)").is_err()); + assert!(check_type_change("FLOAT64", "INT64").is_err()); + assert!(check_type_change("BIGNUMERIC", "NUMERIC").is_err()); + assert!(check_type_change("INT64", "NUMERIC(5, 0)").is_err()); + } + + #[test] + fn foreign_keys_are_not_enforced_and_reject_actions() { + assert_eq!( + foreign_key_sql(PATH, "`my-proj`.`d`.`r`", "fk", "a, b", "x, y", &[Some("NO ACTION"), None]) + .unwrap(), + "ALTER TABLE `my-proj`.`d`.`t` ADD CONSTRAINT `fk` FOREIGN KEY (`a`, `b`) REFERENCES `my-proj`.`d`.`r`(`x`, `y`) NOT ENFORCED" + ); + assert!(foreign_key_sql(PATH, "r", "fk", "a", "x", &[Some("CASCADE")]).is_err()); + assert!(foreign_key_sql(PATH, "r", "fk", "a, b", "x", &[]).is_err()); + } + + #[test] + fn uses_project_only_when_connection_params_are_sent() { + assert_eq!(optional_data_project(&json!({"schema": "d"})), None); + assert_eq!( + optional_data_project(&json!({"params": {"database": "my-proj"}})), + Some("my-proj".to_string()) + ); + // "All databases": no browsed project, names resolve in the billing + // project like the rest of the metadata calls. + assert_eq!( + optional_data_project(&json!({ + "params": {"database": "", "extra": {"billing_project": "bill"}} + })), + Some("bill".to_string()) + ); + } + + #[test] + fn alter_column_sets_and_clears_descriptions() { + let plain = column("value", "INT64", false); + let mut described = plain.clone(); + described.comment = Some("Net amount".into()); + assert_eq!( + alter_column_sql(PATH, &plain, &described).unwrap(), + vec!["ALTER TABLE `my-proj`.`d`.`t` ALTER COLUMN `value` SET OPTIONS(description = 'Net amount')"] + ); + let cleared = + "ALTER TABLE `my-proj`.`d`.`t` ALTER COLUMN `value` SET OPTIONS(description = NULL)"; + assert_eq!( + alter_column_sql(PATH, &described, &plain).unwrap(), + vec![cleared] + ); + let mut empty = plain.clone(); + empty.comment = Some(String::new()); + assert_eq!( + alter_column_sql(PATH, &described, &empty).unwrap(), + vec![cleared] + ); + assert!(alter_column_sql(PATH, &described, &described.clone()).is_err()); + assert!(alter_column_sql(PATH, &plain, &empty).is_err()); + } + + #[test] + fn add_column_carries_the_description() { + let mut described = column("note", "STRING", false); + described.comment = Some("Free text".into()); + assert_eq!( + add_column_sql(PATH, &described).unwrap(), + vec!["ALTER TABLE `my-proj`.`d`.`t` ADD COLUMN `note` STRING OPTIONS(description = 'Free text')"] + ); } } diff --git a/src/handlers/metadata.rs b/src/handlers/metadata.rs index e080414..d9a494e 100644 --- a/src/handlers/metadata.rs +++ b/src/handlers/metadata.rs @@ -1,32 +1,115 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::HashSet; +use std::future::Future; +use std::sync::Arc; -use serde_json::{json, Value}; +use serde_json::{json, Map, Value}; -use crate::bigquery::QueryOutput; -use crate::client::{self, BigQueryClient}; +use crate::catalog::{self, FETCH_CONCURRENCY}; +use crate::client::{self, BigQueryClient, Statement}; use crate::models::{inner_params, ConnectionParams}; use crate::rpc::{error_response, ok_response}; -use crate::sql::{qualified, value_literal}; +use crate::schema_ddl::{self, TableKind}; +use crate::sql::{ + field_type_sql, fully_qualified, information_schema, string_literal, QueryParams, +}; -pub async fn get_databases(id: Value, params: &Value) -> Value { - let connection = ConnectionParams::from_value(inner_params(params)); - match connection.project_id() { - Ok(project) => ok_response(id, json!([project])), - Err(error) => error_response(id, -32602, &error), +/// Above this many objects `get_tables` skips the per-table description +/// lookup, so expanding very large datasets stays fast. +pub const COMMENT_LOOKUP_LIMIT: usize = 100; +const COMMENT_LOOKUP_CONCURRENCY: usize = 16; + +const TABLE_TYPES: &[&str] = &["TABLE", "EXTERNAL", "SNAPSHOT", "CLONE"]; + +/// A connected client plus the project the host is browsing. The client +/// always runs jobs in the billing project; metadata calls use `project`. +pub struct Target { + pub client: Arc, + pub project: String, +} + +impl Target { + pub async fn from_params(params: &Value) -> Result { + let connection = ConnectionParams::from_value(inner_params(params)); + let project = connection.data_project_id()?.to_string(); + let client = client::client_for(&connection).await?; + Ok(Self { client, project }) + } + + /// Location of a dataset, so jobs touching it run in its region even + /// when the connection location differs. + pub async fn location(&self, dataset: &str) -> Option { + catalog::dataset_location(&self.client, &self.project, dataset).await + } + + pub fn path(&self, dataset: &str, object: &str) -> String { + fully_qualified(&self.project, dataset, object) + } + + /// Runs a DDL statement and clears the metadata cache. + pub async fn run_ddl(&self, sql: &str, dataset: &str) -> Result<(), String> { + let statement = Statement::new(sql) + .in_dataset(&self.project, dataset) + .at_location(self.location(dataset).await); + let result = self.client.run_statement(&statement).await; + catalog::invalidate(); + result.map(|_| ()) } } -pub async fn get_schemas(id: Value, params: &Value) -> Value { +pub async fn get_databases(id: Value, params: &Value) -> Value { respond(id, async { let connection = ConnectionParams::from_value(inner_params(params)); + let configured = connection.data_project_id()?.to_string(); let client = client::client_for(&connection).await?; - Ok(json!(client.list_datasets().await?)) + let listed = match client.list_projects().await { + Ok(projects) => projects, + Err(error) => { + eprintln!( + "[bigquery] projects.list failed, using configured projects only: {error}" + ); + Vec::new() + } + }; + Ok(json!(merge_projects( + &configured, + &client.project_id, + listed, + connection.additional_projects() + ))) + }) + .await +} + +/// Visible projects plus the configured, billing and additional projects, +/// de-duplicated and sorted. +pub fn merge_projects( + configured: &str, + billing: &str, + listed: Vec, + additional: Vec, +) -> Vec { + let mut projects = listed; + projects.push(configured.to_string()); + projects.push(billing.to_string()); + projects.extend(additional); + let mut seen = HashSet::new(); + projects.retain(|project| !project.is_empty() && seen.insert(project.clone())); + projects.sort(); + projects +} + +pub async fn get_schemas(id: Value, params: &Value) -> Value { + respond(id, async { + let target = Target::from_params(params).await?; + Ok(json!( + target.client.list_datasets_in(&target.project).await? + )) }) .await } pub async fn get_tables(id: Value, params: &Value) -> Value { - list_objects(id, params, &["TABLE", "EXTERNAL", "SNAPSHOT", "CLONE"]).await + list_objects(id, params, TABLE_TYPES).await } pub async fn get_views(id: Value, params: &Value) -> Value { @@ -37,29 +120,59 @@ pub async fn get_materialized_views(id: Value, params: &Value) -> Value { list_objects(id, params, &["MATERIALIZED_VIEW"]).await } +fn object_names(tables: &[Value], allowed_types: &[&str]) -> Vec { + let mut names = tables + .iter() + .filter(|table| { + table + .get("type") + .and_then(Value::as_str) + .is_some_and(|kind| allowed_types.contains(&kind)) + }) + .filter_map(|table| { + table + .pointer("/tableReference/tableId") + .and_then(Value::as_str) + .map(str::to_owned) + }) + .collect::>(); + names.sort(); + names +} + async fn list_objects(id: Value, params: &Value, allowed_types: &[&str]) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; - let client = client::client_for(&connection).await?; - let mut objects = client - .list_tables(schema) - .await? - .into_iter() - .filter(|table| { - table - .get("type") - .and_then(Value::as_str) - .is_some_and(|kind| allowed_types.contains(&kind)) - }) - .filter_map(|table| { - table - .pointer("/tableReference/tableId") - .and_then(Value::as_str) - .map(|name| json!({"name": name, "schema": schema, "comment": null})) + let target = Target::from_params(params).await?; + let tables = catalog::tables(&target.client, &target.project, schema).await?; + let names = object_names(&tables, allowed_types); + // tables.list does not return descriptions; fetch them for + // reasonably sized datasets (resources are cached for later calls). + let resources = if names.len() <= COMMENT_LOOKUP_LIMIT { + catalog::tables_batch( + target.client.clone(), + &target.project, + schema, + names.clone(), + COMMENT_LOOKUP_CONCURRENCY, + ) + .await + } else { + Vec::new() + }; + let objects = names + .iter() + .map(|name| { + let comment = resources + .iter() + .find(|(resource_name, _)| resource_name == name) + .and_then(|(_, resource)| resource.as_ref().ok()) + .and_then(|resource| resource.get("description")) + .cloned() + .unwrap_or(Value::Null); + json!({"name": name, "schema": schema, "comment": comment}) }) .collect::>(); - objects.sort_by(|left, right| left["name"].as_str().cmp(&right["name"].as_str())); Ok(json!(objects)) }) .await @@ -79,43 +192,35 @@ pub async fn get_materialized_view_columns(id: Value, params: &Value) -> Value { async fn columns_for(id: Value, params: &Value, keys: &[&str]) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; let object = required_any(params, keys)?; - let client = client::client_for(&connection).await?; - Ok(json!(columns_value(&client, schema, object).await?)) + let target = Target::from_params(params).await?; + let resource = catalog::table(&target.client, &target.project, schema, object).await?; + Ok(json!(columns_from_resource(&resource)?)) }) .await } -async fn columns_value( - client: &BigQueryClient, - schema: &str, - table: &str, -) -> Result, String> { - let resource = client.get_table(schema, table).await?; - let fields = client::table_fields(&resource)?; - let primary_keys = resource - .pointer("/tableConstraints/primaryKey/columns") - .and_then(Value::as_array) - .map(|columns| { - columns - .iter() - .filter_map(Value::as_str) - .collect::>() - }) - .unwrap_or_default(); - +/// Host `TableColumn` rows from a `tables.get` resource. +pub fn columns_from_resource(resource: &Value) -> Result, String> { + let fields = client::table_fields(resource)?; + let primary_keys = schema_ddl::primary_key_columns(resource); + let parse = + |value: &Option| value.as_deref().and_then(|value| value.parse::().ok()); Ok(fields .iter() .map(|field| { json!({ "name": field.name, - "data_type": field.display_type(), + "data_type": field_type_sql(field, false), "is_nullable": field.mode.as_deref() != Some("REQUIRED"), "default_value": field.default_value_expression, - "is_pk": primary_keys.contains(field.name.as_str()), + "is_pk": primary_keys.iter().any(|column| column == &field.name), "is_auto_increment": false, + "is_generated": false, + "character_maximum_length": parse(&field.max_length), + "numeric_precision": parse(&field.precision), + "numeric_scale": parse(&field.scale), "comment": field.description }) }) @@ -124,21 +229,17 @@ async fn columns_value( pub async fn get_foreign_keys(id: Value, params: &Value) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; let table = required_any(params, &["table", "table_name"])?; - let client = client::client_for(&connection).await?; - Ok(json!(foreign_keys_value(&client, schema, table).await?)) + let target = Target::from_params(params).await?; + let resource = catalog::table(&target.client, &target.project, schema, table).await?; + Ok(json!(foreign_keys_from_resource(&resource, schema))) }) .await } -async fn foreign_keys_value( - client: &BigQueryClient, - schema: &str, - table: &str, -) -> Result, String> { - let resource = client.get_table(schema, table).await?; +/// Host `ForeignKey` rows (one per column, ordered by `seq_in_fk`). +pub fn foreign_keys_from_resource(resource: &Value, schema: &str) -> Vec { let mut output = Vec::new(); for foreign_key in resource .pointer("/tableConstraints/foreignKeys") @@ -146,266 +247,312 @@ async fn foreign_keys_value( .into_iter() .flatten() { - let name = foreign_key - .get("name") - .and_then(Value::as_str) - .unwrap_or("foreign_key"); let referenced_table = foreign_key .pointer("/referencedTable/tableId") .and_then(Value::as_str) .unwrap_or_default(); + let name = foreign_key + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .map(str::to_owned) + .unwrap_or_else(|| format!("fk_{referenced_table}")); let referenced_schema = foreign_key .pointer("/referencedTable/datasetId") .and_then(Value::as_str) .unwrap_or(schema); - for reference in foreign_key + let referenced_project = foreign_key + .pointer("/referencedTable/projectId") + .and_then(Value::as_str); + for (index, reference) in foreign_key .get("columnReferences") .and_then(Value::as_array) .into_iter() .flatten() + .enumerate() { + let column = reference + .get("referencingColumn") + .and_then(Value::as_str) + .unwrap_or_default(); + let referenced_column = reference + .get("referencedColumn") + .and_then(Value::as_str) + .unwrap_or_default(); output.push(json!({ "name": name, "constraint_name": name, - "column_name": reference.get("referencingColumn").and_then(Value::as_str).unwrap_or_default(), + "column_name": column, "ref_table": referenced_table, "referenced_table": referenced_table, - "ref_column": reference.get("referencedColumn").and_then(Value::as_str).unwrap_or_default(), - "referenced_column": reference.get("referencedColumn").and_then(Value::as_str).unwrap_or_default(), + "ref_column": referenced_column, + "referenced_column": referenced_column, "foreign_schema_name": referenced_schema, + "foreign_project": referenced_project, "on_update": "NO ACTION", - "on_delete": "NO ACTION" + "on_delete": "NO ACTION", + "seq_in_fk": index + 1 })); } } - Ok(output) + output } pub async fn get_indexes(id: Value, params: &Value) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; let table = required_any(params, &["table", "table_name"])?; - let client = client::client_for(&connection).await?; - let resource = client.get_table(schema, table).await?; - let indexes = resource - .pointer("/tableConstraints/primaryKey/columns") - .and_then(Value::as_array) - .map(|columns| { - columns - .iter() - .enumerate() - .filter_map(|(index, column)| { - column.as_str().map(|column| { - json!({ - "name": "PRIMARY", - "index_name": "PRIMARY", - "column_name": column, - "columns": [column], - "is_unique": true, - "is_primary": true, - "seq_in_index": index + 1, - "is_expression": false - }) - }) - }) - .collect::>() + let target = Target::from_params(params).await?; + let resource = catalog::table(&target.client, &target.project, schema, table).await?; + Ok(json!(indexes_from_resource(&resource))) + }) + .await +} + +/// The informational primary key, reported as the `PRIMARY` index. +pub fn indexes_from_resource(resource: &Value) -> Vec { + schema_ddl::primary_key_columns(resource) + .iter() + .enumerate() + .map(|(index, column)| { + json!({ + "name": "PRIMARY", + "index_name": "PRIMARY", + "column_name": column, + "columns": [column], + "is_unique": true, + "is_primary": true, + "seq_in_index": index + 1, + "is_expression": false }) - .unwrap_or_default(); - Ok(json!(indexes)) + }) + .collect() +} + +/// `INFORMATION_SCHEMA.TABLES.ddl` for one object, run in the dataset +/// location with the object name as a query parameter. +async fn information_schema_ddl( + target: &Target, + schema: &str, + name: &str, +) -> Result { + let mut parameters = QueryParams::new(); + let name_parameter = parameters.string(name); + let sql = format!( + "SELECT ddl FROM {} WHERE table_name = {name_parameter}", + information_schema(&target.project, schema, "TABLES") + ); + let location = target.location(schema).await; + let statement = Statement::new(sql) + .with_parameters(parameters.into_vec()) + .at_location(location) + .returning(1); + let output = target.client.run_statement(&statement).await?; + output + .rows + .first() + .and_then(|row| row.first()) + .and_then(Value::as_str) + .filter(|ddl| !ddl.trim().is_empty()) + .map(str::to_owned) + .ok_or_else(|| format!("INFORMATION_SCHEMA returned no DDL for '{schema}.{name}'")) +} + +/// DDL of a table-like object: `INFORMATION_SCHEMA.TABLES.ddl` when it can be +/// read, otherwise rebuilt from the REST resource. +async fn object_ddl(target: &Target, schema: &str, name: &str) -> Result { + let resource = catalog::table(&target.client, &target.project, schema, name).await?; + if schema_ddl::is_legacy_sql_view(&resource) { + return schema_ddl::table_ddl(&resource, &target.project, schema, name); + } + match information_schema_ddl(target, schema, name).await { + Ok(ddl) => Ok(ddl), + Err(error) => { + eprintln!( + "[bigquery] INFORMATION_SCHEMA DDL unavailable, rebuilding from metadata: {error}" + ); + schema_ddl::table_ddl(&resource, &target.project, schema, name) + } + } +} + +pub async fn get_table_ddl(id: Value, params: &Value) -> Value { + respond(id, async { + let schema = required_schema(params)?; + let table = required_any(params, &["table", "table_name"])?; + let target = Target::from_params(params).await?; + Ok(json!(object_ddl(&target, schema, table).await?)) }) .await } pub async fn get_view_definition(id: Value, params: &Value) -> Value { - definition(id, params, false).await + definition(id, params).await } pub async fn get_materialized_view_definition(id: Value, params: &Value) -> Value { - definition(id, params, true).await + definition(id, params).await } -async fn definition(id: Value, params: &Value, materialized: bool) -> Value { +async fn definition(id: Value, params: &Value) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; let name = required_any(params, &["view_name", "view"])?; - let client = client::client_for(&connection).await?; - let resource = client.get_table(schema, name).await?; - let pointer = if materialized { - "/materializedView/query" - } else { - "/view/query" - }; - let query = resource - .pointer(pointer) - .and_then(Value::as_str) - .ok_or_else(|| format!("Definition is unavailable for '{schema}.{name}'"))?; - let kind = if materialized { - "MATERIALIZED VIEW" - } else { - "VIEW" - }; - Ok(json!(format!( - "CREATE {kind} {} AS\n{query}", - qualified(schema, name) - ))) + let target = Target::from_params(params).await?; + Ok(json!(object_ddl(&target, schema, name).await?)) }) .await } pub async fn create_view(id: Value, params: &Value) -> Value { - mutate_view(id, params, false).await -} - -pub async fn alter_view(id: Value, params: &Value) -> Value { - mutate_view(id, params, true).await -} - -async fn mutate_view(id: Value, params: &Value, replace: bool) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; let name = required_any(params, &["view_name", "name"])?; let definition = required_any(params, &["definition"])?; - let prefix = if replace { - "CREATE OR REPLACE VIEW" + let target = Target::from_params(params).await?; + let sql = if starts_with_create(definition) { + definition.to_string() } else { - "CREATE VIEW" + format!("CREATE VIEW {} AS\n{definition}", target.path(schema, name)) }; - let sql = format!("{prefix} {} AS\n{definition}", qualified(schema, name)); - let client = client::client_for(&connection).await?; - client.execute_query(&sql, Some(schema), Some(0), 1).await?; + target.run_ddl(&sql, schema).await?; Ok(Value::Null) }) .await } -pub async fn drop_view(id: Value, params: &Value) -> Value { +pub async fn alter_view(id: Value, params: &Value) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; let name = required_any(params, &["view_name", "name"])?; - let client = client::client_for(&connection).await?; - client - .execute_query( - &format!("DROP VIEW IF EXISTS {}", qualified(schema, name)), - Some(schema), - Some(0), - 1, - ) - .await?; + let definition = required_any(params, &["definition"])?; + let target = Target::from_params(params).await?; + let resource = catalog::table(&target.client, &target.project, schema, name).await?; + let sql = replace_view_sql(&resource, &target.path(schema, name), definition)?; + target.run_ddl(&sql, schema).await?; Ok(Value::Null) }) .await } -pub async fn refresh_materialized_view(id: Value, _params: &Value) -> Value { - error_response( - id, - -32601, - "BigQuery materialized views refresh automatically and do not support manual refresh", - ) +/// `CREATE OR REPLACE VIEW` that keeps the view description, friendly name +/// and labels. Legacy SQL views are refused. +pub fn replace_view_sql(resource: &Value, path: &str, definition: &str) -> Result { + match TableKind::from_resource(resource) { + TableKind::View => {} + TableKind::MaterializedView => { + return Err( + "Materialized views cannot be replaced from the view editor; run CREATE OR REPLACE MATERIALIZED VIEW in a query tab." + .to_string(), + ) + } + _ => return Err("Only views can be altered with the view editor".to_string()), + } + if schema_ddl::is_legacy_sql_view(resource) { + return Err( + "This view uses legacy SQL. Tabularis edits GoogleSQL views only; recreate it in GoogleSQL with CREATE OR REPLACE VIEW in a query tab." + .to_string(), + ); + } + if starts_with_create(definition) { + return Ok(definition.to_string()); + } + schema_ddl::view_ddl(resource, path, Some(definition), true) } -pub async fn get_routines(id: Value, params: &Value) -> Value { - respond(id, async { - let (client, schema) = client_and_schema(params).await?; - let project = &client.project_id; - let sql = format!( - "SELECT routine_name, routine_type FROM {} ORDER BY routine_name", - crate::sql::quote_identifier(&format!( - "{project}.{schema}.INFORMATION_SCHEMA.ROUTINES" - )) - ); - let output = client.execute_query(&sql, Some(schema), None, 1).await?; - Ok(json!(rows_as_objects(&output) - .into_iter() - .map(|row| json!({ - "name": row.get("routine_name").cloned().unwrap_or(Value::Null), - "routine_type": row.get("routine_type").cloned().unwrap_or_else(|| json!("FUNCTION")), - "schema": schema - })) - .collect::>())) - }) - .await +fn starts_with_create(sql: &str) -> bool { + sql.trim_start() + .get(..6) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("CREATE")) } -pub async fn get_routine_parameters(id: Value, params: &Value) -> Value { +pub async fn drop_view(id: Value, params: &Value) -> Value { respond(id, async { - let (client, schema) = client_and_schema(params).await?; - let routine = required_any(params, &["routine_name", "routine"])?; - let project = &client.project_id; - let routine_literal = value_literal(&json!(routine), None)?; - let sql = format!( - "SELECT parameter_name, data_type, parameter_mode FROM {} WHERE specific_name = {routine_literal} ORDER BY ordinal_position", - crate::sql::quote_identifier(&format!( - "{project}.{schema}.INFORMATION_SCHEMA.PARAMETERS" - )) - ); - let output = client.execute_query(&sql, Some(schema), None, 1).await?; - Ok(json!(rows_as_objects(&output) - .into_iter() - .map(|row| json!({ - "name": row.get("parameter_name").cloned().unwrap_or(Value::Null), - "data_type": row.get("data_type").cloned().unwrap_or(Value::Null), - "mode": row.get("parameter_mode").cloned().unwrap_or_else(|| json!("IN")) - })) - .collect::>())) + let schema = required_schema(params)?; + let name = required_any(params, &["view_name", "name"])?; + let target = Target::from_params(params).await?; + let kind = match catalog::table(&target.client, &target.project, schema, name).await { + Ok(resource) => TableKind::from_resource(&resource), + Err(_) => TableKind::View, + }; + let sql = schema_ddl::drop_sql(kind, &target.path(schema, name)); + target.run_ddl(&sql, schema).await?; + Ok(Value::Null) }) .await } -pub async fn get_routine_definition(id: Value, params: &Value) -> Value { +/// `CALL BQ.REFRESH_MATERIALIZED_VIEW('project.dataset.view')`. +pub fn refresh_materialized_view_sql(project: &str, dataset: &str, view: &str) -> String { + format!( + "CALL BQ.REFRESH_MATERIALIZED_VIEW({})", + string_literal(&format!("{project}.{dataset}.{view}")) + ) +} + +pub async fn refresh_materialized_view(id: Value, params: &Value) -> Value { respond(id, async { - let (client, schema) = client_and_schema(params).await?; - let routine = required_any(params, &["routine_name", "routine"])?; - let project = &client.project_id; - let routine_literal = value_literal(&json!(routine), None)?; - let sql = format!( - "SELECT ddl FROM {} WHERE routine_name = {routine_literal} LIMIT 1", - crate::sql::quote_identifier(&format!( - "{project}.{schema}.INFORMATION_SCHEMA.ROUTINES" - )) - ); - let output = client.execute_query(&sql, Some(schema), Some(1), 1).await?; - output - .rows - .first() - .and_then(|row| row.first()) - .cloned() - .ok_or_else(|| format!("Routine '{routine}' was not found")) + let schema = required_schema(params)?; + let name = required_any(params, &["view_name", "view", "name"])?; + let target = Target::from_params(params).await?; + let sql = refresh_materialized_view_sql(&target.project, schema, name); + let location = target.location(schema).await; + target + .client + .run_statement(&Statement::new(sql).at_location(location)) + .await?; + catalog::invalidate(); + Ok(Value::Null) }) .await } +/// Lists table-like objects and fetches each resource once. +async fn table_resources( + target: &Target, + schema: &str, + allowed_types: &[&str], + limit: Option, +) -> Result<(Vec<(String, Value)>, usize), String> { + let tables = catalog::tables(&target.client, &target.project, schema).await?; + let names = object_names(&tables, allowed_types); + let total = names.len(); + let names = match limit { + Some(limit) => names.into_iter().take(limit).collect(), + None => names, + }; + let fetched = catalog::tables_batch( + target.client.clone(), + &target.project, + schema, + names, + FETCH_CONCURRENCY, + ) + .await; + let mut resources = Vec::with_capacity(fetched.len()); + for (name, resource) in fetched { + resources.push((name, resource?)); + } + Ok((resources, total)) +} + +fn table_schema(name: &str, resource: &Value, schema: &str) -> Result { + Ok(json!({ + "name": name, + "columns": columns_from_resource(resource)?, + "foreign_keys": foreign_keys_from_resource(resource, schema) + })) +} + pub async fn get_schema_snapshot(id: Value, params: &Value) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; - let client = client::client_for(&connection).await?; - let tables = client.list_tables(schema).await?; - let mut snapshot = Vec::new(); - for table in tables { - let Some(name) = table - .pointer("/tableReference/tableId") - .and_then(Value::as_str) - else { - continue; - }; - let kind = table - .get("type") - .and_then(Value::as_str) - .unwrap_or_default(); - if !matches!(kind, "TABLE" | "EXTERNAL" | "SNAPSHOT" | "CLONE") { - continue; - } - snapshot.push(json!({ - "name": name, - "columns": columns_value(&client, schema, name).await?, - "foreign_keys": foreign_keys_value(&client, schema, name).await? - })); - } + let target = Target::from_params(params).await?; + let (resources, _) = table_resources(&target, schema, TABLE_TYPES, None).await?; + let snapshot = resources + .iter() + .map(|(name, resource)| table_schema(name, resource, schema)) + .collect::, _>>()?; Ok(json!(snapshot)) }) .await @@ -421,22 +568,54 @@ pub async fn get_all_foreign_keys_batch(id: Value, params: &Value) -> Value { async fn batch_metadata(id: Value, params: &Value, columns: bool) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; - let tables = params + let target = Target::from_params(params).await?; + let requested = params .get("tables") .and_then(Value::as_array) - .map(|tables| tables.iter().filter_map(Value::as_str).collect::>()) - .unwrap_or_default(); - let client = client::client_for(&connection).await?; - let mut result = serde_json::Map::new(); - for table in tables { - let value = if columns { - json!(columns_value(&client, schema, table).await?) - } else { - json!(foreign_keys_value(&client, schema, table).await?) - }; - result.insert(table.to_string(), value); + .map(|tables| { + tables + .iter() + .filter_map(Value::as_str) + .collect::>() + }); + // Only fetch the tables the host asked for. + let tables = catalog::tables(&target.client, &target.project, schema).await?; + let names: Vec = object_names(&tables, TABLE_TYPES) + .into_iter() + .filter(|name| { + requested + .as_ref() + .is_none_or(|requested| requested.contains(name.as_str())) + }) + .collect(); + let fetched = catalog::tables_batch( + target.client.clone(), + &target.project, + schema, + names, + FETCH_CONCURRENCY, + ) + .await; + // One unreadable or concurrently deleted table must not fail the + // whole batch: skip it and keep the others. + let mut result = Map::new(); + for (name, resource) in fetched { + let value = resource.and_then(|resource| { + if columns { + columns_from_resource(&resource).map(|columns| json!(columns)) + } else { + Ok(json!(foreign_keys_from_resource(&resource, schema))) + } + }); + match value { + Ok(value) => { + result.insert(name, value); + } + Err(error) => { + eprintln!("bigquery-plugin: skipping metadata of {schema}.{name}: {error}") + } + } } Ok(Value::Object(result)) }) @@ -445,64 +624,24 @@ async fn batch_metadata(id: Value, params: &Value, columns: bool) -> Value { pub async fn get_ai_schema_context(id: Value, params: &Value) -> Value { respond(id, async { - let connection = ConnectionParams::from_value(inner_params(params)); let schema = required_schema(params)?; let max_tables = params .get("max_tables") .and_then(Value::as_u64) .unwrap_or(20) as usize; - let client = client::client_for(&connection).await?; - let tables = client - .list_tables(schema) - .await? - .into_iter() - .filter(|table| table.get("type").and_then(Value::as_str) == Some("TABLE")) - .filter_map(|table| { - table - .pointer("/tableReference/tableId") - .and_then(Value::as_str) - .map(str::to_owned) - }) - .collect::>(); - let total = tables.len(); - let mut output = Vec::new(); - for table in tables.into_iter().take(max_tables) { - output.push(json!({ - "name": table, - "columns": columns_value(&client, schema, &table).await?, - "foreign_keys": foreign_keys_value(&client, schema, &table).await? - })); - } - Ok(json!({"tables": output, "total_table_count": total})) + let target = Target::from_params(params).await?; + let (resources, total) = + table_resources(&target, schema, TABLE_TYPES, Some(max_tables)).await?; + let tables = resources + .iter() + .map(|(name, resource)| table_schema(name, resource, schema)) + .collect::, _>>()?; + Ok(json!({"tables": tables, "total_table_count": total})) }) .await } -fn rows_as_objects(output: &QueryOutput) -> Vec> { - output - .rows - .iter() - .map(|row| { - output - .columns - .iter() - .cloned() - .zip(row.iter().cloned()) - .collect() - }) - .collect() -} - -async fn client_and_schema( - params: &Value, -) -> Result<(std::sync::Arc, &str), String> { - let connection = ConnectionParams::from_value(inner_params(params)); - let schema = required_schema(params)?; - let client = client::client_for(&connection).await?; - Ok((client, schema)) -} - -fn required_schema(params: &Value) -> Result<&str, String> { +pub fn required_schema(params: &Value) -> Result<&str, String> { params .get("schema") .and_then(Value::as_str) @@ -511,7 +650,7 @@ fn required_schema(params: &Value) -> Result<&str, String> { .ok_or_else(|| "A BigQuery dataset is required as 'schema'".to_string()) } -fn required_any<'a>(params: &'a Value, keys: &[&str]) -> Result<&'a str, String> { +pub fn required_any<'a>(params: &'a Value, keys: &[&str]) -> Result<&'a str, String> { keys.iter() .find_map(|key| params.get(*key).and_then(Value::as_str)) .map(str::trim) @@ -524,12 +663,89 @@ fn required_any<'a>(params: &'a Value, keys: &[&str]) -> Result<&'a str, String> }) } -async fn respond(id: Value, future: F) -> Value +pub async fn respond(id: Value, future: F) -> Value where - F: std::future::Future>, + F: Future>, { match future.await { Ok(value) => ok_response(id, value), Err(error) => error_response(id, -32603, &error), } } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn merges_projects_without_duplicates() { + assert_eq!( + merge_projects( + "my-proj", + "billing-proj", + vec!["z-proj".into(), "my-proj".into()], + vec!["bigquery-public-data".into()] + ), + vec!["bigquery-public-data", "billing-proj", "my-proj", "z-proj"] + ); + assert_eq!(merge_projects("p", "p", vec![], vec![]), vec!["p"]); + } + + #[test] + fn object_names_filter_and_sort_types() { + let tables = vec![ + json!({"type": "VIEW", "tableReference": {"tableId": "v"}}), + json!({"type": "TABLE", "tableReference": {"tableId": "b"}}), + json!({"type": "SNAPSHOT", "tableReference": {"tableId": "a"}}), + json!({"type": "EXTERNAL", "tableReference": {"tableId": "c"}}), + ]; + assert_eq!(object_names(&tables, TABLE_TYPES), vec!["a", "b", "c"]); + assert_eq!(object_names(&tables, &["VIEW"]), vec!["v"]); + } + + #[test] + fn refresh_uses_a_string_path() { + assert_eq!( + refresh_materialized_view_sql("my-proj", "d", "mv'1"), + "CALL BQ.REFRESH_MATERIALIZED_VIEW('my-proj.d.mv\\'1')" + ); + } + + #[test] + fn replacing_views_keeps_options_and_refuses_legacy_sql() { + let view = json!({"type": "VIEW", "description": "Daily", "labels": {"a": "b"}, "view": {"query": "SELECT 1"}}); + assert_eq!( + replace_view_sql(&view, "`p-1`.`d`.`v`", "SELECT 2").unwrap(), + "CREATE OR REPLACE VIEW `p-1`.`d`.`v` OPTIONS(description = 'Daily', labels = [('a', 'b')])\nAS\nSELECT 2" + ); + let legacy = json!({"type": "VIEW", "view": {"query": "SELECT 1", "useLegacySql": true}}); + assert!(replace_view_sql(&legacy, "`p`.`d`.`v`", "SELECT 2") + .unwrap_err() + .contains("legacy SQL")); + let mv = json!({"type": "MATERIALIZED_VIEW"}); + assert!(replace_view_sql(&mv, "`p`.`d`.`v`", "SELECT 2").is_err()); + } + + #[test] + fn columns_include_precision_and_descriptions() { + let resource = json!({ + "schema": {"fields": [ + {"name": "id", "type": "INTEGER", "mode": "REQUIRED"}, + {"name": "price", "type": "NUMERIC", "precision": "10", "scale": "2", "description": "Net"}, + {"name": "code", "type": "STRING", "maxLength": "8"} + ]}, + "tableConstraints": {"primaryKey": {"columns": ["id"]}} + }); + let columns = columns_from_resource(&resource).unwrap(); + assert_eq!(columns[0]["data_type"], "INT64"); + assert_eq!(columns[0]["is_pk"], true); + assert_eq!(columns[0]["is_nullable"], false); + assert_eq!(columns[1]["data_type"], "NUMERIC(10, 2)"); + assert_eq!(columns[1]["numeric_precision"], 10); + assert_eq!(columns[1]["numeric_scale"], 2); + assert_eq!(columns[1]["comment"], "Net"); + assert_eq!(columns[2]["character_maximum_length"], 8); + } +} diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index 54c80e1..ea790bc 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -4,3 +4,4 @@ pub mod crud; pub mod ddl; pub mod metadata; pub mod query; +pub mod routines; diff --git a/src/handlers/query.rs b/src/handlers/query.rs index 83e4c90..761aa27 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -1,43 +1,160 @@ +use std::sync::Arc; use std::time::Instant; use serde_json::{json, Value}; -use crate::client; +use crate::bigquery::QueryOutput; +use crate::catalog; +use crate::client::{self, BigQueryClient, QueryRequest}; +use crate::explain; use crate::models::{inner_params, ConnectionParams}; use crate::rpc::{error_response, ok_response}; +use crate::session::{self, SessionState}; +use crate::statements::{self, TransactionControl}; + +const DEFAULT_PAGE_SIZE: u32 = 10_000; + +struct Paging<'a> { + schema: Option<&'a str>, + page_size: u32, + page: u32, +} + +impl<'a> Paging<'a> { + fn from_params(params: &'a Value) -> Self { + let page_size = params + .get("limit") + .or_else(|| params.get("page_size")) + .and_then(Value::as_u64) + .and_then(|value| u32::try_from(value).ok()) + .filter(|value| *value > 0) + .unwrap_or(DEFAULT_PAGE_SIZE); + let page = params + .get("page") + .and_then(Value::as_u64) + .and_then(|value| u32::try_from(value).ok()) + .unwrap_or(1) + .max(1); + Self { + schema: params + .get("schema") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()), + page_size, + page, + } + } + + fn request<'s>(&self, sql: &'s str) -> QueryRequest<'s> + where + 'a: 's, + { + QueryRequest { + sql, + dataset: self.schema, + page_size: self.page_size, + page: self.page, + fetch_rows: true, + ..QueryRequest::default() + } + } + + fn render(&self, output: &QueryOutput, started: Instant) -> Value { + let mut value = output.to_tabularis(self.page, self.page_size); + value["execution_time_ms"] = json!(started.elapsed().as_secs_f64() * 1000.0); + value + } +} + +fn session_id(params: &Value) -> Option<&str> { + params + .get("session_id") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) +} + +/// Run one statement, inside a session when one is given, and keep the +/// session's transaction flag in step with BEGIN / COMMIT / ROLLBACK. Any +/// statement that is not a plain query may have changed the catalog, so the +/// metadata cache is cleared once it has run. +async fn run_statement( + client: &Arc, + state: Option<&mut SessionState>, + request: &QueryRequest<'_>, +) -> Result { + let result = run_in_session(client, state, request).await; + if !statements::is_read_only_query(request.sql) { + // Also after a failure: a script may have applied part of its DDL. + catalog::invalidate(); + } + result +} + +async fn run_in_session( + client: &Arc, + state: Option<&mut SessionState>, + request: &QueryRequest<'_>, +) -> Result { + if statements::split(request.sql).len() == 1 + && statements::transaction_control(request.sql) == TransactionControl::Savepoint + { + return Err( + "BigQuery does not support savepoints; the statement was not sent to BigQuery" + .to_string(), + ); + } + let Some(state) = state else { + return client.run_query(request, None).await; + }; + state.client = Some(client.clone()); + let result = client.run_query(request, Some(&mut *state)).await; + match &result { + Ok(_) => { + state.in_transaction = + statements::transaction_state_after(request.sql, state.in_transaction) + } + Err(_) if statements::ends_transaction(request.sql) => state.in_transaction = false, + Err(_) => {} + } + result +} pub async fn execute_query(id: Value, params: &Value) -> Value { let connection = ConnectionParams::from_value(inner_params(params)); let Some(query) = params.get("query").and_then(Value::as_str) else { return error_response(id, -32602, "Invalid params: missing 'query'"); }; - let limit = params - .get("limit") - .or_else(|| params.get("page_size")) - .and_then(Value::as_u64) - .and_then(|value| u32::try_from(value).ok()); - let page = params - .get("page") - .and_then(Value::as_u64) - .and_then(|value| u32::try_from(value).ok()) - .unwrap_or(1) - .max(1); - let schema = params.get("schema").and_then(Value::as_str); + let paging = Paging::from_params(params); let started = Instant::now(); + let client = match client::client_for(&connection).await { + Ok(client) => client, + Err(error) => return error_response(id, -32603, &error), + }; + let request = paging.request(query); - let result = async { - let client = client::client_for(&connection).await?; - client.execute_query(query, schema, limit, page).await - } - .await; + let Some(host_session) = session_id(params) else { + return match run_statement(&client, None, &request).await { + Ok(output) => ok_response(id, paging.render(&output, started)), + Err(error) => error_response(id, -32603, &error), + }; + }; + // Session-aware hosts accept `{result, in_transaction}` and report a + // failed statement as `{error, in_transaction}` so the tab's transaction + // badge stays correct. + let shared = session::get_or_create(host_session).await; + let mut state = shared.lock().await; + let result = run_statement(&client, Some(&mut state), &request).await; + let in_transaction = state.in_transaction; match result { - Ok(output) => { - let mut value = output.to_tabularis(page, limit.unwrap_or(10_000)); - value["execution_time_ms"] = json!(started.elapsed().as_secs_f64() * 1000.0); - ok_response(id, value) - } - Err(error) => error_response(id, -32603, &error), + Ok(output) => ok_response( + id, + json!({"result": paging.render(&output, started), "in_transaction": in_transaction}), + ), + Err(error) => ok_response( + id, + json!({"error": error, "in_transaction": in_transaction}), + ), } } @@ -48,28 +165,51 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { .and_then(Value::as_array) .map(|values| values.iter().filter_map(Value::as_str).collect::>()) .unwrap_or_default(); - let limit = params - .get("limit") - .and_then(Value::as_u64) - .and_then(|value| u32::try_from(value).ok()); - let page = params - .get("page") - .and_then(Value::as_u64) - .and_then(|value| u32::try_from(value).ok()) - .unwrap_or(1) - .max(1); - let schema = params.get("schema").and_then(Value::as_str); + let paging = Paging::from_params(params); let client = match client::client_for(&connection).await { Ok(client) => client, Err(error) => return error_response(id, -32603, &error), }; + if let Some(host_session) = session_id(params) { + let shared = session::get_or_create(host_session).await; + let mut state = shared.lock().await; + let results = run_batch(&client, &paging, &queries, Some(&mut state)).await; + let in_transaction = state.in_transaction; + return ok_response( + id, + json!({"results": results, "in_transaction": in_transaction}), + ); + } + + // Without a host session, statements that leave state behind (DECLARE, + // SET, temporary tables, transactions) still need to see each other, so + // the batch runs in a short-lived BigQuery session. + if queries.len() > 1 && queries.iter().any(|query| statements::needs_session(query)) { + let mut state = SessionState::default(); + let results = run_batch(&client, &paging, &queries, Some(&mut state)).await; + if let Err(error) = session::abort(&mut state).await { + eprintln!("bigquery-plugin: unable to abort batch session: {error}"); + } + return ok_response(id, json!(results)); + } + let results = run_batch(&client, &paging, &queries, None).await; + ok_response(id, json!(results)) +} + +async fn run_batch( + client: &Arc, + paging: &Paging<'_>, + queries: &[&str], + mut state: Option<&mut SessionState>, +) -> Vec { let mut results = Vec::with_capacity(queries.len()); for query in queries { let started = Instant::now(); - match client.execute_query(query, schema, limit, page).await { + let request = paging.request(query); + match run_statement(client, state.as_deref_mut(), &request).await { Ok(output) => results.push(json!({ - "result": output.to_tabularis(page, limit.unwrap_or(10_000)), + "result": paging.render(&output, started), "error": null, "execution_time_ms": started.elapsed().as_secs_f64() * 1000.0 })), @@ -80,7 +220,16 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { })), } } - ok_response(id, json!(results)) + results +} + +pub async fn release_session(id: Value, params: &Value) -> Value { + if let Some(host_session) = session_id(params) { + if let Err(error) = session::release(host_session).await { + eprintln!("bigquery-plugin: unable to abort session: {error}"); + } + } + ok_response(id, Value::Null) } pub async fn explain_query(id: Value, params: &Value) -> Value { @@ -88,54 +237,27 @@ pub async fn explain_query(id: Value, params: &Value) -> Value { let Some(query) = params.get("query").and_then(Value::as_str) else { return error_response(id, -32602, "Invalid params: missing 'query'"); }; - let schema = params.get("schema").and_then(Value::as_str); + let analyze = params + .get("analyze") + .and_then(Value::as_bool) + .unwrap_or(false); + let schema = params + .get("schema") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()); let result = async { let client = client::client_for(&connection).await?; - client.dry_run(query, schema).await + if analyze { + let job = client.run_for_statistics(query, schema, None).await?; + Ok::<_, String>(explain::payload("analyze", &job)) + } else { + let job = client.dry_run_job(query, schema, None).await?; + Ok(explain::payload("estimate", &job)) + } } .await; match result { - Ok(output) => { - let total_bytes = output - .get("total_bytes_processed") - .and_then(Value::as_str) - .and_then(|value| value.parse::().ok()); - let raw_output = serde_json::to_string_pretty(&output).ok(); - ok_response( - id, - json!({ - "root": { - "id": "bigquery-dry-run", - "node_type": "BigQuery dry run", - "relation": null, - "startup_cost": null, - "total_cost": null, - "plan_rows": null, - "actual_rows": null, - "actual_time_ms": null, - "actual_loops": null, - "buffers_hit": null, - "buffers_read": null, - "filter": null, - "index_condition": null, - "join_type": null, - "hash_condition": null, - "extra": { - "total_bytes_processed": output.get("total_bytes_processed").cloned().unwrap_or(Value::Null), - "estimated_gibibytes_processed": total_bytes.map(|bytes| bytes as f64 / 1_073_741_824.0), - "job_id": output.get("job_id").cloned().unwrap_or(Value::Null) - }, - "children": [] - }, - "planning_time_ms": null, - "execution_time_ms": null, - "original_query": query, - "driver": "bigquery", - "has_analyze_data": false, - "raw_output": raw_output - }), - ) - } + Ok(payload) => ok_response(id, explain::raw_output(&payload, query)), Err(error) => error_response(id, -32603, &error), } } diff --git a/src/handlers/routines.rs b/src/handlers/routines.rs new file mode 100644 index 0000000..d4e70e8 --- /dev/null +++ b/src/handlers/routines.rs @@ -0,0 +1,115 @@ +//! Routine metadata and management backed by the free `routines.list` and +//! `routines.get` REST calls. + +use serde_json::{json, Value}; + +use crate::catalog; +use crate::handlers::metadata::{required_any, required_schema, respond, Target}; +use crate::models::optional_data_project; +use crate::routines::{self, CallArg, RoutineKind}; + +pub async fn get_routines(id: Value, params: &Value) -> Value { + respond(id, async { + let schema = required_schema(params)?; + let target = Target::from_params(params).await?; + let mut routines = catalog::routines(&target.client, &target.project, schema) + .await? + .iter() + .filter_map(routines::routine_info) + .collect::>(); + routines.sort_by(|left, right| left["name"].as_str().cmp(&right["name"].as_str())); + Ok(json!(routines)) + }) + .await +} + +async fn routine_resource(params: &Value) -> Result<(Target, String, String, Value), String> { + let schema = required_schema(params)?.to_string(); + let name = required_any(params, &["routine_name", "routine"])?.to_string(); + let target = Target::from_params(params).await?; + let resource = catalog::routine(&target.client, &target.project, &schema, &name).await?; + Ok((target, schema, name, resource)) +} + +pub async fn get_routine_parameters(id: Value, params: &Value) -> Value { + respond(id, async { + let (_, _, _, resource) = routine_resource(params).await?; + Ok(json!(routines::routine_parameters(&resource))) + }) + .await +} + +pub async fn get_routine_definition(id: Value, params: &Value) -> Value { + respond(id, async { + let (target, schema, name, resource) = routine_resource(params).await?; + Ok(json!(routines::routine_ddl( + &resource, + &target.path(&schema, &name), + true + )?)) + }) + .await +} + +pub async fn get_routine_edit_script(id: Value, params: &Value) -> Value { + get_routine_definition(id, params).await +} + +pub async fn build_routine_call_sql(id: Value, params: &Value) -> Value { + respond(id, async { + let args: Vec = match params.get("args") { + Some(Value::Null) | None => Vec::new(), + Some(args) => serde_json::from_value(args.clone()) + .map_err(|error| format!("Invalid routine arguments: {error}"))?, + }; + let (target, schema, name, resource) = routine_resource(params).await?; + Ok(json!(routines::routine_call_sql( + &resource, + &target.path(&schema, &name), + &args + )?)) + }) + .await +} + +pub async fn routine_create_template(id: Value, params: &Value) -> Value { + respond(id, async { + let routine_type = params + .get("routine_type") + .and_then(Value::as_str) + .unwrap_or("FUNCTION"); + let schema = params.get("schema").and_then(Value::as_str); + // Older hosts send no `params`; the template is then + // dataset-qualified. + let project = optional_data_project(params); + Ok(json!(routines::create_template( + routine_type, + project.as_deref(), + schema + ))) + }) + .await +} + +pub async fn drop_routine(id: Value, params: &Value) -> Value { + respond(id, async { + let schema = required_schema(params)?; + let name = required_any(params, &["routine_name", "routine"])?; + let target = Target::from_params(params).await?; + // The host only knows FUNCTION / PROCEDURE; table functions need + // their own keyword, so the stored routine type wins. + let kind = match catalog::routine(&target.client, &target.project, schema, name).await { + Ok(resource) => RoutineKind::from_resource(&resource), + Err(_) => RoutineKind::from_name( + params + .get("routine_type") + .and_then(Value::as_str) + .unwrap_or("FUNCTION"), + ), + }; + let sql = routines::drop_routine_sql(kind, &target.path(schema, name)); + target.run_ddl(&sql, schema).await?; + Ok(Value::Null) + }) + .await +} diff --git a/src/lib.rs b/src/lib.rs index 4b8c487..e05f262 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,18 @@ pub mod auth; pub mod bigquery; +pub mod browse; +pub mod cancel; +pub mod catalog; pub mod client; +pub mod endpoint; +pub mod errors; +pub mod explain; pub mod handlers; pub mod models; +pub mod routines; pub mod rpc; +pub mod schema_ddl; +pub mod server; +pub mod session; pub mod sql; +pub mod statements; diff --git a/src/main.rs b/src/main.rs index c1b0a3e..5b30441 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,28 +1,6 @@ -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::io::BufReader; #[tokio::main] async fn main() { - let stdin = tokio::io::stdin(); - let mut lines = BufReader::new(stdin).lines(); - let mut stdout = tokio::io::stdout(); - - while let Ok(Some(line)) = lines.next_line().await { - if line.trim().is_empty() { - continue; - } - - let response = bigquery_plugin::rpc::handle_line(&line).await; - let mut encoded = match serde_json::to_vec(&response) { - Ok(value) => value, - Err(error) => { - eprintln!("failed to serialize JSON-RPC response: {error}"); - continue; - } - }; - encoded.push(b'\n'); - - if stdout.write_all(&encoded).await.is_err() || stdout.flush().await.is_err() { - break; - } - } + bigquery_plugin::server::serve(BufReader::new(tokio::io::stdin()), tokio::io::stdout()).await; } diff --git a/src/models.rs b/src/models.rs index 8fed910..8c5a2e9 100644 --- a/src/models.rs +++ b/src/models.rs @@ -3,6 +3,10 @@ use std::collections::HashMap; use serde::Deserialize; use serde_json::Value; +/// Stays below the host's default 120 s plugin call timeout, so a timed-out +/// query is reported and cancelled by the plugin rather than abandoned. +pub const DEFAULT_TIMEOUT_SECONDS: u64 = 110; + #[derive(Debug, Clone, Default)] pub struct ConnectionParams { pub driver: Option, @@ -62,6 +66,37 @@ impl ConnectionParams { .ok_or_else(|| "Project ID is required in the Database field".to_string()) } + /// Project that owns and pays for query jobs. Multi-project browsing + /// sends the selected project in `database`, so an explicit + /// `billing_project` extra keeps jobs in the configured project. + pub fn billing_project_id(&self) -> Result<&str, String> { + match self.extra_value("billing_project") { + Some(project) => Ok(project), + None => self.project_id(), + } + } + + /// Project whose datasets are browsed: the selected database, or the + /// billing project when the host sends no database ("all databases"). + pub fn data_project_id(&self) -> Result<&str, String> { + self.project_id().or_else(|_| self.billing_project_id()) + } + + /// Projects listed by `get_databases` in addition to the projects the + /// credential can see, e.g. `bigquery-public-data`. + pub fn additional_projects(&self) -> Vec { + self.extra_value("additional_projects") + .map(|value| { + value + .split([',', '\n', ' ']) + .map(str::trim) + .filter(|project| !project.is_empty()) + .map(str::to_owned) + .collect() + }) + .unwrap_or_default() + } + pub fn extra_value(&self, key: &str) -> Option<&str> { self.extra .get(key) @@ -95,7 +130,15 @@ impl ConnectionParams { self.extra_value("query_timeout_seconds") .and_then(|value| value.parse::().ok()) .map(|value| value.clamp(30, 3_600)) - .unwrap_or(300) + .unwrap_or(DEFAULT_TIMEOUT_SECONDS) + } + + pub fn api_endpoint(&self) -> Option<&str> { + self.extra_value("api_endpoint") + } + + pub fn impersonate_service_account(&self) -> Option<&str> { + self.extra_value("impersonate_service_account") } } @@ -103,6 +146,18 @@ pub fn inner_params(value: &Value) -> &Value { value.get("params").unwrap_or(value) } +/// Data project for generated SQL. Current hosts send the connection +/// parameters, with the browsed project applied, to every SQL builder; older +/// hosts send none, and the statements are then dataset-qualified and resolve +/// in the job project. +pub fn optional_data_project(request: &Value) -> Option { + request.get("params")?; + ConnectionParams::from_value(inner_params(request)) + .data_project_id() + .ok() + .map(str::to_owned) +} + #[derive(Debug, Clone, Deserialize, PartialEq)] pub struct ColumnDefinition { pub name: String, @@ -114,6 +169,11 @@ pub struct ColumnDefinition { pub is_auto_increment: bool, #[serde(default)] pub default_value: Option, + /// Column description. Hosts send it only when the user entered one + /// (manifest capability `column_comments`); the DDL builders emit it as + /// `OPTIONS(description = ...)`. + #[serde(default, alias = "description")] + pub comment: Option, } #[cfg(test)] @@ -135,6 +195,24 @@ mod tests { assert!(!params.extra.contains_key("ignored")); } + #[test] + fn billing_project_defaults_to_database_and_honors_override() { + let value = json!({"database": "public-data"}); + let params = ConnectionParams::from_value(&value); + assert_eq!(params.billing_project_id().unwrap(), "public-data"); + let value = json!({ + "database": "bigquery-public-data", + "extra": {"billing_project": "my-billing", "additional_projects": "a-1, b-2\nc"} + }); + let params = ConnectionParams::from_value(&value); + assert_eq!(params.project_id().unwrap(), "bigquery-public-data"); + assert_eq!(params.billing_project_id().unwrap(), "my-billing"); + assert_eq!(params.additional_projects(), vec!["a-1", "b-2", "c"]); + let value = json!({"database": "", "extra": {"billing_project": "my-billing"}}); + let params = ConnectionParams::from_value(&value); + assert_eq!(params.data_project_id().unwrap(), "my-billing"); + } + #[test] fn clamps_timeout() { let value = json!({"database": "p", "extra": {"query_timeout_seconds": "2"}}); diff --git a/src/routines.rs b/src/routines.rs new file mode 100644 index 0000000..913c4c3 --- /dev/null +++ b/src/routines.rs @@ -0,0 +1,921 @@ +//! Routine metadata and SQL built from BigQuery `routines.get` resources. + +use serde::Deserialize; +use serde_json::{json, Value}; + +use crate::bigquery::{FieldSchema, RangeElementType}; +use crate::schema_ddl::options_clause; +use crate::sql::{canonical_type, quote_identifier, string_literal, value_literal}; + +/// Routine kinds as reported by the REST API. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RoutineKind { + ScalarFunction, + AggregateFunction, + TableFunction, + Procedure, +} + +impl RoutineKind { + pub fn from_resource(resource: &Value) -> Self { + Self::from_name( + resource + .get("routineType") + .and_then(Value::as_str) + .unwrap_or("SCALAR_FUNCTION"), + ) + } + + /// Accepts REST names and the host's `FUNCTION` / `PROCEDURE`. + pub fn from_name(name: &str) -> Self { + match name + .trim() + .to_ascii_uppercase() + .replace([' ', '-'], "_") + .as_str() + { + "PROCEDURE" => Self::Procedure, + "TABLE_VALUED_FUNCTION" | "TABLE_FUNCTION" | "TVF" => Self::TableFunction, + "AGGREGATE_FUNCTION" => Self::AggregateFunction, + _ => Self::ScalarFunction, + } + } + + /// The host groups routines into `PROCEDURE` and `FUNCTION`. + pub fn host_type(self) -> &'static str { + match self { + Self::Procedure => "PROCEDURE", + _ => "FUNCTION", + } + } + + pub fn keyword(self) -> &'static str { + match self { + Self::ScalarFunction => "FUNCTION", + Self::AggregateFunction => "AGGREGATE FUNCTION", + Self::TableFunction => "TABLE FUNCTION", + Self::Procedure => "PROCEDURE", + } + } + + pub fn drop_keyword(self) -> &'static str { + match self { + Self::AggregateFunction | Self::ScalarFunction => "FUNCTION", + Self::TableFunction => "TABLE FUNCTION", + Self::Procedure => "PROCEDURE", + } + } +} + +/// `{name, routine_type}` entry for `get_routines`. +pub fn routine_info(resource: &Value) -> Option { + let name = resource + .pointer("/routineReference/routineId") + .and_then(Value::as_str)?; + Some(json!({ + "name": name, + "routine_type": RoutineKind::from_resource(resource).host_type(), + "definition": null + })) +} + +/// Renders a REST `StandardSqlDataType`. +pub fn standard_type_sql(data_type: &Value) -> String { + let kind = data_type + .get("typeKind") + .and_then(Value::as_str) + .unwrap_or("TYPE_KIND_UNSPECIFIED"); + match kind { + "ARRAY" => format!( + "ARRAY<{}>", + data_type + .get("arrayElementType") + .map(standard_type_sql) + .unwrap_or_else(|| "STRING".to_string()) + ), + "STRUCT" => { + let fields = data_type + .pointer("/structType/fields") + .and_then(Value::as_array) + .map(|fields| { + fields + .iter() + .map(|field| { + let field_type = field + .get("type") + .map(standard_type_sql) + .unwrap_or_else(|| "STRING".to_string()); + match field.get("name").and_then(Value::as_str) { + Some(name) if !name.is_empty() => { + format!("{} {field_type}", quote_identifier(name)) + } + _ => field_type, + } + }) + .collect::>() + .join(", ") + }) + .unwrap_or_default(); + format!("STRUCT<{fields}>") + } + "RANGE" => format!( + "RANGE<{}>", + data_type + .get("rangeElementType") + .map(standard_type_sql) + .unwrap_or_else(|| "DATE".to_string()) + ), + "TYPE_KIND_UNSPECIFIED" => "ANY TYPE".to_string(), + other => canonical_type(other), + } +} + +/// Converts a `StandardSqlDataType` into the table field shape used by the +/// literal renderer. +pub fn standard_type_field(name: &str, data_type: &Value) -> FieldSchema { + let kind = data_type + .get("typeKind") + .and_then(Value::as_str) + .unwrap_or("STRING"); + match kind { + "ARRAY" => { + let mut element = data_type + .get("arrayElementType") + .map(|element| standard_type_field(name, element)) + .unwrap_or_else(|| FieldSchema { + name: name.to_string(), + field_type: "STRING".to_string(), + ..Default::default() + }); + element.mode = Some("REPEATED".to_string()); + element + } + "STRUCT" => FieldSchema { + name: name.to_string(), + field_type: "RECORD".to_string(), + fields: data_type + .pointer("/structType/fields") + .and_then(Value::as_array) + .map(|fields| { + fields + .iter() + .map(|field| { + standard_type_field( + field + .get("name") + .and_then(Value::as_str) + .unwrap_or_default(), + field.get("type").unwrap_or(&Value::Null), + ) + }) + .collect() + }) + .unwrap_or_default(), + ..Default::default() + }, + "RANGE" => FieldSchema { + name: name.to_string(), + field_type: "RANGE".to_string(), + range_element_type: Some(RangeElementType { + element_type: data_type + .pointer("/rangeElementType/typeKind") + .and_then(Value::as_str) + .unwrap_or("DATE") + .to_string(), + }), + ..Default::default() + }, + other => FieldSchema { + name: name.to_string(), + field_type: other.to_string(), + ..Default::default() + }, + } +} + +fn arguments(resource: &Value) -> Vec { + resource + .get("arguments") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default() +} + +fn argument_type_sql(argument: &Value) -> String { + if argument.get("argumentKind").and_then(Value::as_str) == Some("ANY_TYPE") { + return "ANY TYPE".to_string(); + } + argument + .get("dataType") + .map(standard_type_sql) + .unwrap_or_else(|| "ANY TYPE".to_string()) +} + +fn return_table_sql(resource: &Value) -> Option { + let columns = resource + .pointer("/returnTableType/columns") + .and_then(Value::as_array)?; + Some(format!( + "TABLE<{}>", + columns + .iter() + .map(|column| { + format!( + "{} {}", + quote_identifier( + column + .get("name") + .and_then(Value::as_str) + .unwrap_or_default() + ), + column + .get("type") + .map(standard_type_sql) + .unwrap_or_else(|| "STRING".to_string()) + ) + }) + .collect::>() + .join(", ") + )) +} + +fn return_type_sql(resource: &Value) -> Option { + resource.get("returnType").map(standard_type_sql) +} + +/// Host `RoutineParameter` rows. Arguments are numbered from 1; a function's +/// return type is reported as the unnamed `OUT` row at position 0. +pub fn routine_parameters(resource: &Value) -> Vec { + let kind = RoutineKind::from_resource(resource); + let mut parameters = Vec::new(); + if kind != RoutineKind::Procedure { + let return_type = match kind { + RoutineKind::TableFunction => return_table_sql(resource), + _ => return_type_sql(resource), + }; + if let Some(data_type) = return_type { + parameters.push(json!({ + "name": "", + "data_type": data_type, + "mode": "OUT", + "ordinal_position": 0 + })); + } + } + for (index, argument) in arguments(resource).iter().enumerate() { + let position = index + 1; + let name = argument + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .map(str::to_owned) + .unwrap_or_else(|| format!("arg{position}")); + let mode = argument + .get("mode") + .and_then(Value::as_str) + .filter(|mode| matches!(*mode, "IN" | "OUT" | "INOUT")) + .unwrap_or("IN"); + parameters.push(json!({ + "name": name, + "data_type": argument_type_sql(argument), + "mode": mode, + "ordinal_position": position + })); + } + parameters +} + +/// A string literal for a routine body, preferring a readable raw string. +pub fn body_literal(body: &str) -> String { + if !body.contains("\"\"\"") && !body.ends_with('\\') && !body.ends_with('"') { + return format!("r\"\"\"{body}\"\"\""); + } + let escaped = body.replace('\\', "\\\\").replace('"', "\\\""); + format!("\"\"\"{escaped}\"\"\"") +} + +fn language(resource: &Value) -> String { + resource + .get("language") + .and_then(Value::as_str) + .unwrap_or("SQL") + .to_ascii_uppercase() +} + +fn routine_options(resource: &Value, kind: RoutineKind) -> Vec<(String, String)> { + let mut options = Vec::new(); + if let Some(description) = resource.get("description").and_then(Value::as_str) { + options.push(("description".to_string(), string_literal(description))); + } + if let Some(libraries) = resource.get("importedLibraries").and_then(Value::as_array) { + let libraries = libraries + .iter() + .filter_map(Value::as_str) + .map(string_literal) + .collect::>(); + if !libraries.is_empty() { + options.push(("library".to_string(), format!("[{}]", libraries.join(", ")))); + } + } + if kind == RoutineKind::Procedure + && resource.get("strictMode").and_then(Value::as_bool) == Some(false) + { + options.push(("strict_mode".to_string(), "FALSE".to_string())); + } + if let Some(governance) = resource.get("dataGovernanceType").and_then(Value::as_str) { + options.push(( + "data_governance_type".to_string(), + string_literal(governance), + )); + } + if let Some(python) = resource.get("pythonOptions") { + if let Some(entry_point) = python.get("entryPoint").and_then(Value::as_str) { + options.push(("entry_point".to_string(), string_literal(entry_point))); + } + if let Some(packages) = python.get("packages").and_then(Value::as_array) { + let packages = packages + .iter() + .filter_map(Value::as_str) + .map(string_literal) + .collect::>(); + options.push(("packages".to_string(), format!("[{}]", packages.join(", ")))); + } + } + if let Some(runtime) = resource + .pointer("/externalRuntimeOptions/runtimeVersion") + .and_then(Value::as_str) + { + options.push(("runtime_version".to_string(), string_literal(runtime))); + } + if let Some(remote) = resource.get("remoteFunctionOptions") { + if let Some(endpoint) = remote.get("endpoint").and_then(Value::as_str) { + options.push(("endpoint".to_string(), string_literal(endpoint))); + } + if let Some(rows) = remote.get("maxBatchingRows").and_then(Value::as_str) { + options.push(("max_batching_rows".to_string(), rows.to_string())); + } + if let Some(context) = remote.get("userDefinedContext").and_then(Value::as_object) { + let mut pairs = context + .iter() + .map(|(key, value)| { + format!( + "({}, {})", + string_literal(key), + string_literal(value.as_str().unwrap_or_default()) + ) + }) + .collect::>(); + pairs.sort(); + options.push(( + "user_defined_context".to_string(), + format!("[{}]", pairs.join(", ")), + )); + } + } + if let Some(spark) = resource.get("sparkOptions") { + options.push(("engine".to_string(), "'SPARK'".to_string())); + for (key, option) in [ + ("runtimeVersion", "runtime_version"), + ("mainFileUri", "main_file_uri"), + ("mainClass", "main_class"), + ("containerImage", "container_image"), + ] { + if let Some(value) = spark.get(key).and_then(Value::as_str) { + options.push((option.to_string(), string_literal(value))); + } + } + } + options +} + +fn connection_clause(connection: &str) -> String { + let parts = connection.split('/').collect::>(); + let path = if parts.len() == 6 && parts[0] == "projects" { + format!("{}.{}.{}", parts[1], parts[3], parts[5]) + } else { + connection.to_string() + }; + format!("WITH CONNECTION {}", quote_identifier(&path)) +} + +/// Rebuilds `CREATE [OR REPLACE] ...` for a routine from its REST resource, +/// including the language, non-SQL bodies and OPTIONS. +pub fn routine_ddl(resource: &Value, path: &str, replace: bool) -> Result { + let kind = RoutineKind::from_resource(resource); + let body = resource + .get("definitionBody") + .and_then(Value::as_str) + .unwrap_or_default(); + let language = language(resource); + let options = routine_options(resource, kind); + let create = if replace { + "CREATE OR REPLACE" + } else { + "CREATE" + }; + let argument_list = arguments(resource) + .iter() + .map(|argument| { + let name = quote_identifier( + argument + .get("name") + .and_then(Value::as_str) + .unwrap_or_default(), + ); + let data_type = argument_type_sql(argument); + match kind { + RoutineKind::Procedure => { + let mode = argument.get("mode").and_then(Value::as_str).unwrap_or("IN"); + format!("{mode} {name} {data_type}") + } + RoutineKind::AggregateFunction + if argument.get("isAggregate").and_then(Value::as_bool) == Some(false) => + { + format!("{name} {data_type} NOT AGGREGATE") + } + _ => format!("{name} {data_type}"), + } + }) + .collect::>() + .join(", "); + let mut sql = format!("{create} {} {path}({argument_list})", kind.keyword()); + + match kind { + RoutineKind::Procedure => { + if language != "SQL" { + if let Some(connection) = resource + .pointer("/sparkOptions/connection") + .and_then(Value::as_str) + { + sql.push('\n'); + sql.push_str(&connection_clause(connection)); + } + sql.push_str(&options_clause(&options)); + sql.push_str(&format!("\nLANGUAGE {language}")); + if !body.is_empty() { + sql.push_str(&format!("\nAS {}", body_literal(body))); + } + return Ok(sql); + } + sql.push_str(&options_clause(&options)); + let trimmed = body.trim(); + if trimmed.to_ascii_uppercase().starts_with("BEGIN") { + sql.push_str(&format!("\n{trimmed}")); + } else { + sql.push_str(&format!("\nBEGIN\n{trimmed}\nEND")); + } + Ok(sql) + } + RoutineKind::TableFunction => { + if let Some(table) = return_table_sql(resource) { + sql.push_str(&format!("\nRETURNS {table}")); + } + sql.push_str(&options_clause(&options)); + sql.push_str(&format!("\nAS\n{}", body.trim())); + Ok(sql) + } + RoutineKind::ScalarFunction | RoutineKind::AggregateFunction => { + if let Some(return_type) = return_type_sql(resource) { + sql.push_str(&format!("\nRETURNS {return_type}")); + } + if let Some(remote) = resource.get("remoteFunctionOptions") { + let connection = remote + .get("connection") + .and_then(Value::as_str) + .unwrap_or_default(); + sql.push_str(&format!("\nREMOTE {}", connection_clause(connection))); + sql.push_str(&options_clause(&options)); + return Ok(sql); + } + if language == "SQL" { + sql.push_str(&format!("\nAS (\n {}\n)", body.trim())); + sql.push_str(&options_clause(&options)); + return Ok(sql); + } + match resource.get("determinismLevel").and_then(Value::as_str) { + Some("DETERMINISTIC") => sql.push_str("\nDETERMINISTIC"), + Some("NOT_DETERMINISTIC") => sql.push_str("\nNOT DETERMINISTIC"), + _ => {} + } + let language_name = match language.as_str() { + "JAVASCRIPT" => "js".to_string(), + "PYTHON" => "python".to_string(), + other => other.to_ascii_lowercase(), + }; + sql.push_str(&format!("\nLANGUAGE {language_name}")); + if let Some(connection) = resource + .pointer("/externalRuntimeOptions/connection") + .and_then(Value::as_str) + { + sql.push('\n'); + sql.push_str(&connection_clause(connection)); + } + sql.push_str(&options_clause(&options)); + sql.push_str(&format!("\nAS {}", body_literal(body))); + Ok(sql) + } + } +} + +pub fn drop_routine_sql(kind: RoutineKind, path: &str) -> String { + format!("DROP {} IF EXISTS {path}", kind.drop_keyword()) +} + +/// Mirror of the host `RoutineCallArg`. +#[derive(Debug, Clone, Deserialize)] +pub struct CallArg { + #[serde(default)] + pub name: String, + #[serde(default = "default_mode")] + pub mode: String, + #[serde(default)] + pub value: Option, + #[serde(default)] + pub is_raw: bool, +} + +fn default_mode() -> String { + "IN".to_string() +} + +fn argument_value(arg: &CallArg, data_type: Option<&Value>) -> Result { + let Some(value) = arg.value.as_deref() else { + return Ok("NULL".to_string()); + }; + if arg.is_raw { + return Ok(value.to_string()); + } + match data_type { + Some(data_type) if data_type.get("typeKind").is_some() => value_literal( + &Value::String(value.to_string()), + Some(&standard_type_field(&arg.name, data_type)), + ) + .map_err(|error| format!("Argument '{}': {error}", arg.name)), + _ => Ok(string_literal(value)), + } +} + +/// Builds the invocation script for the Run Routine dialog: +/// `CALL` for procedures (OUT/INOUT through script variables), +/// `SELECT fn(...)` for scalar and aggregate functions and +/// `SELECT * FROM tvf(...)` for table functions. +pub fn routine_call_sql(resource: &Value, path: &str, args: &[CallArg]) -> Result { + let kind = RoutineKind::from_resource(resource); + let declared = arguments(resource); + let data_type = |index: usize, arg: &CallArg| { + declared + .iter() + .find(|argument| { + argument.get("name").and_then(Value::as_str) == Some(arg.name.as_str()) + }) + .or_else(|| declared.get(index)) + .and_then(|argument| argument.get("dataType")) + .cloned() + }; + + if kind != RoutineKind::Procedure { + let values = args + .iter() + .enumerate() + .map(|(index, arg)| argument_value(arg, data_type(index, arg).as_ref())) + .collect::, _>>()?; + let call = format!("{path}({})", values.join(", ")); + return Ok(match kind { + RoutineKind::TableFunction => format!("SELECT *\nFROM {call};"), + _ => format!("SELECT {call} AS result;"), + }); + } + + let mut declarations = Vec::new(); + let mut call_args = Vec::new(); + let mut outputs = Vec::new(); + for (index, arg) in args.iter().enumerate() { + let mode = arg.mode.to_ascii_uppercase(); + if mode == "IN" { + call_args.push(argument_value(arg, data_type(index, arg).as_ref())?); + continue; + } + let variable = format!("tabularis_arg_{}", index + 1); + let declared_type = data_type(index, arg) + .map(|data_type| standard_type_sql(&data_type)) + .filter(|data_type| data_type != "ANY TYPE") + .unwrap_or_else(|| "STRING".to_string()); + let default = if mode == "INOUT" && arg.value.is_some() { + format!( + " DEFAULT {}", + argument_value(arg, data_type(index, arg).as_ref())? + ) + } else { + String::new() + }; + declarations.push(format!("DECLARE {variable} {declared_type}{default};")); + let label = if arg.name.is_empty() { + variable.clone() + } else { + arg.name.clone() + }; + outputs.push(format!("{variable} AS {}", quote_identifier(&label))); + call_args.push(variable); + } + let mut script = declarations; + script.push(format!("CALL {path}({});", call_args.join(", "))); + if !outputs.is_empty() { + script.push(format!("SELECT {};", outputs.join(", "))); + } + Ok(script.join("\n")) +} + +/// Starter script for a new routine. +pub fn create_template(routine_type: &str, project: Option<&str>, dataset: Option<&str>) -> String { + let path = |name: &str| match ( + project.filter(|value| !value.is_empty()), + dataset.filter(|value| !value.is_empty()), + ) { + (Some(project), Some(dataset)) => crate::sql::fully_qualified(project, dataset, name), + (None, Some(dataset)) => crate::sql::qualified(dataset, name), + _ => quote_identifier(name), + }; + match RoutineKind::from_name(routine_type) { + RoutineKind::Procedure => format!( + "CREATE OR REPLACE PROCEDURE {}(IN input_id INT64, OUT row_count INT64)\n\ + OPTIONS(description = 'Describe the procedure')\n\ + BEGIN\n \ + -- Multi-statement GoogleSQL body\n \ + SET row_count = (\n SELECT COUNT(*)\n FROM {}\n WHERE id = input_id\n );\n\ + END;", + path("my_procedure"), + path("my_table") + ), + RoutineKind::TableFunction => format!( + "CREATE OR REPLACE TABLE FUNCTION {}(min_id INT64)\n\ + OPTIONS(description = 'Describe the table function')\n\ + AS\n\ + SELECT *\nFROM {}\nWHERE id >= min_id;", + path("my_table_function"), + path("my_table") + ), + RoutineKind::ScalarFunction | RoutineKind::AggregateFunction => format!( + "CREATE OR REPLACE FUNCTION {}(x INT64, y INT64)\n\ + RETURNS INT64\n\ + AS (\n x + y\n)\n\ + OPTIONS(description = 'Describe the function');\n\n\ + -- JavaScript UDF:\n\ + -- CREATE OR REPLACE FUNCTION {}(s STRING)\n\ + -- RETURNS STRING\n\ + -- LANGUAGE js\n\ + -- AS r\"\"\"return s.toUpperCase();\"\"\";\n\n\ + -- Table function:\n\ + -- CREATE OR REPLACE TABLE FUNCTION {}(min_id INT64)\n\ + -- AS SELECT * FROM {} WHERE id >= min_id;", + path("my_function"), + path("my_js_function"), + path("my_table_function"), + path("my_table") + ), + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn scalar() -> Value { + json!({ + "routineReference": {"projectId": "my-proj", "datasetId": "d", "routineId": "add_tax"}, + "routineType": "SCALAR_FUNCTION", + "language": "SQL", + "arguments": [ + {"name": "amount", "dataType": {"typeKind": "NUMERIC"}}, + {"name": "rate", "dataType": {"typeKind": "FLOAT64"}} + ], + "returnType": {"typeKind": "NUMERIC"}, + "definitionBody": "amount * (1 + rate)", + "description": "Adds 'tax'" + }) + } + + #[test] + fn maps_routine_kinds_to_host_types() { + assert_eq!(routine_info(&scalar()).unwrap()["routine_type"], "FUNCTION"); + let procedure = json!({"routineReference": {"routineId": "p"}, "routineType": "PROCEDURE"}); + assert_eq!( + routine_info(&procedure).unwrap()["routine_type"], + "PROCEDURE" + ); + let tvf = + json!({"routineReference": {"routineId": "t"}, "routineType": "TABLE_VALUED_FUNCTION"}); + assert_eq!(routine_info(&tvf).unwrap()["routine_type"], "FUNCTION"); + assert_eq!( + RoutineKind::from_name("TABLE FUNCTION"), + RoutineKind::TableFunction + ); + assert_eq!( + drop_routine_sql(RoutineKind::TableFunction, "`p-1`.`d`.`t`"), + "DROP TABLE FUNCTION IF EXISTS `p-1`.`d`.`t`" + ); + assert_eq!( + drop_routine_sql(RoutineKind::AggregateFunction, "`p`.`d`.`a`"), + "DROP FUNCTION IF EXISTS `p`.`d`.`a`" + ); + } + + #[test] + fn parameters_have_positions_and_a_named_return_row() { + let parameters = routine_parameters(&scalar()); + assert_eq!( + parameters, + vec![ + json!({"name": "", "data_type": "NUMERIC", "mode": "OUT", "ordinal_position": 0}), + json!({"name": "amount", "data_type": "NUMERIC", "mode": "IN", "ordinal_position": 1}), + json!({"name": "rate", "data_type": "FLOAT64", "mode": "IN", "ordinal_position": 2}), + ] + ); + let procedure = json!({ + "routineType": "PROCEDURE", + "arguments": [ + {"name": "x", "mode": "INOUT", "dataType": {"typeKind": "ARRAY", "arrayElementType": {"typeKind": "STRING"}}}, + {"argumentKind": "ANY_TYPE", "mode": "OUT"} + ] + }); + assert_eq!( + routine_parameters(&procedure), + vec![ + json!({"name": "x", "data_type": "ARRAY", "mode": "INOUT", "ordinal_position": 1}), + json!({"name": "arg2", "data_type": "ANY TYPE", "mode": "OUT", "ordinal_position": 2}), + ] + ); + let inferred = json!({"routineType": "SCALAR_FUNCTION", "arguments": []}); + assert!(routine_parameters(&inferred).is_empty()); + } + + #[test] + fn rebuilds_sql_and_javascript_functions() { + assert_eq!( + routine_ddl(&scalar(), "`my-proj`.`d`.`add_tax`", true).unwrap(), + "CREATE OR REPLACE FUNCTION `my-proj`.`d`.`add_tax`(`amount` NUMERIC, `rate` FLOAT64)\nRETURNS NUMERIC\nAS (\n amount * (1 + rate)\n) OPTIONS(description = 'Adds \\'tax\\'')" + ); + let js = json!({ + "routineType": "SCALAR_FUNCTION", + "language": "JAVASCRIPT", + "arguments": [{"name": "s", "dataType": {"typeKind": "STRING"}}], + "returnType": {"typeKind": "STRING"}, + "determinismLevel": "DETERMINISTIC", + "importedLibraries": ["gs://bucket/lib.js"], + "definitionBody": "return s + \"\\n\";" + }); + assert_eq!( + routine_ddl(&js, "`p`.`d`.`f`", true).unwrap(), + "CREATE OR REPLACE FUNCTION `p`.`d`.`f`(`s` STRING)\nRETURNS STRING\nDETERMINISTIC\nLANGUAGE js OPTIONS(library = ['gs://bucket/lib.js'])\nAS r\"\"\"return s + \"\\n\";\"\"\"" + ); + } + + #[test] + fn rebuilds_table_functions_procedures_and_aggregates() { + let tvf = json!({ + "routineType": "TABLE_VALUED_FUNCTION", + "arguments": [{"name": "min_id", "dataType": {"typeKind": "INT64"}}], + "returnTableType": {"columns": [{"name": "id", "type": {"typeKind": "INT64"}}]}, + "definitionBody": "SELECT id FROM `p`.`d`.`t` WHERE id >= min_id" + }); + assert_eq!( + routine_ddl(&tvf, "`p`.`d`.`tvf`", true).unwrap(), + "CREATE OR REPLACE TABLE FUNCTION `p`.`d`.`tvf`(`min_id` INT64)\nRETURNS TABLE<`id` INT64>\nAS\nSELECT id FROM `p`.`d`.`t` WHERE id >= min_id" + ); + let procedure = json!({ + "routineType": "PROCEDURE", + "strictMode": false, + "arguments": [ + {"name": "a", "mode": "IN", "dataType": {"typeKind": "INT64"}}, + {"name": "b", "mode": "OUT", "dataType": {"typeKind": "STRUCT", "structType": {"fields": [{"name": "x", "type": {"typeKind": "DATE"}}]}}} + ], + "definitionBody": "BEGIN\n SET b = STRUCT(CURRENT_DATE());\nEND" + }); + assert_eq!( + routine_ddl(&procedure, "`p`.`d`.`proc`", true).unwrap(), + "CREATE OR REPLACE PROCEDURE `p`.`d`.`proc`(IN `a` INT64, OUT `b` STRUCT<`x` DATE>) OPTIONS(strict_mode = FALSE)\nBEGIN\n SET b = STRUCT(CURRENT_DATE());\nEND" + ); + let aggregate = json!({ + "routineType": "AGGREGATE_FUNCTION", + "arguments": [ + {"name": "v", "dataType": {"typeKind": "FLOAT64"}}, + {"name": "w", "isAggregate": false, "dataType": {"typeKind": "FLOAT64"}} + ], + "returnType": {"typeKind": "FLOAT64"}, + "definitionBody": "SUM(v) * w" + }); + assert_eq!( + routine_ddl(&aggregate, "`p`.`d`.`agg`", false).unwrap(), + "CREATE AGGREGATE FUNCTION `p`.`d`.`agg`(`v` FLOAT64, `w` FLOAT64 NOT AGGREGATE)\nRETURNS FLOAT64\nAS (\n SUM(v) * w\n)" + ); + } + + #[test] + fn body_literals_stay_valid() { + assert_eq!(body_literal("return 1;"), "r\"\"\"return 1;\"\"\""); + assert_eq!( + body_literal("x = \"\"\"; y\\"), + "\"\"\"x = \\\"\\\"\\\"; y\\\\\"\"\"" + ); + assert_eq!(body_literal("say \"hi\""), "\"\"\"say \\\"hi\\\"\"\"\""); + } + + #[test] + fn builds_calls_with_typed_arguments() { + let args = vec![ + CallArg { + name: "amount".into(), + mode: "IN".into(), + value: Some("10.5".into()), + is_raw: false, + }, + CallArg { + name: "rate".into(), + mode: "IN".into(), + value: None, + is_raw: false, + }, + ]; + assert_eq!( + routine_call_sql(&scalar(), "`my-proj`.`d`.`add_tax`", &args).unwrap(), + "SELECT `my-proj`.`d`.`add_tax`(NUMERIC '10.5', NULL) AS result;" + ); + let bad = vec![CallArg { + name: "amount".into(), + mode: "IN".into(), + value: Some("ten".into()), + is_raw: false, + }]; + assert!(routine_call_sql(&scalar(), "`p`.`d`.`f`", &bad).is_err()); + + let tvf = json!({ + "routineType": "TABLE_VALUED_FUNCTION", + "arguments": [{"name": "label", "dataType": {"typeKind": "STRING"}}] + }); + let args = vec![CallArg { + name: "label".into(), + mode: "IN".into(), + value: Some("it's\nme".into()), + is_raw: false, + }]; + assert_eq!( + routine_call_sql(&tvf, "`p`.`d`.`tvf`", &args).unwrap(), + "SELECT *\nFROM `p`.`d`.`tvf`('it\\'s\\nme');" + ); + + let procedure = json!({ + "routineType": "PROCEDURE", + "arguments": [ + {"name": "a", "mode": "IN", "dataType": {"typeKind": "INT64"}}, + {"name": "b", "mode": "INOUT", "dataType": {"typeKind": "STRING"}}, + {"name": "c", "mode": "OUT", "dataType": {"typeKind": "INT64"}} + ] + }); + let args = vec![ + CallArg { + name: "a".into(), + mode: "IN".into(), + value: Some("1".into()), + is_raw: false, + }, + CallArg { + name: "b".into(), + mode: "INOUT".into(), + value: Some("x".into()), + is_raw: false, + }, + CallArg { + name: "c".into(), + mode: "OUT".into(), + value: None, + is_raw: false, + }, + ]; + assert_eq!( + routine_call_sql(&procedure, "`p`.`d`.`proc`", &args).unwrap(), + "DECLARE tabularis_arg_2 STRING DEFAULT 'x';\nDECLARE tabularis_arg_3 INT64;\nCALL `p`.`d`.`proc`(1, tabularis_arg_2, tabularis_arg_3);\nSELECT tabularis_arg_2 AS `b`, tabularis_arg_3 AS `c`;" + ); + let raw = vec![CallArg { + name: "a".into(), + mode: "IN".into(), + value: Some("1 + 1".into()), + is_raw: true, + }]; + assert_eq!( + routine_call_sql(&procedure, "`p`.`d`.`proc`", &raw).unwrap(), + "CALL `p`.`d`.`proc`(1 + 1);" + ); + } + + #[test] + fn templates_are_fully_qualified() { + let procedure = create_template("PROCEDURE", Some("my-proj"), Some("d")); + assert!(procedure.starts_with("CREATE OR REPLACE PROCEDURE `my-proj`.`d`.`my_procedure`(")); + assert!(procedure.contains("BEGIN") && procedure.trim_end().ends_with("END;")); + let function = create_template("FUNCTION", None, Some("d")); + assert!(function.starts_with("CREATE OR REPLACE FUNCTION `d`.`my_function`(")); + let tvf = create_template("TABLE_VALUED_FUNCTION", Some("p"), Some("d")); + assert!(tvf.starts_with("CREATE OR REPLACE TABLE FUNCTION `p`.`d`.`my_table_function`")); + } +} diff --git a/src/rpc.rs b/src/rpc.rs index f487c78..5b38b16 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -3,11 +3,13 @@ use serde_json::{json, Value}; use crate::handlers; pub async fn handle_line(line: &str) -> Value { - let request: Value = match serde_json::from_str(line) { - Ok(value) => value, - Err(error) => return error_response(Value::Null, -32700, &format!("parse error: {error}")), - }; + match serde_json::from_str(line) { + Ok(request) => handle_request(request).await, + Err(error) => error_response(Value::Null, -32700, &format!("parse error: {error}")), + } +} +pub async fn handle_request(request: Value) -> Value { let id = request.get("id").cloned().unwrap_or(Value::Null); if request.get("jsonrpc").and_then(Value::as_str) != Some("2.0") { return error_response(id, -32600, "invalid JSON-RPC request"); @@ -44,9 +46,14 @@ pub async fn handle_line(line: &str) -> Value { "refresh_materialized_view" => { handlers::metadata::refresh_materialized_view(id, ¶ms).await } - "get_routines" => handlers::metadata::get_routines(id, ¶ms).await, - "get_routine_parameters" => handlers::metadata::get_routine_parameters(id, ¶ms).await, - "get_routine_definition" => handlers::metadata::get_routine_definition(id, ¶ms).await, + "get_routines" => handlers::routines::get_routines(id, ¶ms).await, + "get_routine_parameters" => handlers::routines::get_routine_parameters(id, ¶ms).await, + "get_routine_definition" => handlers::routines::get_routine_definition(id, ¶ms).await, + "get_routine_edit_script" => handlers::routines::get_routine_edit_script(id, ¶ms).await, + "build_routine_call_sql" => handlers::routines::build_routine_call_sql(id, ¶ms).await, + "routine_create_template" => handlers::routines::routine_create_template(id, ¶ms).await, + "drop_routine" => handlers::routines::drop_routine(id, ¶ms).await, + "get_table_ddl" => handlers::metadata::get_table_ddl(id, ¶ms).await, "get_schema_snapshot" => handlers::metadata::get_schema_snapshot(id, ¶ms).await, "get_all_columns_batch" => handlers::metadata::get_all_columns_batch(id, ¶ms).await, "get_all_foreign_keys_batch" => { @@ -56,6 +63,7 @@ pub async fn handle_line(line: &str) -> Value { "execute_query" => handlers::query::execute_query(id, ¶ms).await, "execute_query_batch" => handlers::query::execute_query_batch(id, ¶ms).await, "explain_query" => handlers::query::explain_query(id, ¶ms).await, + "release_session" => handlers::query::release_session(id, ¶ms).await, "insert_record" => handlers::crud::insert_record(id, ¶ms).await, "update_record" => handlers::crud::update_record(id, ¶ms).await, "delete_record" => handlers::crud::delete_record(id, ¶ms).await, @@ -66,6 +74,7 @@ pub async fn handle_line(line: &str) -> Value { "get_create_foreign_key_sql" => { handlers::ddl::get_create_foreign_key_sql(id, ¶ms).await } + "get_drop_table_sql" => handlers::ddl::get_drop_table_sql(id, ¶ms).await, "drop_index" => handlers::ddl::drop_index(id, ¶ms).await, "drop_foreign_key" => handlers::ddl::drop_foreign_key(id, ¶ms).await, "save_blob_to_file" => handlers::blob::save_blob_to_file(id, ¶ms).await, diff --git a/src/schema_ddl.rs b/src/schema_ddl.rs new file mode 100644 index 0000000..da72f8f --- /dev/null +++ b/src/schema_ddl.rs @@ -0,0 +1,717 @@ +//! GoogleSQL DDL rebuilt from BigQuery REST table resources. +//! +//! Used when `INFORMATION_SCHEMA.TABLES.ddl` is unavailable and to preserve +//! table options when a view is replaced. + +use serde_json::Value; + +use crate::bigquery::FieldSchema; +use crate::client::table_fields; +use crate::sql::{ + base_type, field_type_sql, fully_qualified, is_required, quote_identifier, string_literal, +}; + +/// BigQuery object kinds, as reported by `tables.list` / `tables.get`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TableKind { + Table, + View, + MaterializedView, + External, + Snapshot, +} + +impl TableKind { + pub fn from_resource(resource: &Value) -> Self { + Self::from_type( + resource + .get("type") + .and_then(Value::as_str) + .unwrap_or("TABLE"), + ) + } + + pub fn from_type(kind: &str) -> Self { + match kind { + "VIEW" => Self::View, + "MATERIALIZED_VIEW" => Self::MaterializedView, + "EXTERNAL" => Self::External, + "SNAPSHOT" => Self::Snapshot, + _ => Self::Table, + } + } + + /// Object keyword for `DROP ` and `CREATE `. + pub fn keyword(self) -> &'static str { + match self { + Self::Table => "TABLE", + Self::View => "VIEW", + Self::MaterializedView => "MATERIALIZED VIEW", + Self::External => "EXTERNAL TABLE", + Self::Snapshot => "SNAPSHOT TABLE", + } + } + + pub fn is_table_like(self) -> bool { + matches!(self, Self::Table | Self::External | Self::Snapshot) + } +} + +pub fn drop_sql(kind: TableKind, path: &str) -> String { + format!("DROP {} IF EXISTS {path}", kind.keyword()) +} + +/// Rejects DML on objects BigQuery cannot modify, with an actionable message. +pub fn ensure_dml_allowed(resource: &Value) -> Result<(), String> { + match TableKind::from_resource(resource) { + TableKind::Table => Ok(()), + TableKind::Snapshot => Err( + "Table snapshots are read-only in BigQuery. Restore the snapshot into a table with CREATE TABLE ... CLONE to edit rows." + .to_string(), + ), + TableKind::External => Err( + "External tables are read-only: BigQuery DML does not modify data stored outside BigQuery." + .to_string(), + ), + TableKind::View | TableKind::MaterializedView => { + Err("Views and materialized views cannot be edited with DML.".to_string()) + } + } +} + +/// Rejects schema changes on objects whose columns cannot be altered. +pub fn ensure_schema_editable(resource: &Value) -> Result<(), String> { + match TableKind::from_resource(resource) { + TableKind::Table => Ok(()), + TableKind::Snapshot => { + Err("Table snapshots are read-only: their schema cannot be altered.".to_string()) + } + TableKind::External => Err( + "External table columns cannot be altered with ALTER TABLE; recreate the table with CREATE OR REPLACE EXTERNAL TABLE." + .to_string(), + ), + TableKind::View | TableKind::MaterializedView => { + Err("Edit the view definition instead of altering its columns.".to_string()) + } + } +} + +pub fn primary_key_columns(resource: &Value) -> Vec { + resource + .pointer("/tableConstraints/primaryKey/columns") + .and_then(Value::as_array) + .map(|columns| { + columns + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect() + }) + .unwrap_or_default() +} + +pub fn is_legacy_sql_view(resource: &Value) -> bool { + resource + .pointer("/view/useLegacySql") + .and_then(Value::as_bool) + .unwrap_or(false) +} + +/// `OPTIONS(...)` with a leading space, or an empty string. +pub fn options_clause(options: &[(String, String)]) -> String { + if options.is_empty() { + return String::new(); + } + let body = options + .iter() + .map(|(key, value)| format!("{key} = {value}")) + .collect::>() + .join(", "); + format!(" OPTIONS({body})") +} + +/// Description, friendly name and labels shared by tables and views. +pub fn common_options(resource: &Value) -> Vec<(String, String)> { + let mut options = Vec::new(); + if let Some(name) = resource.get("friendlyName").and_then(Value::as_str) { + options.push(("friendly_name".to_string(), string_literal(name))); + } + if let Some(description) = resource.get("description").and_then(Value::as_str) { + options.push(("description".to_string(), string_literal(description))); + } + if let Some(labels) = resource.get("labels").and_then(Value::as_object) { + if !labels.is_empty() { + let mut pairs = labels + .iter() + .map(|(key, value)| { + format!( + "({}, {})", + string_literal(key), + string_literal(value.as_str().unwrap_or_default()) + ) + }) + .collect::>(); + pairs.sort(); + options.push(("labels".to_string(), format!("[{}]", pairs.join(", ")))); + } + } + options +} + +fn partition_options(resource: &Value) -> Vec<(String, String)> { + let mut options = Vec::new(); + if let Some(ms) = resource + .pointer("/timePartitioning/expirationMs") + .and_then(Value::as_str) + .and_then(|value| value.parse::().ok()) + { + let days = ms / 86_400_000.0; + let rendered = if days.fract() == 0.0 { + format!("{days:.0}") + } else { + format!("{days}") + }; + options.push(("partition_expiration_days".to_string(), rendered)); + } + let require = resource + .get("requirePartitionFilter") + .or_else(|| resource.pointer("/timePartitioning/requirePartitionFilter")) + .and_then(Value::as_bool) + .unwrap_or(false); + if require { + options.push(("require_partition_filter".to_string(), "TRUE".to_string())); + } + options +} + +/// ` PARTITION BY ...` with a leading newline, or an empty string. +pub fn partition_clause(resource: &Value, fields: &[FieldSchema]) -> String { + if let Some(partitioning) = resource.get("timePartitioning") { + let granularity = partitioning + .get("type") + .and_then(Value::as_str) + .unwrap_or("DAY"); + let expression = match partitioning.get("field").and_then(Value::as_str) { + Some(column) => { + let column_type = fields + .iter() + .find(|field| field.name == column) + .map(base_type) + .unwrap_or_else(|| "TIMESTAMP".to_string()); + let quoted = quote_identifier(column); + match (column_type.as_str(), granularity) { + ("DATE", "DAY") => quoted, + ("DATE", _) => format!("DATE_TRUNC({quoted}, {granularity})"), + ("DATETIME", _) => format!("DATETIME_TRUNC({quoted}, {granularity})"), + (_, "DAY") => format!("DATE({quoted})"), + _ => format!("TIMESTAMP_TRUNC({quoted}, {granularity})"), + } + } + None if granularity == "DAY" => "_PARTITIONDATE".to_string(), + None => format!("TIMESTAMP_TRUNC(_PARTITIONTIME, {granularity})"), + }; + return format!("\nPARTITION BY {expression}"); + } + if let Some(range) = resource.get("rangePartitioning") { + let column = range + .get("field") + .and_then(Value::as_str) + .unwrap_or_default(); + let bound = |key: &str| { + range + .pointer(&format!("/range/{key}")) + .and_then(|value| { + value + .as_str() + .map(str::to_owned) + .or_else(|| value.as_i64().map(|v| v.to_string())) + }) + .unwrap_or_else(|| "0".to_string()) + }; + return format!( + "\nPARTITION BY RANGE_BUCKET({}, GENERATE_ARRAY({}, {}, {}))", + quote_identifier(column), + bound("start"), + bound("end"), + bound("interval") + ); + } + String::new() +} + +pub fn cluster_clause(resource: &Value) -> String { + let columns = resource + .pointer("/clustering/fields") + .and_then(Value::as_array) + .map(|fields| { + fields + .iter() + .filter_map(Value::as_str) + .map(quote_identifier) + .collect::>() + }) + .unwrap_or_default(); + if columns.is_empty() { + String::new() + } else { + format!("\nCLUSTER BY {}", columns.join(", ")) + } +} + +/// One column definition for CREATE TABLE / ADD COLUMN. +pub fn column_definition(field: &FieldSchema) -> String { + let mut column = format!( + "{} {}", + quote_identifier(&field.name), + field_type_sql(field, true) + ); + // GoogleSQL column_schema: name type [NOT NULL] [DEFAULT expr] [OPTIONS(...)] + if is_required(field) { + column.push_str(" NOT NULL"); + } + if let Some(default) = field + .default_value_expression + .as_deref() + .filter(|value| !value.trim().is_empty()) + { + column.push_str(" DEFAULT "); + column.push_str(default.trim()); + } + if let Some(description) = field.description.as_deref() { + column.push_str(&format!( + " OPTIONS(description = {})", + string_literal(description) + )); + } + column +} + +fn constraint_lines(resource: &Value, project: &str, dataset: &str) -> Vec { + let mut lines = Vec::new(); + let primary_key = primary_key_columns(resource); + if !primary_key.is_empty() { + lines.push(format!( + "PRIMARY KEY ({}) NOT ENFORCED", + primary_key + .iter() + .map(|column| quote_identifier(column)) + .collect::>() + .join(", ") + )); + } + for foreign_key in resource + .pointer("/tableConstraints/foreignKeys") + .and_then(Value::as_array) + .into_iter() + .flatten() + { + let references = foreign_key + .get("columnReferences") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let column = |key: &str| { + references + .iter() + .filter_map(|reference| reference.get(key).and_then(Value::as_str)) + .map(quote_identifier) + .collect::>() + .join(", ") + }; + let referenced = |key: &str, fallback: &str| { + foreign_key + .pointer(&format!("/referencedTable/{key}")) + .and_then(Value::as_str) + .unwrap_or(fallback) + .to_string() + }; + let target = fully_qualified( + &referenced("projectId", project), + &referenced("datasetId", dataset), + &referenced("tableId", ""), + ); + let prefix = foreign_key + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .map(|name| format!("CONSTRAINT {} ", quote_identifier(name))) + .unwrap_or_default(); + lines.push(format!( + "{prefix}FOREIGN KEY ({}) REFERENCES {target}({}) NOT ENFORCED", + column("referencingColumn"), + column("referencedColumn") + )); + } + lines +} + +fn column_block(lines: &[String]) -> String { + if lines.is_empty() { + String::new() + } else { + format!(" (\n {}\n)", lines.join(",\n ")) + } +} + +/// Rebuilds the CREATE statement of any table-like object from its resource. +pub fn table_ddl( + resource: &Value, + project: &str, + dataset: &str, + name: &str, +) -> Result { + let path = fully_qualified(project, dataset, name); + let kind = TableKind::from_resource(resource); + match kind { + TableKind::View => view_ddl(resource, &path, None, false), + TableKind::MaterializedView => materialized_view_ddl(resource, &path), + TableKind::Snapshot => { + let base = |key: &str| { + resource + .pointer(&format!("/snapshotDefinition/baseTableReference/{key}")) + .and_then(Value::as_str) + .unwrap_or_default() + .to_string() + }; + let mut sql = format!( + "CREATE SNAPSHOT TABLE {path}\nCLONE {}", + fully_qualified(&base("projectId"), &base("datasetId"), &base("tableId")) + ); + if let Some(time) = resource + .pointer("/snapshotDefinition/snapshotTime") + .and_then(Value::as_str) + { + sql.push_str(&format!( + "\nFOR SYSTEM_TIME AS OF TIMESTAMP {}", + string_literal(time) + )); + } + sql.push_str(&options_clause(&common_options(resource))); + Ok(sql) + } + TableKind::External => { + let fields = table_fields(resource)?; + let columns = fields.iter().map(column_definition).collect::>(); + let mut options = Vec::new(); + if let Some(format) = resource + .pointer("/externalDataConfiguration/sourceFormat") + .and_then(Value::as_str) + { + let format = match format { + "NEWLINE_DELIMITED_JSON" => "JSON", + "GOOGLE_SHEETS" => "GOOGLE_SHEETS", + other => other, + }; + options.push(("format".to_string(), string_literal(format))); + } + if let Some(uris) = resource + .pointer("/externalDataConfiguration/sourceUris") + .and_then(Value::as_array) + { + let uris = uris + .iter() + .filter_map(Value::as_str) + .map(string_literal) + .collect::>(); + options.push(("uris".to_string(), format!("[{}]", uris.join(", ")))); + } + options.extend(common_options(resource)); + let connection = resource + .pointer("/externalDataConfiguration/connectionId") + .and_then(Value::as_str) + .map(|connection| format!("\nWITH CONNECTION {}", connection_path(connection))) + .unwrap_or_default(); + Ok(format!( + "CREATE EXTERNAL TABLE {path}{}{connection}{}", + column_block(&columns), + options_clause(&options) + )) + } + TableKind::Table => { + let fields = table_fields(resource)?; + let mut lines = fields.iter().map(column_definition).collect::>(); + lines.extend(constraint_lines(resource, project, dataset)); + let mut options = partition_options(resource); + options.extend(common_options(resource)); + let collation = resource + .get("defaultCollation") + .and_then(Value::as_str) + .map(|collation| format!("\nDEFAULT COLLATE {}", string_literal(collation))) + .unwrap_or_default(); + let clone = resource + .pointer("/cloneDefinition/baseTableReference") + .map(|base| { + let part = + |key: &str| base.get(key).and_then(Value::as_str).unwrap_or_default(); + format!( + "\n-- Cloned from {}", + fully_qualified(part("projectId"), part("datasetId"), part("tableId")) + ) + }) + .unwrap_or_default(); + Ok(format!( + "CREATE TABLE {path}{}{collation}{}{}{}{clone}", + column_block(&lines), + partition_clause(resource, &fields), + cluster_clause(resource), + options_clause(&options) + )) + } + } +} + +/// `projects/p/locations/l/connections/c` or `p.l.c` → `` `p.l.c` ``. +fn connection_path(connection: &str) -> String { + let parts = connection.split('/').collect::>(); + if parts.len() == 6 && parts[0] == "projects" { + quote_identifier(&format!("{}.{}.{}", parts[1], parts[3], parts[5])) + } else { + quote_identifier(connection) + } +} + +/// CREATE [OR REPLACE] VIEW with the resource options. `query` replaces the +/// stored query when an edited definition is saved. +pub fn view_ddl( + resource: &Value, + path: &str, + query: Option<&str>, + replace: bool, +) -> Result { + let stored = resource.pointer("/view/query").and_then(Value::as_str); + let query = query + .or(stored) + .ok_or_else(|| "The view definition is unavailable".to_string())?; + let prefix = if replace { + "CREATE OR REPLACE VIEW" + } else { + "CREATE VIEW" + }; + let mut sql = format!( + "{prefix} {path}{}\nAS\n{}", + options_clause(&common_options(resource)), + query.trim() + ); + if is_legacy_sql_view(resource) && query == stored.unwrap_or_default() { + sql = format!( + "-- Legacy SQL view: BigQuery stores this query as legacy SQL.\n-- Tabularis can display it but cannot replace it as a GoogleSQL view.\n{sql}" + ); + } + Ok(sql) +} + +pub fn materialized_view_ddl(resource: &Value, path: &str) -> Result { + let query = resource + .pointer("/materializedView/query") + .and_then(Value::as_str) + .ok_or_else(|| "The materialized view definition is unavailable".to_string())?; + let fields = table_fields(resource).unwrap_or_default(); + let mut options = Vec::new(); + if let Some(enabled) = resource + .pointer("/materializedView/enableRefresh") + .and_then(Value::as_bool) + { + options.push(( + "enable_refresh".to_string(), + if enabled { "TRUE" } else { "FALSE" }.to_string(), + )); + } + if let Some(ms) = resource + .pointer("/materializedView/refreshIntervalMs") + .and_then(Value::as_str) + .and_then(|value| value.parse::().ok()) + { + options.push(( + "refresh_interval_minutes".to_string(), + (ms / 60_000).to_string(), + )); + } + if let Some(staleness) = resource + .pointer("/materializedView/maxStaleness") + .and_then(Value::as_str) + { + options.push(( + "max_staleness".to_string(), + format!("INTERVAL {}", string_literal(staleness)), + )); + } + if resource + .pointer("/materializedView/allowNonIncrementalDefinition") + .and_then(Value::as_bool) + .unwrap_or(false) + { + options.push(( + "allow_non_incremental_definition".to_string(), + "TRUE".to_string(), + )); + } + options.extend(partition_options(resource)); + options.extend(common_options(resource)); + Ok(format!( + "CREATE MATERIALIZED VIEW {path}{}{}{}\nAS\n{}", + partition_clause(resource, &fields), + cluster_clause(resource), + options_clause(&options), + query.trim() + )) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn table_resource() -> Value { + json!({ + "type": "TABLE", + "description": "Daily \"orders\"\nwith notes", + "labels": {"team": "data", "env": "prod"}, + "schema": {"fields": [ + {"name": "id", "type": "INTEGER", "mode": "REQUIRED"}, + {"name": "created_at", "type": "TIMESTAMP"}, + {"name": "amount", "type": "NUMERIC", "precision": "12", "scale": "2", "description": "O'Brien amount"}, + {"name": "status", "type": "STRING", "maxLength": "10", "defaultValueExpression": "'new'"}, + {"name": "items", "type": "RECORD", "mode": "REPEATED", "fields": [ + {"name": "sku", "type": "STRING", "mode": "REQUIRED"} + ]}, + {"name": "valid", "type": "RANGE", "rangeElementType": {"type": "DATE"}}, + {"name": "customer_id", "type": "INT64"} + ]}, + "timePartitioning": {"type": "DAY", "field": "created_at", "expirationMs": "7776000000"}, + "requirePartitionFilter": true, + "clustering": {"fields": ["status", "customer_id"]}, + "tableConstraints": { + "primaryKey": {"columns": ["id"]}, + "foreignKeys": [{ + "name": "fk_customer", + "referencedTable": {"projectId": "my-proj", "datasetId": "crm", "tableId": "customers"}, + "columnReferences": [{"referencingColumn": "customer_id", "referencedColumn": "id"}] + }] + } + }) + } + + #[test] + fn column_definition_puts_not_null_before_default() { + let field: FieldSchema = serde_json::from_value(json!({ + "name": "status", + "type": "STRING", + "mode": "REQUIRED", + "defaultValueExpression": "'new'", + "description": "State" + })) + .unwrap(); + assert_eq!( + column_definition(&field), + "`status` STRING NOT NULL DEFAULT 'new' OPTIONS(description = 'State')" + ); + } + + #[test] + fn rebuilds_partitioned_clustered_table_ddl() { + let ddl = table_ddl(&table_resource(), "my-proj", "sales", "orders").unwrap(); + assert_eq!( + ddl, + "CREATE TABLE `my-proj`.`sales`.`orders` (\n \ + `id` INT64 NOT NULL,\n \ + `created_at` TIMESTAMP,\n \ + `amount` NUMERIC(12, 2) OPTIONS(description = 'O\\'Brien amount'),\n \ + `status` STRING(10) DEFAULT 'new',\n \ + `items` ARRAY>,\n \ + `valid` RANGE,\n \ + `customer_id` INT64,\n \ + PRIMARY KEY (`id`) NOT ENFORCED,\n \ + CONSTRAINT `fk_customer` FOREIGN KEY (`customer_id`) REFERENCES `my-proj`.`crm`.`customers`(`id`) NOT ENFORCED\n)\n\ + PARTITION BY DATE(`created_at`)\n\ + CLUSTER BY `status`, `customer_id` \ + OPTIONS(partition_expiration_days = 90, require_partition_filter = TRUE, description = 'Daily \"orders\"\\nwith notes', labels = [('env', 'prod'), ('team', 'data')])" + ); + } + + #[test] + fn rebuilds_ingestion_and_range_partitioning() { + let ingestion = json!({"timePartitioning": {"type": "HOUR"}}); + assert_eq!( + partition_clause(&ingestion, &[]), + "\nPARTITION BY TIMESTAMP_TRUNC(_PARTITIONTIME, HOUR)" + ); + let daily = json!({"timePartitioning": {"type": "DAY"}}); + assert_eq!( + partition_clause(&daily, &[]), + "\nPARTITION BY _PARTITIONDATE" + ); + let range = json!({"rangePartitioning": {"field": "bucket", "range": {"start": "0", "end": "100", "interval": "10"}}}); + assert_eq!( + partition_clause(&range, &[]), + "\nPARTITION BY RANGE_BUCKET(`bucket`, GENERATE_ARRAY(0, 100, 10))" + ); + } + + #[test] + fn rebuilds_views_snapshots_and_external_tables() { + let view = json!({"type": "VIEW", "description": "v", "view": {"query": "SELECT 1", "useLegacySql": false}}); + assert_eq!( + table_ddl(&view, "p-1", "d", "v").unwrap(), + "CREATE VIEW `p-1`.`d`.`v` OPTIONS(description = 'v')\nAS\nSELECT 1" + ); + let legacy = json!({"type": "VIEW", "view": {"query": "SELECT 1 FROM [p:d.t]", "useLegacySql": true}}); + assert!(table_ddl(&legacy, "p", "d", "v") + .unwrap() + .starts_with("-- Legacy SQL view")); + let snapshot = json!({ + "type": "SNAPSHOT", + "snapshotDefinition": { + "baseTableReference": {"projectId": "p-1", "datasetId": "d", "tableId": "t"}, + "snapshotTime": "2024-01-01T00:00:00Z" + } + }); + assert_eq!( + table_ddl(&snapshot, "p-1", "d", "s").unwrap(), + "CREATE SNAPSHOT TABLE `p-1`.`d`.`s`\nCLONE `p-1`.`d`.`t`\nFOR SYSTEM_TIME AS OF TIMESTAMP '2024-01-01T00:00:00Z'" + ); + let external = json!({ + "type": "EXTERNAL", + "schema": {"fields": [{"name": "a", "type": "STRING"}]}, + "externalDataConfiguration": {"sourceFormat": "CSV", "sourceUris": ["gs://b/x.csv"]} + }); + assert_eq!( + table_ddl(&external, "p", "d", "e").unwrap(), + "CREATE EXTERNAL TABLE `p`.`d`.`e` (\n `a` STRING\n) OPTIONS(format = 'CSV', uris = ['gs://b/x.csv'])" + ); + } + + #[test] + fn rebuilds_materialized_view_with_refresh_options() { + let mv = json!({ + "type": "MATERIALIZED_VIEW", + "schema": {"fields": [{"name": "day", "type": "DATE"}]}, + "timePartitioning": {"type": "DAY", "field": "day"}, + "clustering": {"fields": ["day"]}, + "materializedView": {"query": "SELECT day FROM t", "enableRefresh": true, "refreshIntervalMs": "1800000"} + }); + assert_eq!( + table_ddl(&mv, "p", "d", "mv").unwrap(), + "CREATE MATERIALIZED VIEW `p`.`d`.`mv`\nPARTITION BY `day`\nCLUSTER BY `day` OPTIONS(enable_refresh = TRUE, refresh_interval_minutes = 30)\nAS\nSELECT day FROM t" + ); + } + + #[test] + fn maps_kinds_to_drop_keywords_and_dml_rules() { + assert_eq!( + drop_sql(TableKind::from_type("SNAPSHOT"), "`p`.`d`.`s`"), + "DROP SNAPSHOT TABLE IF EXISTS `p`.`d`.`s`" + ); + assert_eq!(TableKind::from_type("EXTERNAL").keyword(), "EXTERNAL TABLE"); + assert_eq!( + TableKind::from_type("MATERIALIZED_VIEW").keyword(), + "MATERIALIZED VIEW" + ); + assert!(ensure_dml_allowed(&json!({"type": "TABLE"})).is_ok()); + assert!(ensure_dml_allowed(&json!({"type": "SNAPSHOT"})) + .unwrap_err() + .contains("read-only")); + assert!(ensure_dml_allowed(&json!({"type": "EXTERNAL"})).is_err()); + assert!(ensure_schema_editable(&json!({"type": "SNAPSHOT"})).is_err()); + } +} diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..a7339ee --- /dev/null +++ b/src/server.rs @@ -0,0 +1,177 @@ +//! Concurrent JSON-RPC dispatch over newline-delimited stdio. +//! +//! Each request runs in its own task, so a long query never blocks metadata +//! calls or the host's `cancel` notification. A single writer task owns the +//! output stream, so response lines never interleave. Notifications (messages +//! without an `id`) are never answered. A handler that panics is logged and +//! its request still gets a JSON-RPC internal error, so the host never waits +//! for a reply that cannot come. + +use std::collections::HashMap; +use std::future::Future; + +use serde_json::Value; +use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::sync::mpsc; +use tokio::task::{Id, JoinError, JoinSet}; + +use crate::{cancel, rpc}; + +pub async fn serve(reader: R, writer: W) +where + R: AsyncBufRead + Unpin, + W: AsyncWrite + Unpin + Send + 'static, +{ + serve_with(reader, writer, rpc::handle_request).await; +} + +/// [`serve`] with a custom request handler, so tests can inject one. +pub async fn serve_with(reader: R, writer: W, handler: H) +where + R: AsyncBufRead + Unpin, + W: AsyncWrite + Unpin + Send + 'static, + H: Fn(Value) -> F, + F: Future + Send + 'static, +{ + let (sender, mut receiver) = mpsc::unbounded_channel::(); + let writer_task = tokio::spawn(async move { + let mut writer = writer; + while let Some(response) = receiver.recv().await { + let mut encoded = match serde_json::to_vec(&response) { + Ok(value) => value, + Err(error) => { + eprintln!("failed to serialize JSON-RPC response: {error}"); + continue; + } + }; + encoded.push(b'\n'); + if writer.write_all(&encoded).await.is_err() || writer.flush().await.is_err() { + break; + } + } + }); + + let mut tasks = JoinSet::new(); + // Request id of each running task, to answer it if the task panics. + let mut pending: HashMap = HashMap::new(); + let mut lines = reader.lines(); + while let Ok(Some(line)) = lines.next_line().await { + while let Some(joined) = tasks.try_join_next_with_id() { + settle(joined, &mut pending, &sender); + } + if line.trim().is_empty() { + continue; + } + let message: Value = match serde_json::from_str(&line) { + Ok(value) => value, + Err(error) => { + let _ = sender.send(rpc::error_response( + Value::Null, + -32700, + &format!("parse error: {error}"), + )); + continue; + } + }; + let Some(id) = message.get("id").cloned() else { + handle_notification(&message); + continue; + }; + let task_sender = sender.clone(); + let request = handler(message); + let request_id = id.clone(); + let handle = tasks.spawn(async move { + let response = cancel::scope(request_id, request).await; + let _ = task_sender.send(response); + }); + pending.insert(handle.id(), id); + } + + // Input closed: let running requests answer, then flush the writer. + while let Some(joined) = tasks.join_next_with_id().await { + settle(joined, &mut pending, &sender); + } + drop(sender); + let _ = writer_task.await; +} + +/// Forget a finished task; if it panicked, log it and answer its request. +fn settle( + joined: Result<(Id, ()), JoinError>, + pending: &mut HashMap, + sender: &mpsc::UnboundedSender, +) { + match joined { + Ok((task, ())) => { + pending.remove(&task); + } + Err(error) => { + let id = pending.remove(&error.id()); + eprintln!( + "bigquery-plugin: request {} failed: {error}", + id.as_ref() + .map_or_else(|| "?".to_string(), Value::to_string) + ); + if let Some(id) = id { + let _ = sender.send(rpc::error_response( + id, + -32603, + "Internal error: the BigQuery plugin failed while handling this request", + )); + } + } + } +} + +fn handle_notification(message: &Value) { + if message.get("method").and_then(Value::as_str) == Some("cancel") { + if let Some(id) = message.pointer("/params/id") { + cancel::cancel(id); + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; + + use super::*; + + #[tokio::test] + async fn panicking_handlers_still_answer_with_an_internal_error() { + let (plugin_read, mut host_write) = tokio::io::duplex(4096); + let (host_read, plugin_write) = tokio::io::duplex(4096); + let serving = tokio::spawn(serve_with( + BufReader::new(plugin_read), + plugin_write, + |message: Value| async move { + if message["method"] == "boom" { + panic!("handler exploded"); + } + rpc::error_response(message["id"].clone(), 1, "ok") + }, + )); + host_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"method\":\"boom\",\"id\":7}\n") + .await + .unwrap(); + host_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"method\":\"fine\",\"id\":8}\n") + .await + .unwrap(); + drop(host_write); + serving.await.unwrap(); + + let mut lines = BufReader::new(host_read).lines(); + let mut responses = Vec::new(); + while let Some(line) = lines.next_line().await.unwrap() { + responses.push(serde_json::from_str::(&line).unwrap()); + } + responses.sort_by_key(|response| response["id"].as_i64()); + assert_eq!(responses.len(), 2, "{responses:?}"); + assert_eq!(responses[0]["id"], json!(7)); + assert_eq!(responses[0]["error"]["code"], json!(-32603)); + assert_eq!(responses[1]["id"], json!(8)); + } +} diff --git a/src/session.rs b/src/session.rs new file mode 100644 index 0000000..87f85e5 --- /dev/null +++ b/src/session.rs @@ -0,0 +1,88 @@ +//! Host editor sessions mapped to BigQuery sessions. +//! +//! Tabularis pins an editor tab to a `session_id`. The first statement run for +//! it creates a BigQuery session (`createSession: true`); later statements join +//! it through `connectionProperties`, so variables, temporary tables and +//! multi-statement transactions survive between runs. + +use std::collections::HashMap; +use std::sync::{Arc, OnceLock}; + +use tokio::sync::Mutex; + +use crate::client::BigQueryClient; + +#[derive(Default)] +pub struct SessionState { + /// BigQuery session id, once the first job created it. + pub bigquery_id: Option, + /// Sessions are regional; later jobs must run where the first one ran. + pub location: Option, + pub in_transaction: bool, + pub client: Option>, +} + +type SharedSession = Arc>; + +static SESSIONS: OnceLock>> = OnceLock::new(); + +fn sessions() -> &'static Mutex> { + SESSIONS.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// The state for a host session, created on first use. Statements of one +/// session run one at a time by locking the returned mutex. +pub async fn get_or_create(host_id: &str) -> SharedSession { + sessions() + .lock() + .await + .entry(host_id.to_string()) + .or_default() + .clone() +} + +/// Forget a host session and abort its BigQuery session, if any. +pub async fn release(host_id: &str) -> Result<(), String> { + let Some(session) = sessions().lock().await.remove(host_id) else { + return Ok(()); + }; + let mut state = session.lock().await; + abort(&mut state).await +} + +/// Abort every known session, for plugin shutdown. +pub async fn release_all() { + let all: Vec = sessions().lock().await.drain().map(|(_, s)| s).collect(); + for session in all { + let mut state = session.lock().await; + let _ = abort(&mut state).await; + } +} + +pub async fn abort(state: &mut SessionState) -> Result<(), String> { + // Take the id only once there is a client to abort it with, so it is + // never dropped without aborting the server-side session. + let Some(client) = state.client.clone() else { + return Ok(()); + }; + let Some(id) = state.bigquery_id.take() else { + return Ok(()); + }; + state.in_transaction = false; + client.abort_session(&id, state.location.as_deref()).await +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn abort_without_client_keeps_the_session_id() { + let mut state = SessionState { + bigquery_id: Some("bq-session".into()), + ..SessionState::default() + }; + abort(&mut state).await.unwrap(); + assert_eq!(state.bigquery_id.as_deref(), Some("bq-session")); + } +} diff --git a/src/sql.rs b/src/sql.rs index 96513e6..40203e1 100644 --- a/src/sql.rs +++ b/src/sql.rs @@ -1,79 +1,251 @@ -use serde_json::Value; +//! GoogleSQL identifier quoting, literal rendering and typed query parameters. + +use base64::Engine; +use serde_json::{json, Value}; use crate::bigquery::FieldSchema; +/// The value Tabularis sends when a cell is reset to its column default. +pub const USE_DEFAULT_SENTINEL: &str = "__USE_DEFAULT__"; + +/// Quotes one identifier part with backticks. Every part of a path is quoted +/// separately, so project IDs with dashes and names with dots stay intact. pub fn quote_identifier(identifier: &str) -> String { - format!("`{}`", identifier.replace('`', "\\`")) + let mut quoted = String::with_capacity(identifier.len() + 2); + quoted.push('`'); + for character in identifier.chars() { + match character { + '`' => quoted.push_str("\\`"), + '\\' => quoted.push_str("\\\\"), + '\n' => quoted.push_str("\\n"), + '\r' => quoted.push_str("\\r"), + '\t' => quoted.push_str("\\t"), + other => quoted.push(other), + } + } + quoted.push('`'); + quoted } +/// `` `dataset`.`object` `` (resolved against the job project). pub fn qualified(dataset: &str, object: &str) -> String { - quote_identifier(&format!("{dataset}.{object}")) + format!("{}.{}", quote_identifier(dataset), quote_identifier(object)) } +/// `` `project`.`dataset`.`object` ``. pub fn fully_qualified(project: &str, dataset: &str, object: &str) -> String { - quote_identifier(&format!("{project}.{dataset}.{object}")) + format!( + "{}.{}.{}", + quote_identifier(project), + quote_identifier(dataset), + quote_identifier(object) + ) +} + +/// Fully qualified when the project is known, dataset-qualified otherwise. +pub fn object_path(project: Option<&str>, dataset: &str, object: &str) -> String { + match project.filter(|project| !project.is_empty()) { + Some(project) => fully_qualified(project, dataset, object), + None => qualified(dataset, object), + } } +/// `` `project`.`dataset`.INFORMATION_SCHEMA. ``. +pub fn information_schema(project: &str, dataset: &str, view: &str) -> String { + format!( + "{}.{}.INFORMATION_SCHEMA.{view}", + quote_identifier(project), + quote_identifier(dataset) + ) +} + +/// A single-quoted GoogleSQL string literal. +pub fn string_literal(value: &str) -> String { + format!("'{}'", escape_string(value)) +} + +/// Escapes a string for a quoted GoogleSQL literal. Newlines, tabs and other +/// control characters are written as escapes so the literal stays valid. +pub fn escape_string(value: &str) -> String { + let mut escaped = String::with_capacity(value.len()); + for character in value.chars() { + match character { + '\\' => escaped.push_str("\\\\"), + '\'' => escaped.push_str("\\'"), + '\n' => escaped.push_str("\\n"), + '\r' => escaped.push_str("\\r"), + '\t' => escaped.push_str("\\t"), + control if (control as u32) < 0x20 || control as u32 == 0x7f => { + escaped.push_str(&format!("\\x{:02x}", control as u32)); + } + control if control.is_control() => { + escaped.push_str(&format!("\\u{:04x}", control as u32)); + } + other => escaped.push(other), + } + } + escaped +} + +/// Canonical upper-case type name for a field (legacy names normalized). +pub fn base_type(field: &FieldSchema) -> String { + canonical_type(&field.field_type) +} + +pub fn canonical_type(name: &str) -> String { + match name.trim().to_ascii_uppercase().as_str() { + "INTEGER" => "INT64".to_string(), + "FLOAT" => "FLOAT64".to_string(), + "BOOLEAN" => "BOOL".to_string(), + "RECORD" => "STRUCT".to_string(), + "BIGDECIMAL" => "BIGNUMERIC".to_string(), + "DECIMAL" => "NUMERIC".to_string(), + other => other.to_string(), + } +} + +pub fn is_repeated(field: &FieldSchema) -> bool { + field.mode.as_deref() == Some("REPEATED") +} + +pub fn is_required(field: &FieldSchema) -> bool { + field.mode.as_deref() == Some("REQUIRED") +} + +/// Type of one element of the field, ignoring `REPEATED`. +pub fn element_type_sql(field: &FieldSchema, ddl: bool) -> String { + let base = base_type(field); + match base.as_str() { + "STRUCT" => { + let fields = field + .fields + .iter() + .map(|nested| { + let name = if ddl { + quote_identifier(&nested.name) + } else { + nested.name.clone() + }; + let mut column = format!("{name} {}", field_type_sql(nested, ddl)); + if ddl && is_required(nested) { + column.push_str(" NOT NULL"); + } + column + }) + .collect::>() + .join(", "); + format!("STRUCT<{fields}>") + } + "RANGE" => { + let element = field + .range_element_type + .as_ref() + .map(|range| canonical_type(&range.element_type)) + .unwrap_or_else(|| "DATE".to_string()); + format!("RANGE<{element}>") + } + "STRING" | "BYTES" => match field.max_length.as_deref() { + Some(length) => format!("{base}({length})"), + None => base, + }, + "NUMERIC" | "BIGNUMERIC" => match (field.precision.as_deref(), field.scale.as_deref()) { + (Some(precision), Some(scale)) => format!("{base}({precision}, {scale})"), + (Some(precision), None) => format!("{base}({precision})"), + _ => base, + }, + _ => base, + } +} + +/// Full column type including `ARRAY<...>` for repeated fields. With `ddl` +/// set, nested names are quoted and nested `NOT NULL` is preserved. +pub fn field_type_sql(field: &FieldSchema, ddl: bool) -> String { + let element = element_type_sql(field, ddl); + if is_repeated(field) { + format!("ARRAY<{element}>") + } else { + element + } +} + +/// Renders a JSON cell value as a GoogleSQL literal of the column type. pub fn value_literal(value: &Value, field: Option<&FieldSchema>) -> Result { if value.is_null() { return Ok("NULL".to_string()); } - let field_type = field - .map(|field| field.field_type.to_ascii_uppercase()) - .unwrap_or_else(|| infer_type(value).to_string()); - - if field.and_then(|field| field.mode.as_deref()) == Some("REPEATED") { + if let Some(field) = field.filter(|field| is_repeated(field)) { let values = normalize_array(value)?; - let mut scalar_field = field.cloned().unwrap_or_else(|| FieldSchema { - name: String::new(), - field_type: "STRING".to_string(), - mode: None, - description: None, - fields: vec![], - default_value_expression: None, - }); - scalar_field.mode = None; + let mut element = field.clone(); + element.mode = None; let literals = values .iter() - .map(|value| value_literal(value, Some(&scalar_field))) + .map(|value| { + if value.is_null() { + Err("BigQuery arrays cannot contain NULL elements".to_string()) + } else { + value_literal(value, Some(&element)) + } + }) .collect::, _>>()?; - return Ok(format!("[{}]", literals.join(", "))); + return Ok(format!( + "ARRAY<{}>[{}]", + element_type_sql(&element, true), + literals.join(", ") + )); } + let field_type = field + .map(base_type) + .unwrap_or_else(|| infer_type(value).to_string()); match field_type.as_str() { - "RECORD" | "STRUCT" => struct_literal(value, field), - "BOOL" | "BOOLEAN" => value - .as_bool() - .or_else(|| value.as_str().and_then(parse_bool)) - .map(|value| if value { "TRUE" } else { "FALSE" }.to_string()) - .ok_or_else(|| "Expected a boolean value".to_string()), - "INT64" | "INTEGER" => validated_number(value, true), - "FLOAT" | "FLOAT64" | "NUMERIC" | "BIGNUMERIC" => validated_number(value, false), - "JSON" => { - let serialized = if let Some(value) = value.as_str() { - serde_json::from_str::(value) - .map_err(|error| format!("Invalid JSON value: {error}"))? - .to_string() + "STRUCT" => struct_literal(value, field), + "ARRAY" => { + let literals = normalize_array(value)? + .iter() + .map(|value| value_literal(value, None)) + .collect::, _>>()?; + Ok(format!("[{}]", literals.join(", "))) + } + "BOOL" => bool_text(value).map(|value| value.to_ascii_uppercase()), + "INT64" => integer_text(value), + "FLOAT64" => { + let text = float_text(value)?; + if matches!(text.as_str(), "NaN" | "Infinity" | "-Infinity") { + Ok(format!("CAST('{text}' AS FLOAT64)")) } else { - value.to_string() - }; - Ok(format!("JSON '{}'", escape_string(&serialized))) + Ok(text) + } } - "BYTES" => { - let value = value - .as_str() - .ok_or_else(|| "Expected base64 text for BYTES".to_string())?; - let encoded = value - .strip_prefix("BLOB:") - .and_then(|rest| rest.splitn(4, ':').nth(3)) - .unwrap_or(value); - Ok(format!("FROM_BASE64('{}')", escape_string(encoded))) + "NUMERIC" | "BIGNUMERIC" => Ok(format!("{field_type} '{}'", decimal_text(value)?)), + "JSON" => Ok(format!("JSON {}", string_literal(&json_text(value)?))), + "BYTES" => Ok(format!( + "FROM_BASE64({})", + string_literal(&bytes_text(value)?) + )), + "DATE" | "TIME" | "DATETIME" | "TIMESTAMP" => Ok(format!( + "{field_type} {}", + string_literal(&temporal_text(&field_type, value)?) + )), + "GEOGRAPHY" => { + let text = scalar_text(value)?; + if text.trim_start().starts_with('{') { + Ok(format!("ST_GEOGFROMGEOJSON({})", string_literal(&text))) + } else { + Ok(format!("ST_GEOGFROMTEXT({})", string_literal(&text))) + } } - "DATE" | "TIME" | "DATETIME" | "TIMESTAMP" | "GEOGRAPHY" => { - let value = scalar_text(value)?; - Ok(format!("{field_type} '{}'", escape_string(&value))) + "INTERVAL" => Ok(format!( + "CAST({} AS INTERVAL)", + string_literal(&scalar_text(value)?) + )), + "RANGE" => { + let element = range_element(field); + Ok(format!( + "RANGE<{element}> {}", + string_literal(&range_text(value)?) + )) } - _ => Ok(format!("'{}'", escape_string(&scalar_text(value)?))), + _ => Ok(string_literal(&scalar_text(value)?)), } } @@ -88,8 +260,8 @@ fn struct_literal(value: &Value, field: Option<&FieldSchema>) -> Result return Err("STRUCT value must be a JSON object".to_string()), }; let fields = field.map(|field| field.fields.as_slice()).unwrap_or(&[]); - let values = if fields.is_empty() { - object + if fields.is_empty() { + let values = object .iter() .map(|(name, value)| { Ok(format!( @@ -98,21 +270,301 @@ fn struct_literal(value: &Value, field: Option<&FieldSchema>) -> Result, String>>()? - } else { - fields - .iter() - .map(|nested| { - let value = object.get(&nested.name).unwrap_or(&Value::Null); - Ok(format!( - "{} AS {}", - value_literal(value, Some(nested))?, - quote_identifier(&nested.name) + .collect::, String>>()?; + return Ok(format!("STRUCT({})", values.join(", "))); + } + if let Some(unknown) = object + .keys() + .find(|key| !fields.iter().any(|field| &field.name == *key)) + { + return Err(format!("STRUCT field '{unknown}' does not exist")); + } + let values = fields + .iter() + .map(|nested| { + value_literal( + object.get(&nested.name).unwrap_or(&Value::Null), + Some(nested), + ) + }) + .collect::, String>>()?; + let mut element = field.cloned().unwrap_or_default(); + element.mode = None; + Ok(format!( + "{}({})", + element_type_sql(&element, true), + values.join(", ") + )) +} + +/// Collects named `@pN` query parameters for one statement. +#[derive(Debug, Default)] +pub struct QueryParams { + parameters: Vec, +} + +impl QueryParams { + pub fn new() -> Self { + Self::default() + } + + pub fn into_vec(self) -> Vec { + self.parameters + } + + pub fn len(&self) -> usize { + self.parameters.len() + } + + pub fn is_empty(&self) -> bool { + self.parameters.is_empty() + } + + fn push(&mut self, parameter_type: Value, parameter_value: Value) -> String { + let name = format!("p{}", self.parameters.len()); + self.parameters.push(json!({ + "name": name, + "parameterType": parameter_type, + "parameterValue": parameter_value + })); + format!("@{name}") + } + + /// Binds a STRING parameter, e.g. an object name in a catalog query. + pub fn string(&mut self, value: &str) -> String { + self.scalar("STRING", value.to_string()) + } + + fn scalar(&mut self, type_name: &str, value: String) -> String { + self.push(json!({"type": type_name}), json!({"value": value})) + } + + /// Returns the SQL expression that yields `value` as the column type. + /// Scalars become typed parameters; BYTES, NUMERIC, JSON, GEOGRAPHY, + /// INTERVAL and RANGE travel as STRING parameters converted in SQL, and + /// ARRAY / STRUCT values become validated typed literals. + pub fn bind(&mut self, value: &Value, field: Option<&FieldSchema>) -> Result { + if value.is_null() { + return Ok("NULL".to_string()); + } + // Arrays and structs are rendered as validated typed literals, which + // keeps nested element types exact. + if let Some(field) = field.filter(|field| is_repeated(field)) { + return value_literal(value, Some(field)); + } + let field_type = field + .map(base_type) + .unwrap_or_else(|| infer_type(value).to_string()); + match field_type.as_str() { + "STRUCT" | "ARRAY" => value_literal(value, field), + "JSON" => { + let parameter = self.scalar("STRING", json_text(value)?); + Ok(format!("PARSE_JSON({parameter})")) + } + "GEOGRAPHY" => { + let text = scalar_text(value)?; + let geojson = text.trim_start().starts_with('{'); + let parameter = self.scalar("STRING", text); + Ok(if geojson { + format!("ST_GEOGFROMGEOJSON({parameter})") + } else { + format!("ST_GEOGFROMTEXT({parameter})") + }) + } + "INTERVAL" => { + let parameter = self.scalar("STRING", scalar_text(value)?); + Ok(format!("CAST({parameter} AS INTERVAL)")) + } + "RANGE" => { + let element = range_element(field); + let parameter = self.scalar("STRING", range_text(value)?); + Ok(format!("CAST({parameter} AS RANGE<{element}>)")) + } + "BYTES" => { + let parameter = self.scalar("STRING", bytes_text(value)?); + Ok(format!("FROM_BASE64({parameter})")) + } + // Decimal text is cast in SQL so no precision is lost on the way. + "NUMERIC" | "BIGNUMERIC" => { + let parameter = self.scalar("STRING", decimal_text(value)?); + Ok(format!("CAST({parameter} AS {field_type})")) + } + other if is_simple_parameter_type(other) => { + let text = scalar_parameter_text(other, value)?; + Ok(self.scalar(other, text)) + } + _ => { + let text = scalar_text(value)?; + Ok(self.scalar("STRING", text)) + } + } + } +} + +fn is_simple_parameter_type(type_name: &str) -> bool { + matches!( + type_name, + "STRING" + | "BYTES" + | "INT64" + | "FLOAT64" + | "NUMERIC" + | "BIGNUMERIC" + | "BOOL" + | "DATE" + | "TIME" + | "DATETIME" + | "TIMESTAMP" + ) +} + +fn scalar_parameter_text(type_name: &str, value: &Value) -> Result { + match type_name { + "BOOL" => bool_text(value), + "INT64" => integer_text(value), + "FLOAT64" => float_text(value), + "NUMERIC" | "BIGNUMERIC" => decimal_text(value), + "BYTES" => bytes_text(value), + "DATE" | "TIME" | "DATETIME" | "TIMESTAMP" => temporal_text(type_name, value), + _ => scalar_text(value), + } +} + +fn range_element(field: Option<&FieldSchema>) -> String { + field + .and_then(|field| field.range_element_type.as_ref()) + .map(|range| canonical_type(&range.element_type)) + .unwrap_or_else(|| "DATE".to_string()) +} + +fn range_text(value: &Value) -> Result { + match value { + Value::Object(object) => { + let bound = |key: &str| match object.get(key) { + None | Some(Value::Null) => Ok("UNBOUNDED".to_string()), + Some(value) => scalar_text(value), + }; + Ok(format!("[{}, {})", bound("start")?, bound("end")?)) + } + _ => { + let text = scalar_text(value)?; + let trimmed = text.trim(); + if trimmed.starts_with('[') && trimmed.ends_with(')') && trimmed.contains(',') { + Ok(trimmed.to_string()) + } else { + Err(format!( + "Invalid RANGE value '{text}': expected '[start, end)' with UNBOUNDED for open bounds" )) - }) - .collect::, String>>()? + } + } + } +} + +fn bool_text(value: &Value) -> Result { + value + .as_bool() + .or_else(|| value.as_str().and_then(parse_bool)) + .map(|value| value.to_string()) + .ok_or_else(|| "Expected a boolean value".to_string()) +} + +fn integer_text(value: &Value) -> Result { + let text = scalar_text(value)?; + text.trim() + .parse::() + .map(|number| number.to_string()) + .map_err(|_| format!("Invalid INT64 value '{text}'")) +} + +fn float_text(value: &Value) -> Result { + let text = scalar_text(value)?; + let trimmed = text.trim(); + match trimmed.to_ascii_lowercase().as_str() { + "nan" => return Ok("NaN".to_string()), + "inf" | "+inf" | "infinity" | "+infinity" => return Ok("Infinity".to_string()), + "-inf" | "-infinity" => return Ok("-Infinity".to_string()), + _ => {} + } + if trimmed.parse::().is_ok_and(f64::is_finite) { + Ok(trimmed.to_string()) + } else { + Err(format!("Invalid FLOAT64 value '{text}'")) + } +} + +/// Validates a NUMERIC/BIGNUMERIC decimal without going through f64. +fn decimal_text(value: &Value) -> Result { + let text = scalar_text(value)?; + let trimmed = text.trim(); + let (mantissa, exponent) = match trimmed.find(['e', 'E']) { + Some(index) => (&trimmed[..index], Some(&trimmed[index + 1..])), + None => (trimmed, None), }; - Ok(format!("STRUCT({})", values.join(", "))) + let unsigned = mantissa.strip_prefix(['+', '-']).unwrap_or(mantissa); + let mut parts = unsigned.splitn(2, '.'); + let integer = parts.next().unwrap_or_default(); + let fraction = parts.next().unwrap_or_default(); + let digits_ok = |part: &str| part.chars().all(|character| character.is_ascii_digit()); + let exponent_ok = exponent.is_none_or(|exponent| { + let digits = exponent.strip_prefix(['+', '-']).unwrap_or(exponent); + !digits.is_empty() && digits_ok(digits) + }); + if (integer.is_empty() && fraction.is_empty()) + || !digits_ok(integer) + || !digits_ok(fraction) + || !exponent_ok + { + return Err(format!("Invalid NUMERIC value '{text}'")); + } + Ok(trimmed.to_string()) +} + +fn json_text(value: &Value) -> Result { + match value.as_str() { + Some(text) => serde_json::from_str::(text) + .map(|parsed| parsed.to_string()) + .map_err(|error| format!("Invalid JSON value: {error}")), + None => Ok(value.to_string()), + } +} + +/// Accepts plain base64 or the Tabularis `BLOB:::` wire +/// format and returns validated base64. +fn bytes_text(value: &Value) -> Result { + let text = value + .as_str() + .ok_or_else(|| "Expected base64 text for BYTES".to_string())?; + let encoded = text + .strip_prefix("BLOB:") + .and_then(|rest| rest.splitn(3, ':').nth(2)) + .unwrap_or(text) + .trim(); + base64::engine::general_purpose::STANDARD + .decode(encoded) + .map_err(|error| format!("Invalid base64 BYTES value: {error}"))?; + Ok(encoded.to_string()) +} + +/// Normalizes the RFC 3339 text Tabularis shows for temporal cells into the +/// canonical GoogleSQL format. +fn temporal_text(type_name: &str, value: &Value) -> Result { + let text = scalar_text(value)?; + let mut normalized = text.trim().to_string(); + if matches!(type_name, "DATETIME" | "TIMESTAMP") + && normalized.len() > 10 + && normalized.as_bytes()[10] == b'T' + { + normalized.replace_range(10..11, " "); + } + if type_name == "TIMESTAMP" { + if let Some(stripped) = normalized.strip_suffix('Z') { + normalized = format!("{stripped}+00:00"); + } + } + if normalized.is_empty() { + return Err(format!("Invalid {type_name} value ''")); + } + Ok(normalized) } fn normalize_array(value: &Value) -> Result, String> { @@ -124,18 +576,6 @@ fn normalize_array(value: &Value) -> Result, String> { } } -fn validated_number(value: &Value, integer: bool) -> Result { - let text = scalar_text(value)?; - let valid = if integer { - text.parse::().is_ok() - } else { - text.parse::().is_ok() && !matches!(text.as_str(), "NaN" | "inf" | "-inf") - }; - valid - .then_some(text.clone()) - .ok_or_else(|| format!("Invalid numeric value '{text}'")) -} - fn scalar_text(value: &Value) -> Result { match value { Value::String(value) => Ok(value.clone()), @@ -158,40 +598,290 @@ fn infer_type(value: &Value) -> &'static str { } fn parse_bool(value: &str) -> Option { - match value.to_ascii_lowercase().as_str() { + match value.trim().to_ascii_lowercase().as_str() { "true" | "1" => Some(true), "false" | "0" => Some(false), _ => None, } } -fn escape_string(value: &str) -> String { - value - .replace('\0', "") - .replace('\\', "\\\\") - .replace('\'', "\\'") -} - #[cfg(test)] mod tests { use serde_json::json; - use super::{qualified, value_literal}; + use super::*; + use crate::bigquery::RangeElementType; + + fn field(name: &str, field_type: &str) -> FieldSchema { + FieldSchema { + name: name.into(), + field_type: field_type.into(), + ..Default::default() + } + } #[test] - fn quotes_qualified_identifiers() { + fn quotes_each_path_part() { assert_eq!( - qualified("sales-eu", "order`items"), - "`sales-eu.order\\`items`" + fully_qualified("my-project-123", "sales.eu", "order`items"), + "`my-project-123`.`sales.eu`.`order\\`items`" ); + assert_eq!(qualified("ds", "t"), "`ds`.`t`"); + assert_eq!(object_path(None, "ds", "t"), "`ds`.`t`"); + assert_eq!(object_path(Some("p-1"), "ds", "t"), "`p-1`.`ds`.`t`"); + assert_eq!( + information_schema("bigquery-public-data", "samples", "TABLES"), + "`bigquery-public-data`.`samples`.INFORMATION_SCHEMA.TABLES" + ); + assert_eq!(quote_identifier("a\\b"), "`a\\\\b`"); } #[test] - fn escapes_strings_and_preserves_numbers() { + fn escapes_quotes_backslashes_and_control_characters() { + assert_eq!(string_literal("O'Reilly"), "'O\\'Reilly'"); + assert_eq!(string_literal("a\\b"), "'a\\\\b'"); + assert_eq!( + string_literal("line1\nline2\r\tend"), + "'line1\\nline2\\r\\tend'" + ); + assert_eq!(string_literal("bell\u{7}nul\u{0}"), "'bell\\x07nul\\x00'"); + assert_eq!(string_literal("c1\u{85}"), "'c1\\u0085'"); + assert_eq!( + string_literal("caffè ☕ 東京 \u{1F600}"), + "'caffè ☕ 東京 \u{1F600}'" + ); + assert_eq!(string_literal("say \"hi\""), "'say \"hi\"'"); + } + + #[test] + fn renders_scalar_literals() { assert_eq!(value_literal(&json!(42), None).unwrap(), "42"); assert_eq!( - value_literal(&json!("O'Reilly"), None).unwrap(), - "'O\\'Reilly'" + value_literal(&json!("42"), Some(&field("n", "INTEGER"))).unwrap(), + "42" + ); + assert!(value_literal(&json!("4x"), Some(&field("n", "INT64"))).is_err()); + assert_eq!( + value_literal(&json!("1.50"), Some(&field("n", "NUMERIC"))).unwrap(), + "NUMERIC '1.50'" + ); + assert_eq!( + value_literal( + &json!("-12345678901234567890.123456789"), + Some(&field("n", "BIGNUMERIC")) + ) + .unwrap(), + "BIGNUMERIC '-12345678901234567890.123456789'" + ); + assert!(value_literal(&json!("1.2.3"), Some(&field("n", "NUMERIC"))).is_err()); + assert!(value_literal(&json!("1; DROP"), Some(&field("n", "NUMERIC"))).is_err()); + assert_eq!( + value_literal(&json!("NaN"), Some(&field("f", "FLOAT64"))).unwrap(), + "CAST('NaN' AS FLOAT64)" + ); + assert_eq!( + value_literal(&json!(1.5), Some(&field("f", "FLOAT"))).unwrap(), + "1.5" + ); + assert_eq!( + value_literal(&json!("true"), Some(&field("b", "BOOLEAN"))).unwrap(), + "TRUE" + ); + assert_eq!( + value_literal(&json!("a\nb"), Some(&field("s", "STRING"))).unwrap(), + "'a\\nb'" + ); + } + + #[test] + fn renders_special_type_literals() { + assert_eq!( + value_literal(&json!("POINT(1 2)"), Some(&field("g", "GEOGRAPHY"))).unwrap(), + "ST_GEOGFROMTEXT('POINT(1 2)')" + ); + assert_eq!( + value_literal(&json!("1-2 3 4:5:6"), Some(&field("i", "INTERVAL"))).unwrap(), + "CAST('1-2 3 4:5:6' AS INTERVAL)" + ); + let mut range = field("r", "RANGE"); + range.range_element_type = Some(RangeElementType { + element_type: "DATE".into(), + }); + assert_eq!( + value_literal(&json!("[2024-01-01, UNBOUNDED)"), Some(&range)).unwrap(), + "RANGE '[2024-01-01, UNBOUNDED)'" + ); + assert_eq!( + value_literal(&json!({"start": "2024-01-01", "end": null}), Some(&range)).unwrap(), + "RANGE '[2024-01-01, UNBOUNDED)'" + ); + assert!(value_literal(&json!("2024-01-01"), Some(&range)).is_err()); + assert_eq!( + value_literal( + &json!("BLOB:3:application/octet-stream:AQID"), + Some(&field("b", "BYTES")) + ) + .unwrap(), + "FROM_BASE64('AQID')" + ); + assert!(value_literal(&json!("not base64!"), Some(&field("b", "BYTES"))).is_err()); + assert_eq!( + value_literal(&json!("{\"a\": \"it's\\n\"}"), Some(&field("j", "JSON"))).unwrap(), + "JSON '{\"a\":\"it\\'s\\\\n\"}'" + ); + assert_eq!( + value_literal( + &json!("2024-01-02T03:04:05.123456Z"), + Some(&field("t", "TIMESTAMP")) + ) + .unwrap(), + "TIMESTAMP '2024-01-02 03:04:05.123456+00:00'" + ); + assert_eq!( + value_literal(&json!("2024-01-02T03:04:05"), Some(&field("t", "DATETIME"))).unwrap(), + "DATETIME '2024-01-02 03:04:05'" + ); + } + + #[test] + fn renders_nested_struct_and_array_literals() { + let mut tags = field("tags", "STRING"); + tags.mode = Some("REPEATED".into()); + let mut record = field("item", "RECORD"); + record.fields = vec![ + field("sku", "STRING"), + field("price", "NUMERIC"), + tags.clone(), + ]; + assert_eq!( + value_literal( + &json!({"sku": "a'b", "price": "9.99", "tags": ["x", "y\n"]}), + Some(&record) + ) + .unwrap(), + "STRUCT<`sku` STRING, `price` NUMERIC, `tags` ARRAY>('a\\'b', NUMERIC '9.99', ARRAY['x', 'y\\n'])" + ); + assert!(value_literal(&json!({"nope": 1}), Some(&record)).is_err()); + let mut items = record.clone(); + items.mode = Some("REPEATED".into()); + assert_eq!( + value_literal(&json!("[]"), Some(&items)).unwrap(), + "ARRAY>>[]" + ); + assert!(value_literal(&json!([null]), Some(&tags)).is_err()); + } + + #[test] + fn renders_column_types() { + let mut name = field("name", "STRING"); + name.max_length = Some("50".into()); + assert_eq!(field_type_sql(&name, true), "STRING(50)"); + let mut amount = field("amount", "NUMERIC"); + amount.precision = Some("10".into()); + amount.scale = Some("2".into()); + assert_eq!(field_type_sql(&amount, true), "NUMERIC(10, 2)"); + let mut required = field("id", "INTEGER"); + required.mode = Some("REQUIRED".into()); + let mut record = field("s", "RECORD"); + record.mode = Some("REPEATED".into()); + record.fields = vec![required, field("my field", "FLOAT")]; + assert_eq!( + field_type_sql(&record, true), + "ARRAY>" + ); + assert_eq!( + field_type_sql(&record, false), + "ARRAY>" + ); + } + + #[test] + fn binds_typed_parameters() { + let mut params = QueryParams::new(); + assert_eq!( + params + .bind(&json!("x\n'y"), Some(&field("s", "STRING"))) + .unwrap(), + "@p0" + ); + assert_eq!( + params + .bind(&json!(7), Some(&field("n", "INTEGER"))) + .unwrap(), + "@p1" + ); + assert_eq!( + params + .bind(&json!("POINT(0 0)"), Some(&field("g", "GEOGRAPHY"))) + .unwrap(), + "ST_GEOGFROMTEXT(@p2)" + ); + assert_eq!( + params + .bind(&json!({"a": 1}), Some(&field("j", "JSON"))) + .unwrap(), + "PARSE_JSON(@p3)" + ); + assert_eq!( + params + .bind(&json!("1-0"), Some(&field("i", "INTERVAL"))) + .unwrap(), + "CAST(@p4 AS INTERVAL)" + ); + let mut range = field("r", "RANGE"); + range.range_element_type = Some(RangeElementType { + element_type: "TIMESTAMP".into(), + }); + assert_eq!( + params + .bind(&json!("[UNBOUNDED, 2024-01-01 00:00:00+00)"), Some(&range)) + .unwrap(), + "CAST(@p5 AS RANGE)" + ); + let mut ints = field("ids", "INT64"); + ints.mode = Some("REPEATED".into()); + assert_eq!( + params.bind(&json!([1, "2"]), Some(&ints)).unwrap(), + "ARRAY[1, 2]" + ); + assert_eq!(params.bind(&Value::Null, Some(&ints)).unwrap(), "NULL"); + assert!(params + .bind(&json!("abc"), Some(&field("n", "INT64"))) + .is_err()); + assert_eq!( + params + .bind(&json!("AQID"), Some(&field("b", "BYTES"))) + .unwrap(), + "FROM_BASE64(@p6)" + ); + assert_eq!( + params + .bind(&json!("1.5"), Some(&field("n", "BIGNUMERIC"))) + .unwrap(), + "CAST(@p7 AS BIGNUMERIC)" + ); + let values = params.into_vec(); + assert_eq!(values.len(), 8); + assert_eq!( + values[0], + json!({"name": "p0", "parameterType": {"type": "STRING"}, "parameterValue": {"value": "x\n'y"}}) + ); + assert_eq!(values[1]["parameterType"]["type"], "INT64"); + assert_eq!(values[1]["parameterValue"]["value"], "7"); + assert_eq!(values[3]["parameterValue"]["value"], "{\"a\":1}"); + assert_eq!(values[6]["parameterValue"]["value"], "AQID"); + assert_eq!(values[7]["parameterType"]["type"], "STRING"); + } + + #[test] + fn struct_parameters_fall_back_to_literals() { + let mut record = field("s", "RECORD"); + record.fields = vec![field("a", "INT64")]; + let mut params = QueryParams::new(); + assert_eq!( + params.bind(&json!({"a": 1}), Some(&record)).unwrap(), + "STRUCT<`a` INT64>(1)" ); + assert!(params.is_empty()); } } diff --git a/src/statements.rs b/src/statements.rs new file mode 100644 index 0000000..942b8c7 --- /dev/null +++ b/src/statements.rs @@ -0,0 +1,456 @@ +//! Lightweight GoogleSQL statement classification used by the query runtime. +//! It never rewrites SQL; it only decides how a statement is executed. + +/// Split a script into statements on top-level semicolons, skipping quoted +/// strings, quoted identifiers and comments. +pub fn split(sql: &str) -> Vec { + let bytes = sql.as_bytes(); + let mut statements = Vec::new(); + let mut start = 0; + let mut index = 0; + while index < bytes.len() { + if let Some(next) = skip_quoted_or_comment(bytes, index) { + index = next; + continue; + } + if bytes[index] == b';' { + push_statement(&mut statements, &sql[start..index]); + start = index + 1; + } + index += 1; + } + if start < sql.len() { + push_statement(&mut statements, &sql[start.min(sql.len())..]); + } + statements +} + +/// When `index` starts a quoted string, quoted identifier or comment, the +/// index just past it. +fn skip_quoted_or_comment(bytes: &[u8], mut index: usize) -> Option { + match bytes[index] { + b'\'' | b'"' | b'`' => { + let quote = bytes[index]; + let triple = index + 2 < bytes.len() + && bytes[index + 1] == quote + && bytes[index + 2] == quote + && quote != b'`'; + index += if triple { 3 } else { 1 }; + while index < bytes.len() { + if bytes[index] == b'\\' { + index += 2; + continue; + } + if bytes[index] == quote { + if !triple { + break; + } + if index + 2 < bytes.len() + && bytes[index + 1] == quote + && bytes[index + 2] == quote + { + index += 2; + break; + } + } + index += 1; + } + Some(index + 1) + } + b'-' if bytes.get(index + 1) == Some(&b'-') => Some(line_end(bytes, index)), + b'#' => Some(line_end(bytes, index)), + b'/' if bytes.get(index + 1) == Some(&b'*') => { + index += 2; + while index + 1 < bytes.len() && !(bytes[index] == b'*' && bytes[index + 1] == b'/') { + index += 1; + } + Some(index + 2) + } + _ => None, + } +} + +fn line_end(bytes: &[u8], mut index: usize) -> usize { + while index < bytes.len() && bytes[index] != b'\n' { + index += 1; + } + index +} + +/// Top-level tokens of a script: keywords and identifiers (upper-cased), +/// statement separators and anything else. +#[derive(Debug, Clone, PartialEq, Eq)] +enum Token { + Word(String), + Semicolon, + Open, + Close, + Other, +} + +fn tokens(sql: &str) -> Vec { + let bytes = sql.as_bytes(); + let mut tokens = Vec::new(); + let mut index = 0; + while index < bytes.len() { + if let Some(next) = skip_quoted_or_comment(bytes, index) { + // Comments separate words; strings and quoted names are values. + if matches!(bytes[index], b'\'' | b'"' | b'`') { + tokens.push(Token::Other); + } + index = next; + continue; + } + let byte = bytes[index]; + if byte.is_ascii_alphanumeric() || byte == b'_' { + let start = index; + while index < bytes.len() + && (bytes[index].is_ascii_alphanumeric() || bytes[index] == b'_') + { + index += 1; + } + tokens.push(Token::Word(sql[start..index].to_ascii_uppercase())); + continue; + } + match byte { + b';' => tokens.push(Token::Semicolon), + b'(' => tokens.push(Token::Open), + b')' => tokens.push(Token::Close), + byte if byte.is_ascii_whitespace() => {} + _ => tokens.push(Token::Other), + } + index += 1; + } + tokens +} + +/// BEGIN / COMMIT / ROLLBACK / SAVEPOINT statements of a script in order, +/// including those nested in scripting blocks (`IF ... THEN COMMIT; END IF`). +/// Statement starts are tracked through block depth: after `;`, `BEGIN`, +/// `THEN`, `ELSE`, `DO`, `LOOP` and `REPEAT`, but not inside parentheses or +/// a `CASE` expression. +fn transaction_events(sql: &str) -> Vec { + let tokens = tokens(sql); + let word = |index: usize| match tokens.get(index) { + Some(Token::Word(word)) => Some(word.as_str()), + _ => None, + }; + let mut events = Vec::new(); + // One entry per open block: true for a CASE expression, false for a + // scripting block (BEGIN, IF, LOOP, WHILE, REPEAT, FOR, CASE statement). + let mut blocks: Vec = Vec::new(); + let mut parens = 0usize; + let mut at_start = true; + let mut index = 0; + while index < tokens.len() { + let token = &tokens[index]; + index += 1; + match token { + Token::Open => { + parens += 1; + at_start = false; + } + Token::Close => parens = parens.saturating_sub(1), + Token::Semicolon if parens == 0 => at_start = true, + Token::Word(_) if parens > 0 => {} + Token::Word(word_here) if at_start => { + at_start = false; + match word_here.as_str() { + "COMMIT" | "ROLLBACK" => { + if word(index) == Some("TO") { + events.push(TransactionControl::Savepoint); + } else { + events.push(TransactionControl::End); + if word(index) == Some("TRANSACTION") { + index += 1; + } + } + } + "SAVEPOINT" | "RELEASE" => events.push(TransactionControl::Savepoint), + "BEGIN" => match tokens.get(index) { + None | Some(Token::Semicolon) => events.push(TransactionControl::Begin), + Some(Token::Word(next)) if next == "TRANSACTION" => { + events.push(TransactionControl::Begin); + index += 1; + } + _ => { + blocks.push(false); + at_start = true; + } + }, + "IF" | "WHILE" | "FOR" | "CASE" => blocks.push(false), + "LOOP" | "REPEAT" => { + blocks.push(false); + at_start = true; + } + "ELSE" => at_start = true, + "END" => { + blocks.pop(); + if matches!( + word(index), + Some("IF" | "LOOP" | "WHILE" | "REPEAT" | "FOR" | "CASE") + ) { + index += 1; + } + } + _ => {} + } + } + Token::Word(word_here) => { + let in_expression = blocks.last() == Some(&true); + match word_here.as_str() { + "CASE" => blocks.push(true), + "END" => { + blocks.pop(); + } + "THEN" | "ELSE" | "DO" if !in_expression => at_start = true, + _ => {} + } + } + _ => at_start = false, + } + } + events +} + +fn push_statement(statements: &mut Vec, statement: &str) { + let statement = strip_comments(statement); + if !statement.is_empty() { + statements.push(statement); + } +} + +/// Remove leading comments and surrounding whitespace. +fn strip_comments(statement: &str) -> String { + let mut rest = statement.trim(); + loop { + if rest.starts_with("--") || rest.starts_with('#') { + rest = rest + .split_once('\n') + .map(|(_, tail)| tail) + .unwrap_or("") + .trim(); + } else if rest.starts_with("/*") { + rest = rest + .split_once("*/") + .map(|(_, tail)| tail) + .unwrap_or("") + .trim(); + } else { + return rest.to_string(); + } + } +} + +fn words(statement: &str) -> Vec { + strip_comments(statement) + .split_whitespace() + .take(4) + .map(|word| word.trim_end_matches(';').to_ascii_uppercase()) + .collect() +} + +/// Effect of a statement on an explicit session transaction. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TransactionControl { + Begin, + End, + Savepoint, + None, +} + +pub fn transaction_control(statement: &str) -> TransactionControl { + let words = words(statement); + let words: Vec<&str> = words.iter().map(String::as_str).collect(); + match words.as_slice() { + ["SAVEPOINT", ..] | ["RELEASE", "SAVEPOINT", ..] | ["ROLLBACK", "TO", ..] => { + TransactionControl::Savepoint + } + ["BEGIN"] | ["BEGIN", "TRANSACTION"] => TransactionControl::Begin, + ["COMMIT"] | ["COMMIT", "TRANSACTION"] | ["ROLLBACK"] | ["ROLLBACK", "TRANSACTION"] => { + TransactionControl::End + } + _ => TransactionControl::None, + } +} + +/// Transaction state after running `sql` successfully from `initial`. +pub fn transaction_state_after(sql: &str, initial: bool) -> bool { + transaction_events(sql) + .into_iter() + .fold(initial, |state, event| match event { + TransactionControl::Begin => true, + TransactionControl::End => false, + TransactionControl::Savepoint | TransactionControl::None => state, + }) +} + +/// Whether the SQL contains a COMMIT or ROLLBACK, which ends a transaction +/// even when the statement fails. +pub fn ends_transaction(sql: &str) -> bool { + transaction_events(sql).contains(&TransactionControl::End) +} + +/// Whether BigQuery may run this SQL as a multi-statement script. +pub fn may_be_script(sql: &str) -> bool { + let statements = split(sql); + statements.len() > 1 + || statements.first().is_some_and(|statement| { + transaction_control(statement) == TransactionControl::None + && matches!( + words(statement).first().map(String::as_str), + Some( + "DECLARE" + | "SET" + | "BEGIN" + | "IF" + | "LOOP" + | "WHILE" + | "FOR" + | "REPEAT" + | "CALL" + | "EXECUTE" + ) + ) + }) +} + +/// Whether running the SQL again has no side effects, so an expired cached +/// result may be recomputed for a later page. +pub fn is_read_only_query(sql: &str) -> bool { + let statements = split(sql); + statements.len() == 1 + && matches!( + words(statements[0].trim_start_matches(|c: char| c == '(' || c.is_whitespace())) + .first() + .map(String::as_str), + Some("SELECT" | "WITH") + ) +} + +/// Whether a statement creates session state that later statements need. +pub fn needs_session(statement: &str) -> bool { + let words = words(statement); + let words: Vec<&str> = words.iter().map(String::as_str).collect(); + matches!( + words.as_slice(), + ["DECLARE", ..] + | ["SET", ..] + | ["CREATE", "TEMP", ..] + | ["CREATE", "TEMPORARY", ..] + | ["CREATE", "OR", "REPLACE", "TEMP"] + | ["CREATE", "OR", "REPLACE", "TEMPORARY"] + ) || transaction_control(statement) != TransactionControl::None +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn splits_on_top_level_semicolons_only() { + let statements = + split("SELECT ';' AS a; -- c;\nSELECT `x;y` FROM t /* ; */; SELECT \"\"\"a;b\"\"\";"); + assert_eq!( + statements, + vec![ + "SELECT ';' AS a", + "SELECT `x;y` FROM t /* ; */", + "SELECT \"\"\"a;b\"\"\"" + ] + ); + assert_eq!(split(" ;; "), Vec::::new()); + } + + #[test] + fn tracks_transactions() { + assert_eq!(transaction_control("BEGIN"), TransactionControl::Begin); + assert_eq!( + transaction_control("begin transaction;"), + TransactionControl::Begin + ); + assert_eq!(transaction_control("COMMIT"), TransactionControl::End); + assert_eq!( + transaction_control("ROLLBACK TRANSACTION"), + TransactionControl::End + ); + assert_eq!( + transaction_control("ROLLBACK TO SAVEPOINT tabularis_read"), + TransactionControl::Savepoint + ); + assert_eq!( + transaction_control("SAVEPOINT tabularis_read"), + TransactionControl::Savepoint + ); + assert_eq!( + transaction_control("BEGIN SELECT 1; END"), + TransactionControl::None + ); + assert!(transaction_state_after( + "BEGIN TRANSACTION; UPDATE t SET a = 1", + false + )); + assert!(!transaction_state_after("UPDATE t SET a = 1; COMMIT", true)); + assert!(ends_transaction("-- done\nROLLBACK")); + } + + #[test] + fn tracks_transactions_inside_scripting_blocks() { + assert!(!transaction_state_after( + "IF (SELECT COUNT(*) FROM t) > 0 THEN COMMIT; END IF;", + true + )); + assert!(ends_transaction( + "IF x THEN UPDATE t SET a = 1; ELSE ROLLBACK; END IF" + )); + assert!(ends_transaction( + "BEGIN\n UPDATE t SET a = 1;\nEXCEPTION WHEN ERROR THEN\n ROLLBACK TRANSACTION;\nEND" + )); + assert!(transaction_state_after( + "WHILE x < 3 DO BEGIN TRANSACTION; SET x = x + 1; END WHILE", + false + )); + assert!(transaction_state_after( + "LOOP BEGIN TRANSACTION; LEAVE; END LOOP", + false + )); + // A block BEGIN is not a transaction, and CASE expressions do not + // start statements. + assert!(!transaction_state_after("BEGIN SELECT 1; END", false)); + assert!(!ends_transaction( + "SELECT CASE WHEN a THEN rollback ELSE commit END FROM t" + )); + assert!(!ends_transaction( + "SELECT 'COMMIT'; SELECT `rollback` FROM t" + )); + assert!(!ends_transaction("ROLLBACK TO SAVEPOINT s")); + assert!(transaction_state_after("BEGIN; UPDATE t SET a = 1", false)); + } + + #[test] + fn classifies_scripts_and_rerunnable_reads() { + assert!(may_be_script("DECLARE x INT64")); + assert!(may_be_script("SELECT 1; SELECT 2")); + assert!(!may_be_script("SELECT 1;")); + assert!(!may_be_script("BEGIN TRANSACTION")); + assert!(is_read_only_query( + " -- note\nWITH a AS (SELECT 1) SELECT * FROM a" + )); + assert!(!is_read_only_query("DELETE FROM t WHERE true")); + assert!(!is_read_only_query("SELECT 1; DELETE FROM t WHERE true")); + assert!(is_read_only_query("(SELECT 1) UNION ALL (SELECT 2)")); + assert!(is_read_only_query( + "( (\n WITH a AS (SELECT 1) SELECT * FROM a))" + )); + assert!(!is_read_only_query("(DELETE FROM t WHERE true)")); + } + + #[test] + fn detects_statements_that_need_a_session() { + assert!(needs_session("DECLARE x INT64 DEFAULT 1")); + assert!(needs_session("create temp table t AS SELECT 1")); + assert!(needs_session("BEGIN TRANSACTION")); + assert!(!needs_session("SELECT 1")); + assert!(!needs_session("CREATE TABLE d.t (a INT64)")); + } +} diff --git a/tests/conformance.rs b/tests/conformance.rs new file mode 100644 index 0000000..051ac90 --- /dev/null +++ b/tests/conformance.rs @@ -0,0 +1,225 @@ +//! Query, session and explain conformance: RPC results must deserialize into +//! the structures Tabularis uses (copied from `src-tauri/src/models.rs` and +//! the response readers in `src-tauri/src/plugins/driver.rs`). Metadata, DDL +//! and routine shapes are covered by `conformance_metadata.rs`. + +mod support; + +use serde::Deserialize; +use serde_json::{json, Value}; +use support::{call, complete, mock_connection}; +use wiremock::matchers::{body_partial_json, method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +#[derive(Debug, Deserialize)] +#[allow(dead_code)] +struct Pagination { + page: u32, + page_size: u32, + total_rows: Option, + has_more: bool, +} + +#[derive(Debug, Deserialize)] +#[allow(dead_code)] +struct QueryResult { + columns: Vec, + rows: Vec>, + affected_rows: u64, + #[serde(default)] + truncated: bool, + pagination: Option, + #[serde(default)] + additional_results: Option>, + #[serde(default)] + server_time_ms: Option, + #[serde(default)] + fetch_time_ms: Option, +} + +#[derive(Debug, Deserialize)] +#[allow(dead_code)] +struct BatchStatementResult { + result: Option, + error: Option, + execution_time_ms: Option, +} + +#[derive(Debug, Deserialize)] +#[allow(dead_code)] +struct RawExplainOutput { + engine: String, + format: String, + payload: String, + original_query: String, +} + +/// `parse_query_response` in the host. +fn parse_query_response(value: Value) -> (QueryResult, bool) { + if let Some(object) = value.as_object() { + if let (Some(result), Some(flag)) = (object.get("result"), object.get("in_transaction")) { + return ( + serde_json::from_value(result.clone()).unwrap(), + flag.as_bool().unwrap_or(false), + ); + } + } + (serde_json::from_value(value).unwrap(), false) +} + +/// `parse_session_error` in the host. +fn parse_session_error(value: &Value) -> Option<(String, bool)> { + let object = value.as_object()?; + let error = object.get("error")?.as_str()?; + let in_transaction = object.get("in_transaction")?.as_bool().unwrap_or(false); + Some((error.to_string(), in_transaction)) +} + +/// `parse_batch_response` in the host. +fn parse_batch_response(value: Value) -> (Vec, bool) { + if let Some(object) = value.as_object() { + if let Some(results) = object.get("results") { + return ( + serde_json::from_value(results.clone()).unwrap(), + object + .get("in_transaction") + .and_then(Value::as_bool) + .unwrap_or(false), + ); + } + } + (serde_json::from_value(value).unwrap(), false) +} + +async fn server() -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/bigquery/v2/projects/p/queries")) + .and(body_partial_json(json!({"query": "SELECT nope"}))) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": {"code": 400, "message": "Unrecognized name: nope", "errors": [{"reason": "invalidQuery", "message": "Unrecognized name: nope"}]} + }))) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/bigquery/v2/projects/p/queries")) + .respond_with(ResponseTemplate::new(200).set_body_json({ + let mut response = complete("c", &[1, 2], 2, None); + response["sessionInfo"] = json!({"sessionId": "s"}); + response + })) + .with_priority(5) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/bigquery/v2/projects/p/jobs")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "statistics": {"query": {"totalBytesProcessed": "10"}} + }))) + .mount(&server) + .await; + server +} + +#[tokio::test] +async fn execute_query_without_session_is_a_bare_query_result() { + let server = server().await; + let response = call( + "execute_query", + json!({"params": mock_connection(&server), "query": "SELECT 1", "limit": 100, "page": 1, "schema": null}), + ) + .await; + // `execute_query` without a session deserializes the result directly. + let result: QueryResult = serde_json::from_value(response["result"].clone()).unwrap(); + assert_eq!(result.rows.len(), 2); + assert_eq!(result.pagination.unwrap().page_size, 100); + + // The host's count query sends no limit. + let response = call( + "execute_query", + json!({"params": mock_connection(&server), "query": "SELECT COUNT(*) FROM (SELECT 1) as count_wrapper", "limit": null, "page": 1, "schema": null}), + ) + .await; + let _: QueryResult = serde_json::from_value(response["result"].clone()).unwrap(); +} + +#[tokio::test] +async fn execute_query_in_session_uses_the_wrapped_shapes() { + let server = server().await; + let request = |sql: &str| json!({"params": mock_connection(&server), "query": sql, "limit": 10, "page": 1, "schema": null, "session_id": "conformance-tab"}); + let response = call("execute_query", request("BEGIN")).await; + let value = response["result"].clone(); + assert!(parse_session_error(&value).is_none()); + let (result, in_transaction) = parse_query_response(value); + assert_eq!(result.columns, vec!["n"]); + assert!(in_transaction); + + let response = call("execute_query", request("SELECT nope")).await; + let (error, in_transaction) = parse_session_error(&response["result"]).unwrap(); + assert!(error.contains("Unrecognized name: nope")); + assert!(in_transaction); + + let response = call("release_session", json!({"session_id": "conformance-tab"})).await; + assert_eq!(response, json!({"jsonrpc": "2.0", "result": null, "id": 1})); +} + +#[tokio::test] +async fn execute_query_batch_uses_bare_and_wrapped_shapes() { + let server = server().await; + let response = call( + "execute_query_batch", + json!({"params": mock_connection(&server), "queries": ["SELECT 1", "SELECT nope"], "limit": 10, "page": 1, "schema": null}), + ) + .await; + // Without a session the host reads a bare array. + let results: Vec = + serde_json::from_value(response["result"].clone()).unwrap(); + assert!(results[0].result.is_some() && results[0].error.is_none()); + assert!(results[1].result.is_none() && results[1].error.is_some()); + + let response = call( + "execute_query_batch", + json!({"params": mock_connection(&server), "queries": ["BEGIN TRANSACTION", "SELECT 1"], "limit": 10, "page": 1, "schema": null, "session_id": "batch-tab"}), + ) + .await; + let (results, in_transaction) = parse_batch_response(response["result"].clone()); + assert_eq!(results.len(), 2); + assert!(in_transaction); + call("release_session", json!({"session_id": "batch-tab"})).await; +} + +#[tokio::test] +async fn explain_query_returns_the_raw_shape() { + let server = server().await; + let response = call( + "explain_query", + json!({"params": mock_connection(&server), "query": "SELECT 1", "analyze": false, "schema": null}), + ) + .await; + let raw: RawExplainOutput = serde_json::from_value(response["result"].clone()).unwrap(); + assert_eq!(raw.engine, "bigquery"); + assert_eq!(raw.format, "bigquery-job-json"); + assert_eq!(raw.original_query, "SELECT 1"); + let payload: Value = serde_json::from_str(&raw.payload).unwrap(); + assert_eq!(payload["version"], 1); + + let manifest: Value = serde_json::from_str(include_str!("../.tabularium")).unwrap(); + let parsers = manifest["explain_parsers"].as_array().unwrap(); + assert!(parsers + .iter() + .any(|parser| parser["engine"] == raw.engine && parser["format"] == raw.format)); +} + +#[tokio::test] +async fn errors_use_json_rpc_error_objects() { + let server = server().await; + let response = call( + "execute_query", + json!({"params": mock_connection(&server), "query": "SELECT nope", "limit": 10, "page": 1}), + ) + .await; + assert_eq!(response["error"]["code"], -32603); + assert!(response["error"]["message"].is_string()); + assert_eq!(response["id"], 1); +} diff --git a/tests/conformance_metadata.rs b/tests/conformance_metadata.rs new file mode 100644 index 0000000..827b594 --- /dev/null +++ b/tests/conformance_metadata.rs @@ -0,0 +1,434 @@ +//! Metadata, DDL and routine result conformance against the Tabularis host +//! models. Query, session and explain result shapes are covered by +//! `conformance.rs`. +//! +//! The model definitions below are copied from +//! `tabularis/src-tauri/src/models.rs` at host commit +//! `03bee207fe62a64d482dbbc18e4b688203c7fec3` (doc comments trimmed). The +//! host deserializes plugin results with serde into these exact types, so a +//! result that fails here fails in Tabularis. Inputs are recorded BigQuery +//! REST resources in `tests/fixtures/rest/`. + +#![allow(dead_code)] + +mod support; + +use std::collections::HashMap; + +use bigquery_plugin::handlers::metadata::{ + columns_from_resource, foreign_keys_from_resource, indexes_from_resource, merge_projects, +}; +use bigquery_plugin::routines::{ + create_template, routine_call_sql, routine_ddl, routine_info, routine_parameters, CallArg, +}; +use bigquery_plugin::schema_ddl::table_ddl; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; +use serde_json::{json, Value}; +use support::call as rpc; + +// BEGIN host model definitions. + +#[derive(Debug, Serialize, Deserialize)] +pub struct TableInfo { + pub name: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub comment: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct TableColumn { + pub name: String, + pub data_type: String, + pub is_pk: bool, + pub is_nullable: bool, + pub is_auto_increment: bool, + #[serde(default)] + pub is_generated: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub default_value: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub character_maximum_length: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub numeric_precision: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub numeric_scale: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub comment: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ForeignKey { + pub name: String, + pub column_name: String, + pub ref_table: String, + pub ref_column: String, + pub on_delete: Option, + pub on_update: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub seq_in_fk: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct Index { + pub name: String, + pub column_name: String, + pub is_unique: bool, + pub is_primary: bool, + pub seq_in_index: i32, + #[serde(default)] + pub is_expression: bool, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct TableSchema { + pub name: String, + pub columns: Vec, + pub foreign_keys: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RoutineInfo { + pub name: String, + pub routine_type: String, + pub definition: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub identity_args: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RoutineParameter { + pub name: String, + pub data_type: String, + pub mode: String, + pub ordinal_position: i32, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct RoutineCallArg { + pub name: String, + pub mode: String, + #[serde(default)] + pub value: Option, + #[serde(default)] + pub is_raw: bool, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct ColumnDefinition { + pub name: String, + pub data_type: String, + pub is_nullable: bool, + pub is_pk: bool, + pub is_auto_increment: bool, + pub default_value: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub comment: Option, +} + +// END host model definitions. + +fn fixture(name: &str) -> Value { + let path = format!("{}/tests/fixtures/rest/{name}", env!("CARGO_MANIFEST_DIR")); + serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap() +} + +fn host(value: Value) -> T { + serde_json::from_value(value.clone()) + .unwrap_or_else(|error| panic!("host model rejected {value}: {error}")) +} + +fn routine(name: &str) -> Value { + fixture("routines.json") + .as_array() + .unwrap() + .iter() + .find(|routine| routine["routineReference"]["routineId"] == name) + .cloned() + .unwrap() +} + +#[test] +fn columns_match_host_table_column() { + let columns: Vec = host(json!(columns_from_resource(&fixture( + "table_orders.json" + )) + .unwrap())); + let by_name = columns + .iter() + .map(|column| (column.name.as_str(), column)) + .collect::>(); + assert_eq!(columns.len(), 10); + assert!(by_name["order_id"].is_pk && by_name["line_id"].is_pk); + assert!(!by_name["order_id"].is_nullable); + assert_eq!(by_name["order_id"].comment.as_deref(), Some("Order key")); + assert_eq!(by_name["status"].data_type, "STRING(16)"); + assert_eq!(by_name["status"].character_maximum_length, Some(16)); + assert_eq!(by_name["status"].default_value.as_deref(), Some("'new'")); + assert_eq!(by_name["amount"].numeric_precision, Some(12)); + assert_eq!(by_name["amount"].numeric_scale, Some(2)); + assert_eq!(by_name["validity"].data_type, "RANGE"); + assert_eq!( + by_name["items"].data_type, + "ARRAY>" + ); +} + +#[test] +fn foreign_keys_and_indexes_match_host_models() { + let resource = fixture("table_orders.json"); + let keys: Vec = host(json!(foreign_keys_from_resource(&resource, "sales"))); + assert_eq!(keys.len(), 3); + assert_eq!(keys[0].name, "fk_customer"); + assert_eq!(keys[0].ref_table, "customers"); + assert_eq!(keys[1].name, "fk_order_lines"); + assert_eq!( + keys[1..] + .iter() + .map(|key| key.seq_in_fk) + .collect::>(), + vec![Some(1), Some(2)] + ); + let indexes: Vec = host(json!(indexes_from_resource(&resource))); + assert_eq!(indexes.len(), 2); + assert!(indexes + .iter() + .all(|index| index.is_primary && index.name == "PRIMARY")); + assert_eq!(indexes[1].seq_in_index, 2); +} + +#[test] +fn schema_snapshot_rows_match_host_table_schema() { + let resource = fixture("table_orders.json"); + let snapshot: TableSchema = host(json!({ + "name": "orders", + "columns": columns_from_resource(&resource).unwrap(), + "foreign_keys": foreign_keys_from_resource(&resource, "sales") + })); + assert_eq!(snapshot.columns.len(), 10); +} + +#[test] +fn routines_match_host_routine_info() { + let routines = fixture("routines.json") + .as_array() + .unwrap() + .iter() + .filter_map(routine_info) + .collect::>(); + let routines: Vec = host(json!(routines)); + let types = routines + .iter() + .map(|routine| (routine.name.as_str(), routine.routine_type.as_str())) + .collect::>(); + assert_eq!(types["close_orders"], "PROCEDURE"); + assert_eq!(types["orders_since"], "FUNCTION"); + assert_eq!(types["weighted_avg"], "FUNCTION"); + assert!(routines + .iter() + .all(|routine| routine.identity_args.is_none())); +} + +#[test] +fn routine_parameters_match_host_shape() { + for resource in fixture("routines.json").as_array().unwrap() { + let parameters: Vec = host(json!(routine_parameters(resource))); + let arguments = resource["arguments"].as_array().map_or(0, Vec::len); + let calls = parameters + .iter() + .filter(|parameter| parameter.ordinal_position >= 1) + .collect::>(); + assert_eq!(calls.len(), arguments); + for (index, parameter) in calls.iter().enumerate() { + assert_eq!(parameter.ordinal_position, index as i32 + 1); + assert!(!parameter.name.is_empty()); + } + for parameter in parameters + .iter() + .filter(|parameter| parameter.ordinal_position == 0) + { + assert_eq!(parameter.name, ""); + assert_eq!(parameter.mode, "OUT"); + } + } + let tvf: Vec = host(json!(routine_parameters(&routine("orders_since")))); + assert_eq!( + tvf[0].data_type, + "TABLE<`order_id` INT64, `items` ARRAY>>" + ); + let procedure: Vec = + host(json!(routine_parameters(&routine("close_orders")))); + assert_eq!( + procedure + .iter() + .map(|parameter| parameter.mode.as_str()) + .collect::>(), + vec!["IN", "OUT", "INOUT"] + ); +} + +#[test] +fn routine_scripts_are_fully_qualified() { + let path = "`my-project-1`.`sales`.`close_orders`"; + let ddl = routine_ddl(&routine("close_orders"), path, true).unwrap(); + assert!(ddl.starts_with( + "CREATE OR REPLACE PROCEDURE `my-project-1`.`sales`.`close_orders`(IN `before` DATE, OUT `closed` INT64, INOUT `note` STRING)\nBEGIN" + )); + let host_args: Vec = host(json!([ + {"name": "before", "mode": "IN", "value": "2024-01-01", "is_raw": false}, + {"name": "closed", "mode": "OUT", "value": null, "is_raw": false}, + {"name": "note", "mode": "INOUT", "value": "it's done", "is_raw": false} + ])); + let args: Vec = serde_json::from_value(json!(host_args)).unwrap(); + assert_eq!( + routine_call_sql(&routine("close_orders"), path, &args).unwrap(), + "DECLARE tabularis_arg_2 INT64;\nDECLARE tabularis_arg_3 STRING DEFAULT 'it\\'s done';\nCALL `my-project-1`.`sales`.`close_orders`(DATE '2024-01-01', tabularis_arg_2, tabularis_arg_3);\nSELECT tabularis_arg_2 AS `closed`, tabularis_arg_3 AS `note`;" + ); + let js = routine_ddl( + &routine("upper_js"), + "`my-project-1`.`sales`.`upper_js`", + true, + ) + .unwrap(); + assert!(js.contains("LANGUAGE js") && js.ends_with("AS r\"\"\"return s.toUpperCase();\"\"\"")); + let inferred = routine_ddl(&routine("inferred"), "`p`.`d`.`inferred`", true).unwrap(); + assert!( + inferred.starts_with("CREATE OR REPLACE FUNCTION `p`.`d`.`inferred`(`x` ANY TYPE)\nAS (") + ); + assert!( + create_template("FUNCTION", Some("my-project-1"), Some("sales")) + .starts_with("CREATE OR REPLACE FUNCTION `my-project-1`.`sales`.") + ); +} + +#[test] +fn generated_table_ddl_preserves_layout() { + let ddl = table_ddl( + &fixture("table_orders.json"), + "my-project-1", + "sales", + "orders", + ) + .unwrap(); + assert!(ddl.starts_with("CREATE TABLE `my-project-1`.`sales`.`orders` (")); + assert!(ddl.contains("PRIMARY KEY (`order_id`, `line_id`) NOT ENFORCED")); + assert!(ddl.contains( + "FOREIGN KEY (`order_id`, `line_id`) REFERENCES `my-project-1`.`sales`.`order_lines`(`order_id`, `line_id`) NOT ENFORCED" + )); + assert!(ddl.contains("\nPARTITION BY DATE(`placed_at`)\nCLUSTER BY `customer_id`")); + assert!(ddl.contains("OPTIONS(description = 'Customer orders')")); +} + +#[test] +fn databases_are_a_string_list() { + let databases: Vec = host(json!(merge_projects( + "my-project-1", + "my-project-1", + vec!["other-project".into()], + vec!["bigquery-public-data".into()] + ))); + assert_eq!( + databases, + vec!["bigquery-public-data", "my-project-1", "other-project"] + ); +} + +#[tokio::test] +async fn sql_builders_return_host_shapes() { + let column = ColumnDefinition { + name: "note".into(), + data_type: "STRING".into(), + is_nullable: true, + is_pk: false, + is_auto_increment: false, + default_value: None, + comment: None, + }; + let response = rpc( + "get_create_table_sql", + json!({"table_name": "events", "columns": [column], "schema": "analytics"}), + ) + .await; + let statements: Vec = host(response["result"].clone()); + assert_eq!( + statements, + vec!["CREATE TABLE `analytics`.`events` (\n `note` STRING\n)"] + ); + + let response = rpc( + "get_create_table_sql", + json!({ + "table_name": "events", + "schema": "analytics", + "columns": [column], + "params": {"driver": "bigquery", "database": "my-project-1"} + }), + ) + .await; + let statements: Vec = host(response["result"].clone()); + assert_eq!( + statements, + vec!["CREATE TABLE `my-project-1`.`analytics`.`events` (\n `note` STRING\n)"] + ); + + let response = rpc( + "get_add_column_sql", + json!({"table": "events", "schema": "analytics", "column": column}), + ) + .await; + let statements: Vec = host(response["result"].clone()); + assert_eq!( + statements, + vec!["ALTER TABLE `analytics`.`events` ADD COLUMN `note` STRING"] + ); + + let mut required = column.clone(); + required.is_nullable = false; + let response = rpc( + "get_add_column_sql", + json!({"table": "events", "schema": "analytics", "column": required}), + ) + .await; + assert!(response["error"]["message"] + .as_str() + .unwrap() + .contains("NOT NULL")); + + let mut widened = column.clone(); + widened.data_type = "INT64".into(); + let response = rpc( + "get_alter_column_sql", + json!({"table": "events", "schema": "analytics", "old_column": column, "new_column": widened}), + ) + .await; + assert!(response["error"]["message"] + .as_str() + .unwrap() + .contains("cannot change a column from STRING to INT64")); + + let response = rpc( + "routine_create_template", + json!({"routine_type": "PROCEDURE", "schema": "sales"}), + ) + .await; + let template: String = host(response["result"].clone()); + assert!(template.starts_with("CREATE OR REPLACE PROCEDURE `sales`.`my_procedure`(")); + + let response = rpc("get_create_index_sql", json!({"table": "t", "index_name": "i", "columns": ["a"], "is_unique": false, "schema": "d"})).await; + assert_eq!(response["error"]["code"], -32601); +} + +#[tokio::test] +async fn routine_management_methods_are_routed() { + for method in [ + "get_routine_edit_script", + "build_routine_call_sql", + "drop_routine", + "get_table_ddl", + ] { + let response = rpc(method, json!({"params": {"driver": "bigquery"}})).await; + assert_ne!( + response["error"]["code"], -32601, + "{method} must be implemented" + ); + } +} diff --git a/tests/ddl_context.rs b/tests/ddl_context.rs new file mode 100644 index 0000000..a5635c7 --- /dev/null +++ b/tests/ddl_context.rs @@ -0,0 +1,337 @@ +//! DDL builder requests in the exact shape the Tabularis host sends since the +//! DDL connection context (host commit +//! `03bee207fe62a64d482dbbc18e4b688203c7fec3`, `plugins/driver.rs`): every +//! builder carries `params` with the browsed project applied, columns may +//! carry `comment`, and the explorer asks `get_drop_table_sql` for the drop +//! statement. + +mod support; + +use serde_json::{json, Value}; +use support::{call, result}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +/// `ConnectionParams` as the host serializes it: browsed project in +/// `database`, the billing project and plugin settings in `extra`. +fn host_params(endpoint: &str) -> Value { + json!({ + "driver": "bigquery", + "host": null, + "port": null, + "username": null, + "password": null, + "database": "analytics-project", + "ssl_mode": null, + "ssl_ca": null, + "ssl_cert": null, + "ssl_key": null, + "ssh_enabled": null, + "ssh_connection_id": null, + "save_in_keychain": null, + "k8s_enabled": null, + "k8s_connection_id": null, + "ssm_enabled": null, + "connection_id": "conn-1", + "extra": { + "auth_method": "none", + "api_endpoint": endpoint, + "billing_project": "billing-project" + } + }) +} + +/// Host `ColumnDefinition`: `comment` is omitted when unset. +fn column(name: &str, data_type: &str, comment: Option<&str>) -> Value { + let mut column = json!({ + "name": name, + "data_type": data_type, + "is_nullable": true, + "is_pk": false, + "is_auto_increment": false, + "default_value": null + }); + if let Some(comment) = comment { + column["comment"] = json!(comment); + } + column +} + +fn statements(response: &Value) -> Vec { + serde_json::from_value(result(response).clone()).unwrap() +} + +/// Endpoint for builders that never call BigQuery. +const UNUSED: &str = "http://127.0.0.1:9"; + +#[tokio::test] +async fn create_table_is_project_qualified_with_column_descriptions() { + let response = call( + "get_create_table_sql", + json!({ + "params": host_params(UNUSED), + "table_name": "events", + "columns": [column("id", "INT64", None), column("note", "STRING", Some("Free 'text'"))], + "schema": "sales" + }), + ) + .await; + assert_eq!( + statements(&response), + vec![ + "CREATE TABLE `analytics-project`.`sales`.`events` (\n \ + `id` INT64,\n \ + `note` STRING OPTIONS(description = 'Free \\'text\\'')\n)" + ] + ); +} + +#[tokio::test] +async fn add_column_is_project_qualified_with_description() { + // Column changes look the table up first; an unknown table is left to + // BigQuery. + let server = MockServer::start().await; + let response = call( + "get_add_column_sql", + json!({ + "params": host_params(&server.uri()), + "table": "events", + "column": column("note", "STRING", Some("Free text")), + "schema": "sales" + }), + ) + .await; + assert_eq!( + statements(&response), + vec!["ALTER TABLE `analytics-project`.`sales`.`events` ADD COLUMN `note` STRING OPTIONS(description = 'Free text')"] + ); +} + +#[tokio::test] +async fn alter_column_sets_and_clears_the_description() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path( + "/bigquery/v2/projects/analytics-project/datasets/sales/tables/events", + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "type": "TABLE", + "schema": {"fields": [{"name": "note", "type": "STRING", "mode": "NULLABLE"}]} + }))) + .mount(&server) + .await; + let params = host_params(&server.uri()); + + let response = call( + "get_alter_column_sql", + json!({ + "params": params, + "table": "events", + "old_column": column("note", "STRING", None), + "new_column": column("note", "STRING", Some("Primary contact")), + "schema": "sales" + }), + ) + .await; + assert_eq!( + statements(&response), + vec!["ALTER TABLE `analytics-project`.`sales`.`events` ALTER COLUMN `note` SET OPTIONS(description = 'Primary contact')"] + ); + + let response = call( + "get_alter_column_sql", + json!({ + "params": params, + "table": "events", + "old_column": column("note", "STRING", Some("Primary contact")), + "new_column": column("note", "STRING", None), + "schema": "sales" + }), + ) + .await; + assert_eq!( + statements(&response), + vec!["ALTER TABLE `analytics-project`.`sales`.`events` ALTER COLUMN `note` SET OPTIONS(description = NULL)"] + ); +} + +#[tokio::test] +async fn foreign_keys_reference_the_browsed_project() { + let response = call( + "get_create_foreign_key_sql", + json!({ + "params": host_params(UNUSED), + "table": "orders", + "fk_name": "fk_user", + "column": "user_id", + "ref_table": "users", + "ref_column": "id", + "on_delete": null, + "on_update": null, + "schema": "sales" + }), + ) + .await; + assert_eq!( + statements(&response), + vec!["ALTER TABLE `analytics-project`.`sales`.`orders` ADD CONSTRAINT `fk_user` FOREIGN KEY (`user_id`) REFERENCES `analytics-project`.`sales`.`users`(`id`) NOT ENFORCED"] + ); +} + +#[tokio::test] +async fn routine_template_and_index_requests_carry_the_context() { + let response = call( + "routine_create_template", + json!({"params": host_params(UNUSED), "routine_type": "PROCEDURE", "schema": "sales"}), + ) + .await; + let template = result(&response).as_str().unwrap(); + assert!( + template + .starts_with("CREATE OR REPLACE PROCEDURE `analytics-project`.`sales`.`my_procedure`("), + "{template}" + ); + assert!(template.contains("FROM `analytics-project`.`sales`.`my_table`")); + + let response = call( + "get_create_index_sql", + json!({ + "params": host_params(UNUSED), + "table": "events", + "index_name": "idx", + "columns": ["note"], + "is_unique": false, + "schema": "sales" + }), + ) + .await; + assert_eq!(response["error"]["code"], -32601); +} + +#[tokio::test] +async fn all_databases_mode_falls_back_to_the_billing_project() { + let server = MockServer::start().await; + let mut params = host_params(&server.uri()); + params["database"] = json!(""); + let response = call( + "get_add_column_sql", + json!({"params": params, "table": "events", "column": column("n", "INT64", None), "schema": "sales"}), + ) + .await; + assert_eq!( + statements(&response), + vec!["ALTER TABLE `billing-project`.`sales`.`events` ADD COLUMN `n` INT64"] + ); +} + +async fn mount_kind(server: &MockServer, table: &str, kind: Option<&str>) { + let route = format!("/bigquery/v2/projects/analytics-project/datasets/sales/tables/{table}"); + let response = match kind { + Some(kind) => ResponseTemplate::new(200).set_body_json(json!({ + "type": kind, + "tableReference": {"projectId": "analytics-project", "datasetId": "sales", "tableId": table} + })), + None => ResponseTemplate::new(404).set_body_json(json!({ + "error": {"code": 404, "message": "Not found", "errors": [{"reason": "notFound"}]} + })), + }; + Mock::given(method("GET")) + .and(path(route)) + .respond_with(response) + .mount(server) + .await; +} + +#[tokio::test] +async fn drop_table_uses_the_object_kind_keyword() { + let server = MockServer::start().await; + let cases = [ + ("plain", Some("TABLE"), "DROP TABLE IF EXISTS"), + ("ext", Some("EXTERNAL"), "DROP EXTERNAL TABLE IF EXISTS"), + ( + "daily_snapshot", + Some("SNAPSHOT"), + "DROP SNAPSHOT TABLE IF EXISTS", + ), + ("v", Some("VIEW"), "DROP VIEW IF EXISTS"), + ( + "mv", + Some("MATERIALIZED_VIEW"), + "DROP MATERIALIZED VIEW IF EXISTS", + ), + ("missing", None, "DROP TABLE IF EXISTS"), + ]; + for (table, kind, _) in cases { + mount_kind(&server, table, kind).await; + } + for (table, _, keyword) in cases { + let response = call( + "get_drop_table_sql", + json!({ + "params": host_params(&server.uri()), + "table": table, + "schema": "sales", + "database": "analytics-project" + }), + ) + .await; + assert_eq!( + result(&response), + &json!(format!("{keyword} `analytics-project`.`sales`.`{table}`")), + "{table}" + ); + } +} + +#[tokio::test] +async fn drop_table_without_context_falls_back_to_the_host() { + for request in [ + json!({"table": "t", "schema": "sales", "database": null}), + json!({"params": host_params(UNUSED), "table": "t", "schema": null, "database": null}), + ] { + let response = call("get_drop_table_sql", request).await; + assert_eq!(result(&response), &Value::Null); + } +} + +#[tokio::test] +async fn drop_table_without_credentials_still_qualifies_the_table() { + let mut params = host_params(UNUSED); + params["extra"]["auth_method"] = json!("service_account"); + let response = call( + "get_drop_table_sql", + json!({ + "params": params, + "table": "orders", + "schema": "sales", + "database": "analytics-project" + }), + ) + .await; + assert_eq!( + result(&response), + &json!("DROP TABLE IF EXISTS `analytics-project`.`sales`.`orders`") + ); +} + +#[tokio::test] +async fn drop_table_uses_the_database_field_when_params_have_none() { + let server = MockServer::start().await; + mount_kind(&server, "daily_snapshot", Some("SNAPSHOT")).await; + let mut params = host_params(&server.uri()); + params["database"] = json!(""); + let response = call( + "get_drop_table_sql", + json!({ + "params": params, + "table": "daily_snapshot", + "schema": "sales", + "database": "analytics-project" + }), + ) + .await; + assert_eq!( + result(&response), + &json!("DROP SNAPSHOT TABLE IF EXISTS `analytics-project`.`sales`.`daily_snapshot`") + ); +} diff --git a/tests/dispatch.rs b/tests/dispatch.rs new file mode 100644 index 0000000..fab8055 --- /dev/null +++ b/tests/dispatch.rs @@ -0,0 +1,146 @@ +//! Concurrent stdio dispatch and host `cancel` notifications. + +mod support; + +use std::time::Duration; + +use bigquery_plugin::server::serve; +use serde_json::{json, Value}; +use support::mock_connection; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +async fn read_line( + lines: &mut tokio::io::Lines>>, +) -> Option { + tokio::time::timeout(Duration::from_secs(10), lines.next_line()) + .await + .expect("response within 10 s") + .unwrap() + .map(|line| serde_json::from_str(&line).unwrap()) +} + +#[tokio::test] +async fn cancel_notification_stops_the_job_without_a_reply() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/bigquery/v2/projects/p/queries")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "jobReference": {"projectId": "p", "jobId": "long", "location": "EU"}, + "jobComplete": false + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/bigquery/v2/projects/p/queries/long")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"jobComplete": false})) + .set_delay(Duration::from_secs(30)), + ) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/bigquery/v2/projects/p/jobs/long/cancel")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"job": {}}))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/bigquery/v2/projects/p/datasets")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"datasets": []}))) + .mount(&server) + .await; + + let (host, plugin) = tokio::io::duplex(64 * 1024); + let (plugin_read, plugin_write) = tokio::io::split(plugin); + let serving = tokio::spawn(serve(BufReader::new(plugin_read), plugin_write)); + let (host_read, mut host_write) = tokio::io::split(host); + let mut lines = BufReader::new(host_read).lines(); + + let request = json!({ + "jsonrpc": "2.0", "method": "execute_query", "id": 7, + "params": {"params": mock_connection(&server), "query": "SELECT * FROM huge", "limit": 10, "page": 1} + }); + let ping = json!({ + "jsonrpc": "2.0", "method": "ping", "id": 8, + "params": {"params": mock_connection(&server)} + }); + for message in [request, ping] { + host_write + .write_all(format!("{message}\n").as_bytes()) + .await + .unwrap(); + } + + // The ping answers while the query is still running. + let first = read_line(&mut lines).await.unwrap(); + assert_eq!(first["id"], 8); + assert_eq!(first["result"], Value::Null); + + for notification in [ + json!({"jsonrpc": "2.0", "method": "cancel", "params": {"id": 7}}), + json!({"jsonrpc": "2.0", "method": "cancel", "params": {"id": 12345}}), + json!({"jsonrpc": "2.0", "method": "progress", "params": {}}), + ] { + host_write + .write_all(format!("{notification}\n").as_bytes()) + .await + .unwrap(); + } + + let cancelled = read_line(&mut lines).await.unwrap(); + assert_eq!(cancelled["id"], 7); + assert!(cancelled["error"]["message"] + .as_str() + .unwrap() + .contains("cancelled")); + + host_write.shutdown().await.unwrap(); + drop(host_write); + assert!( + read_line(&mut lines).await.is_none(), + "notifications must never be answered" + ); + serving.await.unwrap(); + // Give the spawned jobs.cancel call time to arrive before verifying. + tokio::time::sleep(Duration::from_millis(200)).await; + let cancels: Vec<_> = server + .received_requests() + .await + .unwrap() + .into_iter() + .filter(|request| request.url.path().ends_with("/jobs/long/cancel")) + .collect(); + assert_eq!(cancels.len(), 1); + assert_eq!( + cancels[0].url.query(), + Some("location=EU"), + "jobs.cancel must carry the job location" + ); +} + +#[tokio::test] +async fn parse_errors_are_reported_with_a_null_id() { + let (host, plugin) = tokio::io::duplex(4096); + let (plugin_read, plugin_write) = tokio::io::split(plugin); + let serving = tokio::spawn(serve(BufReader::new(plugin_read), plugin_write)); + let (host_read, mut host_write) = tokio::io::split(host); + let mut lines = BufReader::new(host_read).lines(); + host_write.write_all(b"{not json\n\n").await.unwrap(); + host_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":3}\n") + .await + .unwrap(); + let parse_error = read_line(&mut lines).await.unwrap(); + assert_eq!(parse_error["error"]["code"], -32700); + assert_eq!(parse_error["id"], Value::Null); + let initialized = read_line(&mut lines).await.unwrap(); + assert_eq!( + initialized, + json!({"jsonrpc": "2.0", "result": null, "id": 3}) + ); + host_write.shutdown().await.unwrap(); + serving.await.unwrap(); +} diff --git a/tests/emulator.rs b/tests/emulator.rs new file mode 100644 index 0000000..1981e75 --- /dev/null +++ b/tests/emulator.rs @@ -0,0 +1,240 @@ +//! Integration tests against a BigQuery emulator. +//! +//! Skipped unless `BIGQUERY_EMULATOR_HOST` is set, for example: +//! +//! ```bash +//! docker run -d -p 9050:9050 ghcr.io/goccy/bigquery-emulator:latest --project=test --dataset=ds +//! BIGQUERY_EMULATOR_HOST=localhost:9050 cargo test --test emulator +//! ``` +//! +//! Metadata, DDL and CRUD round trips live in `emulator_metadata.rs`. +//! +//! The emulator does not implement sessions, dry runs, DML row counts, +//! child jobs or `tabledata.list` paging, so those paths are covered by the +//! mocked tests and only checked here for graceful behaviour. + +mod support; + +use bigquery_plugin::client::{self, JobRef}; +use bigquery_plugin::models::ConnectionParams; +use serde_json::{json, Value}; +use support::{call, result, unique_name}; + +const PROJECT: &str = support::EMULATOR_PROJECT; +const DATASET: &str = support::EMULATOR_DATASET; + +async fn query(params: &Value, sql: &str, limit: u32, page: u32) -> Value { + call( + "execute_query", + json!({"params": params, "query": sql, "limit": limit, "page": page, "schema": DATASET}), + ) + .await +} + +async fn create_numbers(params: &Value, table: &str, rows: u32) { + let response = query( + params, + &format!( + "CREATE TABLE `{DATASET}.{table}` AS SELECT n, CONCAT('row-', CAST(n AS STRING)) AS label FROM UNNEST(GENERATE_ARRAY(1, {rows})) AS n" + ), + 10, + 1, + ) + .await; + result(&response); +} + +#[tokio::test] +async fn connects_and_lists_datasets_and_tables() { + let params = emulator!(); + assert_eq!( + call("test_connection", json!({"params": params})).await["result"], + Value::Null + ); + assert_eq!( + call("ping", json!({"params": params})).await["result"], + Value::Null + ); + + let schemas = call("get_schemas", json!({"params": params})).await; + assert!(result(&schemas) + .as_array() + .unwrap() + .iter() + .any(|schema| schema == DATASET)); + + let table = unique_name("listing"); + create_numbers(¶ms, &table, 3).await; + let tables = call("get_tables", json!({"params": params, "schema": DATASET})).await; + assert!(result(&tables).to_string().contains(&table), "{tables}"); +} + +#[tokio::test] +async fn pages_through_large_results_without_rerunning() { + let params = emulator!(); + let table = unique_name("paging"); + create_numbers(¶ms, &table, 2500).await; + let sql = format!("SELECT n, label FROM `{DATASET}.{table}` ORDER BY n"); + + let mut seen = Vec::new(); + for page in 1..=3 { + let response = query(¶ms, &sql, 1000, page).await; + let value = result(&response); + assert_eq!(value["pagination"]["total_rows"], 2500); + assert_eq!(value["pagination"]["has_more"], page < 3); + seen.extend( + value["rows"] + .as_array() + .unwrap() + .iter() + .map(|row| row[0].as_i64().unwrap()), + ); + } + assert_eq!(seen, (1..=2500).collect::>()); +} + +#[tokio::test] +async fn browses_tables_like_a_table_tab() { + let params = emulator!(); + let table = unique_name("browse"); + create_numbers(¶ms, &table, 30).await; + let sql = format!("SELECT * FROM `{DATASET}`.`{table}` LIMIT 25"); + let first = query(¶ms, &sql, 20, 1).await; + let second = query(¶ms, &sql, 20, 2).await; + assert_eq!(result(&first)["rows"].as_array().unwrap().len(), 20); + assert_eq!(result(&first)["columns"], json!(["n", "label"])); + assert_eq!(result(&second)["rows"].as_array().unwrap().len(), 5); + assert_eq!(result(&second)["pagination"]["has_more"], false); +} + +#[tokio::test] +async fn runs_dml_and_reports_success() { + let params = emulator!(); + let table = unique_name("dml"); + create_numbers(¶ms, &table, 10).await; + let response = query( + ¶ms, + &format!("UPDATE `{DATASET}.{table}` SET label = 'x' WHERE n <= 4"), + 10, + 1, + ) + .await; + let affected = result(&response)["affected_rows"].as_u64().unwrap(); + // The emulator omits numDmlAffectedRows; BigQuery reports 4. + assert!(affected == 0 || affected == 4, "{affected}"); + let check = query( + ¶ms, + &format!("SELECT COUNT(*) AS c FROM `{DATASET}.{table}` WHERE label = 'x'"), + 10, + 1, + ) + .await; + assert_eq!(result(&check)["rows"], json!([[4]])); +} + +#[tokio::test] +async fn formats_timestamps_exactly() { + let params = emulator!(); + let response = query( + ¶ms, + "SELECT TIMESTAMP_MICROS(1700000000123457) AS created, TIMESTAMP '1969-12-31 23:59:59.5 UTC' AS before_epoch", + 10, + 1, + ) + .await; + assert_eq!( + result(&response)["rows"], + json!([["2023-11-14T22:13:20.123457Z", "1969-12-31T23:59:59.500000Z"]]) + ); +} + +#[tokio::test] +async fn returns_the_last_result_of_a_script() { + let params = emulator!(); + let response = query(¶ms, "SELECT 1 AS a; SELECT 2 AS b", 10, 1).await; + assert_eq!(result(&response)["rows"], json!([[2]])); +} + +#[tokio::test] +async fn reports_invalid_queries() { + let params = emulator!(); + let response = query( + ¶ms, + "SELECT * FROM missing_dataset.missing_table", + 10, + 1, + ) + .await; + assert!(response["error"]["message"] + .as_str() + .unwrap() + .contains("missing_table")); +} + +#[tokio::test] +async fn cancels_jobs_through_the_rest_api() { + let params = emulator!(); + let response = query(¶ms, "SELECT 1 AS n", 10, 1).await; + let job_id = result(&response)["metadata"]["job_id"] + .as_str() + .unwrap() + .to_string(); + let client = client::client_for(&ConnectionParams::from_value(¶ms)) + .await + .unwrap(); + // The job has finished; BigQuery and the emulator accept a late cancel. + client + .cancel_job(&JobRef { + project: PROJECT.to_string(), + job_id, + location: None, + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn sessions_keep_state_when_supported() { + let params = emulator!(); + let run = |sql: &str| { + call( + "execute_query", + json!({"params": params, "query": sql, "limit": 10, "page": 1, "schema": DATASET, "session_id": "emulator-tab"}), + ) + }; + let declared = run("CREATE TEMP TABLE session_numbers AS SELECT 1 AS n").await; + let declared = result(&declared); + if declared.get("error").is_some() { + eprintln!("emulator does not support sessions: {}", declared["error"]); + call("release_session", json!({"session_id": "emulator-tab"})).await; + return; + } + let read = run("SELECT n FROM session_numbers").await; + let read = result(&read); + match read.get("result") { + Some(value) => assert_eq!(value["rows"], json!([[1]])), + None => eprintln!("emulator does not keep session state: {}", read["error"]), + } + assert_eq!(read["in_transaction"], false); + let released = call("release_session", json!({"session_id": "emulator-tab"})).await; + assert_eq!(released["result"], Value::Null); +} + +#[tokio::test] +async fn explain_returns_a_raw_plan() { + let params = emulator!(); + for analyze in [false, true] { + let response = call( + "explain_query", + json!({"params": params, "query": "SELECT 1 AS n", "analyze": analyze, "schema": DATASET}), + ) + .await; + let raw = result(&response); + assert_eq!(raw["format"], "bigquery-job-json"); + let payload: Value = serde_json::from_str(raw["payload"].as_str().unwrap()).unwrap(); + assert_eq!( + payload["mode"], + if analyze { "analyze" } else { "estimate" } + ); + } +} diff --git a/tests/emulator_metadata.rs b/tests/emulator_metadata.rs new file mode 100644 index 0000000..76020d9 --- /dev/null +++ b/tests/emulator_metadata.rs @@ -0,0 +1,260 @@ +//! Live metadata, DDL and CRUD checks against the BigQuery emulator +//! (`ghcr.io/goccy/bigquery-emulator`). +//! +//! Skipped unless `BIGQUERY_EMULATOR_HOST` is set, for example: +//! +//! ```bash +//! docker run -d -p 9050:9050 ghcr.io/goccy/bigquery-emulator:latest --project=test --dataset=ds +//! BIGQUERY_EMULATOR_HOST=localhost:9050 cargo test --test emulator_metadata +//! ``` +//! +//! The emulator implements a subset of BigQuery: `@@row_count`, +//! `INFORMATION_SCHEMA.TABLES.ddl`, informational keys and the routines API +//! are missing, so those paths fall back or are not exercised here. + +mod support; + +use serde_json::{json, Value}; +use support::{call, unique_name, EMULATOR_DATASET, EMULATOR_PROJECT}; + +struct Emulator { + params: Value, +} + +impl Emulator { + fn connect() -> Option { + support::emulator_connection().map(|params| Self { params }) + } + + async fn call(&self, method: &str, mut arguments: Value) -> Result { + arguments["params"] = self.params.clone(); + if arguments.get("schema").is_none() { + arguments["schema"] = json!(EMULATOR_DATASET); + } + let response = call(method, arguments).await; + match response.get("error") { + Some(error) => Err(error["message"].as_str().unwrap_or_default().to_string()), + None => Ok(response["result"].clone()), + } + } + + async fn query(&self, sql: &str) -> Result { + self.call( + "execute_query", + json!({"query": sql, "limit": 100, "page": 1}), + ) + .await + } +} + +fn path(table: &str) -> String { + format!("`{EMULATOR_PROJECT}`.`{EMULATOR_DATASET}`.`{table}`") +} + +/// The emulator drops informational primary keys, so row edits take the +/// keyless path, whose transaction guard uses `@@row_count`, which the +/// emulator does not implement. +fn assert_dml(result: Result, expected: u64) { + match result { + Ok(value) => assert_eq!(value.as_u64(), Some(expected)), + Err(error) if error.contains("@@row_count") => { + eprintln!("emulator lacks @@row_count; keyless edit guard not exercised") + } + Err(error) => panic!("DML failed: {error}"), + } +} + +#[tokio::test] +async fn metadata_and_crud_round_trip() { + let Some(emulator) = Emulator::connect() else { + eprintln!("BIGQUERY_EMULATOR_HOST is not set; skipping emulator test"); + return; + }; + let table = unique_name("metadata_items"); + let table = table.as_str(); + let path = path(table); + let _ = emulator + .query(&format!("DROP TABLE IF EXISTS {path}")) + .await; + emulator + .query(&format!( + "CREATE TABLE {path} (id INT64 NOT NULL, name STRING, price NUMERIC, note STRING, PRIMARY KEY (id) NOT ENFORCED)" + )) + .await + .expect("create table"); + + let databases = emulator.call("get_databases", json!({})).await.unwrap(); + assert!(databases + .as_array() + .unwrap() + .iter() + .any(|project| project == EMULATOR_PROJECT)); + + let tables = emulator.call("get_tables", json!({})).await.unwrap(); + assert!(tables + .as_array() + .unwrap() + .iter() + .any(|entry| entry["name"] == table)); + + let columns = emulator + .call("get_columns", json!({"table": table})) + .await + .unwrap(); + assert_eq!(columns.as_array().unwrap().len(), 4); + assert_eq!(columns[0]["name"], "id"); + assert_eq!(columns[2]["data_type"], "NUMERIC"); + + let inserted = emulator + .call( + "insert_record", + json!({ + "table": table, + "data": {"id": 1, "name": "O'Reilly\nline", "price": "12.50", "note": "a\\b"} + }), + ) + .await + .expect("insert"); + assert!(inserted.as_u64().is_some()); + + let rows = emulator + .query(&format!( + "SELECT name, CAST(price AS STRING), note FROM {path} WHERE id = 1" + )) + .await + .unwrap(); + assert_eq!(rows["rows"][0][0], "O'Reilly\nline"); + assert_eq!(rows["rows"][0][2], "a\\b"); + + let identity = json!({"id": 1}); + let updated = emulator + .call( + "update_record", + json!({"table": table, "pk_map": identity, "col_name": "note", "new_val": "tab\there"}), + ) + .await; + assert_dml(updated, 1); + + let snapshot = emulator + .call("get_schema_snapshot", json!({})) + .await + .unwrap(); + assert!(snapshot + .as_array() + .unwrap() + .iter() + .any(|entry| entry["name"] == table)); + let batch = emulator + .call("get_all_columns_batch", json!({})) + .await + .unwrap(); + assert_eq!(batch[table].as_array().unwrap().len(), 4); + + let view = unique_name("metadata_view"); + let view = view.as_str(); + let created = emulator + .call( + "create_view", + json!({"view_name": view, "definition": format!("SELECT id, name FROM {path}")}), + ) + .await; + created.expect("create view"); + let definition = emulator + .call("get_view_definition", json!({"view_name": view})) + .await; + assert!(definition + .expect("view definition") + .as_str() + .unwrap() + .contains(view)); + let _ = emulator.call("drop_view", json!({"view_name": view})).await; + + let deleted = emulator + .call( + "delete_record", + json!({"table": table, "pk_map": {"id": 1}}), + ) + .await; + assert_dml(deleted, 1); + + let ddl = emulator + .call("get_table_ddl", json!({"table": table})) + .await + .expect("table ddl"); + assert!(ddl.as_str().unwrap().contains(table)); + + let drop = emulator + .call( + "get_drop_table_sql", + json!({"table": table, "database": EMULATOR_PROJECT}), + ) + .await + .expect("drop table sql"); + assert_eq!(drop, json!(format!("DROP TABLE IF EXISTS {path}"))); + emulator + .query(drop.as_str().unwrap()) + .await + .expect("drop table"); +} + +#[tokio::test] +async fn typed_parameters_round_trip() { + let Some(emulator) = Emulator::connect() else { + eprintln!("BIGQUERY_EMULATOR_HOST is not set; skipping emulator test"); + return; + }; + let table = unique_name("metadata_types"); + let table = table.as_str(); + let path = path(table); + let _ = emulator + .query(&format!("DROP TABLE IF EXISTS {path}")) + .await; + emulator + .query(&format!( + "CREATE TABLE {path} (id INT64, flag BOOL, day DATE, ts TIMESTAMP, ratio FLOAT64, \ + blob BYTES, doc JSON, place GEOGRAPHY, ids ARRAY, \ + item STRUCT, big BIGNUMERIC)" + )) + .await + .expect("create table"); + emulator + .call( + "insert_record", + json!({ + "table": table, + "data": { + "id": 7, + "flag": true, + "day": "2024-02-29", + "ts": "2024-01-02T03:04:05.123456Z", + "ratio": 0.5, + "blob": "BLOB:3:application/octet-stream:AQID", + "doc": {"k": "v'\n"}, + "place": "POINT(1 2)", + "ids": [1, 2, 3], + "item": {"sku": "a'b", "qty": 2}, + "big": "123456789012345678901234567890.123456789" + } + }), + ) + .await + .expect("insert typed row"); + let rows = emulator + .query(&format!( + "SELECT flag, CAST(day AS STRING), TO_BASE64(blob), JSON_VALUE(doc, '$.k'), \ + ARRAY_LENGTH(ids), item.sku, CAST(big AS STRING) FROM {path} WHERE id = 7" + )) + .await + .unwrap(); + let row = &rows["rows"][0]; + assert_eq!(row[0], true); + assert_eq!(row[1], "2024-02-29"); + assert_eq!(row[2], "AQID"); + assert_eq!(row[3], "v'\n"); + assert_eq!(row[4], 3); + assert_eq!(row[5], "a'b"); + assert_eq!(row[6], "123456789012345678901234567890.123456789"); + let _ = emulator + .query(&format!("DROP TABLE IF EXISTS {path}")) + .await; +} diff --git a/tests/fixtures/rest/routines.json b/tests/fixtures/rest/routines.json new file mode 100644 index 0000000..b10d707 --- /dev/null +++ b/tests/fixtures/rest/routines.json @@ -0,0 +1,61 @@ +[ + { + "routineReference": {"projectId": "my-project-1", "datasetId": "sales", "routineId": "net_amount"}, + "routineType": "SCALAR_FUNCTION", + "language": "SQL", + "arguments": [ + {"name": "amount", "dataType": {"typeKind": "NUMERIC"}}, + {"name": "rate", "dataType": {"typeKind": "FLOAT64"}} + ], + "returnType": {"typeKind": "NUMERIC"}, + "definitionBody": "amount * (1 - rate)" + }, + { + "routineReference": {"projectId": "my-project-1", "datasetId": "sales", "routineId": "upper_js"}, + "routineType": "SCALAR_FUNCTION", + "language": "JAVASCRIPT", + "arguments": [{"name": "s", "dataType": {"typeKind": "STRING"}}], + "returnType": {"typeKind": "STRING"}, + "definitionBody": "return s.toUpperCase();" + }, + { + "routineReference": {"projectId": "my-project-1", "datasetId": "sales", "routineId": "inferred"}, + "routineType": "SCALAR_FUNCTION", + "language": "SQL", + "arguments": [{"name": "x", "argumentKind": "ANY_TYPE"}], + "definitionBody": "x" + }, + { + "routineReference": {"projectId": "my-project-1", "datasetId": "sales", "routineId": "orders_since"}, + "routineType": "TABLE_VALUED_FUNCTION", + "language": "SQL", + "arguments": [{"name": "since", "dataType": {"typeKind": "TIMESTAMP"}}], + "returnTableType": {"columns": [ + {"name": "order_id", "type": {"typeKind": "INT64"}}, + {"name": "items", "type": {"typeKind": "ARRAY", "arrayElementType": {"typeKind": "STRUCT", "structType": {"fields": [{"name": "sku", "type": {"typeKind": "STRING"}}]}}}} + ]}, + "definitionBody": "SELECT order_id, items FROM `my-project-1`.`sales`.`orders` WHERE placed_at >= since" + }, + { + "routineReference": {"projectId": "my-project-1", "datasetId": "sales", "routineId": "close_orders"}, + "routineType": "PROCEDURE", + "language": "SQL", + "arguments": [ + {"name": "before", "mode": "IN", "dataType": {"typeKind": "DATE"}}, + {"name": "closed", "mode": "OUT", "dataType": {"typeKind": "INT64"}}, + {"name": "note", "mode": "INOUT", "dataType": {"typeKind": "STRING"}} + ], + "definitionBody": "BEGIN\n UPDATE `my-project-1`.`sales`.`orders` SET status = 'closed' WHERE DATE(placed_at) < before;\n SET closed = @@row_count;\nEND" + }, + { + "routineReference": {"projectId": "my-project-1", "datasetId": "sales", "routineId": "weighted_avg"}, + "routineType": "AGGREGATE_FUNCTION", + "language": "SQL", + "arguments": [ + {"name": "v", "dataType": {"typeKind": "FLOAT64"}}, + {"name": "w", "dataType": {"typeKind": "FLOAT64"}} + ], + "returnType": {"typeKind": "FLOAT64"}, + "definitionBody": "SUM(v * w) / SUM(w)" + } +] diff --git a/tests/fixtures/rest/table_orders.json b/tests/fixtures/rest/table_orders.json new file mode 100644 index 0000000..254220a --- /dev/null +++ b/tests/fixtures/rest/table_orders.json @@ -0,0 +1,43 @@ +{ + "kind": "bigquery#table", + "id": "my-project-1:sales.orders", + "tableReference": {"projectId": "my-project-1", "datasetId": "sales", "tableId": "orders"}, + "type": "TABLE", + "description": "Customer orders", + "schema": { + "fields": [ + {"name": "order_id", "type": "INTEGER", "mode": "REQUIRED", "description": "Order key"}, + {"name": "customer_id", "type": "INTEGER", "mode": "NULLABLE"}, + {"name": "line_id", "type": "INTEGER", "mode": "REQUIRED"}, + {"name": "status", "type": "STRING", "mode": "NULLABLE", "maxLength": "16", "defaultValueExpression": "'new'"}, + {"name": "amount", "type": "NUMERIC", "precision": "12", "scale": "2"}, + {"name": "placed_at", "type": "TIMESTAMP"}, + {"name": "location", "type": "GEOGRAPHY"}, + {"name": "payload", "type": "JSON"}, + {"name": "validity", "type": "RANGE", "rangeElementType": {"type": "DATE"}}, + {"name": "items", "type": "RECORD", "mode": "REPEATED", "fields": [ + {"name": "sku", "type": "STRING", "mode": "REQUIRED"}, + {"name": "qty", "type": "INTEGER"} + ]} + ] + }, + "timePartitioning": {"type": "DAY", "field": "placed_at"}, + "clustering": {"fields": ["customer_id"]}, + "tableConstraints": { + "primaryKey": {"columns": ["order_id", "line_id"]}, + "foreignKeys": [ + { + "name": "fk_customer", + "referencedTable": {"projectId": "my-project-1", "datasetId": "crm", "tableId": "customers"}, + "columnReferences": [{"referencingColumn": "customer_id", "referencedColumn": "id"}] + }, + { + "referencedTable": {"projectId": "my-project-1", "datasetId": "sales", "tableId": "order_lines"}, + "columnReferences": [ + {"referencingColumn": "order_id", "referencedColumn": "order_id"}, + {"referencingColumn": "line_id", "referencedColumn": "line_id"} + ] + } + ] + } +} diff --git a/tests/manifest_contract.rs b/tests/manifest_contract.rs index d0dc1c1..cc09e26 100644 --- a/tests/manifest_contract.rs +++ b/tests/manifest_contract.rs @@ -5,7 +5,8 @@ use serde_json::Value; #[test] fn manifest_release_files_exist_and_versions_match() { let manifest: Value = serde_json::from_str(include_str!("../.tabularium")).unwrap(); - assert_eq!(manifest["name"], "bigquery"); + assert_eq!(manifest["id"], "bigquery"); + assert_eq!(manifest["name"], "BigQuery"); assert_eq!(manifest["version"], env!("CARGO_PKG_VERSION")); assert_eq!(manifest["executable"], "bigquery-plugin"); let extension = &manifest["ui_extensions"].as_array().unwrap()[0]; @@ -14,6 +15,12 @@ fn manifest_release_files_exist_and_versions_match() { let root = Path::new(env!("CARGO_MANIFEST_DIR")); assert!(root.join("ui/src/connection-fields.tsx").is_file()); assert!(root.join("ui/vite.config.ts").is_file()); + let icon = manifest["icon"].as_str().unwrap(); + let icon_file = icon.rsplit('/').next().unwrap(); + assert!(icon.starts_with( + "https://raw.githubusercontent.com/TabularisDB/tabularis-bigquery-plugin/main/" + )); + assert!(root.join(icon_file).is_file(), "{icon_file}"); } #[test] @@ -32,3 +39,40 @@ fn manifest_declares_safe_bigquery_defaults() { "GCP Project ID" ); } + +#[test] +fn manifest_declares_bigquery_feature_coverage() { + let manifest: Value = serde_json::from_str(include_str!("../.tabularium")).unwrap(); + let capabilities = &manifest["capabilities"]; + assert_eq!(capabilities["routine_management"], true); + assert_eq!(capabilities["materialized_views"], true); + assert_eq!(capabilities["alter_primary_key"], true); + assert_eq!(capabilities["column_comments"], true); + assert_eq!(capabilities["drop_table_sql"], true); + assert_eq!(manifest["min_runtime_version"], "0.29.0"); + assert_eq!(capabilities["sql_dialect"], "mysql"); + let syntax = &capabilities["sql_syntax"]; + assert_eq!(syntax["extends"], "generic"); + assert_eq!(syntax["triple_quoted_strings"], true); + assert_eq!(syntax["hash_comments"], true); + assert_eq!(syntax["script_blocks"], true); + assert_eq!(syntax["string_escape"], "backslash"); + assert_eq!(syntax["like_escape"], false); + assert_eq!(syntax["boolean_literals"], "keywords"); + assert_eq!(syntax["formatter"], "bigquery"); + for (generic, native) in [ + ("JSON", "JSON"), + ("BLOB", "BYTES"), + ("UUID", "STRING"), + ("SERIAL", "INT64"), + ("DECIMAL", "NUMERIC"), + ("TIMESTAMP", "TIMESTAMP"), + ] { + assert_eq!(manifest["type_mappings"][generic], native, "{generic}"); + } + for data_type in manifest["data_types"].as_array().unwrap() { + if matches!(data_type["name"].as_str(), Some("NUMERIC" | "BIGNUMERIC")) { + assert_eq!(data_type["requires_precision"], false); + } + } +} diff --git a/tests/projects_mock.rs b/tests/projects_mock.rs new file mode 100644 index 0000000..b61b594 --- /dev/null +++ b/tests/projects_mock.rs @@ -0,0 +1,200 @@ +//! Billing and data projects, metadata statements on the job runtime and +//! metadata cache invalidation, against a mocked BigQuery REST API. +//! +//! Kept in its own test binary: the metadata cache is process-wide, so +//! statements in other suites would invalidate it while these tests count +//! `tables.get` calls. + +mod support; + +use serde_json::{json, Value}; +use support::{bodies, call, complete, connection, result}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const BILLING_QUERIES: &str = "/bigquery/v2/projects/bill/queries"; +const TABLE: &str = "/bigquery/v2/projects/data/datasets/ds/tables/items"; + +/// A connection browsing `data` whose jobs run in the `bill` project. +fn params(server: &MockServer) -> Value { + let mut params = connection(&server.uri(), "data"); + params["extra"]["billing_project"] = json!("bill"); + params +} + +async fn mount_table(server: &MockServer) { + Mock::given(method("GET")) + .and(path(TABLE)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "type": "TABLE", + "tableReference": {"projectId": "data", "datasetId": "ds", "tableId": "items"}, + "schema": {"fields": [ + {"name": "id", "type": "INT64", "mode": "REQUIRED"}, + {"name": "name", "type": "STRING", "mode": "NULLABLE"} + ]}, + "tableConstraints": {"primaryKey": {"columns": ["id"]}} + }))) + .mount(server) + .await; + Mock::given(method("GET")) + .and(path("/bigquery/v2/projects/data/datasets/ds")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"location": "EU"}))) + .mount(server) + .await; +} + +async fn table_gets(server: &MockServer) -> usize { + server + .received_requests() + .await + .unwrap_or_default() + .iter() + .filter(|request| request.method.as_str() == "GET" && request.url.path() == TABLE) + .count() +} + +#[tokio::test] +async fn jobs_run_in_the_billing_project_and_resolve_names_in_the_data_project() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(BILLING_QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("j", &[1], 1, None))) + .mount(&server) + .await; + + let response = call( + "execute_query", + json!({"params": params(&server), "query": "SELECT n FROM items", "limit": 10, "page": 1, "schema": "ds"}), + ) + .await; + assert_eq!(result(&response)["rows"], json!([[1]])); + let sent = bodies(&server, "POST", BILLING_QUERIES).await; + assert_eq!(sent.len(), 1); + assert_eq!( + sent[0]["defaultDataset"], + json!({"projectId": "data", "datasetId": "ds"}) + ); +} + +#[tokio::test] +async fn row_edits_use_the_job_runtime_with_typed_parameters() { + let server = MockServer::start().await; + mount_table(&server).await; + Mock::given(method("POST")) + .and(path(BILLING_QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "jobReference": {"projectId": "bill", "jobId": "dml", "location": "EU"}, + "jobComplete": true, + "numDmlAffectedRows": "1" + }))) + .mount(&server) + .await; + + let response = call( + "insert_record", + json!({"params": params(&server), "schema": "ds", "table": "items", "data": {"id": 7, "name": "O'Reilly"}}), + ) + .await; + assert_eq!(result(&response), &json!(1)); + + let sent = bodies(&server, "POST", BILLING_QUERIES).await; + assert_eq!(sent.len(), 1); + let body = &sent[0]; + assert!(body["requestId"].as_str().is_some_and(|id| !id.is_empty())); + assert_eq!(body["labels"], json!({"client": "tabularis"})); + assert!(body["jobTimeoutMs"].is_string()); + assert_eq!(body["location"], "EU"); + assert_eq!( + body["defaultDataset"], + json!({"projectId": "data", "datasetId": "ds"}) + ); + assert_eq!(body["parameterMode"], "NAMED"); + let parameters = body["queryParameters"].as_array().unwrap(); + assert_eq!(parameters.len(), 2); + assert!(parameters + .iter() + .any(|parameter| parameter["parameterValue"]["value"] == "O'Reilly")); + assert!(!body["query"].as_str().unwrap().contains("O'Reilly")); +} + +#[tokio::test] +async fn editor_ddl_clears_the_metadata_cache() { + let server = MockServer::start().await; + mount_table(&server).await; + Mock::given(method("POST")) + .and(path(BILLING_QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "jobReference": {"projectId": "bill", "jobId": "ddl", "location": "EU"}, + "jobComplete": true + }))) + .mount(&server) + .await; + let columns = || { + call( + "get_columns", + json!({"params": params(&server), "schema": "ds", "table": "items"}), + ) + }; + let query = |sql: &'static str| { + call( + "execute_query", + json!({"params": params(&server), "query": sql, "limit": 10, "page": 1, "schema": "ds"}), + ) + }; + + result(&columns().await); + result(&columns().await); + assert_eq!(table_gets(&server).await, 1, "second lookup is cached"); + + result(&query("ALTER TABLE items ADD COLUMN note STRING").await); + result(&columns().await); + assert_eq!(table_gets(&server).await, 2, "DDL clears the cache"); +} + +#[tokio::test] +async fn column_batches_fetch_only_requested_tables_and_skip_failures() { + let server = MockServer::start().await; + let tables = "/bigquery/v2/projects/data/datasets/ds/tables"; + let reference = |name: &str| json!({"type": "TABLE", "tableReference": {"projectId": "data", "datasetId": "ds", "tableId": name}}); + Mock::given(method("GET")) + .and(path(tables)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "tables": [reference("a"), reference("gone"), reference("other")] + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{tables}/a"))) + .respond_with(ResponseTemplate::new(200).set_body_json({ + let mut table = reference("a"); + table["schema"] = + json!({"fields": [{"name": "id", "type": "INT64", "mode": "REQUIRED"}]}); + table + })) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{tables}/gone"))) + .respond_with(ResponseTemplate::new(404).set_body_json(json!({ + "error": {"code": 404, "message": "Not found: Table data:ds.gone", "errors": [{"reason": "notFound", "message": "Not found"}]} + }))) + .mount(&server) + .await; + + let response = call( + "get_all_columns_batch", + json!({"params": params(&server), "schema": "ds", "tables": ["a", "gone"]}), + ) + .await; + let value = result(&response); + assert_eq!(value["a"][0]["name"], "id"); + assert!(value.get("gone").is_none(), "{value}"); + assert!(value.get("other").is_none(), "{value}"); + let requests = server.received_requests().await.unwrap_or_default(); + assert!( + requests + .iter() + .all(|request| request.url.path() != format!("{tables}/other")), + "unrequested tables must not be fetched" + ); +} diff --git a/tests/rpc_contract.rs b/tests/rpc_contract.rs index c799e02..0521965 100644 --- a/tests/rpc_contract.rs +++ b/tests/rpc_contract.rs @@ -40,7 +40,7 @@ async fn create_table_ddl_matches_bigquery_constraints() { }); let response = handle_line(&request.to_string()).await; let sql = response["result"][0].as_str().unwrap(); - assert!(sql.starts_with("CREATE TABLE `analytics.events`")); + assert!(sql.starts_with("CREATE TABLE `analytics`.`events`")); assert!(sql.contains("`payload` JSON")); assert!(sql.contains("PRIMARY KEY (`id`) NOT ENFORCED")); } diff --git a/tests/runtime_mock.rs b/tests/runtime_mock.rs new file mode 100644 index 0000000..6f9526b --- /dev/null +++ b/tests/runtime_mock.rs @@ -0,0 +1,748 @@ +//! Query runtime behaviour against a mocked BigQuery REST API. + +mod support; + +use std::time::Duration; + +use serde_json::{json, Value}; +use support::{bodies, call, complete, field, mock_connection, queries, query_value, result}; +use wiremock::matchers::{body_partial_json, method, path, query_param}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const QUERIES: &str = "/bigquery/v2/projects/p/queries"; +const JOBS: &str = "/bigquery/v2/projects/p/jobs"; + +fn query_params(server: &MockServer, sql: &str, limit: u32, page: u32) -> Value { + json!({ + "params": mock_connection(server), + "query": sql, + "limit": limit, + "page": page, + "schema": "ds" + }) +} + +#[tokio::test] +async fn follows_page_tokens_until_the_page_is_full() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete( + "j1", + &[1, 2], + 5, + Some("t1"), + ))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{QUERIES}/j1"))) + .and(query_param("pageToken", "t1")) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("j1", &[3, 4, 5], 5, None))) + .mount(&server) + .await; + + let response = call( + "execute_query", + query_params(&server, "SELECT n FROM t", 5, 1), + ) + .await; + let value = result(&response); + assert_eq!(value["rows"], json!([[1], [2], [3], [4], [5]])); + assert_eq!(value["pagination"]["has_more"], false); + assert_eq!(value["pagination"]["total_rows"], 5); + + let gets = queries(&server, "GET", &format!("{QUERIES}/j1")).await; + assert_eq!(gets.len(), 1); + assert_eq!(query_value(&gets[0], "maxResults"), Some("3")); + assert_eq!(query_value(&gets[0], "location"), Some("US")); +} + +#[tokio::test] +async fn sends_idempotent_labelled_jobs_and_reuses_them_for_later_pages() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete( + "j2", + &[1, 2], + 6, + Some("t"), + ))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{QUERIES}/j2"))) + .and(query_param("startIndex", "2")) + .respond_with(ResponseTemplate::new(200).set_body_json(complete( + "j2", + &[3, 4], + 6, + Some("u"), + ))) + .mount(&server) + .await; + + let sql = "SELECT n FROM t ORDER BY n"; + let first = call("execute_query", query_params(&server, sql, 2, 1)).await; + assert_eq!(result(&first)["rows"], json!([[1], [2]])); + let second = call("execute_query", query_params(&server, sql, 2, 2)).await; + assert_eq!(result(&second)["rows"], json!([[3], [4]])); + assert_eq!(result(&second)["pagination"]["has_more"], true); + + let posts = bodies(&server, "POST", QUERIES).await; + assert_eq!(posts.len(), 1, "page 2 must not run the SQL again"); + let body = &posts[0]; + assert_eq!(body["labels"], json!({"client": "tabularis"})); + assert_eq!(body["jobTimeoutMs"], "110000"); + assert_eq!(body["formatOptions"], json!({"useInt64Timestamp": true})); + assert_eq!(body["maxResults"], 2); + assert_eq!( + body["defaultDataset"], + json!({"projectId": "p", "datasetId": "ds"}) + ); + let request_id = body["requestId"].as_str().unwrap(); + assert_eq!( + uuid::Uuid::parse_str(request_id).unwrap().get_version_num(), + 4 + ); +} + +#[tokio::test] +async fn expired_results_of_side_effecting_statements_are_not_rerun() { + let server = MockServer::start().await; + let response = call( + "execute_query", + query_params(&server, "DELETE FROM t WHERE true", 10, 2), + ) + .await; + assert!(response["error"]["message"] + .as_str() + .unwrap() + .contains("expired")); + assert!(bodies(&server, "POST", QUERIES).await.is_empty()); +} + +#[tokio::test] +async fn retries_transient_failures_with_the_same_request_id() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(503).set_body_json(json!({ + "error": {"code": 503, "message": "backend", "errors": [{"reason": "backendError", "message": "backend"}]} + }))) + .up_to_n_times(1) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(403).set_body_json(json!({ + "error": {"code": 403, "message": "slow down", "errors": [{"reason": "rateLimitExceeded", "message": "slow down"}]} + }))) + .up_to_n_times(1) + .with_priority(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("j3", &[7], 1, None))) + .with_priority(3) + .mount(&server) + .await; + + let response = call("execute_query", query_params(&server, "SELECT 7", 10, 1)).await; + assert_eq!(result(&response)["rows"], json!([[7]])); + let posts = bodies(&server, "POST", QUERIES).await; + assert_eq!(posts.len(), 3); + assert_eq!(posts[0]["requestId"], posts[1]["requestId"]); + assert_eq!(posts[1]["requestId"], posts[2]["requestId"]); +} + +#[tokio::test] +async fn does_not_retry_permanent_errors_and_explains_them() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(403).set_body_json(json!({ + "error": {"code": 403, "message": "Access Denied: Project p", "errors": [{"reason": "accessDenied", "message": "Access Denied: Project p"}]} + }))) + .mount(&server) + .await; + let response = call("execute_query", query_params(&server, "SELECT 1", 10, 1)).await; + let message = response["error"]["message"].as_str().unwrap(); + assert!( + message.contains("Access Denied: Project p [accessDenied, HTTP 403]"), + "{message}" + ); + assert!( + message.contains("Hint: Grant the account BigQuery Job User"), + "{message}" + ); + assert_eq!(bodies(&server, "POST", QUERIES).await.len(), 1); +} + +#[tokio::test] +async fn refreshes_credentials_once_after_unauthorized() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(401).set_body_json(json!({ + "error": {"code": 401, "message": "Request had invalid authentication credentials."} + }))) + .up_to_n_times(1) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("j4", &[1], 1, None))) + .with_priority(2) + .mount(&server) + .await; + let response = call("execute_query", query_params(&server, "SELECT 1", 10, 1)).await; + assert_eq!(result(&response)["rows"], json!([[1]])); + assert_eq!(bodies(&server, "POST", QUERIES).await.len(), 2); +} + +#[tokio::test] +async fn judges_failure_by_the_job_error_result() { + let server = MockServer::start().await; + let mut warning = complete("j5", &[1], 1, None); + warning["errors"] = json!([{"reason": "invalid", "message": "A warning"}]); + Mock::given(method("POST")) + .and(path(QUERIES)) + .and(body_partial_json(json!({"query": "SELECT 1"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(warning)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{JOBS}/j5"))) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"status": {"state": "DONE"}})), + ) + .mount(&server) + .await; + let response = call("execute_query", query_params(&server, "SELECT 1", 10, 1)).await; + assert_eq!(result(&response)["rows"], json!([[1]])); + + let mut failed = complete("j6", &[], 0, None); + failed["errors"] = json!([{"reason": "invalidQuery", "message": "Unrecognized name: x"}]); + Mock::given(method("POST")) + .and(path(QUERIES)) + .and(body_partial_json(json!({"query": "SELECT x"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(failed)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{JOBS}/j6"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "status": {"state": "DONE", "errorResult": {"reason": "invalidQuery", "location": "1:8", "message": "Unrecognized name: x"}} + }))) + .mount(&server) + .await; + let response = call("execute_query", query_params(&server, "SELECT x", 10, 1)).await; + let message = response["error"]["message"].as_str().unwrap(); + assert!( + message.starts_with("Unrecognized name: x (line 1, column 8) [invalidQuery]"), + "{message}" + ); +} + +#[tokio::test] +async fn polls_running_jobs_until_complete() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "jobReference": {"projectId": "p", "jobId": "slow", "location": "EU"}, + "jobComplete": false + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{QUERIES}/slow"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"jobComplete": false}))) + .up_to_n_times(1) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{QUERIES}/slow"))) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("slow", &[1, 2], 2, None))) + .with_priority(2) + .mount(&server) + .await; + let response = call( + "execute_query", + query_params(&server, "SELECT n FROM big", 10, 1), + ) + .await; + assert_eq!(result(&response)["rows"], json!([[1], [2]])); + let gets = queries(&server, "GET", &format!("{QUERIES}/slow")).await; + assert_eq!(gets.len(), 2); + assert!(gets + .iter() + .all(|query| query_value(query, "location") == Some("EU"))); +} + +#[tokio::test] +async fn browses_native_tables_through_tabledata_list() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/bigquery/v2/projects/p/datasets/ds/tables/events")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "type": "TABLE", + "numRows": "10", + "schema": {"fields": [field("n", "INT64"), field("at", "TIMESTAMP")]} + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path( + "/bigquery/v2/projects/p/datasets/ds/tables/events/data", + )) + .and(query_param("startIndex", "5")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "totalRows": "10", + "pageToken": "next", + "rows": [{"f": [{"v": "5"}, {"v": "1700000000123457"}]}] + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path( + "/bigquery/v2/projects/p/datasets/ds/tables/events/data", + )) + .and(query_param("pageToken", "next")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "totalRows": "10", + "rows": [{"f": [{"v": "6"}, {"v": null}]}] + }))) + .mount(&server) + .await; + + // Host table tabs send `SELECT * FROM `ds`.`events` LIMIT n` and page. + let response = call( + "execute_query", + query_params(&server, "SELECT * FROM `ds`.`events` LIMIT 7", 5, 2), + ) + .await; + let value = result(&response); + assert_eq!( + value["rows"], + json!([[5, "2023-11-14T22:13:20.123457Z"], [6, null]]) + ); + assert_eq!(value["pagination"]["total_rows"], 7); + assert_eq!(value["pagination"]["has_more"], false); + assert_eq!(value["metadata"]["source"], "tabledata.list"); + assert!(bodies(&server, "POST", QUERIES).await.is_empty()); + let pages = queries( + &server, + "GET", + "/bigquery/v2/projects/p/datasets/ds/tables/events/data", + ) + .await; + assert_eq!(query_value(&pages[0], "maxResults"), Some("2")); + assert_eq!(query_value(&pages[1], "maxResults"), Some("1")); +} + +#[tokio::test] +async fn views_fall_back_to_query_jobs() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/bigquery/v2/projects/p/datasets/ds/tables/v")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"type": "VIEW"}))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("view", &[1], 1, None))) + .mount(&server) + .await; + let response = call( + "execute_query", + query_params(&server, "SELECT * FROM `ds`.`v`", 10, 1), + ) + .await; + assert_eq!(result(&response)["metadata"]["source"], "job"); + assert_eq!(bodies(&server, "POST", QUERIES).await.len(), 1); +} + +#[tokio::test] +async fn scripts_return_the_last_select_and_sum_dml_rows() { + let server = MockServer::start().await; + let mut parent = complete("script", &[], 0, None); + parent["schema"] = Value::Null; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(parent)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{JOBS}/script"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "statistics": {"query": {"statementType": "SCRIPT"}} + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(JOBS)) + .and(query_param("parentJobId", "script")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"jobs": [ + {"jobReference": {"projectId": "p", "jobId": "c3", "location": "US"}, + "statistics": {"creationTime": "3", "query": {"statementType": "DELETE", "numDmlAffectedRows": "4"}}}, + {"jobReference": {"projectId": "p", "jobId": "c2", "location": "US"}, + "statistics": {"creationTime": "2", "query": {"statementType": "SELECT"}, "scriptStatistics": {"evaluationKind": "STATEMENT"}}}, + {"jobReference": {"projectId": "p", "jobId": "c1", "location": "US"}, + "statistics": {"creationTime": "1", "query": {"statementType": "INSERT", "numDmlAffectedRows": "3"}}} + ]}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{QUERIES}/c2"))) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("c2", &[42], 1, None))) + .mount(&server) + .await; + + let sql = "INSERT INTO t VALUES (1); SELECT 42 AS n; DELETE FROM t WHERE true"; + let response = call("execute_query", query_params(&server, sql, 10, 1)).await; + let value = result(&response); + assert_eq!(value["rows"], json!([[42]])); + assert_eq!(value["affected_rows"], 7); + assert_eq!(bodies(&server, "POST", QUERIES).await[0]["maxResults"], 0); +} + +async fn session_mocks(server: &MockServer) { + Mock::given(method("POST")) + .and(path(QUERIES)) + .and(body_partial_json(json!({"createSession": true}))) + .respond_with(ResponseTemplate::new(200).set_body_json({ + let mut response = complete("s1", &[1], 1, None); + response["sessionInfo"] = json!({"sessionId": "bq-session"}); + response + })) + .with_priority(1) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .and(body_partial_json(json!({"query": "SELECT broken"}))) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": {"code": 400, "message": "Unrecognized name: broken", "errors": [{"reason": "invalidQuery", "message": "Unrecognized name: broken"}]} + }))) + .with_priority(1) + .mount(server) + .await; + Mock::given(method("GET")) + .and(path(format!("{QUERIES}/s1"))) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("s1", &[1], 1, None))) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("s2", &[2], 1, None))) + .with_priority(5) + .mount(server) + .await; +} + +#[tokio::test] +async fn host_sessions_map_to_bigquery_sessions_and_track_transactions() { + let server = MockServer::start().await; + session_mocks(&server).await; + let run = |sql: &str| { + let mut params = query_params(&server, sql, 10, 1); + params["session_id"] = json!("tab-1"); + call("execute_query", params) + }; + + let begin = run("BEGIN TRANSACTION").await; + assert_eq!(result(&begin)["in_transaction"], true); + assert_eq!(result(&begin)["result"]["rows"], json!([[1]])); + + let failed = run("SELECT broken").await; + let failed = result(&failed); + assert!(failed["error"] + .as_str() + .unwrap() + .contains("Unrecognized name")); + assert_eq!(failed["in_transaction"], true); + + let savepoint = run("SAVEPOINT tabularis_read").await; + assert!(result(&savepoint)["error"] + .as_str() + .unwrap() + .contains("savepoints")); + + let commit = run("COMMIT").await; + assert_eq!(result(&commit)["in_transaction"], false); + + let posts = bodies(&server, "POST", QUERIES).await; + assert_eq!(posts.len(), 3, "the savepoint is never sent"); + assert_eq!(posts[0]["createSession"], true); + for joined in &posts[1..] { + assert_eq!( + joined["connectionProperties"], + json!([{"key": "session_id", "value": "bq-session"}]) + ); + assert_eq!(joined["location"], "US"); + assert!(joined.get("createSession").is_none()); + } + + let released = call("release_session", json!({"session_id": "tab-1"})).await; + assert_eq!(released["result"], Value::Null); + let posts = bodies(&server, "POST", QUERIES).await; + assert_eq!( + posts.last().unwrap()["query"], + "CALL BQ.ABORT_SESSION('bq-session')" + ); + let unknown = call("release_session", json!({"session_id": "never-used"})).await; + assert_eq!(unknown["result"], Value::Null); +} + +#[tokio::test] +async fn stateful_batches_share_a_short_lived_session() { + let server = MockServer::start().await; + session_mocks(&server).await; + let response = call( + "execute_query_batch", + json!({ + "params": mock_connection(&server), + "queries": ["DECLARE x INT64 DEFAULT 1", "SELECT x"], + "limit": 10, + "page": 1 + }), + ) + .await; + let results = result(&response) + .as_array() + .expect("bare array without session_id"); + assert_eq!(results.len(), 2); + assert!( + results.iter().all(|item| item["error"].is_null()), + "{results:?}" + ); + let posts = bodies(&server, "POST", QUERIES).await; + assert_eq!(posts[0]["createSession"], true); + assert_eq!(posts[1]["connectionProperties"][0]["value"], "bq-session"); + assert_eq!( + posts.last().unwrap()["query"], + "CALL BQ.ABORT_SESSION('bq-session')" + ); +} + +#[tokio::test] +async fn explain_returns_a_dry_run_estimate_without_running_the_query() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(JOBS)) + .and(body_partial_json(json!({"configuration": {"dryRun": true}}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "jobReference": {"projectId": "p", "location": "US"}, + "statistics": {"query": {"statementType": "SELECT", "totalBytesProcessed": "1099511627776", "cacheHit": false}} + }))) + .mount(&server) + .await; + let response = call( + "explain_query", + json!({"params": mock_connection(&server), "query": "SELECT * FROM ds.t", "analyze": false, "schema": "ds"}), + ) + .await; + let raw = result(&response); + assert_eq!(raw["engine"], "bigquery"); + assert_eq!(raw["format"], "bigquery-job-json"); + let payload: Value = serde_json::from_str(raw["payload"].as_str().unwrap()).unwrap(); + assert_eq!(payload["mode"], "estimate"); + assert_eq!(payload["estimated_cost_usd"], 6.25); + assert!(bodies(&server, "POST", QUERIES).await.is_empty()); +} + +#[tokio::test] +async fn explain_analyze_reads_the_executed_query_plan() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("an", &[], 0, None))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{JOBS}/an"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "jobReference": {"projectId": "p", "jobId": "an"}, + "statistics": {"startTime": "1", "endTime": "11", "query": { + "totalBytesBilled": "10485760", "totalSlotMs": "120", + "queryPlan": [{"id": "0", "name": "S00: Input", "slotMs": "100"}, {"id": "1", "name": "S01: Output", "inputStages": ["0"], "slotMs": "20"}] + }} + }))) + .mount(&server) + .await; + let response = call( + "explain_query", + json!({"params": mock_connection(&server), "query": "SELECT 1", "analyze": true}), + ) + .await; + let payload: Value = + serde_json::from_str(result(&response)["payload"].as_str().unwrap()).unwrap(); + assert_eq!(payload["mode"], "analyze"); + assert_eq!(payload["query_plan"].as_array().unwrap().len(), 2); + assert_eq!(payload["total_slot_ms"], 120); + assert_eq!(payload["elapsed_ms"], 10); +} + +#[tokio::test] +async fn ping_lists_a_single_dataset() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/bigquery/v2/projects/p/datasets")) + .and(query_param("maxResults", "1")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"datasets": []}))) + .expect(1) + .mount(&server) + .await; + let response = call("ping", json!({"params": mock_connection(&server)})).await; + assert_eq!(response["result"], Value::Null); + tokio::time::sleep(Duration::from_millis(10)).await; +} + +#[tokio::test] +async fn later_pages_of_a_session_creating_statement_reuse_its_job() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .and(body_partial_json(json!({"createSession": true}))) + .respond_with(ResponseTemplate::new(200).set_body_json({ + let mut response = complete("sp", &[1, 2], 4, Some("t")); + response["sessionInfo"] = json!({"sessionId": "bq-paged"}); + response + })) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{QUERIES}/sp"))) + .and(query_param("startIndex", "2")) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("sp", &[3, 4], 4, None))) + .mount(&server) + .await; + let run = |page: u32| { + let mut params = query_params(&server, "SELECT n FROM t ORDER BY n", 2, page); + params["session_id"] = json!("tab-paged"); + call("execute_query", params) + }; + let first = run(1).await; + assert_eq!(result(&first)["result"]["rows"], json!([[1], [2]])); + let second = run(2).await; + assert_eq!(result(&second)["result"]["rows"], json!([[3], [4]])); + assert_eq!( + bodies(&server, "POST", QUERIES).await.len(), + 1, + "page 2 must not run the statement again" + ); + call("release_session", json!({"session_id": "tab-paged"})).await; +} + +#[tokio::test] +async fn failed_session_creating_jobs_still_release_their_session() { + let server = MockServer::start().await; + let mut failed = complete("sf", &[], 0, None); + failed["errors"] = json!([{"reason": "invalidQuery", "message": "Division by zero"}]); + Mock::given(method("POST")) + .and(path(QUERIES)) + .and(body_partial_json(json!({"createSession": true}))) + .respond_with(ResponseTemplate::new(200).set_body_json(failed)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{JOBS}/sf"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "status": {"state": "DONE", "errorResult": {"reason": "invalidQuery", "message": "Division by zero"}}, + "statistics": {"sessionInfo": {"sessionId": "bq-failed"}} + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete("ab", &[], 0, None))) + .with_priority(5) + .mount(&server) + .await; + + let mut params = query_params(&server, "SELECT 1 / 0", 10, 1); + params["session_id"] = json!("tab-failed"); + let response = call("execute_query", params).await; + assert!(result(&response)["error"] + .as_str() + .unwrap() + .contains("Division by zero")); + call("release_session", json!({"session_id": "tab-failed"})).await; + let posts = bodies(&server, "POST", QUERIES).await; + assert_eq!( + posts.last().unwrap()["query"], + "CALL BQ.ABORT_SESSION('bq-failed')" + ); +} + +#[tokio::test] +async fn repeated_child_job_page_tokens_fail_instead_of_looping() { + let server = MockServer::start().await; + let mut parent = complete("loop", &[], 0, None); + parent["schema"] = Value::Null; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(parent)) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{JOBS}/loop"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "statistics": {"query": {"statementType": "SCRIPT"}} + }))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(JOBS)) + .and(query_param("parentJobId", "loop")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "jobs": [], "nextPageToken": "same" + }))) + .mount(&server) + .await; + let sql = "SELECT 1; SELECT 2"; + let response = call("execute_query", query_params(&server, sql, 10, 1)).await; + let message = response["error"]["message"].as_str().unwrap(); + assert!(message.contains("twice"), "{message}"); + assert_eq!(queries(&server, "GET", JOBS).await.len(), 2); +} + +#[tokio::test] +async fn repeated_result_page_tokens_fail_instead_of_duplicating_rows() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(QUERIES)) + .respond_with(ResponseTemplate::new(200).set_body_json(complete( + "dup", + &[1], + 10, + Some("same"), + ))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path(format!("{QUERIES}/dup"))) + .and(query_param("pageToken", "same")) + .respond_with(ResponseTemplate::new(200).set_body_json(complete( + "dup", + &[2], + 10, + Some("same"), + ))) + .mount(&server) + .await; + let response = call( + "execute_query", + query_params(&server, "SELECT n FROM t", 5, 1), + ) + .await; + let message = response["error"]["message"].as_str().unwrap(); + assert!(message.contains("twice"), "{message}"); +} diff --git a/tests/support/mod.rs b/tests/support/mod.rs new file mode 100644 index 0000000..68b3105 --- /dev/null +++ b/tests/support/mod.rs @@ -0,0 +1,132 @@ +#![allow(dead_code)] + +use bigquery_plugin::rpc::handle_request; +use serde_json::{json, Value}; +use wiremock::MockServer; + +/// Connection params pointing the plugin at a local mock or emulator. +pub fn connection(endpoint: &str, project: &str) -> Value { + json!({ + "driver": "bigquery", + "database": project, + "extra": {"auth_method": "none", "api_endpoint": endpoint} + }) +} + +pub fn mock_connection(server: &MockServer) -> Value { + connection(&server.uri(), "p") +} + +/// Project and dataset created by the documented emulator command: +/// `bigquery-emulator --project=test --dataset=ds`. +pub const EMULATOR_PROJECT: &str = "test"; +pub const EMULATOR_DATASET: &str = "ds"; + +/// Connection params for the emulator named by `BIGQUERY_EMULATOR_HOST`, or +/// `None` when emulator tests should be skipped. +pub fn emulator_connection() -> Option { + let host = std::env::var("BIGQUERY_EMULATOR_HOST").ok()?; + let host = host.trim(); + if host.is_empty() { + return None; + } + let endpoint = if host.starts_with("http") { + host.to_string() + } else { + format!("http://{host}") + }; + Some(connection(&endpoint, EMULATOR_PROJECT)) +} + +/// Skip the current test unless an emulator is configured. +#[macro_export] +macro_rules! emulator { + () => { + match support::emulator_connection() { + Some(params) => params, + None => { + eprintln!("BIGQUERY_EMULATOR_HOST is not set; skipping emulator test"); + return; + } + } + }; +} + +pub fn unique_name(prefix: &str) -> String { + format!("{prefix}_{}", uuid::Uuid::new_v4().simple()) +} + +pub async fn call(method: &str, params: Value) -> Value { + handle_request(json!({"jsonrpc": "2.0", "method": method, "params": params, "id": 1})).await +} + +pub fn result(response: &Value) -> &Value { + assert!( + response.get("error").is_none(), + "unexpected error: {response}" + ); + &response["result"] +} + +pub fn field(name: &str, field_type: &str) -> Value { + json!({"name": name, "type": field_type, "mode": "NULLABLE"}) +} + +pub fn rows(values: &[i64]) -> Value { + Value::Array( + values + .iter() + .map(|value| json!({"f": [{"v": value.to_string()}]})) + .collect(), + ) +} + +/// A completed jobs.query / getQueryResults response with one INT64 column. +pub fn complete(job_id: &str, values: &[i64], total: u64, page_token: Option<&str>) -> Value { + let mut response = json!({ + "jobReference": {"projectId": "p", "jobId": job_id, "location": "US"}, + "jobComplete": true, + "schema": {"fields": [field("n", "INT64")]}, + "rows": rows(values), + "totalRows": total.to_string() + }); + if let Some(token) = page_token { + response["pageToken"] = json!(token); + } + response +} + +pub async fn bodies(server: &MockServer, method: &str, path: &str) -> Vec { + server + .received_requests() + .await + .unwrap_or_default() + .into_iter() + .filter(|request| request.method.as_str() == method && request.url.path() == path) + .map(|request| serde_json::from_slice(&request.body).unwrap_or(Value::Null)) + .collect() +} + +pub async fn queries(server: &MockServer, method: &str, path: &str) -> Vec> { + server + .received_requests() + .await + .unwrap_or_default() + .into_iter() + .filter(|request| request.method.as_str() == method && request.url.path() == path) + .map(|request| { + request + .url + .query_pairs() + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect() + }) + .collect() +} + +pub fn query_value<'a>(query: &'a [(String, String)], key: &str) -> Option<&'a str> { + query + .iter() + .find(|(name, _)| name == key) + .map(|(_, value)| value.as_str()) +} diff --git a/ui/pnpm-workspace.yaml b/ui/pnpm-workspace.yaml new file mode 100644 index 0000000..3ee2585 --- /dev/null +++ b/ui/pnpm-workspace.yaml @@ -0,0 +1,8 @@ +# Own workspace root so pnpm 9 installs this package alone; the build +# allowlist lets pnpm 10 and later run the esbuild install script. +packages: + - . +allowBuilds: + esbuild: true +onlyBuiltDependencies: + - esbuild diff --git a/ui/src/connection-fields.tsx b/ui/src/connection-fields.tsx index d7989c3..ac57d91 100644 --- a/ui/src/connection-fields.tsx +++ b/ui/src/connection-fields.tsx @@ -20,7 +20,12 @@ type ConnectionContext = TypedSlotProps<"connection-modal.extra_fields">["contex setSecretField?: (key: string, value: string) => void; }; -type AuthMethod = "adc" | "service_account" | "authorized_user" | "access_token"; +type AuthMethod = "adc" | "service_account" | "authorized_user" | "access_token" | "none"; + +// Modes that authenticate without the plugin-owned credential. +const CREDENTIAL_FREE_METHODS: readonly AuthMethod[] = ["adc", "none"]; + +const LOOPBACK_ENDPOINT = /^https?:\/\/(localhost|127\.0\.0\.1|\[::1\])(:\d+)?(\/|$)/i; function Field({ label, children }: PropsWithChildren<{ label: string }>) { return ( @@ -47,6 +52,9 @@ function authenticationHelp(authMethod: AuthMethod, target: string): string { if (authMethod === "access_token") { return `Put a current OAuth access token in ${target}. Access tokens expire and must be replaced manually.`; } + if (authMethod === "none") { + return "Sends requests without credentials. Only allowed with a localhost API endpoint, such as a BigQuery emulator."; + } return "No credential field is used. The plugin follows Application Default Credentials, including GOOGLE_APPLICATION_CREDENTIALS, gcloud ADC and attached service accounts."; } @@ -62,6 +70,17 @@ const BigQueryConnectionFields = defineSlot( const authMethod = (extra.auth_method || "adc") as AuthMethod; const setSecretField = connection.setSecretField; const secureCredentialAvailable = typeof setSecretField === "function"; + const apiEndpoint = extra.api_endpoint || ""; + const endpointIsLocal = LOOPBACK_ENDPOINT.test(apiEndpoint.trim()); + + const changeAuthMethod = (next: AuthMethod) => { + setExtraField("auth_method", next); + // A credential typed or stored for the previous mode must not be kept + // silently when the new mode does not use one. + if (CREDENTIAL_FREE_METHODS.includes(next) && setSecretField) { + setSecretField("credential", ""); + } + }; const credential = connection.secretFields?.credential ?? { value: "", hasStoredValue: false, @@ -95,7 +114,7 @@ const BigQueryConnectionFields = defineSlot(

{authenticationHelp(authMethod, target)}

- {authMethod !== "adc" && secureCredentialAvailable ? ( + {authMethod === "none" && !endpointIsLocal ? ( +

+ Set a localhost API endpoint under Advanced options, or choose another authentication + method. +

+ ) : null} + + {!CREDENTIAL_FREE_METHODS.includes(authMethod) && secureCredentialAvailable ? ( {authMethod === "access_token" ? ( ) => setExtraField("query_timeout_seconds", event.target.value) } @@ -187,6 +214,87 @@ const BigQueryConnectionFields = defineSlot( style={styles.input} /> + +

+ Keep the query timeout at or below the Tabularis plugin call timeout (120 seconds by + default) or raise that setting too. +

+ +
+ Advanced options +
+
+ + ) => + setExtraField("billing_project", event.target.value) + } + placeholder="Defaults to the project ID" + autoComplete="off" + spellCheck={false} + style={styles.input} + /> + + + ) => + setExtraField("additional_projects", event.target.value) + } + placeholder="bigquery-public-data, other-project" + autoComplete="off" + spellCheck={false} + style={styles.input} + /> + +
+

+ Optional. Jobs run and are billed in the billing project. Additional projects are + comma-separated and browsed alongside the project ID. +

+ + ) => + setExtraField("api_endpoint", event.target.value) + } + placeholder="https://bigquery.googleapis.com" + autoComplete="off" + spellCheck={false} + style={styles.input} + /> + +

+ Optional. Use a Private Service Connect endpoint (https) or a local emulator such as + http://localhost:9050. Plain http is only accepted for localhost. +

+ + ) => + setExtraField("impersonate_service_account", event.target.value) + } + placeholder="reader@project.iam.gserviceaccount.com" + autoComplete="off" + spellCheck={false} + style={styles.input} + /> + +

+ Optional. The selected credential requests short-lived tokens for this service + account; it needs the Service Account Token Creator role on it. +

+
+
); }, diff --git a/ui/src/styles.ts b/ui/src/styles.ts index 436cfe5..1965559 100644 --- a/ui/src/styles.ts +++ b/ui/src/styles.ts @@ -62,6 +62,26 @@ export const styles = { fontSize: "11px", padding: 0, } as CSSProperties, + warning: { + color: "var(--color-text-danger, #f87171)", + fontSize: "11px", + lineHeight: 1.5, + margin: "-6px 0 0", + } as CSSProperties, + details: { + border: "1px solid var(--color-border-strong, #475569)", + borderRadius: "8px", + padding: "8px 12px", + } as CSSProperties, + summary: { + color: "var(--color-text-muted, #94a3b8)", + cursor: "pointer", + fontSize: "11px", + fontWeight: 600, + letterSpacing: "0.06em", + textTransform: "uppercase", + marginBottom: "4px", + } as CSSProperties, grid: { display: "grid", gridTemplateColumns: "1fr 1fr",