feat: add location to WriteStream and add WriteStreamView support

google/cloud/bigquery_storage/__init__.py

@@ -59,6 +59,7 @@
 from google.cloud.bigquery_storage_v1.types.stream import ReadStream
 from google.cloud.bigquery_storage_v1.types.stream import WriteStream
 from google.cloud.bigquery_storage_v1.types.stream import DataFormat
+from google.cloud.bigquery_storage_v1.types.stream import WriteStreamView
 from google.cloud.bigquery_storage_v1.types.table import TableFieldSchema
 from google.cloud.bigquery_storage_v1.types.table import TableSchema
 
@@ -99,6 +100,7 @@
  "ReadStream",
  "WriteStream",
  "DataFormat",
+ "WriteStreamView",
  "TableFieldSchema",
  "TableSchema",
 )

google/cloud/bigquery_storage_v1/types/__init__.py

@@ -54,6 +54,7 @@
  ReadStream,
  WriteStream,
  DataFormat,
+ WriteStreamView,
 )
 from .table import (
  TableFieldSchema,
@@ -93,6 +94,7 @@
  "ReadStream",
  "WriteStream",
  "DataFormat",
+ "WriteStreamView",
  "TableFieldSchema",
  "TableSchema",
 )

google/cloud/bigquery_storage_v1/types/storage.py

@@ -554,12 +554,21 @@ class GetWriteStreamRequest(proto.Message):
  name (str):
  Required. Name of the stream to get, in the form of
  ``projects/{project}/datasets/{dataset}/tables/{table}/streams/{stream}``.
+ view (google.cloud.bigquery_storage_v1.types.WriteStreamView):
+ Indicates whether to get full or partial view
+ of the WriteStream. If not set, view returned
+ will be basic.
  """
 
  name = proto.Field(
  proto.STRING,
  number=1,
  )
+ view = proto.Field(
+ proto.ENUM,
+ number=3,
+ enum=stream.WriteStreamView,
+ )
 
 
 class BatchCommitWriteStreamsRequest(proto.Message):

google/cloud/bigquery_storage_v1/types/stream.py

@@ -25,6 +25,7 @@
  package="google.cloud.bigquery.storage.v1",
  manifest={
  "DataFormat",
+ "WriteStreamView",
  "ReadSession",
  "ReadStream",
  "WriteStream",
@@ -39,6 +40,15 @@ class DataFormat(proto.Enum):
  ARROW = 2
 
 
+class WriteStreamView(proto.Enum):
+ r"""WriteStreamView is a view enum that controls what details
+ about a write stream should be returned.
+ """
+ WRITE_STREAM_VIEW_UNSPECIFIED = 0
+ BASIC = 1
+ FULL = 2
+
+
 class ReadSession(proto.Message):
  r"""Information about the ReadSession.
 
@@ -305,6 +315,11 @@ class WriteStream(proto.Message):
  of date during the life time of the stream.
  write_mode (google.cloud.bigquery_storage_v1.types.WriteStream.WriteMode):
  Immutable. Mode of the stream.
+ location (str):
+ Immutable. The geographic location where the
+ stream's dataset resides. See
+ https://cloud.google.com/bigquery/docs/locations
+ for supported locations.
  """
 
  class Type(proto.Enum):
@@ -348,6 +363,10 @@ class WriteMode(proto.Enum):
  number=7,
  enum=WriteMode,
  )
+ location = proto.Field(
+ proto.STRING,
+ number=8,
+ )
 
 
 __all__ = tuple(sorted(__protobuf__.manifest))

scripts/fixup_bigquery_storage_v1_keywords.py

@@ -45,7 +45,7 @@ class bigquery_storageCallTransformer(cst.CSTTransformer):
  'create_write_stream': ('parent', 'write_stream', ),
  'finalize_write_stream': ('name', ),
  'flush_rows': ('write_stream', 'offset', ),
- 'get_write_stream': ('name', ),
+ 'get_write_stream': ('name', 'view', ),
  'read_rows': ('read_stream', 'offset', ),
  'split_read_stream': ('name', 'fraction', ),
  }

tests/unit/gapic/bigquery_storage_v1/test_big_query_read.py

@@ -18,8 +18,8 @@
 # try/except added for compatibility with python < 3.8
 try:
  from unittest import mock
- from unittest.mock import AsyncMock
-except ImportError:
+ from unittest.mock import AsyncMock # pragma: NO COVER
+except ImportError: # pragma: NO COVER
  import mock
 
 import grpc

tests/unit/gapic/bigquery_storage_v1/test_big_query_write.py

@@ -18,8 +18,8 @@
 # try/except added for compatibility with python < 3.8
 try:
  from unittest import mock
- from unittest.mock import AsyncMock
-except ImportError:
+ from unittest.mock import AsyncMock # pragma: NO COVER
+except ImportError: # pragma: NO COVER
  import mock
 
 import grpc
@@ -720,6 +720,7 @@ def test_create_write_stream(request_type, transport: str = "grpc"):
  name="name_value",
  type_=stream.WriteStream.Type.COMMITTED,
  write_mode=stream.WriteStream.WriteMode.INSERT,
+ location="location_value",
  )
  response = client.create_write_stream(request)
 
@@ -733,6 +734,7 @@ def test_create_write_stream(request_type, transport: str = "grpc"):
  assert response.name == "name_value"
  assert response.type_ == stream.WriteStream.Type.COMMITTED
  assert response.write_mode == stream.WriteStream.WriteMode.INSERT
+ assert response.location == "location_value"
 
 
 def test_create_write_stream_empty_call():
@@ -776,6 +778,7 @@ async def test_create_write_stream_async(
  name="name_value",
  type_=stream.WriteStream.Type.COMMITTED,
  write_mode=stream.WriteStream.WriteMode.INSERT,
+ location="location_value",
  )
  )
  response = await client.create_write_stream(request)
@@ -790,6 +793,7 @@ async def test_create_write_stream_async(
  assert response.name == "name_value"
  assert response.type_ == stream.WriteStream.Type.COMMITTED
  assert response.write_mode == stream.WriteStream.WriteMode.INSERT
+ assert response.location == "location_value"
 
 
 @pytest.mark.asyncio
@@ -1050,6 +1054,7 @@ def test_get_write_stream(request_type, transport: str = "grpc"):
  name="name_value",
  type_=stream.WriteStream.Type.COMMITTED,
  write_mode=stream.WriteStream.WriteMode.INSERT,
+ location="location_value",
  )
  response = client.get_write_stream(request)
 
@@ -1063,6 +1068,7 @@ def test_get_write_stream(request_type, transport: str = "grpc"):
  assert response.name == "name_value"
  assert response.type_ == stream.WriteStream.Type.COMMITTED
  assert response.write_mode == stream.WriteStream.WriteMode.INSERT
+ assert response.location == "location_value"
 
 
 def test_get_write_stream_empty_call():
@@ -1102,6 +1108,7 @@ async def test_get_write_stream_async(
  name="name_value",
  type_=stream.WriteStream.Type.COMMITTED,
  write_mode=stream.WriteStream.WriteMode.INSERT,
+ location="location_value",
  )
  )
  response = await client.get_write_stream(request)
@@ -1116,6 +1123,7 @@ async def test_get_write_stream_async(
  assert response.name == "name_value"
  assert response.type_ == stream.WriteStream.Type.COMMITTED
  assert response.write_mode == stream.WriteStream.WriteMode.INSERT
+ assert response.location == "location_value"
 
 
 @pytest.mark.asyncio

tests/unit/gapic/bigquery_storage_v1beta2/test_big_query_read.py

@@ -18,8 +18,8 @@
 # try/except added for compatibility with python < 3.8
 try:
  from unittest import mock
- from unittest.mock import AsyncMock
-except ImportError:
+ from unittest.mock import AsyncMock # pragma: NO COVER
+except ImportError: # pragma: NO COVER
  import mock
 
 import grpc

tests/unit/gapic/bigquery_storage_v1beta2/test_big_query_write.py

@@ -18,8 +18,8 @@
 # try/except added for compatibility with python < 3.8
 try:
  from unittest import mock
- from unittest.mock import AsyncMock
-except ImportError:
+ from unittest.mock import AsyncMock # pragma: NO COVER
+except ImportError: # pragma: NO COVER
  import mock
 
 import grpc