diff --git a/google/cloud/bigquery_storage/__init__.py b/google/cloud/bigquery_storage/__init__.py index 54a40387..0132dbae 100644 --- a/google/cloud/bigquery_storage/__init__.py +++ b/google/cloud/bigquery_storage/__init__.py @@ -20,6 +20,7 @@ from google.cloud.bigquery_storage_v1 import BigQueryReadClient from google.cloud.bigquery_storage_v1 import gapic_types as types +from google.cloud.bigquery_storage_v1.reader import ReadRowsStream from google.cloud.bigquery_storage_v1.services.big_query_write.async_client import ( BigQueryWriteAsyncClient, ) @@ -66,6 +67,7 @@ WriteStreamView, ) from google.cloud.bigquery_storage_v1.types.table import TableFieldSchema, TableSchema +from google.cloud.bigquery_storage_v1.writer import AppendRowsStream __all__ = ( "BigQueryReadClient", @@ -101,6 +103,8 @@ "StreamStats", "ThrottleState", "ReadSession", + "AppendRowsStream", + "ReadRowsStream", "ReadStream", "WriteStream", "DataFormat", diff --git a/owlbot.py b/owlbot.py index b05a7066..cb3f05b8 100644 --- a/owlbot.py +++ b/owlbot.py @@ -87,6 +87,26 @@ '"DataFormat",\n \\g<0>', ) + # Expose handwritten classes AppendRowsStream and ReadRowsStream here. + s.replace( + library / "google/cloud/bigquery_storage/__init__.py", + f"from google.cloud.bigquery_storage_{library.name} import BigQueryReadClient", + ( + f"from google.cloud.bigquery_storage_{library.name}.writer import AppendRowsStream\n" + "\\g<0>" + f"from google.cloud.bigquery_storage_{library.name}.reader import ReadRowsStream\n" + ), + ) + s.replace( + library / f"google/cloud/bigquery_storage_{library.name}*/types/__init__.py", + r"""["']ReadSession["']""", + ( + '"AppendRowsStream",\n' + ' "ReadRowsStream",\n' + ' \\g<0>' + ), + ) + s.move( [library], excludes=[