Skip to content

Commit 21a822a

Browse files
committed
fix python prepare params implicit cause RSS raising
Signed-off-by: ericyuanhui <285521263@qq.com>
1 parent 2a28669 commit 21a822a

2 files changed

Lines changed: 169 additions & 1 deletion

File tree

src_py/connection.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ def __init__(self, database: Database, num_threads: int = 0):
5252
self._query_timeout_ms = 0
5353
self._query_results: WeakSet[QueryResult] = WeakSet()
5454
self._capi_scan_tables: set[str] = set()
55+
self._pybind_implicit_prepared_cache: dict[str, Any] = {}
5556
self.database._register_connection(self)
5657
self.init_connection()
5758

@@ -116,6 +117,7 @@ def close(self) -> None:
116117
for query_result in list(self._query_results):
117118
query_result.close()
118119
self._query_results.clear()
120+
self._pybind_implicit_prepared_cache.clear()
119121

120122
if self._connection is not None and not self.database.is_closed:
121123
self._connection.close()
@@ -460,9 +462,23 @@ def _execute_with_pybind(
460462
return py_connection.query(query)
461463

462464
query, parameters = self._normalize_parameters_for_pybind(query, parameters)
463-
prepared = py_connection.prepare(query, parameters)
465+
prepared = self._get_or_prepare_pybind_statement(
466+
py_connection, query, parameters
467+
)
464468
return py_connection.execute(prepared, parameters)
465469

470+
def _get_or_prepare_pybind_statement(
471+
self,
472+
py_connection: Any,
473+
query: str,
474+
parameters: dict[str, Any],
475+
) -> Any:
476+
prepared = self._pybind_implicit_prepared_cache.get(query)
477+
if prepared is None:
478+
prepared = py_connection.prepare(query, parameters)
479+
self._pybind_implicit_prepared_cache[query] = prepared
480+
return prepared
481+
466482
def _maybe_raise_scan_unsupported_object(self, query: str) -> None:
467483
match = re.search(
468484
r"\bLOAD\s+FROM\s+([A-Za-z_][A-Za-z0-9_]*)\b", query, re.IGNORECASE
Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
from __future__ import annotations
2+
3+
from types import SimpleNamespace
4+
5+
import pytest
6+
7+
import ladybug as lb
8+
import ladybug.connection as lb_connection
9+
10+
11+
class _FakeResult:
12+
def isSuccess(self) -> bool:
13+
return True
14+
15+
def hasNextQueryResult(self) -> bool:
16+
return False
17+
18+
19+
class _FakePreparedStatement:
20+
def __init__(self, query: str, parameters: dict[str, object]):
21+
self.query = query
22+
self.parameters = dict(parameters)
23+
24+
25+
class _FakePybindConnection:
26+
def __init__(self) -> None:
27+
self.prepare_calls: list[tuple[str, dict[str, object]]] = []
28+
self.execute_calls: list[tuple[_FakePreparedStatement, dict[str, object]]] = []
29+
self.query_calls: list[str] = []
30+
self.closed = False
31+
32+
def prepare(
33+
self, query: str, parameters: dict[str, object]
34+
) -> _FakePreparedStatement:
35+
self.prepare_calls.append((query, dict(parameters)))
36+
return _FakePreparedStatement(query, parameters)
37+
38+
def execute(
39+
self, prepared: _FakePreparedStatement, parameters: dict[str, object]
40+
) -> _FakeResult:
41+
self.execute_calls.append((prepared, dict(parameters)))
42+
return _FakeResult()
43+
44+
def query(self, query: str) -> _FakeResult:
45+
self.query_calls.append(query)
46+
return _FakeResult()
47+
48+
def close(self) -> None:
49+
self.closed = True
50+
51+
52+
class _FakeBackendConnection:
53+
def __init__(self) -> None:
54+
self.closed = False
55+
56+
def close(self) -> None:
57+
self.closed = True
58+
59+
60+
class _FakeDatabase:
61+
def __init__(self) -> None:
62+
self._use_pybind_backend = True
63+
self._database = object()
64+
self.is_closed = False
65+
self.registered_connections: list[lb.Connection] = []
66+
67+
def _register_connection(self, connection: lb.Connection) -> None:
68+
self.registered_connections.append(connection)
69+
70+
def _unregister_connection(self, connection: lb.Connection) -> None:
71+
self.registered_connections.remove(connection)
72+
73+
def init_database(self) -> None:
74+
return None
75+
76+
77+
@pytest.fixture
78+
def fake_pybind_connection(monkeypatch: pytest.MonkeyPatch) -> _FakePybindConnection:
79+
fake_pybind = _FakePybindConnection()
80+
81+
monkeypatch.setattr(
82+
lb_connection, "get_pybind_module", lambda: SimpleNamespace()
83+
)
84+
monkeypatch.setattr(
85+
lb.Connection,
86+
"init_connection",
87+
lambda self: setattr(self, "_connection", _FakeBackendConnection()),
88+
)
89+
monkeypatch.setattr(
90+
lb.Connection, "_get_pybind_connection", lambda self: fake_pybind
91+
)
92+
return fake_pybind
93+
94+
95+
def test_pybind_implicit_prepare_reuses_same_query(
96+
fake_pybind_connection: _FakePybindConnection,
97+
) -> None:
98+
conn = lb.Connection(_FakeDatabase())
99+
100+
conn.execute("RETURN $value", {"value": 1})
101+
conn.execute("RETURN $value", {"value": 2})
102+
103+
assert fake_pybind_connection.query_calls == []
104+
assert len(fake_pybind_connection.prepare_calls) == 1
105+
assert [call[0].query for call in fake_pybind_connection.execute_calls] == [
106+
"RETURN $value",
107+
"RETURN $value",
108+
]
109+
assert [call[1] for call in fake_pybind_connection.execute_calls] == [
110+
{"value": 1},
111+
{"value": 2},
112+
]
113+
114+
115+
def test_pybind_implicit_prepare_does_not_share_different_queries(
116+
fake_pybind_connection: _FakePybindConnection,
117+
) -> None:
118+
conn = lb.Connection(_FakeDatabase())
119+
120+
conn.execute("RETURN $value", {"value": 1})
121+
conn.execute("RETURN $other", {"other": 1})
122+
123+
assert [call[0] for call in fake_pybind_connection.prepare_calls] == [
124+
"RETURN $value",
125+
"RETURN $other",
126+
]
127+
128+
129+
def test_pybind_no_parameter_query_skips_prepare_cache(
130+
fake_pybind_connection: _FakePybindConnection,
131+
) -> None:
132+
conn = lb.Connection(_FakeDatabase())
133+
134+
conn.execute("RETURN 1")
135+
136+
assert fake_pybind_connection.prepare_calls == []
137+
assert fake_pybind_connection.query_calls == ["RETURN 1"]
138+
139+
140+
def test_pybind_close_clears_implicit_prepare_cache(
141+
fake_pybind_connection: _FakePybindConnection,
142+
) -> None:
143+
conn = lb.Connection(_FakeDatabase())
144+
145+
conn.execute("RETURN $value", {"value": 1})
146+
147+
assert set(conn._pybind_implicit_prepared_cache) == {"RETURN $value"}
148+
149+
conn.close()
150+
151+
assert conn._pybind_implicit_prepared_cache == {}
152+
assert fake_pybind_connection.closed is True

0 commit comments

Comments
 (0)