systemprompt_database/admin/

introspection.rs

//! Schema introspection service.
//!
//! Part of the documented sqlx allowlist — every query here is built
//! dynamically because the table name is supplied at runtime as a
//! [`SafeIdentifier`].
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use std::sync::Arc;

use sqlx::Row;
use sqlx::postgres::PgPool;

use crate::admin::identifier::SafeIdentifier;
use crate::error::{DatabaseResult, RepositoryError, is_undefined_table};
use crate::models::{ColumnInfo, DatabaseInfo, IndexInfo, TableInfo};

#[derive(Debug)]
pub struct DatabaseAdminService {
    pool: Arc<PgPool>,
}

impl DatabaseAdminService {
    pub const fn new(pool: Arc<PgPool>) -> Self {
        Self { pool }
    }

    pub async fn list_tables(&self) -> DatabaseResult<Vec<TableInfo>> {
        let rows = sqlx::query(
            r"
            SELECT
                t.table_name as name,
                COALESCE(s.n_live_tup, 0) as row_count,
                COALESCE(pg_total_relation_size(to_regclass(quote_ident(t.table_name))), 0) as size_bytes
            FROM information_schema.tables t
            LEFT JOIN pg_stat_user_tables s ON t.table_name = s.relname
            WHERE t.table_schema = 'public'
            ORDER BY t.table_name
            ",
        )
        .fetch_all(&*self.pool)
        .await?;

        let tables = rows
            .iter()
            .map(|row| {
                let name: String = row.get("name");
                let row_count: i64 = row.get("row_count");
                let size_bytes: i64 = row.get("size_bytes");
                TableInfo {
                    name,
                    row_count,
                    size_bytes,
                    columns: vec![],
                }
            })
            .collect();

        Ok(tables)
    }

    pub async fn describe_table(
        &self,
        table_name: &SafeIdentifier,
    ) -> DatabaseResult<(Vec<ColumnInfo>, i64)> {
        let rows = sqlx::query(
            "SELECT column_name, data_type, is_nullable, column_default FROM \
             information_schema.columns WHERE table_schema = 'public' AND table_name = $1 \
             ORDER BY ordinal_position",
        )
        .bind(table_name.as_str())
        .fetch_all(&*self.pool)
        .await?;

        if rows.is_empty() {
            return Err(RepositoryError::not_found("table", table_name));
        }

        let pk_rows = sqlx::query(
            r"
            SELECT a.attname as column_name
            FROM pg_index i
            JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
            WHERE i.indrelid = $1::regclass AND i.indisprimary
            ",
        )
        .bind(table_name.as_str())
        .fetch_all(&*self.pool)
        .await?;

        let pk_columns: Vec<String> = pk_rows
            .iter()
            .map(|row| row.get::<String, _>("column_name"))
            .collect();

        let columns = rows
            .iter()
            .map(|row| {
                let name: String = row.get("column_name");
                let data_type: String = row.get("data_type");
                let nullable_str: String = row.get("is_nullable");
                let nullable = nullable_str.to_uppercase() == "YES";
                let default: Option<String> = row.get("column_default");
                let primary_key = pk_columns.contains(&name);

                ColumnInfo {
                    name,
                    data_type,
                    nullable,
                    primary_key,
                    default,
                }
            })
            .collect();

        let row_count = self.count_rows(table_name).await?;

        Ok((columns, row_count))
    }

    pub async fn list_table_indexes(
        &self,
        table_name: &SafeIdentifier,
    ) -> DatabaseResult<Vec<IndexInfo>> {
        let rows = sqlx::query(
            r"
            SELECT
                i.relname as index_name,
                ix.indisunique as is_unique,
                array_agg(a.attname ORDER BY array_position(ix.indkey, a.attnum)) as columns
            FROM pg_class t
            JOIN pg_index ix ON t.oid = ix.indrelid
            JOIN pg_class i ON i.oid = ix.indexrelid
            JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
            WHERE t.relname = $1 AND t.relkind = 'r'
            GROUP BY i.relname, ix.indisunique
            ORDER BY i.relname
            ",
        )
        .bind(table_name.as_str())
        .fetch_all(&*self.pool)
        .await?;

        let indexes = rows
            .iter()
            .map(|row| {
                let name: String = row.get("index_name");
                let unique: bool = row.get("is_unique");
                let columns: Vec<String> = row.get("columns");
                IndexInfo {
                    name,
                    columns,
                    unique,
                }
            })
            .collect();

        Ok(indexes)
    }

    pub async fn list_tables_counted(&self) -> DatabaseResult<Vec<TableInfo>> {
        let mut counted = Vec::new();
        for mut table in self.list_tables().await? {
            let ident = SafeIdentifier::parse(&table.name)
                .map_err(|e| RepositoryError::decode(format!("table name {}", table.name), e))?;
            let count_query = format!("SELECT COUNT(*) as count FROM {}", ident.quoted());
            match sqlx::query_scalar::<_, i64>(sqlx::AssertSqlSafe(count_query))
                .fetch_one(&*self.pool)
                .await
            {
                Ok(row_count) => table.row_count = row_count,
                Err(e) if is_undefined_table(&e) => continue,
                Err(e) => return Err(e.into()),
            }
            counted.push(table);
        }
        Ok(counted)
    }

    pub async fn count_rows(&self, table_name: &SafeIdentifier) -> DatabaseResult<i64> {
        let quoted_table = table_name.quoted();
        let count_query = format!("SELECT COUNT(*) as count FROM {quoted_table}");
        let row_count: i64 = sqlx::query_scalar(sqlx::AssertSqlSafe(count_query))
            .fetch_one(&*self.pool)
            .await?;

        Ok(row_count)
    }

    pub async fn get_database_info(&self) -> DatabaseResult<DatabaseInfo> {
        let version: String = sqlx::query_scalar("SELECT version()")
            .fetch_one(&*self.pool)
            .await?;

        let size: i64 = sqlx::query_scalar("SELECT pg_database_size(current_database())")
            .fetch_one(&*self.pool)
            .await?;

        let size = u64::try_from(size)
            .map_err(|e| RepositoryError::decode("pg_database_size(current_database())", e))?;

        let tables = self.list_tables().await?;

        Ok(DatabaseInfo {
            path: "PostgreSQL".to_owned(),
            size,
            version,
            tables,
        })
    }
}