teaql_provider_postgres/
lib.rs

1#![allow(warnings)]
2use std::collections::BTreeMap;
3use std::future::Future;
4use std::pin::Pin;
5
6use chrono::{DateTime, NaiveDate, Utc};
7use deadpool_postgres::Pool;
8use rust_decimal::Decimal;
9use std::sync::Arc;
10use teaql_core::{
11    BinaryOp, DataType, EntityDescriptor, Expr, InsertCommand, PropertyDescriptor, Record,
12    SelectQuery, UpdateCommand, Value,
13};
14use teaql_runtime::{GraphNode, InternalIdGenerator, RuntimeError, SchemaProvider, UserContext};
15use teaql_sql::{
16    CompiledQuery, DatabaseKind, SqlCompileError, SqlDialect, SqlTransport,
17    quote_identifier_if_needed,
18};
19use tokio::sync::Mutex;
20
21pub const DEFAULT_ID_SPACE_TABLE: &str = "teaql_id_space";
22
23#[derive(Debug, Default, Clone, Copy)]
24pub struct PostgresDialect;
25
26impl SqlDialect for PostgresDialect {
27    fn kind(&self) -> DatabaseKind {
28        DatabaseKind::PostgreSql
29    }
30
31    fn quote_ident(&self, ident: &str) -> String {
32        quote_ident(ident)
33    }
34
35    fn placeholder(&self, index: usize) -> String {
36        format!("${index}")
37    }
38
39    fn schema_setup_sqls(&self) -> &'static [&'static str] {
40        &[CREATE_SOUNDEX_FUNCTION]
41    }
42
43    fn schema_type_sql(
44        &self,
45        data_type: DataType,
46        _property: &PropertyDescriptor,
47    ) -> Result<&'static str, SqlCompileError> {
48        match data_type {
49            DataType::Bool => Ok("BOOLEAN"),
50            DataType::I64 | DataType::U64 => Ok("BIGINT"),
51            DataType::F64 => Ok("DOUBLE PRECISION"),
52            DataType::Decimal => Ok("NUMERIC"),
53            DataType::Text => Ok("VARCHAR(255)"),
54            DataType::LargeText => Ok("TEXT"),
55            DataType::Json => Ok("JSONB"),
56            DataType::Date => Ok("DATE"),
57            DataType::Timestamp => Ok("TIMESTAMPTZ"),
58        }
59    }
60
61    fn compile_in(
62        &self,
63        entity: &EntityDescriptor,
64        left: &Expr,
65        op: BinaryOp,
66        right: &Expr,
67        params: &mut Vec<Value>,
68    ) -> Result<String, SqlCompileError> {
69        match op {
70            BinaryOp::InLarge | BinaryOp::NotInLarge => {
71                let Expr::Value(Value::List(values)) = right else {
72                    let lhs = self.compile_expr(entity, left, params)?;
73                    let rhs = self.compile_expr(entity, right, params)?;
74                    let operator = match op {
75                        BinaryOp::InLarge => "= ANY",
76                        BinaryOp::NotInLarge => "<> ALL",
77                        _ => unreachable!(),
78                    };
79                    return Ok(format!("({lhs} {operator} ({rhs}))"));
80                };
81                if values.is_empty() {
82                    return Err(SqlCompileError::EmptyInList);
83                }
84                let lhs = self.compile_expr(entity, left, params)?;
85                params.push(Value::List(values.clone()));
86                let placeholder = self.placeholder(params.len());
87                let operator = match op {
88                    BinaryOp::InLarge => "= ANY",
89                    BinaryOp::NotInLarge => "<> ALL",
90                    _ => unreachable!(),
91                };
92                Ok(format!("({lhs} {operator}({placeholder}))"))
93            }
94            _ => {
95                let lhs = self.compile_expr(entity, left, params)?;
96                let operator = match op {
97                    BinaryOp::In => "IN",
98                    BinaryOp::NotIn => "NOT IN",
99                    _ => unreachable!(),
100                };
101                match right {
102                    Expr::Value(Value::List(values)) => {
103                        if values.is_empty() {
104                            return Err(SqlCompileError::EmptyInList);
105                        }
106                        let mut placeholders = Vec::with_capacity(values.len());
107                        for value in values {
108                            params.push(value.clone());
109                            placeholders.push(self.placeholder(params.len()));
110                        }
111                        Ok(format!("({lhs} {operator} ({}))", placeholders.join(", ")))
112                    }
113                    _ => {
114                        let rhs = self.compile_expr(entity, right, params)?;
115                        Ok(format!("({lhs} {operator} ({rhs}))"))
116                    }
117                }
118            }
119        }
120    }
121}
122
123const CREATE_SOUNDEX_FUNCTION: &str = r#"
124CREATE OR REPLACE FUNCTION soundex(input text)
125RETURNS text
126LANGUAGE plpgsql
127IMMUTABLE
128STRICT
129AS $$
130DECLARE
131    normalized text := upper(regexp_replace(input, '[^A-Za-z]', '', 'g'));
132    first_char text;
133    output text;
134    previous_code text;
135    code text;
136    ch text;
137    i integer;
138BEGIN
139    IF normalized = '' THEN
140        RETURN '0000';
141    END IF;
142
143    first_char := substr(normalized, 1, 1);
144    output := first_char;
145    previous_code := CASE
146        WHEN first_char IN ('B', 'F', 'P', 'V') THEN '1'
147        WHEN first_char IN ('C', 'G', 'J', 'K', 'Q', 'S', 'X', 'Z') THEN '2'
148        WHEN first_char IN ('D', 'T') THEN '3'
149        WHEN first_char = 'L' THEN '4'
150        WHEN first_char IN ('M', 'N') THEN '5'
151        WHEN first_char = 'R' THEN '6'
152        ELSE '0'
153    END;
154
155    FOR i IN 2..char_length(normalized) LOOP
156        ch := substr(normalized, i, 1);
157        code := CASE
158            WHEN ch IN ('B', 'F', 'P', 'V') THEN '1'
159            WHEN ch IN ('C', 'G', 'J', 'K', 'Q', 'S', 'X', 'Z') THEN '2'
160            WHEN ch IN ('D', 'T') THEN '3'
161            WHEN ch = 'L' THEN '4'
162            WHEN ch IN ('M', 'N') THEN '5'
163            WHEN ch = 'R' THEN '6'
164            ELSE '0'
165        END;
166
167        IF code <> '0' AND code <> previous_code THEN
168            output := output || code;
169            IF char_length(output) = 4 THEN
170                RETURN output;
171            END IF;
172        END IF;
173        previous_code := code;
174    END LOOP;
175
176    RETURN rpad(output, 4, '0');
177END;
178$$
179"#;
180
181#[derive(Debug)]
182pub enum MutationExecutorError {
183    Driver(tokio_postgres::Error),
184    Pool(String),
185    SqlCompile(SqlCompileError),
186    UnsupportedValue(&'static str),
187    UnsupportedColumnType(String),
188    Bind(String),
189}
190
191impl std::fmt::Display for MutationExecutorError {
192    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
193        match self {
194            Self::Driver(err) => err.fmt(f),
195            Self::Pool(err) => write!(f, "postgres pool error: {err}"),
196            Self::SqlCompile(err) => err.fmt(f),
197            Self::UnsupportedValue(kind) => {
198                write!(f, "unsupported bind value for mutation executor: {kind}")
199            }
200            Self::UnsupportedColumnType(kind) => {
201                write!(f, "unsupported column type for record decoding: {kind}")
202            }
203            Self::Bind(message) => write!(f, "bind error: {message}"),
204        }
205    }
206}
207
208impl std::error::Error for MutationExecutorError {}
209
210impl From<tokio_postgres::Error> for MutationExecutorError {
211    fn from(value: tokio_postgres::Error) -> Self {
212        Self::Driver(value)
213    }
214}
215
216impl From<SqlCompileError> for MutationExecutorError {
217    fn from(value: SqlCompileError) -> Self {
218        Self::SqlCompile(value)
219    }
220}
221
222#[derive(Clone)]
223pub struct PgMutationExecutor {
224    pool: Pool,
225}
226
227impl SqlTransport for PgMutationExecutor {
228    type Error = MutationExecutorError;
229
230    async fn fetch_all_sql(&self, query: &CompiledQuery) -> Result<Vec<Record>, Self::Error> {
231        self.fetch_all(query).await
232    }
233
234    async fn execute_sql(&self, query: &CompiledQuery) -> Result<u64, Self::Error> {
235        self.execute(query).await
236    }
237}
238
239impl teaql_sql::StreamingSqlTransport for PgMutationExecutor {
240    fn stream_sql(
241        &self,
242        query: CompiledQuery,
243        chunk_size: usize,
244    ) -> teaql_data_service::QueryStream<'_, Self::Error> {
245        let pool = self.pool.clone();
246        Box::pin(async_stream::try_stream! {
247            use futures_util::TryStreamExt;
248            let mut args = PgArgs { values: Vec::new() }; for value in &query.params { bind_pg(&mut args, value)?; }
249            let client = pool.get().await.map_err(|e| MutationExecutorError::Pool(e.to_string()))?;
250            let params = args.as_refs();
251            let rows = client.query_raw(&query.sql, params).await?;
252            futures_util::pin_mut!(rows);
253            let mut chunk = Vec::with_capacity(chunk_size); let mut index = 0;
254            while let Some(row) = rows.try_next().await? { chunk.push(decode_pg_row(&row)?); if chunk.len()==chunk_size { yield teaql_data_service::StreamChunk { rows: std::mem::take(&mut chunk), chunk_index:index, is_last:false }; index+=1; } }
255            if !chunk.is_empty() { yield teaql_data_service::StreamChunk { rows:chunk, chunk_index:index, is_last:true }; }
256        })
257    }
258}
259
260impl teaql_sql::SqlTransaction for PgMutationExecutor {
261    type Error = MutationExecutorError;
262
263    async fn commit_sql(self) -> Result<(), Self::Error> {
264        Err(MutationExecutorError::Bind(
265            "Transactions not supported yet".to_string(),
266        ))
267    }
268
269    async fn rollback_sql(self) -> Result<(), Self::Error> {
270        Err(MutationExecutorError::Bind(
271            "Transactions not supported yet".to_string(),
272        ))
273    }
274}
275
276impl teaql_sql::SqlTransactionTransport for PgMutationExecutor {
277    type Tx<'a>
278        = Self
279    where
280        Self: 'a;
281
282    async fn begin_sql(&self) -> Result<Self::Tx<'_>, Self::Error> {
283        Err(MutationExecutorError::Bind(
284            "Transactions not supported yet".to_string(),
285        ))
286    }
287}
288
289impl PgMutationExecutor {
290    pub fn new(pool: Pool) -> Self {
291        Self { pool }
292    }
293
294    pub fn pool(&self) -> Pool {
295        self.pool.clone()
296    }
297
298    pub async fn ensure_schema(
299        &self,
300        dialect: &PostgresDialect,
301        entities: &[&EntityDescriptor],
302    ) -> Result<(), MutationExecutorError> {
303        let client = self
304            .pool
305            .get()
306            .await
307            .map_err(|e| MutationExecutorError::Pool(e.to_string()))?;
308        for sql in dialect.schema_setup_sqls() {
309            client.execute(*sql, &[]).await?;
310        }
311        self.ensure_id_space_table(DEFAULT_ID_SPACE_TABLE).await?;
312
313        for entity in entities {
314            if !self.table_exists(&entity.table_name).await? {
315                let sql = dialect.compile_create_table(entity)?;
316                client.execute(&sql, &[]).await?;
317                continue;
318            }
319
320            let existing_columns = self.table_columns(&entity.table_name).await?;
321            for property in &entity.properties {
322                let bare_column = strip_identifier_quotes(&property.column_name).to_lowercase();
323                if existing_columns.contains(&bare_column) {
324                    continue;
325                }
326                let sql = dialect.compile_add_column(entity, property)?;
327                client.execute(&sql, &[]).await?;
328            }
329
330            for sql in dialect.schema_indexes_sqls(entity)? {
331                client.execute(&sql, &[]).await?;
332            }
333        }
334        Ok(())
335    }
336
337    pub async fn ensure_id_space_table(
338        &self,
339        table_name: &str,
340    ) -> Result<(), MutationExecutorError> {
341        let sql = format!(
342            "CREATE TABLE IF NOT EXISTS {} (type_name VARCHAR(100) PRIMARY KEY, current_level BIGINT NOT NULL)",
343            quote_ident(table_name)
344        );
345        let client = self
346            .pool
347            .get()
348            .await
349            .map_err(|e| MutationExecutorError::Pool(e.to_string()))?;
350        client.execute(&sql, &[]).await?;
351        Ok(())
352    }
353
354    pub async fn execute(&self, query: &CompiledQuery) -> Result<u64, MutationExecutorError> {
355        let mut args = PgArgs { values: Vec::new() };
356        for value in &query.params {
357            bind_pg(&mut args, value)?;
358        }
359        let client = self
360            .pool
361            .get()
362            .await
363            .map_err(|e| MutationExecutorError::Pool(e.to_string()))?;
364        let result = client.execute(&query.sql, &args.as_refs()).await?;
365        Ok(result)
366    }
367
368    pub async fn fetch_all(
369        &self,
370        query: &CompiledQuery,
371    ) -> Result<Vec<Record>, MutationExecutorError> {
372        let mut args = PgArgs { values: Vec::new() };
373        for value in &query.params {
374            bind_pg(&mut args, value)?;
375        }
376        let client = self
377            .pool
378            .get()
379            .await
380            .map_err(|e| MutationExecutorError::Pool(e.to_string()))?;
381        let rows = client.query(&query.sql, &args.as_refs()).await?;
382        rows.iter().map(decode_pg_row).collect()
383    }
384
385    async fn table_exists(&self, table_name: &str) -> Result<bool, MutationExecutorError> {
386        let client = self
387            .pool
388            .get()
389            .await
390            .map_err(|e| MutationExecutorError::Pool(e.to_string()))?;
391        let row = client
392            .query_one(
393                "SELECT COUNT(1)
394             FROM information_schema.tables
395             WHERE table_schema = current_schema()
396               AND table_name = $1",
397                &[&table_name],
398            )
399            .await?;
400        let exists: i64 = row.try_get(0)?;
401        Ok(exists > 0)
402    }
403
404    async fn table_columns(
405        &self,
406        table_name: &str,
407    ) -> Result<std::collections::BTreeSet<String>, MutationExecutorError> {
408        let client = self
409            .pool
410            .get()
411            .await
412            .map_err(|e| MutationExecutorError::Pool(e.to_string()))?;
413        let rows = client
414            .query(
415                "SELECT column_name
416             FROM information_schema.columns
417             WHERE table_schema = current_schema()
418               AND table_name = $1",
419                &[&table_name],
420            )
421            .await?;
422        let mut columns = std::collections::BTreeSet::new();
423        for row in rows {
424            let name: String = row.try_get("column_name")?;
425            columns.insert(name.to_lowercase());
426        }
427        Ok(columns)
428    }
429}
430
431async fn ensure_initial_graphs_postgres(
432    executor: &PgMutationExecutor,
433    dialect: &PostgresDialect,
434    ctx: &UserContext,
435) -> Result<(), MutationExecutorError> {
436    for graph in ctx.initial_graphs() {
437        let entity = ctx.entity(&graph.entity).ok_or_else(|| {
438            MutationExecutorError::Bind(format!("missing entity: {}", graph.entity))
439        })?;
440        if initial_graph_exists_postgres(executor, dialect, entity, graph).await? {
441            if let Some(query) = compile_initial_graph_update(dialect, entity, graph)? {
442                executor.execute(&query).await?;
443            }
444            continue;
445        }
446        let query = compile_initial_graph_insert(dialect, entity, graph)?;
447        executor.execute(&query).await?;
448    }
449    Ok(())
450}
451
452async fn initial_graph_exists_postgres(
453    executor: &PgMutationExecutor,
454    dialect: &PostgresDialect,
455    entity: &EntityDescriptor,
456    graph: &GraphNode,
457) -> Result<bool, MutationExecutorError> {
458    let Some(id) = graph.values.get("id") else {
459        return Ok(false);
460    };
461    let query = dialect.compile_select(
462        entity,
463        &SelectQuery::new(&graph.entity)
464            .project("id")
465            .filter(Expr::eq("id", id.clone()))
466            .limit(1),
467    )?;
468    Ok(!executor.fetch_all(&query).await?.is_empty())
469}
470
471fn compile_initial_graph_insert(
472    dialect: &impl SqlDialect,
473    entity: &EntityDescriptor,
474    graph: &GraphNode,
475) -> Result<CompiledQuery, MutationExecutorError> {
476    let mut command = InsertCommand::new(&graph.entity);
477    for (field, value) in &graph.values {
478        command = command.value(field.clone(), value.clone());
479    }
480    dialect.compile_insert(entity, &command).map_err(Into::into)
481}
482
483fn compile_initial_graph_update(
484    dialect: &impl SqlDialect,
485    entity: &EntityDescriptor,
486    graph: &crate::GraphNode,
487) -> Result<Option<CompiledQuery>, MutationExecutorError> {
488    let Some(id) = graph.values.get("id") else {
489        return Ok(None);
490    };
491    let mut command = UpdateCommand::new(&graph.entity, id.clone());
492    for (field, value) in &graph.values {
493        if field == "id" {
494            continue;
495        }
496        command = command.value(field.clone(), value.clone());
497    }
498    match dialect.compile_update(entity, &command) {
499        Ok(query) => Ok(Some(query)),
500        Err(SqlCompileError::EmptyMutation(_)) => Ok(None),
501        Err(err) => Err(err.into()),
502    }
503}
504
505pub trait PostgresSchemaExt {
506    fn ensure_postgres_schema(
507        &self,
508    ) -> Pin<Box<dyn Future<Output = Result<(), MutationExecutorError>> + '_>>;
509}
510
511pub async fn ensure_postgres_schema_for(ctx: &UserContext) -> Result<(), MutationExecutorError> {
512    let dialect = ctx.get_resource::<PostgresDialect>().ok_or_else(|| {
513        MutationExecutorError::Bind("missing typed resource: PostgresDialect".to_owned())
514    })?;
515    let executor = ctx.get_resource::<PgMutationExecutor>().ok_or_else(|| {
516        MutationExecutorError::Bind("missing typed resource: PgMutationExecutor".to_owned())
517    })?;
518
519    let entities = ctx.all_entities();
520
521    executor.ensure_schema(dialect, &entities).await?;
522    ensure_initial_graphs_postgres(executor, dialect, ctx).await
523}
524
525#[cfg(test)]
526mod streaming_tests {
527    use super::*;
528    use futures_util::StreamExt;
529    use teaql_sql::StreamingSqlTransport;
530
531    #[tokio::test]
532    async fn streams_from_real_postgres_when_configured() {
533        let Ok(url) = std::env::var("TEAQL_TEST_POSTGRES_URL") else {
534            return;
535        };
536        let mut config = deadpool_postgres::Config::new();
537        config.url = Some(url);
538        let pool = config
539            .create_pool(
540                Some(deadpool_postgres::Runtime::Tokio1),
541                tokio_postgres::NoTls,
542            )
543            .unwrap();
544        let executor = PgMutationExecutor::new(pool);
545        let query = CompiledQuery {
546            sql: "SELECT id FROM (VALUES (1), (2), (3), (4), (5)) AS fixture(id) ORDER BY id"
547                .to_owned(),
548            params: vec![],
549            comment: None,
550        };
551        let mut stream = executor.stream_sql(query, 2);
552        let mut sizes = Vec::new();
553        while let Some(chunk) = stream.next().await {
554            sizes.push(chunk.unwrap().rows.len());
555        }
556        assert_eq!(sizes, vec![2, 2, 1]);
557    }
558}
559
560impl PostgresSchemaExt for UserContext {
561    fn ensure_postgres_schema(
562        &self,
563    ) -> Pin<Box<dyn Future<Output = Result<(), MutationExecutorError>> + '_>> {
564        Box::pin(ensure_postgres_schema_for(self))
565    }
566}
567
568#[derive(Debug, Default, Clone, Copy)]
569pub struct PostgresSchemaProvider;
570
571impl SchemaProvider for PostgresSchemaProvider {
572    fn ensure_schema<'a>(
573        &'a self,
574        ctx: &'a UserContext,
575    ) -> Pin<Box<dyn Future<Output = Result<(), RuntimeError>> + Send + 'a>> {
576        Box::pin(async move {
577            ensure_postgres_schema_for(ctx)
578                .await
579                .map_err(|err| RuntimeError::Schema(err.to_string()))
580        })
581    }
582}
583
584pub trait PostgresProviderExt {
585    fn use_postgres_provider(&mut self, executor: PgMutationExecutor) -> &mut Self;
586}
587
588impl PostgresProviderExt for UserContext {
589    fn use_postgres_provider(&mut self, executor: PgMutationExecutor) -> &mut Self {
590        self.insert_resource(PostgresDialect);
591        self.insert_resource(executor);
592        self.set_schema_provider(PostgresSchemaProvider);
593        self
594    }
595}
596
597#[derive(Clone)]
598pub struct PgIdSpaceGenerator {
599    pool: Pool,
600    table_name: String,
601}
602
603impl PgIdSpaceGenerator {
604    pub fn new(pool: Pool) -> Self {
605        Self {
606            pool,
607            table_name: DEFAULT_ID_SPACE_TABLE.to_owned(),
608        }
609    }
610
611    pub fn from_executor(executor: PgMutationExecutor) -> Self {
612        Self::new(executor.pool())
613    }
614
615    pub fn with_table_name(mut self, table_name: impl Into<String>) -> Self {
616        self.table_name = table_name.into();
617        self
618    }
619
620    pub async fn ensure_table(&self) -> Result<(), MutationExecutorError> {
621        PgMutationExecutor::new(self.pool.clone())
622            .ensure_id_space_table(&self.table_name)
623            .await
624    }
625
626    pub async fn next_id(&self, entity: &str) -> Result<u64, MutationExecutorError> {
627        self.ensure_table().await?;
628        let update_sql = format!(
629            "UPDATE {} SET current_level = current_level + 1 WHERE type_name = $1 RETURNING current_level",
630            quote_ident(&self.table_name)
631        );
632        let client = self
633            .pool
634            .get()
635            .await
636            .map_err(|e| MutationExecutorError::Pool(e.to_string()))?;
637        let row = client.query_opt(&update_sql, &[&entity]).await?;
638
639        let id = match row {
640            Some(r) => {
641                let level: i64 = r.try_get(0)?;
642                level
643            }
644            None => {
645                let insert_sql = format!(
646                    "INSERT INTO {} (type_name, current_level) VALUES ($1, 1) RETURNING current_level",
647                    quote_ident(&self.table_name)
648                );
649                let insert_res = client.query_one(&insert_sql, &[&entity]).await;
650                match insert_res {
651                    Ok(r) => {
652                        let level: i64 = r.try_get(0)?;
653                        level
654                    }
655                    Err(_) => {
656                        let row = client.query_one(&update_sql, &[&entity]).await?;
657                        let level: i64 = row.try_get(0)?;
658                        level
659                    }
660                }
661            }
662        };
663
664        u64::try_from(id).map_err(|_| {
665            MutationExecutorError::Bind(format!("generated id {id} cannot be represented as u64"))
666        })
667    }
668}
669
670impl InternalIdGenerator for PgIdSpaceGenerator {
671    fn generate_id(&self, entity: &str) -> Result<u64, RuntimeError> {
672        let generator = self.clone();
673        let entity = entity.to_owned();
674        block_on_id_generation(async move { generator.next_id(&entity).await })
675    }
676}
677
678fn block_on_id_generation<F>(future: F) -> Result<u64, RuntimeError>
679where
680    F: Future<Output = Result<u64, MutationExecutorError>> + Send + 'static,
681{
682    let result = match tokio::runtime::Handle::try_current() {
683        Ok(handle) => tokio::task::block_in_place(|| handle.block_on(future)),
684        Err(_) => tokio::runtime::Builder::new_current_thread()
685            .enable_all()
686            .build()
687            .map_err(|err| RuntimeError::IdGeneration(err.to_string()))?
688            .block_on(future),
689    };
690    result.map_err(|err| RuntimeError::IdGeneration(err.to_string()))
691}
692
693fn quote_ident(ident: &str) -> String {
694    quote_identifier_if_needed(ident, '"')
695}
696
697/// Strip wrapping identifier quotes from a SQL identifier so that bare column
698/// names returned by `information_schema.columns` can be compared with
699/// potentially-quoted `PropertyDescriptor::column_name` values.
700fn strip_identifier_quotes(ident: &str) -> &str {
701    let bytes = ident.as_bytes();
702    if bytes.len() >= 2 {
703        let (first, last) = (bytes[0], bytes[bytes.len() - 1]);
704        if (first == b'"' && last == b'"')
705            || (first == b'`' && last == b'`')
706            || (first == b'[' && last == b']')
707        {
708            return &ident[1..ident.len() - 1];
709        }
710    }
711    ident
712}
713
714fn try_parse_datetime_from_str(s: &str) -> Option<chrono::DateTime<chrono::Utc>> {
715    if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) {
716        return Some(dt.with_timezone(&chrono::Utc));
717    }
718    if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
719        return Some(chrono::DateTime::from_naive_utc_and_offset(
720            ndt,
721            chrono::Utc,
722        ));
723    }
724    if let Ok(nd) = chrono::NaiveDate::parse_from_str(s, "%Y-%m-%d") {
725        let ndt = nd.and_hms_opt(0, 0, 0)?;
726        return Some(chrono::DateTime::from_naive_utc_and_offset(
727            ndt,
728            chrono::Utc,
729        ));
730    }
731    None
732}
733
734#[derive(Debug, Clone, Copy)]
735struct PgNull;
736
737impl tokio_postgres::types::ToSql for PgNull {
738    fn to_sql(
739        &self,
740        ty: &tokio_postgres::types::Type,
741        out: &mut bytes::BytesMut,
742    ) -> Result<tokio_postgres::types::IsNull, Box<dyn std::error::Error + Sync + Send>> {
743        Ok(tokio_postgres::types::IsNull::Yes)
744    }
745
746    fn accepts(ty: &tokio_postgres::types::Type) -> bool {
747        true
748    }
749
750    fn to_sql_checked(
751        &self,
752        ty: &tokio_postgres::types::Type,
753        out: &mut bytes::BytesMut,
754    ) -> Result<tokio_postgres::types::IsNull, Box<dyn std::error::Error + Sync + Send>> {
755        Ok(tokio_postgres::types::IsNull::Yes)
756    }
757}
758
759struct PgArgs {
760    values: Vec<Box<dyn tokio_postgres::types::ToSql + Sync + Send>>,
761}
762impl PgArgs {
763    fn add<T: tokio_postgres::types::ToSql + Sync + Send + 'static>(&mut self, v: T) {
764        self.values.push(Box::new(v));
765    }
766    fn as_refs(&self) -> Vec<&(dyn tokio_postgres::types::ToSql + Sync)> {
767        self.values.iter().map(|b| b.as_ref() as _).collect()
768    }
769}
770
771fn bind_pg(args: &mut PgArgs, value: &Value) -> Result<(), MutationExecutorError> {
772    match value {
773        Value::Null => {
774            args.add(PgNull);
775        }
776        Value::Bool(v) => args.add(*v),
777        Value::I64(v) => args.add(*v),
778        Value::U64(v) => {
779            let v = i64::try_from(*v).map_err(|_| {
780                MutationExecutorError::Bind(format!("u64 value {v} exceeds i64 range"))
781            })?;
782            args.add(v);
783        }
784        Value::F64(v) => args.add(*v),
785        Value::Decimal(v) => args.add(*v),
786        Value::Text(v) => match try_parse_datetime_from_str(v) {
787            Some(dt) => args.add(dt),
788            None => args.add(v.clone()),
789        },
790        Value::Json(v) => {
791            let j_val: serde_json::Value =
792                serde_json::to_value(v).map_err(|e| MutationExecutorError::Bind(e.to_string()))?;
793            args.add(j_val);
794        }
795        Value::Date(v) => args.add(*v),
796        Value::Timestamp(v) => args.add(v.to_datetime()),
797        Value::Object(_) => return Err(MutationExecutorError::UnsupportedValue("object")),
798        Value::List(values) => bind_pg_list(args, values)?,
799        Value::TypedNull(dt) => match dt {
800            DataType::Bool => args.add(Option::<bool>::None),
801            DataType::I64 | DataType::U64 => args.add(Option::<i64>::None),
802            DataType::F64 => args.add(Option::<f64>::None),
803            DataType::Decimal => args.add(Option::<Decimal>::None),
804            DataType::Text | DataType::LargeText => args.add(Option::<String>::None),
805            DataType::Json => args.add(Option::<serde_json::Value>::None),
806            DataType::Date => args.add(Option::<NaiveDate>::None),
807            DataType::Timestamp => args.add(Option::<DateTime<Utc>>::None),
808        },
809    }
810    Ok(())
811}
812
813fn bind_pg_list(args: &mut PgArgs, values: &[Value]) -> Result<(), MutationExecutorError> {
814    let Some(first) = values.first() else {
815        return Err(MutationExecutorError::UnsupportedValue("empty list"));
816    };
817    match first {
818        Value::Bool(_) => {
819            let values = values
820                .iter()
821                .map(|value| match value {
822                    Value::Bool(value) => Ok(*value),
823                    _ => Err(MutationExecutorError::UnsupportedValue("mixed bool list")),
824                })
825                .collect::<Result<Vec<_>, _>>()?;
826            args.add(values);
827        }
828        Value::I64(_) => {
829            let values = values
830                .iter()
831                .map(|value| match value {
832                    Value::I64(value) => Ok(*value),
833                    _ => Err(MutationExecutorError::UnsupportedValue("mixed i64 list")),
834                })
835                .collect::<Result<Vec<_>, _>>()?;
836            args.add(values);
837        }
838        Value::U64(_) => {
839            let values = values
840                .iter()
841                .map(|value| match value {
842                    Value::U64(value) => i64::try_from(*value).map_err(|_| {
843                        MutationExecutorError::Bind(format!("u64 value {value} exceeds i64 range"))
844                    }),
845                    _ => Err(MutationExecutorError::UnsupportedValue("mixed u64 list")),
846                })
847                .collect::<Result<Vec<_>, _>>()?;
848            args.add(values);
849        }
850        Value::F64(_) => {
851            let values = values
852                .iter()
853                .map(|value| match value {
854                    Value::F64(value) => Ok(*value),
855                    _ => Err(MutationExecutorError::UnsupportedValue("mixed f64 list")),
856                })
857                .collect::<Result<Vec<_>, _>>()?;
858            args.add(values);
859        }
860        Value::Decimal(_) => {
861            let values = values
862                .iter()
863                .map(|value| match value {
864                    Value::Decimal(value) => Ok(*value),
865                    _ => Err(MutationExecutorError::UnsupportedValue(
866                        "mixed decimal list",
867                    )),
868                })
869                .collect::<Result<Vec<_>, _>>()?;
870            args.add(values);
871        }
872        Value::Text(_) => {
873            let values = values
874                .iter()
875                .map(|value| match value {
876                    Value::Text(value) => Ok(value.clone()),
877                    _ => Err(MutationExecutorError::UnsupportedValue("mixed text list")),
878                })
879                .collect::<Result<Vec<_>, _>>()?;
880            args.add(values);
881        }
882        Value::Date(_) => {
883            let values = values
884                .iter()
885                .map(|value| match value {
886                    Value::Date(value) => Ok(*value),
887                    _ => Err(MutationExecutorError::UnsupportedValue("mixed date list")),
888                })
889                .collect::<Result<Vec<_>, _>>()?;
890            args.add(values);
891        }
892        Value::Timestamp(_) => {
893            let values = values
894                .iter()
895                .map(|value| match value {
896                    Value::Timestamp(value) => Ok(value.to_datetime()),
897                    _ => Err(MutationExecutorError::UnsupportedValue(
898                        "mixed timestamp list",
899                    )),
900                })
901                .collect::<Result<Vec<_>, _>>()?;
902            args.add(values);
903        }
904        Value::Null => return Err(MutationExecutorError::UnsupportedValue("null list")),
905        Value::Json(_) => return Err(MutationExecutorError::UnsupportedValue("json list")),
906        Value::Object(_) => return Err(MutationExecutorError::UnsupportedValue("object list")),
907        Value::List(_) => return Err(MutationExecutorError::UnsupportedValue("nested list")),
908        Value::TypedNull(_) => return Err(MutationExecutorError::UnsupportedValue("null list")),
909    }
910    Ok(())
911}
912
913fn decode_pg_row(row: &tokio_postgres::Row) -> Result<Record, MutationExecutorError> {
914    let mut record = BTreeMap::new();
915    for (index, column) in row.columns().iter().enumerate() {
916        let name = column.name().to_owned();
917        let type_name = column.type_().name().to_ascii_uppercase();
918
919        let value = match type_name.as_str() {
920            "BOOL" | "BOOLEAN" => {
921                let v: Option<bool> = row.try_get(index)?;
922                match v {
923                    Some(v) => Value::Bool(v),
924                    None => Value::Null,
925                }
926            }
927            "INT2" => {
928                let v: Option<i16> = row.try_get(index)?;
929                match v {
930                    Some(v) => Value::I64(v as i64),
931                    None => Value::Null,
932                }
933            }
934            "INT4" => {
935                let v: Option<i32> = row.try_get(index)?;
936                match v {
937                    Some(v) => Value::I64(v as i64),
938                    None => Value::Null,
939                }
940            }
941            "INT8" => {
942                let v: Option<i64> = row.try_get(index)?;
943                match v {
944                    Some(v) => Value::I64(v),
945                    None => Value::Null,
946                }
947            }
948            "FLOAT4" => {
949                let v: Option<f32> = row.try_get(index)?;
950                match v {
951                    Some(v) => Value::F64(v as f64),
952                    None => Value::Null,
953                }
954            }
955            "FLOAT8" => {
956                let v: Option<f64> = row.try_get(index)?;
957                match v {
958                    Some(v) => Value::F64(v),
959                    None => Value::Null,
960                }
961            }
962            "NUMERIC" => {
963                let v: Option<Decimal> = row.try_get(index)?;
964                match v {
965                    Some(v) => Value::Decimal(v),
966                    None => Value::Null,
967                }
968            }
969            "JSON" | "JSONB" => {
970                let v: Option<serde_json::Value> = row.try_get(index)?;
971                match v {
972                    Some(j) => Value::Json(j.into()),
973                    None => Value::Null,
974                }
975            }
976            "DATE" => {
977                let v: Option<NaiveDate> = row.try_get(index)?;
978                match v {
979                    Some(v) => Value::Date(v),
980                    None => Value::Null,
981                }
982            }
983            "TIMESTAMP" | "TIMESTAMPTZ" => {
984                let v: Option<DateTime<Utc>> = row.try_get(index)?;
985                match v {
986                    Some(v) => Value::Timestamp(teaql_core::time::Timestamp(v.timestamp_millis())),
987                    None => Value::Null,
988                }
989            }
990            "TEXT" | "VARCHAR" | "BPCHAR" | "NAME" | "UUID" => {
991                let v: Option<String> = row.try_get(index)?;
992                match v {
993                    Some(v) => Value::Text(v),
994                    None => Value::Null,
995                }
996            }
997            other => {
998                return Err(MutationExecutorError::UnsupportedColumnType(
999                    other.to_owned(),
1000                ));
1001            }
1002        };
1003        record.insert(name, value);
1004    }
1005    Ok(record)
1006}
1007
1008#[cfg(test)]
1009mod tests {
1010    use super::*;
1011    use teaql_core::{DeleteCommand, RecoverCommand};
1012
1013    fn entity() -> EntityDescriptor {
1014        EntityDescriptor::new("Order")
1015            .table_name("orders")
1016            .property(
1017                PropertyDescriptor::new("id", DataType::U64)
1018                    .column_name("id")
1019                    .id()
1020                    .not_null(),
1021            )
1022            .property(
1023                PropertyDescriptor::new("version", DataType::I64)
1024                    .column_name("version")
1025                    .version()
1026                    .not_null(),
1027            )
1028            .property(PropertyDescriptor::new("name", DataType::Text).column_name("name"))
1029    }
1030
1031    #[test]
1032    fn postgres_dialect_compiles_mutations_with_numbered_placeholders() {
1033        let insert = PostgresDialect
1034            .compile_insert(
1035                &entity(),
1036                &InsertCommand::new("Order")
1037                    .value("id", 1_u64)
1038                    .value("name", "A"),
1039            )
1040            .unwrap();
1041        assert_eq!(insert.sql, "INSERT INTO orders (id, name) VALUES ($1, $2)");
1042
1043        let update = PostgresDialect
1044            .compile_update(
1045                &entity(),
1046                &UpdateCommand::new("Order", 1_u64)
1047                    .expected_version(3)
1048                    .value("name", "B"),
1049            )
1050            .unwrap();
1051        assert_eq!(
1052            update.sql,
1053            "UPDATE orders SET name = $1, version = $2 WHERE id = $3 AND version = $4"
1054        );
1055
1056        let delete = PostgresDialect
1057            .compile_delete(
1058                &entity(),
1059                &DeleteCommand::new("Order", 1_u64).expected_version(3),
1060            )
1061            .unwrap();
1062        let recover = PostgresDialect
1063            .compile_recover(&entity(), &RecoverCommand::new("Order", 1_u64, -4))
1064            .unwrap();
1065        assert_eq!(
1066            delete.sql,
1067            "UPDATE orders SET version = $1 WHERE id = $2 AND version = $3"
1068        );
1069        assert_eq!(
1070            recover.sql,
1071            "UPDATE orders SET version = $1 WHERE id = $2 AND version = $3"
1072        );
1073    }
1074
1075    #[test]
1076    fn postgres_dialect_compiles_schema_and_large_in_array_binds() {
1077        let create = PostgresDialect.compile_create_table(&entity()).unwrap();
1078        assert_eq!(
1079            create,
1080            "CREATE TABLE IF NOT EXISTS orders (id BIGINT PRIMARY KEY NOT NULL, version BIGINT NOT NULL, name VARCHAR(255))"
1081        );
1082        assert!(
1083            PostgresDialect
1084                .schema_setup_sqls()
1085                .iter()
1086                .any(|sql| sql.contains("CREATE OR REPLACE FUNCTION soundex"))
1087        );
1088
1089        let query = PostgresDialect
1090            .compile_select(
1091                &entity(),
1092                &SelectQuery::new("Order")
1093                    .filter(Expr::in_large(
1094                        "id",
1095                        vec![Value::from(1_u64), Value::from(2_u64)],
1096                    ))
1097                    .order_asc("id"),
1098            )
1099            .unwrap();
1100        assert_eq!(
1101            query.sql,
1102            "SELECT id, version, name FROM orders WHERE (id = ANY($1)) ORDER BY id ASC"
1103        );
1104        assert_eq!(
1105            query.params,
1106            vec![Value::List(vec![Value::U64(1), Value::U64(2)])]
1107        );
1108    }
1109}
teaql_provider_postgres/lib.rs

teaql_provider_postgres/
lib.rs
