Skip to content

Commit fdcdc18

Browse files
authored
fix: reduce the number of labels added to query jobs (#1245)
* fix: reduce the number of labels added to query jobs Fixes internal issue 386825477 * fix unit tests
1 parent 4cf8df3 commit fdcdc18

File tree

4 files changed

+23
-7
lines changed

4 files changed

+23
-7
lines changed

‎bigframes/core/log_adapter.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,12 @@
1717
from typing import List
1818

1919
_lock = threading.Lock()
20-
MAX_LABELS_COUNT = 64
20+
21+
# The limit is 64 (https://cloud.google.com/bigquery/docs/labels-intro#requirements),
22+
# but leave a few spare for internal labels to be added.
23+
# See internal issue 386825477.
24+
MAX_LABELS_COUNT = 64 - 8
25+
2126
_api_methods: List = []
2227
_excluded_methods = ["__setattr__", "__getattr__"]
2328

‎bigframes/session/__init__.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -373,10 +373,18 @@ def close(self):
373373
"""Delete resources that were created with this session's session_id.
374374
This includes BigQuery tables, remote functions and cloud functions
375375
serving the remote functions."""
376-
self._temp_storage_manager.clean_up_tables()
377-
self._remote_function_session.clean_up(
378-
self.bqclient, self.cloudfunctionsclient, self.session_id
379-
)
376+
377+
# Protect against failure when the Session is a fake for testing or
378+
# failed to initialize.
379+
temp_storage_manager = getattr(self, "_temp_storage_manager", None)
380+
if temp_storage_manager:
381+
self._temp_storage_manager.clean_up_tables()
382+
383+
remote_function_session = getattr(self, "_remote_function_session", None)
384+
if remote_function_session:
385+
self._remote_function_session.clean_up(
386+
self.bqclient, self.cloudfunctionsclient, self.session_id
387+
)
380388

381389
def read_gbq(
382390
self,

‎tests/unit/core/test_log_adapter.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,10 @@
1616

1717
from bigframes.core import log_adapter
1818

19-
MAX_LABELS_COUNT = 64
19+
# The limit is 64 (https://cloud.google.com/bigquery/docs/labels-intro#requirements),
20+
# but leave a few spare for internal labels to be added.
21+
# See internal issue 386825477.
22+
MAX_LABELS_COUNT = 56
2023

2124

2225
@pytest.fixture

‎tests/unit/session/test_io_bigquery.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ def test_create_job_configs_labels_length_limit_met_and_labels_is_none():
114114
job_configs_labels=None, api_methods=api_methods
115115
)
116116
assert labels is not None
117-
assert len(labels) == 64
117+
assert len(labels) == log_adapter.MAX_LABELS_COUNT
118118
assert "dataframe-head" in labels.values()
119119

120120

0 commit comments

Comments
 (0)