diff --git a/src/llmUtil.jl b/src/llmUtil.jl index 87658d4..5545822 100644 --- a/src/llmUtil.jl +++ b/src/llmUtil.jl @@ -2,7 +2,7 @@ module llmUtil export formatLLMtext, extractthink, checkAgentResponse_JSON, clean_json_response, extract_column_metadata, generate_embedding_payloads, resolve_semantic_cluster, - harvest_entity_catalog, resolve_entity + harvest_entity_catalog, resolve_entity, harvest_db_undirected_schema_graph using UUIDs, JSON, Dates, DataFrames, StringDistances, Graphs, LibPQ using ..util @@ -223,13 +223,35 @@ end -""" - harvest_db_undirected_schema_graph(pg_conn_str::String) +""" Harvest database schema as undirected graph from PostgreSQL. -Connects to a PostgreSQL instance, queries its metadata catalogs, and returns: -1. `g::SimpleDiGraph`: The structural graph where nodes are tables. -2. `id_to_table::Dict{Int, String}`: Maps numerical node IDs to real table names. -3. `table_to_id::Dict{String, Int}`: Maps table names back to graph node IDs. +Extracts table structure and relationships from a PostgreSQL database by querying +system catalogs to build a graph representation of tables and their relationships. + +# Arguments +- `pg_conn_str::String` + PostgreSQL connection string in LibPQ format (e.g., "host=hostname port=5432 dbname=database user=username password=secret") + +# Return +- `g::SimpleGraph`: An undirected graph where nodes represent tables and edges represent + relationships (foreign keys or shared ID column patterns). +- `id_to_table::Dict{Int, String}`: Maps numerical node IDs (1..n) to actual table names. +- `table_to_id::Dict{String, Int}`: Reverse mapping from table names to graph node IDs. + +# Details +The function extracts schema information using two strategies: +1. **Explicit Foreign Keys**: Queries `pg_constraint` for actual foreign key relationships +2. **Fallback Name Matching**: Infers relationships from shared column naming patterns + (e.g., `seller_id`, `product_id` columns across tables) + +# Example +```julia +julia> using GeneralUtils +julia> pg_conn_str = "host=localhost port=5432 dbname=winedb user=admin password=secret" +julia> g, id_to_table, table_to_id = GeneralUtils.harvest_db_undirected_schema_graph(pg_conn_str) +julia> vertices(g) +10 +``` """ function harvest_db_undirected_schema_graph(pg_conn_str) conn = LibPQ.Connection(pg_conn_str) @@ -292,7 +314,7 @@ function harvest_db_undirected_schema_graph(pg_conn_str) sharing_tables = gdf.table_name # Connect all tables that share this ID column - for i in 1:length(sharing_tables), j in (i+1):length(sharing_tables) + for i in eachindex(sharing_tables), j in (i+1):length(sharing_tables) t1 = sharing_tables[i] t2 = sharing_tables[j]