Add initial support for fetching data from remote databases (postgres, motherduck, duckdb) (#2)

* chore: add adbc dependencies

* feat(datafetch): add DataFetchError enum

* feat(datafetch): add TableMetadata and ColumnMetadata types

* feat(datafetch): add DataFetcher trait

* feat(datafetch): add DriverManager for ADBC driver loading

* feat(datafetch): add AdbcFetcher with discovery skeleton

* feat(datafetch): implement get_objects result parsing

* feat(datafetch): implement fetch_table with parquet writing

* feat(catalog): add arrow_schema_json to TableInfo

* feat: integrate DataFetcher with SchemaProvider

* feat(engine): implement connect() with ADBC discovery

* test: add datafetch integration tests

* chore: add driver vendoring directory structure

* fix: implement build.rs for driver vendoring

* fix: use proper Arrow schema serialization

* feat: implement LazyTableProvider for on-demand fetch

* fix: quote SQL identifiers to prevent injection

* fix: allow syncing empty tables

* fix: improve sql_type_to_arrow with NULL handling and warnings

* docs: add documentation to datafetch public APIs

* chore: replace ADBC deps with sqlx

* refactor(datafetch): remove ADBC, add NativeFetcher skeleton

* feat(datafetch): add StreamingParquetWriter

* refactor(catalog): migrate postgres_manager to sqlx

* feat(datafetch): add DuckDB discovery

* feat(datafetch): add DuckDB fetch_table

* feat(datafetch): add PostgreSQL discovery

* feat(datafetch): add PostgreSQL fetch with Arrow conversion

* refactor(catalog): merge add_table and update_table_schema

* test: update tests for native datafetch

* refactor(datafetch): centralize streaming Parquet writer

* test: add comprehensive tests for streaming parquet writer

* log errors

* build connection string from input

* automatically handle ssl for postgres

* simplify to use Source instead of raw json

* fix projection push down

* add support for filters/limits at table scan

* add access to catalog

* use correct size for ordinal position

* duckdb (motherduck) should only return tables for the configured catalog

* first pass at golden path integration test

* refactor to use test harness/fixture abstraction for better re-use

* remove existing integration tests in favor of new "golden path" tests

* remove un-used function

* refactor to match enum, not string

* cargo fmt

* inline postgres arrow logic

* refactor to stream postgres results

* fix log message

* refactor duckdb to write results in a streaming fashion

* use parameterized query for duckdb discovery

Author: zfarrell

Cargo.lock

@@ -14,6 +14,7 @@ path = "src/bin/server.rs"
 [dependencies]
 datafusion = "50.2"
 duckdb = { version = "1.4", features = ["bundled"] }
+sqlx = { version = "0.8", features = ["runtime-tokio", "postgres", "chrono", "tls-rustls"] }
 rustyline = "17.0"
 tokio = { version = "1.47", features = ["full"] }
 serde = { version = "1.0", features = ["derive"] }
@@ -28,10 +29,7 @@ env_logger = "0.11"
 clap = { version = "4.5", features = ["derive"] }
 arrow-csv = "56.2"
 arrow-json = "56.2"
-tokio-postgres = "0.7"
-postgres-native-tls = "0.5"
-native-tls = "0.2"
-deadpool-postgres = "0.14"
+arrow-schema = { version = "56.2", features = ["serde"] }
 object_store = { version = "0.12", features = ["aws"] }
 url = "2.5"
 axum = "0.8.7"
@@ -40,6 +38,9 @@ tower-http = { version = "0.6.6", features = ["trace", "cors"] }
 config = "0.15.19"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+thiserror = "2.0.17"
+uuid = { version = "1.11", features = ["v4"] }
+urlencoding = "2.1"
 
 [dev-dependencies]
 testcontainers = "0.23"

Cargo.toml

@@ -1,5 +1,8 @@
 use anyhow::Result;
 use clap::{Parser, ValueEnum};
+use rivetdb::catalog::CatalogManager;
+use rivetdb::datafusion::{HotDataEngine, QueryResponse};
+use rivetdb::source::Source;
 use rustyline::completion::{Completer, Pair};
 use rustyline::error::ReadlineError;
 use rustyline::highlight::Highlighter;
@@ -9,8 +12,6 @@ use rustyline::{Context, Editor, Helper};
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Duration;
-use rivetdb::catalog::CatalogManager;
-use rivetdb::datafusion::{HotDataEngine, QueryResponse};
 
 #[derive(Parser)]
 #[command(name = "hotdata", about = "HotData Query Engine", version)]
@@ -405,7 +406,7 @@ async fn handle_command(state: &mut ReplState, line: &str) -> Result<()> {
             print_help();
         }
         _ if line.starts_with("connect ") => {
-            handle_connect(state, line)?;
+            handle_connect(state, line).await?;
         }
         "list-connections" => {
             handle_list_connections(state)?;
@@ -462,7 +463,7 @@ fn print_duration_sec(d: Duration) {
     println!("execution time {}", time);
 }
 
-fn handle_connect(state: &mut ReplState, line: &str) -> Result<()> {
+async fn handle_connect(state: &mut ReplState, line: &str) -> Result<()> {
     // Parse: connect <type> <name> <key=value> ...
     let parts: Vec<&str> = line.split_whitespace().collect();
 
@@ -534,8 +535,11 @@ fn handle_connect(state: &mut ReplState, line: &str) -> Result<()> {
 
     let config_value = serde_json::Value::Object(config);
 
+    // Deserialize to Source enum
+    let source: Source = serde_json::from_value(config_value)?;
+
     // Connect through the engine
-    state.engine.connect(source_type, name, config_value)?;
+    state.engine.connect(name, source).await?;
 
     Ok(())
 }

src/bin/repl.rs

@@ -123,6 +123,15 @@ impl DuckdbCatalogManager {
             conn.execute("INSERT INTO schema_migrations (version) VALUES (1)", [])?;
         }
 
+        // Migration 2: Add arrow_schema_json column to tables
+        if current_version < 2 {
+            conn.execute(
+                "ALTER TABLE tables ADD COLUMN IF NOT EXISTS arrow_schema_json TEXT",
+                [],
+            )?;
+            conn.execute("INSERT INTO schema_migrations (version) VALUES (2)", [])?;
+        }
+
         Ok(())
     }
 
@@ -135,6 +144,7 @@ impl DuckdbCatalogManager {
             parquet_path: row.get(4)?,
             state_path: row.get(5)?,
             last_sync: row.get(6)?,
+            arrow_schema_json: row.get(7)?,
         })
     }
 }
@@ -221,18 +231,25 @@ impl CatalogManager for DuckdbCatalogManager {
         }
     }
 
-    fn add_table(&self, connection_id: i32, schema_name: &str, table_name: &str) -> Result<i32> {
+    fn add_table(
+        &self,
+        connection_id: i32,
+        schema_name: &str,
+        table_name: &str,
+        arrow_schema_json: &str,
+    ) -> Result<i32> {
         if self.readonly {
             anyhow::bail!("Cannot add table in readonly mode");
         }
 
         let conn_guard = self.get_connection_guard()?;
         let conn = conn_guard.as_ref().unwrap(); // Safe: get_connection_guard verified it's Some
 
-        // add table to catalog. if we already know about the table (we're reimporting connection), ignore it
+        // Insert or update table with schema
         conn.execute(
-            "INSERT INTO tables (connection_id, schema_name, table_name) VALUES (?, ?, ?) ON CONFLICT DO NOTHING",
-            params![connection_id, schema_name, table_name],
+            "INSERT INTO tables (connection_id, schema_name, table_name, arrow_schema_json) VALUES (?, ?, ?, ?)
+             ON CONFLICT (connection_id, schema_name, table_name) DO UPDATE SET arrow_schema_json = excluded.arrow_schema_json",
+            params![connection_id, schema_name, table_name, arrow_schema_json],
         )?;
 
         let id: i32 = conn.query_row(
@@ -253,7 +270,7 @@ impl CatalogManager for DuckdbCatalogManager {
         if let Some(conn_id) = connection_id {
             let mut stmt = conn.prepare(
                 "SELECT id, connection_id, schema_name, table_name, parquet_path, state_path, \
-                 CAST(last_sync AS VARCHAR) as last_sync \
+                 CAST(last_sync AS VARCHAR) as last_sync, arrow_schema_json \
                  FROM tables WHERE connection_id = ? ORDER BY schema_name, table_name",
             )?;
             let rows = stmt.query_map(params![conn_id], DuckdbCatalogManager::table_mapper)?;
@@ -264,7 +281,7 @@ impl CatalogManager for DuckdbCatalogManager {
         } else {
             let mut stmt = conn.prepare(
                 "SELECT id, connection_id, schema_name, table_name, parquet_path, state_path, \
-                 CAST(last_sync AS VARCHAR) as last_sync \
+                 CAST(last_sync AS VARCHAR) as last_sync, arrow_schema_json \
                  FROM tables ORDER BY schema_name, table_name",
             )?;
             let rows = stmt.query_map([], DuckdbCatalogManager::table_mapper)?;
@@ -288,7 +305,7 @@ impl CatalogManager for DuckdbCatalogManager {
 
         let mut stmt = conn.prepare(
             "SELECT id, connection_id, schema_name, table_name, parquet_path, state_path, \
-             CAST(last_sync AS VARCHAR) as last_sync \
+             CAST(last_sync AS VARCHAR) as last_sync, arrow_schema_json \
              FROM tables WHERE connection_id = ? AND schema_name = ? AND table_name = ?",
         )?;
 

src/catalog/duckdb_manager.rs

@@ -19,6 +19,7 @@ pub struct TableInfo {
     pub parquet_path: Option<String>,
     pub state_path: Option<String>,
     pub last_sync: Option<String>,
+    pub arrow_schema_json: Option<String>,
 }
 
 pub trait CatalogManager: Debug + Send + Sync {
@@ -29,7 +30,13 @@ pub trait CatalogManager: Debug + Send + Sync {
     fn list_connections(&self) -> Result<Vec<ConnectionInfo>>;
     fn add_connection(&self, name: &str, source_type: &str, config_json: &str) -> Result<i32>;
     fn get_connection(&self, name: &str) -> Result<Option<ConnectionInfo>>;
-    fn add_table(&self, connection_id: i32, schema_name: &str, table_name: &str) -> Result<i32>;
+    fn add_table(
+        &self,
+        connection_id: i32,
+        schema_name: &str,
+        table_name: &str,
+        arrow_schema_json: &str,
+    ) -> Result<i32>;
     fn list_tables(&self, connection_id: Option<i32>) -> Result<Vec<TableInfo>>;
     fn get_table(
         &self,