Skip to content

Commit 8a200ae

Browse files
luiscarbonel1991luiscarbonel1991
authored
Merge pull request #33 from luiscarbonel1991/fix/pg-catalog-schema-visibility
fix: replace information_schema with pg_catalog for schema visibility (#32)
2 parents d90ecfd + 3966004 commit 8a200ae

7 files changed

Lines changed: 311 additions & 185 deletions

File tree

‎src/nlp2sql/adapters/postgres_repository.py‎

Lines changed: 192 additions & 148 deletions
Large diffs are not rendered by default.

‎src/nlp2sql/adapters/regex_query_validator.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,9 +48,7 @@ async def validate_columns(self, sql: str, tables: list[Any]) -> list[str]:
4848

4949
# Extract table aliases (FROM/JOIN table alias, table AS alias)
5050
table_aliases: set[str] = set()
51-
for match in re.finditer(
52-
r"\b(?:from|join)\s+(?:\w+\.)?(\w+)\s+(?:AS\s+)?([a-z_]\w*)\b", sql, re.IGNORECASE
53-
):
51+
for match in re.finditer(r"\b(?:from|join)\s+(?:\w+\.)?(\w+)\s+(?:AS\s+)?([a-z_]\w*)\b", sql, re.IGNORECASE):
5452
alias = match.group(2).lower()
5553
if alias not in SQL_KEYWORDS:
5654
table_aliases.add(alias)

‎src/nlp2sql/client.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ def stats(self) -> dict[str, Any]:
9191
# Factory — the ``connect()`` function
9292
# ------------------------------------------------------------------
9393

94+
9495
async def connect(
9596
database_url: str,
9697
*,

‎src/nlp2sql/core/entities.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,5 +90,3 @@ class QueryExample:
9090
complexity: int # 1-5 scale
9191
tags: List[str] = field(default_factory=list)
9292
metadata: Dict[str, Any] = field(default_factory=dict)
93-
94-

‎src/nlp2sql/core/sql_keywords.py‎

Lines changed: 102 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,104 @@
11
"""SQL keywords — shared set for validation and tokenization."""
22

3-
SQL_KEYWORDS: frozenset[str] = frozenset({
4-
"select", "from", "where", "and", "or", "not", "in", "between",
5-
"group", "by", "order", "having", "limit", "offset", "as", "on",
6-
"join", "left", "right", "inner", "outer", "cross", "full",
7-
"case", "when", "then", "else", "end", "null", "true", "false",
8-
"asc", "desc", "distinct", "union", "all", "exists", "any",
9-
"count", "sum", "avg", "min", "max", "abs", "round", "floor", "ceil",
10-
"date_trunc", "dateadd", "datediff", "current_date", "getdate",
11-
"extract", "epoch", "convert_timezone", "to_date", "to_char",
12-
"coalesce", "nullif", "cast", "like", "ilike", "is", "with",
13-
"month", "year", "day", "quarter", "week", "hour", "minute", "second",
14-
"varchar", "int", "integer", "bigint", "numeric", "decimal",
15-
"date", "timestamp", "boolean", "float", "double", "precision",
16-
"over", "partition", "row_number", "rank", "dense_rank",
17-
"lag", "lead", "first_value", "last_value", "listagg",
18-
"approximate", "interval", "explain", "analyze",
19-
})
3+
SQL_KEYWORDS: frozenset[str] = frozenset(
4+
{
5+
"select",
6+
"from",
7+
"where",
8+
"and",
9+
"or",
10+
"not",
11+
"in",
12+
"between",
13+
"group",
14+
"by",
15+
"order",
16+
"having",
17+
"limit",
18+
"offset",
19+
"as",
20+
"on",
21+
"join",
22+
"left",
23+
"right",
24+
"inner",
25+
"outer",
26+
"cross",
27+
"full",
28+
"case",
29+
"when",
30+
"then",
31+
"else",
32+
"end",
33+
"null",
34+
"true",
35+
"false",
36+
"asc",
37+
"desc",
38+
"distinct",
39+
"union",
40+
"all",
41+
"exists",
42+
"any",
43+
"count",
44+
"sum",
45+
"avg",
46+
"min",
47+
"max",
48+
"abs",
49+
"round",
50+
"floor",
51+
"ceil",
52+
"date_trunc",
53+
"dateadd",
54+
"datediff",
55+
"current_date",
56+
"getdate",
57+
"extract",
58+
"epoch",
59+
"convert_timezone",
60+
"to_date",
61+
"to_char",
62+
"coalesce",
63+
"nullif",
64+
"cast",
65+
"like",
66+
"ilike",
67+
"is",
68+
"with",
69+
"month",
70+
"year",
71+
"day",
72+
"quarter",
73+
"week",
74+
"hour",
75+
"minute",
76+
"second",
77+
"varchar",
78+
"int",
79+
"integer",
80+
"bigint",
81+
"numeric",
82+
"decimal",
83+
"date",
84+
"timestamp",
85+
"boolean",
86+
"float",
87+
"double",
88+
"precision",
89+
"over",
90+
"partition",
91+
"row_number",
92+
"rank",
93+
"dense_rank",
94+
"lag",
95+
"lead",
96+
"first_value",
97+
"last_value",
98+
"listagg",
99+
"approximate",
100+
"interval",
101+
"explain",
102+
"analyze",
103+
}
104+
)

‎tests/conftest.py‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,7 @@
2121
POSTGRES_PASS = os.getenv("NLP2SQL_TEST_POSTGRES_PASS", "testpass")
2222
POSTGRES_DB = os.getenv("NLP2SQL_TEST_POSTGRES_DB", "testdb")
2323

24-
POSTGRES_URL = (
25-
f"postgresql://{POSTGRES_USER}:{POSTGRES_PASS}"
26-
f"@{POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB}"
27-
)
24+
POSTGRES_URL = f"postgresql://{POSTGRES_USER}:{POSTGRES_PASS}@{POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB}"
2825

2926
# Expected tables from docker/init-schema.sql
3027
EXPECTED_TABLES = {"users", "categories", "products", "orders", "order_items", "reviews"}

‎tests/test_postgres_integration.py‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,7 @@ async def test_connect_and_discover_tables(self, postgres_available):
2323
tables = await repo.get_tables()
2424
table_names = {t.name.lower() for t in tables}
2525

26-
assert EXPECTED_TABLES.issubset(table_names), (
27-
f"Missing tables: {EXPECTED_TABLES - table_names}"
28-
)
26+
assert EXPECTED_TABLES.issubset(table_names), f"Missing tables: {EXPECTED_TABLES - table_names}"
2927

3028
async def test_column_discovery(self, postgres_available):
3129
repo = PostgreSQLRepository(postgres_available)
@@ -77,9 +75,7 @@ async def test_schema_filters_exclude_tables(self, postgres_available, mock_embe
7775
await repo.initialize()
7876
await manager.initialize(DatabaseType.POSTGRES)
7977

80-
context = await manager.get_optimal_schema_context(
81-
"How many users?", DatabaseType.POSTGRES, max_tokens=4000
82-
)
78+
context = await manager.get_optimal_schema_context("How many users?", DatabaseType.POSTGRES, max_tokens=4000)
8379
assert "reviews" not in context.lower()
8480

8581
async def test_schema_filters_include_tables(self, postgres_available, mock_embedding_provider):
@@ -92,9 +88,7 @@ async def test_schema_filters_include_tables(self, postgres_available, mock_embe
9288
await repo.initialize()
9389
await manager.initialize(DatabaseType.POSTGRES)
9490

95-
context = await manager.get_optimal_schema_context(
96-
"How many users?", DatabaseType.POSTGRES, max_tokens=4000
97-
)
91+
context = await manager.get_optimal_schema_context("How many users?", DatabaseType.POSTGRES, max_tokens=4000)
9892
assert "users" in context.lower()
9993

10094
async def test_embedding_search_against_real_schema(self, postgres_available, mock_embedding_provider, tmp_path):
@@ -110,9 +104,18 @@ async def test_embedding_search_against_real_schema(self, postgres_available, mo
110104

111105
elements = []
112106
for table in tables:
113-
elements.append({"type": "table", "name": table.name, "columns": [{"name": c["name"]} for c in table.columns]})
107+
elements.append(
108+
{"type": "table", "name": table.name, "columns": [{"name": c["name"]} for c in table.columns]}
109+
)
114110
for col in table.columns:
115-
elements.append({"type": "column", "name": col["name"], "table_name": table.name, "data_type": col.get("data_type", "unknown")})
111+
elements.append(
112+
{
113+
"type": "column",
114+
"name": col["name"],
115+
"table_name": table.name,
116+
"data_type": col.get("data_type", "unknown"),
117+
}
118+
)
116119

117120
await emb_manager.add_schema_elements(elements, DatabaseType.POSTGRES)
118121
results = await emb_manager.search_similar("customer email", top_k=5, min_score=0.0)

0 commit comments

Comments
 (0)