Skip to content

Commit a81e432

Browse files
gforsythcpcloud
authored andcommitted
fix(bigquery): fully qualified memtable names in compile
Use `_make_session` to ensure that the `dataset_id` and `project` attributes in `self._session_dataset` are set before we compile to SQL.
1 parent 8a2662b commit a81e432

File tree

2 files changed

+22
-2
lines changed

2 files changed

+22
-2
lines changed

ibis/backends/bigquery/__init__.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -558,9 +558,12 @@ def table(
558558
return rename_partitioned_column(table_expr, bq_table, self.partition_column)
559559

560560
def _make_session(self) -> tuple[str, str]:
561-
if self._session_dataset is None:
561+
if (
562+
self._session_dataset is None
563+
and (client := getattr(self, "client", None)) is not None
564+
):
562565
job_config = bq.QueryJobConfig(use_query_cache=False)
563-
query = self.client.query(
566+
query = client.query(
564567
"SELECT 1", job_config=job_config, project=self.billing_project
565568
)
566569
query.result()
@@ -615,6 +618,7 @@ def compile(
615618
The output of compilation. The type of this value depends on the
616619
backend.
617620
"""
621+
self._make_session()
618622
self._define_udf_translation_rules(expr)
619623
sql = self.compiler.to_ast_ensure_limit(expr, limit, params=params).compile()
620624

ibis/backends/bigquery/tests/system/test_client.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,22 @@ def test_fully_qualified_table_creation(con, project_id, dataset_id, temp_table)
389389
assert t.get_name() == f"{project_id}.{dataset_id}.{temp_table}"
390390

391391

392+
def test_fully_qualified_memtable_compile(project_id, dataset_id):
393+
new_bq_con = ibis.bigquery.connect(project_id=project_id, dataset_id=dataset_id)
394+
# New connection shouldn't have _session_dataset populated after connection
395+
assert new_bq_con._session_dataset is None
396+
397+
t = ibis.memtable(
398+
{"a": [1, 2, 3], "b": [4, 5, 6]},
399+
schema=ibis.schema({"a": "int64", "b": "int64"}),
400+
)
401+
402+
# call to compile should fill in _session_dataset
403+
sql = new_bq_con.compile(t)
404+
assert new_bq_con._session_dataset is not None
405+
assert project_id in sql
406+
407+
392408
def test_create_table_with_options(con):
393409
name = gen_name("bigquery_temp_table")
394410
schema = ibis.schema(dict(a="int64", b="int64", c="array<string>", d="date"))

0 commit comments

Comments
 (0)