Skip to content

Commit 6940527

Browse files
alxmrsclaude
alxmrs
and
claude
committed
Address review: SQL that replaces a table forgets its Dataset; skip comments
- CREATE OR REPLACE TABLE over a name registered from a Dataset forgets the Dataset, as deregister_table does, so a query prepared after it can't write into the old arrays. - Comments before a statement that creates or drops a table no longer hide it. Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent 6719421 commit 6940527

2 files changed

Lines changed: 22 additions & 7 deletions

File tree

‎tests/test_prepare.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -222,11 +222,21 @@ def test_now_is_the_time_a_query_over_a_lazy_table_was_prepared():
222222
assert query().equals(first)
223223

224224

225-
def test_a_table_replaced_in_sql_counts_too():
225+
@pytest.mark.parametrize("comment", ["", "-- refresh\n", "/* refresh */ "])
226+
def test_a_table_replaced_in_sql_counts_too(comment):
226227
ctx = Context()
227228
ctx.sql("CREATE TABLE c AS VALUES (1), (2)").collect()
228229
query = ctx.prepare("SELECT COUNT(*) AS n FROM c")
229230
assert query().column("n").to_pylist() == [2]
230-
ctx.sql("CREATE OR REPLACE TABLE c AS VALUES (1), (2), (3)").collect()
231+
ctx.sql(
232+
comment + "CREATE OR REPLACE TABLE c AS VALUES (1), (2), (3)"
233+
).collect()
231234
with pytest.raises(ValueError, match="registered again"):
232235
query()
236+
237+
238+
def test_a_dataset_replaced_in_sql_takes_no_writes(ctx):
239+
ctx.sql("CREATE OR REPLACE TABLE a AS SELECT * FROM b").collect()
240+
query = ctx.prepare("SELECT SUM(w) AS s FROM a")
241+
with pytest.raises(ValueError, match="no Dataset"):
242+
query(a=_a(np.ones((4, 3))))

‎xarray_sql/sql.py‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -252,15 +252,18 @@ def sql(self, query: str, *args, **kwargs) -> XarrayDataFrame:
252252
# plans it, replacing a table without deregistering it.
253253
replaced = _table_replaced(query)
254254
if replaced is not None:
255-
self._deregistrations[replaced] = (
256-
self._deregistrations.get(replaced, 0) + 1
257-
)
255+
self._forget(replaced)
258256
return self._wrap(super().sql(query, *args, **kwargs))
259257

260258
def deregister_table(self, name: str) -> None:
261259
"""Remove a table from the session, and forget the Dataset
262260
registered under its name."""
263261
super().deregister_table(name)
262+
self._forget(name)
263+
264+
def _forget(self, name: str) -> None:
265+
"""Forget the Dataset registered as ``name``, whose table has been
266+
removed or replaced."""
264267
self._registered_datasets.pop(name, None)
265268
self._deregistrations[name] = self._deregistrations.get(name, 0) + 1
266269

@@ -470,9 +473,11 @@ def _tables_read(plan: LogicalPlan) -> set[str]:
470473

471474

472475
_TABLE_DDL = re.compile(
473-
r"\s*(?:CREATE\s+(?:OR\s+REPLACE\s+)?(?:EXTERNAL\s+)?TABLE|DROP\s+TABLE)"
476+
# Any comments before the statement, then the statement.
477+
r"(?:\s+|--[^\n]*|/\*.*?\*/)*"
478+
r"(?:CREATE\s+(?:OR\s+REPLACE\s+)?(?:EXTERNAL\s+)?TABLE|DROP\s+TABLE)"
474479
r"\s+(?:IF\s+(?:NOT\s+)?EXISTS\s+)?(\"[^\"]+\"|[\w.]+)",
475-
re.IGNORECASE,
480+
re.IGNORECASE | re.DOTALL,
476481
)
477482

478483

0 commit comments

Comments
 (0)