#

# 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.

#

"""Unit tests for BigTable service."""

import logging

import string

import unittest

import uuid

# pytype: skip-file

from datetime import datetime

from datetime import timezone

from random import choice

from mock import MagicMock

from mock import patch

import apache_beam as beam

from apache_beam.internal.metrics.metric import ServiceCallMetric

from apache_beam.io.gcp import bigtableio

from apache_beam.io.gcp import resource_identifiers

from apache_beam.metrics import monitoring_infos

from apache_beam.metrics.execution import MetricsEnvironment

from apache_beam.metrics.metric import Lineage

from apache_beam.testing.test_pipeline import TestPipeline

_LOGGER = logging.getLogger(__name__)

# Protect against environments where bigtable library is not available.

try:

from google.cloud.bigtable import client

from google.cloud.bigtable.batcher import MutationsBatcher

from google.cloud.bigtable.instance import Instance

from google.cloud.bigtable.row import Cell

from google.cloud.bigtable.row import DirectRow

from google.cloud.bigtable.row import PartialRowData

from google.cloud.bigtable.row_filters import TimestampRange

from google.cloud.bigtable.table import Table

from google.rpc.code_pb2 import ALREADY_EXISTS

from google.rpc.code_pb2 import OK

from google.rpc.status_pb2 import Status

except ImportError as e:

client = None

@unittest.skipIf(client is None, 'Bigtable dependencies are not installed')

class TestBeamRowToPartialRowData(unittest.TestCase):

# Beam Row schema:

# - key: bytes

# - column_families: dict[str, dict[str, list[beam.Row(cell_schema)]]]

# cell_schema:

# - value: bytes

# - timestamp_micros: int

def test_beam_row_to_bigtable_row(self):

# create test beam row

families = {

'family_1': {

'col_1': [

beam.Row(value=b'a-1', timestamp_micros=100_000_000),

beam.Row(value=b'b-1', timestamp_micros=200_000_000),

beam.Row(value=b'c-1', timestamp_micros=300_000_000)

],

'col_2': [

beam.Row(value=b'a-2', timestamp_micros=400_000_000),

beam.Row(value=b'b-2', timestamp_micros=500_000_000),

beam.Row(value=b'c-2', timestamp_micros=600_000_000)

],

},

'family_2': {

'column_qualifier': [

beam.Row(value=b'val-1', timestamp_micros=700_000_000),

beam.Row(value=b'val-2', timestamp_micros=800_000_000),

beam.Row(value=b'val-3', timestamp_micros=900_000_000)

]

}

}

beam_row = beam.Row(key=b'key', column_families=families)

# get equivalent bigtable row

doFn = bigtableio.ReadFromBigtable._BeamRowToPartialRowData()

bigtable_row: PartialRowData = next(doFn.process(beam_row))

# using bigtable utils (PartialRowData methods), check that beam row data

# landed in the right cells

self.assertEqual(beam_row.key, bigtable_row.row_key)

self.assertEqual([

Cell(c.value, c.timestamp_micros)

for c in beam_row.column_families['family_1']['col_1']

],

bigtable_row.find_cells('family_1', b'col_1'))

self.assertEqual([

Cell(c.value, c.timestamp_micros)

for c in beam_row.column_families['family_1']['col_2']

],

bigtable_row.find_cells('family_1', b'col_2'))

self.assertEqual([

Cell(c.value, c.timestamp_micros)

for c in beam_row.column_families['family_2']['column_qualifier']

],

bigtable_row.find_cells('family_2', b'column_qualifier'))

@unittest.skipIf(client is None, 'Bigtable dependencies are not installed')

class TestBigtableDirectRowToBeamRow(unittest.TestCase):

doFn = bigtableio.WriteToBigTable._DirectRowMutationsToBeamRow()

def test_set_cell(self):

# create some set cell mutations

direct_row: DirectRow = DirectRow('key-1')

direct_row.set_cell(

'col_fam',

b'col',

b'a',

datetime.fromtimestamp(100_000).replace(tzinfo=timezone.utc))

direct_row.set_cell(

'col_fam',

b'other-col',

b'b',

datetime.fromtimestamp(200_000).replace(tzinfo=timezone.utc))

direct_row.set_cell(

'other_col_fam',

b'col',

b'c',

datetime.fromtimestamp(300_000).replace(tzinfo=timezone.utc))

# get equivalent beam row

beam_row = next(self.doFn.process(direct_row))

# sort both lists of mutations for convenience

beam_row_mutations = sorted(beam_row.mutations, key=lambda m: m['value'])

bt_row_mutations = sorted(

direct_row._get_mutations(), key=lambda m: m.set_cell.value)

self.assertEqual(beam_row.key, direct_row.row_key)

self.assertEqual(len(beam_row_mutations), len(bt_row_mutations))

# check that the values in each beam mutation is equal to the original

# Bigtable direct row mutations

for i in range(len(beam_row_mutations)):

beam_mutation = beam_row_mutations[i]

bt_mutation = bt_row_mutations[i].set_cell

self.assertEqual(beam_mutation['type'], b'SetCell')

self.assertEqual(

beam_mutation['family_name'].decode(), bt_mutation.family_name)

self.assertEqual(

beam_mutation['column_qualifier'], bt_mutation.column_qualifier)

self.assertEqual(beam_mutation['value'], bt_mutation.value)

self.assertEqual(

int.from_bytes(beam_mutation['timestamp_micros'], 'big'),

bt_mutation.timestamp_micros)

def test_delete_cells(self):

# create some delete cell mutations. one with a timestamp range

direct_row: DirectRow = DirectRow('key-1')

direct_row.delete_cell('col_fam', b'col-1')

direct_row.delete_cell(

'other_col_fam',

b'col-2',

time_range=TimestampRange(

start=datetime.fromtimestamp(10_000_000, tz=timezone.utc)))

direct_row.delete_cells(

'another_col_fam', [b'col-3', b'col-4', b'col-5'],

time_range=TimestampRange(

start=datetime.fromtimestamp(50_000_000, tz=timezone.utc),

end=datetime.fromtimestamp(100_000_000, tz=timezone.utc)))

# get equivalent beam row

beam_row = next(self.doFn.process(direct_row))

# sort both lists of mutations for convenience

beam_row_mutations = sorted(

beam_row.mutations, key=lambda m: m['column_qualifier'])

bt_row_mutations = sorted(

direct_row._get_mutations(),

key=lambda m: m.delete_from_column.column_qualifier)

self.assertEqual(beam_row.key, direct_row.row_key)

self.assertEqual(len(beam_row_mutations), len(bt_row_mutations))

# check that the values in each beam mutation is equal to the original

# Bigtable direct row mutations

for i in range(len(beam_row_mutations)):

beam_mutation = beam_row_mutations[i]

bt_mutation = bt_row_mutations[i].delete_from_column

print(bt_mutation)

self.assertEqual(beam_mutation['type'], b'DeleteFromColumn')

self.assertEqual(

beam_mutation['family_name'].decode(), bt_mutation.family_name)

self.assertEqual(

beam_mutation['column_qualifier'], bt_mutation.column_qualifier)

# check we set a timestamp range only when appropriate

if bt_mutation.time_range.start_timestamp_micros:

self.assertEqual(

int.from_bytes(beam_mutation['start_timestamp_micros'], 'big'),

bt_mutation.time_range.start_timestamp_micros)

else:

self.assertTrue('start_timestamp_micros' not in beam_mutation)

if bt_mutation.time_range.end_timestamp_micros:

self.assertEqual(

int.from_bytes(beam_mutation['end_timestamp_micros'], 'big'),

bt_mutation.time_range.end_timestamp_micros)

else:

self.assertTrue('end_timestamp_micros' not in beam_mutation)

def test_delete_column_family(self):

# create mutation to delete column family

direct_row: DirectRow = DirectRow('key-1')

direct_row.delete_cells('col_fam-1', direct_row.ALL_COLUMNS)

direct_row.delete_cells('col_fam-2', direct_row.ALL_COLUMNS)

# get equivalent beam row

beam_row = next(self.doFn.process(direct_row))

# sort both lists of mutations for convenience

beam_row_mutations = sorted(

beam_row.mutations, key=lambda m: m['family_name'])

bt_row_mutations = sorted(

direct_row._get_mutations(),

key=lambda m: m.delete_from_column.family_name)

self.assertEqual(beam_row.key, direct_row.row_key)

self.assertEqual(len(beam_row_mutations), len(bt_row_mutations))

# check that the values in each beam mutation is equal to the original

# Bigtable direct row mutations

for i in range(len(beam_row_mutations)):

beam_mutation = beam_row_mutations[i]

bt_mutation = bt_row_mutations[i].delete_from_family

self.assertEqual(beam_mutation['type'], b'DeleteFromFamily')

self.assertEqual(

beam_mutation['family_name'].decode(), bt_mutation.family_name)

def test_delete_row(self):

# create mutation to delete the Bigtable row

direct_row: DirectRow = DirectRow('key-1')

direct_row.delete()

# get equivalent beam row

beam_row = next(self.doFn.process(direct_row))

self.assertEqual(beam_row.key, direct_row.row_key)

beam_mutation = beam_row.mutations[0]

self.assertEqual(beam_mutation['type'], b'DeleteFromRow')

@unittest.skipIf(client is None, 'Bigtable dependencies are not installed')

class TestWriteBigTable(unittest.TestCase):

TABLE_PREFIX = "python-test"

_PROJECT_ID = TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]

_INSTANCE_ID = TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]

_TABLE_ID = TABLE_PREFIX + "-" + str(uuid.uuid4())[:8]

def setUp(self):

client = MagicMock()

instance = Instance(self._INSTANCE_ID, client)

self.table = Table(self._TABLE_ID, instance)

def test_write(self):

direct_rows = [self.generate_row(i) for i in range(5)]

# TODO(https://github.com/apache/beam/issues/34549): This test relies on

# lineage metrics which Prism doesn't seem to handle correctly. Defaulting

# to FnApiRunner instead.

runner = 'FnApiRunner'

with patch.object(MutationsBatcher, 'mutate'), \

patch.object(MutationsBatcher, 'close'), TestPipeline(runner) as p:

_ = p | beam.Create(direct_rows) | bigtableio.WriteToBigTable(

self._PROJECT_ID, self._INSTANCE_ID, self._TABLE_ID)

self.assertSetEqual(

Lineage.query(p.result.metrics(), Lineage.SINK),

set([

f"bigtable:{self._PROJECT_ID}.{self._INSTANCE_ID}.{self._TABLE_ID}"

]))

def test_write_metrics(self):

MetricsEnvironment.process_wide_container().reset()

write_fn = bigtableio._BigTableWriteFn(

self._PROJECT_ID,

self._INSTANCE_ID,

self._TABLE_ID,

flush_count=1000,

max_row_bytes=5242880)

write_fn.table = self.table

write_fn.start_bundle()

number_of_rows = 2

error = Status()

error.message = 'Entity already exists.'

error.code = ALREADY_EXISTS

success = Status()

success.message = 'Success'

success.code = OK

rows_response = [error, success] * number_of_rows

with patch.object(Table, 'mutate_rows', return_value=rows_response):

direct_rows = [self.generate_row(i) for i in range(number_of_rows * 2)]

for direct_row in direct_rows:

write_fn.process(direct_row)

try:

write_fn.finish_bundle()

except: # pylint: disable=bare-except

# Currently we fail the bundle when there are any failures.

# TODO(https://github.com/apache/beam/issues/21396): remove after

# bigtableio can selectively retry.

pass

self.verify_write_call_metric(

self._PROJECT_ID,

self._INSTANCE_ID,

self._TABLE_ID,

ServiceCallMetric.bigtable_error_code_to_grpc_status_string(

ALREADY_EXISTS),

2)

self.verify_write_call_metric(

self._PROJECT_ID,

self._INSTANCE_ID,

self._TABLE_ID,

ServiceCallMetric.bigtable_error_code_to_grpc_status_string(OK),

2)

def generate_row(self, index=0):

rand = choice(string.ascii_letters + string.digits)

value = ''.join(rand for i in range(100))

column_family_id = 'cf1'

key = "beam_key%s" % ('{0:07}'.format(index))

direct_row = DirectRow(row_key=key)

for column_id in range(10):

direct_row.set_cell(

column_family_id, ('field%s' % column_id).encode('utf-8'),

value,

datetime.now())

return direct_row

def verify_write_call_metric(

self, project_id, instance_id, table_id, status, count):

"""Check if a metric was recorded for the Datastore IO write API call."""

process_wide_monitoring_infos = list(

MetricsEnvironment.process_wide_container().

to_runner_api_monitoring_infos(None).values())

resource = resource_identifiers.BigtableTable(

project_id, instance_id, table_id)

labels = {

monitoring_infos.SERVICE_LABEL: 'BigTable',

monitoring_infos.METHOD_LABEL: 'google.bigtable.v2.MutateRows',

monitoring_infos.RESOURCE_LABEL: resource,

monitoring_infos.BIGTABLE_PROJECT_ID_LABEL: project_id,

monitoring_infos.INSTANCE_ID_LABEL: instance_id,

monitoring_infos.TABLE_ID_LABEL: table_id,

monitoring_infos.STATUS_LABEL: status

}

expected_mi = monitoring_infos.int64_counter(

monitoring_infos.API_REQUEST_COUNT_URN, count, labels=labels)

expected_mi.ClearField("start_time")

found = False

for actual_mi in process_wide_monitoring_infos:

actual_mi.ClearField("start_time")

if expected_mi == actual_mi:

found = True

break

self.assertTrue(

found, "Did not find write call metric with status: %s" % status)

def test_custom_flush_config(self):

direct_rows = [self.generate_row(0)]

with patch.object(

MutationsBatcher, '__init__', return_value=None) as mock_init, \

patch.object(MutationsBatcher, 'mutate'), \

patch.object(MutationsBatcher, 'close'), TestPipeline() as p:

_ = p | beam.Create(direct_rows) | bigtableio.WriteToBigTable(

self._PROJECT_ID,

self._INSTANCE_ID,

self._TABLE_ID,

flush_count=1001,

max_row_bytes=5000001)

mock_init.assert_called_once()

call_args = mock_init.call_args.kwargs

assert call_args['flush_count'] == 1001

assert call_args['max_row_bytes'] == 5000001

if __name__ == '__main__':

unittest.main()