Spaces:
Sleeping
Sleeping
File size: 6,583 Bytes
f0e5200 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | # import os
# import sqlite3
# import threading
# from typing import Dict, List, Set, Tuple
# def get_schema(db_path):
# schema_map = get_table_to_columns(db_path)
# schema_text = ""
# for table, col_names in schema_map.items():
# schema_text += f"{table}({', '.join(col_names)})\n"
# return schema_text
# _SCHEMA_LOCK = threading.Lock()
# _SCHEMA_CACHE: Dict[str, Tuple[str, Dict[str, List[str]]]] = {}
# def _db_state_fingerprint(db_path: str) -> str:
# try:
# st = os.stat(db_path)
# return f"{st.st_mtime_ns}:{st.st_size}"
# except OSError:
# return "missing"
# def _connect_readonly(db_path: str) -> sqlite3.Connection:
# uri = f"file:{os.path.abspath(db_path)}?mode=ro"
# conn = sqlite3.connect(uri, uri=True, check_same_thread=False)
# conn.execute("PRAGMA query_only = ON;")
# conn.execute("PRAGMA foreign_keys = ON;")
# return conn
# def get_table_to_columns(db_path: str) -> Dict[str, List[str]]:
# """
# Return mapping of table -> column names for the SQLite DB at db_path.
# Tables and columns are returned lowercased.
# """
# fp = _db_state_fingerprint(db_path)
# with _SCHEMA_LOCK:
# cached = _SCHEMA_CACHE.get(db_path)
# if cached is not None and cached[0] == fp:
# return cached[1]
# schema: Dict[str, List[str]] = {}
# with _connect_readonly(db_path) as conn:
# cur = conn.execute(
# "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
# )
# tables = [r[0] for r in cur.fetchall() if r and isinstance(r[0], str)]
# for table in tables:
# table_l = table.lower()
# try:
# cur = conn.execute(f'PRAGMA table_info("{table}")')
# cols = [row[1].lower() for row in cur.fetchall() if row and isinstance(row[1], str)]
# schema[table_l] = cols
# except sqlite3.Error:
# schema[table_l] = []
# with _SCHEMA_LOCK:
# _SCHEMA_CACHE[db_path] = (fp, schema)
# return schema
# def get_db_tables_and_columns(db_path: str) -> Tuple[Set[str], Set[str]]:
# schema = get_table_to_columns(db_path)
# tables = set(schema.keys())
# columns: Set[str] = set()
# for cols in schema.values():
# columns.update(cols)
# return tables, columns
import os
import sqlite3
import threading
from typing import Dict, List, Set, Tuple
# ===============================
# 🔥 SCHEMA TEXT (for prompting)
# ===============================
def get_schema(db_path: str) -> str:
schema_map = get_table_to_columns(db_path)
schema_text = ""
for table, col_names in schema_map.items():
schema_text += f"{table}({', '.join(col_names)})\n"
return schema_text
# ===============================
# 🔥 CACHE + LOCK
# ===============================
_SCHEMA_LOCK = threading.Lock()
_SCHEMA_CACHE: Dict[str, Tuple[str, Dict[str, List[str]]]] = {}
def _db_state_fingerprint(db_path: str) -> str:
try:
st = os.stat(db_path)
return f"{st.st_mtime_ns}:{st.st_size}"
except OSError:
return "missing"
def _connect_readonly(db_path: str) -> sqlite3.Connection:
uri = f"file:{os.path.abspath(db_path)}?mode=ro"
conn = sqlite3.connect(uri, uri=True, check_same_thread=False)
conn.execute("PRAGMA query_only = ON;")
conn.execute("PRAGMA foreign_keys = ON;")
return conn
# ===============================
# 🔥 CORE: TABLE → COLUMNS
# ===============================
def get_table_to_columns(db_path: str) -> Dict[str, List[str]]:
"""
Return mapping of table -> column names (ONLY names, no types).
"""
fp = _db_state_fingerprint(db_path)
with _SCHEMA_LOCK:
cached = _SCHEMA_CACHE.get(db_path)
if cached is not None and cached[0] == fp:
return cached[1]
schema: Dict[str, List[str]] = {}
with _connect_readonly(db_path) as conn:
cur = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
)
tables = [r[0] for r in cur.fetchall() if r and isinstance(r[0], str)]
for table in tables:
table_l = table.lower()
try:
cur = conn.execute(f'PRAGMA table_info("{table}")')
cols = []
for row in cur.fetchall():
col_name = row[1].lower()
cols.append(col_name)
schema[table_l] = list(set(cols)) # remove duplicates
except sqlite3.Error:
schema[table_l] = []
with _SCHEMA_LOCK:
_SCHEMA_CACHE[db_path] = (fp, schema)
return schema
# ===============================
# 🔥 TABLE + COLUMN SETS
# ===============================
def get_db_tables_and_columns(db_path: str) -> Tuple[Set[str], Set[str]]:
schema = get_table_to_columns(db_path)
tables = set(schema.keys())
columns: Set[str] = set()
for cols in schema.values():
columns.update(cols)
return tables, columns
# ===============================
# 🔥 FOREIGN KEYS (IMPORTANT)
# ===============================
def get_foreign_keys(db_path: str) -> List[Tuple[str, str, str, str]]:
"""
Returns list of foreign key relations:
(table, column, ref_table, ref_column)
"""
fks = []
with _connect_readonly(db_path) as conn:
cur = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
)
tables = [r[0] for r in cur.fetchall()]
for table in tables:
try:
cur = conn.execute(f'PRAGMA foreign_key_list("{table}")')
for row in cur.fetchall():
fks.append((
table.lower(),
row[3].lower(), # column
row[2].lower(), # referenced table
row[4].lower() # referenced column
))
except sqlite3.Error:
continue
return fks
# ===============================
# 🔥 FINAL: CONSTRAINT GRAPH
# ===============================
def get_constraint_graph(db_path: str):
"""
Build schema constraint graph:
- tables
- columns
- foreign key relations
"""
tables, columns = get_db_tables_and_columns(db_path)
fks = get_foreign_keys(db_path)
return {
"tables": tables,
"columns": columns,
"foreign_keys": fks
} |