Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: avoid failure if closing AppendRowsStream before opening #304

Merged
merged 3 commits into from Sep 20, 2021
Merged

fix: avoid failure if closing AppendRowsStream before opening #304

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
caf5fae
fix: avoid failure if closing AppendRowsStream before opening
tswast Sep 17, 2021
ef572b4
Merge branch 'main' into issue302-appendrowsstream-close
tswast Sep 17, 2021
ffdd126
don't create a thread if we're going to block anyway
tswast Sep 17, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
14 changes: 3 additions & 11 deletions google/cloud/bigquery_storage_v1beta2/writer.py
View file
Open in desktop
Expand Up @@ -259,21 +259,12 @@ def close(self, reason: Optional[Exception] = None):

This method is idempotent. Additional calls will have no effect.

The method does not block, it delegates the shutdown operations to a background
thread.

Args:
reason: The reason to close this. If ``None``, this is considered
an "intentional" shutdown. This is passed to the callbacks
specified via :meth:`add_close_callback`.
"""
self._regular_shutdown_thread = threading.Thread(
name=_REGULAR_SHUTDOWN_THREAD_NAME,
daemon=True,
target=self._shutdown,
kwargs={"reason": reason},
)
self._regular_shutdown_thread.start()
self._shutdown(reason=reason)

def _shutdown(self, reason: Optional[Exception] = None):
"""Run the actual shutdown sequence (stop the stream and all helper threads).
Expand All @@ -293,7 +284,8 @@ def _shutdown(self, reason: Optional[Exception] = None):
self._consumer.stop()
self._consumer = None

self._rpc.close()
if self._rpc is not None:
self._rpc.close()
self._rpc = None
self._closed = True
_LOGGER.debug("Finished stopping manager.")
Expand Down
9 changes: 9 additions & 0 deletions tests/unit/test_writer_v1beta2.py
View file
Open in desktop
Expand Up @@ -20,6 +20,7 @@
from google.api_core import exceptions
from google.cloud.bigquery_storage_v1beta2.services import big_query_write
from google.cloud.bigquery_storage_v1beta2 import types as gapic_types
from google.cloud.bigquery_storage_v1beta2 import exceptions as bqstorage_exceptions
from google.protobuf import descriptor_pb2


Expand All @@ -44,6 +45,14 @@ def test_constructor_and_default_state(module_under_test):
assert manager._client is mock_client


def test_close_before_open(module_under_test):
mock_client = mock.create_autospec(big_query_write.BigQueryWriteClient)
manager = module_under_test.AppendRowsStream(mock_client, REQUEST_TEMPLATE)
manager.close()
with pytest.raises(bqstorage_exceptions.StreamClosedError):
manager.send(object())


@mock.patch("google.api_core.bidi.BidiRpc", autospec=True)
@mock.patch("google.api_core.bidi.BackgroundConsumer", autospec=True)
def test_open(background_consumer, bidi_rpc, module_under_test):
Expand Down