| # | |
| # Licensed to the Apache Software Foundation (ASF) under one or more | |
| # contributor license agreements. See the NOTICE file distributed with | |
| # this work for additional information regarding copyright ownership. | |
| # The ASF licenses this file to You under the Apache License, Version 2.0 | |
| # (the "License"); you may not use this file except in compliance with | |
| # the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| """Bigtable connector | |
| This module implements writing to Bigtable tables. | |
| The default mode is to set row data to write to Bigtable tables. | |
| The syntax supported is described here: | |
| https://cloud.google.com/bigtable/docs/quickstart-cbt | |
| Bigtable connector can be used as main outputs. A main output | |
| (common case) is expected to be massive and will be split into | |
| manageable chunks and processed in parallel. In the example below | |
| we created a list of rows then passed to the GeneratedDirectRows | |
| DoFn to set the Cells and then we call the _BigTableWriteFn to insert | |
| those generated rows in the table. | |
| main_table = (p | |
| | beam.Create(self._generate()) | |
| | WriteToBigTable(project_id, | |
| instance_id, | |
| table_id)) | |
| """ | |
| # pytype: skip-file | |
| import logging | |
| import struct | |
| import apache_beam as beam | |
| from apache_beam import version as beam_version | |
| from apache_beam.internal.metrics.metric import ServiceCallMetric | |
| from apache_beam.io.gcp import resource_identifiers | |
| from apache_beam.metrics import Metrics | |
| from apache_beam.metrics import monitoring_infos | |
| from apache_beam.metrics.metric import Lineage | |
| from apache_beam.transforms import PTransform | |
| from apache_beam.transforms.display import DisplayDataItem | |
| from apache_beam.transforms.external import BeamJarExpansionService | |
| from apache_beam.transforms.external import SchemaAwareExternalTransform | |
| from apache_beam.typehints.row_type import RowTypeConstraint | |
| _LOGGER = logging.getLogger(__name__) | |
| FLUSH_COUNT = 1000 | |
| MAX_ROW_BYTES = 5242880 # 5MB | |
| try: | |
| from google.api_core.gapic_v1 import client_info as client_info_lib | |
| from google.cloud.bigtable import Client | |
| from google.cloud.bigtable.batcher import MutationsBatcher | |
| from google.cloud.bigtable.row import Cell | |
| from google.cloud.bigtable.row import PartialRowData | |
| except ImportError: | |
| _LOGGER.warning( | |
| 'ImportError: from google.cloud.bigtable import Client', exc_info=True) | |
| __all__ = ['WriteToBigTable', 'ReadFromBigtable'] | |
| class _BigTableWriteFn(beam.DoFn): | |
| """ Creates the connector can call and add_row to the batcher using each | |
| row in beam pipe line | |
| Args: | |
| project_id(str): GCP Project ID | |
| instance_id(str): GCP Instance ID | |
| table_id(str): GCP Table ID | |
| flush_count(int): Max number of rows to flush | |
| max_row_bytes(int) Max number of row mutations size to flush | |
| """ | |
| def __init__( | |
| self, project_id, instance_id, table_id, flush_count, max_row_bytes): | |
| """ Constructor of the Write connector of Bigtable | |
| Args: | |
| project_id(str): GCP Project of to write the Rows | |
| instance_id(str): GCP Instance to write the Rows | |
| table_id(str): GCP Table to write the `DirectRows` | |
| flush_count(int): Max number of rows to flush | |
| max_row_bytes(int) Max number of row mutations size to flush | |
| """ | |
| super().__init__() | |
| self.beam_options = { | |
| 'project_id': project_id, | |
| 'instance_id': instance_id, | |
| 'table_id': table_id, | |
| 'flush_count': flush_count, | |
| 'max_row_bytes': max_row_bytes, | |
| } | |
| self.table = None | |
| self.batcher = None | |
| self.service_call_metric = None | |
| self.written = Metrics.counter(self.__class__, 'Written Row') | |
| def __getstate__(self): | |
| return self.beam_options | |
| def __setstate__(self, options): | |
| self.beam_options = options | |
| self.table = None | |
| self.batcher = None | |
| self.service_call_metric = None | |
| self.written = Metrics.counter(self.__class__, 'Written Row') | |
| def write_mutate_metrics(self, status_list): | |
| for status in status_list: | |
| code = status.code if status else None | |
| grpc_status_string = ( | |
| ServiceCallMetric.bigtable_error_code_to_grpc_status_string(code)) | |
| self.service_call_metric.call(grpc_status_string) | |
| def start_service_call_metrics(self, project_id, instance_id, table_id): | |
| resource = resource_identifiers.BigtableTable( | |
| project_id, instance_id, table_id) | |
| labels = { | |
| monitoring_infos.SERVICE_LABEL: 'BigTable', | |
| # TODO(JIRA-11985): Add Ptransform label. | |
| monitoring_infos.METHOD_LABEL: 'google.bigtable.v2.MutateRows', | |
| monitoring_infos.RESOURCE_LABEL: resource, | |
| monitoring_infos.BIGTABLE_PROJECT_ID_LABEL: ( | |
| self.beam_options['project_id']), | |
| monitoring_infos.INSTANCE_ID_LABEL: self.beam_options['instance_id'], | |
| monitoring_infos.TABLE_ID_LABEL: self.beam_options['table_id'] | |
| } | |
| return ServiceCallMetric( | |
| request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, | |
| base_labels=labels) | |
| def start_bundle(self): | |
| if self.table is None: | |
| client = Client( | |
| project=self.beam_options['project_id'], | |
| client_info=client_info_lib.ClientInfo( | |
| user_agent="apache-beam/%s (GPN:Beam)" % | |
| beam_version.__version__)) | |
| instance = client.instance(self.beam_options['instance_id']) | |
| self.table = instance.table(self.beam_options['table_id']) | |
| self.service_call_metric = self.start_service_call_metrics( | |
| self.beam_options['project_id'], | |
| self.beam_options['instance_id'], | |
| self.beam_options['table_id']) | |
| self.batcher = MutationsBatcher( | |
| self.table, | |
| batch_completed_callback=self.write_mutate_metrics, | |
| flush_count=self.beam_options['flush_count'], | |
| max_row_bytes=self.beam_options['max_row_bytes']) | |
| def process(self, row): | |
| self.written.inc() | |
| # You need to set the timestamp in the cells in this row object, | |
| # when we do a retry we will mutating the same object, but, with this | |
| # we are going to set our cell with new values. | |
| # Example: | |
| # direct_row.set_cell('cf1', | |
| # 'field1', | |
| # 'value1', | |
| # timestamp=datetime.now()) | |
| self.batcher.mutate(row) | |
| def finish_bundle(self): | |
| if self.batcher: | |
| self.batcher.close() | |
| self.batcher = None | |
| # Report Lineage metrics on write | |
| Lineage.sinks().add( | |
| 'bigtable', | |
| self.beam_options['project_id'], | |
| self.beam_options['instance_id'], | |
| self.beam_options['table_id']) | |
| def display_data(self): | |
| return { | |
| 'projectId': DisplayDataItem( | |
| self.beam_options['project_id'], label='Bigtable Project Id'), | |
| 'instanceId': DisplayDataItem( | |
| self.beam_options['instance_id'], label='Bigtable Instance Id'), | |
| 'tableId': DisplayDataItem( | |
| self.beam_options['table_id'], label='Bigtable Table Id') | |
| } | |
| class WriteToBigTable(beam.PTransform): | |
| """A transform that writes rows to a Bigtable table. | |
| Takes an input PCollection of `DirectRow` objects containing un-committed | |
| mutations. For more information about this row object, visit | |
| https://cloud.google.com/python/docs/reference/bigtable/latest/row#class-googlecloudbigtablerowdirectrowrowkey-tablenone | |
| If flag `use_cross_language` is set to true, this transform will use the | |
| multi-language transforms framework to inject the Java native write transform | |
| into the pipeline. | |
| """ | |
| URN = "beam:schematransform:org.apache.beam:bigtable_write:v1" | |
| def __init__( | |
| self, | |
| project_id, | |
| instance_id, | |
| table_id, | |
| use_cross_language=False, | |
| expansion_service=None, | |
| flush_count=FLUSH_COUNT, | |
| max_row_bytes=MAX_ROW_BYTES, | |
| ): | |
| """Initialize an WriteToBigTable transform. | |
| :param table_id: | |
| The ID of the table to write to. | |
| :param instance_id: | |
| The ID of the instance where the table resides. | |
| :param project_id: | |
| The GCP project ID. | |
| :param use_cross_language: | |
| If set to True, will use the Java native transform via cross-language. | |
| :param expansion_service: | |
| The address of the expansion service in the case of using cross-language. | |
| If no expansion service is provided, will attempt to run the default GCP | |
| expansion service. | |
| :type flush_count: int | |
| :param flush_count: (Optional) Max number of rows to flush. | |
| Default is FLUSH_COUNT (1000 rows). | |
| :type max_row_bytes: int | |
| :param max_row_bytes: (Optional) Max number of row mutations size to flush. | |
| Default is MAX_ROW_BYTES (5 MB). | |
| """ | |
| super().__init__() | |
| self._table_id = table_id | |
| self._instance_id = instance_id | |
| self._project_id = project_id | |
| self._use_cross_language = use_cross_language | |
| if use_cross_language: | |
| self._expansion_service = ( | |
| expansion_service or BeamJarExpansionService( | |
| 'sdks:java:io:google-cloud-platform:expansion-service:build')) | |
| self.schematransform_config = ( | |
| SchemaAwareExternalTransform.discover_config( | |
| self._expansion_service, self.URN)) | |
| self._flush_count = flush_count | |
| self._max_row_bytes = max_row_bytes | |
| def expand(self, input): | |
| if self._use_cross_language: | |
| external_write = SchemaAwareExternalTransform( | |
| identifier=self.schematransform_config.identifier, | |
| expansion_service=self._expansion_service, | |
| rearrange_based_on_discovery=True, | |
| table_id=self._table_id, | |
| instance_id=self._instance_id, | |
| project_id=self._project_id) | |
| return ( | |
| input | |
| | beam.ParDo(self._DirectRowMutationsToBeamRow()).with_output_types( | |
| RowTypeConstraint.from_fields( | |
| [("key", bytes), ("mutations", list[dict[str, bytes]])])) | |
| | external_write) | |
| else: | |
| return ( | |
| input | |
| | beam.ParDo( | |
| _BigTableWriteFn( | |
| self._project_id, | |
| self._instance_id, | |
| self._table_id, | |
| flush_count=self._flush_count, | |
| max_row_bytes=self._max_row_bytes))) | |
| class _DirectRowMutationsToBeamRow(beam.DoFn): | |
| def process(self, direct_row): | |
| args = {"key": direct_row.row_key, "mutations": []} | |
| # start accumulating mutations in a list | |
| # In google-cloud-bigtable >= 2.44.0, _get_mutations() returns Python | |
| # dataclass objects (RowMutationEntry) instead of protobuf messages. | |
| # Use _get_mutation_pbs() to retrieve Mutation protobuf objects. | |
| mutations = ( | |
| direct_row._get_mutation_pbs() if hasattr( | |
| direct_row, '_get_mutation_pbs') else direct_row._get_mutations()) | |
| for mutation in mutations: | |
| if mutation.__contains__("set_cell"): | |
| mutation_dict = { | |
| "type": b'SetCell', | |
| "family_name": mutation.set_cell.family_name.encode('utf-8'), | |
| "column_qualifier": mutation.set_cell.column_qualifier, | |
| "value": mutation.set_cell.value, | |
| "timestamp_micros": struct.pack( | |
| '>q', mutation.set_cell.timestamp_micros) | |
| } | |
| elif mutation.__contains__("delete_from_column"): | |
| mutation_dict = { | |
| "type": b'DeleteFromColumn', | |
| "family_name": mutation.delete_from_column.family_name.encode( | |
| 'utf-8'), | |
| "column_qualifier": mutation.delete_from_column.column_qualifier | |
| } | |
| time_range = mutation.delete_from_column.time_range | |
| if time_range.start_timestamp_micros: | |
| mutation_dict['start_timestamp_micros'] = struct.pack( | |
| '>q', time_range.start_timestamp_micros) | |
| if time_range.end_timestamp_micros: | |
| mutation_dict['end_timestamp_micros'] = struct.pack( | |
| '>q', time_range.end_timestamp_micros) | |
| elif mutation.__contains__("delete_from_family"): | |
| mutation_dict = { | |
| "type": b'DeleteFromFamily', | |
| "family_name": mutation.delete_from_family.family_name.encode( | |
| 'utf-8') | |
| } | |
| elif mutation.__contains__("delete_from_row"): | |
| mutation_dict = {"type": b'DeleteFromRow'} | |
| else: | |
| raise ValueError("Unexpected mutation") | |
| args["mutations"].append(mutation_dict) | |
| yield beam.Row(**args) | |
| class ReadFromBigtable(PTransform): | |
| """Reads rows from Bigtable. | |
| Returns a PCollection of PartialRowData objects, each representing a | |
| Bigtable row. For more information about this row object, visit | |
| https://cloud.google.com/python/docs/reference/bigtable/latest/row#class-googlecloudbigtablerowpartialrowdatarowkey | |
| """ | |
| URN = "beam:schematransform:org.apache.beam:bigtable_read:v1" | |
| def __init__(self, project_id, instance_id, table_id, expansion_service=None): | |
| """Initialize a ReadFromBigtable transform. | |
| :param table_id: | |
| The ID of the table to read from. | |
| :param instance_id: | |
| The ID of the instance where the table resides. | |
| :param project_id: | |
| The GCP project ID. | |
| :param expansion_service: | |
| The address of the expansion service. If no expansion service is | |
| provided, will attempt to run the default GCP expansion service. | |
| """ | |
| super().__init__() | |
| self._table_id = table_id | |
| self._instance_id = instance_id | |
| self._project_id = project_id | |
| self._expansion_service = ( | |
| expansion_service or BeamJarExpansionService( | |
| 'sdks:java:io:google-cloud-platform:expansion-service:build')) | |
| self.schematransform_config = SchemaAwareExternalTransform.discover_config( | |
| self._expansion_service, self.URN) | |
| def expand(self, input): | |
| external_read = SchemaAwareExternalTransform( | |
| identifier=self.schematransform_config.identifier, | |
| expansion_service=self._expansion_service, | |
| rearrange_based_on_discovery=True, | |
| table_id=self._table_id, | |
| instance_id=self._instance_id, | |
| project_id=self._project_id, | |
| flatten=False) | |
| return ( | |
| input.pipeline | |
| | external_read | |
| | beam.ParDo(self._BeamRowToPartialRowData())) | |
| # PartialRowData has some useful methods for querying data within a row. | |
| # To make use of those methods and to give Python users a more familiar | |
| # object, we process each Beam Row and return a PartialRowData equivalent. | |
| class _BeamRowToPartialRowData(beam.DoFn): | |
| def process(self, row): | |
| key = row.key | |
| families = row.column_families | |
| # initialize PartialRowData object | |
| partial_row: PartialRowData = PartialRowData(key) | |
| for fam_name, col_fam in families.items(): | |
| if fam_name not in partial_row.cells: | |
| partial_row.cells[fam_name] = {} | |
| for col_qualifier, cells in col_fam.items(): | |
| # store column qualifier as bytes to follow PartialRowData behavior | |
| col_qualifier_bytes = col_qualifier.encode() | |
| if col_qualifier not in partial_row.cells[fam_name]: | |
| partial_row.cells[fam_name][col_qualifier_bytes] = [] | |
| for cell in cells: | |
| value = cell.value | |
| timestamp_micros = cell.timestamp_micros | |
| partial_row.cells[fam_name][col_qualifier_bytes].append( | |
| Cell(value, timestamp_micros)) | |
| yield partial_row |