diff --git a/Cargo.lock b/Cargo.lock index 3f2c298ca7..9dbd292e2e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -504,6 +504,27 @@ version = "0.7.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f02882884d3e1bc524fb12c79f107f6ad0e1cfd498c536ffb494301740995dfe" +[[package]] +name = "arrow" +version = "57.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3bd47f2a6ddc39244bd722a27ee5da66c03369d087b9e024eafdb03e98b98ea7" +dependencies = [ + "arrow-arith 57.3.1", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-cast 57.3.1", + "arrow-csv 57.3.1", + "arrow-data 57.3.1", + "arrow-ipc 57.3.1", + "arrow-json 57.3.1", + "arrow-ord 57.3.1", + "arrow-row 57.3.1", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", + "arrow-string 57.3.1", +] + [[package]] name = "arrow" version = "58.3.0" @@ -514,12 +535,12 @@ dependencies = [ "arrow-array 58.3.0", "arrow-buffer 58.3.0", "arrow-cast 58.3.0", - "arrow-csv", + "arrow-csv 58.3.0", "arrow-data 58.3.0", "arrow-ipc 58.3.0", "arrow-json 58.3.0", "arrow-ord 58.3.0", - "arrow-row", + "arrow-row 58.3.0", "arrow-schema 58.3.0", "arrow-select 58.3.0", "arrow-string 58.3.0", @@ -656,6 +677,21 @@ dependencies = [ "ryu", ] +[[package]] +name = "arrow-csv" +version = "57.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27ddb80a4848e03b1655af496d5ac2563a779e5742fcb48f2ca2e089c9cd2197" +dependencies = [ + "arrow-array 57.3.1", + "arrow-cast 57.3.1", + "arrow-schema 57.3.1", + "chrono", + "csv", + "csv-core", + "regex", +] + [[package]] name = "arrow-csv" version = "58.3.0" @@ -800,6 +836,19 @@ dependencies = [ "arrow-select 58.3.0", ] +[[package]] +name = "arrow-row" +version = "57.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a931b520a2a5e22033e01a6f2486b4cdc26f9106b759abeebc320f125e94d7" +dependencies = [ + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", + "half", +] + [[package]] name = "arrow-row" version = "58.3.0" @@ -1408,6 +1457,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" dependencies = [ "aws-lc-sys", + "untrusted 0.7.1", "zeroize", ] @@ -1913,6 +1963,16 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bcder" +version = "0.7.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b593e5aeaf7992d388c08a9831c921cd703718064b3e50ba8e6d666d6cf86ca7" +dependencies = [ + "bytes", + "smallvec", +] + [[package]] name = "bdd" version = "0.0.1" @@ -2403,7 +2463,7 @@ version = "0.22.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2235eb320cd7178862a32dd111bd0c0f71a368e393add4914c50129add478eab" dependencies = [ - "arrow", + "arrow 58.3.0", "buoyant_kernel_derive", "bytes", "chrono", @@ -2721,7 +2781,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" dependencies = [ "chrono", - "phf", + "phf 0.12.1", ] [[package]] @@ -3991,7 +4051,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", @@ -3999,7 +4059,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", @@ -4026,7 +4086,7 @@ dependencies = [ "regex", "serde", "serde_json", - "sqlparser", + "sqlparser 0.61.0", "strum 0.27.2", "thiserror 2.0.18", "tokio", @@ -4840,6 +4900,12 @@ dependencies = [ "ext-trait", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + [[package]] name = "fastbloom" version = "0.14.1" @@ -5279,7 +5345,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -7111,6 +7177,27 @@ dependencies = [ "uuid", ] +[[package]] +name = "iggy_connector_redshift_sink" +version = "0.4.1-edge.1" +dependencies = [ + "arrow 57.3.1", + "async-trait", + "humantime", + "iggy_common", + "iggy_connector_sdk", + "parquet 57.3.1", + "rust-s3", + "secrecy", + "serde", + "serde_json", + "simd-json", + "sqlx", + "tokio", + "tracing", + "uuid", +] + [[package]] name = "iggy_connector_s3_sink" version = "0.4.0" @@ -7382,6 +7469,8 @@ checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" name = "integration" version = "0.0.1" dependencies = [ + "anyhow", + "arrow 57.3.1", "assert_cmd", "async-trait", "base64", @@ -7410,6 +7499,8 @@ dependencies = [ "lazy_static", "libc", "mongodb", + "parquet 57.3.1", + "pgwire", "predicates", "rand 0.10.1", "rcgen", @@ -7424,6 +7515,7 @@ dependencies = [ "serial_test", "server", "socket2 0.6.4", + "sqlparser 0.62.0", "sqlx", "sysinfo 0.39.5", "tempfile", @@ -7431,6 +7523,7 @@ dependencies = [ "testcontainers", "testcontainers-modules", "tokio", + "tokio-postgres", "toml 1.1.2+spec-1.1.0", "tracing", "tracing-subscriber", @@ -7763,6 +7856,29 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4345964bb142484797b161f473a503a434de77149dd8c7427788c6e13379388" +[[package]] +name = "lazy-regex" +version = "3.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bae91019476d3ec7147de9aa291cadb6d870abf2f3015d2da73a90325ac1496" +dependencies = [ + "lazy-regex-proc_macros", + "once_cell", + "regex-lite", +] + +[[package]] +name = "lazy-regex-proc_macros" +version = "3.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4de9c1e1439d8b7b3061b2d209809f447ca33241733d9a3c01eabf2dc8d94358" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "syn 2.0.118", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -8461,7 +8577,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -8968,6 +9084,15 @@ dependencies = [ "objc2-foundation", ] +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.37.3" @@ -9623,13 +9748,67 @@ dependencies = [ "sha2 0.10.9", ] +[[package]] +name = "pg_interval" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c386dd54fce258fc04e668126ae68589a0d92e03a90ea67881d1300f70fd6170" +dependencies = [ + "bytes", + "chrono", + "postgres-types", +] + +[[package]] +name = "pgwire" +version = "0.40.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7981cfde34009be689a05a30c497ad5fbb552531d3d54230b3627264ff1bc384" +dependencies = [ + "async-trait", + "aws-lc-rs", + "base64", + "bytes", + "chrono", + "derive-new", + "futures", + "hex", + "lazy-regex", + "md5", + "pg_interval", + "postgres-types", + "rand 0.10.1", + "rust_decimal", + "rustls-pki-types", + "ryu", + "serde", + "serde_json", + "smol_str", + "stringprep", + "thiserror 2.0.18", + "tokio", + "tokio-rustls", + "tokio-util", + "x509-certificate", +] + [[package]] name = "phf" version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" dependencies = [ - "phf_shared", + "phf_shared 0.12.1", +] + +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_shared 0.13.1", + "serde", ] [[package]] @@ -9641,6 +9820,15 @@ dependencies = [ "siphasher", ] +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + [[package]] name = "pico-args" version = "0.5.0" @@ -9814,6 +10002,39 @@ dependencies = [ "serde", ] +[[package]] +name = "postgres-protocol" +version = "0.6.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" +dependencies = [ + "base64", + "byteorder", + "bytes", + "fallible-iterator", + "hmac 0.13.0", + "md-5 0.11.0", + "memchr", + "rand 0.10.1", + "sha2 0.11.0", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" +dependencies = [ + "array-init", + "bytes", + "chrono", + "fallible-iterator", + "postgres-protocol", + "serde_core", + "serde_json", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -10182,7 +10403,7 @@ dependencies = [ "libc", "once_cell", "raw-cpuid", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "web-sys", "winapi", ] @@ -10820,7 +11041,7 @@ dependencies = [ "cfg-if", "getrandom 0.2.17", "libc", - "untrusted", + "untrusted 0.9.0", "windows-sys 0.52.0", ] @@ -11063,6 +11284,7 @@ dependencies = [ "borsh", "bytes", "num-traits", + "postgres-types", "rand 0.8.6", "rkyv", "serde", @@ -11234,7 +11456,7 @@ dependencies = [ "aws-lc-rs", "ring", "rustls-pki-types", - "untrusted", + "untrusted 0.9.0", ] [[package]] @@ -12161,6 +12383,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" @@ -12270,6 +12502,16 @@ dependencies = [ "recursive", ] +[[package]] +name = "sqlparser" +version = "0.62.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c6d1b651dc4edf07eead2a0c6c78016ce971bc2c10da5266861b13f25e7cec" +dependencies = [ + "log", + "recursive", +] + [[package]] name = "sqlx" version = "0.9.0" @@ -13134,6 +13376,32 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "tokio-postgres" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf 0.13.1", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand 0.10.1", + "socket2 0.6.4", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -13834,6 +14102,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" @@ -14133,6 +14407,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" @@ -14142,6 +14425,15 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.125" @@ -14151,6 +14443,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] @@ -14318,6 +14611,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" @@ -14887,6 +15187,25 @@ dependencies = [ "tap", ] +[[package]] +name = "x509-certificate" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca9eb9a0c822c67129d5b8fcc2806c6bc4f50496b420825069a440669bcfbf7f" +dependencies = [ + "bcder", + "bytes", + "chrono", + "der", + "hex", + "pem", + "ring", + "signature", + "spki", + "thiserror 2.0.18", + "zeroize", +] + [[package]] name = "x509-parser" version = "0.18.1" diff --git a/Cargo.toml b/Cargo.toml index 537ae81024..8af2a9d05e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ members = [ "core/connectors/sinks/mongodb_sink", "core/connectors/sinks/postgres_sink", "core/connectors/sinks/quickwit_sink", + "core/connectors/sinks/redshift_sink", "core/connectors/sinks/s3_sink", "core/connectors/sinks/stdout_sink", "core/connectors/sinks/surrealdb_sink", @@ -244,6 +245,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" @@ -290,6 +292,7 @@ simd-json = { version = "0.17.0", features = ["serde_impl"] } slab = "0.4.12" smallvec = "1.15" socket2 = "0.6.4" +sqlparser = "0.62.0" sqlx = { version = "0.9.0", features = [ "runtime-tokio", "tls-rustls", @@ -315,6 +318,7 @@ testcontainers = { version = "0.27.3", features = ["reusable-containers"] } testcontainers-modules = { version = "0.15.0", features = ["postgres", "http_wait"] } thiserror = "2.0.18" tokio = { version = "1.52.3", features = ["full"] } +tokio-postgres = "0.7.18" tokio-rustls = "0.26.4" tokio-tungstenite = { version = "0.29", features = ["rustls-tls-webpki-roots"] } tokio-util = { version = "0.7.18", features = ["compat"] } diff --git a/core/connectors/README.md b/core/connectors/README.md index 64bfc9a1f8..ca79e18f9c 100644 --- a/core/connectors/README.md +++ b/core/connectors/README.md @@ -88,6 +88,7 @@ Each sink should have its own, custom configuration, which is passed along with - **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 +- **Reshift Sink** - stores messages in Redshift warehouse tables via S3 as staging ## Source diff --git a/core/connectors/sinks/README.md b/core/connectors/sinks/README.md index 57ea055490..1a6006b4a1 100644 --- a/core/connectors/sinks/README.md +++ b/core/connectors/sinks/README.md @@ -17,6 +17,7 @@ Sink connectors are responsible for writing data from Iggy streams to external s | **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 | +| **redshift_sink** | Stores messages in Redshift warehouse tables with configurable schemas vis S3 as staging | The sink is represented by the single `Sink` trait, which defines the basic interface for all sink connectors. It provides methods for initializing the sink, writing data to external destination, and closing the sink. diff --git a/core/connectors/sinks/redshift_sink/Cargo.toml b/core/connectors/sinks/redshift_sink/Cargo.toml new file mode 100644 index 0000000000..d420c3becf --- /dev/null +++ b/core/connectors/sinks/redshift_sink/Cargo.toml @@ -0,0 +1,52 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "iggy_connector_redshift_sink" +version = "0.4.1-edge.1" +description = "Iggy Redshift sink connector for storing stream messages into Redshift warehouse via S3" +edition = "2024" +license = "Apache-2.0" +keywords = ["iggy", "messaging", "streaming", "redshift", "sink"] +categories = ["command-line-utilities", "warehouse", "network-programming"] +homepage = "https://iggy.apache.org" +documentation = "https://iggy.apache.org/docs" +repository = "https://github.com/apache/iggy" +readme = "../../README.md" +publish = false + +[lib] +crate-type = ["cdylib", "lib"] + +[dependencies] +arrow = { workspace = true } +async-trait = { workspace = true } +humantime = { workspace = true } +iggy_common = { workspace = true } +iggy_connector_sdk = { workspace = true } +parquet = { workspace = true } +rust-s3 = { workspace = true } +secrecy = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +sqlx = { workspace = true, features = ["runtime-tokio", "tls-rustls", "postgres", "chrono"] } +tokio = { workspace = true } +tracing = { workspace = true } +uuid = { workspace = true, features = ["v7"] } + +[dev-dependencies] +simd-json = { workspace = true } diff --git a/core/connectors/sinks/redshift_sink/README.md b/core/connectors/sinks/redshift_sink/README.md new file mode 100644 index 0000000000..01a9742706 --- /dev/null +++ b/core/connectors/sinks/redshift_sink/README.md @@ -0,0 +1,119 @@ +# Redshift Sink Connector + +Writes Apache Iggy stream messages into Amazon Redshift via S3-staged Parquet +files and a `COPY` load. + +Each connector batch is serialized to a Parquet file and uploaded to the +configured S3 bucket/prefix, then loaded into the target Redshift table with a +`COPY` statement. This makes S3 a staging area rather than a destination in +its own right — Redshift is the system of record for the data. + +Persistent load failures are at-most-once from the runtime's perspective: +messages may already be committed in Iggy before this connector exhausts its +write attempts, so failed loads are logged but not redelivered. + +## Configuration + +```toml +type = "sink" +key = "redshift" +enabled = true +version = 0 +name = "Redshift sink" +path = "../../target/release/libiggy_connector_redshift_sink" +verbose = false + +[[streams]] +stream = "user_events" +topics = ["users", "orders"] +schema = "json" +batch_length = 100 +poll_interval = "5ms" +consumer_group = "redshift_sink" + +[plugin_config] +connection_string = "postgresql://user:pass@localhost:5439/database" +target_table = "iggy_messages" +batch_size = 100 +max_connections = 10 +include_metadata = true +include_checksum = true +include_origin_timestamp = true +payload_format = "varbyte" +aws_access_key_id = "admin" +aws_secret_access_key = "password" +s3_bucket = "iggystaging" +s3_prefix = "iggy/messages" +s3_endpoint = "http://localhost:9000" +aws_region = "us-east-1" +archive = true +``` + +### Plugin Fields + +| Field | Required | Default | Description | +| --- | --- | --- | --- | +| `connection_string` | yes | — | Postgres-wire connection string used to reach the Redshift cluster and issue the `COPY` command. | +| `target_table` | yes | — | Destination Redshift table that batches are copied into. | +| `batch_size` | no | `100` | Number of messages buffered per Parquet file / `COPY` operation. | +| `max_connections` | no | `5` | Size of the connection pool used against Redshift. | +| `include_metadata` | no | `false` | Stores stream/topic/partition/offset/timestamp/schema fields alongside the payload. | +| `include_checksum` | no | `false` | Stores the Iggy message checksum. | +| `include_origin_timestamp` | no | `false` | Stores the original Iggy origin timestamp. | +| `payload_format` | no | `varbyte` | Encoding used for the payload column in the Parquet file. See **Payload Format** below. | +| `verbose_logging` | no | `false` | Enables verbose logging for debugging purposes. | +| `max_retries` | no | `3` | Maximum number of retries for failed `COPY` operations. `0` disables retries (only one attempt will be made) | +| `retry_delay` | no | `1s` | Delay in seconds between retry attempts. | +| `aws_access_key_id` | yes | — | AWS access key used for S3 staging. | +| `aws_secret_access_key` | yes | — | AWS secret key used for S3 staging. | +| `s3_bucket` | yes | — | S3 bucket that Parquet batch files are staged into before the Redshift `COPY`. | +| `s3_prefix` | yes | — | Key prefix under which staged Parquet files are written, e.g. `iggy/messages`. | +| `s3_endpoint` | no | — | Override endpoint for S3-compatible stores (e.g. MinIO). Omit for AWS S3 itself. | +| `aws_region` | yes | — | AWS region for the S3 bucket. | +| `archive` | no | `false` | See **Archiving Staged Files** below. | + +## Staging via S3 + +Redshift's `COPY` command loads from files, not from a live stream, so every +batch is first written out as a Parquet file and uploaded to +`s3:////...` before the `COPY` into `target_table` runs. +S3 is purely a staging area in this flow — it is not queried directly by +consumers of the data, and its cost is the price of getting bulk data into +Redshift efficiently rather than row-by-row. + +## Archiving Staged Files + +The `archive` field controls what happens to a batch's Parquet file **after** +it has been successfully loaded into Redshift: + +- `archive = true` — the file is kept, moved under an `archive` prefix + (i.e. `s3:///archive/...`) instead of being deleted. + Useful for replay, auditing, or downstream batch jobs that read Parquet + directly. +- `archive = false` — the file is deleted from S3 once the `COPY` succeeds, + since Redshift itself is now the source of truth for that data and the + staged copy has no further purpose. + +## Payload Format + +`payload_format` controls how the payload column is written in the staged +Parquet file, which in turn determines its type once loaded into Redshift: + +- Parquet has no dedicated JSON logical type, so a `payload_format = "json"` + payload is written as a Parquet `VARCHAR` (string), not a structured type. +- As a result, the column lands in Redshift as `VARCHAR`, not `SUPER`. +- To query the payload as structured data downstream, use Redshift's + `JSON_PARSE()` (or equivalent JSON functions) on the `VARCHAR` column at + query time rather than expecting a native `SUPER` column out of the box. + +## Stored Shape + +With metadata enabled, records contain: + +- `id`: original Iggy message id as numeric +- `iggy_stream`, `iggy_topic`, `iggy_partition_id`, `iggy_offset` +- `iggy_timestamp`, `iggy_origin_timestamp`, `iggy_checksum`, +- `payload`: encoded per `payload_format` (see above) + +The `messages_processed` counter reports valid records submitted to Redshift +via `COPY`. diff --git a/core/connectors/sinks/redshift_sink/config.toml b/core/connectors/sinks/redshift_sink/config.toml new file mode 100644 index 0000000000..4e695a0dde --- /dev/null +++ b/core/connectors/sinks/redshift_sink/config.toml @@ -0,0 +1,49 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +type = "sink" +key = "redshift" +enabled = true +version = 0 +name = "Redshift sink" +path = "../../target/release/libiggy_connector_redshift_sink" +verbose = false + +[[streams]] +stream = "user_events" +topics = ["users", "orders"] +schema = "json" +batch_length = 100 +poll_interval = "5ms" +consumer_group = "redshift_sink" + +[plugin_config] +connection_string = "postgresql://user:pass@localhost:5439/database" +target_table = "iggy_messages" +batch_size = 100 +max_connections = 10 +include_metadata = true +include_checksum = true +include_origin_timestamp = true +payload_format = "varbyte" +aws_access_key_id = "admin" +aws_secret_access_key = "password" +s3_bucket = "iggystaging" +s3_prefix = "iggy/messages" +s3_endpoint = "http://localhost:9000" +aws_region = "us-east-1" +archive = true diff --git a/core/connectors/sinks/redshift_sink/src/config.rs b/core/connectors/sinks/redshift_sink/src/config.rs new file mode 100644 index 0000000000..b0a4b30ca9 --- /dev/null +++ b/core/connectors/sinks/redshift_sink/src/config.rs @@ -0,0 +1,127 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::datatypes::DataType; +use iggy_connector_sdk::Error; +use secrecy::{ExposeSecret, SecretString}; + +/// Configuration for the Redshift Sink +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct RedshiftSinkConfig { + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub connection_string: SecretString, + pub target_table: String, + pub batch_size: Option, + pub max_connections: Option, + pub include_metadata: Option, + pub include_checksum: Option, + pub include_origin_timestamp: Option, + pub payload_format: Option, + pub verbose_logging: Option, + pub max_retries: Option, + pub retry_delay: Option, + /// aws_access_key_id and aws_secret_access_key MUST be provided + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub aws_access_key_id: SecretString, + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub aws_secret_access_key: SecretString, + pub s3_bucket: String, + pub s3_prefix: String, + pub s3_endpoint: Option, + pub aws_region: String, + /// Offers the option to archive staged S3 files after COPY + /// Defaults to deletion once COPY completes + /// Files are moved to different prefix within the same bucket + pub archive: Option, +} + +impl RedshiftSinkConfig { + pub fn validate(&self) -> Result<(), Error> { + let mut errors = String::new(); + + if self.connection_string.expose_secret().is_empty() { + errors.push_str("connection_string is empty\n"); + } + + if self.target_table.is_empty() { + errors.push_str(", target_table is empty\n"); + } + + if self.s3_bucket.is_empty() { + errors.push_str(", s3_bucket is empty\n"); + } + + if self.aws_region.is_empty() { + errors.push_str(", aws_region is empty\n"); + } + + // Validate AWS credentials: access keys must be provided + let has_access_key = !self.aws_access_key_id.expose_secret().is_empty(); + + let has_secret_key = !self.aws_secret_access_key.expose_secret().is_empty(); + + if !(has_access_key && has_secret_key) { + errors.push_str(", aws_access_key_id and aws_secret_access_key are empty\n"); + } + + if !errors.is_empty() { + Err(Error::InvalidConfigValue(errors)) + } else { + Ok(()) + } + } +} + +/// This connector supports: +/// 1. Byte -> which has VARBYTE as the Redshift equivalent +/// 2. Text -> which has VARCHAR as the Redshift equivalent +/// +/// We dont have Json because we are using parquet as a means to sink ingestion +/// As at the development of this connector there's no direct parquet type that matches JSON +/// For JSON needs Reshshift has SUPER(VARCHAR can be parsed by JSON_PARSE) +#[allow(unused)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum PayloadFormat { + Json, + Text, + #[default] + Varbyte, +} + +impl PayloadFormat { + pub fn from_config(s: Option<&str>) -> Self { + match s.map(|s| s.to_lowercase()).as_deref() { + Some("text") | Some("json") => PayloadFormat::Text, + _ => PayloadFormat::Varbyte, + } + } + + pub fn sql_type(&self) -> &'static str { + match self { + PayloadFormat::Varbyte => "VARBYTE", + PayloadFormat::Text | PayloadFormat::Json => "VARCHAR", + } + } + + pub fn arrow_type(&self) -> DataType { + match self { + PayloadFormat::Varbyte => DataType::Binary, + PayloadFormat::Text => DataType::Utf8, + PayloadFormat::Json => DataType::Utf8, + } + } +} diff --git a/core/connectors/sinks/redshift_sink/src/lib.rs b/core/connectors/sinks/redshift_sink/src/lib.rs new file mode 100644 index 0000000000..b70803ab5c --- /dev/null +++ b/core/connectors/sinks/redshift_sink/src/lib.rs @@ -0,0 +1,1182 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod config; + +use std::{str::FromStr, sync::Arc, time::Duration}; + +use arrow::{ + array::{ + ArrayRef, BinaryArray, Decimal256Array, Int32Array, Int64Array, RecordBatch, StringArray, + TimestampMicrosecondArray, + }, + datatypes::{DataType, Field, Schema, TimeUnit}, +}; +use async_trait::async_trait; +use humantime::Duration as HumanDuration; +use iggy_connector_sdk::{ + ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata, sink_connector, +}; +use parquet::arrow::ArrowWriter; +use s3::{Bucket, Region, creds::Credentials}; +use secrecy::ExposeSecret; +use sqlx::{AssertSqlSafe, Pool, Postgres, postgres::PgPoolOptions}; +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::config::{PayloadFormat, RedshiftSinkConfig}; + +sink_connector!(RedshiftSink); + +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_RETRY_DELAY: &str = "1s"; +const DEFAULT_MAX_CONNECTIONS: u32 = 5; +const DEFAULT_ARCHIVE_PREFIX: &str = "archive/messages"; + +#[derive(Debug)] +pub struct RedshiftSink { + pub id: u32, + config: RedshiftSinkConfig, + pool: Option>, + state: Mutex, + verbose: bool, + bucket: Option>, +} + +#[async_trait] +impl Sink for RedshiftSink { + async fn open(&mut self) -> Result<(), Error> { + tracing::info!( + sink_id = self.id, + table = %self.config.target_table, "opening Redshift sink connector" + ); + + self.connect().await?; + self.ensure_table_exists().await?; + Ok(()) + } + + async fn consume( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: MessagesMetadata, + messages: Vec, + ) -> Result<(), Error> { + tracing::debug!( + sink_id = self.id, + count = messages.len(), + "consuming messages" + ); + self.process_messages(topic_metadata, &messages_metadata, &messages) + .await + } + + async fn close(&mut self) -> Result<(), Error> { + tracing::info!(sink_id = self.id, "closing Redshift sink connector"); + + if let Some(pool) = self.pool.take() { + pool.close().await; + + tracing::debug!(sink_id = self.id, "database pool closed"); + } + + let state = self.state.lock().await; + + tracing::info!( + sink_id = self.id, + messages_processed = state.messages_processed, + batches_loaded = state.batches_loaded, + insertion_errors = state.insertion_errors, + "Redshift sink connector closed", + ); + + Ok(()) + } +} + +impl RedshiftSink { + pub fn new(id: u32, config: RedshiftSinkConfig) -> Self { + let verbose = config.verbose_logging.unwrap_or(false); + + Self { + id, + config, + pool: None, + state: Mutex::new(State::default()), + verbose, + bucket: None, + } + } + + async fn connect(&mut self) -> Result<(), Error> { + let max_connections = self + .config + .max_connections + .unwrap_or(DEFAULT_MAX_CONNECTIONS); + + let redacted = redact_connection_string(self.config.connection_string.expose_secret()); + + tracing::info!(max_connections, dsn = %redacted, "connecting to Redshift"); + + let pool = PgPoolOptions::new() + .max_connections(max_connections) + .connect(self.config.connection_string.expose_secret()) + .await + .map_err(|e| Error::InitError(format!("Failed to connect to Redshift: {e}")))?; + + sqlx::query("SELECT 1").execute(&pool).await.map_err(|e| { + tracing::error!("Tracing failed: {:#?}", e); + Error::InitError(format!("Warehouse connectivity test failed: {e}")) + })?; + + self.pool = Some(pool); + tracing::debug!("Redshift connection pool established"); + + let region = self.build_region()?; + + let credentials = Credentials::new( + Some(self.config.aws_access_key_id.expose_secret()), + Some(self.config.aws_secret_access_key.expose_secret()), + None, + None, + None, + ) + .map_err(|e| { + tracing::error!("Failed to create S3 credentials: {e}"); + Error::InvalidConfig + })?; + + let mut bucket = Bucket::new(&self.config.s3_bucket, region, credentials).map_err(|e| { + tracing::error!("Failed to create S3 bucket client: {e}"); + Error::InvalidConfig + })?; + + if self.config.s3_endpoint.is_some() { + bucket = bucket.with_path_style(); + } + + self.bucket = Some(bucket); + + tracing::info!("Redshift sink connector ready"); + + Ok(()) + } + + fn build_region(&self) -> Result { + if let Some(endpoint) = &self.config.s3_endpoint { + tracing::debug!(endpoint = %endpoint, "using custom S3 endpoint"); + Ok(Region::Custom { + region: self.config.aws_region.clone(), + endpoint: endpoint.clone(), + }) + } else { + Region::from_str(&self.config.aws_region).map_err(|_| Error::InvalidConfig) + } + } + + async fn ensure_table_exists(&self) -> Result<(), Error> { + let pool = self.get_pool()?; + + let table_name = &self.config.target_table; + let payload_type = self.payload_format().sql_type(); + + let (query, _) = self.build_create_table_sql()?; + + tracing::debug!("ensuring target table exists"); + + sqlx::query(AssertSqlSafe(query)) + .execute(pool) + .await + .map_err(|e| { + tracing::error!(error = %e); + Error::InitError(format!("Failed to create table '{table_name}': {e}")) + })?; + + tracing::info!(table = %table_name, payload_type, "target table ready"); + + Ok(()) + } + + async fn process_messages( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], + ) -> Result<(), Error> { + let batch_size = self.config.batch_size.unwrap_or(100) as usize; + + for batch in messages.chunks(batch_size) { + match self + .insert_batch(batch, topic_metadata, messages_metadata) + .await + { + Ok(()) => { + self.state.lock().await.batches_loaded += 1; + } + Err(e) => { + self.state.lock().await.insertion_errors += batch.len() as u64; + tracing::error!(error = %e, batch_size = batch.len(), "failed to insert batch"); + } + } + } + + let mut state = self.state.lock().await; + state.messages_processed += messages.len() as u64; + + if self.verbose { + tracing::info!( + sink_id = self.id, + total_processed = state.messages_processed, + batch_received = messages.len(), + table = %self.config.target_table, + batches_loaded = state.batches_loaded, + "processed message batch" + ); + } else { + tracing::debug!( + sink_id = self.id, + total_processed = state.messages_processed, + table = %self.config.target_table, + "processed message batch" + ); + } + + Ok(()) + } + + async fn insert_batch( + &self, + messages: &[ConsumedMessage], + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + ) -> Result<(), Error> { + if messages.is_empty() { + return Ok(()); + } + + let include_metadata = self.config.include_metadata.unwrap_or(true); + let include_checksum = self.config.include_checksum.unwrap_or(true); + let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); + let payload_format = self.payload_format(); + + let record_batch = create_record_batch( + topic_metadata, + messages_metadata, + messages, + include_metadata, + include_checksum, + include_origin_timestamp, + payload_format, + )?; + + let content = encode_parquet(&record_batch)?; + + tracing::debug!( + bytes = content.len(), + rows = record_batch.num_rows(), + "encoded parquet batch" + ); + + let s3_path = self.upload_parquet(&content).await?; + self.copy_parquet(&s3_path).await?; + self.archive_parquet(&s3_path).await?; + + tracing::info!(count = messages.len(), path = %s3_path, "batch inserted into Redshift"); + + Ok(()) + } + + async fn copy_parquet(&self, s3_path: &str) -> Result<(), Error> { + let max_retries = self.get_max_retries(); + let retry_delay = self.get_retry_delay(); + let sql = self.build_copy_sql(s3_path); + let pool = self.get_pool()?; + + tracing::debug!(table = %self.config.target_table, s3_path, "issuing Redshift COPY"); + + retry_with_backoff( + "redshift COPY", + max_retries, + retry_delay, + is_transient_error, + || async { + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .map(|_| ()) + }, + ) + .await?; + + tracing::debug!(table = %self.config.target_table, "Redshift COPY completed"); + + Ok(()) + } + + fn build_create_table_sql(&self) -> Result<(String, u32), Error> { + let table_name = &self.config.target_table; + let quoted_table = quote_identifier(table_name)?; + + let include_metadata = self.config.include_metadata.unwrap_or(true); + let include_checksum = self.config.include_checksum.unwrap_or(true); + let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); + let payload_type = self.payload_format().sql_type(); + + let mut params_per_row: u32 = 1; // id + + let mut query = + format!("CREATE TABLE IF NOT EXISTS {quoted_table} (id DECIMAL(39, 0) PRIMARY KEY"); + + if include_metadata { + query.push_str(", iggy_offset BIGINT, iggy_timestamp TIMESTAMPTZ, iggy_stream TEXT, iggy_topic TEXT, iggy_partition_id INTEGER"); + params_per_row += 5; + } + + if include_checksum { + query.push_str(", iggy_checksum VARCHAR"); + params_per_row += 1; + } + + if include_origin_timestamp { + query.push_str(", iggy_origin_timestamp TIMESTAMPTZ"); + params_per_row += 1; + } + + query.push_str(&format!(", payload {payload_type}")); + query.push_str(", created_at TIMESTAMPTZ DEFAULT GETDATE()"); + params_per_row += 2; + + Ok((query, params_per_row)) + } + + fn build_copy_sql(&self, s3_path: &str) -> String { + // Built via format! (not sqlx binds) because the Redshift/Pgwire endpoint here + // uses a Simple Query Handler that doesn't support prepared statements with binds. + let credentials = format!( + "CREDENTIALS 'ACCESS_KEY_ID={};SECRET_ACCESS_KEY={}'", + self.config.aws_access_key_id.expose_secret(), + self.config.aws_secret_access_key.expose_secret() + ); + + format!( + "COPY {} FROM '{}' {} FORMAT AS PARQUET REGION '{}';", + self.config.target_table, s3_path, credentials, self.config.aws_region + ) + } + + async fn upload_parquet(&self, content: &[u8]) -> Result { + let file_id = Uuid::now_v7(); + let key = build_s3_key(&self.config.s3_prefix, &format!("{file_id}.parquet")); + let bucket = self.get_bucket()?; + + tracing::debug!(key = %key, bytes = content.len(), "uploading parquet to S3"); + + let response = bucket.put_object(&key, content).await.map_err(|e| { + tracing::error!("Failed to upload to S3 key '{key}': {e}"); + Error::Storage(format!("S3 upload failed: {e}")) + })?; + + ensure_s3_status(response.status_code(), 200, "S3 upload")?; + + let path = format!("s3://{}{}", bucket.name(), key); + tracing::info!(path = %path, bytes = content.len(), "uploaded parquet to S3"); + + Ok(path) + } + + async fn archive_parquet(&self, key: &str) -> Result<(), Error> { + let old_key = key + .strip_prefix(&format!("s3://{}/", self.config.s3_bucket)) + .unwrap_or(key); + + if !self.get_archive() { + self.delete_object(old_key).await?; + tracing::info!(key = old_key, "deleted parquet file (archiving disabled)"); + return Ok(()); + } + + let bucket = self.get_bucket()?; + let prefix = self.config.s3_prefix.trim_matches('/'); + let archived_key = old_key.replacen(prefix, DEFAULT_ARCHIVE_PREFIX.trim_matches('/'), 1); + + tracing::debug!(from = old_key, to = %archived_key, "archiving parquet file"); + + let status_code = bucket + .copy_object_internal(old_key, &archived_key) + .await + .map_err(|e| { + tracing::error!(key = old_key, error = %e, "failed to copy object for archiving"); + Error::Storage(format!("S3 archiving failed: {e}")) + })?; + + ensure_s3_status(status_code, 200, "S3 archive copy")?; + + self.delete_object(old_key).await?; + tracing::info!(archived_to = %archived_key, "archived parquet file"); + + Ok(()) + } + + async fn delete_object(&self, key: &str) -> Result<(), Error> { + let bucket = self.get_bucket()?; + + let response = bucket.delete_object(key).await.map_err(|e| { + tracing::error!(key, error = %e, "failed to delete S3 object"); + Error::Storage(format!("S3 deleting failed: {e}")) + })?; + ensure_s3_status(response.status_code(), 204, "S3 object deletion")?; + + tracing::debug!(key, "deleted S3 object"); + Ok(()) + } + + fn get_pool(&self) -> Result<&Pool, Error> { + self.pool + .as_ref() + .ok_or_else(|| Error::InitError("Database not connected".to_string())) + } + + fn get_bucket(&self) -> Result<&Bucket, Error> { + let r = self + .bucket + .as_ref() + .ok_or_else(|| Error::InitError("Database not connected".to_string()))?; + + Ok(r) + } + + fn payload_format(&self) -> PayloadFormat { + PayloadFormat::from_config(self.config.payload_format.as_deref()) + } + + fn get_max_retries(&self) -> u32 { + self.config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES) + } + + fn get_retry_delay(&self) -> Duration { + self.config + .retry_delay + .as_deref() + .unwrap_or(DEFAULT_RETRY_DELAY) + .parse::() + .map(Into::into) + .unwrap_or_else(|_| Duration::from_secs(1)) + } + + fn get_archive(&self) -> bool { + self.config.archive.unwrap_or(false) + } +} + +#[derive(Debug, Default)] +struct State { + messages_processed: u64, + batches_loaded: u64, + insertion_errors: u64, +} + +/// Generic retry helper with linear backoff, used for transient warehouse errors. +async fn retry_with_backoff( + operation: &str, + max_retries: u32, + base_delay: Duration, + is_transient: impl Fn(&sqlx::Error) -> bool, + mut op: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempts = 0u32; + + loop { + match op().await { + Ok(value) => return Ok(value), + Err(e) => { + attempts += 1; + let transient = is_transient(&e); + + if !transient || attempts >= max_retries { + tracing::error!(operation, attempts, error = %e, "operation failed permanently"); + return Err(Error::CannotStoreData(format!( + "{operation} failed after {attempts} attempts: {e}" + ))); + } + + tracing::warn!(operation, attempts, max_retries, error = %e, "transient error, retrying"); + tokio::time::sleep(base_delay * attempts).await; + } + } + } +} + +fn encode_parquet(batch: &RecordBatch) -> Result, Error> { + let mut content = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut content, batch.schema(), None).map_err(|e| { + tracing::error!(error = %e, "failed to create parquet writer"); + Error::WriteFailure(format!("Failed to create parquet writer: {e}")) + })?; + + writer.write(batch).map_err(|e| { + tracing::error!(error = %e, "failed to write parquet batch"); + Error::WriteFailure(format!("Failed to write parquet: {e}")) + })?; + + writer.close().map_err(|e| { + tracing::error!(error = %e, "failed to close parquet writer"); + Error::WriteFailure(format!("Failed to close writer: {e}")) + })?; + + Ok(content) +} + +fn build_s3_key(prefix: &str, filename: &str) -> String { + if prefix.is_empty() { + format!("/{filename}") + } else { + format!("/{}/{filename}", prefix.trim_end_matches('/')) + } +} + +fn ensure_s3_status( + status: T, + expected: T, + context: &str, +) -> Result<(), Error> { + if status != expected { + tracing::error!(context, %status, "unexpected S3 response status"); + return Err(Error::Storage(format!( + "{context} failed with status {status}" + ))); + } + Ok(()) +} + +fn create_record_batch( + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], + include_metadata: bool, + include_checksum: bool, + include_origin_timestamp: bool, + payload_format: PayloadFormat, +) -> Result { + let mut fields = vec![Field::new("id", DataType::Decimal256(39, 0), false)]; + let mut columns: Vec = vec![id_column(messages)?]; + + if include_metadata { + let (mut metadata_fields, mut metadata_columns) = + metadata_columns(topic_metadata, messages_metadata, messages); + fields.append(&mut metadata_fields); + columns.append(&mut metadata_columns); + } + + if include_checksum { + fields.push(Field::new("iggy_checksum", DataType::Utf8, false)); + columns.push(checksum_column(messages)); + } + + if include_origin_timestamp { + fields.push(Field::new( + "iggy_origin_timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + )); + columns.push(origin_timestamp_column(messages)); + } + + fields.push(Field::new("payload", payload_format.arrow_type(), false)); + columns.push(payload_column(messages, payload_format)?); + + let schema = Arc::new(Schema::new(fields)); + let batch = + RecordBatch::try_new(schema, columns).map_err(|e| Error::CannotStoreData(e.to_string()))?; + + tracing::debug!( + rows = batch.num_rows(), + columns = batch.num_columns(), + "built record batch" + ); + + Ok(batch) +} + +fn id_column(messages: &[ConsumedMessage]) -> Result { + let ids = Decimal256Array::from_iter_values( + messages + .iter() + .map(|v| arrow::datatypes::i256::from_parts(v.id, 0)), + ) + .with_precision_and_scale(39, 0) + .map_err(|e| Error::CannotStoreData(e.to_string()))?; + + Ok(Arc::new(ids)) +} + +fn metadata_columns( + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], +) -> (Vec, Vec) { + let fields = vec![ + Field::new("iggy_offset", DataType::Int64, false), + Field::new( + "iggy_timestamp", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new("iggy_stream", DataType::Utf8, false), + Field::new("iggy_topic", DataType::Utf8, false), + Field::new("iggy_partition_id", DataType::Int32, false), + ]; + + let columns: Vec = vec![ + Arc::new(Int64Array::from_iter_values( + messages.iter().map(|v| v.offset as i64), + )), + Arc::new(TimestampMicrosecondArray::from_iter_values( + messages.iter().map(|v| v.timestamp as i64), + )), + Arc::new(StringArray::from_iter_values( + (0..messages.len()).map(|_| topic_metadata.stream.clone()), + )), + Arc::new(StringArray::from_iter_values( + (0..messages.len()).map(|_| topic_metadata.topic.clone()), + )), + Arc::new(Int32Array::from_iter_values( + (0..messages.len()).map(|_| messages_metadata.partition_id as i32), + )), + ]; + + (fields, columns) +} + +fn checksum_column(messages: &[ConsumedMessage]) -> ArrayRef { + Arc::new(StringArray::from_iter_values( + messages.iter().map(|v| v.checksum.to_string()), + )) +} + +fn origin_timestamp_column(messages: &[ConsumedMessage]) -> ArrayRef { + Arc::new(TimestampMicrosecondArray::from_iter_values( + messages.iter().map(|v| v.origin_timestamp as i64), + )) +} + +fn payload_column(messages: &[ConsumedMessage], format: PayloadFormat) -> Result { + match format { + PayloadFormat::Varbyte => { + let values: Vec> = messages + .iter() + .map(|v| v.payload.clone().try_to_bytes()) + .collect::>()?; + let slices: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect(); + Ok(Arc::new(BinaryArray::from_vec(slices))) + } + PayloadFormat::Text => { + let values: Vec = messages + .iter() + .map(|v| { + let bytes = v.payload.try_to_bytes()?; + String::from_utf8(bytes).map_err(|_| Error::InvalidTextPayload) + }) + .collect::>()?; + Ok(Arc::new(StringArray::from_iter_values(values.iter()))) + } + PayloadFormat::Json => { + let values: Vec = messages + .iter() + .map(|v| { + let bytes = v.payload.try_to_bytes()?; + + Ok(serde_json::from_slice::(&bytes) + .map_err(|_| Error::InvalidJsonPayload)? + .to_string()) + }) + .collect::>()?; + Ok(Arc::new(StringArray::from_iter_values(values.iter()))) + } + } +} + +fn redact_connection_string(conn_str: &str) -> String { + // Guard against very short strings + const PREVIEW_LEN: usize = 3; + + if let Some(scheme_end) = conn_str.find("://") { + let scheme = &conn_str[..scheme_end + 3]; + let rest = &conn_str[scheme_end + 3..]; + + // Stop preview at the first sensitive boundary + let safe_end = rest + .find([':', '@', '?', '/']) + .unwrap_or(rest.len()) + .min(PREVIEW_LEN); + + let preview = &rest[..safe_end]; + return format!("{scheme}{preview}***"); + } + + let preview: String = conn_str.chars().take(3).collect(); + format!("{preview}***") +} + +fn is_transient_error(e: &sqlx::Error) -> bool { + match e { + sqlx::Error::Io(_) => true, + sqlx::Error::PoolTimedOut => true, + sqlx::Error::PoolClosed => false, + sqlx::Error::Protocol(_) => false, + sqlx::Error::Database(db_err) => db_err.code().is_some_and(|code| { + matches!( + code.as_ref(), + "40001" | "40P01" | "57P01" | "57P02" | "57P03" | "08000" | "08003" | "08006" + ) + }), + _ => false, + } +} + +fn quote_identifier(name: &str) -> Result { + if name.is_empty() { + return Err(Error::InitError("Table name cannot be empty".to_string())); + } + if name.contains('\0') { + return Err(Error::InitError( + "Table name cannot contain null characters".to_string(), + )); + } + let escaped = name.replace('"', "\"\""); + Ok(format!("\"{escaped}\"")) +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use iggy_connector_sdk::{Payload, Schema}; + use secrecy::SecretString; + + use super::*; + + fn test_config( + include_checksum: bool, + include_origin_timestamp: bool, + include_metadata: bool, + ) -> RedshiftSinkConfig { + RedshiftSinkConfig { + connection_string: SecretString::from("postgresql://localhost/db"), + target_table: "messages".to_string(), + batch_size: Some(100), + max_connections: None, + include_metadata: Some(include_metadata), + include_checksum: Some(include_checksum), + include_origin_timestamp: Some(include_origin_timestamp), + payload_format: None, + verbose_logging: None, + max_retries: None, + retry_delay: None, + aws_access_key_id: SecretString::from("admin"), + aws_secret_access_key: SecretString::from("password"), + s3_bucket: "iggymessages".into(), + s3_prefix: "iggy/messages".into(), + s3_endpoint: None, + aws_region: "us-east-1".into(), + archive: None, + } + } + + fn test_topic_metadata() -> TopicMetadata { + TopicMetadata { + stream: "test_stream".to_string(), + topic: "test_topic".to_string(), + } + } + + fn test_messages_metadata() -> MessagesMetadata { + MessagesMetadata { + partition_id: 7, + current_offset: 0, + schema: Schema::Json, + } + } + + fn test_message(payload: Payload) -> ConsumedMessage { + ConsumedMessage { + id: 42, + offset: 9, + checksum: 123, + timestamp: 1_767_225_600_000_000, + origin_timestamp: 1_700_000_000_000_001, + headers: None, + payload, + } + } + + fn json_payload(value: serde_json::Value) -> Payload { + let mut bytes = serde_json::to_vec(&value).expect("Failed to serialize JSON"); + Payload::Json(simd_json::to_owned_value(&mut bytes).expect("Failed to parse JSON")) + } + + #[test] + fn given_empty_connection_string_should_error() { + let mut config = test_config(false, false, false); + config.connection_string = SecretString::default(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_target_table_should_error() { + let mut config = test_config(false, false, false); + config.target_table = String::new(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_s3_bucket_should_error() { + let mut config = test_config(false, false, false); + config.s3_bucket = String::new(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_aws_region_should_error() { + let mut config = test_config(false, false, false); + config.aws_region = String::new(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_aws_access_key_id_should_error() { + let mut config = test_config(false, false, false); + config.aws_access_key_id = SecretString::default(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_empty_aws_secret_access_key_should_error() { + let mut config = test_config(false, false, false); + config.aws_secret_access_key = SecretString::default(); + + assert!(config.validate().is_err()); + } + + #[test] + fn given_json_format_should_return_text() { + assert_eq!( + PayloadFormat::from_config(Some("json")), + PayloadFormat::Text + ); + assert_eq!( + PayloadFormat::from_config(Some("JSON")), + PayloadFormat::Text + ); + } + + #[test] + fn given_text_format_should_return_text() { + assert_eq!( + PayloadFormat::from_config(Some("text")), + PayloadFormat::Text + ); + assert_eq!( + PayloadFormat::from_config(Some("TEXT")), + PayloadFormat::Text + ); + } + + #[test] + fn given_bytea_or_unknown_format_should_return_bytea() { + assert_eq!( + PayloadFormat::from_config(Some("bytea")), + PayloadFormat::Varbyte + ); + assert_eq!( + PayloadFormat::from_config(Some("unknown")), + PayloadFormat::Varbyte + ); + assert_eq!(PayloadFormat::from_config(None), PayloadFormat::Varbyte); + } + + #[test] + fn given_payload_format_should_return_correct_sql_type() { + assert_eq!(PayloadFormat::Varbyte.sql_type(), "VARBYTE"); + assert_eq!(PayloadFormat::Json.sql_type(), "VARCHAR"); + assert_eq!(PayloadFormat::Text.sql_type(), "VARCHAR"); + } + + #[test] + fn given_payload_format_should_return_correct_arrow_type() { + assert_eq!( + PayloadFormat::Varbyte.arrow_type(), + arrow::datatypes::DataType::Binary + ); + assert_eq!( + PayloadFormat::Json.arrow_type(), + arrow::datatypes::DataType::Utf8 + ); + assert_eq!( + PayloadFormat::Text.arrow_type(), + arrow::datatypes::DataType::Utf8 + ); + } + + #[test] + fn given_all_options_enabled_should_build_full_create_query() { + let sink = RedshiftSink::new(1, test_config(true, true, true)); + let (query, param_count) = sink + .build_create_table_sql() + .expect("Failed to build create query"); + + assert!(query.contains("CREATE TABLE IF NOT EXISTS \"messages\"")); + assert!(query.contains("iggy_offset")); + assert!(query.contains("iggy_timestamp")); + assert!(query.contains("iggy_stream")); + assert!(query.contains("iggy_topic")); + assert!(query.contains("iggy_partition_id")); + assert!(query.contains("iggy_checksum")); + assert!(query.contains("iggy_origin_timestamp")); + assert!(query.contains("payload")); + assert!(query.contains("created_at")); + assert_eq!(param_count, 10); + } + + #[test] + fn given_all_options_enabled_should_build_full_parquet() { + let payload = json_payload(serde_json::json!({"name": "Bebeto", "active": true})); + let message = test_message(payload); + let record_batch = create_record_batch( + &test_topic_metadata(), + &test_messages_metadata(), + &[message], + true, + true, + true, + PayloadFormat::Varbyte, + ) + .expect("Failed to create record batch"); + + assert_eq!(record_batch.num_rows(), 1); + assert_eq!(record_batch.num_columns(), 9); + + let columns = record_batch.schema(); + let columns: HashSet<&str> = columns.fields().iter().map(|f| f.name().as_ref()).collect(); + + let expected_columns: HashSet<&str> = HashSet::from([ + "id", + "iggy_offset", + "iggy_timestamp", + "iggy_stream", + "iggy_topic", + "iggy_partition_id", + "iggy_checksum", + "iggy_origin_timestamp", + "payload", + "created_at", + ]); + + assert_eq!(columns.difference(&expected_columns).count(), 0); + } + + #[test] + fn given_metadata_disabled_should_build_minimal_create_query() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + let (query, param_count) = sink + .build_create_table_sql() + .expect("Failed to build create query"); + + assert!(query.contains("CREATE TABLE IF NOT EXISTS \"messages\"")); + assert!(!query.contains("iggy_offset")); + assert!(!query.contains("iggy_timestamp")); + assert!(!query.contains("iggy_stream")); + assert!(!query.contains("iggy_topic")); + assert!(!query.contains("iggy_partition_id")); + assert!(!query.contains("iggy_checksum")); + assert!(!query.contains("iggy_origin_timestamp")); + assert!(query.contains("payload")); + assert!(query.contains("created_at")); + assert_eq!(param_count, 3); + } + + #[test] + fn given_metadata_disabled_should_build_minimal_parquet() { + let payload = json_payload(serde_json::json!({"name": "Bebeto", "active": true})); + let message = test_message(payload); + let record_batch = create_record_batch( + &test_topic_metadata(), + &test_messages_metadata(), + &[message], + false, + false, + false, + PayloadFormat::Varbyte, + ) + .expect("Failed to create record batch"); + + assert_eq!(record_batch.num_rows(), 1); + assert_eq!(record_batch.num_columns(), 2); + + let columns = record_batch.schema(); + let columns: HashSet<&str> = columns.fields().iter().map(|f| f.name().as_ref()).collect(); + + let expected_columns: HashSet<&str> = HashSet::from(["id", "payload", "created_at"]); + + assert_eq!(columns.difference(&expected_columns).count(), 0); + } + + #[test] + fn given_microseconds_should_parse_timestamp_correctly() { + let record_batch = create_record_batch( + &test_topic_metadata(), + &test_messages_metadata(), + &[test_message(json_payload(serde_json::json!({})))], + true, + false, + false, + PayloadFormat::Varbyte, + ) + .expect("Failed to create record batch"); + + let timestamp_col = record_batch + .column(2) + .as_any() + .downcast_ref::() + .expect("Failed to downcast to Timestamp Microsecond array"); + + let timestamp = timestamp_col.value(0); + + assert_eq!(timestamp, 1_767_225_600_000_000); + } + + #[test] + fn given_default_config_should_use_default_archive() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + assert!(sink.get_archive()); + } + + #[test] + fn given_archive_enabled_should_use_archive() { + let mut sink = RedshiftSink::new(1, test_config(false, false, false)); + sink.config.archive = Some(true); + + assert!(sink.get_archive()); + } + + #[test] + fn given_default_config_should_use_default_retries() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + assert_eq!(sink.get_max_retries(), DEFAULT_MAX_RETRIES); + } + + #[test] + fn given_custom_retries_should_use_custom_value() { + let mut config = test_config(false, false, false); + config.max_retries = Some(5); + let sink = RedshiftSink::new(1, config); + assert_eq!(sink.get_max_retries(), 5); + } + + #[test] + fn given_default_config_should_use_default_retry_delay() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + assert_eq!(sink.get_retry_delay(), Duration::from_secs(1)); + } + + #[test] + fn given_custom_retry_delay_should_parse_humantime() { + let mut config = test_config(false, false, false); + config.retry_delay = Some("500ms".to_string()); + let sink = RedshiftSink::new(1, config); + assert_eq!(sink.get_retry_delay(), Duration::from_millis(500)); + } + + #[test] + fn given_verbose_logging_enabled_should_set_verbose_flag() { + let mut config = test_config(false, false, false); + config.verbose_logging = Some(true); + let sink = RedshiftSink::new(1, config); + assert!(sink.verbose); + } + + #[test] + fn given_verbose_logging_disabled_should_not_set_verbose_flag() { + let sink = RedshiftSink::new(1, test_config(false, false, false)); + assert!(!sink.verbose); + } + + #[test] + fn given_connection_string_with_credentials_should_redact() { + let conn = "postgres://redshift:redshift@localhost:5432/db"; + let redacted = redact_connection_string(conn); + assert_eq!(redacted, "postgres://red***"); + } + + #[test] + fn given_connection_string_without_scheme_should_redact() { + let conn = "localhost:5432/db"; + let redacted = redact_connection_string(conn); + assert_eq!(redacted, "loc***"); + } + + #[test] + fn given_postgresql_scheme_should_redact() { + let conn = "postgresql://admin:secret123@db.example.com:5432/mydb"; + let redacted = redact_connection_string(conn); + assert_eq!(redacted, "postgresql://adm***"); + } + + #[test] + fn given_special_chars_in_identifier_should_escape() { + let result = quote_identifier("table\"name").expect("Failed to quote"); + assert_eq!(result, "\"table\"\"name\""); + } + + #[test] + fn given_empty_identifier_should_fail() { + let result = quote_identifier(""); + assert!(result.is_err()); + } + + #[test] + fn given_null_char_in_identifier_should_fail() { + let result = quote_identifier("table\0name"); + assert!(result.is_err()); + } + + #[test] + fn given_normal_identifier_should_quote() { + let result = quote_identifier("my_table").expect("Failed to quote"); + assert_eq!(result, "\"my_table\""); + } + + #[test] + fn given_identifier_with_spaces_should_quote() { + let result = quote_identifier("my table").expect("Failed to quote"); + assert_eq!(result, "\"my table\""); + } + + #[test] + fn given_identifier_with_sql_injection_should_escape() { + let result = quote_identifier("messages\"; DROP TABLE users; --").expect("Failed to quote"); + assert_eq!(result, "\"messages\"\"; DROP TABLE users; --\""); + } +} diff --git a/core/integration/Cargo.toml b/core/integration/Cargo.toml index 62168b3f24..f966a03113 100644 --- a/core/integration/Cargo.toml +++ b/core/integration/Cargo.toml @@ -34,6 +34,7 @@ login-session = ["dep:zbus-secret-service-keyring-store"] vsr = ["iggy/vsr"] [dependencies] +arrow = { workspace = true } assert_cmd = { workspace = true } async-trait = { workspace = true } base64 = { workspace = true } @@ -63,6 +64,8 @@ keyring-core = { workspace = true } lazy_static = { workspace = true } libc = { workspace = true } mongodb = { workspace = true } +parquet = { workspace = true } +pgwire = { workspace = true } predicates = { workspace = true } rand = { workspace = true } rcgen = { workspace = true } @@ -82,6 +85,7 @@ serde_json = { workspace = true } serial_test = { workspace = true } server = { workspace = true } socket2 = { workspace = true } +sqlparser = { workspace = true } sqlx = { workspace = true } sysinfo = { workspace = true } tempfile = { workspace = true } @@ -89,6 +93,7 @@ test-case = { workspace = true } testcontainers = { workspace = true } testcontainers-modules = { workspace = true } tokio = { workspace = true, features = ["full", "test-util"] } +tokio-postgres = { workspace = true } toml = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/core/integration/tests/connectors/fixtures/mod.rs b/core/integration/tests/connectors/fixtures/mod.rs index 885867de96..6af6f886fa 100644 --- a/core/integration/tests/connectors/fixtures/mod.rs +++ b/core/integration/tests/connectors/fixtures/mod.rs @@ -27,6 +27,7 @@ mod influxdb; mod mongodb; mod postgres; mod quickwit; +mod redshift; mod s3; mod surrealdb; mod wiremock; @@ -84,3 +85,8 @@ pub use surrealdb::{ SurrealDbSinkRawFixture, }; pub use wiremock::{WireMockDirectFixture, WireMockWrappedFixture}; + +pub use redshift::{ + RedshiftSinkByteaFixture, RedshiftSinkFixture, RedshiftSinkJsonFixture, + RedshiftSinkNoArchiveFixture, +}; 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..9b7847fe9f --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/container.rs @@ -0,0 +1,267 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use integration::harness::TestBinaryError; +use pgwire::tokio::process_socket; +use sqlx::{Pool, Postgres, postgres::PgPoolOptions}; +use testcontainers::{ + ContainerAsync, GenericImage, ImageExt, + core::{IntoContainerPort, WaitFor, wait::HttpWaitStrategy}, + runners::AsyncRunner, +}; +use testcontainers_modules::postgres; +use tokio::{net::TcpListener, task::JoinHandle}; + +use crate::connectors::fixtures::{ + self, + redshift::redshift_mock::handler::{RedshiftHandler, RedshiftHandlerFactory}, +}; + +const MINIO_IMAGE: &str = "docker.io/minio/minio"; +const MINIO_TAG: &str = "RELEASE.2025-09-07T16-13-09Z"; +const MINIO_PORT: u16 = 9000; +const MINIO_CONSOLE_PORT: u16 = 9001; +const POSTGRES_PORT: u16 = 5432; + +pub const MINIO_ACCESS_KEY: &str = "admin"; +pub const MINIO_SECRET_KEY: &str = "password"; +pub const MINIO_BUCKET: &str = "iggystaging"; +pub const DEFAULT_SINK_TABLE: &str = "iggy_messages"; +pub const STAGING_REGION: &str = "us-east-1"; +pub const STAGING_PREFIX: &str = "iggy/messages"; + +pub const ENV_SINK_CONNECTION_STRING: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_CONNECTION_STRING"; +pub const ENV_SINK_TARGET_TABLE: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_TARGET_TABLE"; +pub const ENV_SINK_PAYLOAD_FORMAT: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_PAYLOAD_FORMAT"; +pub const ENV_SINK_STAGING_ACCESS_KEY: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_AWS_ACCESS_KEY_ID"; +pub const ENV_SINK_STAGING_SECRET: &str = + "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_AWS_SECRET_ACCESS_KEY"; +pub const ENV_SINK_S3_BUCKET: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_S3_BUCKET"; +pub const ENV_SINK_S3_PREFIX: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_S3_PREFIX"; +pub const ENV_SINK_S3_ENDPOINT: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_S3_ENDPOINT"; +pub const ENV_SINK_STAGING_REGION: &str = "IGGY_CONNECTORS_SINK_REDSHIFT_PLUGIN_CONFIG_AWS_REGION"; +pub const TEST_NETWORK: &str = "apache-iggy"; +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(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(TEST_NETWORK) + .with_container_name(container_name) + .with_env_var("MINIO_ROOT_USER", MINIO_ACCESS_KEY) + .with_env_var("MINIO_ROOT_PASSWORD", MINIO_SECRET_KEY) + .with_cmd(vec!["server", "/data", "--console-address", ":9001"]) + .with_mapped_port(0, MINIO_PORT.tcp()) + .with_mapped_port(0, MINIO_CONSOLE_PORT.tcp()) + .start() + .await + .map_err(|error| TestBinaryError::FixtureSetup { + fixture_type: "MinioContainer".to_string(), + message: format!("Failed to start container: {error}"), + })?; + + tracing::info!("Started MinIO container"); + + let mapped_port = container + .ports() + .await + .map_err(|error| TestBinaryError::FixtureSetup { + fixture_type: "MinioContainer".to_string(), + message: format!("Failed to get ports: {error}"), + })? + .map_to_host_port_ipv4(MINIO_PORT) + .ok_or_else(|| TestBinaryError::FixtureSetup { + fixture_type: "MinioContainer".to_string(), + message: "No mapping for MinIO port".to_string(), + })?; + + let endpoint = format!("http://localhost:{mapped_port}"); + tracing::info!("MinIO container available at {endpoint}"); + + Ok(Self { + container, + endpoint, + }) + } +} + +/// Base container management for PostgreSQL fixtures. +pub struct PostgresContainer { + #[allow(dead_code)] + container: ContainerAsync, + pub connection_string: String, +} + +impl PostgresContainer { + pub async fn start() -> Result { + let container = postgres::Postgres::default() + .with_container_name(fixtures::unique_container_name("postgres")) + .start() + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "PostgresContainer".to_string(), + message: format!("Failed to start container: {e}"), + })?; + + let host_port = container + .get_host_port_ipv4(POSTGRES_PORT) + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "PostgresContainer".to_string(), + message: format!("Failed to get port: {e}"), + })?; + + let connection_string = format!("postgres://postgres:postgres@localhost:{host_port}"); + + Ok(Self { + container, + connection_string, + }) + } + + pub async fn create_pool(&self) -> Result, TestBinaryError> { + PgPoolOptions::new() + .max_connections(1) + .connect(&self.connection_string) + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "PostgresContainer".to_string(), + message: format!("Failed to connect: {e}"), + }) + } +} + +pub struct RedshiftContainer { + #[allow(dead_code)] + accept_task: JoinHandle<()>, + pub connection_string: String, +} + +impl RedshiftContainer { + pub async fn start( + target_connection: String, + s3_endpoint: String, + ) -> Result { + let (pg_client, connection) = + tokio_postgres::connect(&target_connection, tokio_postgres::NoTls) + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftContainer".into(), + message: e.to_string(), + })?; + + tokio::spawn(async move { + if let Err(e) = connection.await { + panic!("{}", e.to_string()) + } + }); + + let redshql = RedshiftHandler::new(pg_client, s3_endpoint); + + let factory = Arc::new(RedshiftHandlerFactory { + handler: Arc::new(redshql), + }); + + let listener = + TcpListener::bind("127.0.0.1:0") + .await + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftMockContainer".to_string(), + message: format!("bind failed: {e}"), + })?; + + let host_port = listener + .local_addr() + .map_err(|e| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftMockContainer".to_string(), + message: format!("failed to get local address: {e}"), + })? + .port(); + + let accept_task = tokio::spawn(async move { + loop { + match listener.accept().await { + Ok((incoming_socket, _addr)) => { + let factory_ref = factory.clone(); + + tokio::spawn(async move { + if let Err(e) = process_socket(incoming_socket, None, factory_ref).await + { + panic!("{}", e.to_string()) + } + }); + } + + Err(e) => { + panic!("{}", e.to_string()) + } + } + } + }); + + Ok(Self { + accept_task, + connection_string: format!("postgres://postgres@localhost:{host_port}/postgres"), + }) + } +} + +/// Payload format for sink connector. +#[derive(Debug, Clone, Copy, Default)] +pub enum SinkPayloadFormat { + #[default] + Bytea, + Text, +} + +/// Schema format for message encoding. +#[derive(Debug, Clone, Copy, Default)] +pub enum SinkSchema { + #[default] + Json, + Raw, +} diff --git a/core/integration/tests/connectors/fixtures/redshift/mod.rs b/core/integration/tests/connectors/fixtures/redshift/mod.rs new file mode 100644 index 0000000000..4599ee0d0c --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/mod.rs @@ -0,0 +1,26 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod container; +mod redshift_mock; +mod sink; + +pub use container::{MinioContainer, PostgresContainer, RedshiftContainer}; +pub use sink::{ + RedshiftSinkByteaFixture, RedshiftSinkFixture, RedshiftSinkJsonFixture, + RedshiftSinkNoArchiveFixture, +}; diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/copy.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/copy.rs new file mode 100644 index 0000000000..37114bb718 --- /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(format!("Not a copy statement")))? + } + + 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..7d44430df4 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/create.rs @@ -0,0 +1,270 @@ +// 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 +} + +#[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(format!("Not a create statement")))? + } + + 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; + + // 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)?; + } + 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, + }) +} + +fn parse_ident_list(parser: &mut Parser) -> Result, ParserError> { + parser.expect_token(&Token::LParen)?; + let idents = parser.parse_comma_separated(Parser::parse_identifier)?; + parser.expect_token(&Token::RParen)?; + Ok(idents) +} + +fn parse_column_list(parser: &mut Parser) -> Result, ParserError> { + parser.expect_token(&Token::LParen)?; + let mut columns = Vec::new(); + + loop { + let name = parser.parse_identifier()?; + let data_type = parser.parse_data_type()?; + + let encoding = None; + let mut not_null = false; + let mut default = None; + let identity = None; + let mut primary_key = false; + let mut references = None; + + // Column constraints are unordered too — loop until comma or close paren + loop { + match parser.peek_token().token { + Token::Comma | Token::RParen => break, + _ => {} + } + let word = expect_word(parser)?; + match word.to_uppercase().as_str() { + "NOT" => { + parser.expect_keyword(Keyword::NULL)?; + not_null = true; + } + "NULL" => not_null = false, + "DEFAULT" => default = Some(parser.parse_expr()?), + "PRIMARY" => { + parser.expect_keyword(Keyword::KEY)?; + primary_key = true; + } + "REFERENCES" => { + references = Some(parser.parse_object_name(false)?); + } + // column-level DISTKEY/SORTKEY flags — tolerate, table-level fields win + "DISTKEY" | "SORTKEY" => {} + unknown => { + Err(ParserError::ParserError(format!( + "Unsupported column constraint: {unknown}" + )))?; + } + } + } + + columns.push(RedshiftColumnDef { + name, + data_type, + encoding, + not_null, + default, + identity, + primary_key, + references, + }); + + if parser.consume_token(&Token::Comma) { + continue; + } + parser.expect_token(&Token::RParen)?; + break; + } + + Ok(columns) +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs new file mode 100644 index 0000000000..139bf918f7 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/handler.rs @@ -0,0 +1,305 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use async_trait::async_trait; +use pgwire::{ + api::{ + ClientInfo, ClientPortalStore, PgWireServerHandlers, Type, + portal::Portal, + query::{ExtendedQueryHandler, SimpleQueryHandler}, + results::{DescribePortalResponse, DescribeStatementResponse, FieldInfo, Response, Tag}, + stmt::{NoopQueryParser, StoredStatement}, + store::PortalStore, + }, + error::{ErrorInfo, PgWireError, PgWireResult}, +}; +use sqlparser::{ast::Statement, dialect::RedshiftSqlDialect, parser::Parser}; +use tokio_postgres::{Client as PgClient, Statement as PgStatement}; + +use crate::connectors::fixtures::redshift::redshift_mock::{ + copy::try_parse_redshift_copy, + create::try_parse_redshift_create_table, + load::{ + S3Client, ToPgError, execute_create_table, execute_s3_copy, execute_select, split_s3_uri, + }, +}; + +/// Statement failed to parse under the Redshift dialect. +pub const SYNTAX_ERROR: &str = "26000"; +/// Parsed, but uses a construct we don't implement (custom COPY/CREATE +/// extensions, unsupported statement kinds, etc.). +pub const FEATURE_NOT_SUPPORTED: &str = "42601"; +/// The Postgres connection backing this mock failed outright (as +/// opposed to Postgres returning a well-formed DB error). +pub const CONNECTION_EXCEPTION: &str = "08000"; +/// Empty/missing statement. +pub const INVALID_QUERY: &str = "42601"; +/// Statement kind we recognize but intentionally don't support. +pub const WARNING_UNSUPPORTED: &str = "01000"; + +pub struct RedshiftHandlerFactory { + pub handler: Arc, +} + +impl PgWireServerHandlers for RedshiftHandlerFactory { + fn simple_query_handler(&self) -> Arc { + self.handler.clone() + } + + fn extended_query_handler(&self) -> Arc { + self.handler.clone() + } +} + +pub struct RedshiftHandler { + pg: PgClient, + s3_endpoint: String, +} + +impl RedshiftHandler { + pub fn new(pg: PgClient, s3_endpoint: String) -> Self { + Self { pg, s3_endpoint } + } + + /// Prepares `sql` against the backing Postgres connection for describe + /// purposes. Returns `Ok(None)` for statement kinds (currently just + /// `COPY`) that describe to zero fields rather than going through an + /// unsupported `PREPARE`. + async fn prepare_describable(&self, sql: &str) -> PgWireResult> { + let dialect = RedshiftSqlDialect {}; + let statements = parse_redshift_sql(&dialect, sql)?; + + if matches!(statements.first(), Some(Statement::Copy { .. })) { + return Ok(None); + } + + self.pg + .prepare(sql) + .await + .map(Some) + .map_err(|e| map_pg_client_error(e, CONNECTION_EXCEPTION)) + } +} + +#[async_trait] +impl ExtendedQueryHandler for RedshiftHandler { + type Statement = String; + type QueryParser = NoopQueryParser; + + fn query_parser(&self) -> Arc { + Arc::new(NoopQueryParser::new()) + } + + async fn do_query( + &self, + _client: &mut C, + portal: &Portal, + _max_rows: usize, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + let query = &portal.statement.statement; + + if query.trim().is_empty() { + return Ok(Response::EmptyQuery); + } + + execute_statement(query, &self.pg, &self.s3_endpoint).await + } + + async fn do_describe_statement( + &self, + _client: &mut C, + stmt: &StoredStatement, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + let Some(prepared) = self.prepare_describable(&stmt.statement).await? else { + return Ok(DescribeStatementResponse::new(vec![], vec![])); + }; + + let param_types: Vec = prepared.params().to_vec(); + + let fields: Vec = prepared + .columns() + .iter() + .map(|col| { + FieldInfo::new( + col.name().to_owned(), + None, + None, + col.type_().clone(), + pgwire::api::results::FieldFormat::Text, + ) + }) + .collect(); + + Ok(DescribeStatementResponse::new(param_types, fields)) + } + + async fn do_describe_portal( + &self, + _client: &mut C, + portal: &Portal, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + let Some(prepared) = self + .prepare_describable(&portal.statement.statement) + .await? + else { + return Ok(DescribePortalResponse::new(vec![])); + }; + + let fields: Vec = prepared + .columns() + .iter() + .enumerate() + .map(|(idx, col)| { + FieldInfo::new( + col.name().to_owned(), + None, + None, + col.type_().clone(), + portal.result_column_format.format_for(idx), + ) + }) + .collect(); + + Ok(DescribePortalResponse::new(fields)) + } +} + +#[async_trait] +impl SimpleQueryHandler for RedshiftHandler { + async fn do_query(&self, _client: &mut C, query: &str) -> PgWireResult> + where + C: ClientInfo + ClientPortalStore + Unpin + Send + Sync, + C::PortalStore: PortalStore, + { + if query.trim().is_empty() { + return Ok(vec![Response::EmptyQuery]); + } + + Ok(vec![ + execute_statement(query, &self.pg, &self.s3_endpoint).await?, + ]) + } +} + +/// Dispatches a single SQL statement to the appropriate executor based on +/// its parsed kind. Kept intentionally thin — each branch delegates to a +/// dedicated function so individual statement kinds can be read (and +/// tested) in isolation. +async fn execute_statement( + query: &str, + pg: &PgClient, + s3_endpoint: &str, +) -> PgWireResult { + let dialect = RedshiftSqlDialect {}; + let statements = parse_redshift_sql(&dialect, query)?; + + match statements.first() { + Some(Statement::Query(_)) => execute_select(query, pg).await, + Some(Statement::CreateTable(_)) => execute_create(query, pg).await, + Some(Statement::Copy { .. }) => execute_copy(query, pg, s3_endpoint).await, + Some(other) => Err(pg_warning( + WARNING_UNSUPPORTED, + format!("Unsupported: {other:?}"), + )), + None => Err(pg_error(INVALID_QUERY, "Invalid query")), + } +} + +async fn execute_create(query: &str, pg: &PgClient) -> PgWireResult { + let dialect = RedshiftSqlDialect {}; + let parser = Parser::new(&dialect) + .try_with_sql(query) + .map_err(|e| pg_error(SYNTAX_ERROR, format!("Unsupported: {e:?}")))?; + + let create = try_parse_redshift_create_table(parser) + .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; + + execute_create_table(create, pg).await +} + +async fn execute_copy(query: &str, pg: &PgClient, s3_endpoint: &str) -> PgWireResult { + let dialect = RedshiftSqlDialect {}; + let parser = Parser::new(&dialect) + .try_with_sql(query) + .map_err(|e| pg_error(SYNTAX_ERROR, format!("Unsupported: {e:?}")))?; + + let r_copy = try_parse_redshift_copy(parser) + .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; + + let (bucket_name, prefix) = split_s3_uri(&r_copy.s3_uri).map_err(|e| e.to_pg_wire_error())?; + + let s3_client = S3Client::new( + &bucket_name, + s3_endpoint, + &r_copy.access_key_id, + &r_copy.secret_access_key, + &r_copy.region, + ) + .await + .map_err(|e| e.to_pg_wire_error())?; + + let rows = execute_s3_copy(&r_copy, pg, s3_client, &bucket_name, &prefix) + .await + .map_err(|e| pg_error(FEATURE_NOT_SUPPORTED, format!("Unsupported: {e:?}")))?; + + Ok(Response::Execution(Tag::new("copy").with_rows(rows))) +} + +/// Builds a `PgWireError::UserError` with severity `ERROR`. Replaces the +/// repeated `PgWireError::UserError(Box::new(ErrorInfo::new(...)))` calls. +fn pg_error(code: &str, message: impl std::fmt::Display) -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".into(), + code.into(), + message.to_string(), + ))) +} + +/// Same as [`pg_error`] but with severity `WARNING`, for statement kinds +/// we recognize but choose not to support. +fn pg_warning(code: &str, message: impl std::fmt::Display) -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "WARNING".into(), + code.into(), + message.to_string(), + ))) +} + +/// Maps a `tokio_postgres::Error` to a `PgWireError`, preserving the +/// upstream SQLSTATE/message when Postgres itself produced the error, and +/// falling back to `fallback_code` for connection-level failures. +fn map_pg_client_error(err: tokio_postgres::Error, fallback_code: &str) -> PgWireError { + match err.as_db_error() { + Some(db_err) => pg_error(db_err.code().code(), db_err.message()), + None => pg_error(fallback_code, format!("connection failed: {err}")), + } +} + +fn parse_redshift_sql(dialect: &RedshiftSqlDialect, sql: &str) -> PgWireResult> { + Parser::parse_sql(dialect, sql).map_err(|e| pg_error(SYNTAX_ERROR, e)) +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs new file mode 100644 index 0000000000..857d393c5a --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/load.rs @@ -0,0 +1,603 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::{fmt::Write, sync::Arc}; + +use arrow::{ + array::{ + Array, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Decimal256Array, + Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, TimestampMicrosecondArray, + }, + datatypes::DataType, +}; + +use bytes::Bytes; +use futures::{StreamExt, pin_mut, stream}; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use pgwire::{ + api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag}, + error::{ErrorInfo, PgWireError, PgWireResult}, +}; + +use s3::{Bucket, Region, creds::Credentials}; +use sqlparser::ast::DataType as SDataType; +use sqlx::types::chrono; +use tokio_postgres::{ + Client as PgClient, GenericClient, + binary_copy::BinaryCopyInWriter, + types::{ToSql, Type as PgType}, +}; + +use crate::connectors::fixtures::redshift::redshift_mock::{ + copy::{CopyFormat, RedshiftCopy}, + create::{RedshiftCreateTable, TableKind}, +}; + +pub async fn execute_select(sql: &str, pg: &PgClient) -> PgWireResult { + let value = pg + .client() + .query_one_scalar::(sql, &[]) + .await + .map_err(|e| pg_wire_error(&e, "Query failed"))?; + + let schema = Arc::new(vec![FieldInfo::new( + "column".into(), + None, + None, + PgType::INT4, + FieldFormat::Text, + )]); + + let schema_ref = schema.clone(); + + let row_stream = stream::iter(std::iter::once(value)).map(move |v| { + let mut encoder = DataRowEncoder::new(schema_ref.clone()); + encoder.encode_field(&v)?; + + Ok(encoder.take_row()) + }); + + Ok(Response::Query(QueryResponse::new(schema, row_stream))) +} + +pub async fn execute_create_table( + create: RedshiftCreateTable, + pg: &PgClient, +) -> PgWireResult { + let create_kw = match create.table_kind { + TableKind::Regular => "CREATE TABLE", + TableKind::Temp => "CREATE TEMP TABLE", + // Postgres accepts (and ignores) the LOCAL keyword per the SQL + // standard, so this round-trips fine. + TableKind::LocalTemp => "CREATE LOCAL TEMPORARY TABLE", + }; + + let mut sql = String::from(create_kw); + sql.push(' '); + + if create.if_not_exists { + sql.push_str("IF NOT EXISTS "); + } + + write!(sql, "{} (", create.table.to_string()).unwrap(); + + let fields = create.columns.iter().fold(vec![], |mut acc, v| { + let mut f = format!("{} {}", v.name, map_type(&v.data_type)); + + if let Some(id) = &v.identity { + f.push_str(&format!( + " GENERATED BY DEFAULT AS IDENTITY (START WITH {} INCREMENT BY {})", + id.seed, id.step + )); + } else if let Some(default) = &v.default { + let default = format!("{}", default).replace("GETDATE", "NOW"); + f.push_str(&format!(" DEFAULT {default}")); + } + + if v.primary_key { + f.push_str(" PRIMARY KEY"); + } + + if v.not_null { + f.push_str(" NOT NULL"); + } + + if let Some(target) = &v.references { + f.push_str(&format!(" REFERENCES {target}")); + } + + acc.push(f); + + acc + }); + + sql.push_str(&fields.join(", ")); + sql.push_str(");"); + + let rows_affected = pg + .execute(&sql, &[]) + .await + .map_err(|e| pg_wire_error(&e, "CREATE TABLE failed"))?; + + Ok(Response::Execution( + Tag::new("create table").with_rows(rows_affected as usize), + )) +} + +pub async fn execute_s3_copy( + copy: &RedshiftCopy, + pg: &PgClient, + s3_client: S3Client, + bucket_name: &str, + prefix: &str, +) -> PgWireResult { + let table_name = copy.table.to_string(); + let mut failed_objects = 0usize; + + match ©.format { + CopyFormat::Parquet => { + let existing_columns = fetch_table_columns(pg, ©.table.to_string()) + .await + .map_err(|e| e.to_pg_wire_error())?; + + tracing::info!( + "target table '{table_name}' {}", + if existing_columns.is_some() { + "exists" + } else { + "does not exist" + } + ); + + let bytes = Bytes::from( + s3_client + .get_object(&prefix) + .await + .map_err(|e| e.to_pg_wire_error())?, + ); + + tracing::info!("File '{}' read", prefix); + + let cols: Vec = match existing_columns { + Some(cols) => strip_created_at(cols), + None => { + let inferred = + infer_parquet_schema(bytes.clone()).map_err(|e| e.to_pg_wire_error())?; + + let columns_sql = inferred + .iter() + .map(|f| format!("{} {}", f.name.to_lowercase(), f.pg_type.name())) + .collect::>() + .join(", "); + + create_table(pg, &table_name, &columns_sql) + .await + .map_err(|e| e.to_pg_wire_error())?; + + strip_created_at(inferred) + } + }; + + match load_one_object(pg, ©.table.to_string(), &cols, bytes).await { + Ok(n) => { + tracing::info!("{n} records stored"); + } + Err(e) => { + // Simplification vs real Redshift: MAXERROR there counts bad *rows* + // across the whole load; here each failed *file* counts as one error. + // Good enough for a test double where you're usually asserting + // "load either fully succeeds or trips MAXERROR", not exact counts. + failed_objects += 1; + tracing::error!("[copy] error loading s3://{bucket_name}/{prefix}: {e:#}"); + if failed_objects > copy.max_error as usize { + Err(format!("MAXERROR ({}) exceeded", copy.max_error).to_pg_wire_error())? + } + } + } + } + other => Err(format!("{:?} unsupported", other).to_pg_wire_error())?, + }; + + Ok(1) +} + +async fn fetch_table_columns(pg: &PgClient, table: &str) -> Result>, String> { + let rows = pg + .query( + "SELECT column_name, udt_name FROM information_schema.columns \ + WHERE table_name = $1 ORDER BY ordinal_position", + &[&table], + ) + .await + .map_err(|e| e.to_string())?; + + if rows.is_empty() { + return Ok(None); + } + + let rows: Result, String> = rows + .into_iter() + .map(|row| { + let name: String = row.get(0); + let udt: String = row.get(1); + Ok(ColumnDef { + name, + pg_type: udt_name_to_type(&udt)?, + }) + }) + .collect(); + + Ok(Some(rows?)) +} + +async fn create_table(pg: &PgClient, table: &str, columns: &str) -> Result { + let sql = format!("CREATE TABLE {} ({});", table, columns); + + tracing::info!("{sql}"); + + let result = pg + .client() + .execute(&sql, &[]) + .await + .map_err(|e| e.to_string())?; + + Ok(result as usize) +} + +async fn load_one_object( + pg: &PgClient, + table: &str, + columns: &[ColumnDef], + bytes: Bytes, +) -> Result { + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes) + .map_err(|e| e.to_string())? + .build() + .map_err(|e| e.to_string())?; + + let col_list = columns + .iter() + .map(|v| format!("\"{}\"", v.name)) + .collect::>() + .join(", "); + + let types = columns + .iter() + .map(|v| v.pg_type.clone()) + .collect::>(); + + let copy_sql = format!("COPY {table} ({col_list}) FROM STDIN BINARY"); + let sink = pg.copy_in(©_sql).await.map_err(|e| e.to_string())?; + + tracing::info!("COPY FROM STDIN started"); + + let writer = BinaryCopyInWriter::new(sink, &types); + + pin_mut!(writer); + + let mut n = 0usize; + + for batch in reader { + let batch = batch.map_err(|e| e.to_string())?; + for row_idx in 0..batch.num_rows() { + let row_values = extract_row(&batch, row_idx, columns)?; + + let refs: Vec<&(dyn ToSql + Sync)> = row_values + .iter() + .map(|v| v.as_ref() as &(dyn ToSql + Sync)) + .collect(); + + writer + .as_mut() + .write(&refs) + .await + .map_err(|e| e.to_string())?; + + n += 1; + } + } + + writer.finish().await.map_err(|e| e.to_string())?; + + Ok(n) +} + +struct ColumnDef { + name: String, + pg_type: PgType, +} + +macro_rules! scalar_column { + ($array:expr, $arr_ty:ty, $val_ty:ty, $row:expr, $conv:expr) => {{ + let a = $array + .as_any() + .downcast_ref::<$arr_ty>() + .ok_or_else(|| format!("expected {} array", stringify!($arr_ty)))?; + + if a.is_null($row) { + Box::new(None::<$val_ty>) as Box + } else { + let conv: fn(_) -> $val_ty = $conv; + Box::new(conv(a.value($row))) as Box + } + }}; +} + +/// Only covers common scalar types. Extend as your Parquet exports need more — +/// this deliberately doesn't try to handle structs, lists, or decimals up front. +fn extract_row<'a>( + batch: &'a RecordBatch, + row: usize, + columns: &'a [ColumnDef], +) -> Result>, String> { + let mut out = Vec::with_capacity(columns.len()); + + for (i, col) in columns.iter().enumerate() { + let array = batch.column(i); + + let value: Box = match array.data_type() { + DataType::Utf8 => { + scalar_column!(array, StringArray, String, row, |v: &str| v.to_string()) + } + DataType::Int64 => { + scalar_column!(array, Int64Array, i64, row, |v: i64| v) + } + DataType::Int32 => scalar_column!(array, Int32Array, i32, row, |v: i32| v), + DataType::Float64 => scalar_column!(array, Float64Array, f64, row, |v: f64| v), + DataType::Boolean => scalar_column!(array, BooleanArray, bool, row, |v: bool| v), + DataType::Date32 => scalar_column!(array, Date32Array, i32, row, |v: i32| v), + DataType::Binary => { + scalar_column!(array, BinaryArray, Vec, row, |v: &[u8]| v.to_vec()) + } + DataType::Decimal128(_, _) => { + let a = array + .as_any() + .downcast_ref::() + .ok_or("expected Decimal256Array")?; + if a.is_null(row) { + Box::new(None::) + } else { + Box::new(a.value(row).to_string()) + } + } + // Downcast to fit i256 serialization constraints + // Decimal(39, 0) > becomes Decimal256 on arrow + // Decimal256 serialization will require extra handling + DataType::Decimal256(_, _) => { + let a = array + .as_any() + .downcast_ref::() + .ok_or("expected Decimal256Array")?; + if a.is_null(row) { + Box::new(None::) + } else { + let raw = a.value(row).to_string(); + Box::new(raw) + } + } + DataType::Timestamp(_, _) => { + let a = array + .as_any() + .downcast_ref::() + .ok_or("expected TimestampMicrosecondArray")?; + if a.is_null(row) { + Box::new(None::>) + } else { + let micros = a.value(row); + let dt = chrono::DateTime::::from_timestamp_micros(micros) + .ok_or("Invalid timestamp")?; + Box::new(dt) + } + } + other => Err(format!( + "unsupported parquet column type {other:?} for column {}", + col.name + ))?, + }; + out.push(value); + } + + Ok(out) +} + +fn udt_name_to_type(udt: &str) -> Result { + Ok(match udt { + "int2" => PgType::INT2, + "int4" => PgType::INT4, + "int8" => PgType::INT8, + "float4" => PgType::FLOAT4, + "float8" => PgType::FLOAT8, + // Numeric serialiation requires extra work + // Safe to use VARCHAR + "numeric" => PgType::VARCHAR, + "bool" => PgType::BOOL, + "text" | "varchar" | "bpchar" => PgType::TEXT, + "timestamp" => PgType::TIMESTAMP, + "timestamptz" => PgType::TIMESTAMPTZ, + "date" => PgType::DATE, + "jsonb" => PgType::JSONB, + "bytea" => PgType::BYTEA, + other => Err(format!("unsupported column type for COPY target: {other}"))?, + }) +} + +fn arrow_to_type(a_type: &DataType) -> Result { + match a_type { + DataType::Boolean => Ok(PgType::BOOL), + DataType::Binary | DataType::FixedSizeBinary(_) => Ok(PgType::BYTEA), + DataType::Float64 => Ok(PgType::FLOAT8), + DataType::Float32 | DataType::Float16 => Ok(PgType::FLOAT4), + DataType::Int64 => Ok(PgType::INT8), + DataType::Int32 => Ok(PgType::INT4), + DataType::Decimal128(_, _) => Ok(PgType::VARCHAR), + DataType::Decimal256(_, _) => Ok(PgType::VARCHAR), + DataType::Utf8 => Ok(PgType::TEXT), + DataType::Date32 => Ok(PgType::DATE), + DataType::Timestamp(_, _) => Ok(PgType::TIMESTAMPTZ), + other => Err(format!("Unsuppoerted type: {}", other)), + } +} + +fn map_type(dt: &SDataType) -> String { + match dt { + // Redshift's SUPER (semi-structured) has no PG equivalent for + // your emulator's purposes -> JSON is the closest usable stand-in. + SDataType::Custom(name, _) if name.to_string().eq_ignore_ascii_case("SUPER") => { + "JSON".to_string() + } + SDataType::Custom(name, mods) if name.to_string().eq_ignore_ascii_case("VARBYTE") => { + "BYTEA".to_string() + } + SDataType::Decimal(_) => "VARCHAR".into(), + other => other.to_string(), + } +} + +fn infer_parquet_schema(bytes: Bytes) -> Result, String> { + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).map_err(|e| e.to_string())?; + + reader + .schema() + .fields() + .iter() + .map(|v| { + Ok(ColumnDef { + name: v.name().into(), + pg_type: arrow_to_type(v.data_type())?, + }) + }) + .collect() +} + +fn strip_created_at(cols: Vec) -> Vec { + cols.into_iter() + .filter(|v| v.name != "created_at") + .collect() +} + +/// S3 +#[allow(unused)] +#[derive(Clone)] +pub struct S3Client { + bucket_name: String, + inner: Box, +} + +impl S3Client { + pub async fn new( + bucket_name: &str, + s3_endpoint: &str, + access_key: &str, + secret_key: &str, + region: &str, + ) -> Result { + let region = Region::Custom { + region: region.into(), + endpoint: s3_endpoint.into(), + }; + + let credentials = Credentials::new(Some(access_key), Some(secret_key), None, None, None) + .map_err(|e| e.to_string())?; + + let bucket = Bucket::new(bucket_name, region, credentials) + .map_err(|e| format!("failed to setup bucket: {e}"))? + .with_path_style(); + + Ok(S3Client { + bucket_name: bucket_name.into(), + inner: bucket, + }) + } + + pub async fn get_object(&self, key: &str) -> Result, String> { + tracing::info!( + "Downloading object '{}' from bucket '{}'", + key, + self.bucket_name + ); + + let response = self + .inner + .get_object(key) + .await + .map_err(|e| e.to_string())?; + + if response.status_code() != 200 { + tracing::error!( + "S3 get object returned status {}: {}", + response.status_code(), + String::from_utf8_lossy(response.as_slice()) + ); + return Err(format!( + "S3 get_object failed with status {}", + response.status_code() + )); + } + + tracing::info!( + "Retrieved {} bytes to s3://{}/{}", + response.bytes().len(), + self.inner.name(), + key + ); + + Ok(response.bytes().to_vec()) + } +} + +pub fn split_s3_uri(uri: &str) -> Result<(String, String), String> { + let rest = uri + .strip_prefix("s3://") + .ok_or(format!("not an s3:// URI"))?; + + match rest.split_once('/') { + Some((b, p)) => Ok((b.to_string(), p.to_string())), + None => Err(format!("s3 URI missing key/prefix: {uri}")), + } +} + +pub trait ToPgError { + fn to_pg_wire_error(self) -> PgWireError; +} + +impl ToPgError for T +where + T: Into, +{ + fn to_pg_wire_error(self) -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".into(), + "GG000".into(), + self.into(), + ))) + } +} + +fn pg_wire_error(e: &tokio_postgres::Error, context: &str) -> PgWireError { + match e.as_db_error() { + Some(db_err) => PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".into(), + db_err.code().code().to_string(), + db_err.message().to_string(), + ))), + None => PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".into(), + "XX000".into(), + format!("{context}: {e}"), + ))), + } +} diff --git a/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs new file mode 100644 index 0000000000..1d9ba47977 --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/redshift_mock/mod.rs @@ -0,0 +1,56 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use sqlparser::{ + parser::{Parser, ParserError}, + tokenizer::{Token, Word}, +}; + +pub mod copy; +pub mod create; +pub mod handler; +pub mod load; + +pub fn expect_word(parser: &mut Parser) -> Result { + match parser.next_token().token { + Token::Word(Word { value, .. }) => Ok(value), + Token::SemiColon => Ok("SemiColon".into()), + other => Err(ParserError::ParserError(format!( + "expected identifier, got {other:?}" + ))), + } +} + +pub fn parse_string_literal(parser: &mut Parser) -> Result { + match parser.next_token().token { + Token::SingleQuotedString(s) => Ok(s), + other => Err(ParserError::ParserError(format!( + "expected string literal, got {other:?}" + ))), + } +} + +pub fn parse_number_literal(parser: &mut Parser) -> Result { + match parser.next_token().token { + Token::Number(s, _) => s + .parse() + .map_err(|_| ParserError::ParserError(format!("bad number: {s}"))), + other => Err(ParserError::ParserError(format!( + "expected number, got {other:?}" + ))), + } +} diff --git a/core/integration/tests/connectors/fixtures/redshift/sink.rs b/core/integration/tests/connectors/fixtures/redshift/sink.rs new file mode 100644 index 0000000000..fc4a8a0f6c --- /dev/null +++ b/core/integration/tests/connectors/fixtures/redshift/sink.rs @@ -0,0 +1,423 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::{collections::HashMap, time::Duration}; + +use async_trait::async_trait; +use integration::harness::{TestBinaryError, TestFixture}; +use sqlx::{Pool, Postgres}; + +use crate::connectors::fixtures::{ + self, + redshift::{ + MinioContainer, PostgresContainer, RedshiftContainer, + container::{ + DEFAULT_POLL_ATTEMPTS, DEFAULT_POLL_INTERVAL_MS, DEFAULT_SINK_TABLE, + DEFAULT_TEST_STREAM, DEFAULT_TEST_TOPIC, ENV_SINK_ARCHIVE, ENV_SINK_CONNECTION_STRING, + ENV_SINK_PATH, ENV_SINK_PAYLOAD_FORMAT, ENV_SINK_S3_BUCKET, ENV_SINK_S3_ENDPOINT, + ENV_SINK_S3_PREFIX, ENV_SINK_STAGING_ACCESS_KEY, ENV_SINK_STAGING_REGION, + ENV_SINK_STAGING_SECRET, ENV_SINK_STREAMS_0_CONSUMER_GROUP, ENV_SINK_STREAMS_0_SCHEMA, + ENV_SINK_STREAMS_0_STREAM, ENV_SINK_STREAMS_0_TOPICS, ENV_SINK_TARGET_TABLE, + MINIO_ACCESS_KEY, MINIO_BUCKET, MINIO_SECRET_KEY, STAGING_PREFIX, STAGING_REGION, + SinkPayloadFormat, SinkSchema, + }, + }, +}; + +pub struct RedshiftSinkFixture { + #[allow(dead_code)] + minio: MinioContainer, + #[allow(dead_code)] + redshift: RedshiftContainer, + postgres: PostgresContainer, + payload_format: SinkPayloadFormat, + schema: SinkSchema, + pub minio_endpoint: String, +} + +impl RedshiftSinkFixture { + /// Assertions read from here — never from `redshift`. + pub async fn target_pool(&self) -> Result, TestBinaryError> { + self.postgres.create_pool().await + } + + /// Fetch rows from the sink table with polling until expected count is reached. + /// + /// Returns an error if the expected count is not reached within the poll attempts. + pub async fn fetch_rows_as( + &self, + pool: &Pool, + query: &str, + expected_count: usize, + ) -> Result, TestBinaryError> + where + T: Send + Unpin + for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow>, + { + let mut rows = Vec::new(); + for _ in 0..DEFAULT_POLL_ATTEMPTS { + let result = sqlx::query_as::<_, T>(sqlx::AssertSqlSafe(query)) + .fetch_all(pool) + .await; + if let Ok(fetched) = result { + rows = fetched; + if rows.len() >= expected_count { + return Ok(rows); + } + } else if let Err(e) = result { + tracing::error!("{e}"); + } + + tokio::time::sleep(Duration::from_millis(DEFAULT_POLL_INTERVAL_MS)).await; + } + Err(TestBinaryError::InvalidState { + message: format!( + "Expected {} rows but got {} after {} poll attempts", + expected_count, + rows.len(), + DEFAULT_POLL_ATTEMPTS + ), + }) + } +} + +#[async_trait] +impl TestFixture for RedshiftSinkFixture { + async fn setup() -> Result { + let postgres = PostgresContainer::start().await?; + + let minio_name = fixtures::unique_container_name("minio-redshift"); + + let minio = MinioContainer::start(&minio_name).await?; + + create_bucket(&minio.endpoint)?; + + let redshift = + RedshiftContainer::start(postgres.connection_string.clone(), minio.endpoint.clone()) + .await?; + + Ok(Self { + minio_endpoint: minio.endpoint.clone(), + minio, + redshift, + postgres, + payload_format: SinkPayloadFormat::default(), + schema: SinkSchema::default(), + }) + } + + fn connectors_runtime_envs(&self) -> std::collections::HashMap { + let mut envs = HashMap::new(); + + envs.insert( + ENV_SINK_CONNECTION_STRING.to_string(), + self.redshift.connection_string.clone(), + ); + envs.insert( + ENV_SINK_TARGET_TABLE.to_string(), + DEFAULT_SINK_TABLE.to_string(), + ); + + envs.insert( + ENV_SINK_STAGING_ACCESS_KEY.to_string(), + MINIO_ACCESS_KEY.to_string(), + ); + + envs.insert( + ENV_SINK_STAGING_SECRET.to_string(), + MINIO_SECRET_KEY.to_string(), + ); + + envs.insert(ENV_SINK_S3_BUCKET.to_string(), MINIO_BUCKET.to_string()); + envs.insert(ENV_SINK_S3_PREFIX.to_string(), STAGING_PREFIX.to_string()); + + envs.insert( + ENV_SINK_S3_ENDPOINT.to_string(), + self.minio_endpoint.clone(), + ); + + envs.insert( + ENV_SINK_STAGING_REGION.to_string(), + STAGING_REGION.to_string(), + ); + + envs.insert( + ENV_SINK_STREAMS_0_STREAM.to_string(), + DEFAULT_TEST_STREAM.to_string(), + ); + envs.insert( + ENV_SINK_STREAMS_0_TOPICS.to_string(), + format!("[{}]", DEFAULT_TEST_TOPIC), + ); + envs.insert( + ENV_SINK_STREAMS_0_CONSUMER_GROUP.to_string(), + "test".to_string(), + ); + + envs.insert( + ENV_SINK_PATH.to_string(), + "../../target/debug/libiggy_connector_redshift_sink".to_string(), + ); + + let schema_str = match self.schema { + SinkSchema::Json => "json", + SinkSchema::Raw => "raw", + }; + envs.insert( + ENV_SINK_STREAMS_0_SCHEMA.to_string(), + schema_str.to_string(), + ); + + let format_str = match self.payload_format { + SinkPayloadFormat::Bytea => "bytea", + SinkPayloadFormat::Text => "text", + }; + envs.insert(ENV_SINK_PAYLOAD_FORMAT.to_string(), format_str.to_string()); + + envs + } +} + +/// redshift sink fixture for bytea payload format. +pub struct RedshiftSinkByteaFixture { + inner: RedshiftSinkFixture, +} + +impl std::ops::Deref for RedshiftSinkByteaFixture { + type Target = RedshiftSinkFixture; + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[async_trait] +impl TestFixture for RedshiftSinkByteaFixture { + async fn setup() -> Result { + let postgres = PostgresContainer::start().await?; + + let minio_name = fixtures::unique_container_name("minio-redshift"); + + let minio = MinioContainer::start(&minio_name).await?; + + create_bucket(&minio.endpoint)?; + + let redshift = + RedshiftContainer::start(postgres.connection_string.clone(), minio.endpoint.clone()) + .await?; + + Ok(Self { + inner: RedshiftSinkFixture { + minio_endpoint: minio.endpoint.clone(), + minio, + redshift, + postgres, + payload_format: SinkPayloadFormat::Bytea, + schema: SinkSchema::Raw, + }, + }) + } + + fn connectors_runtime_envs(&self) -> HashMap { + self.inner.connectors_runtime_envs() + } +} + +/// redshift sink fixture for bytea payload format. +pub struct RedshiftSinkJsonFixture { + inner: RedshiftSinkFixture, +} + +impl std::ops::Deref for RedshiftSinkJsonFixture { + type Target = RedshiftSinkFixture; + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[async_trait] +impl TestFixture for RedshiftSinkJsonFixture { + async fn setup() -> Result { + let postgres = PostgresContainer::start().await?; + + let minio_name = fixtures::unique_container_name("minio-redshift"); + + let minio = MinioContainer::start(&minio_name).await?; + + create_bucket(&minio.endpoint)?; + + let redshift = + RedshiftContainer::start(postgres.connection_string.clone(), minio.endpoint.clone()) + .await?; + + Ok(Self { + inner: RedshiftSinkFixture { + minio_endpoint: minio.endpoint.clone(), + minio, + redshift, + postgres, + payload_format: SinkPayloadFormat::Text, + schema: SinkSchema::Json, + }, + }) + } + + fn connectors_runtime_envs(&self) -> HashMap { + self.inner.connectors_runtime_envs() + } +} + +/// redshift sink fixture for bytea payload format. +pub struct RedshiftSinkNoArchiveFixture { + inner: RedshiftSinkFixture, +} + +impl RedshiftSinkNoArchiveFixture { + pub fn confirm_empty_bucket(&self) -> Result { + bucket_empty(&self.minio_endpoint) + } +} + +impl std::ops::Deref for RedshiftSinkNoArchiveFixture { + type Target = RedshiftSinkFixture; + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[async_trait] +impl TestFixture for RedshiftSinkNoArchiveFixture { + async fn setup() -> Result { + let postgres = PostgresContainer::start().await?; + + let minio_name = fixtures::unique_container_name("minio-redshift"); + + let minio = MinioContainer::start(&minio_name).await?; + + create_bucket(&minio.endpoint)?; + + let redshift = + RedshiftContainer::start(postgres.connection_string.clone(), minio.endpoint.clone()) + .await?; + + Ok(Self { + inner: RedshiftSinkFixture { + minio_endpoint: minio.endpoint.clone(), + minio, + redshift, + postgres, + payload_format: SinkPayloadFormat::Text, + schema: SinkSchema::Json, + }, + }) + } + + fn connectors_runtime_envs(&self) -> HashMap { + let mut envs = self.inner.connectors_runtime_envs(); + + let schema_str = match self.schema { + SinkSchema::Json => "json", + SinkSchema::Raw => "raw", + }; + envs.insert( + ENV_SINK_STREAMS_0_SCHEMA.to_string(), + schema_str.to_string(), + ); + + envs.insert(ENV_SINK_ARCHIVE.to_string(), false.to_string()); + + let format_str = match self.payload_format { + SinkPayloadFormat::Bytea => "bytea", + SinkPayloadFormat::Text => "text", + }; + + envs.insert(ENV_SINK_PAYLOAD_FORMAT.to_string(), format_str.to_string()); + + envs + } +} + +fn create_bucket(minio_endpoint: &str) -> Result<(), TestBinaryError> { + use std::process::Command; + + let host = minio_endpoint.trim_start_matches("http://"); + let mc_host = format!("http://{}:{}@{}", MINIO_ACCESS_KEY, MINIO_SECRET_KEY, host); + + let output = Command::new("docker") + .args([ + "run", + "--rm", + "--network=host", + "-e", + &format!("MC_HOST_minio={}", mc_host), + "minio/mc", + "mb", + "--ignore-existing", + &format!("minio/{}", MINIO_BUCKET), + ]) + .output() + .map_err(|error| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftFixture".to_string(), + message: format!("Failed to run mc command: {error}"), + })?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + let stdout = String::from_utf8_lossy(&output.stdout); + return Err(TestBinaryError::FixtureSetup { + fixture_type: "IcebergFixture".to_string(), + message: format!("Failed to create bucket: stderr={stderr}, stdout={stdout}"), + }); + } + + tracing::info!("Created MinIO bucket: {MINIO_BUCKET}"); + Ok(()) +} + +fn bucket_empty(minio_endpoint: &str) -> Result { + use std::process::Command; + + let host = minio_endpoint.trim_start_matches("http://"); + let mc_host = format!("http://{}:{}@{}", MINIO_ACCESS_KEY, MINIO_SECRET_KEY, host); + + let output = Command::new("docker") + .args([ + "run", + "--rm", + "--network=host", + "-e", + &format!("MC_HOST_minio={}", mc_host), + "minio/mc", + "ls", + &format!("minio/{}", MINIO_BUCKET), + ]) + .output() + .map_err(|error| TestBinaryError::FixtureSetup { + fixture_type: "RedshiftFixture".to_string(), + message: format!("Failed to run mc command: {error}"), + })?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + let stdout = String::from_utf8_lossy(&output.stdout); + return Err(TestBinaryError::FixtureSetup { + fixture_type: "IcebergFixture".to_string(), + message: format!("Failed to create bucket: stderr={stderr}, stdout={stdout}"), + }); + } + + tracing::info!("Created MinIO bucket: {MINIO_BUCKET}"); + Ok(output.stdout.is_empty()) +} diff --git a/core/integration/tests/connectors/mod.rs b/core/integration/tests/connectors/mod.rs index a1433160b9..56cccf9309 100644 --- a/core/integration/tests/connectors/mod.rs +++ b/core/integration/tests/connectors/mod.rs @@ -30,6 +30,7 @@ mod postgres; mod quickwit; mod random; mod random_source_liveness; +mod redshift; mod runtime; mod s3; mod stdout; diff --git a/core/integration/tests/connectors/redshift/mod.rs b/core/integration/tests/connectors/redshift/mod.rs new file mode 100644 index 0000000000..1038c093b8 --- /dev/null +++ b/core/integration/tests/connectors/redshift/mod.rs @@ -0,0 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod redshift_sink; + +const TEST_MESSAGE_COUNT: usize = 3; diff --git a/core/integration/tests/connectors/redshift/redshift_sink.rs b/core/integration/tests/connectors/redshift/redshift_sink.rs new file mode 100644 index 0000000000..eefc0653c6 --- /dev/null +++ b/core/integration/tests/connectors/redshift/redshift_sink.rs @@ -0,0 +1,402 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::TEST_MESSAGE_COUNT; +use crate::connectors::fixtures::{ + RedshiftSinkByteaFixture, RedshiftSinkFixture, RedshiftSinkJsonFixture, + RedshiftSinkNoArchiveFixture, +}; +use crate::connectors::{TestMessage, create_test_messages}; +use bytes::Bytes; +use iggy::prelude::{IggyMessage, Partitioning}; +use iggy_common::Identifier; +use iggy_common::MessageClient; +use iggy_connector_sdk::api::SinkInfoResponse; +use integration::harness::seeds; +use integration::iggy_harness; + +use reqwest::Client; + +const SINK_TABLE: &str = "iggy_messages"; +const API_KEY: &str = "test-api-key"; +const ICEBERG_SINK_KEY: &str = "redshift"; + +type SinkRow = (i64, String, String, String); +type SinkRawRow = (i64, String, String, Vec); +type SinkJsonRow = (i64, String); + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn redshift_sink_initializes_and_runs(harness: &TestHarness, fixture: RedshiftSinkFixture) { + let api_address = harness + .connectors_runtime() + .expect("connector runtime should be available") + .http_url(); + + let http_client = Client::new(); + + let response = http_client + .get(format!("{}/sinks", api_address)) + .header("api-key", API_KEY) + .send() + .await + .expect("Failed to get sinks"); + + assert_eq!(response.status(), 200); + let sinks: Vec = response.json().await.expect("Failed to parse sinks"); + + assert_eq!(sinks.len(), 1); + assert_eq!(sinks[0].key, ICEBERG_SINK_KEY); + assert!(sinks[0].enabled); + + drop(fixture); +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn json_messages_sink_stores_as_text( + harness: &TestHarness, + fixture: RedshiftSinkJsonFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture + .target_pool() + .await + .expect("Failed to create target postgres pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let messages_data = create_test_messages(TEST_MESSAGE_COUNT); + + let mut messages: Vec = messages_data + .iter() + .enumerate() + .map(|(i, msg)| { + let payload = serde_json::to_vec(msg).expect("Failed to serialize message"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(payload)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!( + "SELECT iggy_offset, iggy_stream, iggy_topic, payload FROM {SINK_TABLE} ORDER BY iggy_offset" + ); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, stream, topic, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!(stream, seeds::names::STREAM, "Stream mismatch at row {i}"); + assert_eq!(topic, seeds::names::TOPIC, "Topic mismatch at row {i}"); + + let stored: TestMessage = + serde_json::from_str(payload).expect("Failed to deserialize stored payload"); + assert_eq!(stored, messages_data[i], "Message data mismatch at row {i}"); + } +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn binary_messages_sink_stores_as_bytea( + harness: &TestHarness, + fixture: RedshiftSinkByteaFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture.target_pool().await.expect("Failed to create pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let raw_payloads: Vec> = vec![ + b"plain text message".to_vec(), + vec![0x00, 0x01, 0x02, 0xFF, 0xFE, 0xFD], + vec![0xDE, 0xAD, 0xBE, 0xEF], + ]; + + let mut messages: Vec = raw_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(payload.clone())) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!( + "SELECT iggy_offset, iggy_stream, iggy_topic, payload FROM {SINK_TABLE} ORDER BY iggy_offset" + ); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, _, _, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!(payload, &raw_payloads[i], "Payload mismatch at row {i}"); + } +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn json_messages_sink_stores_as_json( + harness: &TestHarness, + fixture: RedshiftSinkJsonFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture.target_pool().await.expect("Failed to create pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let json_payloads: Vec = vec![ + serde_json::json!({"name": "Alice", "age": 30}), + serde_json::json!({"items": [1, 2, 3], "active": true}), + serde_json::json!({"nested": {"key": "value"}, "count": 42}), + ]; + + let mut messages: Vec = json_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + let bytes = serde_json::to_vec(payload).expect("Failed to serialize json"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(bytes)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!("SELECT iggy_offset, payload FROM {SINK_TABLE} ORDER BY iggy_offset"); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!( + payload, + &json_payloads[i].to_string(), + "JSON payload mismatch at row {i}" + ); + } +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn json_messages_sink_stores_as_bytea( + harness: &TestHarness, + fixture: RedshiftSinkByteaFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture.target_pool().await.expect("Failed to create pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let json_payloads: Vec = vec![ + serde_json::json!({"name": "Alice", "age": 30}), + serde_json::json!({"items": [1, 2, 3], "active": true}), + serde_json::json!({"nested": {"key": "value"}, "count": 42}), + ]; + + let mut messages: Vec = json_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + let bytes = serde_json::to_vec(payload).expect("Failed to serialize json"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(bytes)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!( + "SELECT iggy_offset, iggy_stream, iggy_topic, payload FROM {SINK_TABLE} ORDER BY iggy_offset" + ); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, _, _, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!( + serde_json::from_slice::(payload).expect("Failed to parse bytes"), + json_payloads[i], + "Payload mismatch at row {i}" + ); + } +} + +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/redshift/sink.toml")), + seed = seeds::connector_stream +)] +async fn sink_with_no_archive_deletes_s3_artefact( + harness: &TestHarness, + fixture: RedshiftSinkNoArchiveFixture, +) { + let client = harness.root_client().await.unwrap(); + let pool = fixture.target_pool().await.expect("Failed to create pool"); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let json_payloads: Vec = vec![ + serde_json::json!({"name": "Alice", "age": 30}), + serde_json::json!({"items": [1, 2, 3], "active": true}), + serde_json::json!({"nested": {"key": "value"}, "count": 42}), + ]; + + let mut messages: Vec = json_payloads + .iter() + .enumerate() + .map(|(i, payload)| { + let bytes = serde_json::to_vec(payload).expect("Failed to serialize json"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(bytes)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + let query = format!("SELECT iggy_offset, payload FROM {SINK_TABLE} ORDER BY iggy_offset"); + let rows: Vec = fixture + .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT) + .await + .expect("Failed to fetch rows"); + + assert_eq!( + rows.len(), + TEST_MESSAGE_COUNT, + "Expected {TEST_MESSAGE_COUNT} rows in PostgreSQL table" + ); + + for (i, (offset, payload)) in rows.iter().enumerate() { + assert_eq!(*offset, i as i64, "Offset mismatch at row {i}"); + assert_eq!( + payload, + &json_payloads[i].to_string(), + "JSON payload mismatch at row {i}" + ); + + assert_eq!( + fixture + .confirm_empty_bucket() + .expect("Failed to read empty bucket"), + true + ); + } +} diff --git a/core/integration/tests/connectors/redshift/sink.toml b/core/integration/tests/connectors/redshift/sink.toml new file mode 100644 index 0000000000..c46dd3bafa --- /dev/null +++ b/core/integration/tests/connectors/redshift/sink.toml @@ -0,0 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[connectors] +config_type = "local" +config_dir = "../connectors/sinks/redshift_sink"