diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index c50edb98b184..4f88b406161e 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -374,7 +374,6 @@ def chain_after(result): import fastavro from objsize import get_deep_size - import apache_beam as beam from apache_beam import coders from apache_beam import pvalue @@ -422,7 +421,7 @@ def chain_after(result): from apache_beam.typehints.schemas import schema_from_element_type from apache_beam.utils import retry from apache_beam.utils.annotations import deprecated - +from apache_beam.io.gcp.internal.clients.bigquery.bigquery_v2_messages import BigqueryTablesGetRequest try: from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference from apache_beam.io.gcp.internal.clients.bigquery import TableReference @@ -926,7 +925,10 @@ def _export_files(self, bq): else: table_ref = self.table_reference table = bq.get_table( - table_ref.projectId, table_ref.datasetId, table_ref.tableId) + table_ref.projectId, + table_ref.datasetId, + table_ref.tableId, + view=BigqueryTablesGetRequest.ViewValueValuesEnum.BASIC) return table.schema, metadata_list diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py b/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py index 0b4cfe2ecbae..897aecfee275 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py @@ -55,7 +55,7 @@ from apache_beam import Map from apache_beam import ParDo -from apache_beam.io import BigQueryDisposition +from apache_beam.io.gcp.bigquery import BigQueryDisposition from apache_beam.io import BigQuerySource from apache_beam.io import Read from apache_beam.io import WriteToBigQuery diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 081571bfef99..60dd13ca78f6 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -66,6 +66,7 @@ from apache_beam.typehints.typehints import Any from apache_beam.utils import retry from apache_beam.utils.histogram import LinearBucket +from apache_beam.io.gcp.internal.clients.bigquery.bigquery_v2_messages import BigqueryTablesGetRequest # Protect against environments where bigquery library is not available. # pylint: disable=wrong-import-order, wrong-import-position @@ -772,7 +773,12 @@ def _insert_all_rows( @retry.with_exponential_backoff( num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter) - def get_table(self, project_id, dataset_id, table_id): + def get_table( + self, + project_id, + dataset_id, + table_id, + view=BigqueryTablesGetRequest.ViewValueValuesEnum.BASIC): """Lookup a table's metadata object. Args: @@ -780,6 +786,7 @@ def get_table(self, project_id, dataset_id, table_id): project_id: table lookup parameter dataset_id: table lookup parameter table_id: table lookup parameter + view: table information parameter Returns: bigquery.Table instance @@ -787,7 +794,7 @@ def get_table(self, project_id, dataset_id, table_id): HttpError: if lookup failed. """ request = bigquery.BigqueryTablesGetRequest( - projectId=project_id, datasetId=dataset_id, tableId=table_id) + projectId=project_id, datasetId=dataset_id, tableId=table_id, view=view) response = self.client.tables.Get(request) return response diff --git a/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py b/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py index 1aafb1b60a85..a2eabb49302f 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py @@ -55,7 +55,7 @@ from apache_beam import Map from apache_beam import ParDo -from apache_beam.io import BigQueryDisposition +from apache_beam.io.gcp.bigquery import BigQueryDisposition from apache_beam.io import Read from apache_beam.io import WriteToBigQuery from apache_beam.io.gcp.bigquery_tools import parse_table_schema_from_json diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py index fd0bc9694071..268f108883e7 100644 --- a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py +++ b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py @@ -935,12 +935,29 @@ class BigqueryTablesGetRequest(_messages.Message): selectedFields: List of fields to return (comma-separated). If unspecified, all fields are returned tableId: Table ID of the requested table + view: table information parameter """ datasetId = _messages.StringField(1, required=True) projectId = _messages.StringField(2, required=True) selectedFields = _messages.StringField(3) tableId = _messages.StringField(4, required=True) + view = _messages.EnumField('ViewValueValuesEnum', 5, default='BASIC') + # @@protoc_insertion_point(class_scope:google.cloud.bigquery_v2.BigqueryTablesGetRequest) + # @@protoc_insertion_point(enum_scope:google.cloud.bigquery_v2.BigqueryTablesGetRequest.ViewValueValuesEnum) + class ViewValueValuesEnum(_messages.Enum): + r"""table information parameter + + Values: + TABLE_METADATA_VIEW_UNSPECIFIED: The default value. Default to the STORAGE_STATS view. + BASIC: Includes basic table information, such as schema and partitioning specification. This view does not include storage statistics like `numRows` or `numBytes`. It is significantly more efficient and suitable for high query rates. + STORAGE_STATS: Includes all information from the BASIC view, along with additional storage statistics such as numBytes, numLongTermBytes, numRows, and lastModifiedTime. + FULL: Includes all table information, including storage statistics. It returns same information as STORAGE_STATS view, but may contain additional information in the future. + """ + TABLE_METADATA_VIEW_UNSPECIFIED = 0 + BASIC = 1 + STORAGE_STATS = 2 + FULL = 3 class BigqueryTablesInsertRequest(_messages.Message): diff --git a/sdks/python/apache_beam/yaml/yaml_io.py b/sdks/python/apache_beam/yaml/yaml_io.py index 67a5e65bda93..ed2f38e4fff2 100644 --- a/sdks/python/apache_beam/yaml/yaml_io.py +++ b/sdks/python/apache_beam/yaml/yaml_io.py @@ -35,7 +35,7 @@ import apache_beam as beam import apache_beam.io as beam_io from apache_beam import coders -from apache_beam.io import ReadFromBigQuery +from apache_beam.io.gcp.bigquery import ReadFromBigQuery from apache_beam.io import ReadFromTFRecord from apache_beam.io import WriteToBigQuery from apache_beam.io import WriteToTFRecord