From 6ec283d56cdca8427445373735fe2e4cbdc077b1 Mon Sep 17 00:00:00 2001 From: Bebeto Nyamwamu Date: Fri, 17 Jul 2026 22:13:46 -0400 Subject: [PATCH 1/5] feat(connectors): add Redshift sink connector This PR introduces the Redshift sink connector. P.S. a forced redo and push of an earlier commit on the same branch to resolve a divergence that was affecting cli tests(help command) --- Cargo.lock | 344 ++++- Cargo.toml | 4 + core/connectors/README.md | 1 + core/connectors/sinks/README.md | 1 + .../connectors/sinks/redshift_sink/Cargo.toml | 52 + core/connectors/sinks/redshift_sink/README.md | 119 ++ .../sinks/redshift_sink/config.toml | 49 + .../sinks/redshift_sink/src/config.rs | 127 ++ .../connectors/sinks/redshift_sink/src/lib.rs | 1182 +++++++++++++++++ core/integration/Cargo.toml | 5 + .../tests/connectors/fixtures/mod.rs | 5 + .../connectors/fixtures/redshift/container.rs | 265 ++++ .../tests/connectors/fixtures/redshift/mod.rs | 26 + .../fixtures/redshift/redshift_mock/copy.rs | 146 ++ .../fixtures/redshift/redshift_mock/create.rs | 276 ++++ .../redshift/redshift_mock/handler.rs | 305 +++++ .../fixtures/redshift/redshift_mock/load.rs | 603 +++++++++ .../fixtures/redshift/redshift_mock/mod.rs | 56 + .../connectors/fixtures/redshift/sink.rs | 436 ++++++ core/integration/tests/connectors/mod.rs | 1 + .../tests/connectors/redshift/mod.rs | 20 + .../connectors/redshift/redshift_sink.rs | 401 ++++++ .../tests/connectors/redshift/sink.toml | 20 + 23 files changed, 4431 insertions(+), 13 deletions(-) create mode 100644 core/connectors/sinks/redshift_sink/Cargo.toml create mode 100644 core/connectors/sinks/redshift_sink/README.md create mode 100644 core/connectors/sinks/redshift_sink/config.toml create mode 100644 core/connectors/sinks/redshift_sink/src/config.rs create mode 100644 core/connectors/sinks/redshift_sink/src/lib.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/container.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/mod.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/copy.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/create.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/sink.rs create mode 100644 core/integration/tests/connectors/redshift/mod.rs create mode 100644 core/integration/tests/connectors/redshift/redshift_sink.rs create mode 100644 core/integration/tests/connectors/redshift/sink.toml diff --git a/Cargo.lock b/Cargo.lock index 22cbf3cfe2..90908bf8b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -504,6 +504,27 @@ version = "0.7.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f02882884d3e1bc524fb12c79f107f6ad0e1cfd498c536ffb494301740995dfe" +[[package]] +name = "arrow" +version = "57.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3bd47f2a6ddc39244bd722a27ee5da66c03369d087b9e024eafdb03e98b98ea7" +dependencies = [ + "arrow-arith 57.3.1", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-cast 57.3.1", + "arrow-csv 57.3.1", + "arrow-data 57.3.1", + "arrow-ipc 57.3.1", + "arrow-json 57.3.1", + "arrow-ord 57.3.1", + "arrow-row 57.3.1", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", + "arrow-string 57.3.1", +] + [[package]] name = "arrow" version = "58.3.0" @@ -514,12 +535,12 @@ dependencies = [ "arrow-array 58.3.0", "arrow-buffer 58.3.0", "arrow-cast 58.3.0", - "arrow-csv", + "arrow-csv 58.3.0", "arrow-data 58.3.0", "arrow-ipc 58.3.0", "arrow-json 58.3.0", "arrow-ord 58.3.0", - "arrow-row", + "arrow-row 58.3.0", "arrow-schema 58.3.0", "arrow-select 58.3.0", "arrow-string 58.3.0", @@ -656,6 +677,21 @@ dependencies = [ "ryu", ] +[[package]] +name = "arrow-csv" +version = "57.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27ddb80a4848e03b1655af496d5ac2563a779e5742fcb48f2ca2e089c9cd2197" +dependencies = [ + "arrow-array 57.3.1", + "arrow-cast 57.3.1", + "arrow-schema 57.3.1", + "chrono", + "csv", + "csv-core", + "regex", +] + [[package]] name = "arrow-csv" version = "58.3.0" @@ -800,6 +836,19 @@ dependencies = [ "arrow-select 58.3.0", ] +[[package]] +name = "arrow-row" +version = "57.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a931b520a2a5e22033e01a6f2486b4cdc26f9106b759abeebc320f125e94d7" +dependencies = [ + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", + "half", +] + [[package]] name = "arrow-row" version = "58.3.0" @@ -1408,6 +1457,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" dependencies = [ "aws-lc-sys", + "untrusted 0.7.1", "zeroize", ] @@ -1913,6 +1963,16 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bcder" +version = "0.7.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b593e5aeaf7992d388c08a9831c921cd703718064b3e50ba8e6d666d6cf86ca7" +dependencies = [ + "bytes", + "smallvec", +] + [[package]] name = "bdd" version = "0.0.1" @@ -2403,7 +2463,7 @@ version = "0.22.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2235eb320cd7178862a32dd111bd0c0f71a368e393add4914c50129add478eab" dependencies = [ - "arrow", + "arrow 58.3.0", "buoyant_kernel_derive", "bytes", "chrono", @@ -2721,7 +2781,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" dependencies = [ "chrono", - "phf", + "phf 0.12.1", ] [[package]] @@ -3992,7 +4052,7 @@ version = "0.32.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4588e95ff3b2ccdba56d9ec262bd3467c0593000f729402528706f62be8be1ca" dependencies = [ - "arrow", + "arrow 58.3.0", "arrow-arith 58.3.0", "arrow-array 58.3.0", "arrow-buffer 58.3.0", @@ -4000,7 +4060,7 @@ dependencies = [ "arrow-ipc 58.3.0", "arrow-json 58.3.0", "arrow-ord 58.3.0", - "arrow-row", + "arrow-row 58.3.0", "arrow-schema 58.3.0", "arrow-select 58.3.0", "async-trait", @@ -4027,7 +4087,7 @@ dependencies = [ "regex", "serde", "serde_json", - "sqlparser", + "sqlparser 0.61.0", "strum 0.27.2", "thiserror 2.0.18", "tokio", @@ -4841,6 +4901,12 @@ dependencies = [ "ext-trait", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + [[package]] name = "fastbloom" version = "0.14.1" @@ -5280,7 +5346,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -7112,6 +7178,27 @@ dependencies = [ "uuid", ] +[[package]] +name = "iggy_connector_redshift_sink" +version = "0.4.1-edge.1" +dependencies = [ + "arrow 57.3.1", + "async-trait", + "humantime", + "iggy_common", + "iggy_connector_sdk", + "parquet 57.3.1", + "rust-s3", + "secrecy", + "serde", + "serde_json", + "simd-json", + "sqlx", + "tokio", + "tracing", + "uuid", +] + [[package]] name = "iggy_connector_s3_sink" version = "0.4.0" @@ -7383,6 +7470,7 @@ checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" name = "integration" version = "0.0.1" dependencies = [ + "arrow 57.3.1", "assert_cmd", "async-trait", "base64", @@ -7411,6 +7499,8 @@ dependencies = [ "lazy_static", "libc", "mongodb", + "parquet 57.3.1", + "pgwire", "predicates", "rand 0.10.1", "rcgen", @@ -7425,6 +7515,7 @@ dependencies = [ "serial_test", "server", "socket2 0.6.4", + "sqlparser 0.62.0", "sqlx", "sysinfo 0.39.5", "tempfile", @@ -7432,6 +7523,7 @@ dependencies = [ "testcontainers", "testcontainers-modules", "tokio", + "tokio-postgres", "toml 1.1.2+spec-1.1.0", "tracing", "tracing-subscriber", @@ -7764,6 +7856,29 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4345964bb142484797b161f473a503a434de77149dd8c7427788c6e13379388" +[[package]] +name = "lazy-regex" +version = "3.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bae91019476d3ec7147de9aa291cadb6d870abf2f3015d2da73a90325ac1496" +dependencies = [ + "lazy-regex-proc_macros", + "once_cell", + "regex-lite", +] + +[[package]] +name = "lazy-regex-proc_macros" +version = "3.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4de9c1e1439d8b7b3061b2d209809f447ca33241733d9a3c01eabf2dc8d94358" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "syn 2.0.118", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -8462,7 +8577,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -8969,6 +9084,15 @@ dependencies = [ "objc2-foundation", ] +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.37.3" @@ -9624,13 +9748,67 @@ dependencies = [ "sha2 0.10.9", ] +[[package]] +name = "pg_interval" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c386dd54fce258fc04e668126ae68589a0d92e03a90ea67881d1300f70fd6170" +dependencies = [ + "bytes", + "chrono", + "postgres-types", +] + +[[package]] +name = "pgwire" +version = "0.40.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7981cfde34009be689a05a30c497ad5fbb552531d3d54230b3627264ff1bc384" +dependencies = [ + "async-trait", + "aws-lc-rs", + "base64", + "bytes", + "chrono", + "derive-new", + "futures", + "hex", + "lazy-regex", + "md5", + "pg_interval", + "postgres-types", + "rand 0.10.1", + "rust_decimal", + "rustls-pki-types", + "ryu", + "serde", + "serde_json", + "smol_str", + "stringprep", + "thiserror 2.0.18", + "tokio", + "tokio-rustls", + "tokio-util", + "x509-certificate", +] + [[package]] name = "phf" version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" dependencies = [ - "phf_shared", + "phf_shared 0.12.1", +] + +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_shared 0.13.1", + "serde", ] [[package]] @@ -9642,6 +9820,15 @@ dependencies = [ "siphasher", ] +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + [[package]] name = "pico-args" version = "0.5.0" @@ -9815,6 +10002,39 @@ dependencies = [ "serde", ] +[[package]] +name = "postgres-protocol" +version = "0.6.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" +dependencies = [ + "base64", + "byteorder", + "bytes", + "fallible-iterator", + "hmac 0.13.0", + "md-5 0.11.0", + "memchr", + "rand 0.10.1", + "sha2 0.11.0", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" +dependencies = [ + "array-init", + "bytes", + "chrono", + "fallible-iterator", + "postgres-protocol", + "serde_core", + "serde_json", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -10183,7 +10403,7 @@ dependencies = [ "libc", "once_cell", "raw-cpuid", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "web-sys", "winapi", ] @@ -10821,7 +11041,7 @@ dependencies = [ "cfg-if", "getrandom 0.2.17", "libc", - "untrusted", + "untrusted 0.9.0", "windows-sys 0.52.0", ] @@ -11064,6 +11284,7 @@ dependencies = [ "borsh", "bytes", "num-traits", + "postgres-types", "rand 0.8.6", "rkyv", "serde", @@ -11235,7 +11456,7 @@ dependencies = [ "aws-lc-rs", "ring", "rustls-pki-types", - "untrusted", + "untrusted 0.9.0", ] [[package]] @@ -12169,6 +12390,16 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8e2fb0f499abb4d162f2bedad68f5ef91a1682b5a03596ddb67efd37768d100" +[[package]] +name = "smol_str" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4aaa7368fcf4852a4c2dd92df0cace6a71f2091ca0a23391ce7f3a31833f1523" +dependencies = [ + "borsh", + "serde_core", +] + [[package]] name = "snafu" version = "0.8.9" @@ -12278,6 +12509,16 @@ dependencies = [ "recursive", ] +[[package]] +name = "sqlparser" +version = "0.62.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c6d1b651dc4edf07eead2a0c6c78016ce971bc2c10da5266861b13f25e7cec" +dependencies = [ + "log", + "recursive", +] + [[package]] name = "sqlx" version = "0.9.0" @@ -13142,6 +13383,32 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "tokio-postgres" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf 0.13.1", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand 0.10.1", + "socket2 0.6.4", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -13842,6 +14109,12 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" +[[package]] +name = "untrusted" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a" + [[package]] name = "untrusted" version = "0.9.0" @@ -14141,6 +14414,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.4+wasi-0.2.12" @@ -14150,6 +14432,15 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.125" @@ -14159,6 +14450,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] @@ -14326,6 +14618,13 @@ name = "whoami" version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "998767ef88740d1f5b0682a9c53c24431453923962269c2db68ee43788c5a40d" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite", + "web-sys", +] [[package]] name = "widestring" @@ -14895,6 +15194,25 @@ dependencies = [ "tap", ] +[[package]] +name = "x509-certificate" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca9eb9a0c822c67129d5b8fcc2806c6bc4f50496b420825069a440669bcfbf7f" +dependencies = [ + "bcder", + "bytes", + "chrono", + "der", + "hex", + "pem", + "ring", + "signature", + "spki", + "thiserror 2.0.18", + "zeroize", +] + [[package]] name = "x509-parser" version = "0.18.1" diff --git a/Cargo.toml b/Cargo.toml index 87a96add3e..dbb43ad925 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ members = [ "core/connectors/sinks/mongodb_sink", "core/connectors/sinks/postgres_sink", "core/connectors/sinks/quickwit_sink", + "core/connectors/sinks/redshift_sink", "core/connectors/sinks/s3_sink", "core/connectors/sinks/stdout_sink", "core/connectors/sinks/surrealdb_sink", @@ -247,6 +248,7 @@ parquet = "57.3.1" partitions = { path = "core/partitions" } passterm = "2.0.6" paste = "1.0" +pgwire = "0.40.4" postcard = { version = "1.1.3", features = ["alloc"] } predicates = "3.1.4" proc-macro2 = "1" @@ -293,6 +295,7 @@ simd-json = { version = "0.17.0", features = ["serde_impl"] } slab = "0.4.12" smallvec = "1.15" socket2 = "0.6.4" +sqlparser = "0.62.0" sqlx = { version = "0.9.0", features = [ "runtime-tokio", "tls-rustls", @@ -318,6 +321,7 @@ testcontainers = { version = "0.27.3", features = ["reusable-containers"] } testcontainers-modules = { version = "0.15.0", features = ["postgres", "http_wait"] } thiserror = "2.0.18" tokio = { version = "1.52.3", features = ["full"] } +tokio-postgres = "0.7.18" tokio-rustls = "0.26.4" tokio-tungstenite = { version = "0.29", features = ["rustls-tls-webpki-roots"] } tokio-util = { version = "0.7.18", features = ["compat"] } diff --git a/core/connectors/README.md b/core/connectors/README.md index 64bfc9a1f8..28932d9bc1 100644 --- a/core/connectors/README.md +++ b/core/connectors/README.md @@ -85,6 +85,7 @@ Each sink should have its own, custom configuration, which is passed along with - **Iceberg Sink** - writes data to Apache Iceberg tables via REST catalog - **PostgreSQL Sink** - stores messages in PostgreSQL database tables - **Quickwit Sink** - indexes messages in Quickwit search engine +- **Reshift Sink** - stores messages in Redshift warehouse tables via S3 as staging - **S3 Sink** - writes messages to Amazon S3 and S3-compatible stores (MinIO, R2, B2, DO Spaces) - **Stdout Sink** - prints messages to standard output (useful for debugging/development) - **SurrealDB Sink** - writes messages into SurrealDB with deterministic record IDs for idempotent replay diff --git a/core/connectors/sinks/README.md b/core/connectors/sinks/README.md index 57ea055490..212d30aae5 100644 --- a/core/connectors/sinks/README.md +++ b/core/connectors/sinks/README.md @@ -14,6 +14,7 @@ Sink connectors are responsible for writing data from Iggy streams to external s | **influxdb_sink** | Writes messages to InfluxDB as line-protocol points; supports both V2 (org/bucket, Flux) and V3 (db, SQL) | | **postgres_sink** | Stores messages in PostgreSQL database tables with configurable schemas | | **quickwit_sink** | Indexes messages in Quickwit search engine for log analytics | +| **redshift_sink** | Stores messages in Redshift warehouse tables with configurable schemas vis S3 as staging | | **s3_sink** | Writes messages to Amazon S3 and S3-compatible stores (MinIO, R2, B2, DO Spaces) | | **stdout_sink** | Prints messages to standard output (useful for debugging and development) | | **surrealdb_sink** | Writes messages into SurrealDB with deterministic record IDs for idempotent replay | diff --git a/core/connectors/sinks/redshift_sink/Cargo.toml b/core/connectors/sinks/redshift_sink/Cargo.toml new file mode 100644 index 0000000000..d420c3becf --- /dev/null +++ b/core/connectors/sinks/redshift_sink/Cargo.toml @@ -0,0 +1,52 @@ +# 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_redshift_sink" +version = "0.4.1-edge.1" +description = "Iggy Redshift sink connector for storing stream messages into Redshift warehouse via S3" +edition = "2024" +license = "Apache-2.0" +keywords = ["iggy", "messaging", "streaming", "redshift", "sink"] +categories = ["command-line-utilities", "warehouse", "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 } +humantime = { workspace = true } +iggy_common = { workspace = true } +iggy_connector_sdk = { workspace = true } +parquet = { workspace = true } +rust-s3 = { workspace = true } +secrecy = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +sqlx = { workspace = true, features = ["runtime-tokio", "tls-rustls", "postgres", "chrono"] } +tokio = { workspace = true } +tracing = { workspace = true } +uuid = { workspace = true, features = ["v7"] } + +[dev-dependencies] +simd-json = { workspace = true } diff --git a/core/connectors/sinks/redshift_sink/README.md b/core/connectors/sinks/redshift_sink/README.md new file mode 100644 index 0000000000..01a9742706 --- /dev/null +++ b/core/connectors/sinks/redshift_sink/README.md @@ -0,0 +1,119 @@ +# Redshift Sink Connector + +Writes Apache Iggy stream messages into Amazon Redshift via S3-staged Parquet +files and a `COPY` load. + +Each connector batch is serialized to a Parquet file and uploaded to the +configured S3 bucket/prefix, then loaded into the target Redshift table with a +`COPY` statement. This makes S3 a staging area rather than a destination in +its own right — Redshift is the system of record for the data. + +Persistent load failures are at-most-once from the runtime's perspective: +messages may already be committed in Iggy before this connector exhausts its +write attempts, so failed loads are logged but not redelivered. + +## Configuration + +```toml +type = "sink" +key = "redshift" +enabled = true +version = 0 +name = "Redshift sink" +path = "../../target/release/libiggy_connector_redshift_sink" +verbose = false + +[[streams]] +stream = "user_events" +topics = ["users", "orders"] +schema = "json" +batch_length = 100 +poll_interval = "5ms" +consumer_group = "redshift_sink" + +[plugin_config] +connection_string = "postgresql://user:pass@localhost:5439/database" +target_table = "iggy_messages" +batch_size = 100 +max_connections = 10 +include_metadata = true +include_checksum = true +include_origin_timestamp = true +payload_format = "varbyte" +aws_access_key_id = "admin" +aws_secret_access_key = "password" +s3_bucket = "iggystaging" +s3_prefix = "iggy/messages" +s3_endpoint = "http://localhost:9000" +aws_region = "us-east-1" +archive = true +``` + +### Plugin Fields + +| Field | Required | Default | Description | +| --- | --- | --- | --- | +| `connection_string` | yes | — | Postgres-wire connection string used to reach the Redshift cluster and issue the `COPY` command. | +| `target_table` | yes | — | Destination Redshift table that batches are copied into. | +| `batch_size` | no | `100` | Number of messages buffered per Parquet file / `COPY` operation. | +| `max_connections` | no | `5` | Size of the connection pool used against Redshift. | +| `include_metadata` | no | `false` | Stores stream/topic/partition/offset/timestamp/schema fields alongside the payload. | +| `include_checksum` | no | `false` | Stores the Iggy message checksum. | +| `include_origin_timestamp` | no | `false` | Stores the original Iggy origin timestamp. | +| `payload_format` | no | `varbyte` | Encoding used for the payload column in the Parquet file. See **Payload Format** below. | +| `verbose_logging` | no | `false` | Enables verbose logging for debugging purposes. | +| `max_retries` | no | `3` | Maximum number of retries for failed `COPY` operations. `0` disables retries (only one attempt will be made) | +| `retry_delay` | no | `1s` | Delay in seconds between retry attempts. | +| `aws_access_key_id` | yes | — | AWS access key used for S3 staging. | +| `aws_secret_access_key` | yes | — | AWS secret key used for S3 staging. | +| `s3_bucket` | yes | — | S3 bucket that Parquet batch files are staged into before the Redshift `COPY`. | +| `s3_prefix` | yes | — | Key prefix under which staged Parquet files are written, e.g. `iggy/messages`. | +| `s3_endpoint` | no | — | Override endpoint for S3-compatible stores (e.g. MinIO). Omit for AWS S3 itself. | +| `aws_region` | yes | — | AWS region for the S3 bucket. | +| `archive` | no | `false` | See **Archiving Staged Files** below. | + +## Staging via S3 + +Redshift's `COPY` command loads from files, not from a live stream, so every +batch is first written out as a Parquet file and uploaded to +`s3:////...` before the `COPY` into `target_table` runs. +S3 is purely a staging area in this flow — it is not queried directly by +consumers of the data, and its cost is the price of getting bulk data into +Redshift efficiently rather than row-by-row. + +## Archiving Staged Files + +The `archive` field controls what happens to a batch's Parquet file **after** +it has been successfully loaded into Redshift: + +- `archive = true` — the file is kept, moved under an `archive` prefix + (i.e. `s3:///archive/...`) instead of being deleted. + Useful for replay, auditing, or downstream batch jobs that read Parquet + directly. +- `archive = false` — the file is deleted from S3 once the `COPY` succeeds, + since Redshift itself is now the source of truth for that data and the + staged copy has no further purpose. + +## Payload Format + +`payload_format` controls how the payload column is written in the staged +Parquet file, which in turn determines its type once loaded into Redshift: + +- Parquet has no dedicated JSON logical type, so a `payload_format = "json"` + payload is written as a Parquet `VARCHAR` (string), not a structured type. +- As a result, the column lands in Redshift as `VARCHAR`, not `SUPER`. +- To query the payload as structured data downstream, use Redshift's + `JSON_PARSE()` (or equivalent JSON functions) on the `VARCHAR` column at + query time rather than expecting a native `SUPER` column out of the box. + +## Stored Shape + +With metadata enabled, records contain: + +- `id`: original Iggy message id as numeric +- `iggy_stream`, `iggy_topic`, `iggy_partition_id`, `iggy_offset` +- `iggy_timestamp`, `iggy_origin_timestamp`, `iggy_checksum`, +- `payload`: encoded per `payload_format` (see above) + +The `messages_processed` counter reports valid records submitted to Redshift +via `COPY`. diff --git a/core/connectors/sinks/redshift_sink/config.toml b/core/connectors/sinks/redshift_sink/config.toml new file mode 100644 index 0000000000..4e695a0dde --- /dev/null +++ b/core/connectors/sinks/redshift_sink/config.toml @@ -0,0 +1,49 @@ +# 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 = "redshift" +enabled = true +version = 0 +name = "Redshift sink" +path = "../../target/release/libiggy_connector_redshift_sink" +verbose = false + +[[streams]] +stream = "user_events" +topics = ["users", "orders"] +schema = "json" +batch_length = 100 +poll_interval = "5ms" +consumer_group = "redshift_sink" + +[plugin_config] +connection_string = "postgresql://user:pass@localhost:5439/database" +target_table = "iggy_messages" +batch_size = 100 +max_connections = 10 +include_metadata = true +include_checksum = true +include_origin_timestamp = true +payload_format = "varbyte" +aws_access_key_id = "admin" +aws_secret_access_key = "password" +s3_bucket = "iggystaging" +s3_prefix = "iggy/messages" +s3_endpoint = "http://localhost:9000" +aws_region = "us-east-1" +archive = true diff --git a/core/connectors/sinks/redshift_sink/src/config.rs b/core/connectors/sinks/redshift_sink/src/config.rs new file mode 100644 index 0000000000..b0a4b30ca9 --- /dev/null +++ b/core/connectors/sinks/redshift_sink/src/config.rs @@ -0,0 +1,127 @@ +// 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 arrow::datatypes::DataType; +use iggy_connector_sdk::Error; +use secrecy::{ExposeSecret, SecretString}; + +/// Configuration for the Redshift Sink +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct RedshiftSinkConfig { + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub connection_string: SecretString, + pub target_table: String, + pub batch_size: Option, + pub max_connections: Option, + pub include_metadata: Option, + pub include_checksum: Option, + pub include_origin_timestamp: Option, + pub payload_format: Option, + pub verbose_logging: Option, + pub max_retries: Option, + pub retry_delay: Option, + /// aws_access_key_id and aws_secret_access_key MUST be provided + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub aws_access_key_id: SecretString, + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub aws_secret_access_key: SecretString, + pub s3_bucket: String, + pub s3_prefix: String, + pub s3_endpoint: Option, + pub aws_region: String, + /// Offers the option to archive staged S3 files after COPY + /// Defaults to deletion once COPY completes + /// Files are moved to different prefix within the same bucket + pub archive: Option, +} + +impl RedshiftSinkConfig { + pub fn validate(&self) -> Result<(), Error> { + let mut errors = String::new(); + + if self.connection_string.expose_secret().is_empty() { + errors.push_str("connection_string is empty\n"); + } + + if self.target_table.is_empty() { + errors.push_str(", target_table is empty\n"); + } + + if self.s3_bucket.is_empty() { + errors.push_str(", s3_bucket is empty\n"); + } + + if self.aws_region.is_empty() { + errors.push_str(", aws_region is empty\n"); + } + + // Validate AWS credentials: access keys must be provided + let has_access_key = !self.aws_access_key_id.expose_secret().is_empty(); + + let has_secret_key = !self.aws_secret_access_key.expose_secret().is_empty(); + + if !(has_access_key && has_secret_key) { + errors.push_str(", aws_access_key_id and aws_secret_access_key are empty\n"); + } + + if !errors.is_empty() { + Err(Error::InvalidConfigValue(errors)) + } else { + Ok(()) + } + } +} + +/// This connector supports: +/// 1. Byte -> which has VARBYTE as the Redshift equivalent +/// 2. Text -> which has VARCHAR as the Redshift equivalent +/// +/// We dont have Json because we are using parquet as a means to sink ingestion +/// As at the development of this connector there's no direct parquet type that matches JSON +/// For JSON needs Reshshift has SUPER(VARCHAR can be parsed by JSON_PARSE) +#[allow(unused)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum PayloadFormat { + Json, + Text, + #[default] + Varbyte, +} + +impl PayloadFormat { + pub fn from_config(s: Option<&str>) -> Self { + match s.map(|s| s.to_lowercase()).as_deref() { + Some("text") | Some("json") => PayloadFormat::Text, + _ => PayloadFormat::Varbyte, + } + } + + pub fn sql_type(&self) -> &'static str { + match self { + PayloadFormat::Varbyte => "VARBYTE", + PayloadFormat::Text | PayloadFormat::Json => "VARCHAR", + } + } + + pub fn arrow_type(&self) -> DataType { + match self { + PayloadFormat::Varbyte => DataType::Binary, + PayloadFormat::Text => DataType::Utf8, + PayloadFormat::Json => DataType::Utf8, + } + } +} diff --git a/core/connectors/sinks/redshift_sink/src/lib.rs b/core/connectors/sinks/redshift_sink/src/lib.rs new file mode 100644 index 0000000000..211e354c20 --- /dev/null +++ b/core/connectors/sinks/redshift_sink/src/lib.rs @@ -0,0 +1,1182 @@ +// 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. + +mod config; + +use std::{str::FromStr, sync::Arc, time::Duration}; + +use arrow::{ + array::{ + ArrayRef, BinaryArray, Decimal256Array, Int32Array, Int64Array, RecordBatch, StringArray, + TimestampMicrosecondArray, + }, + datatypes::{DataType, Field, Schema, TimeUnit}, +}; +use async_trait::async_trait; +use humantime::Duration as HumanDuration; +use iggy_connector_sdk::{ + ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata, sink_connector, +}; +use parquet::arrow::ArrowWriter; +use s3::{Bucket, Region, creds::Credentials}; +use secrecy::ExposeSecret; +use sqlx::{AssertSqlSafe, Pool, Postgres, postgres::PgPoolOptions}; +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::config::{PayloadFormat, RedshiftSinkConfig}; + +sink_connector!(RedshiftSink); + +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_RETRY_DELAY: &str = "1s"; +const DEFAULT_MAX_CONNECTIONS: u32 = 5; +const DEFAULT_ARCHIVE_PREFIX: &str = "archive/messages"; + +#[derive(Debug)] +pub struct RedshiftSink { + pub id: u32, + config: RedshiftSinkConfig, + pool: Option>, + state: Mutex, + verbose: bool, + bucket: Option>, +} + +#[async_trait] +impl Sink for RedshiftSink { + async fn open(&mut self) -> Result<(), Error> { + tracing::info!( + sink_id = self.id, + table = %self.config.target_table, "opening Redshift sink connector" + ); + + self.connect().await?; + self.ensure_table_exists().await?; + Ok(()) + } + + async fn consume( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: MessagesMetadata, + messages: Vec, + ) -> Result<(), Error> { + tracing::debug!( + sink_id = self.id, + count = messages.len(), + "consuming messages" + ); + self.process_messages(topic_metadata, &messages_metadata, &messages) + .await + } + + async fn close(&mut self) -> Result<(), Error> { + tracing::info!(sink_id = self.id, "closing Redshift sink connector"); + + if let Some(pool) = self.pool.take() { + pool.close().await; + + tracing::debug!(sink_id = self.id, "database pool closed"); + } + + let state = self.state.lock().await; + + tracing::info!( + sink_id = self.id, + messages_processed = state.messages_processed, + batches_loaded = state.batches_loaded, + insertion_errors = state.insertion_errors, + "Redshift sink connector closed", + ); + + Ok(()) + } +} + +impl RedshiftSink { + pub fn new(id: u32, config: RedshiftSinkConfig) -> Self { + let verbose = config.verbose_logging.unwrap_or(false); + + Self { + id, + config, + pool: None, + state: Mutex::new(State::default()), + verbose, + bucket: None, + } + } + + async fn connect(&mut self) -> Result<(), Error> { + let max_connections = self + .config + .max_connections + .unwrap_or(DEFAULT_MAX_CONNECTIONS); + + let redacted = redact_connection_string(self.config.connection_string.expose_secret()); + + tracing::info!(max_connections, dsn = %redacted, "connecting to Redshift"); + + let pool = PgPoolOptions::new() + .max_connections(max_connections) + .connect(self.config.connection_string.expose_secret()) + .await + .map_err(|e| Error::InitError(format!("Failed to connect to Redshift: {e}")))?; + + sqlx::query("SELECT 1").execute(&pool).await.map_err(|e| { + tracing::error!("Tracing failed: {:#?}", e); + Error::InitError(format!("Warehouse connectivity test failed: {e}")) + })?; + + self.pool = Some(pool); + tracing::debug!("Redshift connection pool established"); + + let region = self.build_region()?; + + let credentials = Credentials::new( + Some(self.config.aws_access_key_id.expose_secret()), + Some(self.config.aws_secret_access_key.expose_secret()), + None, + None, + None, + ) + .map_err(|e| { + tracing::error!("Failed to create S3 credentials: {e}"); + Error::InvalidConfig + })?; + + let mut bucket = Bucket::new(&self.config.s3_bucket, region, credentials).map_err(|e| { + tracing::error!("Failed to create S3 bucket client: {e}"); + Error::InvalidConfig + })?; + + if self.config.s3_endpoint.is_some() { + bucket = bucket.with_path_style(); + } + + self.bucket = Some(bucket); + + tracing::info!("Redshift sink connector ready"); + + Ok(()) + } + + fn build_region(&self) -> Result { + if let Some(endpoint) = &self.config.s3_endpoint { + tracing::debug!(endpoint = %endpoint, "using custom S3 endpoint"); + Ok(Region::Custom { + region: self.config.aws_region.clone(), + endpoint: endpoint.clone(), + }) + } else { + Region::from_str(&self.config.aws_region).map_err(|_| Error::InvalidConfig) + } + } + + async fn ensure_table_exists(&self) -> Result<(), Error> { + let pool = self.get_pool()?; + + let table_name = &self.config.target_table; + let payload_type = self.payload_format().sql_type(); + + let (query, _) = self.build_create_table_sql()?; + + tracing::debug!("ensuring target table exists"); + + sqlx::query(AssertSqlSafe(query)) + .execute(pool) + .await + .map_err(|e| { + tracing::error!(error = %e); + Error::InitError(format!("Failed to create table '{table_name}': {e}")) + })?; + + tracing::info!(table = %table_name, payload_type, "target table ready"); + + Ok(()) + } + + async fn process_messages( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], + ) -> Result<(), Error> { + let batch_size = self.config.batch_size.unwrap_or(100) as usize; + + for batch in messages.chunks(batch_size) { + match self + .insert_batch(batch, topic_metadata, messages_metadata) + .await + { + Ok(()) => { + self.state.lock().await.batches_loaded += 1; + } + Err(e) => { + self.state.lock().await.insertion_errors += batch.len() as u64; + tracing::error!(error = %e, batch_size = batch.len(), "failed to insert batch"); + } + } + } + + let mut state = self.state.lock().await; + state.messages_processed += messages.len() as u64; + + if self.verbose { + tracing::info!( + sink_id = self.id, + total_processed = state.messages_processed, + batch_received = messages.len(), + table = %self.config.target_table, + batches_loaded = state.batches_loaded, + "processed message batch" + ); + } else { + tracing::debug!( + sink_id = self.id, + total_processed = state.messages_processed, + table = %self.config.target_table, + "processed message batch" + ); + } + + Ok(()) + } + + async fn insert_batch( + &self, + messages: &[ConsumedMessage], + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + ) -> Result<(), Error> { + if messages.is_empty() { + return Ok(()); + } + + let include_metadata = self.config.include_metadata.unwrap_or(true); + let include_checksum = self.config.include_checksum.unwrap_or(true); + let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); + let payload_format = self.payload_format(); + + let record_batch = create_record_batch( + topic_metadata, + messages_metadata, + messages, + include_metadata, + include_checksum, + include_origin_timestamp, + payload_format, + )?; + + let content = encode_parquet(&record_batch)?; + + tracing::debug!( + bytes = content.len(), + rows = record_batch.num_rows(), + "encoded parquet batch" + ); + + let s3_path = self.upload_parquet(&content).await?; + self.copy_parquet(&s3_path).await?; + self.archive_parquet(&s3_path).await?; + + tracing::info!(count = messages.len(), path = %s3_path, "batch inserted into Redshift"); + + Ok(()) + } + + async fn copy_parquet(&self, s3_path: &str) -> Result<(), Error> { + let max_retries = self.get_max_retries(); + let retry_delay = self.get_retry_delay(); + let sql = self.build_copy_sql(s3_path); + let pool = self.get_pool()?; + + tracing::debug!(table = %self.config.target_table, s3_path, "issuing Redshift COPY"); + + retry_with_backoff( + "redshift COPY", + max_retries, + retry_delay, + is_transient_error, + || async { + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .map(|_| ()) + }, + ) + .await?; + + tracing::debug!(table = %self.config.target_table, "Redshift COPY completed"); + + Ok(()) + } + + fn build_create_table_sql(&self) -> Result<(String, u32), Error> { + let table_name = &self.config.target_table; + let quoted_table = quote_identifier(table_name)?; + + let include_metadata = self.config.include_metadata.unwrap_or(true); + let include_checksum = self.config.include_checksum.unwrap_or(true); + let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); + let payload_type = self.payload_format().sql_type(); + + let mut params_per_row: u32 = 1; // id + + let mut query = + format!("CREATE TABLE IF NOT EXISTS {quoted_table} (id DECIMAL(39, 0) PRIMARY KEY"); + + if include_metadata { + query.push_str(", iggy_offset BIGINT, iggy_timestamp TIMESTAMPTZ, iggy_stream TEXT, iggy_topic TEXT, iggy_partition_id INTEGER"); + params_per_row += 5; + } + + if include_checksum { + query.push_str(", iggy_checksum VARCHAR"); + params_per_row += 1; + } + + if include_origin_timestamp { + query.push_str(", iggy_origin_timestamp TIMESTAMPTZ"); + params_per_row += 1; + } + + query.push_str(&format!(", payload {payload_type}")); + query.push_str(", created_at TIMESTAMPTZ DEFAULT GETDATE());"); + params_per_row += 2; + + Ok((query, params_per_row)) + } + + fn build_copy_sql(&self, s3_path: &str) -> String { + // Built via format! (not sqlx binds) because the Redshift/Pgwire endpoint here + // uses a Simple Query Handler that doesn't support prepared statements with binds. + let credentials = format!( + "CREDENTIALS 'ACCESS_KEY_ID={};SECRET_ACCESS_KEY={}'", + self.config.aws_access_key_id.expose_secret(), + self.config.aws_secret_access_key.expose_secret() + ); + + format!( + "COPY {} FROM '{}' {} FORMAT AS PARQUET REGION '{}';", + self.config.target_table, s3_path, credentials, self.config.aws_region + ) + } + + async fn upload_parquet(&self, content: &[u8]) -> Result { + let file_id = Uuid::now_v7(); + let key = build_s3_key(&self.config.s3_prefix, &format!("{file_id}.parquet")); + let bucket = self.get_bucket()?; + + tracing::debug!(key = %key, bytes = content.len(), "uploading parquet to S3"); + + let response = bucket.put_object(&key, content).await.map_err(|e| { + tracing::error!("Failed to upload to S3 key '{key}': {e}"); + Error::Storage(format!("S3 upload failed: {e}")) + })?; + + ensure_s3_status(response.status_code(), 200, "S3 upload")?; + + let path = format!("s3://{}{}", bucket.name(), key); + tracing::info!(path = %path, bytes = content.len(), "uploaded parquet to S3"); + + Ok(path) + } + + async fn archive_parquet(&self, key: &str) -> Result<(), Error> { + let old_key = key + .strip_prefix(&format!("s3://{}/", self.config.s3_bucket)) + .unwrap_or(key); + + if !self.get_archive() { + self.delete_object(old_key).await?; + tracing::info!(key = old_key, "deleted parquet file (archiving disabled)"); + return Ok(()); + } + + let bucket = self.get_bucket()?; + let prefix = self.config.s3_prefix.trim_matches('/'); + let archived_key = old_key.replacen(prefix, DEFAULT_ARCHIVE_PREFIX.trim_matches('/'), 1); + + tracing::debug!(from = old_key, to = %archived_key, "archiving parquet file"); + + let status_code = bucket + .copy_object_internal(old_key, &archived_key) + .await + .map_err(|e| { + tracing::error!(key = old_key, error = %e, "failed to copy object for archiving"); + Error::Storage(format!("S3 archiving failed: {e}")) + })?; + + ensure_s3_status(status_code, 200, "S3 archive copy")?; + + self.delete_object(old_key).await?; + tracing::info!(archived_to = %archived_key, "archived parquet file"); + + Ok(()) + } + + async fn delete_object(&self, key: &str) -> Result<(), Error> { + let bucket = self.get_bucket()?; + + let response = bucket.delete_object(key).await.map_err(|e| { + tracing::error!(key, error = %e, "failed to delete S3 object"); + Error::Storage(format!("S3 deleting failed: {e}")) + })?; + ensure_s3_status(response.status_code(), 204, "S3 object deletion")?; + + tracing::debug!(key, "deleted S3 object"); + Ok(()) + } + + fn get_pool(&self) -> Result<&Pool, Error> { + self.pool + .as_ref() + .ok_or_else(|| Error::InitError("Database not connected".to_string())) + } + + fn get_bucket(&self) -> Result<&Bucket, Error> { + let r = self + .bucket + .as_ref() + .ok_or_else(|| Error::InitError("Database not connected".to_string()))?; + + Ok(r) + } + + fn payload_format(&self) -> PayloadFormat { + PayloadFormat::from_config(self.config.payload_format.as_deref()) + } + + fn get_max_retries(&self) -> u32 { + self.config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES) + } + + fn get_retry_delay(&self) -> Duration { + self.config + .retry_delay + .as_deref() + .unwrap_or(DEFAULT_RETRY_DELAY) + .parse::() + .map(Into::into) + .unwrap_or_else(|_| Duration::from_secs(1)) + } + + fn get_archive(&self) -> bool { + self.config.archive.unwrap_or(false) + } +} + +#[derive(Debug, Default)] +struct State { + messages_processed: u64, + batches_loaded: u64, + insertion_errors: u64, +} + +/// Generic retry helper with linear backoff, used for transient warehouse errors. +async fn retry_with_backoff( + operation: &str, + max_retries: u32, + base_delay: Duration, + is_transient: impl Fn(&sqlx::Error) -> bool, + mut op: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempts = 0u32; + + loop { + match op().await { + Ok(value) => return Ok(value), + Err(e) => { + attempts += 1; + let transient = is_transient(&e); + + if !transient || attempts >= max_retries { + tracing::error!(operation, attempts, error = %e, "operation failed permanently"); + return Err(Error::CannotStoreData(format!( + "{operation} failed after {attempts} attempts: {e}" + ))); + } + + tracing::warn!(operation, attempts, max_retries, error = %e, "transient error, retrying"); + tokio::time::sleep(base_delay * attempts).await; + } + } + } +} + +fn encode_parquet(batch: &RecordBatch) -> Result, Error> { + let mut content = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut content, batch.schema(), None).map_err(|e| { + tracing::error!(error = %e, "failed to create parquet writer"); + Error::WriteFailure(format!("Failed to create parquet writer: {e}")) + })?; + + writer.write(batch).map_err(|e| { + tracing::error!(error = %e, "failed to write parquet batch"); + Error::WriteFailure(format!("Failed to write parquet: {e}")) + })?; + + writer.close().map_err(|e| { + tracing::error!(error = %e, "failed to close parquet writer"); + Error::WriteFailure(format!("Failed to close writer: {e}")) + })?; + + Ok(content) +} + +fn build_s3_key(prefix: &str, filename: &str) -> String { + if prefix.is_empty() { + format!("/{filename}") + } else { + format!("/{}/{filename}", prefix.trim_end_matches('/')) + } +} + +fn ensure_s3_status( + status: T, + expected: T, + context: &str, +) -> Result<(), Error> { + if status != expected { + tracing::error!(context, %status, "unexpected S3 response status"); + return Err(Error::Storage(format!( + "{context} failed with status {status}" + ))); + } + Ok(()) +} + +fn create_record_batch( + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], + include_metadata: bool, + include_checksum: bool, + include_origin_timestamp: bool, + payload_format: PayloadFormat, +) -> Result { + let mut fields = vec![Field::new("id", DataType::Decimal256(39, 0), false)]; + let mut columns: Vec = vec![id_column(messages)?]; + + if include_metadata { + let (mut metadata_fields, mut metadata_columns) = + metadata_columns(topic_metadata, messages_metadata, messages); + fields.append(&mut metadata_fields); + columns.append(&mut metadata_columns); + } + + if include_checksum { + fields.push(Field::new("iggy_checksum", DataType::Utf8, false)); + columns.push(checksum_column(messages)); + } + + if include_origin_timestamp { + fields.push(Field::new( + "iggy_origin_timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + )); + columns.push(origin_timestamp_column(messages)); + } + + fields.push(Field::new("payload", payload_format.arrow_type(), false)); + columns.push(payload_column(messages, payload_format)?); + + let schema = Arc::new(Schema::new(fields)); + let batch = + RecordBatch::try_new(schema, columns).map_err(|e| Error::CannotStoreData(e.to_string()))?; + + tracing::debug!( + rows = batch.num_rows(), + columns = batch.num_columns(), + "built record batch" + ); + + Ok(batch) +} + +fn id_column(messages: &[ConsumedMessage]) -> Result { + let ids = Decimal256Array::from_iter_values( + messages + .iter() + .map(|v| arrow::datatypes::i256::from_parts(v.id, 0)), + ) + .with_precision_and_scale(39, 0) + .map_err(|e| Error::CannotStoreData(e.to_string()))?; + + Ok(Arc::new(ids)) +} + +fn metadata_columns( + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], +) -> (Vec, Vec) { + let fields = vec![ + Field::new("iggy_offset", DataType::Int64, false), + Field::new( + "iggy_timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new("iggy_stream", DataType::Utf8, false), + Field::new("iggy_topic", DataType::Utf8, false), + Field::new("iggy_partition_id", DataType::Int32, false), + ]; + + let columns: Vec = vec![ + Arc::new(Int64Array::from_iter_values( + messages.iter().map(|v| v.offset as i64), + )), + Arc::new(TimestampMicrosecondArray::from_iter_values( + messages.iter().map(|v| v.timestamp as i64), + )), + Arc::new(StringArray::from_iter_values( + (0..messages.len()).map(|_| topic_metadata.stream.clone()), + )), + Arc::new(StringArray::from_iter_values( + (0..messages.len()).map(|_| topic_metadata.topic.clone()), + )), + Arc::new(Int32Array::from_iter_values( + (0..messages.len()).map(|_| messages_metadata.partition_id as i32), + )), + ]; + + (fields, columns) +} + +fn checksum_column(messages: &[ConsumedMessage]) -> ArrayRef { + Arc::new(StringArray::from_iter_values( + messages.iter().map(|v| v.checksum.to_string()), + )) +} + +fn origin_timestamp_column(messages: &[ConsumedMessage]) -> ArrayRef { + Arc::new(TimestampMicrosecondArray::from_iter_values( + messages.iter().map(|v| v.origin_timestamp as i64), + )) +} + +fn payload_column(messages: &[ConsumedMessage], format: PayloadFormat) -> Result { + match format { + PayloadFormat::Varbyte => { + let values: Vec> = messages + .iter() + .map(|v| v.payload.clone().try_to_bytes()) + .collect::>()?; + let slices: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect(); + Ok(Arc::new(BinaryArray::from_vec(slices))) + } + PayloadFormat::Text => { + let values: Vec = messages + .iter() + .map(|v| { + let bytes = v.payload.try_to_bytes()?; + String::from_utf8(bytes).map_err(|_| Error::InvalidTextPayload) + }) + .collect::>()?; + Ok(Arc::new(StringArray::from_iter_values(values.iter()))) + } + PayloadFormat::Json => { + let values: Vec = messages + .iter() + .map(|v| { + let bytes = v.payload.try_to_bytes()?; + + Ok(serde_json::from_slice::(&bytes) + .map_err(|_| Error::InvalidJsonPayload)? + .to_string()) + }) + .collect::>()?; + Ok(Arc::new(StringArray::from_iter_values(values.iter()))) + } + } +} + +fn redact_connection_string(conn_str: &str) -> String { + // Guard against very short strings + const PREVIEW_LEN: usize = 3; + + if let Some(scheme_end) = conn_str.find("://") { + let scheme = &conn_str[..scheme_end + 3]; + let rest = &conn_str[scheme_end + 3..]; + + // Stop preview at the first sensitive boundary + let safe_end = rest + .find([':', '@', '?', '/']) + .unwrap_or(rest.len()) + .min(PREVIEW_LEN); + + let preview = &rest[..safe_end]; + return format!("{scheme}{preview}***"); + } + + let preview: String = conn_str.chars().take(3).collect(); + format!("{preview}***") +} + +fn is_transient_error(e: &sqlx::Error) -> bool { + match e { + sqlx::Error::Io(_) => true, + sqlx::Error::PoolTimedOut => true, + sqlx::Error::PoolClosed => false, + sqlx::Error::Protocol(_) => false, + sqlx::Error::Database(db_err) => db_err.code().is_some_and(|code| { + matches!( + code.as_ref(), + "40001" | "40P01" | "57P01" | "57P02" | "57P03" | "08000" | "08003" | "08006" + ) + }), + _ => false, + } +} + +fn quote_identifier(name: &str) -> Result { + if name.is_empty() { + return Err(Error::InitError("Table name cannot be empty".to_string())); + } + if name.contains('\0') { + return Err(Error::InitError( + "Table name cannot contain null characters".to_string(), + )); + } + let escaped = name.replace('"', "\"\""); + Ok(format!("\"{escaped}\"")) +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use iggy_connector_sdk::{Payload, Schema}; + use secrecy::SecretString; + + use super::*; + + fn test_config( + include_checksum: bool, + include_origin_timestamp: bool, + include_metadata: bool, + ) -> RedshiftSinkConfig { + RedshiftSinkConfig { + connection_string: SecretString::from("postgresql://localhost/db"), + target_table: "messages".to_string(), + batch_size: Some(100), + max_connections: None, + include_metadata: Some(include_metadata), + include_checksum: Some(include_checksum), + include_origin_timestamp: Some(include_origin_timestamp), + payload_format: None, + verbose_logging: None, + max_retries: None, + retry_delay: None, + aws_access_key_id: SecretString::from("admin"), + aws_secret_access_key: SecretString::from("password"), + s3_bucket: "iggymessages".into(), + s3_prefix: "iggy/messages".into(), + s3_endpoint: None, + aws_region: "us-east-1".into(), + archive: None, + } + } + + fn test_topic_metadata() -> TopicMetadata { + TopicMetadata { + stream: "test_stream".to_string(), + topic: "test_topic".to_string(), + } + } + + fn test_messages_metadata() -> MessagesMetadata { + MessagesMetadata { + partition_id: 7, + current_offset: 0, + schema: Schema::Json, + } + } + + fn test_message(payload: Payload) -> ConsumedMessage { + ConsumedMessage { + id: 42, + offset: 9, + checksum: 123, + timestamp: 1_767_225_600_000_000, + origin_timestamp: 1_700_000_000_000_001, + headers: None, + payload, + } + } + + fn json_payload(value: serde_json::Value) -> Payload { + let mut bytes = serde_json::to_vec(&value).expect("Failed to serialize JSON"); + Payload::Json(simd_json::to_owned_value(&mut bytes).expect("Failed to parse JSON")) + } + + #[test] + fn given_empty_connection_string_should_error() { + let mut config = test_config(false, false, false); + config.connection_string = SecretString::default(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_target_table_should_error() { + let mut config = test_config(false, false, false); + config.target_table = String::new(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_s3_bucket_should_error() { + let mut config = test_config(false, false, false); + config.s3_bucket = String::new(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_aws_region_should_error() { + let mut config = test_config(false, false, false); + config.aws_region = String::new(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_aws_access_key_id_should_error() { + let mut config = test_config(false, false, false); + config.aws_access_key_id = SecretString::default(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_aws_secret_access_key_should_error() { + let mut config = test_config(false, false, false); + config.aws_secret_access_key = SecretString::default(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_json_format_should_return_text() { + assert_eq!( + PayloadFormat::from_config(Some("json")), + PayloadFormat::Text + ); + assert_eq!( + PayloadFormat::from_config(Some("JSON")), + PayloadFormat::Text + ); + } + + #[test] + fn given_text_format_should_return_text() { + assert_eq!( + PayloadFormat::from_config(Some("text")), + PayloadFormat::Text + ); + assert_eq!( + PayloadFormat::from_config(Some("TEXT")), + PayloadFormat::Text + ); + } + + #[test] + fn given_bytea_or_unknown_format_should_return_bytea() { + assert_eq!( + PayloadFormat::from_config(Some("bytea")), + PayloadFormat::Varbyte + ); + assert_eq!( + PayloadFormat::from_config(Some("unknown")), + PayloadFormat::Varbyte + ); + assert_eq!(PayloadFormat::from_config(None), PayloadFormat::Varbyte); + } + + #[test] + fn given_payload_format_should_return_correct_sql_type() { + assert_eq!(PayloadFormat::Varbyte.sql_type(), "VARBYTE"); + assert_eq!(PayloadFormat::Json.sql_type(), "VARCHAR"); + assert_eq!(PayloadFormat::Text.sql_type(), "VARCHAR"); + } + + #[test] + fn given_payload_format_should_return_correct_arrow_type() { + assert_eq!( + PayloadFormat::Varbyte.arrow_type(), + arrow::datatypes::DataType::Binary + ); + assert_eq!( + PayloadFormat::Json.arrow_type(), + arrow::datatypes::DataType::Utf8 + ); + assert_eq!( + PayloadFormat::Text.arrow_type(), + arrow::datatypes::DataType::Utf8 + ); + } + + #[test] + fn given_all_options_enabled_should_build_full_create_query() { + let sink = RedshiftSink::new(1, test_config(true, true, true)); + let (query, param_count) = sink + .build_create_table_sql() + .expect("Failed to build create query"); + + assert!(query.contains("CREATE TABLE IF NOT EXISTS \"messages\"")); + assert!(query.contains("iggy_offset")); + assert!(query.contains("iggy_timestamp")); + assert!(query.contains("iggy_stream")); + assert!(query.contains("iggy_topic")); + assert!(query.contains("iggy_partition_id")); + assert!(query.contains("iggy_checksum")); + assert!(query.contains("iggy_origin_timestamp")); + assert!(query.contains("payload")); + assert!(query.contains("created_at")); + assert_eq!(param_count, 10); + } + + #[test] + fn given_all_options_enabled_should_build_full_parquet() { + let payload = json_payload(serde_json::json!({"name": "Bebeto", "active": true})); + let message = test_message(payload); + let record_batch = create_record_batch( + &test_topic_metadata(), + &test_messages_metadata(), + &[message], + true, + true, + true, + PayloadFormat::Varbyte, + ) + .expect("Failed to create record batch"); + + assert_eq!(record_batch.num_rows(), 1); + assert_eq!(record_batch.num_columns(), 9); + + let columns = record_batch.schema(); + let columns: HashSet<&str> = columns.fields().iter().map(|f| f.name().as_ref()).collect(); + + let expected_columns: HashSet<&str> = HashSet::from([ + "id", + "iggy_offset", + "iggy_timestamp", + "iggy_stream", + "iggy_topic", + "iggy_partition_id", + "iggy_checksum", + "iggy_origin_timestamp", + "payload", + "created_at", + ]); + + assert_eq!(columns.difference(&expected_columns).count(), 0); + } + + #[test] + fn given_metadata_disabled_should_build_minimal_create_query() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + let (query, param_count) = sink + .build_create_table_sql() + .expect("Failed to build create query"); + + assert!(query.contains("CREATE TABLE IF NOT EXISTS \"messages\"")); + assert!(!query.contains("iggy_offset")); + assert!(!query.contains("iggy_timestamp")); + assert!(!query.contains("iggy_stream")); + assert!(!query.contains("iggy_topic")); + assert!(!query.contains("iggy_partition_id")); + assert!(!query.contains("iggy_checksum")); + assert!(!query.contains("iggy_origin_timestamp")); + assert!(query.contains("payload")); + assert!(query.contains("created_at")); + assert_eq!(param_count, 3); + } + + #[test] + fn given_metadata_disabled_should_build_minimal_parquet() { + let payload = json_payload(serde_json::json!({"name": "Bebeto", "active": true})); + let message = test_message(payload); + let record_batch = create_record_batch( + &test_topic_metadata(), + &test_messages_metadata(), + &[message], + false, + false, + false, + PayloadFormat::Varbyte, + ) + .expect("Failed to create record batch"); + + assert_eq!(record_batch.num_rows(), 1); + assert_eq!(record_batch.num_columns(), 2); + + let columns = record_batch.schema(); + let columns: HashSet<&str> = columns.fields().iter().map(|f| f.name().as_ref()).collect(); + + let expected_columns: HashSet<&str> = HashSet::from(["id", "payload", "created_at"]); + + assert_eq!(columns.difference(&expected_columns).count(), 0); + } + + #[test] + fn given_microseconds_should_parse_timestamp_correctly() { + let record_batch = create_record_batch( + &test_topic_metadata(), + &test_messages_metadata(), + &[test_message(json_payload(serde_json::json!({})))], + true, + false, + false, + PayloadFormat::Varbyte, + ) + .expect("Failed to create record batch"); + + let timestamp_col = record_batch + .column(2) + .as_any() + .downcast_ref::() + .expect("Failed to downcast to Timestamp Microsecond array"); + + let timestamp = timestamp_col.value(0); + + assert_eq!(timestamp, 1_767_225_600_000_000); + } + + #[test] + fn given_default_config_should_use_default_archive() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + assert!(!sink.get_archive()); + } + + #[test] + fn given_archive_enabled_should_use_archive() { + let mut sink = RedshiftSink::new(1, test_config(false, false, false)); + sink.config.archive = Some(true); + + assert!(sink.get_archive()); + } + + #[test] + fn given_default_config_should_use_default_retries() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + assert_eq!(sink.get_max_retries(), DEFAULT_MAX_RETRIES); + } + + #[test] + fn given_custom_retries_should_use_custom_value() { + let mut config = test_config(false, false, false); + config.max_retries = Some(5); + let sink = RedshiftSink::new(1, config); + assert_eq!(sink.get_max_retries(), 5); + } + + #[test] + fn given_default_config_should_use_default_retry_delay() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + assert_eq!(sink.get_retry_delay(), Duration::from_secs(1)); + } + + #[test] + fn given_custom_retry_delay_should_parse_humantime() { + let mut config = test_config(false, false, false); + config.retry_delay = Some("500ms".to_string()); + let sink = RedshiftSink::new(1, config); + assert_eq!(sink.get_retry_delay(), Duration::from_millis(500)); + } + + #[test] + fn given_verbose_logging_enabled_should_set_verbose_flag() { + let mut config = test_config(false, false, false); + config.verbose_logging = Some(true); + let sink = RedshiftSink::new(1, config); + assert!(sink.verbose); + } + + #[test] + fn given_verbose_logging_disabled_should_not_set_verbose_flag() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + assert!(!sink.verbose); + } + + #[test] + fn given_connection_string_with_credentials_should_redact() { + let conn = "postgres://redshift:redshift@localhost:5432/db"; + let redacted = redact_connection_string(conn); + assert_eq!(redacted, "postgres://red***"); + } + + #[test] + fn given_connection_string_without_scheme_should_redact() { + let conn = "localhost:5432/db"; + let redacted = redact_connection_string(conn); + assert_eq!(redacted, "loc***"); + } + + #[test] + fn given_postgresql_scheme_should_redact() { + let conn = "postgresql://admin:secret123@db.example.com:5432/mydb"; + let redacted = redact_connection_string(conn); + assert_eq!(redacted, "postgresql://adm***"); + } + + #[test] + fn given_special_chars_in_identifier_should_escape() { + let result = quote_identifier("table\"name").expect("Failed to quote"); + assert_eq!(result, "\"table\"\"name\""); + } + + #[test] + fn given_empty_identifier_should_fail() { + let result = quote_identifier(""); + assert!(result.is_err()); + } + + #[test] + fn given_null_char_in_identifier_should_fail() { + let result = quote_identifier("table\0name"); + assert!(result.is_err()); + } + + #[test] + fn given_normal_identifier_should_quote() { + let result = quote_identifier("my_table").expect("Failed to quote"); + assert_eq!(result, "\"my_table\""); + } + + #[test] + fn given_identifier_with_spaces_should_quote() { + let result = quote_identifier("my table").expect("Failed to quote"); + assert_eq!(result, "\"my table\""); + } + + #[test] + fn given_identifier_with_sql_injection_should_escape() { + let result = quote_identifier("messages\"; DROP TABLE users; --").expect("Failed to quote"); + assert_eq!(result, "\"messages\"\"; DROP TABLE users; --\""); + } +} diff --git a/core/integration/Cargo.toml b/core/integration/Cargo.toml index 9569d4d1d2..0f95f4ffa7 100644 --- a/core/integration/Cargo.toml +++ b/core/integration/Cargo.toml @@ -34,6 +34,7 @@ login-session = ["dep:zbus-secret-service-keyring-store"] vsr = ["iggy/vsr"] [dependencies] +arrow = { workspace = true } assert_cmd = { workspace = true } async-trait = { workspace = true } base64 = { workspace = true } @@ -63,6 +64,8 @@ keyring-core = { workspace = true } lazy_static = { workspace = true } libc = { workspace = true } mongodb = { workspace = true } +parquet = { workspace = true } +pgwire = { workspace = true } predicates = { workspace = true } rand = { workspace = true } rcgen = { workspace = true } @@ -82,6 +85,7 @@ serde_json = { workspace = true } serial_test = { workspace = true } server = { workspace = true } socket2 = { workspace = true } +sqlparser = { workspace = true } sqlx = { workspace = true } sysinfo = { workspace = true } tempfile = { workspace = true } @@ -89,6 +93,7 @@ test-case = { workspace = true } testcontainers = { workspace = true } testcontainers-modules = { workspace = true } tokio = { workspace = true, features = ["full", "test-util"] } +tokio-postgres = { workspace = true } toml = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/core/integration/tests/connectors/fixtures/mod.rs b/core/integration/tests/connectors/fixtures/mod.rs index 885867de96..f3bf6d17f0 100644 --- a/core/integration/tests/connectors/fixtures/mod.rs +++ b/core/integration/tests/connectors/fixtures/mod.rs @@ -27,6 +27,7 @@ mod influxdb; mod mongodb; mod postgres; mod quickwit; +mod redshift; mod s3; mod surrealdb; mod wiremock; @@ -78,6 +79,10 @@ pub use postgres::{ PostgresSourceJsonbFixture, PostgresSourceMarkFixture, PostgresSourceOps, }; pub use quickwit::{QuickwitFixture, QuickwitOps, QuickwitPreCreatedFixture}; +pub use redshift::{ + RedshiftSinkByteaFixture, RedshiftSinkFixture, RedshiftSinkJsonFixture, + RedshiftSinkNoArchiveFixture, +}; pub use s3::{S3SinkFixture, S3SinkOps, S3SinkRotationFixture}; pub use surrealdb::{ SurrealDbOps, SurrealDbSinkBatchFixture, SurrealDbSinkFixture, SurrealDbSinkJsonFixture, diff --git a/core/integration/tests/connectors/fixtures/redshift/container.rs b/core/integration/tests/connectors/fixtures/redshift/container.rs new file mode 100644 index 0000000000..fb696e8135 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/container.rs @@ -0,0 +1,265 @@ +// 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 std::sync::Arc; + +use integration::harness::TestBinaryError; +use pgwire::tokio::process_socket; +use sqlx::{Pool, Postgres, postgres::PgPoolOptions}; +use testcontainers::{ + ContainerAsync, GenericImage, ImageExt, + core::{IntoContainerPort, WaitFor, wait::HttpWaitStrategy}, + runners::AsyncRunner, +}; +use testcontainers_modules::postgres; +use tokio::{net::TcpListener, task::JoinHandle}; + +use crate::connectors::fixtures::{ + self, + redshift::redshift_mock::handler::{RedshiftHandler, RedshiftHandlerFactory}, +}; + +const MINIO_IMAGE: &str = "docker.io/minio/minio"; +const MINIO_TAG: &str = "RELEASE.2025-09-07T16-13-09Z"; +const MINIO_PORT: u16 = 9000; +const MINIO_CONSOLE_PORT: u16 = 9001; +const POSTGRES_PORT: u16 = 5432; + +pub const MINIO_ACCESS_KEY: &str = "admin"; +pub const MINIO_SECRET_KEY: &str = "password"; +pub const MINIO_BUCKET: &str = "iggystaging"; +pub const DEFAULT_SINK_TABLE: &str = "iggy_messages"; +pub const STAGING_REGION: &str = "us-east-1"; +pub const STAGING_PREFIX: &str = "iggy/messages"; + +pub const ENV_SINK_CONNECTION_STRING: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_CONNECTION_STRING"; +pub const ENV_SINK_TARGET_TABLE: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_TARGET_TABLE"; +pub const ENV_SINK_PAYLOAD_FORMAT: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_PAYLOAD_FORMAT"; +pub const ENV_SINK_STAGING_ACCESS_KEY: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_AWS_ACCESS_KEY_ID"; +pub const ENV_SINK_STAGING_SECRET: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_AWS_SECRET_ACCESS_KEY"; +pub const ENV_SINK_S3_BUCKET: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_S3_BUCKET"; +pub const ENV_SINK_S3_PREFIX: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_S3_PREFIX"; +pub const ENV_SINK_S3_ENDPOINT: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_S3_ENDPOINT"; +pub const ENV_SINK_STAGING_REGION: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_AWS_REGION"; +pub const ENV_SINK_PATH: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PATH"; +pub const ENV_SINK_STREAMS_0_STREAM: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_STREAMS_0_STREAM"; +pub const ENV_SINK_STREAMS_0_TOPICS: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_STREAMS_0_TOPICS"; +pub const ENV_SINK_STREAMS_0_SCHEMA: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_STREAMS_0_SCHEMA"; +pub const ENV_SINK_STREAMS_0_CONSUMER_GROUP: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_STREAMS_0_CONSUMER_GROUP"; +pub const ENV_SINK_ARCHIVE: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_ARCHIVE"; +pub const DEFAULT_TEST_STREAM: &str = "test_stream"; +pub const DEFAULT_TEST_TOPIC: &str = "test_topic"; + +pub const DEFAULT_POLL_ATTEMPTS: usize = 100; +pub const DEFAULT_POLL_INTERVAL_MS: u64 = 50; + +pub struct MinioContainer { + #[allow(dead_code)] + container: ContainerAsync, + pub endpoint: String, +} + +impl MinioContainer { + pub async fn start(network: &str, container_name: &str) -> Result { + let container = GenericImage::new(MINIO_IMAGE, MINIO_TAG) + .with_exposed_port(MINIO_PORT.tcp()) + .with_exposed_port(MINIO_CONSOLE_PORT.tcp()) + .with_wait_for(WaitFor::http( + HttpWaitStrategy::new("/minio/health/live") + .with_port(MINIO_PORT.tcp()) + .with_expected_status_code(200u16), + )) + .with_network(network) + .with_container_name(container_name) + .with_env_var("MINIO_ROOT_USER", MINIO_ACCESS_KEY) + .with_env_var("MINIO_ROOT_PASSWORD", MINIO_SECRET_KEY) + .with_cmd(vec!["server", "/data", "--console-address", ":9001"]) + .with_mapped_port(0, MINIO_PORT.tcp()) + .with_mapped_port(0, MINIO_CONSOLE_PORT.tcp()) + .start() + .await + .map_err(|error| TestBinaryError::FixtureSetup { + fixture_type: "MinioContainer".to_string(), + message: format!("Failed to start container: {error}"), + })?; + + tracing::info!("Started MinIO container"); + + let mapped_port = container + .ports() + .await + .map_err(|error| TestBinaryError::FixtureSetup { + fixture_type: "MinioContainer".to_string(), + message: format!("Failed to get ports: {error}"), + })? + .map_to_host_port_ipv4(MINIO_PORT) + .ok_or_else(|| TestBinaryError::FixtureSetup { + fixture_type: "MinioContainer".to_string(), + message: "No mapping for MinIO port".to_string(), + })?; + + let endpoint = format!("http://localhost:{mapped_port}"); + tracing::info!("MinIO container available at {endpoint}"); + + Ok(Self { + container, + endpoint, + }) + } +} + +/// Base container management for PostgreSQL fixtures. +pub struct PostgresContainer { + #[allow(dead_code)] + container: ContainerAsync, + pub connection_string: String, +} + +impl PostgresContainer { + pub async fn start() -> Result { + let container = postgres::Postgres::default() + .with_container_name(fixtures::unique_container_name("postgres")) + .start() + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "PostgresContainer".to_string(), + message: format!("Failed to start container: {e}"), + })?; + + let host_port = container + .get_host_port_ipv4(POSTGRES_PORT) + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "PostgresContainer".to_string(), + message: format!("Failed to get port: {e}"), + })?; + + let connection_string = format!("postgres://postgres:postgres@localhost:{host_port}"); + + Ok(Self { + container, + connection_string, + }) + } + + pub async fn create_pool(&self) -> Result, TestBinaryError> { + PgPoolOptions::new() + .max_connections(1) + .connect(&self.connection_string) + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "PostgresContainer".to_string(), + message: format!("Failed to connect: {e}"), + }) + } +} + +pub struct RedshiftContainer { + #[allow(dead_code)] + accept_task: JoinHandle<()>, + pub connection_string: String, +} + +impl RedshiftContainer { + pub async fn start( + target_connection: String, + s3_endpoint: String, + ) -> Result { + let (pg_client, connection) = + tokio_postgres::connect(&target_connection, tokio_postgres::NoTls) + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftContainer".into(), + message: e.to_string(), + })?; + + tokio::spawn(async move { + if let Err(e) = connection.await { + panic!("{}", e.to_string()) + } + }); + + let redshql = RedshiftHandler::new(pg_client, s3_endpoint); + + let factory = Arc::new(RedshiftHandlerFactory { + handler: Arc::new(redshql), + }); + + let listener = + TcpListener::bind("127.0.0.1:0") + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftMockContainer".to_string(), + message: format!("bind failed: {e}"), + })?; + + let host_port = listener + .local_addr() + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftMockContainer".to_string(), + message: format!("failed to get local address: {e}"), + })? + .port(); + + let accept_task = tokio::spawn(async move { + loop { + match listener.accept().await { + Ok((incoming_socket, _addr)) => { + let factory_ref = factory.clone(); + + tokio::spawn(async move { + if let Err(e) = process_socket(incoming_socket, None, factory_ref).await + { + panic!("{}", e.to_string()) + } + }); + } + + Err(e) => { + panic!("{}", e.to_string()) + } + } + } + }); + + Ok(Self { + accept_task, + connection_string: format!("postgres://postgres@localhost:{host_port}/postgres"), + }) + } +} + +/// Payload format for sink connector. +#[derive(Debug, Clone, Copy, Default)] +pub enum SinkPayloadFormat { + #[default] + Bytea, + Text, +} + +/// Schema format for message encoding. +#[derive(Debug, Clone, Copy, Default)] +pub enum SinkSchema { + #[default] + Json, + Raw, +} diff --git a/core/integration/tests/connectors/fixtures/redshift/mod.rs b/core/integration/tests/connectors/fixtures/redshift/mod.rs new file mode 100644 index 0000000000..4599ee0d0c --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/mod.rs @@ -0,0 +1,26 @@ +// 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. + +mod container; +mod redshift_mock; +mod sink; + +pub use container::{MinioContainer, PostgresContainer, RedshiftContainer}; +pub use sink::{ + RedshiftSinkByteaFixture, RedshiftSinkFixture, RedshiftSinkJsonFixture, + RedshiftSinkNoArchiveFixture, +}; diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/copy.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/copy.rs new file mode 100644 index 0000000000..a5650d1f42 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/copy.rs @@ -0,0 +1,146 @@ +// 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 std::collections::HashMap; + +use sqlparser::{ + ast::ObjectName, + keywords::Keyword, + parser::{Parser, ParserError}, + tokenizer::Token, +}; + +use crate::connectors::fixtures::redshift::redshift_mock::{ + expect_word, parse_number_literal, parse_string_literal, +}; + +#[allow(unused)] +#[derive(Debug, Clone)] +pub struct RedshiftCopy { + pub table: ObjectName, + pub s3_uri: String, + #[allow(dead_code)] + pub access_key_id: String, + pub secret_access_key: String, + pub format: CopyFormat, + pub max_error: u32, + #[allow(dead_code)] + pub region: String, + pub terminator: bool, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum CopyFormat { + Parquet, + Other(String), +} + +pub fn try_parse_redshift_copy(mut parser: Parser) -> Result { + if !parser.parse_keyword(Keyword::COPY) { + Err(ParserError::ParserError("Not a copy statement".to_string()))? + } + + let table = parser.parse_object_name(false)?; + + parser.expect_keyword(Keyword::FROM)?; + + let s3_uri = parse_string_literal(&mut parser)?; + + if !s3_uri.starts_with("s3://") { + Err(ParserError::ParserError(format!( + "expected s3:// URI, got {s3_uri}" + )))? + } + + let mut access_key_id = String::new(); + let mut secret_access_key = String::new(); + let mut format = CopyFormat::Parquet; + let mut max_error = 0u32; + let mut region = String::new(); + let mut terminator = false; + + // Real Redshift COPY options are unordered after FROM — loop until EOF + while parser.peek_token() != Token::EOF { + let word = expect_word(&mut parser)?; + + match word.to_uppercase().as_str() { + "CREDENTIALS" | "IAM_ROLE" => { + let credentials = parse_string_literal(&mut parser)?; + let mut credentials = parse_credentials(&credentials); + access_key_id = credentials + .remove("ACCESS_KEY_ID") + .ok_or_else(|| ParserError::ParserError("Missing access_key_id".into()))?; + secret_access_key = credentials + .remove("SECRET_ACCESS_KEY") + .ok_or_else(|| ParserError::ParserError("Missing access_key_id".into()))?; + } + "FORMAT" => { + let _ = parser.parse_keyword(Keyword::AS); // "FORMAT AS X" or bare "FORMAT X" + format = match expect_word(&mut parser)?.to_uppercase().as_str() { + "PARQUET" => CopyFormat::Parquet, + other => CopyFormat::Other(other.to_string()), + }; + } + "MAXERROR" => max_error = parse_number_literal(&mut parser)?, + "REGION" => region = parse_string_literal(&mut parser)?, + // clauses you don't emit but want to tolerate rather than error on + "GZIP" | "COMPUPDATE" | "STATUPDATE" => { + let _ = parser.parse_one_of_keywords(&[Keyword::ON, Keyword::OFF]); + } + "IGNOREHEADER" => { + parse_number_literal(&mut parser)?; + } + "DELIMITER" => { + parse_string_literal(&mut parser)?; + } + "SEMICOLON" => terminator = true, + unknown => { + return Err(ParserError::ParserError(format!( + "unsupported COPY clause: {unknown}" + ))); + } + } + } + + Ok(RedshiftCopy { + table, + s3_uri, + access_key_id, + secret_access_key, + format, + max_error, + region, + terminator, + }) +} + +/// Parses a Redshift-style `CREDENTIALS '...'` value into key-value pairs. +/// Input example: "ACCESS_KEY_ID=admin; SECRET_ACCESS_KEY=1234" +fn parse_credentials(raw: &str) -> HashMap { + raw.split(';') + .filter_map(|pair| { + let pair = pair.trim(); + if pair.is_empty() { + return None; + } + let mut parts = pair.splitn(2, '='); + let key = parts.next()?.trim().to_string().to_uppercase(); + let value = parts.next()?.trim().to_string(); + Some((key, value)) + }) + .collect() +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/create.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/create.rs new file mode 100644 index 0000000000..6bf91191c0 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/create.rs @@ -0,0 +1,276 @@ +// 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 sqlparser::{ + ast::{DataType, Expr, Ident, ObjectName}, + keywords::Keyword, + parser::{Parser, ParserError}, + tokenizer::Token, +}; + +use crate::connectors::fixtures::redshift::redshift_mock::expect_word; + +#[allow(unused)] +#[derive(Debug, Clone)] +pub struct RedshiftCreateTable { + pub table: ObjectName, + pub if_not_exists: bool, + pub columns: Vec, + pub table_kind: TableKind, // TEMP / LOCAL TEMP / regular + pub dist_style: Option, + pub dist_key: Option, // column name, only valid when dist_style == Key + pub sort_key: Option, + pub backup: Option, // BACKUP YES | NO + pub terminator: bool, +} + +#[allow(unused)] +#[derive(Debug, Clone)] +pub struct RedshiftColumnDef { + pub name: Ident, + pub data_type: DataType, + // ENCODE ZSTD, LZO, RAW, etc. + pub encoding: Option, + pub not_null: bool, + pub default: Option, + // IDENTITY(seed, step) + pub identity: Option, + pub primary_key: bool, + // simplified FK target + pub references: Option, +} + +#[derive(Debug, Clone)] +pub enum TableKind { + Regular, + Temp, + LocalTemp, +} + +#[derive(Debug, Clone)] +pub enum DistStyle { + Even, + Key, + All, + Auto, +} + +#[allow(unused)] +#[derive(Debug, Clone)] +pub enum SortKey { + Compound(Vec), + Interleaved(Vec), +} + +#[allow(unused)] +#[derive(Debug, Clone, PartialEq)] +pub enum ColumnEncoding { + Raw, + Bytedict, + Delta, + Delta32k, + Lzo, + Mostly8, + Mostly16, + Mostly32, + Runlength, + Text255, + Text32k, + Zstd, + Az64, +} + +#[derive(Debug, Clone)] +pub struct IdentitySpec { + pub seed: i64, + pub step: i64, +} + +// Parse CREATE +pub fn try_parse_redshift_create_table( + mut parser: Parser, +) -> Result { + if !parser.parse_keyword(Keyword::CREATE) { + Err(ParserError::ParserError( + "Not a create statement".to_string(), + ))? + } + + let table_kind = + if parser.parse_keyword(Keyword::TEMPORARY) || parser.parse_keyword(Keyword::TEMP) { + TableKind::Temp + } else if parser.parse_keywords(&[Keyword::LOCAL, Keyword::TEMPORARY]) + || parser.parse_keywords(&[Keyword::LOCAL, Keyword::TEMP]) + { + TableKind::LocalTemp + } else { + TableKind::Regular + }; + + parser.expect_keyword(Keyword::TABLE)?; + + let if_not_exists = parser.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]); + + let table = parser.parse_object_name(false)?; + + let columns = parse_column_list(&mut parser)?; + + let mut dist_style = None; + let mut dist_key = None; + let mut sort_key = None; + let mut backup = None; + let mut terminator = false; + + // Table-level clauses after the column list are unordered, same as COPY options + while parser.peek_token() != Token::EOF { + let word = expect_word(&mut parser)?; + match word.to_uppercase().as_str() { + "DISTSTYLE" => { + dist_style = Some(match expect_word(&mut parser)?.to_uppercase().as_str() { + "EVEN" => DistStyle::Even, + "KEY" => DistStyle::Key, + "ALL" => DistStyle::All, + "AUTO" => DistStyle::Auto, + other => Err(ParserError::ParserError(format!( + "unknown DISTSTYLE: {other}" + )))?, + }); + } + "DISTKEY" => { + parser.expect_token(&Token::LParen)?; + dist_key = Some(parser.parse_identifier()?); + parser.expect_token(&Token::RParen)?; + } + "SORTKEY" => { + sort_key = Some(SortKey::Compound(parse_ident_list(&mut parser)?)); + } + "COMPOUND" => { + parser.expect_keyword(Keyword::SORTKEY)?; + sort_key = Some(SortKey::Compound(parse_ident_list(&mut parser)?)); + } + "INTERLEAVED" => { + parser.expect_keyword(Keyword::SORTKEY)?; + sort_key = Some(SortKey::Interleaved(parse_ident_list(&mut parser)?)); + } + "BACKUP" => { + backup = Some(match expect_word(&mut parser)?.to_uppercase().as_str() { + "YES" => true, + "NO" => false, + other => Err(ParserError::ParserError(format!( + "expected YES|NO after BACKUP, got {other}" + )))?, + }); + } + "ENCODE" => { + // table-level ENCODE AUTO|NONE — tolerate, not modeled per-table yet + let _ = expect_word(&mut parser)?; + } + "SEMICOLON" => terminator = true, + unknown => { + return Err(ParserError::ParserError(format!( + "unsupported CREATE TABLE clause: {unknown}" + ))); + } + } + } + + Ok(RedshiftCreateTable { + table, + if_not_exists, + columns, + table_kind, + dist_style, + dist_key, + sort_key, + backup, + terminator, + }) +} + +fn parse_ident_list(parser: &mut Parser) -> Result, ParserError> { + parser.expect_token(&Token::LParen)?; + let idents = parser.parse_comma_separated(Parser::parse_identifier)?; + parser.expect_token(&Token::RParen)?; + Ok(idents) +} + +fn parse_column_list(parser: &mut Parser) -> Result, ParserError> { + parser.expect_token(&Token::LParen)?; + let mut columns = Vec::new(); + + loop { + let name = parser.parse_identifier()?; + let data_type = parser.parse_data_type()?; + + let encoding = None; + let mut not_null = false; + let mut default = None; + let identity = None; + let mut primary_key = false; + let mut references = None; + + // Column constraints are unordered too — loop until comma or close paren + loop { + match parser.peek_token().token { + Token::Comma | Token::RParen => break, + _ => {} + } + let word = expect_word(parser)?; + match word.to_uppercase().as_str() { + "NOT" => { + parser.expect_keyword(Keyword::NULL)?; + not_null = true; + } + "NULL" => not_null = false, + "DEFAULT" => default = Some(parser.parse_expr()?), + "PRIMARY" => { + parser.expect_keyword(Keyword::KEY)?; + primary_key = true; + } + "REFERENCES" => { + references = Some(parser.parse_object_name(false)?); + } + // column-level DISTKEY/SORTKEY flags — tolerate, table-level fields win + "DISTKEY" | "SORTKEY" => {} + unknown => { + Err(ParserError::ParserError(format!( + "Unsupported column constraint: {unknown}" + )))?; + } + } + } + + columns.push(RedshiftColumnDef { + name, + data_type, + encoding, + not_null, + default, + identity, + primary_key, + references, + }); + + if parser.consume_token(&Token::Comma) { + continue; + } + parser.expect_token(&Token::RParen)?; + break; + } + + Ok(columns) +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs new file mode 100644 index 0000000000..139bf918f7 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs @@ -0,0 +1,305 @@ +// 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 std::sync::Arc; + +use async_trait::async_trait; +use pgwire::{ + api::{ + ClientInfo, ClientPortalStore, PgWireServerHandlers, Type, + portal::Portal, + query::{ExtendedQueryHandler, SimpleQueryHandler}, + results::{DescribePortalResponse, DescribeStatementResponse, FieldInfo, Response, Tag}, + stmt::{NoopQueryParser, StoredStatement}, + store::PortalStore, + }, + error::{ErrorInfo, PgWireError, PgWireResult}, +}; +use sqlparser::{ast::Statement, dialect::RedshiftSqlDialect, parser::Parser}; +use tokio_postgres::{Client as PgClient, Statement as PgStatement}; + +use crate::connectors::fixtures::redshift::redshift_mock::{ + copy::try_parse_redshift_copy, + create::try_parse_redshift_create_table, + load::{ + S3Client, ToPgError, execute_create_table, execute_s3_copy, execute_select, split_s3_uri, + }, +}; + +/// Statement failed to parse under the Redshift dialect. +pub const SYNTAX_ERROR: &str = "26000"; +/// Parsed, but uses a construct we don't implement (custom COPY/CREATE +/// extensions, unsupported statement kinds, etc.). +pub const FEATURE_NOT_SUPPORTED: &str = "42601"; +/// The Postgres connection backing this mock failed outright (as +/// opposed to Postgres returning a well-formed DB error). +pub const CONNECTION_EXCEPTION: &str = "08000"; +/// Empty/missing statement. +pub const INVALID_QUERY: &str = "42601"; +/// Statement kind we recognize but intentionally don't support. +pub const WARNING_UNSUPPORTED: &str = "01000"; + +pub struct RedshiftHandlerFactory { + pub handler: Arc, +} + +impl PgWireServerHandlers for RedshiftHandlerFactory { + fn simple_query_handler(&self) -> Arc { + self.handler.clone() + } + + fn extended_query_handler(&self) -> Arc { + self.handler.clone() + } +} + +pub struct RedshiftHandler { + pg: PgClient, + s3_endpoint: String, +} + +impl RedshiftHandler { + pub fn new(pg: PgClient, s3_endpoint: String) -> Self { + Self { pg, s3_endpoint } + } + + /// Prepares `sql` against the backing Postgres connection for describe + /// purposes. Returns `Ok(None)` for statement kinds (currently just + /// `COPY`) that describe to zero fields rather than going through an + /// unsupported `PREPARE`. + async fn prepare_describable(&self, sql: &str) -> PgWireResult> { + let dialect = RedshiftSqlDialect {}; + let statements = parse_redshift_sql(&dialect, sql)?; + + if matches!(statements.first(), Some(Statement::Copy { .. })) { + return Ok(None); + } + + self.pg + .prepare(sql) + .await + .map(Some) + .map_err(|e| map_pg_client_error(e, CONNECTION_EXCEPTION)) + } +} + +#[async_trait] +impl ExtendedQueryHandler for RedshiftHandler { + type Statement = String; + type QueryParser = NoopQueryParser; + + fn query_parser(&self) -> Arc { + Arc::new(NoopQueryParser::new()) + } + + async fn do_query( + &self, + _client: &mut C, + portal: &Portal, + _max_rows: usize, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + let query = &portal.statement.statement; + + if query.trim().is_empty() { + return Ok(Response::EmptyQuery); + } + + execute_statement(query, &self.pg, &self.s3_endpoint).await + } + + async fn do_describe_statement( + &self, + _client: &mut C, + stmt: &StoredStatement, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + let Some(prepared) = self.prepare_describable(&stmt.statement).await? else { + return Ok(DescribeStatementResponse::new(vec![], vec![])); + }; + + let param_types: Vec = prepared.params().to_vec(); + + let fields: Vec = prepared + .columns() + .iter() + .map(|col| { + FieldInfo::new( + col.name().to_owned(), + None, + None, + col.type_().clone(), + pgwire::api::results::FieldFormat::Text, + ) + }) + .collect(); + + Ok(DescribeStatementResponse::new(param_types, fields)) + } + + async fn do_describe_portal( + &self, + _client: &mut C, + portal: &Portal, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + let Some(prepared) = self + .prepare_describable(&portal.statement.statement) + .await? + else { + return Ok(DescribePortalResponse::new(vec![])); + }; + + let fields: Vec = prepared + .columns() + .iter() + .enumerate() + .map(|(idx, col)| { + FieldInfo::new( + col.name().to_owned(), + None, + None, + col.type_().clone(), + portal.result_column_format.format_for(idx), + ) + }) + .collect(); + + Ok(DescribePortalResponse::new(fields)) + } +} + +#[async_trait] +impl SimpleQueryHandler for RedshiftHandler { + async fn do_query(&self, _client: &mut C, query: &str) -> PgWireResult> + where + C: ClientInfo + ClientPortalStore + Unpin + Send + Sync, + C::PortalStore: PortalStore, + { + if query.trim().is_empty() { + return Ok(vec![Response::EmptyQuery]); + } + + Ok(vec![ + execute_statement(query, &self.pg, &self.s3_endpoint).await?, + ]) + } +} + +/// Dispatches a single SQL statement to the appropriate executor based on +/// its parsed kind. Kept intentionally thin — each branch delegates to a +/// dedicated function so individual statement kinds can be read (and +/// tested) in isolation. +async fn execute_statement( + query: &str, + pg: &PgClient, + s3_endpoint: &str, +) -> PgWireResult { + let dialect = RedshiftSqlDialect {}; + let statements = parse_redshift_sql(&dialect, query)?; + + match statements.first() { + Some(Statement::Query(_)) => execute_select(query, pg).await, + Some(Statement::CreateTable(_)) => execute_create(query, pg).await, + Some(Statement::Copy { .. }) => execute_copy(query, pg, s3_endpoint).await, + Some(other) => Err(pg_warning( + WARNING_UNSUPPORTED, + format!("Unsupported: {other:?}"), + )), + None => Err(pg_error(INVALID_QUERY, "Invalid query")), + } +} + +async fn execute_create(query: &str, pg: &PgClient) -> PgWireResult { + let dialect = RedshiftSqlDialect {}; + let parser = Parser::new(&dialect) + .try_with_sql(query) + .map_err(|e| pg_error(SYNTAX_ERROR, format!("Unsupported: {e:?}")))?; + + let create = try_parse_redshift_create_table(parser) + .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; + + execute_create_table(create, pg).await +} + +async fn execute_copy(query: &str, pg: &PgClient, s3_endpoint: &str) -> PgWireResult { + let dialect = RedshiftSqlDialect {}; + let parser = Parser::new(&dialect) + .try_with_sql(query) + .map_err(|e| pg_error(SYNTAX_ERROR, format!("Unsupported: {e:?}")))?; + + let r_copy = try_parse_redshift_copy(parser) + .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; + + let (bucket_name, prefix) = split_s3_uri(&r_copy.s3_uri).map_err(|e| e.to_pg_wire_error())?; + + let s3_client = S3Client::new( + &bucket_name, + s3_endpoint, + &r_copy.access_key_id, + &r_copy.secret_access_key, + &r_copy.region, + ) + .await + .map_err(|e| e.to_pg_wire_error())?; + + let rows = execute_s3_copy(&r_copy, pg, s3_client, &bucket_name, &prefix) + .await + .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; + + Ok(Response::Execution(Tag::new("copy").with_rows(rows))) +} + +/// Builds a `PgWireError::UserError` with severity `ERROR`. Replaces the +/// repeated `PgWireError::UserError(Box::new(ErrorInfo::new(...)))` calls. +fn pg_error(code: &str, message: impl std::fmt::Display) -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".into(), + code.into(), + message.to_string(), + ))) +} + +/// Same as [`pg_error`] but with severity `WARNING`, for statement kinds +/// we recognize but choose not to support. +fn pg_warning(code: &str, message: impl std::fmt::Display) -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "WARNING".into(), + code.into(), + message.to_string(), + ))) +} + +/// Maps a `tokio_postgres::Error` to a `PgWireError`, preserving the +/// upstream SQLSTATE/message when Postgres itself produced the error, and +/// falling back to `fallback_code` for connection-level failures. +fn map_pg_client_error(err: tokio_postgres::Error, fallback_code: &str) -> PgWireError { + match err.as_db_error() { + Some(db_err) => pg_error(db_err.code().code(), db_err.message()), + None => pg_error(fallback_code, format!("connection failed: {err}")), + } +} + +fn parse_redshift_sql(dialect: &RedshiftSqlDialect, sql: &str) -> PgWireResult> { + Parser::parse_sql(dialect, sql).map_err(|e| pg_error(SYNTAX_ERROR, e)) +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs new file mode 100644 index 0000000000..8de525b15d --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs @@ -0,0 +1,603 @@ +// 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 std::{fmt::Write, sync::Arc}; + +use arrow::{ + array::{ + Array, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Decimal256Array, + Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, + }, + datatypes::DataType, +}; + +use bytes::Bytes; +use futures::{StreamExt, pin_mut, stream}; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use pgwire::{ + api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag}, + error::{ErrorInfo, PgWireError, PgWireResult}, +}; + +use s3::{Bucket, Region, creds::Credentials}; +use sqlparser::ast::DataType as SDataType; +use sqlx::types::chrono; +use tokio_postgres::{ + Client as PgClient, GenericClient, + binary_copy::BinaryCopyInWriter, + types::{ToSql, Type as PgType}, +}; + +use crate::connectors::fixtures::redshift::redshift_mock::{ + copy::{CopyFormat, RedshiftCopy}, + create::{RedshiftCreateTable, TableKind}, +}; + +pub async fn execute_select(sql: &str, pg: &PgClient) -> PgWireResult { + let value = pg + .client() + .query_one_scalar::(sql, &[]) + .await + .map_err(|e| pg_wire_error(&e, "Query failed"))?; + + let schema = Arc::new(vec![FieldInfo::new( + "column".into(), + None, + None, + PgType::INT4, + FieldFormat::Text, + )]); + + let schema_ref = schema.clone(); + + let row_stream = stream::iter(std::iter::once(value)).map(move |v| { + let mut encoder = DataRowEncoder::new(schema_ref.clone()); + encoder.encode_field(&v)?; + + Ok(encoder.take_row()) + }); + + Ok(Response::Query(QueryResponse::new(schema, row_stream))) +} + +pub async fn execute_create_table( + create: RedshiftCreateTable, + pg: &PgClient, +) -> PgWireResult { + let create_kw = match create.table_kind { + TableKind::Regular => "CREATE TABLE", + TableKind::Temp => "CREATE TEMP TABLE", + // Postgres accepts (and ignores) the LOCAL keyword per the SQL + // standard, so this round-trips fine. + TableKind::LocalTemp => "CREATE LOCAL TEMPORARY TABLE", + }; + + let mut sql = String::from(create_kw); + sql.push(' '); + + if create.if_not_exists { + sql.push_str("IF NOT EXISTS "); + } + + write!(sql, "{} (", create.table).unwrap(); + + let fields = create.columns.iter().fold(vec![], |mut acc, v| { + let mut f = format!("{} {}", v.name, map_type(&v.data_type)); + + if let Some(id) = &v.identity { + f.push_str(&format!( + " GENERATED BY DEFAULT AS IDENTITY (START WITH {} INCREMENT BY {})", + id.seed, id.step + )); + } else if let Some(default) = &v.default { + let default = format!("{}", default).replace("GETDATE", "NOW"); + f.push_str(&format!(" DEFAULT {default}")); + } + + if v.primary_key { + f.push_str(" PRIMARY KEY"); + } + + if v.not_null { + f.push_str(" NOT NULL"); + } + + if let Some(target) = &v.references { + f.push_str(&format!(" REFERENCES {target}")); + } + + acc.push(f); + + acc + }); + + sql.push_str(&fields.join(", ")); + sql.push_str(");"); + + let rows_affected = pg + .execute(&sql, &[]) + .await + .map_err(|e| pg_wire_error(&e, "CREATE TABLE failed"))?; + + Ok(Response::Execution( + Tag::new("create table").with_rows(rows_affected as usize), + )) +} + +pub async fn execute_s3_copy( + copy: &RedshiftCopy, + pg: &PgClient, + s3_client: S3Client, + bucket_name: &str, + prefix: &str, +) -> PgWireResult { + let table_name = copy.table.to_string(); + let mut failed_objects = 0usize; + + match ©.format { + CopyFormat::Parquet => { + let existing_columns = fetch_table_columns(pg, ©.table.to_string()) + .await + .map_err(|e| e.to_pg_wire_error())?; + + tracing::info!( + "target table '{table_name}' {}", + if existing_columns.is_some() { + "exists" + } else { + "does not exist" + } + ); + + let bytes = Bytes::from( + s3_client + .get_object(prefix) + .await + .map_err(|e| e.to_pg_wire_error())?, + ); + + tracing::info!("File '{}' read", prefix); + + let cols: Vec = match existing_columns { + Some(cols) => strip_created_at(cols), + None => { + let inferred = + infer_parquet_schema(bytes.clone()).map_err(|e| e.to_pg_wire_error())?; + + let columns_sql = inferred + .iter() + .map(|f| format!("{} {}", f.name.to_lowercase(), f.pg_type.name())) + .collect::>() + .join(", "); + + create_table(pg, &table_name, &columns_sql) + .await + .map_err(|e| e.to_pg_wire_error())?; + + strip_created_at(inferred) + } + }; + + match load_one_object(pg, ©.table.to_string(), &cols, bytes).await { + Ok(n) => { + tracing::info!("{n} records stored"); + } + Err(e) => { + // Simplification vs real Redshift: MAXERROR there counts bad *rows* + // across the whole load; here each failed *file* counts as one error. + // Good enough for a test double where you're usually asserting + // "load either fully succeeds or trips MAXERROR", not exact counts. + failed_objects += 1; + tracing::error!("[copy] error loading s3://{bucket_name}/{prefix}: {e:#}"); + if failed_objects > copy.max_error as usize { + Err(format!("MAXERROR ({}) exceeded", copy.max_error).to_pg_wire_error())? + } + } + } + } + other => Err(format!("{:?} unsupported", other).to_pg_wire_error())?, + }; + + Ok(1) +} + +async fn fetch_table_columns(pg: &PgClient, table: &str) -> Result>, String> { + let rows = pg + .query( + "SELECT column_name, udt_name FROM information_schema.columns \ + WHERE table_name = $1 ORDER BY ordinal_position", + &[&table], + ) + .await + .map_err(|e| e.to_string())?; + + if rows.is_empty() { + return Ok(None); + } + + let rows: Result, String> = rows + .into_iter() + .map(|row| { + let name: String = row.get(0); + let udt: String = row.get(1); + Ok(ColumnDef { + name, + pg_type: udt_name_to_type(&udt)?, + }) + }) + .collect(); + + Ok(Some(rows?)) +} + +async fn create_table(pg: &PgClient, table: &str, columns: &str) -> Result { + let sql = format!("CREATE TABLE {} ({});", table, columns); + + tracing::info!("{sql}"); + + let result = pg + .client() + .execute(&sql, &[]) + .await + .map_err(|e| e.to_string())?; + + Ok(result as usize) +} + +async fn load_one_object( + pg: &PgClient, + table: &str, + columns: &[ColumnDef], + bytes: Bytes, +) -> Result { + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes) + .map_err(|e| e.to_string())? + .build() + .map_err(|e| e.to_string())?; + + let col_list = columns + .iter() + .map(|v| format!("\"{}\"", v.name)) + .collect::>() + .join(", "); + + let types = columns + .iter() + .map(|v| v.pg_type.clone()) + .collect::>(); + + let copy_sql = format!("COPY {table} ({col_list}) FROM STDIN BINARY"); + let sink = pg.copy_in(©_sql).await.map_err(|e| e.to_string())?; + + tracing::info!("COPY FROM STDIN started"); + + let writer = BinaryCopyInWriter::new(sink, &types); + + pin_mut!(writer); + + let mut n = 0usize; + + for batch in reader { + let batch = batch.map_err(|e| e.to_string())?; + for row_idx in 0..batch.num_rows() { + let row_values = extract_row(&batch, row_idx, columns)?; + + let refs: Vec<&(dyn ToSql + Sync)> = row_values + .iter() + .map(|v| v.as_ref() as &(dyn ToSql + Sync)) + .collect(); + + writer + .as_mut() + .write(&refs) + .await + .map_err(|e| e.to_string())?; + + n += 1; + } + } + + writer.finish().await.map_err(|e| e.to_string())?; + + Ok(n) +} + +struct ColumnDef { + name: String, + pg_type: PgType, +} + +macro_rules! scalar_column { + ($array:expr, $arr_ty:ty, $val_ty:ty, $row:expr, $conv:expr) => {{ + let a = $array + .as_any() + .downcast_ref::<$arr_ty>() + .ok_or_else(|| format!("expected {} array", stringify!($arr_ty)))?; + + if a.is_null($row) { + Box::new(None::<$val_ty>) as Box + } else { + let conv: fn(_) -> $val_ty = $conv; + Box::new(conv(a.value($row))) as Box + } + }}; +} + +/// Only covers common scalar types. Extend as your Parquet exports need more — +/// this deliberately doesn't try to handle structs, lists, or decimals up front. +fn extract_row<'a>( + batch: &'a RecordBatch, + row: usize, + columns: &'a [ColumnDef], +) -> Result>, String> { + let mut out = Vec::with_capacity(columns.len()); + + for (i, col) in columns.iter().enumerate() { + let array = batch.column(i); + + let value: Box = match array.data_type() { + DataType::Utf8 => { + scalar_column!(array, StringArray, String, row, |v: &str| v.to_string()) + } + DataType::Int64 => { + scalar_column!(array, Int64Array, i64, row, |v: i64| v) + } + DataType::Int32 => scalar_column!(array, Int32Array, i32, row, |v: i32| v), + DataType::Float64 => scalar_column!(array, Float64Array, f64, row, |v: f64| v), + DataType::Boolean => scalar_column!(array, BooleanArray, bool, row, |v: bool| v), + DataType::Date32 => scalar_column!(array, Date32Array, i32, row, |v: i32| v), + DataType::Binary => { + scalar_column!(array, BinaryArray, Vec, row, |v: &[u8]| v.to_vec()) + } + DataType::Decimal128(_, _) => { + let a = array + .as_any() + .downcast_ref::() + .ok_or("expected Decimal256Array")?; + if a.is_null(row) { + Box::new(None::) + } else { + Box::new(a.value(row).to_string()) + } + } + // Downcast to fit i256 serialization constraints + // Decimal(39, 0) > becomes Decimal256 on arrow + // Decimal256 serialization will require extra handling + DataType::Decimal256(_, _) => { + let a = array + .as_any() + .downcast_ref::() + .ok_or("expected Decimal256Array")?; + if a.is_null(row) { + Box::new(None::) + } else { + let raw = a.value(row).to_string(); + Box::new(raw) + } + } + DataType::Timestamp(_, _) => { + let a = array + .as_any() + .downcast_ref::() + .ok_or("expected TimestampMicrosecondArray")?; + if a.is_null(row) { + Box::new(None::>) + } else { + let micros = a.value(row); + let dt = chrono::DateTime::::from_timestamp_micros(micros) + .ok_or("Invalid timestamp")?; + Box::new(dt) + } + } + other => Err(format!( + "unsupported parquet column type {other:?} for column {}", + col.name + ))?, + }; + out.push(value); + } + + Ok(out) +} + +fn udt_name_to_type(udt: &str) -> Result { + Ok(match udt { + "int2" => PgType::INT2, + "int4" => PgType::INT4, + "int8" => PgType::INT8, + "float4" => PgType::FLOAT4, + "float8" => PgType::FLOAT8, + // Numeric serialiation requires extra work + // Safe to use VARCHAR + "numeric" => PgType::VARCHAR, + "bool" => PgType::BOOL, + "text" | "varchar" | "bpchar" => PgType::TEXT, + "timestamp" => PgType::TIMESTAMP, + "timestamptz" => PgType::TIMESTAMPTZ, + "date" => PgType::DATE, + "jsonb" => PgType::JSONB, + "bytea" => PgType::BYTEA, + other => Err(format!("unsupported column type for COPY target: {other}"))?, + }) +} + +fn arrow_to_type(a_type: &DataType) -> Result { + match a_type { + DataType::Boolean => Ok(PgType::BOOL), + DataType::Binary | DataType::FixedSizeBinary(_) => Ok(PgType::BYTEA), + DataType::Float64 => Ok(PgType::FLOAT8), + DataType::Float32 | DataType::Float16 => Ok(PgType::FLOAT4), + DataType::Int64 => Ok(PgType::INT8), + DataType::Int32 => Ok(PgType::INT4), + DataType::Decimal128(_, _) => Ok(PgType::VARCHAR), + DataType::Decimal256(_, _) => Ok(PgType::VARCHAR), + DataType::Utf8 => Ok(PgType::TEXT), + DataType::Date32 => Ok(PgType::DATE), + DataType::Timestamp(_, _) => Ok(PgType::TIMESTAMPTZ), + other => Err(format!("Unsuppoerted type: {}", other)), + } +} + +fn map_type(dt: &SDataType) -> String { + match dt { + // Redshift's SUPER (semi-structured) has no PG equivalent for + // your emulator's purposes -> JSON is the closest usable stand-in. + SDataType::Custom(name, _) if name.to_string().eq_ignore_ascii_case("SUPER") => { + "JSON".to_string() + } + SDataType::Custom(name, mods) if name.to_string().eq_ignore_ascii_case("VARBYTE") => { + "BYTEA".to_string() + } + SDataType::Decimal(_) => "VARCHAR".into(), + other => other.to_string(), + } +} + +fn infer_parquet_schema(bytes: Bytes) -> Result, String> { + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).map_err(|e| e.to_string())?; + + reader + .schema() + .fields() + .iter() + .map(|v| { + Ok(ColumnDef { + name: v.name().into(), + pg_type: arrow_to_type(v.data_type())?, + }) + }) + .collect() +} + +fn strip_created_at(cols: Vec) -> Vec { + cols.into_iter() + .filter(|v| v.name != "created_at") + .collect() +} + +/// S3 +#[allow(unused)] +#[derive(Clone)] +pub struct S3Client { + bucket_name: String, + inner: Box, +} + +impl S3Client { + pub async fn new( + bucket_name: &str, + s3_endpoint: &str, + access_key: &str, + secret_key: &str, + region: &str, + ) -> Result { + let region = Region::Custom { + region: region.into(), + endpoint: s3_endpoint.into(), + }; + + let credentials = Credentials::new(Some(access_key), Some(secret_key), None, None, None) + .map_err(|e| e.to_string())?; + + let bucket = Bucket::new(bucket_name, region, credentials) + .map_err(|e| format!("failed to setup bucket: {e}"))? + .with_path_style(); + + Ok(S3Client { + bucket_name: bucket_name.into(), + inner: bucket, + }) + } + + pub async fn get_object(&self, key: &str) -> Result, String> { + tracing::info!( + "Downloading object '{}' from bucket '{}'", + key, + self.bucket_name + ); + + let response = self + .inner + .get_object(key) + .await + .map_err(|e| e.to_string())?; + + if response.status_code() != 200 { + tracing::error!( + "S3 get object returned status {}: {}", + response.status_code(), + String::from_utf8_lossy(response.as_slice()) + ); + return Err(format!( + "S3 get_object failed with status {}", + response.status_code() + )); + } + + tracing::info!( + "Retrieved {} bytes to s3://{}/{}", + response.bytes().len(), + self.inner.name(), + key + ); + + Ok(response.bytes().to_vec()) + } +} + +pub fn split_s3_uri(uri: &str) -> Result<(String, String), String> { + let rest = uri + .strip_prefix("s3://") + .ok_or("not an s3:// URI".to_string())?; + + match rest.split_once('/') { + Some((b, p)) => Ok((b.to_string(), p.to_string())), + None => Err(format!("s3 URI missing key/prefix: {uri}")), + } +} + +pub trait ToPgError { + fn to_pg_wire_error(self) -> PgWireError; +} + +impl ToPgError for T +where + T: Into, +{ + fn to_pg_wire_error(self) -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".into(), + "GG000".into(), + self.into(), + ))) + } +} + +fn pg_wire_error(e: &tokio_postgres::Error, context: &str) -> PgWireError { + match e.as_db_error() { + Some(db_err) => PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".into(), + db_err.code().code().to_string(), + db_err.message().to_string(), + ))), + None => PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".into(), + "XX000".into(), + format!("{context}: {e}"), + ))), + } +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs new file mode 100644 index 0000000000..1d9ba47977 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs @@ -0,0 +1,56 @@ +// 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 sqlparser::{ + parser::{Parser, ParserError}, + tokenizer::{Token, Word}, +}; + +pub mod copy; +pub mod create; +pub mod handler; +pub mod load; + +pub fn expect_word(parser: &mut Parser) -> Result { + match parser.next_token().token { + Token::Word(Word { value, .. }) => Ok(value), + Token::SemiColon => Ok("SemiColon".into()), + other => Err(ParserError::ParserError(format!( + "expected identifier, got {other:?}" + ))), + } +} + +pub fn parse_string_literal(parser: &mut Parser) -> Result { + match parser.next_token().token { + Token::SingleQuotedString(s) => Ok(s), + other => Err(ParserError::ParserError(format!( + "expected string literal, got {other:?}" + ))), + } +} + +pub fn parse_number_literal(parser: &mut Parser) -> Result { + match parser.next_token().token { + Token::Number(s, _) => s + .parse() + .map_err(|_| ParserError::ParserError(format!("bad number: {s}"))), + other => Err(ParserError::ParserError(format!( + "expected number, got {other:?}" + ))), + } +} diff --git a/core/integration/tests/connectors/fixtures/redshift/sink.rs b/core/integration/tests/connectors/fixtures/redshift/sink.rs new file mode 100644 index 0000000000..df062c9dd9 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/sink.rs @@ -0,0 +1,436 @@ +// 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 std::{collections::HashMap, time::Duration}; + +use async_trait::async_trait; +use integration::harness::{TestBinaryError, TestFixture}; +use sqlx::{Pool, Postgres}; +use uuid::Uuid; + +use crate::connectors::fixtures::{ + self, + redshift::{ + MinioContainer, PostgresContainer, RedshiftContainer, + container::{ + DEFAULT_POLL_ATTEMPTS, DEFAULT_POLL_INTERVAL_MS, DEFAULT_SINK_TABLE, + DEFAULT_TEST_STREAM, DEFAULT_TEST_TOPIC, ENV_SINK_ARCHIVE, ENV_SINK_CONNECTION_STRING, + ENV_SINK_PATH, ENV_SINK_PAYLOAD_FORMAT, ENV_SINK_S3_BUCKET, ENV_SINK_S3_ENDPOINT, + ENV_SINK_S3_PREFIX, ENV_SINK_STAGING_ACCESS_KEY, ENV_SINK_STAGING_REGION, + ENV_SINK_STAGING_SECRET, ENV_SINK_STREAMS_0_CONSUMER_GROUP, ENV_SINK_STREAMS_0_SCHEMA, + ENV_SINK_STREAMS_0_STREAM, ENV_SINK_STREAMS_0_TOPICS, ENV_SINK_TARGET_TABLE, + MINIO_ACCESS_KEY, MINIO_BUCKET, MINIO_SECRET_KEY, STAGING_PREFIX, STAGING_REGION, + SinkPayloadFormat, SinkSchema, + }, + }, +}; + +pub struct RedshiftSinkFixture { + #[allow(dead_code)] + minio: MinioContainer, + #[allow(dead_code)] + redshift: RedshiftContainer, + postgres: PostgresContainer, + payload_format: SinkPayloadFormat, + schema: SinkSchema, + pub minio_endpoint: String, +} + +impl RedshiftSinkFixture { + /// Assertions read from here — never from `redshift`. + pub async fn target_pool(&self) -> Result, TestBinaryError> { + self.postgres.create_pool().await + } + + /// Fetch rows from the sink table with polling until expected count is reached. + /// + /// Returns an error if the expected count is not reached within the poll attempts. + pub async fn fetch_rows_as( + &self, + pool: &Pool, + query: &str, + expected_count: usize, + ) -> Result, TestBinaryError> + where + T: Send + Unpin + for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow>, + { + let mut rows = Vec::new(); + for _ in 0..DEFAULT_POLL_ATTEMPTS { + let result = sqlx::query_as::<_, T>(sqlx::AssertSqlSafe(query)) + .fetch_all(pool) + .await; + if let Ok(fetched) = result { + rows = fetched; + if rows.len() >= expected_count { + return Ok(rows); + } + } else if let Err(e) = result { + tracing::error!("{e}"); + } + + tokio::time::sleep(Duration::from_millis(DEFAULT_POLL_INTERVAL_MS)).await; + } + Err(TestBinaryError::InvalidState { + message: format!( + "Expected {} rows but got {} after {} poll attempts", + expected_count, + rows.len(), + DEFAULT_POLL_ATTEMPTS + ), + }) + } +} + +#[async_trait] +impl TestFixture for RedshiftSinkFixture { + async fn setup() -> Result { + let postgres = PostgresContainer::start().await?; + + let id = Uuid::now_v7(); + let network = format!("iggy-redshift-{id}"); + + let minio_name = fixtures::unique_container_name("minio-redshift"); + + let minio = MinioContainer::start(&network, &minio_name).await?; + + create_bucket(&minio.endpoint)?; + + let redshift = + RedshiftContainer::start(postgres.connection_string.clone(), minio.endpoint.clone()) + .await?; + + Ok(Self { + minio_endpoint: minio.endpoint.clone(), + minio, + redshift, + postgres, + payload_format: SinkPayloadFormat::default(), + schema: SinkSchema::default(), + }) + } + + fn connectors_runtime_envs(&self) -> std::collections::HashMap { + let mut envs = HashMap::new(); + + envs.insert( + ENV_SINK_CONNECTION_STRING.to_string(), + self.redshift.connection_string.clone(), + ); + envs.insert( + ENV_SINK_TARGET_TABLE.to_string(), + DEFAULT_SINK_TABLE.to_string(), + ); + + envs.insert( + ENV_SINK_STAGING_ACCESS_KEY.to_string(), + MINIO_ACCESS_KEY.to_string(), + ); + + envs.insert( + ENV_SINK_STAGING_SECRET.to_string(), + MINIO_SECRET_KEY.to_string(), + ); + + envs.insert(ENV_SINK_S3_BUCKET.to_string(), MINIO_BUCKET.to_string()); + envs.insert(ENV_SINK_S3_PREFIX.to_string(), STAGING_PREFIX.to_string()); + + envs.insert( + ENV_SINK_S3_ENDPOINT.to_string(), + self.minio_endpoint.clone(), + ); + + envs.insert( + ENV_SINK_STAGING_REGION.to_string(), + STAGING_REGION.to_string(), + ); + + envs.insert( + ENV_SINK_STREAMS_0_STREAM.to_string(), + DEFAULT_TEST_STREAM.to_string(), + ); + envs.insert( + ENV_SINK_STREAMS_0_TOPICS.to_string(), + format!("[{}]", DEFAULT_TEST_TOPIC), + ); + envs.insert( + ENV_SINK_STREAMS_0_CONSUMER_GROUP.to_string(), + "test".to_string(), + ); + + envs.insert( + ENV_SINK_PATH.to_string(), + "../../target/debug/libiggy_connector_redshift_sink".to_string(), + ); + + let schema_str = match self.schema { + SinkSchema::Json => "json", + SinkSchema::Raw => "raw", + }; + envs.insert( + ENV_SINK_STREAMS_0_SCHEMA.to_string(), + schema_str.to_string(), + ); + + let format_str = match self.payload_format { + SinkPayloadFormat::Bytea => "bytea", + SinkPayloadFormat::Text => "text", + }; + envs.insert(ENV_SINK_PAYLOAD_FORMAT.to_string(), format_str.to_string()); + + envs + } +} + +/// redshift sink fixture for bytea payload format. +pub struct RedshiftSinkByteaFixture { + inner: RedshiftSinkFixture, +} + +impl std::ops::Deref for RedshiftSinkByteaFixture { + type Target = RedshiftSinkFixture; + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[async_trait] +impl TestFixture for RedshiftSinkByteaFixture { + async fn setup() -> Result { + let postgres = PostgresContainer::start().await?; + + let id = Uuid::now_v7(); + let network = format!("iggy-redshift-{id}"); + + let minio_name = fixtures::unique_container_name("minio-redshift"); + + let minio = MinioContainer::start(&network, &minio_name).await?; + + create_bucket(&minio.endpoint)?; + + let redshift = + RedshiftContainer::start(postgres.connection_string.clone(), minio.endpoint.clone()) + .await?; + + Ok(Self { + inner: RedshiftSinkFixture { + minio_endpoint: minio.endpoint.clone(), + minio, + redshift, + postgres, + payload_format: SinkPayloadFormat::Bytea, + schema: SinkSchema::Raw, + }, + }) + } + + fn connectors_runtime_envs(&self) -> HashMap { + self.inner.connectors_runtime_envs() + } +} + +/// redshift sink fixture for bytea payload format. +pub struct RedshiftSinkJsonFixture { + inner: RedshiftSinkFixture, +} + +impl std::ops::Deref for RedshiftSinkJsonFixture { + type Target = RedshiftSinkFixture; + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[async_trait] +impl TestFixture for RedshiftSinkJsonFixture { + async fn setup() -> Result { + let postgres = PostgresContainer::start().await?; + + let id = Uuid::now_v7(); + let network = format!("iggy-redshift-{id}"); + + let minio_name = fixtures::unique_container_name("minio-redshift"); + + let minio = MinioContainer::start(&network, &minio_name).await?; + + create_bucket(&minio.endpoint)?; + + let redshift = + RedshiftContainer::start(postgres.connection_string.clone(), minio.endpoint.clone()) + .await?; + + Ok(Self { + inner: RedshiftSinkFixture { + minio_endpoint: minio.endpoint.clone(), + minio, + redshift, + postgres, + payload_format: SinkPayloadFormat::Text, + schema: SinkSchema::Json, + }, + }) + } + + fn connectors_runtime_envs(&self) -> HashMap { + self.inner.connectors_runtime_envs() + } +} + +/// redshift sink fixture for bytea payload format. +pub struct RedshiftSinkNoArchiveFixture { + inner: RedshiftSinkFixture, +} + +impl RedshiftSinkNoArchiveFixture { + pub fn confirm_empty_bucket(&self) -> Result { + bucket_empty(&self.minio_endpoint) + } +} + +impl std::ops::Deref for RedshiftSinkNoArchiveFixture { + type Target = RedshiftSinkFixture; + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[async_trait] +impl TestFixture for RedshiftSinkNoArchiveFixture { + async fn setup() -> Result { + let postgres = PostgresContainer::start().await?; + + let id = Uuid::now_v7(); + let network = format!("iggy-redshift-{id}"); + + let minio_name = fixtures::unique_container_name("minio-redshift"); + + let minio = MinioContainer::start(&network, &minio_name).await?; + + create_bucket(&minio.endpoint)?; + + let redshift = + RedshiftContainer::start(postgres.connection_string.clone(), minio.endpoint.clone()) + .await?; + + Ok(Self { + inner: RedshiftSinkFixture { + minio_endpoint: minio.endpoint.clone(), + minio, + redshift, + postgres, + payload_format: SinkPayloadFormat::Text, + schema: SinkSchema::Json, + }, + }) + } + + fn connectors_runtime_envs(&self) -> HashMap { + let mut envs = self.inner.connectors_runtime_envs(); + + let schema_str = match self.schema { + SinkSchema::Json => "json", + SinkSchema::Raw => "raw", + }; + envs.insert( + ENV_SINK_STREAMS_0_SCHEMA.to_string(), + schema_str.to_string(), + ); + + envs.insert(ENV_SINK_ARCHIVE.to_string(), false.to_string()); + + let format_str = match self.payload_format { + SinkPayloadFormat::Bytea => "bytea", + SinkPayloadFormat::Text => "text", + }; + + envs.insert(ENV_SINK_PAYLOAD_FORMAT.to_string(), format_str.to_string()); + + envs + } +} + +fn create_bucket(minio_endpoint: &str) -> Result<(), TestBinaryError> { + use std::process::Command; + + let host = minio_endpoint.trim_start_matches("http://"); + let mc_host = format!("http://{}:{}@{}", MINIO_ACCESS_KEY, MINIO_SECRET_KEY, host); + + let output = Command::new("docker") + .args([ + "run", + "--rm", + "--network=host", + "-e", + &format!("MC_HOST_minio={}", mc_host), + "minio/mc", + "mb", + "--ignore-existing", + &format!("minio/{}", MINIO_BUCKET), + ]) + .output() + .map_err(|error| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftFixture".to_string(), + message: format!("Failed to run mc command: {error}"), + })?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + let stdout = String::from_utf8_lossy(&output.stdout); + return Err(TestBinaryError::FixtureSetup { + fixture_type: "IcebergFixture".to_string(), + message: format!("Failed to create bucket: stderr={stderr}, stdout={stdout}"), + }); + } + + tracing::info!("Created MinIO bucket: {MINIO_BUCKET}"); + Ok(()) +} + +fn bucket_empty(minio_endpoint: &str) -> Result { + use std::process::Command; + + let host = minio_endpoint.trim_start_matches("http://"); + let mc_host = format!("http://{}:{}@{}", MINIO_ACCESS_KEY, MINIO_SECRET_KEY, host); + + let output = Command::new("docker") + .args([ + "run", + "--rm", + "--network=host", + "-e", + &format!("MC_HOST_minio={}", mc_host), + "minio/mc", + "ls", + &format!("minio/{}", MINIO_BUCKET), + ]) + .output() + .map_err(|error| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftFixture".to_string(), + message: format!("Failed to run mc command: {error}"), + })?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + let stdout = String::from_utf8_lossy(&output.stdout); + return Err(TestBinaryError::FixtureSetup { + fixture_type: "IcebergFixture".to_string(), + message: format!("Failed to create bucket: stderr={stderr}, stdout={stdout}"), + }); + } + + tracing::info!("Created MinIO bucket: {MINIO_BUCKET}"); + Ok(output.stdout.is_empty()) +} diff --git a/core/integration/tests/connectors/mod.rs b/core/integration/tests/connectors/mod.rs index a1433160b9..56cccf9309 100644 --- a/core/integration/tests/connectors/mod.rs +++ b/core/integration/tests/connectors/mod.rs @@ -30,6 +30,7 @@ mod postgres; mod quickwit; mod random; mod random_source_liveness; +mod redshift; mod runtime; mod s3; mod stdout; diff --git a/core/integration/tests/connectors/redshift/mod.rs b/core/integration/tests/connectors/redshift/mod.rs new file mode 100644 index 0000000000..1038c093b8 --- /dev/null +++ b/core/integration/tests/connectors/redshift/mod.rs @@ -0,0 +1,20 @@ +// 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. + +mod redshift_sink; + +const TEST_MESSAGE_COUNT: usize = 3; diff --git a/core/integration/tests/connectors/redshift/redshift_sink.rs b/core/integration/tests/connectors/redshift/redshift_sink.rs new file mode 100644 index 0000000000..7f6e50d9cb --- /dev/null +++ b/core/integration/tests/connectors/redshift/redshift_sink.rs @@ -0,0 +1,401 @@ +// 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 super::TEST_MESSAGE_COUNT; +use crate::connectors::fixtures::{ + RedshiftSinkByteaFixture, RedshiftSinkFixture, RedshiftSinkJsonFixture, + RedshiftSinkNoArchiveFixture, +}; +use crate::connectors::{TestMessage, create_test_messages}; +use bytes::Bytes; +use iggy::prelude::{IggyMessage, Partitioning}; +use iggy_common::Identifier; +use iggy_common::MessageClient; +use iggy_connector_sdk::api::SinkInfoResponse; +use integration::harness::seeds; +use integration::iggy_harness; + +use reqwest::Client; + +const SINK_TABLE: &str = "iggy_messages"; +const API_KEY: &str = "test-api-key"; +const ICEBERG_SINK_KEY: &str = "redshift"; + +type SinkRow = (i64, String, String, String); +type SinkRawRow = (i64, String, String, Vec); +type SinkJsonRow = (i64, String); + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn redshift_sink_initializes_and_runs(harness: &TestHarness, fixture: RedshiftSinkFixture) { + let api_address = harness + .connectors_runtime() + .expect("connector runtime should be available") + .http_url(); + + let http_client = Client::new(); + + let response = http_client + .get(format!("{}/sinks", api_address)) + .header("api-key", API_KEY) + .send() + .await + .expect("Failed to get sinks"); + + assert_eq!(response.status(), 200); + let sinks: Vec = response.json().await.expect("Failed to parse sinks"); + + assert_eq!(sinks.len(), 1); + assert_eq!(sinks[0].key, ICEBERG_SINK_KEY); + assert!(sinks[0].enabled); + + drop(fixture); +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn json_messages_sink_stores_as_text( + harness: &TestHarness, + fixture: RedshiftSinkJsonFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture + .target_pool() + .await + .expect("Failed to create target postgres pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let messages_data = create_test_messages(TEST_MESSAGE_COUNT); + + let mut messages: Vec = messages_data + .iter() + .enumerate() + .map(|(i, msg)| { + let payload = serde_json::to_vec(msg).expect("Failed to serialize message"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(payload)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!( + "SELECT iggy_offset, iggy_stream, iggy_topic, payload FROM {SINK_TABLE} ORDER BY iggy_offset" + ); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, stream, topic, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!(stream, seeds::names::STREAM, "Stream mismatch at row {i}"); + assert_eq!(topic, seeds::names::TOPIC, "Topic mismatch at row {i}"); + + let stored: TestMessage = + serde_json::from_str(payload).expect("Failed to deserialize stored payload"); + assert_eq!(stored, messages_data[i], "Message data mismatch at row {i}"); + } +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn binary_messages_sink_stores_as_bytea( + harness: &TestHarness, + fixture: RedshiftSinkByteaFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture.target_pool().await.expect("Failed to create pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let raw_payloads: Vec> = vec![ + b"plain text message".to_vec(), + vec![0x00, 0x01, 0x02, 0xFF, 0xFE, 0xFD], + vec![0xDE, 0xAD, 0xBE, 0xEF], + ]; + + let mut messages: Vec = raw_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(payload.clone())) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!( + "SELECT iggy_offset, iggy_stream, iggy_topic, payload FROM {SINK_TABLE} ORDER BY iggy_offset" + ); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, _, _, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!(payload, &raw_payloads[i], "Payload mismatch at row {i}"); + } +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn json_messages_sink_stores_as_json( + harness: &TestHarness, + fixture: RedshiftSinkJsonFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture.target_pool().await.expect("Failed to create pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let json_payloads: Vec = vec![ + serde_json::json!({"name": "Alice", "age": 30}), + serde_json::json!({"items": [1, 2, 3], "active": true}), + serde_json::json!({"nested": {"key": "value"}, "count": 42}), + ]; + + let mut messages: Vec = json_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + let bytes = serde_json::to_vec(payload).expect("Failed to serialize json"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(bytes)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!("SELECT iggy_offset, payload FROM {SINK_TABLE} ORDER BY iggy_offset"); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!( + payload, + &json_payloads[i].to_string(), + "JSON payload mismatch at row {i}" + ); + } +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn json_messages_sink_stores_as_bytea( + harness: &TestHarness, + fixture: RedshiftSinkByteaFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture.target_pool().await.expect("Failed to create pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let json_payloads: Vec = vec![ + serde_json::json!({"name": "Alice", "age": 30}), + serde_json::json!({"items": [1, 2, 3], "active": true}), + serde_json::json!({"nested": {"key": "value"}, "count": 42}), + ]; + + let mut messages: Vec = json_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + let bytes = serde_json::to_vec(payload).expect("Failed to serialize json"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(bytes)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!( + "SELECT iggy_offset, iggy_stream, iggy_topic, payload FROM {SINK_TABLE} ORDER BY iggy_offset" + ); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, _, _, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!( + serde_json::from_slice::(payload).expect("Failed to parse bytes"), + json_payloads[i], + "Payload mismatch at row {i}" + ); + } +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn sink_with_no_archive_deletes_s3_artefact( + harness: &TestHarness, + fixture: RedshiftSinkNoArchiveFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture.target_pool().await.expect("Failed to create pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let json_payloads: Vec = vec![ + serde_json::json!({"name": "Alice", "age": 30}), + serde_json::json!({"items": [1, 2, 3], "active": true}), + serde_json::json!({"nested": {"key": "value"}, "count": 42}), + ]; + + let mut messages: Vec = json_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + let bytes = serde_json::to_vec(payload).expect("Failed to serialize json"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(bytes)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!("SELECT iggy_offset, payload FROM {SINK_TABLE} ORDER BY iggy_offset"); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!( + payload, + &json_payloads[i].to_string(), + "JSON payload mismatch at row {i}" + ); + + assert!( + fixture + .confirm_empty_bucket() + .expect("Failed to read empty bucket") + ); + } +} diff --git a/core/integration/tests/connectors/redshift/sink.toml b/core/integration/tests/connectors/redshift/sink.toml new file mode 100644 index 0000000000..c46dd3bafa --- /dev/null +++ b/core/integration/tests/connectors/redshift/sink.toml @@ -0,0 +1,20 @@ +# 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. + +[connectors] +config_type = "local" +config_dir = "../connectors/sinks/redshift_sink" From 3708d133d5c0e2b1f06cf62ca2ba0af42d7b6d7c Mon Sep 17 00:00:00 2001 From: Bebeto Nyamwamu Date: Thu, 6 Aug 2026 13:05:35 -0400 Subject: [PATCH 2/5] fix(redshift): improved correctness and reliability This change improves the Redshift sink connector across correctness and reliability. The connector now emits explicit column lists, quotes SQL identifiers, validates schemas, verifies created_at metadata, and uses merge semantics for uniqueness handling. Runtime errors are propagated correctly, and COPY operations are atomic while cleanup remains best-effort(there is attempt but failure does not stop overall operation). AWS access key/secret key authentication has been replaced with IAM roles for COPY operations. Safer connection string redaction, mock behavior and metadata flag are documented, JSON payload format logging is imple, and unnecessary message materialization has been removed to reduce allocations. --- Cargo.lock | 29 +- Cargo.toml | 2 +- core/connectors/sinks/redshift_sink/README.md | 108 +++- .../sinks/redshift_sink/src/config.rs | 51 +- .../connectors/sinks/redshift_sink/src/lib.rs | 550 +++++++++++++----- core/integration/Cargo.toml | 1 + .../tests/connectors/fixtures/mod.rs | 4 +- .../connectors/fixtures/redshift/container.rs | 55 +- .../tests/connectors/fixtures/redshift/mod.rs | 4 +- .../fixtures/redshift/redshift_mock/ddl.rs | 77 +++ .../fixtures/redshift/redshift_mock/dml.rs | 209 +++++++ .../fixtures/redshift/redshift_mock/dql.rs | 92 +++ .../redshift/redshift_mock/handler.rs | 293 ++++------ .../fixtures/redshift/redshift_mock/load.rs | 363 ++---------- .../fixtures/redshift/redshift_mock/mod.rs | 42 +- .../fixtures/redshift/redshift_mock/parser.rs | 549 +++++++++++++++++ .../fixtures/redshift/redshift_mock/util.rs | 157 +++++ .../connectors/fixtures/redshift/sink.rs | 30 +- .../connectors/redshift/redshift_sink.rs | 64 +- 19 files changed, 1915 insertions(+), 765 deletions(-) create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/ddl.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/dml.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/dql.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/parser.rs create mode 100644 core/integration/tests/connectors/fixtures/redshift/redshift_mock/util.rs diff --git a/Cargo.lock b/Cargo.lock index 86d1432cb9..f1676b150b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7540,6 +7540,7 @@ dependencies = [ "bytemuck", "bytes", "cfg_aliases", + "chrono", "compio", "configs", "configs_derive", @@ -7578,7 +7579,7 @@ dependencies = [ "serde_json", "serial_test", "server", - "socket2 0.6.4", + "socket2 0.6.5", "sqlparser 0.62.0", "sqlx", "sysinfo 0.39.6", @@ -7938,7 +7939,7 @@ dependencies = [ "proc-macro2", "quote", "regex", - "syn 2.0.118", + "syn 2.0.119", ] [[package]] @@ -9885,7 +9886,7 @@ dependencies = [ "md5", "pg_interval", "postgres-types", - "rand 0.10.1", + "rand 0.10.2", "rust_decimal", "rustls-pki-types", "ryu", @@ -9893,7 +9894,7 @@ dependencies = [ "serde_json", "smol_str", "stringprep", - "thiserror 2.0.18", + "thiserror 2.0.19", "tokio", "tokio-rustls", "tokio-util", @@ -10128,7 +10129,7 @@ dependencies = [ "hmac 0.13.0", "md-5 0.11.0", "memchr", - "rand 0.10.1", + "rand 0.10.2", "sha2 0.11.0", "stringprep", ] @@ -12649,6 +12650,18 @@ checksum = "13c6d1b651dc4edf07eead2a0c6c78016ce971bc2c10da5266861b13f25e7cec" dependencies = [ "log", "recursive", + "sqlparser_derive", +] + +[[package]] +name = "sqlparser_derive" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -13545,8 +13558,8 @@ dependencies = [ "pin-project-lite", "postgres-protocol", "postgres-types", - "rand 0.10.1", - "socket2 0.6.4", + "rand 0.10.2", + "socket2 0.6.5", "tokio", "tokio-util", "whoami", @@ -15304,7 +15317,7 @@ dependencies = [ "ring", "signature", "spki", - "thiserror 2.0.18", + "thiserror 2.0.19", "zeroize", ] diff --git a/Cargo.toml b/Cargo.toml index 8ce63f00b4..ef236bfc2c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -305,7 +305,7 @@ simd-json = { version = "0.17.3", features = ["serde_impl"] } slab = "0.4.12" smallvec = "1.15" socket2 = "0.6.5" -sqlparser = "0.62.0" +sqlparser = { version = "0.62.0", features = ["visitor"] } sqlx = { version = "0.9.0", features = [ "runtime-tokio", "tls-rustls", diff --git a/core/connectors/sinks/redshift_sink/README.md b/core/connectors/sinks/redshift_sink/README.md index 01a9742706..5600e2b3ff 100644 --- a/core/connectors/sinks/redshift_sink/README.md +++ b/core/connectors/sinks/redshift_sink/README.md @@ -57,15 +57,16 @@ archive = true | `target_table` | yes | — | Destination Redshift table that batches are copied into. | | `batch_size` | no | `100` | Number of messages buffered per Parquet file / `COPY` operation. | | `max_connections` | no | `5` | Size of the connection pool used against Redshift. | -| `include_metadata` | no | `false` | Stores stream/topic/partition/offset/timestamp/schema fields alongside the payload. | +| `include_metadata` | no | `true` | Stores stream/topic/partition/offset/timestamp/schema fields alongside the payload. | | `include_checksum` | no | `false` | Stores the Iggy message checksum. | | `include_origin_timestamp` | no | `false` | Stores the original Iggy origin timestamp. | | `payload_format` | no | `varbyte` | Encoding used for the payload column in the Parquet file. See **Payload Format** below. | | `verbose_logging` | no | `false` | Enables verbose logging for debugging purposes. | | `max_retries` | no | `3` | Maximum number of retries for failed `COPY` operations. `0` disables retries (only one attempt will be made) | | `retry_delay` | no | `1s` | Delay in seconds between retry attempts. | -| `aws_access_key_id` | yes | — | AWS access key used for S3 staging. | -| `aws_secret_access_key` | yes | — | AWS secret key used for S3 staging. | +| `aws_iam_role` | yes | — | AWS IAM role with S3-Redshift write privileges used for S3 staging. | +| `aws_access_key_id` | no | — | AWS access key used for S3 staging. | +| `aws_secret_access_key` | no | — | AWS secret key used for S3 staging. | | `s3_bucket` | yes | — | S3 bucket that Parquet batch files are staged into before the Redshift `COPY`. | | `s3_prefix` | yes | — | Key prefix under which staged Parquet files are written, e.g. `iggy/messages`. | | `s3_endpoint` | no | — | Override endpoint for S3-compatible stores (e.g. MinIO). Omit for AWS S3 itself. | @@ -117,3 +118,104 @@ With metadata enabled, records contain: The `messages_processed` counter reports valid records submitted to Redshift via `COPY`. + +## Test Suite Setup + +Six queries validate connector behavior end-to-end. Each is shown in its **production (Redshift)** form; where the pgwire-postgres test harness diverges, the substitution is noted inline. + +### 1. Connection check + +```sql +SELECT 1 +``` + +Confirms warehouse connectivity. No dialect differences. + +## 2. Staging/target table creation + +```sql +CREATE TABLE IF NOT EXISTS {table_name} ( + id VARCHAR(40), + iggy_offset VARCHAR(20), + iggy_timestamp VARCHAR(20), + iggy_stream TEXT, + iggy_topic TEXT, + iggy_partition_id BIGINT, + iggy_checksum VARCHAR, + iggy_origin_timestamp VARCHAR(20), + payload {payload_type}, + created_at TIMESTAMPTZ DEFAULT GETDATE() +); +``` + +- Staging table name = `staging_` + `{table_name}`. +- **pgwire test substitution:** `GETDATE()` → `NOW()`. +- **pgwire test substitution:** `VARBYTE` → `BYTEA`. This affects the `column` when we have `VARBYTE` as the type. +- `iggy_offset`, `iggy_timestamp`, and `iggy_origin_timestamp` are u64 values in Iggy but are stored as `VARCHAR` rather than `BIGINT`. `BIGINT` is signed and tops out below `u64::MAX`, so a `VARCHAR` column sidesteps the overflow risk on the upper half of the u64 range without pulling in `DECIMAL`'s added precision/rounding handling. +- `iggy_partition_id` is u32 in Iggy but is stored as `BIGINT` rather than `INTEGER`. `INTEGER` is signed and tops out below `u32::MAX`, so a `BIGINT` column sidesteps the overflow risk. + +## 3. Schema drift check + +**Redshift:** + +```sql +SELECT "column", type +FROM pg_table_def +WHERE tablename = 'target_table'; +``` + +**pgwire test equivalent:** + +```sql +SELECT column_name, type +FROM information_schema.columns +WHERE table_name = 'target_table'; +``` + +Substitutions: `pg_table_def` → `information_schema.columns`, `"column"` → `column_name`, `udt_name`* → `type`. + +## 4. S3 → staging load + +**Redshift:** + +```sql +COPY {staging_table} ({columns}) +FROM '{s3_path}' +CREDENTIALS 'AWS_IAM_ROLE={iam_role}' +FORMAT AS PARQUET +REGION '{region}'; +``` + +**pgwire test equivalent:** + +```sql +COPY {staging_table} ({columns}) +FROM STDIN BINARY +``` + +The `s3_path` is parsed and used to fetch the object from the MinIO instance backing the mock container, with access key and secret key supplied to the container via environment variables rather than an IAM role. Instead of Redshift pulling directly from S3, the connector reads the object itself and streams it into the mock over `COPY ... FROM STDIN BINARY`, so the `CREDENTIALS`, `FORMAT AS PARQUET`, and `REGION` clauses have no equivalent here. + +## 5. Staging → target merge (idempotent upsert) + +```sql +MERGE INTO "target_table" AS t +USING staging_target_table AS sm +ON t.id = sm.id +WHEN NOT MATCHED THEN INSERT ( + id, iggy_offset, iggy_timestamp, iggy_stream, iggy_topic, + iggy_partition_id, iggy_checksum, iggy_origin_timestamp, payload, created_at +) VALUES ( + sm.id, sm.iggy_offset, sm.iggy_timestamp, sm.iggy_stream, sm.iggy_topic, + sm.iggy_partition_id, sm.iggy_checksum, sm.iggy_origin_timestamp, sm.payload, sm.created_at +); +``` + +Uniqueness enforced on `id` — no update branch by design (insert-only merge). No dialect differences. + +## 6. Staging table reset + +```sql +TRUNCATE staging_target_table; +``` + +Clears staging ahead of the next load cycle. No dialect differences. diff --git a/core/connectors/sinks/redshift_sink/src/config.rs b/core/connectors/sinks/redshift_sink/src/config.rs index b0a4b30ca9..bc0627fc5d 100644 --- a/core/connectors/sinks/redshift_sink/src/config.rs +++ b/core/connectors/sinks/redshift_sink/src/config.rs @@ -35,10 +35,11 @@ pub struct RedshiftSinkConfig { pub max_retries: Option, pub retry_delay: Option, /// aws_access_key_id and aws_secret_access_key MUST be provided - #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] - pub aws_access_key_id: SecretString, - #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] - pub aws_secret_access_key: SecretString, + #[serde(serialize_with = "iggy_common::serde_secret::serialize_optional_secret")] + pub aws_access_key_id: Option, + #[serde(serialize_with = "iggy_common::serde_secret::serialize_optional_secret")] + pub aws_secret_access_key: Option, + pub aws_iam_role: String, pub s3_bucket: String, pub s3_prefix: String, pub s3_endpoint: Option, @@ -69,13 +70,20 @@ impl RedshiftSinkConfig { errors.push_str(", aws_region is empty\n"); } - // Validate AWS credentials: access keys must be provided - let has_access_key = !self.aws_access_key_id.expose_secret().is_empty(); + if self.aws_iam_role.is_empty() { + errors.push_str(", aws_iam_role is empty\n"); + } + + if let (Some(access), Some(secret)) = (&self.aws_access_key_id, &self.aws_secret_access_key) + { + // Validate AWS credentials: access keys must be provided + let has_access_key = !access.expose_secret().is_empty(); - let has_secret_key = !self.aws_secret_access_key.expose_secret().is_empty(); + let has_secret_key = !secret.expose_secret().is_empty(); - if !(has_access_key && has_secret_key) { - errors.push_str(", aws_access_key_id and aws_secret_access_key are empty\n"); + if !(has_access_key && has_secret_key) { + errors.push_str(", aws_access_key_id and aws_secret_access_key are empty\n"); + } } if !errors.is_empty() { @@ -92,11 +100,9 @@ impl RedshiftSinkConfig { /// /// We dont have Json because we are using parquet as a means to sink ingestion /// As at the development of this connector there's no direct parquet type that matches JSON -/// For JSON needs Reshshift has SUPER(VARCHAR can be parsed by JSON_PARSE) -#[allow(unused)] +/// For JSON needs Redshift has SUPER(VARCHAR can be parsed by JSON_PARSE)ß #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum PayloadFormat { - Json, Text, #[default] Varbyte, @@ -105,15 +111,29 @@ pub enum PayloadFormat { impl PayloadFormat { pub fn from_config(s: Option<&str>) -> Self { match s.map(|s| s.to_lowercase()).as_deref() { - Some("text") | Some("json") => PayloadFormat::Text, - _ => PayloadFormat::Varbyte, + Some("text") => PayloadFormat::Text, + Some("json") => { + tracing::warn!("Json is not supported, falling back to Text"); + PayloadFormat::Text + } + + other => { + if other.is_some() { + tracing::warn!( + "Unrecognized payload_format {:?}, falling back to VARBYTE", + other + ); + } + + PayloadFormat::Varbyte + } } } pub fn sql_type(&self) -> &'static str { match self { PayloadFormat::Varbyte => "VARBYTE", - PayloadFormat::Text | PayloadFormat::Json => "VARCHAR", + PayloadFormat::Text => "VARCHAR", } } @@ -121,7 +141,6 @@ impl PayloadFormat { match self { PayloadFormat::Varbyte => DataType::Binary, PayloadFormat::Text => DataType::Utf8, - PayloadFormat::Json => DataType::Utf8, } } } diff --git a/core/connectors/sinks/redshift_sink/src/lib.rs b/core/connectors/sinks/redshift_sink/src/lib.rs index 211e354c20..1bb2869f0b 100644 --- a/core/connectors/sinks/redshift_sink/src/lib.rs +++ b/core/connectors/sinks/redshift_sink/src/lib.rs @@ -17,24 +17,25 @@ mod config; -use std::{str::FromStr, sync::Arc, time::Duration}; +use std::{collections::BTreeMap, str::FromStr, sync::Arc, time::Duration}; use arrow::{ - array::{ - ArrayRef, BinaryArray, Decimal256Array, Int32Array, Int64Array, RecordBatch, StringArray, - TimestampMicrosecondArray, - }, - datatypes::{DataType, Field, Schema, TimeUnit}, + array::{ArrayRef, BinaryBuilder, Int64Array, RecordBatch, StringArray, StringBuilder}, + datatypes::{DataType, Field, Schema}, }; use async_trait::async_trait; use humantime::Duration as HumanDuration; use iggy_connector_sdk::{ ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata, sink_connector, }; -use parquet::arrow::ArrowWriter; +use parquet::{ + arrow::ArrowWriter, + basic::{Compression, ZstdLevel}, + file::properties::WriterProperties, +}; use s3::{Bucket, Region, creds::Credentials}; use secrecy::ExposeSecret; -use sqlx::{AssertSqlSafe, Pool, Postgres, postgres::PgPoolOptions}; +use sqlx::{AssertSqlSafe, Pool, Postgres, Row, postgres::PgPoolOptions}; use tokio::sync::Mutex; use uuid::Uuid; @@ -66,7 +67,10 @@ impl Sink for RedshiftSink { ); self.connect().await?; - self.ensure_table_exists().await?; + // Ensuring tables exist + self.ensure_tables_exist().await?; + // Checking for schema drift + self.ensure_schema_match().await?; Ok(()) } @@ -149,8 +153,14 @@ impl RedshiftSink { let region = self.build_region()?; let credentials = Credentials::new( - Some(self.config.aws_access_key_id.expose_secret()), - Some(self.config.aws_secret_access_key.expose_secret()), + self.config + .aws_access_key_id + .as_ref() + .map(|v| v.expose_secret()), + self.config + .aws_secret_access_key + .as_ref() + .map(|v| v.expose_secret()), None, None, None, @@ -188,29 +198,172 @@ impl RedshiftSink { } } - async fn ensure_table_exists(&self) -> Result<(), Error> { + async fn ensure_tables_exist(&self) -> Result<(), Error> { let pool = self.get_pool()?; - let table_name = &self.config.target_table; + let target_table = quote_identifier(&self.config.target_table)?; + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + let payload_type = self.payload_format().sql_type(); - let (query, _) = self.build_create_table_sql()?; + let target_query = self.build_create_table_sql(&target_table)?; + + let staging_query = self.build_create_table_sql(&staging_table)?; + + tracing::debug!("ensuring staging and target tables exist"); + + sqlx::query(AssertSqlSafe(staging_query)) + .execute(pool) + .await + .map_err(|e| { + tracing::error!(error = %e); + Error::InitError(format!("Failed to create table '{staging_table}': {e}")) + })?; - tracing::debug!("ensuring target table exists"); + tracing::debug!("Staging table created"); - sqlx::query(AssertSqlSafe(query)) + sqlx::query(AssertSqlSafe(target_query)) .execute(pool) .await .map_err(|e| { tracing::error!(error = %e); - Error::InitError(format!("Failed to create table '{table_name}': {e}")) + Error::InitError(format!("Failed to create table '{target_table}': {e}")) })?; - tracing::info!(table = %table_name, payload_type, "target table ready"); + tracing::info!( + staging_table = staging_table, + target_table = target_table, + payload_type, + "staging and target tables ready" + ); + + Ok(()) + } + + // This method ensures that the target table schema matches the expected schema. + // it also verifies there is a created_at column + async fn ensure_schema_match(&self) -> Result<(), Error> { + let include_metadata = self.config.include_metadata.unwrap_or(true); + let include_checksum = self.config.include_checksum.unwrap_or(true); + let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); + let target_table = quote_identifier(&self.config.target_table)?; + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + let payload_type = self.payload_format().sql_type(); + let pool = self.get_pool()?; + + let mut expected_cols: BTreeMap<&str, &str> = BTreeMap::new(); + expected_cols.insert("id", "VARCHAR"); + if include_metadata { + expected_cols.insert("iggy_offset", "VARCHAR"); + expected_cols.insert("iggy_timestamp", "VARCHAR"); + expected_cols.insert("iggy_stream", "TEXT"); + expected_cols.insert("iggy_topic", "TEXT"); + expected_cols.insert("iggy_partition_id", "BIGINT"); + } + if include_checksum { + expected_cols.insert("iggy_checksum", "VARCHAR"); + } + if include_origin_timestamp { + expected_cols.insert("iggy_origin_timestamp", "VARCHAR"); + } + expected_cols.insert("payload", payload_type); + expected_cols.insert("created_at", "TIMESTAMPTZ"); + + let target_cols = Self::load_columns(pool, &target_table).await?; + let staging_cols = Self::load_columns(pool, &staging_table).await?; + + let mut mismatches = Self::diff_schema(&target_table, &target_cols, &expected_cols); + mismatches.extend(Self::diff_schema( + &staging_table, + &staging_cols, + &expected_cols, + )); + + tracing::info!("Mismatches: {:?}", mismatches); + + if !mismatches.is_empty() { + return Err(Error::InitError(format!( + "Schema mismatch detected:\n{}", + mismatches.join("\n") + ))); + } Ok(()) } + fn diff_schema( + table_name: &str, + actual_cols: &BTreeMap, + expected_cols: &BTreeMap<&str, &str>, + ) -> Vec { + let mut errors = Vec::new(); + + for (col_name, expected_type) in expected_cols { + match actual_cols.get(*col_name) { + None => errors.push(format!( + "{table_name}: missing column '{col_name}' (expected {expected_type})" + )), + Some(actual_type) if !Self::type_matches(actual_type, expected_type) => errors.push(format!( + "{table_name}: column '{col_name}' type mismatch — expected {expected_type}, found {actual_type}" + )), + _ => {} + } + } + errors + } + + fn type_matches(actual: &str, expected: &str) -> bool { + Self::normalize_type(actual) == Self::normalize_type(expected) + } + + async fn load_columns( + pool: &sqlx::PgPool, + table: &str, + ) -> Result, Error> { + let query = format!( + "SELECT \"column\", type FROM pg_table_def WHERE tablename = '{}'", + table.replace('"', "") + ); + + let rows = sqlx::query(AssertSqlSafe(query)) + .fetch_all(pool) + .await + .map_err(|e| Error::InitError(format!("Failed to read schema for '{table}': {e}")))?; + + if rows.is_empty() { + return Err(Error::InitError(format!( + "Table '{table}' was not found or has no visible columns" + ))); + } + + rows.into_iter() + .map(|row| { + Ok(( + row.try_get::("column") + .map_err(|e| Error::InitError(e.to_string()))?, + Self::normalize_type( + &row.try_get::("type") + .map_err(|e| Error::InitError(e.to_string()))?, + ) + .to_string(), + )) + }) + .collect() + } + + fn normalize_type(t: &str) -> &'static str { + match t.to_ascii_uppercase().as_str() { + "INTEGER" | "INT" | "INT4" => "INTEGER", + "INT8" | "BIGINT" => "BIGINT", + "VARCHAR" | "CHARACTER VARYING" => "VARCHAR", + // Having bytea because of the Postgres Test + "BYTEA" | "VARBYTE" | "VARBINARY" | "BINARY VARYING" => "VARBYTE", + "TEXT" => "TEXT", + "TIMESTAMPTZ" | "TIMESTAMP WITH TIME ZONE" => "TIMESTAMPTZ", + _ => "UNKNOWN", + } + } + async fn process_messages( &self, topic_metadata: &TopicMetadata, @@ -224,12 +377,28 @@ impl RedshiftSink { .insert_batch(batch, topic_metadata, messages_metadata) .await { - Ok(()) => { - self.state.lock().await.batches_loaded += 1; + Ok(path) => { + // Messages were received and ingested to Redshift + if let Some(s3_path) = path { + // Truncate the staging table + if let Err(e) = self.staging_cleanup().await { + tracing::warn!(error = %e, "failed to cleanup staging table"); + } + + // Handle archiving + if let Err(e) = self.archive_parquet(&s3_path).await { + tracing::warn!(error = %e, "failed to archive parquet file: {}", s3_path); + } + + self.state.lock().await.batches_loaded += 1 + } else { + tracing::info!("Zero messages found for processing"); + } } Err(e) => { self.state.lock().await.insertion_errors += batch.len() as u64; tracing::error!(error = %e, batch_size = batch.len(), "failed to insert batch"); + return Err(e); } } } @@ -258,14 +427,21 @@ impl RedshiftSink { Ok(()) } + // This function builds a parquet from messages and metadata + // It uploads the parquet to S3 and returns the path + // It then copies the parquet to Redshift target table via + // a staging table by means of a MERGE statement + // This function treats parquet-generation, uploading to s3, + // copying to staging and merging to target as atomic + // process of focus for this sink connector async fn insert_batch( &self, messages: &[ConsumedMessage], topic_metadata: &TopicMetadata, messages_metadata: &MessagesMetadata, - ) -> Result<(), Error> { + ) -> Result, Error> { if messages.is_empty() { - return Ok(()); + return Ok(None); } let include_metadata = self.config.include_metadata.unwrap_or(true); @@ -292,24 +468,80 @@ impl RedshiftSink { ); let s3_path = self.upload_parquet(&content).await?; - self.copy_parquet(&s3_path).await?; - self.archive_parquet(&s3_path).await?; + + let schema = record_batch.schema(); + let cols = schema + .fields + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + + // Copy the parquet file to Redshift staging + // Cleanup + tracing::info!("copying parquet to Redshift staging"); + if let Err(e) = self.copy_parquet(&s3_path, &cols).await { + let key = s3_path + .strip_prefix(&format!("s3://{}/", self.config.s3_bucket)) + .ok_or(Error::InvalidConfigValue("Missing Cleanup S3 path".into()))?; + + self.delete_object(key).await?; + + Err(e)? + } + + tracing::info!("Redshift stging COPY completed"); + + // Do a merge into Redshift target table + self.merge_into_target(&cols).await?; + + tracing::info!("Redshift target table merge completed"); tracing::info!(count = messages.len(), path = %s3_path, "batch inserted into Redshift"); - Ok(()) + Ok(Some(s3_path)) } - async fn copy_parquet(&self, s3_path: &str) -> Result<(), Error> { + async fn copy_parquet(&self, s3_path: &str, cols: &[&str]) -> Result<(), Error> { let max_retries = self.get_max_retries(); let retry_delay = self.get_retry_delay(); - let sql = self.build_copy_sql(s3_path); + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + + let sql = self.build_copy_sql(&staging_table, s3_path, &cols.join(", "))?; let pool = self.get_pool()?; tracing::debug!(table = %self.config.target_table, s3_path, "issuing Redshift COPY"); retry_with_backoff( - "redshift COPY", + "Redshift COPY", + max_retries, + retry_delay, + is_transient_error, + || async { + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .map(|_| ()) + }, + ) + .await?; + + tracing::debug!(staging_table = staging_table, "Redshift COPY completed"); + + Ok(()) + } + + async fn merge_into_target(&self, cols: &[&str]) -> Result<(), Error> { + let max_retries = self.get_max_retries(); + let retry_delay = self.get_retry_delay(); + let target_table = quote_identifier(&self.config.target_table)?; + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + let sql = self.build_merge_sql(cols, &staging_table, &target_table); + let pool = self.get_pool()?; + + tracing::debug!(table = %self.config.target_table, "issuing Redshift MERGE"); + + retry_with_backoff( + "Redshift MERGE", max_retries, retry_delay, is_transient_error, @@ -322,62 +554,102 @@ impl RedshiftSink { ) .await?; - tracing::debug!(table = %self.config.target_table, "Redshift COPY completed"); + tracing::debug!(staging_table = %staging_table, target_table = %target_table, "Redshift MERGE completed"); Ok(()) } - fn build_create_table_sql(&self) -> Result<(String, u32), Error> { - let table_name = &self.config.target_table; - let quoted_table = quote_identifier(table_name)?; + async fn staging_cleanup(&self) -> Result<(), Error> { + let max_retries = self.get_max_retries(); + let retry_delay = self.get_retry_delay(); + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + let sql = self.build_truncate_sql(&staging_table); + let pool = self.get_pool()?; + + tracing::debug!(table = %self.config.target_table, "issuing Redshift TRUNCATE"); + + retry_with_backoff( + "Redshift TRUNCATE", + max_retries, + retry_delay, + is_transient_error, + || async { + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .map(|_| ()) + }, + ) + .await?; + + tracing::debug!(table = %self.config.target_table, "Redshift TRUNCATE completed"); + Ok(()) + } + + fn build_create_table_sql(&self, table_name: &str) -> Result { let include_metadata = self.config.include_metadata.unwrap_or(true); let include_checksum = self.config.include_checksum.unwrap_or(true); let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); let payload_type = self.payload_format().sql_type(); - let mut params_per_row: u32 = 1; // id - - let mut query = - format!("CREATE TABLE IF NOT EXISTS {quoted_table} (id DECIMAL(39, 0) PRIMARY KEY"); + let mut query = format!("CREATE TABLE IF NOT EXISTS {table_name} (id VARCHAR(40)"); if include_metadata { - query.push_str(", iggy_offset BIGINT, iggy_timestamp TIMESTAMPTZ, iggy_stream TEXT, iggy_topic TEXT, iggy_partition_id INTEGER"); - params_per_row += 5; + query.push_str(", iggy_offset VARCHAR(20), iggy_timestamp VARCHAR(20), iggy_stream TEXT, iggy_topic TEXT, iggy_partition_id BIGINT"); } if include_checksum { query.push_str(", iggy_checksum VARCHAR"); - params_per_row += 1; } if include_origin_timestamp { - query.push_str(", iggy_origin_timestamp TIMESTAMPTZ"); - params_per_row += 1; + query.push_str(", iggy_origin_timestamp VARCHAR(20)"); } query.push_str(&format!(", payload {payload_type}")); query.push_str(", created_at TIMESTAMPTZ DEFAULT GETDATE());"); - params_per_row += 2; - Ok((query, params_per_row)) + Ok(query) } - fn build_copy_sql(&self, s3_path: &str) -> String { - // Built via format! (not sqlx binds) because the Redshift/Pgwire endpoint here - // uses a Simple Query Handler that doesn't support prepared statements with binds. - let credentials = format!( - "CREDENTIALS 'ACCESS_KEY_ID={};SECRET_ACCESS_KEY={}'", - self.config.aws_access_key_id.expose_secret(), - self.config.aws_secret_access_key.expose_secret() - ); + fn build_copy_sql( + &self, + staging_table: &str, + s3_path: &str, + cols: &str, + ) -> Result { + // Redshift allows this from the docs + // https://docs.aws.amazon.com/redshift/latest/dg/r_COPY_command_examples.html + let iam_role = quote_identifier(&self.config.aws_iam_role)?; + + let region = quote_identifier(&self.config.aws_region)?; + + Ok(format!( + "COPY {} ({}) FROM '{}' CREDENTIALS 'AWS_IAM_ROLE={}' FORMAT AS PARQUET REGION '{}';", + staging_table, cols, s3_path, iam_role, region + )) + } + + fn build_merge_sql(&self, cols: &[&str], staging: &str, target: &str) -> String { + let t_cols = cols.join(", "); + + let s_cols = cols + .iter() + .map(|v| format!("sm.{v}")) + .collect::>() + .join(", "); format!( - "COPY {} FROM '{}' {} FORMAT AS PARQUET REGION '{}';", - self.config.target_table, s3_path, credentials, self.config.aws_region + "MERGE INTO {} AS t USING {} AS sm ON t.id = sm.id WHEN NOT MATCHED THEN INSERT ({}) VALUES ({});", + target, staging, t_cols, s_cols ) } + fn build_truncate_sql(&self, table: &str) -> String { + format!("TRUNCATE {};", table) + } + async fn upload_parquet(&self, content: &[u8]) -> Result { let file_id = Uuid::now_v7(); let key = build_s3_key(&self.config.s3_prefix, &format!("{file_id}.parquet")); @@ -411,7 +683,22 @@ impl RedshiftSink { let bucket = self.get_bucket()?; let prefix = self.config.s3_prefix.trim_matches('/'); - let archived_key = old_key.replacen(prefix, DEFAULT_ARCHIVE_PREFIX.trim_matches('/'), 1); + let archive_prefix = DEFAULT_ARCHIVE_PREFIX.trim_matches('/'); + + let suffix = if prefix.is_empty() { + old_key + } else { + old_key + .strip_prefix(prefix) + .map(|s| s.trim_start_matches('/')) + .unwrap_or(old_key) + }; + + let archived_key = if archive_prefix.is_empty() { + suffix.to_string() + } else { + format!("{}/{}", archive_prefix, suffix) + }; tracing::debug!(from = old_key, to = %archived_key, "archiving parquet file"); @@ -511,7 +798,7 @@ where let transient = is_transient(&e); if !transient || attempts >= max_retries { - tracing::error!(operation, attempts, error = %e, "operation failed permanently"); + tracing::error!(operation = operation, attempts = attempts, error = %e, "operation failed permanently"); return Err(Error::CannotStoreData(format!( "{operation} failed after {attempts} attempts: {e}" ))); @@ -525,11 +812,16 @@ where } fn encode_parquet(batch: &RecordBatch) -> Result, Error> { + let props = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .build(); + let mut content = Vec::new(); - let mut writer = ArrowWriter::try_new(&mut content, batch.schema(), None).map_err(|e| { - tracing::error!(error = %e, "failed to create parquet writer"); - Error::WriteFailure(format!("Failed to create parquet writer: {e}")) - })?; + let mut writer = + ArrowWriter::try_new(&mut content, batch.schema(), Some(props)).map_err(|e| { + tracing::error!(error = %e, "failed to create parquet writer"); + Error::WriteFailure(format!("Failed to create parquet writer: {e}")) + })?; writer.write(batch).map_err(|e| { tracing::error!(error = %e, "failed to write parquet batch"); @@ -575,8 +867,8 @@ fn create_record_batch( include_origin_timestamp: bool, payload_format: PayloadFormat, ) -> Result { - let mut fields = vec![Field::new("id", DataType::Decimal256(39, 0), false)]; - let mut columns: Vec = vec![id_column(messages)?]; + let mut fields = vec![Field::new("id", DataType::Utf8, false)]; + let mut columns: Vec = vec![id_column(messages)]; if include_metadata { let (mut metadata_fields, mut metadata_columns) = @@ -591,11 +883,7 @@ fn create_record_batch( } if include_origin_timestamp { - fields.push(Field::new( - "iggy_origin_timestamp", - DataType::Timestamp(TimeUnit::Microsecond, None), - false, - )); + fields.push(Field::new("iggy_origin_timestamp", DataType::Utf8, false)); columns.push(origin_timestamp_column(messages)); } @@ -615,16 +903,10 @@ fn create_record_batch( Ok(batch) } -fn id_column(messages: &[ConsumedMessage]) -> Result { - let ids = Decimal256Array::from_iter_values( - messages - .iter() - .map(|v| arrow::datatypes::i256::from_parts(v.id, 0)), - ) - .with_precision_and_scale(39, 0) - .map_err(|e| Error::CannotStoreData(e.to_string()))?; - - Ok(Arc::new(ids)) +fn id_column(messages: &[ConsumedMessage]) -> ArrayRef { + Arc::new(StringArray::from_iter_values( + messages.iter().map(|v| v.id.to_string()), + )) } fn metadata_columns( @@ -633,23 +915,19 @@ fn metadata_columns( messages: &[ConsumedMessage], ) -> (Vec, Vec) { let fields = vec![ - Field::new("iggy_offset", DataType::Int64, false), - Field::new( - "iggy_timestamp", - DataType::Timestamp(TimeUnit::Microsecond, None), - false, - ), + Field::new("iggy_offset", DataType::Utf8, false), + Field::new("iggy_timestamp", DataType::Utf8, false), Field::new("iggy_stream", DataType::Utf8, false), Field::new("iggy_topic", DataType::Utf8, false), - Field::new("iggy_partition_id", DataType::Int32, false), + Field::new("iggy_partition_id", DataType::Int64, false), ]; let columns: Vec = vec![ - Arc::new(Int64Array::from_iter_values( - messages.iter().map(|v| v.offset as i64), + Arc::new(StringArray::from_iter_values( + messages.iter().map(|v| v.offset.to_string()), )), - Arc::new(TimestampMicrosecondArray::from_iter_values( - messages.iter().map(|v| v.timestamp as i64), + Arc::new(StringArray::from_iter_values( + messages.iter().map(|v| v.timestamp.to_string()), )), Arc::new(StringArray::from_iter_values( (0..messages.len()).map(|_| topic_metadata.stream.clone()), @@ -657,8 +935,8 @@ fn metadata_columns( Arc::new(StringArray::from_iter_values( (0..messages.len()).map(|_| topic_metadata.topic.clone()), )), - Arc::new(Int32Array::from_iter_values( - (0..messages.len()).map(|_| messages_metadata.partition_id as i32), + Arc::new(Int64Array::from_iter_values( + (0..messages.len()).map(|_| messages_metadata.partition_id as i64), )), ]; @@ -672,43 +950,33 @@ fn checksum_column(messages: &[ConsumedMessage]) -> ArrayRef { } fn origin_timestamp_column(messages: &[ConsumedMessage]) -> ArrayRef { - Arc::new(TimestampMicrosecondArray::from_iter_values( - messages.iter().map(|v| v.origin_timestamp as i64), + Arc::new(StringArray::from_iter_values( + messages.iter().map(|v| v.origin_timestamp.to_string()), )) } fn payload_column(messages: &[ConsumedMessage], format: PayloadFormat) -> Result { match format { PayloadFormat::Varbyte => { - let values: Vec> = messages - .iter() - .map(|v| v.payload.clone().try_to_bytes()) - .collect::>()?; - let slices: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect(); - Ok(Arc::new(BinaryArray::from_vec(slices))) + let mut builder = BinaryBuilder::with_capacity(messages.len(), 0); + + for m in messages { + builder.append_value(m.payload.try_to_bytes()?); + } + + Ok(Arc::new(builder.finish())) } PayloadFormat::Text => { - let values: Vec = messages - .iter() - .map(|v| { - let bytes = v.payload.try_to_bytes()?; - String::from_utf8(bytes).map_err(|_| Error::InvalidTextPayload) - }) - .collect::>()?; - Ok(Arc::new(StringArray::from_iter_values(values.iter()))) - } - PayloadFormat::Json => { - let values: Vec = messages - .iter() - .map(|v| { - let bytes = v.payload.try_to_bytes()?; - - Ok(serde_json::from_slice::(&bytes) - .map_err(|_| Error::InvalidJsonPayload)? - .to_string()) - }) - .collect::>()?; - Ok(Arc::new(StringArray::from_iter_values(values.iter()))) + let mut builder = StringBuilder::with_capacity(messages.len(), 0); + + for m in messages { + let bytes = m.payload.try_to_bytes()?; + let s = std::str::from_utf8(&bytes).map_err(|_| Error::InvalidTextPayload)?; + + builder.append_value(s); + } + + Ok(Arc::new(builder.finish())) } } } @@ -721,11 +989,16 @@ fn redact_connection_string(conn_str: &str) -> String { let scheme = &conn_str[..scheme_end + 3]; let rest = &conn_str[scheme_end + 3..]; + let bound_end = rest.find([':', '@', '?', '/']).unwrap_or(rest.len()); + // Stop preview at the first sensitive boundary let safe_end = rest - .find([':', '@', '?', '/']) - .unwrap_or(rest.len()) - .min(PREVIEW_LEN); + .char_indices() + .map(|(i, _)| i) + .chain(std::iter::once(rest.len())) + .take_while(|&i| i <= bound_end) + .nth(PREVIEW_LEN) + .unwrap_or(bound_end); let preview = &rest[..safe_end]; return format!("{scheme}{preview}***"); @@ -790,8 +1063,9 @@ mod tests { verbose_logging: None, max_retries: None, retry_delay: None, - aws_access_key_id: SecretString::from("admin"), - aws_secret_access_key: SecretString::from("password"), + aws_access_key_id: Some(SecretString::from("admin")), + aws_secret_access_key: Some(SecretString::from("password")), + aws_iam_role: "arn:aws:iam::123456789012:role/Iggy".into(), s3_bucket: "iggymessages".into(), s3_prefix: "iggy/messages".into(), s3_endpoint: None, @@ -867,7 +1141,7 @@ mod tests { #[test] fn given_empty_aws_access_key_id_should_error() { let mut config = test_config(false, false, false); - config.aws_access_key_id = SecretString::default(); + config.aws_access_key_id = Some(SecretString::default()); assert!(config.validate().is_err()); } @@ -875,7 +1149,7 @@ mod tests { #[test] fn given_empty_aws_secret_access_key_should_error() { let mut config = test_config(false, false, false); - config.aws_secret_access_key = SecretString::default(); + config.aws_secret_access_key = Some(SecretString::default()); assert!(config.validate().is_err()); } @@ -920,7 +1194,6 @@ mod tests { #[test] fn given_payload_format_should_return_correct_sql_type() { assert_eq!(PayloadFormat::Varbyte.sql_type(), "VARBYTE"); - assert_eq!(PayloadFormat::Json.sql_type(), "VARCHAR"); assert_eq!(PayloadFormat::Text.sql_type(), "VARCHAR"); } @@ -930,10 +1203,6 @@ mod tests { PayloadFormat::Varbyte.arrow_type(), arrow::datatypes::DataType::Binary ); - assert_eq!( - PayloadFormat::Json.arrow_type(), - arrow::datatypes::DataType::Utf8 - ); assert_eq!( PayloadFormat::Text.arrow_type(), arrow::datatypes::DataType::Utf8 @@ -943,8 +1212,12 @@ mod tests { #[test] fn given_all_options_enabled_should_build_full_create_query() { let sink = RedshiftSink::new(1, test_config(true, true, true)); - let (query, param_count) = sink - .build_create_table_sql() + + let target_table = + quote_identifier(&sink.config.target_table).expect("Failed to quote table identifier"); + + let query = sink + .build_create_table_sql(&target_table) .expect("Failed to build create query"); assert!(query.contains("CREATE TABLE IF NOT EXISTS \"messages\"")); @@ -957,7 +1230,6 @@ mod tests { assert!(query.contains("iggy_origin_timestamp")); assert!(query.contains("payload")); assert!(query.contains("created_at")); - assert_eq!(param_count, 10); } #[test] @@ -1000,8 +1272,11 @@ mod tests { #[test] fn given_metadata_disabled_should_build_minimal_create_query() { let sink = RedshiftSink::new(1, test_config(false, false, false)); - let (query, param_count) = sink - .build_create_table_sql() + + let target_table = + quote_identifier(&sink.config.target_table).expect("Failed to quote table identifier"); + let query = sink + .build_create_table_sql(&target_table) .expect("Failed to build create query"); assert!(query.contains("CREATE TABLE IF NOT EXISTS \"messages\"")); @@ -1014,7 +1289,6 @@ mod tests { assert!(!query.contains("iggy_origin_timestamp")); assert!(query.contains("payload")); assert!(query.contains("created_at")); - assert_eq!(param_count, 3); } #[test] @@ -1059,12 +1333,12 @@ mod tests { let timestamp_col = record_batch .column(2) .as_any() - .downcast_ref::() + .downcast_ref::() .expect("Failed to downcast to Timestamp Microsecond array"); let timestamp = timestamp_col.value(0); - assert_eq!(timestamp, 1_767_225_600_000_000); + assert_eq!(timestamp, "1767225600000000"); } #[test] diff --git a/core/integration/Cargo.toml b/core/integration/Cargo.toml index 77b0acb2bf..3a85893efd 100644 --- a/core/integration/Cargo.toml +++ b/core/integration/Cargo.toml @@ -44,6 +44,7 @@ base64 = { workspace = true } bon = { workspace = true } bytemuck = { workspace = true } bytes = { workspace = true } +chrono = { workspace = true } compio = { workspace = true } configs = { workspace = true } configs_derive = { workspace = true } diff --git a/core/integration/tests/connectors/fixtures/mod.rs b/core/integration/tests/connectors/fixtures/mod.rs index 0b03d957eb..dc208a858c 100644 --- a/core/integration/tests/connectors/fixtures/mod.rs +++ b/core/integration/tests/connectors/fixtures/mod.rs @@ -83,8 +83,8 @@ pub use postgres::{ }; pub use quickwit::{QuickwitFixture, QuickwitOps, QuickwitPreCreatedFixture}; pub use redshift::{ - RedshiftSinkByteaFixture, RedshiftSinkFixture, RedshiftSinkJsonFixture, - RedshiftSinkNoArchiveFixture, + RedshiftSinkFixture, RedshiftSinkJsonFixture, RedshiftSinkNoArchiveFixture, + RedshiftSinkVarbyteFixture, }; pub use s3::{S3SinkFixture, S3SinkOps, S3SinkRotationFixture}; pub use surrealdb::{ diff --git a/core/integration/tests/connectors/fixtures/redshift/container.rs b/core/integration/tests/connectors/fixtures/redshift/container.rs index fb696e8135..962c2378c0 100644 --- a/core/integration/tests/connectors/fixtures/redshift/container.rs +++ b/core/integration/tests/connectors/fixtures/redshift/container.rs @@ -25,19 +25,23 @@ use testcontainers::{ core::{IntoContainerPort, WaitFor, wait::HttpWaitStrategy}, runners::AsyncRunner, }; -use testcontainers_modules::postgres; use tokio::{net::TcpListener, task::JoinHandle}; use crate::connectors::fixtures::{ self, - redshift::redshift_mock::handler::{RedshiftHandler, RedshiftHandlerFactory}, + redshift::redshift_mock::{handler::RedshiftHandlerFactory, load::S3Client}, }; +const POSTGRES_IMAGE: &str = "postgres"; +const POSTGRES_TAG: &str = "15-alpine"; +const POSTGRES_PORT: u16 = 5432; +const POSTGRES_DB: &str = "postgres"; +const POSTGRES_USER: &str = "postgres"; +const POSTGRES_PASSWORD: &str = "postgres"; const MINIO_IMAGE: &str = "docker.io/minio/minio"; const MINIO_TAG: &str = "RELEASE.2025-09-07T16-13-09Z"; const MINIO_PORT: u16 = 9000; const MINIO_CONSOLE_PORT: u16 = 9001; -const POSTGRES_PORT: u16 = 5432; pub const MINIO_ACCESS_KEY: &str = "admin"; pub const MINIO_SECRET_KEY: &str = "password"; @@ -45,12 +49,14 @@ pub const MINIO_BUCKET: &str = "iggystaging"; pub const DEFAULT_SINK_TABLE: &str = "iggy_messages"; pub const STAGING_REGION: &str = "us-east-1"; pub const STAGING_PREFIX: &str = "iggy/messages"; +pub const AWS_IAM_ROLE: &str = "arn:aws:iam::0123456789012:role/iggyRole"; pub const ENV_SINK_CONNECTION_STRING: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_CONNECTION_STRING"; pub const ENV_SINK_TARGET_TABLE: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_TARGET_TABLE"; pub const ENV_SINK_PAYLOAD_FORMAT: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_PAYLOAD_FORMAT"; +pub const ENV_SINK_AWS_IAM_ROLE: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_AWS_IAM_ROLE"; pub const ENV_SINK_STAGING_ACCESS_KEY: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_AWS_ACCESS_KEY_ID"; pub const ENV_SINK_STAGING_SECRET: &str = @@ -130,13 +136,20 @@ impl MinioContainer { /// Base container management for PostgreSQL fixtures. pub struct PostgresContainer { #[allow(dead_code)] - container: ContainerAsync, + container: ContainerAsync, pub connection_string: String, } impl PostgresContainer { pub async fn start() -> Result { - let container = postgres::Postgres::default() + let container = GenericImage::new(POSTGRES_IMAGE, POSTGRES_TAG) + .with_exposed_port(POSTGRES_PORT.tcp()) + .with_wait_for(WaitFor::message_on_stdout( + "database system is ready to accept connections", + )) + .with_env_var("POSTGRES_DB", POSTGRES_DB) + .with_env_var("POSTGRES_USER", POSTGRES_USER) + .with_env_var("POSTGRES_PASSWORD", POSTGRES_PASSWORD) .with_container_name(fixtures::unique_container_name("postgres")) .start() .await @@ -184,24 +197,22 @@ impl RedshiftContainer { target_connection: String, s3_endpoint: String, ) -> Result { - let (pg_client, connection) = - tokio_postgres::connect(&target_connection, tokio_postgres::NoTls) - .await - .map_err(|e| TestBinaryError::FixtureSetup { - fixture_type: "RedshiftContainer".into(), - message: e.to_string(), - })?; - - tokio::spawn(async move { - if let Err(e) = connection.await { - panic!("{}", e.to_string()) - } - }); - - let redshql = RedshiftHandler::new(pg_client, s3_endpoint); + let s3_client = S3Client::new( + MINIO_BUCKET, + &s3_endpoint, + MINIO_ACCESS_KEY, + MINIO_SECRET_KEY, + STAGING_REGION, + ) + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftContainer".to_string(), + message: format!("failed to create S3 client: {e}"), + })?; let factory = Arc::new(RedshiftHandlerFactory { - handler: Arc::new(redshql), + pg_dsn: target_connection, + s3_client: Arc::new(s3_client), }); let listener = @@ -252,7 +263,7 @@ impl RedshiftContainer { #[derive(Debug, Clone, Copy, Default)] pub enum SinkPayloadFormat { #[default] - Bytea, + Varbyte, Text, } diff --git a/core/integration/tests/connectors/fixtures/redshift/mod.rs b/core/integration/tests/connectors/fixtures/redshift/mod.rs index 4599ee0d0c..0bfaa2f485 100644 --- a/core/integration/tests/connectors/fixtures/redshift/mod.rs +++ b/core/integration/tests/connectors/fixtures/redshift/mod.rs @@ -21,6 +21,6 @@ mod sink; pub use container::{MinioContainer, PostgresContainer, RedshiftContainer}; pub use sink::{ - RedshiftSinkByteaFixture, RedshiftSinkFixture, RedshiftSinkJsonFixture, - RedshiftSinkNoArchiveFixture, + RedshiftSinkFixture, RedshiftSinkJsonFixture, RedshiftSinkNoArchiveFixture, + RedshiftSinkVarbyteFixture, }; diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/ddl.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/ddl.rs new file mode 100644 index 0000000000..a2bc75749c --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/ddl.rs @@ -0,0 +1,77 @@ +// 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 pgwire::error::PgWireResult; +use sqlparser::ast::Statement as SqlStatement; +use tokio_postgres::Client; + +use super::{parser::ParsedStatement, util::backend_err}; + +pub async fn execute_create(client: &Client, raw_sql: &str) -> PgWireResult { + // Serialize concurrent DDL on the same backend session pool using a + // Postgres advisory lock, keyed by a hash of statement text. This does + // NOT protect against DDL issued from other proxies/paths outside this + // service — pair it with `lock_timeout`/`statement_timeout` GUCs set on + // the pooled connection so a stuck CREATE can't wedge the pool. + let lock_key = ddl_lock_key(raw_sql); + client + .execute("SELECT pg_advisory_lock($1)", &[&lock_key]) + .await + .map_err(backend_err)?; + + tracing::debug!(sql = raw_sql, "executing DDL: CREATE"); + let result = client.execute(raw_sql, &[]).await.map_err(|e| { + tracing::error!("{}", e); + e + }); + + client + .execute("SELECT pg_advisory_unlock($1)", &[&lock_key]) + .await + .map_err(backend_err)?; + + result.map_err(backend_err) +} + +pub async fn execute_truncate(client: &Client, stmt: &ParsedStatement) -> PgWireResult { + let table_names = truncate_targets(&stmt.ast); + + // Even when allowed, log loudly before it happens — this is the one + // statement class where "log after success" is useless (there's nothing + // to roll back to reconstruct intent from). + tracing::info!(tables = ?table_names, sql = stmt.raw_sql, "executing TRUNCATE"); + + client + .execute(&stmt.raw_sql, &[]) + .await + .map_err(backend_err) +} + +fn truncate_targets(ast: &SqlStatement) -> Vec { + if let SqlStatement::Truncate(trunc) = ast { + trunc.table_names.iter().map(|t| t.to_string()).collect() + } else { + vec![] + } +} + +fn ddl_lock_key(sql: &str) -> i64 { + use std::hash::{Hash, Hasher}; + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + sql.hash(&mut hasher); + hasher.finish() as i64 +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/dml.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/dml.rs new file mode 100644 index 0000000000..804b46908e --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/dml.rs @@ -0,0 +1,209 @@ +// 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 std::{collections::HashSet, sync::Arc}; + +use bytes::Bytes; +use futures::stream; +use pgwire::{api::results::Response, error::PgWireResult}; +use sqlparser::ast::{CopyLegacyOption, CopySource, Statement as SqlStatement}; +use tokio_postgres::Client; + +use crate::connectors::fixtures::redshift::redshift_mock::util::backend_err; + +use super::{ + handler::ExecCtx, + load::{S3Client, fetch_table_columns, infer_parquet_schema, load_one_object}, + util::{columns_to_field_info, decode_param, row_to_data_row, user_err}, +}; + +/// Shared by INSERT and MERGE +pub async fn execute_dml<'a>(client: &Client, ctx: ExecCtx<'a>) -> PgWireResult { + let portal = ctx.portal().ok_or(user_err("Missing portal"))?; + + let raw_sql = &portal.statement.statement.raw_sql; + + let prepared = client + .prepare(raw_sql) + .await + .map_err(|e| pgwire::error::PgWireError::ApiError(Box::new(e)))?; + + let param_types = prepared.params(); + let mut bound_params: Vec> = + Vec::with_capacity(param_types.len()); + + for (i, ty) in param_types.iter().enumerate() { + bound_params.push(decode_param(portal, ty, i)?); + } + + let param_refs: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = + bound_params.iter().map(|b| b.as_ref() as &_).collect(); + + let has_returning = portal + .statement + .statement + .raw_sql + .to_ascii_uppercase() + .contains("RETURNING"); + + if has_returning { + let rows = client + .query(&prepared, ¶m_refs) + .await + .map_err(backend_err)?; + + let fields = Arc::new(columns_to_field_info(prepared.columns())); + + let fields_c = fields.clone(); + + let data_rows = stream::iter( + rows.into_iter() + .map(move |r| row_to_data_row(&r, &fields_c)), + ); + + Ok(Response::Query(pgwire::api::results::QueryResponse::new( + fields, data_rows, + ))) + } else { + let affected = client + .execute(&prepared, ¶m_refs) + .await + .map_err(backend_err)?; + + Ok(Response::Execution( + pgwire::api::results::Tag::new("INSERT").with_rows(affected as usize), + )) + } +} + +pub async fn execute_copy<'a>( + client: &Client, + s3_client: &S3Client, + ctx: ExecCtx<'a>, +) -> PgWireResult { + let portal = ctx.portal().ok_or_else(|| user_err("Missing portal"))?; + + let SqlStatement::Copy { + ref legacy_options, + ref source, + .. + } = portal.statement.statement.ast + else { + return Ok(0); + }; + + let is_parquet = legacy_options + .iter() + .any(|v| matches!(v, CopyLegacyOption::Parquet)); + + if !is_parquet { + return Err(user_err("Expected parquet")); + } + + execute_parquet_copy( + client, + s3_client, + source, + &portal.statement.statement.raw_sql, + ) + .await +} + +async fn execute_parquet_copy( + client: &Client, + s3_client: &S3Client, + source: &CopySource, + raw_sql: &str, +) -> PgWireResult { + let (table_name, cols) = match source { + CopySource::Table { + table_name, + columns, + } => (table_name, columns), + CopySource::Query(_) => return Err(user_err("Unsupported")), + }; + + let s3_uri = extract_s3_path(raw_sql)?; + + let (bucket_name, prefix) = split_s3_uri(&s3_uri)?; + + let existing_cols = fetch_table_columns(client, &table_name.to_string()) + .await + .map_err(|e| user_err(e.to_string()))? + .ok_or_else(|| user_err(format!("A required table is missing: {}", table_name)))?; + + let existing_names: HashSet<&str> = existing_cols.iter().map(|v| v.name.as_str()).collect(); + if !cols + .iter() + .all(|v| existing_names.contains(v.value.as_str())) + { + return Err(user_err(format!( + "Column mismatch for table: {}", + table_name + ))); + } + + let bytes = Bytes::from( + s3_client + .get_object(&prefix) + .await + .map_err(|e| user_err(e.to_string()))?, + ); + tracing::info!("File '{}' read", prefix); + + let inferred = infer_parquet_schema(bytes.clone()).map_err(user_err)?; + + let requested_cols: HashSet<&str> = cols.iter().map(|v| v.value.as_str()).collect(); + let in_cols: Vec<_> = inferred + .into_iter() + .filter(|c| requested_cols.contains(c.name.as_str())) + .collect(); + + let n = load_one_object(client, &format!("{}", table_name), &in_cols, bytes) + .await + .map_err(|e| { + tracing::error!("[copy] error loading s3://{bucket_name}/{prefix}: {e:#}"); + user_err(format!("Failed to load s3://{bucket_name}/{prefix}: {e}")) + })?; + + tracing::info!("{n} records stored"); + Ok(n as u64) +} + +pub fn extract_s3_path(copy_sql: &str) -> PgWireResult { + let start = copy_sql.find("s3://").ok_or(user_err(format!( + "Invalid query - Missing s3:// prefix: {copy_sql}" + )))?; + + let rest = ©_sql[start..]; + let end = rest.find('\'').ok_or(user_err(format!( + "Invalid query - Missing s3 link end: {rest}" + )))?; + + Ok(rest[..end].to_string()) +} + +pub fn split_s3_uri(uri: &str) -> PgWireResult<(String, String)> { + let rest = uri.strip_prefix("s3://").ok_or(user_err(format!( + "Invalid query - Missing s3:// prefix: {uri}" + )))?; + + match rest.split_once('/') { + Some((b, p)) => Ok((b.to_string(), p.to_string())), + None => Err(user_err(format!("Invalid query - Missing s3 key: {uri}"))), + } +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/dql.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/dql.rs new file mode 100644 index 0000000000..f5edb1ebc3 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/dql.rs @@ -0,0 +1,92 @@ +// 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 std::sync::Arc; + +use futures::stream; +use pgwire::api::results::{QueryResponse, Response}; +use pgwire::error::PgWireResult; +use tokio_postgres::Client; + +use super::{ + handler::ExecCtx, + util::{backend_err, columns_to_field_info, decode_param, row_to_data_row, user_err}, +}; + +pub async fn execute_select<'a>( + client: &Client, + ctx: ExecCtx<'a>, + max_rows: usize, +) -> PgWireResult { + let portal = ctx.portal().ok_or(user_err("Missing portal"))?; + + let raw_sql = &portal.statement.statement.raw_sql; + + let prepared = client.prepare(raw_sql).await.map_err(backend_err)?; + + let mut bound_params: Vec> = Vec::new(); + for (i, ty) in prepared.params().iter().enumerate() { + bound_params.push(decode_param(portal, ty, i)?); + } + let param_refs: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = + bound_params.iter().map(|b| b.as_ref() as &_).collect(); + + let fields = Arc::new(columns_to_field_info(prepared.columns())); + + if max_rows == 0 { + // 0 means "no limit" per the wire protocol: fetch everything. + // Stream via query_raw + try_next rather than query() so you're not + // buffering a huge result set in one Vec before encoding it — + // pgwire's Response::Query can take a Stream, not just a Vec. + let rows = client + .query(&prepared, ¶m_refs) + .await + .map_err(backend_err)?; + + let fields_c = Arc::clone(&fields); + + let data_rows = stream::iter( + rows.into_iter() + .map(move |r| row_to_data_row(&r, &fields_c.clone())), + ); + + Ok(Response::Query(QueryResponse::new( + fields.clone(), + data_rows, + ))) + } else { + let rows = client + .query(&prepared, ¶m_refs) + .await + .map_err(backend_err)?; + + let fields_c = Arc::clone(&fields); + + let truncated: Vec<_> = rows.into_iter().take(max_rows).collect(); + + let data_rows = stream::iter( + truncated + .into_iter() + .map(move |r| row_to_data_row(&r, &fields_c)), + ); + + Ok(Response::Query(QueryResponse::new( + fields.clone(), + data_rows, + ))) + } +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs index 139bf918f7..e19b96a82f 100644 --- a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs @@ -23,87 +23,88 @@ use pgwire::{ ClientInfo, ClientPortalStore, PgWireServerHandlers, Type, portal::Portal, query::{ExtendedQueryHandler, SimpleQueryHandler}, - results::{DescribePortalResponse, DescribeStatementResponse, FieldInfo, Response, Tag}, - stmt::{NoopQueryParser, StoredStatement}, + results::{DescribePortalResponse, DescribeStatementResponse, FieldInfo, Response}, + stmt::{QueryParser, StoredStatement}, store::PortalStore, }, - error::{ErrorInfo, PgWireError, PgWireResult}, + error::{PgWireError, PgWireResult}, }; -use sqlparser::{ast::Statement, dialect::RedshiftSqlDialect, parser::Parser}; -use tokio_postgres::{Client as PgClient, Statement as PgStatement}; +use tokio_postgres::Client as PgClient; use crate::connectors::fixtures::redshift::redshift_mock::{ - copy::try_parse_redshift_copy, - create::try_parse_redshift_create_table, - load::{ - S3Client, ToPgError, execute_create_table, execute_s3_copy, execute_select, split_s3_uri, - }, + load::S3Client, + util::{backend_err, columns_to_field_info, user_err}, }; -/// Statement failed to parse under the Redshift dialect. -pub const SYNTAX_ERROR: &str = "26000"; -/// Parsed, but uses a construct we don't implement (custom COPY/CREATE -/// extensions, unsupported statement kinds, etc.). -pub const FEATURE_NOT_SUPPORTED: &str = "42601"; -/// The Postgres connection backing this mock failed outright (as -/// opposed to Postgres returning a well-formed DB error). -pub const CONNECTION_EXCEPTION: &str = "08000"; -/// Empty/missing statement. -pub const INVALID_QUERY: &str = "42601"; -/// Statement kind we recognize but intentionally don't support. -pub const WARNING_UNSUPPORTED: &str = "01000"; +use super::{ + ddl, dml, dql, + parser::{ParsedStatement, QueryClass, RedshiftQueryParser}, +}; pub struct RedshiftHandlerFactory { - pub handler: Arc, + pub pg_dsn: String, + pub s3_client: Arc, } impl PgWireServerHandlers for RedshiftHandlerFactory { fn simple_query_handler(&self) -> Arc { - self.handler.clone() + Arc::new(RedshiftHandler::new( + self.pg_dsn.clone(), + self.s3_client.clone(), + )) } fn extended_query_handler(&self) -> Arc { - self.handler.clone() + Arc::new(RedshiftHandler::new( + self.pg_dsn.clone(), + self.s3_client.clone(), + )) } } -pub struct RedshiftHandler { - pg: PgClient, - s3_endpoint: String, +struct RedshiftHandler { + pg_dsn: String, + pg: tokio::sync::OnceCell, + s3_client: Arc, + query_parser: Arc, } impl RedshiftHandler { - pub fn new(pg: PgClient, s3_endpoint: String) -> Self { - Self { pg, s3_endpoint } - } - - /// Prepares `sql` against the backing Postgres connection for describe - /// purposes. Returns `Ok(None)` for statement kinds (currently just - /// `COPY`) that describe to zero fields rather than going through an - /// unsupported `PREPARE`. - async fn prepare_describable(&self, sql: &str) -> PgWireResult> { - let dialect = RedshiftSqlDialect {}; - let statements = parse_redshift_sql(&dialect, sql)?; - - if matches!(statements.first(), Some(Statement::Copy { .. })) { - return Ok(None); + pub fn new(pg_dsn: String, s3_client: Arc) -> Self { + Self { + pg_dsn, + pg: tokio::sync::OnceCell::new(), + s3_client, + query_parser: Arc::new(RedshiftQueryParser), } + } + async fn pg_client(&self) -> Result<&tokio_postgres::Client, PgWireError> { self.pg - .prepare(sql) + .get_or_try_init(|| async { + let (pg_client, pg_conn) = + tokio_postgres::connect(&self.pg_dsn, tokio_postgres::NoTls) + .await + .map_err(backend_err)?; + + tokio::spawn(async move { + if let Err(e) = pg_conn.await { + tracing::error!("Postgres connection error: {e}"); + } + }); + Ok::<_, PgWireError>(pg_client) + }) .await - .map(Some) - .map_err(|e| map_pg_client_error(e, CONNECTION_EXCEPTION)) } } #[async_trait] impl ExtendedQueryHandler for RedshiftHandler { - type Statement = String; - type QueryParser = NoopQueryParser; + type Statement = ParsedStatement; + type QueryParser = RedshiftQueryParser; fn query_parser(&self) -> Arc { - Arc::new(NoopQueryParser::new()) + self.query_parser.clone() } async fn do_query( @@ -115,13 +116,15 @@ impl ExtendedQueryHandler for RedshiftHandler { where C: ClientInfo + Unpin + Send + Sync, { - let query = &portal.statement.statement; + let stmt = portal.statement.statement.clone(); - if query.trim().is_empty() { + if stmt.raw_sql.trim().is_empty() { return Ok(Response::EmptyQuery); } - execute_statement(query, &self.pg, &self.s3_endpoint).await + let pg = self.pg_client().await?; + + execute_statement(stmt, ExecCtx::Bound(portal), pg, &self.s3_client).await } async fn do_describe_statement( @@ -132,25 +135,20 @@ impl ExtendedQueryHandler for RedshiftHandler { where C: ClientInfo + Unpin + Send + Sync, { - let Some(prepared) = self.prepare_describable(&stmt.statement).await? else { + if matches!(stmt.statement.class, QueryClass::DmlCopy) { return Ok(DescribeStatementResponse::new(vec![], vec![])); - }; + } + + let prepared = self + .pg_client() + .await? + .prepare(&stmt.statement.raw_sql) + .await + .map_err(backend_err)?; let param_types: Vec = prepared.params().to_vec(); - let fields: Vec = prepared - .columns() - .iter() - .map(|col| { - FieldInfo::new( - col.name().to_owned(), - None, - None, - col.type_().clone(), - pgwire::api::results::FieldFormat::Text, - ) - }) - .collect(); + let fields: Vec = columns_to_field_info(prepared.columns()); Ok(DescribeStatementResponse::new(param_types, fields)) } @@ -163,27 +161,18 @@ impl ExtendedQueryHandler for RedshiftHandler { where C: ClientInfo + Unpin + Send + Sync, { - let Some(prepared) = self - .prepare_describable(&portal.statement.statement) - .await? - else { + if matches!(portal.statement.statement.class, QueryClass::DmlCopy) { return Ok(DescribePortalResponse::new(vec![])); - }; - - let fields: Vec = prepared - .columns() - .iter() - .enumerate() - .map(|(idx, col)| { - FieldInfo::new( - col.name().to_owned(), - None, - None, - col.type_().clone(), - portal.result_column_format.format_for(idx), - ) - }) - .collect(); + } + + let prepared = self + .pg_client() + .await? + .prepare(&portal.statement.statement.raw_sql) + .await + .map_err(backend_err)?; + + let fields: Vec = columns_to_field_info(prepared.columns()); Ok(DescribePortalResponse::new(fields)) } @@ -200,106 +189,68 @@ impl SimpleQueryHandler for RedshiftHandler { return Ok(vec![Response::EmptyQuery]); } + let stmt = self.query_parser.parse_sql(_client, query, &[]).await?; + + let pg = self.pg_client().await?; + Ok(vec![ - execute_statement(query, &self.pg, &self.s3_endpoint).await?, + execute_statement(stmt, ExecCtx::Unbound, pg, &self.s3_client).await?, ]) } } -/// Dispatches a single SQL statement to the appropriate executor based on -/// its parsed kind. Kept intentionally thin — each branch delegates to a -/// dedicated function so individual statement kinds can be read (and -/// tested) in isolation. -async fn execute_statement( - query: &str, +async fn execute_statement<'a>( + stmt: ParsedStatement, + ctx: ExecCtx<'a>, pg: &PgClient, - s3_endpoint: &str, + s3_client: &S3Client, ) -> PgWireResult { - let dialect = RedshiftSqlDialect {}; - let statements = parse_redshift_sql(&dialect, query)?; - - match statements.first() { - Some(Statement::Query(_)) => execute_select(query, pg).await, - Some(Statement::CreateTable(_)) => execute_create(query, pg).await, - Some(Statement::Copy { .. }) => execute_copy(query, pg, s3_endpoint).await, - Some(other) => Err(pg_warning( - WARNING_UNSUPPORTED, - format!("Unsupported: {other:?}"), - )), - None => Err(pg_error(INVALID_QUERY, "Invalid query")), - } -} - -async fn execute_create(query: &str, pg: &PgClient) -> PgWireResult { - let dialect = RedshiftSqlDialect {}; - let parser = Parser::new(&dialect) - .try_with_sql(query) - .map_err(|e| pg_error(SYNTAX_ERROR, format!("Unsupported: {e:?}")))?; - - let create = try_parse_redshift_create_table(parser) - .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; - - execute_create_table(create, pg).await -} - -async fn execute_copy(query: &str, pg: &PgClient, s3_endpoint: &str) -> PgWireResult { - let dialect = RedshiftSqlDialect {}; - let parser = Parser::new(&dialect) - .try_with_sql(query) - .map_err(|e| pg_error(SYNTAX_ERROR, format!("Unsupported: {e:?}")))?; - - let r_copy = try_parse_redshift_copy(parser) - .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; + match stmt.class { + QueryClass::DdlCreate => { + let affected = ddl::execute_create(pg, &stmt.raw_sql).await?; - let (bucket_name, prefix) = split_s3_uri(&r_copy.s3_uri).map_err(|e| e.to_pg_wire_error())?; + Ok(Response::Execution( + pgwire::api::results::Tag::new("CREATE").with_rows(affected as usize), + )) + } + QueryClass::DdlTruncate => { + let affected = ddl::execute_truncate(pg, &stmt).await?; - let s3_client = S3Client::new( - &bucket_name, - s3_endpoint, - &r_copy.access_key_id, - &r_copy.secret_access_key, - &r_copy.region, - ) - .await - .map_err(|e| e.to_pg_wire_error())?; + Ok(Response::Execution( + pgwire::api::results::Tag::new("TRUNCATE").with_rows(affected as usize), + )) + } + QueryClass::Dql => { + let response = dql::execute_select(pg, ctx, 0).await?; - let rows = execute_s3_copy(&r_copy, pg, s3_client, &bucket_name, &prefix) - .await - .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; + Ok(response) + } + QueryClass::DmlCopy => { + let affected = dml::execute_copy(pg, s3_client, ctx).await?; - Ok(Response::Execution(Tag::new("copy").with_rows(rows))) -} + Ok(Response::Execution( + pgwire::api::results::Tag::new("COPY").with_rows(affected as usize), + )) + } + QueryClass::DmlMerge => { + let response = dml::execute_dml(pg, ctx).await?; -/// Builds a `PgWireError::UserError` with severity `ERROR`. Replaces the -/// repeated `PgWireError::UserError(Box::new(ErrorInfo::new(...)))` calls. -fn pg_error(code: &str, message: impl std::fmt::Display) -> PgWireError { - PgWireError::UserError(Box::new(ErrorInfo::new( - "ERROR".into(), - code.into(), - message.to_string(), - ))) + Ok(response) + } + QueryClass::Other => Err(user_err(format!("Unsupported: {:?}", stmt.raw_sql))), + } } -/// Same as [`pg_error`] but with severity `WARNING`, for statement kinds -/// we recognize but choose not to support. -fn pg_warning(code: &str, message: impl std::fmt::Display) -> PgWireError { - PgWireError::UserError(Box::new(ErrorInfo::new( - "WARNING".into(), - code.into(), - message.to_string(), - ))) +pub enum ExecCtx<'a> { + Bound(&'a Portal), + Unbound, } -/// Maps a `tokio_postgres::Error` to a `PgWireError`, preserving the -/// upstream SQLSTATE/message when Postgres itself produced the error, and -/// falling back to `fallback_code` for connection-level failures. -fn map_pg_client_error(err: tokio_postgres::Error, fallback_code: &str) -> PgWireError { - match err.as_db_error() { - Some(db_err) => pg_error(db_err.code().code(), db_err.message()), - None => pg_error(fallback_code, format!("connection failed: {err}")), +impl<'a> ExecCtx<'a> { + pub fn portal(&self) -> Option<&'a Portal> { + match self { + ExecCtx::Bound(p) => Some(p), + ExecCtx::Unbound => None, + } } } - -fn parse_redshift_sql(dialect: &RedshiftSqlDialect, sql: &str) -> PgWireResult> { - Parser::parse_sql(dialect, sql).map_err(|e| pg_error(SYNTAX_ERROR, e)) -} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs index 8de525b15d..d34024c50b 100644 --- a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs @@ -15,215 +15,37 @@ // specific language governing permissions and limitations // under the License. -use std::{fmt::Write, sync::Arc}; - use arrow::{ array::{ - Array, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Decimal256Array, - Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, + Array, BinaryArray, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, + StringArray, }, datatypes::DataType, }; use bytes::Bytes; -use futures::{StreamExt, pin_mut, stream}; +use futures::pin_mut; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; -use pgwire::{ - api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag}, - error::{ErrorInfo, PgWireError, PgWireResult}, -}; - use s3::{Bucket, Region, creds::Credentials}; -use sqlparser::ast::DataType as SDataType; -use sqlx::types::chrono; use tokio_postgres::{ - Client as PgClient, GenericClient, + Client as PgClient, binary_copy::BinaryCopyInWriter, types::{ToSql, Type as PgType}, }; -use crate::connectors::fixtures::redshift::redshift_mock::{ - copy::{CopyFormat, RedshiftCopy}, - create::{RedshiftCreateTable, TableKind}, -}; - -pub async fn execute_select(sql: &str, pg: &PgClient) -> PgWireResult { - let value = pg - .client() - .query_one_scalar::(sql, &[]) - .await - .map_err(|e| pg_wire_error(&e, "Query failed"))?; - - let schema = Arc::new(vec![FieldInfo::new( - "column".into(), - None, - None, - PgType::INT4, - FieldFormat::Text, - )]); - - let schema_ref = schema.clone(); - - let row_stream = stream::iter(std::iter::once(value)).map(move |v| { - let mut encoder = DataRowEncoder::new(schema_ref.clone()); - encoder.encode_field(&v)?; - - Ok(encoder.take_row()) - }); - - Ok(Response::Query(QueryResponse::new(schema, row_stream))) -} - -pub async fn execute_create_table( - create: RedshiftCreateTable, - pg: &PgClient, -) -> PgWireResult { - let create_kw = match create.table_kind { - TableKind::Regular => "CREATE TABLE", - TableKind::Temp => "CREATE TEMP TABLE", - // Postgres accepts (and ignores) the LOCAL keyword per the SQL - // standard, so this round-trips fine. - TableKind::LocalTemp => "CREATE LOCAL TEMPORARY TABLE", - }; - - let mut sql = String::from(create_kw); - sql.push(' '); - - if create.if_not_exists { - sql.push_str("IF NOT EXISTS "); - } - - write!(sql, "{} (", create.table).unwrap(); - - let fields = create.columns.iter().fold(vec![], |mut acc, v| { - let mut f = format!("{} {}", v.name, map_type(&v.data_type)); - - if let Some(id) = &v.identity { - f.push_str(&format!( - " GENERATED BY DEFAULT AS IDENTITY (START WITH {} INCREMENT BY {})", - id.seed, id.step - )); - } else if let Some(default) = &v.default { - let default = format!("{}", default).replace("GETDATE", "NOW"); - f.push_str(&format!(" DEFAULT {default}")); - } - - if v.primary_key { - f.push_str(" PRIMARY KEY"); - } - - if v.not_null { - f.push_str(" NOT NULL"); - } - - if let Some(target) = &v.references { - f.push_str(&format!(" REFERENCES {target}")); - } - - acc.push(f); - - acc - }); - - sql.push_str(&fields.join(", ")); - sql.push_str(");"); - - let rows_affected = pg - .execute(&sql, &[]) - .await - .map_err(|e| pg_wire_error(&e, "CREATE TABLE failed"))?; - - Ok(Response::Execution( - Tag::new("create table").with_rows(rows_affected as usize), - )) -} - -pub async fn execute_s3_copy( - copy: &RedshiftCopy, +pub async fn fetch_table_columns( pg: &PgClient, - s3_client: S3Client, - bucket_name: &str, - prefix: &str, -) -> PgWireResult { - let table_name = copy.table.to_string(); - let mut failed_objects = 0usize; - - match ©.format { - CopyFormat::Parquet => { - let existing_columns = fetch_table_columns(pg, ©.table.to_string()) - .await - .map_err(|e| e.to_pg_wire_error())?; - - tracing::info!( - "target table '{table_name}' {}", - if existing_columns.is_some() { - "exists" - } else { - "does not exist" - } - ); - - let bytes = Bytes::from( - s3_client - .get_object(prefix) - .await - .map_err(|e| e.to_pg_wire_error())?, - ); - - tracing::info!("File '{}' read", prefix); - - let cols: Vec = match existing_columns { - Some(cols) => strip_created_at(cols), - None => { - let inferred = - infer_parquet_schema(bytes.clone()).map_err(|e| e.to_pg_wire_error())?; - - let columns_sql = inferred - .iter() - .map(|f| format!("{} {}", f.name.to_lowercase(), f.pg_type.name())) - .collect::>() - .join(", "); - - create_table(pg, &table_name, &columns_sql) - .await - .map_err(|e| e.to_pg_wire_error())?; - - strip_created_at(inferred) - } - }; - - match load_one_object(pg, ©.table.to_string(), &cols, bytes).await { - Ok(n) => { - tracing::info!("{n} records stored"); - } - Err(e) => { - // Simplification vs real Redshift: MAXERROR there counts bad *rows* - // across the whole load; here each failed *file* counts as one error. - // Good enough for a test double where you're usually asserting - // "load either fully succeeds or trips MAXERROR", not exact counts. - failed_objects += 1; - tracing::error!("[copy] error loading s3://{bucket_name}/{prefix}: {e:#}"); - if failed_objects > copy.max_error as usize { - Err(format!("MAXERROR ({}) exceeded", copy.max_error).to_pg_wire_error())? - } - } - } - } - other => Err(format!("{:?} unsupported", other).to_pg_wire_error())?, - }; + table: &str, +) -> Result>, String> { + let query = format!( + "SELECT column_name, udt_name FROM information_schema.columns WHERE table_name = '{}' ORDER BY ordinal_position", + table.replace('"', "") + ); - Ok(1) -} - -async fn fetch_table_columns(pg: &PgClient, table: &str) -> Result>, String> { - let rows = pg - .query( - "SELECT column_name, udt_name FROM information_schema.columns \ - WHERE table_name = $1 ORDER BY ordinal_position", - &[&table], - ) - .await - .map_err(|e| e.to_string())?; + let rows = pg.query(&query, &[]).await.map_err(|e| { + tracing::error!("{:?}", e); + e.to_string() + })?; if rows.is_empty() { return Ok(None); @@ -244,21 +66,7 @@ async fn fetch_table_columns(pg: &PgClient, table: &str) -> Result Result { - let sql = format!("CREATE TABLE {} ({});", table, columns); - - tracing::info!("{sql}"); - - let result = pg - .client() - .execute(&sql, &[]) - .await - .map_err(|e| e.to_string())?; - - Ok(result as usize) -} - -async fn load_one_object( +pub async fn load_one_object( pg: &PgClient, table: &str, columns: &[ColumnDef], @@ -281,7 +89,11 @@ async fn load_one_object( .collect::>(); let copy_sql = format!("COPY {table} ({col_list}) FROM STDIN BINARY"); - let sink = pg.copy_in(©_sql).await.map_err(|e| e.to_string())?; + let sink = pg.copy_in(©_sql).await.map_err(|e| { + tracing::error!("{:?}", e.as_db_error()); + + e.to_string() + })?; tracing::info!("COPY FROM STDIN started"); @@ -293,6 +105,7 @@ async fn load_one_object( for batch in reader { let batch = batch.map_err(|e| e.to_string())?; + for row_idx in 0..batch.num_rows() { let row_values = extract_row(&batch, row_idx, columns)?; @@ -301,24 +114,29 @@ async fn load_one_object( .map(|v| v.as_ref() as &(dyn ToSql + Sync)) .collect(); - writer - .as_mut() - .write(&refs) - .await - .map_err(|e| e.to_string())?; + writer.as_mut().write(&refs).await.map_err(|e| { + tracing::error!("{:?}", e.as_db_error()); + + e.to_string() + })?; n += 1; } } - writer.finish().await.map_err(|e| e.to_string())?; + writer.finish().await.map_err(|e| { + tracing::error!("{:?}", e.as_db_error()); + + e.to_string() + })?; Ok(n) } -struct ColumnDef { - name: String, - pg_type: PgType, +#[derive(Debug)] +pub struct ColumnDef { + pub name: String, + pub pg_type: PgType, } macro_rules! scalar_column { @@ -359,50 +177,9 @@ fn extract_row<'a>( DataType::Int32 => scalar_column!(array, Int32Array, i32, row, |v: i32| v), DataType::Float64 => scalar_column!(array, Float64Array, f64, row, |v: f64| v), DataType::Boolean => scalar_column!(array, BooleanArray, bool, row, |v: bool| v), - DataType::Date32 => scalar_column!(array, Date32Array, i32, row, |v: i32| v), DataType::Binary => { scalar_column!(array, BinaryArray, Vec, row, |v: &[u8]| v.to_vec()) } - DataType::Decimal128(_, _) => { - let a = array - .as_any() - .downcast_ref::() - .ok_or("expected Decimal256Array")?; - if a.is_null(row) { - Box::new(None::) - } else { - Box::new(a.value(row).to_string()) - } - } - // Downcast to fit i256 serialization constraints - // Decimal(39, 0) > becomes Decimal256 on arrow - // Decimal256 serialization will require extra handling - DataType::Decimal256(_, _) => { - let a = array - .as_any() - .downcast_ref::() - .ok_or("expected Decimal256Array")?; - if a.is_null(row) { - Box::new(None::) - } else { - let raw = a.value(row).to_string(); - Box::new(raw) - } - } - DataType::Timestamp(_, _) => { - let a = array - .as_any() - .downcast_ref::() - .ok_or("expected TimestampMicrosecondArray")?; - if a.is_null(row) { - Box::new(None::>) - } else { - let micros = a.value(row); - let dt = chrono::DateTime::::from_timestamp_micros(micros) - .ok_or("Invalid timestamp")?; - Box::new(dt) - } - } other => Err(format!( "unsupported parquet column type {other:?} for column {}", col.name @@ -446,28 +223,11 @@ fn arrow_to_type(a_type: &DataType) -> Result { DataType::Decimal128(_, _) => Ok(PgType::VARCHAR), DataType::Decimal256(_, _) => Ok(PgType::VARCHAR), DataType::Utf8 => Ok(PgType::TEXT), - DataType::Date32 => Ok(PgType::DATE), - DataType::Timestamp(_, _) => Ok(PgType::TIMESTAMPTZ), other => Err(format!("Unsuppoerted type: {}", other)), } } -fn map_type(dt: &SDataType) -> String { - match dt { - // Redshift's SUPER (semi-structured) has no PG equivalent for - // your emulator's purposes -> JSON is the closest usable stand-in. - SDataType::Custom(name, _) if name.to_string().eq_ignore_ascii_case("SUPER") => { - "JSON".to_string() - } - SDataType::Custom(name, mods) if name.to_string().eq_ignore_ascii_case("VARBYTE") => { - "BYTEA".to_string() - } - SDataType::Decimal(_) => "VARCHAR".into(), - other => other.to_string(), - } -} - -fn infer_parquet_schema(bytes: Bytes) -> Result, String> { +pub fn infer_parquet_schema(bytes: Bytes) -> Result, String> { let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).map_err(|e| e.to_string())?; reader @@ -483,12 +243,6 @@ fn infer_parquet_schema(bytes: Bytes) -> Result, String> { .collect() } -fn strip_created_at(cols: Vec) -> Vec { - cols.into_iter() - .filter(|v| v.name != "created_at") - .collect() -} - /// S3 #[allow(unused)] #[derive(Clone)] @@ -558,46 +312,3 @@ impl S3Client { Ok(response.bytes().to_vec()) } } - -pub fn split_s3_uri(uri: &str) -> Result<(String, String), String> { - let rest = uri - .strip_prefix("s3://") - .ok_or("not an s3:// URI".to_string())?; - - match rest.split_once('/') { - Some((b, p)) => Ok((b.to_string(), p.to_string())), - None => Err(format!("s3 URI missing key/prefix: {uri}")), - } -} - -pub trait ToPgError { - fn to_pg_wire_error(self) -> PgWireError; -} - -impl ToPgError for T -where - T: Into, -{ - fn to_pg_wire_error(self) -> PgWireError { - PgWireError::UserError(Box::new(ErrorInfo::new( - "ERROR".into(), - "GG000".into(), - self.into(), - ))) - } -} - -fn pg_wire_error(e: &tokio_postgres::Error, context: &str) -> PgWireError { - match e.as_db_error() { - Some(db_err) => PgWireError::UserError(Box::new(ErrorInfo::new( - "ERROR".into(), - db_err.code().code().to_string(), - db_err.message().to_string(), - ))), - None => PgWireError::UserError(Box::new(ErrorInfo::new( - "ERROR".into(), - "XX000".into(), - format!("{context}: {e}"), - ))), - } -} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs index 1d9ba47977..8631084aeb 100644 --- a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs @@ -15,42 +15,10 @@ // specific language governing permissions and limitations // under the License. -use sqlparser::{ - parser::{Parser, ParserError}, - tokenizer::{Token, Word}, -}; - -pub mod copy; -pub mod create; +pub mod ddl; +pub mod dml; +pub mod dql; pub mod handler; pub mod load; - -pub fn expect_word(parser: &mut Parser) -> Result { - match parser.next_token().token { - Token::Word(Word { value, .. }) => Ok(value), - Token::SemiColon => Ok("SemiColon".into()), - other => Err(ParserError::ParserError(format!( - "expected identifier, got {other:?}" - ))), - } -} - -pub fn parse_string_literal(parser: &mut Parser) -> Result { - match parser.next_token().token { - Token::SingleQuotedString(s) => Ok(s), - other => Err(ParserError::ParserError(format!( - "expected string literal, got {other:?}" - ))), - } -} - -pub fn parse_number_literal(parser: &mut Parser) -> Result { - match parser.next_token().token { - Token::Number(s, _) => s - .parse() - .map_err(|_| ParserError::ParserError(format!("bad number: {s}"))), - other => Err(ParserError::ParserError(format!( - "expected number, got {other:?}" - ))), - } -} +pub mod parser; +pub mod util; diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/parser.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/parser.rs new file mode 100644 index 0000000000..cf0fef569e --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/parser.rs @@ -0,0 +1,549 @@ +// 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 std::ops::ControlFlow; + +use async_trait::async_trait; +use pgwire::{ + api::{ + ClientInfo, Type as PgWireType, + portal::Format, + results::{FieldFormat, FieldInfo}, + stmt::QueryParser, + }, + error::PgWireResult, +}; +use sqlparser::{ + ast::{ + CreateTable, DataType, Expr, HiveDistributionStyle, Ident, ObjectName, ObjectNamePart, + Select, SelectFlavor, SelectItem, SetExpr, Statement as SqlStatement, TableFactor, Value, + VisitMut, VisitorMut, + }, + dialect::{PostgreSqlDialect, RedshiftSqlDialect}, + parser::Parser as SqlParser, +}; + +use super::util::user_err; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QueryClass { + DdlCreate, + DdlTruncate, + DmlMerge, + DmlCopy, + Dql, + /// Anything we don't special-case: passthrough with no rewriting, + /// still logged + Other, +} + +#[derive(Debug, Clone)] +pub struct ParsedStatement { + pub raw_sql: String, + pub ast: SqlStatement, + // Captured during Parse and reused during Describe. + pub parameter_types: Vec, + // Fields + pub result_columns: Vec, + pub class: QueryClass, +} + +#[derive(Clone)] +pub struct RedshiftQueryParser; + +#[async_trait] +impl QueryParser for RedshiftQueryParser { + type Statement = ParsedStatement; + + async fn parse_sql( + &self, + _client: &C, + sql: &str, + param_types: &[Option], + ) -> PgWireResult + where + C: ClientInfo + Send + Sync, + { + tracing::debug!("Parsing sql"); + let dialect = RedshiftSqlDialect {}; + + let mut asts = SqlParser::parse_sql(&dialect, sql) + .map_err(|e| user_err(format!("sql parse error: {e}")))?; + + if asts.len() != 1 { + // Reject multi-statement Parse messages outright. + return Err(user_err( + "only a single statement is permitted per Parse message", + )); + } + + let _ = asts.visit(&mut RedshiftExprRewriter); + + tracing::debug!("Query rewritten"); + + let ast = asts.remove(0); + let class = classify(&ast); + let result_columns = select_schema(&ast); + + tracing::debug!("Done parsing"); + + let mut p_stmt = ParsedStatement { + raw_sql: ast.to_string(), + ast, + class, + parameter_types: param_types + .iter() + .clone() + .map(|ty| ty.clone().unwrap_or(PgWireType::UNKNOWN)) + .collect(), + result_columns, + }; + + p_stmt.rewrite_to_postgres().map_err(user_err)?; + + tracing::debug!("Postgres rewrite, {}", p_stmt.raw_sql); + + Ok(p_stmt) + } + + fn get_parameter_types(&self, stmt: &Self::Statement) -> PgWireResult> { + Ok(stmt.parameter_types.clone()) + } + + fn get_result_schema( + &self, + stmt: &Self::Statement, + _column_format: Option<&Format>, + ) -> PgWireResult> { + Ok(stmt.result_columns.clone()) + } +} + +fn classify(stmt: &SqlStatement) -> QueryClass { + match stmt { + SqlStatement::CreateTable { .. } => QueryClass::DdlCreate, + + SqlStatement::Truncate { .. } => QueryClass::DdlTruncate, + + SqlStatement::Merge { .. } => QueryClass::DmlMerge, + + SqlStatement::Copy { .. } => QueryClass::DmlCopy, + + SqlStatement::Query(_) => QueryClass::Dql, + + _ => QueryClass::Other, + } +} + +fn select_schema(stmt: &SqlStatement) -> Vec { + let SqlStatement::Query(query) = stmt else { + return vec![]; + }; + let SetExpr::Select(select) = query.body.as_ref() else { + return vec![]; + }; + + select + .projection + .iter() + .filter_map(|item| match item { + SelectItem::ExprWithAlias { expr, alias } => { + let field_info = FieldInfo::new( + alias.value.clone(), + None, + None, + expression_type(expr), + FieldFormat::Text, + ); + + Some(field_info) + } + SelectItem::UnnamedExpr(expr) => { + tracing::info!(?expr, resolved = ?expression_type(expr)); + let name = match expr { + Expr::Identifier(ident) => ident.value.clone(), + Expr::CompoundIdentifier(parts) => parts + .last() + .map(|ident| ident.value.clone()) + .unwrap_or_else(|| expr.to_string()), + _ => expr.to_string(), + }; + + let field_info = + FieldInfo::new(name, None, None, expression_type(expr), FieldFormat::Text); + + Some(field_info) + } + + _ => None, + }) + .collect() +} + +fn expression_type(expr: &Expr) -> PgWireType { + match expr { + Expr::Value(value) => match &value.value { + Value::Boolean(_) => PgWireType::BOOL, + Value::Number(number, _) if number.contains(['.', 'e', 'E']) => PgWireType::NUMERIC, + Value::Number(number, _) if number.parse::().is_ok() => PgWireType::INT4, + Value::Number(number, _) if number.parse::().is_ok() => PgWireType::INT8, + Value::Number(_, _) => PgWireType::NUMERIC, + Value::SingleQuotedString(_) + | Value::DollarQuotedString(_) + | Value::EscapedStringLiteral(_) + | Value::UnicodeStringLiteral(_) => PgWireType::TEXT, + Value::Null | Value::Placeholder(_) => PgWireType::UNKNOWN, + _ => PgWireType::UNKNOWN, + }, + + Expr::Cast { data_type, .. } => match data_type.to_string().to_uppercase().as_str() { + "BOOL" | "BOOLEAN" => PgWireType::BOOL, + "SMALLINT" | "INT2" => PgWireType::INT2, + "INTEGER" | "INT" | "INT4" => PgWireType::INT4, + "BIGINT" | "INT8" => PgWireType::INT8, + "REAL" | "FLOAT4" => PgWireType::FLOAT4, + "DOUBLE PRECISION" | "FLOAT8" => PgWireType::FLOAT8, + "NUMERIC" | "DECIMAL" => PgWireType::NUMERIC, + "TEXT" => PgWireType::TEXT, + "VARCHAR" | "CHARACTER VARYING" => PgWireType::VARCHAR, + "DATE" => PgWireType::DATE, + "TIMESTAMP" => PgWireType::TIMESTAMP, + "TIMESTAMP WITH TIME ZONE" => PgWireType::TIMESTAMPTZ, + _ => PgWireType::UNKNOWN, + }, + + Expr::BinaryOp { + op: + sqlparser::ast::BinaryOperator::Eq + | sqlparser::ast::BinaryOperator::NotEq + | sqlparser::ast::BinaryOperator::Lt + | sqlparser::ast::BinaryOperator::LtEq + | sqlparser::ast::BinaryOperator::Gt + | sqlparser::ast::BinaryOperator::GtEq + | sqlparser::ast::BinaryOperator::And + | sqlparser::ast::BinaryOperator::Or, + .. + } => PgWireType::BOOL, + + Expr::UnaryOp { op, expr } => match op { + sqlparser::ast::UnaryOperator::Not => PgWireType::BOOL, + sqlparser::ast::UnaryOperator::Minus | sqlparser::ast::UnaryOperator::Plus => { + expression_type(expr) + } + _ => PgWireType::UNKNOWN, + }, + + _ => PgWireType::UNKNOWN, + } +} + +impl ParsedStatement { + pub fn rewrite_to_postgres(&mut self) -> Result<(), String> { + match &mut self.ast { + SqlStatement::CreateTable(create_table) => { + redshift_create_table_to_postgres(create_table)?; + } + SqlStatement::Query(query) if matches!(query.body.as_ref(), SetExpr::Select(_)) => { + let SetExpr::Select(select) = query.body.as_mut() else { + return Err("No select body found".into()); + }; + + redshift_select_to_postgres(select)?; + } + _ => {} + } + + self.raw_sql = self.ast.to_string(); + + Ok(()) + } +} + +pub fn redshift_create_table_to_postgres(create: &mut CreateTable) -> Result<(), String> { + // These cannot be expressed as PostgreSQL CREATE TABLE. + let unsupported = [ + ("OR REPLACE", create.or_replace), + ("EXTERNAL", create.external), + ("TRANSIENT", create.transient), + ("ICEBERG", create.iceberg), + ("SNAPSHOT", create.snapshot), + ("DYNAMIC", create.dynamic), + ("WITHOUT ROWID", create.without_rowid), + ("COPY GRANTS", create.copy_grants), + ("REQUIRE USER", create.require_user), + ("STRICT", create.strict), + ]; + + if let Some((feature, _)) = unsupported.into_iter().find(|(_, present)| *present) { + return Err(format!( + "Redshift CREATE TABLE uses {feature}, which has no PostgreSQL CREATE TABLE equivalent" + )); + } + + if create.file_format.is_some() + || create.location.is_some() + || create.hive_formats.is_some() + || create.hive_distribution != HiveDistributionStyle::NONE + { + return Err( + "External/Hive storage options require a PostgreSQL foreign-table migration, \ + not CREATE TABLE transpilation." + .into(), + ); + } + + if create.clone.is_some() || create.version.is_some() { + return Err( + "CLONE / table-version syntax has no PostgreSQL CREATE TABLE equivalent".into(), + ); + } + + // Redshift physical-design directives have no PostgreSQL DDL equivalent. + if create.diststyle.take().is_some() { + tracing::warn!("Dropped Redshift DISTSTYLE."); + } + if create.distkey.take().is_some() { + tracing::warn!("Dropped Redshift DISTKEY."); + } + if create.sortkey.take().is_some() { + tracing::warn!("Dropped Redshift SORTKEY; create a PostgreSQL index separately if needed."); + } + if create.backup.take().is_some() { + tracing::warn!("Dropped Redshift BACKUP setting; configure PostgreSQL backups externally."); + } + + // `VOLATILE` is not PostgreSQL CREATE TABLE syntax. Treat it as TEMPORARY. + if create.volatile { + create.volatile = false; + create.temporary = true; + tracing::warn!("Translated VOLATILE to TEMPORARY."); + } + + for column in &mut create.columns { + match &column.data_type { + // PostgreSQL BYTEA has no length modifier. + DataType::Varbinary(_) => { + column.data_type = DataType::Bytea; + } + + // Fallback as sqlparser version parses VARBYTE + // as a custom type instead. + DataType::Custom(name, _) if name.to_string().eq_ignore_ascii_case("VARBYTE") => { + column.data_type = DataType::Bytea; + } + + _ => {} + } + } + + let sql = create.to_string(); + + // Syntax validation only; sqlparser deliberately does not perform full + // PostgreSQL semantic validation. + SqlParser::parse_sql(&PostgreSqlDialect {}, &sql) + .map_err(|error| format!("Generated SQL is not PostgreSQL syntax: {error}"))?; + + Ok(()) +} + +pub fn redshift_select_to_postgres(select: &mut Box