Skip to content

Commit

Permalink
fix(bigquery): fully qualified memtable names in compile
Browse files Browse the repository at this point in the history
Use `_make_session` to ensure that the `dataset_id` and `project`
attributes in `self._session_dataset` are set before we compile to SQL.
  • Loading branch information
gforsyth authored and cpcloud committed Nov 30, 2023
1 parent 8a2662b commit a81e432
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 2 deletions.
8 changes: 6 additions & 2 deletions ibis/backends/bigquery/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -558,9 +558,12 @@ def table(
return rename_partitioned_column(table_expr, bq_table, self.partition_column)

def _make_session(self) -> tuple[str, str]:
if self._session_dataset is None:
if (
self._session_dataset is None
and (client := getattr(self, "client", None)) is not None
):
job_config = bq.QueryJobConfig(use_query_cache=False)
query = self.client.query(
query = client.query(
"SELECT 1", job_config=job_config, project=self.billing_project
)
query.result()
Expand Down Expand Up @@ -615,6 +618,7 @@ def compile(
The output of compilation. The type of this value depends on the
backend.
"""
self._make_session()
self._define_udf_translation_rules(expr)
sql = self.compiler.to_ast_ensure_limit(expr, limit, params=params).compile()

Expand Down
16 changes: 16 additions & 0 deletions ibis/backends/bigquery/tests/system/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,22 @@ def test_fully_qualified_table_creation(con, project_id, dataset_id, temp_table)
assert t.get_name() == f"{project_id}.{dataset_id}.{temp_table}"


def test_fully_qualified_memtable_compile(project_id, dataset_id):
new_bq_con = ibis.bigquery.connect(project_id=project_id, dataset_id=dataset_id)
# New connection shouldn't have _session_dataset populated after connection
assert new_bq_con._session_dataset is None

t = ibis.memtable(
{"a": [1, 2, 3], "b": [4, 5, 6]},
schema=ibis.schema({"a": "int64", "b": "int64"}),
)

# call to compile should fill in _session_dataset
sql = new_bq_con.compile(t)
assert new_bq_con._session_dataset is not None
assert project_id in sql


def test_create_table_with_options(con):
name = gen_name("bigquery_temp_table")
schema = ibis.schema(dict(a="int64", b="int64", c="array<string>", d="date"))
Expand Down

0 comments on commit a81e432

Please sign in to comment.