| # |
| # 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 |
| from typing import Dict |
| from typing import List |
| |
| import apache_beam as beam |
| 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.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__) |
| |
| try: |
| from google.cloud.bigtable import Client |
| from google.cloud.bigtable.row import Cell, PartialRowData |
| from google.cloud.bigtable.batcher import MutationsBatcher |
| |
| FLUSH_COUNT = 1000 |
| MAX_ROW_BYTES = 5242880 # 5MB |
| |
| 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 |
| |
| """ |
| def __init__(self, project_id, instance_id, table_id): |
| """ 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` |
| """ |
| super().__init__() |
| self.beam_options = { |
| 'project_id': project_id, |
| 'instance_id': instance_id, |
| 'table_id': table_id |
| } |
| 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']) |
| 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) |
| |
| 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 |
| |
| 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): |
| """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. |
| """ |
| 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)) |
| |
| 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, |
| tableId=self._table_id, |
| instanceId=self._instance_id, |
| projectId=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))) |
| |
| class _DirectRowMutationsToBeamRow(beam.DoFn): |
| def process(self, direct_row): |
| args = {"key": direct_row.row_key, "mutations": []} |
| # start accumulating mutations in a list |
| for mutation in direct_row._get_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, |
| tableId=self._table_id, |
| instanceId=self._instance_id, |
| projectId=self._project_id) |
| |
| 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 |
