feat: bigframes.options and bigframes.option_context now uses thread-local variables to prevent context managers in separate threads from affecting each other
#652
Conversation
…ead-local variables
to prevent context managers in separate threads from affecting each other
In our tests, this allows us to actually test things like
`bf.option_context("display.repr_mode", "deferred"):` without always
having some other test change the display mode and break the test.
Fixes internal issue 308657813
product-auto-label
bot
added
the
api: bigquery
bigframes/_config/__init__.py
Outdated
| self._compute_options = compute_options.ComputeOptions() | ||
| self._local.bigquery_options = None | ||
|
|
||
| def _init_biquery_thread_local(self): |
There was a problem hiding this comment.
Thanks. Fixed in 4fa1a92
|
e2e failure of |
|
Cover failed.
Looks like I now have some dead code in that test. Will investigate. |
|
Messaged offline. @tswast will fix the failing test, and then I will review. |
|
|
|
I'm able to reproduce locally. The thread local variable is working how I expect, but the test is making an assumption that the DataFrame doesn't have a I'll update the test. |
bigframes/_config/__init__.py
Outdated
| # BigQuery options are special because they can only be set once per | ||
| # session, so we need an indicator as to whether we are using the | ||
| # thread-local session or the global session. | ||
| self._local.is_bigquery_thread_local = False |
There was a problem hiding this comment.
If I'm reading this correctly, would checking for self._local.bigquery_options is None be able to replace this boolean state? It seems a little more DRY to do it that way, if so.
There was a problem hiding this comment.
I like it! Done in a7f3c96
bigframes/_config/__init__.py
Outdated
| # Already thread-local, so don't reset any options that have been set | ||
| # already. No locks needed since this only modifies thread-local | ||
| # variables. | ||
| if self._local.is_bigquery_thread_local: |
There was a problem hiding this comment.
i.e., this could be "if self._local.bigquery_options is not None"
bigframes/_config/__init__.py
Outdated
|
|
||
| @property | ||
| def bigquery(self) -> bigquery_options.BigQueryOptions: | ||
| """Options to use with the BigQuery engine.""" | ||
| if self._local.is_bigquery_thread_local: |
bigframes/_config/__init__.py
Outdated
|
|
||
| This is set to True by using `option_context` with a "bigquery" option. | ||
| """ | ||
| return self._local.is_bigquery_thread_local |
There was a problem hiding this comment.
This could also check none (and the comment would need to be tweaked.)
| # sessions if we allow that. | ||
| if bigframes.options.is_bigquery_thread_local: | ||
| bigframes.close_session() | ||
| bigframes.options._local.is_bigquery_thread_local = False |
There was a problem hiding this comment.
close_session could set _local.bigquery_options to None
There was a problem hiding this comment.
Someone could call close_session() inside of a context manager. We actually have some tests that start and stop multiple sessions in this thread-local context. By using the options as the thread locality indicator, we need to make sure that object stays around.
There was a problem hiding this comment.
Oh, right, makes sense.
In our tests, this allows us to actually test things like
bf.option_context("display.repr_mode", "deferred"):without always having some other test change the display mode and break the test.Fixes internal issue 308657813
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
Fixes #<issue_number_goes_here> 🦕