diff --git a/Cargo.lock b/Cargo.lock index 7e44d56d6f..75a1e0d334 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -515,6 +515,27 @@ version = "0.7.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" +[[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" @@ -525,12 +546,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", @@ -667,6 +688,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" @@ -811,6 +847,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" @@ -1419,6 +1468,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00bdb5da18dac48ca2cc7cd4a98e533e8635a58e2361d13a1a4ee3888e0d72f1" dependencies = [ "aws-lc-sys", + "untrusted 0.7.1", "zeroize", ] @@ -1932,6 +1982,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" @@ -2432,7 +2492,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", @@ -2744,7 +2804,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" dependencies = [ "chrono", - "phf", + "phf 0.12.1", ] [[package]] @@ -4032,7 +4092,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", @@ -4040,7 +4100,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", @@ -4067,7 +4127,7 @@ dependencies = [ "regex", "serde", "serde_json", - "sqlparser", + "sqlparser 0.61.0", "strum 0.27.2", "thiserror 2.0.19", "tokio", @@ -4877,6 +4937,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.17.0" @@ -5316,7 +5382,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -7166,6 +7232,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "iggy_connector_redshift_sink" +version = "0.4.1-edge.1" +dependencies = [ + "arrow 57.3.1", + "async-trait", + "chrono", + "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.5.0-edge.2" @@ -7446,6 +7534,7 @@ checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" name = "integration" version = "0.0.1" dependencies = [ + "arrow 57.3.1", "assert_cmd", "async-trait", "base64", @@ -7476,6 +7565,8 @@ dependencies = [ "lazy_static", "libc", "mongodb", + "parquet 57.3.1", + "pgwire", "predicates", "rand 0.10.2", "rcgen", @@ -7488,6 +7579,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "sqlparser 0.62.0", "sqlx", "sysinfo 0.39.6", "tempfile", @@ -7495,6 +7587,7 @@ dependencies = [ "testcontainers", "testcontainers-modules", "tokio", + "tokio-postgres", "toml 1.1.3+spec-1.1.0", "tracing", "tracing-subscriber", @@ -7825,6 +7918,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.119", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -8576,7 +8692,7 @@ checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -9083,6 +9199,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" @@ -9734,13 +9859,67 @@ dependencies = [ "pest", ] +[[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.2", + "rust_decimal", + "rustls-pki-types", + "ryu", + "serde", + "serde_json", + "smol_str", + "stringprep", + "thiserror 2.0.19", + "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]] @@ -9752,6 +9931,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" @@ -9930,6 +10118,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.2", + "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" @@ -10321,7 +10542,7 @@ dependencies = [ "libc", "once_cell", "raw-cpuid", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "web-sys", "winapi", ] @@ -10975,7 +11196,7 @@ dependencies = [ "cfg-if", "getrandom 0.2.17", "libc", - "untrusted", + "untrusted 0.9.0", "windows-sys 0.52.0", ] @@ -11219,6 +11440,7 @@ dependencies = [ "borsh", "bytes", "num-traits", + "postgres-types", "rand 0.8.7", "rkyv", "serde", @@ -11369,7 +11591,7 @@ dependencies = [ "aws-lc-rs", "ring", "rustls-pki-types", - "untrusted", + "untrusted 0.9.0", ] [[package]] @@ -12239,6 +12461,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" @@ -12348,6 +12580,28 @@ dependencies = [ "recursive", ] +[[package]] +name = "sqlparser" +version = "0.62.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +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]] name = "sqlx" version = "0.9.0" @@ -13223,6 +13477,32 @@ dependencies = [ "syn 2.0.119", ] +[[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.2", + "socket2 0.6.5", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -13941,6 +14221,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" @@ -14240,6 +14526,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" @@ -14249,6 +14544,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.126" @@ -14258,6 +14562,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] @@ -14425,6 +14730,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" @@ -14928,6 +15240,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.19", + "zeroize", +] + [[package]] name = "x509-parser" version = "0.18.1" diff --git a/Cargo.toml b/Cargo.toml index 31e37bacf1..dc9ec0dec9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,6 +44,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", @@ -256,6 +257,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" @@ -301,6 +303,7 @@ simd-json = { version = "0.17.3", features = ["serde_impl"] } slab = "0.4.12" smallvec = "1.15" socket2 = "0.6.5" +sqlparser = { version = "0.62.0", features = ["visitor"] } sqlx = { version = "0.9.0", features = [ "runtime-tokio", "tls-rustls", @@ -327,6 +330,7 @@ testcontainers = { version = "0.27.3", features = ["reusable-containers"] } testcontainers-modules = { version = "0.15.0", features = ["postgres", "http_wait"] } thiserror = "2.0.19" tokio = { version = "1.53.1", features = ["full"] } +tokio-postgres = "0.7.18" tokio-rustls = "0.26.4" tokio-tungstenite = { version = "0.30", 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 2690a97756..890a263b5f 100644 --- a/core/connectors/README.md +++ b/core/connectors/README.md @@ -86,6 +86,7 @@ Each sink should have its own, custom configuration, which is passed along with - **Meilisearch Sink** - indexes messages in Meilisearch - **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 e23e1ace9c..91de1fd084 100644 --- a/core/connectors/sinks/README.md +++ b/core/connectors/sinks/README.md @@ -15,6 +15,7 @@ Sink connectors are responsible for writing data from Iggy streams to external s | **meilisearch_sink** | Indexes messages in Meilisearch for full-text search | | **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..fb465701af --- /dev/null +++ b/core/connectors/sinks/redshift_sink/Cargo.toml @@ -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. + +[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 + +[package.metadata.cargo-machete] +ignored = ["rust-s3"] + +[lib] +crate-type = ["cdylib", "lib"] + +[dependencies] +arrow = { workspace = true } +async-trait = { workspace = true } +chrono = { 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..676c23a2c6 --- /dev/null +++ b/core/connectors/sinks/redshift_sink/README.md @@ -0,0 +1,215 @@ +# 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 | `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_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. | +| `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`. + +## 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_iggy_messages" (id, iggy_offset, iggy_timestamp, iggy_stream, iggy_topic, iggy_partition_id, iggy_checksum, iggy_origin_timestamp, payload, created_at) FROM 's3://iggystaging/iggy/messages/019ff3d5-a06f-7921-b084-0c67cabfefed.parquet' +CREDENTIALS 'aws_iam_role=arn:aws:iam::0123456789012:role/iggyRole' +FORMAT AS PARQUET +REGION 'us-east-1'; +``` + +**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 insert (idempotent upsert) + +```sql +INSERT INTO "iggy_messages" (id, iggy_offset, iggy_timestamp, iggy_stream, iggy_topic, iggy_partition_id, iggy_checksum, iggy_origin_timestamp, payload, created_at) +SELECT s.id, s.iggy_offset, s.iggy_timestamp, s.iggy_stream, s.iggy_topic, s.iggy_partition_id, s.iggy_checksum, s.iggy_origin_timestamp, s.payload, s.created_at +FROM (SELECT sm.*, ROW_NUMBER() OVER (PARTITION BY sm.id ORDER BY sm.created_at) AS rn FROM "staging_iggy_messages" sm) s +WHERE s.rn = 1 +AND NOT EXISTS (SELECT 1 FROM "iggy_messages" t WHERE t.id = s.id); +``` + +Uniqueness enforced on `id` — no update branch by design. 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/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..bc0627fc5d --- /dev/null +++ b/core/connectors/sinks/redshift_sink/src/config.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 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_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, + 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"); + } + + 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 = !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 !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 Redshift has SUPER(VARCHAR can be parsed by JSON_PARSE)ß +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum PayloadFormat { + Text, + #[default] + Varbyte, +} + +impl PayloadFormat { + pub fn from_config(s: Option<&str>) -> Self { + match s.map(|s| s.to_lowercase()).as_deref() { + 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 => "VARCHAR", + } + } + + pub fn arrow_type(&self) -> DataType { + match self { + PayloadFormat::Varbyte => DataType::Binary, + PayloadFormat::Text => 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..507213f21f --- /dev/null +++ b/core/connectors/sinks/redshift_sink/src/lib.rs @@ -0,0 +1,1483 @@ +// 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::{collections::BTreeMap, str::FromStr, sync::Arc, time::Duration}; + +use arrow::{ + 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, + basic::{Compression, ZstdLevel}, + file::properties::WriterProperties, +}; +use s3::{Bucket, Region, creds::Credentials}; +use secrecy::ExposeSecret; +use sqlx::{AssertSqlSafe, Pool, Postgres, Row, 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?; + // Ensuring tables exist + self.ensure_tables_exist().await?; + // Checking for schema drift + self.ensure_schema_match().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( + 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, + ) + .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_tables_exist(&self) -> Result<(), Error> { + let pool = self.get_pool()?; + + 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 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!("Staging table created"); + + sqlx::query(AssertSqlSafe(target_query)) + .execute(pool) + .await + .map_err(|e| { + tracing::error!(error = %e); + Error::InitError(format!("Failed to create table '{target_table}': {e}")) + })?; + + 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", "VARCHAR"); + + 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 { + let base = t.split('(').next().unwrap_or(t).trim(); + + match base.to_ascii_uppercase().as_str() { + "INTEGER" | "INT" | "INT4" => "INTEGER", + // e.g. "bigint" + "INT8" | "BIGINT" => "BIGINT", + // e.g "character varying(40)", "character varying(20)", "character varying(256)" + "VARCHAR" | "CHARACTER VARYING" => "VARCHAR", + // Having bytea because of the Postgres Test + // e.g. "binary varying(64000)" + "BYTEA" | "VARBYTE" | "VARBINARY" | "BINARY VARYING" => "VARBYTE", + "TEXT" => "TEXT", + // e.g. "timestamp with time zone" + "TIMESTAMPTZ" | "TIMESTAMP WITH TIME ZONE" => "TIMESTAMPTZ", + _ => "UNKNOWN", + } + } + + 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(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); + } + } + } + + 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(()) + } + + // 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> { + if messages.is_empty() { + return Ok(None); + } + + 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?; + + 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.insert_into_target(&cols).await?; + + tracing::info!("Redshift target table merge completed"); + + tracing::info!(count = messages.len(), path = %s3_path, "batch inserted into Redshift"); + + Ok(Some(s3_path)) + } + + 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 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", + 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 insert_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_insert_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 INSERT", + 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, target_table = %target_table, "Redshift INSERT completed"); + + Ok(()) + } + + 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 query = format!("CREATE TABLE IF NOT EXISTS {table_name} (id VARCHAR(40)"); + + if include_metadata { + 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"); + } + + if include_origin_timestamp { + query.push_str(", iggy_origin_timestamp VARCHAR(20)"); + } + + query.push_str(&format!(", payload {payload_type}")); + query.push_str(", created_at VARCHAR);"); + + Ok(query) + } + + 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)?.replace('"', ""); + + let region = quote_identifier(&self.config.aws_region)?.replace('"', ""); + + Ok(format!( + "COPY {} ({}) FROM '{}' CREDENTIALS 'aws_iam_role={}' FORMAT AS PARQUET REGION '{}';", + staging_table, cols, s3_path, iam_role, region + )) + } + + fn build_insert_sql(&self, cols: &[&str], staging: &str, target: &str) -> String { + let t_cols = cols.join(", "); + + let s_cols = cols + .iter() + .map(|v| format!("s.{v}")) + .collect::>() + .join(", "); + + format!( + " + INSERT INTO {} ({}) + SELECT {} + FROM (SELECT sm.*, ROW_NUMBER() OVER (PARTITION BY sm.id ORDER BY sm.created_at) AS rn FROM {} sm) s + WHERE s.rn = 1 + AND NOT EXISTS (SELECT 1 FROM {} t WHERE t.id = s.id);", + target, t_cols, s_cols, staging, target + ) + } + + 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")); + 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 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"); + + 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 = operation, attempts = 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 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(), 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"); + 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::Utf8, 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::Utf8, false)); + columns.push(origin_timestamp_column(messages)); + } + + fields.push(Field::new("payload", payload_format.arrow_type(), false)); + columns.push(payload_column(messages, payload_format)?); + + fields.push(Field::new("created_at", DataType::Utf8, false)); + columns.push(created_at_column(messages.len())?); + + 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]) -> ArrayRef { + Arc::new(StringArray::from_iter_values( + messages.iter().map(|v| v.id.to_string()), + )) +} + +fn metadata_columns( + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], +) -> (Vec, Vec) { + let fields = vec![ + 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::Int64, false), + ]; + + let columns: Vec = vec![ + Arc::new(StringArray::from_iter_values( + messages.iter().map(|v| v.offset.to_string()), + )), + 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()), + )), + Arc::new(StringArray::from_iter_values( + (0..messages.len()).map(|_| topic_metadata.topic.clone()), + )), + Arc::new(Int64Array::from_iter_values( + (0..messages.len()).map(|_| messages_metadata.partition_id as i64), + )), + ]; + + (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(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 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 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())) + } + } +} + +fn created_at_column(size: usize) -> Result { + let now = chrono::Utc::now().to_rfc3339(); + + let mut builder = StringBuilder::with_capacity(size, 0); + + for _ in 0..size { + builder.append_value(now.as_str()); + } + + Ok(Arc::new(builder.finish())) +} + +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..]; + + let bound_end = rest.find([':', '@', '?', '/']).unwrap_or(rest.len()); + + // Stop preview at the first sensitive boundary + let safe_end = rest + .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}***"); + } + + 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: 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, + 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 = Some(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 = Some(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::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::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 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\"")); + 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")); + } + + #[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(), 10); + + 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 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\"")); + 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")); + } + + #[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(), 3); + + 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, "1767225600000000"); + } + + #[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 bb7bdc51b2..8da11d1b5a 100644 --- a/core/integration/Cargo.toml +++ b/core/integration/Cargo.toml @@ -36,6 +36,7 @@ default = ["login-session"] login-session = ["dep:zbus-secret-service-keyring-store"] [dependencies] +arrow = { workspace = true } assert_cmd = { workspace = true } async-trait = { workspace = true } base64 = { workspace = true } @@ -71,6 +72,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 } @@ -88,6 +91,7 @@ secrecy = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } serial_test = { workspace = true } +sqlparser = { workspace = true } sqlx = { workspace = true } sysinfo = { workspace = true } tempfile = { workspace = true } @@ -95,6 +99,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 e4992d6785..dc208a858c 100644 --- a/core/integration/tests/connectors/fixtures/mod.rs +++ b/core/integration/tests/connectors/fixtures/mod.rs @@ -28,6 +28,7 @@ mod meilisearch; mod mongodb; mod postgres; mod quickwit; +mod redshift; mod s3; mod surrealdb; mod wiremock; @@ -81,6 +82,10 @@ pub use postgres::{ PostgresSourceOps, }; pub use quickwit::{QuickwitFixture, QuickwitOps, QuickwitPreCreatedFixture}; +pub use redshift::{ + RedshiftSinkFixture, RedshiftSinkJsonFixture, RedshiftSinkNoArchiveFixture, + RedshiftSinkVarbyteFixture, +}; 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..962c2378c0 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/container.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 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 tokio::{net::TcpListener, task::JoinHandle}; + +use crate::connectors::fixtures::{ + self, + 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; + +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 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 = + "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 = 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 + .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 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 { + pg_dsn: target_connection, + s3_client: Arc::new(s3_client), + }); + + 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] + Varbyte, + 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..0bfaa2f485 --- /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::{ + RedshiftSinkFixture, RedshiftSinkJsonFixture, RedshiftSinkNoArchiveFixture, + RedshiftSinkVarbyteFixture, +}; 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/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 new file mode 100644 index 0000000000..8fc6314f24 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs @@ -0,0 +1,256 @@ +// 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}, + stmt::{QueryParser, StoredStatement}, + store::PortalStore, + }, + error::{PgWireError, PgWireResult}, +}; +use tokio_postgres::Client as PgClient; + +use crate::connectors::fixtures::redshift::redshift_mock::{ + load::S3Client, + util::{backend_err, columns_to_field_info, user_err}, +}; + +use super::{ + ddl, dml, dql, + parser::{ParsedStatement, QueryClass, RedshiftQueryParser}, +}; + +pub struct RedshiftHandlerFactory { + pub pg_dsn: String, + pub s3_client: Arc, +} + +impl PgWireServerHandlers for RedshiftHandlerFactory { + fn simple_query_handler(&self) -> Arc { + Arc::new(RedshiftHandler::new( + self.pg_dsn.clone(), + self.s3_client.clone(), + )) + } + + fn extended_query_handler(&self) -> Arc { + Arc::new(RedshiftHandler::new( + self.pg_dsn.clone(), + self.s3_client.clone(), + )) + } +} + +struct RedshiftHandler { + pg_dsn: String, + pg: tokio::sync::OnceCell, + s3_client: Arc, + query_parser: Arc, +} + +impl RedshiftHandler { + 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 + .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 + } +} + +#[async_trait] +impl ExtendedQueryHandler for RedshiftHandler { + type Statement = ParsedStatement; + type QueryParser = RedshiftQueryParser; + + fn query_parser(&self) -> Arc { + self.query_parser.clone() + } + + async fn do_query( + &self, + _client: &mut C, + portal: &Portal, + _max_rows: usize, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + let stmt = portal.statement.statement.clone(); + + if stmt.raw_sql.trim().is_empty() { + return Ok(Response::EmptyQuery); + } + + let pg = self.pg_client().await?; + + execute_statement(stmt, ExecCtx::Bound(portal), pg, &self.s3_client).await + } + + async fn do_describe_statement( + &self, + _client: &mut C, + stmt: &StoredStatement, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + 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 = columns_to_field_info(prepared.columns()); + + Ok(DescribeStatementResponse::new(param_types, fields)) + } + + async fn do_describe_portal( + &self, + _client: &mut C, + portal: &Portal, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + if matches!(portal.statement.statement.class, QueryClass::DmlCopy) { + return Ok(DescribePortalResponse::new(vec![])); + } + + 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)) + } +} + +#[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]); + } + + let stmt = self.query_parser.parse_sql(_client, query, &[]).await?; + + let pg = self.pg_client().await?; + + Ok(vec![ + execute_statement(stmt, ExecCtx::Unbound, pg, &self.s3_client).await?, + ]) + } +} + +async fn execute_statement<'a>( + stmt: ParsedStatement, + ctx: ExecCtx<'a>, + pg: &PgClient, + s3_client: &S3Client, +) -> PgWireResult { + match stmt.class { + QueryClass::DdlCreate => { + let affected = ddl::execute_create(pg, &stmt.raw_sql).await?; + + Ok(Response::Execution( + pgwire::api::results::Tag::new("CREATE").with_rows(affected as usize), + )) + } + QueryClass::DdlTruncate => { + let affected = ddl::execute_truncate(pg, &stmt).await?; + + 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?; + + Ok(response) + } + QueryClass::DmlCopy => { + let affected = dml::execute_copy(pg, s3_client, ctx).await?; + + Ok(Response::Execution( + pgwire::api::results::Tag::new("COPY").with_rows(affected as usize), + )) + } + QueryClass::DmlInsert => { + let response = dml::execute_dml(pg, ctx).await?; + + Ok(response) + } + QueryClass::Other => Err(user_err(format!("Unsupported: {:?}", stmt.raw_sql))), + } +} + +pub enum ExecCtx<'a> { + Bound(&'a Portal), + Unbound, +} + +impl<'a> ExecCtx<'a> { + pub fn portal(&self) -> Option<&'a Portal> { + match self { + ExecCtx::Bound(p) => Some(p), + ExecCtx::Unbound => None, + } + } +} 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..d34024c50b --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs @@ -0,0 +1,314 @@ +// 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::{ + array::{ + Array, BinaryArray, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, + StringArray, + }, + datatypes::DataType, +}; + +use bytes::Bytes; +use futures::pin_mut; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use s3::{Bucket, Region, creds::Credentials}; +use tokio_postgres::{ + Client as PgClient, + binary_copy::BinaryCopyInWriter, + types::{ToSql, Type as PgType}, +}; + +pub async fn fetch_table_columns( + pg: &PgClient, + 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('"', "") + ); + + let rows = pg.query(&query, &[]).await.map_err(|e| { + tracing::error!("{:?}", 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?)) +} + +pub 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| { + tracing::error!("{:?}", e.as_db_error()); + + 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| { + tracing::error!("{:?}", e.as_db_error()); + + e.to_string() + })?; + + n += 1; + } + } + + writer.finish().await.map_err(|e| { + tracing::error!("{:?}", e.as_db_error()); + + e.to_string() + })?; + + Ok(n) +} + +#[derive(Debug)] +pub struct ColumnDef { + pub name: String, + pub 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::Binary => { + scalar_column!(array, BinaryArray, Vec, row, |v: &[u8]| v.to_vec()) + } + 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), + other => Err(format!("Unsuppoerted type: {}", other)), + } +} + +pub 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() +} + +/// 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()) + } +} 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..8631084aeb --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs @@ -0,0 +1,24 @@ +// 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. + +pub mod ddl; +pub mod dml; +pub mod dql; +pub mod handler; +pub mod load; +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..95a811992f --- /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, + DmlInsert, + 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::Insert { .. } => QueryClass::DmlInsert, + + 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