diff --git a/config.yaml b/config.yaml index e69de29..c5c42f6 100644 --- a/config.yaml +++ b/config.yaml @@ -0,0 +1,28 @@ +connection_string: postgresql://postgres:mysecretpassword@localhost:5433/postgres + +# The query to execute. This MUST be a COPY...TO STDOUT (FORMAT binary) command. +# We cast all columns to match the binary types our parser expects. +query: "COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score::REAL, is_active::BOOLEAN, last_login::TIMESTAMP, notes::TEXT, course_id::INT, start_date::DATE, rating::FLOAT8 FROM benchmark_table) TO STDOUT (FORMAT binary)" + +# 'schema' is the ordered list of columns *exactly* as they appear in the query's SELECT statement. +schema: + - column_name: id + arrow_type: Int64 + - column_name: uuid + arrow_type: Utf8 + - column_name: username + arrow_type: Utf8 + - column_name: score + arrow_type: Float32 + - column_name: is_active + arrow_type: Boolean + - column_name: last_login + arrow_type: Timestamp(Nanosecond, None) + - column_name: notes + arrow_type: Utf8 + - column_name: course_id + arrow_type: Int32 + - column_name: start_date + arrow_type: Date32 + - column_name: rating + arrow_type: Float64 \ No newline at end of file diff --git a/setup_db.sql b/setup_db.sql index e69de29..6b0d918 100644 --- a/setup_db.sql +++ b/setup_db.sql @@ -0,0 +1,127 @@ +-- 1. Drop the old table if it exists +DROP TABLE IF EXISTS benchmark_table; + +-- 2. Create the new 10-column table +CREATE TABLE benchmark_table ( + id BIGINT, + uuid TEXT, -- Storing as TEXT for simplicity, can be native UUID + username TEXT, + score REAL, + is_active BOOLEAN, + last_login TIMESTAMP, + notes TEXT, + course_id INT, + start_date DATE, + rating FLOAT8 +); + +-- 3. Use COPY FROM STDIN to bulk-load 100 rows of data +-- This is the fastest way to insert data into Postgres. +COPY benchmark_table (id, uuid, username, score, is_active, last_login, notes, course_id, start_date, rating) FROM STDIN WITH (FORMAT csv); +1,a1b2c3d4-e5f6-7890-1234-567890abcdef,user_1,95.5,true,2025-01-10 10:30:00,First note,101,2024-01-01,4.5 +2,b2c3d4e5-f6a7-8901-2345-678901bcdef0,user_2,88.0,true,2025-02-15 11:45:10,Second note,102,2024-01-02,4.2 +3,c3d4e5f6-a7b8-9012-3456-789012cdef01,user_3,72.1,false,2025-03-20 12:00:05,,101,2024-01-03,3.8 +4,d4e5f6a7-b8c9-0123-4567-890123def012,user_4,99.9,true,2025-04-25 14:15:30,High score,103,2024-01-04,4.9 +5,e5f6a7b8-c9d0-1234-5678-901234ef0123,user_5,65.0,false,2025-05-30 09:10:20,Needs review,102,2024-01-05,3.1 +6,f6a7b8c9-d0e1-2345-6789-012345f01234,user_6,78.5,true,2025-06-01 18:05:00,Passed,101,2024-01-06,4.0 +7,a7b8c9d0-e1f2-3456-7890-123456012345,user_7,82.3,true,2025-07-07 13:20:45,,104,2024-01-07,4.1 +8,b8c9d0e1-f2a3-4567-8901-234567123456,user_8,91.0,true,2025-08-12 10:00:00,Excellent,103,2024-01-08,4.7 +9,c9d0e1f2-a3b4-5678-9012-345678234567,user_9,55.0,false,2025-09-18 20:30:15,Failed,102,2024-01-09,2.5 +10,d0e1f2a3-b4c5-6789-0123-456789345678,user_10,89.5,true,2025-10-22 11:11:11,Good work,101,2024-01-10,4.6 +11,e1f2a3b4-c5d6-7890-1234-567890456789,user_11,92.0,true,2025-11-28 14:55:00,Solid,103,2024-01-11,4.8 +12,f2a3b4c5-d6e7-8901-2345-678901567890,user_12,75.0,true,2025-12-01 08:30:00,,104,2024-01-12,3.9 +13,a3b4c5d6-e7f8-9012-3456-789012678901,user_13,68.2,false,2025-01-15 17:45:30,Re-test,102,2024-01-13,3.3 +14,b4c5d6e7-f8a9-0123-4567-890123789012,user_14,81.0,true,2025-02-20 12:01:00,Passed,101,2024-01-14,4.3 +15,c5d6e7f8-a9b0-1234-5678-901234890123,user_15,77.7,true,2025-03-25 10:10:10,Average,103,2024-01-15,3.7 +16,d6e7f8a9-b0c1-2345-6789-012345901234,user_16,94.5,true,2025-04-30 16:20:25,,104,2024-01-16,4.8 +17,e7f8a9b0-c1d2-3456-7890-123456012345,user_17,83.0,true,2025-05-05 11:30:00,Good,101,2024-01-17,4.2 +18,f8a9b0c1-d2e3-4567-8901-234567123456,user_18,60.5,false,2025-06-10 19:00:50,Needs help,102,2024-01-18,2.9 +19,a9b0c1d2-e3f4-5678-9012-345678234567,user_19,88.5,true,2025-07-15 14:45:00,Great,103,2024-01-19,4.6 +20,b0c1d2e3-f4a5-6789-0123-456789345678,user_20,90.0,true,2025-08-20 09:00:00,,104,2024-01-20,4.7 +21,c1d2e3f4-a5b6-7890-1234-567890456789,user_21,93.2,true,2025-09-25 10:30:00,Top score,101,2024-01-21,4.9 +22,d2e3f4a5-b6c7-8901-2345-678901567890,user_22,76.8,true,2025-10-30 11:45:10,Passed,102,2024-01-22,3.9 +23,e3f4a5b6-c7d8-9012-3456-789012678901,user_23,69.9,false,2025-11-04 12:00:05,,101,2024-01-23,3.5 +24,f4a5b6c7-d8e9-0123-4567-890123789012,user_24,85.0,true,2025-12-10 14:15:30,Good,103,2024-01-24,4.4 +25,a5b6c7d8-e9f0-1234-5678-901234890123,user_25,70.0,true,2025-01-15 09:10:20,Review,102,2024-01-25,3.6 +26,b6c7d8e9-f0a1-2345-6789-012345901234,user_26,81.5,true,2025-02-20 18:05:00,Passed,101,2024-01-26,4.1 +27,c7d8e9f0-a1b2-3456-7890-123456012345,user_27,87.3,true,2025-03-27 13:20:45,,104,2024-01-27,4.5 +28,d8e9f0a1-b2c3-4567-8901-234567123456,user_28,94.0,true,2025-04-01 10:00:00,Excellent,103,2024-01-28,4.8 +29,e9f0a1b2-c3d4-5678-9012-345678234567,user_29,58.0,false,2025-05-07 20:30:15,Failed,102,2024-01-29,2.8 +30,f0a1b2c3-d4e5-6789-0123-456789345678,user_30,90.5,true,2025-06-12 11:11:11,Great work,101,2024-01-30,4.7 +31,0a1b2c3d-4e5f-6789-0123-456789abcdef,user_31,96.0,true,2025-07-18 14:55:00,Top,103,2024-01-31,4.9 +32,1b2c3d4e-5f6a-7890-1234-567890bcdef0,user_32,74.0,true,2025-08-23 08:30:00,,104,2024-02-01,3.8 +33,2c3d4e5f-6a7b-8901-2345-678901cdef01,user_33,67.0,false,2025-09-28 17:45:30,Re-test,102,2024-02-02,3.2 +34,3d4e5f6a-7b8c-9012-3456-789012def012,user_34,80.0,true,2025-10-03 12:01:00,Passed,101,2024-02-03,4.2 +35,4e5f6a7b-8c9d-0123-4567-890123ef0123,user_35,79.0,true,2025-11-08 10:10:10,Average,103,2024-02-04,3.8 +36,5f6a7b8c-9d0e-1234-5678-901234f01234,user_36,95.0,true,2025-12-14 16:20:25,,104,2024-02-05,4.9 +37,6a7b8c9d-0e1f-2345-6789-012345012345,user_37,84.5,true,2025-01-20 11:30:00,Good,101,2024-02-06,4.3 +38,7b8c9d0e-1f2a-3456-7890-123456123456,user_38,62.0,false,2025-02-25 19:00:50,Needs help,102,2024-02-07,3.0 +39,8c9d0e1f-2a3b-4567-8901-234567234567,user_39,87.0,true,2025-03-02 14:45:00,Great,103,2024-02-08,4.5 +40,9d0e1f2a-3b4c-5678-9012-345678345678,user_40,91.8,true,2025-04-07 09:00:00,,104,2024-02-09,4.7 +41,0e1f2a3b-4c5d-6789-0123-456789456789,user_41,94.2,true,2025-05-13 10:30:00,Top score,101,2024-02-10,4.8 +42,1f2a3b4c-5d6e-7890-1234-567890567890,user_42,78.0,true,2025-06-18 11:45:10,Passed,102,2024-02-11,4.0 +43,2a3b4c5d-6e7f-8901-2345-678901678901,user_43,71.5,false,2025-07-24 12:00:05,,101,2024-02-12,3.6 +44,3b4c5d6e-7f8a-9012-3456-789012789012,user_44,86.3,true,2025-08-29 14:15:30,Good,103,2024-02-13,4.5 +45,4c5d6e7f-8a9b-0123-4567-890123890123,user_45,72.0,true,2025-09-03 09:10:20,Review,102,2024-02-14,3.7 +46,5d6e7f8a-9b0c-1234-5678-901234901234,user_46,83.0,true,2025-10-09 18:05:00,Passed,101,2024-02-15,4.2 +47,6e7f8a9b-0c1d-2345-6789-012345012345,user_47,89.0,true,2025-11-14 13:20:45,,104,2024-02-16,4.6 +48,7f8a9b0c-1d2e-3456-7890-123456123456,user_48,96.5,true,2025-12-20 10:00:00,Excellent,103,2024-02-17,4.9 +49,8a9b0c1d-2e3f-4567-8901-234567234567,user_49,60.0,false,2025-01-25 20:30:15,Failed,102,2024-02-18,2.9 +50,9b0c1d2e-3f4a-5678-9012-345678345678,user_50,92.5,true,2025-02-01 11:11:11,Great work,101,2024-02-19,4.8 +51,a1b2c3d4-e5f6-7890-1234-567890abcdef,user_51,95.5,true,2025-01-10 10:30:00,First note,101,2024-01-01,4.5 +52,b2c3d4e5-f6a7-8901-2345-678901bcdef0,user_52,88.0,true,2025-02-15 11:45:10,Second note,102,2024-01-02,4.2 +53,c3d4e5f6-a7b8-9012-3456-789012cdef01,user_53,72.1,false,2025-03-20 12:00:05,,101,2024-01-03,3.8 +54,d4e5f6a7-b8c9-0123-4567-890123def012,user_54,99.9,true,2025-04-25 14:15:30,High score,103,2024-01-04,4.9 +55,e5f6a7b8-c9d0-1234-5678-901234ef0123,user_55,65.0,false,2025-05-30 09:10:20,Needs review,102,2024-01-05,3.1 +56,f6a7b8c9-d0e1-2345-6789-012345f01234,user_56,78.5,true,2025-06-01 18:05:00,Passed,101,2024-01-06,4.0 +57,a7b8c9d0-e1f2-3456-7890-123456012345,user_57,82.3,true,2025-07-07 13:20:45,,104,2024-01-07,4.1 +58,b8c9d0e1-f2a3-4567-8901-234567123456,user_58,91.0,true,2025-08-12 10:00:00,Excellent,103,2024-01-08,4.7 +59,c9d0e1f2-a3b4-5678-9012-345678234567,user_59,55.0,false,2025-09-18 20:30:15,Failed,102,2024-01-09,2.5 +60,d0e1f2a3-b4c5-6789-0123-456789345678,user_60,89.5,true,2025-10-22 11:11:11,Good work,101,2024-01-10,4.6 +61,e1f2a3b4-c5d6-7890-1234-567890456789,user_61,92.0,true,2025-11-28 14:55:00,Solid,103,2024-01-11,4.8 +62,f2a3b4c5-d6e7-8901-2345-678901567890,user_62,75.0,true,2025-12-01 08:30:00,,104,2024-01-12,3.9 +63,a3b4c5d6-e7f8-9012-3456-789012678901,user_63,68.2,false,2025-01-15 17:45:30,Re-test,102,2024-01-13,3.3 +64,b4c5d6e7-f8a9-0123-4567-890123789012,user_64,81.0,true,2025-02-20 12:01:00,Passed,101,2024-01-14,4.3 +65,c5d6e7f8-a9b0-1234-5678-901234890123,user_65,77.7,true,2025-03-25 10:10:10,Average,103,2024-01-15,3.7 +66,d6e7f8a9-b0c1-2345-6789-012345901234,user_66,94.5,true,2025-04-30 16:20:25,,104,2024-01-16,4.8 +67,e7f8a9b0-c1d2-3456-7890-123456012345,user_67,83.0,true,2025-05-05 11:30:00,Good,101,2024-01-12,4.2 +68,f8a9b0c1-d2e3-4567-8901-234567123456,user_68,60.5,false,2025-06-10 19:00:50,Needs help,102,2024-01-18,2.9 +69,a9b0c1d2-e3f4-5678-9012-345678234567,user_69,88.5,true,2025-07-15 14:45:00,Great,103,2024-01-19,4.6 +70,b0c1d2e3-f4a5-6789-0123-456789345678,user_70,90.0,true,2025-08-20 09:00:00,,104,2024-01-20,4.7 +71,c1d2e3f4-a5b6-7890-1234-567890456789,user_71,93.2,true,2025-09-25 10:30:00,Top score,101,2024-01-21,4.9 +72,d2e3f4a5-b6c7-8901-2345-678901567890,user_72,76.8,true,2025-10-30 11:45:10,Passed,102,2024-01-22,3.9 +73,e3f4a5b6-c7d8-9012-3456-789012678901,user_73,69.9,false,2025-11-04 12:00:05,,101,2024-01-23,3.5 +74,f4a5b6c7-d8e9-0123-4567-890123789012,user_74,85.0,true,2025-12-10 14:15:30,Good,103,2024-01-24,4.4 +75,a5b6c7d8-e9f0-1234-5678-901234890123,user_75,70.0,true,2025-01-15 09:10:20,Review,102,2024-01-25,3.6 +76,b6c7d8e9-f0a1-2345-6789-012345901234,user_76,81.5,true,2025-02-20 18:05:00,Passed,101,2024-01-26,4.1 +77,c7d8e9f0-a1b2-3456-7890-123456012345,user_77,87.3,true,2025-03-27 13:20:45,,104,2024-01-27,4.5 +78,d8e9f0a1-b2c3-4567-8901-234567123456,user_78,94.0,true,2025-04-01 10:00:00,Excellent,103,2024-01-28,4.8 +79,e9f0a1b2-c3d4-5678-9012-345678234567,user_79,58.0,false,2025-05-07 20:30:15,Failed,102,2024-01-29,2.8 +80,f0a1b2c3-d4e5-6789-0123-456789345678,user_80,90.5,true,2025-06-12 11:11:11,Great work,101,2024-01-30,4.7 +81,0a1b2c3d-4e5f-6789-0123-456789abcdef,user_81,96.0,true,2025-07-18 14:55:00,Top,103,2024-01-31,4.9 +82,1b2c3d4e-5f6a-7890-1234-567890bcdef0,user_82,74.0,true,2025-08-23 08:30:00,,104,2024-02-01,3.8 +83,2c3d4e5f-6a7b-8901-2345-678901cdef01,user_83,67.0,false,2025-09-28 17:45:30,Re-test,102,2024-02-02,3.2 +84,3d4e5f6a-7b8c-9012-3456-789012def012,user_84,80.0,true,2025-10-03 12:01:00,Passed,101,2024-02-03,4.2 +85,4e5f6a7b-8c9d-0123-4567-890123ef0123,user_85,79.0,true,2025-11-08 10:10:10,Average,103,2024-02-04,3.8 +86,5f6a7b8c-9d0e-1234-5678-901234f01234,user_86,95.0,true,2025-12-14 16:20:25,,104,2024-02-05,4.9 +87,6a7b8c9d-0e1f-2345-6789-012345012345,user_87,84.5,true,2025-01-20 11:30:00,Good,101,2024-02-06,4.3 +88,7b8c9d0e-1f2a-3456-7890-123456123456,user_88,62.0,false,2025-02-25 19:00:50,Needs help,102,2024-02-07,3.0 +89,8c9d0e1f-2a3b-4567-8901-234567234567,user_89,87.0,true,2025-03-02 14:45:00,Great,103,2024-02-08,4.5 +90,9d0e1f2a-3b4c-5678-9012-345678345678,user_90,91.8,true,2025-04-07 09:00:00,,104,2024-02-09,4.7 +91,0e1f2a3b-4c5d-6789-0123-456789456789,user_91,94.2,true,2025-05-13 10:30:00,Top score,101,2024-02-10,4.8 +92,1f2a3b4c-5d6e-7890-1234-567890567890,user_92,78.0,true,2025-06-18 11:45:10,Passed,102,2024-02-11,4.0 +93,2a3b4c5d-6e7f-8901-2345-678901678901,user_93,71.5,false,2025-07-24 12:00:05,,101,2024-02-12,3.6 +94,3b4c5d6e-7f8a-9012-3456-789012789012,user_94,86.3,true,2025-08-29 14:15:30,Good,103,2024-02-13,4.5 +95,4c5d6e7f-8a9b-0123-4567-890123890123,user_95,72.0,true,2025-09-03 09:10:20,Review,102,2024-02-14,3.7 +96,5d6e7f8a-9b0c-1234-5678-901234901234,user_96,83.0,true,2025-10-09 18:05:00,Passed,101,2024-02-15,4.2 +97,6e7f8a9b-0c1d-2345-6789-012345012345,user_97,89.0,true,2025-11-14 13:20:45,,104,2024-02-16,4.6 +98,7f8a9b0c-1d2e-3456-7890-123456123456,user_98,96.5,true,2025-12-20 10:00:00,Excellent,103,2024-02-17,4.9 +99,8a9b0c1d-2e3f-4567-8901-234567234567,user_99,60.0,false,2025-01-25 20:30:15,Failed,102,2024-02-18,2.9 +100,9b0c1d2e-3f4a-5678-9012-345678345678,user_100,92.5,true,2025-02-01 11:11:11,Great work,101,2024-02-19,4.8 +\. +-- ``` +-- **To run this script:** +-- 1. Save it as `setup_db.sql`. +-- 2. Use the `psql` command-line tool (which comes with Postgres) to execute it against your database (the one running in Docker): +-- ```bash +-- psql "postgresql://postgres:mysecretpassword@localhost:5433/postgres" -f setup_db.sql \ No newline at end of file diff --git a/src/config.rs b/src/config.rs index e69de29..a652fde 100644 --- a/src/config.rs +++ b/src/config.rs @@ -0,0 +1,61 @@ +// --- External Crates --- +use serde::Deserialize; +use std::collections::HashMap; +use anyhow::{Context, Result, anyhow}; +use std::fmt::{self, Display}; + +// --- 1. CUSTOM ERROR DEFINITION --- + +#[derive(Debug)] +pub struct ConfigError(String); + +impl Display for ConfigError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "Configuration Error: {}", self.0) + } +} +impl std::error::Error for ConfigError {} + + +// --- 2. CONFIGURATION STRUCTS (The "Configured Opinion") --- +// We make these 'pub' (public) so src/lib.rs can use them. + +#[derive(Debug, Deserialize, Clone)] +pub struct ColumnConfig { + pub column_name: String, + pub arrow_type: String, +} + +#[derive(Debug, Deserialize)] +pub struct ConnectorConfig { + pub connection_string: String, + pub query: String, // <--- ADD THIS FIELD + pub schema: Vec, +} + +// --- 3. CONFIG LOADING FUNCTION --- +// This is also 'pub' so src/lib.rs can call it. +pub fn load_and_validate_config(path: &str) -> Result { + let file_content = std::fs::read_to_string(path) + .context(format!("Failed to read config file at path: {}", path))?; + + let config: ConnectorConfig = serde_yaml::from_str(&file_content) + .context("Failed to deserialize YAML configuration")?; + + if config.connection_string.is_empty() { + return Err(ConfigError("Connection string cannot be empty.".to_string())).map_err(anyhow::Error::from)?; + } + if config.query.is_empty() { + return Err(ConfigError("Query cannot be empty.".to_string())).map_err(anyhow::Error::from)?; + } + // FIX: Make the COPY check more flexible (allows newlines) + let uppercase_query = config.query.trim().to_uppercase(); + if !uppercase_query.starts_with("COPY") || !uppercase_query.contains("TO STDOUT") || !uppercase_query.contains("FORMAT BINARY") { + return Err(ConfigError("Query must be a 'COPY ... TO STDOUT (FORMAT binary)' command.".to_string())).map_err(anyhow::Error::from)?; + } + if config.schema.is_empty() { + return Err(ConfigError("Schema cannot be empty.".to_string())).map_err(anyhow::Error::from)?; + } + + Ok(config) +} \ No newline at end of file diff --git a/src/parser.rs b/src/parser.rs index e69de29..5185656 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -0,0 +1,271 @@ +// --- External Crates --- +use std::pin::Pin; +use std::sync::Arc; +use tokio_postgres::{NoTls, CopyOutStream}; +use anyhow::{Context, Result, anyhow}; +// FIX: Import new builders +use arrow::array::{ + ArrayBuilder, ArrayRef, + Int64Builder, Float64Builder, Float32Builder, StringBuilder, BooleanBuilder, + TimestampNanosecondBuilder, Date32Builder, Int32Builder +}; +// FIX: Import new types +use arrow::datatypes::{ + DataType, Field, Schema, + Float64Type, Float32Type, Int64Type, Int32Type, Utf8Type, BooleanType, TimestampNanosecondType, + Date32Type +}; +use arrow::record_batch::RecordBatch; +use futures_util::stream::StreamExt; +use bytes::Bytes; +use byteorder::{BigEndian, ReadBytesExt}; +use std::io::{Cursor, Read}; +use std::str; +use chrono::{NaiveDateTime, NaiveDate}; + +// --- Internal Crates --- +use crate::config::ConnectorConfig; + + +// --- 1. CORE DATABASE LOGIC --- +pub async fn run_db_logic(config: ConnectorConfig) -> Result<()> { + // 1. Establish the connection + println!("UncheckedIO: Attempting connection..."); + + let (client, connection) = tokio_postgres::connect(&config.connection_string, NoTls).await + .context("Failed to connect to PostgreSQL database")?; + + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("Postgres connection error: {}", e); + } + }); + + // 2. Execute the COPY TO STDOUT command + // FIX: We are now using the query from the config file! + let copy_query = &config.query; + println!("UncheckedIO: Executing user-defined query..."); + + let copy_stream = client.copy_out(copy_query.as_str()).await + .context("Failed to execute COPY TO STDOUT protocol. Check your query syntax and permissions.")?; + + // 3. Handle the Binary Stream + let pinned_stream: Pin> = Box::pin(copy_stream); + + // Build the Arrow Schema from the config + let schema_fields: Vec = config.schema.iter().map(|col_cfg| { + // We now allow "nullable" to be controlled by the column name, a temporary "hack" + let nullable = col_cfg.column_name == "notes"; + + let arrow_type = match col_cfg.arrow_type.as_str() { + "Int64" => DataType::Int64, + "Int32" => DataType::Int32, // NEW + "Float64" => DataType::Float64, // NEW + "Float32" => DataType::Float32, + "Utf8" | "String" => DataType::Utf8, + "Boolean" => DataType::Boolean, + "Timestamp(Nanosecond, None)" => DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + "Date32" => DataType::Date32, // NEW + _ => panic!("Unsupported type in config: {}", col_cfg.arrow_type), + }; + Field::new(&col_cfg.column_name, arrow_type, nullable) + }).collect(); + let arrow_schema = Arc::new(Schema::new(schema_fields)); + + // Call the parser + let (rows_processed, record_batch) = handle_binary_copy(pinned_stream, arrow_schema.clone()).await?; + + println!("UncheckedIO: Successfully parsed {} rows via binary stream.", rows_processed); + + // 4. Final Output Confirmation + println!("UncheckedIO: Built RecordBatch with {} rows and {} columns.", + record_batch.num_rows(), record_batch.num_columns()); + + println!("UncheckedIO: Data transfer complete. We lived."); + Ok(()) +} + + +// --- 2. BINARY STREAM HANDLER (THE DYNAMIC PARSER) --- + +// This enum will hold our different builder types +enum DynamicBuilder { + Int64(Box), + Int32(Box), // NEW + Float64(Box), // NEW + Float32(Box), + String(Box), + Boolean(Box), + Timestamp(Box), + Date32(Box), // NEW +} + +// Postgres Epoch for timestamps +const POSTGRES_EPOCH_NAIVE: NaiveDateTime = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap().and_hms_opt(0, 0, 0).unwrap(); +// Unix Epoch for dates +const UNIX_EPOCH_NAIVE_DATE: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap(); + +// This function is now private to this module (it's not 'pub') +async fn handle_binary_copy( + mut stream: Pin>, + arrow_schema: Arc +) -> Result<(usize, RecordBatch)> { + + let mut binary_buffer: Vec = Vec::new(); + + // 1. Aggregate all bytes from the stream into our buffer + while let Some(segment_result) = stream.next().await { + let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?; + binary_buffer.extend_from_slice(&segment); + } + + let bytes_received = binary_buffer.len(); + if bytes_received == 0 { + return Err(anyhow!("Received zero bytes from COPY stream. Check query and permissions.")); + } + + println!("UncheckedIO: Received {} total bytes from stream.", bytes_received); + + // --- PARSE POSTGRES BINARY HEADER --- + let mut cursor = Cursor::new(&binary_buffer[..]); + let mut magic_signature = [0u8; 11]; + cursor.read_exact(&mut magic_signature).context("Failed to read magic signature")?; + if &magic_signature != b"PGCOPY\n\xff\r\n\0" { + return Err(anyhow!("Invalid Postgres COPY binary signature.")); + } + let _flags = cursor.read_u32::().context("Failed to read flags")?; + let _header_ext_len = cursor.read_u32::().context("Failed to read header extension length")?; + println!("UncheckedIO: Postgres binary header validated."); + + // --- DYNAMIC FIELD DESERIALIZATION --- + + // 1. Create a dynamic list of builders based on the schema + let mut builders: Vec = arrow_schema.fields().iter().map(|field| { + match field.data_type() { + DataType::Int64 => DynamicBuilder::Int64(Box::new(Int64Builder::new())), + DataType::Int32 => DynamicBuilder::Int32(Box::new(Int32Builder::new())), + DataType::Float64 => DynamicBuilder::Float64(Box::new(Float64Builder::new())), + DataType::Float32 => DynamicBuilder::Float32(Box::new(Float32Builder::new())), + DataType::Utf8 => DynamicBuilder::String(Box::new(StringBuilder::new())), + DataType::Boolean => DynamicBuilder::Boolean(Box::new(BooleanBuilder::new())), + DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None) => { + DynamicBuilder::Timestamp(Box::new(TimestampNanosecondBuilder::new())) + }, + DataType::Date32 => DynamicBuilder::Date32(Box::new(Date32Builder::new())), + _ => panic!("Unsupported type in builder creation!"), + } + }).collect(); + + let mut rows_processed = 0; + + // 2. Loop through the rest of the buffer until we hit the 2-byte trailer + loop { + let col_count = match cursor.read_i16::() { + Ok(count) => count, + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { + break; // Reached end of buffer *before* trailer, assume it's the end + } + Err(e) => return Err(e.into()), + }; + + // Check for the 2-byte trailer (-1) which signals the end of the data + if col_count == -1 { + println!("UncheckedIO: Reached end-of-stream trailer."); + break; // End of stream + } + + rows_processed += 1; + + // 3. Dynamic Row-Parsing Loop + for (i, builder) in builders.iter_mut().enumerate() { + let field_len_i32 = cursor.read_i32::() + .context(format!("Failed to read field length for col {}", i))?; + + if field_len_i32 == -1 { + // Handle NULLs + match builder { + DynamicBuilder::Int64(b) => b.append_null(), + DynamicBuilder::Float64(b) => b.append_null(), + DynamicBuilder::Int32(b) => b.append_null(), + DynamicBuilder::Float32(b) => b.append_null(), + DynamicBuilder::String(b) => b.append_null(), + DynamicBuilder::Boolean(b) => b.append_null(), + DynamicBuilder::Timestamp(b) => b.append_null(), + DynamicBuilder::Date32(b) => b.append_null(), + } + continue; // Go to the next field + } + + let field_len_usize = field_len_i32 as usize; + + match builder { + DynamicBuilder::Int64(b) => { + let val = cursor.read_i64::()?; + b.append_value(val); + } + DynamicBuilder::Int32(b) => { + let val = cursor.read_i32::()?; + b.append_value(val); + } + DynamicBuilder::Float64(b) => { + let val = cursor.read_f64::()?; + b.append_value(val); + } + DynamicBuilder::Float32(b) => { + let val = cursor.read_f32::()?; + b.append_value(val); + } + DynamicBuilder::String(b) => { + let mut str_buf = vec![0; field_len_usize]; + cursor.read_exact(&mut str_buf)?; + let val_str = str::from_utf8(&str_buf) + .context(format!("Failed to parse UTF-8 string for col {}", i))?; + b.append_value(val_str); + } + DynamicBuilder::Boolean(b) => { + let val_bool = cursor.read_u8()?; + b.append_value(val_bool != 0); + } + DynamicBuilder::Timestamp(b) => { + let pg_micros = cursor.read_i64::()?; + let unix_epoch = NaiveDateTime::from_timestamp_opt(0, 0).unwrap(); + let pg_epoch = POSTGRES_EPOCH_NAIVE; + let epoch_delta_micros = (pg_epoch - unix_epoch).num_microseconds().unwrap(); + let unix_micros = epoch_delta_micros + pg_micros; + let unix_nanos = unix_micros * 1000; + b.append_value(unix_nanos); + } + DynamicBuilder::Date32(b) => { + // Postgres binary DATE is i32 days since 2000-01-01 + let pg_days = cursor.read_i32::()?; + // Arrow Date32 is i32 days since 1970-01-01 (Unix Epoch) + let epoch_delta_days = (POSTGRES_EPOCH_NAIVE.date() - UNIX_EPOCH_NAIVE_DATE).num_days() as i32; + let unix_days = epoch_delta_days + pg_days; + b.append_value(unix_days); + } + } + } + } + + // 4. Finalize the Arrays + let final_columns: Vec = builders.into_iter().map(|builder| { + match builder { + DynamicBuilder::Int64(mut b) => Arc::new(b.finish()) as ArrayRef, + DynamicBuilder::Int32(mut b) => Arc::new(b.finish()) as ArrayRef, + DynamicBuilder::Float64(mut b) => Arc::new(b.finish()) as ArrayRef, + DynamicBuilder::Float32(mut b) => Arc::new(b.finish()) as ArrayRef, + DynamicBuilder::String(mut b) => Arc::new(b.finish()) as ArrayRef, + DynamicBuilder::Boolean(mut b) => Arc::new(b.finish()) as ArrayRef, + DynamicBuilder::Timestamp(mut b) => Arc::new(b.finish()) as ArrayRef, + DynamicBuilder::Date32(mut b) => Arc::new(b.finish()) as ArrayRef, + } + }).collect(); + + // 5. Build the Final RecordBatch + let record_batch = RecordBatch::try_new( + arrow_schema.clone(), + final_columns, + ).context("Failed to create final Arrow RecordBatch")?; + + Ok((rows_processed, record_batch)) +} \ No newline at end of file