-
-
Notifications
You must be signed in to change notification settings - Fork 17
Expand file tree
/
Copy pathtest_postgres_schema.py
More file actions
193 lines (160 loc) · 6.65 KB
/
Copy pathtest_postgres_schema.py
File metadata and controls
193 lines (160 loc) · 6.65 KB
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
import json
import sys
import types
from engraphis.backends import postgres_schema
from engraphis.core.interfaces import SchemaSnapshot, SearchFilter
from engraphis.service import MemoryService
class _Cursor:
def __init__(self):
self.calls = []
self.result = []
def __enter__(self):
return self
def __exit__(self, *args):
return False
def execute(self, query, params=()):
normalized = " ".join(query.split())
self.calls.append((normalized, tuple(params)))
if "current_database()" in normalized:
self.result = [("appdb",)]
elif "information_schema.tables" in normalized:
self.result = [
("public", "users", "BASE TABLE"),
("public", "orders", "BASE TABLE"),
("auth", "accounts", "BASE TABLE"),
]
elif "information_schema.columns" in normalized:
self.result = [
("public", "users", "id", 1, "integer", "NO", None),
("public", "users", "tenant_id", 2, "integer", "NO", None),
("public", "orders", "account_id", 1, "integer", "NO", None),
("auth", "accounts", "id", 1, "integer", "NO", None),
]
else:
self.result = [
("PRIMARY KEY", "public", "users", "id",
"public", "users", "id", "shared_name"),
("PRIMARY KEY", "public", "orders", "account_id",
"public", "orders", "account_id", "shared_name"),
("FOREIGN KEY", "public", "orders", "account_id",
"auth", "accounts", "id", "orders_account_fk"),
]
def fetchone(self):
return self.result[0]
def fetchall(self):
return list(self.result)
class _Connection:
def __init__(self):
self.cursor_obj = _Cursor()
self.closed = False
def __enter__(self):
return self
def __exit__(self, *args):
return False
def cursor(self):
return self.cursor_obj
def close(self):
self.closed = True
def test_postgres_connect_and_statement_timeouts_are_bounded(monkeypatch):
captured = {}
connection = _Connection()
def connect(dsn, **kwargs):
captured["dsn"] = dsn
captured.update(kwargs)
return connection
monkeypatch.setitem(sys.modules, "psycopg", types.SimpleNamespace(connect=connect))
monkeypatch.setenv("ENGRAPHIS_POSTGRES_CONNECT_TIMEOUT", "9999")
monkeypatch.setenv("ENGRAPHIS_POSTGRES_STATEMENT_TIMEOUT_MS", "45000")
snapshot = postgres_schema.PostgresSchemaIntrospector().inspect(
"postgresql://local/appdb"
)
assert snapshot.metadata["database"] == "appdb"
assert captured["connect_timeout"] == postgres_schema._MAX_CONNECT_TIMEOUT_SECONDS
timeout_call = connection.cursor_obj.calls[0]
assert "set_config('statement_timeout'" in timeout_call[0]
assert timeout_call[1] == ("45000",)
def test_postgres_introspection_is_filtered_bounded_and_cross_schema_safe(monkeypatch):
connection = _Connection()
monkeypatch.setattr(postgres_schema, "_connect", lambda dsn: connection)
snapshot = postgres_schema.PostgresSchemaIntrospector().inspect(
dsn, schemas=["public", "auth"]
)
assert connection.closed is True
assert dsn not in snapshot.text
assert dsn not in json.dumps(snapshot.metadata)
assert snapshot.metadata["source_digest"]
ids = {entity["id"] for entity in snapshot.entities}
assert "constraint:public.users.shared_name" in ids
assert "constraint:public.orders.shared_name" in ids
assert {
("table:public.orders", "table:auth.accounts", "references")
} <= {
(relation["source"], relation["target"], relation["relation"])
for relation in snapshot.relations
}
constraint_query, params = connection.cursor_obj.calls[-1]
assert "tc.constraint_schema=ccu.constraint_schema" in constraint_query
assert "tc.table_schema=ccu.table_schema" not in constraint_query
assert params[:2] == (["auth", "public"], ["auth", "public"])
def test_service_never_persists_postgres_dsn(monkeypatch):
snapshot = SchemaSnapshot(
title="PostgreSQL schema: appdb",
text="# PostgreSQL schema: appdb\n\n## public.users\n\n- `id`: integer not null",
entities=[
{"id": "database:appdb", "name": "appdb", "kind": "database"},
{"id": "table:public.users", "name": "public.users", "kind": "table"},
],
relations=[{
"source": "database:appdb",
"target": "table:public.users",
"relation": "contains",
}],
metadata={"database": "appdb", "tables": 1, "source_digest": "abc123"},
)
class _Introspector:
def inspect(self, supplied, *, schemas=None):
assert supplied == dsn
return snapshot
monkeypatch.setattr(
postgres_schema, "get_postgres_introspector", lambda: _Introspector()
)
service = MemoryService.create(":memory:")
result = service.import_postgres_schema(dsn, workspace="acme")
wid = service.store.get_or_create_workspace("acme")
memories = service.store.list_memories(
SearchFilter(workspace_id=wid), include_invalid=True
)
audit = service.store.conn.execute("SELECT * FROM audit").fetchall()
receipts = service.store.list_receipts(workspace_id=wid)
serialized = json.dumps({
"result": result,
"memories": [memory.content for memory in memories],
"metadata": [memory.metadata for memory in memories],
"audit": [dict(row) for row in audit],
"receipts": receipts,
}, default=str)
assert dsn not in serialized
assert "secret" not in serialized
def test_large_postgres_snapshot_keeps_every_chunk_distinct(monkeypatch):
snapshot = SchemaSnapshot(
title="PostgreSQL schema: large",
text="\n\n".join(
f"## public.table_{index}\n\n- `id`: integer not null"
for index in range(5_000)
),
metadata={"database": "large", "tables": 5_000, "source_digest": "digest"},
)
class _Introspector:
def inspect(self, supplied, *, schemas=None):
return snapshot
monkeypatch.setattr(
postgres_schema, "get_postgres_introspector", lambda: _Introspector()
)
service = MemoryService.create(":memory:", graph_extractor="none")
result = service.import_postgres_schema(
"postgresql://local/large", workspace="acme"
)
assert len(result["memory_ids"]) > 1
assert len(set(result["memory_ids"])) == len(result["memory_ids"])