vrtb_pg/
engine.rs

1// PostgreSQL Engine adapter — implements vrtb_core::engine::Engine.
2//
3// tokio-postgres is async, but the core `Engine` trait is synchronous (so the
4// engine-agnostic conformance logic in vrtb-core stays driver-free). We bridge
5// the two by giving each engine its OWN current-thread Tokio runtime: the
6// connection task is spawned onto it, and every sync call `block_on`s a query,
7// which also drives the spawned connection task. Callers must therefore NOT be
8// inside another Tokio runtime (CLI integration tests use plain `#[test]`).
9
10use tokio::runtime::Runtime;
11use tokio_postgres::{Client, NoTls, Row};
12
13use vrtb_core::engine::{ColumnSchema, Dialect, Engine, LogicalType, TableRef, TableSchema};
14use vrtb_core::error::{Result as CoreResult, VeritableError};
15use vrtb_utils::sql::from_table;
16
17use crate::dialect::PgDialect;
18
19pub struct PostgresEngine {
20    runtime: Runtime,
21    client: Client,
22    dialect: PgDialect,
23}
24
25impl PostgresEngine {
26    // Connect using any tokio-postgres config string — libpq `key=value` form
27    // or a `postgres://user:pass@host:port/db` URL. The conn string is NOT
28    // echoed into errors — it can carry a password.
29    pub fn connect(conn_str: &str) -> CoreResult<Self> {
30        let runtime = tokio::runtime::Builder::new_current_thread()
31            .enable_all()
32            .build()
33            .map_err(|e| VeritableError::Connectivity(format!("tokio runtime: {e}")))?;
34
35        let client = runtime.block_on(async {
36            let (client, connection) = tokio_postgres::connect(conn_str, NoTls)
37                .await
38                .map_err(|e| VeritableError::Connectivity(format!("postgres connect: {e}")))?;
39            // Drive the connection on this engine's runtime; it makes progress
40            // whenever a later block_on(query) polls the runtime.
41            tokio::spawn(async move {
42                if let Err(e) = connection.await {
43                    eprintln!("postgres connection error: {e}");
44                }
45            });
46            Ok::<_, VeritableError>(client)
47        })?;
48
49        Ok(PostgresEngine {
50            runtime,
51            client,
52            dialect: PgDialect,
53        })
54    }
55
56    // Convenience constructor mirroring common discrete parameters.
57    pub fn new(
58        host: &str,
59        port: u16,
60        user: &str,
61        password: &str,
62        dbname: &str,
63    ) -> CoreResult<Self> {
64        Self::connect(&format!(
65            "host={host} port={port} user={user} password={password} dbname={dbname}"
66        ))
67    }
68}
69
70fn query_err(e: tokio_postgres::Error) -> VeritableError {
71    VeritableError::Query(e.to_string())
72}
73
74// Stringify a Postgres cell by column type. Covers the types the checksum query
75// returns (int8 cnt, text sums) plus common scalars; NULL → "". Unsupported
76// types (e.g. numeric, timestamp) error out — a documented limit of generic
77// `execute`, not hit by the conformance path.
78fn pg_value_to_string(row: &Row, i: usize) -> CoreResult<String> {
79    let ty = row.columns()[i].type_();
80    let s = match ty.name() {
81        "bool" => row
82            .get::<_, Option<bool>>(i)
83            .map(|v| v.to_string())
84            .unwrap_or_default(),
85        "int2" => row
86            .get::<_, Option<i16>>(i)
87            .map(|v| v.to_string())
88            .unwrap_or_default(),
89        "int4" => row
90            .get::<_, Option<i32>>(i)
91            .map(|v| v.to_string())
92            .unwrap_or_default(),
93        "int8" => row
94            .get::<_, Option<i64>>(i)
95            .map(|v| v.to_string())
96            .unwrap_or_default(),
97        "float4" => row
98            .get::<_, Option<f32>>(i)
99            .map(|v| v.to_string())
100            .unwrap_or_default(),
101        "float8" => row
102            .get::<_, Option<f64>>(i)
103            .map(|v| v.to_string())
104            .unwrap_or_default(),
105        "text" | "varchar" | "bpchar" | "name" => {
106            row.get::<_, Option<String>>(i).unwrap_or_default()
107        }
108        other => {
109            return Err(VeritableError::Query(format!(
110                "unsupported Postgres column type {other:?} for generic stringify"
111            )));
112        }
113    };
114    Ok(s)
115}
116
117// Parse the scale out of a `numeric(p, s)` spelling.
118fn parse_scale(s: &str) -> Option<u8> {
119    let inside = s.split('(').nth(1)?.trim_end_matches(')');
120    inside.split(',').nth(1)?.trim().parse().ok()
121}
122
123// Map a Postgres `format_type` string to a Veritable [`LogicalType`].
124fn map_pg_type(raw: &str) -> CoreResult<LogicalType> {
125    let s = raw.trim().to_lowercase();
126    let ty = if s.starts_with("integer")
127        || s.starts_with("bigint")
128        || s.starts_with("smallint")
129        || s.starts_with("int")
130    {
131        LogicalType::Int
132    } else if s.starts_with("numeric") || s.starts_with("decimal") {
133        LogicalType::Decimal {
134            scale: parse_scale(&s).unwrap_or(0),
135        }
136    } else if s.starts_with("timestamp") {
137        LogicalType::Timestamp { precision: 6 }
138    } else if s.starts_with("character varying")
139        || s.starts_with("varchar")
140        || s == "text"
141        || s.starts_with("character")
142        || s.starts_with("char")
143        || s.starts_with("bpchar")
144        || s == "name"
145    {
146        LogicalType::String
147    } else if s.starts_with("boolean") || s == "bool" {
148        LogicalType::Boolean
149    } else if s.starts_with("bytea") {
150        LogicalType::Binary
151    } else {
152        return Err(VeritableError::Schema(format!(
153            "unsupported Postgres type: {raw}"
154        )));
155    };
156    Ok(ty)
157}
158
159const INTROSPECT_SQL: &str = "SELECT
160    a.attname AS name,
161    format_type(a.atttypid, a.atttypmod) AS data_type,
162    a.attnotnull AS not_null,
163    pg_get_expr(d.adbin, d.adrelid) AS default_value,
164    COALESCE(i.indisprimary, false) AS primary_key
165FROM pg_attribute a
166LEFT JOIN pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
167LEFT JOIN pg_index i ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) AND i.indisprimary
168WHERE a.attrelid = $1::text::regclass AND a.attnum > 0 AND NOT a.attisdropped
169ORDER BY a.attnum";
170
171impl Engine for PostgresEngine {
172    fn name(&self) -> &str {
173        "Postgres"
174    }
175
176    fn introspect(&self, table: &TableRef) -> CoreResult<TableSchema> {
177        let qualified = from_table(table);
178        let rows = self
179            .runtime
180            .block_on(self.client.query(INTROSPECT_SQL, &[&qualified]))
181            .map_err(query_err)?;
182
183        let mut columns = Vec::with_capacity(rows.len());
184        for row in &rows {
185            let name: String = row.get("name");
186            let data_type: String = row.get("data_type");
187            let not_null: bool = row.get("not_null");
188            let default_value: Option<String> = row.get("default_value");
189            let primary_key: bool = row.get("primary_key");
190            columns.push(ColumnSchema {
191                name,
192                ty: map_pg_type(&data_type)?,
193                nullable: !not_null,
194                default_value,
195                primary_key,
196            });
197        }
198        Ok(TableSchema { columns })
199    }
200
201    fn dialect(&self) -> &dyn Dialect {
202        &self.dialect
203    }
204
205    fn execute(&self, sql: &str) -> CoreResult<Vec<Vec<String>>> {
206        let rows = self
207            .runtime
208            .block_on(self.client.query(sql, &[]))
209            .map_err(query_err)?;
210        let mut out = Vec::with_capacity(rows.len());
211        for row in &rows {
212            let mut r = Vec::with_capacity(row.len());
213            for i in 0..row.len() {
214                r.push(pg_value_to_string(row, i)?);
215            }
216            out.push(r);
217        }
218        Ok(out)
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use vrtb_core::engine::LogicalType;
226
227    #[test]
228    fn maps_common_pg_types() {
229        assert_eq!(map_pg_type("integer").unwrap(), LogicalType::Int);
230        assert_eq!(map_pg_type("bigint").unwrap(), LogicalType::Int);
231        assert_eq!(
232            map_pg_type("numeric(12,2)").unwrap(),
233            LogicalType::Decimal { scale: 2 }
234        );
235        assert_eq!(
236            map_pg_type("timestamp without time zone").unwrap(),
237            LogicalType::Timestamp { precision: 6 }
238        );
239        assert_eq!(
240            map_pg_type("character varying(255)").unwrap(),
241            LogicalType::String
242        );
243        assert_eq!(map_pg_type("text").unwrap(), LogicalType::String);
244        assert_eq!(map_pg_type("boolean").unwrap(), LogicalType::Boolean);
245        assert_eq!(map_pg_type("bytea").unwrap(), LogicalType::Binary);
246        assert!(map_pg_type("json").is_err());
247    }
248}
vrtb_pg/engine.rs

vrtb_pg/
engine.rs
