ducklake 0.0.17

mod arrow;
mod dialects;
pub(crate) mod sea_query_ext;
mod types;

use std::sync::OnceLock;

use arrow_array::RecordBatch;
use arrow_schema::Schema;
pub(crate) use dialects::Dialect;
use dialects::SqlConvertible;
use sea_query::Expr;
use sqlx::prelude::*;
pub(crate) use types::chrono::UtcDateTime;
pub(crate) use types::uuid::UuidText;

use crate::{DucklakeError, DucklakeResult};

/* ------------------------------------------ DISPATCH ----------------------------------------- */

/// Dispatches an operation over the concrete database backend behind a [`Pool`].
macro_rules! dispatch_pool {
    ($self:ident, $pool:ident => $call:expr) => {
        match &$self.0 {
            #[cfg(feature = "postgres")]
            AnyPool::Postgres($pool) => $call,
            #[cfg(feature = "mysql")]
            AnyPool::MySql($pool) => $call,
            #[cfg(feature = "sqlite")]
            AnyPool::Sqlite($pool) => $call,
        }
    };
}

/// Dispatches an operation over the concrete database backend behind a [`Transaction`].
macro_rules! dispatch_tx {
    ($self:ident, $tx:ident => $call:expr) => {
        match &mut $self.0 {
            #[cfg(feature = "postgres")]
            AnyTransaction::Postgres($tx) => $call,
            #[cfg(feature = "mysql")]
            AnyTransaction::MySql($tx) => $call,
            #[cfg(feature = "sqlite")]
            AnyTransaction::Sqlite($tx) => $call,
        }
    };
}

/* -------------------------------------------- POOL ------------------------------------------- */

/// Single-connection pool to a dynamic database backend (Postgres, MySQL, SQLite).
#[derive(Clone)]
pub(crate) struct Pool(AnyPool);

#[derive(Clone)]
enum AnyPool {
    #[cfg(feature = "postgres")]
    Postgres(sqlx::Pool<sqlx::Postgres>),
    #[cfg(feature = "mysql")]
    MySql(sqlx::Pool<sqlx::MySql>),
    #[cfg(feature = "sqlite")]
    Sqlite(sqlx::Pool<sqlx::Sqlite>),
}

impl Pool {
    pub(crate) fn dialect(&self) -> Dialect {
        match self.0 {
            #[cfg(feature = "postgres")]
            AnyPool::Postgres(_) => Dialect::Postgres,
            #[cfg(feature = "mysql")]
            AnyPool::MySql(_) => Dialect::MySql,
            #[cfg(feature = "sqlite")]
            AnyPool::Sqlite(_) => Dialect::Sqlite,
        }
    }

    /// Whether two pools connect to the same catalog database.
    pub(crate) fn is_same_catalog(&self, other: &Self) -> bool {
        match (&self.0, &other.0) {
            #[cfg(feature = "postgres")]
            (AnyPool::Postgres(this), AnyPool::Postgres(other)) => {
                let this = this.connect_options();
                let other = other.connect_options();
                ServerCatalogKey::from_postgres(this.as_ref())
                    == ServerCatalogKey::from_postgres(other.as_ref())
            }
            #[cfg(feature = "mysql")]
            (AnyPool::MySql(this), AnyPool::MySql(other)) => {
                let this = this.connect_options();
                let other = other.connect_options();
                ServerCatalogKey::from_mysql(this.as_ref())
                    == ServerCatalogKey::from_mysql(other.as_ref())
            }
            #[cfg(feature = "sqlite")]
            (AnyPool::Sqlite(this), AnyPool::Sqlite(other)) => {
                let this = this.connect_options();
                let other = other.connect_options();
                normalized_path(this.get_filename()) == normalized_path(other.get_filename())
            }
            _ => false,
        }
    }

    pub(crate) async fn new(url: &str) -> DucklakeResult<Self> {
        // NOTE: Choose 8 because this allows the highest concurrency query in this
        //  repo to send all queries simultaneously.
        #[cfg(any(feature = "postgres", feature = "mysql"))]
        const POOL_SIZE: u32 = 8;

        let pool = if url.starts_with("postgresql://") || url.starts_with("postgres://") {
            #[cfg(feature = "postgres")]
            {
                let pool = sqlx::postgres::PgPoolOptions::new()
                    .max_connections(POOL_SIZE)
                    .connect(url)
                    .await?;
                AnyPool::Postgres(pool)
            }
            #[cfg(not(feature = "postgres"))]
            panic!("Postgres support is not enabled. Enable the 'postgres' feature.");
        } else if url.starts_with("mysql://") {
            #[cfg(feature = "mysql")]
            {
                let pool = sqlx::mysql::MySqlPoolOptions::new()
                    .max_connections(POOL_SIZE)
                    .connect(url)
                    .await?;
                AnyPool::MySql(pool)
            }
            #[cfg(not(feature = "mysql"))]
            panic!("MySQL support is not enabled. Enable the 'mysql' feature.");
        } else if url.starts_with("sqlite://") {
            #[cfg(feature = "sqlite")]
            {
                use sqlx::sqlite::SqliteConnectOptions;

                let connect_options = url.parse::<SqliteConnectOptions>()?.create_if_missing(true);
                let pool = sqlx::sqlite::SqlitePoolOptions::new()
                    .max_connections(1)
                    .connect_with(connect_options)
                    .await?;
                AnyPool::Sqlite(pool)
            }
            #[cfg(not(feature = "sqlite"))]
            panic!("SQLite support is not enabled. Enable the 'sqlite' feature.");
        } else {
            return Err(DucklakeError::UnsupportedDatabase(url.to_string()));
        };
        Ok(Pool(pool))
    }

    pub(crate) async fn close(&self) {
        dispatch_pool!(self, pool => {
            pool.close().await
        })
    }

    pub(crate) async fn table_exists(&self, table_name: &str) -> DucklakeResult<bool> {
        let result: (bool,) = match &self.0 {
            #[cfg(feature = "postgres")]
            AnyPool::Postgres(pool) => {
                let sql = "SELECT to_regclass($1) IS NOT NULL";
                log_sql(sql, None);
                sqlx::query_as(sql).bind(table_name).fetch_one(pool).await?
            }
            #[cfg(feature = "mysql")]
            AnyPool::MySql(pool) => {
                let sql = r#"SELECT COUNT(*) > 0
                   FROM information_schema.tables
                   WHERE table_schema = DATABASE() AND table_name = ?"#;
                log_sql(sql, None);
                sqlx::query_as(sql).bind(table_name).fetch_one(pool).await?
            }
            #[cfg(feature = "sqlite")]
            AnyPool::Sqlite(pool) => {
                let sql = r#"SELECT COUNT(*) > 0
                   FROM sqlite_master
                   WHERE type = 'table' AND name = ?"#;
                log_sql(sql, None);
                sqlx::query_as(sql).bind(table_name).fetch_one(pool).await?
            }
        };
        Ok(result.0)
    }

    pub(crate) async fn fetch_one<O>(&self, query: &impl SqlConvertible) -> DucklakeResult<O>
    where
        O: RowType,
    {
        let (sql, values) = query.to_sql(self.dialect());
        log_sql(sql.as_str(), Some(&values));
        let result = dispatch_pool!(self, pool => {
            sqlx::query_as_with(sql, values).fetch_one(pool).await?
        });
        Ok(result)
    }

    pub(crate) async fn fetch_all<O>(&self, query: &impl SqlConvertible) -> DucklakeResult<Vec<O>>
    where
        O: RowType,
    {
        let (sql, values) = query.to_sql(self.dialect());
        log_sql(sql.as_str(), Some(&values));
        let result = dispatch_pool!(self, pool => {
            sqlx::query_as_with(sql, values).fetch_all(pool).await?
        });
        Ok(result)
    }

    pub(crate) async fn fetch_optional<O>(
        &self,
        query: &impl SqlConvertible,
    ) -> DucklakeResult<Option<O>>
    where
        O: RowType,
    {
        let (sql, values) = query.to_sql(self.dialect());
        log_sql(sql.as_str(), Some(&values));
        let result = dispatch_pool!(self, pool => {
            sqlx::query_as_with(sql, values).fetch_optional(pool).await?
        });
        Ok(result)
    }

    pub(crate) async fn fetch_all_arrow(
        &self,
        query: &impl SqlConvertible,
        schema: &Schema,
    ) -> DucklakeResult<RecordBatch> {
        let (sql, values) = query.to_sql(self.dialect());
        log_sql(sql.as_str(), Some(&values));
        match &self.0 {
            #[cfg(feature = "postgres")]
            AnyPool::Postgres(pool) => {
                let rows = sqlx::query_with(sql, values).fetch(pool);
                arrow::decode_rows(rows, schema).await
            }
            #[cfg(feature = "mysql")]
            AnyPool::MySql(_) => unimplemented!("data inlining is not yet implemented for MySQL"),
            #[cfg(feature = "sqlite")]
            AnyPool::Sqlite(pool) => {
                let rows = sqlx::query_with(sql, values).fetch(pool);
                arrow::decode_rows(rows, schema).await
            }
        }
    }

    pub(crate) async fn begin(&self) -> DucklakeResult<Transaction> {
        let tx = match &self.0 {
            #[cfg(feature = "postgres")]
            AnyPool::Postgres(pool) => {
                let sql = "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ";
                log_sql(sql, None);
                AnyTransaction::Postgres(pool.begin_with(sql).await?)
            }
            #[cfg(feature = "mysql")]
            AnyPool::MySql(pool) => {
                let sql = "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ; START TRANSACTION";
                log_sql(sql, None);
                AnyTransaction::MySql(pool.begin_with(sql).await?)
            }
            #[cfg(feature = "sqlite")]
            AnyPool::Sqlite(pool) => {
                let sql = "BEGIN IMMEDIATE";
                log_sql(sql, None);
                AnyTransaction::Sqlite(pool.begin_with(sql).await?)
            }
        };
        Ok(Transaction(tx))
    }
}

/* ---------------------------------------- TRANSACTION ---------------------------------------- */

pub(crate) struct Transaction(AnyTransaction);

enum AnyTransaction {
    #[cfg(feature = "postgres")]
    Postgres(sqlx::Transaction<'static, sqlx::Postgres>),
    #[cfg(feature = "mysql")]
    MySql(sqlx::Transaction<'static, sqlx::MySql>),
    #[cfg(feature = "sqlite")]
    Sqlite(sqlx::Transaction<'static, sqlx::Sqlite>),
}

impl Transaction {
    pub(crate) fn dialect(&self) -> Dialect {
        match self.0 {
            #[cfg(feature = "postgres")]
            AnyTransaction::Postgres(_) => Dialect::Postgres,
            #[cfg(feature = "mysql")]
            AnyTransaction::MySql(_) => Dialect::MySql,
            #[cfg(feature = "sqlite")]
            AnyTransaction::Sqlite(_) => Dialect::Sqlite,
        }
    }

    pub(crate) async fn execute(&mut self, query: &impl SqlConvertible) -> DucklakeResult<()> {
        let (sql, values) = query.to_sql(self.dialect());
        log_sql(sql.as_str(), Some(&values));
        dispatch_tx!(self, tx => {
            sqlx::query_with(sql, values).execute(&mut **tx).await?;
        });
        Ok(())
    }

    /// Insert a single entity into its backing table.
    pub(crate) async fn insert_entity(
        &mut self,
        entity: impl sea_query_ext::InsertableEntity,
    ) -> DucklakeResult<()> {
        let query = entity.insert_into_table();
        self.execute(&query).await
    }

    /// Insert the given entities into their backing table.
    ///
    /// Insertions are automatically batched into multiple statements to prevent exhausting the
    /// underlying database's bind parameter limit.
    pub(crate) async fn insert_entities<E>(
        &mut self,
        entities: impl IntoIterator<Item = E>,
    ) -> DucklakeResult<()>
    where
        E: sea_query_ext::InsertableEntity,
    {
        self.insert_entities_multi_row_values(None, entities).await
    }

    /// Insert the given entities into a dynamically named table.
    pub(crate) async fn insert_entities_into<E>(
        &mut self,
        table: &str,
        entities: impl IntoIterator<Item = E>,
    ) -> DucklakeResult<()>
    where
        E: sea_query_ext::InsertableEntity,
    {
        self.insert_entities_multi_row_values(Some(table), entities)
            .await
    }

    /// Insert the given entities using one or more multi-row `VALUES` statements, batching them
    /// such that the backend's bind parameter limit is respected.
    async fn insert_entities_multi_row_values<E>(
        &mut self,
        table: Option<&str>,
        entities: impl IntoIterator<Item = E>,
    ) -> DucklakeResult<()>
    where
        E: sea_query_ext::InsertableEntity,
    {
        let chunk_size = self.dialect().max_bind_params() / E::NUM_COLUMNS;
        // NOTE: We materialize the entities into an owned `vec::IntoIter` up front. Some call
        //  sites pass borrowing iterators (e.g. `slice.iter().map(...)`); holding such an
        //  iterator across the `await` below would make the resulting future non-`Send`.
        let mut entities = entities.into_iter().collect::<Vec<_>>().into_iter();
        // NOTE: Unfortunately, we cannot use `itertools.chunks` because the resulting future
        //  would not be `Send`.
        loop {
            let chunk: Vec<E> = entities.by_ref().take(chunk_size).collect();
            if chunk.is_empty() {
                break;
            }
            let mut query = E::insert_all_into_table(chunk);
            if let Some(table) = table {
                query.into_table(table.to_string());
            }
            self.execute(&query).await?;
        }
        Ok(())
    }

    pub(crate) async fn insert_all_arrow(
        &mut self,
        table: &str,
        data: RecordBatch,
    ) -> DucklakeResult<()> {
        if data.num_rows() == 0 || data.num_columns() == 0 {
            return Ok(());
        }

        // Build the insertion query
        let mut stmt = sea_query::Query::insert();
        stmt.into_table(table.to_string())
            .columns(data.schema().fields().iter().map(|f| f.name().clone()));
        // NOTE: We use dummy values for the placeholders here and replace them with the Arrow
        //  data below. This way, we are not dependent on data types supported by sea-query.
        //  For example,
        (0..data.num_rows()).for_each(|_| {
            let row = (0..data.num_columns())
                .map(|_| Expr::value(false))
                .collect::<Vec<_>>();
            stmt.values_panic(row);
        });
        let (sql, _) = stmt.to_sql(self.dialect());
        log_sql(sql.as_str(), None);

        // Execute the insertion query with the appropriate arguments built from the Arrow data
        match &mut self.0 {
            #[cfg(feature = "postgres")]
            AnyTransaction::Postgres(tx) => {
                let args: sqlx::postgres::PgArguments = arrow::encode_record_batch(&data)?;
                sqlx::query_with(sql, args).execute(&mut **tx).await?;
            }
            #[cfg(feature = "mysql")]
            AnyTransaction::MySql(_) => {
                unimplemented!("data inlining is not yet implemented for MySQL")
            }
            #[cfg(feature = "sqlite")]
            AnyTransaction::Sqlite(tx) => {
                let args: sqlx::sqlite::SqliteArguments = arrow::encode_record_batch(&data)?;
                sqlx::query_with(sql, args).execute(&mut **tx).await?;
            }
        };
        Ok(())
    }

    pub(crate) async fn fetch_one<O>(&mut self, query: &impl SqlConvertible) -> DucklakeResult<O>
    where
        O: RowType,
    {
        let (sql, values) = query.to_sql(self.dialect());
        log_sql(sql.as_str(), Some(&values));
        let result = dispatch_tx!(self, tx => {
            sqlx::query_as_with(sql, values).fetch_one(&mut **tx).await?
        });
        Ok(result)
    }

    pub(crate) async fn fetch_all<O>(
        &mut self,
        query: &impl SqlConvertible,
    ) -> DucklakeResult<Vec<O>>
    where
        O: RowType,
    {
        let (sql, values) = query.to_sql(self.dialect());
        log_sql(sql.as_str(), Some(&values));
        let result = dispatch_tx!(self, tx => {
            sqlx::query_as_with(sql, values).fetch_all(&mut **tx).await?
        });
        Ok(result)
    }

    pub(crate) async fn commit(self) -> DucklakeResult<()> {
        log_sql("COMMIT", None);
        match self.0 {
            #[cfg(feature = "postgres")]
            AnyTransaction::Postgres(tx) => tx.commit().await?,
            #[cfg(feature = "mysql")]
            AnyTransaction::MySql(tx) => tx.commit().await?,
            #[cfg(feature = "sqlite")]
            AnyTransaction::Sqlite(tx) => tx.commit().await?,
        };
        Ok(())
    }

    pub(crate) async fn rollback(self) -> DucklakeResult<()> {
        log_sql("ROLLBACK", None);
        match self.0 {
            #[cfg(feature = "postgres")]
            AnyTransaction::Postgres(tx) => tx.rollback().await?,
            #[cfg(feature = "mysql")]
            AnyTransaction::MySql(tx) => tx.rollback().await?,
            #[cfg(feature = "sqlite")]
            AnyTransaction::Sqlite(tx) => tx.rollback().await?,
        };
        Ok(())
    }
}

/* --------------------------------------------------------------------------------------------- */
/*                                             UTILS                                             */
/* --------------------------------------------------------------------------------------------- */

/* ------------------------------------------ LOGGING ------------------------------------------ */

#[allow(clippy::print_stdout)]
fn log_sql(sql: &str, values: Option<&sea_query_sqlx::SqlxValues>) {
    static VERBOSE: OnceLock<bool> = OnceLock::new();
    let verbose =
        *VERBOSE.get_or_init(|| std::env::var("DUCKLAKE_SQL_VERBOSE").as_deref() == Ok("1"));
    if verbose {
        match values {
            Some(values) if !values.0.0.is_empty() => {
                println!("[ducklake sql] {sql} -- values: {:?}", values.0.0)
            }
            _ => println!("[ducklake sql] {sql}"),
        }
    }
}

/* ----------------------------------------- CONNECTION ---------------------------------------- */

#[cfg(any(feature = "postgres", feature = "mysql"))]
#[derive(PartialEq, Eq)]
struct ServerCatalogKey<'a> {
    endpoint: ServerEndpoint,
    port: u16,
    database: Option<&'a str>,
}

#[cfg(any(feature = "postgres", feature = "mysql"))]
#[derive(PartialEq, Eq)]
enum ServerEndpoint {
    Host(String),
    Socket(std::path::PathBuf),
}

#[cfg(any(feature = "postgres", feature = "mysql"))]
impl<'a> ServerCatalogKey<'a> {
    fn new(
        host: &str,
        port: u16,
        socket: Option<&std::path::PathBuf>,
        database: Option<&'a str>,
    ) -> Self {
        let socket = socket
            .map(std::path::PathBuf::as_path)
            .or_else(|| host.starts_with('/').then(|| std::path::Path::new(host)));
        let endpoint = match socket {
            Some(socket) => ServerEndpoint::Socket(normalized_path(socket)),
            None => ServerEndpoint::Host(host.to_ascii_lowercase()),
        };
        Self {
            endpoint,
            port,
            database,
        }
    }

    #[cfg(feature = "postgres")]
    fn from_postgres(options: &'a sqlx::postgres::PgConnectOptions) -> Self {
        Self::new(
            options.get_host(),
            options.get_port(),
            options.get_socket(),
            options.get_database().or(Some(options.get_username())),
        )
    }

    #[cfg(feature = "mysql")]
    fn from_mysql(options: &'a sqlx::mysql::MySqlConnectOptions) -> Self {
        Self::new(
            options.get_host(),
            options.get_port(),
            options.get_socket(),
            options.get_database(),
        )
    }
}

fn normalized_path(path: &std::path::Path) -> std::path::PathBuf {
    std::fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())
}

/* ------------------------------------------ ROW TYPE ----------------------------------------- */

#[cfg(not(any(feature = "postgres", feature = "mysql", feature = "sqlite")))]
pub(crate) trait RowType = Send + Unpin;

#[cfg(all(feature = "postgres", not(feature = "mysql"), not(feature = "sqlite")))]
pub(crate) trait RowType =
    Send + Unpin + for<'r> FromRow<'r, <sqlx::Postgres as sqlx::Database>::Row>;

#[cfg(all(not(feature = "postgres"), feature = "mysql", not(feature = "sqlite")))]
pub(crate) trait RowType =
    Send + Unpin + for<'r> FromRow<'r, <sqlx::MySql as sqlx::Database>::Row>;

#[cfg(all(not(feature = "postgres"), not(feature = "mysql"), feature = "sqlite"))]
pub(crate) trait RowType =
    Send + Unpin + for<'r> FromRow<'r, <sqlx::Sqlite as sqlx::Database>::Row>;

#[cfg(all(feature = "postgres", feature = "mysql", not(feature = "sqlite")))]
pub(crate) trait RowType = Send
    + Unpin
    + for<'r> FromRow<'r, <sqlx::Postgres as sqlx::Database>::Row>
    + for<'r> FromRow<'r, <sqlx::MySql as sqlx::Database>::Row>;

#[cfg(all(feature = "postgres", not(feature = "mysql"), feature = "sqlite"))]
pub(crate) trait RowType = Send
    + Unpin
    + for<'r> FromRow<'r, <sqlx::Postgres as sqlx::Database>::Row>
    + for<'r> FromRow<'r, <sqlx::Sqlite as sqlx::Database>::Row>;

#[cfg(all(not(feature = "postgres"), feature = "mysql", feature = "sqlite"))]
pub(crate) trait RowType = Send
    + Unpin
    + for<'r> FromRow<'r, <sqlx::MySql as sqlx::Database>::Row>
    + for<'r> FromRow<'r, <sqlx::Sqlite as sqlx::Database>::Row>;

#[cfg(all(feature = "postgres", feature = "mysql", feature = "sqlite"))]
pub(crate) trait RowType = Send
    + Unpin
    + for<'r> FromRow<'r, <sqlx::Postgres as sqlx::Database>::Row>
    + for<'r> FromRow<'r, <sqlx::MySql as sqlx::Database>::Row>
    + for<'r> FromRow<'r, <sqlx::Sqlite as sqlx::Database>::Row>;