nyx/src/surface/datastore.rs

//! Data-store detection.
//!
//! Walks the post-pass-2 [`GlobalSummaries`] looking for callees whose
//! name is a known database / cache / blob-store driver entry point,
//! and emits one [`SurfaceNode::DataStore`] per resolved store.
//!
//! The detector is name-based on purpose: the receiver's full type is
//! often unknown after pass 2, but the leaf name of a driver call
//! (`psycopg2.connect`, `mysql.createConnection`, `gorm.Open`,
//! `Eloquent::find`, `ActiveRecord::Base.connection`) carries enough
//! signal for surface-level chain composition.  False positives here
//! are forgiving — the surface map is informational, not a finding
//! that fires on its own.

use super::{DataStore, DataStoreKind, SourceLocation, SurfaceNode};
use crate::summary::{CalleeSite, FuncSummary, GlobalSummaries};

/// One detection rule: leaf-name pattern → store kind + label.  Stored
/// as a flat list so adding a new ORM / driver is a one-line edit.
struct DriverRule {
    /// Substring to match against the callee's leaf name (case-insensitive).
    leaf: &'static str,
    kind: DataStoreKind,
    /// Human-readable label attached to the emitted node.  Used by the
    /// chain composer and the `nyx surface` CLI tree.
    label: &'static str,
}

const DRIVER_RULES: &[DriverRule] = &[
    // Python — relational
    DriverRule { leaf: "psycopg2.connect", kind: DataStoreKind::Sql, label: "PostgreSQL (psycopg2)" },
    DriverRule { leaf: "psycopg.connect",  kind: DataStoreKind::Sql, label: "PostgreSQL (psycopg3)" },
    DriverRule { leaf: "mysql.connector.connect", kind: DataStoreKind::Sql, label: "MySQL (mysql.connector)" },
    DriverRule { leaf: "MySQLdb.connect",  kind: DataStoreKind::Sql, label: "MySQL (MySQLdb)" },
    DriverRule { leaf: "pymysql.connect",  kind: DataStoreKind::Sql, label: "MySQL (PyMySQL)" },
    DriverRule { leaf: "sqlite3.connect",  kind: DataStoreKind::Sql, label: "SQLite (sqlite3)" },
    DriverRule { leaf: "sqlalchemy.create_engine", kind: DataStoreKind::Sql, label: "SQLAlchemy" },
    DriverRule { leaf: "django.db.connection", kind: DataStoreKind::Sql, label: "Django ORM" },
    // Python — kv / doc
    DriverRule { leaf: "redis.Redis",      kind: DataStoreKind::KeyValue, label: "Redis" },
    DriverRule { leaf: "redis.from_url",   kind: DataStoreKind::KeyValue, label: "Redis" },
    DriverRule { leaf: "pymongo.MongoClient", kind: DataStoreKind::Document, label: "MongoDB" },
    DriverRule { leaf: "boto3.client",     kind: DataStoreKind::BlobStore, label: "AWS (boto3)" },
    DriverRule { leaf: "boto3.resource",   kind: DataStoreKind::BlobStore, label: "AWS (boto3)" },

    // JavaScript / TypeScript — relational
    DriverRule { leaf: "knex",             kind: DataStoreKind::Sql, label: "Knex.js" },
    DriverRule { leaf: "createConnection", kind: DataStoreKind::Sql, label: "MySQL/Postgres (mysql/pg)" },
    DriverRule { leaf: "Sequelize",        kind: DataStoreKind::Sql, label: "Sequelize" },
    DriverRule { leaf: "TypeORM.createConnection", kind: DataStoreKind::Sql, label: "TypeORM" },
    DriverRule { leaf: "PrismaClient",     kind: DataStoreKind::Sql, label: "Prisma" },
    DriverRule { leaf: "pool.query",       kind: DataStoreKind::Sql, label: "pg/mysql pool" },
    DriverRule { leaf: "client.query",     kind: DataStoreKind::Sql, label: "pg client" },
    DriverRule { leaf: "db.query",         kind: DataStoreKind::Sql, label: "Generic SQL driver" },
    // JS — kv / doc
    DriverRule { leaf: "redis.createClient", kind: DataStoreKind::KeyValue, label: "Redis (node-redis)" },
    DriverRule { leaf: "ioredis",          kind: DataStoreKind::KeyValue, label: "ioredis" },
    DriverRule { leaf: "MongoClient.connect", kind: DataStoreKind::Document, label: "MongoDB (node)" },
    DriverRule { leaf: "AWS.S3",           kind: DataStoreKind::BlobStore, label: "AWS S3" },

    // Java — JDBC / Hibernate
    DriverRule { leaf: "DriverManager.getConnection", kind: DataStoreKind::Sql, label: "JDBC" },
    DriverRule { leaf: "JdbcTemplate",     kind: DataStoreKind::Sql, label: "Spring JdbcTemplate" },
    DriverRule { leaf: "EntityManager",    kind: DataStoreKind::Sql, label: "JPA EntityManager" },
    DriverRule { leaf: "SessionFactory.openSession", kind: DataStoreKind::Sql, label: "Hibernate" },
    DriverRule { leaf: "Jedis",            kind: DataStoreKind::KeyValue, label: "Jedis (Redis)" },
    DriverRule { leaf: "MongoClients.create", kind: DataStoreKind::Document, label: "MongoDB (java-driver)" },

    // Go — sql + ORM
    DriverRule { leaf: "sql.Open",         kind: DataStoreKind::Sql, label: "database/sql" },
    DriverRule { leaf: "gorm.Open",        kind: DataStoreKind::Sql, label: "GORM" },
    DriverRule { leaf: "sqlx.Connect",     kind: DataStoreKind::Sql, label: "sqlx" },
    DriverRule { leaf: "sqlx.Open",        kind: DataStoreKind::Sql, label: "sqlx" },
    DriverRule { leaf: "redis.NewClient",  kind: DataStoreKind::KeyValue, label: "go-redis" },
    DriverRule { leaf: "mongo.Connect",    kind: DataStoreKind::Document, label: "MongoDB (go-driver)" },

    // PHP — Eloquent / PDO
    DriverRule { leaf: "PDO",              kind: DataStoreKind::Sql, label: "PDO" },
    DriverRule { leaf: "Eloquent::find",   kind: DataStoreKind::Sql, label: "Laravel Eloquent" },
    DriverRule { leaf: "Eloquent::where",  kind: DataStoreKind::Sql, label: "Laravel Eloquent" },
    DriverRule { leaf: "DB::connection",   kind: DataStoreKind::Sql, label: "Laravel DB" },
    DriverRule { leaf: "Doctrine",         kind: DataStoreKind::Sql, label: "Doctrine ORM" },

    // Ruby — ActiveRecord
    DriverRule { leaf: "ActiveRecord::Base.connection", kind: DataStoreKind::Sql, label: "ActiveRecord" },
    DriverRule { leaf: "ActiveRecord::Base.find",       kind: DataStoreKind::Sql, label: "ActiveRecord" },
    DriverRule { leaf: ".find_by_sql",     kind: DataStoreKind::Sql, label: "ActiveRecord raw SQL" },

    // Rust — sqlx / diesel
    DriverRule { leaf: "sqlx::query",      kind: DataStoreKind::Sql, label: "sqlx" },
    DriverRule { leaf: "sqlx::query_as",   kind: DataStoreKind::Sql, label: "sqlx" },
    DriverRule { leaf: "diesel::sql_query", kind: DataStoreKind::Sql, label: "Diesel" },
    DriverRule { leaf: "PgConnection::establish", kind: DataStoreKind::Sql, label: "Diesel" },

    // Type-qualified — fires when the SSA type-fact engine resolves a
    // receiver to `TypeKind::DatabaseConnection` regardless of the bare
    // callee name (e.g. `conn = psycopg2.connect(); conn.cursor()` →
    // typed_call_receivers maps the `.cursor` ordinal to "DatabaseConnection").
    DriverRule { leaf: "DatabaseConnection.cursor",  kind: DataStoreKind::Sql, label: "Database connection" },
    DriverRule { leaf: "DatabaseConnection.execute", kind: DataStoreKind::Sql, label: "Database connection" },
    DriverRule { leaf: "DatabaseConnection.query",   kind: DataStoreKind::Sql, label: "Database connection" },
    DriverRule { leaf: "DatabaseConnection.exec",    kind: DataStoreKind::Sql, label: "Database connection" },
    DriverRule { leaf: "DatabaseConnection.prepare", kind: DataStoreKind::Sql, label: "Database connection" },
    DriverRule { leaf: "DatabaseConnection.commit",  kind: DataStoreKind::Sql, label: "Database connection" },
    DriverRule { leaf: "FileHandle.read",  kind: DataStoreKind::Filesystem, label: "Filesystem" },
    DriverRule { leaf: "FileHandle.write", kind: DataStoreKind::Filesystem, label: "Filesystem" },
    DriverRule { leaf: "FileHandle.close", kind: DataStoreKind::Filesystem, label: "Filesystem" },

    // Filesystem (best-effort: language-agnostic open()-family)
    DriverRule { leaf: "open",             kind: DataStoreKind::Filesystem, label: "Filesystem" },
];

/// Walk every function summary's callee list and emit one
/// [`SurfaceNode::DataStore`] per matched driver call.  De-duped on
/// `(file, line, label)`.
///
/// When the bare callee name does not hit a rule, the type-fact engine's
/// per-call `typed_call_receivers` map (read off the matching
/// [`crate::summary::SsaFuncSummary`]) is consulted: a callee whose
/// receiver was resolved to `TypeKind::DatabaseConnection` or
/// `TypeKind::FileHandle` is retried under the type-qualified name
/// `"DatabaseConnection.<method>"` / `"FileHandle.<method>"`, picking up
/// the bound-receiver call shapes (`conn.cursor()` after
/// `conn = psycopg2.connect()`) that the name-only matcher misses.
pub fn detect_data_stores(summaries: &GlobalSummaries) -> Vec<SurfaceNode> {
    let mut out: Vec<SurfaceNode> = Vec::new();
    let mut seen: std::collections::HashSet<(String, u32, String)> =
        std::collections::HashSet::new();
    for (key, summary) in summaries.iter() {
        let typed = summaries.get_ssa(key).map(|s| s.typed_call_receivers.as_slice());
        for callee in &summary.callees {
            let rule = match_rule(&callee.name).or_else(|| {
                typed
                    .and_then(|t| container_for_ordinal(t, callee.ordinal))
                    .and_then(|c| match_rule(&qualify(c, &callee.name)))
            });
            let Some(rule) = rule else { continue };
            let location = call_site_location(summary, callee);
            let dedup = (
                location.file.clone(),
                location.line,
                rule.label.to_string(),
            );
            if !seen.insert(dedup) {
                continue;
            }
            out.push(SurfaceNode::DataStore(DataStore {
                location,
                kind: rule.kind,
                label: rule.label.to_string(),
            }));
        }
    }
    out
}

/// Last segment of a callee text after the final `.` or `::`.
fn leaf_segment(name: &str) -> &str {
    let after_colon = name.rsplit("::").next().unwrap_or(name);
    after_colon.rsplit('.').next().unwrap_or(after_colon)
}

/// Build a type-qualified callee name (`"{container}.{method}"`) for
/// retry-matching when the bare callee text did not hit any rule.
fn qualify(container: &str, callee_name: &str) -> String {
    format!("{}.{}", container, leaf_segment(callee_name))
}

/// Linear-scan helper since `typed_call_receivers` is a small
/// `Vec<(ordinal, container)>` per function. Typical lengths are 0 to a
/// few dozen; a HashMap-per-summary would be wasteful.
fn container_for_ordinal(typed: &[(u32, String)], ordinal: u32) -> Option<&str> {
    typed.iter().find(|(o, _)| *o == ordinal).map(|(_, c)| c.as_str())
}

fn match_rule(callee: &str) -> Option<&'static DriverRule> {
    let cl = callee.trim().to_ascii_lowercase();
    // Normalize `::` → `.` so segment-split treats both as separators.
    let cl_segments = cl.replace("::", ".");
    DRIVER_RULES.iter().find(|r| {
        let rl = r.leaf.to_ascii_lowercase();
        if r.leaf.contains('.') || r.leaf.contains("::") {
            // Qualified pattern (e.g. `psycopg2.connect`, `Eloquent::find`):
            // substring on the full callee text.  Qualified shapes are
            // unambiguous so substring is precise enough.
            cl.contains(&rl)
        } else {
            // Bare leaf (e.g. `open`, `fetch`, `PrismaClient`): require a
            // whole-segment match.  Prevents `fopen` / `OpenSearch` /
            // `getPrismaClient` from FP-matching short bare leaves.
            cl_segments.split('.').any(|seg| seg == rl)
        }
    })
}

/// Source location of a call site.  Reads the 1-based `(line, col)`
/// recorded on the [`CalleeSite`] at CFG-build time (populated for every
/// summary produced after the span field landed); for legacy summaries
/// loaded from SQLite with no span, falls back to the function's host
/// file with line 0.
fn call_site_location(summary: &FuncSummary, callee: &CalleeSite) -> SourceLocation {
    let (line, col) = callee.span.unwrap_or((0, 0));
    SourceLocation {
        file: summary.file_path.clone(),
        line,
        col,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::symbol::{FuncKey, Lang};

    fn summary_with_callees(name: &str, file: &str, callees: &[&str]) -> (FuncKey, FuncSummary) {
        let key = FuncKey::new_function(Lang::Python, file, name, None);
        let summary = FuncSummary {
            name: name.to_string(),
            file_path: file.to_string(),
            lang: "python".to_string(),
            param_count: 0,
            callees: callees
                .iter()
                .map(|c| CalleeSite::bare(c.to_string()))
                .collect(),
            ..Default::default()
        };
        (key, summary)
    }

    #[test]
    fn datastore_carries_callee_span_when_present() {
        // When the CFG populates `CalleeSite.span`, the detected datastore
        // node's `SourceLocation` must reflect that 1-based `(line, col)`
        // — not the legacy `(0, 0)` fallback.
        let mut gs = GlobalSummaries::new();
        let key = FuncKey::new_function(Lang::Python, "app.py", "init", None);
        let mut callee = CalleeSite::bare("psycopg2.connect");
        callee.span = Some((42, 13));
        let summary = FuncSummary {
            name: "init".into(),
            file_path: "app.py".into(),
            lang: "python".into(),
            param_count: 0,
            callees: vec![callee],
            ..Default::default()
        };
        gs.insert(key, summary);
        let nodes = detect_data_stores(&gs);
        assert_eq!(nodes.len(), 1);
        let SurfaceNode::DataStore(ds) = &nodes[0] else {
            panic!()
        };
        assert_eq!(ds.location.line, 42);
        assert_eq!(ds.location.col, 13);
    }

    #[test]
    fn detects_psycopg2_connect() {
        let mut gs = GlobalSummaries::new();
        let (k, s) = summary_with_callees("init", "app.py", &["psycopg2.connect"]);
        gs.insert(k, s);
        let nodes = detect_data_stores(&gs);
        assert_eq!(nodes.len(), 1);
        let SurfaceNode::DataStore(ds) = &nodes[0] else {
            panic!()
        };
        assert_eq!(ds.kind, DataStoreKind::Sql);
        assert_eq!(ds.label, "PostgreSQL (psycopg2)");
    }

    #[test]
    fn detects_gorm_open() {
        let mut gs = GlobalSummaries::new();
        let (k, s) = summary_with_callees("init", "main.go", &["gorm.Open"]);
        gs.insert(k, s);
        let nodes = detect_data_stores(&gs);
        assert_eq!(nodes.len(), 1);
        let SurfaceNode::DataStore(ds) = &nodes[0] else {
            panic!()
        };
        assert_eq!(ds.label, "GORM");
    }

    #[test]
    fn dedup_collapses_repeats_in_same_file() {
        let mut gs = GlobalSummaries::new();
        let (k, s) = summary_with_callees(
            "init",
            "app.py",
            &["psycopg2.connect", "psycopg2.connect"],
        );
        gs.insert(k, s);
        let nodes = detect_data_stores(&gs);
        assert_eq!(nodes.len(), 1);
    }

    #[test]
    fn bare_open_rule_does_not_match_fopen_or_opensearch() {
        let mut gs = GlobalSummaries::new();
        let (k, s) = summary_with_callees(
            "init",
            "app.py",
            &[
                "fopen",
                "popen",
                "OpenSearch",
                "openssl_encrypt",
                "MongoClient.openSession",
            ],
        );
        gs.insert(k, s);
        let nodes = detect_data_stores(&gs);
        assert!(
            nodes.is_empty(),
            "bare `open` rule should not FP on {nodes:?}",
        );
    }

    #[test]
    fn bare_open_rule_still_matches_real_open() {
        let mut gs = GlobalSummaries::new();
        let (k, s) = summary_with_callees("loader", "app.py", &["open"]);
        gs.insert(k, s);
        let nodes = detect_data_stores(&gs);
        assert_eq!(nodes.len(), 1);
        let SurfaceNode::DataStore(ds) = &nodes[0] else {
            panic!()
        };
        assert_eq!(ds.kind, DataStoreKind::Filesystem);

        let mut gs = GlobalSummaries::new();
        let (k, s) = summary_with_callees("loader", "app.py", &["builtins.open"]);
        gs.insert(k, s);
        let nodes = detect_data_stores(&gs);
        assert_eq!(nodes.len(), 1);
    }

    #[test]
    fn typed_receiver_database_connection_resolves_bound_cursor() {
        // `conn = psycopg2.connect(); conn.cursor()` — the bare callee
        // `conn.cursor` is not in DRIVER_RULES, but the SSA type-fact
        // engine populates `typed_call_receivers` with
        // `(ordinal, "DatabaseConnection")` for the `.cursor` ordinal.
        // The detector retries under `DatabaseConnection.cursor` and
        // emits a Sql datastore node.
        use crate::summary::ssa_summary::SsaFuncSummary;
        let mut gs = GlobalSummaries::new();
        let key = FuncKey::new_function(Lang::Python, "app.py", "load", None);
        let summary = FuncSummary {
            name: "load".into(),
            file_path: "app.py".into(),
            lang: "python".into(),
            param_count: 0,
            callees: vec![
                {
                    let mut c = CalleeSite::bare("conn.cursor");
                    c.ordinal = 7;
                    c.span = Some((4, 8));
                    c
                },
            ],
            ..Default::default()
        };
        gs.insert(key.clone(), summary);
        let mut ssa = SsaFuncSummary::default();
        ssa.typed_call_receivers
            .push((7, "DatabaseConnection".into()));
        gs.insert_ssa(key, ssa);
        let nodes = detect_data_stores(&gs);
        assert_eq!(nodes.len(), 1, "expected typed retry to hit; got {nodes:?}");
        let SurfaceNode::DataStore(ds) = &nodes[0] else {
            panic!()
        };
        assert_eq!(ds.kind, DataStoreKind::Sql);
        assert_eq!(ds.label, "Database connection");
        assert_eq!(ds.location.line, 4);
    }

    #[test]
    fn typed_receiver_without_ssa_summary_falls_through() {
        // No SsaFuncSummary inserted → bare `client.cursor` does not match
        // any rule and `typed_call_receivers` is unreachable. Detector
        // emits zero nodes (no panic on missing SSA side).
        let mut gs = GlobalSummaries::new();
        let (k, s) = summary_with_callees("load", "app.py", &["client.cursor"]);
        gs.insert(k, s);
        assert!(detect_data_stores(&gs).is_empty());
    }
}