This commit is contained in:
2026-07-15 07:11:09 +07:00
parent 95954249ce
commit 1ad46c6e18
2 changed files with 419 additions and 33 deletions
+126 -33
View File
@@ -177,6 +177,7 @@ function checkAgentResponse_JSON(responsedict::T1, requiredKeys::T2
return (ispass, errormsg)
end
""" Convert a plain text string containing key-value pairs into a JSON-formatted string.
This function takes text containing key-value pairs (typically extracted from LLM responses)
@@ -221,6 +222,96 @@ function clean_json_response(text::String)
end
"""
harvest_db_undirected_schema_graph(pg_conn_str::String)
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.
"""
function harvest_db_undirected_schema_graph(pg_conn_str)
conn = LibPQ.Connection(pg_conn_str)
try
# 1. Fetch all user tables
table_query = """
SELECT c.relname AS table_name, c.oid
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = 'public' AND c.relkind = 'r';
"""
table_df = DataFrame(execute(conn, table_query))
table_names = table_df.table_name
num_tables = length(table_names)
table_to_id = Dict{String, Int}(name => i for (i, name) in enumerate(table_names))
id_to_table = Dict{Int, String}(i => name for (i, name) in enumerate(table_names))
# CRITICAL: Use SimpleGraph (Undirected) so pathfinding can traverse both ways
g = SimpleGraph(num_tables)
# 2. Strategy A: Extract Explicit Foreign Keys
fk_query = """
SELECT
conrelid::regclass::text AS source_table,
confrelid::regclass::text AS target_table
FROM pg_constraint c
JOIN pg_namespace n ON n.oid = c.connamespace
WHERE c.contype = 'f' AND n.nspname = 'public';
"""
fk_df = DataFrame(execute(conn, fk_query))
for row in eachrow(fk_df)
src = split(replace(row.source_table, "\"" => ""), '.')[end]
tgt = split(replace(row.target_table, "\"" => ""), '.')[end]
if haskey(table_to_id, src) && haskey(table_to_id, tgt)
add_edge!(g, table_to_id[src], table_to_id[tgt])
end
end
# 3. Strategy B: Fallback Name Matching (For missing explicit constraints)
# Fetch all columns for all tables
col_query = """
SELECT a.attname AS column_name, c.relname AS table_name
FROM pg_attribute a
JOIN pg_class c ON c.oid = a.attrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = 'public' AND c.relkind = 'r' AND a.attnum > 0 AND NOT a.attisdropped;
"""
col_df = DataFrame(execute(conn, col_query))
# Group columns by name to see who shares what fields
for gdf in groupby(col_df, :column_name)
col_name = gdf.column_name[1]
# We only infer connections on ID fields (e.g., seller_id, product_id)
# Avoid matching generic names like 'id', 'created_at', or 'name'
if endswith(lowercase(col_name), "_id") && col_name != "id"
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)
t1 = sharing_tables[i]
t2 = sharing_tables[j]
if haskey(table_to_id, t1) && haskey(table_to_id, t2)
add_edge!(g, table_to_id[t1], table_to_id[t2])
end
end
end
end
return g, id_to_table, table_to_id
finally
close(conn)
end
end
""" Extract vector metadata from PostgreSQL database.
Queries PostgreSQL system catalogs to extract column metadata including table names,
@@ -261,36 +352,38 @@ products user_id integer Reference to user false true
"""
function extract_column_metadata(pg_conn_str::String)::DataFrame
conn = LibPQ.Connection(pg_conn_str)
return extract_column_metadata(conn)
end
function extract_column_metadata(conn::LibPQ.Connection)::DataFrame
# This direct SQL query pulls the column specifications along with column-level descriptions
query =
"""
SELECT
c.relname AS column_name,
format_type(a.atttypid, a.atttypmod) AS data_type,
COALESCE(d.description, '') AS column_description,
CASE WHEN pk.contype = 'p' THEN true ELSE false END AS is_primary_key,
CASE WHEN AS table_name,
a.attnamefk.contype = 'f' THEN true ELSE false END AS is_foreign_key
FROM pg_attribute a
JOIN pg_class c ON c.oid = a.attrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
-- Join to fetch column comments/descriptions
LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
-- Check if column is part of a Primary Key
LEFT JOIN pg_constraint pk ON pk.conrelid = c.oid
AND pk.contype = 'p'
AND a.attnum = ANY(pk.conkey)
-- Check if column is part of a Foreign Key
LEFT JOIN pg_constraint fk ON fk.conrelid = c.oid
AND fk.contype = 'f'
AND a.attnum = ANY(fk.conkey)
WHERE
n.nspname = 'public' -- Only user schemas
AND c.relkind = 'r' -- Only standard tables
AND a.attnum > 0 -- Skip system hidden columns
AND NOT a.attisdropped; -- Skip dropped columns
"""
query = """
SELECT
c.relname AS table_name,
a.attname AS column_name,
format_type(a.atttypid, a.atttypmod) AS data_type,
COALESCE(d.description, '') AS column_description,
CASE WHEN pk.contype = 'p' THEN true ELSE false END AS is_primary_key,
CASE WHEN fk.contype = 'f' THEN true ELSE false END AS is_foreign_key
FROM pg_attribute a
JOIN pg_class c ON c.oid = a.attrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
-- Join to fetch column comments/descriptions
LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
-- Check if column is part of a Primary Key
LEFT JOIN pg_constraint pk ON pk.conrelid = c.oid
AND pk.contype = 'p'
AND a.attnum = ANY(pk.conkey)
-- Check if column is part of a Foreign Key
LEFT JOIN pg_constraint fk ON fk.conrelid = c.oid
AND fk.contype = 'f'
AND a.attnum = ANY(fk.conkey)
WHERE
n.nspname = 'public' -- Only user schemas
AND c.relkind = 'r' -- Only standard tables
AND a.attnum > 0 -- Skip system hidden columns
AND NOT a.attisdropped; -- Skip dropped columns
"""
try
# Execute and format into a clean DataFrame
@@ -503,7 +596,7 @@ Extracts unique, non-null values from a specific column to build a local index f
semantic search or entity resolution.
# Arguments
- `conn_str::String`
- `pg_conn_str::String`
PostgreSQL connection string in LibPQ format (e.g., "host=hostname port=5432 dbname=database user=username password=secret")
- `table::String`
Table name to query
@@ -531,12 +624,12 @@ julia> fruits = GeneralUtils.harvest_entity_catalog(conn, "products", "fruit_nam
["Apple", "Banana", "Orange", "Mango"]
```
"""
function harvest_entity_catalog(conn_str::String, table::String, column::String)::Vector{String}
conn = LibPQ.Connection(conn_str)
function harvest_entity_catalog(pg_conn_str::String, table::String, column::String)::Vector{String}
conn = LibPQ.Connection(pg_conn_str)
return harvest_entity_catalog(conn, table, column)
end
function harvest_entity_catalog(conn, table::String, column::String)::Vector{String}
function harvest_entity_catalog(conn::LibPQ.Connection, table::String, column::String)::Vector{String}
# We only care about unique, non-null values to keep the index fast and dense
query = "SELECT DISTINCT $(column) FROM $(table) WHERE $(column) IS NOT NULL;"