From a81e4328e8baeb9f4138a811c60f9b4c2bea7614 Mon Sep 17 00:00:00 2001
From: Gil Forsyth <gil@forsyth.dev>
Date: Wed, 29 Nov 2023 10:21:02 -0500
Subject: [PATCH] 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.
---
 ibis/backends/bigquery/__init__.py               |  8 ++++++--
 .../bigquery/tests/system/test_client.py         | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/ibis/backends/bigquery/__init__.py b/ibis/backends/bigquery/__init__.py
index e8f180e46dee..84909085ed11 100644
--- a/ibis/backends/bigquery/__init__.py
+++ b/ibis/backends/bigquery/__init__.py
@@ -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()
@@ -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()
 
diff --git a/ibis/backends/bigquery/tests/system/test_client.py b/ibis/backends/bigquery/tests/system/test_client.py
index aa2e5d1db4bc..3d033f89c2c1 100644
--- a/ibis/backends/bigquery/tests/system/test_client.py
+++ b/ibis/backends/bigquery/tests/system/test_client.py
@@ -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"))