use sqry_core::graph::unified::build::staging::{StagingGraph, StagingOp};
use sqry_core::graph::unified::edge::EdgeKind;
use sqry_core::graph::unified::node::NodeKind;
use sqry_core::graph::unified::storage::NodeEntry;
use sqry_core::plugin::LanguagePlugin;
use sqry_lang_sql::SqlPlugin;
use std::collections::HashMap;
use std::path::PathBuf;
fn build_string_lookup(staging: &StagingGraph) -> HashMap<u32, String> {
let mut lookup = HashMap::new();
for op in staging.operations() {
if let StagingOp::InternString { local_id, value } = op {
lookup.insert(local_id.index(), value.clone());
}
}
lookup
}
fn find_node_entry<'a>(
staging: &'a StagingGraph,
name: &str,
kind: NodeKind,
) -> Option<&'a NodeEntry> {
let strings = build_string_lookup(staging);
for op in staging.operations() {
if let StagingOp::AddNode { entry, .. } = op
&& entry.kind == kind
{
let node_name = strings.get(&entry.name.index());
if node_name.is_some_and(|n| n == name) {
return Some(entry);
}
}
}
None
}
fn count_edges_of_kind(staging: &StagingGraph, kind: &EdgeKind) -> usize {
staging
.operations()
.iter()
.filter(|op| {
if let StagingOp::AddEdge {
kind: edge_kind, ..
} = op
{
std::mem::discriminant(edge_kind) == std::mem::discriminant(kind)
} else {
false
}
})
.count()
}
fn build_graph(source: &[u8]) -> StagingGraph {
let plugin = SqlPlugin::default();
let file = PathBuf::from("test.sql");
let tree = plugin.parse_ast(source).expect("parse failed");
let mut staging = StagingGraph::new();
let builder = plugin.graph_builder().expect("graph builder");
builder
.build_graph(&tree, source, &file, &mut staging)
.expect("build graph");
staging
}
#[test]
fn test_create_table_uses_table_node_kind() {
let source = br#"
CREATE TABLE customers (
id INT PRIMARY KEY,
name VARCHAR(100) NOT NULL,
email VARCHAR(255) UNIQUE,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
"#;
let staging = build_graph(source);
let entry = find_node_entry(&staging, "customers", NodeKind::Variable);
assert!(
entry.is_some(),
"Expected customers table with NodeKind::Variable"
);
}
#[test]
fn test_create_view_uses_view_node_kind() {
let source = br#"
CREATE TABLE users (id INT, status VARCHAR(20));
CREATE VIEW active_users AS
SELECT id FROM users WHERE status = 'active';
"#;
let staging = build_graph(source);
let table_entry = find_node_entry(&staging, "users", NodeKind::Variable);
assert!(
table_entry.is_some(),
"Expected users with NodeKind::Variable"
);
let view_entry = find_node_entry(&staging, "active_users", NodeKind::Variable);
assert!(
view_entry.is_some(),
"Expected active_users with NodeKind::Variable"
);
}
#[test]
fn test_create_materialized_view_uses_view_node_kind() {
let source = br#"
CREATE TABLE logs (id INT, message TEXT);
CREATE MATERIALIZED VIEW log_summary AS
SELECT COUNT(*) as total FROM logs;
"#;
let staging = build_graph(source);
let view_entry = find_node_entry(&staging, "log_summary", NodeKind::Variable);
assert!(
view_entry.is_some(),
"Expected log_summary materialized view with NodeKind::Variable"
);
}
#[test]
fn test_create_function_uses_procedure_node_kind() {
let source = br#"
CREATE FUNCTION calculate_discount(price DECIMAL, rate DECIMAL)
RETURNS DECIMAL AS $$
BEGIN
RETURN price * rate;
END;
$$ LANGUAGE plpgsql;
"#;
let staging = build_graph(source);
let entry = find_node_entry(&staging, "calculate_discount", NodeKind::Function);
assert!(
entry.is_some(),
"Expected calculate_discount with NodeKind::Function"
);
}
#[test]
fn test_create_procedure_uses_procedure_node_kind() {
let source = br#"
CREATE FUNCTION archive_old_records(cutoff_date DATE) RETURNS VOID AS $$
BEGIN
DELETE FROM records WHERE created_at < cutoff_date;
END;
$$ LANGUAGE plpgsql;
"#;
let staging = build_graph(source);
let entry = find_node_entry(&staging, "archive_old_records", NodeKind::Function);
assert!(
entry.is_some(),
"Expected archive_old_records with NodeKind::Function"
);
}
#[test]
fn test_create_trigger_uses_trigger_node_kind() {
let source = br#"
CREATE FUNCTION notify_change() RETURNS TRIGGER AS $$
BEGIN
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
CREATE TRIGGER user_update_trigger
AFTER UPDATE ON users
FOR EACH ROW
EXECUTE FUNCTION notify_change();
"#;
let staging = build_graph(source);
let trigger_entry = find_node_entry(&staging, "user_update_trigger", NodeKind::Function);
assert!(
trigger_entry.is_some(),
"Expected user_update_trigger with NodeKind::Function"
);
}
#[test]
fn test_view_creates_table_read_edge() {
let source = br#"
CREATE TABLE products (id INT, name VARCHAR(100), price DECIMAL);
CREATE VIEW expensive_products AS
SELECT * FROM products WHERE price > 1000;
"#;
let staging = build_graph(source);
assert!(
find_node_entry(&staging, "products", NodeKind::Variable).is_some(),
"products table should exist"
);
assert!(
find_node_entry(&staging, "expensive_products", NodeKind::Variable).is_some(),
"expensive_products view should exist"
);
}
#[test]
fn test_procedure_creates_table_read_edge() {
let source = br#"
CREATE TABLE accounts (id INT, balance DECIMAL);
CREATE FUNCTION get_total_balance() RETURNS DECIMAL AS $$
SELECT SUM(balance) FROM accounts;
$$ LANGUAGE sql;
"#;
let staging = build_graph(source);
assert!(
find_node_entry(&staging, "accounts", NodeKind::Variable).is_some(),
"accounts table should exist"
);
assert!(
find_node_entry(&staging, "get_total_balance", NodeKind::Function).is_some(),
"get_total_balance procedure should exist"
);
let table_read_count = count_edges_of_kind(
&staging,
&EdgeKind::TableRead {
table_name: Default::default(),
schema: None,
},
);
assert!(
table_read_count >= 1,
"Expected at least 1 TableRead edge from procedure to table"
);
}
#[test]
fn test_procedure_creates_table_write_edge() {
let source = br#"
CREATE TABLE audit_log (id INT, message TEXT);
CREATE FUNCTION log_event(msg TEXT) RETURNS VOID AS $$
BEGIN
INSERT INTO audit_log (message) VALUES (msg);
END;
$$ LANGUAGE plpgsql;
"#;
let staging = build_graph(source);
assert!(
find_node_entry(&staging, "audit_log", NodeKind::Variable).is_some(),
"audit_log table should exist"
);
assert!(
find_node_entry(&staging, "log_event", NodeKind::Function).is_some(),
"log_event procedure should exist"
);
let table_write_count = count_edges_of_kind(
&staging,
&EdgeKind::TableWrite {
table_name: Default::default(),
schema: None,
operation: sqry_core::graph::unified::edge::TableWriteOp::Insert,
},
);
assert!(
table_write_count >= 1,
"Expected at least 1 TableWrite edge from procedure to table"
);
}
#[test]
fn test_trigger_creates_triggered_by_edge() {
let source = br#"
CREATE TABLE orders (id INT, status VARCHAR(20));
CREATE FUNCTION update_timestamp() RETURNS TRIGGER AS $$
BEGIN
NEW.updated_at = NOW();
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
CREATE TRIGGER orders_update_trigger
BEFORE UPDATE ON orders
FOR EACH ROW
EXECUTE FUNCTION update_timestamp();
"#;
let staging = build_graph(source);
assert!(
find_node_entry(&staging, "orders", NodeKind::Variable).is_some(),
"orders table should exist"
);
assert!(
find_node_entry(&staging, "orders_update_trigger", NodeKind::Function).is_some(),
"orders_update_trigger trigger should exist"
);
let triggered_by_count = count_edges_of_kind(
&staging,
&EdgeKind::TriggeredBy {
trigger_name: Default::default(),
schema: None,
},
);
assert!(
triggered_by_count >= 1,
"Expected at least 1 TriggeredBy edge from trigger to table"
);
}
#[test]
fn test_all_four_database_node_kinds_together() {
let source = br#"
CREATE TABLE inventory (
id INT PRIMARY KEY,
product_name VARCHAR(100),
quantity INT,
last_updated TIMESTAMP
);
CREATE VIEW low_stock AS
SELECT product_name, quantity FROM inventory WHERE quantity < 10;
CREATE FUNCTION restock_item(item_id INT, amount INT) RETURNS VOID AS $$
BEGIN
UPDATE inventory SET quantity = quantity + amount WHERE id = item_id;
END;
$$ LANGUAGE plpgsql;
CREATE TRIGGER track_inventory_changes
AFTER UPDATE ON inventory
FOR EACH ROW
EXECUTE FUNCTION log_inventory_change();
"#;
let staging = build_graph(source);
assert!(
find_node_entry(&staging, "inventory", NodeKind::Variable).is_some(),
"inventory table with NodeKind::Variable should exist"
);
assert!(
find_node_entry(&staging, "low_stock", NodeKind::Variable).is_some(),
"low_stock view with NodeKind::Variable should exist"
);
assert!(
find_node_entry(&staging, "restock_item", NodeKind::Function).is_some(),
"restock_item procedure with NodeKind::Function should exist"
);
assert!(
find_node_entry(&staging, "track_inventory_changes", NodeKind::Function).is_some(),
"track_inventory_changes trigger with NodeKind::Function should exist"
);
}
#[test]
fn test_schema_qualified_table_uses_table_node_kind() {
let source = br#"
CREATE TABLE public.employees (
id INT PRIMARY KEY,
name VARCHAR(100)
);
"#;
let staging = build_graph(source);
let entry = find_node_entry(&staging, "employees", NodeKind::Variable);
assert!(
entry.is_some(),
"Expected employees table (schema-qualified) with NodeKind::Variable"
);
}
#[test]
fn test_multiple_tables_all_use_table_node_kind() {
let source = br#"
CREATE TABLE users (id INT, name VARCHAR(100));
CREATE TABLE posts (id INT, user_id INT, title VARCHAR(200));
CREATE TABLE comments (id INT, post_id INT, content TEXT);
"#;
let staging = build_graph(source);
assert!(
find_node_entry(&staging, "users", NodeKind::Variable).is_some(),
"users should use NodeKind::Variable"
);
assert!(
find_node_entry(&staging, "posts", NodeKind::Variable).is_some(),
"posts should use NodeKind::Variable"
);
assert!(
find_node_entry(&staging, "comments", NodeKind::Variable).is_some(),
"comments should use NodeKind::Variable"
);
}
#[test]
fn test_database_lineage_chain() {
let source = br#"
CREATE TABLE raw_data (id INT, value INT);
CREATE VIEW processed_data AS
SELECT id, value * 2 as doubled FROM raw_data;
CREATE FUNCTION get_processed_count() RETURNS INT AS $$
SELECT COUNT(*) FROM raw_data;
$$ LANGUAGE sql;
CREATE TRIGGER validate_raw_data
BEFORE INSERT ON raw_data
FOR EACH ROW
EXECUTE FUNCTION validate_value();
"#;
let staging = build_graph(source);
assert!(
find_node_entry(&staging, "raw_data", NodeKind::Variable).is_some(),
"raw_data table should exist"
);
assert!(
find_node_entry(&staging, "processed_data", NodeKind::Variable).is_some(),
"processed_data view should exist"
);
assert!(
find_node_entry(&staging, "get_processed_count", NodeKind::Function).is_some(),
"get_processed_count procedure should exist"
);
assert!(
find_node_entry(&staging, "validate_raw_data", NodeKind::Function).is_some(),
"validate_raw_data trigger should exist"
);
let table_read_count = count_edges_of_kind(
&staging,
&EdgeKind::TableRead {
table_name: Default::default(),
schema: None,
},
);
assert!(
table_read_count >= 1,
"Expected TableRead edges in lineage chain"
);
let triggered_by_count = count_edges_of_kind(
&staging,
&EdgeKind::TriggeredBy {
trigger_name: Default::default(),
schema: None,
},
);
assert!(
triggered_by_count >= 1,
"Expected TriggeredBy edge in lineage chain"
);
}