indexlake_catalog_postgres/
catalog.rs

1use arrow::array::RecordBatch;
2use arrow::datatypes::Schema;
3use bb8::Pool;
4use bb8_postgres::PostgresConnectionManager;
5use bb8_postgres::tokio_postgres::NoTls;
6use bb8_postgres::tokio_postgres::binary_copy::BinaryCopyInWriter;
7use bb8_postgres::tokio_postgres::types::ToSql;
8use futures::StreamExt;
9use indexlake::catalog::{
10    Catalog, CatalogDataType, CatalogDatabase, CatalogSchemaRef, Row, RowStream, Scalar,
11    Transaction,
12};
13use indexlake::expr::{BinaryExpr, Expr};
14use indexlake::{ILError, ILResult};
15use log::{error, trace};
16use uuid::Uuid;
17
18use crate::SqlDisplay;
19use crate::types::{catalog_data_type_to_pg_type, scalar_to_pg_value};
20
21#[derive(Debug, Clone)]
22pub struct PostgresCatalog {
23    pool: Pool<PostgresConnectionManager<NoTls>>,
24}
25
26impl PostgresCatalog {
27    pub fn new(pool: Pool<PostgresConnectionManager<NoTls>>) -> Self {
28        Self { pool }
29    }
30}
31
32#[async_trait::async_trait]
33impl Catalog for PostgresCatalog {
34    fn database(&self) -> CatalogDatabase {
35        CatalogDatabase::Postgres
36    }
37
38    async fn query(&self, sql: &str, schema: CatalogSchemaRef) -> ILResult<RowStream<'static>> {
39        trace!("postgres query: {sql}");
40        let conn = self
41            .pool
42            .get_owned()
43            .await
44            .map_err(|e| ILError::catalog(format!("failed to get postgres connection: {e}")))?;
45        let pg_row_stream = conn
46            .query_raw(sql, Vec::<String>::new())
47            .await
48            .map_err(|e| {
49                ILError::catalog(format!("failed to query postgres: {} {e}", SqlDisplay(sql)))
50            })?;
51
52        let stream = pg_row_stream.map(move |row| {
53            let pg_row =
54                row.map_err(|e| ILError::catalog(format!("failed to get postgres row: {e}")))?;
55            pg_row_to_row(&pg_row, &schema)
56        });
57        Ok(Box::pin(stream))
58    }
59
60    async fn transaction(&self) -> ILResult<Box<dyn Transaction>> {
61        let conn = self
62            .pool
63            .get_owned()
64            .await
65            .map_err(|e| ILError::catalog(format!("failed to get postgres connection: {e}")))?;
66        conn.batch_execute("START TRANSACTION")
67            .await
68            .map_err(|e| ILError::catalog(format!("failed to start postgres txn: {e}")))?;
69        Ok(Box::new(PostgresTransaction {
70            conn: Some(conn),
71            done: false,
72        }))
73    }
74
75    async fn truncate(&self, table_name: &str) -> ILResult<()> {
76        let conn = self
77            .pool
78            .get()
79            .await
80            .map_err(|e| ILError::catalog(format!("failed to get postgres connection: {e}")))?;
81        conn.execute(
82            &format!("TRUNCATE TABLE {}", self.sql_identifier(table_name)),
83            &[],
84        )
85        .await
86        .map(|r| r as usize)
87        .map_err(|e| {
88            ILError::catalog(format!(
89                "failed to truncate table {table_name} on postgres: {e}"
90            ))
91        })?;
92        Ok(())
93    }
94
95    async fn size(&self, table_name: &str) -> ILResult<usize> {
96        let conn = self
97            .pool
98            .get()
99            .await
100            .map_err(|e| ILError::catalog(format!("failed to get postgres connection: {e}")))?;
101        let row = conn
102            .query_one(
103                &format!("SELECT pg_total_relation_size('{table_name}')"),
104                &[],
105            )
106            .await
107            .map_err(|e| {
108                ILError::catalog(format!(
109                    "failed to get size of table {table_name} on postgres: {e}"
110                ))
111            })?;
112        let size = row
113            .try_get::<_, Option<i64>>(0)
114            .map_err(|e| ILError::storage(format!("Failed to get size from row: {e}")))?;
115        Ok(size.unwrap_or(0) as usize)
116    }
117
118    fn sql_identifier(&self, ident: &str) -> String {
119        format!("\"{ident}\"")
120    }
121
122    fn sql_binary_literal(&self, value: &[u8]) -> String {
123        format!("E'\\\\x{}'", hex::encode(value))
124    }
125
126    fn sql_uuid_literal(&self, value: &Uuid) -> String {
127        format!("'{value}'")
128    }
129
130    fn sql_string_literal(&self, value: &str) -> String {
131        let value = value.replace("'", "''");
132        format!("'{value}'")
133    }
134
135    // TODO impl this
136    fn supports_filter(&self, filter: &Expr, _schema: &Schema) -> ILResult<bool> {
137        match filter {
138            Expr::Function(_) => Ok(false),
139            Expr::BinaryExpr(BinaryExpr { left, right, .. }) => {
140                if let Expr::Column(_) = left.as_ref()
141                    && let Expr::Literal(lit) = right.as_ref()
142                    && matches!(lit.value, Scalar::List(_))
143                {
144                    Ok(false)
145                } else if let Expr::Literal(lit) = left.as_ref()
146                    && let Expr::Column(_) = right.as_ref()
147                    && matches!(lit.value, Scalar::List(_))
148                {
149                    Ok(false)
150                } else {
151                    Ok(true)
152                }
153            }
154            _ => Ok(true),
155        }
156    }
157
158    fn unparse_expr(&self, expr: &Expr) -> ILResult<String> {
159        match expr {
160            Expr::Column(name) => Ok(self.sql_identifier(name)),
161            Expr::Literal(literal) => literal.value.to_sql(self),
162            Expr::BinaryExpr(binary_expr) => {
163                let left = self.unparse_expr(&binary_expr.left)?;
164                let right = self.unparse_expr(&binary_expr.right)?;
165                Ok(format!("({} {} {})", left, binary_expr.op, right))
166            }
167            Expr::Not(expr) => Ok(format!("NOT {}", self.unparse_expr(expr)?)),
168            Expr::IsNull(expr) => Ok(format!("{} IS NULL", self.unparse_expr(expr)?)),
169            Expr::IsNotNull(expr) => Ok(format!("{} IS NOT NULL", self.unparse_expr(expr)?)),
170            Expr::InList(in_list) => {
171                let list = in_list
172                    .list
173                    .iter()
174                    .map(|expr| self.unparse_expr(expr))
175                    .collect::<ILResult<Vec<_>>>()?
176                    .join(", ");
177                Ok(format!(
178                    "{} IN ({})",
179                    self.unparse_expr(&in_list.expr)?,
180                    list
181                ))
182            }
183            Expr::Function(_) => Err(ILError::invalid_input(
184                "Function can only be used for index",
185            )),
186            Expr::Like(like) => {
187                let expr = self.unparse_expr(&like.expr)?;
188                let pattern = self.unparse_expr(&like.pattern)?;
189                match (like.negated, like.case_insensitive) {
190                    (true, true) => Ok(format!("{expr} NOT ILIKE {pattern}")),
191                    (true, false) => Ok(format!("{expr} NOT LIKE {pattern}")),
192                    (false, true) => Ok(format!("{expr} ILIKE {pattern}")),
193                    (false, false) => Ok(format!("{expr} LIKE {pattern}")),
194                }
195            }
196            Expr::Cast(cast) => {
197                let catalog_datatype = CatalogDataType::from_arrow(&cast.cast_type)?;
198                let expr_sql = self.unparse_expr(&cast.expr)?;
199                Ok(format!(
200                    "CAST({} AS {})",
201                    expr_sql,
202                    self.unparse_catalog_data_type(catalog_datatype),
203                ))
204            }
205            Expr::TryCast(_) => Err(ILError::invalid_input("TRY_CAST is not supported in SQL")),
206            Expr::Negative(expr) => Ok(format!("-{}", self.unparse_expr(expr)?)),
207            Expr::Case(case) => {
208                let mut sql = String::new();
209                sql.push_str("CASE");
210                for (when, then) in &case.when_then {
211                    sql.push_str(&format!(
212                        " WHEN {} THEN {}",
213                        self.unparse_expr(when)?,
214                        self.unparse_expr(then)?,
215                    ));
216                }
217                if let Some(else_expr) = &case.else_expr {
218                    sql.push_str(&format!(" ELSE {}", self.unparse_expr(else_expr)?));
219                }
220                sql.push_str(" END");
221                Ok(sql)
222            }
223        }
224    }
225
226    fn unparse_catalog_data_type(&self, data_type: CatalogDataType) -> String {
227        match data_type {
228            CatalogDataType::Boolean => "BOOLEAN".to_string(),
229            CatalogDataType::Int8 => "SMALLINT".to_string(),
230            CatalogDataType::Int16 => "SMALLINT".to_string(),
231            CatalogDataType::Int32 => "INTEGER".to_string(),
232            CatalogDataType::Int64 => "BIGINT".to_string(),
233            CatalogDataType::UInt8 => "SMALLINT".to_string(),
234            CatalogDataType::UInt16 => "INTEGER".to_string(),
235            CatalogDataType::UInt32 => "BIGINT".to_string(),
236            CatalogDataType::UInt64 => "FLOAT4".to_string(),
237            CatalogDataType::Float32 => "FLOAT4".to_string(),
238            CatalogDataType::Float64 => "FLOAT8".to_string(),
239            CatalogDataType::Utf8 => "VARCHAR".to_string(),
240            CatalogDataType::Binary => "BYTEA".to_string(),
241            CatalogDataType::Uuid => "UUID".to_string(),
242        }
243    }
244}
245
246#[derive(Debug)]
247pub struct PostgresTransaction {
248    conn: Option<bb8::PooledConnection<'static, PostgresConnectionManager<NoTls>>>,
249    done: bool,
250}
251
252impl PostgresTransaction {
253    fn check_done(&self) -> ILResult<()> {
254        if self.done {
255            return Err(ILError::catalog(
256                "Transaction already committed or rolled back",
257            ));
258        }
259        Ok(())
260    }
261
262    fn conn_mut(
263        &mut self,
264    ) -> ILResult<&mut bb8::PooledConnection<'static, PostgresConnectionManager<NoTls>>> {
265        self.conn
266            .as_mut()
267            .ok_or_else(|| ILError::catalog("connection already taken"))
268    }
269}
270
271#[async_trait::async_trait]
272impl Transaction for PostgresTransaction {
273    async fn query(&mut self, sql: &str, schema: CatalogSchemaRef) -> ILResult<RowStream> {
274        trace!("postgres txn query: {sql}");
275        self.check_done()?;
276
277        let pg_row_stream = self
278            .conn_mut()?
279            .query_raw(sql, Vec::<String>::new())
280            .await
281            .map_err(|e| {
282                ILError::catalog(format!("failed to query postgres: {} {e}", SqlDisplay(sql)))
283            })?;
284
285        let stream = pg_row_stream.map(move |row| {
286            let pg_row =
287                row.map_err(|e| ILError::catalog(format!("failed to get postgres row: {e}")))?;
288            pg_row_to_row(&pg_row, &schema)
289        });
290        Ok(Box::pin(stream))
291    }
292
293    async fn execute(&mut self, sql: &str) -> ILResult<usize> {
294        trace!("postgres txn execute: {sql}");
295        self.check_done()?;
296        self.conn_mut()?
297            .execute(sql, &[])
298            .await
299            .map(|r| r as usize)
300            .map_err(|e| {
301                ILError::catalog(format!(
302                    "failed to execute postgres: {} {e}",
303                    SqlDisplay(sql)
304                ))
305            })
306    }
307
308    async fn execute_batch(&mut self, sqls: &[String]) -> ILResult<()> {
309        trace!("postgres txn execute batch: {:?}", sqls);
310        self.check_done()?;
311        let sql = sqls.join(";");
312        self.conn_mut()?.batch_execute(&sql).await.map_err(|e| {
313            ILError::catalog(format!(
314                "failed to execute batch postgres: {} {e}",
315                SqlDisplay(&sql)
316            ))
317        })
318    }
319
320    async fn insert_rows(
321        &mut self,
322        table_name: &str,
323        field_names: &[String],
324        batches: &[RecordBatch],
325    ) -> ILResult<()> {
326        trace!("postgres txn insert rows: {table_name}");
327        self.check_done()?;
328
329        let copy_sql = format!(
330            "COPY {} ({}) FROM STDIN BINARY",
331            table_name,
332            field_names.join(", ")
333        );
334        trace!("postgres txn binary copy_in: {copy_sql}");
335
336        let sink = self
337            .conn_mut()?
338            .copy_in(&copy_sql)
339            .await
340            .map_err(|e| ILError::catalog(format!("failed to start COPY: {e}")))?;
341        let pg_types = batches[0]
342            .schema()
343            .fields()
344            .iter()
345            .map(|f| CatalogDataType::from_arrow(f.data_type()).map(catalog_data_type_to_pg_type))
346            .collect::<ILResult<Vec<_>>>()?;
347        let mut writer = std::pin::pin!(BinaryCopyInWriter::new(sink, &pg_types));
348
349        for batch in batches {
350            for row_idx in 0..batch.num_rows() {
351                let values = batch
352                    .columns()
353                    .iter()
354                    .map(|col| {
355                        let scalar = Scalar::try_from_array(col.as_ref(), row_idx)?;
356                        scalar_to_pg_value(scalar)
357                    })
358                    .collect::<ILResult<Vec<_>>>()?;
359                let refs: Vec<&(dyn ToSql + Sync)> =
360                    values.iter().map(|v| v as &(dyn ToSql + Sync)).collect();
361                writer
362                    .as_mut()
363                    .write(&refs)
364                    .await
365                    .map_err(|e| ILError::catalog(format!("failed to write COPY row: {e}")))?;
366            }
367        }
368
369        let rows = writer
370            .finish()
371            .await
372            .map_err(|e| ILError::catalog(format!("failed to finish COPY: {e}")))?;
373        trace!("postgres COPY inserted {} rows", rows);
374
375        Ok(())
376    }
377
378    async fn commit(&mut self) -> ILResult<()> {
379        trace!("postgres txn commit");
380        self.check_done()?;
381        self.conn_mut()?
382            .batch_execute("COMMIT")
383            .await
384            .map_err(|e| ILError::catalog(format!("failed to commit postgres txn: {e}")))?;
385        self.done = true;
386        Ok(())
387    }
388
389    async fn rollback(&mut self) -> ILResult<()> {
390        trace!("postgres txn rollback");
391        self.check_done()?;
392        self.conn_mut()?
393            .batch_execute("ROLLBACK")
394            .await
395            .map_err(|e| ILError::catalog(format!("failed to rollback postgres txn: {e}")))?;
396        self.done = true;
397        Ok(())
398    }
399}
400
401impl Drop for PostgresTransaction {
402    fn drop(&mut self) {
403        if self.done {
404            return;
405        }
406        if let Some(conn) = self.conn.take() {
407            tokio::spawn(async move {
408                if let Err(e) = conn.batch_execute("ROLLBACK").await {
409                    error!("[indexlake] failed to rollback postgres txn: {e}");
410                }
411            });
412        }
413    }
414}
415
416fn pg_row_to_row(
417    pg_row: &bb8_postgres::tokio_postgres::Row,
418    schema: &CatalogSchemaRef,
419) -> ILResult<Row> {
420    let mut values = Vec::new();
421    let err_mapping = |e: bb8_postgres::tokio_postgres::error::Error| {
422        ILError::catalog(format!("failed to get row value: {e}"))
423    };
424    for (idx, field) in schema.columns.iter().enumerate() {
425        let scalar = match field.data_type {
426            CatalogDataType::Boolean => {
427                let v: Option<bool> = pg_row.try_get(idx).map_err(err_mapping)?;
428                Scalar::Boolean(v)
429            }
430            CatalogDataType::Int8 => {
431                let v: Option<i16> = pg_row.try_get(idx).map_err(err_mapping)?;
432                Scalar::Int8(v.map(|v| v as i8))
433            }
434            CatalogDataType::Int16 => {
435                let v: Option<i16> = pg_row.try_get(idx).map_err(err_mapping)?;
436                Scalar::Int16(v)
437            }
438            CatalogDataType::Int32 => {
439                let v: Option<i32> = pg_row.try_get(idx).map_err(err_mapping)?;
440                Scalar::Int32(v)
441            }
442            CatalogDataType::Int64 => {
443                let v: Option<i64> = pg_row.try_get(idx).map_err(err_mapping)?;
444                Scalar::Int64(v)
445            }
446            CatalogDataType::UInt8 => {
447                let v: Option<i16> = pg_row.try_get(idx).map_err(err_mapping)?;
448                Scalar::UInt8(v.map(|v| v as u8))
449            }
450            CatalogDataType::UInt16 => {
451                let v: Option<i32> = pg_row.try_get(idx).map_err(err_mapping)?;
452                Scalar::UInt16(v.map(|v| v as u16))
453            }
454            CatalogDataType::UInt32 => {
455                let v: Option<i64> = pg_row.try_get(idx).map_err(err_mapping)?;
456                Scalar::UInt32(v.map(|v| v as u32))
457            }
458            CatalogDataType::UInt64 => {
459                let v: Option<f32> = pg_row.try_get(idx).map_err(err_mapping)?;
460                Scalar::UInt64(v.map(|v| v as u64))
461            }
462            CatalogDataType::Float32 => {
463                let v: Option<f32> = pg_row.try_get(idx).map_err(err_mapping)?;
464                Scalar::Float32(v)
465            }
466            CatalogDataType::Float64 => {
467                let v: Option<f64> = pg_row.try_get(idx).map_err(err_mapping)?;
468                Scalar::Float64(v)
469            }
470            CatalogDataType::Utf8 => {
471                let v: Option<String> = pg_row.try_get(idx).map_err(err_mapping)?;
472                Scalar::Utf8(v)
473            }
474            CatalogDataType::Binary => {
475                let v: Option<Vec<u8>> = pg_row.try_get(idx).map_err(err_mapping)?;
476                Scalar::Binary(v)
477            }
478            CatalogDataType::Uuid => {
479                let v: Option<uuid::Uuid> = pg_row.try_get(idx).map_err(err_mapping)?;
480                Scalar::Binary(v.map(|v| v.as_bytes().to_vec()))
481            }
482        };
483        if !field.nullable && scalar.is_null() {
484            return Err(ILError::catalog(format!(
485                "column {} is not nullable but got null value",
486                field.name
487            )));
488        }
489        values.push(scalar);
490    }
491    Ok(Row::new(schema.clone(), values))
492}
indexlake_catalog_postgres/catalog.rs

indexlake_catalog_postgres/
catalog.rs
