feat: pg_dump support + schema evolution

docs/analytics/quickstart.mdx

@@ -64,12 +64,12 @@ will become resolved as `deltalake` tables become production-ready.
 - [ ] `UPDATE` statements
 - [ ] Nested `DELETE` statements
 - [ ] Partitioning tables by column
-- [ ] Some Postgres types like arrays, JSON, time, and timestamp with time zone
+- [ ] Some Postgres types like JSON, time, and timestamp with time zone
 - [ ] User-defined functions, aggregations, or types
 - [ ] Referencing `deltalake` and regular Postgres `heap` tables in the same query
 - [ ] Write-ahead-log (WAL) support and `ROLLBACK`
 - [ ] Foreign keys
 - [ ] Index scans
-- [ ] `TEMP` tables
+- [ ] Collations
 - [ ] Using an external data lake as a table storage provider
 - [ ] Full text search over `deltalake` tables with `pg_bm25`

pg_analytics/README.md

@@ -65,6 +65,7 @@ Parquet files belonging to dropped data.
 - [x] `deltalake` tables behave like regular Postgres tables and support most Postgres queries (JOINs, CTEs, window functions, etc.)
 - [x] Vacuum and Parquet storage optimization
 - [x] `INSERT`, `TRUNCATE`, `DELETE`, `COPY`
+- [x] Physical backups with `pg_dump`
 
 ### Known Limitations
 
@@ -73,13 +74,13 @@ As `pg_analytics` becomes production-ready, many of these will be resolved.
 - [ ] `UPDATE` statements
 - [ ] Nested `DELETE` statements
 - [ ] Partitioning tables by column
-- [ ] Some Postgres types like arrays, JSON, time, and timestamp with time zone
+- [ ] Some Postgres types like JSON, time, and timestamp with time zone
 - [ ] User-defined functions, aggregations, or types
 - [ ] Referencing `deltalake` and regular Postgres `heap` tables in the same query
-- [ ] Write-ahead-log (WAL) support and `ROLLBACK`
+- [ ] Write-ahead-log (WAL) support/`ROLLBACK`/logical replication
 - [ ] Foreign keys
 - [ ] Index scans
-- [ ] `TEMP` tables
+- [ ] Collations
 - [ ] Using an external data lake as a table storage provider
 - [ ] Full text search over `deltalake` tables with `pg_bm25`
 

pg_analytics/src/datafusion/datatype.rs

@@ -68,14 +68,25 @@
 
     fn from_sql_data_type(sql_data_type: SQLDataType) -> Result<DataType, ParadeError> {
         let result = match sql_data_type {
-            SQLDataType::Boolean => DataType::Boolean,
-            SQLDataType::Text => DataType::Utf8,
-            SQLDataType::Int2(_) => DataType::Int16,
-            SQLDataType::Int4(_) => DataType::Int32,
-            SQLDataType::Int8(_) => DataType::Int64,
-            SQLDataType::UnsignedInt4(_) => DataType::UInt32,
-            SQLDataType::Float4 => DataType::Float32,
-            SQLDataType::Float8 => DataType::Float64,
+            SQLDataType::Boolean | SQLDataType::Bool => DataType::Boolean,
+            SQLDataType::Text
+            | SQLDataType::Char(_)
+            | SQLDataType::Varchar(_)
+            | SQLDataType::String(_) => DataType::Utf8,
+            SQLDataType::TinyInt(_) => DataType::Int8,
+            SQLDataType::SmallInt(_) | SQLDataType::Int2(_) => DataType::Int16,
+            SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_) => DataType::Int32,
+            SQLDataType::BigInt(_) | SQLDataType::Int8(_) => DataType::Int64,
+            SQLDataType::UnsignedTinyInt(_) => DataType::UInt8,
+            SQLDataType::UnsignedSmallInt(_) | SQLDataType::UnsignedInt2(_) => DataType::UInt16,
+            SQLDataType::UnsignedInt(_)
+            | SQLDataType::UnsignedInteger(_)
+            | SQLDataType::UnsignedInt4(_) => DataType::UInt32,
+            SQLDataType::UnsignedBigInt(_) | SQLDataType::UnsignedInt8(_) => DataType::UInt64,
+            SQLDataType::Float4 | SQLDataType::Float(_) => DataType::Float32,
+            SQLDataType::Double | SQLDataType::DoublePrecision | SQLDataType::Float8 => {
+                DataType::Float64
+            }
             SQLDataType::Bytea => DataType::Binary,
             SQLDataType::Numeric(ExactNumberInfo::PrecisionAndScale(precision, scale)) => {
                 let casted_precision = precision as u8;

pg_analytics/src/datafusion/schema.rs

@@ -1,11 +1,15 @@
+use async_std::stream::StreamExt;
 use async_std::task;
 use async_trait::async_trait;
 use deltalake::datafusion::arrow::datatypes::Schema as ArrowSchema;
 use deltalake::datafusion::arrow::record_batch::RecordBatch;
 use deltalake::datafusion::catalog::schema::SchemaProvider;
 use deltalake::datafusion::datasource::TableProvider;
 use deltalake::datafusion::error::Result;
+use deltalake::datafusion::execution::context::SessionState;
+use deltalake::datafusion::execution::TaskContext;
 use deltalake::datafusion::logical_expr::Expr;
+use deltalake::datafusion::physical_plan::SendableRecordBatchStream;
 use deltalake::kernel::Action;
 use deltalake::kernel::Schema as DeltaSchema;
 use deltalake::operations::create::CreateBuilder;
@@ -18,6 +22,8 @@
 use deltalake::protocol::{DeltaOperation, SaveMode};
 use deltalake::storage::ObjectStoreRef;
 use deltalake::table::state::DeltaTableState;
+use deltalake::writer::record_batch::RecordBatchWriter;
+use deltalake::writer::{DeltaWriter as DeltaWriterTrait, WriteMode};
 use deltalake::DeltaTable;
 use parking_lot::{Mutex, RwLock};
 use pgrx::*;
@@ -39,6 +45,7 @@
     schema_name: String,
     tables: RwLock<HashMap<String, Arc<dyn TableProvider>>>,
     writers: Mutex<HashMap<String, DeltaWriter>>,
+    streams: Mutex<HashMap<String, SendableRecordBatchStream>>,
     dir: PathBuf,
 }
 
@@ -49,6 +56,7 @@
             schema_name: schema_name.to_string(),
             tables: RwLock::new(HashMap::new()),
             writers: Mutex::new(HashMap::new()),
+            streams: Mutex::new(HashMap::new()),
             dir,
         })
     }
@@ -250,7 +258,7 @@
         let mut writer_lock = self.writers.lock();
         let writer = writer_lock
             .get_mut(table_name)
-            .ok_or(NotFound::Table(table_name.to_string()))?;
+            .ok_or(NotFound::Writer(table_name.to_string()))?;
 
         task::block_on(writer.write(&batch))?;
 
@@ -270,7 +278,7 @@
         let mut writer_lock = self.writers.lock();
         let writer = writer_lock
             .remove(table_name)
-            .ok_or(NotFound::Table(table_name.to_string()))?;
+            .ok_or(NotFound::Writer(table_name.to_string()))?;
 
         // Generate commit actions by closing the writer and commit to delta logs
         let actions = task::block_on(writer.close())?;
@@ -356,6 +364,41 @@
         Ok(metrics)
     }
 
+    pub async fn merge_schema(
+        &self,
+        table_name: &str,
+        batch: RecordBatch,
+    ) -> Result<(), ParadeError> {
+        // Open the DeltaTable
+        let mut delta_table = Self::get_delta_table(self, table_name).await?;
+
+        // Write the RecordBatch to the DeltaTable
+        let mut writer = RecordBatchWriter::for_table(&delta_table)?;
+        writer
+            .write_with_mode(batch, WriteMode::MergeSchema)
+            .await?;
+        writer.flush_and_commit(&mut delta_table).await?;
+
+        // Update the DeltaTable
+        delta_table.load().await?;
+
+        // Create and register a new writer
+        Self::register_writer(
+            self,
+            table_name,
+            Self::create_writer(self, delta_table.object_store(), writer.arrow_schema())?,
+        )?;
+
+        // Commit the table
+        Self::register_table(
+            self,
+            table_name.to_string(),
+            Arc::new(delta_table) as Arc<dyn TableProvider>,
+        )?;
+
+        Ok(())
+    }
+
     // SchemaProvider stores immutable TableProviders, whereas many DeltaOps methods
     // require a mutable DeltaTable. This function gets a mutable DeltaTable from
     // a TableProvider using the DeltaOps UpdateBuilder.
@@ -410,6 +453,49 @@
         Ok(delta_table)
     }
 
+    pub fn register_stream(
+        &self,
+        name: &str,
+        stream: SendableRecordBatchStream,
+    ) -> Result<(), ParadeError> {
+        let mut streams = self.streams.lock();
+        streams.insert(name.to_string(), stream);
+
+        Ok(())
+    }
+
+    pub async fn create_stream(
+        &self,
+        name: &str,
+        state: &SessionState,
+        task_context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream, ParadeError> {
+        let delta_table = Self::get_delta_table(self, name).await?;
+
+        Ok(delta_table
+            .scan(state, None, &[], None)
+            .await
+            .map(|plan| plan.execute(0, task_context))??)
+    }
+
+    pub fn get_next_streamed_batch(&self, name: &str) -> Result<Option<RecordBatch>, ParadeError> {
+        let mut streams = self.streams.lock();
+        let stream = streams
+            .get_mut(name)
+            .ok_or(NotFound::Stream(name.to_string()))?;
+
+        let batch = task::block_on(stream.next());
+
+        match batch {
+            Some(Ok(b)) => Ok(Some(b)),
+            None => {
+                streams.remove(name);
+                Ok(None)
+            }
+            Some(Err(err)) => Err(ParadeError::DataFusion(err)),
+        }
+    }
+
     // Helper function to register a table writer
     fn register_writer(&self, name: &str, writer: DeltaWriter) -> Result<(), ParadeError> {
         let mut writers = self.writers.lock();

pg_analytics/src/errors.rs

@@ -51,6 +51,12 @@
     #[error("No table registered with name {0}")]
     Table(String),
 
+    #[error("No writer found for table {0}")]
+    Writer(String),
+
+    #[error("No stream found for table {0}")]
+    Stream(String),
+
     #[error("Failed to convert to datum {0}")]
     Datum(String),
 
@@ -81,8 +87,14 @@
     #[error("Custom Postgres types are not supported")]
     CustomPostgresType,
 
-    #[error("ALTER TABLE is not yet supported. Please DROP and CREATE the table instead.")]
-    AlterTable,
+    #[error("DROP COLUMN is not yet supported. Please recreate the table instead.")]
+    DropColumn,
+
+    #[error("ALTER COLUMN is not yet supported. Please recreate the table instead.")]
+    AlterColumn,
+
+    #[error("RENAME COLUMN is not yet supported. Please recreate the table instead.")]
+    RenameColumn,
 
     #[error("UPDATE is not yet supported for deltalake tables")]
     Update,

pg_analytics/src/hooks/alter.rs

@@ -1,9 +1,20 @@
+use async_std::task;
+use deltalake::datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
+use deltalake::datafusion::arrow::record_batch::RecordBatch;
+use deltalake::datafusion::sql::parser;
+use deltalake::datafusion::sql::sqlparser::ast::{AlterTableOperation::*, ColumnOption, Statement};
 use pgrx::*;
+use std::sync::Arc;
 
+use crate::datafusion::context::DatafusionContext;
+use crate::datafusion::datatype::DatafusionTypeTranslator;
 use crate::errors::{NotSupported, ParadeError};
 use crate::hooks::handler::DeltaHandler;
 
-pub unsafe fn alter(alter_stmt: *mut pg_sys::AlterTableStmt) -> Result<(), ParadeError> {
+pub unsafe fn alter(
+    alter_stmt: *mut pg_sys::AlterTableStmt,
+    statement: &parser::Statement,
+) -> Result<(), ParadeError> {
     let rangevar = (*alter_stmt).relation;
     let rangevar_oid = pg_sys::RangeVarGetRelidExtended(
         rangevar,
@@ -23,5 +34,46 @@
         return Ok(());
     }
 
-    Err(NotSupported::AlterTable.into())
+    let pg_relation = unsafe { PgRelation::from_pg_owned(relation) };
+    let table_name = pg_relation.name();
+    let schema_name = pg_relation.namespace();
+    let mut fields_to_add = vec![];
+
+    if let parser::Statement::Statement(inner_statement) = statement {
+        if let Statement::AlterTable { operations, .. } = inner_statement.as_ref() {
+            for operation in operations {
+                match operation {
+                    AddColumn { column_def, .. } => {
+                        let options = &column_def.options;
+                        let nullability = options
+                            .iter()
+                            .any(|opt| matches!(opt.option, ColumnOption::Null));
+                        fields_to_add.push(Field::new(
+                            column_def.name.value.clone(),
+                            DataType::from_sql_data_type(column_def.data_type.clone())?,
+                            !nullability,
+                        ));
+                    }
+                    DropColumn { .. } => {
+                        return Err(NotSupported::DropColumn.into());
+                    }
+                    AlterColumn { .. } | ChangeColumn { .. } => {
+                        return Err(NotSupported::AlterColumn.into());
+                    }
+                    _ => {}
+                }
+            }
+        }
+    }
+
+    if !fields_to_add.is_empty() {
+        let schema = Arc::new(ArrowSchema::new(fields_to_add));
+        let batch = RecordBatch::new_empty(schema);
+
+        DatafusionContext::with_schema_provider(schema_name, |provider| {
+            task::block_on(provider.merge_schema(table_name, batch))
+        })?;
+    }
+
+    Ok(())
 }

pg_analytics/src/hooks/executor.rs

@@ -1,9 +1,7 @@
 use deltalake::datafusion::error::DataFusionError;
 use deltalake::datafusion::logical_expr::LogicalPlan;
-
 use deltalake::datafusion::sql::parser::DFParser;
 use deltalake::datafusion::sql::planner::SqlToRel;
-
 use deltalake::datafusion::sql::sqlparser::dialect::PostgreSqlDialect;
 use pgrx::*;
 use std::ffi::CStr;
@@ -30,12 +28,15 @@
     unsafe {
         let ps = query_desc.plannedstmt;
         let rtable = (*ps).rtable;
+        let query = CStr::from_ptr(query_desc.sourceText).to_str()?;
 
         // Only use this hook for deltalake tables
         // Allow INSERTs to go through
         if rtable.is_null()
             || query_desc.operation == pg_sys::CmdType_CMD_INSERT
             || !DeltaHandler::rtable_is_delta(rtable)?
+            // Tech Debt: Find a less hacky way to let COPY go through
+            || query.to_lowercase().starts_with("copy")
         {
             prev_hook(query_desc, direction, count, execute_once);
             return Ok(());
@@ -44,11 +45,11 @@
         // Execute SELECT, DELETE, UPDATE
         match query_desc.operation {
             pg_sys::CmdType_CMD_DELETE => {
-                let logical_plan = create_logical_plan(query_desc.clone())?;
+                let logical_plan = create_logical_plan(query)?;
                 delete(rtable, query_desc, logical_plan)
             }
             pg_sys::CmdType_CMD_SELECT => {
-                let logical_plan = create_logical_plan(query_desc.clone())?;
+                let logical_plan = create_logical_plan(query)?;
                 select(query_desc, logical_plan)
             }
             pg_sys::CmdType_CMD_UPDATE => Err(NotSupported::Update.into()),
@@ -61,9 +62,8 @@
 }
 
 #[inline]
-fn create_logical_plan(query_desc: PgBox<pg_sys::QueryDesc>) -> Result<LogicalPlan, ParadeError> {
+fn create_logical_plan(query: &str) -> Result<LogicalPlan, ParadeError> {
     let dialect = PostgreSqlDialect {};
-    let query = unsafe { CStr::from_ptr(query_desc.sourceText).to_str()? };
     let ast = DFParser::parse_sql_with_dialect(query, &dialect)
         .map_err(|err| ParadeError::DataFusion(DataFusionError::SQL(err, None)))?;
     let statement = &ast[0];