feat: add location to WriteStream and add WriteStreamView support (#507)

* chore: use gapic-generator-python 1.3.1

PiperOrigin-RevId: 472772457

Source-Link: googleapis/googleapis@855b74d

Source-Link: googleapis/googleapis-gen@b64b1e7
Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiYjY0YjFlN2RhM2UxMzhmMTVjYTM2MTU1MmVmMDU0NWU1NDg5MWI0ZiJ9

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* feat: add location to WriteStream and add WriteStreamView support

PiperOrigin-RevId: 472835596

Source-Link: googleapis/googleapis@1cf9407

Source-Link: googleapis/googleapis-gen@59d0f10
Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiNTlkMGYxMGJhZGVjODY3OTEzZWJjMmViNGJhYzAyMWY3OTkwNDUxOSJ9

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: gcf-owl-bot[bot]

packages/google-cloud-bigquery-storage/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",
 )

packages/google-cloud-bigquery-storage/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",
 )

packages/google-cloud-bigquery-storage/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):

packages/google-cloud-bigquery-storage/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))

packages/google-cloud-bigquery-storage/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', ),
     }

packages/google-cloud-bigquery-storage/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

packages/google-cloud-bigquery-storage/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

packages/google-cloud-bigquery-storage/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

packages/google-cloud-bigquery-storage/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