diff --git a/Cargo.lock b/Cargo.lock index d32301556d..37daab2710 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5419,6 +5419,96 @@ dependencies = [ "slab", ] +[[package]] +name = "gcloud-auth" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b43924e3df02cb3b846ca66a7ee58e8c13eb2556d0308c71f6154083f6980365" +dependencies = [ + "async-trait", + "base64 0.22.1", + "gcloud-metadata", + "home", + "jsonwebtoken 10.4.0", + "reqwest 0.13.5", + "serde", + "serde_json", + "thiserror 2.0.20", + "time", + "token-source", + "tokio", + "tracing", + "urlencoding", +] + +[[package]] +name = "gcloud-bigquery" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17c9e8ff7894d451cf5b704772326b5fbf112f75b8f229700197c923a763d99f" +dependencies = [ + "anyhow", + "arrow", + "async-stream", + "async-trait", + "backon", + "base64 0.22.1", + "bigdecimal", + "gcloud-auth", + "gcloud-gax", + "gcloud-googleapis", + "num-bigint", + "prost-types", + "reqwest 0.13.5", + "reqwest-middleware", + "serde", + "serde_json", + "thiserror 2.0.20", + "time", + "token-source", + "tokio", + "tracing", +] + +[[package]] +name = "gcloud-gax" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "558634081d1cf458d1bc0166f69475a80f445b5b991f162028b293a38307059d" +dependencies = [ + "http 1.5.0", + "thiserror 2.0.20", + "token-source", + "tokio", + "tokio-retry2", + "tonic", + "tower", + "tracing", +] + +[[package]] +name = "gcloud-googleapis" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d79fc4ec36213a21b734a2ce204176eb39c439fa528f9a388c9a1349ca344b0d" +dependencies = [ + "prost", + "prost-types", + "tonic", + "tonic-prost", +] + +[[package]] +name = "gcloud-metadata" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd3152612316be627be52fe9ca72331eb48425059b3a6a700e7adde223e061d5" +dependencies = [ + "reqwest 0.13.5", + "thiserror 2.0.20", + "tokio", +] + [[package]] name = "generator" version = "0.8.9" @@ -7125,6 +7215,36 @@ dependencies = [ "uuid", ] +[[package]] +name = "iggy_connector_bigquery_sink" +version = "0.1.0" +dependencies = [ + "arrow", + "async-trait", + "axum", + "base64 0.23.1", + "gcloud-bigquery", + "gcloud-gax", + "gcloud-googleapis", + "hex", + "humantime", + "iggy_common", + "iggy_connector_sdk", + "reqwest 0.13.5", + "rustls", + "secrecy", + "serde", + "serde_json", + "simd-json", + "token-source", + "tokio", + "tokio-stream", + "toml 1.1.6+spec-1.1.0", + "tonic", + "tonic-prost", + "tracing", +] + [[package]] name = "iggy_connector_clickhouse_sink" version = "0.2.1-edge.1" @@ -11658,6 +11778,7 @@ dependencies = [ "js-sys", "log", "mime", + "mime_guess", "percent-encoding", "pin-project-lite", "quinn", @@ -14086,6 +14207,15 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fd3ca314f692efd6c868f8408f53fe444634a845f96c028b97d35f6a1f79f0ee" +[[package]] +name = "token-source" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75746ae15bef509f21039a652383104424208fdae172a964a8930858b9a78412" +dependencies = [ + "async-trait", +] + [[package]] name = "tokio" version = "1.53.1" @@ -14140,6 +14270,16 @@ dependencies = [ "whoami", ] +[[package]] +name = "tokio-retry2" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e05e40fd65880de11383bbc693d3c636045ed1b5010ef1265a1d2b41d73334a1" +dependencies = [ + "pin-project", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.5" @@ -14311,6 +14451,7 @@ dependencies = [ "axum", "base64 0.22.1", "bytes", + "flate2", "h2 0.4.19", "http 1.5.0", "http-body 1.1.0", @@ -14323,11 +14464,13 @@ dependencies = [ "socket2 0.6.5", "sync_wrapper", "tokio", + "tokio-rustls", "tokio-stream", "tower", "tower-layer", "tower-service", "tracing", + "webpki-roots 1.0.9", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 82267292f2..5c54ec1767 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ members = [ "core/configs_derive", "core/connectors/runtime", "core/connectors/sdk", + "core/connectors/sinks/bigquery_sink", "core/connectors/sinks/clickhouse_sink", "core/connectors/sinks/delta_sink", "core/connectors/sinks/doris_sink", @@ -186,6 +187,13 @@ futures-core = { version = "0.3.34", default-features = false } # Keep on the minor compio-tls 0.10 uses; the message_bus tripwire test catches drift. futures-rustls = "0.26.0" futures-util = "0.3.34" +gcloud-bigquery = { version = "1.7.0", default-features = false, features = [ + "auth", + "jwt-rust-crypto", + "rustls-tls", +] } +gcloud-gax = "1.4.0" +gcloud-googleapis = { version = "1.3.0", features = ["bigquery"] } getrandom = { version = "0.4", features = ["wasm_js"] } git2 = { version = "0.21.0", default-features = false, features = ["vendored-libgit2"] } gloo = "0.12" @@ -341,12 +349,16 @@ test-case = "3.3.1" testcontainers = { version = "0.27.3", features = ["reusable-containers"] } testcontainers-modules = { version = "0.15.0", features = ["postgres", "http_wait"] } thiserror = "2.0.20" +token-source = "1.0.0" tokio = { version = "1.53.1", features = ["full"] } tokio-postgres = "0.7.18" tokio-rustls = "0.26.5" +tokio-stream = { version = "0.1.19", features = ["net"] } tokio-tungstenite = { version = "0.30", features = ["rustls-tls-webpki-roots"] } tokio-util = { version = "0.7.19", features = ["compat"] } toml = "1.1.6" +tonic = { version = "0.14.6", default-features = false, features = ["codegen", "router", "server"] } +tonic-prost = "0.14.6" tower-http = { version = "0.7.1", features = ["add-extension", "cors", "trace"] } tracing = "0.1.44" tracing-appender = "0.2.5" diff --git a/core/connectors/runtime/example_config/connectors/bigquery_sink.toml b/core/connectors/runtime/example_config/connectors/bigquery_sink.toml new file mode 100644 index 0000000000..188b118f67 --- /dev/null +++ b/core/connectors/runtime/example_config/connectors/bigquery_sink.toml @@ -0,0 +1,63 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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. + +type = "sink" +key = "bigquery" +enabled = true +version = 0 +name = "BigQuery sink" +path = "target/release/libiggy_connector_bigquery_sink" +verbose = false + +[[streams]] +stream = "example_stream" +topics = ["example_topic"] +schema = "json" +batch_length = 1000 +poll_interval = "5ms" +consumer_group = "bigquery_sink_connector" + +[plugin_config] +project_id = "my-project" +dataset = "iggy" +table = "events" + +# "mapped" (default): top-level JSON fields map to table columns by name. +# "raw": the whole payload goes into payload_column (JSON, STRING or BYTES). +mode = "mapped" +# payload_column = "payload" + +# Credentials: a key file, an inline key, or neither for Application +# Default Credentials (GOOGLE_APPLICATION_CREDENTIALS, gcloud, metadata server). +# credentials_path = "/secrets/bigquery-writer.json" +# credentials_json = "{...}" + +# Writes iggy_stream, iggy_topic, iggy_partition_id, iggy_offset, +# iggy_timestamp and iggy_id. The table must have these columns. +include_metadata = true +# Writes message headers into an iggy_headers JSON or STRING column. +include_headers = false + +# Columns no row in a request sets: "default" (column default, else NULL) or "null". +missing_value = "default" + +max_request_bytes = 8388608 +max_retries = 3 +retry_delay = "1s" +max_retry_delay = "30s" +timeout = "30s" +verbose_logging = false diff --git a/core/connectors/sinks/README.md b/core/connectors/sinks/README.md index 83f30a5aca..70a28940ac 100644 --- a/core/connectors/sinks/README.md +++ b/core/connectors/sinks/README.md @@ -8,6 +8,7 @@ Sink connectors are responsible for writing data from Iggy streams to external s | Sink | Description | | ---- | ----------- | +| **bigquery_sink** | Writes messages to Google BigQuery tables through the Storage Write API | | **doris_sink** | Loads JSON messages into Apache Doris tables via the Stream Load HTTP API | | **elasticsearch_sink** | Sends messages to Elasticsearch indices for full-text search and analytics | | **iceberg_sink** | Writes data to Apache Iceberg tables via REST catalog with S3/GCS/Azure storage | diff --git a/core/connectors/sinks/bigquery_sink/Cargo.toml b/core/connectors/sinks/bigquery_sink/Cargo.toml new file mode 100644 index 0000000000..2f6f23063d --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/Cargo.toml @@ -0,0 +1,64 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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. + +[package] +name = "iggy_connector_bigquery_sink" +version = "0.1.0" +description = "Iggy BigQuery sink connector for streaming messages into Google BigQuery via the Storage Write API" +edition = "2024" +license = "Apache-2.0" +keywords = ["iggy", "messaging", "streaming", "bigquery", "sink"] +categories = ["command-line-utilities", "database", "network-programming"] +homepage = "https://iggy.apache.org" +documentation = "https://iggy.apache.org/docs" +repository = "https://github.com/apache/iggy" +readme = "../../README.md" +publish = false + +[lib] +crate-type = ["cdylib", "lib"] + +[dependencies] +arrow = { workspace = true } +async-trait = { workspace = true } +base64 = { workspace = true } +gcloud-bigquery = { workspace = true } +gcloud-gax = { workspace = true } +gcloud-googleapis = { workspace = true } +hex = { workspace = true } +humantime = { workspace = true } +iggy_connector_sdk = { workspace = true } +reqwest = { workspace = true } +rustls = { workspace = true } +secrecy = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +simd-json = { workspace = true } +token-source = { workspace = true } +tracing = { workspace = true } + +[dev-dependencies] +axum = { workspace = true } +iggy_common = { workspace = true } +tokio = { workspace = true } +tokio-stream = { workspace = true } +toml = { workspace = true } +tonic = { workspace = true } +tonic-prost = { workspace = true } + +[lints] +workspace = true diff --git a/core/connectors/sinks/bigquery_sink/README.md b/core/connectors/sinks/bigquery_sink/README.md new file mode 100644 index 0000000000..d2095713c0 --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/README.md @@ -0,0 +1,263 @@ +# BigQuery Sink Connector + +The BigQuery sink connector consumes messages from Iggy topics and writes them into a Google BigQuery table through the [Storage Write API](https://cloud.google.com/bigquery/docs/write-api) `_default` stream, with rows encoded as Arrow record batches. + +## Features + +- **Storage Write API**: gRPC `AppendRows` on the table's `_default` stream. No legacy `insertAll`. +- **Two modes**: `mapped` maps top-level JSON fields to table columns by name. `raw` stores the whole payload in one JSON, STRING or BYTES column. +- **Schema validation at startup**: the table schema is read with `tables.get` and checked before any data flows. Unsupported column types fail `open()`. +- **Iggy metadata columns**: stream, topic, partition, offset, timestamp and message id, plus optional headers, for lineage and deduplication. +- **Per-row error handling**: a row that cannot be written is dropped and logged with its offset. The rest of the batch is still written. +- **Request sizing**: batches are split below the 10 MB `AppendRows` limit. +- **Retries**: transient gRPC failures are retried with jittered exponential backoff. + +## Configuration + +```toml +type = "sink" +key = "bigquery" +enabled = true +version = 0 +name = "BigQuery sink" +path = "target/release/libiggy_connector_bigquery_sink" + +[[streams]] +stream = "example_stream" +topics = ["example_topic"] +schema = "json" +batch_length = 1000 +poll_interval = "5ms" +consumer_group = "bigquery_sink_connector" + +[plugin_config] +project_id = "my-project" +dataset = "iggy" +table = "events" +mode = "mapped" +credentials_path = "/secrets/bigquery-writer.json" +include_metadata = true +``` + +## Configuration Options + +| Option | Type | Default | Description | +| ------ | ---- | ------- | ----------- | +| `project_id` | string | required | Google Cloud project that owns the dataset | +| `dataset` | string | required | BigQuery dataset | +| `table` | string | required | Target table. It must already exist | +| `mode` | string | `"mapped"` | `mapped` or `raw` | +| `payload_column` | string | `"payload"` | Payload column in `raw` mode | +| `credentials_path` | string | none | Service account key file | +| `credentials_json` | string | none | Inline service account key. Mutually exclusive with `credentials_path` | +| `include_metadata` | bool | `true` | Write the `iggy_*` metadata columns listed below | +| `include_headers` | bool | `false` | Write message headers into `iggy_headers` | +| `missing_value` | string | `"default"` | How BigQuery fills columns no row in a request sets: `default` or `null` | +| `max_request_bytes` | usize | `8388608` | Upper bound for one `AppendRows` request, clamped to 64 KiB..9 MiB | +| `max_retries` | u32 | `3` | Total attempts per call, including the first | +| `retry_delay` | string | `"1s"` | Base retry delay | +| `max_retry_delay` | string | `"30s"` | Upper bound for one retry delay | +| `timeout` | string | `"30s"` | Timeout for `tables.get` and each `AppendRows` call | +| `verbose_logging` | bool | `false` | Log each batch at info level instead of debug | +| `endpoint` | string | none | REST endpoint override, for local fakes and emulators. Must be set with `grpc_endpoint`; no credentials are used | +| `grpc_endpoint` | string | none | Storage Write API `host:port` override | + +## Authentication + +The connector picks the first source that applies: + +1. `endpoint` and `grpc_endpoint`: no credentials, plain-text connections. For tests only. +2. `credentials_path`: a service account key file. +3. `credentials_json`: an inline service account key. +4. Application Default Credentials: `GOOGLE_APPLICATION_CREDENTIALS`, `gcloud auth application-default login`, or the metadata server on GCE, GKE and Cloud Run. + +```toml +# Key file +credentials_path = "/secrets/bigquery-writer.json" + +# Inline key, for example injected through an environment override: +# IGGY_CONNECTORS_SINK_BIGQUERY_PLUGIN_CONFIG_CREDENTIALS_JSON='{"type":"service_account",...}' +``` + +### IAM permissions + +The identity needs `bigquery.tables.get` (read the schema) and `bigquery.tables.updateData` (append rows) on the target table. The predefined role `roles/bigquery.dataEditor`, granted on the table or its dataset, includes both. + +## Modes + +### `mapped` (default) + +Each payload must be a JSON object. Top-level fields map to table columns by name. Fields that match no column are ignored. + +- **Payload types**: `Json`, `Text` and `Proto` text are accepted when they hold a JSON object. `Raw`, `FlatBuffer` and `Avro` payloads are rejected. +- **Missing columns**: a column that no row in a request sets is left out of the request, and BigQuery fills it according to `missing_value`. A REQUIRED column without a default value must be present in every row. +- **Defaults apply per request, not per row**: Arrow has no "unset" marker, so once one row in a request sets a column, the other rows send NULL for it. For a REQUIRED column with a default, rows that omit it while other rows in the same batch set it are rejected. Set such columns in every message or in none. +- **REPEATED** columns that are absent or `null` are written as empty arrays. + +### `raw` + +The payload is stored in `payload_column`. Every other table column is left to its default, so a REQUIRED column without a default fails `open()`. + +| Payload column type | Accepted payloads | +| ------------------- | ----------------- | +| `JSON` | `Json`, plus `Text`, `Proto` and UTF-8 `Raw` payloads that parse as JSON | +| `STRING` | `Json` (serialized), `Text`, `Proto`, UTF-8 `Raw` | +| `BYTES` | every payload type, as its raw bytes | + +## Metadata Columns + +With `include_metadata = true` the table must contain these columns. A missing column fails `open()` with the DDL to add it. + +| Column | BigQuery type | Value | +| ------ | ------------- | ----- | +| `iggy_stream` | `STRING` | Stream name | +| `iggy_topic` | `STRING` | Topic name | +| `iggy_partition_id` | `INT64` | Partition id | +| `iggy_offset` | `INT64` | Message offset | +| `iggy_timestamp` | `TIMESTAMP` | Message timestamp | +| `iggy_id` | `STRING` | Message id, as a decimal string (it is a u128) | +| `iggy_headers` | `JSON` or `STRING` | Headers, only with `include_headers = true`. Binary values are written as `{"data": "", "iggy_header_encoding": "base64"}` | + +## Type Mapping + +| BigQuery | Accepted JSON values | +| -------- | -------------------- | +| `STRING`, `GEOGRAPHY` (WKT) | string | +| `BYTES` | base64 string | +| `INT64` | integer, or a string holding one | +| `FLOAT64` | number, or a string holding one | +| `BOOL` | `true` / `false` | +| `NUMERIC`, `BIGNUMERIC` | number or string. Prefer strings to avoid float rounding | +| `TIMESTAMP` | RFC 3339 string, or an integer in microseconds since the epoch | +| `DATETIME` | `YYYY-MM-DDTHH:MM:SS[.ffffff]` string | +| `DATE` | `YYYY-MM-DD` string | +| `TIME` | `HH:MM:SS[.ffffff]` string | +| `JSON` | any JSON value, stored as that value | +| `RECORD` | object | +| `REPEATED` | array | + +`INTERVAL` and `RANGE` columns are not supported and fail `open()`. + +## Reliability + +### Errors and retries + +- `UNAVAILABLE`, `DEADLINE_EXCEEDED`, `INTERNAL`, `ABORTED` and `RESOURCE_EXHAUSTED` are retried up to `max_retries` attempts in total. +- Every other gRPC error, including `INVALID_ARGUMENT`, `PERMISSION_DENIED` and `UNAUTHENTICATED`, fails the request without a retry. +- `tables.get` retries HTTP 429, 5xx and transport errors. A 403 or 404 fails `open()` immediately. + +### Bad rows + +A message that cannot become a row (wrong payload type, a value that does not fit its column, a missing REQUIRED value, a row larger than `max_request_bytes`) is dropped before the request is sent. When BigQuery reports row-level errors, it appends nothing from that request: the connector drops the reported rows and appends the remaining rows once more. If that second append also reports row errors, the remaining rows of that request are dropped and the batch returns an error. + +Every dropped row is logged at `warn` with its stream, topic, partition and offset. There is no dead-letter queue. On shutdown the connector logs rows written, rejected and failed. + +### Delivery semantics + +The connector runtime commits consumer group offsets when messages are polled, before `consume()` runs, and a batch for which `consume()` returns an error is not redelivered. As a result: + +- A batch that still fails after the last retry, or hits a permanent error, is lost. +- A process crash between the offset commit and a successful append also loses that batch. +- The `_default` stream has no append offsets. When an append result is lost and the request is retried, rows can be written twice. Deduplicate downstream on `iggy_stream`, `iggy_topic`, `iggy_partition_id` and `iggy_offset`, for example: + +```sql +SELECT * EXCEPT(rn) FROM ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY iggy_stream, iggy_topic, iggy_partition_id, iggy_offset + ORDER BY iggy_timestamp + ) AS rn + FROM `my-project.iggy.events` +) WHERE rn = 1 +``` + +### Schema changes + +The schema is read once in `open()`. Restart the connector after changing the table. + +## Testing + +Unit tests cover configuration, credential source selection, schema mapping, Arrow encoding, request splitting and error classification. + +`tests/bigquery_sink.rs` runs the sink against an in-process fake (`tests/common/mod.rs`): an axum `tables.get` endpoint and a tonic `BigQueryWrite` service that records every `AppendRows` request as an Arrow batch and answers from a script. It covers successful writes, request splitting, local and BigQuery-reported row errors, retryable and permanent failures, and `open()` failures. No Google Cloud access is needed. + +The fake cannot check what only BigQuery knows: that it accepts the Arrow encoding, its type coercion, and real credentials. Verify those manually as follows. + +### Manual verification against BigQuery + +1. Create a dataset and table (requires `gcloud` and `bq`): + + ```bash + export PROJECT=my-project + gcloud auth application-default login + gcloud services enable bigquery.googleapis.com bigquerystorage.googleapis.com --project "$PROJECT" + bq --project_id="$PROJECT" mk --dataset iggy + bq --project_id="$PROJECT" query --use_legacy_sql=false " + CREATE TABLE iggy.events ( + user_id INT64 NOT NULL, + event STRING, + amount NUMERIC, + created TIMESTAMP DEFAULT CURRENT_TIMESTAMP(), + attrs JSON, + iggy_stream STRING, + iggy_topic STRING, + iggy_partition_id INT64, + iggy_offset INT64, + iggy_timestamp TIMESTAMP, + iggy_id STRING + )" + ``` + +2. Build the binaries and the plugin, and write a runtime config that loads only this connector: + + ```bash + cargo build --bin iggy-server --bin iggy --bin iggy-connectors + cargo build -p iggy_connector_bigquery_sink + mkdir -p /tmp/bq-sink-test/connectors + sed 's#^config_dir = .*#config_dir = "/tmp/bq-sink-test/connectors"#' \ + core/connectors/runtime/example_config/config.toml > /tmp/bq-sink-test/config.toml + cat > /tmp/bq-sink-test/connectors/bigquery.toml <), +} + +pub(crate) struct BigQueryClient { + http: reqwest::Client, + table_url: String, + token: Option>, + stream: Option, + write_client: Client, + table_path: String, + missing_value: i32, +} + +impl CredentialSource { + pub(crate) fn from_config(config: &BigQuerySinkConfig) -> Self { + if let (Some(rest), Some(grpc)) = (&config.endpoint, &config.grpc_endpoint) { + return CredentialSource::Emulator { + rest: rest.trim_end_matches('/').to_owned(), + grpc: grpc.clone(), + }; + } + if let Some(path) = &config.credentials_path { + return CredentialSource::KeyFile(path.clone()); + } + if let Some(json) = &config.credentials_json { + return CredentialSource::InlineKey(json.clone()); + } + CredentialSource::ApplicationDefault + } + + pub(crate) fn label(&self) -> &'static str { + match self { + CredentialSource::Emulator { .. } => "none (endpoint override)", + CredentialSource::KeyFile(_) => "service account key file", + CredentialSource::InlineKey(_) => "inline service account key", + CredentialSource::ApplicationDefault => "application default credentials", + } + } +} + +impl BigQueryClient { + /// Build credentials and the gRPC channel pool. No BigQuery API call is + /// made yet. + pub(crate) async fn connect( + config: &BigQuerySinkConfig, + settings: &Settings, + ) -> Result { + install_crypto_provider(); + let source = CredentialSource::from_config(config); + info!("BigQuery sink: using {}", source.label()); + + let write_config = StreamingWriteConfig::default() + .with_channel_config(ChannelConfig::default().with_timeout(settings.timeout)); + + let (client_config, rest_base, token) = match &source { + CredentialSource::Emulator { rest, grpc } => ( + ClientConfig::new_with_emulator(grpc, rest.clone()), + rest.clone(), + None, + ), + CredentialSource::KeyFile(path) => { + let credentials = + CredentialsFile::new_from_file(path.clone()) + .await + .map_err(|e| { + Error::InitError(format!("cannot load credentials_path '{path}': {e}")) + })?; + with_credentials(credentials).await? + } + CredentialSource::InlineKey(json) => { + let credentials = CredentialsFile::new_from_str(json.expose_secret()) + .await + .map_err(|e| Error::InitError(format!("cannot parse credentials_json: {e}")))?; + with_credentials(credentials).await? + } + CredentialSource::ApplicationDefault => { + let (client_config, _) = ClientConfig::new_with_auth().await.map_err(|e| { + Error::InitError(format!("application default credentials: {e}")) + })?; + let provider = HttpClientConfig::default_token_provider() + .await + .map_err(|e| { + Error::InitError(format!("application default credentials: {e}")) + })?; + ( + client_config, + BIGQUERY_REST_ENDPOINT.to_owned(), + Some(provider.token_source()), + ) + } + }; + + let write_client = Client::new(client_config.with_streaming_write_config(write_config)) + .await + .map_err(|e| { + Error::InitError(format!("cannot connect to the Storage Write API: {e}")) + })?; + let http = reqwest::Client::builder() + .timeout(settings.timeout) + .build() + .map_err(|e| Error::InitError(format!("cannot build HTTP client: {e}")))?; + + Ok(BigQueryClient { + http, + table_url: format!( + "{rest_base}/bigquery/v2/projects/{}/datasets/{}/tables/{}?fields=schema", + config.project_id, config.dataset, config.table + ), + token, + stream: None, + write_client, + table_path: format!( + "projects/{}/datasets/{}/tables/{}", + config.project_id, config.dataset, config.table + ), + missing_value: settings.missing_value.as_proto(), + }) + } + + /// `tables.get`, returning the raw response body. + pub(crate) async fn fetch_table(&self) -> Result, TableError> { + let mut request = self.http.get(&self.table_url); + if let Some(token) = &self.token { + let value = token + .token() + .await + .map_err(|e| TableError::Token(e.to_string()))?; + request = request.header(AUTHORIZATION, value); + } + let response = request + .send() + .await + .map_err(|e| TableError::Transport(e.to_string()))?; + let status = response.status(); + let body = response + .bytes() + .await + .map_err(|e| TableError::Transport(e.to_string()))?; + if !status.is_success() { + return Err(TableError::Http { + status: status.as_u16(), + body: error_message(&body), + }); + } + Ok(body.to_vec()) + } + + /// Resolve the table's `_default` write stream. Also proves the + /// credentials can reach the Storage Write API. + pub(crate) async fn resolve_write_stream(&self) -> Result { + Ok(self + .write_client + .default_storage_writer() + .create_write_stream(&self.table_path) + .await?) + } + + pub(crate) fn set_write_stream(&mut self, stream: DefaultStream) { + self.stream = Some(stream); + } + + /// One `AppendRows` request for one chunk. + pub(crate) async fn append(&self, chunk: &Chunk) -> Result { + let stream = self.stream.as_ref().ok_or(AppendError::NoResponse)?; + let request = AppendRowsRequestBuilder::new_arrow( + chunk.schema_bytes.clone(), + chunk.batch_bytes.clone(), + ) + .with_default_missing_value_interpretation(self.missing_value); + let mut responses = stream.append_rows(vec![request]).await?; + let response = responses.message().await?.ok_or(AppendError::NoResponse)?; + + if !response.row_errors.is_empty() { + let rows = response + .row_errors + .into_iter() + .map(|row| { + ( + usize::try_from(row.index).unwrap_or(usize::MAX), + row.message, + ) + }) + .collect(); + return Ok(AppendOutcome::RowErrors(rows)); + } + match response.response { + Some(Response::Error(status)) => Err(AppendError::Rpc { + code: status.code.into(), + message: status.message, + }), + Some(Response::AppendResult(_)) | None => Ok(AppendOutcome::Appended), + } + } +} + +impl fmt::Debug for BigQueryClient { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BigQueryClient") + .field("table_path", &self.table_path) + .field("stream_open", &self.stream.is_some()) + .finish_non_exhaustive() + } +} + +async fn with_credentials( + credentials: CredentialsFile, +) -> Result<(ClientConfig, String, Option>), Error> { + let (client_config, _) = ClientConfig::new_with_credentials(credentials.clone()) + .await + .map_err(|e| Error::InitError(format!("service account credentials: {e}")))?; + let provider: DefaultTokenSourceProvider = + HttpClientConfig::default_token_provider_with(credentials) + .await + .map_err(|e| Error::InitError(format!("service account credentials: {e}")))?; + Ok(( + client_config, + BIGQUERY_REST_ENDPOINT.to_owned(), + Some(provider.token_source()), + )) +} + +/// Google APIs answer errors with a pretty-printed JSON document. Keep the +/// log on one line by pulling out `error.message`, falling back to the body +/// with line breaks collapsed. +fn error_message(body: &[u8]) -> String { + serde_json::from_slice::(body) + .ok() + .and_then(|value| { + value + .pointer("/error/message") + .and_then(serde_json::Value::as_str) + .map(str::to_owned) + }) + .unwrap_or_else(|| { + String::from_utf8_lossy(body) + .split_whitespace() + .collect::>() + .join(" ") + }) +} + +/// tonic builds its TLS config from rustls' process-wide default provider. +/// With more than one provider compiled in there is no default and the +/// first TLS connection panics, so pick one explicitly. The plugin is its +/// own shared library, so this only affects its own copy of rustls. +fn install_crypto_provider() { + let _ = rustls::crypto::ring::default_provider().install_default(); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::config; + + #[test] + fn given_no_credentials_should_use_application_default() { + assert!(matches!( + CredentialSource::from_config(&config()), + CredentialSource::ApplicationDefault + )); + } + + #[test] + fn given_key_file_should_use_it() { + let mut config = config(); + config.credentials_path = Some("/secrets/key.json".into()); + assert!(matches!( + CredentialSource::from_config(&config), + CredentialSource::KeyFile(path) if path == "/secrets/key.json" + )); + } + + #[test] + fn given_inline_key_should_use_it() { + let mut config = config(); + config.credentials_json = Some(SecretString::from("{\"type\":\"service_account\"}")); + let CredentialSource::InlineKey(secret) = CredentialSource::from_config(&config) else { + panic!("expected inline key"); + }; + assert_eq!(secret.expose_secret(), "{\"type\":\"service_account\"}"); + } + + #[test] + fn given_endpoint_overrides_should_skip_credentials() { + let mut config = config(); + config.endpoint = Some("http://127.0.0.1:9050/".into()); + config.grpc_endpoint = Some("127.0.0.1:9060".into()); + config.credentials_path = Some("/ignored.json".into()); + let CredentialSource::Emulator { rest, grpc } = CredentialSource::from_config(&config) + else { + panic!("expected endpoint override"); + }; + assert_eq!(rest, "http://127.0.0.1:9050"); + assert_eq!(grpc, "127.0.0.1:9060"); + } + + #[test] + fn given_google_error_body_should_extract_single_line_message() { + let body = br#"{ + "error": { + "code": 404, + "message": "Not found: Table proj:iggy.missing", + "status": "NOT_FOUND" + } +}"#; + assert_eq!(error_message(body), "Not found: Table proj:iggy.missing"); + assert_eq!(error_message(b"bad\ngateway"), "bad gateway"); + } + + #[tokio::test] + async fn given_missing_key_file_should_fail_to_connect() { + let mut config = config(); + config.credentials_path = Some("/definitely/not/here.json".into()); + let settings = Settings::from_config(&config); + let error = BigQueryClient::connect(&config, &settings) + .await + .unwrap_err(); + assert!( + matches!(&error, Error::InitError(reason) if reason.contains("credentials_path")), + "{error:?}" + ); + } + + #[tokio::test] + async fn given_invalid_inline_key_should_fail_to_connect() { + let mut config = config(); + config.credentials_json = Some(SecretString::from("not json")); + let settings = Settings::from_config(&config); + let error = BigQueryClient::connect(&config, &settings) + .await + .unwrap_err(); + let Error::InitError(reason) = &error else { + panic!("expected InitError, got {error:?}"); + }; + assert!(reason.contains("credentials_json")); + assert!(!reason.contains("not json"), "secret leaked: {reason}"); + } +} diff --git a/core/connectors/sinks/bigquery_sink/src/encode.rs b/core/connectors/sinks/bigquery_sink/src/encode.rs new file mode 100644 index 0000000000..344491fa64 --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/src/encode.rs @@ -0,0 +1,1289 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! Turns a run of `ConsumedMessage`s into Arrow record batches that fit in +//! one `AppendRows` request each. +//! +//! A message that cannot become a row (wrong payload type, a value that does +//! not fit its column, a row larger than the request budget) is rejected on +//! its own and reported back with its offset. The rest of the batch is still +//! written. +//! +//! Mapped mode decodes JSON objects with `arrow-json`. The fast path decodes +//! the whole run at once. When that fails, bisection isolates the offending +//! rows before the survivors are decoded together. +//! +//! The writer schema of a request only contains the columns that at least one +//! row in the run sets. Columns no row mentions are left out, so BigQuery +//! fills them according to `missing_value` (the column default, or NULL). + +use crate::schema::{BqType, Column, MetaColumn, Mode, RawKind, TableLayout, UTC}; +use arrow::array::{ + ArrayRef, BinaryArray, Int64Array, RecordBatch, StringArray, StringBuilder, + TimestampMicrosecondArray, +}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; +use arrow::error::ArrowError; +use arrow::ipc::writer::{ + CompressionContext, DictionaryTracker, IpcDataGenerator, IpcWriteOptions, write_message, +}; +use arrow::json::ReaderBuilder; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64; +use iggy_connector_sdk::{ConsumedMessage, MessagesMetadata, Payload, TopicMetadata}; +use simd_json::{OwnedValue, StaticNode}; +use std::fmt::Write as _; +use std::sync::Arc; + +const HEADER_ENCODING_BASE64: &str = "base64"; +const UTC_OFFSET: &str = "+00:00"; +const IPC_BATCH_OVERHEAD: usize = 1024; +const IPC_COLUMN_OVERHEAD: usize = 128; + +/// One `AppendRows` request worth of rows. +#[derive(Debug)] +pub(crate) struct Chunk { + pub batch: RecordBatch, + /// Iggy offset of each row, same order as `batch`. + pub offsets: Vec, + pub schema_bytes: Vec, + pub batch_bytes: Vec, +} + +/// A message that was not turned into a row. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Rejected { + pub offset: u64, + pub reason: String, +} + +#[derive(Debug, Default)] +pub(crate) struct Encoded { + pub chunks: Vec, + pub rejected: Vec, +} + +/// Per-run context shared by every row. +pub(crate) struct RunContext<'a> { + pub layout: &'a TableLayout, + pub topic: &'a TopicMetadata, + pub messages: &'a MessagesMetadata, + pub max_request_bytes: usize, +} + +/// Encode a run of messages. Fails only on an internal Arrow error, which +/// means the batch as a whole cannot be built. Per-row problems end up in +/// `Encoded::rejected`. +pub(crate) fn encode( + ctx: &RunContext<'_>, + messages: Vec, +) -> Result { + let mut rejected = Vec::new(); + let batch = match &ctx.layout.raw { + Some((kind, field)) => encode_raw(ctx, *kind, field, messages, &mut rejected)?, + None => encode_mapped(ctx, messages, &mut rejected)?, + }; + let mut chunks = Vec::new(); + if let Some((batch, offsets)) = batch { + split_into_chunks( + batch, + offsets, + ctx.max_request_bytes, + &mut chunks, + &mut rejected, + )?; + } + Ok(Encoded { chunks, rejected }) +} + +/// Build a new chunk holding only the rows of `chunk` whose index is not in +/// `drop`. Used after BigQuery reports row-level errors. +pub(crate) fn without_rows(chunk: &Chunk, drop: &[usize]) -> Result, ArrowError> { + let keep: Vec = (0..chunk.batch.num_rows()) + .map(|row| !drop.contains(&row)) + .collect(); + let offsets: Vec = chunk + .offsets + .iter() + .zip(&keep) + .filter_map(|(offset, keep)| keep.then_some(*offset)) + .collect(); + if offsets.is_empty() { + return Ok(None); + } + let mask = arrow::array::BooleanArray::from(keep); + let batch = arrow::compute::filter_record_batch(&chunk.batch, &mask)?; + let (schema_bytes, batch_bytes) = encode_ipc(&batch)?; + Ok(Some(Chunk { + batch, + offsets, + schema_bytes, + batch_bytes, + })) +} + +// ─── Mapped mode ───────────────────────────────────────────────────────────── + +fn encode_mapped( + ctx: &RunContext<'_>, + messages: Vec, + rejected: &mut Vec, +) -> Result)>, ArrowError> { + let layout = ctx.layout; + let mut rows = Vec::with_capacity(messages.len()); + let mut kept = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match json_object(payload).and_then(|mut row| { + normalize_row(&mut row, layout)?; + Ok(row) + }) { + Ok(row) => { + rows.push(row); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + if rows.is_empty() { + return Ok(None); + } + + let present = present_columns(&rows, layout); + if present.is_empty() && layout.metadata.is_empty() { + rejected.extend(kept.iter().map(|m| Rejected { + offset: m.offset, + reason: "payload has no field matching a table column".to_owned(), + })); + return Ok(None); + } + let payload_fields: Vec = present.iter().map(|&i| layout.fields[i].clone()).collect(); + + let (decoded, survivors) = decode_rows(&payload_fields, &rows, |row, reason| { + rejected.push(Rejected { + offset: kept[row].offset, + reason, + }); + })?; + if survivors.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = survivors.iter().map(|&row| &kept[row]).collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + + let mut fields = payload_fields; + let mut columns: Vec = decoded.map(|b| b.columns().to_vec()).unwrap_or_default(); + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Take the JSON object out of a payload. Text and proto text must parse as +/// a JSON object. +fn json_object(payload: Payload) -> Result { + let value = match payload { + Payload::Json(value) => value, + Payload::Text(text) | Payload::Proto(text) => { + let mut bytes = text.into_bytes(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload text is not valid JSON: {e}"))? + } + other => { + return Err(format!( + "{} payload is not supported in mapped mode", + payload_kind(&other) + )); + } + }; + match value { + OwnedValue::Object(_) => Ok(value), + _ => Err("payload is not a JSON object".to_owned()), + } +} + +/// Prepare one row for `arrow-json`: +/// - a REQUIRED column without a default must be present and non-null, +/// - REPEATED columns that are absent or null become `[]`, +/// - JSON column values become their JSON text, +/// - BYTES column values are base64 in JSON and become hex, which is what +/// `arrow-json` decodes into binary. +fn normalize_row(row: &mut OwnedValue, layout: &TableLayout) -> Result<(), String> { + let OwnedValue::Object(object) = row else { + return Err("payload is not a JSON object".to_owned()); + }; + for column in &layout.columns { + let value = object.get_mut(column.name.as_str()); + match (value, column.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(column.name.clone(), empty_array()); + } + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Required) + if !column.has_default => + { + return Err(format!( + "missing value for REQUIRED column '{}'", + column.name + )); + } + (Some(value), _) => normalize_value(value, column)?, + (None, _) => {} + } + } + Ok(()) +} + +fn normalize_value(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if column.mode == Mode::Repeated { + let OwnedValue::Array(items) = value else { + return Err(format!("column '{}' expects an array", column.name)); + }; + for item in items.iter_mut() { + normalize_scalar(item, column)?; + } + return Ok(()); + } + normalize_scalar(value, column) +} + +fn normalize_scalar(value: &mut OwnedValue, column: &Column) -> Result<(), String> { + if matches!(value, OwnedValue::Static(StaticNode::Null)) { + return Ok(()); + } + match &column.bq_type { + BqType::Json => { + let text = simd_json::to_string(value) + .map_err(|e| format!("column '{}': cannot serialize JSON: {e}", column.name))?; + *value = OwnedValue::String(text); + } + BqType::Bytes => { + let OwnedValue::String(text) = value else { + return Err(format!("column '{}' expects a base64 string", column.name)); + }; + let bytes = BASE64 + .decode(text.as_bytes()) + .map_err(|e| format!("column '{}': invalid base64: {e}", column.name))?; + *value = OwnedValue::String(hex::encode(bytes)); + } + BqType::Record(children) => { + let OwnedValue::Object(object) = value else { + return Ok(()); + }; + for child in children { + match (object.get_mut(child.name.as_str()), child.mode) { + (None | Some(OwnedValue::Static(StaticNode::Null)), Mode::Repeated) => { + object.insert(child.name.clone(), empty_array()); + } + (Some(child_value), _) => normalize_value(child_value, child)?, + (None, _) => {} + } + } + } + _ => {} + } + Ok(()) +} + +/// Indexes of the layout columns that at least one row sets, in table order. +fn present_columns(rows: &[OwnedValue], layout: &TableLayout) -> Vec { + let mut present = vec![false; layout.columns.len()]; + for row in rows { + if let OwnedValue::Object(object) = row { + for key in object.keys() { + if let Some(&index) = layout.index.get(key.as_str()) { + present[index] = true; + } + } + } + } + present + .iter() + .enumerate() + .filter_map(|(index, set)| set.then_some(index)) + .collect() +} + +/// Decode rows into one batch. Returns the batch (`None` when there are no +/// payload fields) and the indexes of the rows it holds. +fn decode_rows( + fields: &[FieldRef], + rows: &[OwnedValue], + on_reject: impl FnMut(usize, String), +) -> Result<(Option, Vec), ArrowError> { + decode_rows_with(fields, rows, on_reject, decode) +} + +fn decode_rows_with<'a>( + fields: &[FieldRef], + rows: &'a [OwnedValue], + mut on_reject: impl FnMut(usize, String), + mut decoder: impl FnMut(&SchemaRef, &[&'a OwnedValue]) -> Result, +) -> Result<(Option, Vec), ArrowError> { + if fields.is_empty() || rows.is_empty() { + return Ok((None, (0..rows.len()).collect())); + } + let target: SchemaRef = Arc::new(Schema::new(fields.to_vec())); + let schema: SchemaRef = Arc::new(Schema::new( + fields + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect::>(), + )); + let row_refs: Vec<&OwnedValue> = rows.iter().collect(); + let first_error = match decoder(&schema, &row_refs) { + Ok(batch) => { + return Ok((Some(to_target(batch, &target)?), (0..rows.len()).collect())); + } + Err(error) => error, + }; + + let mut valid = vec![true; rows.len()]; + isolate_invalid_rows( + &schema, + &row_refs, + 0, + first_error, + &mut valid, + &mut on_reject, + &mut decoder, + )?; + + let mut valid_rows = Vec::with_capacity(rows.len()); + let mut survivors = Vec::with_capacity(rows.len()); + for (index, row) in row_refs.into_iter().enumerate() { + if valid[index] { + valid_rows.push(row); + survivors.push(index); + } + } + if valid_rows.is_empty() { + return Ok((None, survivors)); + } + let batch = decoder(&schema, &valid_rows)?; + Ok((Some(to_target(batch, &target)?), survivors)) +} + +fn isolate_invalid_rows<'a>( + schema: &SchemaRef, + rows: &[&'a OwnedValue], + first_index: usize, + error: ArrowError, + valid: &mut [bool], + on_reject: &mut impl FnMut(usize, String), + decoder: &mut impl FnMut(&SchemaRef, &[&'a OwnedValue]) -> Result, +) -> Result<(), ArrowError> { + if rows.len() == 1 { + valid[first_index] = false; + on_reject( + first_index, + format!("row does not match the table schema: {error}"), + ); + return Ok(()); + } + + let midpoint = rows.len() / 2; + let (left, right) = rows.split_at(midpoint); + if let Err(error) = decoder(schema, left) { + isolate_invalid_rows(schema, left, first_index, error, valid, on_reject, decoder)?; + } + if let Err(error) = decoder(schema, right) { + isolate_invalid_rows( + schema, + right, + first_index + midpoint, + error, + valid, + on_reject, + decoder, + )?; + } + Ok(()) +} + +/// The type `arrow-json` decodes into. Without the `chrono-tz` feature it +/// only parses offset timezones, so `"UTC"` is decoded as `"+00:00"` and +/// relabelled afterwards. Both name the same instant, so the relabel is a +/// metadata change. +fn decode_type(data_type: &DataType) -> DataType { + match data_type { + DataType::Timestamp(unit, Some(_)) => DataType::Timestamp(*unit, Some(UTC_OFFSET.into())), + DataType::List(item) => DataType::List(Arc::new( + item.as_ref() + .clone() + .with_data_type(decode_type(item.data_type())), + )), + DataType::Struct(children) => DataType::Struct( + children + .iter() + .map(|f| { + f.as_ref() + .clone() + .with_data_type(decode_type(f.data_type())) + }) + .collect(), + ), + other => other.clone(), + } +} + +fn to_target(batch: RecordBatch, target: &SchemaRef) -> Result { + let columns = batch + .columns() + .iter() + .zip(target.fields()) + .map(|(column, field)| { + if column.data_type() == field.data_type() { + Ok(column.clone()) + } else { + cast(column, field.data_type()) + } + }) + .collect::, _>>()?; + RecordBatch::try_new(target.clone(), columns) +} + +fn decode(schema: &SchemaRef, rows: &[S]) -> Result { + let mut decoder = ReaderBuilder::new(schema.clone()) + .with_batch_size(rows.len().max(1)) + .build_decoder()?; + decoder.serialize(rows)?; + decoder + .flush()? + .ok_or_else(|| ArrowError::JsonError("decoder produced no rows".to_owned())) +} + +// ─── Raw mode ──────────────────────────────────────────────────────────────── + +fn encode_raw( + ctx: &RunContext<'_>, + kind: RawKind, + payload_field: &FieldRef, + messages: Vec, + rejected: &mut Vec, +) -> Result)>, ArrowError> { + let mut kept = Vec::with_capacity(messages.len()); + let payload_column: ArrayRef = match kind { + RawKind::Bytes => { + let mut values = Vec::with_capacity(messages.len()); + for message in messages { + match message.payload.try_to_bytes() { + Ok(bytes) => { + values.push(bytes); + kept.push(message); + } + Err(e) => rejected.push(Rejected { + offset: message.offset, + reason: e.to_string(), + }), + } + } + Arc::new(BinaryArray::from_iter_values(values)) + } + RawKind::Json | RawKind::String => { + let mut values = Vec::with_capacity(messages.len()); + for mut message in messages { + let payload = std::mem::replace(&mut message.payload, Payload::Raw(Vec::new())); + match raw_text(payload, kind) { + Ok(text) => { + values.push(text); + kept.push(message); + } + Err(reason) => rejected.push(Rejected { + offset: message.offset, + reason, + }), + } + } + Arc::new(StringArray::from_iter_values(values)) + } + }; + if kept.is_empty() { + return Ok(None); + } + let kept: Vec<&ConsumedMessage> = kept.iter().collect(); + let offsets = kept.iter().map(|m| m.offset).collect(); + let mut fields = vec![payload_field.clone()]; + let mut columns = vec![payload_column]; + append_metadata(ctx, &kept, &mut fields, &mut columns)?; + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)?; + Ok(Some((batch, offsets))) +} + +/// Payload as text for a JSON or STRING column. A JSON column only takes +/// text that parses as JSON. +fn raw_text(payload: Payload, kind: RawKind) -> Result { + let text = match payload { + Payload::Json(value) => { + return simd_json::to_string(&value).map_err(|e| format!("cannot serialize JSON: {e}")); + } + Payload::Text(text) | Payload::Proto(text) => text, + Payload::Raw(bytes) => { + String::from_utf8(bytes).map_err(|_| "raw payload is not valid UTF-8".to_owned())? + } + other => { + return Err(format!( + "{} payload cannot be written to a text column", + payload_kind(&other) + )); + } + }; + if kind == RawKind::Json { + let mut bytes = text.as_bytes().to_vec(); + simd_json::to_owned_value(&mut bytes) + .map_err(|e| format!("payload is not valid JSON: {e}"))?; + } + Ok(text) +} + +// ─── Metadata ──────────────────────────────────────────────────────────────── + +fn append_metadata( + ctx: &RunContext<'_>, + messages: &[&ConsumedMessage], + fields: &mut Vec, + columns: &mut Vec, +) -> Result<(), ArrowError> { + let rows = messages.len(); + for (meta, field) in &ctx.layout.metadata { + let column: ArrayRef = match meta { + MetaColumn::Stream => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + ctx.topic.stream.as_str(), + rows, + ))), + MetaColumn::Topic => Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + ctx.topic.topic.as_str(), + rows, + ))), + MetaColumn::PartitionId => Arc::new(Int64Array::from_iter_values(std::iter::repeat_n( + i64::from(ctx.messages.partition_id), + rows, + ))), + MetaColumn::Offset => Arc::new(Int64Array::from_iter_values( + messages.iter().map(|m| saturating_i64(m.offset)), + )), + MetaColumn::Timestamp => Arc::new( + TimestampMicrosecondArray::from_iter_values( + messages.iter().map(|m| saturating_i64(m.timestamp)), + ) + .with_timezone(UTC), + ), + MetaColumn::Id => { + let mut builder = StringBuilder::with_capacity(rows, rows.saturating_mul(39)); + for message in messages { + write!(&mut builder, "{}", message.id).map_err(|_| { + ArrowError::ComputeError("cannot format message ID".to_owned()) + })?; + builder.append_value(""); + } + Arc::new(builder.finish()) + } + MetaColumn::Headers => { + let required = !field.is_nullable(); + Arc::new(StringArray::from_iter( + messages.iter().map(|m| headers_json(m, required)), + )) + } + }; + fields.push(field.clone()); + columns.push(column); + } + Ok(()) +} + +/// Headers as a JSON object. Text values are written as strings, binary +/// values as `{"data": , "iggy_header_encoding": "base64"}`, the same +/// shape `http_sink` uses. +fn headers_json(message: &ConsumedMessage, required: bool) -> Option { + let headers = message.headers.as_ref().filter(|h| !h.is_empty()); + let Some(headers) = headers else { + return required.then(|| "{}".to_owned()); + }; + let mut object = serde_json::Map::with_capacity(headers.len()); + for (key, value) in headers { + let encoded = match value.as_raw() { + Ok(raw) => serde_json::json!({ + "data": BASE64.encode(raw), + "iggy_header_encoding": HEADER_ENCODING_BASE64, + }), + Err(_) => serde_json::Value::String(value.to_string_value()), + }; + object.insert(key.to_string_value(), encoded); + } + Some(serde_json::Value::Object(object).to_string()) +} + +// ─── Request sizing ────────────────────────────────────────────────────────── + +/// Estimate the rows per request from the Arrow allocation, then encode each +/// candidate chunk. A single row that does not fit is rejected. +fn split_into_chunks( + batch: RecordBatch, + offsets: Vec, + max_bytes: usize, + chunks: &mut Vec, + rejected: &mut Vec, +) -> Result<(), ArrowError> { + let schema_bytes = encode_schema(batch.schema_ref())?; + split_into_chunks_with( + batch, + offsets, + max_bytes, + schema_bytes, + chunks, + rejected, + encode_batch, + ) +} + +fn split_into_chunks_with( + batch: RecordBatch, + offsets: Vec, + max_bytes: usize, + schema_bytes: Vec, + chunks: &mut Vec, + rejected: &mut Vec, + mut encoder: impl FnMut(&RecordBatch) -> Result, ArrowError>, +) -> Result<(), ArrowError> { + let rows_per_chunk = estimated_rows_per_chunk(&batch, schema_bytes.len(), max_bytes); + let mut start = 0; + while start < batch.num_rows() { + let remaining = batch.num_rows() - start; + let mut row_count = rows_per_chunk.min(remaining); + loop { + let candidate = batch.slice(start, row_count); + let batch_bytes = encoder(&candidate)?; + let size = schema_bytes.len() + batch_bytes.len(); + if size <= max_bytes { + chunks.push(Chunk { + batch: candidate, + offsets: offsets[start..start + row_count].to_vec(), + schema_bytes: schema_bytes.clone(), + batch_bytes, + }); + start += row_count; + break; + } + if row_count == 1 { + rejected.push(Rejected { + offset: offsets[start], + reason: format!( + "row is {size} bytes encoded, above max_request_bytes ({max_bytes})" + ), + }); + start += 1; + break; + } + + let batch_budget = max_bytes.saturating_sub(schema_bytes.len()); + let smaller = row_count.saturating_mul(batch_budget) / batch_bytes.len().max(1); + row_count = smaller.clamp(1, row_count - 1); + } + } + Ok(()) +} + +fn estimated_rows_per_chunk(batch: &RecordBatch, schema_bytes: usize, max_bytes: usize) -> usize { + if batch.num_rows() == 0 { + return 0; + } + let overhead = + IPC_BATCH_OVERHEAD.saturating_add(batch.num_columns().saturating_mul(IPC_COLUMN_OVERHEAD)); + let batch_budget = max_bytes.saturating_sub(schema_bytes.saturating_add(overhead)); + let bytes_per_row = batch + .get_array_memory_size() + .div_ceil(batch.num_rows()) + .max(1); + (batch_budget / bytes_per_row).max(1).min(batch.num_rows()) +} + +/// Serialize the schema and the batch as Arrow IPC stream messages, the +/// format `AppendRowsRequest.arrow_rows` carries. +fn encode_ipc(batch: &RecordBatch) -> Result<(Vec, Vec), ArrowError> { + Ok((encode_schema(batch.schema_ref())?, encode_batch(batch)?)) +} + +fn encode_schema(schema: &SchemaRef) -> Result, ArrowError> { + let options = IpcWriteOptions::default(); + let generator = IpcDataGenerator::default(); + let mut tracker = DictionaryTracker::new(true); + let schema = generator.schema_to_bytes_with_dictionary_tracker(schema, &mut tracker, &options); + let mut schema_bytes = Vec::new(); + write_message(&mut schema_bytes, schema, &options)?; + Ok(schema_bytes) +} + +fn encode_batch(batch: &RecordBatch) -> Result, ArrowError> { + let options = IpcWriteOptions::default(); + let generator = IpcDataGenerator::default(); + let mut tracker = DictionaryTracker::new(true); + let mut compression = CompressionContext::default(); + let _ = generator.schema_to_bytes_with_dictionary_tracker( + batch.schema_ref(), + &mut tracker, + &options, + ); + + let (dictionaries, encoded) = + generator.encode(batch, &mut tracker, &options, &mut compression)?; + let mut batch_bytes = Vec::new(); + for dictionary in dictionaries { + write_message(&mut batch_bytes, dictionary, &options)?; + } + write_message(&mut batch_bytes, encoded, &options)?; + Ok(batch_bytes) +} + +// ─── Helpers ───────────────────────────────────────────────────────────────── + +fn payload_kind(payload: &Payload) -> &'static str { + match payload { + Payload::Json(_) => "JSON", + Payload::Raw(_) => "raw", + Payload::Text(_) => "text", + Payload::Proto(_) => "proto", + Payload::FlatBuffer(_) => "FlatBuffer", + Payload::Avro(_) => "Avro", + } +} + +fn empty_array() -> OwnedValue { + OwnedValue::Array(Box::default()) +} + +fn saturating_i64(value: u64) -> i64 { + i64::try_from(value).unwrap_or(i64::MAX) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::WriteMode; + use crate::schema::parse_table_schema; + use crate::test_support::settings; + use arrow::array::{Array, AsArray, ListArray}; + use arrow::datatypes::Int64Type; + use iggy_common::{HeaderKey, HeaderValue}; + use iggy_connector_sdk::Schema as PayloadSchema; + use std::collections::BTreeMap; + use std::str::FromStr; + + const METADATA_FIELDS: &str = r#" + {"name":"iggy_stream","type":"STRING"}, + {"name":"iggy_topic","type":"STRING"}, + {"name":"iggy_partition_id","type":"INTEGER"}, + {"name":"iggy_offset","type":"INTEGER"}, + {"name":"iggy_timestamp","type":"TIMESTAMP"}, + {"name":"iggy_id","type":"STRING"}"#; + + fn layout(fields: &str, mode: WriteMode, include_metadata: bool) -> TableLayout { + let body = format!(r#"{{"schema":{{"fields":[{fields}]}}}}"#); + let columns = parse_table_schema(body.as_bytes()).unwrap(); + TableLayout::build(columns, &settings(mode, include_metadata)).unwrap() + } + + fn topic() -> TopicMetadata { + TopicMetadata { + stream: "orders".into(), + topic: "created".into(), + } + } + + fn messages_metadata() -> MessagesMetadata { + MessagesMetadata { + partition_id: 3, + current_offset: 100, + schema: PayloadSchema::Json, + } + } + + fn message(offset: u64, payload: Payload) -> ConsumedMessage { + ConsumedMessage { + id: u128::from(offset) + 1000, + offset, + checksum: 0, + timestamp: 1_700_000_000_000_000 + offset, + origin_timestamp: 0, + headers: None, + payload, + } + } + + fn json(offset: u64, text: &str) -> ConsumedMessage { + let mut bytes = text.as_bytes().to_vec(); + message( + offset, + Payload::Json(simd_json::to_owned_value(&mut bytes).unwrap()), + ) + } + + fn run(layout: &TableLayout, messages: Vec) -> Encoded { + run_with_budget(layout, messages, 8 * 1024 * 1024) + } + + fn run_with_budget( + layout: &TableLayout, + messages: Vec, + max_request_bytes: usize, + ) -> Encoded { + let topic = topic(); + let metadata = messages_metadata(); + let ctx = RunContext { + layout, + topic: &topic, + messages: &metadata, + max_request_bytes, + }; + encode(&ctx, messages).unwrap() + } + + fn column_names(batch: &RecordBatch) -> Vec { + batch + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect() + } + + #[test] + fn given_json_rows_should_encode_payload_and_metadata_columns() { + let layout = layout( + &format!( + r#"{{"name":"user_id","type":"INT64"}},{{"name":"event","type":"STRING"}},{METADATA_FIELDS}"# + ), + WriteMode::Mapped, + true, + ); + let encoded = run( + &layout, + vec![ + json(10, r#"{"user_id": 1, "event": "signup", "ignored": true}"#), + json(11, r#"{"user_id": "2", "event": "login"}"#), + ], + ); + assert!(encoded.rejected.is_empty()); + assert_eq!(encoded.chunks.len(), 1); + let chunk = &encoded.chunks[0]; + assert_eq!(chunk.offsets, vec![10, 11]); + assert_eq!( + column_names(&chunk.batch), + vec![ + "user_id", + "event", + "iggy_stream", + "iggy_topic", + "iggy_partition_id", + "iggy_offset", + "iggy_timestamp", + "iggy_id" + ] + ); + let user_ids = chunk.batch.column(0).as_primitive::(); + assert_eq!(user_ids.values(), &[1, 2]); + let stream = chunk.batch.column(2).as_string::(); + assert_eq!(stream.value(1), "orders"); + let partition = chunk.batch.column(4).as_primitive::(); + assert_eq!(partition.value(0), 3); + let offset = chunk.batch.column(5).as_primitive::(); + assert_eq!(offset.values(), &[10, 11]); + let id = chunk.batch.column(7).as_string::(); + assert_eq!(id.value(0), "1010"); + assert!(!chunk.schema_bytes.is_empty()); + assert!(!chunk.batch_bytes.is_empty()); + } + + #[test] + fn given_row_with_wrong_type_should_reject_only_that_row() { + let layout = layout( + r#"{"name":"user_id","type":"INT64"}"#, + WriteMode::Mapped, + false, + ); + let encoded = run( + &layout, + vec![ + json(1, r#"{"user_id": 1}"#), + json(2, r#"{"user_id": "not-a-number"}"#), + json(3, r#"{"user_id": 3}"#), + ], + ); + assert_eq!(encoded.rejected.len(), 1); + assert_eq!(encoded.rejected[0].offset, 2); + let chunk = &encoded.chunks[0]; + assert_eq!(chunk.offsets, vec![1, 3]); + assert_eq!( + chunk.batch.column(0).as_primitive::().values(), + &[1, 3] + ); + } + + #[test] + fn given_one_invalid_row_should_isolate_it_with_logarithmic_decodes() { + let layout = layout( + r#"{"name":"user_id","type":"INT64"}"#, + WriteMode::Mapped, + false, + ); + let rows: Vec = (0..32) + .map(|index| { + let text = if index == 17 { + r#"{"user_id":"not-a-number"}"#.to_owned() + } else { + format!(r#"{{"user_id":{index}}}"#) + }; + let mut bytes = text.into_bytes(); + simd_json::to_owned_value(&mut bytes).unwrap() + }) + .collect(); + let mut rejected = Vec::new(); + let mut decode_calls = 0; + + let (batch, survivors) = decode_rows_with( + &layout.fields, + &rows, + |index, _| rejected.push(index), + |schema, rows| { + decode_calls += 1; + decode(schema, rows) + }, + ) + .unwrap(); + + assert_eq!(rejected, vec![17]); + assert_eq!(survivors.len(), 31); + assert_eq!(batch.unwrap().num_rows(), 31); + assert!(decode_calls <= 12, "used {decode_calls} decode calls"); + } + + #[test] + fn given_text_payload_should_parse_as_json_and_reject_non_json() { + let layout = layout( + r#"{"name":"name","type":"STRING"}"#, + WriteMode::Mapped, + false, + ); + let encoded = run( + &layout, + vec![ + message(1, Payload::Text(r#"{"name":"a"}"#.into())), + message(2, Payload::Text("plain text".into())), + message(3, Payload::Raw(vec![1, 2, 3])), + json(4, r#"[1, 2]"#), + ], + ); + let rejected: Vec = encoded.rejected.iter().map(|r| r.offset).collect(); + assert_eq!(rejected, vec![2, 3, 4]); + assert_eq!(encoded.chunks[0].offsets, vec![1]); + } + + #[test] + fn given_column_absent_from_every_row_should_leave_it_out_of_writer_schema() { + let layout = layout( + r#"{"name":"a","type":"STRING"},{"name":"b","type":"STRING"},{"name":"c","type":"STRING"}"#, + WriteMode::Mapped, + false, + ); + let encoded = run( + &layout, + vec![json(1, r#"{"c": "x"}"#), json(2, r#"{"a": "y"}"#)], + ); + assert_eq!(column_names(&encoded.chunks[0].batch), vec!["a", "c"]); + } + + #[test] + fn given_required_column_missing_should_reject_unless_it_has_default() { + let layout = layout( + r#"{"name":"id","type":"INT64","mode":"REQUIRED"}, + {"name":"created","type":"TIMESTAMP","mode":"REQUIRED","defaultValueExpression":"CURRENT_TIMESTAMP()"}"#, + WriteMode::Mapped, + false, + ); + let encoded = run(&layout, vec![json(1, r#"{"id": 7}"#), json(2, r#"{}"#)]); + assert_eq!(encoded.rejected.len(), 1); + assert_eq!(encoded.rejected[0].offset, 2); + assert!(encoded.rejected[0].reason.contains("REQUIRED column 'id'")); + assert_eq!(column_names(&encoded.chunks[0].batch), vec!["id"]); + } + + #[test] + fn given_null_repeated_column_should_write_empty_array() { + let layout = layout( + r#"{"name":"tags","type":"STRING","mode":"REPEATED"}"#, + WriteMode::Mapped, + false, + ); + let encoded = run( + &layout, + vec![ + json(1, r#"{"tags": ["a", "b"]}"#), + json(2, r#"{"tags": null}"#), + ], + ); + assert!(encoded.rejected.is_empty()); + let tags = encoded.chunks[0] + .batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_eq!(tags.value_length(0), 2); + assert_eq!(tags.value_length(1), 0); + assert!(!tags.is_null(1)); + } + + #[test] + fn given_json_and_bytes_columns_should_convert_values() { + let layout = layout( + r#"{"name":"attrs","type":"JSON"},{"name":"blob","type":"BYTES"}"#, + WriteMode::Mapped, + false, + ); + let encoded = run( + &layout, + vec![ + json(1, r#"{"attrs": {"k": [1, 2]}, "blob": "aGVsbG8="}"#), + json(2, r#"{"blob": "not base64!"}"#), + ], + ); + assert_eq!(encoded.rejected.len(), 1); + assert_eq!(encoded.rejected[0].offset, 2); + let batch = &encoded.chunks[0].batch; + assert_eq!( + batch.column(0).as_string::().value(0), + r#"{"k":[1,2]}"# + ); + assert_eq!(batch.column(1).as_binary::().value(0), b"hello"); + } + + #[test] + fn given_timestamp_and_numeric_strings_should_decode() { + let layout = layout( + r#"{"name":"at","type":"TIMESTAMP"},{"name":"amount","type":"NUMERIC"},{"name":"day","type":"DATE"}"#, + WriteMode::Mapped, + false, + ); + let encoded = run( + &layout, + vec![json( + 1, + r#"{"at": "2024-01-02T03:04:05.123456Z", "amount": "12.345", "day": "2024-01-02"}"#, + )], + ); + assert!(encoded.rejected.is_empty(), "{:?}", encoded.rejected); + let batch = &encoded.chunks[0].batch; + assert_eq!( + batch.column(0).data_type(), + &DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some(UTC.into())) + ); + let at = batch + .column(0) + .as_primitive::(); + assert_eq!(at.value(0), 1_704_164_645_123_456); + } + + #[test] + fn given_nested_record_with_missing_required_child_should_reject_row() { + let layout = layout( + r#"{"name":"customer","type":"RECORD","fields":[ + {"name":"id","type":"INT64","mode":"REQUIRED"}, + {"name":"tags","type":"STRING","mode":"REPEATED"} + ]}"#, + WriteMode::Mapped, + false, + ); + let encoded = run( + &layout, + vec![ + json(1, r#"{"customer": {"id": 1}}"#), + json(2, r#"{"customer": {"tags": ["x"]}}"#), + ], + ); + assert_eq!(encoded.rejected.len(), 1); + assert_eq!(encoded.rejected[0].offset, 2); + assert_eq!(encoded.chunks[0].offsets, vec![1]); + } + + #[test] + fn given_raw_json_column_should_accept_json_and_reject_invalid_text() { + let layout = layout( + &format!(r#"{{"name":"payload","type":"JSON"}},{METADATA_FIELDS}"#), + WriteMode::Raw, + true, + ); + let encoded = run( + &layout, + vec![ + json(1, r#"{"a": 1}"#), + message(2, Payload::Text("not json".into())), + message(3, Payload::Raw(br#"{"b":2}"#.to_vec())), + ], + ); + assert_eq!(encoded.rejected.len(), 1); + assert_eq!(encoded.rejected[0].offset, 2); + let batch = &encoded.chunks[0].batch; + assert_eq!(column_names(batch)[0], "payload"); + let payload = batch.column(0).as_string::(); + assert_eq!(payload.value(0), r#"{"a":1}"#); + assert_eq!(payload.value(1), r#"{"b":2}"#); + } + + #[test] + fn given_raw_bytes_column_should_accept_every_payload_type() { + let layout = layout( + r#"{"name":"payload","type":"BYTES"}"#, + WriteMode::Raw, + false, + ); + let encoded = run( + &layout, + vec![ + message(1, Payload::Raw(vec![0, 159, 146, 150])), + message(2, Payload::Text("hi".into())), + ], + ); + assert!(encoded.rejected.is_empty()); + let payload = encoded.chunks[0].batch.column(0).as_binary::(); + assert_eq!(payload.value(0), &[0, 159, 146, 150]); + assert_eq!(payload.value(1), b"hi"); + } + + #[test] + fn given_raw_string_column_should_reject_non_utf8_bytes() { + let layout = layout( + r#"{"name":"payload","type":"STRING"}"#, + WriteMode::Raw, + false, + ); + let encoded = run(&layout, vec![message(1, Payload::Raw(vec![0xff, 0xfe]))]); + assert_eq!(encoded.rejected.len(), 1); + assert!(encoded.chunks.is_empty()); + } + + #[test] + fn given_batch_over_budget_should_split_and_keep_row_order() { + let layout = layout( + r#"{"name":"payload","type":"STRING"}"#, + WriteMode::Raw, + false, + ); + let messages = (0..40) + .map(|offset| message(offset, Payload::Text("x".repeat(4096)))) + .collect(); + let mut encoded = run(&layout, messages); + let source = encoded.chunks.pop().unwrap(); + let schema_bytes = encode_schema(source.batch.schema_ref()).unwrap(); + let mut chunks = Vec::new(); + let mut rejected = Vec::new(); + let mut encode_calls = 0; + split_into_chunks_with( + source.batch, + source.offsets, + 64 * 1024, + schema_bytes, + &mut chunks, + &mut rejected, + |batch| { + encode_calls += 1; + encode_batch(batch) + }, + ) + .unwrap(); + + assert!(rejected.is_empty()); + assert!(chunks.len() > 1); + assert_eq!(encode_calls, chunks.len()); + for chunk in &chunks { + assert!(chunk.schema_bytes.len() + chunk.batch_bytes.len() <= 64 * 1024); + assert_eq!(chunk.batch.num_rows(), chunk.offsets.len()); + } + let offsets: Vec = chunks.iter().flat_map(|c| c.offsets.clone()).collect(); + assert_eq!(offsets, (0..40).collect::>()); + } + + #[test] + fn given_single_row_over_budget_should_reject_it() { + let layout = layout( + r#"{"name":"payload","type":"STRING"}"#, + WriteMode::Raw, + false, + ); + let encoded = run_with_budget( + &layout, + vec![ + message(1, Payload::Text("small".into())), + message(2, Payload::Text("y".repeat(200 * 1024))), + ], + 64 * 1024, + ); + assert_eq!(encoded.rejected.len(), 1); + assert_eq!(encoded.rejected[0].offset, 2); + assert_eq!(encoded.chunks.len(), 1); + assert_eq!(encoded.chunks[0].offsets, vec![1]); + } + + #[test] + fn given_row_errors_should_rebuild_chunk_without_those_rows() { + let layout = layout( + r#"{"name":"payload","type":"STRING"}"#, + WriteMode::Raw, + false, + ); + let encoded = run( + &layout, + (0..4) + .map(|offset| message(offset, Payload::Text(format!("row-{offset}")))) + .collect(), + ); + let chunk = without_rows(&encoded.chunks[0], &[1, 3]).unwrap().unwrap(); + assert_eq!(chunk.offsets, vec![0, 2]); + let payload = chunk.batch.column(0).as_string::(); + assert_eq!(payload.value(1), "row-2"); + assert!(without_rows(&chunk, &[0, 1]).unwrap().is_none()); + } + + #[test] + fn given_headers_should_encode_text_and_binary_values() { + let mut headers = BTreeMap::new(); + headers.insert( + HeaderKey::from_str("trace").unwrap(), + HeaderValue::from_str("abc").unwrap(), + ); + headers.insert( + HeaderKey::from_str("blob").unwrap(), + HeaderValue::try_from(&[1u8, 2][..]).unwrap(), + ); + let mut msg = message(1, Payload::Text("x".into())); + msg.headers = Some(headers); + let text = headers_json(&msg, false).unwrap(); + let value: serde_json::Value = serde_json::from_str(&text).unwrap(); + assert_eq!(value["trace"], "abc"); + assert_eq!(value["blob"]["data"], "AQI="); + assert_eq!(value["blob"]["iggy_header_encoding"], "base64"); + + let empty = message(2, Payload::Text("x".into())); + assert_eq!(headers_json(&empty, false), None); + assert_eq!(headers_json(&empty, true).as_deref(), Some("{}")); + } +} diff --git a/core/connectors/sinks/bigquery_sink/src/error.rs b/core/connectors/sinks/bigquery_sink/src/error.rs new file mode 100644 index 0000000000..00ce7768a5 --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/src/error.rs @@ -0,0 +1,207 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! Failure classification for BigQuery calls. +//! +//! Only transient failures are retried. Retrying a permanent one (bad +//! schema, missing permission) cannot succeed and only delays the error. + +use crate::schema::SchemaError; +use gcloud_gax::grpc::{Code, Status}; +use iggy_connector_sdk::Error; +use std::fmt; + +/// An `AppendRows` call that did not append. Row-level errors are not part +/// of this type: they are an expected outcome the caller handles by +/// dropping the reported rows. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum AppendError { + /// The call failed, or the response carried an error status. + Rpc { code: Code, message: String }, + /// The response stream closed without a response. + NoResponse, +} + +/// A `tables.get` call that did not return a usable schema. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum TableError { + Http { status: u16, body: String }, + Transport(String), + Token(String), + Schema(SchemaError), +} + +impl AppendError { + pub(crate) fn is_retryable(&self) -> bool { + match self { + AppendError::Rpc { code, .. } => is_retryable_code(*code), + // The append may or may not have landed. Retrying risks a + // duplicate, which the delivery guarantees allow. + AppendError::NoResponse => true, + } + } +} + +impl TableError { + pub(crate) fn is_retryable(&self) -> bool { + match self { + TableError::Http { status, .. } => *status == 429 || *status >= 500, + TableError::Transport(_) | TableError::Token(_) => true, + TableError::Schema(_) => false, + } + } +} + +/// gRPC codes BigQuery documents as safe to retry for `AppendRows`. +pub(crate) fn is_retryable_code(code: Code) -> bool { + matches!( + code, + Code::Unavailable + | Code::DeadlineExceeded + | Code::Internal + | Code::Aborted + | Code::ResourceExhausted + ) +} + +impl From for AppendError { + fn from(status: Status) -> Self { + AppendError::Rpc { + code: status.code(), + message: status.message().to_owned(), + } + } +} + +impl From for Error { + fn from(error: AppendError) -> Self { + if error.is_retryable() { + Error::CannotStoreData(error.to_string()) + } else { + Error::PermanentHttpError(error.to_string()) + } + } +} + +impl From for Error { + fn from(error: TableError) -> Self { + Error::InitError(error.to_string()) + } +} + +impl fmt::Display for AppendError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + AppendError::Rpc { code, message } => write!(f, "AppendRows {code:?}: {message}"), + AppendError::NoResponse => write!(f, "AppendRows stream closed without a response"), + } + } +} + +impl fmt::Display for TableError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + TableError::Http { status, body } => { + write!(f, "tables.get returned HTTP {status}: {body}") + } + TableError::Transport(reason) => write!(f, "tables.get failed: {reason}"), + TableError::Token(reason) => write!(f, "cannot obtain an access token: {reason}"), + TableError::Schema(reason) => write!(f, "{reason}"), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rpc(code: Code) -> AppendError { + AppendError::Rpc { + code, + message: "boom".into(), + } + } + + #[test] + fn given_transient_codes_should_be_retryable() { + for code in [ + Code::Unavailable, + Code::DeadlineExceeded, + Code::Internal, + Code::Aborted, + Code::ResourceExhausted, + ] { + assert!(rpc(code).is_retryable(), "{code:?}"); + } + } + + #[test] + fn given_permanent_codes_should_not_be_retryable() { + for code in [ + Code::InvalidArgument, + Code::PermissionDenied, + Code::Unauthenticated, + Code::NotFound, + Code::FailedPrecondition, + ] { + assert!(!rpc(code).is_retryable(), "{code:?}"); + } + } + + #[test] + fn given_missing_response_should_be_retryable() { + assert!(AppendError::NoResponse.is_retryable()); + } + + #[test] + fn given_status_should_keep_code_and_message() { + let error = AppendError::from(Status::permission_denied("no access")); + assert_eq!( + error, + AppendError::Rpc { + code: Code::PermissionDenied, + message: "no access".into() + } + ); + } + + #[test] + fn given_append_error_should_map_to_sdk_error_by_retryability() { + assert!(matches!( + Error::from(rpc(Code::Unavailable)), + Error::CannotStoreData(_) + )); + assert!(matches!( + Error::from(rpc(Code::InvalidArgument)), + Error::PermanentHttpError(_) + )); + } + + #[test] + fn given_table_failures_should_retry_only_transient_errors() { + let http = |status| TableError::Http { + status, + body: String::new(), + }; + assert!(http(429).is_retryable()); + assert!(http(503).is_retryable()); + assert!(!http(403).is_retryable()); + assert!(!http(404).is_retryable()); + assert!(TableError::Transport("reset".into()).is_retryable()); + assert!(TableError::Token("token service unavailable".into()).is_retryable()); + } +} diff --git a/core/connectors/sinks/bigquery_sink/src/lib.rs b/core/connectors/sinks/bigquery_sink/src/lib.rs new file mode 100644 index 0000000000..ac5bde1f8f --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/src/lib.rs @@ -0,0 +1,461 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! Google BigQuery sink connector. +//! +//! Writes Iggy messages into a BigQuery table through the Storage Write API +//! `_default` stream, with rows encoded as Arrow record batches. +//! +//! Module map: +//! - `schema`: BigQuery table schema (from `tables.get`) to Arrow schema. +//! - `encode`: `ConsumedMessage`s to Arrow record batches, split by size. +//! - `client`: credentials, `tables.get`, and `AppendRows` calls. +//! - `error`: gRPC status classification into retryable and permanent. +//! - `sink`: the `Sink` trait implementation. + +use humantime::Duration as HumanDuration; +use iggy_connector_sdk::{Error, sink_connector}; +use secrecy::SecretString; +use serde::Deserialize; +use std::str::FromStr; +use std::sync::atomic::AtomicU64; +use std::time::Duration; +use tracing::warn; + +mod client; +mod encode; +mod error; +mod schema; +mod sink; + +sink_connector!(BigQuerySink); + +const DEFAULT_PAYLOAD_COLUMN: &str = "payload"; +const DEFAULT_MAX_REQUEST_BYTES: usize = 8 * 1024 * 1024; +/// Hard ceiling below the 10 MB `AppendRows` limit, leaving room for the +/// request envelope (stream name, writer schema, protobuf framing). +const MAX_REQUEST_BYTES_CEILING: usize = 9 * 1024 * 1024; +/// Smallest accepted request budget. Anything lower makes most real rows +/// unwritable. +const MIN_REQUEST_BYTES: usize = 64 * 1024; +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_RETRY_DELAY: &str = "1s"; +const DEFAULT_MAX_RETRY_DELAY: &str = "30s"; +const DEFAULT_TIMEOUT: &str = "30s"; + +/// Plugin configuration, deserialized from `[plugin_config]`. +/// +/// `Deserialize` only. Nothing re-serializes a plugin config, and leaving +/// `Serialize` off is what keeps `credentials_json` unserializable. +#[derive(Debug, Clone, Deserialize)] +pub struct BigQuerySinkConfig { + pub project_id: String, + pub dataset: String, + pub table: String, + /// `mapped` (default) or `raw`. + pub mode: Option, + /// Path to a service account key file. Mutually exclusive with + /// `credentials_json`. When neither is set, Application Default + /// Credentials are used. + pub credentials_path: Option, + /// Inline service account key JSON. + pub credentials_json: Option, + /// Column holding the payload in `raw` mode (default `payload`). + pub payload_column: Option, + /// Write `iggy_stream`, `iggy_topic`, `iggy_partition_id`, `iggy_offset`, + /// `iggy_timestamp` and `iggy_id` (default `true`). + pub include_metadata: Option, + /// Write message headers into an `iggy_headers` column (default `false`). + pub include_headers: Option, + /// How BigQuery fills columns absent from a request: `default` (column + /// default value, else NULL) or `null`. + pub missing_value: Option, + /// Upper bound for one `AppendRows` request (default 8 MiB). + pub max_request_bytes: Option, + /// Total attempts per append, including the first (default 3). + pub max_retries: Option, + /// Base retry delay, e.g. `"1s"` (default `1s`). + pub retry_delay: Option, + /// Upper bound for a single retry delay (default `30s`). + pub max_retry_delay: Option, + /// Per-request timeout for `tables.get` and `AppendRows` (default `30s`). + pub timeout: Option, + /// Log every batch at info level instead of debug (default `false`). + pub verbose_logging: Option, + /// REST endpoint override, e.g. `http://127.0.0.1:9050`. Must be set + /// together with `grpc_endpoint`; when both are set no credentials are + /// used. Intended for local fakes and emulators. + pub endpoint: Option, + /// gRPC `host:port` override for the Storage Write API. + pub grpc_endpoint: Option, +} + +/// How message payloads map onto table columns. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WriteMode { + /// Top-level JSON fields map to table columns by name. + #[default] + Mapped, + /// The whole payload goes into one column. + Raw, +} + +/// Maps to `AppendRowsRequest.default_missing_value_interpretation`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MissingValue { + #[default] + Default, + Null, +} + +impl MissingValue { + /// Wire value of `AppendRowsRequest.MissingValueInterpretation`. + pub(crate) fn as_proto(self) -> i32 { + match self { + MissingValue::Null => 1, + MissingValue::Default => 2, + } + } +} + +/// Resolved, validated runtime settings. Built once in `new()`. +#[derive(Debug, Clone)] +pub(crate) struct Settings { + pub mode: WriteMode, + pub payload_column: String, + pub include_metadata: bool, + pub include_headers: bool, + pub missing_value: MissingValue, + pub max_request_bytes: usize, + pub max_retries: u32, + pub retry_delay: Duration, + pub max_retry_delay: Duration, + pub timeout: Duration, + pub verbose: bool, +} + +/// Counters updated on the hot path without taking a lock. +#[derive(Debug, Default)] +pub(crate) struct Counters { + pub rows_written: AtomicU64, + pub rows_rejected: AtomicU64, + pub rows_failed: AtomicU64, +} + +#[derive(Debug)] +struct OpenState { + client: client::BigQueryClient, + layout: schema::TableLayout, +} + +#[derive(Debug)] +pub struct BigQuerySink { + id: u32, + config: BigQuerySinkConfig, + settings: Settings, + /// `project.dataset.table`, safe to log. + target: String, + state: Option, + counters: Counters, +} + +impl BigQuerySink { + pub fn new(id: u32, config: BigQuerySinkConfig) -> Self { + let settings = Settings::from_config(&config); + let target = format!("{}.{}.{}", config.project_id, config.dataset, config.table); + BigQuerySink { + id, + config, + settings, + target, + state: None, + counters: Counters::default(), + } + } +} + +impl Settings { + fn from_config(config: &BigQuerySinkConfig) -> Self { + let mut retry_delay = parse_duration(config.retry_delay.as_deref(), DEFAULT_RETRY_DELAY); + let mut max_retry_delay = + parse_duration(config.max_retry_delay.as_deref(), DEFAULT_MAX_RETRY_DELAY); + if retry_delay > max_retry_delay { + warn!( + "BigQuery sink: retry_delay ({retry_delay:?}) is greater than max_retry_delay ({max_retry_delay:?}), swapping them" + ); + std::mem::swap(&mut retry_delay, &mut max_retry_delay); + } + + let requested = config + .max_request_bytes + .unwrap_or(DEFAULT_MAX_REQUEST_BYTES); + let max_request_bytes = requested.clamp(MIN_REQUEST_BYTES, MAX_REQUEST_BYTES_CEILING); + if max_request_bytes != requested { + warn!( + "BigQuery sink: max_request_bytes {requested} is outside [{MIN_REQUEST_BYTES}, {MAX_REQUEST_BYTES_CEILING}], using {max_request_bytes}" + ); + } + + Settings { + mode: config.mode.unwrap_or_default(), + payload_column: config + .payload_column + .clone() + .unwrap_or_else(|| DEFAULT_PAYLOAD_COLUMN.to_owned()), + include_metadata: config.include_metadata.unwrap_or(true), + include_headers: config.include_headers.unwrap_or(false), + missing_value: config.missing_value.unwrap_or_default(), + max_request_bytes, + max_retries: config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES), + retry_delay, + max_retry_delay, + timeout: parse_duration(config.timeout.as_deref(), DEFAULT_TIMEOUT), + verbose: config.verbose_logging.unwrap_or(false), + } + } +} + +impl BigQuerySinkConfig { + /// Structural checks that need no network. Called from `open()` because + /// `new()` cannot fail. + pub(crate) fn validate(&self) -> Result<(), Error> { + for (name, value) in [ + ("project_id", &self.project_id), + ("dataset", &self.dataset), + ("table", &self.table), + ] { + if value.trim().is_empty() { + return Err(Error::InvalidConfigValue(format!( + "{name} must not be empty" + ))); + } + } + if self.credentials_path.is_some() && self.credentials_json.is_some() { + return Err(Error::InvalidConfigValue( + "set either credentials_path or credentials_json, not both".into(), + )); + } + if self.endpoint.is_some() != self.grpc_endpoint.is_some() { + return Err(Error::InvalidConfigValue( + "endpoint and grpc_endpoint must be set together".into(), + )); + } + if let Some(column) = &self.payload_column + && column.trim().is_empty() + { + return Err(Error::InvalidConfigValue( + "payload_column must not be empty".into(), + )); + } + Ok(()) + } +} + +fn parse_duration(input: Option<&str>, default: &str) -> Duration { + let raw = input.unwrap_or(default); + HumanDuration::from_str(raw) + .map(|d| *d) + .unwrap_or_else(|e| { + warn!("BigQuery sink: invalid duration '{raw}': {e}, using default '{default}'"); + HumanDuration::from_str(default) + .map(|d| *d) + .unwrap_or(Duration::from_secs(1)) + }) +} + +#[cfg(test)] +pub(crate) mod test_support { + use super::*; + + pub(crate) fn config() -> BigQuerySinkConfig { + BigQuerySinkConfig { + project_id: "proj".into(), + dataset: "ds".into(), + table: "events".into(), + mode: None, + credentials_path: None, + credentials_json: None, + payload_column: None, + include_metadata: None, + include_headers: None, + missing_value: None, + max_request_bytes: None, + max_retries: None, + retry_delay: None, + max_retry_delay: None, + timeout: None, + verbose_logging: None, + endpoint: None, + grpc_endpoint: None, + } + } + + pub(crate) fn settings(mode: WriteMode, include_metadata: bool) -> Settings { + let mut config = config(); + config.mode = Some(mode); + config.include_metadata = Some(include_metadata); + Settings::from_config(&config) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::config as test_config; + + #[test] + fn given_minimal_config_should_apply_defaults() { + let settings = Settings::from_config(&test_config()); + assert_eq!(settings.mode, WriteMode::Mapped); + assert_eq!(settings.payload_column, "payload"); + assert!(settings.include_metadata); + assert!(!settings.include_headers); + assert_eq!(settings.missing_value, MissingValue::Default); + assert_eq!(settings.max_request_bytes, DEFAULT_MAX_REQUEST_BYTES); + assert_eq!(settings.max_retries, 3); + assert_eq!(settings.retry_delay, Duration::from_secs(1)); + assert_eq!(settings.max_retry_delay, Duration::from_secs(30)); + assert_eq!(settings.timeout, Duration::from_secs(30)); + assert!(!settings.verbose); + } + + #[test] + fn given_minimal_toml_should_deserialize() { + let config: BigQuerySinkConfig = toml::from_str( + r#" + project_id = "proj" + dataset = "ds" + table = "events" + "#, + ) + .expect("minimal config should parse"); + assert!(config.validate().is_ok()); + } + + #[test] + fn given_raw_mode_and_null_missing_value_should_deserialize() { + let config: BigQuerySinkConfig = toml::from_str( + r#" + project_id = "proj" + dataset = "ds" + table = "events" + mode = "raw" + missing_value = "null" + "#, + ) + .expect("config should parse"); + assert_eq!(config.mode, Some(WriteMode::Raw)); + assert_eq!(config.missing_value, Some(MissingValue::Null)); + } + + #[test] + fn given_unknown_mode_should_fail_to_deserialize() { + let result = toml::from_str::( + r#" + project_id = "proj" + dataset = "ds" + table = "events" + mode = "upsert" + "#, + ); + assert!(result.is_err()); + } + + #[test] + fn given_missing_table_should_fail_to_deserialize() { + let result = toml::from_str::( + r#" + project_id = "proj" + dataset = "ds" + "#, + ); + assert!(result.is_err()); + } + + #[test] + fn given_both_credential_sources_should_fail_validation() { + let mut config = test_config(); + config.credentials_path = Some("/tmp/key.json".into()); + config.credentials_json = Some(SecretString::from("{}")); + assert!(matches!( + config.validate(), + Err(Error::InvalidConfigValue(_)) + )); + } + + #[test] + fn given_only_rest_endpoint_should_fail_validation() { + let mut config = test_config(); + config.endpoint = Some("http://127.0.0.1:9050".into()); + assert!(matches!( + config.validate(), + Err(Error::InvalidConfigValue(_)) + )); + } + + #[test] + fn given_empty_table_should_fail_validation() { + let mut config = test_config(); + config.table = " ".into(); + assert!(matches!( + config.validate(), + Err(Error::InvalidConfigValue(_)) + )); + } + + #[test] + fn given_reversed_retry_delays_should_swap_them() { + let mut config = test_config(); + config.retry_delay = Some("10s".into()); + config.max_retry_delay = Some("2s".into()); + let settings = Settings::from_config(&config); + assert_eq!(settings.retry_delay, Duration::from_secs(2)); + assert_eq!(settings.max_retry_delay, Duration::from_secs(10)); + } + + #[test] + fn given_oversized_request_budget_should_clamp_below_api_limit() { + let mut config = test_config(); + config.max_request_bytes = Some(50 * 1024 * 1024); + let settings = Settings::from_config(&config); + assert_eq!(settings.max_request_bytes, MAX_REQUEST_BYTES_CEILING); + } + + #[test] + fn given_invalid_duration_should_fall_back_to_default() { + let mut config = test_config(); + config.timeout = Some("soon".into()); + let settings = Settings::from_config(&config); + assert_eq!(settings.timeout, Duration::from_secs(30)); + } + + #[test] + fn given_debug_format_should_not_leak_inline_credentials() { + let mut config = test_config(); + config.credentials_json = Some(SecretString::from("super-secret-key")); + let rendered = format!("{config:?}"); + assert!(!rendered.contains("super-secret-key")); + } + + #[test] + fn given_missing_value_should_map_to_proto_enum() { + assert_eq!(MissingValue::Null.as_proto(), 1); + assert_eq!(MissingValue::Default.as_proto(), 2); + } +} diff --git a/core/connectors/sinks/bigquery_sink/src/schema.rs b/core/connectors/sinks/bigquery_sink/src/schema.rs new file mode 100644 index 0000000000..74eaabdcc4 --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/src/schema.rs @@ -0,0 +1,769 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! BigQuery table schema model and its mapping to Arrow. +//! +//! The schema comes from the `tables.get` REST call as JSON. It is parsed +//! with local serde types rather than a client library model so that every +//! BigQuery type (GEOGRAPHY, RANGE, ...) deserializes, and the unsupported +//! ones can be rejected with a clear message in `open()`. +//! +//! Arrow types follow the Storage Write API Arrow mapping: +//! +//! | BigQuery | Arrow | +//! |------------|--------------------------------| +//! | STRING | Utf8 | +//! | BYTES | Binary | +//! | INT64 | Int64 | +//! | FLOAT64 | Float64 | +//! | BOOL | Boolean | +//! | NUMERIC | Decimal128(38, 9) | +//! | BIGNUMERIC | Decimal256(76, 38) | +//! | TIMESTAMP | Timestamp(Microsecond, "UTC") | +//! | DATETIME | Timestamp(Microsecond, none) | +//! | DATE | Date32 | +//! | TIME | Time64(Microsecond) | +//! | GEOGRAPHY | Utf8 (WKT) | +//! | JSON | Utf8 (JSON text) | +//! | RECORD | Struct | +//! | REPEATED | List | + +use crate::{Settings, WriteMode}; +use arrow::datatypes::{DataType, Field, FieldRef, Fields, TimeUnit}; +use serde::Deserialize; +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +pub(crate) const UTC: &str = "UTC"; +const LIST_ITEM: &str = "item"; + +/// One table column, recursively for RECORD types. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Column { + pub name: String, + pub bq_type: BqType, + pub mode: Mode, + pub has_default: bool, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum BqType { + String, + Bytes, + Int64, + Float64, + Bool, + Numeric, + BigNumeric, + Timestamp, + Datetime, + Date, + Time, + Geography, + Json, + Record(Vec), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Mode { + Nullable, + Required, + Repeated, +} + +/// Iggy metadata written next to the payload. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum MetaColumn { + Stream, + Topic, + PartitionId, + Offset, + Timestamp, + Id, + Headers, +} + +/// What the `raw` mode payload column holds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RawKind { + Json, + String, + Bytes, +} + +/// The validated mapping from messages to the target table, built once in +/// `open()`. +#[derive(Debug)] +pub(crate) struct TableLayout { + /// Payload columns in table order. Empty in `raw` mode. + pub columns: Vec, + /// Arrow field for each entry of `columns`, same index. + pub fields: Vec, + /// Column name to index into `columns`. + pub index: HashMap, + /// Enabled metadata columns with their Arrow fields. + pub metadata: Vec<(MetaColumn, FieldRef)>, + /// `raw` mode payload column. + pub raw: Option<(RawKind, FieldRef)>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum SchemaError { + Malformed(String), + EmptySchema, + UnsupportedType { + column: String, + bq_type: String, + }, + MissingColumn { + column: String, + hint: String, + }, + WrongType { + column: String, + expected: String, + actual: String, + }, + UnfillableColumn(String), +} + +impl fmt::Display for SchemaError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + SchemaError::Malformed(reason) => write!(f, "malformed table schema: {reason}"), + SchemaError::EmptySchema => write!(f, "table has no columns"), + SchemaError::UnsupportedType { column, bq_type } => write!( + f, + "column '{column}' has type {bq_type}, which the BigQuery sink does not support" + ), + SchemaError::MissingColumn { column, hint } => { + write!(f, "column '{column}' is missing from the table. {hint}") + } + SchemaError::WrongType { + column, + expected, + actual, + } => write!(f, "column '{column}' must be {expected}, found {actual}"), + SchemaError::UnfillableColumn(column) => write!( + f, + "column '{column}' is REQUIRED without a default value and raw mode never writes it" + ), + } + } +} + +impl MetaColumn { + pub(crate) const ALL: [MetaColumn; 7] = [ + MetaColumn::Stream, + MetaColumn::Topic, + MetaColumn::PartitionId, + MetaColumn::Offset, + MetaColumn::Timestamp, + MetaColumn::Id, + MetaColumn::Headers, + ]; + + pub(crate) fn name(self) -> &'static str { + match self { + MetaColumn::Stream => "iggy_stream", + MetaColumn::Topic => "iggy_topic", + MetaColumn::PartitionId => "iggy_partition_id", + MetaColumn::Offset => "iggy_offset", + MetaColumn::Timestamp => "iggy_timestamp", + MetaColumn::Id => "iggy_id", + MetaColumn::Headers => "iggy_headers", + } + } + + fn ddl_type(self) -> &'static str { + match self { + MetaColumn::Stream | MetaColumn::Topic | MetaColumn::Id => "STRING", + MetaColumn::PartitionId | MetaColumn::Offset => "INT64", + MetaColumn::Timestamp => "TIMESTAMP", + MetaColumn::Headers => "JSON", + } + } + + fn accepts(self, bq_type: &BqType) -> bool { + match self { + MetaColumn::Stream | MetaColumn::Topic | MetaColumn::Id => { + matches!(bq_type, BqType::String) + } + MetaColumn::PartitionId | MetaColumn::Offset => matches!(bq_type, BqType::Int64), + MetaColumn::Timestamp => matches!(bq_type, BqType::Timestamp), + MetaColumn::Headers => matches!(bq_type, BqType::Json | BqType::String), + } + } + + fn enabled(self, settings: &Settings) -> bool { + match self { + MetaColumn::Headers => settings.include_headers, + _ => settings.include_metadata, + } + } +} + +/// Parse the body of a `tables.get` response. +pub(crate) fn parse_table_schema(body: &[u8]) -> Result, SchemaError> { + let resource: TableResource = + serde_json::from_slice(body).map_err(|e| SchemaError::Malformed(e.to_string()))?; + let fields = resource + .schema + .map(|schema| schema.fields) + .unwrap_or_default(); + if fields.is_empty() { + return Err(SchemaError::EmptySchema); + } + fields.iter().map(Column::try_from).collect() +} + +impl TableLayout { + pub(crate) fn build(columns: Vec, settings: &Settings) -> Result { + let mut by_name: HashMap<&str, &Column> = + columns.iter().map(|c| (c.name.as_str(), c)).collect(); + + let mut metadata = Vec::new(); + for meta in MetaColumn::ALL { + if !meta.enabled(settings) { + continue; + } + let column = by_name + .remove(meta.name()) + .ok_or_else(|| SchemaError::MissingColumn { + column: meta.name().to_owned(), + hint: format!( + "Add `{} {}` to the table or disable it in the connector config.", + meta.name(), + meta.ddl_type() + ), + })?; + if column.mode == Mode::Repeated || !meta.accepts(&column.bq_type) { + return Err(SchemaError::WrongType { + column: column.name.clone(), + expected: meta.ddl_type().to_owned(), + actual: column.type_label(), + }); + } + metadata.push((meta, Arc::new(meta.arrow_field(column.mode)))); + } + + match settings.mode { + WriteMode::Raw => Self::build_raw(&columns, by_name, metadata, settings), + WriteMode::Mapped => Ok(Self::build_mapped(&columns, &metadata)), + } + } + + fn build_raw( + columns: &[Column], + mut remaining: HashMap<&str, &Column>, + metadata: Vec<(MetaColumn, FieldRef)>, + settings: &Settings, + ) -> Result { + let name = settings.payload_column.as_str(); + let column = remaining + .remove(name) + .ok_or_else(|| SchemaError::MissingColumn { + column: name.to_owned(), + hint: "Add a JSON, STRING or BYTES column for the payload or set payload_column." + .to_owned(), + })?; + let kind = match (&column.bq_type, column.mode) { + (BqType::Json, Mode::Nullable | Mode::Required) => RawKind::Json, + (BqType::String, Mode::Nullable | Mode::Required) => RawKind::String, + (BqType::Bytes, Mode::Nullable | Mode::Required) => RawKind::Bytes, + _ => { + return Err(SchemaError::WrongType { + column: column.name.clone(), + expected: "JSON, STRING or BYTES".to_owned(), + actual: column.type_label(), + }); + } + }; + + // Every other column is left out of the writer schema, so BigQuery + // fills it from its default. A REQUIRED column without one would + // reject every row. + if let Some(unfillable) = columns.iter().find(|c| { + remaining.contains_key(c.name.as_str()) && c.mode == Mode::Required && !c.has_default + }) { + return Err(SchemaError::UnfillableColumn(unfillable.name.clone())); + } + + let field = Arc::new(column.arrow_field()); + Ok(TableLayout { + columns: Vec::new(), + fields: Vec::new(), + index: HashMap::new(), + metadata, + raw: Some((kind, field)), + }) + } + + fn build_mapped(columns: &[Column], metadata: &[(MetaColumn, FieldRef)]) -> Self { + let payload_columns: Vec = columns + .iter() + .filter(|c| !metadata.iter().any(|(meta, _)| meta.name() == c.name)) + .cloned() + .collect(); + let fields = payload_columns + .iter() + .map(|c| Arc::new(c.arrow_field())) + .collect(); + let index = payload_columns + .iter() + .enumerate() + .map(|(i, c)| (c.name.clone(), i)) + .collect(); + TableLayout { + columns: payload_columns, + fields, + index, + metadata: metadata.to_vec(), + raw: None, + } + } +} + +impl Column { + /// Arrow field as the Storage Write API expects it. REQUIRED maps to a + /// non-nullable field. REPEATED maps to a list whose items are + /// non-nullable, because BigQuery arrays cannot hold NULL. + pub(crate) fn arrow_field(&self) -> Field { + let element = self.bq_type.arrow_type(); + match self.mode { + Mode::Nullable => Field::new(&self.name, element, true), + Mode::Required => Field::new(&self.name, element, false), + Mode::Repeated => Field::new( + &self.name, + DataType::List(Arc::new(Field::new(LIST_ITEM, element, false))), + true, + ), + } + } + + fn type_label(&self) -> String { + let base = self.bq_type.label(); + match self.mode { + Mode::Repeated => format!("REPEATED {base}"), + _ => base.to_owned(), + } + } +} + +impl BqType { + fn arrow_type(&self) -> DataType { + match self { + BqType::String | BqType::Geography | BqType::Json => DataType::Utf8, + BqType::Bytes => DataType::Binary, + BqType::Int64 => DataType::Int64, + BqType::Float64 => DataType::Float64, + BqType::Bool => DataType::Boolean, + BqType::Numeric => DataType::Decimal128(38, 9), + BqType::BigNumeric => DataType::Decimal256(76, 38), + BqType::Timestamp => DataType::Timestamp(TimeUnit::Microsecond, Some(UTC.into())), + BqType::Datetime => DataType::Timestamp(TimeUnit::Microsecond, None), + BqType::Date => DataType::Date32, + BqType::Time => DataType::Time64(TimeUnit::Microsecond), + BqType::Record(children) => DataType::Struct(Fields::from( + children.iter().map(Column::arrow_field).collect::>(), + )), + } + } + + fn label(&self) -> &'static str { + match self { + BqType::String => "STRING", + BqType::Bytes => "BYTES", + BqType::Int64 => "INT64", + BqType::Float64 => "FLOAT64", + BqType::Bool => "BOOL", + BqType::Numeric => "NUMERIC", + BqType::BigNumeric => "BIGNUMERIC", + BqType::Timestamp => "TIMESTAMP", + BqType::Datetime => "DATETIME", + BqType::Date => "DATE", + BqType::Time => "TIME", + BqType::Geography => "GEOGRAPHY", + BqType::Json => "JSON", + BqType::Record(_) => "RECORD", + } + } +} + +impl MetaColumn { + fn arrow_field(self, mode: Mode) -> Field { + let data_type = match self { + MetaColumn::Stream | MetaColumn::Topic | MetaColumn::Id | MetaColumn::Headers => { + DataType::Utf8 + } + MetaColumn::PartitionId | MetaColumn::Offset => DataType::Int64, + MetaColumn::Timestamp => DataType::Timestamp(TimeUnit::Microsecond, Some(UTC.into())), + }; + Field::new(self.name(), data_type, mode != Mode::Required) + } +} + +// ─── REST representation ───────────────────────────────────────────────────── + +#[derive(Debug, Deserialize)] +struct TableResource { + schema: Option, +} + +#[derive(Debug, Deserialize)] +struct RestSchema { + #[serde(default)] + fields: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct RestField { + name: String, + #[serde(rename = "type")] + field_type: String, + mode: Option, + #[serde(default)] + fields: Vec, + default_value_expression: Option, +} + +impl TryFrom<&RestField> for Column { + type Error = SchemaError; + + fn try_from(field: &RestField) -> Result { + let bq_type = match field.field_type.to_ascii_uppercase().as_str() { + "STRING" => BqType::String, + "BYTES" => BqType::Bytes, + "INTEGER" | "INT64" => BqType::Int64, + "FLOAT" | "FLOAT64" => BqType::Float64, + "BOOLEAN" | "BOOL" => BqType::Bool, + "NUMERIC" | "DECIMAL" => BqType::Numeric, + "BIGNUMERIC" | "BIGDECIMAL" => BqType::BigNumeric, + "TIMESTAMP" => BqType::Timestamp, + "DATETIME" => BqType::Datetime, + "DATE" => BqType::Date, + "TIME" => BqType::Time, + "GEOGRAPHY" => BqType::Geography, + "JSON" => BqType::Json, + "RECORD" | "STRUCT" => { + if field.fields.is_empty() { + return Err(SchemaError::Malformed(format!( + "RECORD column '{}' has no fields", + field.name + ))); + } + BqType::Record( + field + .fields + .iter() + .map(Column::try_from) + .collect::>()?, + ) + } + other => { + return Err(SchemaError::UnsupportedType { + column: field.name.clone(), + bq_type: other.to_owned(), + }); + } + }; + let mode = match field + .mode + .as_deref() + .map(str::to_ascii_uppercase) + .as_deref() + { + None | Some("NULLABLE") => Mode::Nullable, + Some("REQUIRED") => Mode::Required, + Some("REPEATED") => Mode::Repeated, + Some(other) => { + return Err(SchemaError::Malformed(format!( + "column '{}' has unknown mode {other}", + field.name + ))); + } + }; + Ok(Column { + name: field.name.clone(), + bq_type, + mode, + has_default: field.default_value_expression.is_some(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::settings; + + fn schema_json(fields: &str) -> Vec { + format!(r#"{{"schema":{{"fields":[{fields}]}}}}"#).into_bytes() + } + + const METADATA_FIELDS: &str = r#" + {"name":"iggy_stream","type":"STRING"}, + {"name":"iggy_topic","type":"STRING"}, + {"name":"iggy_partition_id","type":"INTEGER"}, + {"name":"iggy_offset","type":"INTEGER","mode":"REQUIRED"}, + {"name":"iggy_timestamp","type":"TIMESTAMP"}, + {"name":"iggy_id","type":"STRING"}"#; + + #[test] + fn given_every_scalar_type_should_map_to_storage_write_arrow_types() { + let body = schema_json( + r#" + {"name":"s","type":"STRING"}, + {"name":"b","type":"BYTES"}, + {"name":"i","type":"INTEGER"}, + {"name":"f","type":"FLOAT"}, + {"name":"ok","type":"BOOLEAN"}, + {"name":"n","type":"NUMERIC"}, + {"name":"bn","type":"BIGNUMERIC"}, + {"name":"ts","type":"TIMESTAMP"}, + {"name":"dt","type":"DATETIME"}, + {"name":"d","type":"DATE"}, + {"name":"tm","type":"TIME"}, + {"name":"g","type":"GEOGRAPHY"}, + {"name":"j","type":"JSON"}"#, + ); + let columns = parse_table_schema(&body).unwrap(); + let types: Vec = columns + .iter() + .map(|c| c.arrow_field().data_type().clone()) + .collect(); + assert_eq!( + types, + vec![ + DataType::Utf8, + DataType::Binary, + DataType::Int64, + DataType::Float64, + DataType::Boolean, + DataType::Decimal128(38, 9), + DataType::Decimal256(76, 38), + DataType::Timestamp(TimeUnit::Microsecond, Some(UTC.into())), + DataType::Timestamp(TimeUnit::Microsecond, None), + DataType::Date32, + DataType::Time64(TimeUnit::Microsecond), + DataType::Utf8, + DataType::Utf8, + ] + ); + } + + #[test] + fn given_nested_repeated_record_should_map_to_list_of_struct() { + let body = schema_json( + r#"{"name":"items","type":"RECORD","mode":"REPEATED","fields":[ + {"name":"sku","type":"STRING","mode":"REQUIRED"}, + {"name":"qty","type":"INT64"} + ]}"#, + ); + let columns = parse_table_schema(&body).unwrap(); + let field = columns[0].arrow_field(); + let DataType::List(item) = field.data_type() else { + panic!("expected list, got {:?}", field.data_type()); + }; + assert!(!item.is_nullable()); + let DataType::Struct(children) = item.data_type() else { + panic!("expected struct item"); + }; + assert_eq!(children[0].name(), "sku"); + assert!(!children[0].is_nullable()); + assert!(children[1].is_nullable()); + } + + #[test] + fn given_required_column_should_be_non_nullable() { + let columns = parse_table_schema(&schema_json( + r#"{"name":"id","type":"INT64","mode":"REQUIRED"}"#, + )) + .unwrap(); + assert!(!columns[0].arrow_field().is_nullable()); + } + + #[test] + fn given_interval_column_should_be_rejected() { + let result = parse_table_schema(&schema_json(r#"{"name":"span","type":"INTERVAL"}"#)); + assert!(matches!( + result, + Err(SchemaError::UnsupportedType { ref column, .. }) if column == "span" + )); + } + + #[test] + fn given_nested_range_column_should_be_rejected() { + let result = parse_table_schema(&schema_json( + r#"{"name":"r","type":"RECORD","fields":[{"name":"window","type":"RANGE"}]}"#, + )); + assert!(matches!( + result, + Err(SchemaError::UnsupportedType { ref column, .. }) if column == "window" + )); + } + + #[test] + fn given_table_without_schema_should_fail() { + assert_eq!(parse_table_schema(b"{}"), Err(SchemaError::EmptySchema)); + } + + #[test] + fn given_default_value_expression_should_mark_column_with_default() { + let columns = parse_table_schema(&schema_json( + r#"{"name":"created","type":"TIMESTAMP","mode":"REQUIRED","defaultValueExpression":"CURRENT_TIMESTAMP()"}"#, + )) + .unwrap(); + assert!(columns[0].has_default); + } + + #[test] + fn given_mapped_mode_with_metadata_should_split_payload_and_metadata_columns() { + let body = schema_json(&format!( + r#"{{"name":"user_id","type":"INT64"}},{METADATA_FIELDS},{{"name":"event","type":"STRING"}}"# + )); + let layout = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Mapped, true), + ) + .unwrap(); + let payload: Vec<&str> = layout.columns.iter().map(|c| c.name.as_str()).collect(); + assert_eq!(payload, vec!["user_id", "event"]); + assert_eq!(layout.metadata.len(), 6); + let offset = layout + .metadata + .iter() + .find(|(meta, _)| *meta == MetaColumn::Offset) + .unwrap(); + assert!(!offset.1.is_nullable()); + assert_eq!(layout.index["event"], 1); + } + + #[test] + fn given_metadata_enabled_but_column_missing_should_fail_with_hint() { + let body = schema_json(r#"{"name":"user_id","type":"INT64"}"#); + let result = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Mapped, true), + ); + let Err(SchemaError::MissingColumn { column, hint }) = result else { + panic!("expected MissingColumn"); + }; + assert_eq!(column, "iggy_stream"); + assert!(hint.contains("iggy_stream STRING")); + } + + #[test] + fn given_metadata_column_with_wrong_type_should_fail() { + let body = schema_json( + r#"{"name":"iggy_stream","type":"STRING"}, + {"name":"iggy_topic","type":"STRING"}, + {"name":"iggy_partition_id","type":"STRING"}"#, + ); + let result = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Mapped, true), + ); + assert!(matches!( + result, + Err(SchemaError::WrongType { ref column, .. }) if column == "iggy_partition_id" + )); + } + + #[test] + fn given_metadata_disabled_should_treat_iggy_columns_as_payload() { + let body = schema_json(r#"{"name":"iggy_offset","type":"INT64"}"#); + let layout = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Mapped, false), + ) + .unwrap(); + assert!(layout.metadata.is_empty()); + assert_eq!(layout.columns.len(), 1); + } + + #[test] + fn given_raw_mode_with_json_payload_column_should_build() { + let body = schema_json(&format!( + r#"{{"name":"payload","type":"JSON"}},{METADATA_FIELDS}"# + )); + let layout = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Raw, true), + ) + .unwrap(); + let (kind, field) = layout.raw.as_ref().unwrap(); + assert_eq!(*kind, RawKind::Json); + assert_eq!(field.data_type(), &DataType::Utf8); + assert!(layout.columns.is_empty()); + } + + #[test] + fn given_raw_mode_without_payload_column_should_fail() { + let body = schema_json(r#"{"name":"data","type":"BYTES"}"#); + let result = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Raw, false), + ); + assert!(matches!( + result, + Err(SchemaError::MissingColumn { ref column, .. }) if column == "payload" + )); + } + + #[test] + fn given_raw_mode_with_numeric_payload_column_should_fail() { + let body = schema_json(r#"{"name":"payload","type":"INT64"}"#); + let result = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Raw, false), + ); + assert!(matches!(result, Err(SchemaError::WrongType { .. }))); + } + + #[test] + fn given_raw_mode_with_required_column_without_default_should_fail() { + let body = schema_json( + r#"{"name":"payload","type":"BYTES"}, + {"name":"tenant","type":"STRING","mode":"REQUIRED"}"#, + ); + let result = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Raw, false), + ); + assert_eq!( + result.unwrap_err(), + SchemaError::UnfillableColumn("tenant".into()) + ); + } + + #[test] + fn given_raw_mode_with_required_column_with_default_should_build() { + let body = schema_json( + r#"{"name":"payload","type":"BYTES"}, + {"name":"tenant","type":"STRING","mode":"REQUIRED","defaultValueExpression":"'acme'"}"#, + ); + let layout = TableLayout::build( + parse_table_schema(&body).unwrap(), + &settings(WriteMode::Raw, false), + ) + .unwrap(); + assert_eq!(layout.raw.unwrap().0, RawKind::Bytes); + } +} diff --git a/core/connectors/sinks/bigquery_sink/src/sink.rs b/core/connectors/sinks/bigquery_sink/src/sink.rs new file mode 100644 index 0000000000..a13368d30d --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/src/sink.rs @@ -0,0 +1,303 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +use crate::client::{AppendOutcome, BigQueryClient}; +use crate::encode::{self, Chunk, RunContext}; +use crate::error::{AppendError, TableError}; +use crate::schema::{TableLayout, parse_table_schema}; +use crate::{BigQuerySink, OpenState}; +use async_trait::async_trait; +use iggy_connector_sdk::retry::{RetryPolicy, retry_async}; +use iggy_connector_sdk::{ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata}; +use std::sync::atomic::Ordering; +use tracing::{debug, error, info, warn}; + +#[async_trait] +impl Sink for BigQuerySink { + async fn open(&mut self) -> Result<(), Error> { + info!( + "Opening BigQuery sink connector ID: {} for table {} (mode: {:?})", + self.id, self.target, self.settings.mode + ); + // The SDK turns a failed open into a bare status code, so this is the + // only place the reason reaches the runtime log. + self.try_open().await.inspect_err(|e| { + error!( + "Failed to open BigQuery sink connector ID: {} for table {}: {e}", + self.id, self.target + ); + }) + } + + async fn consume( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: MessagesMetadata, + messages: Vec, + ) -> Result<(), Error> { + if messages.is_empty() { + return Ok(()); + } + let Some(state) = &self.state else { + return Err(Error::InitError("BigQuery sink is not open".into())); + }; + let received = messages.len(); + debug!( + "BigQuery sink ID: {} received {received} messages from {}/{} partition {} current_offset {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + messages_metadata.current_offset + ); + + let ctx = RunContext { + layout: &state.layout, + topic: topic_metadata, + messages: &messages_metadata, + max_request_bytes: self.settings.max_request_bytes, + }; + let encoded = encode::encode(&ctx, messages).map_err(|e| { + self.counters + .rows_failed + .fetch_add(received as u64, Ordering::Relaxed); + error!( + "BigQuery sink ID: {} cannot encode {received} messages from {}/{} partition {}: {e}", + self.id, topic_metadata.stream, topic_metadata.topic, messages_metadata.partition_id + ); + Error::Serialization(e.to_string()) + })?; + + for rejected in &encoded.rejected { + warn!( + "BigQuery sink ID: {} dropped message stream={} topic={} partition={} offset={}: {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + rejected.offset, + rejected.reason + ); + } + self.counters + .rows_rejected + .fetch_add(encoded.rejected.len() as u64, Ordering::Relaxed); + + let mut written = 0u64; + let mut last_error = None; + for chunk in &encoded.chunks { + match self + .write_chunk(&state.client, chunk, topic_metadata, &messages_metadata) + .await + { + Ok(rows) => written += rows, + Err(e) => last_error = Some(e), + } + } + self.counters + .rows_written + .fetch_add(written, Ordering::Relaxed); + + if self.settings.verbose { + info!( + "BigQuery sink ID: {} wrote {written} of {received} messages from {}/{} partition {} to {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + self.target + ); + } else { + debug!( + "BigQuery sink ID: {} wrote {written} of {received} messages from {}/{} partition {} to {}", + self.id, + topic_metadata.stream, + topic_metadata.topic, + messages_metadata.partition_id, + self.target + ); + } + + match last_error { + Some(e) => Err(e), + None => Ok(()), + } + } + + async fn close(&mut self) -> Result<(), Error> { + self.state = None; + info!( + "Closed BigQuery sink connector ID: {}, rows written: {}, rows rejected: {}, rows failed: {}", + self.id, + self.counters.rows_written.load(Ordering::Relaxed), + self.counters.rows_rejected.load(Ordering::Relaxed), + self.counters.rows_failed.load(Ordering::Relaxed) + ); + Ok(()) + } +} + +impl BigQuerySink { + async fn try_open(&mut self) -> Result<(), Error> { + self.config.validate()?; + + let mut client = BigQueryClient::connect(&self.config, &self.settings).await?; + let policy = self.retry_policy(); + + let context = format!("BigQuery sink ID: {} tables.get", self.id); + let body = retry_async(policy, &context, TableError::is_retryable, || { + client.fetch_table() + }) + .await + .map_err(|failure| Error::InitError(format!("cannot read the table schema: {failure}")))?; + + let columns = parse_table_schema(&body).map_err(|e| Error::from(TableError::Schema(e)))?; + let layout = TableLayout::build(columns, &self.settings) + .map_err(|e| Error::from(TableError::Schema(e)))?; + + let context = format!("BigQuery sink ID: {} write stream setup", self.id); + let stream = retry_async(policy, &context, AppendError::is_retryable, || { + client.resolve_write_stream() + }) + .await + .map_err(|failure| Error::InitError(format!("cannot open the write stream: {failure}")))?; + client.set_write_stream(stream); + + info!( + "Opened BigQuery sink connector ID: {} for table {}: {} payload column(s), {} metadata column(s)", + self.id, + self.target, + layout.raw.as_ref().map_or(layout.columns.len(), |_| 1), + layout.metadata.len() + ); + self.state = Some(OpenState { client, layout }); + Ok(()) + } + + fn retry_policy(&self) -> RetryPolicy { + RetryPolicy { + max_attempts: self.settings.max_retries, + base_delay: self.settings.retry_delay, + max_delay: self.settings.max_retry_delay, + } + } + + /// Append one chunk. When BigQuery reports row errors, nothing was + /// appended: drop the reported rows and append the rest once more. + /// Returns the number of rows written. + async fn write_chunk( + &self, + client: &BigQueryClient, + chunk: &Chunk, + topic: &TopicMetadata, + messages: &MessagesMetadata, + ) -> Result { + let row_errors = match self + .append_with_retry(client, chunk, topic, messages) + .await? + { + AppendOutcome::Appended => return Ok(chunk.offsets.len() as u64), + AppendOutcome::RowErrors(row_errors) => row_errors, + }; + + let mut dropped = Vec::with_capacity(row_errors.len()); + for (row, reason) in row_errors { + let offset = chunk.offsets.get(row).copied(); + warn!( + "BigQuery sink ID: {} dropped row rejected by BigQuery stream={} topic={} partition={} offset={}: {reason}", + self.id, + topic.stream, + topic.topic, + messages.partition_id, + offset.map_or_else(|| format!("unknown (row {row})"), |o| o.to_string()) + ); + dropped.push(row); + } + self.counters + .rows_rejected + .fetch_add(dropped.len() as u64, Ordering::Relaxed); + + let Some(retry) = encode::without_rows(chunk, &dropped) + .map_err(|e| Error::Serialization(e.to_string()))? + else { + return Ok(0); + }; + match self + .append_with_retry(client, &retry, topic, messages) + .await? + { + AppendOutcome::Appended => Ok(retry.offsets.len() as u64), + AppendOutcome::RowErrors(again) => { + self.counters + .rows_failed + .fetch_add(retry.offsets.len() as u64, Ordering::Relaxed); + error!( + "BigQuery sink ID: {} dropped {} rows (offsets {}..={}) of {}/{} partition {}: BigQuery reported {} more row error(s) after the bad rows were removed", + self.id, + retry.offsets.len(), + retry.offsets.first().copied().unwrap_or_default(), + retry.offsets.last().copied().unwrap_or_default(), + topic.stream, + topic.topic, + messages.partition_id, + again.len() + ); + Err(Error::PermanentHttpError(format!( + "{} row error(s) persisted after removing rejected rows", + again.len() + ))) + } + } + } + + async fn append_with_retry( + &self, + client: &BigQueryClient, + chunk: &Chunk, + topic: &TopicMetadata, + messages: &MessagesMetadata, + ) -> Result { + let context = format!( + "BigQuery sink ID: {} AppendRows ({} rows)", + self.id, + chunk.offsets.len() + ); + retry_async( + self.retry_policy(), + &context, + AppendError::is_retryable, + || client.append(chunk), + ) + .await + .map_err(|failure| { + self.counters + .rows_failed + .fetch_add(chunk.offsets.len() as u64, Ordering::Relaxed); + error!( + "BigQuery sink ID: {} lost {} rows (offsets {}..={}) of {}/{} partition {}: {failure}", + self.id, + chunk.offsets.len(), + chunk.offsets.first().copied().unwrap_or_default(), + chunk.offsets.last().copied().unwrap_or_default(), + topic.stream, + topic.topic, + messages.partition_id + ); + Error::from(failure.error) + }) + } +} diff --git a/core/connectors/sinks/bigquery_sink/tests/bigquery_sink.rs b/core/connectors/sinks/bigquery_sink/tests/bigquery_sink.rs new file mode 100644 index 0000000000..67501176e5 --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/tests/bigquery_sink.rs @@ -0,0 +1,424 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! End-to-end tests of the sink against the in-process fake BigQuery. + +mod common; + +use arrow::array::{Array, AsArray, RecordBatch}; +use arrow::datatypes::Int64Type; +use common::{AppendScript, FakeBigQuery}; +use iggy_connector_bigquery_sink::{BigQuerySink, BigQuerySinkConfig}; +use iggy_connector_sdk::{ + ConsumedMessage, Error, MessagesMetadata, Payload, Schema, Sink, TopicMetadata, +}; +use tonic::Code; + +const EVENTS_TABLE: &str = r#" + {"name":"user_id","type":"INT64","mode":"REQUIRED"}, + {"name":"event","type":"STRING"}, + {"name":"created","type":"TIMESTAMP","defaultValueExpression":"CURRENT_TIMESTAMP()"}, + {"name":"iggy_stream","type":"STRING"}, + {"name":"iggy_topic","type":"STRING"}, + {"name":"iggy_partition_id","type":"INT64"}, + {"name":"iggy_offset","type":"INT64"}, + {"name":"iggy_timestamp","type":"TIMESTAMP"}, + {"name":"iggy_id","type":"STRING"}"#; + +fn config(fake: &FakeBigQuery, extra: &str) -> BigQuerySinkConfig { + toml::from_str(&format!( + r#" + project_id = "proj" + dataset = "ds" + table = "events" + endpoint = "{}" + grpc_endpoint = "{}" + retry_delay = "5ms" + max_retry_delay = "20ms" + {extra} + "#, + fake.rest_url, fake.grpc_addr + )) + .expect("valid test config") +} + +async fn open_sink(fake: &FakeBigQuery, extra: &str) -> BigQuerySink { + let mut sink = BigQuerySink::new(1, config(fake, extra)); + sink.open().await.expect("sink should open"); + sink +} + +fn topic() -> TopicMetadata { + TopicMetadata { + stream: "shop".into(), + topic: "events".into(), + } +} + +fn metadata() -> MessagesMetadata { + MessagesMetadata { + partition_id: 2, + current_offset: 99, + schema: Schema::Json, + } +} + +fn json_message(offset: u64, text: &str) -> ConsumedMessage { + let mut bytes = text.as_bytes().to_vec(); + ConsumedMessage { + id: u128::from(offset), + offset, + checksum: 0, + timestamp: 1_700_000_000_000_000, + origin_timestamp: 0, + headers: None, + payload: Payload::Json(simd_json::to_owned_value(&mut bytes).unwrap()), + } +} + +fn events(offsets: std::ops::Range) -> Vec { + offsets + .map(|offset| { + json_message( + offset, + &format!(r#"{{"user_id": {offset}, "event": "e{offset}"}}"#), + ) + }) + .collect() +} + +fn offsets(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("iggy_offset") + .expect("iggy_offset column") + .as_primitive::() + .values() + .to_vec() + }) + .collect() +} + +#[tokio::test] +async fn given_json_messages_should_append_rows_with_metadata() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + let sink = open_sink(&fake, "").await; + + sink.consume(&topic(), metadata(), events(0..3)) + .await + .expect("consume should succeed"); + + let appends = fake.appends(); + assert_eq!(appends.len(), 1); + let append = &appends[0]; + assert_eq!( + append.write_stream, + "projects/proj/datasets/ds/tables/events/streams/_default" + ); + assert_eq!(append.default_missing_value_interpretation, 2); + + let batch = &append.batch; + let names: Vec<&str> = batch + .schema_ref() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect(); + assert!(!names.contains(&"created"), "unset column must be omitted"); + assert_eq!( + batch + .column_by_name("user_id") + .unwrap() + .as_primitive::() + .values(), + &[0, 1, 2] + ); + assert_eq!( + batch + .column_by_name("iggy_partition_id") + .unwrap() + .as_primitive::() + .value(0), + 2 + ); + assert_eq!( + batch + .column_by_name("iggy_stream") + .unwrap() + .as_string::() + .value(0), + "shop" + ); + assert_eq!(offsets(&fake.stored()), vec![0, 1, 2]); +} + +#[tokio::test] +async fn given_null_missing_value_should_send_null_interpretation() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + let sink = open_sink(&fake, r#"missing_value = "null""#).await; + sink.consume(&topic(), metadata(), events(0..1)) + .await + .unwrap(); + assert_eq!(fake.appends()[0].default_missing_value_interpretation, 1); +} + +#[tokio::test] +async fn given_batch_over_request_budget_should_split_into_several_appends() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + let sink = open_sink(&fake, "max_request_bytes = 65536").await; + let messages = (0..60) + .map(|offset| { + json_message( + offset, + &format!( + r#"{{"user_id": {offset}, "event": "{}"}}"#, + "x".repeat(4096) + ), + ) + }) + .collect(); + + sink.consume(&topic(), metadata(), messages).await.unwrap(); + + assert!(fake.appends().len() > 1); + assert_eq!(offsets(&fake.stored()), (0..60).collect::>()); +} + +#[tokio::test] +async fn given_invalid_rows_should_drop_them_and_write_the_rest() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + let sink = open_sink(&fake, "").await; + let messages = vec![ + json_message(0, r#"{"user_id": 0}"#), + json_message(1, r#"{"event": "no user"}"#), + json_message(2, r#"{"user_id": "abc"}"#), + json_message(3, r#"{"user_id": 3}"#), + ]; + + sink.consume(&topic(), metadata(), messages) + .await + .expect("rejected rows do not fail the batch"); + + assert_eq!(offsets(&fake.stored()), vec![0, 3]); +} + +#[tokio::test] +async fn given_row_errors_should_reappend_without_rejected_rows() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.script(&[AppendScript::RowErrors(vec![1, 3])]); + let sink = open_sink(&fake, "").await; + + sink.consume(&topic(), metadata(), events(10..15)) + .await + .expect("row errors are handled by dropping rows"); + + let appends = fake.appends(); + assert_eq!(appends.len(), 2); + assert!(!appends[0].stored); + assert_eq!(offsets(&fake.stored()), vec![10, 12, 14]); +} + +#[tokio::test] +async fn given_row_errors_twice_should_fail_the_chunk() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.script(&[ + AppendScript::RowErrors(vec![0]), + AppendScript::RowErrors(vec![0]), + ]); + let sink = open_sink(&fake, "").await; + + let result = sink.consume(&topic(), metadata(), events(0..3)).await; + + assert!( + matches!(result, Err(Error::PermanentHttpError(_))), + "{result:?}" + ); + assert!(fake.stored().is_empty()); +} + +#[tokio::test] +async fn given_transient_failure_should_retry_and_succeed() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.script(&[ + AppendScript::CallError(Code::Unavailable), + AppendScript::ResponseError(Code::ResourceExhausted), + ]); + let sink = open_sink(&fake, "").await; + + sink.consume(&topic(), metadata(), events(0..2)) + .await + .expect("should succeed on the third attempt"); + + assert_eq!(offsets(&fake.stored()), vec![0, 1]); +} + +#[tokio::test] +async fn given_transient_failure_beyond_retry_budget_should_fail() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.script(&[ + AppendScript::CallError(Code::Unavailable), + AppendScript::CallError(Code::Unavailable), + AppendScript::CallError(Code::Unavailable), + ]); + let sink = open_sink(&fake, "max_retries = 2").await; + + let result = sink.consume(&topic(), metadata(), events(0..2)).await; + + assert!( + matches!(result, Err(Error::CannotStoreData(_))), + "{result:?}" + ); + assert!(fake.stored().is_empty()); + assert_eq!(fake.pending_script(), 1, "only two attempts were made"); +} + +#[tokio::test] +async fn given_invalid_argument_should_fail_without_retry() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.script(&[AppendScript::ResponseError(Code::InvalidArgument)]); + let sink = open_sink(&fake, "").await; + + let result = sink.consume(&topic(), metadata(), events(0..2)).await; + + assert!( + matches!(result, Err(Error::PermanentHttpError(_))), + "{result:?}" + ); + assert_eq!(fake.appends().len(), 1); +} + +#[tokio::test] +async fn given_permission_and_auth_errors_should_be_permanent() { + for code in [Code::PermissionDenied, Code::Unauthenticated] { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.script(&[AppendScript::CallError(code), AppendScript::Ok]); + let sink = open_sink(&fake, "").await; + + let result = sink.consume(&topic(), metadata(), events(0..1)).await; + + assert!( + matches!(result, Err(Error::PermanentHttpError(_))), + "{code:?}: {result:?}" + ); + assert!(fake.stored().is_empty(), "{code:?} must not be retried"); + } +} + +#[tokio::test] +async fn given_unsupported_column_type_should_fail_open() { + let fake = FakeBigQuery::start(r#"{"name":"span","type":"INTERVAL"}"#).await; + let mut sink = BigQuerySink::new(1, config(&fake, "include_metadata = false")); + + let result = sink.open().await; + + let Err(Error::InitError(reason)) = result else { + panic!("expected InitError, got {result:?}"); + }; + assert!(reason.contains("span"), "{reason}"); +} + +#[tokio::test] +async fn given_missing_metadata_column_should_fail_open() { + let fake = FakeBigQuery::start(r#"{"name":"user_id","type":"INT64"}"#).await; + let mut sink = BigQuerySink::new(1, config(&fake, "")); + + let result = sink.open().await; + + let Err(Error::InitError(reason)) = result else { + panic!("expected InitError, got {result:?}"); + }; + assert!(reason.contains("iggy_stream STRING"), "{reason}"); +} + +#[tokio::test] +async fn given_table_not_found_should_fail_open_without_retry() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.table_responses(&[(404, r#"{"error":{"message":"Not found: Table"}}"#)]); + let mut sink = BigQuerySink::new(1, config(&fake, "")); + + let result = sink.open().await; + + assert!(matches!(result, Err(Error::InitError(_))), "{result:?}"); + assert_eq!(fake.table_calls(), 1); +} + +#[tokio::test] +async fn given_table_lookup_unavailable_once_should_retry_and_open() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.table_responses(&[(503, "busy")]); + let mut sink = BigQuerySink::new(1, config(&fake, "")); + + sink.open().await.expect("second tables.get succeeds"); + + assert_eq!(fake.table_calls(), 2); +} + +#[tokio::test] +async fn given_write_stream_permission_denied_should_fail_open() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + fake.fail_get_write_stream(Code::PermissionDenied); + let mut sink = BigQuerySink::new(1, config(&fake, "")); + + let result = sink.open().await; + + assert!(matches!(result, Err(Error::InitError(_))), "{result:?}"); +} + +#[tokio::test] +async fn given_raw_mode_should_store_payload_column() { + let fake = FakeBigQuery::start(r#"{"name":"payload","type":"JSON"}"#).await; + let sink = open_sink(&fake, "mode = \"raw\"\ninclude_metadata = false").await; + let messages = vec![ + json_message(0, r#"{"a": 1}"#), + ConsumedMessage { + payload: Payload::Text("not json".into()), + ..json_message(1, "{}") + }, + ]; + + sink.consume(&topic(), metadata(), messages).await.unwrap(); + + let stored = fake.stored(); + assert_eq!(stored.len(), 1); + let payload = stored[0].column(0).as_string::(); + assert_eq!(payload.len(), 1); + assert_eq!(payload.value(0), r#"{"a":1}"#); +} + +#[tokio::test] +async fn given_empty_batch_should_not_call_bigquery() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + let sink = open_sink(&fake, "").await; + + sink.consume(&topic(), metadata(), Vec::new()) + .await + .unwrap(); + + assert!(fake.appends().is_empty()); +} + +#[tokio::test] +async fn given_consume_before_open_should_fail() { + let fake = FakeBigQuery::start(EVENTS_TABLE).await; + let sink = BigQuerySink::new(1, config(&fake, "")); + + let result = sink.consume(&topic(), metadata(), events(0..1)).await; + + assert!(matches!(result, Err(Error::InitError(_))), "{result:?}"); +} diff --git a/core/connectors/sinks/bigquery_sink/tests/common/mod.rs b/core/connectors/sinks/bigquery_sink/tests/common/mod.rs new file mode 100644 index 0000000000..1d63549310 --- /dev/null +++ b/core/connectors/sinks/bigquery_sink/tests/common/mod.rs @@ -0,0 +1,387 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! In-process fake of the two BigQuery surfaces the sink uses: +//! +//! - REST `tables.get`, served by axum, returning a configurable schema or +//! HTTP error. +//! - gRPC `BigQueryWrite` (`GetWriteStream` and `AppendRows`), served by +//! tonic. No server stubs are published for this service, so the routing +//! below is written by hand in the shape `tonic-build` generates. +//! +//! Each `AppendRows` call takes the next scripted outcome (default: success) +//! and records the decoded Arrow rows, so tests can assert on exactly what +//! the sink sent. + +use arrow::array::RecordBatch; +use arrow::ipc::reader::StreamReader; +use axum::Router; +use axum::extract::State; +use axum::http::StatusCode; +use axum::routing::get; +use gcloud_googleapis::cloud::bigquery::storage::v1::append_rows_request::Rows; +use gcloud_googleapis::cloud::bigquery::storage::v1::append_rows_response::{ + AppendResult, Response, +}; +use gcloud_googleapis::cloud::bigquery::storage::v1::{ + AppendRowsRequest, AppendRowsResponse, GetWriteStreamRequest, RowError, WriteStream, +}; +use gcloud_googleapis::rpc::Status as RpcStatus; +use std::collections::VecDeque; +use std::convert::Infallible; +use std::io::Cursor; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; +use tokio::net::TcpListener; +use tokio_stream::StreamExt; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::codegen::{BoxFuture, BoxStream, Context, Poll, Service, http}; +use tonic::server::{Grpc, NamedService, StreamingService, UnaryService}; +use tonic::{Code, Status}; +use tonic_prost::ProstCodec; + +const GET_WRITE_STREAM: &str = "/google.cloud.bigquery.storage.v1.BigQueryWrite/GetWriteStream"; +const APPEND_ROWS: &str = "/google.cloud.bigquery.storage.v1.BigQueryWrite/AppendRows"; + +/// What the next `AppendRows` call answers. +#[derive(Debug, Clone)] +pub enum AppendScript { + Ok, + /// Response with row errors at these request indexes. Nothing is stored. + RowErrors(Vec), + /// Response carrying an error status. Nothing is stored. + ResponseError(Code), + /// The RPC itself fails. + CallError(Code), +} + +/// One `AppendRows` request as the fake saw it. +#[derive(Debug, Clone)] +pub struct ReceivedAppend { + pub write_stream: String, + pub default_missing_value_interpretation: i32, + pub batch: RecordBatch, + pub stored: bool, +} + +#[derive(Debug)] +struct TableResponse { + status: u16, + body: String, +} + +#[derive(Debug, Default)] +struct FakeState { + table: Mutex>, + table_calls: Mutex, + get_stream_error: Mutex>, + script: Mutex>, + appends: Mutex>, +} + +pub struct FakeBigQuery { + pub rest_url: String, + pub grpc_addr: String, + state: Arc, +} + +impl FakeBigQuery { + /// Start both servers with `schema_fields` (the JSON array body of + /// `schema.fields`) as the table schema. + pub async fn start(schema_fields: &str) -> Self { + let state = Arc::new(FakeState::default()); + state + .table + .lock() + .unwrap() + .push_back(TableResponse::schema(schema_fields)); + + let rest = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let rest_addr = rest.local_addr().unwrap(); + let router = Router::new() + .route( + "/bigquery/v2/projects/{project}/datasets/{dataset}/tables/{table}", + get(tables_get), + ) + .with_state(state.clone()); + tokio::spawn(async move { + axum::serve(rest, router).await.unwrap(); + }); + + let grpc = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let grpc_addr: SocketAddr = grpc.local_addr().unwrap(); + let service = BigQueryWriteService { + state: state.clone(), + }; + tokio::spawn(async move { + tonic::transport::Server::builder() + .add_service(service) + .serve_with_incoming(TcpListenerStream::new(grpc)) + .await + .unwrap(); + }); + + FakeBigQuery { + rest_url: format!("http://{rest_addr}"), + grpc_addr: grpc_addr.to_string(), + state, + } + } + + /// Queue `tables.get` responses ahead of the schema given to `start`. + /// Once the queue holds one entry it is repeated forever. + pub fn table_responses(&self, responses: &[(u16, &str)]) { + let mut table = self.state.table.lock().unwrap(); + let last = table.pop_back(); + table.clear(); + for (status, body) in responses { + table.push_back(TableResponse { + status: *status, + body: (*body).to_owned(), + }); + } + if let Some(last) = last { + table.push_back(last); + } + } + + pub fn fail_get_write_stream(&self, code: Code) { + *self.state.get_stream_error.lock().unwrap() = Some(code); + } + + pub fn script(&self, outcomes: &[AppendScript]) { + self.state + .script + .lock() + .unwrap() + .extend(outcomes.iter().cloned()); + } + + pub fn pending_script(&self) -> usize { + self.state.script.lock().unwrap().len() + } + + pub fn table_calls(&self) -> usize { + *self.state.table_calls.lock().unwrap() + } + + pub fn appends(&self) -> Vec { + self.state.appends.lock().unwrap().clone() + } + + /// Batches the fake accepted, in order. + pub fn stored(&self) -> Vec { + self.appends() + .into_iter() + .filter(|a| a.stored) + .map(|a| a.batch) + .collect() + } +} + +impl TableResponse { + fn schema(fields: &str) -> Self { + TableResponse { + status: 200, + body: format!(r#"{{"schema":{{"fields":[{fields}]}}}}"#), + } + } +} + +async fn tables_get(State(state): State>) -> (StatusCode, String) { + *state.table_calls.lock().unwrap() += 1; + let mut table = state.table.lock().unwrap(); + let response = if table.len() > 1 { + table.pop_front().unwrap() + } else { + let only = table.front().unwrap(); + TableResponse { + status: only.status, + body: only.body.clone(), + } + }; + ( + StatusCode::from_u16(response.status).unwrap(), + response.body, + ) +} + +// ─── gRPC service ──────────────────────────────────────────────────────────── + +#[derive(Clone)] +struct BigQueryWriteService { + state: Arc, +} + +impl NamedService for BigQueryWriteService { + const NAME: &'static str = "google.cloud.bigquery.storage.v1.BigQueryWrite"; +} + +impl Service> for BigQueryWriteService +where + B: tonic::codegen::Body + Send + 'static, + B::Error: Into + Send + 'static, +{ + type Response = http::Response; + type Error = Infallible; + type Future = BoxFuture; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + let state = self.state.clone(); + match request.uri().path() { + GET_WRITE_STREAM => Box::pin(async move { + let mut grpc = Grpc::new(ProstCodec::default()); + Ok(grpc.unary(GetWriteStreamHandler { state }, request).await) + }), + APPEND_ROWS => Box::pin(async move { + let mut grpc = Grpc::new(ProstCodec::default()); + Ok(grpc.streaming(AppendRowsHandler { state }, request).await) + }), + _ => Box::pin(async move { Ok(Status::unimplemented("not faked").into_http()) }), + } + } +} + +struct GetWriteStreamHandler { + state: Arc, +} + +impl UnaryService for GetWriteStreamHandler { + type Response = WriteStream; + type Future = BoxFuture, Status>; + + fn call(&mut self, request: tonic::Request) -> Self::Future { + let error = *self.state.get_stream_error.lock().unwrap(); + Box::pin(async move { + if let Some(code) = error { + return Err(Status::new(code, "scripted GetWriteStream failure")); + } + Ok(tonic::Response::new(WriteStream { + name: request.into_inner().name, + ..Default::default() + })) + }) + } +} + +struct AppendRowsHandler { + state: Arc, +} + +impl StreamingService for AppendRowsHandler { + type Response = AppendRowsResponse; + type ResponseStream = BoxStream; + type Future = BoxFuture, Status>; + + fn call( + &mut self, + request: tonic::Request>, + ) -> Self::Future { + let state = self.state.clone(); + Box::pin(async move { + let mut requests = request.into_inner(); + let mut responses = Vec::new(); + while let Some(append) = requests.next().await { + let append = append?; + let script = state + .script + .lock() + .unwrap() + .pop_front() + .unwrap_or(AppendScript::Ok); + if let AppendScript::CallError(code) = script { + return Err(Status::new(code, "scripted AppendRows failure")); + } + let batch = decode_arrow(&append); + let stored = matches!(script, AppendScript::Ok); + state.appends.lock().unwrap().push(ReceivedAppend { + write_stream: append.write_stream.clone(), + default_missing_value_interpretation: append + .default_missing_value_interpretation, + batch, + stored, + }); + responses.push(Ok(response_for(script, append.write_stream))); + } + let stream: Self::ResponseStream = Box::pin(tokio_stream::iter(responses)); + Ok(tonic::Response::new(stream)) + }) + } +} + +fn response_for(script: AppendScript, write_stream: String) -> AppendRowsResponse { + match script { + AppendScript::Ok | AppendScript::CallError(_) => AppendRowsResponse { + write_stream, + response: Some(Response::AppendResult(AppendResult { offset: None })), + ..Default::default() + }, + AppendScript::RowErrors(indexes) => AppendRowsResponse { + write_stream, + row_errors: indexes + .into_iter() + .map(|index| RowError { + index, + code: 1, + message: format!("scripted error for row {index}"), + }) + .collect(), + response: Some(Response::Error(RpcStatus { + code: Code::InvalidArgument as i32, + message: "rows rejected".into(), + details: Vec::new(), + })), + ..Default::default() + }, + AppendScript::ResponseError(code) => AppendRowsResponse { + write_stream, + response: Some(Response::Error(RpcStatus { + code: code as i32, + message: "scripted response error".into(), + details: Vec::new(), + })), + ..Default::default() + }, + } +} + +fn decode_arrow(append: &AppendRowsRequest) -> RecordBatch { + let Some(Rows::ArrowRows(arrow)) = &append.rows else { + panic!("sink must send Arrow rows"); + }; + let mut bytes = arrow + .writer_schema + .as_ref() + .expect("writer schema") + .serialized_schema + .clone(); + bytes.extend_from_slice( + &arrow + .rows + .as_ref() + .expect("record batch") + .serialized_record_batch, + ); + let reader = StreamReader::try_new(Cursor::new(bytes), None).expect("valid Arrow IPC"); + let batches: Vec = reader.collect::>().expect("valid batch"); + assert_eq!(batches.len(), 1, "one record batch per request"); + batches.into_iter().next().unwrap() +}