soma-schema 0.2.0

use async_trait::async_trait;
use sqlx::{postgres::PgConnection, Connection, PgPool};
use std::sync::Arc;
use tokio::sync::Mutex;

use crate::driver::{AppliedMigration, LockGuard, MigrationDriver};
use crate::error::{Error, Result};
use crate::migration::Migration;

/// Validate that an identifier contains only [A-Za-z0-9_].
/// A leading digit is allowed (e.g. `00_schema_migrations`).
fn validate_ident(s: &str) -> Result<()> {
    if s.is_empty() || !s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
        return Err(Error::InvalidIdentifier(s.to_owned()));
    }
    Ok(())
}

/// Split a SQL script into individual statements by semicolons that appear outside
/// single-quoted string literals and `--` line comments.
///
/// Handles:
/// - Single-quoted strings (`'...'`), including escaped quotes (`''`)
/// - `--` line comments (to end of line)
/// - Dollar-quoted strings are not parsed; avoid bare `;` inside `$$...$$` blocks.
///
/// ponytail: handles every pattern present in our DDL/DML migrations; a full SQL parser
/// would be a heavy dependency for this narrow use case. Ceiling: dollar-quoted bodies
/// containing bare semicolons — upgrade path is a real parser (e.g. `pg_query`).
fn split_statements(sql: &str) -> Vec<String> {
    let mut stmts: Vec<String> = Vec::new();
    let mut current = String::new();
    let chars: Vec<char> = sql.chars().collect();
    let mut i = 0;

    while i < chars.len() {
        let c = chars[i];

        if c == '-' && i + 1 < chars.len() && chars[i + 1] == '-' {
            // Line comment: consume until newline.
            while i < chars.len() && chars[i] != '\n' {
                current.push(chars[i]);
                i += 1;
            }
        } else if c == '\'' {
            // Single-quoted string: consume until closing quote, handling '' escapes.
            current.push(c);
            i += 1;
            loop {
                if i >= chars.len() {
                    break;
                }
                let qc = chars[i];
                current.push(qc);
                i += 1;
                if qc == '\'' {
                    // Check for escaped quote ('').
                    if i < chars.len() && chars[i] == '\'' {
                        current.push(chars[i]);
                        i += 1;
                    } else {
                        break;
                    }
                }
            }
        } else if c == ';' {
            let trimmed = current.trim().to_owned();
            if !trimmed.is_empty()
                && !trimmed
                    .lines()
                    .all(|l| l.trim().is_empty() || l.trim().starts_with("--"))
            {
                stmts.push(trimmed);
            }
            current.clear();
            i += 1;
        } else {
            current.push(c);
            i += 1;
        }
    }

    // Trailing content after the last semicolon (should be empty for well-formed SQL).
    let trimmed = current.trim().to_owned();
    if !trimmed.is_empty()
        && !trimmed
            .lines()
            .all(|l| l.trim().is_empty() || l.trim().starts_with("--"))
    {
        stmts.push(trimmed);
    }

    stmts
}

/// Configuration for the Postgres driver.
///
/// # Construction
///
/// `PostgresConfig` is intentionally constructable with struct-literal syntax so callers
/// can write `PostgresConfig { schema: Some("app".into()), ..Default::default() }`.
/// It is therefore NOT marked `#[non_exhaustive]` — doing so would break that ergonomic.
/// New fields added in future versions will always have defaults provided via `Default`.
#[derive(Debug, Clone)]
pub struct PostgresConfig {
    /// Target schema. `None` means use the connection's search_path default.
    pub schema: Option<String>,
    /// Tracking table name. Defaults to `00_schema_migrations`.
    pub table: String,
    /// Advisory lock key. Defaults to 918273645.
    pub advisory_lock_key: i64,
}

impl Default for PostgresConfig {
    fn default() -> Self {
        Self {
            schema: None,
            table: "00_schema_migrations".to_owned(),
            advisory_lock_key: 918273645,
        }
    }
}

/// A held advisory lock. Keeps a dedicated connection alive.
/// Releases the lock on Drop via a background task.
pub struct PgLockGuard {
    /// The connection is wrapped in Arc<Mutex<>> so we can move it into the drop task.
    conn: Arc<Mutex<Option<PgConnection>>>,
    key: i64,
}

impl LockGuard for PgLockGuard {}

impl Drop for PgLockGuard {
    fn drop(&mut self) {
        let conn_arc = Arc::clone(&self.conn);
        let key = self.key;
        // ponytail: async cleanup in Drop is inherently unreliable — if the runtime is
        // draining at shutdown, spawn() silently drops the future. In practice Postgres
        // releases session-level advisory locks when the connection closes (either via
        // conn.close() or TCP teardown), so this is best-effort cleanup: the lock will
        // be released eventually, but an operator-visible log line on failure is better
        // than silence.
        if let Ok(handle) = tokio::runtime::Handle::try_current() {
            handle.spawn(async move {
                let mut guard = conn_arc.lock().await;
                if let Some(mut conn) = guard.take() {
                    if let Err(e) = sqlx::query("SELECT pg_advisory_unlock($1)")
                        .bind(key)
                        .execute(&mut conn)
                        .await
                    {
                        // ponytail: can't propagate errors from Drop; log so operators
                        // can see it. Postgres releases session locks on TCP close anyway.
                        eprintln!("WARN soma-schema: pg_advisory_unlock(key={key}) failed: {e}; lock will be released when the connection closes");
                    }
                    let _ = conn.close().await;
                }
            });
        }
    }
}

/// Postgres implementation of `MigrationDriver`.
///
/// The pool must have at least 2 connections: one is reserved for the advisory lock,
/// and at least one more is needed for migration queries. `PgPoolOptions::new().max_connections(2)`
/// is the recommended minimum for CLI use.
#[derive(Debug)]
pub struct PostgresDriver {
    pool: PgPool,
    config: PostgresConfig,
}

impl PostgresDriver {
    /// Create a new `PostgresDriver`.
    ///
    /// # Errors
    ///
    /// Returns `Error::InvalidIdentifier` if the table or schema name contains
    /// characters outside `[A-Za-z0-9_]`.
    ///
    /// Returns `Error::PoolTooSmall` if the pool has fewer than 2 connections.
    /// One connection is reserved for the advisory lock; all migration queries
    /// need at least one more. With `max_connections == 1` every `pool.begin()`
    /// call would wait forever.
    pub fn new(pool: PgPool, config: PostgresConfig) -> Result<Self> {
        validate_ident(&config.table)?;
        if let Some(ref s) = config.schema {
            validate_ident(s)?;
        }
        // One connection is permanently held for the advisory lock.
        // Every migration query needs a second connection from the pool.
        if pool.options().get_max_connections() < 2 {
            return Err(Error::PoolTooSmall);
        }
        Ok(Self { pool, config })
    }

    /// Return `"schema"."table"` or just `"table"` depending on whether schema is set.
    fn qualified_table(&self) -> String {
        match self.config.schema.as_deref() {
            Some(s) => format!("\"{s}\".\"{}\"", &self.config.table),
            None => format!("\"{}\"", &self.config.table),
        }
    }

    fn set_search_path_sql(&self) -> Option<String> {
        // SET LOCAL is transaction-scoped and cannot escape the transaction boundary,
        // preventing search_path leaking to other pooled connections after commit.
        self.config
            .schema
            .as_deref()
            .map(|s| format!("SET LOCAL search_path TO \"{s}\""))
    }
}

#[async_trait]
impl MigrationDriver for PostgresDriver {
    async fn acquire_lock(&self) -> Result<Box<dyn LockGuard>> {
        let mut conn = self.pool.acquire().await?.detach();
        sqlx::query("SELECT pg_advisory_lock($1)")
            .bind(self.config.advisory_lock_key)
            .execute(&mut conn)
            .await?;
        let guard = PgLockGuard {
            conn: Arc::new(Mutex::new(Some(conn))),
            key: self.config.advisory_lock_key,
        };
        Ok(Box::new(guard))
    }

    async fn run_setup_sql(&self, name: &str, sql: &str) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        if let Some(sp) = self.set_search_path_sql() {
            sqlx::query(&sp)
                .execute(&mut *tx)
                .await
                .map_err(|e| Error::SetupFailed {
                    file: name.to_owned(),
                    source: e,
                })?;
        }
        for stmt in split_statements(sql) {
            sqlx::query(&stmt)
                .execute(&mut *tx)
                .await
                .map_err(|e| Error::SetupFailed {
                    file: name.to_owned(),
                    source: e,
                })?;
        }
        tx.commit().await.map_err(|e| Error::SetupFailed {
            file: name.to_owned(),
            source: e,
        })?;
        Ok(())
    }

    async fn ensure_tracking_table(&self) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        // Create schema if configured.
        if let Some(schema) = &self.config.schema {
            let create_schema = format!("CREATE SCHEMA IF NOT EXISTS \"{schema}\"");
            sqlx::query(&create_schema).execute(&mut *tx).await?;
        }
        // Set search_path so the table creation lands in the right schema.
        if let Some(sp) = self.set_search_path_sql() {
            sqlx::query(&sp).execute(&mut *tx).await?;
        }
        let qt = self.qualified_table();
        let create_table = format!(
            r#"CREATE TABLE IF NOT EXISTS {qt} (
    version        INTEGER      NOT NULL,
    file           VARCHAR(255) NOT NULL,
    name           VARCHAR(255) NOT NULL,
    checksum       TEXT         NOT NULL,
    description    TEXT,
    batch          INTEGER      NOT NULL,
    applied_at     TIMESTAMPTZ  NOT NULL DEFAULT now(),
    applied_by     TEXT         NOT NULL DEFAULT current_user,
    execution_ms   INTEGER,
    PRIMARY KEY (version, file)
)"#
        );
        sqlx::query(&create_table).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }

    async fn applied(&self) -> Result<Vec<AppliedMigration>> {
        use chrono::{DateTime, Utc};
        let qt = self.qualified_table();
        let sql = format!(
            "SELECT version, file, name, checksum, description, batch, applied_at, applied_by, execution_ms \
             FROM {qt} ORDER BY version ASC, file ASC"
        );
        let rows = sqlx::query_as::<
            _,
            (
                i32,
                String,
                String,
                String,
                Option<String>,
                i32,
                DateTime<Utc>,
                String,
                Option<i32>,
            ),
        >(&sql)
        .fetch_all(&self.pool)
        .await?;
        Ok(rows
            .into_iter()
            .map(
                |(
                    version,
                    file,
                    name,
                    checksum,
                    description,
                    batch,
                    applied_at,
                    applied_by,
                    execution_ms,
                )| {
                    AppliedMigration {
                        version: version as u32,
                        file,
                        name,
                        checksum,
                        description,
                        batch,
                        applied_at,
                        applied_by,
                        execution_ms,
                    }
                },
            )
            .collect())
    }

    async fn apply(&self, migration: &Migration, up_sql: &str, batch: i32) -> Result<()> {
        let start = std::time::Instant::now();
        let mut tx = self.pool.begin().await?;
        // Set search_path so DDL lands in the right schema.
        if let Some(sp) = self.set_search_path_sql() {
            sqlx::query(&sp).execute(&mut *tx).await?;
        }
        for stmt in split_statements(up_sql) {
            sqlx::query(&stmt).execute(&mut *tx).await?;
        }
        let elapsed_ms = start.elapsed().as_millis() as i32;
        let qt = self.qualified_table();
        let insert = format!(
            "INSERT INTO {qt} (version, file, name, checksum, description, batch, applied_at, applied_by, execution_ms) \
             VALUES ($1, $2, $3, $4, $5, $6, now(), current_user, $7)"
        );
        sqlx::query(&insert)
            .bind(migration.version as i32)
            .bind(&migration.file)
            .bind(&migration.name)
            .bind(&migration.checksum)
            .bind(migration.why.as_deref())
            .bind(batch)
            .bind(elapsed_ms)
            .execute(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(())
    }

    async fn revert(&self, applied: &AppliedMigration, down_sql: &str) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        if let Some(sp) = self.set_search_path_sql() {
            sqlx::query(&sp).execute(&mut *tx).await?;
        }
        for stmt in split_statements(down_sql) {
            sqlx::query(&stmt).execute(&mut *tx).await?;
        }
        let qt = self.qualified_table();
        let delete = format!("DELETE FROM {qt} WHERE version = $1 AND file = $2");
        sqlx::query(&delete)
            .bind(applied.version as i32)
            .bind(&applied.file)
            .execute(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn valid_ident() {
        assert!(validate_ident("00_schema_migrations").is_ok());
        assert!(validate_ident("my_schema").is_ok());
        assert!(validate_ident("ABC123").is_ok());
    }

    #[test]
    fn invalid_ident_rejects_special_chars() {
        assert!(matches!(
            validate_ident("bad-name"),
            Err(Error::InvalidIdentifier(_))
        ));
        assert!(matches!(
            validate_ident("has space"),
            Err(Error::InvalidIdentifier(_))
        ));
        assert!(matches!(
            validate_ident(""),
            Err(Error::InvalidIdentifier(_))
        ));
        assert!(matches!(
            validate_ident("semi;colon"),
            Err(Error::InvalidIdentifier(_))
        ));
    }

    #[test]
    fn leading_digit_allowed() {
        // e.g. "00_schema_migrations" — leading digit is valid
        assert!(validate_ident("00_schema_migrations").is_ok());
        assert!(validate_ident("01_vault").is_ok());
    }

    #[test]
    fn qualified_table_format() {
        // Verify schema.table ordering without needing a live pool.
        let with_schema = {
            let s = "myschema";
            let t = "00_schema_migrations";
            format!("\"{s}\".\"{t}\"")
        };
        assert_eq!(with_schema, "\"myschema\".\"00_schema_migrations\"");

        let without_schema = {
            let t = "00_schema_migrations";
            format!("\"{t}\"")
        };
        assert_eq!(without_schema, "\"00_schema_migrations\"");
    }

    // --- split_statements tests ---

    #[test]
    fn split_single_statement() {
        let stmts = split_statements("CREATE TABLE t (id INT)");
        assert_eq!(stmts, vec!["CREATE TABLE t (id INT)"]);
    }

    #[test]
    fn split_multiple_statements() {
        let sql = "CREATE TABLE a (id INT);\nCREATE TABLE b (id INT);";
        let stmts = split_statements(sql);
        assert_eq!(stmts.len(), 2);
        assert_eq!(stmts[0], "CREATE TABLE a (id INT)");
        assert_eq!(stmts[1], "CREATE TABLE b (id INT)");
    }

    #[test]
    fn split_single_quoted_string_with_semicolon_stays_one_statement() {
        // The semicolon inside 'hello;world' must NOT split the statement.
        let sql = "INSERT INTO t (name) VALUES ('hello;world');";
        let stmts = split_statements(sql);
        assert_eq!(stmts.len(), 1);
        assert!(stmts[0].contains("'hello;world'"));
    }

    #[test]
    fn split_single_quoted_escaped_quote() {
        // '' inside a string is an escaped single quote, not a string close.
        let sql = "INSERT INTO t (name) VALUES ('it''s fine');";
        let stmts = split_statements(sql);
        assert_eq!(stmts.len(), 1);
        assert!(stmts[0].contains("'it''s fine'"));
    }

    #[test]
    fn split_line_comment_semicolon_does_not_split() {
        // A semicolon in a -- comment must NOT be treated as a statement terminator.
        let sql = "CREATE TABLE t (id INT); -- trailing; comment\nCREATE TABLE u (id INT);";
        let stmts = split_statements(sql);
        // The first real statement ends at the `;` before `--`, the second after the newline.
        assert_eq!(stmts.len(), 2, "got: {stmts:?}");
    }

    #[test]
    fn split_trailing_statement_without_semicolon() {
        // A final statement without a trailing `;` must still be captured.
        let sql = "CREATE TABLE a (id INT);\nCREATE TABLE b (id INT)";
        let stmts = split_statements(sql);
        assert_eq!(stmts.len(), 2);
        assert_eq!(stmts[1], "CREATE TABLE b (id INT)");
    }

    #[test]
    fn split_empty_input_returns_empty() {
        assert!(split_statements("").is_empty());
    }

    #[test]
    fn split_only_comments_returns_empty() {
        let sql = "-- just a comment\n-- another comment";
        assert!(
            split_statements(sql).is_empty(),
            "comment-only input should produce no statements"
        );
    }

    #[test]
    fn split_dollar_quote_limitation_documented() {
        // ponytail: dollar-quoted strings ($$...$$) are NOT parsed; a semicolon inside
        // $$ ... $$ will incorrectly split the statement. This is the known ceiling —
        // upgrade path is a real SQL parser (e.g. pg_query crate). The test documents
        // the ACTUAL current behavior rather than the desired behavior, so CI catches
        // any accidental change to the status quo.
        let sql = "CREATE FUNCTION f() RETURNS void LANGUAGE plpgsql AS $$ BEGIN NULL; END $$;";
        let stmts = split_statements(sql);
        // With the current implementation, the bare `;` inside $$ splits the statement.
        // This assertion intentionally documents the limitation — it will need updating
        // if/when dollar-quote support is added.
        assert!(
            stmts.len() > 1,
            "expected dollar-quote limitation to cause a split (got {} stmt(s)): {:?}",
            stmts.len(),
            stmts
        );
    }
}