fraiseql_db/postgres/

introspector.rs

/// ! PostgreSQL database introspection for fact tables.
use deadpool_postgres::Pool;
use fraiseql_error::{FraiseQLError, Result};
use tokio_postgres::Row;

use super::pg_detail;
use crate::{
    DatabaseType,
    introspector::{DatabaseIntrospector, RelationInfo, RelationKind},
};

/// PostgreSQL introspector for fact table metadata.
pub struct PostgresIntrospector {
    pool: Pool,
}

/// Split `schema.relation` into its parts; an unqualified name yields `None` for
/// the schema, which the catalog queries resolve against the search path.
fn split_relation(relation: &str) -> (Option<&str>, &str) {
    relation
        .find('.')
        .map_or((None, relation), |dot| (Some(&relation[..dot]), &relation[dot + 1..]))
}

/// One index on a relation, with its key list in **declared order** (#1307).
///
/// Order is the whole point: an index on `(status, pk_invoice)` serves
/// `ORDER BY status, pk_invoice` without a sort, while separate indexes on
/// `status` and on `pk_invoice` leave an incremental sort on every page. The flat
/// column set `get_indexed_columns` returns cannot tell those two apart — it
/// reports `{pk_invoice, status}` for both, which is why #1307 could not be built
/// on it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IndexInfo {
    /// The index's own relation name.
    pub name:   String,
    /// Whether the index is unique.
    pub unique: bool,
    /// Key definitions in declared order. A key is a bare column name (`status`)
    /// or an expression as Postgres renders it (`(data ->> 'id'::text)`), so an
    /// expression index is visible rather than silently dropped.
    pub keys:   Vec<String>,
}

impl PostgresIntrospector {
    /// Create new PostgreSQL introspector from connection pool.
    #[must_use]
    pub const fn new(pool: Pool) -> Self {
        Self { pool }
    }
}

impl DatabaseIntrospector for PostgresIntrospector {
    async fn list_fact_tables(&self) -> Result<Vec<String>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // Query information_schema for tables matching tf_* pattern across all schemas
        // in the current search path.
        let query = r"
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = ANY(current_schemas(false))
              AND table_type = 'BASE TABLE'
              AND table_name LIKE 'tf_%'
            ORDER BY table_name
        ";

        let rows: Vec<Row> =
            client.query(query, &[]).await.map_err(|e| FraiseQLError::Database {
                message:   format!("Failed to list fact tables: {}", pg_detail(&e)),
                sql_state: e.code().map(|c| c.code().to_string()),
            })?;

        let tables = rows
            .into_iter()
            .map(|row| {
                let name: String = row.get(0);
                name
            })
            .collect();

        Ok(tables)
    }

    async fn get_columns(&self, table_name: &str) -> Result<Vec<(String, String, bool)>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // Support schema-qualified names like "benchmark.tv_post".
        // When no schema prefix is present, search across all schemas in the current
        // search path so that non-public schemas work without explicit qualification.
        let rows: Vec<Row> = if let Some(dot) = table_name.find('.') {
            let (schema, name) = (&table_name[..dot], &table_name[dot + 1..]);
            client
                .query(
                    r"SELECT column_name,
                             -- An extension type (ltree, vector) reads as USER-DEFINED
                             -- here; its own name is in udt_name (#1498).
                             CASE WHEN data_type = 'USER-DEFINED' THEN udt_name ELSE data_type END,
                             is_nullable = 'YES' as is_nullable
                      FROM information_schema.columns
                      WHERE table_name = $1 AND table_schema = $2
                      ORDER BY ordinal_position",
                    &[&name, &schema],
                )
                .await
        } else {
            client
                .query(
                    r"SELECT column_name,
                             -- An extension type (ltree, vector) reads as USER-DEFINED
                             -- here; its own name is in udt_name (#1498).
                             CASE WHEN data_type = 'USER-DEFINED' THEN udt_name ELSE data_type END,
                             is_nullable = 'YES' as is_nullable
                      FROM information_schema.columns
                      WHERE table_name = $1
                        AND table_schema = ANY(current_schemas(false))
                      ORDER BY ordinal_position",
                    &[&table_name],
                )
                .await
        }
        .map_err(|e| FraiseQLError::Database {
            message:   format!("Failed to query column information: {}", pg_detail(&e)),
            sql_state: e.code().map(|c| c.code().to_string()),
        })?;

        let columns = rows
            .into_iter()
            .map(|row| {
                let name: String = row.get(0);
                let data_type: String = row.get(1);
                let is_nullable: bool = row.get(2);
                (name, data_type, is_nullable)
            })
            .collect();

        Ok(columns)
    }

    async fn get_indexed_columns(&self, table_name: &str) -> Result<Vec<String>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // Support schema-qualified names like "benchmark.tv_post".
        // When no schema prefix is present, search across all schemas in the current
        // search path so that non-public schemas work without explicit qualification.
        let rows: Vec<Row> = if let Some(dot) = table_name.find('.') {
            let (schema, name) = (&table_name[..dot], &table_name[dot + 1..]);
            client
                .query(
                    r"SELECT DISTINCT a.attname AS column_name
                      FROM pg_index i
                      JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                      JOIN pg_class t ON t.oid = i.indrelid
                      JOIN pg_namespace n ON n.oid = t.relnamespace
                      WHERE t.relname = $1 AND n.nspname = $2 AND a.attnum > 0
                      ORDER BY column_name",
                    &[&name, &schema],
                )
                .await
        } else {
            client
                .query(
                    r"SELECT DISTINCT a.attname AS column_name
                      FROM pg_index i
                      JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                      JOIN pg_class t ON t.oid = i.indrelid
                      JOIN pg_namespace n ON n.oid = t.relnamespace
                      WHERE t.relname = $1
                        AND n.nspname = ANY(current_schemas(false))
                        AND a.attnum > 0
                      ORDER BY column_name",
                    &[&table_name],
                )
                .await
        }
        .map_err(|e| FraiseQLError::Database {
            message:   format!("Failed to query index information: {}", pg_detail(&e)),
            sql_state: e.code().map(|c| c.code().to_string()),
        })?;

        let indexed_columns = rows
            .into_iter()
            .map(|row| {
                let name: String = row.get(0);
                name
            })
            .collect();

        Ok(indexed_columns)
    }

    fn database_type(&self) -> DatabaseType {
        DatabaseType::PostgreSQL
    }

    async fn list_relations(&self) -> Result<Vec<RelationInfo>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // List all tables and views visible in the current search path, excluding
        // system schemas. Uses current_schemas(false) to match the same scope as
        // get_columns / get_indexed_columns so that relation existence checks are
        // consistent with column lookups.
        let rows: Vec<Row> = client
            .query(
                r"SELECT table_schema, table_name,
                         CASE table_type WHEN 'VIEW' THEN 'view' ELSE 'table' END AS kind
                  FROM information_schema.tables
                  WHERE table_schema = ANY(current_schemas(false))
                  ORDER BY table_schema, table_name",
                &[],
            )
            .await
            .map_err(|e| FraiseQLError::Database {
                message:   format!("Failed to list relations: {}", pg_detail(&e)),
                sql_state: e.code().map(|c| c.code().to_string()),
            })?;

        let relations = rows
            .into_iter()
            .map(|row| {
                let schema: String = row.get(0);
                let name: String = row.get(1);
                let kind_str: String = row.get(2);
                let kind = if kind_str == "view" {
                    RelationKind::View
                } else {
                    RelationKind::Table
                };
                RelationInfo { schema, name, kind }
            })
            .collect();

        Ok(relations)
    }

    async fn get_sample_jsonb(
        &self,
        table_name: &str,
        column_name: &str,
    ) -> Result<Option<serde_json::Value>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // Query for a sample row with non-null JSON data
        // Use format! for identifiers (safe because we validate table_name pattern)
        let query = format!(
            r#"
            SELECT "{column}"::text
            FROM "{table}"
            WHERE "{column}" IS NOT NULL
            LIMIT 1
            "#,
            table = table_name,
            column = column_name
        );

        let rows: Vec<Row> =
            client.query(&query, &[]).await.map_err(|e| FraiseQLError::Database {
                message:   format!("Failed to query sample JSONB: {}", pg_detail(&e)),
                sql_state: e.code().map(|c| c.code().to_string()),
            })?;

        if rows.is_empty() {
            return Ok(None);
        }

        let json_text: Option<String> = rows[0].get(0);
        if let Some(text) = json_text {
            let value: serde_json::Value =
                serde_json::from_str(&text).map_err(|e| FraiseQLError::Parse {
                    message:  format!("Failed to parse JSONB sample: {e}"),
                    location: format!("{table_name}.{column_name}"),
                })?;
            Ok(Some(value))
        } else {
            Ok(None)
        }
    }

    async fn get_sample_json_rows(
        &self,
        table_name: &str,
        column_name: &str,
        limit: usize,
    ) -> Result<Vec<serde_json::Value>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // Quote each dot-separated part (with `"` doubled) so a schema-qualified
        // source like `analytics.v_thing` resolves as schema + relation instead
        // of one bogus identifier.
        let table = table_name
            .split('.')
            .map(|part| format!("\"{}\"", part.replace('"', "\"\"")))
            .collect::<Vec<_>>()
            .join(".");
        let column = format!("\"{}\"", column_name.replace('"', "\"\""));
        let query =
            format!("SELECT {column}::text FROM {table} WHERE {column} IS NOT NULL LIMIT {limit}");

        let rows: Vec<Row> =
            client.query(&query, &[]).await.map_err(|e| FraiseQLError::Database {
                message:   format!("Failed to query sample JSON rows: {}", pg_detail(&e)),
                sql_state: e.code().map(|c| c.code().to_string()),
            })?;

        rows.iter()
            .filter_map(|row| row.get::<_, Option<String>>(0))
            .map(|text| {
                serde_json::from_str(&text).map_err(|e| FraiseQLError::Parse {
                    message:  format!("Failed to parse JSON sample: {e}"),
                    location: format!("{table_name}.{column_name}"),
                })
            })
            .collect()
    }

    async fn function_exists(&self, schema: Option<&str>, name: &str) -> Result<Option<bool>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // Mirror `pg_catalog::resolve_functions` name/schema matching (verbatim,
        // schema-qualified or `current_schemas(false)`-scoped). Restrict to kinds a
        // mutation can actually call — regular function `f` or procedure `p` —
        // excluding aggregate `a` / window `w`, which `resolve_functions` does not
        // filter (a minor over-match there).
        let row = if let Some(schema) = schema {
            client
                .query_one(
                    "SELECT EXISTS(SELECT 1 FROM pg_proc p \
                       JOIN pg_namespace n ON n.oid = p.pronamespace \
                       WHERE n.nspname = $1 AND p.proname = $2 AND p.prokind IN ('f','p'))",
                    &[&schema, &name],
                )
                .await
        } else {
            client
                .query_one(
                    "SELECT EXISTS(SELECT 1 FROM pg_proc p \
                       JOIN pg_namespace n ON n.oid = p.pronamespace \
                       WHERE p.proname = $1 AND n.nspname = ANY(current_schemas(false)) \
                       AND p.prokind IN ('f','p'))",
                    &[&name],
                )
                .await
        }
        .map_err(|e| FraiseQLError::Database {
            message:   format!("Failed to probe function existence: {}", pg_detail(&e)),
            sql_state: e.code().map(|c| c.code().to_string()),
        })?;

        Ok(Some(row.get(0)))
    }

    async fn qualified_relation_exists(&self, schema: &str, name: &str) -> Result<Option<bool>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // Resolve exactly as the runtime does (`postgres/adapter/mod.rs` quotes each
        // component case-sensitively), so a mixed-case relation in an off-search_path
        // schema the runtime *can* serve is not falsely reported missing. `to_regclass`
        // takes the quoted identifier as a string and ignores `search_path` for a
        // schema-qualified name, returning NULL when the relation is genuinely absent.
        let quoted = format!(
            "{}.{}",
            crate::identifier::quote_postgres_identifier(schema),
            crate::identifier::quote_postgres_identifier(name),
        );
        let row = client
            .query_one("SELECT to_regclass($1) IS NOT NULL", &[&quoted])
            .await
            .map_err(|e| FraiseQLError::Database {
                message:   format!("Failed to probe relation existence: {}", pg_detail(&e)),
                sql_state: e.code().map(|c| c.code().to_string()),
            })?;

        Ok(Some(row.get(0)))
    }

    async fn get_view_definition(&self, relation: &str) -> Result<Option<String>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        // `to_regclass` resolves a bare name through `search_path` and a qualified
        // one verbatim, exactly as `qualified_relation_exists` does, and yields NULL
        // rather than raising when the relation is absent. `pg_get_viewdef` of a
        // non-view is NULL too, so a table and a missing relation are the same
        // answer here — which is correct, because neither has a body to read.
        let row = client
            .query_one("SELECT pg_get_viewdef(to_regclass($1), true)", &[&relation])
            .await
            .map_err(|e| FraiseQLError::Database {
                message:   format!("Failed to read view definition: {}", pg_detail(&e)),
                sql_state: e.code().map(|c| c.code().to_string()),
            })?;

        Ok(row.get(0))
    }
}
impl PostgresIntrospector {
    /// Every index on `relation`, keys in declared order (#1307).
    ///
    /// Keys come from `pg_get_indexdef(indexrelid, n, true)` rather than from
    /// `pg_attribute`, because an expression key has no `attnum`: the JSONB
    /// identity a page is tie-broken on (`data->>'id'`) is indexed as an
    /// expression, and reading `indkey` alone would report that index as having
    /// fewer keys than it has.
    ///
    /// Returns an empty vector for a view — a view carries no indexes, so a
    /// caller holding a `sql_source` must resolve it with
    /// [`resolve_base_relations`](Self::resolve_base_relations) first.
    ///
    /// # Errors
    ///
    /// Returns [`FraiseQLError::Database`] if the catalog query fails.
    pub async fn get_index_definitions(&self, relation: &str) -> Result<Vec<IndexInfo>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        let (schema, name) = split_relation(relation);
        let query = r"
            SELECT i.relname                       AS index_name,
                   ix.indisunique                  AS is_unique,
                   array_agg(pg_get_indexdef(ix.indexrelid, k.ord::int, true)
                             ORDER BY k.ord)       AS keys
            FROM pg_index ix
            JOIN pg_class i ON i.oid = ix.indexrelid
            JOIN pg_class t ON t.oid = ix.indrelid
            JOIN pg_namespace n ON n.oid = t.relnamespace
            JOIN LATERAL generate_series(1, ix.indnkeyatts) AS k(ord) ON true
            WHERE t.relname = $1
              AND ($2::text IS NULL OR n.nspname = $2)
              AND ($2::text IS NOT NULL OR n.nspname = ANY(current_schemas(false)))
            GROUP BY i.relname, ix.indisunique
            ORDER BY i.relname
        ";

        let rows: Vec<Row> =
            client
                .query(query, &[&name, &schema])
                .await
                .map_err(|e| FraiseQLError::Database {
                    message:   format!("Failed to query index definitions: {}", pg_detail(&e)),
                    sql_state: e.code().map(|c| c.code().to_string()),
                })?;

        Ok(rows
            .into_iter()
            .map(|row| IndexInfo {
                name:   row.get(0),
                unique: row.get(1),
                keys:   row.get::<_, Vec<String>>(2),
            })
            .collect())
    }

    /// The tables and materialized views a view reads from (#1307).
    ///
    /// A query binds to `v_invoice`, and `v_invoice` has no indexes — the indexes
    /// that decide whether a page needs a sort live on `tb_invoice`. Resolved
    /// through `pg_rewrite`/`pg_depend` rather than by assuming a `v_`/`tb_`
    /// naming convention, which is a project convention and not a guarantee about
    /// any particular database.
    ///
    /// Returns every base relation the view depends on. More than one means the
    /// view is a join, and a caller cannot attribute an ordering to a single
    /// table without reading the view body.
    ///
    /// # Errors
    ///
    /// Returns [`FraiseQLError::Database`] if the catalog query fails.
    pub async fn resolve_base_relations(&self, view: &str) -> Result<Vec<String>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        let (schema, name) = split_relation(view);
        let query = r"
            SELECT DISTINCT base.relname
            FROM pg_class v
            JOIN pg_namespace vn ON vn.oid = v.relnamespace
            JOIN pg_rewrite r ON r.ev_class = v.oid
            JOIN pg_depend d ON d.objid = r.oid
            JOIN pg_class base ON base.oid = d.refobjid
            WHERE v.relname = $1
              AND ($2::text IS NULL OR vn.nspname = $2)
              AND ($2::text IS NOT NULL OR vn.nspname = ANY(current_schemas(false)))
              AND base.relkind IN ('r', 'm', 'p')
              AND base.oid <> v.oid
            ORDER BY base.relname
        ";

        let rows: Vec<Row> =
            client
                .query(query, &[&name, &schema])
                .await
                .map_err(|e| FraiseQLError::Database {
                    message:   format!("Failed to resolve view base relations: {}", pg_detail(&e)),
                    sql_state: e.code().map(|c| c.code().to_string()),
                })?;

        Ok(rows.into_iter().map(|row| row.get(0)).collect())
    }

    /// Get all column names for a view/table.
    ///
    /// # Arguments
    ///
    /// * `view_name` - Name of the view or table to introspect
    ///
    /// # Returns
    ///
    /// List of all column names in the view/table.
    ///
    /// # Errors
    ///
    /// Returns `FraiseQLError::ConnectionPool` if a connection cannot be acquired,
    /// or `FraiseQLError::Database` if the query fails.
    pub async fn get_view_columns(&self, view_name: &str) -> Result<Vec<String>> {
        let client = self.pool.get().await.map_err(|e| FraiseQLError::ConnectionPool {
            message: format!("Failed to acquire connection: {e}"),
        })?;

        let query = r"
            SELECT column_name
            FROM information_schema.columns
            WHERE table_name = $1
              AND table_schema = ANY(current_schemas(false))
            ORDER BY ordinal_position
        ";

        let rows: Vec<Row> =
            client.query(query, &[&view_name]).await.map_err(|e| FraiseQLError::Database {
                message:   format!("Failed to query view columns: {}", pg_detail(&e)),
                sql_state: e.code().map(|c| c.code().to_string()),
            })?;

        let columns = rows
            .into_iter()
            .map(|row| {
                let name: String = row.get(0);
                name
            })
            .collect();

        Ok(columns)
    }
}