ahri_tre_app/

sql_source.rs

use std::path::PathBuf;

use ahri_tre_pgmeta::{PgMetadataAdapter, PgMetadataConnection};
use ahri_tre_types::{TreValueType, ValueTypeId};
use duckdb::Connection;

use crate::{
    AnalyzeSqlSourceRequest, AppError, OpenDuckDbSourceRequest, OpenMsSqlSourceRequest,
    OpenPostgresSourceRequest, OpenSqliteSourceRequest, SqlSourceAnalysis, SqlSourceColumnAnalysis,
    SqlSourceVocabularyAnalysis, SqlSourceVocabularyItemAnalysis,
};

const SOURCE_ALIAS: &str = "__tre_sql_source";

pub enum SqlSourceConnection {
    DuckDb {
        path: PathBuf,
        connection: Connection,
        description: String,
    },
    Sqlite {
        path: PathBuf,
        description: String,
    },
    Postgres {
        connection_info: String,
        description: String,
        schema: Option<String>,
        connection: Box<PgMetadataConnection>,
    },
    MsSql {
        connection_string: String,
        description: String,
        schema: Option<String>,
    },
}

impl std::fmt::Debug for SqlSourceConnection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SqlSourceConnection")
            .field("kind", &self.kind())
            .field("description", &self.description())
            .finish_non_exhaustive()
    }
}

impl SqlSourceConnection {
    pub fn open_duckdb(request: OpenDuckDbSourceRequest) -> Result<Self, AppError> {
        let connection = Connection::open(&request.path)
            .map_err(|error| infrastructure_error("open DuckDB SQL source", error))?;
        Ok(Self::DuckDb {
            description: format!("duckdb:{}", request.path.display()),
            path: request.path,
            connection,
        })
    }

    pub fn open_sqlite(request: OpenSqliteSourceRequest) -> Result<Self, AppError> {
        if !request.path.is_file() {
            return Err(AppError::Validation(format!(
                "SQLite SQL source not found: {}",
                request.path.display()
            )));
        }
        Ok(Self::Sqlite {
            description: format!("sqlite:{}", request.path.display()),
            path: request.path,
        })
    }

    pub fn open_postgres(request: OpenPostgresSourceRequest) -> Result<Self, AppError> {
        let adapter = PgMetadataAdapter::direct(request.config.clone());
        let connection = adapter
            .connect()
            .map_err(|error| infrastructure_error("open PostgreSQL SQL source", error))?;
        Ok(Self::Postgres {
            connection_info: request.config.conninfo(),
            description: request.config.safe_conninfo(),
            schema: request.schema,
            connection: Box::new(connection),
        })
    }

    pub fn open_mssql(request: OpenMsSqlSourceRequest) -> Result<Self, AppError> {
        if request.connection_string.trim().is_empty() {
            return Err(AppError::Validation(
                "MSSQL source connection string must not be empty".to_string(),
            ));
        }
        Ok(Self::MsSql {
            description: request
                .safe_description
                .unwrap_or_else(|| "mssql:<redacted>".to_string()),
            connection_string: request.connection_string,
            schema: request.schema,
        })
    }

    pub fn kind(&self) -> &'static str {
        match self {
            Self::DuckDb { .. } => "duckdb",
            Self::Sqlite { .. } => "sqlite",
            Self::Postgres { .. } => "postgresql",
            Self::MsSql { .. } => "mssql",
        }
    }

    pub fn description(&self) -> &str {
        match self {
            Self::DuckDb { description, .. }
            | Self::Sqlite { description, .. }
            | Self::Postgres { description, .. }
            | Self::MsSql { description, .. } => description,
        }
    }

    pub fn analyze(
        &mut self,
        request: AnalyzeSqlSourceRequest,
    ) -> Result<SqlSourceAnalysis, AppError> {
        let sql = read_only_select(&request.sql)?;
        let shape = QueryShape::parse(&sql);
        let analysis_connection = Connection::open_in_memory()
            .map_err(|error| infrastructure_error("open SQL analysis DuckDB", error))?;
        let prepared = self.prepare_query_connection(&analysis_connection, &shape)?;
        let columns = describe_query(&analysis_connection, &sql)?;
        let mut warnings = shape.warnings.clone();
        warnings.extend(prepared.warnings);
        let mut analyzed = Vec::with_capacity(columns.len());
        let single_table = shape.tables.len() == 1;

        for described in columns {
            let lineage = shape.lineage_for(&described.name, single_table);
            let mut value_type_id = value_type_id_for_sql(&described.sql_type);
            let mut value_type = value_type_name(value_type_id)?.to_string();
            let mut vocabulary = None;
            if let Some(lineage) = &lineage {
                match infer_vocabulary(
                    &analysis_connection,
                    self,
                    &lineage.table_name,
                    &described.name,
                ) {
                    Ok(Some(inferred)) => {
                        value_type_id = TreValueType::Enumeration.value_type_id();
                        value_type = TreValueType::Enumeration.datastore_name().to_string();
                        vocabulary = Some(inferred);
                    }
                    Ok(None) => {}
                    Err(message) => warnings.push(message),
                }
            } else {
                warnings.push(format!(
                    "metadata lineage for column {} is uncertain; registered variable will use name and value type only",
                    described.name
                ));
            }

            let description = lineage.as_ref().and_then(|lineage| {
                self.column_description(&lineage.table_name, &lineage.column_name)
                    .ok()
                    .flatten()
            });
            analyzed.push(SqlSourceColumnAnalysis {
                name: described.name,
                sql_type: described.sql_type,
                value_type_id,
                value_type,
                description,
                source_table: lineage.as_ref().map(|lineage| lineage.table_name.clone()),
                source_column: lineage.as_ref().map(|lineage| lineage.column_name.clone()),
                vocabulary,
            });
        }

        Ok(SqlSourceAnalysis {
            source_kind: self.kind().to_string(),
            source_description: self.description().to_string(),
            columns: analyzed,
            warnings,
        })
    }

    pub fn prepare_lake_query(
        &self,
        connection: &Connection,
        sql: &str,
    ) -> Result<PreparedSqlSourceQuery, AppError> {
        let sql = read_only_select(sql)?;
        let shape = QueryShape::parse(&sql);
        let prepared = self.prepare_query_connection(connection, &shape)?;
        let mut warnings = shape.warnings;
        warnings.extend(prepared.warnings);
        Ok(PreparedSqlSourceQuery { sql, warnings })
    }

    fn prepare_query_connection(
        &self,
        connection: &Connection,
        shape: &QueryShape,
    ) -> Result<PreparedSqlSourceQuery, AppError> {
        self.attach(connection)?;
        let mut warnings = Vec::new();
        for table in &shape.tables {
            if table.name.contains('.') {
                warnings.push(format!(
                    "query table {} is schema-qualified; automatic temporary view creation is skipped",
                    table.name
                ));
                continue;
            }
            let view_sql = quote_identifier(&table.name);
            let source_sql = source_relation_sql(SOURCE_ALIAS, &table.name);
            connection
                .execute_batch(&format!(
                    "CREATE OR REPLACE TEMP VIEW {view_sql} AS SELECT * FROM {source_sql};"
                ))
                .map_err(|error| {
                    infrastructure_error("prepare SQL source temporary view", error)
                })?;
        }
        Ok(PreparedSqlSourceQuery {
            sql: String::new(),
            warnings,
        })
    }

    fn attach(&self, connection: &Connection) -> Result<(), AppError> {
        match self {
            Self::DuckDb { path, .. } => connection
                .execute_batch(&format!(
                    "ATTACH OR REPLACE {} AS {} (READ_ONLY);",
                    sql_literal(&path.display().to_string()),
                    quote_identifier(SOURCE_ALIAS)
                ))
                .map_err(|error| infrastructure_error("attach DuckDB SQL source", error))?,
            Self::Sqlite { path, .. } => connection
                .execute_batch(&format!(
                    "INSTALL sqlite; LOAD sqlite; ATTACH OR REPLACE {} AS {} (TYPE sqlite, READ_ONLY);",
                    sql_literal(&path.display().to_string()),
                    quote_identifier(SOURCE_ALIAS)
                ))
                .map_err(|error| infrastructure_error("attach SQLite SQL source", error))?,
            Self::Postgres {
                connection_info,
                schema,
                ..
            } => {
                let schema_sql = schema
                    .as_deref()
                    .filter(|value| !value.trim().is_empty())
                    .map(|value| format!(", SCHEMA {}", sql_literal(value)))
                    .unwrap_or_default();
                connection
                    .execute_batch(&format!(
                        "INSTALL postgres; LOAD postgres; ATTACH OR REPLACE {} AS {} (TYPE postgres, READ_ONLY{});",
                        sql_literal(connection_info),
                        quote_identifier(SOURCE_ALIAS),
                        schema_sql
                    ))
                    .map_err(|error| infrastructure_error("attach PostgreSQL SQL source", error))?;
            }
            Self::MsSql {
                connection_string,
                schema,
                ..
            } => {
                let schema_sql = schema
                    .as_deref()
                    .filter(|value| !value.trim().is_empty())
                    .map(|value| format!(", SCHEMA {}", sql_literal(value)))
                    .unwrap_or_default();
                connection
                    .execute_batch(&format!(
                        "INSTALL sqlserver FROM community; LOAD sqlserver; ATTACH OR REPLACE {} AS {} (TYPE sqlserver, READ_ONLY{});",
                        sql_literal(connection_string),
                        quote_identifier(SOURCE_ALIAS),
                        schema_sql
                    ))
                    .map_err(|error| infrastructure_error("attach MSSQL SQL source", error))?;
            }
        }
        Ok(())
    }

    fn column_description(
        &mut self,
        table_name: &str,
        column_name: &str,
    ) -> Result<Option<String>, AppError> {
        match self {
            Self::Postgres {
                connection, schema, ..
            } => {
                let schema = schema.as_deref().unwrap_or("public");
                let row = connection
                    .client()
                    .query_opt(
                        "
                        SELECT pg_catalog.col_description(
                                   format('%I.%I', table_schema, table_name)::regclass::oid,
                                   ordinal_position
                               ) AS description
                          FROM information_schema.columns
                         WHERE table_schema = $1
                           AND table_name = $2
                           AND column_name = $3
                        ",
                        &[&schema, &table_name, &column_name],
                    )
                    .map_err(|error| {
                        infrastructure_error("read PostgreSQL column description", error)
                    })?;
                Ok(row.and_then(|row| row.get::<_, Option<String>>("description")))
            }
            _ => Ok(None),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PreparedSqlSourceQuery {
    pub sql: String,
    pub warnings: Vec<String>,
}

#[derive(Debug, Clone)]
struct DescribedColumn {
    name: String,
    sql_type: String,
}

#[derive(Debug, Clone)]
struct QueryShape {
    tables: Vec<QueryTable>,
    columns: Vec<QueryColumn>,
    warnings: Vec<String>,
}

#[derive(Debug, Clone)]
struct QueryTable {
    name: String,
    alias: Option<String>,
}

#[derive(Debug, Clone)]
struct QueryColumn {
    output_name: String,
    source_name: Option<String>,
    qualifier: Option<String>,
    wildcard: bool,
}

#[derive(Debug, Clone)]
struct ColumnLineage {
    table_name: String,
    column_name: String,
}

impl QueryShape {
    fn parse(sql: &str) -> Self {
        let mut warnings = Vec::new();
        let tokens = tokenize(sql);
        let from_index = tokens
            .iter()
            .position(|token| token.eq_ignore_ascii_case("from"));
        let Some(from_index) = from_index else {
            return Self {
                tables: Vec::new(),
                columns: Vec::new(),
                warnings: vec![
                    "query metadata discovery could not find a top-level FROM clause".to_string(),
                ],
            };
        };
        let select_text = sql
            .get(6..sql.to_ascii_lowercase().find(" from ").unwrap_or(sql.len()))
            .unwrap_or("")
            .trim();
        let columns = parse_select_columns(select_text, &mut warnings);
        let mut tables = Vec::new();
        let mut index = from_index;
        while index < tokens.len() {
            if tokens[index].eq_ignore_ascii_case("from")
                || tokens[index].eq_ignore_ascii_case("join")
            {
                index += 1;
                if index >= tokens.len() {
                    break;
                }
                let name = tokens[index].clone();
                if name == "(" {
                    warnings.push(
                        "query contains a nested source; rich metadata discovery is limited"
                            .to_string(),
                    );
                    index += 1;
                    continue;
                }
                let alias = tokens.get(index + 1).and_then(|candidate| {
                    if is_table_alias_stop(candidate) {
                        None
                    } else if candidate.eq_ignore_ascii_case("as") {
                        tokens.get(index + 2).cloned()
                    } else {
                        Some(candidate.clone())
                    }
                });
                tables.push(QueryTable { name, alias });
            }
            index += 1;
        }
        Self {
            tables,
            columns,
            warnings,
        }
    }

    fn lineage_for(&self, column_name: &str, single_table: bool) -> Option<ColumnLineage> {
        if single_table {
            return self.tables.first().map(|table| ColumnLineage {
                table_name: table.name.clone(),
                column_name: column_name.to_string(),
            });
        }
        self.columns.iter().find_map(|column| {
            if column.wildcard || column.output_name != column_name {
                return None;
            }
            let qualifier = column.qualifier.as_deref()?;
            let table = self.tables.iter().find(|table| {
                table.alias.as_deref() == Some(qualifier) || table.name == qualifier
            })?;
            Some(ColumnLineage {
                table_name: table.name.clone(),
                column_name: column
                    .source_name
                    .clone()
                    .unwrap_or_else(|| column_name.to_string()),
            })
        })
    }
}

fn read_only_select(sql: &str) -> Result<String, AppError> {
    let trimmed = sql.trim().trim_end_matches(';').trim();
    let lowered = trimmed.to_ascii_lowercase();
    let forbidden = [
        ";",
        "insert ",
        "update ",
        "delete ",
        "drop ",
        "create ",
        "alter ",
        "merge ",
        "truncate ",
        "grant ",
        "revoke ",
        "exec ",
        "execute ",
        "call ",
        "copy ",
        "attach ",
        "detach ",
    ];
    if trimmed.is_empty()
        || !lowered.starts_with("select ")
        || forbidden.iter().any(|needle| lowered.contains(needle))
    {
        return Err(AppError::Validation(
            "sql-to-dataset requires a single read-only SELECT statement".to_string(),
        ));
    }
    Ok(trimmed.to_string())
}

fn describe_query(connection: &Connection, sql: &str) -> Result<Vec<DescribedColumn>, AppError> {
    let mut statement = connection
        .prepare(&format!("DESCRIBE SELECT * FROM ({sql}) src"))
        .map_err(|error| infrastructure_error("describe SQL source query", error))?;
    let rows = statement
        .query_map([], |row| {
            Ok(DescribedColumn {
                name: row.get::<_, String>(0)?,
                sql_type: row.get::<_, String>(1)?,
            })
        })
        .map_err(|error| infrastructure_error("read SQL source query description", error))?;
    let mut columns = Vec::new();
    for row in rows {
        columns.push(
            row.map_err(|error| {
                infrastructure_error("decode SQL source query description", error)
            })?,
        );
    }
    Ok(columns)
}

fn infer_vocabulary(
    connection: &Connection,
    source: &SqlSourceConnection,
    _source_table: &str,
    column_name: &str,
) -> Result<Option<SqlSourceVocabularyAnalysis>, String> {
    let candidates = code_table_candidates(column_name);
    for candidate in candidates {
        let relation = source_relation_sql(SOURCE_ALIAS, &candidate);
        let count =
            match connection.query_row(&format!("SELECT count(*) FROM {relation}"), [], |row| {
                row.get::<_, i64>(0)
            }) {
                Ok(count) => count,
                Err(_) => continue,
            };
        if !(0..500).contains(&count) {
            continue;
        }
        let columns = describe_relation_columns(connection, &relation)
            .map_err(|error| format!("could not inspect code table {candidate}: {error}"))?;
        let Some(value_column) = columns.first().map(|column| column.name.clone()) else {
            continue;
        };
        let code_column = columns
            .iter()
            .find(|column| matches!(column.name.as_str(), "label" | "name" | "code"))
            .map(|column| column.name.clone())
            .unwrap_or_else(|| value_column.clone());
        let description_column = columns
            .iter()
            .find(|column| column.name == "description")
            .map(|column| column.name.clone());
        let description_sql = description_column
            .as_ref()
            .map(|column| quote_identifier(column))
            .unwrap_or_else(|| "NULL".to_string());
        let mut statement = connection
            .prepare(&format!(
                "SELECT {}, CAST({} AS VARCHAR), CAST({} AS VARCHAR) FROM {relation} ORDER BY 1",
                quote_identifier(&value_column),
                quote_identifier(&code_column),
                description_sql
            ))
            .map_err(|error| format!("could not read code table {candidate}: {error}"))?;
        let rows = statement
            .query_map([], |row| {
                Ok(SqlSourceVocabularyItemAnalysis {
                    value: row.get::<_, i32>(0)?,
                    code: row.get::<_, String>(1)?,
                    description: row.get::<_, Option<String>>(2)?,
                })
            })
            .map_err(|error| format!("could not query code table {candidate}: {error}"))?;
        let mut items = Vec::new();
        for row in rows {
            items.push(
                row.map_err(|error| format!("could not decode code table {candidate}: {error}"))?,
            );
        }
        if !items.is_empty() {
            return Ok(Some(SqlSourceVocabularyAnalysis {
                name: candidate,
                description: Some(format!("Inferred from {} source code table", source.kind())),
                items,
            }));
        }
    }
    Ok(None)
}

fn describe_relation_columns(
    connection: &Connection,
    relation: &str,
) -> Result<Vec<DescribedColumn>, duckdb::Error> {
    let mut statement = connection.prepare(&format!("DESCRIBE SELECT * FROM {relation}"))?;
    let rows = statement.query_map([], |row| {
        Ok(DescribedColumn {
            name: row.get::<_, String>(0)?,
            sql_type: row.get::<_, String>(1)?,
        })
    })?;
    let mut columns = Vec::new();
    for row in rows {
        columns.push(row?);
    }
    Ok(columns)
}

fn code_table_candidates(column_name: &str) -> Vec<String> {
    let mut candidates = vec![column_name.to_string()];
    if !column_name.ends_with('s') {
        candidates.push(format!("{column_name}s"));
    }
    if let Some(stem) = column_name.strip_suffix('y') {
        candidates.push(format!("{stem}ies"));
    }
    candidates
}

fn parse_select_columns(select_text: &str, warnings: &mut Vec<String>) -> Vec<QueryColumn> {
    split_top_level_commas(select_text)
        .into_iter()
        .map(|item| {
            let item = item.trim();
            if item == "*" {
                return QueryColumn {
                    output_name: "*".to_string(),
                    source_name: None,
                    qualifier: None,
                    wildcard: true,
                };
            }
            let before_alias = item
                .split_once(" as ")
                .map(|(left, _)| left)
                .unwrap_or(item)
                .trim();
            let output_name = item
                .split_once(" as ")
                .map(|(_, right)| right.trim().trim_matches('"').to_string())
                .unwrap_or_else(|| {
                    before_alias
                        .rsplit('.')
                        .next()
                        .unwrap_or(before_alias)
                        .trim_matches('"')
                        .to_string()
                });
            if before_alias.contains('(') {
                warnings.push(format!(
                    "query column {output_name} is computed; rich metadata discovery is limited"
                ));
                return QueryColumn {
                    output_name,
                    source_name: None,
                    qualifier: None,
                    wildcard: false,
                };
            }
            let (qualifier, source_name) = before_alias
                .rsplit_once('.')
                .map(|(qualifier, name)| {
                    (
                        Some(qualifier.trim_matches('"').to_string()),
                        Some(name.trim_matches('"').to_string()),
                    )
                })
                .unwrap_or((None, Some(before_alias.trim_matches('"').to_string())));
            QueryColumn {
                output_name,
                source_name,
                qualifier,
                wildcard: false,
            }
        })
        .collect()
}

fn split_top_level_commas(value: &str) -> Vec<String> {
    let mut items = Vec::new();
    let mut depth = 0_i32;
    let mut start = 0_usize;
    for (index, character) in value.char_indices() {
        match character {
            '(' => depth += 1,
            ')' => depth -= 1,
            ',' if depth == 0 => {
                items.push(value[start..index].to_string());
                start = index + 1;
            }
            _ => {}
        }
    }
    if start < value.len() {
        items.push(value[start..].to_string());
    }
    items
}

fn tokenize(sql: &str) -> Vec<String> {
    sql.replace(',', " ")
        .replace('(', " ( ")
        .replace(')', " ) ")
        .split_whitespace()
        .map(|token| token.trim_matches('"').to_string())
        .collect()
}

fn is_table_alias_stop(token: &str) -> bool {
    matches!(
        token.to_ascii_lowercase().as_str(),
        "on" | "where"
            | "join"
            | "left"
            | "right"
            | "inner"
            | "outer"
            | "full"
            | "cross"
            | "group"
            | "order"
            | "limit"
            | "offset"
            | "union"
    )
}

fn value_type_id_for_sql(sql_type: &str) -> ValueTypeId {
    let lowered = sql_type.to_ascii_lowercase();
    if lowered.contains("int") {
        ValueTypeId(1)
    } else if lowered.contains("float")
        || lowered.contains("double")
        || lowered.contains("real")
        || lowered.contains("decimal")
        || lowered.contains("numeric")
    {
        ValueTypeId(2)
    } else if lowered == "date" {
        ValueTypeId(4)
    } else if lowered.contains("timestamp") || lowered.contains("datetime") {
        ValueTypeId(5)
    } else if lowered == "time" || lowered.starts_with("time(") {
        ValueTypeId(6)
    } else if lowered.contains("bool") || lowered == "bit" {
        ValueTypeId(1)
    } else {
        ValueTypeId(3)
    }
}

pub fn value_type_name(value_type_id: ValueTypeId) -> Result<&'static str, AppError> {
    TreValueType::from_value_type_id(value_type_id)
        .map(TreValueType::datastore_name)
        .map_err(|error| AppError::Validation(error.to_string()))
}

fn source_relation_sql(alias: &str, table_name: &str) -> String {
    format!(
        "{}.{}",
        quote_identifier(alias),
        quote_identifier(table_name)
    )
}

fn quote_identifier(value: &str) -> String {
    format!("\"{}\"", value.replace('"', "\"\""))
}

fn sql_literal(value: &str) -> String {
    format!("'{}'", value.replace('\'', "''"))
}

fn infrastructure_error(operation: &str, error: impl std::fmt::Display) -> AppError {
    AppError::Infrastructure(format!("{operation} failed: {error}"))
}