update

src/llmUtil.jl

@@ -2,7 +2,8 @@ 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_db_undirected_schema_graph
+        harvest_entity_catalog, resolve_entity, harvest_db_undirected_schema_graph,
+        get_db_table_schema
 
 using UUIDs, JSON, Dates, DataFrames, StringDistances, Graphs, LibPQ
 using ..util
@@ -964,6 +965,55 @@ function resolve_entity(messy_input::String, catalog::Vector{String}; threshold=
 end
 
 
+function get_db_table_schema(pg_conn_str::String, table_name::String)::DataFrame
+  conn = LibPQ.Connection(pg_conn_str)
+  return get_db_table_schema(conn, table_name)
+end
+
+function get_db_table_schema(conn::LibPQ.Connection, table_name::String)::DataFrame
+  # This direct SQL query pulls the column specifications along with column-level descriptions
+  query = """
+        SELECT 
+            a.attname AS column_name,
+            format_type(a.atttypid, a.atttypmod) AS data_type,
+            COALESCE(d.description, '') AS column_comment,
+            CASE 
+                WHEN p.contype = 'p' THEN 'PRIMARY KEY'
+                WHEN p.contype = 'u' THEN 'UNIQUE'
+                WHEN p.contype = 'f' THEN 'FOREIGN KEY'
+                WHEN p.contype = 'c' THEN 'CHECK'
+                ELSE ''
+            END AS constraint_type,
+            COALESCE(p.conname, '') AS constraint_name
+        FROM 
+            pg_catalog.pg_attribute a
+        JOIN 
+            pg_catalog.pg_class c ON a.attrelid = c.oid
+        JOIN 
+            pg_catalog.pg_namespace n ON c.relnamespace = n.oid
+        LEFT JOIN 
+            pg_catalog.pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
+        LEFT JOIN 
+            pg_catalog.pg_constraint p ON p.conrelid = c.oid AND a.attnum = ANY(p.conkey)
+        WHERE 
+            c.relname = '$table_name' -- <-- Put your table name here
+            AND n.nspname = 'public'      -- <-- Change schema if not 'public'
+            AND a.attnum > 0 
+            AND NOT a.attisdropped
+        ORDER BY 
+            a.attnum;
+    """
+  
+  try
+      # Execute and format into a clean DataFrame
+      result = LibPQ.execute(conn, query)
+      df = DataFrame(result)
+      
+      return df
+  finally
+      close(conn)
+  end
+end