From b5b62c820e2d4f05b3f96fe4d4995b44d06209d3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 14 Mar 2023 15:47:41 -0700 Subject: [PATCH 01/56] feat!: add new v3.0.0 API skeleton (#745) --- .coveragerc | 2 +- docs/app-profile.rst | 2 +- docs/backup.rst | 2 +- docs/client-intro.rst | 18 +- docs/client.rst | 2 +- docs/cluster.rst | 2 +- docs/column-family.rst | 22 +- docs/data-api.rst | 82 +- docs/encryption-info.rst | 2 +- docs/instance-api.rst | 32 +- docs/instance.rst | 2 +- docs/row-data.rst | 2 +- docs/row-filters.rst | 12 +- docs/row-set.rst | 2 +- docs/row.rst | 2 +- docs/snippets.py | 118 +-- docs/snippets_table.py | 154 ++-- docs/table-api.rst | 40 +- docs/table.rst | 2 +- docs/usage.rst | 16 +- google/cloud/bigtable/__init__.py | 44 +- google/cloud/bigtable/client.py | 868 ++++++++---------- google/cloud/bigtable/deprecated/__init__.py | 25 + .../bigtable/{ => deprecated}/app_profile.py | 8 +- .../cloud/bigtable/{ => deprecated}/backup.py | 20 +- .../bigtable/{ => deprecated}/batcher.py | 8 +- google/cloud/bigtable/deprecated/client.py | 521 +++++++++++ .../bigtable/{ => deprecated}/cluster.py | 22 +- .../{ => deprecated}/column_family.py | 2 +- .../{ => deprecated}/encryption_info.py | 4 +- .../cloud/bigtable/{ => deprecated}/enums.py | 0 .../cloud/bigtable/{ => deprecated}/error.py | 0 .../bigtable/{ => deprecated}/instance.py | 68 +- .../cloud/bigtable/{ => deprecated}/policy.py | 0 .../cloud/bigtable/{ => deprecated}/py.typed | 0 google/cloud/bigtable/{ => deprecated}/row.py | 18 +- .../bigtable/{ => deprecated}/row_data.py | 6 +- .../cloud/bigtable/deprecated/row_filters.py | 838 +++++++++++++++++ .../bigtable/{ => deprecated}/row_merger.py | 2 +- .../bigtable/{ => deprecated}/row_set.py | 0 .../cloud/bigtable/{ => deprecated}/table.py | 64 +- google/cloud/bigtable/exceptions.py | 46 + google/cloud/bigtable/mutations.py | 54 ++ google/cloud/bigtable/mutations_batcher.py | 104 +++ .../cloud/bigtable/read_modify_write_rules.py | 37 + google/cloud/bigtable/read_rows_query.py | 56 ++ google/cloud/bigtable/row_response.py | 130 +++ noxfile.py | 18 +- owlbot.py | 18 +- tests/system/__init__.py | 2 +- tests/system/v2_client/__init__.py | 15 + tests/system/{ => v2_client}/_helpers.py | 0 tests/system/{ => v2_client}/conftest.py | 2 +- tests/system/{ => v2_client}/test_data_api.py | 20 +- .../{ => v2_client}/test_instance_admin.py | 6 +- .../{ => v2_client}/test_table_admin.py | 12 +- tests/unit/__init__.py | 2 +- tests/unit/v2_client/__init__.py | 15 + tests/unit/{ => v2_client}/_testing.py | 0 .../read-rows-acceptance-test.json | 0 .../unit/{ => v2_client}/test_app_profile.py | 44 +- tests/unit/{ => v2_client}/test_backup.py | 36 +- tests/unit/{ => v2_client}/test_batcher.py | 10 +- tests/unit/{ => v2_client}/test_client.py | 59 +- tests/unit/{ => v2_client}/test_cluster.py | 68 +- .../{ => v2_client}/test_column_family.py | 68 +- .../{ => v2_client}/test_encryption_info.py | 8 +- tests/unit/{ => v2_client}/test_error.py | 2 +- tests/unit/{ => v2_client}/test_instance.py | 50 +- tests/unit/{ => v2_client}/test_policy.py | 28 +- tests/unit/{ => v2_client}/test_row.py | 34 +- tests/unit/{ => v2_client}/test_row_data.py | 58 +- .../unit/{ => v2_client}/test_row_filters.py | 214 ++--- tests/unit/{ => v2_client}/test_row_merger.py | 8 +- tests/unit/{ => v2_client}/test_row_set.py | 60 +- tests/unit/{ => v2_client}/test_table.py | 176 ++-- 76 files changed, 3165 insertions(+), 1329 deletions(-) create mode 100644 google/cloud/bigtable/deprecated/__init__.py rename google/cloud/bigtable/{ => deprecated}/app_profile.py (97%) rename google/cloud/bigtable/{ => deprecated}/backup.py (96%) rename google/cloud/bigtable/{ => deprecated}/batcher.py (94%) create mode 100644 google/cloud/bigtable/deprecated/client.py rename google/cloud/bigtable/{ => deprecated}/cluster.py (95%) rename google/cloud/bigtable/{ => deprecated}/column_family.py (99%) rename google/cloud/bigtable/{ => deprecated}/encryption_info.py (93%) rename google/cloud/bigtable/{ => deprecated}/enums.py (100%) rename google/cloud/bigtable/{ => deprecated}/error.py (100%) rename google/cloud/bigtable/{ => deprecated}/instance.py (91%) rename google/cloud/bigtable/{ => deprecated}/policy.py (100%) rename google/cloud/bigtable/{ => deprecated}/py.typed (100%) rename google/cloud/bigtable/{ => deprecated}/row.py (98%) rename google/cloud/bigtable/{ => deprecated}/row_data.py (98%) create mode 100644 google/cloud/bigtable/deprecated/row_filters.py rename google/cloud/bigtable/{ => deprecated}/row_merger.py (99%) rename google/cloud/bigtable/{ => deprecated}/row_set.py (100%) rename google/cloud/bigtable/{ => deprecated}/table.py (95%) create mode 100644 google/cloud/bigtable/exceptions.py create mode 100644 google/cloud/bigtable/mutations.py create mode 100644 google/cloud/bigtable/mutations_batcher.py create mode 100644 google/cloud/bigtable/read_modify_write_rules.py create mode 100644 google/cloud/bigtable/read_rows_query.py create mode 100644 google/cloud/bigtable/row_response.py create mode 100644 tests/system/v2_client/__init__.py rename tests/system/{ => v2_client}/_helpers.py (100%) rename tests/system/{ => v2_client}/conftest.py (98%) rename tests/system/{ => v2_client}/test_data_api.py (94%) rename tests/system/{ => v2_client}/test_instance_admin.py (99%) rename tests/system/{ => v2_client}/test_table_admin.py (96%) create mode 100644 tests/unit/v2_client/__init__.py rename tests/unit/{ => v2_client}/_testing.py (100%) rename tests/unit/{ => v2_client}/read-rows-acceptance-test.json (100%) rename tests/unit/{ => v2_client}/test_app_profile.py (94%) rename tests/unit/{ => v2_client}/test_backup.py (96%) rename tests/unit/{ => v2_client}/test_batcher.py (91%) rename tests/unit/{ => v2_client}/test_client.py (93%) rename tests/unit/{ => v2_client}/test_cluster.py (94%) rename tests/unit/{ => v2_client}/test_column_family.py (87%) rename tests/unit/{ => v2_client}/test_encryption_info.py (94%) rename tests/unit/{ => v2_client}/test_error.py (97%) rename tests/unit/{ => v2_client}/test_instance.py (95%) rename tests/unit/{ => v2_client}/test_policy.py (89%) rename tests/unit/{ => v2_client}/test_row.py (95%) rename tests/unit/{ => v2_client}/test_row_data.py (94%) rename tests/unit/{ => v2_client}/test_row_filters.py (77%) rename tests/unit/{ => v2_client}/test_row_merger.py (97%) rename tests/unit/{ => v2_client}/test_row_set.py (79%) rename tests/unit/{ => v2_client}/test_table.py (91%) diff --git a/.coveragerc b/.coveragerc index 3128ad99e..702b85681 100644 --- a/.coveragerc +++ b/.coveragerc @@ -24,7 +24,7 @@ omit = google/cloud/bigtable_admin/gapic_version.py [report] -fail_under = 100 +fail_under = 99 show_missing = True exclude_lines = # Re-enable the standard pragma diff --git a/docs/app-profile.rst b/docs/app-profile.rst index 5c9d426c2..50e57c179 100644 --- a/docs/app-profile.rst +++ b/docs/app-profile.rst @@ -1,6 +1,6 @@ App Profile ~~~~~~~~~~~ -.. automodule:: google.cloud.bigtable.app_profile +.. automodule:: google.cloud.bigtable.deprecated.app_profile :members: :show-inheritance: diff --git a/docs/backup.rst b/docs/backup.rst index e75abd431..46c32c91b 100644 --- a/docs/backup.rst +++ b/docs/backup.rst @@ -1,6 +1,6 @@ Backup ~~~~~~~~ -.. automodule:: google.cloud.bigtable.backup +.. automodule:: google.cloud.bigtable.deprecated.backup :members: :show-inheritance: diff --git a/docs/client-intro.rst b/docs/client-intro.rst index 242068499..d75cf5f96 100644 --- a/docs/client-intro.rst +++ b/docs/client-intro.rst @@ -1,21 +1,21 @@ Base for Everything =================== -To use the API, the :class:`Client ` +To use the API, the :class:`Client ` class defines a high-level interface which handles authorization and creating other objects: .. code:: python - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client client = Client() Long-lived Defaults ------------------- -When creating a :class:`Client `, the +When creating a :class:`Client `, the ``user_agent`` argument has sensible a default -(:data:`DEFAULT_USER_AGENT `). +(:data:`DEFAULT_USER_AGENT `). However, you may over-ride it and the value will be used throughout all API requests made with the ``client`` you create. @@ -38,14 +38,14 @@ Configuration .. code:: - >>> from google.cloud import bigtable + >>> import google.cloud.deprecated as bigtable >>> client = bigtable.Client() or pass in ``credentials`` and ``project`` explicitly .. code:: - >>> from google.cloud import bigtable + >>> import google.cloud.deprecated as bigtable >>> client = bigtable.Client(project='my-project', credentials=creds) .. tip:: @@ -73,15 +73,15 @@ you can pass the ``read_only`` argument: client = bigtable.Client(read_only=True) This will ensure that the -:data:`READ_ONLY_SCOPE ` is used +:data:`READ_ONLY_SCOPE ` is used for API requests (so any accidental requests that would modify data will fail). Next Step --------- -After a :class:`Client `, the next highest-level -object is an :class:`Instance `. You'll need +After a :class:`Client `, the next highest-level +object is an :class:`Instance `. You'll need one before you can interact with tables or data. Head next to learn about the :doc:`instance-api`. diff --git a/docs/client.rst b/docs/client.rst index c48595c8a..df92a9861 100644 --- a/docs/client.rst +++ b/docs/client.rst @@ -1,6 +1,6 @@ Client ~~~~~~ -.. automodule:: google.cloud.bigtable.client +.. automodule:: google.cloud.bigtable.deprecated.client :members: :show-inheritance: diff --git a/docs/cluster.rst b/docs/cluster.rst index ad33aae5e..9747b226f 100644 --- a/docs/cluster.rst +++ b/docs/cluster.rst @@ -1,6 +1,6 @@ Cluster ~~~~~~~ -.. automodule:: google.cloud.bigtable.cluster +.. automodule:: google.cloud.bigtable.deprecated.cluster :members: :show-inheritance: diff --git a/docs/column-family.rst b/docs/column-family.rst index de6c1eb1f..39095000d 100644 --- a/docs/column-family.rst +++ b/docs/column-family.rst @@ -2,7 +2,7 @@ Column Families =============== When creating a -:class:`ColumnFamily `, it is +:class:`ColumnFamily `, it is possible to set garbage collection rules for expired data. By setting a rule, cells in the table matching the rule will be deleted @@ -10,19 +10,19 @@ during periodic garbage collection (which executes opportunistically in the background). The types -:class:`MaxAgeGCRule `, -:class:`MaxVersionsGCRule `, -:class:`GarbageCollectionRuleUnion ` and -:class:`GarbageCollectionRuleIntersection ` +:class:`MaxAgeGCRule `, +:class:`MaxVersionsGCRule `, +:class:`GarbageCollectionRuleUnion ` and +:class:`GarbageCollectionRuleIntersection ` can all be used as the optional ``gc_rule`` argument in the -:class:`ColumnFamily ` +:class:`ColumnFamily ` constructor. This value is then used in the -:meth:`create() ` and -:meth:`update() ` methods. +:meth:`create() ` and +:meth:`update() ` methods. These rules can be nested arbitrarily, with a -:class:`MaxAgeGCRule ` or -:class:`MaxVersionsGCRule ` +:class:`MaxAgeGCRule ` or +:class:`MaxVersionsGCRule ` at the lowest level of the nesting: .. code:: python @@ -44,6 +44,6 @@ at the lowest level of the nesting: ---- -.. automodule:: google.cloud.bigtable.column_family +.. automodule:: google.cloud.bigtable.deprecated.column_family :members: :show-inheritance: diff --git a/docs/data-api.rst b/docs/data-api.rst index 01a49178f..e68835d1a 100644 --- a/docs/data-api.rst +++ b/docs/data-api.rst @@ -1,7 +1,7 @@ Data API ======== -After creating a :class:`Table ` and some +After creating a :class:`Table ` and some column families, you are ready to store and retrieve data. Cells vs. Columns vs. Column Families @@ -27,7 +27,7 @@ Modifying Data Since data is stored in cells, which are stored in rows, we use the metaphor of a **row** in classes that are used to modify (write, update, delete) data in a -:class:`Table `. +:class:`Table `. Direct vs. Conditional vs. Append --------------------------------- @@ -38,26 +38,26 @@ methods. * The **direct** way is via `MutateRow`_ which involves simply adding, overwriting or deleting cells. The - :class:`DirectRow ` class + :class:`DirectRow ` class handles direct mutations. * The **conditional** way is via `CheckAndMutateRow`_. This method first checks if some filter is matched in a given row, then applies one of two sets of mutations, depending on if a match occurred or not. (These mutation sets are called the "true mutations" and "false mutations".) The - :class:`ConditionalRow ` class + :class:`ConditionalRow ` class handles conditional mutations. * The **append** way is via `ReadModifyWriteRow`_. This simply appends (as bytes) or increments (as an integer) data in a presumed existing cell in a row. The - :class:`AppendRow ` class + :class:`AppendRow ` class handles append mutations. Row Factory ----------- A single factory can be used to create any of the three row types. -To create a :class:`DirectRow `: +To create a :class:`DirectRow `: .. code:: python @@ -66,15 +66,15 @@ To create a :class:`DirectRow `: Unlike the previous string values we've used before, the row key must be ``bytes``. -To create a :class:`ConditionalRow `, -first create a :class:`RowFilter ` and +To create a :class:`ConditionalRow `, +first create a :class:`RowFilter ` and then .. code:: python cond_row = table.row(row_key, filter_=filter_) -To create an :class:`AppendRow ` +To create an :class:`AppendRow ` .. code:: python @@ -95,7 +95,7 @@ Direct Mutations Direct mutations can be added via one of four methods -* :meth:`set_cell() ` allows a +* :meth:`set_cell() ` allows a single value to be written to a column .. code:: python @@ -109,7 +109,7 @@ Direct mutations can be added via one of four methods The value can either be bytes or an integer, which will be converted to bytes as a signed 64-bit integer. -* :meth:`delete_cell() ` deletes +* :meth:`delete_cell() ` deletes all cells (i.e. for all timestamps) in a given column .. code:: python @@ -119,7 +119,7 @@ Direct mutations can be added via one of four methods Remember, this only happens in the ``row`` we are using. If we only want to delete cells from a limited range of time, a - :class:`TimestampRange ` can + :class:`TimestampRange ` can be used .. code:: python @@ -127,9 +127,9 @@ Direct mutations can be added via one of four methods row.delete_cell(column_family_id, column, time_range=time_range) -* :meth:`delete_cells() ` does +* :meth:`delete_cells() ` does the same thing as - :meth:`delete_cell() `, + :meth:`delete_cell() `, but accepts a list of columns in a column family rather than a single one. .. code:: python @@ -138,7 +138,7 @@ Direct mutations can be added via one of four methods time_range=time_range) In addition, if we want to delete cells from every column in a column family, - the special :attr:`ALL_COLUMNS ` + the special :attr:`ALL_COLUMNS ` value can be used .. code:: python @@ -146,7 +146,7 @@ Direct mutations can be added via one of four methods row.delete_cells(column_family_id, row.ALL_COLUMNS, time_range=time_range) -* :meth:`delete() ` will delete the +* :meth:`delete() ` will delete the entire row .. code:: python @@ -177,14 +177,14 @@ Append Mutations Append mutations can be added via one of two methods -* :meth:`append_cell_value() ` +* :meth:`append_cell_value() ` appends a bytes value to an existing cell: .. code:: python append_row.append_cell_value(column_family_id, column, bytes_value) -* :meth:`increment_cell_value() ` +* :meth:`increment_cell_value() ` increments an integer value in an existing cell: .. code:: python @@ -217,7 +217,7 @@ Read Single Row from a Table ---------------------------- To make a `ReadRows`_ API request for a single row key, use -:meth:`Table.read_row() `: +:meth:`Table.read_row() `: .. code:: python @@ -226,34 +226,34 @@ To make a `ReadRows`_ API request for a single row key, use { u'fam1': { b'col1': [ - , - , + , + , ], b'col2': [ - , + , ], }, u'fam2': { b'col3': [ - , - , - , + , + , + , ], }, } >>> cell = row_data.cells[u'fam1'][b'col1'][0] >>> cell - + >>> cell.value b'val1' >>> cell.timestamp datetime.datetime(2016, 2, 27, 3, 41, 18, 122823, tzinfo=) -Rather than returning a :class:`DirectRow ` +Rather than returning a :class:`DirectRow ` or similar class, this method returns a -:class:`PartialRowData ` +:class:`PartialRowData ` instance. This class is used for reading and parsing data rather than for -modifying data (as :class:`DirectRow ` is). +modifying data (as :class:`DirectRow ` is). A filter can also be applied to the results: @@ -262,15 +262,15 @@ A filter can also be applied to the results: row_data = table.read_row(row_key, filter_=filter_val) The allowable ``filter_`` values are the same as those used for a -:class:`ConditionalRow `. For +:class:`ConditionalRow `. For more information, see the -:meth:`Table.read_row() ` documentation. +:meth:`Table.read_row() ` documentation. Stream Many Rows from a Table ----------------------------- To make a `ReadRows`_ API request for a stream of rows, use -:meth:`Table.read_rows() `: +:meth:`Table.read_rows() `: .. code:: python @@ -279,32 +279,32 @@ To make a `ReadRows`_ API request for a stream of rows, use Using gRPC over HTTP/2, a continual stream of responses will be delivered. In particular -* :meth:`consume_next() ` +* :meth:`consume_next() ` pulls the next result from the stream, parses it and stores it on the - :class:`PartialRowsData ` instance -* :meth:`consume_all() ` + :class:`PartialRowsData ` instance +* :meth:`consume_all() ` pulls results from the stream until there are no more -* :meth:`cancel() ` closes +* :meth:`cancel() ` closes the stream -See the :class:`PartialRowsData ` +See the :class:`PartialRowsData ` documentation for more information. As with -:meth:`Table.read_row() `, an optional +:meth:`Table.read_row() `, an optional ``filter_`` can be applied. In addition a ``start_key`` and / or ``end_key`` can be supplied for the stream, a ``limit`` can be set and a boolean ``allow_row_interleaving`` can be specified to allow faster streamed results at the potential cost of non-sequential reads. -See the :meth:`Table.read_rows() ` +See the :meth:`Table.read_rows() ` documentation for more information on the optional arguments. Sample Keys in a Table ---------------------- Make a `SampleRowKeys`_ API request with -:meth:`Table.sample_row_keys() `: +:meth:`Table.sample_row_keys() `: .. code:: python @@ -315,7 +315,7 @@ approximately equal size, which can be used to break up the data for distributed tasks like mapreduces. As with -:meth:`Table.read_rows() `, the +:meth:`Table.read_rows() `, the returned ``keys_iterator`` is connected to a cancellable HTTP/2 stream. The next key in the result can be accessed via diff --git a/docs/encryption-info.rst b/docs/encryption-info.rst index 46f19880f..62b77ea0c 100644 --- a/docs/encryption-info.rst +++ b/docs/encryption-info.rst @@ -1,6 +1,6 @@ Encryption Info ~~~~~~~~~~~~~~~ -.. automodule:: google.cloud.bigtable.encryption_info +.. automodule:: google.cloud.bigtable.deprecated.encryption_info :members: :show-inheritance: diff --git a/docs/instance-api.rst b/docs/instance-api.rst index 88b4eb4dc..78123e8ca 100644 --- a/docs/instance-api.rst +++ b/docs/instance-api.rst @@ -1,7 +1,7 @@ Instance Admin API ================== -After creating a :class:`Client `, you can +After creating a :class:`Client `, you can interact with individual instances for a project. List Instances @@ -9,7 +9,7 @@ List Instances If you want a comprehensive list of all existing instances, make a `ListInstances`_ API request with -:meth:`Client.list_instances() `: +:meth:`Client.list_instances() `: .. code:: python @@ -18,7 +18,7 @@ If you want a comprehensive list of all existing instances, make a Instance Factory ---------------- -To create an :class:`Instance ` object: +To create an :class:`Instance ` object: .. code:: python @@ -40,7 +40,7 @@ Create a new Instance --------------------- After creating the instance object, make a `CreateInstance`_ API request -with :meth:`create() `: +with :meth:`create() `: .. code:: python @@ -54,14 +54,14 @@ Check on Current Operation When modifying an instance (via a `CreateInstance`_ request), the Bigtable API will return a `long-running operation`_ and a corresponding - :class:`Operation ` object + :class:`Operation ` object will be returned by - :meth:`create() `. + :meth:`create() `. You can check if a long-running operation (for a -:meth:`create() ` has finished +:meth:`create() ` has finished by making a `GetOperation`_ request with -:meth:`Operation.finished() `: +:meth:`Operation.finished() `: .. code:: python @@ -71,18 +71,18 @@ by making a `GetOperation`_ request with .. note:: - Once an :class:`Operation ` object + Once an :class:`Operation ` object has returned :data:`True` from - :meth:`finished() `, the + :meth:`finished() `, the object should not be re-used. Subsequent calls to - :meth:`finished() ` + :meth:`finished() ` will result in a :class:`ValueError `. Get metadata for an existing Instance ------------------------------------- After creating the instance object, make a `GetInstance`_ API request -with :meth:`reload() `: +with :meth:`reload() `: .. code:: python @@ -94,7 +94,7 @@ Update an existing Instance --------------------------- After creating the instance object, make an `UpdateInstance`_ API request -with :meth:`update() `: +with :meth:`update() `: .. code:: python @@ -105,7 +105,7 @@ Delete an existing Instance --------------------------- Make a `DeleteInstance`_ API request with -:meth:`delete() `: +:meth:`delete() `: .. code:: python @@ -115,8 +115,8 @@ Next Step --------- Now we go down the hierarchy from -:class:`Instance ` to a -:class:`Table `. +:class:`Instance ` to a +:class:`Table `. Head next to learn about the :doc:`table-api`. diff --git a/docs/instance.rst b/docs/instance.rst index f9be9672f..3a61faf1c 100644 --- a/docs/instance.rst +++ b/docs/instance.rst @@ -1,6 +1,6 @@ Instance ~~~~~~~~ -.. automodule:: google.cloud.bigtable.instance +.. automodule:: google.cloud.bigtable.deprecated.instance :members: :show-inheritance: diff --git a/docs/row-data.rst b/docs/row-data.rst index 503f9b1cb..b9013ebf5 100644 --- a/docs/row-data.rst +++ b/docs/row-data.rst @@ -1,6 +1,6 @@ Row Data ~~~~~~~~ -.. automodule:: google.cloud.bigtable.row_data +.. automodule:: google.cloud.bigtable.deprecated.row_data :members: :show-inheritance: diff --git a/docs/row-filters.rst b/docs/row-filters.rst index 9884ce400..8d1fac46b 100644 --- a/docs/row-filters.rst +++ b/docs/row-filters.rst @@ -2,11 +2,11 @@ Bigtable Row Filters ==================== It is possible to use a -:class:`RowFilter ` +:class:`RowFilter ` when adding mutations to a -:class:`ConditionalRow ` and when -reading row data with :meth:`read_row() ` -or :meth:`read_rows() `. +:class:`ConditionalRow ` and when +reading row data with :meth:`read_row() ` +or :meth:`read_rows() `. As laid out in the `RowFilter definition`_, the following basic filters are provided: @@ -60,8 +60,8 @@ level. For example: ---- -.. automodule:: google.cloud.bigtable.row_filters +.. automodule:: google.cloud.bigtable.deprecated.row_filters :members: :show-inheritance: -.. _RowFilter definition: https://googleapis.dev/python/bigtable/latest/row-filters.html?highlight=rowfilter#google.cloud.bigtable.row_filters.RowFilter +.. _RowFilter definition: https://googleapis.dev/python/bigtable/latest/row-filters.html?highlight=rowfilter#google.cloud.bigtable.deprecated.row_filters.RowFilter diff --git a/docs/row-set.rst b/docs/row-set.rst index 5f7a16a02..92cd107e8 100644 --- a/docs/row-set.rst +++ b/docs/row-set.rst @@ -1,6 +1,6 @@ Row Set ~~~~~~~~ -.. automodule:: google.cloud.bigtable.row_set +.. automodule:: google.cloud.bigtable.deprecated.row_set :members: :show-inheritance: diff --git a/docs/row.rst b/docs/row.rst index 33686608b..e8fa48cdd 100644 --- a/docs/row.rst +++ b/docs/row.rst @@ -1,7 +1,7 @@ Bigtable Row ============ -.. automodule:: google.cloud.bigtable.row +.. automodule:: google.cloud.bigtable.deprecated.row :members: :show-inheritance: :inherited-members: diff --git a/docs/snippets.py b/docs/snippets.py index 1d93fdf12..084f10270 100644 --- a/docs/snippets.py +++ b/docs/snippets.py @@ -16,7 +16,7 @@ """Testable usage examples for Google Cloud Bigtable API wrapper Each example function takes a ``client`` argument (which must be an instance -of :class:`google.cloud.bigtable.client.Client`) and uses it to perform a task +of :class:`google.cloud.bigtable.deprecated.client.Client`) and uses it to perform a task with the API. To facilitate running the examples as system tests, each example is also passed @@ -40,8 +40,8 @@ from test_utils.retry import RetryErrors from google.cloud._helpers import UTC -from google.cloud.bigtable import Client -from google.cloud.bigtable import enums +from google.cloud.bigtable.deprecated import Client +from google.cloud.bigtable.deprecated import enums UNIQUE_SUFFIX = unique_resource_id("-") @@ -110,8 +110,8 @@ def teardown_module(): def test_bigtable_create_instance(): # [START bigtable_api_create_prod_instance] - from google.cloud.bigtable import Client - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import enums my_instance_id = "inst-my-" + UNIQUE_SUFFIX my_cluster_id = "clus-my-" + UNIQUE_SUFFIX @@ -144,8 +144,8 @@ def test_bigtable_create_instance(): def test_bigtable_create_additional_cluster(): # [START bigtable_api_create_cluster] - from google.cloud.bigtable import Client - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import enums # Assuming that there is an existing instance with `INSTANCE_ID` # on the server already. @@ -181,8 +181,8 @@ def test_bigtable_create_reload_delete_app_profile(): import re # [START bigtable_api_create_app_profile] - from google.cloud.bigtable import Client - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import enums routing_policy_type = enums.RoutingPolicyType.ANY @@ -202,7 +202,7 @@ def test_bigtable_create_reload_delete_app_profile(): # [END bigtable_api_create_app_profile] # [START bigtable_api_app_profile_name] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -219,7 +219,7 @@ def test_bigtable_create_reload_delete_app_profile(): assert _profile_name_re.match(app_profile_name) # [START bigtable_api_app_profile_exists] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -230,7 +230,7 @@ def test_bigtable_create_reload_delete_app_profile(): assert app_profile_exists # [START bigtable_api_reload_app_profile] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -241,7 +241,7 @@ def test_bigtable_create_reload_delete_app_profile(): assert app_profile.routing_policy_type == ROUTING_POLICY_TYPE # [START bigtable_api_update_app_profile] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -255,7 +255,7 @@ def test_bigtable_create_reload_delete_app_profile(): assert app_profile.description == description # [START bigtable_api_delete_app_profile] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -269,7 +269,7 @@ def test_bigtable_create_reload_delete_app_profile(): def test_bigtable_list_instances(): # [START bigtable_api_list_instances] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) (instances_list, failed_locations_list) = client.list_instances() @@ -280,7 +280,7 @@ def test_bigtable_list_instances(): def test_bigtable_list_clusters_on_instance(): # [START bigtable_api_list_clusters_on_instance] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -292,7 +292,7 @@ def test_bigtable_list_clusters_on_instance(): def test_bigtable_list_clusters_in_project(): # [START bigtable_api_list_clusters_in_project] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) (clusters_list, failed_locations_list) = client.list_clusters() @@ -309,7 +309,7 @@ def test_bigtable_list_app_profiles(): app_profile = app_profile.create(ignore_warnings=True) # [START bigtable_api_list_app_profiles] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -325,7 +325,7 @@ def test_bigtable_list_app_profiles(): def test_bigtable_instance_exists(): # [START bigtable_api_check_instance_exists] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -337,7 +337,7 @@ def test_bigtable_instance_exists(): def test_bigtable_cluster_exists(): # [START bigtable_api_check_cluster_exists] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -350,7 +350,7 @@ def test_bigtable_cluster_exists(): def test_bigtable_reload_instance(): # [START bigtable_api_reload_instance] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -362,7 +362,7 @@ def test_bigtable_reload_instance(): def test_bigtable_reload_cluster(): # [START bigtable_api_reload_cluster] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -375,7 +375,7 @@ def test_bigtable_reload_cluster(): def test_bigtable_update_instance(): # [START bigtable_api_update_instance] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -389,7 +389,7 @@ def test_bigtable_update_instance(): def test_bigtable_update_cluster(): # [START bigtable_api_update_cluster] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -403,7 +403,7 @@ def test_bigtable_update_cluster(): def test_bigtable_cluster_disable_autoscaling(): # [START bigtable_api_cluster_disable_autoscaling] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -424,8 +424,8 @@ def test_bigtable_create_table(): # [START bigtable_api_create_table] from google.api_core import exceptions from google.api_core import retry - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -450,7 +450,7 @@ def test_bigtable_create_table(): def test_bigtable_list_tables(): # [START bigtable_api_list_tables] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -463,7 +463,7 @@ def test_bigtable_list_tables(): def test_bigtable_delete_cluster(): - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -480,7 +480,7 @@ def test_bigtable_delete_cluster(): operation.result(timeout=1000) # [START bigtable_api_delete_cluster] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -493,7 +493,7 @@ def test_bigtable_delete_cluster(): def test_bigtable_delete_instance(): - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) @@ -515,7 +515,7 @@ def test_bigtable_delete_instance(): INSTANCES_TO_DELETE.append(instance) # [START bigtable_api_delete_instance] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) @@ -531,7 +531,7 @@ def test_bigtable_delete_instance(): def test_bigtable_test_iam_permissions(): # [START bigtable_api_test_iam_permissions] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -547,9 +547,9 @@ def test_bigtable_set_iam_policy_then_get_iam_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_set_iam_policy] - from google.cloud.bigtable import Client - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -563,7 +563,7 @@ def test_bigtable_set_iam_policy_then_get_iam_policy(): assert len(policy_latest.bigtable_admins) > 0 # [START bigtable_api_get_iam_policy] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -577,7 +577,7 @@ def test_bigtable_project_path(): import re # [START bigtable_api_project_path] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) project_path = client.project_path @@ -586,7 +586,7 @@ def test_bigtable_project_path(): def test_bigtable_table_data_client(): # [START bigtable_api_table_data_client] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) table_data_client = client.table_data_client @@ -595,7 +595,7 @@ def test_bigtable_table_data_client(): def test_bigtable_table_admin_client(): # [START bigtable_api_table_admin_client] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) table_admin_client = client.table_admin_client @@ -604,7 +604,7 @@ def test_bigtable_table_admin_client(): def test_bigtable_instance_admin_client(): # [START bigtable_api_instance_admin_client] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance_admin_client = client.instance_admin_client @@ -615,9 +615,9 @@ def test_bigtable_admins_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_admins_policy] - from google.cloud.bigtable import Client - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -636,9 +636,9 @@ def test_bigtable_readers_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_readers_policy] - from google.cloud.bigtable import Client - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_READER_ROLE + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_READER_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -657,9 +657,9 @@ def test_bigtable_users_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_users_policy] - from google.cloud.bigtable import Client - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_USER_ROLE + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_USER_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -678,9 +678,9 @@ def test_bigtable_viewers_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_viewers_policy] - from google.cloud.bigtable import Client - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_VIEWER_ROLE + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_VIEWER_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -699,7 +699,7 @@ def test_bigtable_instance_name(): import re # [START bigtable_api_instance_name] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -711,7 +711,7 @@ def test_bigtable_cluster_name(): import re # [START bigtable_api_cluster_name] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -722,7 +722,7 @@ def test_bigtable_cluster_name(): def test_bigtable_instance_from_pb(): # [START bigtable_api_instance_from_pb] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 client = Client(admin=True) @@ -741,7 +741,7 @@ def test_bigtable_instance_from_pb(): def test_bigtable_cluster_from_pb(): # [START bigtable_api_cluster_from_pb] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 client = Client(admin=True) @@ -767,7 +767,7 @@ def test_bigtable_cluster_from_pb(): def test_bigtable_instance_state(): # [START bigtable_api_instance_state] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -779,7 +779,7 @@ def test_bigtable_instance_state(): def test_bigtable_cluster_state(): # [START bigtable_api_cluster_state] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) diff --git a/docs/snippets_table.py b/docs/snippets_table.py index f27260425..72c342907 100644 --- a/docs/snippets_table.py +++ b/docs/snippets_table.py @@ -16,7 +16,7 @@ """Testable usage examples for Google Cloud Bigtable API wrapper Each example function takes a ``client`` argument (which must be an instance -of :class:`google.cloud.bigtable.client.Client`) and uses it to perform a task +of :class:`google.cloud.bigtable.deprecated.client.Client`) and uses it to perform a task with the API. To facilitate running the examples as system tests, each example is also passed @@ -38,9 +38,9 @@ from test_utils.retry import RetryErrors from google.cloud._helpers import UTC -from google.cloud.bigtable import Client -from google.cloud.bigtable import enums -from google.cloud.bigtable import column_family +from google.cloud.bigtable.deprecated import Client +from google.cloud.bigtable.deprecated import enums +from google.cloud.bigtable.deprecated import column_family INSTANCE_ID = "snippet" + unique_resource_id("-") @@ -113,8 +113,8 @@ def teardown_module(): def test_bigtable_create_table(): # [START bigtable_api_create_table] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -143,7 +143,7 @@ def test_bigtable_sample_row_keys(): assert table_sample.exists() # [START bigtable_api_sample_row_keys] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -159,7 +159,7 @@ def test_bigtable_sample_row_keys(): def test_bigtable_write_read_drop_truncate(): # [START bigtable_api_mutate_rows] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -190,7 +190,7 @@ def test_bigtable_write_read_drop_truncate(): # [END bigtable_api_mutate_rows] assert len(response) == len(rows) # [START bigtable_api_read_row] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -200,7 +200,7 @@ def test_bigtable_write_read_drop_truncate(): # [END bigtable_api_read_row] assert row.row_key.decode("utf-8") == row_key # [START bigtable_api_read_rows] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -218,7 +218,7 @@ def test_bigtable_write_read_drop_truncate(): # [END bigtable_api_read_rows] assert len(total_rows) == len(rows) # [START bigtable_api_drop_by_prefix] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -231,7 +231,7 @@ def test_bigtable_write_read_drop_truncate(): assert row.row_key.decode("utf-8") not in dropped_row_keys # [START bigtable_api_truncate_table] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -246,7 +246,7 @@ def test_bigtable_write_read_drop_truncate(): def test_bigtable_mutations_batcher(): # [START bigtable_api_mutations_batcher] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -297,7 +297,7 @@ def test_bigtable_mutations_batcher(): def test_bigtable_table_column_family(): # [START bigtable_api_table_column_family] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -311,7 +311,7 @@ def test_bigtable_table_column_family(): def test_bigtable_list_tables(): # [START bigtable_api_list_tables] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -324,7 +324,7 @@ def test_bigtable_table_name(): import re # [START bigtable_api_table_name] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -342,7 +342,7 @@ def test_bigtable_table_name(): def test_bigtable_list_column_families(): # [START bigtable_api_list_column_families] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -356,7 +356,7 @@ def test_bigtable_list_column_families(): def test_bigtable_get_cluster_states(): # [START bigtable_api_get_cluster_states] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -374,7 +374,7 @@ def test_bigtable_table_test_iam_permissions(): assert table_policy.exists # [START bigtable_api_table_test_iam_permissions] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -392,9 +392,9 @@ def test_bigtable_table_set_iam_policy_then_get_iam_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_table_set_iam_policy] - from google.cloud.bigtable import Client - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -407,7 +407,7 @@ def test_bigtable_table_set_iam_policy_then_get_iam_policy(): assert len(policy_latest.bigtable_admins) > 0 # [START bigtable_api_table_get_iam_policy] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -419,7 +419,7 @@ def test_bigtable_table_set_iam_policy_then_get_iam_policy(): def test_bigtable_table_exists(): # [START bigtable_api_check_table_exists] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -435,7 +435,7 @@ def test_bigtable_delete_table(): assert table_del.exists() # [START bigtable_api_delete_table] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -448,7 +448,7 @@ def test_bigtable_delete_table(): def test_bigtable_table_row(): # [START bigtable_api_table_row] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -475,7 +475,7 @@ def test_bigtable_table_row(): def test_bigtable_table_append_row(): # [START bigtable_api_table_append_row] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -502,7 +502,7 @@ def test_bigtable_table_append_row(): def test_bigtable_table_direct_row(): # [START bigtable_api_table_direct_row] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -529,8 +529,8 @@ def test_bigtable_table_direct_row(): def test_bigtable_table_conditional_row(): # [START bigtable_api_table_conditional_row] - from google.cloud.bigtable import Client - from google.cloud.bigtable.row_filters import PassAllFilter + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.row_filters import PassAllFilter client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -558,7 +558,7 @@ def test_bigtable_table_conditional_row(): def test_bigtable_column_family_name(): # [START bigtable_api_column_family_name] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -581,8 +581,8 @@ def test_bigtable_column_family_name(): def test_bigtable_create_update_delete_column_family(): # [START bigtable_api_create_column_family] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -598,8 +598,8 @@ def test_bigtable_create_update_delete_column_family(): assert column_families[column_family_id].gc_rule == gc_rule # [START bigtable_api_update_column_family] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -617,8 +617,8 @@ def test_bigtable_create_update_delete_column_family(): assert updated_families[column_family_id].gc_rule == max_age_rule # [START bigtable_api_delete_column_family] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -653,8 +653,8 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): Config.TABLE.mutate_rows(rows) # [START bigtable_api_add_row_key] - from google.cloud.bigtable import Client - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.row_set import RowSet client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -670,9 +670,9 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): assert found_row_keys == expected_row_keys # [START bigtable_api_add_row_range] - from google.cloud.bigtable import Client - from google.cloud.bigtable.row_set import RowSet - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -688,8 +688,8 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): assert found_row_keys == expected_row_keys # [START bigtable_api_row_range_from_keys] - from google.cloud.bigtable import Client - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.row_set import RowSet client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -723,8 +723,8 @@ def test_bigtable_add_row_range_with_prefix(): Config.TABLE.mutate_rows(rows) # [START bigtable_api_add_row_range_with_prefix] - from google.cloud.bigtable import Client - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated.row_set import RowSet client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -747,7 +747,7 @@ def test_bigtable_add_row_range_with_prefix(): def test_bigtable_batcher_mutate_flush_mutate_rows(): # [START bigtable_api_batcher_mutate] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -769,7 +769,7 @@ def test_bigtable_batcher_mutate_flush_mutate_rows(): # [END bigtable_api_batcher_mutate] # [START bigtable_api_batcher_flush] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -795,7 +795,7 @@ def test_bigtable_batcher_mutate_flush_mutate_rows(): table.truncate(timeout=200) # [START bigtable_api_batcher_mutate_rows] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -829,8 +829,8 @@ def test_bigtable_batcher_mutate_flush_mutate_rows(): def test_bigtable_create_family_gc_max_age(): # [START bigtable_api_create_family_gc_max_age] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -851,8 +851,8 @@ def test_bigtable_create_family_gc_max_age(): def test_bigtable_create_family_gc_max_versions(): # [START bigtable_api_create_family_gc_max_versions] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -872,8 +872,8 @@ def test_bigtable_create_family_gc_max_versions(): def test_bigtable_create_family_gc_union(): # [START bigtable_api_create_family_gc_union] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -898,8 +898,8 @@ def test_bigtable_create_family_gc_union(): def test_bigtable_create_family_gc_intersection(): # [START bigtable_api_create_family_gc_intersection] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -927,8 +927,8 @@ def test_bigtable_create_family_gc_intersection(): def test_bigtable_create_family_gc_nested(): # [START bigtable_api_create_family_gc_nested] - from google.cloud.bigtable import Client - from google.cloud.bigtable import column_family + from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable.deprecated import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -978,7 +978,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): row.commit() # [START bigtable_api_row_data_cells] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -993,7 +993,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): assert actual_cell_value == value # [START bigtable_api_row_cell_value] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1006,7 +1006,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): assert cell_value == value # [START bigtable_api_row_cell_values] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1025,7 +1025,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): row.commit() # [START bigtable_api_row_find_cells] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1042,7 +1042,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): def test_bigtable_row_setcell_rowkey(): # [START bigtable_api_row_set_cell] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1061,7 +1061,7 @@ def test_bigtable_row_setcell_rowkey(): assert status.code == 0 # [START bigtable_api_row_row_key] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1073,7 +1073,7 @@ def test_bigtable_row_setcell_rowkey(): assert row_key == ROW_KEY1 # [START bigtable_api_row_table] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1098,7 +1098,7 @@ def test_bigtable_row_delete(): assert written_row_keys == [b"row_key_1"] # [START bigtable_api_row_delete] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1130,7 +1130,7 @@ def test_bigtable_row_delete_cell(): assert written_row_keys == [row_key1] # [START bigtable_api_row_delete_cell] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1163,7 +1163,7 @@ def test_bigtable_row_delete_cells(): assert written_row_keys == [row_key1] # [START bigtable_api_row_delete_cells] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1189,7 +1189,7 @@ def test_bigtable_row_clear(): assert mutation_size > 0 # [START bigtable_api_row_clear] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1208,7 +1208,7 @@ def test_bigtable_row_clear(): def test_bigtable_row_clear_get_mutations_size(): # [START bigtable_api_row_get_mutations_size] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1230,7 +1230,7 @@ def test_bigtable_row_clear_get_mutations_size(): def test_bigtable_row_setcell_commit_rowkey(): # [START bigtable_api_row_set_cell] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1244,7 +1244,7 @@ def test_bigtable_row_setcell_commit_rowkey(): row_obj.commit() # [START bigtable_api_row_commit] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1264,7 +1264,7 @@ def test_bigtable_row_setcell_commit_rowkey(): assert written_row_keys == [b"row_key_1", b"row_key_2"] # [START bigtable_api_row_row_key] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1286,7 +1286,7 @@ def test_bigtable_row_append_cell_value(): row.commit() # [START bigtable_api_row_append_cell_value] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1303,7 +1303,7 @@ def test_bigtable_row_append_cell_value(): assert actual_value == cell_val1 + cell_val2 # [START bigtable_api_row_commit] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1315,7 +1315,7 @@ def test_bigtable_row_append_cell_value(): # [END bigtable_api_row_commit] # [START bigtable_api_row_increment_cell_value] - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) diff --git a/docs/table-api.rst b/docs/table-api.rst index 1bbf85146..ce05a3419 100644 --- a/docs/table-api.rst +++ b/docs/table-api.rst @@ -1,7 +1,7 @@ Table Admin API =============== -After creating an :class:`Instance `, you can +After creating an :class:`Instance `, you can interact with individual tables, groups of tables or column families within a table. @@ -10,33 +10,33 @@ List Tables If you want a comprehensive list of all existing tables in a instance, make a `ListTables`_ API request with -:meth:`Instance.list_tables() `: +:meth:`Instance.list_tables() `: .. code:: python >>> instance.list_tables() - [, - ] + [, + ] Table Factory ------------- -To create a :class:`Table ` object: +To create a :class:`Table ` object: .. code:: python table = instance.table(table_id) -Even if this :class:`Table ` already +Even if this :class:`Table ` already has been created with the API, you'll want this object to use as a -parent of a :class:`ColumnFamily ` -or :class:`Row `. +parent of a :class:`ColumnFamily ` +or :class:`Row `. Create a new Table ------------------ After creating the table object, make a `CreateTable`_ API request -with :meth:`create() `: +with :meth:`create() `: .. code:: python @@ -53,7 +53,7 @@ Delete an existing Table ------------------------ Make a `DeleteTable`_ API request with -:meth:`delete() `: +:meth:`delete() `: .. code:: python @@ -67,7 +67,7 @@ associated with a table, the `GetTable`_ API method returns a table object with the names of the column families. To retrieve the list of column families use -:meth:`list_column_families() `: +:meth:`list_column_families() `: .. code:: python @@ -77,7 +77,7 @@ Column Family Factory --------------------- To create a -:class:`ColumnFamily ` object: +:class:`ColumnFamily ` object: .. code:: python @@ -87,7 +87,7 @@ There is no real reason to use this factory unless you intend to create or delete a column family. In addition, you can specify an optional ``gc_rule`` (a -:class:`GarbageCollectionRule ` +:class:`GarbageCollectionRule ` or similar): .. code:: python @@ -99,7 +99,7 @@ This rule helps the backend determine when and how to clean up old cells in the column family. See :doc:`column-family` for more information about -:class:`GarbageCollectionRule ` +:class:`GarbageCollectionRule ` and related classes. Create a new Column Family @@ -107,7 +107,7 @@ Create a new Column Family After creating the column family object, make a `CreateColumnFamily`_ API request with -:meth:`ColumnFamily.create() ` +:meth:`ColumnFamily.create() ` .. code:: python @@ -117,7 +117,7 @@ Delete an existing Column Family -------------------------------- Make a `DeleteColumnFamily`_ API request with -:meth:`ColumnFamily.delete() ` +:meth:`ColumnFamily.delete() ` .. code:: python @@ -127,7 +127,7 @@ Update an existing Column Family -------------------------------- Make an `UpdateColumnFamily`_ API request with -:meth:`ColumnFamily.delete() ` +:meth:`ColumnFamily.delete() ` .. code:: python @@ -137,9 +137,9 @@ Next Step --------- Now we go down the final step of the hierarchy from -:class:`Table ` to -:class:`Row ` as well as streaming -data directly via a :class:`Table `. +:class:`Table ` to +:class:`Row ` as well as streaming +data directly via a :class:`Table `. Head next to learn about the :doc:`data-api`. diff --git a/docs/table.rst b/docs/table.rst index c230725d1..0d938e0af 100644 --- a/docs/table.rst +++ b/docs/table.rst @@ -1,6 +1,6 @@ Table ~~~~~ -.. automodule:: google.cloud.bigtable.table +.. automodule:: google.cloud.bigtable.deprecated.table :members: :show-inheritance: diff --git a/docs/usage.rst b/docs/usage.rst index 33bf7bb7f..80fb65898 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -21,12 +21,12 @@ Using the API In the hierarchy of API concepts -* a :class:`Client ` owns an - :class:`Instance ` -* an :class:`Instance ` owns a - :class:`Table ` -* a :class:`Table ` owns a - :class:`ColumnFamily ` -* a :class:`Table ` owns a - :class:`Row ` +* a :class:`Client ` owns an + :class:`Instance ` +* an :class:`Instance ` owns a + :class:`Table ` +* a :class:`Table ` owns a + :class:`ColumnFamily ` +* a :class:`Table ` owns a + :class:`Row ` (and all the cells in the row) diff --git a/google/cloud/bigtable/__init__.py b/google/cloud/bigtable/__init__.py index 7331ff241..daa562c0c 100644 --- a/google/cloud/bigtable/__init__.py +++ b/google/cloud/bigtable/__init__.py @@ -1,4 +1,5 @@ -# Copyright 2015 Google LLC +# -*- coding: utf-8 -*- +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,15 +12,44 @@ # 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. +# -"""Google Cloud Bigtable API package.""" - -from google.cloud.bigtable.client import Client +from typing import List, Tuple from google.cloud.bigtable import gapic_version as package_version -__version__: str +from google.cloud.bigtable.client import BigtableDataClient +from google.cloud.bigtable.client import Table + +from google.cloud.bigtable.read_rows_query import ReadRowsQuery +from google.cloud.bigtable.row_response import RowResponse +from google.cloud.bigtable.row_response import CellResponse + +from google.cloud.bigtable.mutations_batcher import MutationsBatcher +from google.cloud.bigtable.mutations import Mutation +from google.cloud.bigtable.mutations import BulkMutationsEntry +from google.cloud.bigtable.mutations import SetCell +from google.cloud.bigtable.mutations import DeleteRangeFromColumn +from google.cloud.bigtable.mutations import DeleteAllFromFamily +from google.cloud.bigtable.mutations import DeleteAllFromRow + +# Type alias for the output of sample_keys +RowKeySamples = List[Tuple[bytes, int]] -__version__ = package_version.__version__ +__version__: str = package_version.__version__ -__all__ = ["__version__", "Client"] +__all__ = ( + "BigtableDataClient", + "Table", + "RowKeySamples", + "ReadRowsQuery", + "MutationsBatcher", + "Mutation", + "BulkMutationsEntry", + "SetCell", + "DeleteRangeFromColumn", + "DeleteAllFromFamily", + "DeleteAllFromRow", + "RowResponse", + "CellResponse", +) diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index c82a268c6..df4bf308f 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -1,4 +1,4 @@ -# Copyright 2015 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,503 +11,427 @@ # 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. +# + +from __future__ import annotations + +from typing import Any, AsyncIterable, TYPE_CHECKING + +from google.cloud.client import ClientWithProject + -"""Parent client for calling the Google Cloud Bigtable API. - -This is the base from which all interactions with the API occur. - -In the hierarchy of API concepts - -* a :class:`~google.cloud.bigtable.client.Client` owns an - :class:`~google.cloud.bigtable.instance.Instance` -* an :class:`~google.cloud.bigtable.instance.Instance` owns a - :class:`~google.cloud.bigtable.table.Table` -* a :class:`~google.cloud.bigtable.table.Table` owns a - :class:`~.column_family.ColumnFamily` -* a :class:`~google.cloud.bigtable.table.Table` owns a - :class:`~google.cloud.bigtable.row.Row` (and all the cells in the row) -""" -import os -import warnings -import grpc # type: ignore - -from google.api_core.gapic_v1 import client_info as client_info_lib -import google.auth # type: ignore -from google.auth.credentials import AnonymousCredentials # type: ignore - -from google.cloud import bigtable_v2 -from google.cloud import bigtable_admin_v2 -from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport -from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin.transports import ( - BigtableInstanceAdminGrpcTransport, -) -from google.cloud.bigtable_admin_v2.services.bigtable_table_admin.transports import ( - BigtableTableAdminGrpcTransport, -) - -from google.cloud import bigtable -from google.cloud.bigtable.instance import Instance -from google.cloud.bigtable.cluster import Cluster - -from google.cloud.client import ClientWithProject # type: ignore - -from google.cloud.bigtable_admin_v2.types import instance -from google.cloud.bigtable.cluster import _CLUSTER_NAME_RE -from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore - - -INSTANCE_TYPE_PRODUCTION = instance.Instance.Type.PRODUCTION -INSTANCE_TYPE_DEVELOPMENT = instance.Instance.Type.DEVELOPMENT -INSTANCE_TYPE_UNSPECIFIED = instance.Instance.Type.TYPE_UNSPECIFIED -SPANNER_ADMIN_SCOPE = "https://www.googleapis.com/auth/spanner.admin" -ADMIN_SCOPE = "https://www.googleapis.com/auth/bigtable.admin" -"""Scope for interacting with the Cluster Admin and Table Admin APIs.""" -DATA_SCOPE = "https://www.googleapis.com/auth/bigtable.data" -"""Scope for reading and writing table data.""" -READ_ONLY_SCOPE = "https://www.googleapis.com/auth/bigtable.data.readonly" -"""Scope for reading table data.""" - -_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator" -_GRPC_CHANNEL_OPTIONS = ( - ("grpc.max_send_message_length", -1), - ("grpc.max_receive_message_length", -1), - ("grpc.keepalive_time_ms", 30000), - ("grpc.keepalive_timeout_ms", 10000), -) - - -def _create_gapic_client(client_class, client_options=None, transport=None): - def inner(self): - return client_class( - credentials=None, - client_info=self._client_info, - client_options=client_options, - transport=transport, - ) - - return inner - - -class Client(ClientWithProject): - """Client for interacting with Google Cloud Bigtable API. - - .. note:: - - Since the Cloud Bigtable API requires the gRPC transport, no - ``_http`` argument is accepted by this class. - - :type project: :class:`str` or :func:`unicode ` - :param project: (Optional) The ID of the project which owns the - instances, tables and data. If not provided, will - attempt to determine from the environment. - - :type credentials: :class:`~google.auth.credentials.Credentials` - :param credentials: (Optional) The OAuth2 Credentials to use for this - client. If not passed, falls back to the default - inferred from the environment. - - :type read_only: bool - :param read_only: (Optional) Boolean indicating if the data scope should be - for reading only (or for writing as well). Defaults to - :data:`False`. - - :type admin: bool - :param admin: (Optional) Boolean indicating if the client will be used to - interact with the Instance Admin or Table Admin APIs. This - requires the :const:`ADMIN_SCOPE`. Defaults to :data:`False`. - - :type: client_info: :class:`google.api_core.gapic_v1.client_info.ClientInfo` - :param client_info: - The client info used to send a user-agent string along with API - requests. If ``None``, then default info will be used. Generally, - you only need to set this if you're developing your own library - or partner tool. - - :type client_options: :class:`~google.api_core.client_options.ClientOptions` - or :class:`dict` - :param client_options: (Optional) Client options used to set user options - on the client. API Endpoint should be set through client_options. - - :type admin_client_options: - :class:`~google.api_core.client_options.ClientOptions` or :class:`dict` - :param admin_client_options: (Optional) Client options used to set user - options on the client. API Endpoint for admin operations should be set - through admin_client_options. - - :type channel: :instance: grpc.Channel - :param channel (grpc.Channel): (Optional) DEPRECATED: - A ``Channel`` instance through which to make calls. - This argument is mutually exclusive with ``credentials``; - providing both will raise an exception. No longer used. - - :raises: :class:`ValueError ` if both ``read_only`` - and ``admin`` are :data:`True` +import google.auth.credentials + +if TYPE_CHECKING: + from google.cloud.bigtable.mutations import Mutation, BulkMutationsEntry + from google.cloud.bigtable.mutations_batcher import MutationsBatcher + from google.cloud.bigtable.row_response import RowResponse + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable import RowKeySamples + from google.cloud.bigtable.row_filters import RowFilter + from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule + + +class BigtableDataClient(ClientWithProject): + def __init__( + self, + *, + project: str | None = None, + pool_size: int = 3, + credentials: google.auth.credentials.Credentials | None = None, + client_options: dict[str, Any] + | "google.api_core.client_options.ClientOptions" + | None = None, + metadata: list[tuple[str, str]] | None = None, + ): + """ + Create a client instance for the Bigtable Data API + + Args: + project: the project which the client acts on behalf of. + If not passed, falls back to the default inferred + from the environment. + pool_size: The number of grpc channels to maintain + in the internal channel pool. + credentials: + Thehe OAuth2 Credentials to use for this + client. If not passed (and if no ``_http`` object is + passed), falls back to the default inferred from the + environment. + client_options (Optional[Union[dict, google.api_core.client_options.ClientOptions]]): + Client options used to set user options + on the client. API Endpoint should be set through client_options. + metadata: a list of metadata headers to be attached to all calls with this client + """ + raise NotImplementedError + + def get_table( + self, instance_id: str, table_id: str, app_profile_id: str | None = None + ) -> Table: + """ + Return a Table instance to make API requests for a specific table. + + Args: + instance_id: The ID of the instance that owns the table. + table_id: The ID of the table. + app_profile_id: (Optional) The app profile to associate with requests. + https://cloud.google.com/bigtable/docs/app-profiles + """ + raise NotImplementedError + + +class Table: """ + Main Data API surface - _table_data_client = None - _table_admin_client = None - _instance_admin_client = None + Table object maintains instance_id, table_id, and app_profile_id context, and passes them with + each call + """ def __init__( self, - project=None, - credentials=None, - read_only=False, - admin=False, - client_info=None, - client_options=None, - admin_client_options=None, - channel=None, + client: BigtableDataClient, + instance_id: str, + table_id: str, + app_profile_id: str | None = None, ): - if client_info is None: - client_info = client_info_lib.ClientInfo( - client_library_version=bigtable.__version__, - ) - if read_only and admin: - raise ValueError( - "A read-only client cannot also perform" "administrative actions." - ) - - # NOTE: We set the scopes **before** calling the parent constructor. - # It **may** use those scopes in ``with_scopes_if_required``. - self._read_only = bool(read_only) - self._admin = bool(admin) - self._client_info = client_info - self._emulator_host = os.getenv(BIGTABLE_EMULATOR) - - if self._emulator_host is not None: - if credentials is None: - credentials = AnonymousCredentials() - if project is None: - project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT - - if channel is not None: - warnings.warn( - "'channel' is deprecated and no longer used.", - DeprecationWarning, - stacklevel=2, - ) - - self._client_options = client_options - self._admin_client_options = admin_client_options - self._channel = channel - self.SCOPE = self._get_scopes() - super(Client, self).__init__( - project=project, - credentials=credentials, - client_options=client_options, - ) - - def _get_scopes(self): - """Get the scopes corresponding to admin / read-only state. + raise NotImplementedError + + async def read_rows_stream( + self, + query: ReadRowsQuery | dict[str, Any], + *, + shard: bool = False, + limit: int | None, + cache_size_limit: int | None = None, + operation_timeout: int | float | None = 60, + per_row_timeout: int | float | None = 10, + idle_timeout: int | float | None = 300, + per_request_timeout: int | float | None = None, + metadata: list[tuple[str, str]] | None = None, + ) -> AsyncIterable[RowResponse]: + """ + Returns a generator to asynchronously stream back row data. + + Failed requests within operation_timeout and operation_deadline policies will be retried. + + By default, row data is streamed eagerly over the network, and fully cached in memory + in the generator, which can be consumed as needed. The size of the generator cache can + be configured with cache_size_limit. When the cache is full, the read_rows_stream will pause + the network stream until space is available + + Args: + - query: contains details about which rows to return + - shard: if True, will attempt to split up and distribute query to multiple + backend nodes in parallel + - limit: a limit on the number of rows to return. Actual limit will be + min(limit, query.limit) + - cache_size: the number of rows to cache in memory. If None, no limits. + Defaults to None + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + time is only counted while actively waiting on the network. + Completed and cached results can still be accessed after the deadline is complete, + with a DeadlineExceeded exception only raised after cached results are exhausted + - per_row_timeout: the time budget for a single row read, in seconds. If a row takes + longer than per_row_timeout to complete, the ongoing network request will be with a + DeadlineExceeded exception, and a retry may be attempted + Applies only to the underlying network call. + - idle_timeout: the number of idle seconds before an active generator is marked as + stale and the cache is drained. The idle count is reset each time the generator + is yielded from + raises DeadlineExceeded on future yields + - per_request_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted + - metadata: Strings which should be sent along with the request as metadata headers. Returns: - Tuple[str, ...]: The tuple of scopes. + - an asynchronous generator that yields rows returned by the query + Raises: + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions + from any retries that failed + - IdleTimeout: if generator was abandoned + """ + raise NotImplementedError + + async def read_rows( + self, + query: ReadRowsQuery | dict[str, Any], + *, + shard: bool = False, + limit: int | None, + operation_timeout: int | float | None = 60, + per_row_timeout: int | float | None = 10, + per_request_timeout: int | float | None = None, + metadata: list[tuple[str, str]] | None = None, + ) -> list[RowResponse]: """ - if self._read_only: - scopes = (READ_ONLY_SCOPE,) - else: - scopes = (DATA_SCOPE,) + Helper function that returns a full list instead of a generator - if self._admin: - scopes += (ADMIN_SCOPE,) + See read_rows_stream - return scopes + Returns: + - a list of the rows returned by the query + """ + raise NotImplementedError - def _emulator_channel(self, transport, options): - """Create a channel using self._credentials + async def read_row( + self, + row_key: str | bytes, + *, + operation_timeout: int | float | None = 60, + per_request_timeout: int | float | None = None, + metadata: list[tuple[str, str]] | None = None, + ) -> RowResponse: + """ + Helper function to return a single row - Works in a similar way to ``grpc.secure_channel`` but using - ``grpc.local_channel_credentials`` rather than - ``grpc.ssh_channel_credentials`` to allow easy connection to a - local emulator. + See read_rows_stream Returns: - grpc.Channel or grpc.aio.Channel + - the individual row requested """ - # TODO: Implement a special credentials type for emulator and use - # "transport.create_channel" to create gRPC channels once google-auth - # extends it's allowed credentials types. - # Note: this code also exists in the firestore client. - if "GrpcAsyncIOTransport" in str(transport.__name__): - return grpc.aio.secure_channel( - self._emulator_host, - self._local_composite_credentials(), - options=options, - ) - else: - return grpc.secure_channel( - self._emulator_host, - self._local_composite_credentials(), - options=options, - ) - - def _local_composite_credentials(self): - """Create credentials for the local emulator channel. - - :return: grpc.ChannelCredentials + raise NotImplementedError + + async def read_rows_sharded( + self, + query_list: list[ReadRowsQuery] | list[dict[str, Any]], + *, + limit: int | None, + cache_size_limit: int | None = None, + operation_timeout: int | float | None = 60, + per_row_timeout: int | float | None = 10, + idle_timeout: int | float | None = 300, + per_request_timeout: int | float | None = None, + metadata: list[tuple[str, str]] | None = None, + ) -> AsyncIterable[RowResponse]: """ - credentials = google.auth.credentials.with_scopes_if_required( - self._credentials, None - ) - request = google.auth.transport.requests.Request() - - # Create the metadata plugin for inserting the authorization header. - metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin( - credentials, request - ) - - # Create a set of grpc.CallCredentials using the metadata plugin. - google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin) - - # Using the local_credentials to allow connection to emulator - local_credentials = grpc.local_channel_credentials() - - # Combine the local credentials and the authorization credentials. - return grpc.composite_channel_credentials( - local_credentials, google_auth_credentials - ) - - def _create_gapic_client_channel(self, client_class, grpc_transport): - if self._emulator_host is not None: - api_endpoint = self._emulator_host - elif self._client_options and self._client_options.api_endpoint: - api_endpoint = self._client_options.api_endpoint - else: - api_endpoint = client_class.DEFAULT_ENDPOINT - - if self._emulator_host is not None: - channel = self._emulator_channel( - transport=grpc_transport, - options=_GRPC_CHANNEL_OPTIONS, - ) - else: - channel = grpc_transport.create_channel( - host=api_endpoint, - credentials=self._credentials, - options=_GRPC_CHANNEL_OPTIONS, - ) - return grpc_transport(channel=channel, host=api_endpoint) - - @property - def project_path(self): - """Project name to be used with Instance Admin API. - - .. note:: - - This property will not change if ``project`` does not, but the - return value is not cached. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_project_path] - :end-before: [END bigtable_api_project_path] - :dedent: 4 - - The project name is of the form - - ``"projects/{project}"`` - - :rtype: str - :returns: Return a fully-qualified project string. + Runs a sharded query in parallel + + Each query in query list will be run concurrently, with results yielded as they are ready + yielded results may be out of order + + Args: + - query_list: a list of queries to run in parallel + """ + raise NotImplementedError + + async def row_exists( + self, + row_key: str | bytes, + *, + operation_timeout: int | float | None = 60, + per_request_timeout: int | float | None = None, + metadata: list[tuple[str, str]] | None = None, + ) -> bool: + """ + Helper function to determine if a row exists + + uses the filters: chain(limit cells per row = 1, strip value) + + Returns: + - a bool indicating whether the row exists + """ + raise NotImplementedError + + async def sample_keys( + self, + *, + operation_timeout: int | float | None = 60, + per_sample_timeout: int | float | None = 10, + per_request_timeout: int | float | None = None, + metadata: list[tuple[str, str]] | None = None, + ) -> RowKeySamples: + """ + Return a set of RowKeySamples that delimit contiguous sections of the table of + approximately equal size + + RowKeySamples output can be used with ReadRowsQuery.shard() to create a sharded query that + can be parallelized across multiple backend nodes read_rows and read_rows_stream + requests will call sample_keys internally for this purpose when sharding is enabled + + RowKeySamples is simply a type alias for list[tuple[bytes, int]]; a list of + row_keys, along with offset positions in the table + + Returns: + - a set of RowKeySamples the delimit contiguous sections of the table + Raises: + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing all GoogleAPIError + exceptions from any retries that failed """ - return self.instance_admin_client.common_project_path(self.project) + raise NotImplementedError - @property - def table_data_client(self): - """Getter for the gRPC stub used for the Table Admin API. + def mutations_batcher(self, **kwargs) -> MutationsBatcher: + """ + Returns a new mutations batcher instance. + + Can be used to iteratively add mutations that are flushed as a group, + to avoid excess network calls - For example: + Returns: + - a MutationsBatcher context manager that can batch requests + """ + return MutationsBatcher(self, **kwargs) - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_table_data_client] - :end-before: [END bigtable_api_table_data_client] - :dedent: 4 + async def mutate_row( + self, + row_key: str | bytes, + mutations: list[Mutation] | Mutation, + *, + operation_timeout: int | float | None = 60, + per_request_timeout: int | float | None = None, + metadata: list[tuple[str, str]] | None = None, + ): + """ + Mutates a row atomically. + + Cells already present in the row are left unchanged unless explicitly changed + by ``mutation``. + + Idempotent operations (i.e, all mutations have an explicit timestamp) will be + retried on server failure. Non-idempotent operations will not. + + Args: + - row_key: the row to apply mutations to + - mutations: the set of mutations to apply to the row + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + time is only counted while actively waiting on the network. + DeadlineExceeded exception raised after timeout + - per_request_timeout: the time budget for an individual network request, + in seconds. If it takes longer than this time to complete, the request + will be cancelled with a DeadlineExceeded exception, and a retry will be + attempted if within operation_timeout budget + - metadata: Strings which should be sent along with the request as metadata headers. + + Raises: + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing all + GoogleAPIError exceptions from any retries that failed + - GoogleAPIError: raised on non-idempotent operations that cannot be + safely retried. + """ + raise NotImplementedError - :rtype: :class:`.bigtable_v2.BigtableClient` - :returns: A BigtableClient object. + async def bulk_mutate_rows( + self, + mutation_entries: list[BulkMutationsEntry], + *, + operation_timeout: int | float | None = 60, + per_request_timeout: int | float | None = None, + metadata: list[tuple[str, str]] | None = None, + ): """ - if self._table_data_client is None: - transport = self._create_gapic_client_channel( - bigtable_v2.BigtableClient, - BigtableGrpcTransport, - ) - klass = _create_gapic_client( - bigtable_v2.BigtableClient, - client_options=self._client_options, - transport=transport, - ) - self._table_data_client = klass(self) - return self._table_data_client - - @property - def table_admin_client(self): - """Getter for the gRPC stub used for the Table Admin API. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_table_admin_client] - :end-before: [END bigtable_api_table_admin_client] - :dedent: 4 - - :rtype: :class:`.bigtable_admin_pb2.BigtableTableAdmin` - :returns: A BigtableTableAdmin instance. - :raises: :class:`ValueError ` if the current - client is not an admin client or if it has not been - :meth:`start`-ed. + Applies mutations for multiple rows in a single batched request. + + Each individual BulkMutationsEntry is applied atomically, but separate entries + may be applied in arbitrary order (even for entries targetting the same row) + In total, the row_mutations can contain at most 100000 individual mutations + across all entries + + Idempotent entries (i.e., entries with mutations with explicit timestamps) + will be retried on failure. Non-idempotent will not, and will reported in a + raised exception group + + Args: + - mutation_entries: the batches of mutations to apply + Each entry will be applied atomically, but entries will be applied + in arbitrary order + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + time is only counted while actively waiting on the network. + DeadlineExceeded exception raised after timeout + - per_request_timeout: the time budget for an individual network request, + in seconds. If it takes longer than this time to complete, the request + will be cancelled with a DeadlineExceeded exception, and a retry will + be attempted if within operation_timeout budget + - metadata: Strings which should be sent along with the request as metadata headers. + + Raises: + - MutationsExceptionGroup if one or more mutations fails + Contains details about any failed entries in .exceptions """ - if self._table_admin_client is None: - if not self._admin: - raise ValueError("Client is not an admin client.") - - transport = self._create_gapic_client_channel( - bigtable_admin_v2.BigtableTableAdminClient, - BigtableTableAdminGrpcTransport, - ) - klass = _create_gapic_client( - bigtable_admin_v2.BigtableTableAdminClient, - client_options=self._admin_client_options, - transport=transport, - ) - self._table_admin_client = klass(self) - return self._table_admin_client - - @property - def instance_admin_client(self): - """Getter for the gRPC stub used for the Table Admin API. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_instance_admin_client] - :end-before: [END bigtable_api_instance_admin_client] - :dedent: 4 - - :rtype: :class:`.bigtable_admin_pb2.BigtableInstanceAdmin` - :returns: A BigtableInstanceAdmin instance. - :raises: :class:`ValueError ` if the current - client is not an admin client or if it has not been - :meth:`start`-ed. + raise NotImplementedError + + async def check_and_mutate_row( + self, + row_key: str | bytes, + predicate: RowFilter | None, + true_case_mutations: Mutation | list[Mutation] | None = None, + false_case_mutations: Mutation | list[Mutation] | None = None, + operation_timeout: int | float | None = 60, + metadata: list[tuple[str, str]] | None = None, + ) -> bool: """ - if self._instance_admin_client is None: - if not self._admin: - raise ValueError("Client is not an admin client.") - - transport = self._create_gapic_client_channel( - bigtable_admin_v2.BigtableInstanceAdminClient, - BigtableInstanceAdminGrpcTransport, - ) - klass = _create_gapic_client( - bigtable_admin_v2.BigtableInstanceAdminClient, - client_options=self._admin_client_options, - transport=transport, - ) - self._instance_admin_client = klass(self) - return self._instance_admin_client - - def instance(self, instance_id, display_name=None, instance_type=None, labels=None): - """Factory to create a instance associated with this client. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_create_prod_instance] - :end-before: [END bigtable_api_create_prod_instance] - :dedent: 4 - - :type instance_id: str - :param instance_id: The ID of the instance. - - :type display_name: str - :param display_name: (Optional) The display name for the instance in - the Cloud Console UI. (Must be between 4 and 30 - characters.) If this value is not set in the - constructor, will fall back to the instance ID. - - :type instance_type: int - :param instance_type: (Optional) The type of the instance. - Possible values are represented - by the following constants: - :data:`google.cloud.bigtable.instance.InstanceType.PRODUCTION`. - :data:`google.cloud.bigtable.instance.InstanceType.DEVELOPMENT`, - Defaults to - :data:`google.cloud.bigtable.instance.InstanceType.UNSPECIFIED`. - - :type labels: dict - :param labels: (Optional) Labels are a flexible and lightweight - mechanism for organizing cloud resources into groups - that reflect a customer's organizational needs and - deployment strategies. They can be used to filter - resources and aggregate metrics. Label keys must be - between 1 and 63 characters long. Maximum 64 labels can - be associated with a given resource. Label values must - be between 0 and 63 characters long. Keys and values - must both be under 128 bytes. - - :rtype: :class:`~google.cloud.bigtable.instance.Instance` - :returns: an instance owned by this client. + Mutates a row atomically based on the output of a predicate filter + + Non-idempotent operation: will not be retried + + Args: + - row_key: the key of the row to mutate + - predicate: the filter to be applied to the contents of the specified row. + Depending on whether or not any results are yielded, + either true_case_mutations or false_case_mutations will be executed. + If None, checks that the row contains any values at all. + - true_case_mutations: + Changes to be atomically applied to the specified row if + predicate yields at least one cell when + applied to row_key. Entries are applied in order, + meaning that earlier mutations can be masked by later + ones. Must contain at least one entry if + false_case_mutations is empty, and at most 100000. + - false_case_mutations: + Changes to be atomically applied to the specified row if + predicate_filter does not yield any cells when + applied to row_key. Entries are applied in order, + meaning that earlier mutations can be masked by later + ones. Must contain at least one entry if + `true_case_mutations is empty, and at most 100000. + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will not be retried. + - metadata: Strings which should be sent along with the request as metadata headers. + Returns: + - bool indicating whether the predicate was true or false + Raises: + - GoogleAPIError exceptions from grpc call """ - return Instance( - instance_id, - self, - display_name=display_name, - instance_type=instance_type, - labels=labels, - ) - - def list_instances(self): - """List instances owned by the project. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_list_instances] - :end-before: [END bigtable_api_list_instances] - :dedent: 4 - - :rtype: tuple - :returns: - (instances, failed_locations), where 'instances' is list of - :class:`google.cloud.bigtable.instance.Instance`, and - 'failed_locations' is a list of locations which could not - be resolved. + raise NotImplementedError + + async def read_modify_write_row( + self, + row_key: str | bytes, + rules: ReadModifyWriteRule + | list[ReadModifyWriteRule] + | dict[str, Any] + | list[dict[str, Any]], + *, + operation_timeout: int | float | None = 60, + metadata: list[tuple[str, str]] | None = None, + ) -> RowResponse: """ - resp = self.instance_admin_client.list_instances( - request={"parent": self.project_path} - ) - instances = [Instance.from_pb(instance, self) for instance in resp.instances] - return instances, resp.failed_locations - - def list_clusters(self): - """List the clusters in the project. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_list_clusters_in_project] - :end-before: [END bigtable_api_list_clusters_in_project] - :dedent: 4 - - :rtype: tuple - :returns: - (clusters, failed_locations), where 'clusters' is list of - :class:`google.cloud.bigtable.instance.Cluster`, and - 'failed_locations' is a list of strings representing - locations which could not be resolved. + Reads and modifies a row atomically according to input ReadModifyWriteRules, + and returns the contents of all modified cells + + The new value for the timestamp is the greater of the existing timestamp or + the current server time. + + Non-idempotent operation: will not be retried + + Args: + - row_key: the key of the row to apply read/modify/write rules to + - rules: A rule or set of rules to apply to the row. + Rules are applied in order, meaning that earlier rules will affect the + results of later ones. + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will not be retried. + - metadata: Strings which should be sent along with the request as metadata headers. + Returns: + - RowResponse: containing cell data that was modified as part of the + operation + Raises: + - GoogleAPIError exceptions from grpc call """ - resp = self.instance_admin_client.list_clusters( - request={ - "parent": self.instance_admin_client.instance_path(self.project, "-") - } - ) - clusters = [] - instances = {} - for cluster in resp.clusters: - match_cluster_name = _CLUSTER_NAME_RE.match(cluster.name) - instance_id = match_cluster_name.group("instance") - if instance_id not in instances: - instances[instance_id] = self.instance(instance_id) - clusters.append(Cluster.from_pb(cluster, instances[instance_id])) - return clusters, resp.failed_locations + raise NotImplementedError diff --git a/google/cloud/bigtable/deprecated/__init__.py b/google/cloud/bigtable/deprecated/__init__.py new file mode 100644 index 000000000..a54fffdf1 --- /dev/null +++ b/google/cloud/bigtable/deprecated/__init__.py @@ -0,0 +1,25 @@ +# Copyright 2015 Google LLC +# +# Licensed 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. + +"""Google Cloud Bigtable API package.""" + +from google.cloud.bigtable.deprecated.client import Client + +from google.cloud.bigtable import gapic_version as package_version + +__version__: str + +__version__ = package_version.__version__ + +__all__ = ["__version__", "Client"] diff --git a/google/cloud/bigtable/app_profile.py b/google/cloud/bigtable/deprecated/app_profile.py similarity index 97% rename from google/cloud/bigtable/app_profile.py rename to google/cloud/bigtable/deprecated/app_profile.py index 8cde66146..a5c3df356 100644 --- a/google/cloud/bigtable/app_profile.py +++ b/google/cloud/bigtable/deprecated/app_profile.py @@ -17,7 +17,7 @@ import re -from google.cloud.bigtable.enums import RoutingPolicyType +from google.cloud.bigtable.deprecated.enums import RoutingPolicyType from google.cloud.bigtable_admin_v2.types import instance from google.protobuf import field_mask_pb2 from google.api_core.exceptions import NotFound @@ -47,8 +47,8 @@ class AppProfile(object): :param: routing_policy_type: (Optional) The type of the routing policy. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.enums.RoutingPolicyType.ANY` - :data:`google.cloud.bigtable.enums.RoutingPolicyType.SINGLE` + :data:`google.cloud.bigtable.deprecated.enums.RoutingPolicyType.ANY` + :data:`google.cloud.bigtable.deprecated.enums.RoutingPolicyType.SINGLE` :type: description: str :param: description: (Optional) Long form description of the use @@ -148,7 +148,7 @@ def from_pb(cls, app_profile_pb, instance): :type app_profile_pb: :class:`instance.app_profile_pb` :param app_profile_pb: An instance protobuf object. - :type instance: :class:`google.cloud.bigtable.instance.Instance` + :type instance: :class:`google.cloud.bigtable.deprecated.instance.Instance` :param instance: The instance that owns the cluster. :rtype: :class:`AppProfile` diff --git a/google/cloud/bigtable/backup.py b/google/cloud/bigtable/deprecated/backup.py similarity index 96% rename from google/cloud/bigtable/backup.py rename to google/cloud/bigtable/deprecated/backup.py index 6986d730a..fc15318bc 100644 --- a/google/cloud/bigtable/backup.py +++ b/google/cloud/bigtable/deprecated/backup.py @@ -19,8 +19,8 @@ from google.cloud._helpers import _datetime_to_pb_timestamp # type: ignore from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient from google.cloud.bigtable_admin_v2.types import table -from google.cloud.bigtable.encryption_info import EncryptionInfo -from google.cloud.bigtable.policy import Policy +from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo +from google.cloud.bigtable.deprecated.policy import Policy from google.cloud.exceptions import NotFound # type: ignore from google.protobuf import field_mask_pb2 @@ -50,7 +50,7 @@ class Backup(object): :type backup_id: str :param backup_id: The ID of the backup. - :type instance: :class:`~google.cloud.bigtable.instance.Instance` + :type instance: :class:`~google.cloud.bigtable.deprecated.instance.Instance` :param instance: The Instance that owns this Backup. :type cluster_id: str @@ -188,7 +188,7 @@ def expire_time(self, new_expire_time): def encryption_info(self): """Encryption info for this Backup. - :rtype: :class:`google.cloud.bigtable.encryption.EncryptionInfo` + :rtype: :class:`google.cloud.bigtable.deprecated.encryption.EncryptionInfo` :returns: The encryption information for this backup. """ return self._encryption_info @@ -238,10 +238,10 @@ def from_pb(cls, backup_pb, instance): :type backup_pb: :class:`table.Backup` :param backup_pb: A Backup protobuf object. - :type instance: :class:`Instance ` + :type instance: :class:`Instance ` :param instance: The Instance that owns the Backup. - :rtype: :class:`~google.cloud.bigtable.backup.Backup` + :rtype: :class:`~google.cloud.bigtable.deprecated.backup.Backup` :returns: The backup parsed from the protobuf response. :raises: ValueError: If the backup name does not match the expected format or the parsed project ID does not match the @@ -440,7 +440,7 @@ def restore(self, table_id, instance_id=None): def get_iam_policy(self): """Gets the IAM access control policy for this backup. - :rtype: :class:`google.cloud.bigtable.policy.Policy` + :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` :returns: The current IAM policy of this backup. """ table_api = self._instance._client.table_admin_client @@ -452,13 +452,13 @@ def set_iam_policy(self, policy): existing policy. For more information about policy, please see documentation of - class `google.cloud.bigtable.policy.Policy` + class `google.cloud.bigtable.deprecated.policy.Policy` - :type policy: :class:`google.cloud.bigtable.policy.Policy` + :type policy: :class:`google.cloud.bigtable.deprecated.policy.Policy` :param policy: A new IAM policy to replace the current IAM policy of this backup. - :rtype: :class:`google.cloud.bigtable.policy.Policy` + :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` :returns: The current IAM policy of this backup. """ table_api = self._instance._client.table_admin_client diff --git a/google/cloud/bigtable/batcher.py b/google/cloud/bigtable/deprecated/batcher.py similarity index 94% rename from google/cloud/bigtable/batcher.py rename to google/cloud/bigtable/deprecated/batcher.py index 3c23f4436..58cf6b6e3 100644 --- a/google/cloud/bigtable/batcher.py +++ b/google/cloud/bigtable/deprecated/batcher.py @@ -42,7 +42,7 @@ class MutationsBatcher(object): capability of asynchronous, parallel RPCs. :type table: class - :param table: class:`~google.cloud.bigtable.table.Table`. + :param table: class:`~google.cloud.bigtable.deprecated.table.Table`. :type flush_count: int :param flush_count: (Optional) Max number of rows to flush. If it @@ -76,7 +76,7 @@ def mutate(self, row): :dedent: 4 :type row: class - :param row: class:`~google.cloud.bigtable.row.DirectRow`. + :param row: class:`~google.cloud.bigtable.deprecated.row.DirectRow`. :raises: One of the following: * :exc:`~.table._BigtableRetryableError` if any @@ -115,8 +115,8 @@ def mutate_rows(self, rows): :end-before: [END bigtable_api_batcher_mutate_rows] :dedent: 4 - :type rows: list:[`~google.cloud.bigtable.row.DirectRow`] - :param rows: list:[`~google.cloud.bigtable.row.DirectRow`]. + :type rows: list:[`~google.cloud.bigtable.deprecated.row.DirectRow`] + :param rows: list:[`~google.cloud.bigtable.deprecated.row.DirectRow`]. :raises: One of the following: * :exc:`~.table._BigtableRetryableError` if any diff --git a/google/cloud/bigtable/deprecated/client.py b/google/cloud/bigtable/deprecated/client.py new file mode 100644 index 000000000..c13e5f0da --- /dev/null +++ b/google/cloud/bigtable/deprecated/client.py @@ -0,0 +1,521 @@ +# Copyright 2015 Google LLC +# +# Licensed 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. + +"""Parent client for calling the Google Cloud Bigtable API. + +This is the base from which all interactions with the API occur. + +In the hierarchy of API concepts + +* a :class:`~google.cloud.bigtable.deprecated.client.Client` owns an + :class:`~google.cloud.bigtable.deprecated.instance.Instance` +* an :class:`~google.cloud.bigtable.deprecated.instance.Instance` owns a + :class:`~google.cloud.bigtable.deprecated.table.Table` +* a :class:`~google.cloud.bigtable.deprecated.table.Table` owns a + :class:`~.column_family.ColumnFamily` +* a :class:`~google.cloud.bigtable.deprecated.table.Table` owns a + :class:`~google.cloud.bigtable.deprecated.row.Row` (and all the cells in the row) +""" +import os +import warnings +import grpc # type: ignore + +from google.api_core.gapic_v1 import client_info as client_info_lib +import google.auth # type: ignore +from google.auth.credentials import AnonymousCredentials # type: ignore + +from google.cloud import bigtable_v2 +from google.cloud import bigtable_admin_v2 +from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport +from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin.transports import ( + BigtableInstanceAdminGrpcTransport, +) +from google.cloud.bigtable_admin_v2.services.bigtable_table_admin.transports import ( + BigtableTableAdminGrpcTransport, +) + +from google.cloud import bigtable +from google.cloud.bigtable.deprecated.instance import Instance +from google.cloud.bigtable.deprecated.cluster import Cluster + +from google.cloud.client import ClientWithProject # type: ignore + +from google.cloud.bigtable_admin_v2.types import instance +from google.cloud.bigtable.deprecated.cluster import _CLUSTER_NAME_RE +from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore + + +INSTANCE_TYPE_PRODUCTION = instance.Instance.Type.PRODUCTION +INSTANCE_TYPE_DEVELOPMENT = instance.Instance.Type.DEVELOPMENT +INSTANCE_TYPE_UNSPECIFIED = instance.Instance.Type.TYPE_UNSPECIFIED +SPANNER_ADMIN_SCOPE = "https://www.googleapis.com/auth/spanner.admin" +ADMIN_SCOPE = "https://www.googleapis.com/auth/bigtable.admin" +"""Scope for interacting with the Cluster Admin and Table Admin APIs.""" +DATA_SCOPE = "https://www.googleapis.com/auth/bigtable.data" +"""Scope for reading and writing table data.""" +READ_ONLY_SCOPE = "https://www.googleapis.com/auth/bigtable.data.readonly" +"""Scope for reading table data.""" + +_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator" +_GRPC_CHANNEL_OPTIONS = ( + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ("grpc.keepalive_time_ms", 30000), + ("grpc.keepalive_timeout_ms", 10000), +) + + +def _create_gapic_client(client_class, client_options=None, transport=None): + def inner(self): + return client_class( + credentials=None, + client_info=self._client_info, + client_options=client_options, + transport=transport, + ) + + return inner + + +class Client(ClientWithProject): + """Client for interacting with Google Cloud Bigtable API. + + DEPRECATED: This class is deprecated and may be removed in a future version + Please use `google.cloud.bigtable.BigtableDataClient` instead. + + .. note:: + + Since the Cloud Bigtable API requires the gRPC transport, no + ``_http`` argument is accepted by this class. + + :type project: :class:`str` or :func:`unicode ` + :param project: (Optional) The ID of the project which owns the + instances, tables and data. If not provided, will + attempt to determine from the environment. + + :type credentials: :class:`~google.auth.credentials.Credentials` + :param credentials: (Optional) The OAuth2 Credentials to use for this + client. If not passed, falls back to the default + inferred from the environment. + + :type read_only: bool + :param read_only: (Optional) Boolean indicating if the data scope should be + for reading only (or for writing as well). Defaults to + :data:`False`. + + :type admin: bool + :param admin: (Optional) Boolean indicating if the client will be used to + interact with the Instance Admin or Table Admin APIs. This + requires the :const:`ADMIN_SCOPE`. Defaults to :data:`False`. + + :type: client_info: :class:`google.api_core.gapic_v1.client_info.ClientInfo` + :param client_info: + The client info used to send a user-agent string along with API + requests. If ``None``, then default info will be used. Generally, + you only need to set this if you're developing your own library + or partner tool. + + :type client_options: :class:`~google.api_core.client_options.ClientOptions` + or :class:`dict` + :param client_options: (Optional) Client options used to set user options + on the client. API Endpoint should be set through client_options. + + :type admin_client_options: + :class:`~google.api_core.client_options.ClientOptions` or :class:`dict` + :param admin_client_options: (Optional) Client options used to set user + options on the client. API Endpoint for admin operations should be set + through admin_client_options. + + :type channel: :instance: grpc.Channel + :param channel (grpc.Channel): (Optional) DEPRECATED: + A ``Channel`` instance through which to make calls. + This argument is mutually exclusive with ``credentials``; + providing both will raise an exception. No longer used. + + :raises: :class:`ValueError ` if both ``read_only`` + and ``admin`` are :data:`True` + """ + + _table_data_client = None + _table_admin_client = None + _instance_admin_client = None + + def __init__( + self, + project=None, + credentials=None, + read_only=False, + admin=False, + client_info=None, + client_options=None, + admin_client_options=None, + channel=None, + ): + warnings.warn( + "'Client' is deprecated. Please use 'google.cloud.bigtable.BigtableDataClient' instead.", + DeprecationWarning, + stacklevel=2, + ) + if client_info is None: + client_info = client_info_lib.ClientInfo( + client_library_version=bigtable.__version__, + ) + if read_only and admin: + raise ValueError( + "A read-only client cannot also perform" "administrative actions." + ) + + # NOTE: We set the scopes **before** calling the parent constructor. + # It **may** use those scopes in ``with_scopes_if_required``. + self._read_only = bool(read_only) + self._admin = bool(admin) + self._client_info = client_info + self._emulator_host = os.getenv(BIGTABLE_EMULATOR) + + if self._emulator_host is not None: + if credentials is None: + credentials = AnonymousCredentials() + if project is None: + project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT + + if channel is not None: + warnings.warn( + "'channel' is deprecated and no longer used.", + DeprecationWarning, + stacklevel=2, + ) + + self._client_options = client_options + self._admin_client_options = admin_client_options + self._channel = channel + self.SCOPE = self._get_scopes() + super(Client, self).__init__( + project=project, + credentials=credentials, + client_options=client_options, + ) + + def _get_scopes(self): + """Get the scopes corresponding to admin / read-only state. + + Returns: + Tuple[str, ...]: The tuple of scopes. + """ + if self._read_only: + scopes = (READ_ONLY_SCOPE,) + else: + scopes = (DATA_SCOPE,) + + if self._admin: + scopes += (ADMIN_SCOPE,) + + return scopes + + def _emulator_channel(self, transport, options): + """Create a channel using self._credentials + + Works in a similar way to ``grpc.secure_channel`` but using + ``grpc.local_channel_credentials`` rather than + ``grpc.ssh_channel_credentials`` to allow easy connection to a + local emulator. + + Returns: + grpc.Channel or grpc.aio.Channel + """ + # TODO: Implement a special credentials type for emulator and use + # "transport.create_channel" to create gRPC channels once google-auth + # extends it's allowed credentials types. + # Note: this code also exists in the firestore client. + if "GrpcAsyncIOTransport" in str(transport.__name__): + return grpc.aio.secure_channel( + self._emulator_host, + self._local_composite_credentials(), + options=options, + ) + else: + return grpc.secure_channel( + self._emulator_host, + self._local_composite_credentials(), + options=options, + ) + + def _local_composite_credentials(self): + """Create credentials for the local emulator channel. + + :return: grpc.ChannelCredentials + """ + credentials = google.auth.credentials.with_scopes_if_required( + self._credentials, None + ) + request = google.auth.transport.requests.Request() + + # Create the metadata plugin for inserting the authorization header. + metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin( + credentials, request + ) + + # Create a set of grpc.CallCredentials using the metadata plugin. + google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin) + + # Using the local_credentials to allow connection to emulator + local_credentials = grpc.local_channel_credentials() + + # Combine the local credentials and the authorization credentials. + return grpc.composite_channel_credentials( + local_credentials, google_auth_credentials + ) + + def _create_gapic_client_channel(self, client_class, grpc_transport): + if self._emulator_host is not None: + api_endpoint = self._emulator_host + elif self._client_options and self._client_options.api_endpoint: + api_endpoint = self._client_options.api_endpoint + else: + api_endpoint = client_class.DEFAULT_ENDPOINT + + if self._emulator_host is not None: + channel = self._emulator_channel( + transport=grpc_transport, + options=_GRPC_CHANNEL_OPTIONS, + ) + else: + channel = grpc_transport.create_channel( + host=api_endpoint, + credentials=self._credentials, + options=_GRPC_CHANNEL_OPTIONS, + ) + return grpc_transport(channel=channel, host=api_endpoint) + + @property + def project_path(self): + """Project name to be used with Instance Admin API. + + .. note:: + + This property will not change if ``project`` does not, but the + return value is not cached. + + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_project_path] + :end-before: [END bigtable_api_project_path] + :dedent: 4 + + The project name is of the form + + ``"projects/{project}"`` + + :rtype: str + :returns: Return a fully-qualified project string. + """ + return self.instance_admin_client.common_project_path(self.project) + + @property + def table_data_client(self): + """Getter for the gRPC stub used for the Table Admin API. + + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_table_data_client] + :end-before: [END bigtable_api_table_data_client] + :dedent: 4 + + :rtype: :class:`.bigtable_v2.BigtableClient` + :returns: A BigtableClient object. + """ + if self._table_data_client is None: + transport = self._create_gapic_client_channel( + bigtable_v2.BigtableClient, + BigtableGrpcTransport, + ) + klass = _create_gapic_client( + bigtable_v2.BigtableClient, + client_options=self._client_options, + transport=transport, + ) + self._table_data_client = klass(self) + return self._table_data_client + + @property + def table_admin_client(self): + """Getter for the gRPC stub used for the Table Admin API. + + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_table_admin_client] + :end-before: [END bigtable_api_table_admin_client] + :dedent: 4 + + :rtype: :class:`.bigtable_admin_pb2.BigtableTableAdmin` + :returns: A BigtableTableAdmin instance. + :raises: :class:`ValueError ` if the current + client is not an admin client or if it has not been + :meth:`start`-ed. + """ + if self._table_admin_client is None: + if not self._admin: + raise ValueError("Client is not an admin client.") + + transport = self._create_gapic_client_channel( + bigtable_admin_v2.BigtableTableAdminClient, + BigtableTableAdminGrpcTransport, + ) + klass = _create_gapic_client( + bigtable_admin_v2.BigtableTableAdminClient, + client_options=self._admin_client_options, + transport=transport, + ) + self._table_admin_client = klass(self) + return self._table_admin_client + + @property + def instance_admin_client(self): + """Getter for the gRPC stub used for the Table Admin API. + + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_instance_admin_client] + :end-before: [END bigtable_api_instance_admin_client] + :dedent: 4 + + :rtype: :class:`.bigtable_admin_pb2.BigtableInstanceAdmin` + :returns: A BigtableInstanceAdmin instance. + :raises: :class:`ValueError ` if the current + client is not an admin client or if it has not been + :meth:`start`-ed. + """ + if self._instance_admin_client is None: + if not self._admin: + raise ValueError("Client is not an admin client.") + + transport = self._create_gapic_client_channel( + bigtable_admin_v2.BigtableInstanceAdminClient, + BigtableInstanceAdminGrpcTransport, + ) + klass = _create_gapic_client( + bigtable_admin_v2.BigtableInstanceAdminClient, + client_options=self._admin_client_options, + transport=transport, + ) + self._instance_admin_client = klass(self) + return self._instance_admin_client + + def instance(self, instance_id, display_name=None, instance_type=None, labels=None): + """Factory to create a instance associated with this client. + + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_create_prod_instance] + :end-before: [END bigtable_api_create_prod_instance] + :dedent: 4 + + :type instance_id: str + :param instance_id: The ID of the instance. + + :type display_name: str + :param display_name: (Optional) The display name for the instance in + the Cloud Console UI. (Must be between 4 and 30 + characters.) If this value is not set in the + constructor, will fall back to the instance ID. + + :type instance_type: int + :param instance_type: (Optional) The type of the instance. + Possible values are represented + by the following constants: + :data:`google.cloud.bigtable.deprecated.instance.InstanceType.PRODUCTION`. + :data:`google.cloud.bigtable.deprecated.instance.InstanceType.DEVELOPMENT`, + Defaults to + :data:`google.cloud.bigtable.deprecated.instance.InstanceType.UNSPECIFIED`. + + :type labels: dict + :param labels: (Optional) Labels are a flexible and lightweight + mechanism for organizing cloud resources into groups + that reflect a customer's organizational needs and + deployment strategies. They can be used to filter + resources and aggregate metrics. Label keys must be + between 1 and 63 characters long. Maximum 64 labels can + be associated with a given resource. Label values must + be between 0 and 63 characters long. Keys and values + must both be under 128 bytes. + + :rtype: :class:`~google.cloud.bigtable.deprecated.instance.Instance` + :returns: an instance owned by this client. + """ + return Instance( + instance_id, + self, + display_name=display_name, + instance_type=instance_type, + labels=labels, + ) + + def list_instances(self): + """List instances owned by the project. + + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_list_instances] + :end-before: [END bigtable_api_list_instances] + :dedent: 4 + + :rtype: tuple + :returns: + (instances, failed_locations), where 'instances' is list of + :class:`google.cloud.bigtable.deprecated.instance.Instance`, and + 'failed_locations' is a list of locations which could not + be resolved. + """ + resp = self.instance_admin_client.list_instances( + request={"parent": self.project_path} + ) + instances = [Instance.from_pb(instance, self) for instance in resp.instances] + return instances, resp.failed_locations + + def list_clusters(self): + """List the clusters in the project. + + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_list_clusters_in_project] + :end-before: [END bigtable_api_list_clusters_in_project] + :dedent: 4 + + :rtype: tuple + :returns: + (clusters, failed_locations), where 'clusters' is list of + :class:`google.cloud.bigtable.deprecated.instance.Cluster`, and + 'failed_locations' is a list of strings representing + locations which could not be resolved. + """ + resp = self.instance_admin_client.list_clusters( + request={ + "parent": self.instance_admin_client.instance_path(self.project, "-") + } + ) + clusters = [] + instances = {} + for cluster in resp.clusters: + match_cluster_name = _CLUSTER_NAME_RE.match(cluster.name) + instance_id = match_cluster_name.group("instance") + if instance_id not in instances: + instances[instance_id] = self.instance(instance_id) + clusters.append(Cluster.from_pb(cluster, instances[instance_id])) + return clusters, resp.failed_locations diff --git a/google/cloud/bigtable/cluster.py b/google/cloud/bigtable/deprecated/cluster.py similarity index 95% rename from google/cloud/bigtable/cluster.py rename to google/cloud/bigtable/deprecated/cluster.py index 11fb5492d..b60d3503c 100644 --- a/google/cloud/bigtable/cluster.py +++ b/google/cloud/bigtable/deprecated/cluster.py @@ -42,7 +42,7 @@ class Cluster(object): :type cluster_id: str :param cluster_id: The ID of the cluster. - :type instance: :class:`~google.cloud.bigtable.instance.Instance` + :type instance: :class:`~google.cloud.bigtable.deprecated.instance.Instance` :param instance: The instance where the cluster resides. :type location_id: str @@ -62,10 +62,10 @@ class Cluster(object): :param default_storage_type: (Optional) The type of storage Possible values are represented by the following constants: - :data:`google.cloud.bigtable.enums.StorageType.SSD`. - :data:`google.cloud.bigtable.enums.StorageType.HDD`, + :data:`google.cloud.bigtable.deprecated.enums.StorageType.SSD`. + :data:`google.cloud.bigtable.deprecated.enums.StorageType.HDD`, Defaults to - :data:`google.cloud.bigtable.enums.StorageType.UNSPECIFIED`. + :data:`google.cloud.bigtable.deprecated.enums.StorageType.UNSPECIFIED`. :type kms_key_name: str :param kms_key_name: (Optional, Creation Only) The name of the KMS customer managed @@ -84,11 +84,11 @@ class Cluster(object): :param _state: (`OutputOnly`) The current state of the cluster. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.enums.Cluster.State.NOT_KNOWN`. - :data:`google.cloud.bigtable.enums.Cluster.State.READY`. - :data:`google.cloud.bigtable.enums.Cluster.State.CREATING`. - :data:`google.cloud.bigtable.enums.Cluster.State.RESIZING`. - :data:`google.cloud.bigtable.enums.Cluster.State.DISABLED`. + :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.NOT_KNOWN`. + :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.READY`. + :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.CREATING`. + :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.RESIZING`. + :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.DISABLED`. :type min_serve_nodes: int :param min_serve_nodes: (Optional) The minimum number of nodes to be set in the cluster for autoscaling. @@ -150,7 +150,7 @@ def from_pb(cls, cluster_pb, instance): :type cluster_pb: :class:`instance.Cluster` :param cluster_pb: An instance protobuf object. - :type instance: :class:`google.cloud.bigtable.instance.Instance` + :type instance: :class:`google.cloud.bigtable.deprecated.instance.Instance` :param instance: The instance that owns the cluster. :rtype: :class:`Cluster` @@ -236,7 +236,7 @@ def name(self): @property def state(self): - """google.cloud.bigtable.enums.Cluster.State: state of cluster. + """google.cloud.bigtable.deprecated.enums.Cluster.State: state of cluster. For example: diff --git a/google/cloud/bigtable/column_family.py b/google/cloud/bigtable/deprecated/column_family.py similarity index 99% rename from google/cloud/bigtable/column_family.py rename to google/cloud/bigtable/deprecated/column_family.py index 80232958d..3d4c1a642 100644 --- a/google/cloud/bigtable/column_family.py +++ b/google/cloud/bigtable/deprecated/column_family.py @@ -195,7 +195,7 @@ class ColumnFamily(object): :param column_family_id: The ID of the column family. Must be of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - :type table: :class:`Table ` + :type table: :class:`Table ` :param table: The table that owns the column family. :type gc_rule: :class:`GarbageCollectionRule` diff --git a/google/cloud/bigtable/encryption_info.py b/google/cloud/bigtable/deprecated/encryption_info.py similarity index 93% rename from google/cloud/bigtable/encryption_info.py rename to google/cloud/bigtable/deprecated/encryption_info.py index 1757297bc..daa0d9232 100644 --- a/google/cloud/bigtable/encryption_info.py +++ b/google/cloud/bigtable/deprecated/encryption_info.py @@ -14,7 +14,7 @@ """Class for encryption info for tables and backups.""" -from google.cloud.bigtable.error import Status +from google.cloud.bigtable.deprecated.error import Status class EncryptionInfo: @@ -27,7 +27,7 @@ class EncryptionInfo: :type encryption_type: int :param encryption_type: See :class:`enums.EncryptionInfo.EncryptionType` - :type encryption_status: google.cloud.bigtable.encryption.Status + :type encryption_status: google.cloud.bigtable.deprecated.encryption.Status :param encryption_status: The encryption status. :type kms_key_version: str diff --git a/google/cloud/bigtable/enums.py b/google/cloud/bigtable/deprecated/enums.py similarity index 100% rename from google/cloud/bigtable/enums.py rename to google/cloud/bigtable/deprecated/enums.py diff --git a/google/cloud/bigtable/error.py b/google/cloud/bigtable/deprecated/error.py similarity index 100% rename from google/cloud/bigtable/error.py rename to google/cloud/bigtable/deprecated/error.py diff --git a/google/cloud/bigtable/instance.py b/google/cloud/bigtable/deprecated/instance.py similarity index 91% rename from google/cloud/bigtable/instance.py rename to google/cloud/bigtable/deprecated/instance.py index 6d092cefd..33475d261 100644 --- a/google/cloud/bigtable/instance.py +++ b/google/cloud/bigtable/deprecated/instance.py @@ -16,9 +16,9 @@ import re -from google.cloud.bigtable.app_profile import AppProfile -from google.cloud.bigtable.cluster import Cluster -from google.cloud.bigtable.table import Table +from google.cloud.bigtable.deprecated.app_profile import AppProfile +from google.cloud.bigtable.deprecated.cluster import Cluster +from google.cloud.bigtable.deprecated.table import Table from google.protobuf import field_mask_pb2 @@ -28,7 +28,7 @@ from google.api_core.exceptions import NotFound -from google.cloud.bigtable.policy import Policy +from google.cloud.bigtable.deprecated.policy import Policy import warnings @@ -61,7 +61,7 @@ class Instance(object): :type instance_id: str :param instance_id: The ID of the instance. - :type client: :class:`Client ` + :type client: :class:`Client ` :param client: The client that owns the instance. Provides authorization and a project ID. @@ -75,10 +75,10 @@ class Instance(object): :param instance_type: (Optional) The type of the instance. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.enums.Instance.Type.PRODUCTION`. - :data:`google.cloud.bigtable.enums.Instance.Type.DEVELOPMENT`, + :data:`google.cloud.bigtable.deprecated.enums.Instance.Type.PRODUCTION`. + :data:`google.cloud.bigtable.deprecated.enums.Instance.Type.DEVELOPMENT`, Defaults to - :data:`google.cloud.bigtable.enums.Instance.Type.UNSPECIFIED`. + :data:`google.cloud.bigtable.deprecated.enums.Instance.Type.UNSPECIFIED`. :type labels: dict :param labels: (Optional) Labels are a flexible and lightweight @@ -95,9 +95,9 @@ class Instance(object): :param _state: (`OutputOnly`) The current state of the instance. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.enums.Instance.State.STATE_NOT_KNOWN`. - :data:`google.cloud.bigtable.enums.Instance.State.READY`. - :data:`google.cloud.bigtable.enums.Instance.State.CREATING`. + :data:`google.cloud.bigtable.deprecated.enums.Instance.State.STATE_NOT_KNOWN`. + :data:`google.cloud.bigtable.deprecated.enums.Instance.State.READY`. + :data:`google.cloud.bigtable.deprecated.enums.Instance.State.CREATING`. """ def __init__( @@ -141,7 +141,7 @@ def from_pb(cls, instance_pb, client): :type instance_pb: :class:`instance.Instance` :param instance_pb: An instance protobuf object. - :type client: :class:`Client ` + :type client: :class:`Client ` :param client: The client that owns the instance. :rtype: :class:`Instance` @@ -196,7 +196,7 @@ def name(self): @property def state(self): - """google.cloud.bigtable.enums.Instance.State: state of Instance. + """google.cloud.bigtable.deprecated.enums.Instance.State: state of Instance. For example: @@ -272,12 +272,12 @@ def create( persisting Bigtable data. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.enums.StorageType.SSD`. - :data:`google.cloud.bigtable.enums.StorageType.HDD`, + :data:`google.cloud.bigtable.deprecated.enums.StorageType.SSD`. + :data:`google.cloud.bigtable.deprecated.enums.StorageType.HDD`, Defaults to - :data:`google.cloud.bigtable.enums.StorageType.UNSPECIFIED`. + :data:`google.cloud.bigtable.deprecated.enums.StorageType.UNSPECIFIED`. - :type clusters: class:`~[~google.cloud.bigtable.cluster.Cluster]` + :type clusters: class:`~[~google.cloud.bigtable.deprecated.cluster.Cluster]` :param clusters: List of clusters to be created. :rtype: :class:`~google.api_core.operation.Operation` @@ -478,7 +478,7 @@ def get_iam_policy(self, requested_policy_version=None): than the one that was requested, based on the feature syntax in the policy fetched. - :rtype: :class:`google.cloud.bigtable.policy.Policy` + :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` :returns: The current IAM policy of this instance """ args = {"resource": self.name} @@ -497,7 +497,7 @@ def set_iam_policy(self, policy): existing policy. For more information about policy, please see documentation of - class `google.cloud.bigtable.policy.Policy` + class `google.cloud.bigtable.deprecated.policy.Policy` For example: @@ -506,11 +506,11 @@ class `google.cloud.bigtable.policy.Policy` :end-before: [END bigtable_api_set_iam_policy] :dedent: 4 - :type policy: :class:`google.cloud.bigtable.policy.Policy` + :type policy: :class:`google.cloud.bigtable.deprecated.policy.Policy` :param policy: A new IAM policy to replace the current IAM policy of this instance - :rtype: :class:`google.cloud.bigtable.policy.Policy` + :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` :returns: The current IAM policy of this instance. """ instance_admin_client = self._client.instance_admin_client @@ -586,12 +586,12 @@ def cluster( :param default_storage_type: (Optional) The type of storage Possible values are represented by the following constants: - :data:`google.cloud.bigtable.enums.StorageType.SSD`. - :data:`google.cloud.bigtable.enums.StorageType.HDD`, + :data:`google.cloud.bigtable.deprecated.enums.StorageType.SSD`. + :data:`google.cloud.bigtable.deprecated.enums.StorageType.HDD`, Defaults to - :data:`google.cloud.bigtable.enums.StorageType.UNSPECIFIED`. + :data:`google.cloud.bigtable.deprecated.enums.StorageType.UNSPECIFIED`. - :rtype: :class:`~google.cloud.bigtable.instance.Cluster` + :rtype: :class:`~google.cloud.bigtable.deprecated.instance.Cluster` :returns: a cluster owned by this instance. :type kms_key_name: str @@ -635,7 +635,7 @@ def list_clusters(self): :rtype: tuple :returns: (clusters, failed_locations), where 'clusters' is list of - :class:`google.cloud.bigtable.instance.Cluster`, and + :class:`google.cloud.bigtable.deprecated.instance.Cluster`, and 'failed_locations' is a list of locations which could not be resolved. """ @@ -664,7 +664,7 @@ def table(self, table_id, mutation_timeout=None, app_profile_id=None): :type app_profile_id: str :param app_profile_id: (Optional) The unique name of the AppProfile. - :rtype: :class:`Table ` + :rtype: :class:`Table ` :returns: The table owned by this instance. """ return Table( @@ -684,7 +684,7 @@ def list_tables(self): :end-before: [END bigtable_api_list_tables] :dedent: 4 - :rtype: list of :class:`Table ` + :rtype: list of :class:`Table ` :returns: The list of tables owned by the instance. :raises: :class:`ValueError ` if one of the returned tables has a name that is not of the expected format. @@ -731,8 +731,8 @@ def app_profile( :param: routing_policy_type: The type of the routing policy. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.enums.RoutingPolicyType.ANY` - :data:`google.cloud.bigtable.enums.RoutingPolicyType.SINGLE` + :data:`google.cloud.bigtable.deprecated.enums.RoutingPolicyType.ANY` + :data:`google.cloud.bigtable.deprecated.enums.RoutingPolicyType.SINGLE` :type: description: str :param: description: (Optional) Long form description of the use @@ -753,7 +753,7 @@ def app_profile( transactional writes for ROUTING_POLICY_TYPE_SINGLE. - :rtype: :class:`~google.cloud.bigtable.app_profile.AppProfile>` + :rtype: :class:`~google.cloud.bigtable.deprecated.app_profile.AppProfile>` :returns: AppProfile for this instance. """ return AppProfile( @@ -776,10 +776,10 @@ def list_app_profiles(self): :end-before: [END bigtable_api_list_app_profiles] :dedent: 4 - :rtype: :list:[`~google.cloud.bigtable.app_profile.AppProfile`] - :returns: A :list:[`~google.cloud.bigtable.app_profile.AppProfile`]. + :rtype: :list:[`~google.cloud.bigtable.deprecated.app_profile.AppProfile`] + :returns: A :list:[`~google.cloud.bigtable.deprecated.app_profile.AppProfile`]. By default, this is a list of - :class:`~google.cloud.bigtable.app_profile.AppProfile` + :class:`~google.cloud.bigtable.deprecated.app_profile.AppProfile` instances. """ resp = self._client.instance_admin_client.list_app_profiles( diff --git a/google/cloud/bigtable/policy.py b/google/cloud/bigtable/deprecated/policy.py similarity index 100% rename from google/cloud/bigtable/policy.py rename to google/cloud/bigtable/deprecated/policy.py diff --git a/google/cloud/bigtable/py.typed b/google/cloud/bigtable/deprecated/py.typed similarity index 100% rename from google/cloud/bigtable/py.typed rename to google/cloud/bigtable/deprecated/py.typed diff --git a/google/cloud/bigtable/row.py b/google/cloud/bigtable/deprecated/row.py similarity index 98% rename from google/cloud/bigtable/row.py rename to google/cloud/bigtable/deprecated/row.py index 752458a08..3b114a74a 100644 --- a/google/cloud/bigtable/row.py +++ b/google/cloud/bigtable/deprecated/row.py @@ -51,7 +51,7 @@ class Row(object): :type row_key: bytes :param row_key: The key for the current row. - :type table: :class:`Table ` + :type table: :class:`Table ` :param table: (Optional) The table that owns the row. """ @@ -86,7 +86,7 @@ def table(self): :end-before: [END bigtable_api_row_table] :dedent: 4 - :rtype: table: :class:`Table ` + :rtype: table: :class:`Table ` :returns: table: The table that owns the row. """ return self._table @@ -105,7 +105,7 @@ class _SetDeleteRow(Row): :type row_key: bytes :param row_key: The key for the current row. - :type table: :class:`Table ` + :type table: :class:`Table ` :param table: The table that owns the row. """ @@ -275,11 +275,11 @@ class DirectRow(_SetDeleteRow): :type row_key: bytes :param row_key: The key for the current row. - :type table: :class:`Table ` + :type table: :class:`Table ` :param table: (Optional) The table that owns the row. This is used for the :meth: `commit` only. Alternatively, DirectRows can be persisted via - :meth:`~google.cloud.bigtable.table.Table.mutate_rows`. + :meth:`~google.cloud.bigtable.deprecated.table.Table.mutate_rows`. """ def __init__(self, row_key, table=None): @@ -519,7 +519,7 @@ class ConditionalRow(_SetDeleteRow): :type row_key: bytes :param row_key: The key for the current row. - :type table: :class:`Table ` + :type table: :class:`Table ` :param table: The table that owns the row. :type filter_: :class:`.RowFilter` @@ -791,7 +791,7 @@ class AppendRow(Row): :type row_key: bytes :param row_key: The key for the current row. - :type table: :class:`Table ` + :type table: :class:`Table ` :param table: The table that owns the row. """ @@ -1107,7 +1107,7 @@ def find_cells(self, column_family_id, column): are located. Returns: - List[~google.cloud.bigtable.row_data.Cell]: The cells stored in the + List[~google.cloud.bigtable.deprecated.row_data.Cell]: The cells stored in the specified column. Raises: @@ -1147,7 +1147,7 @@ def cell_value(self, column_family_id, column, index=0): not specified, will return the first cell. Returns: - ~google.cloud.bigtable.row_data.Cell value: The cell value stored + ~google.cloud.bigtable.deprecated.row_data.Cell value: The cell value stored in the specified column and specified index. Raises: diff --git a/google/cloud/bigtable/row_data.py b/google/cloud/bigtable/deprecated/row_data.py similarity index 98% rename from google/cloud/bigtable/row_data.py rename to google/cloud/bigtable/deprecated/row_data.py index a50fab1ee..9daa1ed8f 100644 --- a/google/cloud/bigtable/row_data.py +++ b/google/cloud/bigtable/deprecated/row_data.py @@ -23,10 +23,10 @@ from google.api_core import retry from google.cloud._helpers import _to_bytes # type: ignore -from google.cloud.bigtable.row_merger import _RowMerger, _State +from google.cloud.bigtable.deprecated.row_merger import _RowMerger, _State from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 from google.cloud.bigtable_v2.types import data as data_v2_pb2 -from google.cloud.bigtable.row import Cell, InvalidChunk, PartialRowData +from google.cloud.bigtable.deprecated.row import Cell, InvalidChunk, PartialRowData # Some classes need to be re-exported here to keep backwards @@ -98,7 +98,7 @@ def _retry_read_rows_exception(exc): """The default retry strategy to be used on retry-able errors. Used by -:meth:`~google.cloud.bigtable.row_data.PartialRowsData._read_next_response`. +:meth:`~google.cloud.bigtable.deprecated.row_data.PartialRowsData._read_next_response`. """ diff --git a/google/cloud/bigtable/deprecated/row_filters.py b/google/cloud/bigtable/deprecated/row_filters.py new file mode 100644 index 000000000..53192acc8 --- /dev/null +++ b/google/cloud/bigtable/deprecated/row_filters.py @@ -0,0 +1,838 @@ +# Copyright 2016 Google LLC +# +# Licensed 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. + +"""Filters for Google Cloud Bigtable Row classes.""" + +import struct + + +from google.cloud._helpers import _microseconds_from_datetime # type: ignore +from google.cloud._helpers import _to_bytes # type: ignore +from google.cloud.bigtable_v2.types import data as data_v2_pb2 + +_PACK_I64 = struct.Struct(">q").pack + + +class RowFilter(object): + """Basic filter to apply to cells in a row. + + These values can be combined via :class:`RowFilterChain`, + :class:`RowFilterUnion` and :class:`ConditionalRowFilter`. + + .. note:: + + This class is a do-nothing base class for all row filters. + """ + + +class _BoolFilter(RowFilter): + """Row filter that uses a boolean flag. + + :type flag: bool + :param flag: An indicator if a setting is turned on or off. + """ + + def __init__(self, flag): + self.flag = flag + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other.flag == self.flag + + def __ne__(self, other): + return not self == other + + +class SinkFilter(_BoolFilter): + """Advanced row filter to skip parent filters. + + :type flag: bool + :param flag: ADVANCED USE ONLY. Hook for introspection into the row filter. + Outputs all cells directly to the output of the read rather + than to any parent filter. Cannot be used within the + ``predicate_filter``, ``true_filter``, or ``false_filter`` + of a :class:`ConditionalRowFilter`. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(sink=self.flag) + + +class PassAllFilter(_BoolFilter): + """Row filter equivalent to not filtering at all. + + :type flag: bool + :param flag: Matches all cells, regardless of input. Functionally + equivalent to leaving ``filter`` unset, but included for + completeness. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(pass_all_filter=self.flag) + + +class BlockAllFilter(_BoolFilter): + """Row filter that doesn't match any cells. + + :type flag: bool + :param flag: Does not match any cells, regardless of input. Useful for + temporarily disabling just part of a filter. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(block_all_filter=self.flag) + + +class _RegexFilter(RowFilter): + """Row filter that uses a regular expression. + + The ``regex`` must be valid RE2 patterns. See Google's + `RE2 reference`_ for the accepted syntax. + + .. _RE2 reference: https://github.com/google/re2/wiki/Syntax + + :type regex: bytes or str + :param regex: + A regular expression (RE2) for some row filter. String values + will be encoded as ASCII. + """ + + def __init__(self, regex): + self.regex = _to_bytes(regex) + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other.regex == self.regex + + def __ne__(self, other): + return not self == other + + +class RowKeyRegexFilter(_RegexFilter): + """Row filter for a row key regular expression. + + The ``regex`` must be valid RE2 patterns. See Google's + `RE2 reference`_ for the accepted syntax. + + .. _RE2 reference: https://github.com/google/re2/wiki/Syntax + + .. note:: + + Special care need be used with the expression used. Since + each of these properties can contain arbitrary bytes, the ``\\C`` + escape sequence must be used if a true wildcard is desired. The ``.`` + character will not match the new line character ``\\n``, which may be + present in a binary value. + + :type regex: bytes + :param regex: A regular expression (RE2) to match cells from rows with row + keys that satisfy this regex. For a + ``CheckAndMutateRowRequest``, this filter is unnecessary + since the row key is already specified. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(row_key_regex_filter=self.regex) + + +class RowSampleFilter(RowFilter): + """Matches all cells from a row with probability p. + + :type sample: float + :param sample: The probability of matching a cell (must be in the + interval ``(0, 1)`` The end points are excluded). + """ + + def __init__(self, sample): + self.sample = sample + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other.sample == self.sample + + def __ne__(self, other): + return not self == other + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(row_sample_filter=self.sample) + + +class FamilyNameRegexFilter(_RegexFilter): + """Row filter for a family name regular expression. + + The ``regex`` must be valid RE2 patterns. See Google's + `RE2 reference`_ for the accepted syntax. + + .. _RE2 reference: https://github.com/google/re2/wiki/Syntax + + :type regex: str + :param regex: A regular expression (RE2) to match cells from columns in a + given column family. For technical reasons, the regex must + not contain the ``':'`` character, even if it is not being + used as a literal. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(family_name_regex_filter=self.regex) + + +class ColumnQualifierRegexFilter(_RegexFilter): + """Row filter for a column qualifier regular expression. + + The ``regex`` must be valid RE2 patterns. See Google's + `RE2 reference`_ for the accepted syntax. + + .. _RE2 reference: https://github.com/google/re2/wiki/Syntax + + .. note:: + + Special care need be used with the expression used. Since + each of these properties can contain arbitrary bytes, the ``\\C`` + escape sequence must be used if a true wildcard is desired. The ``.`` + character will not match the new line character ``\\n``, which may be + present in a binary value. + + :type regex: bytes + :param regex: A regular expression (RE2) to match cells from column that + match this regex (irrespective of column family). + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(column_qualifier_regex_filter=self.regex) + + +class TimestampRange(object): + """Range of time with inclusive lower and exclusive upper bounds. + + :type start: :class:`datetime.datetime` + :param start: (Optional) The (inclusive) lower bound of the timestamp + range. If omitted, defaults to Unix epoch. + + :type end: :class:`datetime.datetime` + :param end: (Optional) The (exclusive) upper bound of the timestamp + range. If omitted, no upper bound is used. + """ + + def __init__(self, start=None, end=None): + self.start = start + self.end = end + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other.start == self.start and other.end == self.end + + def __ne__(self, other): + return not self == other + + def to_pb(self): + """Converts the :class:`TimestampRange` to a protobuf. + + :rtype: :class:`.data_v2_pb2.TimestampRange` + :returns: The converted current object. + """ + timestamp_range_kwargs = {} + if self.start is not None: + timestamp_range_kwargs["start_timestamp_micros"] = ( + _microseconds_from_datetime(self.start) // 1000 * 1000 + ) + if self.end is not None: + end_time = _microseconds_from_datetime(self.end) + if end_time % 1000 != 0: + end_time = end_time // 1000 * 1000 + 1000 + timestamp_range_kwargs["end_timestamp_micros"] = end_time + return data_v2_pb2.TimestampRange(**timestamp_range_kwargs) + + +class TimestampRangeFilter(RowFilter): + """Row filter that limits cells to a range of time. + + :type range_: :class:`TimestampRange` + :param range_: Range of time that cells should match against. + """ + + def __init__(self, range_): + self.range_ = range_ + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other.range_ == self.range_ + + def __ne__(self, other): + return not self == other + + def to_pb(self): + """Converts the row filter to a protobuf. + + First converts the ``range_`` on the current object to a protobuf and + then uses it in the ``timestamp_range_filter`` field. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(timestamp_range_filter=self.range_.to_pb()) + + +class ColumnRangeFilter(RowFilter): + """A row filter to restrict to a range of columns. + + Both the start and end column can be included or excluded in the range. + By default, we include them both, but this can be changed with optional + flags. + + :type column_family_id: str + :param column_family_id: The column family that contains the columns. Must + be of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type start_column: bytes + :param start_column: The start of the range of columns. If no value is + used, the backend applies no upper bound to the + values. + + :type end_column: bytes + :param end_column: The end of the range of columns. If no value is used, + the backend applies no upper bound to the values. + + :type inclusive_start: bool + :param inclusive_start: Boolean indicating if the start column should be + included in the range (or excluded). Defaults + to :data:`True` if ``start_column`` is passed and + no ``inclusive_start`` was given. + + :type inclusive_end: bool + :param inclusive_end: Boolean indicating if the end column should be + included in the range (or excluded). Defaults + to :data:`True` if ``end_column`` is passed and + no ``inclusive_end`` was given. + + :raises: :class:`ValueError ` if ``inclusive_start`` + is set but no ``start_column`` is given or if ``inclusive_end`` + is set but no ``end_column`` is given + """ + + def __init__( + self, + column_family_id, + start_column=None, + end_column=None, + inclusive_start=None, + inclusive_end=None, + ): + self.column_family_id = column_family_id + + if inclusive_start is None: + inclusive_start = True + elif start_column is None: + raise ValueError( + "Inclusive start was specified but no " "start column was given." + ) + self.start_column = start_column + self.inclusive_start = inclusive_start + + if inclusive_end is None: + inclusive_end = True + elif end_column is None: + raise ValueError( + "Inclusive end was specified but no " "end column was given." + ) + self.end_column = end_column + self.inclusive_end = inclusive_end + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return ( + other.column_family_id == self.column_family_id + and other.start_column == self.start_column + and other.end_column == self.end_column + and other.inclusive_start == self.inclusive_start + and other.inclusive_end == self.inclusive_end + ) + + def __ne__(self, other): + return not self == other + + def to_pb(self): + """Converts the row filter to a protobuf. + + First converts to a :class:`.data_v2_pb2.ColumnRange` and then uses it + in the ``column_range_filter`` field. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + column_range_kwargs = {"family_name": self.column_family_id} + if self.start_column is not None: + if self.inclusive_start: + key = "start_qualifier_closed" + else: + key = "start_qualifier_open" + column_range_kwargs[key] = _to_bytes(self.start_column) + if self.end_column is not None: + if self.inclusive_end: + key = "end_qualifier_closed" + else: + key = "end_qualifier_open" + column_range_kwargs[key] = _to_bytes(self.end_column) + + column_range = data_v2_pb2.ColumnRange(**column_range_kwargs) + return data_v2_pb2.RowFilter(column_range_filter=column_range) + + +class ValueRegexFilter(_RegexFilter): + """Row filter for a value regular expression. + + The ``regex`` must be valid RE2 patterns. See Google's + `RE2 reference`_ for the accepted syntax. + + .. _RE2 reference: https://github.com/google/re2/wiki/Syntax + + .. note:: + + Special care need be used with the expression used. Since + each of these properties can contain arbitrary bytes, the ``\\C`` + escape sequence must be used if a true wildcard is desired. The ``.`` + character will not match the new line character ``\\n``, which may be + present in a binary value. + + :type regex: bytes or str + :param regex: A regular expression (RE2) to match cells with values that + match this regex. String values will be encoded as ASCII. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(value_regex_filter=self.regex) + + +class ExactValueFilter(ValueRegexFilter): + """Row filter for an exact value. + + + :type value: bytes or str or int + :param value: + a literal string encodable as ASCII, or the + equivalent bytes, or an integer (which will be packed into 8-bytes). + """ + + def __init__(self, value): + if isinstance(value, int): + value = _PACK_I64(value) + super(ExactValueFilter, self).__init__(value) + + +class ValueRangeFilter(RowFilter): + """A range of values to restrict to in a row filter. + + Will only match cells that have values in this range. + + Both the start and end value can be included or excluded in the range. + By default, we include them both, but this can be changed with optional + flags. + + :type start_value: bytes + :param start_value: The start of the range of values. If no value is used, + the backend applies no lower bound to the values. + + :type end_value: bytes + :param end_value: The end of the range of values. If no value is used, + the backend applies no upper bound to the values. + + :type inclusive_start: bool + :param inclusive_start: Boolean indicating if the start value should be + included in the range (or excluded). Defaults + to :data:`True` if ``start_value`` is passed and + no ``inclusive_start`` was given. + + :type inclusive_end: bool + :param inclusive_end: Boolean indicating if the end value should be + included in the range (or excluded). Defaults + to :data:`True` if ``end_value`` is passed and + no ``inclusive_end`` was given. + + :raises: :class:`ValueError ` if ``inclusive_start`` + is set but no ``start_value`` is given or if ``inclusive_end`` + is set but no ``end_value`` is given + """ + + def __init__( + self, start_value=None, end_value=None, inclusive_start=None, inclusive_end=None + ): + if inclusive_start is None: + inclusive_start = True + elif start_value is None: + raise ValueError( + "Inclusive start was specified but no " "start value was given." + ) + if isinstance(start_value, int): + start_value = _PACK_I64(start_value) + self.start_value = start_value + self.inclusive_start = inclusive_start + + if inclusive_end is None: + inclusive_end = True + elif end_value is None: + raise ValueError( + "Inclusive end was specified but no " "end value was given." + ) + if isinstance(end_value, int): + end_value = _PACK_I64(end_value) + self.end_value = end_value + self.inclusive_end = inclusive_end + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return ( + other.start_value == self.start_value + and other.end_value == self.end_value + and other.inclusive_start == self.inclusive_start + and other.inclusive_end == self.inclusive_end + ) + + def __ne__(self, other): + return not self == other + + def to_pb(self): + """Converts the row filter to a protobuf. + + First converts to a :class:`.data_v2_pb2.ValueRange` and then uses + it to create a row filter protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + value_range_kwargs = {} + if self.start_value is not None: + if self.inclusive_start: + key = "start_value_closed" + else: + key = "start_value_open" + value_range_kwargs[key] = _to_bytes(self.start_value) + if self.end_value is not None: + if self.inclusive_end: + key = "end_value_closed" + else: + key = "end_value_open" + value_range_kwargs[key] = _to_bytes(self.end_value) + + value_range = data_v2_pb2.ValueRange(**value_range_kwargs) + return data_v2_pb2.RowFilter(value_range_filter=value_range) + + +class _CellCountFilter(RowFilter): + """Row filter that uses an integer count of cells. + + The cell count is used as an offset or a limit for the number + of results returned. + + :type num_cells: int + :param num_cells: An integer count / offset / limit. + """ + + def __init__(self, num_cells): + self.num_cells = num_cells + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other.num_cells == self.num_cells + + def __ne__(self, other): + return not self == other + + +class CellsRowOffsetFilter(_CellCountFilter): + """Row filter to skip cells in a row. + + :type num_cells: int + :param num_cells: Skips the first N cells of the row. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(cells_per_row_offset_filter=self.num_cells) + + +class CellsRowLimitFilter(_CellCountFilter): + """Row filter to limit cells in a row. + + :type num_cells: int + :param num_cells: Matches only the first N cells of the row. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(cells_per_row_limit_filter=self.num_cells) + + +class CellsColumnLimitFilter(_CellCountFilter): + """Row filter to limit cells in a column. + + :type num_cells: int + :param num_cells: Matches only the most recent N cells within each column. + This filters a (family name, column) pair, based on + timestamps of each cell. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(cells_per_column_limit_filter=self.num_cells) + + +class StripValueTransformerFilter(_BoolFilter): + """Row filter that transforms cells into empty string (0 bytes). + + :type flag: bool + :param flag: If :data:`True`, replaces each cell's value with the empty + string. As the name indicates, this is more useful as a + transformer than a generic query / filter. + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(strip_value_transformer=self.flag) + + +class ApplyLabelFilter(RowFilter): + """Filter to apply labels to cells. + + Intended to be used as an intermediate filter on a pre-existing filtered + result set. This way if two sets are combined, the label can tell where + the cell(s) originated.This allows the client to determine which results + were produced from which part of the filter. + + .. note:: + + Due to a technical limitation of the backend, it is not currently + possible to apply multiple labels to a cell. + + :type label: str + :param label: Label to apply to cells in the output row. Values must be + at most 15 characters long, and match the pattern + ``[a-z0-9\\-]+``. + """ + + def __init__(self, label): + self.label = label + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other.label == self.label + + def __ne__(self, other): + return not self == other + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(apply_label_transformer=self.label) + + +class _FilterCombination(RowFilter): + """Chain of row filters. + + Sends rows through several filters in sequence. The filters are "chained" + together to process a row. After the first filter is applied, the second + is applied to the filtered output and so on for subsequent filters. + + :type filters: list + :param filters: List of :class:`RowFilter` + """ + + def __init__(self, filters=None): + if filters is None: + filters = [] + self.filters = filters + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other.filters == self.filters + + def __ne__(self, other): + return not self == other + + +class RowFilterChain(_FilterCombination): + """Chain of row filters. + + Sends rows through several filters in sequence. The filters are "chained" + together to process a row. After the first filter is applied, the second + is applied to the filtered output and so on for subsequent filters. + + :type filters: list + :param filters: List of :class:`RowFilter` + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + chain = data_v2_pb2.RowFilter.Chain( + filters=[row_filter.to_pb() for row_filter in self.filters] + ) + return data_v2_pb2.RowFilter(chain=chain) + + +class RowFilterUnion(_FilterCombination): + """Union of row filters. + + Sends rows through several filters simultaneously, then + merges / interleaves all the filtered results together. + + If multiple cells are produced with the same column and timestamp, + they will all appear in the output row in an unspecified mutual order. + + :type filters: list + :param filters: List of :class:`RowFilter` + """ + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + interleave = data_v2_pb2.RowFilter.Interleave( + filters=[row_filter.to_pb() for row_filter in self.filters] + ) + return data_v2_pb2.RowFilter(interleave=interleave) + + +class ConditionalRowFilter(RowFilter): + """Conditional row filter which exhibits ternary behavior. + + Executes one of two filters based on another filter. If the ``base_filter`` + returns any cells in the row, then ``true_filter`` is executed. If not, + then ``false_filter`` is executed. + + .. note:: + + The ``base_filter`` does not execute atomically with the true and false + filters, which may lead to inconsistent or unexpected results. + + Additionally, executing a :class:`ConditionalRowFilter` has poor + performance on the server, especially when ``false_filter`` is set. + + :type base_filter: :class:`RowFilter` + :param base_filter: The filter to condition on before executing the + true/false filters. + + :type true_filter: :class:`RowFilter` + :param true_filter: (Optional) The filter to execute if there are any cells + matching ``base_filter``. If not provided, no results + will be returned in the true case. + + :type false_filter: :class:`RowFilter` + :param false_filter: (Optional) The filter to execute if there are no cells + matching ``base_filter``. If not provided, no results + will be returned in the false case. + """ + + def __init__(self, base_filter, true_filter=None, false_filter=None): + self.base_filter = base_filter + self.true_filter = true_filter + self.false_filter = false_filter + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return ( + other.base_filter == self.base_filter + and other.true_filter == self.true_filter + and other.false_filter == self.false_filter + ) + + def __ne__(self, other): + return not self == other + + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + condition_kwargs = {"predicate_filter": self.base_filter.to_pb()} + if self.true_filter is not None: + condition_kwargs["true_filter"] = self.true_filter.to_pb() + if self.false_filter is not None: + condition_kwargs["false_filter"] = self.false_filter.to_pb() + condition = data_v2_pb2.RowFilter.Condition(**condition_kwargs) + return data_v2_pb2.RowFilter(condition=condition) diff --git a/google/cloud/bigtable/row_merger.py b/google/cloud/bigtable/deprecated/row_merger.py similarity index 99% rename from google/cloud/bigtable/row_merger.py rename to google/cloud/bigtable/deprecated/row_merger.py index 515b91df7..d29d64eb2 100644 --- a/google/cloud/bigtable/row_merger.py +++ b/google/cloud/bigtable/deprecated/row_merger.py @@ -1,6 +1,6 @@ from enum import Enum from collections import OrderedDict -from google.cloud.bigtable.row import Cell, PartialRowData, InvalidChunk +from google.cloud.bigtable.deprecated.row import Cell, PartialRowData, InvalidChunk _MISSING_COLUMN_FAMILY = "Column family {} is not among the cells stored in this row." _MISSING_COLUMN = ( diff --git a/google/cloud/bigtable/row_set.py b/google/cloud/bigtable/deprecated/row_set.py similarity index 100% rename from google/cloud/bigtable/row_set.py rename to google/cloud/bigtable/deprecated/row_set.py diff --git a/google/cloud/bigtable/table.py b/google/cloud/bigtable/deprecated/table.py similarity index 95% rename from google/cloud/bigtable/table.py rename to google/cloud/bigtable/deprecated/table.py index 8605992ba..cf60b066e 100644 --- a/google/cloud/bigtable/table.py +++ b/google/cloud/bigtable/deprecated/table.py @@ -28,24 +28,24 @@ from google.api_core.retry import if_exception_type from google.api_core.retry import Retry from google.cloud._helpers import _to_bytes # type: ignore -from google.cloud.bigtable.backup import Backup -from google.cloud.bigtable.column_family import _gc_rule_from_pb -from google.cloud.bigtable.column_family import ColumnFamily -from google.cloud.bigtable.batcher import MutationsBatcher -from google.cloud.bigtable.batcher import FLUSH_COUNT, MAX_ROW_BYTES -from google.cloud.bigtable.encryption_info import EncryptionInfo -from google.cloud.bigtable.policy import Policy -from google.cloud.bigtable.row import AppendRow -from google.cloud.bigtable.row import ConditionalRow -from google.cloud.bigtable.row import DirectRow -from google.cloud.bigtable.row_data import ( +from google.cloud.bigtable.deprecated.backup import Backup +from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb +from google.cloud.bigtable.deprecated.column_family import ColumnFamily +from google.cloud.bigtable.deprecated.batcher import MutationsBatcher +from google.cloud.bigtable.deprecated.batcher import FLUSH_COUNT, MAX_ROW_BYTES +from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo +from google.cloud.bigtable.deprecated.policy import Policy +from google.cloud.bigtable.deprecated.row import AppendRow +from google.cloud.bigtable.deprecated.row import ConditionalRow +from google.cloud.bigtable.deprecated.row import DirectRow +from google.cloud.bigtable.deprecated.row_data import ( PartialRowsData, _retriable_internal_server_error, ) -from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS -from google.cloud.bigtable.row_set import RowSet -from google.cloud.bigtable.row_set import RowRange -from google.cloud.bigtable import enums +from google.cloud.bigtable.deprecated.row_data import DEFAULT_RETRY_READ_ROWS +from google.cloud.bigtable.deprecated.row_set import RowSet +from google.cloud.bigtable.deprecated.row_set import RowRange +from google.cloud.bigtable.deprecated import enums from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient from google.cloud.bigtable_admin_v2.types import table as admin_messages_v2_pb2 @@ -88,7 +88,7 @@ class _BigtableRetryableError(Exception): ) """The default retry strategy to be used on retry-able errors. -Used by :meth:`~google.cloud.bigtable.table.Table.mutate_rows`. +Used by :meth:`~google.cloud.bigtable.deprecated.table.Table.mutate_rows`. """ @@ -119,7 +119,7 @@ class Table(object): :type table_id: str :param table_id: The ID of the table. - :type instance: :class:`~google.cloud.bigtable.instance.Instance` + :type instance: :class:`~google.cloud.bigtable.deprecated.instance.Instance` :param instance: The instance that owns the table. :type app_profile_id: str @@ -172,7 +172,7 @@ def get_iam_policy(self): :end-before: [END bigtable_api_table_get_iam_policy] :dedent: 4 - :rtype: :class:`google.cloud.bigtable.policy.Policy` + :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` :returns: The current IAM policy of this table. """ table_client = self._instance._client.table_admin_client @@ -184,7 +184,7 @@ def set_iam_policy(self, policy): existing policy. For more information about policy, please see documentation of - class `google.cloud.bigtable.policy.Policy` + class `google.cloud.bigtable.deprecated.policy.Policy` For example: @@ -193,11 +193,11 @@ class `google.cloud.bigtable.policy.Policy` :end-before: [END bigtable_api_table_set_iam_policy] :dedent: 4 - :type policy: :class:`google.cloud.bigtable.policy.Policy` + :type policy: :class:`google.cloud.bigtable.deprecated.policy.Policy` :param policy: A new IAM policy to replace the current IAM policy of this table. - :rtype: :class:`google.cloud.bigtable.policy.Policy` + :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` :returns: The current IAM policy of this table. """ table_client = self._instance._client.table_admin_client @@ -271,7 +271,7 @@ def row(self, row_key, filter_=None, append=False): .. warning:: At most one of ``filter_`` and ``append`` can be used in a - :class:`~google.cloud.bigtable.row.Row`. + :class:`~google.cloud.bigtable.deprecated.row.Row`. :type row_key: bytes :param row_key: The key for the row being created. @@ -284,7 +284,7 @@ def row(self, row_key, filter_=None, append=False): :param append: (Optional) Flag to determine if the row should be used for append mutations. - :rtype: :class:`~google.cloud.bigtable.row.Row` + :rtype: :class:`~google.cloud.bigtable.deprecated.row.Row` :returns: A row owned by this table. :raises: :class:`ValueError ` if both ``filter_`` and ``append`` are used. @@ -307,7 +307,7 @@ def row(self, row_key, filter_=None, append=False): return DirectRow(row_key, self) def append_row(self, row_key): - """Create a :class:`~google.cloud.bigtable.row.AppendRow` associated with this table. + """Create a :class:`~google.cloud.bigtable.deprecated.row.AppendRow` associated with this table. For example: @@ -325,7 +325,7 @@ def append_row(self, row_key): return AppendRow(row_key, self) def direct_row(self, row_key): - """Create a :class:`~google.cloud.bigtable.row.DirectRow` associated with this table. + """Create a :class:`~google.cloud.bigtable.deprecated.row.DirectRow` associated with this table. For example: @@ -343,7 +343,7 @@ def direct_row(self, row_key): return DirectRow(row_key, self) def conditional_row(self, row_key, filter_): - """Create a :class:`~google.cloud.bigtable.row.ConditionalRow` associated with this table. + """Create a :class:`~google.cloud.bigtable.deprecated.row.ConditionalRow` associated with this table. For example: @@ -515,7 +515,7 @@ def get_encryption_info(self): :rtype: dict :returns: Dictionary of encryption info for this table. Keys are cluster ids and - values are tuples of :class:`google.cloud.bigtable.encryption.EncryptionInfo` instances. + values are tuples of :class:`google.cloud.bigtable.deprecated.encryption.EncryptionInfo` instances. """ ENCRYPTION_VIEW = enums.Table.View.ENCRYPTION_VIEW table_client = self._instance._client.table_admin_client @@ -967,7 +967,7 @@ def list_backups(self, cluster_id=None, filter_=None, order_by=None, page_size=0 number of resources in a page. :rtype: :class:`~google.api_core.page_iterator.Iterator` - :returns: Iterator of :class:`~google.cloud.bigtable.backup.Backup` + :returns: Iterator of :class:`~google.cloud.bigtable.deprecated.backup.Backup` resources within the current Instance. :raises: :class:`ValueError ` if one of the returned Backups' name is not of the expected format. @@ -1367,8 +1367,8 @@ def _check_row_table_name(table_name, row): :type table_name: str :param table_name: The name of the table. - :type row: :class:`~google.cloud.bigtable.row.Row` - :param row: An instance of :class:`~google.cloud.bigtable.row.Row` + :type row: :class:`~google.cloud.bigtable.deprecated.row.Row` + :param row: An instance of :class:`~google.cloud.bigtable.deprecated.row.Row` subclasses. :raises: :exc:`~.table.TableMismatchError` if the row does not belong to @@ -1384,8 +1384,8 @@ def _check_row_table_name(table_name, row): def _check_row_type(row): """Checks that a row is an instance of :class:`.DirectRow`. - :type row: :class:`~google.cloud.bigtable.row.Row` - :param row: An instance of :class:`~google.cloud.bigtable.row.Row` + :type row: :class:`~google.cloud.bigtable.deprecated.row.Row` + :param row: An instance of :class:`~google.cloud.bigtable.deprecated.row.Row` subclasses. :raises: :class:`TypeError ` if the row is not an diff --git a/google/cloud/bigtable/exceptions.py b/google/cloud/bigtable/exceptions.py new file mode 100644 index 000000000..86bfe9247 --- /dev/null +++ b/google/cloud/bigtable/exceptions.py @@ -0,0 +1,46 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# + +import sys + + +is_311_plus = sys.version_info >= (3, 11) + + +class BigtableExceptionGroup(ExceptionGroup if is_311_plus else Exception): # type: ignore # noqa: F821 + """ + Represents one or more exceptions that occur during a bulk Bigtable operation + + In Python 3.11+, this is an unmodified exception group. In < 3.10, it is a + custom exception with some exception group functionality backported, but does + Not implement the full API + """ + + def __init__(self, message, excs): + raise NotImplementedError() + + +class MutationsExceptionGroup(BigtableExceptionGroup): + """ + Represents one or more exceptions that occur during a bulk mutation operation + """ + + pass + + +class RetryExceptionGroup(BigtableExceptionGroup): + """Represents one or more exceptions that occur during a retryable operation""" + + pass diff --git a/google/cloud/bigtable/mutations.py b/google/cloud/bigtable/mutations.py new file mode 100644 index 000000000..ed3c2f065 --- /dev/null +++ b/google/cloud/bigtable/mutations.py @@ -0,0 +1,54 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from dataclasses import dataclass +from google.cloud.bigtable.row_response import family_id, qualifier, row_key + + +class Mutation: + pass + + +@dataclass +class SetCell(Mutation): + family: family_id + column_qualifier: qualifier + new_value: bytes | str | int + timestamp_ms: int | None = None + + +@dataclass +class DeleteRangeFromColumn(Mutation): + family: family_id + column_qualifier: qualifier + start_timestamp_ms: int + end_timestamp_ms: int + + +@dataclass +class DeleteAllFromFamily(Mutation): + family_to_delete: family_id + + +@dataclass +class DeleteAllFromRow(Mutation): + pass + + +@dataclass +class BulkMutationsEntry: + row: row_key + mutations: list[Mutation] | Mutation diff --git a/google/cloud/bigtable/mutations_batcher.py b/google/cloud/bigtable/mutations_batcher.py new file mode 100644 index 000000000..2e393cc7e --- /dev/null +++ b/google/cloud/bigtable/mutations_batcher.py @@ -0,0 +1,104 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from google.cloud.bigtable.mutations import Mutation +from google.cloud.bigtable.row_response import row_key +from google.cloud.bigtable.row_filters import RowFilter + +if TYPE_CHECKING: + from google.cloud.bigtable.client import Table # pragma: no cover + + +class MutationsBatcher: + """ + Allows users to send batches using context manager API: + + Runs mutate_row, mutate_rows, and check_and_mutate_row internally, combining + to use as few network requests as required + + Flushes: + - manually + - every flush_interval seconds + - after queue reaches flush_count in quantity + - after queue reaches flush_size_bytes in storage size + - when batcher is closed or destroyed + + async with table.mutations_batcher() as batcher: + for i in range(10): + batcher.add(row, mut) + """ + + queue: asyncio.Queue[tuple[row_key, list[Mutation]]] + conditional_queues: dict[RowFilter, tuple[list[Mutation], list[Mutation]]] + + MB_SIZE = 1024 * 1024 + + def __init__( + self, + table: "Table", + flush_count: int = 100, + flush_size_bytes: int = 100 * MB_SIZE, + max_mutation_bytes: int = 20 * MB_SIZE, + flush_interval: int = 5, + metadata: list[tuple[str, str]] | None = None, + ): + raise NotImplementedError + + async def append(self, row_key: str | bytes, mutation: Mutation | list[Mutation]): + """ + Add a new mutation to the internal queue + """ + raise NotImplementedError + + async def append_conditional( + self, + predicate_filter: RowFilter, + row_key: str | bytes, + if_true_mutations: Mutation | list[Mutation] | None = None, + if_false_mutations: Mutation | list[Mutation] | None = None, + ): + """ + Apply a different set of mutations based on the outcome of predicate_filter + + Calls check_and_mutate_row internally on flush + """ + raise NotImplementedError + + async def flush(self): + """ + Send queue over network in as few calls as possible + + Raises: + - MutationsExceptionGroup if any mutation in the batch fails + """ + raise NotImplementedError + + async def __aenter__(self): + """For context manager API""" + raise NotImplementedError + + async def __aexit__(self, exc_type, exc, tb): + """For context manager API""" + raise NotImplementedError + + async def close(self): + """ + Flush queue and clean up resources + """ + raise NotImplementedError diff --git a/google/cloud/bigtable/read_modify_write_rules.py b/google/cloud/bigtable/read_modify_write_rules.py new file mode 100644 index 000000000..a9b0885f2 --- /dev/null +++ b/google/cloud/bigtable/read_modify_write_rules.py @@ -0,0 +1,37 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from dataclasses import dataclass + +from google.cloud.bigtable.row_response import family_id, qualifier + + +class ReadModifyWriteRule: + pass + + +@dataclass +class IncrementRule(ReadModifyWriteRule): + increment_amount: int + family: family_id + column_qualifier: qualifier + + +@dataclass +class AppendValueRule(ReadModifyWriteRule): + append_value: bytes | str + family: family_id + column_qualifier: qualifier diff --git a/google/cloud/bigtable/read_rows_query.py b/google/cloud/bigtable/read_rows_query.py new file mode 100644 index 000000000..64583b2d7 --- /dev/null +++ b/google/cloud/bigtable/read_rows_query.py @@ -0,0 +1,56 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from google.cloud.bigtable.row_filters import RowFilter + from google.cloud.bigtable import RowKeySamples + + +class ReadRowsQuery: + """ + Class to encapsulate details of a read row request + """ + + def __init__( + self, row_keys: list[str | bytes] | str | bytes | None = None, limit=None + ): + pass + + def set_limit(self, limit: int) -> ReadRowsQuery: + raise NotImplementedError + + def set_filter(self, filter: "RowFilter") -> ReadRowsQuery: + raise NotImplementedError + + def add_rows(self, row_id_list: list[str]) -> ReadRowsQuery: + raise NotImplementedError + + def add_range( + self, start_key: str | bytes | None = None, end_key: str | bytes | None = None + ) -> ReadRowsQuery: + raise NotImplementedError + + def shard(self, shard_keys: "RowKeySamples" | None = None) -> list[ReadRowsQuery]: + """ + Split this query into multiple queries that can be evenly distributed + across nodes and be run in parallel + + Returns: + - a list of queries that represent a sharded version of the original + query (if possible) + """ + raise NotImplementedError diff --git a/google/cloud/bigtable/row_response.py b/google/cloud/bigtable/row_response.py new file mode 100644 index 000000000..be6d8c505 --- /dev/null +++ b/google/cloud/bigtable/row_response.py @@ -0,0 +1,130 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from collections import OrderedDict +from typing import Sequence + +# Type aliases used internally for readability. +row_key = bytes +family_id = str +qualifier = bytes +row_value = bytes + + +class RowResponse(Sequence["CellResponse"]): + """ + Model class for row data returned from server + + Does not represent all data contained in the row, only data returned by a + query. + Expected to be read-only to users, and written by backend + + Can be indexed: + cells = row["family", "qualifier"] + """ + + def __init__(self, key: row_key, cells: list[CellResponse]): + self.row_key = key + self.cells: OrderedDict[ + family_id, OrderedDict[qualifier, list[CellResponse]] + ] = OrderedDict() + """Expected to be used internally only""" + pass + + def get_cells( + self, family: str | None, qualifer: str | bytes | None + ) -> list[CellResponse]: + """ + Returns cells sorted in Bigtable native order: + - Family lexicographically ascending + - Qualifier lexicographically ascending + - Timestamp in reverse chronological order + + If family or qualifier not passed, will include all + + Syntactic sugar: cells = row["family", "qualifier"] + """ + raise NotImplementedError + + def get_index(self) -> dict[family_id, list[qualifier]]: + """ + Returns a list of family and qualifiers for the object + """ + raise NotImplementedError + + def __str__(self): + """ + Human-readable string representation + + (family, qualifier) cells + (ABC, XYZ) [b"123", b"456" ...(+5)] + (DEF, XYZ) [b"123"] + (GHI, XYZ) [b"123", b"456" ...(+2)] + """ + raise NotImplementedError + + +class CellResponse: + """ + Model class for cell data + + Does not represent all data contained in the cell, only data returned by a + query. + Expected to be read-only to users, and written by backend + """ + + def __init__( + self, + value: row_value, + row: row_key, + family: family_id, + column_qualifier: qualifier, + labels: list[str] | None = None, + timestamp: int | None = None, + ): + self.value = value + self.row_key = row + self.family = family + self.column_qualifier = column_qualifier + self.labels = labels + self.timestamp = timestamp + + def decode_value(self, encoding="UTF-8", errors=None) -> str: + """decode bytes to string""" + return self.value.decode(encoding, errors) + + def __int__(self) -> int: + """ + Allows casting cell to int + Interprets value as a 64-bit big-endian signed integer, as expected by + ReadModifyWrite increment rule + """ + return int.from_bytes(self.value, byteorder="big", signed=True) + + def __str__(self) -> str: + """ + Allows casting cell to str + Prints encoded byte string, same as printing value directly. + """ + return str(self.value) + + """For Bigtable native ordering""" + + def __lt__(self, other) -> bool: + raise NotImplementedError + + def __eq__(self, other) -> bool: + raise NotImplementedError diff --git a/noxfile.py b/noxfile.py index 47415385a..8ce6d5d95 100644 --- a/noxfile.py +++ b/noxfile.py @@ -128,8 +128,20 @@ def mypy(session): session.install("-e", ".") session.install("mypy", "types-setuptools", "types-protobuf", "types-mock") session.install("google-cloud-testutils") - # TODO: also verify types on tests, all of google package - session.run("mypy", "google/", "tests/") + session.run( + "mypy", + "google/cloud/bigtable", + "tests/", + "--check-untyped-defs", + "--warn-unreachable", + "--disallow-any-generics", + "--exclude", + "google/cloud/bigtable/deprecated", + "--exclude", + "tests/system/v2_client", + "--exclude", + "tests/unit/v2_client", + ) @nox.session(python=DEFAULT_PYTHON_VERSION) @@ -306,7 +318,7 @@ def cover(session): test runs (not system test runs), and then erases coverage data. """ session.install("coverage", "pytest-cov") - session.run("coverage", "report", "--show-missing", "--fail-under=100") + session.run("coverage", "report", "--show-missing", "--fail-under=99") session.run("coverage", "erase") diff --git a/owlbot.py b/owlbot.py index b6aa2f8a2..d7eb3eaf2 100644 --- a/owlbot.py +++ b/owlbot.py @@ -89,7 +89,7 @@ def get_staging_dirs( samples=True, # set to True only if there are samples split_system_tests=True, microgenerator=True, - cov_level=100, + cov_level=99, ) s.move(templated_files, excludes=[".coveragerc", "README.rst", ".github/release-please.yml"]) @@ -168,8 +168,20 @@ def mypy(session): session.install("-e", ".") session.install("mypy", "types-setuptools", "types-protobuf", "types-mock") session.install("google-cloud-testutils") - # TODO: also verify types on tests, all of google package - session.run("mypy", "google/", "tests/") + session.run( + "mypy", + "google/cloud/bigtable", + "tests/", + "--check-untyped-defs", + "--warn-unreachable", + "--disallow-any-generics", + "--exclude", + "google/cloud/bigtable/deprecated", + "--exclude", + "tests/system/v2_client", + "--exclude", + "tests/unit/v2_client", + ) @nox.session(python=DEFAULT_PYTHON_VERSION) diff --git a/tests/system/__init__.py b/tests/system/__init__.py index 4de65971c..89a37dc92 100644 --- a/tests/system/__init__.py +++ b/tests/system/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/system/v2_client/__init__.py b/tests/system/v2_client/__init__.py new file mode 100644 index 000000000..4de65971c --- /dev/null +++ b/tests/system/v2_client/__init__.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +# Copyright 2020 Google LLC +# +# Licensed 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. +# diff --git a/tests/system/_helpers.py b/tests/system/v2_client/_helpers.py similarity index 100% rename from tests/system/_helpers.py rename to tests/system/v2_client/_helpers.py diff --git a/tests/system/conftest.py b/tests/system/v2_client/conftest.py similarity index 98% rename from tests/system/conftest.py rename to tests/system/v2_client/conftest.py index f39fcba88..bb4f54b41 100644 --- a/tests/system/conftest.py +++ b/tests/system/v2_client/conftest.py @@ -17,7 +17,7 @@ import pytest from test_utils.system import unique_resource_id -from google.cloud.bigtable.client import Client +from google.cloud.bigtable.deprecated.client import Client from google.cloud.environment_vars import BIGTABLE_EMULATOR from . import _helpers diff --git a/tests/system/test_data_api.py b/tests/system/v2_client/test_data_api.py similarity index 94% rename from tests/system/test_data_api.py rename to tests/system/v2_client/test_data_api.py index 2ca7e1504..551a221ee 100644 --- a/tests/system/test_data_api.py +++ b/tests/system/v2_client/test_data_api.py @@ -60,7 +60,7 @@ def rows_to_delete(): def test_table_read_rows_filter_millis(data_table): - from google.cloud.bigtable import row_filters + from google.cloud.bigtable.deprecated import row_filters end = datetime.datetime.now() start = end - datetime.timedelta(minutes=60) @@ -158,8 +158,8 @@ def test_table_drop_by_prefix(data_table, rows_to_delete): def test_table_read_rows_w_row_set(data_table, rows_to_delete): - from google.cloud.bigtable.row_set import RowSet - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange row_keys = [ b"row_key_1", @@ -189,7 +189,7 @@ def test_table_read_rows_w_row_set(data_table, rows_to_delete): def test_rowset_add_row_range_w_pfx(data_table, rows_to_delete): - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowSet row_keys = [ b"row_key_1", @@ -234,7 +234,7 @@ def _write_to_row(row1, row2, row3, row4): from google.cloud._helpers import _datetime_from_microseconds from google.cloud._helpers import _microseconds_from_datetime from google.cloud._helpers import UTC - from google.cloud.bigtable.row_data import Cell + from google.cloud.bigtable.deprecated.row_data import Cell timestamp1 = datetime.datetime.utcnow().replace(tzinfo=UTC) timestamp1_micros = _microseconds_from_datetime(timestamp1) @@ -290,7 +290,7 @@ def test_table_read_row(data_table, rows_to_delete): def test_table_read_rows(data_table, rows_to_delete): - from google.cloud.bigtable.row_data import PartialRowData + from google.cloud.bigtable.deprecated.row_data import PartialRowData row = data_table.direct_row(ROW_KEY) rows_to_delete.append(row) @@ -326,10 +326,10 @@ def test_table_read_rows(data_table, rows_to_delete): def test_read_with_label_applied(data_table, rows_to_delete, skip_on_emulator): - from google.cloud.bigtable.row_filters import ApplyLabelFilter - from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnQualifierRegexFilter + from google.cloud.bigtable.deprecated.row_filters import RowFilterChain + from google.cloud.bigtable.deprecated.row_filters import RowFilterUnion row = data_table.direct_row(ROW_KEY) rows_to_delete.append(row) diff --git a/tests/system/test_instance_admin.py b/tests/system/v2_client/test_instance_admin.py similarity index 99% rename from tests/system/test_instance_admin.py rename to tests/system/v2_client/test_instance_admin.py index e5e311213..debe1ab56 100644 --- a/tests/system/test_instance_admin.py +++ b/tests/system/v2_client/test_instance_admin.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud.bigtable import enums -from google.cloud.bigtable.table import ClusterState +from google.cloud.bigtable.deprecated import enums +from google.cloud.bigtable.deprecated.table import ClusterState from . import _helpers @@ -149,7 +149,7 @@ def test_instance_create_prod( instances_to_delete, skip_on_emulator, ): - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums alt_instance_id = f"ndef{unique_suffix}" instance = admin_client.instance(alt_instance_id, labels=instance_labels) diff --git a/tests/system/test_table_admin.py b/tests/system/v2_client/test_table_admin.py similarity index 96% rename from tests/system/test_table_admin.py rename to tests/system/v2_client/test_table_admin.py index c50189013..107ed41bf 100644 --- a/tests/system/test_table_admin.py +++ b/tests/system/v2_client/test_table_admin.py @@ -97,7 +97,7 @@ def test_table_create_w_families( data_instance_populated, tables_to_delete, ): - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule temp_table_id = "test-create-table-with-failies" column_family_id = "col-fam-id1" @@ -134,7 +134,7 @@ def test_table_create_w_split_keys( def test_column_family_create(data_instance_populated, tables_to_delete): - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule temp_table_id = "test-create-column-family" temp_table = data_instance_populated.table(temp_table_id) @@ -158,7 +158,7 @@ def test_column_family_create(data_instance_populated, tables_to_delete): def test_column_family_update(data_instance_populated, tables_to_delete): - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule temp_table_id = "test-update-column-family" temp_table = data_instance_populated.table(temp_table_id) @@ -219,8 +219,8 @@ def test_table_get_iam_policy( def test_table_set_iam_policy( service_account, data_instance_populated, tables_to_delete, skip_on_emulator ): - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE - from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import Policy temp_table_id = "test-set-iam-policy-table" temp_table = data_instance_populated.table(temp_table_id) @@ -264,7 +264,7 @@ def test_table_backup( skip_on_emulator, ): from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums temp_table_id = "test-backup-table" temp_table = data_instance_populated.table(temp_table_id) diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py index e8e1c3845..89a37dc92 100644 --- a/tests/unit/__init__.py +++ b/tests/unit/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/unit/v2_client/__init__.py b/tests/unit/v2_client/__init__.py new file mode 100644 index 000000000..e8e1c3845 --- /dev/null +++ b/tests/unit/v2_client/__init__.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +# Copyright 2022 Google LLC +# +# Licensed 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. +# diff --git a/tests/unit/_testing.py b/tests/unit/v2_client/_testing.py similarity index 100% rename from tests/unit/_testing.py rename to tests/unit/v2_client/_testing.py diff --git a/tests/unit/read-rows-acceptance-test.json b/tests/unit/v2_client/read-rows-acceptance-test.json similarity index 100% rename from tests/unit/read-rows-acceptance-test.json rename to tests/unit/v2_client/read-rows-acceptance-test.json diff --git a/tests/unit/test_app_profile.py b/tests/unit/v2_client/test_app_profile.py similarity index 94% rename from tests/unit/test_app_profile.py rename to tests/unit/v2_client/test_app_profile.py index 660ee7899..575f25194 100644 --- a/tests/unit/test_app_profile.py +++ b/tests/unit/v2_client/test_app_profile.py @@ -32,19 +32,19 @@ def _make_app_profile(*args, **kwargs): - from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.deprecated.app_profile import AppProfile return AppProfile(*args, **kwargs) def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client return Client(*args, **kwargs) def test_app_profile_constructor_defaults(): - from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.deprecated.app_profile import AppProfile client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -60,7 +60,7 @@ def test_app_profile_constructor_defaults(): def test_app_profile_constructor_explicit(): - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType ANY = RoutingPolicyType.ANY DESCRIPTION_1 = "routing policy any" @@ -99,7 +99,7 @@ def test_app_profile_constructor_explicit(): def test_app_profile_constructor_multi_cluster_ids(): - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType ANY = RoutingPolicyType.ANY DESCRIPTION_1 = "routing policy any" @@ -166,8 +166,8 @@ def test_app_profile___ne__(): def test_app_profile_from_pb_success_w_routing_any(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.app_profile import AppProfile - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -195,8 +195,8 @@ def test_app_profile_from_pb_success_w_routing_any(): def test_app_profile_from_pb_success_w_routing_any_multi_cluster_ids(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.app_profile import AppProfile - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -226,8 +226,8 @@ def test_app_profile_from_pb_success_w_routing_any_multi_cluster_ids(): def test_app_profile_from_pb_success_w_routing_single(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.app_profile import AppProfile - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -259,7 +259,7 @@ def test_app_profile_from_pb_success_w_routing_single(): def test_app_profile_from_pb_w_bad_app_profile_name(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.deprecated.app_profile import AppProfile bad_app_profile_name = "BAD_NAME" @@ -271,7 +271,7 @@ def test_app_profile_from_pb_w_bad_app_profile_name(): def test_app_profile_from_pb_w_instance_id_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.deprecated.app_profile import AppProfile ALT_INSTANCE_ID = "ALT_INSTANCE_ID" client = _Client(PROJECT) @@ -286,7 +286,7 @@ def test_app_profile_from_pb_w_instance_id_mistmatch(): def test_app_profile_from_pb_w_project_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.deprecated.app_profile import AppProfile ALT_PROJECT = "ALT_PROJECT" client = _Client(project=ALT_PROJECT) @@ -304,7 +304,7 @@ def test_app_profile_reload_w_routing_any(): BigtableInstanceAdminClient, ) from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType api = mock.create_autospec(BigtableInstanceAdminClient) credentials = _make_credentials() @@ -400,8 +400,8 @@ def test_app_profile_create_w_routing_any(): from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) - from google.cloud.bigtable.app_profile import AppProfile - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -461,8 +461,8 @@ def test_app_profile_create_w_routing_single(): from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) - from google.cloud.bigtable.app_profile import AppProfile - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -533,7 +533,7 @@ def test_app_profile_update_w_routing_any(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -608,7 +608,7 @@ def test_app_profile_update_w_routing_any_multi_cluster_ids(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -684,7 +684,7 @@ def test_app_profile_update_w_routing_single(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) diff --git a/tests/unit/test_backup.py b/tests/unit/v2_client/test_backup.py similarity index 96% rename from tests/unit/test_backup.py rename to tests/unit/v2_client/test_backup.py index 9882ca339..34cc8823a 100644 --- a/tests/unit/test_backup.py +++ b/tests/unit/v2_client/test_backup.py @@ -48,7 +48,7 @@ def _make_table_admin_client(): def _make_backup(*args, **kwargs): - from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.deprecated.backup import Backup return Backup(*args, **kwargs) @@ -102,7 +102,7 @@ def test_backup_constructor_explicit(): def test_backup_from_pb_w_project_mismatch(): from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.deprecated.backup import Backup alt_project_id = "alt-project-id" client = _Client(project=alt_project_id) @@ -115,7 +115,7 @@ def test_backup_from_pb_w_project_mismatch(): def test_backup_from_pb_w_instance_mismatch(): from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.deprecated.backup import Backup alt_instance = "/projects/%s/instances/alt-instance" % PROJECT_ID client = _Client() @@ -128,7 +128,7 @@ def test_backup_from_pb_w_instance_mismatch(): def test_backup_from_pb_w_bad_name(): from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.deprecated.backup import Backup client = _Client() instance = _Instance(INSTANCE_NAME, client) @@ -139,10 +139,10 @@ def test_backup_from_pb_w_bad_name(): def test_backup_from_pb_success(): - from google.cloud.bigtable.encryption_info import EncryptionInfo - from google.cloud.bigtable.error import Status + from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo + from google.cloud.bigtable.deprecated.error import Status from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.deprecated.backup import Backup from google.cloud._helpers import _datetime_to_pb_timestamp from google.rpc.code_pb2 import Code @@ -190,7 +190,7 @@ def test_backup_from_pb_success(): def test_backup_name(): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -225,7 +225,7 @@ def test_backup_parent_none(): def test_backup_parent_w_cluster(): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -242,7 +242,7 @@ def test_backup_parent_w_cluster(): def test_backup_source_table_none(): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -258,7 +258,7 @@ def test_backup_source_table_none(): def test_backup_source_table_valid(): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -473,7 +473,7 @@ def test_backup_create_w_expire_time_not_set(): def test_backup_create_success(): from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable import Client + from google.cloud.bigtable.deprecated import Client op_future = object() credentials = _make_credentials() @@ -806,12 +806,12 @@ def test_backup_restore_to_another_instance(): def test_backup_get_iam_policy(): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = Client(project=PROJECT_ID, credentials=credentials, admin=True) @@ -842,13 +842,13 @@ def test_backup_get_iam_policy(): def test_backup_set_iam_policy(): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = Client(project=PROJECT_ID, credentials=credentials, admin=True) @@ -887,7 +887,7 @@ def test_backup_set_iam_policy(): def test_backup_test_iam_permissions(): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) diff --git a/tests/unit/test_batcher.py b/tests/unit/v2_client/test_batcher.py similarity index 91% rename from tests/unit/test_batcher.py rename to tests/unit/v2_client/test_batcher.py index 9ae6ed175..0793ed480 100644 --- a/tests/unit/test_batcher.py +++ b/tests/unit/v2_client/test_batcher.py @@ -16,14 +16,14 @@ import mock import pytest -from google.cloud.bigtable.row import DirectRow +from google.cloud.bigtable.deprecated.row import DirectRow TABLE_ID = "table-id" TABLE_NAME = "/tables/" + TABLE_ID def _make_mutation_batcher(table, **kw): - from google.cloud.bigtable.batcher import MutationsBatcher + from google.cloud.bigtable.deprecated.batcher import MutationsBatcher return MutationsBatcher(table, **kw) @@ -92,9 +92,9 @@ def test_mutation_batcher_mutate_w_max_flush_count(): assert table.mutation_calls == 1 -@mock.patch("google.cloud.bigtable.batcher.MAX_MUTATIONS", new=3) +@mock.patch("google.cloud.bigtable.deprecated.batcher.MAX_MUTATIONS", new=3) def test_mutation_batcher_mutate_with_max_mutations_failure(): - from google.cloud.bigtable.batcher import MaxMutationsError + from google.cloud.bigtable.deprecated.batcher import MaxMutationsError table = _Table(TABLE_NAME) mutation_batcher = _make_mutation_batcher(table=table) @@ -109,7 +109,7 @@ def test_mutation_batcher_mutate_with_max_mutations_failure(): mutation_batcher.mutate(row) -@mock.patch("google.cloud.bigtable.batcher.MAX_MUTATIONS", new=3) +@mock.patch("google.cloud.bigtable.deprecated.batcher.MAX_MUTATIONS", new=3) def test_mutation_batcher_mutate_w_max_mutations(): table = _Table(TABLE_NAME) mutation_batcher = _make_mutation_batcher(table=table) diff --git a/tests/unit/test_client.py b/tests/unit/v2_client/test_client.py similarity index 93% rename from tests/unit/test_client.py rename to tests/unit/v2_client/test_client.py index 5944c58a3..9deac6a25 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/v2_client/test_client.py @@ -25,7 +25,7 @@ def _invoke_client_factory(client_class, **kw): - from google.cloud.bigtable.client import _create_gapic_client + from google.cloud.bigtable.deprecated.client import _create_gapic_client return _create_gapic_client(client_class, **kw) @@ -101,23 +101,27 @@ def __init__(self, credentials, emulator_host=None, emulator_channel=None): def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client return Client(*args, **kwargs) @mock.patch("os.environ", {}) def test_client_constructor_defaults(): + import warnings from google.api_core import client_info - from google.cloud.bigtable import __version__ - from google.cloud.bigtable.client import DATA_SCOPE + from google.cloud.bigtable.deprecated import __version__ + from google.cloud.bigtable.deprecated.client import DATA_SCOPE credentials = _make_credentials() - with mock.patch("google.auth.default") as mocked: - mocked.return_value = credentials, PROJECT - client = _make_client() + with warnings.catch_warnings(record=True) as warned: + with mock.patch("google.auth.default") as mocked: + mocked.return_value = credentials, PROJECT + client = _make_client() + # warn about client deprecation + assert len(warned) == 1 assert client.project == PROJECT assert client._credentials is credentials.with_scopes.return_value assert not client._read_only @@ -131,8 +135,8 @@ def test_client_constructor_defaults(): def test_client_constructor_explicit(): import warnings - from google.cloud.bigtable.client import ADMIN_SCOPE - from google.cloud.bigtable.client import DATA_SCOPE + from google.cloud.bigtable.deprecated.client import ADMIN_SCOPE + from google.cloud.bigtable.deprecated.client import DATA_SCOPE credentials = _make_credentials() client_info = mock.Mock() @@ -147,7 +151,8 @@ def test_client_constructor_explicit(): channel=mock.sentinel.channel, ) - assert len(warned) == 1 + # deprecationw arnning for channel and Client deprecation + assert len(warned) == 2 assert client.project == PROJECT assert client._credentials is credentials.with_scopes.return_value @@ -171,8 +176,10 @@ def test_client_constructor_w_both_admin_and_read_only(): def test_client_constructor_w_emulator_host(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT - from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.deprecated.client import ( + _DEFAULT_BIGTABLE_EMULATOR_CLIENT, + ) + from google.cloud.bigtable.deprecated.client import _GRPC_CHANNEL_OPTIONS emulator_host = "localhost:8081" with mock.patch("os.environ", {BIGTABLE_EMULATOR: emulator_host}): @@ -195,7 +202,7 @@ def test_client_constructor_w_emulator_host(): def test_client_constructor_w_emulator_host_w_project(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.deprecated.client import _GRPC_CHANNEL_OPTIONS emulator_host = "localhost:8081" with mock.patch("os.environ", {BIGTABLE_EMULATOR: emulator_host}): @@ -216,8 +223,10 @@ def test_client_constructor_w_emulator_host_w_project(): def test_client_constructor_w_emulator_host_w_credentials(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT - from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.deprecated.client import ( + _DEFAULT_BIGTABLE_EMULATOR_CLIENT, + ) + from google.cloud.bigtable.deprecated.client import _GRPC_CHANNEL_OPTIONS emulator_host = "localhost:8081" credentials = _make_credentials() @@ -238,15 +247,15 @@ def test_client_constructor_w_emulator_host_w_credentials(): def test_client__get_scopes_default(): - from google.cloud.bigtable.client import DATA_SCOPE + from google.cloud.bigtable.deprecated.client import DATA_SCOPE client = _make_client(project=PROJECT, credentials=_make_credentials()) assert client._get_scopes() == (DATA_SCOPE,) def test_client__get_scopes_w_admin(): - from google.cloud.bigtable.client import ADMIN_SCOPE - from google.cloud.bigtable.client import DATA_SCOPE + from google.cloud.bigtable.deprecated.client import ADMIN_SCOPE + from google.cloud.bigtable.deprecated.client import DATA_SCOPE client = _make_client(project=PROJECT, credentials=_make_credentials(), admin=True) expected_scopes = (DATA_SCOPE, ADMIN_SCOPE) @@ -254,7 +263,7 @@ def test_client__get_scopes_w_admin(): def test_client__get_scopes_w_read_only(): - from google.cloud.bigtable.client import READ_ONLY_SCOPE + from google.cloud.bigtable.deprecated.client import READ_ONLY_SCOPE client = _make_client( project=PROJECT, credentials=_make_credentials(), read_only=True @@ -344,7 +353,7 @@ def test_client__local_composite_credentials(): def _create_gapic_client_channel_helper(endpoint=None, emulator_host=None): - from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.deprecated.client import _GRPC_CHANNEL_OPTIONS client_class = mock.Mock(spec=["DEFAULT_ENDPOINT"]) credentials = _make_credentials() @@ -618,7 +627,7 @@ def test_client_instance_admin_client_initialized(): def test_client_instance_factory_defaults(): - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials) @@ -634,8 +643,8 @@ def test_client_instance_factory_defaults(): def test_client_instance_factory_non_defaults(): - from google.cloud.bigtable.instance import Instance - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.deprecated import enums instance_type = enums.Instance.Type.DEVELOPMENT labels = {"foo": "bar"} @@ -665,7 +674,7 @@ def test_client_list_instances(): from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance FAILED_LOCATION = "FAILED" INSTANCE_ID1 = "instance-id1" @@ -717,7 +726,7 @@ def test_client_list_clusters(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.instance import Cluster + from google.cloud.bigtable.deprecated.instance import Cluster instance_api = mock.create_autospec(BigtableInstanceAdminClient) diff --git a/tests/unit/test_cluster.py b/tests/unit/v2_client/test_cluster.py similarity index 94% rename from tests/unit/test_cluster.py rename to tests/unit/v2_client/test_cluster.py index cb0312b0c..e667c2af4 100644 --- a/tests/unit/test_cluster.py +++ b/tests/unit/v2_client/test_cluster.py @@ -42,13 +42,13 @@ def _make_cluster(*args, **kwargs): - from google.cloud.bigtable.cluster import Cluster + from google.cloud.bigtable.deprecated.cluster import Cluster return Cluster(*args, **kwargs) def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client return Client(*args, **kwargs) @@ -72,8 +72,8 @@ def test_cluster_constructor_defaults(): def test_cluster_constructor_explicit(): - from google.cloud.bigtable.enums import StorageType - from google.cloud.bigtable.enums import Cluster + from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.deprecated.enums import Cluster STATE = Cluster.State.READY STORAGE_TYPE_SSD = StorageType.SSD @@ -126,8 +126,8 @@ def test_cluster_kms_key_name_setter(): def test_cluster_from_pb_success(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.cluster import Cluster - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated.cluster import Cluster + from google.cloud.bigtable.deprecated import enums client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -162,7 +162,7 @@ def test_cluster_from_pb_success(): def test_cluster_from_pb_w_bad_cluster_name(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.cluster import Cluster + from google.cloud.bigtable.deprecated.cluster import Cluster bad_cluster_name = "BAD_NAME" @@ -174,7 +174,7 @@ def test_cluster_from_pb_w_bad_cluster_name(): def test_cluster_from_pb_w_instance_id_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.cluster import Cluster + from google.cloud.bigtable.deprecated.cluster import Cluster ALT_INSTANCE_ID = "ALT_INSTANCE_ID" client = _Client(PROJECT) @@ -189,7 +189,7 @@ def test_cluster_from_pb_w_instance_id_mistmatch(): def test_cluster_from_pb_w_project_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.cluster import Cluster + from google.cloud.bigtable.deprecated.cluster import Cluster ALT_PROJECT = "ALT_PROJECT" client = _Client(project=ALT_PROJECT) @@ -204,8 +204,8 @@ def test_cluster_from_pb_w_project_mistmatch(): def test_cluster_from_pb_w_autoscaling(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.cluster import Cluster - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated.cluster import Cluster + from google.cloud.bigtable.deprecated import enums client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -292,8 +292,8 @@ def _make_instance_admin_client(): def test_cluster_reload(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.enums import StorageType - from google.cloud.bigtable.enums import Cluster + from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.deprecated.enums import Cluster credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -349,7 +349,7 @@ def test_cluster_reload(): def test_cluster_exists_hit(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -371,7 +371,7 @@ def test_cluster_exists_hit(): def test_cluster_exists_miss(): - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance from google.api_core import exceptions credentials = _make_credentials() @@ -390,7 +390,7 @@ def test_cluster_exists_miss(): def test_cluster_exists_w_error(): - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance from google.api_core import exceptions credentials = _make_credentials() @@ -416,9 +416,9 @@ def test_cluster_create(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2 - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -471,9 +471,9 @@ def test_cluster_create_w_cmek(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2 - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -531,9 +531,9 @@ def test_cluster_create_w_autoscaling(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2 - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -600,7 +600,7 @@ def test_cluster_update(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -667,7 +667,7 @@ def test_cluster_update_w_autoscaling(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -726,7 +726,7 @@ def test_cluster_update_w_partial_autoscaling_config(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -811,7 +811,7 @@ def test_cluster_update_w_both_manual_and_autoscaling(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -871,8 +871,8 @@ def test_cluster_disable_autoscaling(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.instance import Instance - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.deprecated.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -928,8 +928,8 @@ def test_cluster_disable_autoscaling(): def test_create_cluster_with_both_manual_and_autoscaling(): - from google.cloud.bigtable.instance import Instance - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.deprecated.enums import StorageType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -956,8 +956,8 @@ def test_create_cluster_with_both_manual_and_autoscaling(): def test_create_cluster_with_partial_autoscaling_config(): - from google.cloud.bigtable.instance import Instance - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.deprecated.enums import StorageType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -997,8 +997,8 @@ def test_create_cluster_with_partial_autoscaling_config(): def test_create_cluster_with_no_scaling_config(): - from google.cloud.bigtable.instance import Instance - from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.deprecated.enums import StorageType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) diff --git a/tests/unit/test_column_family.py b/tests/unit/v2_client/test_column_family.py similarity index 87% rename from tests/unit/test_column_family.py rename to tests/unit/v2_client/test_column_family.py index b464024a7..d16d2b20c 100644 --- a/tests/unit/test_column_family.py +++ b/tests/unit/v2_client/test_column_family.py @@ -19,7 +19,7 @@ def _make_max_versions_gc_rule(*args, **kwargs): - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule return MaxVersionsGCRule(*args, **kwargs) @@ -51,7 +51,7 @@ def test_max_versions_gc_rule_to_pb(): def _make_max_age_gc_rule(*args, **kwargs): - from google.cloud.bigtable.column_family import MaxAgeGCRule + from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule return MaxAgeGCRule(*args, **kwargs) @@ -89,7 +89,7 @@ def test_max_age_gc_rule_to_pb(): def _make_gc_rule_union(*args, **kwargs): - from google.cloud.bigtable.column_family import GCRuleUnion + from google.cloud.bigtable.deprecated.column_family import GCRuleUnion return GCRuleUnion(*args, **kwargs) @@ -124,8 +124,8 @@ def test_gc_rule_union___ne__same_value(): def test_gc_rule_union_to_pb(): import datetime from google.protobuf import duration_pb2 - from google.cloud.bigtable.column_family import MaxAgeGCRule - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule max_num_versions = 42 rule1 = MaxVersionsGCRule(max_num_versions) @@ -145,8 +145,8 @@ def test_gc_rule_union_to_pb(): def test_gc_rule_union_to_pb_nested(): import datetime from google.protobuf import duration_pb2 - from google.cloud.bigtable.column_family import MaxAgeGCRule - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule max_num_versions1 = 42 rule1 = MaxVersionsGCRule(max_num_versions1) @@ -171,7 +171,7 @@ def test_gc_rule_union_to_pb_nested(): def _make_gc_rule_intersection(*args, **kwargs): - from google.cloud.bigtable.column_family import GCRuleIntersection + from google.cloud.bigtable.deprecated.column_family import GCRuleIntersection return GCRuleIntersection(*args, **kwargs) @@ -206,8 +206,8 @@ def test_gc_rule_intersection___ne__same_value(): def test_gc_rule_intersection_to_pb(): import datetime from google.protobuf import duration_pb2 - from google.cloud.bigtable.column_family import MaxAgeGCRule - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule max_num_versions = 42 rule1 = MaxVersionsGCRule(max_num_versions) @@ -227,8 +227,8 @@ def test_gc_rule_intersection_to_pb(): def test_gc_rule_intersection_to_pb_nested(): import datetime from google.protobuf import duration_pb2 - from google.cloud.bigtable.column_family import MaxAgeGCRule - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule max_num_versions1 = 42 rule1 = MaxVersionsGCRule(max_num_versions1) @@ -253,13 +253,13 @@ def test_gc_rule_intersection_to_pb_nested(): def _make_column_family(*args, **kwargs): - from google.cloud.bigtable.column_family import ColumnFamily + from google.cloud.bigtable.deprecated.column_family import ColumnFamily return ColumnFamily(*args, **kwargs) def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client return Client(*args, **kwargs) @@ -323,7 +323,7 @@ def test_column_family_to_pb_no_rules(): def test_column_family_to_pb_with_rule(): - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule gc_rule = MaxVersionsGCRule(1) column_family = _make_column_family("column_family_id", None, gc_rule=gc_rule) @@ -336,7 +336,7 @@ def _create_test_helper(gc_rule=None): from google.cloud.bigtable_admin_v2.types import ( bigtable_table_admin as table_admin_v2_pb2, ) - from tests.unit._testing import _FakeStub + from ._testing import _FakeStub from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) @@ -397,14 +397,14 @@ def test_column_family_create(): def test_column_family_create_with_gc_rule(): - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule gc_rule = MaxVersionsGCRule(1337) _create_test_helper(gc_rule=gc_rule) def _update_test_helper(gc_rule=None): - from tests.unit._testing import _FakeStub + from ._testing import _FakeStub from google.cloud.bigtable_admin_v2.types import ( bigtable_table_admin as table_admin_v2_pb2, ) @@ -467,7 +467,7 @@ def test_column_family_update(): def test_column_family_update_with_gc_rule(): - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule gc_rule = MaxVersionsGCRule(1337) _update_test_helper(gc_rule=gc_rule) @@ -478,7 +478,7 @@ def test_column_family_delete(): from google.cloud.bigtable_admin_v2.types import ( bigtable_table_admin as table_admin_v2_pb2, ) - from tests.unit._testing import _FakeStub + from ._testing import _FakeStub from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) @@ -530,15 +530,15 @@ def test_column_family_delete(): def test__gc_rule_from_pb_empty(): - from google.cloud.bigtable.column_family import _gc_rule_from_pb + from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb gc_rule_pb = _GcRulePB() assert _gc_rule_from_pb(gc_rule_pb) is None def test__gc_rule_from_pb_max_num_versions(): - from google.cloud.bigtable.column_family import _gc_rule_from_pb - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule orig_rule = MaxVersionsGCRule(1) gc_rule_pb = orig_rule.to_pb() @@ -549,8 +549,8 @@ def test__gc_rule_from_pb_max_num_versions(): def test__gc_rule_from_pb_max_age(): import datetime - from google.cloud.bigtable.column_family import _gc_rule_from_pb - from google.cloud.bigtable.column_family import MaxAgeGCRule + from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb + from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule orig_rule = MaxAgeGCRule(datetime.timedelta(seconds=1)) gc_rule_pb = orig_rule.to_pb() @@ -561,10 +561,10 @@ def test__gc_rule_from_pb_max_age(): def test__gc_rule_from_pb_union(): import datetime - from google.cloud.bigtable.column_family import _gc_rule_from_pb - from google.cloud.bigtable.column_family import GCRuleUnion - from google.cloud.bigtable.column_family import MaxAgeGCRule - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb + from google.cloud.bigtable.deprecated.column_family import GCRuleUnion + from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule rule1 = MaxVersionsGCRule(1) rule2 = MaxAgeGCRule(datetime.timedelta(seconds=1)) @@ -577,10 +577,10 @@ def test__gc_rule_from_pb_union(): def test__gc_rule_from_pb_intersection(): import datetime - from google.cloud.bigtable.column_family import _gc_rule_from_pb - from google.cloud.bigtable.column_family import GCRuleIntersection - from google.cloud.bigtable.column_family import MaxAgeGCRule - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb + from google.cloud.bigtable.deprecated.column_family import GCRuleIntersection + from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule rule1 = MaxVersionsGCRule(1) rule2 = MaxAgeGCRule(datetime.timedelta(seconds=1)) @@ -592,7 +592,7 @@ def test__gc_rule_from_pb_intersection(): def test__gc_rule_from_pb_unknown_field_name(): - from google.cloud.bigtable.column_family import _gc_rule_from_pb + from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb class MockProto(object): diff --git a/tests/unit/test_encryption_info.py b/tests/unit/v2_client/test_encryption_info.py similarity index 94% rename from tests/unit/test_encryption_info.py rename to tests/unit/v2_client/test_encryption_info.py index 8b92a83ed..0b6a93e9e 100644 --- a/tests/unit/test_encryption_info.py +++ b/tests/unit/v2_client/test_encryption_info.py @@ -14,7 +14,7 @@ import mock -from google.cloud.bigtable import enums +from google.cloud.bigtable.deprecated import enums EncryptionType = enums.EncryptionInfo.EncryptionType @@ -30,7 +30,7 @@ def _make_status_pb(code=_STATUS_CODE, message=_STATUS_MESSAGE): def _make_status(code=_STATUS_CODE, message=_STATUS_MESSAGE): - from google.cloud.bigtable.error import Status + from google.cloud.bigtable.deprecated.error import Status status_pb = _make_status_pb(code=code, message=message) return Status(status_pb) @@ -54,7 +54,7 @@ def _make_info_pb( def _make_encryption_info(*args, **kwargs): - from google.cloud.bigtable.encryption_info import EncryptionInfo + from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo return EncryptionInfo(*args, **kwargs) @@ -70,7 +70,7 @@ def _make_encryption_info_defaults( def test_encryption_info__from_pb(): - from google.cloud.bigtable.encryption_info import EncryptionInfo + from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo info_pb = _make_info_pb() diff --git a/tests/unit/test_error.py b/tests/unit/v2_client/test_error.py similarity index 97% rename from tests/unit/test_error.py rename to tests/unit/v2_client/test_error.py index 8b148473c..072a3b3c3 100644 --- a/tests/unit/test_error.py +++ b/tests/unit/v2_client/test_error.py @@ -20,7 +20,7 @@ def _make_status_pb(**kwargs): def _make_status(status_pb): - from google.cloud.bigtable.error import Status + from google.cloud.bigtable.deprecated.error import Status return Status(status_pb) diff --git a/tests/unit/test_instance.py b/tests/unit/v2_client/test_instance.py similarity index 95% rename from tests/unit/test_instance.py rename to tests/unit/v2_client/test_instance.py index c577adca5..b43e8bb38 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/v2_client/test_instance.py @@ -17,7 +17,7 @@ import pytest from ._testing import _make_credentials -from google.cloud.bigtable.cluster import Cluster +from google.cloud.bigtable.deprecated.cluster import Cluster PROJECT = "project" INSTANCE_ID = "instance-id" @@ -47,7 +47,7 @@ def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client return Client(*args, **kwargs) @@ -61,7 +61,7 @@ def _make_instance_admin_api(): def _make_instance(*args, **kwargs): - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance return Instance(*args, **kwargs) @@ -79,7 +79,7 @@ def test_instance_constructor_defaults(): def test_instance_constructor_non_default(): - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums instance_type = enums.Instance.Type.DEVELOPMENT state = enums.Instance.State.READY @@ -104,7 +104,7 @@ def test_instance_constructor_non_default(): def test_instance__update_from_pb_success(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums instance_type = data_v2_pb2.Instance.Type.PRODUCTION state = enums.Instance.State.READY @@ -129,7 +129,7 @@ def test_instance__update_from_pb_success(): def test_instance__update_from_pb_success_defaults(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums instance_pb = data_v2_pb2.Instance(display_name=DISPLAY_NAME) @@ -156,8 +156,8 @@ def test_instance__update_from_pb_wo_display_name(): def test_instance_from_pb_success(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable import enums - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable.deprecated.instance import Instance credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -184,7 +184,7 @@ def test_instance_from_pb_success(): def test_instance_from_pb_bad_instance_name(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance instance_name = "INCORRECT_FORMAT" instance_pb = data_v2_pb2.Instance(name=instance_name) @@ -195,7 +195,7 @@ def test_instance_from_pb_bad_instance_name(): def test_instance_from_pb_project_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance ALT_PROJECT = "ALT_PROJECT" credentials = _make_credentials() @@ -304,7 +304,7 @@ def _instance_api_response_for_create(): def test_instance_create(): - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums from google.cloud.bigtable_admin_v2.types import Instance from google.cloud.bigtable_admin_v2.types import Cluster import warnings @@ -353,8 +353,8 @@ def test_instance_create(): def test_instance_create_w_clusters(): - from google.cloud.bigtable import enums - from google.cloud.bigtable.cluster import Cluster + from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable.deprecated.cluster import Cluster from google.cloud.bigtable_admin_v2.types import Cluster as cluster_pb from google.cloud.bigtable_admin_v2.types import Instance as instance_pb @@ -473,7 +473,7 @@ def test_instance_exists_w_error(): def test_instance_reload(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums DISPLAY_NAME = "hey-hi-hello" credentials = _make_credentials() @@ -527,7 +527,7 @@ def _instance_api_response_for_update(): def test_instance_update(): - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums from google.protobuf import field_mask_pb2 from google.cloud.bigtable_admin_v2.types import Instance @@ -603,7 +603,7 @@ def test_instance_delete(): def test_instance_get_iam_policy(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -631,7 +631,7 @@ def test_instance_get_iam_policy(): def test_instance_get_iam_policy_w_requested_policy_version(): from google.iam.v1 import policy_pb2, options_pb2 - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -665,8 +665,8 @@ def test_instance_get_iam_policy_w_requested_policy_version(): def test_instance_set_iam_policy(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -721,7 +721,7 @@ def test_instance_test_iam_permissions(): def test_instance_cluster_factory(): - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums CLUSTER_ID = "{}-cluster".format(INSTANCE_ID) LOCATION_ID = "us-central1-c" @@ -749,8 +749,8 @@ def test_instance_list_clusters(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.instance import Instance - from google.cloud.bigtable.instance import Cluster + from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.deprecated.instance import Cluster credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -788,7 +788,7 @@ def test_instance_list_clusters(): def test_instance_table_factory(): - from google.cloud.bigtable.table import Table + from google.cloud.bigtable.deprecated.table import Table app_profile_id = "appProfileId1262094415" instance = _make_instance(INSTANCE_ID, None) @@ -857,7 +857,7 @@ def test_instance_list_tables_failure_name_bad_before(): def test_instance_app_profile_factory(): - from google.cloud.bigtable.enums import RoutingPolicyType + from google.cloud.bigtable.deprecated.enums import RoutingPolicyType instance = _make_instance(INSTANCE_ID, None) @@ -890,7 +890,7 @@ def test_instance_list_app_profiles(): from google.api_core.page_iterator import Iterator from google.api_core.page_iterator import Page from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.deprecated.app_profile import AppProfile class _Iterator(Iterator): def __init__(self, pages): diff --git a/tests/unit/test_policy.py b/tests/unit/v2_client/test_policy.py similarity index 89% rename from tests/unit/test_policy.py rename to tests/unit/v2_client/test_policy.py index 77674517e..ef3df2d2b 100644 --- a/tests/unit/test_policy.py +++ b/tests/unit/v2_client/test_policy.py @@ -14,7 +14,7 @@ def _make_policy(*args, **kw): - from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.deprecated.policy import Policy return Policy(*args, **kw) @@ -48,7 +48,7 @@ def test_policy_ctor_explicit(): def test_policy_bigtable_admins(): - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE MEMBER = "user:phred@example.com" expected = frozenset([MEMBER]) @@ -58,7 +58,7 @@ def test_policy_bigtable_admins(): def test_policy_bigtable_readers(): - from google.cloud.bigtable.policy import BIGTABLE_READER_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_READER_ROLE MEMBER = "user:phred@example.com" expected = frozenset([MEMBER]) @@ -68,7 +68,7 @@ def test_policy_bigtable_readers(): def test_policy_bigtable_users(): - from google.cloud.bigtable.policy import BIGTABLE_USER_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_USER_ROLE MEMBER = "user:phred@example.com" expected = frozenset([MEMBER]) @@ -78,7 +78,7 @@ def test_policy_bigtable_users(): def test_policy_bigtable_viewers(): - from google.cloud.bigtable.policy import BIGTABLE_VIEWER_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_VIEWER_ROLE MEMBER = "user:phred@example.com" expected = frozenset([MEMBER]) @@ -89,7 +89,7 @@ def test_policy_bigtable_viewers(): def test_policy_from_pb_w_empty(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.deprecated.policy import Policy empty = frozenset() message = policy_pb2.Policy() @@ -106,8 +106,8 @@ def test_policy_from_pb_w_empty(): def test_policy_from_pb_w_non_empty(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE - from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import Policy ETAG = b"ETAG" VERSION = 1 @@ -133,8 +133,8 @@ def test_policy_from_pb_w_condition(): import pytest from google.iam.v1 import policy_pb2 from google.api_core.iam import InvalidOperationException, _DICT_ACCESS_MSG - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE - from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import Policy ETAG = b"ETAG" VERSION = 3 @@ -184,7 +184,7 @@ def test_policy_to_pb_empty(): def test_policy_to_pb_explicit(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE VERSION = 1 ETAG = b"ETAG" @@ -204,7 +204,7 @@ def test_policy_to_pb_explicit(): def test_policy_to_pb_w_condition(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE VERSION = 3 ETAG = b"ETAG" @@ -234,7 +234,7 @@ def test_policy_to_pb_w_condition(): def test_policy_from_api_repr_wo_etag(): - from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.deprecated.policy import Policy VERSION = 1 empty = frozenset() @@ -252,7 +252,7 @@ def test_policy_from_api_repr_wo_etag(): def test_policy_from_api_repr_w_etag(): import base64 - from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.deprecated.policy import Policy ETAG = b"ETAG" empty = frozenset() diff --git a/tests/unit/test_row.py b/tests/unit/v2_client/test_row.py similarity index 95% rename from tests/unit/test_row.py rename to tests/unit/v2_client/test_row.py index 49bbfc45c..4850b18c3 100644 --- a/tests/unit/test_row.py +++ b/tests/unit/v2_client/test_row.py @@ -20,13 +20,13 @@ def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client return Client(*args, **kwargs) def _make_row(*args, **kwargs): - from google.cloud.bigtable.row import Row + from google.cloud.bigtable.deprecated.row import Row return Row(*args, **kwargs) @@ -42,7 +42,7 @@ def test_row_table_getter(): def _make__set_delete_row(*args, **kwargs): - from google.cloud.bigtable.row import _SetDeleteRow + from google.cloud.bigtable.deprecated.row import _SetDeleteRow return _SetDeleteRow(*args, **kwargs) @@ -54,7 +54,7 @@ def test__set_detlete_row__get_mutations_virtual(): def _make_direct_row(*args, **kwargs): - from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.deprecated.row import DirectRow return DirectRow(*args, **kwargs) @@ -193,7 +193,7 @@ def test_direct_row_delete(): def test_direct_row_delete_cell(): - from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.deprecated.row import DirectRow class MockRow(DirectRow): def __init__(self, *args, **kwargs): @@ -237,7 +237,7 @@ def test_direct_row_delete_cells_non_iterable(): def test_direct_row_delete_cells_all_columns(): - from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.deprecated.row import DirectRow row_key = b"row_key" column_family_id = "column_family_id" @@ -293,7 +293,7 @@ def test_direct_row_delete_cells_no_time_range(): def test_direct_row_delete_cells_with_time_range(): import datetime from google.cloud._helpers import _EPOCH - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.deprecated.row_filters import TimestampRange microseconds = 30871000 # Makes sure already milliseconds granularity start = _EPOCH + datetime.timedelta(microseconds=microseconds) @@ -386,7 +386,7 @@ def test_direct_row_commit_with_exception(): def _make_conditional_row(*args, **kwargs): - from google.cloud.bigtable.row import ConditionalRow + from google.cloud.bigtable.deprecated.row import ConditionalRow return ConditionalRow(*args, **kwargs) @@ -417,7 +417,7 @@ def test_conditional_row__get_mutations(): def test_conditional_row_commit(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter from google.cloud.bigtable_v2.services.bigtable import BigtableClient project_id = "project-id" @@ -466,7 +466,7 @@ def test_conditional_row_commit(): def test_conditional_row_commit_too_many_mutations(): from google.cloud._testing import _Monkey - from google.cloud.bigtable import row as MUT + from google.cloud.bigtable.deprecated import row as MUT row_key = b"row_key" table = object() @@ -480,7 +480,7 @@ def test_conditional_row_commit_too_many_mutations(): def test_conditional_row_commit_no_mutations(): - from tests.unit._testing import _FakeStub + from ._testing import _FakeStub project_id = "project-id" row_key = b"row_key" @@ -504,7 +504,7 @@ def test_conditional_row_commit_no_mutations(): def _make_append_row(*args, **kwargs): - from google.cloud.bigtable.row import AppendRow + from google.cloud.bigtable.deprecated.row import AppendRow return AppendRow(*args, **kwargs) @@ -564,7 +564,7 @@ def test_append_row_increment_cell_value(): def test_append_row_commit(): from google.cloud._testing import _Monkey - from google.cloud.bigtable import row as MUT + from google.cloud.bigtable.deprecated import row as MUT from google.cloud.bigtable_v2.services.bigtable import BigtableClient project_id = "project-id" @@ -607,7 +607,7 @@ def mock_parse_rmw_row_response(row_response): def test_append_row_commit_no_rules(): - from tests.unit._testing import _FakeStub + from ._testing import _FakeStub project_id = "project-id" row_key = b"row_key" @@ -630,7 +630,7 @@ def test_append_row_commit_no_rules(): def test_append_row_commit_too_many_mutations(): from google.cloud._testing import _Monkey - from google.cloud.bigtable import row as MUT + from google.cloud.bigtable.deprecated import row as MUT row_key = b"row_key" table = object() @@ -644,7 +644,7 @@ def test_append_row_commit_too_many_mutations(): def test__parse_rmw_row_response(): from google.cloud._helpers import _datetime_from_microseconds - from google.cloud.bigtable.row import _parse_rmw_row_response + from google.cloud.bigtable.deprecated.row import _parse_rmw_row_response col_fam1 = "col-fam-id" col_fam2 = "col-fam-id2" @@ -700,7 +700,7 @@ def test__parse_rmw_row_response(): def test__parse_family_pb(): from google.cloud._helpers import _datetime_from_microseconds - from google.cloud.bigtable.row import _parse_family_pb + from google.cloud.bigtable.deprecated.row import _parse_family_pb col_fam1 = "col-fam-id" col_name1 = b"col-name1" diff --git a/tests/unit/test_row_data.py b/tests/unit/v2_client/test_row_data.py similarity index 94% rename from tests/unit/test_row_data.py rename to tests/unit/v2_client/test_row_data.py index 382a81ef1..ee9b065c8 100644 --- a/tests/unit/test_row_data.py +++ b/tests/unit/v2_client/test_row_data.py @@ -27,7 +27,7 @@ def _make_cell(*args, **kwargs): - from google.cloud.bigtable.row_data import Cell + from google.cloud.bigtable.deprecated.row_data import Cell return Cell(*args, **kwargs) @@ -36,7 +36,7 @@ def _cell_from_pb_test_helper(labels=None): import datetime from google.cloud._helpers import _EPOCH from google.cloud.bigtable_v2.types import data as data_v2_pb2 - from google.cloud.bigtable.row_data import Cell + from google.cloud.bigtable.deprecated.row_data import Cell timestamp = _EPOCH + datetime.timedelta(microseconds=TIMESTAMP_MICROS) value = b"value-bytes" @@ -100,7 +100,7 @@ def test_cell___ne__(): def _make_partial_row_data(*args, **kwargs): - from google.cloud.bigtable.row_data import PartialRowData + from google.cloud.bigtable.deprecated.row_data import PartialRowData return PartialRowData(*args, **kwargs) @@ -288,7 +288,7 @@ def trailing_metadata(self): def test__retry_read_rows_exception_miss(): from google.api_core.exceptions import Conflict - from google.cloud.bigtable.row_data import _retry_read_rows_exception + from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception exception = Conflict("testing") assert not _retry_read_rows_exception(exception) @@ -296,7 +296,7 @@ def test__retry_read_rows_exception_miss(): def test__retry_read_rows_exception_service_unavailable(): from google.api_core.exceptions import ServiceUnavailable - from google.cloud.bigtable.row_data import _retry_read_rows_exception + from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception exception = ServiceUnavailable("testing") assert _retry_read_rows_exception(exception) @@ -304,7 +304,7 @@ def test__retry_read_rows_exception_service_unavailable(): def test__retry_read_rows_exception_deadline_exceeded(): from google.api_core.exceptions import DeadlineExceeded - from google.cloud.bigtable.row_data import _retry_read_rows_exception + from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception exception = DeadlineExceeded("testing") assert _retry_read_rows_exception(exception) @@ -312,7 +312,7 @@ def test__retry_read_rows_exception_deadline_exceeded(): def test__retry_read_rows_exception_internal_server_not_retriable(): from google.api_core.exceptions import InternalServerError - from google.cloud.bigtable.row_data import ( + from google.cloud.bigtable.deprecated.row_data import ( _retry_read_rows_exception, RETRYABLE_INTERNAL_ERROR_MESSAGES, ) @@ -325,7 +325,7 @@ def test__retry_read_rows_exception_internal_server_not_retriable(): def test__retry_read_rows_exception_internal_server_retriable(): from google.api_core.exceptions import InternalServerError - from google.cloud.bigtable.row_data import ( + from google.cloud.bigtable.deprecated.row_data import ( _retry_read_rows_exception, RETRYABLE_INTERNAL_ERROR_MESSAGES, ) @@ -337,7 +337,7 @@ def test__retry_read_rows_exception_internal_server_retriable(): def test__retry_read_rows_exception_miss_wrapped_in_grpc(): from google.api_core.exceptions import Conflict - from google.cloud.bigtable.row_data import _retry_read_rows_exception + from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception wrapped = Conflict("testing") exception = _make_grpc_call_error(wrapped) @@ -346,7 +346,7 @@ def test__retry_read_rows_exception_miss_wrapped_in_grpc(): def test__retry_read_rows_exception_service_unavailable_wrapped_in_grpc(): from google.api_core.exceptions import ServiceUnavailable - from google.cloud.bigtable.row_data import _retry_read_rows_exception + from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception wrapped = ServiceUnavailable("testing") exception = _make_grpc_call_error(wrapped) @@ -355,7 +355,7 @@ def test__retry_read_rows_exception_service_unavailable_wrapped_in_grpc(): def test__retry_read_rows_exception_deadline_exceeded_wrapped_in_grpc(): from google.api_core.exceptions import DeadlineExceeded - from google.cloud.bigtable.row_data import _retry_read_rows_exception + from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception wrapped = DeadlineExceeded("testing") exception = _make_grpc_call_error(wrapped) @@ -363,7 +363,7 @@ def test__retry_read_rows_exception_deadline_exceeded_wrapped_in_grpc(): def _make_partial_rows_data(*args, **kwargs): - from google.cloud.bigtable.row_data import PartialRowsData + from google.cloud.bigtable.deprecated.row_data import PartialRowsData return PartialRowsData(*args, **kwargs) @@ -373,13 +373,13 @@ def _partial_rows_data_consume_all(yrd): def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client return Client(*args, **kwargs) def test_partial_rows_data_constructor(): - from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS + from google.cloud.bigtable.deprecated.row_data import DEFAULT_RETRY_READ_ROWS client = _Client() client._data_stub = mock.MagicMock() @@ -436,7 +436,7 @@ def fake_read(*args, **kwargs): def test_partial_rows_data_constructor_with_retry(): - from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS + from google.cloud.bigtable.deprecated.row_data import DEFAULT_RETRY_READ_ROWS client = _Client() client._data_stub = mock.MagicMock() @@ -644,7 +644,7 @@ def test_partial_rows_data_valid_last_scanned_row_key_on_start(): def test_partial_rows_data_invalid_empty_chunk(): - from google.cloud.bigtable.row_data import InvalidChunk + from google.cloud.bigtable.deprecated.row_data import InvalidChunk from google.cloud.bigtable_v2.services.bigtable import BigtableClient client = _Client() @@ -755,14 +755,14 @@ def test_partial_rows_data_yield_retry_rows_data(): def _make_read_rows_request_manager(*args, **kwargs): - from google.cloud.bigtable.row_data import _ReadRowsRequestManager + from google.cloud.bigtable.deprecated.row_data import _ReadRowsRequestManager return _ReadRowsRequestManager(*args, **kwargs) @pytest.fixture(scope="session") def rrrm_data(): - from google.cloud.bigtable import row_set + from google.cloud.bigtable.deprecated import row_set row_range1 = row_set.RowRange(b"row_key21", b"row_key29") row_range2 = row_set.RowRange(b"row_key31", b"row_key39") @@ -851,7 +851,7 @@ def test_RRRM__filter_row_ranges_all_ranges_already_read(rrrm_data): def test_RRRM__filter_row_ranges_all_ranges_already_read_open_closed(): - from google.cloud.bigtable import row_set + from google.cloud.bigtable.deprecated import row_set last_scanned_key = b"row_key54" @@ -895,7 +895,7 @@ def test_RRRM__filter_row_ranges_some_ranges_already_read(rrrm_data): def test_RRRM_build_updated_request(rrrm_data): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter from google.cloud.bigtable_v2 import types row_range1 = rrrm_data["row_range1"] @@ -944,7 +944,7 @@ def test_RRRM_build_updated_request_full_table(): def test_RRRM_build_updated_request_no_start_key(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter from google.cloud.bigtable_v2 import types row_filter = RowSampleFilter(0.33) @@ -972,7 +972,7 @@ def test_RRRM_build_updated_request_no_start_key(): def test_RRRM_build_updated_request_no_end_key(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter from google.cloud.bigtable_v2 import types row_filter = RowSampleFilter(0.33) @@ -998,7 +998,7 @@ def test_RRRM_build_updated_request_no_end_key(): def test_RRRM_build_updated_request_rows(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter row_filter = RowSampleFilter(0.33) last_scanned_key = b"row_key4" @@ -1046,7 +1046,7 @@ def test_RRRM__key_already_read(): def test_RRRM__rows_limit_reached(): - from google.cloud.bigtable.row_data import InvalidRetryRequest + from google.cloud.bigtable.deprecated.row_data import InvalidRetryRequest last_scanned_key = b"row_key14" request = _ReadRowsRequestPB(table_name=TABLE_NAME) @@ -1059,7 +1059,7 @@ def test_RRRM__rows_limit_reached(): def test_RRRM_build_updated_request_last_row_read_raises_invalid_retry_request(): - from google.cloud.bigtable.row_data import InvalidRetryRequest + from google.cloud.bigtable.deprecated.row_data import InvalidRetryRequest last_scanned_key = b"row_key4" request = _ReadRowsRequestPB(table_name=TABLE_NAME) @@ -1073,8 +1073,8 @@ def test_RRRM_build_updated_request_last_row_read_raises_invalid_retry_request() def test_RRRM_build_updated_request_row_ranges_read_raises_invalid_retry_request(): - from google.cloud.bigtable.row_data import InvalidRetryRequest - from google.cloud.bigtable import row_set + from google.cloud.bigtable.deprecated.row_data import InvalidRetryRequest + from google.cloud.bigtable.deprecated import row_set row_range1 = row_set.RowRange(b"row_key21", b"row_key29") @@ -1095,7 +1095,7 @@ def test_RRRM_build_updated_request_row_ranges_read_raises_invalid_retry_request def test_RRRM_build_updated_request_row_ranges_valid(): - from google.cloud.bigtable import row_set + from google.cloud.bigtable.deprecated import row_set row_range1 = row_set.RowRange(b"row_key21", b"row_key29") @@ -1179,7 +1179,7 @@ def _ReadRowsResponseCellChunkPB(*args, **kw): def _make_cell_pb(value): - from google.cloud.bigtable import row_data + from google.cloud.bigtable.deprecated import row_data return row_data.Cell(value, TIMESTAMP_MICROS) diff --git a/tests/unit/test_row_filters.py b/tests/unit/v2_client/test_row_filters.py similarity index 77% rename from tests/unit/test_row_filters.py rename to tests/unit/v2_client/test_row_filters.py index b312cb942..dfb16ba16 100644 --- a/tests/unit/test_row_filters.py +++ b/tests/unit/v2_client/test_row_filters.py @@ -17,7 +17,7 @@ def test_bool_filter_constructor(): - from google.cloud.bigtable.row_filters import _BoolFilter + from google.cloud.bigtable.deprecated.row_filters import _BoolFilter flag = object() row_filter = _BoolFilter(flag) @@ -25,7 +25,7 @@ def test_bool_filter_constructor(): def test_bool_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import _BoolFilter + from google.cloud.bigtable.deprecated.row_filters import _BoolFilter flag = object() row_filter1 = _BoolFilter(flag) @@ -34,7 +34,7 @@ def test_bool_filter___eq__type_differ(): def test_bool_filter___eq__same_value(): - from google.cloud.bigtable.row_filters import _BoolFilter + from google.cloud.bigtable.deprecated.row_filters import _BoolFilter flag = object() row_filter1 = _BoolFilter(flag) @@ -43,7 +43,7 @@ def test_bool_filter___eq__same_value(): def test_bool_filter___ne__same_value(): - from google.cloud.bigtable.row_filters import _BoolFilter + from google.cloud.bigtable.deprecated.row_filters import _BoolFilter flag = object() row_filter1 = _BoolFilter(flag) @@ -52,7 +52,7 @@ def test_bool_filter___ne__same_value(): def test_sink_filter_to_pb(): - from google.cloud.bigtable.row_filters import SinkFilter + from google.cloud.bigtable.deprecated.row_filters import SinkFilter flag = True row_filter = SinkFilter(flag) @@ -62,7 +62,7 @@ def test_sink_filter_to_pb(): def test_pass_all_filter_to_pb(): - from google.cloud.bigtable.row_filters import PassAllFilter + from google.cloud.bigtable.deprecated.row_filters import PassAllFilter flag = True row_filter = PassAllFilter(flag) @@ -72,7 +72,7 @@ def test_pass_all_filter_to_pb(): def test_block_all_filter_to_pb(): - from google.cloud.bigtable.row_filters import BlockAllFilter + from google.cloud.bigtable.deprecated.row_filters import BlockAllFilter flag = True row_filter = BlockAllFilter(flag) @@ -82,7 +82,7 @@ def test_block_all_filter_to_pb(): def test_regex_filterconstructor(): - from google.cloud.bigtable.row_filters import _RegexFilter + from google.cloud.bigtable.deprecated.row_filters import _RegexFilter regex = b"abc" row_filter = _RegexFilter(regex) @@ -90,7 +90,7 @@ def test_regex_filterconstructor(): def test_regex_filterconstructor_non_bytes(): - from google.cloud.bigtable.row_filters import _RegexFilter + from google.cloud.bigtable.deprecated.row_filters import _RegexFilter regex = "abc" row_filter = _RegexFilter(regex) @@ -98,7 +98,7 @@ def test_regex_filterconstructor_non_bytes(): def test_regex_filter__eq__type_differ(): - from google.cloud.bigtable.row_filters import _RegexFilter + from google.cloud.bigtable.deprecated.row_filters import _RegexFilter regex = b"def-rgx" row_filter1 = _RegexFilter(regex) @@ -107,7 +107,7 @@ def test_regex_filter__eq__type_differ(): def test_regex_filter__eq__same_value(): - from google.cloud.bigtable.row_filters import _RegexFilter + from google.cloud.bigtable.deprecated.row_filters import _RegexFilter regex = b"trex-regex" row_filter1 = _RegexFilter(regex) @@ -116,7 +116,7 @@ def test_regex_filter__eq__same_value(): def test_regex_filter__ne__same_value(): - from google.cloud.bigtable.row_filters import _RegexFilter + from google.cloud.bigtable.deprecated.row_filters import _RegexFilter regex = b"abc" row_filter1 = _RegexFilter(regex) @@ -125,7 +125,7 @@ def test_regex_filter__ne__same_value(): def test_row_key_regex_filter_to_pb(): - from google.cloud.bigtable.row_filters import RowKeyRegexFilter + from google.cloud.bigtable.deprecated.row_filters import RowKeyRegexFilter regex = b"row-key-regex" row_filter = RowKeyRegexFilter(regex) @@ -135,7 +135,7 @@ def test_row_key_regex_filter_to_pb(): def test_row_sample_filter_constructor(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter sample = object() row_filter = RowSampleFilter(sample) @@ -143,7 +143,7 @@ def test_row_sample_filter_constructor(): def test_row_sample_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter sample = object() row_filter1 = RowSampleFilter(sample) @@ -152,7 +152,7 @@ def test_row_sample_filter___eq__type_differ(): def test_row_sample_filter___eq__same_value(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter sample = object() row_filter1 = RowSampleFilter(sample) @@ -161,7 +161,7 @@ def test_row_sample_filter___eq__same_value(): def test_row_sample_filter___ne__(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter sample = object() other_sample = object() @@ -171,7 +171,7 @@ def test_row_sample_filter___ne__(): def test_row_sample_filter_to_pb(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter sample = 0.25 row_filter = RowSampleFilter(sample) @@ -181,7 +181,7 @@ def test_row_sample_filter_to_pb(): def test_family_name_regex_filter_to_pb(): - from google.cloud.bigtable.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.deprecated.row_filters import FamilyNameRegexFilter regex = "family-regex" row_filter = FamilyNameRegexFilter(regex) @@ -191,7 +191,7 @@ def test_family_name_regex_filter_to_pb(): def test_column_qualifier_regext_filter_to_pb(): - from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnQualifierRegexFilter regex = b"column-regex" row_filter = ColumnQualifierRegexFilter(regex) @@ -201,7 +201,7 @@ def test_column_qualifier_regext_filter_to_pb(): def test_timestamp_range_constructor(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.deprecated.row_filters import TimestampRange start = object() end = object() @@ -211,7 +211,7 @@ def test_timestamp_range_constructor(): def test_timestamp_range___eq__(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.deprecated.row_filters import TimestampRange start = object() end = object() @@ -221,7 +221,7 @@ def test_timestamp_range___eq__(): def test_timestamp_range___eq__type_differ(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.deprecated.row_filters import TimestampRange start = object() end = object() @@ -231,7 +231,7 @@ def test_timestamp_range___eq__type_differ(): def test_timestamp_range___ne__same_value(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.deprecated.row_filters import TimestampRange start = object() end = object() @@ -243,7 +243,7 @@ def test_timestamp_range___ne__same_value(): def _timestamp_range_to_pb_helper(pb_kwargs, start=None, end=None): import datetime from google.cloud._helpers import _EPOCH - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.deprecated.row_filters import TimestampRange if start is not None: start = _EPOCH + datetime.timedelta(microseconds=start) @@ -291,7 +291,7 @@ def test_timestamp_range_to_pb_end_only(): def test_timestamp_range_filter_constructor(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter range_ = object() row_filter = TimestampRangeFilter(range_) @@ -299,7 +299,7 @@ def test_timestamp_range_filter_constructor(): def test_timestamp_range_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter range_ = object() row_filter1 = TimestampRangeFilter(range_) @@ -308,7 +308,7 @@ def test_timestamp_range_filter___eq__type_differ(): def test_timestamp_range_filter___eq__same_value(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter range_ = object() row_filter1 = TimestampRangeFilter(range_) @@ -317,7 +317,7 @@ def test_timestamp_range_filter___eq__same_value(): def test_timestamp_range_filter___ne__(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter range_ = object() other_range_ = object() @@ -327,8 +327,8 @@ def test_timestamp_range_filter___ne__(): def test_timestamp_range_filter_to_pb(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter + from google.cloud.bigtable.deprecated.row_filters import TimestampRange range_ = TimestampRange() row_filter = TimestampRangeFilter(range_) @@ -338,7 +338,7 @@ def test_timestamp_range_filter_to_pb(): def test_column_range_filter_constructor_defaults(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = object() row_filter = ColumnRangeFilter(column_family_id) @@ -350,7 +350,7 @@ def test_column_range_filter_constructor_defaults(): def test_column_range_filter_constructor_explicit(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = object() start_column = object() @@ -372,7 +372,7 @@ def test_column_range_filter_constructor_explicit(): def test_column_range_filter_constructor_bad_start(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = object() with pytest.raises(ValueError): @@ -380,7 +380,7 @@ def test_column_range_filter_constructor_bad_start(): def test_column_range_filter_constructor_bad_end(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = object() with pytest.raises(ValueError): @@ -388,7 +388,7 @@ def test_column_range_filter_constructor_bad_end(): def test_column_range_filter___eq__(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = object() start_column = object() @@ -413,7 +413,7 @@ def test_column_range_filter___eq__(): def test_column_range_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = object() row_filter1 = ColumnRangeFilter(column_family_id) @@ -422,7 +422,7 @@ def test_column_range_filter___eq__type_differ(): def test_column_range_filter___ne__(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = object() other_column_family_id = object() @@ -448,7 +448,7 @@ def test_column_range_filter___ne__(): def test_column_range_filter_to_pb(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = "column-family-id" row_filter = ColumnRangeFilter(column_family_id) @@ -458,7 +458,7 @@ def test_column_range_filter_to_pb(): def test_column_range_filter_to_pb_inclusive_start(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = "column-family-id" column = b"column" @@ -471,7 +471,7 @@ def test_column_range_filter_to_pb_inclusive_start(): def test_column_range_filter_to_pb_exclusive_start(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = "column-family-id" column = b"column" @@ -486,7 +486,7 @@ def test_column_range_filter_to_pb_exclusive_start(): def test_column_range_filter_to_pb_inclusive_end(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = "column-family-id" column = b"column" @@ -499,7 +499,7 @@ def test_column_range_filter_to_pb_inclusive_end(): def test_column_range_filter_to_pb_exclusive_end(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter column_family_id = "column-family-id" column = b"column" @@ -514,7 +514,7 @@ def test_column_range_filter_to_pb_exclusive_end(): def test_value_regex_filter_to_pb_w_bytes(): - from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRegexFilter value = regex = b"value-regex" row_filter = ValueRegexFilter(value) @@ -524,7 +524,7 @@ def test_value_regex_filter_to_pb_w_bytes(): def test_value_regex_filter_to_pb_w_str(): - from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRegexFilter value = "value-regex" regex = value.encode("ascii") @@ -535,7 +535,7 @@ def test_value_regex_filter_to_pb_w_str(): def test_exact_value_filter_to_pb_w_bytes(): - from google.cloud.bigtable.row_filters import ExactValueFilter + from google.cloud.bigtable.deprecated.row_filters import ExactValueFilter value = regex = b"value-regex" row_filter = ExactValueFilter(value) @@ -545,7 +545,7 @@ def test_exact_value_filter_to_pb_w_bytes(): def test_exact_value_filter_to_pb_w_str(): - from google.cloud.bigtable.row_filters import ExactValueFilter + from google.cloud.bigtable.deprecated.row_filters import ExactValueFilter value = "value-regex" regex = value.encode("ascii") @@ -557,7 +557,7 @@ def test_exact_value_filter_to_pb_w_str(): def test_exact_value_filter_to_pb_w_int(): import struct - from google.cloud.bigtable.row_filters import ExactValueFilter + from google.cloud.bigtable.deprecated.row_filters import ExactValueFilter value = 1 regex = struct.Struct(">q").pack(value) @@ -568,7 +568,7 @@ def test_exact_value_filter_to_pb_w_int(): def test_value_range_filter_constructor_defaults(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter row_filter = ValueRangeFilter() @@ -579,7 +579,7 @@ def test_value_range_filter_constructor_defaults(): def test_value_range_filter_constructor_explicit(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter start_value = object() end_value = object() @@ -600,7 +600,7 @@ def test_value_range_filter_constructor_explicit(): def test_value_range_filter_constructor_w_int_values(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter import struct start_value = 1 @@ -618,21 +618,21 @@ def test_value_range_filter_constructor_w_int_values(): def test_value_range_filter_constructor_bad_start(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter with pytest.raises(ValueError): ValueRangeFilter(inclusive_start=True) def test_value_range_filter_constructor_bad_end(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter with pytest.raises(ValueError): ValueRangeFilter(inclusive_end=True) def test_value_range_filter___eq__(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter start_value = object() end_value = object() @@ -654,7 +654,7 @@ def test_value_range_filter___eq__(): def test_value_range_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter row_filter1 = ValueRangeFilter() row_filter2 = object() @@ -662,7 +662,7 @@ def test_value_range_filter___eq__type_differ(): def test_value_range_filter___ne__(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter start_value = object() other_start_value = object() @@ -685,7 +685,7 @@ def test_value_range_filter___ne__(): def test_value_range_filter_to_pb(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter row_filter = ValueRangeFilter() expected_pb = _RowFilterPB(value_range_filter=_ValueRangePB()) @@ -693,7 +693,7 @@ def test_value_range_filter_to_pb(): def test_value_range_filter_to_pb_inclusive_start(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(start_value=value) @@ -703,7 +703,7 @@ def test_value_range_filter_to_pb_inclusive_start(): def test_value_range_filter_to_pb_exclusive_start(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(start_value=value, inclusive_start=False) @@ -713,7 +713,7 @@ def test_value_range_filter_to_pb_exclusive_start(): def test_value_range_filter_to_pb_inclusive_end(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(end_value=value) @@ -723,7 +723,7 @@ def test_value_range_filter_to_pb_inclusive_end(): def test_value_range_filter_to_pb_exclusive_end(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(end_value=value, inclusive_end=False) @@ -733,7 +733,7 @@ def test_value_range_filter_to_pb_exclusive_end(): def test_cell_count_constructor(): - from google.cloud.bigtable.row_filters import _CellCountFilter + from google.cloud.bigtable.deprecated.row_filters import _CellCountFilter num_cells = object() row_filter = _CellCountFilter(num_cells) @@ -741,7 +741,7 @@ def test_cell_count_constructor(): def test_cell_count___eq__type_differ(): - from google.cloud.bigtable.row_filters import _CellCountFilter + from google.cloud.bigtable.deprecated.row_filters import _CellCountFilter num_cells = object() row_filter1 = _CellCountFilter(num_cells) @@ -750,7 +750,7 @@ def test_cell_count___eq__type_differ(): def test_cell_count___eq__same_value(): - from google.cloud.bigtable.row_filters import _CellCountFilter + from google.cloud.bigtable.deprecated.row_filters import _CellCountFilter num_cells = object() row_filter1 = _CellCountFilter(num_cells) @@ -759,7 +759,7 @@ def test_cell_count___eq__same_value(): def test_cell_count___ne__same_value(): - from google.cloud.bigtable.row_filters import _CellCountFilter + from google.cloud.bigtable.deprecated.row_filters import _CellCountFilter num_cells = object() row_filter1 = _CellCountFilter(num_cells) @@ -768,7 +768,7 @@ def test_cell_count___ne__same_value(): def test_cells_row_offset_filter_to_pb(): - from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.deprecated.row_filters import CellsRowOffsetFilter num_cells = 76 row_filter = CellsRowOffsetFilter(num_cells) @@ -778,7 +778,7 @@ def test_cells_row_offset_filter_to_pb(): def test_cells_row_limit_filter_to_pb(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.deprecated.row_filters import CellsRowLimitFilter num_cells = 189 row_filter = CellsRowLimitFilter(num_cells) @@ -788,7 +788,7 @@ def test_cells_row_limit_filter_to_pb(): def test_cells_column_limit_filter_to_pb(): - from google.cloud.bigtable.row_filters import CellsColumnLimitFilter + from google.cloud.bigtable.deprecated.row_filters import CellsColumnLimitFilter num_cells = 10 row_filter = CellsColumnLimitFilter(num_cells) @@ -798,7 +798,7 @@ def test_cells_column_limit_filter_to_pb(): def test_strip_value_transformer_filter_to_pb(): - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter flag = True row_filter = StripValueTransformerFilter(flag) @@ -808,7 +808,7 @@ def test_strip_value_transformer_filter_to_pb(): def test_apply_label_filter_constructor(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter label = object() row_filter = ApplyLabelFilter(label) @@ -816,7 +816,7 @@ def test_apply_label_filter_constructor(): def test_apply_label_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter label = object() row_filter1 = ApplyLabelFilter(label) @@ -825,7 +825,7 @@ def test_apply_label_filter___eq__type_differ(): def test_apply_label_filter___eq__same_value(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter label = object() row_filter1 = ApplyLabelFilter(label) @@ -834,7 +834,7 @@ def test_apply_label_filter___eq__same_value(): def test_apply_label_filter___ne__(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter label = object() other_label = object() @@ -844,7 +844,7 @@ def test_apply_label_filter___ne__(): def test_apply_label_filter_to_pb(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter label = "label" row_filter = ApplyLabelFilter(label) @@ -854,14 +854,14 @@ def test_apply_label_filter_to_pb(): def test_filter_combination_constructor_defaults(): - from google.cloud.bigtable.row_filters import _FilterCombination + from google.cloud.bigtable.deprecated.row_filters import _FilterCombination row_filter = _FilterCombination() assert row_filter.filters == [] def test_filter_combination_constructor_explicit(): - from google.cloud.bigtable.row_filters import _FilterCombination + from google.cloud.bigtable.deprecated.row_filters import _FilterCombination filters = object() row_filter = _FilterCombination(filters=filters) @@ -869,7 +869,7 @@ def test_filter_combination_constructor_explicit(): def test_filter_combination___eq__(): - from google.cloud.bigtable.row_filters import _FilterCombination + from google.cloud.bigtable.deprecated.row_filters import _FilterCombination filters = object() row_filter1 = _FilterCombination(filters=filters) @@ -878,7 +878,7 @@ def test_filter_combination___eq__(): def test_filter_combination___eq__type_differ(): - from google.cloud.bigtable.row_filters import _FilterCombination + from google.cloud.bigtable.deprecated.row_filters import _FilterCombination filters = object() row_filter1 = _FilterCombination(filters=filters) @@ -887,7 +887,7 @@ def test_filter_combination___eq__type_differ(): def test_filter_combination___ne__(): - from google.cloud.bigtable.row_filters import _FilterCombination + from google.cloud.bigtable.deprecated.row_filters import _FilterCombination filters = object() other_filters = object() @@ -897,9 +897,9 @@ def test_filter_combination___ne__(): def test_row_filter_chain_to_pb(): - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.deprecated.row_filters import RowFilterChain + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() @@ -917,10 +917,10 @@ def test_row_filter_chain_to_pb(): def test_row_filter_chain_to_pb_nested(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.deprecated.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.deprecated.row_filters import RowFilterChain + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -941,9 +941,9 @@ def test_row_filter_chain_to_pb_nested(): def test_row_filter_union_to_pb(): - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.deprecated.row_filters import RowFilterUnion + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() @@ -961,10 +961,10 @@ def test_row_filter_union_to_pb(): def test_row_filter_union_to_pb_nested(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.deprecated.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.deprecated.row_filters import RowFilterUnion + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -985,7 +985,7 @@ def test_row_filter_union_to_pb_nested(): def test_conditional_row_filter_constructor(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter base_filter = object() true_filter = object() @@ -999,7 +999,7 @@ def test_conditional_row_filter_constructor(): def test_conditional_row_filter___eq__(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter base_filter = object() true_filter = object() @@ -1014,7 +1014,7 @@ def test_conditional_row_filter___eq__(): def test_conditional_row_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter base_filter = object() true_filter = object() @@ -1027,7 +1027,7 @@ def test_conditional_row_filter___eq__type_differ(): def test_conditional_row_filter___ne__(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter base_filter = object() other_base_filter = object() @@ -1043,10 +1043,10 @@ def test_conditional_row_filter___ne__(): def test_conditional_row_filter_to_pb(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import CellsRowOffsetFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter + from google.cloud.bigtable.deprecated.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() @@ -1073,9 +1073,9 @@ def test_conditional_row_filter_to_pb(): def test_conditional_row_filter_to_pb_true_only(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() @@ -1095,9 +1095,9 @@ def test_conditional_row_filter_to_pb_true_only(): def test_conditional_row_filter_to_pb_false_only(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() diff --git a/tests/unit/test_row_merger.py b/tests/unit/v2_client/test_row_merger.py similarity index 97% rename from tests/unit/test_row_merger.py rename to tests/unit/v2_client/test_row_merger.py index 483c04536..26cedb34d 100644 --- a/tests/unit/test_row_merger.py +++ b/tests/unit/v2_client/test_row_merger.py @@ -5,9 +5,13 @@ import proto import pytest -from google.cloud.bigtable.row_data import PartialRowsData, PartialRowData, InvalidChunk +from google.cloud.bigtable.deprecated.row_data import ( + PartialRowsData, + PartialRowData, + InvalidChunk, +) from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse -from google.cloud.bigtable.row_merger import _RowMerger +from google.cloud.bigtable.deprecated.row_merger import _RowMerger # TODO: autogenerate protos from diff --git a/tests/unit/test_row_set.py b/tests/unit/v2_client/test_row_set.py similarity index 79% rename from tests/unit/test_row_set.py rename to tests/unit/v2_client/test_row_set.py index 1a33be720..ce0e9bfea 100644 --- a/tests/unit/test_row_set.py +++ b/tests/unit/v2_client/test_row_set.py @@ -14,7 +14,7 @@ def test_row_set_constructor(): - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowSet row_set = RowSet() assert [] == row_set.row_keys @@ -22,8 +22,8 @@ def test_row_set_constructor(): def test_row_set__eq__(): - from google.cloud.bigtable.row_set import RowRange - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowSet row_key1 = b"row_key1" row_key2 = b"row_key1" @@ -42,7 +42,7 @@ def test_row_set__eq__(): def test_row_set__eq__type_differ(): - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowSet row_set1 = RowSet() row_set2 = object() @@ -50,7 +50,7 @@ def test_row_set__eq__type_differ(): def test_row_set__eq__len_row_keys_differ(): - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowSet row_key1 = b"row_key1" row_key2 = b"row_key1" @@ -66,8 +66,8 @@ def test_row_set__eq__len_row_keys_differ(): def test_row_set__eq__len_row_ranges_differ(): - from google.cloud.bigtable.row_set import RowRange - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowSet row_range1 = RowRange(b"row_key4", b"row_key9") row_range2 = RowRange(b"row_key4", b"row_key9") @@ -83,7 +83,7 @@ def test_row_set__eq__len_row_ranges_differ(): def test_row_set__eq__row_keys_differ(): - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowSet row_set1 = RowSet() row_set2 = RowSet() @@ -99,8 +99,8 @@ def test_row_set__eq__row_keys_differ(): def test_row_set__eq__row_ranges_differ(): - from google.cloud.bigtable.row_set import RowRange - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowSet row_range1 = RowRange(b"row_key4", b"row_key9") row_range2 = RowRange(b"row_key14", b"row_key19") @@ -119,8 +119,8 @@ def test_row_set__eq__row_ranges_differ(): def test_row_set__ne__(): - from google.cloud.bigtable.row_set import RowRange - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowSet row_key1 = b"row_key1" row_key2 = b"row_key1" @@ -139,8 +139,8 @@ def test_row_set__ne__(): def test_row_set__ne__same_value(): - from google.cloud.bigtable.row_set import RowRange - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowSet row_key1 = b"row_key1" row_key2 = b"row_key1" @@ -159,7 +159,7 @@ def test_row_set__ne__same_value(): def test_row_set_add_row_key(): - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowSet row_set = RowSet() row_set.add_row_key("row_key1") @@ -168,8 +168,8 @@ def test_row_set_add_row_key(): def test_row_set_add_row_range(): - from google.cloud.bigtable.row_set import RowRange - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowSet row_set = RowSet() row_range1 = RowRange(b"row_key1", b"row_key9") @@ -181,7 +181,7 @@ def test_row_set_add_row_range(): def test_row_set_add_row_range_from_keys(): - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowSet row_set = RowSet() row_set.add_row_range_from_keys( @@ -194,7 +194,7 @@ def test_row_set_add_row_range_from_keys(): def test_row_set_add_row_range_with_prefix(): - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowSet row_set = RowSet() row_set.add_row_range_with_prefix("row") @@ -203,8 +203,8 @@ def test_row_set_add_row_range_with_prefix(): def test_row_set__update_message_request(): from google.cloud._helpers import _to_bytes - from google.cloud.bigtable.row_set import RowRange - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowSet row_set = RowSet() table_name = "table_name" @@ -224,7 +224,7 @@ def test_row_set__update_message_request(): def test_row_range_constructor(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange start_key = "row_key1" end_key = "row_key9" @@ -236,7 +236,7 @@ def test_row_range_constructor(): def test_row_range___hash__set_equality(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange row_range1 = RowRange("row_key1", "row_key9") row_range2 = RowRange("row_key1", "row_key9") @@ -246,7 +246,7 @@ def test_row_range___hash__set_equality(): def test_row_range___hash__not_equals(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange row_range1 = RowRange("row_key1", "row_key9") row_range2 = RowRange("row_key1", "row_key19") @@ -256,7 +256,7 @@ def test_row_range___hash__not_equals(): def test_row_range__eq__(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -266,7 +266,7 @@ def test_row_range__eq__(): def test_row_range___eq__type_differ(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -276,7 +276,7 @@ def test_row_range___eq__type_differ(): def test_row_range__ne__(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -286,7 +286,7 @@ def test_row_range__ne__(): def test_row_range__ne__same_value(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -296,7 +296,7 @@ def test_row_range__ne__same_value(): def test_row_range_get_range_kwargs_closed_open(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -307,7 +307,7 @@ def test_row_range_get_range_kwargs_closed_open(): def test_row_range_get_range_kwargs_open_closed(): - from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.deprecated.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" diff --git a/tests/unit/test_table.py b/tests/unit/v2_client/test_table.py similarity index 91% rename from tests/unit/test_table.py rename to tests/unit/v2_client/test_table.py index e66a8f0f6..ad31e8bc9 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/v2_client/test_table.py @@ -50,11 +50,11 @@ STATUS_INTERNAL = StatusCode.INTERNAL.value[0] -@mock.patch("google.cloud.bigtable.table._MAX_BULK_MUTATIONS", new=3) +@mock.patch("google.cloud.bigtable.deprecated.table._MAX_BULK_MUTATIONS", new=3) def test__compile_mutation_entries_w_too_many_mutations(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import TooManyMutationsError - from google.cloud.bigtable.table import _compile_mutation_entries + from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.deprecated.table import TooManyMutationsError + from google.cloud.bigtable.deprecated.table import _compile_mutation_entries table = mock.Mock(name="table", spec=["name"]) table.name = "table" @@ -72,8 +72,8 @@ def test__compile_mutation_entries_w_too_many_mutations(): def test__compile_mutation_entries_normal(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _compile_mutation_entries + from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.deprecated.table import _compile_mutation_entries from google.cloud.bigtable_v2.types import MutateRowsRequest from google.cloud.bigtable_v2.types import data @@ -109,9 +109,9 @@ def test__compile_mutation_entries_normal(): def test__check_row_table_name_w_wrong_table_name(): - from google.cloud.bigtable.table import _check_row_table_name - from google.cloud.bigtable.table import TableMismatchError - from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.deprecated.table import _check_row_table_name + from google.cloud.bigtable.deprecated.table import TableMismatchError + from google.cloud.bigtable.deprecated.row import DirectRow table = mock.Mock(name="table", spec=["name"]) table.name = "table" @@ -122,8 +122,8 @@ def test__check_row_table_name_w_wrong_table_name(): def test__check_row_table_name_w_right_table_name(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _check_row_table_name + from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.deprecated.table import _check_row_table_name table = mock.Mock(name="table", spec=["name"]) table.name = "table" @@ -133,8 +133,8 @@ def test__check_row_table_name_w_right_table_name(): def test__check_row_type_w_wrong_row_type(): - from google.cloud.bigtable.row import ConditionalRow - from google.cloud.bigtable.table import _check_row_type + from google.cloud.bigtable.deprecated.row import ConditionalRow + from google.cloud.bigtable.deprecated.table import _check_row_type row = ConditionalRow(row_key=b"row_key", table="table", filter_=None) with pytest.raises(TypeError): @@ -142,21 +142,21 @@ def test__check_row_type_w_wrong_row_type(): def test__check_row_type_w_right_row_type(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _check_row_type + from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.deprecated.table import _check_row_type row = DirectRow(row_key=b"row_key", table="table") assert not _check_row_type(row) def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client + from google.cloud.bigtable.deprecated.client import Client return Client(*args, **kwargs) def _make_table(*args, **kwargs): - from google.cloud.bigtable.table import Table + from google.cloud.bigtable.deprecated.table import Table return Table(*args, **kwargs) @@ -219,7 +219,7 @@ def _table_row_methods_helper(): def test_table_row_factory_direct(): - from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.deprecated.row import DirectRow table, row_key = _table_row_methods_helper() with warnings.catch_warnings(record=True) as warned: @@ -234,7 +234,7 @@ def test_table_row_factory_direct(): def test_table_row_factory_conditional(): - from google.cloud.bigtable.row import ConditionalRow + from google.cloud.bigtable.deprecated.row import ConditionalRow table, row_key = _table_row_methods_helper() filter_ = object() @@ -251,7 +251,7 @@ def test_table_row_factory_conditional(): def test_table_row_factory_append(): - from google.cloud.bigtable.row import AppendRow + from google.cloud.bigtable.deprecated.row import AppendRow table, row_key = _table_row_methods_helper() @@ -278,7 +278,7 @@ def test_table_row_factory_failure(): def test_table_direct_row(): - from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.deprecated.row import DirectRow table, row_key = _table_row_methods_helper() row = table.direct_row(row_key) @@ -289,7 +289,7 @@ def test_table_direct_row(): def test_table_conditional_row(): - from google.cloud.bigtable.row import ConditionalRow + from google.cloud.bigtable.deprecated.row import ConditionalRow table, row_key = _table_row_methods_helper() filter_ = object() @@ -301,7 +301,7 @@ def test_table_conditional_row(): def test_table_append_row(): - from google.cloud.bigtable.row import AppendRow + from google.cloud.bigtable.deprecated.row import AppendRow table, row_key = _table_row_methods_helper() row = table.append_row(row_key) @@ -357,7 +357,7 @@ def _create_table_helper(split_keys=[], column_families={}): from google.cloud.bigtable_admin_v2.types import ( bigtable_table_admin as table_admin_messages_v2_pb2, ) - from google.cloud.bigtable.column_family import ColumnFamily + from google.cloud.bigtable.deprecated.column_family import ColumnFamily credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -391,7 +391,7 @@ def test_table_create(): def test_table_create_with_families(): - from google.cloud.bigtable.column_family import MaxVersionsGCRule + from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule families = {"family": MaxVersionsGCRule(5)} _create_table_helper(column_families=families) @@ -404,7 +404,7 @@ def test_table_create_with_split_keys(): def test_table_exists_hit(): from google.cloud.bigtable_admin_v2.types import ListTablesResponse from google.cloud.bigtable_admin_v2.types import Table - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -426,7 +426,7 @@ def test_table_exists_hit(): def test_table_exists_miss(): from google.api_core.exceptions import NotFound - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -447,7 +447,7 @@ def test_table_exists_miss(): def test_table_exists_error(): from google.api_core.exceptions import BadRequest - from google.cloud.bigtable import enums + from google.cloud.bigtable.deprecated import enums credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -512,8 +512,8 @@ def test_table_list_column_families(): def test_table_get_cluster_states(): - from google.cloud.bigtable.enums import Table as enum_table - from google.cloud.bigtable.table import ClusterState + from google.cloud.bigtable.deprecated.enums import Table as enum_table + from google.cloud.bigtable.deprecated.table import ClusterState INITIALIZING = enum_table.ReplicationState.INITIALIZING PLANNED_MAINTENANCE = enum_table.ReplicationState.PLANNED_MAINTENANCE @@ -557,10 +557,10 @@ def test_table_get_cluster_states(): def test_table_get_encryption_info(): from google.rpc.code_pb2 import Code - from google.cloud.bigtable.encryption_info import EncryptionInfo - from google.cloud.bigtable.enums import EncryptionInfo as enum_crypto - from google.cloud.bigtable.enums import Table as enum_table - from google.cloud.bigtable.error import Status + from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo + from google.cloud.bigtable.deprecated.enums import EncryptionInfo as enum_crypto + from google.cloud.bigtable.deprecated.enums import Table as enum_table + from google.cloud.bigtable.deprecated.error import Status ENCRYPTION_TYPE_UNSPECIFIED = enum_crypto.EncryptionType.ENCRYPTION_TYPE_UNSPECIFIED GOOGLE_DEFAULT_ENCRYPTION = enum_crypto.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION @@ -640,9 +640,9 @@ def _make_data_api(): def _table_read_row_helper(chunks, expected_result, app_profile_id=None): from google.cloud._testing import _Monkey - from google.cloud.bigtable import table as MUT - from google.cloud.bigtable.row_set import RowSet - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated import table as MUT + from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -704,8 +704,8 @@ def test_table_read_row_miss_no_chunks_in_response(): def test_table_read_row_complete(): - from google.cloud.bigtable.row_data import Cell - from google.cloud.bigtable.row_data import PartialRowData + from google.cloud.bigtable.deprecated.row_data import Cell + from google.cloud.bigtable.deprecated.row_data import PartialRowData app_profile_id = "app-profile-id" chunk = _ReadRowsResponseCellChunkPB( @@ -768,7 +768,7 @@ def _table_mutate_rows_helper( mutation_timeout=None, app_profile_id=None, retry=None, timeout=None ): from google.rpc.status_pb2 import Status - from google.cloud.bigtable.table import DEFAULT_RETRY + from google.cloud.bigtable.deprecated.table import DEFAULT_RETRY credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -787,7 +787,7 @@ def _table_mutate_rows_helper( response = [Status(code=0), Status(code=1)] instance_mock = mock.Mock(return_value=response) klass_mock = mock.patch( - "google.cloud.bigtable.table._RetryableMutateRowsWorker", + "google.cloud.bigtable.deprecated.table._RetryableMutateRowsWorker", new=mock.MagicMock(return_value=instance_mock), ) @@ -854,9 +854,9 @@ def test_table_mutate_rows_w_mutation_timeout_and_timeout_arg(): def test_table_read_rows(): from google.cloud._testing import _Monkey - from google.cloud.bigtable.row_data import PartialRowsData - from google.cloud.bigtable import table as MUT - from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS + from google.cloud.bigtable.deprecated.row_data import PartialRowsData + from google.cloud.bigtable.deprecated import table as MUT + from google.cloud.bigtable.deprecated.row_data import DEFAULT_RETRY_READ_ROWS credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1017,7 +1017,7 @@ def test_table_read_retry_rows_no_full_table_scan(): def test_table_yield_retry_rows(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1079,9 +1079,9 @@ def test_table_yield_retry_rows(): def test_table_yield_rows_with_row_set(): - from google.cloud.bigtable.row_set import RowSet - from google.cloud.bigtable.row_set import RowRange - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.deprecated.table import _create_row_request credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1174,7 +1174,9 @@ def test_table_truncate(): table = _make_table(TABLE_ID, instance) table_api = client._table_admin_client = _make_table_api() - with mock.patch("google.cloud.bigtable.table.Table.name", new=TABLE_NAME): + with mock.patch( + "google.cloud.bigtable.deprecated.table.Table.name", new=TABLE_NAME + ): result = table.truncate() assert result is None @@ -1255,7 +1257,7 @@ def test_table_mutations_batcher_factory(): def test_table_get_iam_policy(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1286,8 +1288,8 @@ def test_table_get_iam_policy(): def test_table_set_iam_policy(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.policy import Policy - from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1349,7 +1351,7 @@ def test_table_test_iam_permissions(): def test_table_backup_factory_defaults(): - from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.deprecated.backup import Backup instance = _make_table(INSTANCE_ID, None) table = _make_table(TABLE_ID, instance) @@ -1373,8 +1375,8 @@ def test_table_backup_factory_defaults(): def test_table_backup_factory_non_defaults(): import datetime from google.cloud._helpers import UTC - from google.cloud.bigtable.backup import Backup - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.backup import Backup + from google.cloud.bigtable.deprecated.instance import Instance instance = Instance(INSTANCE_ID, None) table = _make_table(TABLE_ID, instance) @@ -1404,7 +1406,7 @@ def _table_list_backups_helper(cluster_id=None, filter_=None, **kwargs): Backup as backup_pb, bigtable_table_admin, ) - from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.deprecated.backup import Backup client = _make_client( project=PROJECT_ID, credentials=_make_credentials(), admin=True @@ -1466,7 +1468,7 @@ def test_table_list_backups_w_options(): def _table_restore_helper(backup_name=None): - from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.deprecated.instance import Instance op_future = object() credentials = _make_credentials() @@ -1502,7 +1504,7 @@ def test_table_restore_table_w_backup_name(): def _make_worker(*args, **kwargs): - from google.cloud.bigtable.table import _RetryableMutateRowsWorker + from google.cloud.bigtable.deprecated.table import _RetryableMutateRowsWorker return _RetryableMutateRowsWorker(*args, **kwargs) @@ -1543,7 +1545,7 @@ def test_rmrw_callable_empty_rows(): def test_rmrw_callable_no_retry_strategy(): - from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.deprecated.row import DirectRow # Setup: # - Mutate 3 rows. @@ -1585,8 +1587,8 @@ def test_rmrw_callable_no_retry_strategy(): def test_rmrw_callable_retry(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import DEFAULT_RETRY + from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.deprecated.table import DEFAULT_RETRY # Setup: # - Mutate 3 rows. @@ -1640,8 +1642,8 @@ def _do_mutate_retryable_rows_helper( mutate_rows_side_effect=None, ): from google.api_core.exceptions import ServiceUnavailable - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _BigtableRetryableError + from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.deprecated.table import _BigtableRetryableError from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 # Setup: @@ -1797,7 +1799,9 @@ def test_rmrw_do_mutate_retryable_rows_w_retryable_error_internal_rst_stream_err # Raise internal server error with RST STREAM error messages # There should be no error raised and that the request is retried from google.api_core.exceptions import InternalServerError - from google.cloud.bigtable.row_data import RETRYABLE_INTERNAL_ERROR_MESSAGES + from google.cloud.bigtable.deprecated.row_data import ( + RETRYABLE_INTERNAL_ERROR_MESSAGES, + ) row_cells = [ (b"row_key_1", ("cf", b"col", b"value1")), @@ -2003,7 +2007,7 @@ def test_rmrw_do_mutate_retryable_rows_mismatch_num_responses(): def test__create_row_request_table_name_only(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request table_name = "table_name" result = _create_row_request(table_name) @@ -2012,14 +2016,14 @@ def test__create_row_request_table_name_only(): def test__create_row_request_row_range_row_set_conflict(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request with pytest.raises(ValueError): _create_row_request(None, end_key=object(), row_set=object()) def test__create_row_request_row_range_start_key(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange table_name = "table_name" @@ -2032,7 +2036,7 @@ def test__create_row_request_row_range_start_key(): def test__create_row_request_row_range_end_key(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange table_name = "table_name" @@ -2045,7 +2049,7 @@ def test__create_row_request_row_range_end_key(): def test__create_row_request_row_range_both_keys(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange table_name = "table_name" @@ -2059,7 +2063,7 @@ def test__create_row_request_row_range_both_keys(): def test__create_row_request_row_range_both_keys_inclusive(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange table_name = "table_name" @@ -2075,8 +2079,8 @@ def test__create_row_request_row_range_both_keys_inclusive(): def test__create_row_request_with_filter(): - from google.cloud.bigtable.table import _create_row_request - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter table_name = "table_name" row_filter = RowSampleFilter(0.33) @@ -2088,7 +2092,7 @@ def test__create_row_request_with_filter(): def test__create_row_request_with_limit(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request table_name = "table_name" limit = 1337 @@ -2098,8 +2102,8 @@ def test__create_row_request_with_limit(): def test__create_row_request_with_row_set(): - from google.cloud.bigtable.table import _create_row_request - from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.deprecated.row_set import RowSet table_name = "table_name" row_set = RowSet() @@ -2109,7 +2113,7 @@ def test__create_row_request_with_row_set(): def test__create_row_request_with_app_profile_id(): - from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.deprecated.table import _create_row_request table_name = "table_name" limit = 1337 @@ -2128,8 +2132,8 @@ def _ReadRowsRequestPB(*args, **kw): def test_cluster_state___eq__(): - from google.cloud.bigtable.enums import Table as enum_table - from google.cloud.bigtable.table import ClusterState + from google.cloud.bigtable.deprecated.enums import Table as enum_table + from google.cloud.bigtable.deprecated.table import ClusterState READY = enum_table.ReplicationState.READY state1 = ClusterState(READY) @@ -2138,8 +2142,8 @@ def test_cluster_state___eq__(): def test_cluster_state___eq__type_differ(): - from google.cloud.bigtable.enums import Table as enum_table - from google.cloud.bigtable.table import ClusterState + from google.cloud.bigtable.deprecated.enums import Table as enum_table + from google.cloud.bigtable.deprecated.table import ClusterState READY = enum_table.ReplicationState.READY state1 = ClusterState(READY) @@ -2148,8 +2152,8 @@ def test_cluster_state___eq__type_differ(): def test_cluster_state___ne__same_value(): - from google.cloud.bigtable.enums import Table as enum_table - from google.cloud.bigtable.table import ClusterState + from google.cloud.bigtable.deprecated.enums import Table as enum_table + from google.cloud.bigtable.deprecated.table import ClusterState READY = enum_table.ReplicationState.READY state1 = ClusterState(READY) @@ -2158,8 +2162,8 @@ def test_cluster_state___ne__same_value(): def test_cluster_state___ne__(): - from google.cloud.bigtable.enums import Table as enum_table - from google.cloud.bigtable.table import ClusterState + from google.cloud.bigtable.deprecated.enums import Table as enum_table + from google.cloud.bigtable.deprecated.table import ClusterState READY = enum_table.ReplicationState.READY INITIALIZING = enum_table.ReplicationState.INITIALIZING @@ -2169,8 +2173,8 @@ def test_cluster_state___ne__(): def test_cluster_state__repr__(): - from google.cloud.bigtable.enums import Table as enum_table - from google.cloud.bigtable.table import ClusterState + from google.cloud.bigtable.deprecated.enums import Table as enum_table + from google.cloud.bigtable.deprecated.table import ClusterState STATE_NOT_KNOWN = enum_table.ReplicationState.STATE_NOT_KNOWN INITIALIZING = enum_table.ReplicationState.INITIALIZING From 507da99b1fe8abf2d999553fb87bc9b42a8eadb1 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 23 Mar 2023 11:22:34 -0700 Subject: [PATCH 02/56] feat: improve rows filters (#751) --- google/cloud/bigtable/row_filters.py | 482 +++--- noxfile.py | 1 - owlbot.py | 1 - tests/unit/test_row_filters.py | 2009 ++++++++++++++++++++++++++ 4 files changed, 2297 insertions(+), 196 deletions(-) create mode 100644 tests/unit/test_row_filters.py diff --git a/google/cloud/bigtable/row_filters.py b/google/cloud/bigtable/row_filters.py index 53192acc8..48a1d4d8a 100644 --- a/google/cloud/bigtable/row_filters.py +++ b/google/cloud/bigtable/row_filters.py @@ -13,18 +13,25 @@ # limitations under the License. """Filters for Google Cloud Bigtable Row classes.""" +from __future__ import annotations import struct +from typing import Any, Sequence, TYPE_CHECKING, overload +from abc import ABC, abstractmethod from google.cloud._helpers import _microseconds_from_datetime # type: ignore from google.cloud._helpers import _to_bytes # type: ignore from google.cloud.bigtable_v2.types import data as data_v2_pb2 +if TYPE_CHECKING: + # import dependencies when type checking + from datetime import datetime + _PACK_I64 = struct.Struct(">q").pack -class RowFilter(object): +class RowFilter(ABC): """Basic filter to apply to cells in a row. These values can be combined via :class:`RowFilterChain`, @@ -35,15 +42,30 @@ class RowFilter(object): This class is a do-nothing base class for all row filters. """ + def _to_pb(self) -> data_v2_pb2.RowFilter: + """Converts the row filter to a protobuf. + + Returns: The converted current object. + """ + return data_v2_pb2.RowFilter(**self.to_dict()) + + @abstractmethod + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + pass + + def __repr__(self) -> str: + return f"{self.__class__.__name__}()" -class _BoolFilter(RowFilter): + +class _BoolFilter(RowFilter, ABC): """Row filter that uses a boolean flag. :type flag: bool :param flag: An indicator if a setting is turned on or off. """ - def __init__(self, flag): + def __init__(self, flag: bool): self.flag = flag def __eq__(self, other): @@ -54,6 +76,9 @@ def __eq__(self, other): def __ne__(self, other): return not self == other + def __repr__(self) -> str: + return f"{self.__class__.__name__}(flag={self.flag})" + class SinkFilter(_BoolFilter): """Advanced row filter to skip parent filters. @@ -66,13 +91,9 @@ class SinkFilter(_BoolFilter): of a :class:`ConditionalRowFilter`. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(sink=self.flag) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"sink": self.flag} class PassAllFilter(_BoolFilter): @@ -84,13 +105,9 @@ class PassAllFilter(_BoolFilter): completeness. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(pass_all_filter=self.flag) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"pass_all_filter": self.flag} class BlockAllFilter(_BoolFilter): @@ -101,16 +118,12 @@ class BlockAllFilter(_BoolFilter): temporarily disabling just part of a filter. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(block_all_filter=self.flag) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"block_all_filter": self.flag} -class _RegexFilter(RowFilter): +class _RegexFilter(RowFilter, ABC): """Row filter that uses a regular expression. The ``regex`` must be valid RE2 patterns. See Google's @@ -124,8 +137,8 @@ class _RegexFilter(RowFilter): will be encoded as ASCII. """ - def __init__(self, regex): - self.regex = _to_bytes(regex) + def __init__(self, regex: str | bytes): + self.regex: bytes = _to_bytes(regex) def __eq__(self, other): if not isinstance(other, self.__class__): @@ -135,6 +148,9 @@ def __eq__(self, other): def __ne__(self, other): return not self == other + def __repr__(self) -> str: + return f"{self.__class__.__name__}(regex={self.regex!r})" + class RowKeyRegexFilter(_RegexFilter): """Row filter for a row key regular expression. @@ -159,13 +175,9 @@ class RowKeyRegexFilter(_RegexFilter): since the row key is already specified. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(row_key_regex_filter=self.regex) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"row_key_regex_filter": self.regex} class RowSampleFilter(RowFilter): @@ -176,8 +188,8 @@ class RowSampleFilter(RowFilter): interval ``(0, 1)`` The end points are excluded). """ - def __init__(self, sample): - self.sample = sample + def __init__(self, sample: float): + self.sample: float = sample def __eq__(self, other): if not isinstance(other, self.__class__): @@ -187,13 +199,12 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): - """Converts the row filter to a protobuf. + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"row_sample_filter": self.sample} - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(row_sample_filter=self.sample) + def __repr__(self) -> str: + return f"{self.__class__.__name__}(sample={self.sample})" class FamilyNameRegexFilter(_RegexFilter): @@ -211,13 +222,9 @@ class FamilyNameRegexFilter(_RegexFilter): used as a literal. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(family_name_regex_filter=self.regex) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"family_name_regex_filter": self.regex} class ColumnQualifierRegexFilter(_RegexFilter): @@ -241,13 +248,9 @@ class ColumnQualifierRegexFilter(_RegexFilter): match this regex (irrespective of column family). """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(column_qualifier_regex_filter=self.regex) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"column_qualifier_regex_filter": self.regex} class TimestampRange(object): @@ -262,9 +265,9 @@ class TimestampRange(object): range. If omitted, no upper bound is used. """ - def __init__(self, start=None, end=None): - self.start = start - self.end = end + def __init__(self, start: "datetime" | None = None, end: "datetime" | None = None): + self.start: "datetime" | None = start + self.end: "datetime" | None = end def __eq__(self, other): if not isinstance(other, self.__class__): @@ -274,23 +277,29 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.TimestampRange: """Converts the :class:`TimestampRange` to a protobuf. - :rtype: :class:`.data_v2_pb2.TimestampRange` - :returns: The converted current object. + Returns: The converted current object. """ + return data_v2_pb2.TimestampRange(**self.to_dict()) + + def to_dict(self) -> dict[str, int]: + """Converts the timestamp range to a dict representation.""" timestamp_range_kwargs = {} if self.start is not None: - timestamp_range_kwargs["start_timestamp_micros"] = ( - _microseconds_from_datetime(self.start) // 1000 * 1000 - ) + start_time = _microseconds_from_datetime(self.start) // 1000 * 1000 + timestamp_range_kwargs["start_timestamp_micros"] = start_time if self.end is not None: end_time = _microseconds_from_datetime(self.end) if end_time % 1000 != 0: + # if not a whole milisecond value, round up end_time = end_time // 1000 * 1000 + 1000 timestamp_range_kwargs["end_timestamp_micros"] = end_time - return data_v2_pb2.TimestampRange(**timestamp_range_kwargs) + return timestamp_range_kwargs + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(start={self.start}, end={self.end})" class TimestampRangeFilter(RowFilter): @@ -300,8 +309,8 @@ class TimestampRangeFilter(RowFilter): :param range_: Range of time that cells should match against. """ - def __init__(self, range_): - self.range_ = range_ + def __init__(self, start: "datetime" | None = None, end: "datetime" | None = None): + self.range_: TimestampRange = TimestampRange(start, end) def __eq__(self, other): if not isinstance(other, self.__class__): @@ -311,16 +320,22 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. First converts the ``range_`` on the current object to a protobuf and then uses it in the ``timestamp_range_filter`` field. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ - return data_v2_pb2.RowFilter(timestamp_range_filter=self.range_.to_pb()) + return data_v2_pb2.RowFilter(timestamp_range_filter=self.range_._to_pb()) + + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"timestamp_range_filter": self.range_.to_dict()} + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(start={self.range_.start!r}, end={self.range_.end!r})" class ColumnRangeFilter(RowFilter): @@ -330,71 +345,72 @@ class ColumnRangeFilter(RowFilter): By default, we include them both, but this can be changed with optional flags. - :type column_family_id: str - :param column_family_id: The column family that contains the columns. Must + :type family_id: str + :param family_id: The column family that contains the columns. Must be of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - :type start_column: bytes - :param start_column: The start of the range of columns. If no value is + :type start_qualifier: bytes + :param start_qualifier: The start of the range of columns. If no value is used, the backend applies no upper bound to the values. - :type end_column: bytes - :param end_column: The end of the range of columns. If no value is used, + :type end_qualifier: bytes + :param end_qualifier: The end of the range of columns. If no value is used, the backend applies no upper bound to the values. :type inclusive_start: bool :param inclusive_start: Boolean indicating if the start column should be included in the range (or excluded). Defaults - to :data:`True` if ``start_column`` is passed and + to :data:`True` if ``start_qualifier`` is passed and no ``inclusive_start`` was given. :type inclusive_end: bool :param inclusive_end: Boolean indicating if the end column should be included in the range (or excluded). Defaults - to :data:`True` if ``end_column`` is passed and + to :data:`True` if ``end_qualifier`` is passed and no ``inclusive_end`` was given. :raises: :class:`ValueError ` if ``inclusive_start`` - is set but no ``start_column`` is given or if ``inclusive_end`` - is set but no ``end_column`` is given + is set but no ``start_qualifier`` is given or if ``inclusive_end`` + is set but no ``end_qualifier`` is given """ def __init__( self, - column_family_id, - start_column=None, - end_column=None, - inclusive_start=None, - inclusive_end=None, + family_id: str, + start_qualifier: bytes | None = None, + end_qualifier: bytes | None = None, + inclusive_start: bool | None = None, + inclusive_end: bool | None = None, ): - self.column_family_id = column_family_id - if inclusive_start is None: inclusive_start = True - elif start_column is None: + elif start_qualifier is None: raise ValueError( - "Inclusive start was specified but no " "start column was given." + "inclusive_start was specified but no start_qualifier was given." ) - self.start_column = start_column - self.inclusive_start = inclusive_start - if inclusive_end is None: inclusive_end = True - elif end_column is None: + elif end_qualifier is None: raise ValueError( - "Inclusive end was specified but no " "end column was given." + "inclusive_end was specified but no end_qualifier was given." ) - self.end_column = end_column + + self.family_id = family_id + + self.start_qualifier = start_qualifier + self.inclusive_start = inclusive_start + + self.end_qualifier = end_qualifier self.inclusive_end = inclusive_end def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented return ( - other.column_family_id == self.column_family_id - and other.start_column == self.start_column - and other.end_column == self.end_column + other.family_id == self.family_id + and other.start_qualifier == self.start_qualifier + and other.end_qualifier == self.end_qualifier and other.inclusive_start == self.inclusive_start and other.inclusive_end == self.inclusive_end ) @@ -402,31 +418,41 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. First converts to a :class:`.data_v2_pb2.ColumnRange` and then uses it in the ``column_range_filter`` field. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ - column_range_kwargs = {"family_name": self.column_family_id} - if self.start_column is not None: + column_range = data_v2_pb2.ColumnRange(**self.range_to_dict()) + return data_v2_pb2.RowFilter(column_range_filter=column_range) + + def range_to_dict(self) -> dict[str, str | bytes]: + """Converts the column range range to a dict representation.""" + column_range_kwargs: dict[str, str | bytes] = {} + column_range_kwargs["family_name"] = self.family_id + if self.start_qualifier is not None: if self.inclusive_start: key = "start_qualifier_closed" else: key = "start_qualifier_open" - column_range_kwargs[key] = _to_bytes(self.start_column) - if self.end_column is not None: + column_range_kwargs[key] = _to_bytes(self.start_qualifier) + if self.end_qualifier is not None: if self.inclusive_end: key = "end_qualifier_closed" else: key = "end_qualifier_open" - column_range_kwargs[key] = _to_bytes(self.end_column) + column_range_kwargs[key] = _to_bytes(self.end_qualifier) + return column_range_kwargs - column_range = data_v2_pb2.ColumnRange(**column_range_kwargs) - return data_v2_pb2.RowFilter(column_range_filter=column_range) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"column_range_filter": self.range_to_dict()} + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(family_id='{self.family_id}', start_qualifier={self.start_qualifier!r}, end_qualifier={self.end_qualifier!r}, inclusive_start={self.inclusive_start}, inclusive_end={self.inclusive_end})" class ValueRegexFilter(_RegexFilter): @@ -450,13 +476,9 @@ class ValueRegexFilter(_RegexFilter): match this regex. String values will be encoded as ASCII. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(value_regex_filter=self.regex) + def to_dict(self) -> dict[str, bytes]: + """Converts the row filter to a dict representation.""" + return {"value_regex_filter": self.regex} class ExactValueFilter(ValueRegexFilter): @@ -469,11 +491,14 @@ class ExactValueFilter(ValueRegexFilter): equivalent bytes, or an integer (which will be packed into 8-bytes). """ - def __init__(self, value): + def __init__(self, value: bytes | str | int): if isinstance(value, int): value = _PACK_I64(value) super(ExactValueFilter, self).__init__(value) + def __repr__(self) -> str: + return f"{self.__class__.__name__}(value={self.regex!r})" + class ValueRangeFilter(RowFilter): """A range of values to restrict to in a row filter. @@ -510,25 +535,29 @@ class ValueRangeFilter(RowFilter): """ def __init__( - self, start_value=None, end_value=None, inclusive_start=None, inclusive_end=None + self, + start_value: bytes | int | None = None, + end_value: bytes | int | None = None, + inclusive_start: bool | None = None, + inclusive_end: bool | None = None, ): if inclusive_start is None: inclusive_start = True elif start_value is None: raise ValueError( - "Inclusive start was specified but no " "start value was given." + "inclusive_start was specified but no start_value was given." ) - if isinstance(start_value, int): - start_value = _PACK_I64(start_value) - self.start_value = start_value - self.inclusive_start = inclusive_start - if inclusive_end is None: inclusive_end = True elif end_value is None: raise ValueError( - "Inclusive end was specified but no " "end value was given." + "inclusive_end was specified but no end_qualifier was given." ) + if isinstance(start_value, int): + start_value = _PACK_I64(start_value) + self.start_value = start_value + self.inclusive_start = inclusive_start + if isinstance(end_value, int): end_value = _PACK_I64(end_value) self.end_value = end_value @@ -547,15 +576,19 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. First converts to a :class:`.data_v2_pb2.ValueRange` and then uses it to create a row filter protobuf. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ + value_range = data_v2_pb2.ValueRange(**self.range_to_dict()) + return data_v2_pb2.RowFilter(value_range_filter=value_range) + + def range_to_dict(self) -> dict[str, bytes]: + """Converts the value range range to a dict representation.""" value_range_kwargs = {} if self.start_value is not None: if self.inclusive_start: @@ -569,12 +602,17 @@ def to_pb(self): else: key = "end_value_open" value_range_kwargs[key] = _to_bytes(self.end_value) + return value_range_kwargs - value_range = data_v2_pb2.ValueRange(**value_range_kwargs) - return data_v2_pb2.RowFilter(value_range_filter=value_range) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"value_range_filter": self.range_to_dict()} + def __repr__(self) -> str: + return f"{self.__class__.__name__}(start_value={self.start_value!r}, end_value={self.end_value!r}, inclusive_start={self.inclusive_start}, inclusive_end={self.inclusive_end})" -class _CellCountFilter(RowFilter): + +class _CellCountFilter(RowFilter, ABC): """Row filter that uses an integer count of cells. The cell count is used as an offset or a limit for the number @@ -584,7 +622,7 @@ class _CellCountFilter(RowFilter): :param num_cells: An integer count / offset / limit. """ - def __init__(self, num_cells): + def __init__(self, num_cells: int): self.num_cells = num_cells def __eq__(self, other): @@ -595,6 +633,9 @@ def __eq__(self, other): def __ne__(self, other): return not self == other + def __repr__(self) -> str: + return f"{self.__class__.__name__}(num_cells={self.num_cells})" + class CellsRowOffsetFilter(_CellCountFilter): """Row filter to skip cells in a row. @@ -603,13 +644,9 @@ class CellsRowOffsetFilter(_CellCountFilter): :param num_cells: Skips the first N cells of the row. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(cells_per_row_offset_filter=self.num_cells) + def to_dict(self) -> dict[str, int]: + """Converts the row filter to a dict representation.""" + return {"cells_per_row_offset_filter": self.num_cells} class CellsRowLimitFilter(_CellCountFilter): @@ -619,13 +656,9 @@ class CellsRowLimitFilter(_CellCountFilter): :param num_cells: Matches only the first N cells of the row. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(cells_per_row_limit_filter=self.num_cells) + def to_dict(self) -> dict[str, int]: + """Converts the row filter to a dict representation.""" + return {"cells_per_row_limit_filter": self.num_cells} class CellsColumnLimitFilter(_CellCountFilter): @@ -637,13 +670,9 @@ class CellsColumnLimitFilter(_CellCountFilter): timestamps of each cell. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(cells_per_column_limit_filter=self.num_cells) + def to_dict(self) -> dict[str, int]: + """Converts the row filter to a dict representation.""" + return {"cells_per_column_limit_filter": self.num_cells} class StripValueTransformerFilter(_BoolFilter): @@ -655,13 +684,9 @@ class StripValueTransformerFilter(_BoolFilter): transformer than a generic query / filter. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(strip_value_transformer=self.flag) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"strip_value_transformer": self.flag} class ApplyLabelFilter(RowFilter): @@ -683,7 +708,7 @@ class ApplyLabelFilter(RowFilter): ``[a-z0-9\\-]+``. """ - def __init__(self, label): + def __init__(self, label: str): self.label = label def __eq__(self, other): @@ -694,16 +719,15 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): - """Converts the row filter to a protobuf. + def to_dict(self) -> dict[str, str]: + """Converts the row filter to a dict representation.""" + return {"apply_label_transformer": self.label} - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(apply_label_transformer=self.label) + def __repr__(self) -> str: + return f"{self.__class__.__name__}(label={self.label})" -class _FilterCombination(RowFilter): +class _FilterCombination(RowFilter, Sequence[RowFilter], ABC): """Chain of row filters. Sends rows through several filters in sequence. The filters are "chained" @@ -714,10 +738,10 @@ class _FilterCombination(RowFilter): :param filters: List of :class:`RowFilter` """ - def __init__(self, filters=None): + def __init__(self, filters: list[RowFilter] | None = None): if filters is None: filters = [] - self.filters = filters + self.filters: list[RowFilter] = filters def __eq__(self, other): if not isinstance(other, self.__class__): @@ -727,6 +751,38 @@ def __eq__(self, other): def __ne__(self, other): return not self == other + def __len__(self) -> int: + return len(self.filters) + + @overload + def __getitem__(self, index: int) -> RowFilter: + # overload signature for type checking + pass + + @overload + def __getitem__(self, index: slice) -> list[RowFilter]: + # overload signature for type checking + pass + + def __getitem__(self, index): + return self.filters[index] + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(filters={self.filters})" + + def __str__(self) -> str: + """ + Returns a string representation of the filter chain. + + Adds line breaks between each sub-filter for readability. + """ + output = [f"{self.__class__.__name__}(["] + for filter_ in self.filters: + filter_lines = f"{filter_},".splitlines() + output.extend([f" {line}" for line in filter_lines]) + output.append("])") + return "\n".join(output) + class RowFilterChain(_FilterCombination): """Chain of row filters. @@ -739,17 +795,20 @@ class RowFilterChain(_FilterCombination): :param filters: List of :class:`RowFilter` """ - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ chain = data_v2_pb2.RowFilter.Chain( - filters=[row_filter.to_pb() for row_filter in self.filters] + filters=[row_filter._to_pb() for row_filter in self.filters] ) return data_v2_pb2.RowFilter(chain=chain) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"chain": {"filters": [f.to_dict() for f in self.filters]}} + class RowFilterUnion(_FilterCombination): """Union of row filters. @@ -764,50 +823,58 @@ class RowFilterUnion(_FilterCombination): :param filters: List of :class:`RowFilter` """ - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ interleave = data_v2_pb2.RowFilter.Interleave( - filters=[row_filter.to_pb() for row_filter in self.filters] + filters=[row_filter._to_pb() for row_filter in self.filters] ) return data_v2_pb2.RowFilter(interleave=interleave) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"interleave": {"filters": [f.to_dict() for f in self.filters]}} + class ConditionalRowFilter(RowFilter): """Conditional row filter which exhibits ternary behavior. - Executes one of two filters based on another filter. If the ``base_filter`` + Executes one of two filters based on another filter. If the ``predicate_filter`` returns any cells in the row, then ``true_filter`` is executed. If not, then ``false_filter`` is executed. .. note:: - The ``base_filter`` does not execute atomically with the true and false + The ``predicate_filter`` does not execute atomically with the true and false filters, which may lead to inconsistent or unexpected results. Additionally, executing a :class:`ConditionalRowFilter` has poor performance on the server, especially when ``false_filter`` is set. - :type base_filter: :class:`RowFilter` - :param base_filter: The filter to condition on before executing the + :type predicate_filter: :class:`RowFilter` + :param predicate_filter: The filter to condition on before executing the true/false filters. :type true_filter: :class:`RowFilter` :param true_filter: (Optional) The filter to execute if there are any cells - matching ``base_filter``. If not provided, no results + matching ``predicate_filter``. If not provided, no results will be returned in the true case. :type false_filter: :class:`RowFilter` :param false_filter: (Optional) The filter to execute if there are no cells - matching ``base_filter``. If not provided, no results + matching ``predicate_filter``. If not provided, no results will be returned in the false case. """ - def __init__(self, base_filter, true_filter=None, false_filter=None): - self.base_filter = base_filter + def __init__( + self, + predicate_filter: RowFilter, + true_filter: RowFilter | None = None, + false_filter: RowFilter | None = None, + ): + self.predicate_filter = predicate_filter self.true_filter = true_filter self.false_filter = false_filter @@ -815,7 +882,7 @@ def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented return ( - other.base_filter == self.base_filter + other.predicate_filter == self.predicate_filter and other.true_filter == self.true_filter and other.false_filter == self.false_filter ) @@ -823,16 +890,43 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ - condition_kwargs = {"predicate_filter": self.base_filter.to_pb()} + condition_kwargs = {"predicate_filter": self.predicate_filter._to_pb()} if self.true_filter is not None: - condition_kwargs["true_filter"] = self.true_filter.to_pb() + condition_kwargs["true_filter"] = self.true_filter._to_pb() if self.false_filter is not None: - condition_kwargs["false_filter"] = self.false_filter.to_pb() + condition_kwargs["false_filter"] = self.false_filter._to_pb() condition = data_v2_pb2.RowFilter.Condition(**condition_kwargs) return data_v2_pb2.RowFilter(condition=condition) + + def condition_to_dict(self) -> dict[str, Any]: + """Converts the condition to a dict representation.""" + condition_kwargs = {"predicate_filter": self.predicate_filter.to_dict()} + if self.true_filter is not None: + condition_kwargs["true_filter"] = self.true_filter.to_dict() + if self.false_filter is not None: + condition_kwargs["false_filter"] = self.false_filter.to_dict() + return condition_kwargs + + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"condition": self.condition_to_dict()} + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(predicate_filter={self.predicate_filter!r}, true_filter={self.true_filter!r}, false_filter={self.false_filter!r})" + + def __str__(self) -> str: + output = [f"{self.__class__.__name__}("] + for filter_type in ("predicate_filter", "true_filter", "false_filter"): + filter_ = getattr(self, filter_type) + if filter_ is None: + continue + # add the new filter set, adding indentations for readability + filter_lines = f"{filter_type}={filter_},".splitlines() + output.extend(f" {line}" for line in filter_lines) + output.append(")") + return "\n".join(output) diff --git a/noxfile.py b/noxfile.py index 94982fcfd..ed69bf85e 100644 --- a/noxfile.py +++ b/noxfile.py @@ -133,7 +133,6 @@ def mypy(session): session.run( "mypy", "google/cloud/bigtable", - "tests/", "--check-untyped-defs", "--warn-unreachable", "--disallow-any-generics", diff --git a/owlbot.py b/owlbot.py index f56a0bd9e..b542b3246 100644 --- a/owlbot.py +++ b/owlbot.py @@ -171,7 +171,6 @@ def mypy(session): session.run( "mypy", "google/cloud/bigtable", - "tests/", "--check-untyped-defs", "--warn-unreachable", "--disallow-any-generics", diff --git a/tests/unit/test_row_filters.py b/tests/unit/test_row_filters.py new file mode 100644 index 000000000..d0fbad42f --- /dev/null +++ b/tests/unit/test_row_filters.py @@ -0,0 +1,2009 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + + +import pytest + + +def test_abstract_class_constructors(): + from google.cloud.bigtable.row_filters import RowFilter + from google.cloud.bigtable.row_filters import _BoolFilter + from google.cloud.bigtable.row_filters import _FilterCombination + from google.cloud.bigtable.row_filters import _CellCountFilter + + with pytest.raises(TypeError): + RowFilter() + with pytest.raises(TypeError): + _BoolFilter(False) + with pytest.raises(TypeError): + _FilterCombination([]) + with pytest.raises(TypeError): + _CellCountFilter(0) + + +def test_bool_filter_constructor(): + for FilterType in _get_bool_filters(): + flag = True + row_filter = FilterType(flag) + assert row_filter.flag is flag + + +def test_bool_filter___eq__type_differ(): + for FilterType in _get_bool_filters(): + flag = object() + row_filter1 = FilterType(flag) + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_bool_filter___eq__same_value(): + for FilterType in _get_bool_filters(): + flag = object() + row_filter1 = FilterType(flag) + row_filter2 = FilterType(flag) + assert row_filter1 == row_filter2 + + +def test_bool_filter___ne__same_value(): + for FilterType in _get_bool_filters(): + flag = object() + row_filter1 = FilterType(flag) + row_filter2 = FilterType(flag) + assert not (row_filter1 != row_filter2) + + +def test_sink_filter_to_pb(): + from google.cloud.bigtable.row_filters import SinkFilter + + flag = True + row_filter = SinkFilter(flag) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(sink=flag) + assert pb_val == expected_pb + + +def test_sink_filter_to_dict(): + from google.cloud.bigtable.row_filters import SinkFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + flag = True + row_filter = SinkFilter(flag) + expected_dict = {"sink": flag} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_sink_filter___repr__(): + from google.cloud.bigtable.row_filters import SinkFilter + + flag = True + row_filter = SinkFilter(flag) + assert repr(row_filter) == "SinkFilter(flag={})".format(flag) + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_pass_all_filter_to_pb(): + from google.cloud.bigtable.row_filters import PassAllFilter + + flag = True + row_filter = PassAllFilter(flag) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(pass_all_filter=flag) + assert pb_val == expected_pb + + +def test_pass_all_filter_to_dict(): + from google.cloud.bigtable.row_filters import PassAllFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + flag = True + row_filter = PassAllFilter(flag) + expected_dict = {"pass_all_filter": flag} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_pass_all_filter___repr__(): + from google.cloud.bigtable.row_filters import PassAllFilter + + flag = True + row_filter = PassAllFilter(flag) + assert repr(row_filter) == "PassAllFilter(flag={})".format(flag) + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_block_all_filter_to_pb(): + from google.cloud.bigtable.row_filters import BlockAllFilter + + flag = True + row_filter = BlockAllFilter(flag) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(block_all_filter=flag) + assert pb_val == expected_pb + + +def test_block_all_filter_to_dict(): + from google.cloud.bigtable.row_filters import BlockAllFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + flag = True + row_filter = BlockAllFilter(flag) + expected_dict = {"block_all_filter": flag} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_block_all_filter___repr__(): + from google.cloud.bigtable.row_filters import BlockAllFilter + + flag = True + row_filter = BlockAllFilter(flag) + assert repr(row_filter) == "BlockAllFilter(flag={})".format(flag) + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_regex_filterconstructor(): + for FilterType in _get_regex_filters(): + regex = b"abc" + row_filter = FilterType(regex) + assert row_filter.regex == regex + + +def test_regex_filterconstructor_non_bytes(): + for FilterType in _get_regex_filters(): + regex = "abc" + row_filter = FilterType(regex) + assert row_filter.regex == b"abc" + + +def test_regex_filter__eq__type_differ(): + for FilterType in _get_regex_filters(): + regex = b"def-rgx" + row_filter1 = FilterType(regex) + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_regex_filter__eq__same_value(): + for FilterType in _get_regex_filters(): + regex = b"trex-regex" + row_filter1 = FilterType(regex) + row_filter2 = FilterType(regex) + assert row_filter1 == row_filter2 + + +def test_regex_filter__ne__same_value(): + for FilterType in _get_regex_filters(): + regex = b"abc" + row_filter1 = FilterType(regex) + row_filter2 = FilterType(regex) + assert not (row_filter1 != row_filter2) + + +def test_row_key_regex_filter_to_pb(): + from google.cloud.bigtable.row_filters import RowKeyRegexFilter + + regex = b"row-key-regex" + row_filter = RowKeyRegexFilter(regex) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(row_key_regex_filter=regex) + assert pb_val == expected_pb + + +def test_row_key_regex_filter_to_dict(): + from google.cloud.bigtable.row_filters import RowKeyRegexFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + regex = b"row-key-regex" + row_filter = RowKeyRegexFilter(regex) + expected_dict = {"row_key_regex_filter": regex} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_row_key_regex_filter___repr__(): + from google.cloud.bigtable.row_filters import RowKeyRegexFilter + + regex = b"row-key-regex" + row_filter = RowKeyRegexFilter(regex) + assert repr(row_filter) == "RowKeyRegexFilter(regex={})".format(regex) + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_row_sample_filter_constructor(): + from google.cloud.bigtable.row_filters import RowSampleFilter + + sample = object() + row_filter = RowSampleFilter(sample) + assert row_filter.sample is sample + + +def test_row_sample_filter___eq__type_differ(): + from google.cloud.bigtable.row_filters import RowSampleFilter + + sample = object() + row_filter1 = RowSampleFilter(sample) + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_row_sample_filter___eq__same_value(): + from google.cloud.bigtable.row_filters import RowSampleFilter + + sample = object() + row_filter1 = RowSampleFilter(sample) + row_filter2 = RowSampleFilter(sample) + assert row_filter1 == row_filter2 + + +def test_row_sample_filter___ne__(): + from google.cloud.bigtable.row_filters import RowSampleFilter + + sample = object() + other_sample = object() + row_filter1 = RowSampleFilter(sample) + row_filter2 = RowSampleFilter(other_sample) + assert row_filter1 != row_filter2 + + +def test_row_sample_filter_to_pb(): + from google.cloud.bigtable.row_filters import RowSampleFilter + + sample = 0.25 + row_filter = RowSampleFilter(sample) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(row_sample_filter=sample) + assert pb_val == expected_pb + + +def test_row_sample_filter___repr__(): + from google.cloud.bigtable.row_filters import RowSampleFilter + + sample = 0.25 + row_filter = RowSampleFilter(sample) + assert repr(row_filter) == "RowSampleFilter(sample={})".format(sample) + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_family_name_regex_filter_to_pb(): + from google.cloud.bigtable.row_filters import FamilyNameRegexFilter + + regex = "family-regex" + row_filter = FamilyNameRegexFilter(regex) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(family_name_regex_filter=regex) + assert pb_val == expected_pb + + +def test_family_name_regex_filter_to_dict(): + from google.cloud.bigtable.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + regex = "family-regex" + row_filter = FamilyNameRegexFilter(regex) + expected_dict = {"family_name_regex_filter": regex.encode()} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_family_name_regex_filter___repr__(): + from google.cloud.bigtable.row_filters import FamilyNameRegexFilter + + regex = "family-regex" + row_filter = FamilyNameRegexFilter(regex) + expected = "FamilyNameRegexFilter(regex=b'family-regex')" + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_column_qualifier_regex_filter_to_pb(): + from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter + + regex = b"column-regex" + row_filter = ColumnQualifierRegexFilter(regex) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(column_qualifier_regex_filter=regex) + assert pb_val == expected_pb + + +def test_column_qualifier_regex_filter_to_dict(): + from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + regex = b"column-regex" + row_filter = ColumnQualifierRegexFilter(regex) + expected_dict = {"column_qualifier_regex_filter": regex} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_column_qualifier_regex_filter___repr__(): + from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter + + regex = b"column-regex" + row_filter = ColumnQualifierRegexFilter(regex) + assert repr(row_filter) == "ColumnQualifierRegexFilter(regex={})".format(regex) + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_timestamp_range_constructor(): + from google.cloud.bigtable.row_filters import TimestampRange + + start = object() + end = object() + time_range = TimestampRange(start=start, end=end) + assert time_range.start is start + assert time_range.end is end + + +def test_timestamp_range___eq__(): + from google.cloud.bigtable.row_filters import TimestampRange + + start = object() + end = object() + time_range1 = TimestampRange(start=start, end=end) + time_range2 = TimestampRange(start=start, end=end) + assert time_range1 == time_range2 + + +def test_timestamp_range___eq__type_differ(): + from google.cloud.bigtable.row_filters import TimestampRange + + start = object() + end = object() + time_range1 = TimestampRange(start=start, end=end) + time_range2 = object() + assert not (time_range1 == time_range2) + + +def test_timestamp_range___ne__same_value(): + from google.cloud.bigtable.row_filters import TimestampRange + + start = object() + end = object() + time_range1 = TimestampRange(start=start, end=end) + time_range2 = TimestampRange(start=start, end=end) + assert not (time_range1 != time_range2) + + +def _timestamp_range_to_pb_helper(pb_kwargs, start=None, end=None): + import datetime + from google.cloud._helpers import _EPOCH + from google.cloud.bigtable.row_filters import TimestampRange + + if start is not None: + start = _EPOCH + datetime.timedelta(microseconds=start) + if end is not None: + end = _EPOCH + datetime.timedelta(microseconds=end) + time_range = TimestampRange(start=start, end=end) + expected_pb = _TimestampRangePB(**pb_kwargs) + time_pb = time_range._to_pb() + assert time_pb.start_timestamp_micros == expected_pb.start_timestamp_micros + assert time_pb.end_timestamp_micros == expected_pb.end_timestamp_micros + assert time_pb == expected_pb + + +def test_timestamp_range_to_pb(): + start_micros = 30871234 + end_micros = 12939371234 + start_millis = start_micros // 1000 * 1000 + assert start_millis == 30871000 + end_millis = end_micros // 1000 * 1000 + 1000 + assert end_millis == 12939372000 + pb_kwargs = {} + pb_kwargs["start_timestamp_micros"] = start_millis + pb_kwargs["end_timestamp_micros"] = end_millis + _timestamp_range_to_pb_helper(pb_kwargs, start=start_micros, end=end_micros) + + +def test_timestamp_range_to_dict(): + from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + import datetime + + row_filter = TimestampRange( + start=datetime.datetime(2019, 1, 1), end=datetime.datetime(2019, 1, 2) + ) + expected_dict = { + "start_timestamp_micros": 1546300800000000, + "end_timestamp_micros": 1546387200000000, + } + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.TimestampRange(**expected_dict) == expected_pb_value + + +def test_timestamp_range_to_pb_start_only(): + # Makes sure already milliseconds granularity + start_micros = 30871000 + start_millis = start_micros // 1000 * 1000 + assert start_millis == 30871000 + pb_kwargs = {} + pb_kwargs["start_timestamp_micros"] = start_millis + _timestamp_range_to_pb_helper(pb_kwargs, start=start_micros, end=None) + + +def test_timestamp_range_to_dict_start_only(): + from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + import datetime + + row_filter = TimestampRange(start=datetime.datetime(2019, 1, 1)) + expected_dict = {"start_timestamp_micros": 1546300800000000} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.TimestampRange(**expected_dict) == expected_pb_value + + +def test_timestamp_range_to_pb_end_only(): + # Makes sure already milliseconds granularity + end_micros = 12939371000 + end_millis = end_micros // 1000 * 1000 + assert end_millis == 12939371000 + pb_kwargs = {} + pb_kwargs["end_timestamp_micros"] = end_millis + _timestamp_range_to_pb_helper(pb_kwargs, start=None, end=end_micros) + + +def test_timestamp_range_to_dict_end_only(): + from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + import datetime + + row_filter = TimestampRange(end=datetime.datetime(2019, 1, 2)) + expected_dict = {"end_timestamp_micros": 1546387200000000} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.TimestampRange(**expected_dict) == expected_pb_value + + +def timestamp_range___repr__(): + from google.cloud.bigtable.row_filters import TimestampRange + + start = object() + end = object() + time_range = TimestampRange(start=start, end=end) + assert repr(time_range) == "TimestampRange(start={}, end={})".format(start, end) + assert repr(time_range) == str(time_range) + assert eval(repr(time_range)) == time_range + + +def test_timestamp_range_filter___eq__type_differ(): + from google.cloud.bigtable.row_filters import TimestampRangeFilter + + range_ = object() + row_filter1 = TimestampRangeFilter(range_) + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_timestamp_range_filter___eq__same_value(): + from google.cloud.bigtable.row_filters import TimestampRangeFilter + + range_ = object() + row_filter1 = TimestampRangeFilter(range_) + row_filter2 = TimestampRangeFilter(range_) + assert row_filter1 == row_filter2 + + +def test_timestamp_range_filter___ne__(): + from google.cloud.bigtable.row_filters import TimestampRangeFilter + + range_ = object() + other_range_ = object() + row_filter1 = TimestampRangeFilter(range_) + row_filter2 = TimestampRangeFilter(other_range_) + assert row_filter1 != row_filter2 + + +def test_timestamp_range_filter_to_pb(): + from google.cloud.bigtable.row_filters import TimestampRangeFilter + + row_filter = TimestampRangeFilter() + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(timestamp_range_filter=_TimestampRangePB()) + assert pb_val == expected_pb + + +def test_timestamp_range_filter_to_dict(): + from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + import datetime + + row_filter = TimestampRangeFilter( + start=datetime.datetime(2019, 1, 1), end=datetime.datetime(2019, 1, 2) + ) + expected_dict = { + "timestamp_range_filter": { + "start_timestamp_micros": 1546300800000000, + "end_timestamp_micros": 1546387200000000, + } + } + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_timestamp_range_filter_empty_to_dict(): + from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter = TimestampRangeFilter() + expected_dict = {"timestamp_range_filter": {}} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_timestamp_range_filter___repr__(): + from google.cloud.bigtable.row_filters import TimestampRangeFilter + import datetime + + start = datetime.datetime(2019, 1, 1) + end = datetime.datetime(2019, 1, 2) + row_filter = TimestampRangeFilter(start, end) + assert ( + repr(row_filter) + == f"TimestampRangeFilter(start={repr(start)}, end={repr(end)})" + ) + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_column_range_filter_constructor_defaults(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = object() + row_filter = ColumnRangeFilter(family_id) + assert row_filter.family_id is family_id + assert row_filter.start_qualifier is None + assert row_filter.end_qualifier is None + assert row_filter.inclusive_start + assert row_filter.inclusive_end + + +def test_column_range_filter_constructor_explicit(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = object() + start_qualifier = object() + end_qualifier = object() + inclusive_start = object() + inclusive_end = object() + row_filter = ColumnRangeFilter( + family_id, + start_qualifier=start_qualifier, + end_qualifier=end_qualifier, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + assert row_filter.family_id is family_id + assert row_filter.start_qualifier is start_qualifier + assert row_filter.end_qualifier is end_qualifier + assert row_filter.inclusive_start is inclusive_start + assert row_filter.inclusive_end is inclusive_end + + +def test_column_range_filter_constructor_(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = object() + with pytest.raises(ValueError): + ColumnRangeFilter(family_id, inclusive_start=True) + + +def test_column_range_filter_constructor_bad_end(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = object() + with pytest.raises(ValueError): + ColumnRangeFilter(family_id, inclusive_end=True) + + +def test_column_range_filter___eq__(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = object() + start_qualifier = object() + end_qualifier = object() + inclusive_start = object() + inclusive_end = object() + row_filter1 = ColumnRangeFilter( + family_id, + start_qualifier=start_qualifier, + end_qualifier=end_qualifier, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + row_filter2 = ColumnRangeFilter( + family_id, + start_qualifier=start_qualifier, + end_qualifier=end_qualifier, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + assert row_filter1 == row_filter2 + + +def test_column_range_filter___eq__type_differ(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = object() + row_filter1 = ColumnRangeFilter(family_id) + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_column_range_filter___ne__(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = object() + other_family_id = object() + start_qualifier = object() + end_qualifier = object() + inclusive_start = object() + inclusive_end = object() + row_filter1 = ColumnRangeFilter( + family_id, + start_qualifier=start_qualifier, + end_qualifier=end_qualifier, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + row_filter2 = ColumnRangeFilter( + other_family_id, + start_qualifier=start_qualifier, + end_qualifier=end_qualifier, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + assert row_filter1 != row_filter2 + + +def test_column_range_filter_to_pb(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = "column-family-id" + row_filter = ColumnRangeFilter(family_id) + col_range_pb = _ColumnRangePB(family_name=family_id) + expected_pb = _RowFilterPB(column_range_filter=col_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_column_range_filter_to_dict(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + family_id = "column-family-id" + row_filter = ColumnRangeFilter(family_id) + expected_dict = {"column_range_filter": {"family_name": family_id}} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_column_range_filter_to_pb_inclusive_start(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = "column-family-id" + column = b"column" + row_filter = ColumnRangeFilter(family_id, start_qualifier=column) + col_range_pb = _ColumnRangePB(family_name=family_id, start_qualifier_closed=column) + expected_pb = _RowFilterPB(column_range_filter=col_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_column_range_filter_to_pb_exclusive_start(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = "column-family-id" + column = b"column" + row_filter = ColumnRangeFilter( + family_id, start_qualifier=column, inclusive_start=False + ) + col_range_pb = _ColumnRangePB(family_name=family_id, start_qualifier_open=column) + expected_pb = _RowFilterPB(column_range_filter=col_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_column_range_filter_to_pb_inclusive_end(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = "column-family-id" + column = b"column" + row_filter = ColumnRangeFilter(family_id, end_qualifier=column) + col_range_pb = _ColumnRangePB(family_name=family_id, end_qualifier_closed=column) + expected_pb = _RowFilterPB(column_range_filter=col_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_column_range_filter_to_pb_exclusive_end(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = "column-family-id" + column = b"column" + row_filter = ColumnRangeFilter(family_id, end_qualifier=column, inclusive_end=False) + col_range_pb = _ColumnRangePB(family_name=family_id, end_qualifier_open=column) + expected_pb = _RowFilterPB(column_range_filter=col_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_column_range_filter___repr__(): + from google.cloud.bigtable.row_filters import ColumnRangeFilter + + family_id = "column-family-id" + start_qualifier = b"column" + end_qualifier = b"column2" + row_filter = ColumnRangeFilter(family_id, start_qualifier, end_qualifier) + expected = "ColumnRangeFilter(family_id='column-family-id', start_qualifier=b'column', end_qualifier=b'column2', inclusive_start=True, inclusive_end=True)" + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_value_regex_filter_to_pb_w_bytes(): + from google.cloud.bigtable.row_filters import ValueRegexFilter + + value = regex = b"value-regex" + row_filter = ValueRegexFilter(value) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(value_regex_filter=regex) + assert pb_val == expected_pb + + +def test_value_regex_filter_to_dict_w_bytes(): + from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + value = regex = b"value-regex" + row_filter = ValueRegexFilter(value) + expected_dict = {"value_regex_filter": regex} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_value_regex_filter_to_pb_w_str(): + from google.cloud.bigtable.row_filters import ValueRegexFilter + + value = "value-regex" + regex = value.encode("ascii") + row_filter = ValueRegexFilter(value) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(value_regex_filter=regex) + assert pb_val == expected_pb + + +def test_value_regex_filter_to_dict_w_str(): + from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + value = "value-regex" + regex = value.encode("ascii") + row_filter = ValueRegexFilter(value) + expected_dict = {"value_regex_filter": regex} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_value_regex_filter___repr__(): + from google.cloud.bigtable.row_filters import ValueRegexFilter + + value = "value-regex" + row_filter = ValueRegexFilter(value) + expected = "ValueRegexFilter(regex=b'value-regex')" + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_exact_value_filter_to_pb_w_bytes(): + from google.cloud.bigtable.row_filters import ExactValueFilter + + value = regex = b"value-regex" + row_filter = ExactValueFilter(value) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(value_regex_filter=regex) + assert pb_val == expected_pb + + +def test_exact_value_filter_to_dict_w_bytes(): + from google.cloud.bigtable.row_filters import ExactValueFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + value = regex = b"value-regex" + row_filter = ExactValueFilter(value) + expected_dict = {"value_regex_filter": regex} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_exact_value_filter_to_pb_w_str(): + from google.cloud.bigtable.row_filters import ExactValueFilter + + value = "value-regex" + regex = value.encode("ascii") + row_filter = ExactValueFilter(value) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(value_regex_filter=regex) + assert pb_val == expected_pb + + +def test_exact_value_filter_to_dict_w_str(): + from google.cloud.bigtable.row_filters import ExactValueFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + value = "value-regex" + regex = value.encode("ascii") + row_filter = ExactValueFilter(value) + expected_dict = {"value_regex_filter": regex} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_exact_value_filter_to_pb_w_int(): + import struct + from google.cloud.bigtable.row_filters import ExactValueFilter + + value = 1 + regex = struct.Struct(">q").pack(value) + row_filter = ExactValueFilter(value) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(value_regex_filter=regex) + assert pb_val == expected_pb + + +def test_exact_value_filter_to_dict_w_int(): + import struct + from google.cloud.bigtable.row_filters import ExactValueFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + value = 1 + regex = struct.Struct(">q").pack(value) + row_filter = ExactValueFilter(value) + expected_dict = {"value_regex_filter": regex} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_exact_value_filter___repr__(): + from google.cloud.bigtable.row_filters import ExactValueFilter + + value = "value-regex" + row_filter = ExactValueFilter(value) + expected = "ExactValueFilter(value=b'value-regex')" + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_value_range_filter_constructor_defaults(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + row_filter = ValueRangeFilter() + + assert row_filter.start_value is None + assert row_filter.end_value is None + assert row_filter.inclusive_start + assert row_filter.inclusive_end + + +def test_value_range_filter_constructor_explicit(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + start_value = object() + end_value = object() + inclusive_start = object() + inclusive_end = object() + + row_filter = ValueRangeFilter( + start_value=start_value, + end_value=end_value, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + + assert row_filter.start_value is start_value + assert row_filter.end_value is end_value + assert row_filter.inclusive_start is inclusive_start + assert row_filter.inclusive_end is inclusive_end + + +def test_value_range_filter_constructor_w_int_values(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + import struct + + start_value = 1 + end_value = 10 + + row_filter = ValueRangeFilter(start_value=start_value, end_value=end_value) + + expected_start_value = struct.Struct(">q").pack(start_value) + expected_end_value = struct.Struct(">q").pack(end_value) + + assert row_filter.start_value == expected_start_value + assert row_filter.end_value == expected_end_value + assert row_filter.inclusive_start + assert row_filter.inclusive_end + + +def test_value_range_filter_constructor_bad_start(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + with pytest.raises(ValueError): + ValueRangeFilter(inclusive_start=True) + + +def test_value_range_filter_constructor_bad_end(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + with pytest.raises(ValueError): + ValueRangeFilter(inclusive_end=True) + + +def test_value_range_filter___eq__(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + start_value = object() + end_value = object() + inclusive_start = object() + inclusive_end = object() + row_filter1 = ValueRangeFilter( + start_value=start_value, + end_value=end_value, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + row_filter2 = ValueRangeFilter( + start_value=start_value, + end_value=end_value, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + assert row_filter1 == row_filter2 + + +def test_value_range_filter___eq__type_differ(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + row_filter1 = ValueRangeFilter() + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_value_range_filter___ne__(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + start_value = object() + other_start_value = object() + end_value = object() + inclusive_start = object() + inclusive_end = object() + row_filter1 = ValueRangeFilter( + start_value=start_value, + end_value=end_value, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + row_filter2 = ValueRangeFilter( + start_value=other_start_value, + end_value=end_value, + inclusive_start=inclusive_start, + inclusive_end=inclusive_end, + ) + assert row_filter1 != row_filter2 + + +def test_value_range_filter_to_pb(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + row_filter = ValueRangeFilter() + expected_pb = _RowFilterPB(value_range_filter=_ValueRangePB()) + assert row_filter._to_pb() == expected_pb + + +def test_value_range_filter_to_dict(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter = ValueRangeFilter() + expected_dict = {"value_range_filter": {}} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_value_range_filter_to_pb_inclusive_start(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + value = b"some-value" + row_filter = ValueRangeFilter(start_value=value) + val_range_pb = _ValueRangePB(start_value_closed=value) + expected_pb = _RowFilterPB(value_range_filter=val_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_value_range_filter_to_pb_exclusive_start(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + value = b"some-value" + row_filter = ValueRangeFilter(start_value=value, inclusive_start=False) + val_range_pb = _ValueRangePB(start_value_open=value) + expected_pb = _RowFilterPB(value_range_filter=val_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_value_range_filter_to_pb_inclusive_end(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + value = b"some-value" + row_filter = ValueRangeFilter(end_value=value) + val_range_pb = _ValueRangePB(end_value_closed=value) + expected_pb = _RowFilterPB(value_range_filter=val_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_value_range_filter_to_pb_exclusive_end(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + value = b"some-value" + row_filter = ValueRangeFilter(end_value=value, inclusive_end=False) + val_range_pb = _ValueRangePB(end_value_open=value) + expected_pb = _RowFilterPB(value_range_filter=val_range_pb) + assert row_filter._to_pb() == expected_pb + + +def test_value_range_filter___repr__(): + from google.cloud.bigtable.row_filters import ValueRangeFilter + + start_value = b"some-value" + end_value = b"some-other-value" + row_filter = ValueRangeFilter( + start_value=start_value, end_value=end_value, inclusive_end=False + ) + expected = "ValueRangeFilter(start_value=b'some-value', end_value=b'some-other-value', inclusive_start=True, inclusive_end=False)" + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_cell_count_constructor(): + for FilerType in _get_cell_count_filters(): + num_cells = object() + row_filter = FilerType(num_cells) + assert row_filter.num_cells is num_cells + + +def test_cell_count___eq__type_differ(): + for FilerType in _get_cell_count_filters(): + num_cells = object() + row_filter1 = FilerType(num_cells) + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_cell_count___eq__same_value(): + for FilerType in _get_cell_count_filters(): + num_cells = object() + row_filter1 = FilerType(num_cells) + row_filter2 = FilerType(num_cells) + assert row_filter1 == row_filter2 + + +def test_cell_count___ne__same_value(): + for FilerType in _get_cell_count_filters(): + num_cells = object() + row_filter1 = FilerType(num_cells) + row_filter2 = FilerType(num_cells) + assert not (row_filter1 != row_filter2) + + +def test_cells_row_offset_filter_to_pb(): + from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + + num_cells = 76 + row_filter = CellsRowOffsetFilter(num_cells) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(cells_per_row_offset_filter=num_cells) + assert pb_val == expected_pb + + +def test_cells_row_offset_filter_to_dict(): + from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + num_cells = 76 + row_filter = CellsRowOffsetFilter(num_cells) + expected_dict = {"cells_per_row_offset_filter": num_cells} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_cells_row_offset_filter___repr__(): + from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + + num_cells = 76 + row_filter = CellsRowOffsetFilter(num_cells) + expected = "CellsRowOffsetFilter(num_cells={})".format(num_cells) + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_cells_row_limit_filter_to_pb(): + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + + num_cells = 189 + row_filter = CellsRowLimitFilter(num_cells) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(cells_per_row_limit_filter=num_cells) + assert pb_val == expected_pb + + +def test_cells_row_limit_filter_to_dict(): + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + num_cells = 189 + row_filter = CellsRowLimitFilter(num_cells) + expected_dict = {"cells_per_row_limit_filter": num_cells} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_cells_row_limit_filter___repr__(): + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + + num_cells = 189 + row_filter = CellsRowLimitFilter(num_cells) + expected = "CellsRowLimitFilter(num_cells={})".format(num_cells) + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_cells_column_limit_filter_to_pb(): + from google.cloud.bigtable.row_filters import CellsColumnLimitFilter + + num_cells = 10 + row_filter = CellsColumnLimitFilter(num_cells) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(cells_per_column_limit_filter=num_cells) + assert pb_val == expected_pb + + +def test_cells_column_limit_filter_to_dict(): + from google.cloud.bigtable.row_filters import CellsColumnLimitFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + num_cells = 10 + row_filter = CellsColumnLimitFilter(num_cells) + expected_dict = {"cells_per_column_limit_filter": num_cells} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_cells_column_limit_filter___repr__(): + from google.cloud.bigtable.row_filters import CellsColumnLimitFilter + + num_cells = 10 + row_filter = CellsColumnLimitFilter(num_cells) + expected = "CellsColumnLimitFilter(num_cells={})".format(num_cells) + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_strip_value_transformer_filter_to_pb(): + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + flag = True + row_filter = StripValueTransformerFilter(flag) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(strip_value_transformer=flag) + assert pb_val == expected_pb + + +def test_strip_value_transformer_filter_to_dict(): + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + flag = True + row_filter = StripValueTransformerFilter(flag) + expected_dict = {"strip_value_transformer": flag} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_strip_value_transformer_filter___repr__(): + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + flag = True + row_filter = StripValueTransformerFilter(flag) + expected = "StripValueTransformerFilter(flag={})".format(flag) + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_apply_label_filter_constructor(): + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + label = object() + row_filter = ApplyLabelFilter(label) + assert row_filter.label is label + + +def test_apply_label_filter___eq__type_differ(): + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + label = object() + row_filter1 = ApplyLabelFilter(label) + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_apply_label_filter___eq__same_value(): + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + label = object() + row_filter1 = ApplyLabelFilter(label) + row_filter2 = ApplyLabelFilter(label) + assert row_filter1 == row_filter2 + + +def test_apply_label_filter___ne__(): + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + label = object() + other_label = object() + row_filter1 = ApplyLabelFilter(label) + row_filter2 = ApplyLabelFilter(other_label) + assert row_filter1 != row_filter2 + + +def test_apply_label_filter_to_pb(): + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + label = "label" + row_filter = ApplyLabelFilter(label) + pb_val = row_filter._to_pb() + expected_pb = _RowFilterPB(apply_label_transformer=label) + assert pb_val == expected_pb + + +def test_apply_label_filter_to_dict(): + from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + label = "label" + row_filter = ApplyLabelFilter(label) + expected_dict = {"apply_label_transformer": label} + assert row_filter.to_dict() == expected_dict + expected_pb_value = row_filter._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_apply_label_filter___repr__(): + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + label = "label" + row_filter = ApplyLabelFilter(label) + expected = "ApplyLabelFilter(label={})".format(label) + assert repr(row_filter) == expected + assert repr(row_filter) == str(row_filter) + assert eval(repr(row_filter)) == row_filter + + +def test_filter_combination_constructor_defaults(): + for FilterType in _get_filter_combination_filters(): + row_filter = FilterType() + assert row_filter.filters == [] + + +def test_filter_combination_constructor_explicit(): + for FilterType in _get_filter_combination_filters(): + filters = object() + row_filter = FilterType(filters=filters) + assert row_filter.filters is filters + + +def test_filter_combination___eq__(): + for FilterType in _get_filter_combination_filters(): + filters = object() + row_filter1 = FilterType(filters=filters) + row_filter2 = FilterType(filters=filters) + assert row_filter1 == row_filter2 + + +def test_filter_combination___eq__type_differ(): + for FilterType in _get_filter_combination_filters(): + filters = object() + row_filter1 = FilterType(filters=filters) + row_filter2 = object() + assert not (row_filter1 == row_filter2) + + +def test_filter_combination___ne__(): + for FilterType in _get_filter_combination_filters(): + filters = object() + other_filters = object() + row_filter1 = FilterType(filters=filters) + row_filter2 = FilterType(filters=other_filters) + assert row_filter1 != row_filter2 + + +def test_filter_combination_len(): + for FilterType in _get_filter_combination_filters(): + filters = [object(), object()] + row_filter = FilterType(filters=filters) + assert len(row_filter) == len(filters) + + +def test_filter_combination_iter(): + for FilterType in _get_filter_combination_filters(): + filters = [object(), object()] + row_filter = FilterType(filters=filters) + assert list(iter(row_filter)) == filters + for filter_, expected in zip(row_filter, filters): + assert filter_ is expected + + +def test_filter_combination___getitem__(): + for FilterType in _get_filter_combination_filters(): + filters = [object(), object()] + row_filter = FilterType(filters=filters) + row_filter[0] is filters[0] + row_filter[1] is filters[1] + with pytest.raises(IndexError): + row_filter[2] + row_filter[:] is filters[:] + + +def test_filter_combination___str__(): + from google.cloud.bigtable.row_filters import PassAllFilter + + for FilterType in _get_filter_combination_filters(): + filters = [PassAllFilter(True), PassAllFilter(False)] + row_filter = FilterType(filters=filters) + expected = ( + "([\n PassAllFilter(flag=True),\n PassAllFilter(flag=False),\n])" + ) + assert expected in str(row_filter) + + +def test_row_filter_chain_to_pb(): + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_pb = row_filter1._to_pb() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_pb = row_filter2._to_pb() + + row_filter3 = RowFilterChain(filters=[row_filter1, row_filter2]) + filter_pb = row_filter3._to_pb() + + expected_pb = _RowFilterPB( + chain=_RowFilterChainPB(filters=[row_filter1_pb, row_filter2_pb]) + ) + assert filter_pb == expected_pb + + +def test_row_filter_chain_to_dict(): + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_dict = row_filter1.to_dict() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_dict = row_filter2.to_dict() + + row_filter3 = RowFilterChain(filters=[row_filter1, row_filter2]) + filter_dict = row_filter3.to_dict() + + expected_dict = {"chain": {"filters": [row_filter1_dict, row_filter2_dict]}} + assert filter_dict == expected_dict + expected_pb_value = row_filter3._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_row_filter_chain_to_pb_nested(): + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter2 = RowSampleFilter(0.25) + + row_filter3 = RowFilterChain(filters=[row_filter1, row_filter2]) + row_filter3_pb = row_filter3._to_pb() + + row_filter4 = CellsRowLimitFilter(11) + row_filter4_pb = row_filter4._to_pb() + + row_filter5 = RowFilterChain(filters=[row_filter3, row_filter4]) + filter_pb = row_filter5._to_pb() + + expected_pb = _RowFilterPB( + chain=_RowFilterChainPB(filters=[row_filter3_pb, row_filter4_pb]) + ) + assert filter_pb == expected_pb + + +def test_row_filter_chain_to_dict_nested(): + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter1 = StripValueTransformerFilter(True) + + row_filter2 = RowSampleFilter(0.25) + + row_filter3 = RowFilterChain(filters=[row_filter1, row_filter2]) + row_filter3_dict = row_filter3.to_dict() + + row_filter4 = CellsRowLimitFilter(11) + row_filter4_dict = row_filter4.to_dict() + + row_filter5 = RowFilterChain(filters=[row_filter3, row_filter4]) + filter_dict = row_filter5.to_dict() + + expected_dict = {"chain": {"filters": [row_filter3_dict, row_filter4_dict]}} + assert filter_dict == expected_dict + expected_pb_value = row_filter5._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_row_filter_chain___repr__(): + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter2 = RowSampleFilter(0.25) + + row_filter3 = RowFilterChain(filters=[row_filter1, row_filter2]) + expected = f"RowFilterChain(filters={[row_filter1, row_filter2]})" + assert repr(row_filter3) == expected + assert eval(repr(row_filter3)) == row_filter3 + + +def test_row_filter_chain___str__(): + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter2 = RowSampleFilter(0.25) + + row_filter3 = RowFilterChain(filters=[row_filter1, row_filter2]) + expected = "RowFilterChain([\n StripValueTransformerFilter(flag=True),\n RowSampleFilter(sample=0.25),\n])" + assert str(row_filter3) == expected + # test nested + row_filter4 = RowFilterChain(filters=[row_filter3]) + expected = "RowFilterChain([\n RowFilterChain([\n StripValueTransformerFilter(flag=True),\n RowSampleFilter(sample=0.25),\n ]),\n])" + assert str(row_filter4) == expected + + +def test_row_filter_union_to_pb(): + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_pb = row_filter1._to_pb() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_pb = row_filter2._to_pb() + + row_filter3 = RowFilterUnion(filters=[row_filter1, row_filter2]) + filter_pb = row_filter3._to_pb() + + expected_pb = _RowFilterPB( + interleave=_RowFilterInterleavePB(filters=[row_filter1_pb, row_filter2_pb]) + ) + assert filter_pb == expected_pb + + +def test_row_filter_union_to_dict(): + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_dict = row_filter1.to_dict() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_dict = row_filter2.to_dict() + + row_filter3 = RowFilterUnion(filters=[row_filter1, row_filter2]) + filter_dict = row_filter3.to_dict() + + expected_dict = {"interleave": {"filters": [row_filter1_dict, row_filter2_dict]}} + assert filter_dict == expected_dict + expected_pb_value = row_filter3._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_row_filter_union_to_pb_nested(): + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter2 = RowSampleFilter(0.25) + + row_filter3 = RowFilterUnion(filters=[row_filter1, row_filter2]) + row_filter3_pb = row_filter3._to_pb() + + row_filter4 = CellsRowLimitFilter(11) + row_filter4_pb = row_filter4._to_pb() + + row_filter5 = RowFilterUnion(filters=[row_filter3, row_filter4]) + filter_pb = row_filter5._to_pb() + + expected_pb = _RowFilterPB( + interleave=_RowFilterInterleavePB(filters=[row_filter3_pb, row_filter4_pb]) + ) + assert filter_pb == expected_pb + + +def test_row_filter_union_to_dict_nested(): + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter1 = StripValueTransformerFilter(True) + + row_filter2 = RowSampleFilter(0.25) + + row_filter3 = RowFilterUnion(filters=[row_filter1, row_filter2]) + row_filter3_dict = row_filter3.to_dict() + + row_filter4 = CellsRowLimitFilter(11) + row_filter4_dict = row_filter4.to_dict() + + row_filter5 = RowFilterUnion(filters=[row_filter3, row_filter4]) + filter_dict = row_filter5.to_dict() + + expected_dict = {"interleave": {"filters": [row_filter3_dict, row_filter4_dict]}} + assert filter_dict == expected_dict + expected_pb_value = row_filter5._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_row_filter_union___repr__(): + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter2 = RowSampleFilter(0.25) + + row_filter3 = RowFilterUnion(filters=[row_filter1, row_filter2]) + expected = "RowFilterUnion(filters=[StripValueTransformerFilter(flag=True), RowSampleFilter(sample=0.25)])" + assert repr(row_filter3) == expected + assert eval(repr(row_filter3)) == row_filter3 + + +def test_row_filter_union___str__(): + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter2 = RowSampleFilter(0.25) + + row_filter3 = RowFilterUnion(filters=[row_filter1, row_filter2]) + expected = "RowFilterUnion([\n StripValueTransformerFilter(flag=True),\n RowSampleFilter(sample=0.25),\n])" + assert str(row_filter3) == expected + # test nested + row_filter4 = RowFilterUnion(filters=[row_filter3]) + expected = "RowFilterUnion([\n RowFilterUnion([\n StripValueTransformerFilter(flag=True),\n RowSampleFilter(sample=0.25),\n ]),\n])" + assert str(row_filter4) == expected + + +def test_conditional_row_filter_constructor(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + + predicate_filter = object() + true_filter = object() + false_filter = object() + cond_filter = ConditionalRowFilter( + predicate_filter, true_filter=true_filter, false_filter=false_filter + ) + assert cond_filter.predicate_filter is predicate_filter + assert cond_filter.true_filter is true_filter + assert cond_filter.false_filter is false_filter + + +def test_conditional_row_filter___eq__(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + + predicate_filter = object() + true_filter = object() + false_filter = object() + cond_filter1 = ConditionalRowFilter( + predicate_filter, true_filter=true_filter, false_filter=false_filter + ) + cond_filter2 = ConditionalRowFilter( + predicate_filter, true_filter=true_filter, false_filter=false_filter + ) + assert cond_filter1 == cond_filter2 + + +def test_conditional_row_filter___eq__type_differ(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + + predicate_filter = object() + true_filter = object() + false_filter = object() + cond_filter1 = ConditionalRowFilter( + predicate_filter, true_filter=true_filter, false_filter=false_filter + ) + cond_filter2 = object() + assert not (cond_filter1 == cond_filter2) + + +def test_conditional_row_filter___ne__(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + + predicate_filter = object() + other_predicate_filter = object() + true_filter = object() + false_filter = object() + cond_filter1 = ConditionalRowFilter( + predicate_filter, true_filter=true_filter, false_filter=false_filter + ) + cond_filter2 = ConditionalRowFilter( + other_predicate_filter, true_filter=true_filter, false_filter=false_filter + ) + assert cond_filter1 != cond_filter2 + + +def test_conditional_row_filter_to_pb(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_pb = row_filter1._to_pb() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_pb = row_filter2._to_pb() + + row_filter3 = CellsRowOffsetFilter(11) + row_filter3_pb = row_filter3._to_pb() + + row_filter4 = ConditionalRowFilter( + row_filter1, true_filter=row_filter2, false_filter=row_filter3 + ) + filter_pb = row_filter4._to_pb() + + expected_pb = _RowFilterPB( + condition=_RowFilterConditionPB( + predicate_filter=row_filter1_pb, + true_filter=row_filter2_pb, + false_filter=row_filter3_pb, + ) + ) + assert filter_pb == expected_pb + + +def test_conditional_row_filter_to_dict(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_dict = row_filter1.to_dict() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_dict = row_filter2.to_dict() + + row_filter3 = CellsRowOffsetFilter(11) + row_filter3_dict = row_filter3.to_dict() + + row_filter4 = ConditionalRowFilter( + row_filter1, true_filter=row_filter2, false_filter=row_filter3 + ) + filter_dict = row_filter4.to_dict() + + expected_dict = { + "condition": { + "predicate_filter": row_filter1_dict, + "true_filter": row_filter2_dict, + "false_filter": row_filter3_dict, + } + } + assert filter_dict == expected_dict + expected_pb_value = row_filter4._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_conditional_row_filter_to_pb_true_only(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_pb = row_filter1._to_pb() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_pb = row_filter2._to_pb() + + row_filter3 = ConditionalRowFilter(row_filter1, true_filter=row_filter2) + filter_pb = row_filter3._to_pb() + + expected_pb = _RowFilterPB( + condition=_RowFilterConditionPB( + predicate_filter=row_filter1_pb, true_filter=row_filter2_pb + ) + ) + assert filter_pb == expected_pb + + +def test_conditional_row_filter_to_dict_true_only(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_dict = row_filter1.to_dict() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_dict = row_filter2.to_dict() + + row_filter3 = ConditionalRowFilter(row_filter1, true_filter=row_filter2) + filter_dict = row_filter3.to_dict() + + expected_dict = { + "condition": { + "predicate_filter": row_filter1_dict, + "true_filter": row_filter2_dict, + } + } + assert filter_dict == expected_dict + expected_pb_value = row_filter3._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_conditional_row_filter_to_pb_false_only(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_pb = row_filter1._to_pb() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_pb = row_filter2._to_pb() + + row_filter3 = ConditionalRowFilter(row_filter1, false_filter=row_filter2) + filter_pb = row_filter3._to_pb() + + expected_pb = _RowFilterPB( + condition=_RowFilterConditionPB( + predicate_filter=row_filter1_pb, false_filter=row_filter2_pb + ) + ) + assert filter_pb == expected_pb + + +def test_conditional_row_filter_to_dict_false_only(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + row_filter1 = StripValueTransformerFilter(True) + row_filter1_dict = row_filter1.to_dict() + + row_filter2 = RowSampleFilter(0.25) + row_filter2_dict = row_filter2.to_dict() + + row_filter3 = ConditionalRowFilter(row_filter1, false_filter=row_filter2) + filter_dict = row_filter3.to_dict() + + expected_dict = { + "condition": { + "predicate_filter": row_filter1_dict, + "false_filter": row_filter2_dict, + } + } + assert filter_dict == expected_dict + expected_pb_value = row_filter3._to_pb() + assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + + +def test_conditional_row_filter___repr__(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter2 = RowSampleFilter(0.25) + row_filter3 = ConditionalRowFilter(row_filter1, true_filter=row_filter2) + expected = ( + "ConditionalRowFilter(predicate_filter=StripValueTransformerFilter(" + "flag=True), true_filter=RowSampleFilter(sample=0.25), false_filter=None)" + ) + assert repr(row_filter3) == expected + assert eval(repr(row_filter3)) == row_filter3 + # test nested + row_filter4 = ConditionalRowFilter(row_filter3, true_filter=row_filter2) + expected = "ConditionalRowFilter(predicate_filter=ConditionalRowFilter(predicate_filter=StripValueTransformerFilter(flag=True), true_filter=RowSampleFilter(sample=0.25), false_filter=None), true_filter=RowSampleFilter(sample=0.25), false_filter=None)" + assert repr(row_filter4) == expected + assert eval(repr(row_filter4)) == row_filter4 + + +def test_conditional_row_filter___str__(): + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import StripValueTransformerFilter + + row_filter1 = StripValueTransformerFilter(True) + row_filter2 = RowSampleFilter(0.25) + row_filter3 = ConditionalRowFilter(row_filter1, true_filter=row_filter2) + expected = "ConditionalRowFilter(\n predicate_filter=StripValueTransformerFilter(flag=True),\n true_filter=RowSampleFilter(sample=0.25),\n)" + assert str(row_filter3) == expected + # test nested + row_filter4 = ConditionalRowFilter( + row_filter3, + true_filter=row_filter2, + false_filter=RowFilterUnion([row_filter1, row_filter2]), + ) + expected = "ConditionalRowFilter(\n predicate_filter=ConditionalRowFilter(\n predicate_filter=StripValueTransformerFilter(flag=True),\n true_filter=RowSampleFilter(sample=0.25),\n ),\n true_filter=RowSampleFilter(sample=0.25),\n false_filter=RowFilterUnion([\n StripValueTransformerFilter(flag=True),\n RowSampleFilter(sample=0.25),\n ]),\n)" + assert str(row_filter4) == expected + + +def _ColumnRangePB(*args, **kw): + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + return data_v2_pb2.ColumnRange(*args, **kw) + + +def _RowFilterPB(*args, **kw): + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + return data_v2_pb2.RowFilter(*args, **kw) + + +def _RowFilterChainPB(*args, **kw): + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + return data_v2_pb2.RowFilter.Chain(*args, **kw) + + +def _RowFilterConditionPB(*args, **kw): + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + return data_v2_pb2.RowFilter.Condition(*args, **kw) + + +def _RowFilterInterleavePB(*args, **kw): + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + return data_v2_pb2.RowFilter.Interleave(*args, **kw) + + +def _TimestampRangePB(*args, **kw): + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + return data_v2_pb2.TimestampRange(*args, **kw) + + +def _ValueRangePB(*args, **kw): + from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + return data_v2_pb2.ValueRange(*args, **kw) + + +def _get_regex_filters(): + from google.cloud.bigtable.row_filters import ( + RowKeyRegexFilter, + FamilyNameRegexFilter, + ColumnQualifierRegexFilter, + ValueRegexFilter, + ExactValueFilter, + ) + + return [ + RowKeyRegexFilter, + FamilyNameRegexFilter, + ColumnQualifierRegexFilter, + ValueRegexFilter, + ExactValueFilter, + ] + + +def _get_bool_filters(): + from google.cloud.bigtable.row_filters import ( + SinkFilter, + PassAllFilter, + BlockAllFilter, + StripValueTransformerFilter, + ) + + return [ + SinkFilter, + PassAllFilter, + BlockAllFilter, + StripValueTransformerFilter, + ] + + +def _get_cell_count_filters(): + from google.cloud.bigtable.row_filters import ( + CellsRowLimitFilter, + CellsRowOffsetFilter, + CellsColumnLimitFilter, + ) + + return [ + CellsRowLimitFilter, + CellsRowOffsetFilter, + CellsColumnLimitFilter, + ] + + +def _get_filter_combination_filters(): + from google.cloud.bigtable.row_filters import ( + RowFilterChain, + RowFilterUnion, + ) + + return [ + RowFilterChain, + RowFilterUnion, + ] From 71b03128a5075d044df9120d3ee9156a00bb74e4 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 3 Apr 2023 14:10:26 -0700 Subject: [PATCH 03/56] feat: read rows query model class (#752) --- google/cloud/bigtable/__init__.py | 2 + google/cloud/bigtable/read_rows_query.py | 206 ++++++++++++- tests/unit/test_read_rows_query.py | 359 +++++++++++++++++++++++ 3 files changed, 554 insertions(+), 13 deletions(-) create mode 100644 tests/unit/test_read_rows_query.py diff --git a/google/cloud/bigtable/__init__.py b/google/cloud/bigtable/__init__.py index daa562c0c..251e41e42 100644 --- a/google/cloud/bigtable/__init__.py +++ b/google/cloud/bigtable/__init__.py @@ -22,6 +22,7 @@ from google.cloud.bigtable.client import Table from google.cloud.bigtable.read_rows_query import ReadRowsQuery +from google.cloud.bigtable.read_rows_query import RowRange from google.cloud.bigtable.row_response import RowResponse from google.cloud.bigtable.row_response import CellResponse @@ -43,6 +44,7 @@ "Table", "RowKeySamples", "ReadRowsQuery", + "RowRange", "MutationsBatcher", "Mutation", "BulkMutationsEntry", diff --git a/google/cloud/bigtable/read_rows_query.py b/google/cloud/bigtable/read_rows_query.py index 64583b2d7..9fd349d5f 100644 --- a/google/cloud/bigtable/read_rows_query.py +++ b/google/cloud/bigtable/read_rows_query.py @@ -13,36 +13,192 @@ # limitations under the License. # from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any +from .row_response import row_key +from dataclasses import dataclass +from google.cloud.bigtable.row_filters import RowFilter if TYPE_CHECKING: - from google.cloud.bigtable.row_filters import RowFilter from google.cloud.bigtable import RowKeySamples +@dataclass +class _RangePoint: + """Model class for a point in a row range""" + + key: row_key + is_inclusive: bool + + +@dataclass +class RowRange: + start: _RangePoint | None + end: _RangePoint | None + + def __init__( + self, + start_key: str | bytes | None = None, + end_key: str | bytes | None = None, + start_is_inclusive: bool | None = None, + end_is_inclusive: bool | None = None, + ): + # check for invalid combinations of arguments + if start_is_inclusive is None: + start_is_inclusive = True + elif start_key is None: + raise ValueError("start_is_inclusive must be set with start_key") + if end_is_inclusive is None: + end_is_inclusive = False + elif end_key is None: + raise ValueError("end_is_inclusive must be set with end_key") + # ensure that start_key and end_key are bytes + if isinstance(start_key, str): + start_key = start_key.encode() + elif start_key is not None and not isinstance(start_key, bytes): + raise ValueError("start_key must be a string or bytes") + if isinstance(end_key, str): + end_key = end_key.encode() + elif end_key is not None and not isinstance(end_key, bytes): + raise ValueError("end_key must be a string or bytes") + + self.start = ( + _RangePoint(start_key, start_is_inclusive) + if start_key is not None + else None + ) + self.end = ( + _RangePoint(end_key, end_is_inclusive) if end_key is not None else None + ) + + def _to_dict(self) -> dict[str, bytes]: + """Converts this object to a dictionary""" + output = {} + if self.start is not None: + key = "start_key_closed" if self.start.is_inclusive else "start_key_open" + output[key] = self.start.key + if self.end is not None: + key = "end_key_closed" if self.end.is_inclusive else "end_key_open" + output[key] = self.end.key + return output + + class ReadRowsQuery: """ Class to encapsulate details of a read row request """ def __init__( - self, row_keys: list[str | bytes] | str | bytes | None = None, limit=None + self, + row_keys: list[str | bytes] | str | bytes | None = None, + row_ranges: list[RowRange] | RowRange | None = None, + limit: int | None = None, + row_filter: RowFilter | None = None, ): - pass + """ + Create a new ReadRowsQuery - def set_limit(self, limit: int) -> ReadRowsQuery: - raise NotImplementedError + Args: + - row_keys: row keys to include in the query + a query can contain multiple keys, but ranges should be preferred + - row_ranges: ranges of rows to include in the query + - limit: the maximum number of rows to return. None or 0 means no limit + default: None (no limit) + - row_filter: a RowFilter to apply to the query + """ + self.row_keys: set[bytes] = set() + self.row_ranges: list[RowRange | dict[str, bytes]] = [] + if row_ranges: + if isinstance(row_ranges, RowRange): + row_ranges = [row_ranges] + for r in row_ranges: + self.add_range(r) + if row_keys: + if not isinstance(row_keys, list): + row_keys = [row_keys] + for k in row_keys: + self.add_key(k) + self.limit: int | None = limit + self.filter: RowFilter | dict[str, Any] | None = row_filter - def set_filter(self, filter: "RowFilter") -> ReadRowsQuery: - raise NotImplementedError + @property + def limit(self) -> int | None: + return self._limit - def add_rows(self, row_id_list: list[str]) -> ReadRowsQuery: - raise NotImplementedError + @limit.setter + def limit(self, new_limit: int | None): + """ + Set the maximum number of rows to return by this query. + + None or 0 means no limit + + Args: + - new_limit: the new limit to apply to this query + Returns: + - a reference to this query for chaining + Raises: + - ValueError if new_limit is < 0 + """ + if new_limit is not None and new_limit < 0: + raise ValueError("limit must be >= 0") + self._limit = new_limit + + @property + def filter(self) -> RowFilter | dict[str, Any] | None: + return self._filter + + @filter.setter + def filter(self, row_filter: RowFilter | dict[str, Any] | None): + """ + Set a RowFilter to apply to this query + + Args: + - row_filter: a RowFilter to apply to this query + Can be a RowFilter object or a dict representation + Returns: + - a reference to this query for chaining + """ + if not ( + isinstance(row_filter, dict) + or isinstance(row_filter, RowFilter) + or row_filter is None + ): + raise ValueError("row_filter must be a RowFilter or dict") + self._filter = row_filter + + def add_key(self, row_key: str | bytes): + """ + Add a row key to this query + + A query can contain multiple keys, but ranges should be preferred + + Args: + - row_key: a key to add to this query + Returns: + - a reference to this query for chaining + Raises: + - ValueError if an input is not a string or bytes + """ + if isinstance(row_key, str): + row_key = row_key.encode() + elif not isinstance(row_key, bytes): + raise ValueError("row_key must be string or bytes") + self.row_keys.add(row_key) def add_range( - self, start_key: str | bytes | None = None, end_key: str | bytes | None = None - ) -> ReadRowsQuery: - raise NotImplementedError + self, + row_range: RowRange | dict[str, bytes], + ): + """ + Add a range of row keys to this query. + + Args: + - row_range: a range of row keys to add to this query + Can be a RowRange object or a dict representation in + RowRange proto format + """ + if not (isinstance(row_range, dict) or isinstance(row_range, RowRange)): + raise ValueError("row_range must be a RowRange or dict") + self.row_ranges.append(row_range) def shard(self, shard_keys: "RowKeySamples" | None = None) -> list[ReadRowsQuery]: """ @@ -54,3 +210,27 @@ def shard(self, shard_keys: "RowKeySamples" | None = None) -> list[ReadRowsQuery query (if possible) """ raise NotImplementedError + + def _to_dict(self) -> dict[str, Any]: + """ + Convert this query into a dictionary that can be used to construct a + ReadRowsRequest protobuf + """ + row_ranges = [] + for r in self.row_ranges: + dict_range = r._to_dict() if isinstance(r, RowRange) else r + row_ranges.append(dict_range) + row_keys = list(self.row_keys) + row_keys.sort() + row_set = {"row_keys": row_keys, "row_ranges": row_ranges} + final_dict: dict[str, Any] = { + "rows": row_set, + } + dict_filter = ( + self.filter.to_dict() if isinstance(self.filter, RowFilter) else self.filter + ) + if dict_filter: + final_dict["filter"] = dict_filter + if self.limit is not None: + final_dict["rows_limit"] = self.limit + return final_dict diff --git a/tests/unit/test_read_rows_query.py b/tests/unit/test_read_rows_query.py new file mode 100644 index 000000000..aa690bc86 --- /dev/null +++ b/tests/unit/test_read_rows_query.py @@ -0,0 +1,359 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +import unittest + +TEST_ROWS = [ + "row_key_1", + b"row_key_2", +] + + +class TestRowRange(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigtable.read_rows_query import RowRange + + return RowRange + + def _make_one(self, *args, **kwargs): + return self._get_target_class()(*args, **kwargs) + + def test_ctor_start_end(self): + row_range = self._make_one("test_row", "test_row2") + self.assertEqual(row_range.start.key, "test_row".encode()) + self.assertEqual(row_range.end.key, "test_row2".encode()) + self.assertEqual(row_range.start.is_inclusive, True) + self.assertEqual(row_range.end.is_inclusive, False) + + def test_ctor_start_only(self): + row_range = self._make_one("test_row3") + self.assertEqual(row_range.start.key, "test_row3".encode()) + self.assertEqual(row_range.start.is_inclusive, True) + self.assertEqual(row_range.end, None) + + def test_ctor_end_only(self): + row_range = self._make_one(end_key="test_row4") + self.assertEqual(row_range.end.key, "test_row4".encode()) + self.assertEqual(row_range.end.is_inclusive, False) + self.assertEqual(row_range.start, None) + + def test_ctor_inclusive_flags(self): + row_range = self._make_one("test_row5", "test_row6", False, True) + self.assertEqual(row_range.start.key, "test_row5".encode()) + self.assertEqual(row_range.end.key, "test_row6".encode()) + self.assertEqual(row_range.start.is_inclusive, False) + self.assertEqual(row_range.end.is_inclusive, True) + + def test_ctor_defaults(self): + row_range = self._make_one() + self.assertEqual(row_range.start, None) + self.assertEqual(row_range.end, None) + + def test_ctor_flags_only(self): + with self.assertRaises(ValueError) as exc: + self._make_one(start_is_inclusive=True, end_is_inclusive=True) + self.assertEqual( + exc.exception.args, + ("start_is_inclusive must be set with start_key",), + ) + with self.assertRaises(ValueError) as exc: + self._make_one(start_is_inclusive=False, end_is_inclusive=False) + self.assertEqual( + exc.exception.args, + ("start_is_inclusive must be set with start_key",), + ) + with self.assertRaises(ValueError) as exc: + self._make_one(start_is_inclusive=False) + self.assertEqual( + exc.exception.args, + ("start_is_inclusive must be set with start_key",), + ) + with self.assertRaises(ValueError) as exc: + self._make_one(end_is_inclusive=True) + self.assertEqual( + exc.exception.args, ("end_is_inclusive must be set with end_key",) + ) + + def test_ctor_invalid_keys(self): + # test with invalid keys + with self.assertRaises(ValueError) as exc: + self._make_one(1, "2") + self.assertEqual(exc.exception.args, ("start_key must be a string or bytes",)) + with self.assertRaises(ValueError) as exc: + self._make_one("1", 2) + self.assertEqual(exc.exception.args, ("end_key must be a string or bytes",)) + + def test__to_dict_defaults(self): + row_range = self._make_one("test_row", "test_row2") + expected = { + "start_key_closed": b"test_row", + "end_key_open": b"test_row2", + } + self.assertEqual(row_range._to_dict(), expected) + + def test__to_dict_inclusive_flags(self): + row_range = self._make_one("test_row", "test_row2", False, True) + expected = { + "start_key_open": b"test_row", + "end_key_closed": b"test_row2", + } + self.assertEqual(row_range._to_dict(), expected) + + +class TestReadRowsQuery(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + + return ReadRowsQuery + + def _make_one(self, *args, **kwargs): + return self._get_target_class()(*args, **kwargs) + + def test_ctor_defaults(self): + query = self._make_one() + self.assertEqual(query.row_keys, set()) + self.assertEqual(query.row_ranges, []) + self.assertEqual(query.filter, None) + self.assertEqual(query.limit, None) + + def test_ctor_explicit(self): + from google.cloud.bigtable.row_filters import RowFilterChain + + filter_ = RowFilterChain() + query = self._make_one(["row_key_1", "row_key_2"], limit=10, row_filter=filter_) + self.assertEqual(len(query.row_keys), 2) + self.assertIn("row_key_1".encode(), query.row_keys) + self.assertIn("row_key_2".encode(), query.row_keys) + self.assertEqual(query.row_ranges, []) + self.assertEqual(query.filter, filter_) + self.assertEqual(query.limit, 10) + + def test_ctor_invalid_limit(self): + with self.assertRaises(ValueError) as exc: + self._make_one(limit=-1) + self.assertEqual(exc.exception.args, ("limit must be >= 0",)) + + def test_set_filter(self): + from google.cloud.bigtable.row_filters import RowFilterChain + + filter1 = RowFilterChain() + query = self._make_one() + self.assertEqual(query.filter, None) + query.filter = filter1 + self.assertEqual(query.filter, filter1) + filter2 = RowFilterChain() + query.filter = filter2 + self.assertEqual(query.filter, filter2) + query.filter = None + self.assertEqual(query.filter, None) + query.filter = RowFilterChain() + self.assertEqual(query.filter, RowFilterChain()) + with self.assertRaises(ValueError) as exc: + query.filter = 1 + self.assertEqual( + exc.exception.args, ("row_filter must be a RowFilter or dict",) + ) + + def test_set_filter_dict(self): + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable_v2.types.bigtable import ReadRowsRequest + + filter1 = RowSampleFilter(0.5) + filter1_dict = filter1.to_dict() + query = self._make_one() + self.assertEqual(query.filter, None) + query.filter = filter1_dict + self.assertEqual(query.filter, filter1_dict) + output = query._to_dict() + self.assertEqual(output["filter"], filter1_dict) + proto_output = ReadRowsRequest(**output) + self.assertEqual(proto_output.filter, filter1._to_pb()) + + query.filter = None + self.assertEqual(query.filter, None) + + def test_set_limit(self): + query = self._make_one() + self.assertEqual(query.limit, None) + query.limit = 10 + self.assertEqual(query.limit, 10) + query.limit = 9 + self.assertEqual(query.limit, 9) + query.limit = 0 + self.assertEqual(query.limit, 0) + with self.assertRaises(ValueError) as exc: + query.limit = -1 + self.assertEqual(exc.exception.args, ("limit must be >= 0",)) + with self.assertRaises(ValueError) as exc: + query.limit = -100 + self.assertEqual(exc.exception.args, ("limit must be >= 0",)) + + def test_add_key_str(self): + query = self._make_one() + self.assertEqual(query.row_keys, set()) + input_str = "test_row" + query.add_key(input_str) + self.assertEqual(len(query.row_keys), 1) + self.assertIn(input_str.encode(), query.row_keys) + input_str2 = "test_row2" + query.add_key(input_str2) + self.assertEqual(len(query.row_keys), 2) + self.assertIn(input_str.encode(), query.row_keys) + self.assertIn(input_str2.encode(), query.row_keys) + + def test_add_key_bytes(self): + query = self._make_one() + self.assertEqual(query.row_keys, set()) + input_bytes = b"test_row" + query.add_key(input_bytes) + self.assertEqual(len(query.row_keys), 1) + self.assertIn(input_bytes, query.row_keys) + input_bytes2 = b"test_row2" + query.add_key(input_bytes2) + self.assertEqual(len(query.row_keys), 2) + self.assertIn(input_bytes, query.row_keys) + self.assertIn(input_bytes2, query.row_keys) + + def test_add_rows_batch(self): + query = self._make_one() + self.assertEqual(query.row_keys, set()) + input_batch = ["test_row", b"test_row2", "test_row3"] + for k in input_batch: + query.add_key(k) + self.assertEqual(len(query.row_keys), 3) + self.assertIn(b"test_row", query.row_keys) + self.assertIn(b"test_row2", query.row_keys) + self.assertIn(b"test_row3", query.row_keys) + # test adding another batch + for k in ["test_row4", b"test_row5"]: + query.add_key(k) + self.assertEqual(len(query.row_keys), 5) + self.assertIn(input_batch[0].encode(), query.row_keys) + self.assertIn(input_batch[1], query.row_keys) + self.assertIn(input_batch[2].encode(), query.row_keys) + self.assertIn(b"test_row4", query.row_keys) + self.assertIn(b"test_row5", query.row_keys) + + def test_add_key_invalid(self): + query = self._make_one() + with self.assertRaises(ValueError) as exc: + query.add_key(1) + self.assertEqual(exc.exception.args, ("row_key must be string or bytes",)) + with self.assertRaises(ValueError) as exc: + query.add_key(["s"]) + self.assertEqual(exc.exception.args, ("row_key must be string or bytes",)) + + def test_duplicate_rows(self): + # should only hold one of each input key + key_1 = b"test_row" + key_2 = b"test_row2" + query = self._make_one(row_keys=[key_1, key_1, key_2]) + self.assertEqual(len(query.row_keys), 2) + self.assertIn(key_1, query.row_keys) + self.assertIn(key_2, query.row_keys) + key_3 = "test_row3" + for i in range(10): + query.add_key(key_3) + self.assertEqual(len(query.row_keys), 3) + + def test_add_range(self): + from google.cloud.bigtable.read_rows_query import RowRange + + query = self._make_one() + self.assertEqual(query.row_ranges, []) + input_range = RowRange(start_key=b"test_row") + query.add_range(input_range) + self.assertEqual(len(query.row_ranges), 1) + self.assertEqual(query.row_ranges[0], input_range) + input_range2 = RowRange(start_key=b"test_row2") + query.add_range(input_range2) + self.assertEqual(len(query.row_ranges), 2) + self.assertEqual(query.row_ranges[0], input_range) + self.assertEqual(query.row_ranges[1], input_range2) + + def test_add_range_dict(self): + query = self._make_one() + self.assertEqual(query.row_ranges, []) + input_range = {"start_key_closed": b"test_row"} + query.add_range(input_range) + self.assertEqual(len(query.row_ranges), 1) + self.assertEqual(query.row_ranges[0], input_range) + + def test_to_dict_rows_default(self): + # dictionary should be in rowset proto format + from google.cloud.bigtable_v2.types.bigtable import ReadRowsRequest + + query = self._make_one() + output = query._to_dict() + self.assertTrue(isinstance(output, dict)) + self.assertEqual(len(output.keys()), 1) + expected = {"rows": {"row_keys": [], "row_ranges": []}} + self.assertEqual(output, expected) + + request_proto = ReadRowsRequest(**output) + self.assertEqual(request_proto.rows.row_keys, []) + self.assertEqual(request_proto.rows.row_ranges, []) + self.assertFalse(request_proto.filter) + self.assertEqual(request_proto.rows_limit, 0) + + def test_to_dict_rows_populated(self): + # dictionary should be in rowset proto format + from google.cloud.bigtable_v2.types.bigtable import ReadRowsRequest + from google.cloud.bigtable.row_filters import PassAllFilter + from google.cloud.bigtable.read_rows_query import RowRange + + row_filter = PassAllFilter(False) + query = self._make_one(limit=100, row_filter=row_filter) + query.add_range(RowRange("test_row", "test_row2")) + query.add_range(RowRange("test_row3")) + query.add_range(RowRange(start_key=None, end_key="test_row5")) + query.add_range(RowRange(b"test_row6", b"test_row7", False, True)) + query.add_range(RowRange()) + query.add_key("test_row") + query.add_key(b"test_row2") + query.add_key("test_row3") + query.add_key(b"test_row3") + query.add_key(b"test_row4") + output = query._to_dict() + self.assertTrue(isinstance(output, dict)) + request_proto = ReadRowsRequest(**output) + rowset_proto = request_proto.rows + # check rows + self.assertEqual(len(rowset_proto.row_keys), 4) + self.assertEqual(rowset_proto.row_keys[0], b"test_row") + self.assertEqual(rowset_proto.row_keys[1], b"test_row2") + self.assertEqual(rowset_proto.row_keys[2], b"test_row3") + self.assertEqual(rowset_proto.row_keys[3], b"test_row4") + # check ranges + self.assertEqual(len(rowset_proto.row_ranges), 5) + self.assertEqual(rowset_proto.row_ranges[0].start_key_closed, b"test_row") + self.assertEqual(rowset_proto.row_ranges[0].end_key_open, b"test_row2") + self.assertEqual(rowset_proto.row_ranges[1].start_key_closed, b"test_row3") + self.assertEqual(rowset_proto.row_ranges[1].end_key_open, b"") + self.assertEqual(rowset_proto.row_ranges[2].start_key_closed, b"") + self.assertEqual(rowset_proto.row_ranges[2].end_key_open, b"test_row5") + self.assertEqual(rowset_proto.row_ranges[3].start_key_open, b"test_row6") + self.assertEqual(rowset_proto.row_ranges[3].end_key_closed, b"test_row7") + self.assertEqual(rowset_proto.row_ranges[4].start_key_closed, b"") + self.assertEqual(rowset_proto.row_ranges[4].end_key_open, b"") + # check limit + self.assertEqual(request_proto.rows_limit, 100) + # check filter + filter_proto = request_proto.filter + self.assertEqual(filter_proto, row_filter._to_pb()) + + def test_shard(self): + pass From c55099fd1a4274d66b165bc5793f25eaf7ee6e64 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 5 Apr 2023 14:11:47 -0700 Subject: [PATCH 04/56] feat: implement row and cell model classes (#753) --- google/cloud/bigtable/__init__.py | 8 +- google/cloud/bigtable/client.py | 14 +- google/cloud/bigtable/mutations.py | 2 +- google/cloud/bigtable/mutations_batcher.py | 2 +- .../cloud/bigtable/read_modify_write_rules.py | 2 +- google/cloud/bigtable/read_rows_query.py | 2 +- google/cloud/bigtable/row.py | 413 +++++++++++ google/cloud/bigtable/row_response.py | 130 ---- tests/unit/test_row.py | 692 ++++++++++++++++++ 9 files changed, 1120 insertions(+), 145 deletions(-) create mode 100644 google/cloud/bigtable/row.py delete mode 100644 google/cloud/bigtable/row_response.py create mode 100644 tests/unit/test_row.py diff --git a/google/cloud/bigtable/__init__.py b/google/cloud/bigtable/__init__.py index 251e41e42..c5581f813 100644 --- a/google/cloud/bigtable/__init__.py +++ b/google/cloud/bigtable/__init__.py @@ -22,9 +22,9 @@ from google.cloud.bigtable.client import Table from google.cloud.bigtable.read_rows_query import ReadRowsQuery +from google.cloud.bigtable.row import Row +from google.cloud.bigtable.row import Cell from google.cloud.bigtable.read_rows_query import RowRange -from google.cloud.bigtable.row_response import RowResponse -from google.cloud.bigtable.row_response import CellResponse from google.cloud.bigtable.mutations_batcher import MutationsBatcher from google.cloud.bigtable.mutations import Mutation @@ -52,6 +52,6 @@ "DeleteRangeFromColumn", "DeleteAllFromFamily", "DeleteAllFromRow", - "RowResponse", - "CellResponse", + "Row", + "Cell", ) diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index df4bf308f..23f0cb6fe 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -25,7 +25,7 @@ if TYPE_CHECKING: from google.cloud.bigtable.mutations import Mutation, BulkMutationsEntry from google.cloud.bigtable.mutations_batcher import MutationsBatcher - from google.cloud.bigtable.row_response import RowResponse + from google.cloud.bigtable.row import Row from google.cloud.bigtable.read_rows_query import ReadRowsQuery from google.cloud.bigtable import RowKeySamples from google.cloud.bigtable.row_filters import RowFilter @@ -109,7 +109,7 @@ async def read_rows_stream( idle_timeout: int | float | None = 300, per_request_timeout: int | float | None = None, metadata: list[tuple[str, str]] | None = None, - ) -> AsyncIterable[RowResponse]: + ) -> AsyncIterable[Row]: """ Returns a generator to asynchronously stream back row data. @@ -166,7 +166,7 @@ async def read_rows( per_row_timeout: int | float | None = 10, per_request_timeout: int | float | None = None, metadata: list[tuple[str, str]] | None = None, - ) -> list[RowResponse]: + ) -> list[Row]: """ Helper function that returns a full list instead of a generator @@ -184,7 +184,7 @@ async def read_row( operation_timeout: int | float | None = 60, per_request_timeout: int | float | None = None, metadata: list[tuple[str, str]] | None = None, - ) -> RowResponse: + ) -> Row: """ Helper function to return a single row @@ -206,7 +206,7 @@ async def read_rows_sharded( idle_timeout: int | float | None = 300, per_request_timeout: int | float | None = None, metadata: list[tuple[str, str]] | None = None, - ) -> AsyncIterable[RowResponse]: + ) -> AsyncIterable[Row]: """ Runs a sharded query in parallel @@ -410,7 +410,7 @@ async def read_modify_write_row( *, operation_timeout: int | float | None = 60, metadata: list[tuple[str, str]] | None = None, - ) -> RowResponse: + ) -> Row: """ Reads and modifies a row atomically according to input ReadModifyWriteRules, and returns the contents of all modified cells @@ -429,7 +429,7 @@ async def read_modify_write_row( Failed requests will not be retried. - metadata: Strings which should be sent along with the request as metadata headers. Returns: - - RowResponse: containing cell data that was modified as part of the + - Row: containing cell data that was modified as part of the operation Raises: - GoogleAPIError exceptions from grpc call diff --git a/google/cloud/bigtable/mutations.py b/google/cloud/bigtable/mutations.py index ed3c2f065..4ff59bff9 100644 --- a/google/cloud/bigtable/mutations.py +++ b/google/cloud/bigtable/mutations.py @@ -15,7 +15,7 @@ from __future__ import annotations from dataclasses import dataclass -from google.cloud.bigtable.row_response import family_id, qualifier, row_key +from google.cloud.bigtable.row import family_id, qualifier, row_key class Mutation: diff --git a/google/cloud/bigtable/mutations_batcher.py b/google/cloud/bigtable/mutations_batcher.py index 2e393cc7e..582786ee4 100644 --- a/google/cloud/bigtable/mutations_batcher.py +++ b/google/cloud/bigtable/mutations_batcher.py @@ -18,7 +18,7 @@ from typing import TYPE_CHECKING from google.cloud.bigtable.mutations import Mutation -from google.cloud.bigtable.row_response import row_key +from google.cloud.bigtable.row import row_key from google.cloud.bigtable.row_filters import RowFilter if TYPE_CHECKING: diff --git a/google/cloud/bigtable/read_modify_write_rules.py b/google/cloud/bigtable/read_modify_write_rules.py index a9b0885f2..839262ea2 100644 --- a/google/cloud/bigtable/read_modify_write_rules.py +++ b/google/cloud/bigtable/read_modify_write_rules.py @@ -16,7 +16,7 @@ from dataclasses import dataclass -from google.cloud.bigtable.row_response import family_id, qualifier +from google.cloud.bigtable.row import family_id, qualifier class ReadModifyWriteRule: diff --git a/google/cloud/bigtable/read_rows_query.py b/google/cloud/bigtable/read_rows_query.py index 9fd349d5f..559b47f04 100644 --- a/google/cloud/bigtable/read_rows_query.py +++ b/google/cloud/bigtable/read_rows_query.py @@ -14,7 +14,7 @@ # from __future__ import annotations from typing import TYPE_CHECKING, Any -from .row_response import row_key +from .row import row_key from dataclasses import dataclass from google.cloud.bigtable.row_filters import RowFilter diff --git a/google/cloud/bigtable/row.py b/google/cloud/bigtable/row.py new file mode 100644 index 000000000..8231e324e --- /dev/null +++ b/google/cloud/bigtable/row.py @@ -0,0 +1,413 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from collections import OrderedDict +from typing import Sequence, Generator, overload, Any +from functools import total_ordering + +# Type aliases used internally for readability. +row_key = bytes +family_id = str +qualifier = bytes +row_value = bytes + + +class Row(Sequence["Cell"]): + """ + Model class for row data returned from server + + Does not represent all data contained in the row, only data returned by a + query. + Expected to be read-only to users, and written by backend + + Can be indexed: + cells = row["family", "qualifier"] + """ + + def __init__( + self, + key: row_key, + cells: list[Cell], + ): + """ + Initializes a Row object + + Row objects are not intended to be created by users. + They are returned by the Bigtable backend. + """ + self.row_key = key + self._cells_map: dict[family_id, dict[qualifier, list[Cell]]] = OrderedDict() + self._cells_list: list[Cell] = [] + # add cells to internal stores using Bigtable native ordering + for cell in cells: + if cell.family not in self._cells_map: + self._cells_map[cell.family] = OrderedDict() + if cell.column_qualifier not in self._cells_map[cell.family]: + self._cells_map[cell.family][cell.column_qualifier] = [] + self._cells_map[cell.family][cell.column_qualifier].append(cell) + self._cells_list.append(cell) + + def get_cells( + self, family: str | None = None, qualifier: str | bytes | None = None + ) -> list[Cell]: + """ + Returns cells sorted in Bigtable native order: + - Family lexicographically ascending + - Qualifier ascending + - Timestamp in reverse chronological order + + If family or qualifier not passed, will include all + + Can also be accessed through indexing: + cells = row["family", "qualifier"] + cells = row["family"] + """ + if family is None: + if qualifier is not None: + # get_cells(None, "qualifier") is not allowed + raise ValueError("Qualifier passed without family") + else: + # return all cells on get_cells() + return self._cells_list + if qualifier is None: + # return all cells in family on get_cells(family) + return list(self._get_all_from_family(family)) + if isinstance(qualifier, str): + qualifier = qualifier.encode("utf-8") + # return cells in family and qualifier on get_cells(family, qualifier) + if family not in self._cells_map: + raise ValueError(f"Family '{family}' not found in row '{self.row_key!r}'") + if qualifier not in self._cells_map[family]: + raise ValueError( + f"Qualifier '{qualifier!r}' not found in family '{family}' in row '{self.row_key!r}'" + ) + return self._cells_map[family][qualifier] + + def _get_all_from_family(self, family: family_id) -> Generator[Cell, None, None]: + """ + Returns all cells in the row for the family_id + """ + if family not in self._cells_map: + raise ValueError(f"Family '{family}' not found in row '{self.row_key!r}'") + qualifier_dict = self._cells_map.get(family, {}) + for cell_batch in qualifier_dict.values(): + for cell in cell_batch: + yield cell + + def __str__(self) -> str: + """ + Human-readable string representation + + { + (family='fam', qualifier=b'col'): [b'value', (+1 more),], + (family='fam', qualifier=b'col2'): [b'other'], + } + """ + output = ["{"] + for family, qualifier in self.get_column_components(): + cell_list = self[family, qualifier] + line = [f" (family={family!r}, qualifier={qualifier!r}): "] + if len(cell_list) == 0: + line.append("[],") + elif len(cell_list) == 1: + line.append(f"[{cell_list[0]}],") + else: + line.append(f"[{cell_list[0]}, (+{len(cell_list)-1} more)],") + output.append("".join(line)) + output.append("}") + return "\n".join(output) + + def __repr__(self): + cell_str_buffer = ["{"] + for family, qualifier in self.get_column_components(): + cell_list = self[family, qualifier] + repr_list = [cell.to_dict() for cell in cell_list] + cell_str_buffer.append(f" ('{family}', {qualifier}): {repr_list},") + cell_str_buffer.append("}") + cell_str = "\n".join(cell_str_buffer) + output = f"Row(key={self.row_key!r}, cells={cell_str})" + return output + + def to_dict(self) -> dict[str, Any]: + """ + Returns a dictionary representation of the cell in the Bigtable Row + proto format + + https://cloud.google.com/bigtable/docs/reference/data/rpc/google.bigtable.v2#row + """ + families_list: list[dict[str, Any]] = [] + for family in self._cells_map: + column_list: list[dict[str, Any]] = [] + for qualifier in self._cells_map[family]: + cells_list: list[dict[str, Any]] = [] + for cell in self._cells_map[family][qualifier]: + cells_list.append(cell.to_dict()) + column_list.append({"qualifier": qualifier, "cells": cells_list}) + families_list.append({"name": family, "columns": column_list}) + return {"key": self.row_key, "families": families_list} + + # Sequence and Mapping methods + def __iter__(self): + """ + Allow iterating over all cells in the row + """ + # iterate as a sequence; yield all cells + for cell in self._cells_list: + yield cell + + def __contains__(self, item): + """ + Implements `in` operator + + Works for both cells in the internal list, and `family` or + `(family, qualifier)` pairs associated with the cells + """ + if isinstance(item, family_id): + # check if family key is in Row + return item in self._cells_map + elif ( + isinstance(item, tuple) + and isinstance(item[0], family_id) + and isinstance(item[1], (qualifier, str)) + ): + # check if (family, qualifier) pair is in Row + qualifer = item[1] if isinstance(item[1], bytes) else item[1].encode() + return item[0] in self._cells_map and qualifer in self._cells_map[item[0]] + # check if Cell is in Row + return item in self._cells_list + + @overload + def __getitem__( + self, + index: family_id | tuple[family_id, qualifier | str], + ) -> list[Cell]: + # overload signature for type checking + pass + + @overload + def __getitem__(self, index: int) -> Cell: + # overload signature for type checking + pass + + @overload + def __getitem__(self, index: slice) -> list[Cell]: + # overload signature for type checking + pass + + def __getitem__(self, index): + """ + Implements [] indexing + + Supports indexing by family, (family, qualifier) pair, + numerical index, and index slicing + """ + if isinstance(index, family_id): + return self.get_cells(family=index) + elif ( + isinstance(index, tuple) + and isinstance(index[0], family_id) + and isinstance(index[1], (qualifier, str)) + ): + return self.get_cells(family=index[0], qualifier=index[1]) + elif isinstance(index, int) or isinstance(index, slice): + # index is int or slice + return self._cells_list[index] + else: + raise TypeError( + "Index must be family_id, (family_id, qualifier), int, or slice" + ) + + def __len__(self): + """ + Implements `len()` operator + """ + return len(self._cells_list) + + def get_column_components(self): + """ + Returns a list of (family, qualifier) pairs associated with the cells + + Pairs can be used for indexing + """ + key_list = [] + for family in self._cells_map: + for qualifier in self._cells_map[family]: + key_list.append((family, qualifier)) + return key_list + + def __eq__(self, other): + """ + Implements `==` operator + """ + # for performance reasons, check row metadata + # before checking individual cells + if not isinstance(other, Row): + return False + if self.row_key != other.row_key: + return False + if len(self._cells_list) != len(other._cells_list): + return False + components = self.get_column_components() + other_components = other.get_column_components() + if len(components) != len(other_components): + return False + if components != other_components: + return False + for family, qualifier in components: + if len(self[family, qualifier]) != len(other[family, qualifier]): + return False + # compare individual cell lists + if self._cells_list != other._cells_list: + return False + return True + + def __ne__(self, other) -> bool: + """ + Implements `!=` operator + """ + return not self == other + + +@total_ordering +class Cell: + """ + Model class for cell data + + Does not represent all data contained in the cell, only data returned by a + query. + Expected to be read-only to users, and written by backend + """ + + def __init__( + self, + value: row_value, + row: row_key, + family: family_id, + column_qualifier: qualifier | str, + timestamp_micros: int, + labels: list[str] | None = None, + ): + """ + Cell constructor + + Cell objects are not intended to be constructed by users. + They are returned by the Bigtable backend. + """ + self.value = value + self.row_key = row + self.family = family + if isinstance(column_qualifier, str): + column_qualifier = column_qualifier.encode() + self.column_qualifier = column_qualifier + self.timestamp_micros = timestamp_micros + self.labels = labels if labels is not None else [] + + def __int__(self) -> int: + """ + Allows casting cell to int + Interprets value as a 64-bit big-endian signed integer, as expected by + ReadModifyWrite increment rule + """ + return int.from_bytes(self.value, byteorder="big", signed=True) + + def to_dict(self) -> dict[str, Any]: + """ + Returns a dictionary representation of the cell in the Bigtable Cell + proto format + + https://cloud.google.com/bigtable/docs/reference/data/rpc/google.bigtable.v2#cell + """ + cell_dict: dict[str, Any] = { + "value": self.value, + } + cell_dict["timestamp_micros"] = self.timestamp_micros + if self.labels: + cell_dict["labels"] = self.labels + return cell_dict + + def __str__(self) -> str: + """ + Allows casting cell to str + Prints encoded byte string, same as printing value directly. + """ + return str(self.value) + + def __repr__(self): + """ + Returns a string representation of the cell + """ + return f"Cell(value={self.value!r}, row={self.row_key!r}, family='{self.family}', column_qualifier={self.column_qualifier!r}, timestamp_micros={self.timestamp_micros}, labels={self.labels})" + + """For Bigtable native ordering""" + + def __lt__(self, other) -> bool: + """ + Implements `<` operator + """ + if not isinstance(other, Cell): + return NotImplemented + this_ordering = ( + self.family, + self.column_qualifier, + -self.timestamp_micros, + self.value, + self.labels, + ) + other_ordering = ( + other.family, + other.column_qualifier, + -other.timestamp_micros, + other.value, + other.labels, + ) + return this_ordering < other_ordering + + def __eq__(self, other) -> bool: + """ + Implements `==` operator + """ + if not isinstance(other, Cell): + return NotImplemented + return ( + self.row_key == other.row_key + and self.family == other.family + and self.column_qualifier == other.column_qualifier + and self.value == other.value + and self.timestamp_micros == other.timestamp_micros + and len(self.labels) == len(other.labels) + and all([label in other.labels for label in self.labels]) + ) + + def __ne__(self, other) -> bool: + """ + Implements `!=` operator + """ + return not self == other + + def __hash__(self): + """ + Implements `hash()` function to fingerprint cell + """ + return hash( + ( + self.row_key, + self.family, + self.column_qualifier, + self.value, + self.timestamp_micros, + tuple(self.labels), + ) + ) diff --git a/google/cloud/bigtable/row_response.py b/google/cloud/bigtable/row_response.py deleted file mode 100644 index be6d8c505..000000000 --- a/google/cloud/bigtable/row_response.py +++ /dev/null @@ -1,130 +0,0 @@ -# Copyright 2023 Google LLC -# -# Licensed 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. -# -from __future__ import annotations - -from collections import OrderedDict -from typing import Sequence - -# Type aliases used internally for readability. -row_key = bytes -family_id = str -qualifier = bytes -row_value = bytes - - -class RowResponse(Sequence["CellResponse"]): - """ - Model class for row data returned from server - - Does not represent all data contained in the row, only data returned by a - query. - Expected to be read-only to users, and written by backend - - Can be indexed: - cells = row["family", "qualifier"] - """ - - def __init__(self, key: row_key, cells: list[CellResponse]): - self.row_key = key - self.cells: OrderedDict[ - family_id, OrderedDict[qualifier, list[CellResponse]] - ] = OrderedDict() - """Expected to be used internally only""" - pass - - def get_cells( - self, family: str | None, qualifer: str | bytes | None - ) -> list[CellResponse]: - """ - Returns cells sorted in Bigtable native order: - - Family lexicographically ascending - - Qualifier lexicographically ascending - - Timestamp in reverse chronological order - - If family or qualifier not passed, will include all - - Syntactic sugar: cells = row["family", "qualifier"] - """ - raise NotImplementedError - - def get_index(self) -> dict[family_id, list[qualifier]]: - """ - Returns a list of family and qualifiers for the object - """ - raise NotImplementedError - - def __str__(self): - """ - Human-readable string representation - - (family, qualifier) cells - (ABC, XYZ) [b"123", b"456" ...(+5)] - (DEF, XYZ) [b"123"] - (GHI, XYZ) [b"123", b"456" ...(+2)] - """ - raise NotImplementedError - - -class CellResponse: - """ - Model class for cell data - - Does not represent all data contained in the cell, only data returned by a - query. - Expected to be read-only to users, and written by backend - """ - - def __init__( - self, - value: row_value, - row: row_key, - family: family_id, - column_qualifier: qualifier, - labels: list[str] | None = None, - timestamp: int | None = None, - ): - self.value = value - self.row_key = row - self.family = family - self.column_qualifier = column_qualifier - self.labels = labels - self.timestamp = timestamp - - def decode_value(self, encoding="UTF-8", errors=None) -> str: - """decode bytes to string""" - return self.value.decode(encoding, errors) - - def __int__(self) -> int: - """ - Allows casting cell to int - Interprets value as a 64-bit big-endian signed integer, as expected by - ReadModifyWrite increment rule - """ - return int.from_bytes(self.value, byteorder="big", signed=True) - - def __str__(self) -> str: - """ - Allows casting cell to str - Prints encoded byte string, same as printing value directly. - """ - return str(self.value) - - """For Bigtable native ordering""" - - def __lt__(self, other) -> bool: - raise NotImplementedError - - def __eq__(self, other) -> bool: - raise NotImplementedError diff --git a/tests/unit/test_row.py b/tests/unit/test_row.py new file mode 100644 index 000000000..32afe92f5 --- /dev/null +++ b/tests/unit/test_row.py @@ -0,0 +1,692 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +import unittest + +import time + +TEST_VALUE = b"1234" +TEST_ROW_KEY = b"row" +TEST_FAMILY_ID = "cf1" +TEST_QUALIFIER = b"col" +TEST_TIMESTAMP = time.time_ns() // 1000 +TEST_LABELS = ["label1", "label2"] + + +class TestRow(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigtable.row import Row + + return Row + + def _make_one(self, *args, **kwargs): + if len(args) == 0: + args = (TEST_ROW_KEY, [self._make_cell()]) + return self._get_target_class()(*args, **kwargs) + + def _make_cell( + self, + value=TEST_VALUE, + row_key=TEST_ROW_KEY, + family_id=TEST_FAMILY_ID, + qualifier=TEST_QUALIFIER, + timestamp=TEST_TIMESTAMP, + labels=TEST_LABELS, + ): + from google.cloud.bigtable.row import Cell + + return Cell(value, row_key, family_id, qualifier, timestamp, labels) + + def test_ctor(self): + cells = [self._make_cell(), self._make_cell()] + row_response = self._make_one(TEST_ROW_KEY, cells) + self.assertEqual(list(row_response), cells) + self.assertEqual(row_response.row_key, TEST_ROW_KEY) + + def test_get_cells(self): + cell_list = [] + for family_id in ["1", "2"]: + for qualifier in [b"a", b"b"]: + cell = self._make_cell(family_id=family_id, qualifier=qualifier) + cell_list.append(cell) + # test getting all cells + row_response = self._make_one(TEST_ROW_KEY, cell_list) + self.assertEqual(row_response.get_cells(), cell_list) + # test getting cells in a family + output = row_response.get_cells(family="1") + self.assertEqual(len(output), 2) + self.assertEqual(output[0].family, "1") + self.assertEqual(output[1].family, "1") + self.assertEqual(output[0], cell_list[0]) + # test getting cells in a family/qualifier + # should accept bytes or str for qualifier + for q in [b"a", "a"]: + output = row_response.get_cells(family="1", qualifier=q) + self.assertEqual(len(output), 1) + self.assertEqual(output[0].family, "1") + self.assertEqual(output[0].column_qualifier, b"a") + self.assertEqual(output[0], cell_list[0]) + # calling with just qualifier should raise an error + with self.assertRaises(ValueError): + row_response.get_cells(qualifier=b"a") + # test calling with bad family or qualifier + with self.assertRaises(ValueError): + row_response.get_cells(family="3", qualifier=b"a") + with self.assertRaises(ValueError): + row_response.get_cells(family="3") + with self.assertRaises(ValueError): + row_response.get_cells(family="1", qualifier=b"c") + + def test___repr__(self): + + cell_str = ( + "{'value': b'1234', 'timestamp_micros': %d, 'labels': ['label1', 'label2']}" + % (TEST_TIMESTAMP) + ) + expected_prefix = "Row(key=b'row', cells=" + row = self._make_one(TEST_ROW_KEY, [self._make_cell()]) + self.assertIn(expected_prefix, repr(row)) + self.assertIn(cell_str, repr(row)) + expected_full = ( + "Row(key=b'row', cells={\n ('cf1', b'col'): [{'value': b'1234', 'timestamp_micros': %d, 'labels': ['label1', 'label2']}],\n})" + % (TEST_TIMESTAMP) + ) + self.assertEqual(expected_full, repr(row)) + # try with multiple cells + row = self._make_one(TEST_ROW_KEY, [self._make_cell(), self._make_cell()]) + self.assertIn(expected_prefix, repr(row)) + self.assertIn(cell_str, repr(row)) + + def test___str__(self): + cells = [ + self._make_cell(value=b"1234", family_id="1", qualifier=b"col"), + self._make_cell(value=b"5678", family_id="3", qualifier=b"col"), + self._make_cell(value=b"1", family_id="3", qualifier=b"col"), + self._make_cell(value=b"2", family_id="3", qualifier=b"col"), + ] + + row_response = self._make_one(TEST_ROW_KEY, cells) + expected = ( + "{\n" + + " (family='1', qualifier=b'col'): [b'1234'],\n" + + " (family='3', qualifier=b'col'): [b'5678', (+2 more)],\n" + + "}" + ) + self.assertEqual(expected, str(row_response)) + + def test_to_dict(self): + from google.cloud.bigtable_v2.types import Row + + cell1 = self._make_cell() + cell2 = self._make_cell() + cell2.value = b"other" + row = self._make_one(TEST_ROW_KEY, [cell1, cell2]) + row_dict = row.to_dict() + expected_dict = { + "key": TEST_ROW_KEY, + "families": [ + { + "name": TEST_FAMILY_ID, + "columns": [ + { + "qualifier": TEST_QUALIFIER, + "cells": [ + { + "value": TEST_VALUE, + "timestamp_micros": TEST_TIMESTAMP, + "labels": TEST_LABELS, + }, + { + "value": b"other", + "timestamp_micros": TEST_TIMESTAMP, + "labels": TEST_LABELS, + }, + ], + } + ], + }, + ], + } + self.assertEqual(len(row_dict), len(expected_dict)) + for key, value in expected_dict.items(): + self.assertEqual(row_dict[key], value) + # should be able to construct a Cell proto from the dict + row_proto = Row(**row_dict) + self.assertEqual(row_proto.key, TEST_ROW_KEY) + self.assertEqual(len(row_proto.families), 1) + family = row_proto.families[0] + self.assertEqual(family.name, TEST_FAMILY_ID) + self.assertEqual(len(family.columns), 1) + column = family.columns[0] + self.assertEqual(column.qualifier, TEST_QUALIFIER) + self.assertEqual(len(column.cells), 2) + self.assertEqual(column.cells[0].value, TEST_VALUE) + self.assertEqual(column.cells[0].timestamp_micros, TEST_TIMESTAMP) + self.assertEqual(column.cells[0].labels, TEST_LABELS) + self.assertEqual(column.cells[1].value, cell2.value) + self.assertEqual(column.cells[1].timestamp_micros, TEST_TIMESTAMP) + self.assertEqual(column.cells[1].labels, TEST_LABELS) + + def test_iteration(self): + from types import GeneratorType + from google.cloud.bigtable.row import Cell + + # should be able to iterate over the Row as a list + cell1 = self._make_cell(value=b"1") + cell2 = self._make_cell(value=b"2") + cell3 = self._make_cell(value=b"3") + row_response = self._make_one(TEST_ROW_KEY, [cell1, cell2, cell3]) + self.assertEqual(len(row_response), 3) + # should create generator object + self.assertIsInstance(iter(row_response), GeneratorType) + result_list = list(row_response) + self.assertEqual(len(result_list), 3) + # should be able to iterate over all cells + idx = 0 + for cell in row_response: + self.assertIsInstance(cell, Cell) + self.assertEqual(cell.value, result_list[idx].value) + self.assertEqual(cell.value, str(idx + 1).encode()) + idx += 1 + + def test_contains_cell(self): + cell3 = self._make_cell(value=b"3") + cell1 = self._make_cell(value=b"1") + cell2 = self._make_cell(value=b"2") + cell4 = self._make_cell(value=b"4") + row_response = self._make_one(TEST_ROW_KEY, [cell3, cell1, cell2]) + self.assertIn(cell1, row_response) + self.assertIn(cell2, row_response) + self.assertNotIn(cell4, row_response) + cell3_copy = self._make_cell(value=b"3") + self.assertIn(cell3_copy, row_response) + + def test_contains_family_id(self): + new_family_id = "new_family_id" + cell = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + cell2 = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + new_family_id, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + row_response = self._make_one(TEST_ROW_KEY, [cell, cell2]) + self.assertIn(TEST_FAMILY_ID, row_response) + self.assertIn("new_family_id", row_response) + self.assertIn(new_family_id, row_response) + self.assertNotIn("not_a_family_id", row_response) + self.assertNotIn(None, row_response) + + def test_contains_family_qualifier_tuple(self): + new_family_id = "new_family_id" + new_qualifier = b"new_qualifier" + cell = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + cell2 = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + new_family_id, + new_qualifier, + TEST_TIMESTAMP, + TEST_LABELS, + ) + row_response = self._make_one(TEST_ROW_KEY, [cell, cell2]) + self.assertIn((TEST_FAMILY_ID, TEST_QUALIFIER), row_response) + self.assertIn(("new_family_id", "new_qualifier"), row_response) + self.assertIn(("new_family_id", b"new_qualifier"), row_response) + self.assertIn((new_family_id, new_qualifier), row_response) + + self.assertNotIn(("not_a_family_id", TEST_QUALIFIER), row_response) + self.assertNotIn((TEST_FAMILY_ID, "not_a_qualifier"), row_response) + self.assertNotIn((TEST_FAMILY_ID, new_qualifier), row_response) + self.assertNotIn(("not_a_family_id", "not_a_qualifier"), row_response) + self.assertNotIn((None, None), row_response) + self.assertNotIn(None, row_response) + + def test_int_indexing(self): + # should be able to index into underlying list with an index number directly + cell_list = [self._make_cell(value=str(i).encode()) for i in range(10)] + sorted(cell_list) + row_response = self._make_one(TEST_ROW_KEY, cell_list) + self.assertEqual(len(row_response), 10) + for i in range(10): + self.assertEqual(row_response[i].value, str(i).encode()) + # backwards indexing should work + self.assertEqual(row_response[-i - 1].value, str(9 - i).encode()) + with self.assertRaises(IndexError): + row_response[10] + with self.assertRaises(IndexError): + row_response[-11] + + def test_slice_indexing(self): + # should be able to index with a range of indices + cell_list = [self._make_cell(value=str(i).encode()) for i in range(10)] + sorted(cell_list) + row_response = self._make_one(TEST_ROW_KEY, cell_list) + self.assertEqual(len(row_response), 10) + self.assertEqual(len(row_response[0:10]), 10) + self.assertEqual(row_response[0:10], cell_list) + self.assertEqual(len(row_response[0:]), 10) + self.assertEqual(row_response[0:], cell_list) + self.assertEqual(len(row_response[:10]), 10) + self.assertEqual(row_response[:10], cell_list) + self.assertEqual(len(row_response[0:10:1]), 10) + self.assertEqual(row_response[0:10:1], cell_list) + self.assertEqual(len(row_response[0:10:2]), 5) + self.assertEqual(row_response[0:10:2], [cell_list[i] for i in range(0, 10, 2)]) + self.assertEqual(len(row_response[0:10:3]), 4) + self.assertEqual(row_response[0:10:3], [cell_list[i] for i in range(0, 10, 3)]) + self.assertEqual(len(row_response[10:0:-1]), 9) + self.assertEqual(len(row_response[10:0:-2]), 5) + self.assertEqual(row_response[10:0:-3], cell_list[10:0:-3]) + self.assertEqual(len(row_response[0:100]), 10) + + def test_family_indexing(self): + # should be able to retrieve cells in a family + new_family_id = "new_family_id" + cell = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + cell2 = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + cell3 = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + new_family_id, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + row_response = self._make_one(TEST_ROW_KEY, [cell, cell2, cell3]) + + self.assertEqual(len(row_response[TEST_FAMILY_ID]), 2) + self.assertEqual(row_response[TEST_FAMILY_ID][0], cell) + self.assertEqual(row_response[TEST_FAMILY_ID][1], cell2) + self.assertEqual(len(row_response[new_family_id]), 1) + self.assertEqual(row_response[new_family_id][0], cell3) + with self.assertRaises(ValueError): + row_response["not_a_family_id"] + with self.assertRaises(TypeError): + row_response[None] + with self.assertRaises(TypeError): + row_response[b"new_family_id"] + + def test_family_qualifier_indexing(self): + # should be able to retrieve cells in a family/qualifier tuplw + new_family_id = "new_family_id" + new_qualifier = b"new_qualifier" + cell = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + cell2 = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + cell3 = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + new_family_id, + new_qualifier, + TEST_TIMESTAMP, + TEST_LABELS, + ) + row_response = self._make_one(TEST_ROW_KEY, [cell, cell2, cell3]) + + self.assertEqual(len(row_response[TEST_FAMILY_ID, TEST_QUALIFIER]), 2) + self.assertEqual(row_response[TEST_FAMILY_ID, TEST_QUALIFIER][0], cell) + self.assertEqual(row_response[TEST_FAMILY_ID, TEST_QUALIFIER][1], cell2) + self.assertEqual(len(row_response[new_family_id, new_qualifier]), 1) + self.assertEqual(row_response[new_family_id, new_qualifier][0], cell3) + self.assertEqual(len(row_response["new_family_id", "new_qualifier"]), 1) + self.assertEqual(len(row_response["new_family_id", b"new_qualifier"]), 1) + with self.assertRaises(ValueError): + row_response[new_family_id, "not_a_qualifier"] + with self.assertRaises(ValueError): + row_response["not_a_family_id", new_qualifier] + with self.assertRaises(TypeError): + row_response[None, None] + with self.assertRaises(TypeError): + row_response[b"new_family_id", b"new_qualifier"] + + def test_get_column_components(self): + # should be able to retrieve (family,qualifier) tuples as keys + new_family_id = "new_family_id" + new_qualifier = b"new_qualifier" + cell = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + cell2 = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + cell3 = self._make_cell( + TEST_VALUE, + TEST_ROW_KEY, + new_family_id, + new_qualifier, + TEST_TIMESTAMP, + TEST_LABELS, + ) + row_response = self._make_one(TEST_ROW_KEY, [cell, cell2, cell3]) + + self.assertEqual(len(row_response.get_column_components()), 2) + self.assertEqual( + row_response.get_column_components(), + [(TEST_FAMILY_ID, TEST_QUALIFIER), (new_family_id, new_qualifier)], + ) + + row_response = self._make_one(TEST_ROW_KEY, []) + self.assertEqual(len(row_response.get_column_components()), 0) + self.assertEqual(row_response.get_column_components(), []) + + row_response = self._make_one(TEST_ROW_KEY, [cell]) + self.assertEqual(len(row_response.get_column_components()), 1) + self.assertEqual( + row_response.get_column_components(), [(TEST_FAMILY_ID, TEST_QUALIFIER)] + ) + + def test_index_of(self): + # given a cell, should find index in underlying list + cell_list = [self._make_cell(value=str(i).encode()) for i in range(10)] + sorted(cell_list) + row_response = self._make_one(TEST_ROW_KEY, cell_list) + + self.assertEqual(row_response.index(cell_list[0]), 0) + self.assertEqual(row_response.index(cell_list[5]), 5) + self.assertEqual(row_response.index(cell_list[9]), 9) + with self.assertRaises(ValueError): + row_response.index(self._make_cell()) + with self.assertRaises(ValueError): + row_response.index(None) + + +class TestCell(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigtable.row import Cell + + return Cell + + def _make_one(self, *args, **kwargs): + if len(args) == 0: + args = ( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + return self._get_target_class()(*args, **kwargs) + + def test_ctor(self): + cell = self._make_one( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + self.assertEqual(cell.value, TEST_VALUE) + self.assertEqual(cell.row_key, TEST_ROW_KEY) + self.assertEqual(cell.family, TEST_FAMILY_ID) + self.assertEqual(cell.column_qualifier, TEST_QUALIFIER) + self.assertEqual(cell.timestamp_micros, TEST_TIMESTAMP) + self.assertEqual(cell.labels, TEST_LABELS) + + def test_to_dict(self): + from google.cloud.bigtable_v2.types import Cell + + cell = self._make_one() + cell_dict = cell.to_dict() + expected_dict = { + "value": TEST_VALUE, + "timestamp_micros": TEST_TIMESTAMP, + "labels": TEST_LABELS, + } + self.assertEqual(len(cell_dict), len(expected_dict)) + for key, value in expected_dict.items(): + self.assertEqual(cell_dict[key], value) + # should be able to construct a Cell proto from the dict + cell_proto = Cell(**cell_dict) + self.assertEqual(cell_proto.value, TEST_VALUE) + self.assertEqual(cell_proto.timestamp_micros, TEST_TIMESTAMP) + self.assertEqual(cell_proto.labels, TEST_LABELS) + + def test_to_dict_no_labels(self): + from google.cloud.bigtable_v2.types import Cell + + cell_no_labels = self._make_one( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + None, + ) + cell_dict = cell_no_labels.to_dict() + expected_dict = { + "value": TEST_VALUE, + "timestamp_micros": TEST_TIMESTAMP, + } + self.assertEqual(len(cell_dict), len(expected_dict)) + for key, value in expected_dict.items(): + self.assertEqual(cell_dict[key], value) + # should be able to construct a Cell proto from the dict + cell_proto = Cell(**cell_dict) + self.assertEqual(cell_proto.value, TEST_VALUE) + self.assertEqual(cell_proto.timestamp_micros, TEST_TIMESTAMP) + self.assertEqual(cell_proto.labels, []) + + def test_int_value(self): + test_int = 1234 + bytes_value = test_int.to_bytes(4, "big", signed=True) + cell = self._make_one( + bytes_value, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + self.assertEqual(int(cell), test_int) + # ensure string formatting works + formatted = "%d" % cell + self.assertEqual(formatted, str(test_int)) + self.assertEqual(int(formatted), test_int) + + def test_int_value_negative(self): + test_int = -99999 + bytes_value = test_int.to_bytes(4, "big", signed=True) + cell = self._make_one( + bytes_value, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + self.assertEqual(int(cell), test_int) + # ensure string formatting works + formatted = "%d" % cell + self.assertEqual(formatted, str(test_int)) + self.assertEqual(int(formatted), test_int) + + def test___str__(self): + test_value = b"helloworld" + cell = self._make_one( + test_value, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + self.assertEqual(str(cell), "b'helloworld'") + self.assertEqual(str(cell), str(test_value)) + + def test___repr__(self): + from google.cloud.bigtable.row import Cell # type: ignore # noqa: F401 + + cell = self._make_one() + expected = ( + "Cell(value=b'1234', row=b'row', " + + "family='cf1', column_qualifier=b'col', " + + f"timestamp_micros={TEST_TIMESTAMP}, labels=['label1', 'label2'])" + ) + self.assertEqual(repr(cell), expected) + # should be able to construct instance from __repr__ + result = eval(repr(cell)) + self.assertEqual(result, cell) + + def test___repr___no_labels(self): + from google.cloud.bigtable.row import Cell # type: ignore # noqa: F401 + + cell_no_labels = self._make_one( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + None, + ) + expected = ( + "Cell(value=b'1234', row=b'row', " + + "family='cf1', column_qualifier=b'col', " + + f"timestamp_micros={TEST_TIMESTAMP}, labels=[])" + ) + self.assertEqual(repr(cell_no_labels), expected) + # should be able to construct instance from __repr__ + result = eval(repr(cell_no_labels)) + self.assertEqual(result, cell_no_labels) + + def test_equality(self): + cell1 = self._make_one() + cell2 = self._make_one() + self.assertEqual(cell1, cell2) + self.assertTrue(cell1 == cell2) + args = ( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + for i in range(0, len(args)): + # try changing each argument + modified_cell = self._make_one(*args[:i], args[i] + args[i], *args[i + 1 :]) + self.assertNotEqual(cell1, modified_cell) + self.assertFalse(cell1 == modified_cell) + self.assertTrue(cell1 != modified_cell) + + def test_hash(self): + # class should be hashable + cell1 = self._make_one() + d = {cell1: 1} + cell2 = self._make_one() + self.assertEqual(d[cell2], 1) + + args = ( + TEST_VALUE, + TEST_ROW_KEY, + TEST_FAMILY_ID, + TEST_QUALIFIER, + TEST_TIMESTAMP, + TEST_LABELS, + ) + for i in range(0, len(args)): + # try changing each argument + modified_cell = self._make_one(*args[:i], args[i] + args[i], *args[i + 1 :]) + with self.assertRaises(KeyError): + d[modified_cell] + + def test_ordering(self): + # create cell list in order from lowest to highest + higher_cells = [] + i = 0 + # families; alphebetical order + for family in ["z", "y", "x"]: + # qualifiers; lowest byte value first + for qualifier in [b"z", b"y", b"x"]: + # timestamps; newest first + for timestamp in [ + TEST_TIMESTAMP, + TEST_TIMESTAMP + 1, + TEST_TIMESTAMP + 2, + ]: + cell = self._make_one( + TEST_VALUE, + TEST_ROW_KEY, + family, + qualifier, + timestamp, + TEST_LABELS, + ) + # cell should be the highest priority encountered so far + self.assertEqual(i, len(higher_cells)) + i += 1 + for other in higher_cells: + self.assertLess(cell, other) + higher_cells.append(cell) + # final order should be reverse of sorted order + expected_order = higher_cells + expected_order.reverse() + self.assertEqual(expected_order, sorted(higher_cells)) From f9a1907ce180571f7b972a7ada83f02d2dba4c13 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 24 Apr 2023 10:30:33 -0700 Subject: [PATCH 05/56] feat: add pooled grpc transport (#748) --- .gitmodules | 3 + gapic-generator-fork | 1 + google/cloud/bigtable/client.py | 325 ++++++- .../bigtable_v2/services/bigtable/client.py | 2 + .../services/bigtable/transports/__init__.py | 3 + .../transports/pooled_grpc_asyncio.py | 426 +++++++++ noxfile.py | 2 + tests/system/test_system.py | 156 ++++ tests/unit/gapic/bigtable_v2/test_bigtable.py | 505 ++++++++++- tests/unit/test_client.py | 814 ++++++++++++++++++ 10 files changed, 2207 insertions(+), 30 deletions(-) create mode 100644 .gitmodules create mode 160000 gapic-generator-fork create mode 100644 google/cloud/bigtable_v2/services/bigtable/transports/pooled_grpc_asyncio.py create mode 100644 tests/system/test_system.py create mode 100644 tests/unit/test_client.py diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 000000000..4186187f4 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "gapic-generator-fork"] + path = gapic-generator-fork + url = git@github.com:googleapis/gapic-generator-python.git diff --git a/gapic-generator-fork b/gapic-generator-fork new file mode 160000 index 000000000..b26cda7d1 --- /dev/null +++ b/gapic-generator-fork @@ -0,0 +1 @@ +Subproject commit b26cda7d163d6e0d45c9684f328ca32fb49b799a diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index 23f0cb6fe..dfd8b16cd 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -15,12 +15,29 @@ from __future__ import annotations -from typing import Any, AsyncIterable, TYPE_CHECKING - +from typing import cast, Any, Optional, AsyncIterable, Set, TYPE_CHECKING + +import asyncio +import grpc +import time +import warnings +import sys +import random + +from google.cloud.bigtable_v2.services.bigtable.client import BigtableClientMeta +from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient +from google.cloud.bigtable_v2.services.bigtable.async_client import DEFAULT_CLIENT_INFO +from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( + PooledBigtableGrpcAsyncIOTransport, +) from google.cloud.client import ClientWithProject +from google.api_core.exceptions import GoogleAPICallError import google.auth.credentials +import google.auth._default +from google.api_core import client_options as client_options_lib + if TYPE_CHECKING: from google.cloud.bigtable.mutations import Mutation, BulkMutationsEntry @@ -42,11 +59,12 @@ def __init__( client_options: dict[str, Any] | "google.api_core.client_options.ClientOptions" | None = None, - metadata: list[tuple[str, str]] | None = None, ): """ Create a client instance for the Bigtable Data API + Client should be created within an async context (running event loop) + Args: project: the project which the client acts on behalf of. If not passed, falls back to the default inferred @@ -61,30 +79,241 @@ def __init__( client_options (Optional[Union[dict, google.api_core.client_options.ClientOptions]]): Client options used to set user options on the client. API Endpoint should be set through client_options. - metadata: a list of metadata headers to be attached to all calls with this client + Raises: + - RuntimeError if called outside of an async context (no running event loop) + - ValueError if pool_size is less than 1 """ - raise NotImplementedError + # set up transport in registry + transport_str = f"pooled_grpc_asyncio_{pool_size}" + transport = PooledBigtableGrpcAsyncIOTransport.with_fixed_size(pool_size) + BigtableClientMeta._transport_registry[transport_str] = transport + # set up client info headers for veneer library + client_info = DEFAULT_CLIENT_INFO + client_info.client_library_version = client_info.gapic_version + # parse client options + if type(client_options) is dict: + client_options = client_options_lib.from_dict(client_options) + client_options = cast( + Optional[client_options_lib.ClientOptions], client_options + ) + # initialize client + ClientWithProject.__init__( + self, + credentials=credentials, + project=project, + client_options=client_options, + ) + self._gapic_client = BigtableAsyncClient( + transport=transport_str, + credentials=credentials, + client_options=client_options, + client_info=client_info, + ) + self.transport = cast( + PooledBigtableGrpcAsyncIOTransport, self._gapic_client.transport + ) + # keep track of active instances to for warmup on channel refresh + self._active_instances: Set[str] = set() + # keep track of table objects associated with each instance + # only remove instance from _active_instances when all associated tables remove it + self._instance_owners: dict[str, Set[int]] = {} + # attempt to start background tasks + self._channel_init_time = time.time() + self._channel_refresh_tasks: list[asyncio.Task[None]] = [] + try: + self.start_background_channel_refresh() + except RuntimeError: + warnings.warn( + f"{self.__class__.__name__} should be started in an " + "asyncio event loop. Channel refresh will not be started", + RuntimeWarning, + stacklevel=2, + ) + + def start_background_channel_refresh(self) -> None: + """ + Starts a background task to ping and warm each channel in the pool + Raises: + - RuntimeError if not called in an asyncio event loop + """ + if not self._channel_refresh_tasks: + # raise RuntimeError if there is no event loop + asyncio.get_running_loop() + for channel_idx in range(self.transport.pool_size): + refresh_task = asyncio.create_task(self._manage_channel(channel_idx)) + if sys.version_info >= (3, 8): + # task names supported in Python 3.8+ + refresh_task.set_name( + f"{self.__class__.__name__} channel refresh {channel_idx}" + ) + self._channel_refresh_tasks.append(refresh_task) + + async def close(self, timeout: float = 2.0): + """ + Cancel all background tasks + """ + for task in self._channel_refresh_tasks: + task.cancel() + group = asyncio.gather(*self._channel_refresh_tasks, return_exceptions=True) + await asyncio.wait_for(group, timeout=timeout) + await self.transport.close() + self._channel_refresh_tasks = [] + + async def _ping_and_warm_instances( + self, channel: grpc.aio.Channel + ) -> list[GoogleAPICallError | None]: + """ + Prepares the backend for requests on a channel + + Pings each Bigtable instance registered in `_active_instances` on the client + + Args: + channel: grpc channel to ping + Returns: + - sequence of results or exceptions from the ping requests + """ + ping_rpc = channel.unary_unary( + "/google.bigtable.v2.Bigtable/PingAndWarmChannel" + ) + tasks = [ping_rpc({"name": n}) for n in self._active_instances] + return await asyncio.gather(*tasks, return_exceptions=True) + + async def _manage_channel( + self, + channel_idx: int, + refresh_interval_min: float = 60 * 35, + refresh_interval_max: float = 60 * 45, + grace_period: float = 60 * 10, + ) -> None: + """ + Background coroutine that periodically refreshes and warms a grpc channel + + The backend will automatically close channels after 60 minutes, so + `refresh_interval` + `grace_period` should be < 60 minutes + + Runs continuously until the client is closed + + Args: + channel_idx: index of the channel in the transport's channel pool + refresh_interval_min: minimum interval before initiating refresh + process in seconds. Actual interval will be a random value + between `refresh_interval_min` and `refresh_interval_max` + refresh_interval_max: maximum interval before initiating refresh + process in seconds. Actual interval will be a random value + between `refresh_interval_min` and `refresh_interval_max` + grace_period: time to allow previous channel to serve existing + requests before closing, in seconds + """ + first_refresh = self._channel_init_time + random.uniform( + refresh_interval_min, refresh_interval_max + ) + next_sleep = max(first_refresh - time.time(), 0) + if next_sleep > 0: + # warm the current channel immediately + channel = self.transport.channels[channel_idx] + await self._ping_and_warm_instances(channel) + # continuously refresh the channel every `refresh_interval` seconds + while True: + await asyncio.sleep(next_sleep) + # prepare new channel for use + new_channel = self.transport.grpc_channel._create_channel() + await self._ping_and_warm_instances(new_channel) + # cycle channel out of use, with long grace window before closure + start_timestamp = time.time() + await self.transport.replace_channel( + channel_idx, grace=grace_period, swap_sleep=10, new_channel=new_channel + ) + # subtract the time spent waiting for the channel to be replaced + next_refresh = random.uniform(refresh_interval_min, refresh_interval_max) + next_sleep = next_refresh - (time.time() - start_timestamp) + + async def _register_instance(self, instance_id: str, owner: Table) -> None: + """ + Registers an instance with the client, and warms the channel pool + for the instance + The client will periodically refresh grpc channel pool used to make + requests, and new channels will be warmed for each registered instance + Channels will not be refreshed unless at least one instance is registered + + Args: + - instance_id: id of the instance to register. + - owner: table that owns the instance. Owners will be tracked in + _instance_owners, and instances will only be unregistered when all + owners call _remove_instance_registration + """ + instance_name = self._gapic_client.instance_path(self.project, instance_id) + self._instance_owners.setdefault(instance_name, set()).add(id(owner)) + if instance_name not in self._active_instances: + self._active_instances.add(instance_name) + if self._channel_refresh_tasks: + # refresh tasks already running + # call ping and warm on all existing channels + for channel in self.transport.channels: + await self._ping_and_warm_instances(channel) + else: + # refresh tasks aren't active. start them as background tasks + self.start_background_channel_refresh() + + async def _remove_instance_registration( + self, instance_id: str, owner: Table + ) -> bool: + """ + Removes an instance from the client's registered instances, to prevent + warming new channels for the instance + + If instance_id is not registered, or is still in use by other tables, returns False + + Args: + - instance_id: id of the instance to remove + - owner: table that owns the instance. Owners will be tracked in + _instance_owners, and instances will only be unregistered when all + owners call _remove_instance_registration + Returns: + - True if instance was removed + """ + instance_name = self._gapic_client.instance_path(self.project, instance_id) + owner_list = self._instance_owners.get(instance_name, set()) + try: + owner_list.remove(id(owner)) + if len(owner_list) == 0: + self._active_instances.remove(instance_name) + return True + except KeyError: + return False def get_table( - self, instance_id: str, table_id: str, app_profile_id: str | None = None + self, + instance_id: str, + table_id: str, + app_profile_id: str | None = None, ) -> Table: """ - Return a Table instance to make API requests for a specific table. + Returns a table instance for making data API requests Args: - instance_id: The ID of the instance that owns the table. + instance_id: The Bigtable instance ID to associate with this client. + instance_id is combined with the client's project to fully + specify the instance table_id: The ID of the table. app_profile_id: (Optional) The app profile to associate with requests. https://cloud.google.com/bigtable/docs/app-profiles """ - raise NotImplementedError + return Table(self, instance_id, table_id, app_profile_id) + + async def __aenter__(self): + self.start_background_channel_refresh() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.close() + await self._gapic_client.__aexit__(exc_type, exc_val, exc_tb) class Table: """ Main Data API surface - Table object maintains instance_id, table_id, and app_profile_id context, and passes them with + Table object maintains table_id, and app_profile_id context, and passes them with each call """ @@ -95,7 +324,41 @@ def __init__( table_id: str, app_profile_id: str | None = None, ): - raise NotImplementedError + """ + Initialize a Table instance + + Must be created within an async context (running event loop) + + Args: + instance_id: The Bigtable instance ID to associate with this client. + instance_id is combined with the client's project to fully + specify the instance + table_id: The ID of the table. table_id is combined with the + instance_id and the client's project to fully specify the table + app_profile_id: (Optional) The app profile to associate with requests. + https://cloud.google.com/bigtable/docs/app-profiles + Raises: + - RuntimeError if called outside of an async context (no running event loop) + """ + self.client = client + self.instance_id = instance_id + self.instance_name = self.client._gapic_client.instance_path( + self.client.project, instance_id + ) + self.table_id = table_id + self.table_name = self.client._gapic_client.table_path( + self.client.project, instance_id, table_id + ) + self.app_profile_id = app_profile_id + # raises RuntimeError if called outside of an async context (no running event loop) + try: + self._register_instance_task = asyncio.create_task( + self.client._register_instance(instance_id, self) + ) + except RuntimeError as e: + raise RuntimeError( + f"{self.__class__.__name__} must be created within an async event loop context." + ) from e async def read_rows_stream( self, @@ -108,7 +371,6 @@ async def read_rows_stream( per_row_timeout: int | float | None = 10, idle_timeout: int | float | None = 300, per_request_timeout: int | float | None = None, - metadata: list[tuple[str, str]] | None = None, ) -> AsyncIterable[Row]: """ Returns a generator to asynchronously stream back row data. @@ -144,7 +406,6 @@ async def read_rows_stream( - per_request_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted - - metadata: Strings which should be sent along with the request as metadata headers. Returns: - an asynchronous generator that yields rows returned by the query @@ -165,7 +426,6 @@ async def read_rows( operation_timeout: int | float | None = 60, per_row_timeout: int | float | None = 10, per_request_timeout: int | float | None = None, - metadata: list[tuple[str, str]] | None = None, ) -> list[Row]: """ Helper function that returns a full list instead of a generator @@ -183,7 +443,6 @@ async def read_row( *, operation_timeout: int | float | None = 60, per_request_timeout: int | float | None = None, - metadata: list[tuple[str, str]] | None = None, ) -> Row: """ Helper function to return a single row @@ -205,7 +464,6 @@ async def read_rows_sharded( per_row_timeout: int | float | None = 10, idle_timeout: int | float | None = 300, per_request_timeout: int | float | None = None, - metadata: list[tuple[str, str]] | None = None, ) -> AsyncIterable[Row]: """ Runs a sharded query in parallel @@ -224,7 +482,6 @@ async def row_exists( *, operation_timeout: int | float | None = 60, per_request_timeout: int | float | None = None, - metadata: list[tuple[str, str]] | None = None, ) -> bool: """ Helper function to determine if a row exists @@ -242,7 +499,6 @@ async def sample_keys( operation_timeout: int | float | None = 60, per_sample_timeout: int | float | None = 10, per_request_timeout: int | float | None = None, - metadata: list[tuple[str, str]] | None = None, ) -> RowKeySamples: """ Return a set of RowKeySamples that delimit contiguous sections of the table of @@ -283,7 +539,6 @@ async def mutate_row( *, operation_timeout: int | float | None = 60, per_request_timeout: int | float | None = None, - metadata: list[tuple[str, str]] | None = None, ): """ Mutates a row atomically. @@ -305,7 +560,6 @@ async def mutate_row( in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted if within operation_timeout budget - - metadata: Strings which should be sent along with the request as metadata headers. Raises: - DeadlineExceeded: raised after operation timeout @@ -322,7 +576,6 @@ async def bulk_mutate_rows( *, operation_timeout: int | float | None = 60, per_request_timeout: int | float | None = None, - metadata: list[tuple[str, str]] | None = None, ): """ Applies mutations for multiple rows in a single batched request. @@ -348,7 +601,6 @@ async def bulk_mutate_rows( in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted if within operation_timeout budget - - metadata: Strings which should be sent along with the request as metadata headers. Raises: - MutationsExceptionGroup if one or more mutations fails @@ -363,7 +615,6 @@ async def check_and_mutate_row( true_case_mutations: Mutation | list[Mutation] | None = None, false_case_mutations: Mutation | list[Mutation] | None = None, operation_timeout: int | float | None = 60, - metadata: list[tuple[str, str]] | None = None, ) -> bool: """ Mutates a row atomically based on the output of a predicate filter @@ -392,7 +643,6 @@ async def check_and_mutate_row( `true_case_mutations is empty, and at most 100000. - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will not be retried. - - metadata: Strings which should be sent along with the request as metadata headers. Returns: - bool indicating whether the predicate was true or false Raises: @@ -409,7 +659,6 @@ async def read_modify_write_row( | list[dict[str, Any]], *, operation_timeout: int | float | None = 60, - metadata: list[tuple[str, str]] | None = None, ) -> Row: """ Reads and modifies a row atomically according to input ReadModifyWriteRules, @@ -427,7 +676,6 @@ async def read_modify_write_row( results of later ones. - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will not be retried. - - metadata: Strings which should be sent along with the request as metadata headers. Returns: - Row: containing cell data that was modified as part of the operation @@ -435,3 +683,28 @@ async def read_modify_write_row( - GoogleAPIError exceptions from grpc call """ raise NotImplementedError + + async def close(self): + """ + Called to close the Table instance and release any resources held by it. + """ + await self.client._remove_instance_registration(self.instance_id, self) + + async def __aenter__(self): + """ + Implement async context manager protocol + + Register this instance with the client, so that + grpc channels will be warmed for the specified instance + """ + await self.client._register_instance(self.instance_id, self) + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """ + Implement async context manager protocol + + Unregister this instance with the client, so that + grpc channels will no longer be warmed + """ + await self.close() diff --git a/google/cloud/bigtable_v2/services/bigtable/client.py b/google/cloud/bigtable_v2/services/bigtable/client.py index 37ab65fe2..60622509a 100644 --- a/google/cloud/bigtable_v2/services/bigtable/client.py +++ b/google/cloud/bigtable_v2/services/bigtable/client.py @@ -53,6 +53,7 @@ from .transports.base import BigtableTransport, DEFAULT_CLIENT_INFO from .transports.grpc import BigtableGrpcTransport from .transports.grpc_asyncio import BigtableGrpcAsyncIOTransport +from .transports.pooled_grpc_asyncio import PooledBigtableGrpcAsyncIOTransport from .transports.rest import BigtableRestTransport @@ -67,6 +68,7 @@ class BigtableClientMeta(type): _transport_registry = OrderedDict() # type: Dict[str, Type[BigtableTransport]] _transport_registry["grpc"] = BigtableGrpcTransport _transport_registry["grpc_asyncio"] = BigtableGrpcAsyncIOTransport + _transport_registry["pooled_grpc_asyncio"] = PooledBigtableGrpcAsyncIOTransport _transport_registry["rest"] = BigtableRestTransport def get_transport_class( diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/__init__.py b/google/cloud/bigtable_v2/services/bigtable/transports/__init__.py index 1b03919f6..e8796bb8c 100644 --- a/google/cloud/bigtable_v2/services/bigtable/transports/__init__.py +++ b/google/cloud/bigtable_v2/services/bigtable/transports/__init__.py @@ -19,6 +19,7 @@ from .base import BigtableTransport from .grpc import BigtableGrpcTransport from .grpc_asyncio import BigtableGrpcAsyncIOTransport +from .pooled_grpc_asyncio import PooledBigtableGrpcAsyncIOTransport from .rest import BigtableRestTransport from .rest import BigtableRestInterceptor @@ -27,12 +28,14 @@ _transport_registry = OrderedDict() # type: Dict[str, Type[BigtableTransport]] _transport_registry["grpc"] = BigtableGrpcTransport _transport_registry["grpc_asyncio"] = BigtableGrpcAsyncIOTransport +_transport_registry["pooled_grpc_asyncio"] = PooledBigtableGrpcAsyncIOTransport _transport_registry["rest"] = BigtableRestTransport __all__ = ( "BigtableTransport", "BigtableGrpcTransport", "BigtableGrpcAsyncIOTransport", + "PooledBigtableGrpcAsyncIOTransport", "BigtableRestTransport", "BigtableRestInterceptor", ) diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/pooled_grpc_asyncio.py b/google/cloud/bigtable_v2/services/bigtable/transports/pooled_grpc_asyncio.py new file mode 100644 index 000000000..372e5796d --- /dev/null +++ b/google/cloud/bigtable_v2/services/bigtable/transports/pooled_grpc_asyncio.py @@ -0,0 +1,426 @@ +# -*- coding: utf-8 -*- +# Copyright 2022 Google LLC +# +# Licensed 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. +# +import asyncio +import warnings +from functools import partialmethod +from functools import partial +from typing import ( + Awaitable, + Callable, + Dict, + Optional, + Sequence, + Tuple, + Union, + List, + Type, +) + +from google.api_core import gapic_v1 +from google.api_core import grpc_helpers_async +from google.auth import credentials as ga_credentials # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore + +import grpc # type: ignore +from grpc.experimental import aio # type: ignore + +from google.cloud.bigtable_v2.types import bigtable +from .base import BigtableTransport, DEFAULT_CLIENT_INFO +from .grpc_asyncio import BigtableGrpcAsyncIOTransport + + +class PooledMultiCallable: + def __init__(self, channel_pool: "PooledChannel", *args, **kwargs): + self._init_args = args + self._init_kwargs = kwargs + self.next_channel_fn = channel_pool.next_channel + + +class PooledUnaryUnaryMultiCallable(PooledMultiCallable, aio.UnaryUnaryMultiCallable): + def __call__(self, *args, **kwargs) -> aio.UnaryUnaryCall: + return self.next_channel_fn().unary_unary( + *self._init_args, **self._init_kwargs + )(*args, **kwargs) + + +class PooledUnaryStreamMultiCallable(PooledMultiCallable, aio.UnaryStreamMultiCallable): + def __call__(self, *args, **kwargs) -> aio.UnaryStreamCall: + return self.next_channel_fn().unary_stream( + *self._init_args, **self._init_kwargs + )(*args, **kwargs) + + +class PooledStreamUnaryMultiCallable(PooledMultiCallable, aio.StreamUnaryMultiCallable): + def __call__(self, *args, **kwargs) -> aio.StreamUnaryCall: + return self.next_channel_fn().stream_unary( + *self._init_args, **self._init_kwargs + )(*args, **kwargs) + + +class PooledStreamStreamMultiCallable( + PooledMultiCallable, aio.StreamStreamMultiCallable +): + def __call__(self, *args, **kwargs) -> aio.StreamStreamCall: + return self.next_channel_fn().stream_stream( + *self._init_args, **self._init_kwargs + )(*args, **kwargs) + + +class PooledChannel(aio.Channel): + def __init__( + self, + pool_size: int = 3, + host: str = "bigtable.googleapis.com", + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + quota_project_id: Optional[str] = None, + default_scopes: Optional[Sequence[str]] = None, + scopes: Optional[Sequence[str]] = None, + default_host: Optional[str] = None, + insecure: bool = False, + **kwargs, + ): + self._pool: List[aio.Channel] = [] + self._next_idx = 0 + if insecure: + self._create_channel = partial(aio.insecure_channel, host) + else: + self._create_channel = partial( + grpc_helpers_async.create_channel, + target=host, + credentials=credentials, + credentials_file=credentials_file, + quota_project_id=quota_project_id, + default_scopes=default_scopes, + scopes=scopes, + default_host=default_host, + **kwargs, + ) + for i in range(pool_size): + self._pool.append(self._create_channel()) + + def next_channel(self) -> aio.Channel: + channel = self._pool[self._next_idx] + self._next_idx = (self._next_idx + 1) % len(self._pool) + return channel + + def unary_unary(self, *args, **kwargs) -> grpc.aio.UnaryUnaryMultiCallable: + return PooledUnaryUnaryMultiCallable(self, *args, **kwargs) + + def unary_stream(self, *args, **kwargs) -> grpc.aio.UnaryStreamMultiCallable: + return PooledUnaryStreamMultiCallable(self, *args, **kwargs) + + def stream_unary(self, *args, **kwargs) -> grpc.aio.StreamUnaryMultiCallable: + return PooledStreamUnaryMultiCallable(self, *args, **kwargs) + + def stream_stream(self, *args, **kwargs) -> grpc.aio.StreamStreamMultiCallable: + return PooledStreamStreamMultiCallable(self, *args, **kwargs) + + async def close(self, grace=None): + close_fns = [channel.close(grace=grace) for channel in self._pool] + return await asyncio.gather(*close_fns) + + async def channel_ready(self): + ready_fns = [channel.channel_ready() for channel in self._pool] + return asyncio.gather(*ready_fns) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.close() + + def get_state(self, try_to_connect: bool = False) -> grpc.ChannelConnectivity: + raise NotImplementedError() + + async def wait_for_state_change(self, last_observed_state): + raise NotImplementedError() + + async def replace_channel( + self, channel_idx, grace=None, swap_sleep=1, new_channel=None + ) -> aio.Channel: + """ + Replaces a channel in the pool with a fresh one. + + The `new_channel` will start processing new requests immidiately, + but the old channel will continue serving existing clients for `grace` seconds + + Args: + channel_idx(int): the channel index in the pool to replace + grace(Optional[float]): The time to wait until all active RPCs are + finished. If a grace period is not specified (by passing None for + grace), all existing RPCs are cancelled immediately. + swap_sleep(Optional[float]): The number of seconds to sleep in between + replacing channels and closing the old one + new_channel(grpc.aio.Channel): a new channel to insert into the pool + at `channel_idx`. If `None`, a new channel will be created. + """ + if channel_idx >= len(self._pool) or channel_idx < 0: + raise ValueError( + f"invalid channel_idx {channel_idx} for pool size {len(self._pool)}" + ) + if new_channel is None: + new_channel = self._create_channel() + old_channel = self._pool[channel_idx] + self._pool[channel_idx] = new_channel + await asyncio.sleep(swap_sleep) + await old_channel.close(grace=grace) + return new_channel + + +class PooledBigtableGrpcAsyncIOTransport(BigtableGrpcAsyncIOTransport): + """Pooled gRPC AsyncIO backend transport for Bigtable. + + Service for reading from and writing to existing Bigtable + tables. + + This class defines the same methods as the primary client, so the + primary client can load the underlying transport implementation + and call it. + + It sends protocol buffers over the wire using gRPC (which is built on + top of HTTP/2); the ``grpcio`` package must be installed. + + This class allows channel pooling, so multiple channels can be used concurrently + when making requests. Channels are rotated in a round-robin fashion. + """ + + @classmethod + def with_fixed_size(cls, pool_size) -> Type["PooledBigtableGrpcAsyncIOTransport"]: + """ + Creates a new class with a fixed channel pool size. + + A fixed channel pool makes compatibility with other transports easier, + as the initializer signature is the same. + """ + + class PooledTransportFixed(cls): + __init__ = partialmethod(cls.__init__, pool_size=pool_size) + + PooledTransportFixed.__name__ = f"{cls.__name__}_{pool_size}" + PooledTransportFixed.__qualname__ = PooledTransportFixed.__name__ + return PooledTransportFixed + + @classmethod + def create_channel( + cls, + pool_size: int = 3, + host: str = "bigtable.googleapis.com", + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + quota_project_id: Optional[str] = None, + **kwargs, + ) -> aio.Channel: + """Create and return a PooledChannel object, representing a pool of gRPC AsyncIO channels + Args: + pool_size (int): The number of channels in the pool. + host (Optional[str]): The host for the channel to use. + credentials (Optional[~.Credentials]): The + authorization credentials to attach to requests. These + credentials identify this application to the service. If + none are specified, the client will attempt to ascertain + the credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + kwargs (Optional[dict]): Keyword arguments, which are passed to the + channel creation. + Returns: + PooledChannel: a channel pool object + """ + + return PooledChannel( + pool_size, + host, + credentials=credentials, + credentials_file=credentials_file, + quota_project_id=quota_project_id, + default_scopes=cls.AUTH_SCOPES, + scopes=scopes, + default_host=cls.DEFAULT_HOST, + **kwargs, + ) + + def __init__( + self, + *, + pool_size: int = 3, + host: str = "bigtable.googleapis.com", + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + api_mtls_endpoint: Optional[str] = None, + client_cert_source: Optional[Callable[[], Tuple[bytes, bytes]]] = None, + ssl_channel_credentials: Optional[grpc.ChannelCredentials] = None, + client_cert_source_for_mtls: Optional[Callable[[], Tuple[bytes, bytes]]] = None, + quota_project_id: Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + always_use_jwt_access: Optional[bool] = False, + api_audience: Optional[str] = None, + ) -> None: + """Instantiate the transport. + + Args: + pool_size (int): the number of grpc channels to maintain in a pool + host (Optional[str]): + The hostname to connect to. + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + This argument is ignored if ``channel`` is provided. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + api_mtls_endpoint (Optional[str]): Deprecated. The mutual TLS endpoint. + If provided, it overrides the ``host`` argument and tries to create + a mutual TLS channel with client SSL credentials from + ``client_cert_source`` or application default SSL credentials. + client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]): + Deprecated. A callback to provide client SSL certificate bytes and + private key bytes, both in PEM format. It is ignored if + ``api_mtls_endpoint`` is None. + ssl_channel_credentials (grpc.ChannelCredentials): SSL credentials + for the grpc channel. It is ignored if ``channel`` is provided. + client_cert_source_for_mtls (Optional[Callable[[], Tuple[bytes, bytes]]]): + A callback to provide client certificate bytes and private key bytes, + both in PEM format. It is used to configure a mutual TLS channel. It is + ignored if ``channel`` or ``ssl_channel_credentials`` is provided. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + always_use_jwt_access (Optional[bool]): Whether self signed JWT should + be used for service account credentials. + + Raises: + google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport + creation failed for any reason. + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + ValueError: if ``pool_size`` <= 0 + """ + if pool_size <= 0: + raise ValueError(f"invalid pool_size: {pool_size}") + self._ssl_channel_credentials = ssl_channel_credentials + self._stubs: Dict[str, Callable] = {} + + if api_mtls_endpoint: + warnings.warn("api_mtls_endpoint is deprecated", DeprecationWarning) + if client_cert_source: + warnings.warn("client_cert_source is deprecated", DeprecationWarning) + + if api_mtls_endpoint: + host = api_mtls_endpoint + + # Create SSL credentials with client_cert_source or application + # default SSL credentials. + if client_cert_source: + cert, key = client_cert_source() + self._ssl_channel_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + else: + self._ssl_channel_credentials = SslCredentials().ssl_credentials + + else: + if client_cert_source_for_mtls and not ssl_channel_credentials: + cert, key = client_cert_source_for_mtls() + self._ssl_channel_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + + # The base transport sets the host, credentials and scopes + BigtableTransport.__init__( + self, + host=host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes, + quota_project_id=quota_project_id, + client_info=client_info, + always_use_jwt_access=always_use_jwt_access, + api_audience=api_audience, + ) + self._quota_project_id = quota_project_id + self._grpc_channel = type(self).create_channel( + pool_size, + self._host, + # use the credentials which are saved + credentials=self._credentials, + # Set ``credentials_file`` to ``None`` here as + # the credentials that we saved earlier should be used. + credentials_file=None, + scopes=self._scopes, + ssl_credentials=self._ssl_channel_credentials, + quota_project_id=self._quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + # Wrap messages. This must be done after self._grpc_channel exists + self._prep_wrapped_messages(client_info) + + @property + def pool_size(self) -> int: + """The number of grpc channels in the pool.""" + return len(self._grpc_channel._pool) + + @property + def channels(self) -> List[grpc.Channel]: + """Acccess the internal list of grpc channels.""" + return self._grpc_channel._pool + + async def replace_channel( + self, channel_idx, grace=None, swap_sleep=1, new_channel=None + ) -> aio.Channel: + """ + Replaces a channel in the pool with a fresh one. + + The `new_channel` will start processing new requests immidiately, + but the old channel will continue serving existing clients for `grace` seconds + + Args: + channel_idx(int): the channel index in the pool to replace + grace(Optional[float]): The time to wait until all active RPCs are + finished. If a grace period is not specified (by passing None for + grace), all existing RPCs are cancelled immediately. + swap_sleep(Optional[float]): The number of seconds to sleep in between + replacing channels and closing the old one + new_channel(grpc.aio.Channel): a new channel to insert into the pool + at `channel_idx`. If `None`, a new channel will be created. + """ + return await self._grpc_channel.replace_channel( + channel_idx, grace, swap_sleep, new_channel + ) + + +__all__ = ("PooledBigtableGrpcAsyncIOTransport",) diff --git a/noxfile.py b/noxfile.py index ed69bf85e..035599844 100644 --- a/noxfile.py +++ b/noxfile.py @@ -49,6 +49,7 @@ SYSTEM_TEST_STANDARD_DEPENDENCIES = [ "mock", "pytest", + "pytest-asyncio", "google-cloud-testutils", ] SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [] @@ -306,6 +307,7 @@ def system(session): "py.test", "--quiet", f"--junitxml=system_{session.python}_sponge_log.xml", + "--ignore=tests/system/v2_client", system_test_folder_path, *session.posargs, ) diff --git a/tests/system/test_system.py b/tests/system/test_system.py new file mode 100644 index 000000000..05633ac91 --- /dev/null +++ b/tests/system/test_system.py @@ -0,0 +1,156 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +import pytest +import pytest_asyncio +import os +import asyncio + +TEST_FAMILY = "test-family" +TEST_FAMILY_2 = "test-family-2" + + +@pytest.fixture(scope="session") +def event_loop(): + return asyncio.get_event_loop() + + +@pytest.fixture(scope="session") +def instance_admin_client(): + """Client for interacting with the Instance Admin API.""" + from google.cloud.bigtable_admin_v2 import BigtableInstanceAdminClient + + with BigtableInstanceAdminClient() as client: + yield client + + +@pytest.fixture(scope="session") +def table_admin_client(): + """Client for interacting with the Table Admin API.""" + from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient + + with BigtableTableAdminClient() as client: + yield client + + +@pytest.fixture(scope="session") +def instance_id(instance_admin_client, project_id): + """ + Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session + """ + from google.cloud.bigtable_admin_v2 import types + from google.api_core import exceptions + + # use user-specified instance if available + user_specified_instance = os.getenv("BIGTABLE_TEST_INSTANCE") + if user_specified_instance: + print("Using user-specified instance: {}".format(user_specified_instance)) + yield user_specified_instance + return + + # create a new temporary test instance + instance_id = "test-instance" + try: + operation = instance_admin_client.create_instance( + parent=f"projects/{project_id}", + instance_id=instance_id, + instance=types.Instance( + display_name="Test Instance", + labels={"python-system-test": "true"}, + ), + clusters={ + "test-cluster": types.Cluster( + location=f"projects/{project_id}/locations/us-central1-b", + serve_nodes=3, + ) + }, + ) + operation.result(timeout=240) + except exceptions.AlreadyExists: + pass + yield instance_id + instance_admin_client.delete_instance( + name=f"projects/{project_id}/instances/{instance_id}" + ) + + +@pytest.fixture(scope="session") +def table_id(table_admin_client, project_id, instance_id): + """ + Returns BIGTABLE_TEST_TABLE if set, otherwise creates a new temporary table for the test session + """ + from google.cloud.bigtable_admin_v2 import types + from google.api_core import exceptions + from google.api_core import retry + + # use user-specified instance if available + user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") + if user_specified_table: + print("Using user-specified table: {}".format(user_specified_table)) + yield user_specified_table + return + + table_id = "test-table" + retry = retry.Retry( + predicate=retry.if_exception_type(exceptions.FailedPrecondition) + ) + try: + table_admin_client.create_table( + parent=f"projects/{project_id}/instances/{instance_id}", + table_id=table_id, + table=types.Table( + column_families={ + TEST_FAMILY: types.ColumnFamily(), + TEST_FAMILY_2: types.ColumnFamily(), + }, + ), + retry=retry, + ) + except exceptions.AlreadyExists: + pass + yield table_id + table_admin_client.delete_table( + name=f"projects/{project_id}/instances/{instance_id}/tables/{table_id}" + ) + + +@pytest_asyncio.fixture(scope="session") +async def client(): + from google.cloud.bigtable import BigtableDataClient + + project = os.getenv("GOOGLE_CLOUD_PROJECT") or None + async with BigtableDataClient(project=project) as client: + yield client + + +@pytest.fixture(scope="session") +def project_id(client): + """Returns the project ID from the client.""" + yield client.project + + +@pytest_asyncio.fixture(scope="session") +async def table(client, table_id, instance_id): + async with client.get_table(instance_id, table_id) as table: + yield table + + +@pytest.mark.asyncio +async def test_ping_and_warm_gapic(client, table): + """ + Simple ping rpc test + This test ensures channels are able to authenticate with backend + """ + request = {"name": table.instance_name} + await client._gapic_client.ping_and_warm(request) diff --git a/tests/unit/gapic/bigtable_v2/test_bigtable.py b/tests/unit/gapic/bigtable_v2/test_bigtable.py index 03ba3044f..b1500aa48 100644 --- a/tests/unit/gapic/bigtable_v2/test_bigtable.py +++ b/tests/unit/gapic/bigtable_v2/test_bigtable.py @@ -100,6 +100,7 @@ def test__get_default_mtls_endpoint(): [ (BigtableClient, "grpc"), (BigtableAsyncClient, "grpc_asyncio"), + (BigtableAsyncClient, "pooled_grpc_asyncio"), (BigtableClient, "rest"), ], ) @@ -116,7 +117,7 @@ def test_bigtable_client_from_service_account_info(client_class, transport_name) assert client.transport._host == ( "bigtable.googleapis.com:443" - if transport_name in ["grpc", "grpc_asyncio"] + if transport_name in ["grpc", "grpc_asyncio", "pooled_grpc_asyncio"] else "https://bigtable.googleapis.com" ) @@ -126,6 +127,7 @@ def test_bigtable_client_from_service_account_info(client_class, transport_name) [ (transports.BigtableGrpcTransport, "grpc"), (transports.BigtableGrpcAsyncIOTransport, "grpc_asyncio"), + (transports.PooledBigtableGrpcAsyncIOTransport, "pooled_grpc_asyncio"), (transports.BigtableRestTransport, "rest"), ], ) @@ -152,6 +154,7 @@ def test_bigtable_client_service_account_always_use_jwt( [ (BigtableClient, "grpc"), (BigtableAsyncClient, "grpc_asyncio"), + (BigtableAsyncClient, "pooled_grpc_asyncio"), (BigtableClient, "rest"), ], ) @@ -175,7 +178,7 @@ def test_bigtable_client_from_service_account_file(client_class, transport_name) assert client.transport._host == ( "bigtable.googleapis.com:443" - if transport_name in ["grpc", "grpc_asyncio"] + if transport_name in ["grpc", "grpc_asyncio", "pooled_grpc_asyncio"] else "https://bigtable.googleapis.com" ) @@ -197,6 +200,11 @@ def test_bigtable_client_get_transport_class(): [ (BigtableClient, transports.BigtableGrpcTransport, "grpc"), (BigtableAsyncClient, transports.BigtableGrpcAsyncIOTransport, "grpc_asyncio"), + ( + BigtableAsyncClient, + transports.PooledBigtableGrpcAsyncIOTransport, + "pooled_grpc_asyncio", + ), (BigtableClient, transports.BigtableRestTransport, "rest"), ], ) @@ -332,6 +340,12 @@ def test_bigtable_client_client_options(client_class, transport_class, transport "grpc_asyncio", "true", ), + ( + BigtableAsyncClient, + transports.PooledBigtableGrpcAsyncIOTransport, + "pooled_grpc_asyncio", + "true", + ), (BigtableClient, transports.BigtableGrpcTransport, "grpc", "false"), ( BigtableAsyncClient, @@ -339,6 +353,12 @@ def test_bigtable_client_client_options(client_class, transport_class, transport "grpc_asyncio", "false", ), + ( + BigtableAsyncClient, + transports.PooledBigtableGrpcAsyncIOTransport, + "pooled_grpc_asyncio", + "false", + ), (BigtableClient, transports.BigtableRestTransport, "rest", "true"), (BigtableClient, transports.BigtableRestTransport, "rest", "false"), ], @@ -530,6 +550,11 @@ def test_bigtable_client_get_mtls_endpoint_and_cert_source(client_class): [ (BigtableClient, transports.BigtableGrpcTransport, "grpc"), (BigtableAsyncClient, transports.BigtableGrpcAsyncIOTransport, "grpc_asyncio"), + ( + BigtableAsyncClient, + transports.PooledBigtableGrpcAsyncIOTransport, + "pooled_grpc_asyncio", + ), (BigtableClient, transports.BigtableRestTransport, "rest"), ], ) @@ -566,6 +591,12 @@ def test_bigtable_client_client_options_scopes( "grpc_asyncio", grpc_helpers_async, ), + ( + BigtableAsyncClient, + transports.PooledBigtableGrpcAsyncIOTransport, + "pooled_grpc_asyncio", + grpc_helpers_async, + ), (BigtableClient, transports.BigtableRestTransport, "rest", None), ], ) @@ -712,6 +743,35 @@ def test_read_rows(request_type, transport: str = "grpc"): assert isinstance(message, bigtable.ReadRowsResponse) +def test_read_rows_pooled_rotation(transport: str = "pooled_grpc_asyncio"): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.read_rows(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.read_rows(request) + assert next_channel.call_count == i + + def test_read_rows_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -931,6 +991,35 @@ def test_sample_row_keys(request_type, transport: str = "grpc"): assert isinstance(message, bigtable.SampleRowKeysResponse) +def test_sample_row_keys_pooled_rotation(transport: str = "pooled_grpc_asyncio"): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.sample_row_keys(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.sample_row_keys(request) + assert next_channel.call_count == i + + def test_sample_row_keys_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -1149,6 +1238,35 @@ def test_mutate_row(request_type, transport: str = "grpc"): assert isinstance(response, bigtable.MutateRowResponse) +def test_mutate_row_pooled_rotation(transport: str = "pooled_grpc_asyncio"): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.mutate_row(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.mutate_row(request) + assert next_channel.call_count == i + + def test_mutate_row_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -1412,6 +1530,35 @@ def test_mutate_rows(request_type, transport: str = "grpc"): assert isinstance(message, bigtable.MutateRowsResponse) +def test_mutate_rows_pooled_rotation(transport: str = "pooled_grpc_asyncio"): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.mutate_rows(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.mutate_rows(request) + assert next_channel.call_count == i + + def test_mutate_rows_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -1645,6 +1792,35 @@ def test_check_and_mutate_row(request_type, transport: str = "grpc"): assert response.predicate_matched is True +def test_check_and_mutate_row_pooled_rotation(transport: str = "pooled_grpc_asyncio"): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.check_and_mutate_row(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.check_and_mutate_row(request) + assert next_channel.call_count == i + + def test_check_and_mutate_row_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -2022,6 +2198,35 @@ def test_ping_and_warm(request_type, transport: str = "grpc"): assert isinstance(response, bigtable.PingAndWarmResponse) +def test_ping_and_warm_pooled_rotation(transport: str = "pooled_grpc_asyncio"): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.ping_and_warm(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.ping_and_warm(request) + assert next_channel.call_count == i + + def test_ping_and_warm_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -2242,6 +2447,35 @@ def test_read_modify_write_row(request_type, transport: str = "grpc"): assert isinstance(response, bigtable.ReadModifyWriteRowResponse) +def test_read_modify_write_row_pooled_rotation(transport: str = "pooled_grpc_asyncio"): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.read_modify_write_row(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.read_modify_write_row(request) + assert next_channel.call_count == i + + def test_read_modify_write_row_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -2501,6 +2735,37 @@ def test_generate_initial_change_stream_partitions( ) +def test_generate_initial_change_stream_partitions_pooled_rotation( + transport: str = "pooled_grpc_asyncio", +): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.generate_initial_change_stream_partitions(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.generate_initial_change_stream_partitions(request) + assert next_channel.call_count == i + + def test_generate_initial_change_stream_partitions_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -2760,6 +3025,35 @@ def test_read_change_stream(request_type, transport: str = "grpc"): assert isinstance(message, bigtable.ReadChangeStreamResponse) +def test_read_change_stream_pooled_rotation(transport: str = "pooled_grpc_asyncio"): + with mock.patch.object( + transports.pooled_grpc_asyncio.PooledChannel, "next_channel" + ) as next_channel: + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = {} + + channel = client.transport._grpc_channel._pool[ + client.transport._grpc_channel._next_idx + ] + next_channel.return_value = channel + + response = client.read_change_stream(request) + + # Establish that next_channel was called + next_channel.assert_called_once() + # Establish that subsequent calls all call next_channel + starting_idx = client.transport._grpc_channel._next_idx + for i in range(2, 10): + response = client.read_change_stream(request) + assert next_channel.call_count == i + + def test_read_change_stream_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -5663,6 +5957,7 @@ def test_transport_get_channel(): [ transports.BigtableGrpcTransport, transports.BigtableGrpcAsyncIOTransport, + transports.PooledBigtableGrpcAsyncIOTransport, transports.BigtableRestTransport, ], ) @@ -5810,6 +6105,7 @@ def test_bigtable_auth_adc(): [ transports.BigtableGrpcTransport, transports.BigtableGrpcAsyncIOTransport, + transports.PooledBigtableGrpcAsyncIOTransport, ], ) def test_bigtable_transport_auth_adc(transport_class): @@ -5837,6 +6133,7 @@ def test_bigtable_transport_auth_adc(transport_class): [ transports.BigtableGrpcTransport, transports.BigtableGrpcAsyncIOTransport, + transports.PooledBigtableGrpcAsyncIOTransport, transports.BigtableRestTransport, ], ) @@ -5939,6 +6236,61 @@ def test_bigtable_grpc_transport_client_cert_source_for_mtls(transport_class): ) +@pytest.mark.parametrize( + "transport_class", [transports.PooledBigtableGrpcAsyncIOTransport] +) +def test_bigtable_pooled_grpc_transport_client_cert_source_for_mtls(transport_class): + cred = ga_credentials.AnonymousCredentials() + + # test with invalid pool size + with pytest.raises(ValueError): + transport_class( + host="squid.clam.whelk", + credentials=cred, + pool_size=0, + ) + + # Check ssl_channel_credentials is used if provided. + for pool_num in range(1, 5): + with mock.patch.object( + transport_class, "create_channel" + ) as mock_create_channel: + mock_ssl_channel_creds = mock.Mock() + transport_class( + host="squid.clam.whelk", + credentials=cred, + ssl_channel_credentials=mock_ssl_channel_creds, + pool_size=pool_num, + ) + mock_create_channel.assert_called_with( + pool_num, + "squid.clam.whelk:443", + credentials=cred, + credentials_file=None, + scopes=None, + ssl_credentials=mock_ssl_channel_creds, + quota_project_id=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + assert mock_create_channel.call_count == 1 + + # Check if ssl_channel_credentials is not provided, then client_cert_source_for_mtls + # is used. + with mock.patch.object(transport_class, "create_channel", return_value=mock.Mock()): + with mock.patch("grpc.ssl_channel_credentials") as mock_ssl_cred: + transport_class( + credentials=cred, + client_cert_source_for_mtls=client_cert_source_callback, + ) + expected_cert, expected_key = client_cert_source_callback() + mock_ssl_cred.assert_called_once_with( + certificate_chain=expected_cert, private_key=expected_key + ) + + def test_bigtable_http_transport_client_cert_source_for_mtls(): cred = ga_credentials.AnonymousCredentials() with mock.patch( @@ -5955,6 +6307,7 @@ def test_bigtable_http_transport_client_cert_source_for_mtls(): [ "grpc", "grpc_asyncio", + "pooled_grpc_asyncio", "rest", ], ) @@ -5968,7 +6321,7 @@ def test_bigtable_host_no_port(transport_name): ) assert client.transport._host == ( "bigtable.googleapis.com:443" - if transport_name in ["grpc", "grpc_asyncio"] + if transport_name in ["grpc", "grpc_asyncio", "pooled_grpc_asyncio"] else "https://bigtable.googleapis.com" ) @@ -5978,6 +6331,7 @@ def test_bigtable_host_no_port(transport_name): [ "grpc", "grpc_asyncio", + "pooled_grpc_asyncio", "rest", ], ) @@ -5991,7 +6345,7 @@ def test_bigtable_host_with_port(transport_name): ) assert client.transport._host == ( "bigtable.googleapis.com:8000" - if transport_name in ["grpc", "grpc_asyncio"] + if transport_name in ["grpc", "grpc_asyncio", "pooled_grpc_asyncio"] else "https://bigtable.googleapis.com:8000" ) @@ -6347,6 +6701,24 @@ async def test_transport_close_async(): async with client: close.assert_not_called() close.assert_called_once() + close.assert_awaited() + + +@pytest.mark.asyncio +async def test_pooled_transport_close_async(): + client = BigtableAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="pooled_grpc_asyncio", + ) + num_channels = len(client.transport._grpc_channel._pool) + with mock.patch.object( + type(client.transport._grpc_channel._pool[0]), "close" + ) as close: + async with client: + close.assert_not_called() + close.assert_called() + assert close.call_count == num_channels + close.assert_awaited() def test_transport_close(): @@ -6413,3 +6785,128 @@ def test_api_key_credentials(client_class, transport_class): always_use_jwt_access=True, api_audience=None, ) + + +@pytest.mark.asyncio +async def test_pooled_transport_replace_default(): + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="pooled_grpc_asyncio", + ) + num_channels = len(client.transport._grpc_channel._pool) + for replace_idx in range(num_channels): + prev_pool = [channel for channel in client.transport._grpc_channel._pool] + grace_period = 4 + with mock.patch.object( + type(client.transport._grpc_channel._pool[0]), "close" + ) as close: + await client.transport.replace_channel(replace_idx, grace=grace_period) + close.assert_called_once() + close.assert_awaited() + close.assert_called_with(grace=grace_period) + assert isinstance( + client.transport._grpc_channel._pool[replace_idx], grpc.aio.Channel + ) + # only the specified channel should be replaced + for i in range(num_channels): + if i == replace_idx: + assert client.transport._grpc_channel._pool[i] != prev_pool[i] + else: + assert client.transport._grpc_channel._pool[i] == prev_pool[i] + with pytest.raises(ValueError): + await client.transport.replace_channel(num_channels + 1) + with pytest.raises(ValueError): + await client.transport.replace_channel(-1) + + +@pytest.mark.asyncio +async def test_pooled_transport_replace_explicit(): + client = BigtableClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="pooled_grpc_asyncio", + ) + num_channels = len(client.transport._grpc_channel._pool) + for replace_idx in range(num_channels): + prev_pool = [channel for channel in client.transport._grpc_channel._pool] + grace_period = 0 + with mock.patch.object( + type(client.transport._grpc_channel._pool[0]), "close" + ) as close: + new_channel = grpc.aio.insecure_channel("localhost:8080") + await client.transport.replace_channel( + replace_idx, grace=grace_period, new_channel=new_channel + ) + close.assert_called_once() + close.assert_awaited() + close.assert_called_with(grace=grace_period) + assert client.transport._grpc_channel._pool[replace_idx] == new_channel + # only the specified channel should be replaced + for i in range(num_channels): + if i == replace_idx: + assert client.transport._grpc_channel._pool[i] != prev_pool[i] + else: + assert client.transport._grpc_channel._pool[i] == prev_pool[i] + + +def test_pooled_transport_next_channel(): + num_channels = 10 + transport = transports.PooledBigtableGrpcAsyncIOTransport( + credentials=ga_credentials.AnonymousCredentials(), + pool_size=num_channels, + ) + assert len(transport._grpc_channel._pool) == num_channels + transport._grpc_channel._next_idx = 0 + # rotate through all channels multiple times + num_cycles = 4 + for _ in range(num_cycles): + for i in range(num_channels - 1): + assert transport._grpc_channel._next_idx == i + got_channel = transport._grpc_channel.next_channel() + assert got_channel == transport._grpc_channel._pool[i] + assert transport._grpc_channel._next_idx == (i + 1) + # test wrap around + assert transport._grpc_channel._next_idx == num_channels - 1 + got_channel = transport._grpc_channel.next_channel() + assert got_channel == transport._grpc_channel._pool[num_channels - 1] + assert transport._grpc_channel._next_idx == 0 + + +def test_pooled_transport_pool_unique_channels(): + num_channels = 50 + + transport = transports.PooledBigtableGrpcAsyncIOTransport( + credentials=ga_credentials.AnonymousCredentials(), + pool_size=num_channels, + ) + channel_list = [channel for channel in transport._grpc_channel._pool] + channel_set = set(channel_list) + assert len(channel_list) == num_channels + assert len(channel_set) == num_channels + for channel in channel_list: + assert isinstance(channel, grpc.aio.Channel) + + +def test_pooled_transport_pool_creation(): + # channels should be created with the specified options + num_channels = 50 + creds = ga_credentials.AnonymousCredentials() + scopes = ["test1", "test2"] + quota_project_id = "test3" + host = "testhost:8080" + with mock.patch( + "google.api_core.grpc_helpers_async.create_channel" + ) as create_channel: + transport = transports.PooledBigtableGrpcAsyncIOTransport( + credentials=creds, + pool_size=num_channels, + scopes=scopes, + quota_project_id=quota_project_id, + host=host, + ) + assert create_channel.call_count == num_channels + for i in range(num_channels): + kwargs = create_channel.call_args_list[i][1] + assert kwargs["target"] == host + assert kwargs["credentials"] == creds + assert kwargs["scopes"] == scopes + assert kwargs["quota_project_id"] == quota_project_id diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py new file mode 100644 index 000000000..ca7220800 --- /dev/null +++ b/tests/unit/test_client.py @@ -0,0 +1,814 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + + +import grpc +import asyncio +import re +import sys + +from google.auth.credentials import AnonymousCredentials +import pytest + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock + from unittest.mock import AsyncMock # type: ignore +except ImportError: # pragma: NO COVER + import mock # type: ignore + from mock import AsyncMock # type: ignore + +VENEER_HEADER_REGEX = re.compile( + r"gapic\/[0-9]+\.[\w.-]+ gax\/[0-9]+\.[\w.-]+ gccl\/[0-9]+\.[\w.-]+ gl-python\/[0-9]+\.[\w.-]+ grpc\/[0-9]+\.[\w.-]+" +) + + +class TestBigtableDataClient: + def _get_target_class(self): + from google.cloud.bigtable.client import BigtableDataClient + + return BigtableDataClient + + def _make_one(self, *args, **kwargs): + return self._get_target_class()(*args, **kwargs) + + @pytest.mark.asyncio + async def test_ctor(self): + expected_project = "project-id" + expected_pool_size = 11 + expected_credentials = AnonymousCredentials() + client = self._make_one( + project="project-id", + pool_size=expected_pool_size, + credentials=expected_credentials, + ) + await asyncio.sleep(0.1) + assert client.project == expected_project + assert len(client.transport._grpc_channel._pool) == expected_pool_size + assert not client._active_instances + assert len(client._channel_refresh_tasks) == expected_pool_size + assert client.transport._credentials == expected_credentials + await client.close() + + @pytest.mark.asyncio + async def test_ctor_super_inits(self): + from google.cloud.bigtable_v2.services.bigtable.async_client import ( + BigtableAsyncClient, + ) + from google.cloud.client import ClientWithProject + from google.api_core import client_options as client_options_lib + + project = "project-id" + pool_size = 11 + credentials = AnonymousCredentials() + client_options = {"api_endpoint": "foo.bar:1234"} + options_parsed = client_options_lib.from_dict(client_options) + transport_str = f"pooled_grpc_asyncio_{pool_size}" + with mock.patch.object(BigtableAsyncClient, "__init__") as bigtable_client_init: + bigtable_client_init.return_value = None + with mock.patch.object( + ClientWithProject, "__init__" + ) as client_project_init: + client_project_init.return_value = None + try: + self._make_one( + project=project, + pool_size=pool_size, + credentials=credentials, + client_options=options_parsed, + ) + except AttributeError: + pass + # test gapic superclass init was called + assert bigtable_client_init.call_count == 1 + kwargs = bigtable_client_init.call_args[1] + assert kwargs["transport"] == transport_str + assert kwargs["credentials"] == credentials + assert kwargs["client_options"] == options_parsed + # test mixin superclass init was called + assert client_project_init.call_count == 1 + kwargs = client_project_init.call_args[1] + assert kwargs["project"] == project + assert kwargs["credentials"] == credentials + assert kwargs["client_options"] == options_parsed + + @pytest.mark.asyncio + async def test_ctor_dict_options(self): + from google.cloud.bigtable_v2.services.bigtable.async_client import ( + BigtableAsyncClient, + ) + from google.api_core.client_options import ClientOptions + from google.cloud.bigtable.client import BigtableDataClient + + client_options = {"api_endpoint": "foo.bar:1234"} + with mock.patch.object(BigtableAsyncClient, "__init__") as bigtable_client_init: + try: + self._make_one(client_options=client_options) + except TypeError: + pass + bigtable_client_init.assert_called_once() + kwargs = bigtable_client_init.call_args[1] + called_options = kwargs["client_options"] + assert called_options.api_endpoint == "foo.bar:1234" + assert isinstance(called_options, ClientOptions) + with mock.patch.object( + BigtableDataClient, "start_background_channel_refresh" + ) as start_background_refresh: + client = self._make_one(client_options=client_options) + start_background_refresh.assert_called_once() + await client.close() + + @pytest.mark.asyncio + async def test_veneer_grpc_headers(self): + # client_info should be populated with headers to + # detect as a veneer client + patch = mock.patch("google.api_core.gapic_v1.method.wrap_method") + with patch as gapic_mock: + client = self._make_one(project="project-id") + wrapped_call_list = gapic_mock.call_args_list + assert len(wrapped_call_list) > 0 + # each wrapped call should have veneer headers + for call in wrapped_call_list: + client_info = call.kwargs["client_info"] + assert client_info is not None, f"{call} has no client_info" + wrapped_user_agent_sorted = " ".join( + sorted(client_info.to_user_agent().split(" ")) + ) + assert VENEER_HEADER_REGEX.match( + wrapped_user_agent_sorted + ), f"'{wrapped_user_agent_sorted}' does not match {VENEER_HEADER_REGEX}" + await client.close() + + @pytest.mark.asyncio + async def test_channel_pool_creation(self): + pool_size = 14 + with mock.patch( + "google.api_core.grpc_helpers_async.create_channel" + ) as create_channel: + create_channel.return_value = AsyncMock() + client = self._make_one(project="project-id", pool_size=pool_size) + assert create_channel.call_count == pool_size + await client.close() + # channels should be unique + client = self._make_one(project="project-id", pool_size=pool_size) + pool_list = list(client.transport._grpc_channel._pool) + pool_set = set(client.transport._grpc_channel._pool) + assert len(pool_list) == len(pool_set) + await client.close() + + @pytest.mark.asyncio + async def test_channel_pool_rotation(self): + from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( + PooledChannel, + ) + + pool_size = 7 + + with mock.patch.object(PooledChannel, "next_channel") as next_channel: + client = self._make_one(project="project-id", pool_size=pool_size) + assert len(client.transport._grpc_channel._pool) == pool_size + next_channel.reset_mock() + with mock.patch.object( + type(client.transport._grpc_channel._pool[0]), "unary_unary" + ) as unary_unary: + # calling an rpc `pool_size` times should use a different channel each time + channel_next = None + for i in range(pool_size): + channel_last = channel_next + channel_next = client.transport.grpc_channel._pool[i] + assert channel_last != channel_next + next_channel.return_value = channel_next + client.transport.ping_and_warm() + assert next_channel.call_count == i + 1 + unary_unary.assert_called_once() + unary_unary.reset_mock() + await client.close() + + @pytest.mark.asyncio + async def test_channel_pool_replace(self): + with mock.patch.object(asyncio, "sleep"): + pool_size = 7 + client = self._make_one(project="project-id", pool_size=pool_size) + for replace_idx in range(pool_size): + start_pool = [ + channel for channel in client.transport._grpc_channel._pool + ] + grace_period = 9 + with mock.patch.object( + type(client.transport._grpc_channel._pool[0]), "close" + ) as close: + new_channel = grpc.aio.insecure_channel("localhost:8080") + await client.transport.replace_channel( + replace_idx, grace=grace_period, new_channel=new_channel + ) + close.assert_called_once_with(grace=grace_period) + close.assert_awaited_once() + assert client.transport._grpc_channel._pool[replace_idx] == new_channel + for i in range(pool_size): + if i != replace_idx: + assert client.transport._grpc_channel._pool[i] == start_pool[i] + else: + assert client.transport._grpc_channel._pool[i] != start_pool[i] + await client.close() + + @pytest.mark.filterwarnings("ignore::RuntimeWarning") + def test_start_background_channel_refresh_sync(self): + # should raise RuntimeError if called in a sync context + client = self._make_one(project="project-id") + with pytest.raises(RuntimeError): + client.start_background_channel_refresh() + + @pytest.mark.asyncio + async def test_start_background_channel_refresh_tasks_exist(self): + # if tasks exist, should do nothing + client = self._make_one(project="project-id") + with mock.patch.object(asyncio, "create_task") as create_task: + client.start_background_channel_refresh() + create_task.assert_not_called() + await client.close() + + @pytest.mark.asyncio + @pytest.mark.parametrize("pool_size", [1, 3, 7]) + async def test_start_background_channel_refresh(self, pool_size): + # should create background tasks for each channel + client = self._make_one(project="project-id", pool_size=pool_size) + ping_and_warm = AsyncMock() + client._ping_and_warm_instances = ping_and_warm + client.start_background_channel_refresh() + assert len(client._channel_refresh_tasks) == pool_size + for task in client._channel_refresh_tasks: + assert isinstance(task, asyncio.Task) + await asyncio.sleep(0.1) + assert ping_and_warm.call_count == pool_size + for channel in client.transport._grpc_channel._pool: + ping_and_warm.assert_any_call(channel) + await client.close() + + @pytest.mark.asyncio + @pytest.mark.skipif( + sys.version_info < (3, 8), reason="Task.name requires python3.8 or higher" + ) + async def test_start_background_channel_refresh_tasks_names(self): + # if tasks exist, should do nothing + pool_size = 3 + client = self._make_one(project="project-id", pool_size=pool_size) + for i in range(pool_size): + name = client._channel_refresh_tasks[i].get_name() + assert str(i) in name + assert "BigtableDataClient channel refresh " in name + await client.close() + + @pytest.mark.asyncio + async def test__ping_and_warm_instances(self): + # test with no instances + with mock.patch.object(asyncio, "gather", AsyncMock()) as gather: + client = self._make_one(project="project-id", pool_size=1) + channel = client.transport._grpc_channel._pool[0] + await client._ping_and_warm_instances(channel) + gather.assert_called_once() + gather.assert_awaited_once() + assert not gather.call_args.args + assert gather.call_args.kwargs == {"return_exceptions": True} + # test with instances + client._active_instances = [ + "instance-1", + "instance-2", + "instance-3", + "instance-4", + ] + with mock.patch.object(asyncio, "gather", AsyncMock()) as gather: + await client._ping_and_warm_instances(channel) + gather.assert_called_once() + gather.assert_awaited_once() + assert len(gather.call_args.args) == 4 + assert gather.call_args.kwargs == {"return_exceptions": True} + for idx, call in enumerate(gather.call_args.args): + assert isinstance(call, grpc.aio.UnaryUnaryCall) + call._request["name"] = client._active_instances[idx] + await client.close() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "refresh_interval, wait_time, expected_sleep", + [ + (0, 0, 0), + (0, 1, 0), + (10, 0, 10), + (10, 5, 5), + (10, 10, 0), + (10, 15, 0), + ], + ) + async def test__manage_channel_first_sleep( + self, refresh_interval, wait_time, expected_sleep + ): + # first sleep time should be `refresh_interval` seconds after client init + import time + + with mock.patch.object(time, "time") as time: + time.return_value = 0 + with mock.patch.object(asyncio, "sleep") as sleep: + sleep.side_effect = asyncio.CancelledError + try: + client = self._make_one(project="project-id") + client._channel_init_time = -wait_time + await client._manage_channel(0, refresh_interval, refresh_interval) + except asyncio.CancelledError: + pass + sleep.assert_called_once() + call_time = sleep.call_args[0][0] + assert ( + abs(call_time - expected_sleep) < 0.1 + ), f"refresh_interval: {refresh_interval}, wait_time: {wait_time}, expected_sleep: {expected_sleep}" + await client.close() + + @pytest.mark.asyncio + async def test__manage_channel_ping_and_warm(self): + from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( + PooledBigtableGrpcAsyncIOTransport, + ) + + # should ping an warm all new channels, and old channels if sleeping + client = self._make_one(project="project-id") + new_channel = grpc.aio.insecure_channel("localhost:8080") + with mock.patch.object(asyncio, "sleep"): + create_channel = mock.Mock() + create_channel.return_value = new_channel + client.transport.grpc_channel._create_channel = create_channel + with mock.patch.object( + PooledBigtableGrpcAsyncIOTransport, "replace_channel" + ) as replace_channel: + replace_channel.side_effect = asyncio.CancelledError + # should ping and warm old channel then new if sleep > 0 + with mock.patch.object( + type(self._make_one()), "_ping_and_warm_instances" + ) as ping_and_warm: + try: + channel_idx = 2 + old_channel = client.transport._grpc_channel._pool[channel_idx] + await client._manage_channel(channel_idx, 10) + except asyncio.CancelledError: + pass + assert ping_and_warm.call_count == 2 + assert old_channel != new_channel + called_with = [call[0][0] for call in ping_and_warm.call_args_list] + assert old_channel in called_with + assert new_channel in called_with + # should ping and warm instantly new channel only if not sleeping + with mock.patch.object( + type(self._make_one()), "_ping_and_warm_instances" + ) as ping_and_warm: + try: + await client._manage_channel(0, 0, 0) + except asyncio.CancelledError: + pass + ping_and_warm.assert_called_once_with(new_channel) + await client.close() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "refresh_interval, num_cycles, expected_sleep", + [ + (None, 1, 60 * 35), + (10, 10, 100), + (10, 1, 10), + ], + ) + async def test__manage_channel_sleeps( + self, refresh_interval, num_cycles, expected_sleep + ): + # make sure that sleeps work as expected + import time + import random + + channel_idx = 1 + with mock.patch.object(random, "uniform") as uniform: + uniform.side_effect = lambda min_, max_: min_ + with mock.patch.object(time, "time") as time: + time.return_value = 0 + with mock.patch.object(asyncio, "sleep") as sleep: + sleep.side_effect = [None for i in range(num_cycles - 1)] + [ + asyncio.CancelledError + ] + try: + client = self._make_one(project="project-id") + if refresh_interval is not None: + await client._manage_channel( + channel_idx, refresh_interval, refresh_interval + ) + else: + await client._manage_channel(channel_idx) + except asyncio.CancelledError: + pass + assert sleep.call_count == num_cycles + total_sleep = sum([call[0][0] for call in sleep.call_args_list]) + assert ( + abs(total_sleep - expected_sleep) < 0.1 + ), f"refresh_interval={refresh_interval}, num_cycles={num_cycles}, expected_sleep={expected_sleep}" + await client.close() + + @pytest.mark.asyncio + async def test__manage_channel_random(self): + import random + + with mock.patch.object(asyncio, "sleep") as sleep: + with mock.patch.object(random, "uniform") as uniform: + uniform.return_value = 0 + try: + uniform.side_effect = asyncio.CancelledError + client = self._make_one(project="project-id", pool_size=1) + except asyncio.CancelledError: + uniform.side_effect = None + uniform.reset_mock() + sleep.reset_mock() + min_val = 200 + max_val = 205 + uniform.side_effect = lambda min_, max_: min_ + sleep.side_effect = [None, None, asyncio.CancelledError] + try: + await client._manage_channel(0, min_val, max_val) + except asyncio.CancelledError: + pass + assert uniform.call_count == 2 + uniform_args = [call[0] for call in uniform.call_args_list] + for found_min, found_max in uniform_args: + assert found_min == min_val + assert found_max == max_val + + @pytest.mark.asyncio + @pytest.mark.parametrize("num_cycles", [0, 1, 10, 100]) + async def test__manage_channel_refresh(self, num_cycles): + # make sure that channels are properly refreshed + from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( + PooledBigtableGrpcAsyncIOTransport, + ) + from google.api_core import grpc_helpers_async + + expected_grace = 9 + expected_refresh = 0.5 + channel_idx = 1 + new_channel = grpc.aio.insecure_channel("localhost:8080") + + with mock.patch.object( + PooledBigtableGrpcAsyncIOTransport, "replace_channel" + ) as replace_channel: + with mock.patch.object(asyncio, "sleep") as sleep: + sleep.side_effect = [None for i in range(num_cycles)] + [ + asyncio.CancelledError + ] + with mock.patch.object( + grpc_helpers_async, "create_channel" + ) as create_channel: + create_channel.return_value = new_channel + client = self._make_one(project="project-id") + create_channel.reset_mock() + try: + await client._manage_channel( + channel_idx, + refresh_interval_min=expected_refresh, + refresh_interval_max=expected_refresh, + grace_period=expected_grace, + ) + except asyncio.CancelledError: + pass + assert sleep.call_count == num_cycles + 1 + assert create_channel.call_count == num_cycles + assert replace_channel.call_count == num_cycles + for call in replace_channel.call_args_list: + args, kwargs = call + assert args[0] == channel_idx + assert kwargs["grace"] == expected_grace + assert kwargs["new_channel"] == new_channel + await client.close() + + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore::RuntimeWarning") + async def test__register_instance(self): + # create the client without calling start_background_channel_refresh + with mock.patch.object(asyncio, "get_running_loop") as get_event_loop: + get_event_loop.side_effect = RuntimeError("no event loop") + client = self._make_one(project="project-id") + assert not client._channel_refresh_tasks + # first call should start background refresh + assert client._active_instances == set() + await client._register_instance("instance-1", mock.Mock()) + assert len(client._active_instances) == 1 + assert client._active_instances == {"projects/project-id/instances/instance-1"} + assert client._channel_refresh_tasks + # next call should not + with mock.patch.object( + type(self._make_one()), "start_background_channel_refresh" + ) as refresh_mock: + await client._register_instance("instance-2", mock.Mock()) + assert len(client._active_instances) == 2 + assert client._active_instances == { + "projects/project-id/instances/instance-1", + "projects/project-id/instances/instance-2", + } + refresh_mock.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore::RuntimeWarning") + async def test__register_instance_ping_and_warm(self): + # should ping and warm each new instance + pool_size = 7 + with mock.patch.object(asyncio, "get_running_loop") as get_event_loop: + get_event_loop.side_effect = RuntimeError("no event loop") + client = self._make_one(project="project-id", pool_size=pool_size) + # first call should start background refresh + assert not client._channel_refresh_tasks + await client._register_instance("instance-1", mock.Mock()) + client = self._make_one(project="project-id", pool_size=pool_size) + assert len(client._channel_refresh_tasks) == pool_size + assert not client._active_instances + # next calls should trigger ping and warm + with mock.patch.object( + type(self._make_one()), "_ping_and_warm_instances" + ) as ping_mock: + # new instance should trigger ping and warm + await client._register_instance("instance-2", mock.Mock()) + assert ping_mock.call_count == pool_size + await client._register_instance("instance-3", mock.Mock()) + assert ping_mock.call_count == pool_size * 2 + # duplcate instances should not trigger ping and warm + await client._register_instance("instance-3", mock.Mock()) + assert ping_mock.call_count == pool_size * 2 + await client.close() + + @pytest.mark.asyncio + async def test__remove_instance_registration(self): + client = self._make_one(project="project-id") + table = mock.Mock() + await client._register_instance("instance-1", table) + await client._register_instance("instance-2", table) + assert len(client._active_instances) == 2 + assert len(client._instance_owners.keys()) == 2 + instance_1_path = client._gapic_client.instance_path( + client.project, "instance-1" + ) + instance_2_path = client._gapic_client.instance_path( + client.project, "instance-2" + ) + assert len(client._instance_owners[instance_1_path]) == 1 + assert list(client._instance_owners[instance_1_path])[0] == id(table) + assert len(client._instance_owners[instance_2_path]) == 1 + assert list(client._instance_owners[instance_2_path])[0] == id(table) + success = await client._remove_instance_registration("instance-1", table) + assert success + assert len(client._active_instances) == 1 + assert len(client._instance_owners[instance_1_path]) == 0 + assert len(client._instance_owners[instance_2_path]) == 1 + assert client._active_instances == {"projects/project-id/instances/instance-2"} + success = await client._remove_instance_registration("nonexistant", table) + assert not success + assert len(client._active_instances) == 1 + await client.close() + + @pytest.mark.asyncio + async def test__multiple_table_registration(self): + async with self._make_one(project="project-id") as client: + async with client.get_table("instance_1", "table_1") as table_1: + instance_1_path = client._gapic_client.instance_path( + client.project, "instance_1" + ) + assert len(client._instance_owners[instance_1_path]) == 1 + assert len(client._active_instances) == 1 + assert id(table_1) in client._instance_owners[instance_1_path] + async with client.get_table("instance_1", "table_2") as table_2: + assert len(client._instance_owners[instance_1_path]) == 2 + assert len(client._active_instances) == 1 + assert id(table_1) in client._instance_owners[instance_1_path] + assert id(table_2) in client._instance_owners[instance_1_path] + # table_2 should be unregistered, but instance should still be active + assert len(client._active_instances) == 1 + assert instance_1_path in client._active_instances + assert id(table_2) not in client._instance_owners[instance_1_path] + # both tables are gone. instance should be unregistered + assert len(client._active_instances) == 0 + assert instance_1_path not in client._active_instances + assert len(client._instance_owners[instance_1_path]) == 0 + + @pytest.mark.asyncio + async def test__multiple_instance_registration(self): + async with self._make_one(project="project-id") as client: + async with client.get_table("instance_1", "table_1") as table_1: + async with client.get_table("instance_2", "table_2") as table_2: + instance_1_path = client._gapic_client.instance_path( + client.project, "instance_1" + ) + instance_2_path = client._gapic_client.instance_path( + client.project, "instance_2" + ) + assert len(client._instance_owners[instance_1_path]) == 1 + assert len(client._instance_owners[instance_2_path]) == 1 + assert len(client._active_instances) == 2 + assert id(table_1) in client._instance_owners[instance_1_path] + assert id(table_2) in client._instance_owners[instance_2_path] + # instance2 should be unregistered, but instance1 should still be active + assert len(client._active_instances) == 1 + assert instance_1_path in client._active_instances + assert len(client._instance_owners[instance_2_path]) == 0 + assert len(client._instance_owners[instance_1_path]) == 1 + assert id(table_1) in client._instance_owners[instance_1_path] + # both tables are gone. instances should both be unregistered + assert len(client._active_instances) == 0 + assert len(client._instance_owners[instance_1_path]) == 0 + assert len(client._instance_owners[instance_2_path]) == 0 + + @pytest.mark.asyncio + async def test_get_table(self): + from google.cloud.bigtable.client import Table + + client = self._make_one(project="project-id") + assert not client._active_instances + expected_table_id = "table-id" + expected_instance_id = "instance-id" + expected_app_profile_id = "app-profile-id" + table = client.get_table( + expected_instance_id, + expected_table_id, + expected_app_profile_id, + ) + await asyncio.sleep(0) + assert isinstance(table, Table) + assert table.table_id == expected_table_id + assert ( + table.table_name + == f"projects/{client.project}/instances/{expected_instance_id}/tables/{expected_table_id}" + ) + assert table.instance_id == expected_instance_id + assert ( + table.instance_name + == f"projects/{client.project}/instances/{expected_instance_id}" + ) + assert table.app_profile_id == expected_app_profile_id + assert table.client is client + assert table.instance_name in client._active_instances + await client.close() + + @pytest.mark.asyncio + async def test_get_table_context_manager(self): + from google.cloud.bigtable.client import Table + + expected_table_id = "table-id" + expected_instance_id = "instance-id" + expected_app_profile_id = "app-profile-id" + expected_project_id = "project-id" + + with mock.patch.object(Table, "close") as close_mock: + async with self._make_one(project=expected_project_id) as client: + async with client.get_table( + expected_instance_id, + expected_table_id, + expected_app_profile_id, + ) as table: + await asyncio.sleep(0) + assert isinstance(table, Table) + assert table.table_id == expected_table_id + assert ( + table.table_name + == f"projects/{expected_project_id}/instances/{expected_instance_id}/tables/{expected_table_id}" + ) + assert table.instance_id == expected_instance_id + assert ( + table.instance_name + == f"projects/{expected_project_id}/instances/{expected_instance_id}" + ) + assert table.app_profile_id == expected_app_profile_id + assert table.client is client + assert table.instance_name in client._active_instances + assert close_mock.call_count == 1 + + @pytest.mark.asyncio + async def test_multiple_pool_sizes(self): + # should be able to create multiple clients with different pool sizes without issue + pool_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256] + for pool_size in pool_sizes: + client = self._make_one(project="project-id", pool_size=pool_size) + assert len(client._channel_refresh_tasks) == pool_size + client_duplicate = self._make_one(project="project-id", pool_size=pool_size) + assert len(client_duplicate._channel_refresh_tasks) == pool_size + assert str(pool_size) in str(client.transport) + await client.close() + await client_duplicate.close() + + @pytest.mark.asyncio + async def test_close(self): + from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( + PooledBigtableGrpcAsyncIOTransport, + ) + + pool_size = 7 + client = self._make_one(project="project-id", pool_size=pool_size) + assert len(client._channel_refresh_tasks) == pool_size + tasks_list = list(client._channel_refresh_tasks) + for task in client._channel_refresh_tasks: + assert not task.done() + with mock.patch.object( + PooledBigtableGrpcAsyncIOTransport, "close", AsyncMock() + ) as close_mock: + await client.close() + close_mock.assert_called_once() + close_mock.assert_awaited() + for task in tasks_list: + assert task.done() + assert task.cancelled() + assert client._channel_refresh_tasks == [] + + @pytest.mark.asyncio + async def test_close_with_timeout(self): + pool_size = 7 + expected_timeout = 19 + client = self._make_one(project="project-id", pool_size=pool_size) + tasks = list(client._channel_refresh_tasks) + with mock.patch.object(asyncio, "wait_for", AsyncMock()) as wait_for_mock: + await client.close(timeout=expected_timeout) + wait_for_mock.assert_called_once() + wait_for_mock.assert_awaited() + assert wait_for_mock.call_args[1]["timeout"] == expected_timeout + client._channel_refresh_tasks = tasks + await client.close() + + @pytest.mark.asyncio + async def test_context_manager(self): + # context manager should close the client cleanly + close_mock = AsyncMock() + true_close = None + async with self._make_one(project="project-id") as client: + true_close = client.close() + client.close = close_mock + for task in client._channel_refresh_tasks: + assert not task.done() + assert client.project == "project-id" + assert client._active_instances == set() + close_mock.assert_not_called() + close_mock.assert_called_once() + close_mock.assert_awaited() + # actually close the client + await true_close + + def test_client_ctor_sync(self): + # initializing client in a sync context should raise RuntimeError + from google.cloud.bigtable.client import BigtableDataClient + + with pytest.warns(RuntimeWarning) as warnings: + client = BigtableDataClient(project="project-id") + expected_warning = [w for w in warnings if "client.py" in w.filename] + assert len(expected_warning) == 1 + assert "BigtableDataClient should be started in an asyncio event loop." in str( + expected_warning[0].message + ) + assert client.project == "project-id" + assert client._channel_refresh_tasks == [] + + +class TestTable: + @pytest.mark.asyncio + async def test_table_ctor(self): + from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.client import Table + + expected_table_id = "table-id" + expected_instance_id = "instance-id" + expected_app_profile_id = "app-profile-id" + client = BigtableDataClient() + assert not client._active_instances + + table = Table( + client, + expected_instance_id, + expected_table_id, + expected_app_profile_id, + ) + await asyncio.sleep(0) + assert table.table_id == expected_table_id + assert table.instance_id == expected_instance_id + assert table.app_profile_id == expected_app_profile_id + assert table.client is client + assert table.instance_name in client._active_instances + # ensure task reaches completion + await table._register_instance_task + assert table._register_instance_task.done() + assert not table._register_instance_task.cancelled() + assert table._register_instance_task.exception() is None + await client.close() + + def test_table_ctor_sync(self): + # initializing client in a sync context should raise RuntimeError + from google.cloud.bigtable.client import Table + + client = mock.Mock() + with pytest.raises(RuntimeError) as e: + Table(client, "instance-id", "table-id") + assert e.match("Table must be created within an async event loop context.") From 3de7a68f80d3098b1b9b20c6b6fcf03eeea1d05c Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 24 May 2023 13:28:51 -0700 Subject: [PATCH 06/56] feat: implement read_rows (#762) --- .gitmodules | 3 + .kokoro/trampoline.sh | 2 +- google/cloud/bigtable/__init__.py | 2 +- google/cloud/bigtable/_read_rows.py | 637 +++++++ google/cloud/bigtable/client.py | 149 +- google/cloud/bigtable/exceptions.py | 84 +- google/cloud/bigtable/iterators.py | 129 ++ google/cloud/bigtable/mutations.py | 13 +- google/cloud/bigtable/mutations_batcher.py | 6 +- .../cloud/bigtable/read_modify_write_rules.py | 12 +- google/cloud/bigtable/read_rows_query.py | 3 +- google/cloud/bigtable/row.py | 176 +- noxfile.py | 38 +- python-api-core | 1 + testing/constraints-3.7.txt | 3 +- tests/system/test_system.py | 164 ++ tests/unit/read-rows-acceptance-test.json | 1665 +++++++++++++++++ tests/unit/test__read_rows.py | 1001 ++++++++++ tests/unit/test_client.py | 486 ++++- tests/unit/test_iterators.py | 251 +++ tests/unit/test_read_rows_acceptance.py | 314 ++++ tests/unit/test_row.py | 16 +- 22 files changed, 4985 insertions(+), 170 deletions(-) create mode 100644 google/cloud/bigtable/_read_rows.py create mode 100644 google/cloud/bigtable/iterators.py create mode 160000 python-api-core create mode 100644 tests/unit/read-rows-acceptance-test.json create mode 100644 tests/unit/test__read_rows.py create mode 100644 tests/unit/test_iterators.py create mode 100644 tests/unit/test_read_rows_acceptance.py diff --git a/.gitmodules b/.gitmodules index 4186187f4..5fa9b1ed5 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,6 @@ +[submodule "python-api-core"] + path = python-api-core + url = git@github.com:googleapis/python-api-core.git [submodule "gapic-generator-fork"] path = gapic-generator-fork url = git@github.com:googleapis/gapic-generator-python.git diff --git a/.kokoro/trampoline.sh b/.kokoro/trampoline.sh index f39236e94..a4241db23 100755 --- a/.kokoro/trampoline.sh +++ b/.kokoro/trampoline.sh @@ -25,4 +25,4 @@ function cleanup() { trap cleanup EXIT $(dirname $0)/populate-secrets.sh # Secret Manager secrets. -python3 "${KOKORO_GFILE_DIR}/trampoline_v1.py" \ No newline at end of file +python3 "${KOKORO_GFILE_DIR}/trampoline_v1.py" diff --git a/google/cloud/bigtable/__init__.py b/google/cloud/bigtable/__init__.py index c5581f813..de46f8a75 100644 --- a/google/cloud/bigtable/__init__.py +++ b/google/cloud/bigtable/__init__.py @@ -22,9 +22,9 @@ from google.cloud.bigtable.client import Table from google.cloud.bigtable.read_rows_query import ReadRowsQuery +from google.cloud.bigtable.read_rows_query import RowRange from google.cloud.bigtable.row import Row from google.cloud.bigtable.row import Cell -from google.cloud.bigtable.read_rows_query import RowRange from google.cloud.bigtable.mutations_batcher import MutationsBatcher from google.cloud.bigtable.mutations import Mutation diff --git a/google/cloud/bigtable/_read_rows.py b/google/cloud/bigtable/_read_rows.py new file mode 100644 index 000000000..1c9e02d5a --- /dev/null +++ b/google/cloud/bigtable/_read_rows.py @@ -0,0 +1,637 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from typing import ( + List, + Any, + AsyncIterable, + AsyncIterator, + AsyncGenerator, + Callable, + Awaitable, + Type, +) + +import asyncio +import time +from functools import partial +from grpc.aio import RpcContext + +from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse +from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient +from google.cloud.bigtable.row import Row, Cell, _LastScannedRow +from google.cloud.bigtable.exceptions import InvalidChunk +from google.cloud.bigtable.exceptions import _RowSetComplete +from google.api_core import retry_async as retries +from google.api_core import exceptions as core_exceptions + +""" +This module provides a set of classes for merging ReadRowsResponse chunks +into Row objects. + +- ReadRowsOperation is the highest level class, providing an interface for asynchronous + merging end-to-end +- StateMachine is used internally to track the state of the merge, including + the current row key and the keys of the rows that have been processed. + It processes a stream of chunks, and will raise InvalidChunk if it reaches + an invalid state. +- State classes track the current state of the StateMachine, and define what + to do on the next chunk. +- RowBuilder is used by the StateMachine to build a Row object. +""" + + +class _ReadRowsOperation(AsyncIterable[Row]): + """ + ReadRowsOperation handles the logic of merging chunks from a ReadRowsResponse stream + into a stream of Row objects. + + ReadRowsOperation.merge_row_response_stream takes in a stream of ReadRowsResponse + and turns them into a stream of Row objects using an internal + StateMachine. + + ReadRowsOperation(request, client) handles row merging logic end-to-end, including + performing retries on stream errors. + """ + + def __init__( + self, + request: dict[str, Any], + client: BigtableAsyncClient, + *, + operation_timeout: float = 600.0, + per_request_timeout: float | None = None, + ): + """ + Args: + - request: the request dict to send to the Bigtable API + - client: the Bigtable client to use to make the request + - operation_timeout: the timeout to use for the entire operation, in seconds + - per_request_timeout: the timeout to use when waiting for each individual grpc request, in seconds + If not specified, defaults to operation_timeout + """ + self._last_emitted_row_key: bytes | None = None + self._emit_count = 0 + self._request = request + self.operation_timeout = operation_timeout + deadline = operation_timeout + time.monotonic() + row_limit = request.get("rows_limit", 0) + if per_request_timeout is None: + per_request_timeout = operation_timeout + # lock in paramters for retryable wrapper + self._partial_retryable = partial( + self._read_rows_retryable_attempt, + client.read_rows, + per_request_timeout, + deadline, + row_limit, + ) + predicate = retries.if_exception_type( + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + core_exceptions.Aborted, + ) + + def on_error_fn(exc): + if predicate(exc): + self.transient_errors.append(exc) + + retry = retries.AsyncRetry( + predicate=predicate, + timeout=self.operation_timeout, + initial=0.01, + multiplier=2, + maximum=60, + on_error=on_error_fn, + is_stream=True, + ) + self._stream: AsyncGenerator[Row, None] | None = retry( + self._partial_retryable + )() + # contains the list of errors that were retried + self.transient_errors: List[Exception] = [] + + def __aiter__(self) -> AsyncIterator[Row]: + """Implements the AsyncIterable interface""" + return self + + async def __anext__(self) -> Row: + """Implements the AsyncIterator interface""" + if self._stream is not None: + return await self._stream.__anext__() + else: + raise asyncio.InvalidStateError("stream is closed") + + async def aclose(self): + """Close the stream and release resources""" + if self._stream is not None: + await self._stream.aclose() + self._stream = None + self._emitted_seen_row_key = None + + async def _read_rows_retryable_attempt( + self, + gapic_fn: Callable[..., Awaitable[AsyncIterable[ReadRowsResponse]]], + per_request_timeout: float, + operation_deadline: float, + total_row_limit: int, + ) -> AsyncGenerator[Row, None]: + """ + Retryable wrapper for merge_rows. This function is called each time + a retry is attempted. + + Some fresh state is created on each retry: + - grpc network stream + - state machine to hold merge chunks received from stream + Some state is shared between retries: + - _last_emitted_row_key is used to ensure that + duplicate rows are not emitted + - request is stored and (potentially) modified on each retry + """ + if self._last_emitted_row_key is not None: + # if this is a retry, try to trim down the request to avoid ones we've already processed + try: + self._request["rows"] = _ReadRowsOperation._revise_request_rowset( + row_set=self._request.get("rows", None), + last_seen_row_key=self._last_emitted_row_key, + ) + except _RowSetComplete: + # if there are no rows left to process, we're done + # This is not expected to happen often, but could occur if + # a retry is triggered quickly after the last row is emitted + return + # revise next request's row limit based on number emitted + if total_row_limit: + new_limit = total_row_limit - self._emit_count + if new_limit == 0: + # we have hit the row limit, so we're done + return + elif new_limit < 0: + raise RuntimeError("unexpected state: emit count exceeds row limit") + else: + self._request["rows_limit"] = new_limit + params_str = f'table_name={self._request.get("table_name", "")}' + app_profile_id = self._request.get("app_profile_id", None) + if app_profile_id: + params_str = f"{params_str},app_profile_id={app_profile_id}" + time_to_deadline = operation_deadline - time.monotonic() + gapic_timeout = max(0, min(time_to_deadline, per_request_timeout)) + new_gapic_stream: RpcContext = await gapic_fn( + self._request, + timeout=gapic_timeout, + metadata=[("x-goog-request-params", params_str)], + ) + try: + state_machine = _StateMachine() + stream = _ReadRowsOperation.merge_row_response_stream( + new_gapic_stream, state_machine + ) + # run until we get a timeout or the stream is exhausted + async for new_item in stream: + if ( + self._last_emitted_row_key is not None + and new_item.row_key <= self._last_emitted_row_key + ): + raise InvalidChunk("Last emitted row key out of order") + # don't yeild _LastScannedRow markers; they + # should only update last_seen_row_key + if not isinstance(new_item, _LastScannedRow): + yield new_item + self._emit_count += 1 + self._last_emitted_row_key = new_item.row_key + if total_row_limit and self._emit_count >= total_row_limit: + return + except (Exception, GeneratorExit) as exc: + # ensure grpc stream is closed + new_gapic_stream.cancel() + raise exc + + @staticmethod + def _revise_request_rowset( + row_set: dict[str, Any] | None, + last_seen_row_key: bytes, + ) -> dict[str, Any]: + """ + Revise the rows in the request to avoid ones we've already processed. + + Args: + - row_set: the row set from the request + - last_seen_row_key: the last row key encountered + Raises: + - _RowSetComplete: if there are no rows left to process after the revision + """ + # if user is doing a whole table scan, start a new one with the last seen key + if row_set is None or ( + len(row_set.get("row_ranges", [])) == 0 + and len(row_set.get("row_keys", [])) == 0 + ): + last_seen = last_seen_row_key + return { + "row_keys": [], + "row_ranges": [{"start_key_open": last_seen}], + } + # remove seen keys from user-specific key list + row_keys: list[bytes] = row_set.get("row_keys", []) + adjusted_keys = [k for k in row_keys if k > last_seen_row_key] + # adjust ranges to ignore keys before last seen + row_ranges: list[dict[str, Any]] = row_set.get("row_ranges", []) + adjusted_ranges = [] + for row_range in row_ranges: + end_key = row_range.get("end_key_closed", None) or row_range.get( + "end_key_open", None + ) + if end_key is None or end_key > last_seen_row_key: + # end range is after last seen key + new_range = row_range.copy() + start_key = row_range.get("start_key_closed", None) or row_range.get( + "start_key_open", None + ) + if start_key is None or start_key <= last_seen_row_key: + # replace start key with last seen + new_range["start_key_open"] = last_seen_row_key + new_range.pop("start_key_closed", None) + adjusted_ranges.append(new_range) + if len(adjusted_keys) == 0 and len(adjusted_ranges) == 0: + # if the query is empty after revision, raise an exception + # this will avoid an unwanted full table scan + raise _RowSetComplete() + return {"row_keys": adjusted_keys, "row_ranges": adjusted_ranges} + + @staticmethod + async def merge_row_response_stream( + response_generator: AsyncIterable[ReadRowsResponse], + state_machine: _StateMachine, + ) -> AsyncGenerator[Row, None]: + """ + Consume chunks from a ReadRowsResponse stream into a set of Rows + + Args: + - response_generator: AsyncIterable of ReadRowsResponse objects. Typically + this is a stream of chunks from the Bigtable API + Returns: + - AsyncGenerator of Rows + Raises: + - InvalidChunk: if the chunk stream is invalid + """ + async for row_response in response_generator: + # unwrap protoplus object for increased performance + response_pb = row_response._pb + last_scanned = response_pb.last_scanned_row_key + # if the server sends a scan heartbeat, notify the state machine. + if last_scanned: + yield state_machine.handle_last_scanned_row(last_scanned) + # process new chunks through the state machine. + for chunk in response_pb.chunks: + complete_row = state_machine.handle_chunk(chunk) + if complete_row is not None: + yield complete_row + # TODO: handle request stats + if not state_machine.is_terminal_state(): + # read rows is complete, but there's still data in the merger + raise InvalidChunk("read_rows completed with partial state remaining") + + +class _StateMachine: + """ + State Machine converts chunks into Rows + + Chunks are added to the state machine via handle_chunk, which + transitions the state machine through the various states. + + When a row is complete, it will be returned from handle_chunk, + and the state machine will reset to AWAITING_NEW_ROW + + If an unexpected chunk is received for the current state, + the state machine will raise an InvalidChunk exception + + The server may send a heartbeat message indicating that it has + processed a particular row, to facilitate retries. This will be passed + to the state machine via handle_last_scanned_row, which emit a + _LastScannedRow marker to the stream. + """ + + __slots__ = ( + "current_state", + "current_family", + "current_qualifier", + "last_seen_row_key", + "adapter", + ) + + def __init__(self): + # represents either the last row emitted, or the last_scanned_key sent from backend + # all future rows should have keys > last_seen_row_key + self.last_seen_row_key: bytes | None = None + self.adapter = _RowBuilder() + self._reset_row() + + def _reset_row(self) -> None: + """ + Drops the current row and transitions to AWAITING_NEW_ROW to start a fresh one + """ + self.current_state: Type[_State] = AWAITING_NEW_ROW + self.current_family: str | None = None + self.current_qualifier: bytes | None = None + self.adapter.reset() + + def is_terminal_state(self) -> bool: + """ + Returns true if the state machine is in a terminal state (AWAITING_NEW_ROW) + + At the end of the read_rows stream, if the state machine is not in a terminal + state, an exception should be raised + """ + return self.current_state == AWAITING_NEW_ROW + + def handle_last_scanned_row(self, last_scanned_row_key: bytes) -> Row: + """ + Called by ReadRowsOperation to notify the state machine of a scan heartbeat + + Returns an empty row with the last_scanned_row_key + """ + if self.last_seen_row_key and self.last_seen_row_key >= last_scanned_row_key: + raise InvalidChunk("Last scanned row key is out of order") + if not self.current_state == AWAITING_NEW_ROW: + raise InvalidChunk("Last scanned row key received in invalid state") + scan_marker = _LastScannedRow(last_scanned_row_key) + self._handle_complete_row(scan_marker) + return scan_marker + + def handle_chunk(self, chunk: ReadRowsResponse.CellChunk) -> Row | None: + """ + Called by ReadRowsOperation to process a new chunk + + Returns a Row if the chunk completes a row, otherwise returns None + """ + if ( + self.last_seen_row_key + and chunk.row_key + and self.last_seen_row_key >= chunk.row_key + ): + raise InvalidChunk("row keys should be strictly increasing") + if chunk.reset_row: + # reset row if requested + self._handle_reset_chunk(chunk) + return None + + # process the chunk and update the state + self.current_state = self.current_state.handle_chunk(self, chunk) + if chunk.commit_row: + # check if row is complete, and return it if so + if not self.current_state == AWAITING_NEW_CELL: + raise InvalidChunk("Commit chunk received in invalid state") + complete_row = self.adapter.finish_row() + self._handle_complete_row(complete_row) + return complete_row + else: + # row is not complete, return None + return None + + def _handle_complete_row(self, complete_row: Row) -> None: + """ + Complete row, update seen keys, and move back to AWAITING_NEW_ROW + + Called by StateMachine when a commit_row flag is set on a chunk, + or when a scan heartbeat is received + """ + self.last_seen_row_key = complete_row.row_key + self._reset_row() + + def _handle_reset_chunk(self, chunk: ReadRowsResponse.CellChunk): + """ + Drop all buffers and reset the row in progress + + Called by StateMachine when a reset_row flag is set on a chunk + """ + # ensure reset chunk matches expectations + if self.current_state == AWAITING_NEW_ROW: + raise InvalidChunk("Reset chunk received when not processing row") + if chunk.row_key: + raise InvalidChunk("Reset chunk has a row key") + if _chunk_has_field(chunk, "family_name"): + raise InvalidChunk("Reset chunk has a family name") + if _chunk_has_field(chunk, "qualifier"): + raise InvalidChunk("Reset chunk has a qualifier") + if chunk.timestamp_micros: + raise InvalidChunk("Reset chunk has a timestamp") + if chunk.labels: + raise InvalidChunk("Reset chunk has labels") + if chunk.value: + raise InvalidChunk("Reset chunk has a value") + self._reset_row() + + +class _State: + """ + Represents a state the state machine can be in + + Each state is responsible for handling the next chunk, and then + transitioning to the next state + """ + + @staticmethod + def handle_chunk( + owner: _StateMachine, chunk: ReadRowsResponse.CellChunk + ) -> Type["_State"]: + raise NotImplementedError + + +class AWAITING_NEW_ROW(_State): + """ + Default state + Awaiting a chunk to start a new row + Exit states: + - AWAITING_NEW_CELL: when a chunk with a row_key is received + """ + + @staticmethod + def handle_chunk( + owner: _StateMachine, chunk: ReadRowsResponse.CellChunk + ) -> Type["_State"]: + if not chunk.row_key: + raise InvalidChunk("New row is missing a row key") + owner.adapter.start_row(chunk.row_key) + # the first chunk signals both the start of a new row and the start of a new cell, so + # force the chunk processing in the AWAITING_CELL_VALUE. + return AWAITING_NEW_CELL.handle_chunk(owner, chunk) + + +class AWAITING_NEW_CELL(_State): + """ + Represents a cell boundary witin a row + + Exit states: + - AWAITING_NEW_CELL: when the incoming cell is complete and ready for another + - AWAITING_CELL_VALUE: when the value is split across multiple chunks + """ + + @staticmethod + def handle_chunk( + owner: _StateMachine, chunk: ReadRowsResponse.CellChunk + ) -> Type["_State"]: + is_split = chunk.value_size > 0 + # track latest cell data. New chunks won't send repeated data + has_family = _chunk_has_field(chunk, "family_name") + has_qualifier = _chunk_has_field(chunk, "qualifier") + if has_family: + owner.current_family = chunk.family_name.value + if not has_qualifier: + raise InvalidChunk("New family must specify qualifier") + if has_qualifier: + owner.current_qualifier = chunk.qualifier.value + if owner.current_family is None: + raise InvalidChunk("Family not found") + + # ensure that all chunks after the first one are either missing a row + # key or the row is the same + if chunk.row_key and chunk.row_key != owner.adapter.current_key: + raise InvalidChunk("Row key changed mid row") + + if owner.current_family is None: + raise InvalidChunk("Missing family for new cell") + if owner.current_qualifier is None: + raise InvalidChunk("Missing qualifier for new cell") + + owner.adapter.start_cell( + family=owner.current_family, + qualifier=owner.current_qualifier, + labels=list(chunk.labels), + timestamp_micros=chunk.timestamp_micros, + ) + owner.adapter.cell_value(chunk.value) + # transition to new state + if is_split: + return AWAITING_CELL_VALUE + else: + # cell is complete + owner.adapter.finish_cell() + return AWAITING_NEW_CELL + + +class AWAITING_CELL_VALUE(_State): + """ + State that represents a split cell's continuation + + Exit states: + - AWAITING_NEW_CELL: when the cell is complete + - AWAITING_CELL_VALUE: when additional value chunks are required + """ + + @staticmethod + def handle_chunk( + owner: _StateMachine, chunk: ReadRowsResponse.CellChunk + ) -> Type["_State"]: + # ensure reset chunk matches expectations + if chunk.row_key: + raise InvalidChunk("In progress cell had a row key") + if _chunk_has_field(chunk, "family_name"): + raise InvalidChunk("In progress cell had a family name") + if _chunk_has_field(chunk, "qualifier"): + raise InvalidChunk("In progress cell had a qualifier") + if chunk.timestamp_micros: + raise InvalidChunk("In progress cell had a timestamp") + if chunk.labels: + raise InvalidChunk("In progress cell had labels") + is_last = chunk.value_size == 0 + owner.adapter.cell_value(chunk.value) + # transition to new state + if not is_last: + return AWAITING_CELL_VALUE + else: + # cell is complete + owner.adapter.finish_cell() + return AWAITING_NEW_CELL + + +class _RowBuilder: + """ + called by state machine to build rows + State machine makes the following guarantees: + Exactly 1 `start_row` for each row. + Exactly 1 `start_cell` for each cell. + At least 1 `cell_value` for each cell. + Exactly 1 `finish_cell` for each cell. + Exactly 1 `finish_row` for each row. + `reset` can be called at any point and can be invoked multiple times in + a row. + """ + + __slots__ = "current_key", "working_cell", "working_value", "completed_cells" + + def __init__(self): + # initialize state + self.reset() + + def reset(self) -> None: + """called when the current in progress row should be dropped""" + self.current_key: bytes | None = None + self.working_cell: Cell | None = None + self.working_value: bytearray | None = None + self.completed_cells: List[Cell] = [] + + def start_row(self, key: bytes) -> None: + """Called to start a new row. This will be called once per row""" + self.current_key = key + + def start_cell( + self, + family: str, + qualifier: bytes, + timestamp_micros: int, + labels: List[str], + ) -> None: + """called to start a new cell in a row.""" + if self.current_key is None: + raise InvalidChunk("start_cell called without a row") + self.working_value = bytearray() + self.working_cell = Cell( + b"", self.current_key, family, qualifier, timestamp_micros, labels + ) + + def cell_value(self, value: bytes) -> None: + """called multiple times per cell to concatenate the cell value""" + if self.working_value is None: + raise InvalidChunk("Cell value received before start_cell") + self.working_value.extend(value) + + def finish_cell(self) -> None: + """called once per cell to signal the end of the value (unless reset)""" + if self.working_cell is None or self.working_value is None: + raise InvalidChunk("finish_cell called before start_cell") + self.working_cell.value = bytes(self.working_value) + self.completed_cells.append(self.working_cell) + self.working_cell = None + self.working_value = None + + def finish_row(self) -> Row: + """called once per row to signal that all cells have been processed (unless reset)""" + if self.current_key is None: + raise InvalidChunk("No row in progress") + new_row = Row(self.current_key, self.completed_cells) + self.reset() + return new_row + + +def _chunk_has_field(chunk: ReadRowsResponse.CellChunk, field: str) -> bool: + """ + Returns true if the field is set on the chunk + + Required to disambiguate between empty strings and unset values + """ + try: + return chunk.HasField(field) + except ValueError: + return False diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index dfd8b16cd..e8c7b5e04 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -15,7 +15,13 @@ from __future__ import annotations -from typing import cast, Any, Optional, AsyncIterable, Set, TYPE_CHECKING +from typing import ( + cast, + Any, + Optional, + Set, + TYPE_CHECKING, +) import asyncio import grpc @@ -32,18 +38,18 @@ ) from google.cloud.client import ClientWithProject from google.api_core.exceptions import GoogleAPICallError - +from google.cloud.bigtable._read_rows import _ReadRowsOperation import google.auth.credentials import google.auth._default from google.api_core import client_options as client_options_lib - +from google.cloud.bigtable.row import Row +from google.cloud.bigtable.read_rows_query import ReadRowsQuery +from google.cloud.bigtable.iterators import ReadRowsIterator if TYPE_CHECKING: from google.cloud.bigtable.mutations import Mutation, BulkMutationsEntry from google.cloud.bigtable.mutations_batcher import MutationsBatcher - from google.cloud.bigtable.row import Row - from google.cloud.bigtable.read_rows_query import ReadRowsQuery from google.cloud.bigtable import RowKeySamples from google.cloud.bigtable.row_filters import RowFilter from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule @@ -281,11 +287,14 @@ async def _remove_instance_registration( except KeyError: return False + # TODO: revisit timeouts https://github.com/googleapis/python-bigtable/issues/782 def get_table( self, instance_id: str, table_id: str, app_profile_id: str | None = None, + default_operation_timeout: float = 600, + default_per_request_timeout: float | None = None, ) -> Table: """ Returns a table instance for making data API requests @@ -298,7 +307,14 @@ def get_table( app_profile_id: (Optional) The app profile to associate with requests. https://cloud.google.com/bigtable/docs/app-profiles """ - return Table(self, instance_id, table_id, app_profile_id) + return Table( + self, + instance_id, + table_id, + app_profile_id, + default_operation_timeout=default_operation_timeout, + default_per_request_timeout=default_per_request_timeout, + ) async def __aenter__(self): self.start_background_channel_refresh() @@ -323,6 +339,9 @@ def __init__( instance_id: str, table_id: str, app_profile_id: str | None = None, + *, + default_operation_timeout: float = 600, + default_per_request_timeout: float | None = None, ): """ Initialize a Table instance @@ -337,9 +356,24 @@ def __init__( instance_id and the client's project to fully specify the table app_profile_id: (Optional) The app profile to associate with requests. https://cloud.google.com/bigtable/docs/app-profiles + default_operation_timeout: (Optional) The default timeout, in seconds + default_per_request_timeout: (Optional) The default timeout for individual + rpc requests, in seconds Raises: - RuntimeError if called outside of an async context (no running event loop) """ + # validate timeouts + if default_operation_timeout <= 0: + raise ValueError("default_operation_timeout must be greater than 0") + if default_per_request_timeout is not None and default_per_request_timeout <= 0: + raise ValueError("default_per_request_timeout must be greater than 0") + if ( + default_per_request_timeout is not None + and default_per_request_timeout > default_operation_timeout + ): + raise ValueError( + "default_per_request_timeout must be less than default_operation_timeout" + ) self.client = client self.instance_id = instance_id self.instance_name = self.client._gapic_client.instance_path( @@ -350,6 +384,10 @@ def __init__( self.client.project, instance_id, table_id ) self.app_profile_id = app_profile_id + + self.default_operation_timeout = default_operation_timeout + self.default_per_request_timeout = default_per_request_timeout + # raises RuntimeError if called outside of an async context (no running event loop) try: self._register_instance_task = asyncio.create_task( @@ -364,68 +402,76 @@ async def read_rows_stream( self, query: ReadRowsQuery | dict[str, Any], *, - shard: bool = False, - limit: int | None, - cache_size_limit: int | None = None, - operation_timeout: int | float | None = 60, - per_row_timeout: int | float | None = 10, - idle_timeout: int | float | None = 300, - per_request_timeout: int | float | None = None, - ) -> AsyncIterable[Row]: + operation_timeout: float | None = None, + per_request_timeout: float | None = None, + ) -> ReadRowsIterator: """ - Returns a generator to asynchronously stream back row data. + Returns an iterator to asynchronously stream back row data. Failed requests within operation_timeout and operation_deadline policies will be retried. - By default, row data is streamed eagerly over the network, and fully cached in memory - in the generator, which can be consumed as needed. The size of the generator cache can - be configured with cache_size_limit. When the cache is full, the read_rows_stream will pause - the network stream until space is available - Args: - query: contains details about which rows to return - - shard: if True, will attempt to split up and distribute query to multiple - backend nodes in parallel - - limit: a limit on the number of rows to return. Actual limit will be - min(limit, query.limit) - - cache_size: the number of rows to cache in memory. If None, no limits. - Defaults to None - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. time is only counted while actively waiting on the network. - Completed and cached results can still be accessed after the deadline is complete, - with a DeadlineExceeded exception only raised after cached results are exhausted - - per_row_timeout: the time budget for a single row read, in seconds. If a row takes - longer than per_row_timeout to complete, the ongoing network request will be with a - DeadlineExceeded exception, and a retry may be attempted - Applies only to the underlying network call. - - idle_timeout: the number of idle seconds before an active generator is marked as - stale and the cache is drained. The idle count is reset each time the generator - is yielded from - raises DeadlineExceeded on future yields + If None, defaults to the Table's default_operation_timeout - per_request_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with - a DeadlineExceeded exception, and a retry will be attempted + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_per_request_timeout Returns: - - an asynchronous generator that yields rows returned by the query + - an asynchronous iterator that yields rows returned by the query Raises: - DeadlineExceeded: raised after operation timeout will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions from any retries that failed - - IdleTimeout: if generator was abandoned + - GoogleAPIError: raised if the request encounters an unrecoverable error + - IdleTimeout: if iterator was abandoned """ - raise NotImplementedError + + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = per_request_timeout or self.default_per_request_timeout + + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout <= 0: + raise ValueError("per_request_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout > operation_timeout: + raise ValueError( + "per_request_timeout must not be greater than operation_timeout" + ) + if per_request_timeout is None: + per_request_timeout = operation_timeout + request = query._to_dict() if isinstance(query, ReadRowsQuery) else query + request["table_name"] = self.table_name + if self.app_profile_id: + request["app_profile_id"] = self.app_profile_id + + # read_rows smart retries is implemented using a series of iterators: + # - client.read_rows: outputs raw ReadRowsResponse objects from backend. Has per_request_timeout + # - ReadRowsOperation.merge_row_response_stream: parses chunks into rows + # - ReadRowsOperation.retryable_merge_rows: adds retries, caching, revised requests, per_request_timeout + # - ReadRowsIterator: adds idle_timeout, moves stats out of stream and into attribute + row_merger = _ReadRowsOperation( + request, + self.client._gapic_client, + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + output_generator = ReadRowsIterator(row_merger) + # add idle timeout to clear resources if generator is abandoned + idle_timeout_seconds = 300 + await output_generator._start_idle_timer(idle_timeout_seconds) + return output_generator async def read_rows( self, query: ReadRowsQuery | dict[str, Any], *, - shard: bool = False, - limit: int | None, - operation_timeout: int | float | None = 60, - per_row_timeout: int | float | None = 10, - per_request_timeout: int | float | None = None, + operation_timeout: float | None = None, + per_request_timeout: float | None = None, ) -> list[Row]: """ Helper function that returns a full list instead of a generator @@ -435,7 +481,13 @@ async def read_rows( Returns: - a list of the rows returned by the query """ - raise NotImplementedError + row_generator = await self.read_rows_stream( + query, + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + results = [row async for row in row_generator] + return results async def read_row( self, @@ -459,12 +511,9 @@ async def read_rows_sharded( query_list: list[ReadRowsQuery] | list[dict[str, Any]], *, limit: int | None, - cache_size_limit: int | None = None, operation_timeout: int | float | None = 60, - per_row_timeout: int | float | None = 10, - idle_timeout: int | float | None = 300, per_request_timeout: int | float | None = None, - ) -> AsyncIterable[Row]: + ) -> ReadRowsIterator: """ Runs a sharded query in parallel diff --git a/google/cloud/bigtable/exceptions.py b/google/cloud/bigtable/exceptions.py index 86bfe9247..975feb101 100644 --- a/google/cloud/bigtable/exceptions.py +++ b/google/cloud/bigtable/exceptions.py @@ -12,13 +12,90 @@ # See the License for the specific language governing permissions and # limitations under the License. # +from __future__ import annotations import sys +from inspect import iscoroutinefunction +from typing import Callable, Any + +from google.api_core import exceptions as core_exceptions is_311_plus = sys.version_info >= (3, 11) +def _convert_retry_deadline( + func: Callable[..., Any], + timeout_value: float | None = None, + retry_errors: list[Exception] | None = None, +): + """ + Decorator to convert RetryErrors raised by api_core.retry into + DeadlineExceeded exceptions, indicating that the underlying retries have + exhaused the timeout value. + Optionally attaches a RetryExceptionGroup to the DeadlineExceeded.__cause__, + detailing the failed exceptions associated with each retry. + + Supports both sync and async function wrapping. + + Args: + - func: The function to decorate + - timeout_value: The timeout value to display in the DeadlineExceeded error message + - retry_errors: An optional list of exceptions to attach as a RetryExceptionGroup to the DeadlineExceeded.__cause__ + """ + timeout_str = f" of {timeout_value:.1f}s" if timeout_value is not None else "" + error_str = f"operation_timeout{timeout_str} exceeded" + + def handle_error(): + new_exc = core_exceptions.DeadlineExceeded( + error_str, + ) + source_exc = None + if retry_errors: + source_exc = RetryExceptionGroup( + f"{len(retry_errors)} failed attempts", retry_errors + ) + new_exc.__cause__ = source_exc + raise new_exc from source_exc + + # separate wrappers for async and sync functions + async def wrapper_async(*args, **kwargs): + try: + return await func(*args, **kwargs) + except core_exceptions.RetryError: + handle_error() + + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except core_exceptions.RetryError: + handle_error() + + return wrapper_async if iscoroutinefunction(func) else wrapper + + +class IdleTimeout(core_exceptions.DeadlineExceeded): + """ + Exception raised by ReadRowsIterator when the generator + has been idle for longer than the internal idle_timeout. + """ + + pass + + +class InvalidChunk(core_exceptions.GoogleAPICallError): + """Exception raised to invalid chunk data from back-end.""" + + +class _RowSetComplete(Exception): + """ + Internal exception for _ReadRowsOperation + Raised in revise_request_rowset when there are no rows left to process when starting a retry attempt + """ + + pass + + class BigtableExceptionGroup(ExceptionGroup if is_311_plus else Exception): # type: ignore # noqa: F821 """ Represents one or more exceptions that occur during a bulk Bigtable operation @@ -29,7 +106,12 @@ class BigtableExceptionGroup(ExceptionGroup if is_311_plus else Exception): # t """ def __init__(self, message, excs): - raise NotImplementedError() + if is_311_plus: + super().__init__(message, excs) + else: + self.exceptions = excs + revised_message = f"{message} ({len(excs)} sub-exceptions)" + super().__init__(revised_message) class MutationsExceptionGroup(BigtableExceptionGroup): diff --git a/google/cloud/bigtable/iterators.py b/google/cloud/bigtable/iterators.py new file mode 100644 index 000000000..169bbc3f3 --- /dev/null +++ b/google/cloud/bigtable/iterators.py @@ -0,0 +1,129 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from typing import AsyncIterable + +import asyncio +import time +import sys + +from google.cloud.bigtable._read_rows import _ReadRowsOperation +from google.cloud.bigtable.exceptions import IdleTimeout +from google.cloud.bigtable.exceptions import _convert_retry_deadline +from google.cloud.bigtable.row import Row + + +class ReadRowsIterator(AsyncIterable[Row]): + """ + Async iterator for ReadRows responses. + """ + + def __init__(self, merger: _ReadRowsOperation): + self._merger: _ReadRowsOperation = merger + self._error: Exception | None = None + self.last_interaction_time = time.time() + self._idle_timeout_task: asyncio.Task[None] | None = None + # wrap merger with a wrapper that properly formats exceptions + self._next_fn = _convert_retry_deadline( + self._merger.__anext__, + self._merger.operation_timeout, + self._merger.transient_errors, + ) + + async def _start_idle_timer(self, idle_timeout: float): + """ + Start a coroutine that will cancel a stream if no interaction + with the iterator occurs for the specified number of seconds. + + Subsequent access to the iterator will raise an IdleTimeout exception. + + Args: + - idle_timeout: number of seconds of inactivity before cancelling the stream + """ + self.last_interaction_time = time.time() + if self._idle_timeout_task is not None: + self._idle_timeout_task.cancel() + self._idle_timeout_task = asyncio.create_task( + self._idle_timeout_coroutine(idle_timeout) + ) + if sys.version_info >= (3, 8): + self._idle_timeout_task.name = "ReadRowsIterator._idle_timeout" + + @property + def active(self): + """ + Returns True if the iterator is still active and has not been closed + """ + return self._error is None + + async def _idle_timeout_coroutine(self, idle_timeout: float): + """ + Coroutine that will cancel a stream if no interaction with the iterator + in the last `idle_timeout` seconds. + """ + while self.active: + next_timeout = self.last_interaction_time + idle_timeout + await asyncio.sleep(next_timeout - time.time()) + if self.last_interaction_time + idle_timeout < time.time() and self.active: + # idle timeout has expired + await self._finish_with_error( + IdleTimeout( + ( + "Timed out waiting for next Row to be consumed. " + f"(idle_timeout={idle_timeout:0.1f}s)" + ) + ) + ) + + def __aiter__(self): + """Implement the async iterator protocol.""" + return self + + async def __anext__(self) -> Row: + """ + Implement the async iterator potocol. + + Return the next item in the stream if active, or + raise an exception if the stream has been closed. + """ + if self._error is not None: + raise self._error + try: + self.last_interaction_time = time.time() + return await self._next_fn() + except Exception as e: + await self._finish_with_error(e) + raise e + + async def _finish_with_error(self, e: Exception): + """ + Helper function to close the stream and clean up resources + after an error has occurred. + """ + if self.active: + await self._merger.aclose() + self._error = e + if self._idle_timeout_task is not None: + self._idle_timeout_task.cancel() + self._idle_timeout_task = None + + async def aclose(self): + """ + Support closing the stream with an explicit call to aclose() + """ + await self._finish_with_error( + StopAsyncIteration(f"{self.__class__.__name__} closed") + ) diff --git a/google/cloud/bigtable/mutations.py b/google/cloud/bigtable/mutations.py index 4ff59bff9..3bb5b2ed6 100644 --- a/google/cloud/bigtable/mutations.py +++ b/google/cloud/bigtable/mutations.py @@ -15,7 +15,6 @@ from __future__ import annotations from dataclasses import dataclass -from google.cloud.bigtable.row import family_id, qualifier, row_key class Mutation: @@ -24,23 +23,23 @@ class Mutation: @dataclass class SetCell(Mutation): - family: family_id - column_qualifier: qualifier + family: str + column_qualifier: bytes new_value: bytes | str | int timestamp_ms: int | None = None @dataclass class DeleteRangeFromColumn(Mutation): - family: family_id - column_qualifier: qualifier + family: str + column_qualifier: bytes start_timestamp_ms: int end_timestamp_ms: int @dataclass class DeleteAllFromFamily(Mutation): - family_to_delete: family_id + family_to_delete: str @dataclass @@ -50,5 +49,5 @@ class DeleteAllFromRow(Mutation): @dataclass class BulkMutationsEntry: - row: row_key + row: bytes mutations: list[Mutation] | Mutation diff --git a/google/cloud/bigtable/mutations_batcher.py b/google/cloud/bigtable/mutations_batcher.py index 582786ee4..9681f4382 100644 --- a/google/cloud/bigtable/mutations_batcher.py +++ b/google/cloud/bigtable/mutations_batcher.py @@ -18,12 +18,14 @@ from typing import TYPE_CHECKING from google.cloud.bigtable.mutations import Mutation -from google.cloud.bigtable.row import row_key from google.cloud.bigtable.row_filters import RowFilter if TYPE_CHECKING: from google.cloud.bigtable.client import Table # pragma: no cover +# Type alias used internally for readability. +_row_key_type = bytes + class MutationsBatcher: """ @@ -44,7 +46,7 @@ class MutationsBatcher: batcher.add(row, mut) """ - queue: asyncio.Queue[tuple[row_key, list[Mutation]]] + queue: asyncio.Queue[tuple[_row_key_type, list[Mutation]]] conditional_queues: dict[RowFilter, tuple[list[Mutation], list[Mutation]]] MB_SIZE = 1024 * 1024 diff --git a/google/cloud/bigtable/read_modify_write_rules.py b/google/cloud/bigtable/read_modify_write_rules.py index 839262ea2..cd6b370df 100644 --- a/google/cloud/bigtable/read_modify_write_rules.py +++ b/google/cloud/bigtable/read_modify_write_rules.py @@ -16,8 +16,6 @@ from dataclasses import dataclass -from google.cloud.bigtable.row import family_id, qualifier - class ReadModifyWriteRule: pass @@ -26,12 +24,12 @@ class ReadModifyWriteRule: @dataclass class IncrementRule(ReadModifyWriteRule): increment_amount: int - family: family_id - column_qualifier: qualifier + family: str + qualifier: bytes @dataclass class AppendValueRule(ReadModifyWriteRule): - append_value: bytes | str - family: family_id - column_qualifier: qualifier + append_value: bytes + family: str + qualifier: bytes diff --git a/google/cloud/bigtable/read_rows_query.py b/google/cloud/bigtable/read_rows_query.py index 559b47f04..e26f99d34 100644 --- a/google/cloud/bigtable/read_rows_query.py +++ b/google/cloud/bigtable/read_rows_query.py @@ -14,7 +14,6 @@ # from __future__ import annotations from typing import TYPE_CHECKING, Any -from .row import row_key from dataclasses import dataclass from google.cloud.bigtable.row_filters import RowFilter @@ -26,7 +25,7 @@ class _RangePoint: """Model class for a point in a row range""" - key: row_key + key: bytes is_inclusive: bool diff --git a/google/cloud/bigtable/row.py b/google/cloud/bigtable/row.py index 8231e324e..a5fb033e6 100644 --- a/google/cloud/bigtable/row.py +++ b/google/cloud/bigtable/row.py @@ -19,10 +19,8 @@ from functools import total_ordering # Type aliases used internally for readability. -row_key = bytes -family_id = str -qualifier = bytes -row_value = bytes +_family_type = str +_qualifier_type = bytes class Row(Sequence["Cell"]): @@ -37,9 +35,11 @@ class Row(Sequence["Cell"]): cells = row["family", "qualifier"] """ + __slots__ = ("row_key", "cells", "_index_data") + def __init__( self, - key: row_key, + key: bytes, cells: list[Cell], ): """ @@ -49,16 +49,28 @@ def __init__( They are returned by the Bigtable backend. """ self.row_key = key - self._cells_map: dict[family_id, dict[qualifier, list[Cell]]] = OrderedDict() - self._cells_list: list[Cell] = [] - # add cells to internal stores using Bigtable native ordering - for cell in cells: - if cell.family not in self._cells_map: - self._cells_map[cell.family] = OrderedDict() - if cell.column_qualifier not in self._cells_map[cell.family]: - self._cells_map[cell.family][cell.column_qualifier] = [] - self._cells_map[cell.family][cell.column_qualifier].append(cell) - self._cells_list.append(cell) + self.cells: list[Cell] = cells + # index is lazily created when needed + self._index_data: OrderedDict[ + _family_type, OrderedDict[_qualifier_type, list[Cell]] + ] | None = None + + @property + def _index( + self, + ) -> OrderedDict[_family_type, OrderedDict[_qualifier_type, list[Cell]]]: + """ + Returns an index of cells associated with each family and qualifier. + + The index is lazily created when needed + """ + if self._index_data is None: + self._index_data = OrderedDict() + for cell in self.cells: + self._index_data.setdefault(cell.family, OrderedDict()).setdefault( + cell.qualifier, [] + ).append(cell) + return self._index_data def get_cells( self, family: str | None = None, qualifier: str | bytes | None = None @@ -81,31 +93,29 @@ def get_cells( raise ValueError("Qualifier passed without family") else: # return all cells on get_cells() - return self._cells_list + return self.cells if qualifier is None: # return all cells in family on get_cells(family) return list(self._get_all_from_family(family)) if isinstance(qualifier, str): qualifier = qualifier.encode("utf-8") # return cells in family and qualifier on get_cells(family, qualifier) - if family not in self._cells_map: + if family not in self._index: raise ValueError(f"Family '{family}' not found in row '{self.row_key!r}'") - if qualifier not in self._cells_map[family]: + if qualifier not in self._index[family]: raise ValueError( f"Qualifier '{qualifier!r}' not found in family '{family}' in row '{self.row_key!r}'" ) - return self._cells_map[family][qualifier] + return self._index[family][qualifier] - def _get_all_from_family(self, family: family_id) -> Generator[Cell, None, None]: + def _get_all_from_family(self, family: str) -> Generator[Cell, None, None]: """ Returns all cells in the row for the family_id """ - if family not in self._cells_map: + if family not in self._index: raise ValueError(f"Family '{family}' not found in row '{self.row_key!r}'") - qualifier_dict = self._cells_map.get(family, {}) - for cell_batch in qualifier_dict.values(): - for cell in cell_batch: - yield cell + for qualifier in self._index[family]: + yield from self._index[family][qualifier] def __str__(self) -> str: """ @@ -135,7 +145,7 @@ def __repr__(self): for family, qualifier in self.get_column_components(): cell_list = self[family, qualifier] repr_list = [cell.to_dict() for cell in cell_list] - cell_str_buffer.append(f" ('{family}', {qualifier}): {repr_list},") + cell_str_buffer.append(f" ('{family}', {qualifier!r}): {repr_list},") cell_str_buffer.append("}") cell_str = "\n".join(cell_str_buffer) output = f"Row(key={self.row_key!r}, cells={cell_str})" @@ -148,25 +158,23 @@ def to_dict(self) -> dict[str, Any]: https://cloud.google.com/bigtable/docs/reference/data/rpc/google.bigtable.v2#row """ - families_list: list[dict[str, Any]] = [] - for family in self._cells_map: - column_list: list[dict[str, Any]] = [] - for qualifier in self._cells_map[family]: - cells_list: list[dict[str, Any]] = [] - for cell in self._cells_map[family][qualifier]: - cells_list.append(cell.to_dict()) - column_list.append({"qualifier": qualifier, "cells": cells_list}) - families_list.append({"name": family, "columns": column_list}) - return {"key": self.row_key, "families": families_list} + family_list = [] + for family_name, qualifier_dict in self._index.items(): + qualifier_list = [] + for qualifier_name, cell_list in qualifier_dict.items(): + cell_dicts = [cell.to_dict() for cell in cell_list] + qualifier_list.append( + {"qualifier": qualifier_name, "cells": cell_dicts} + ) + family_list.append({"name": family_name, "columns": qualifier_list}) + return {"key": self.row_key, "families": family_list} # Sequence and Mapping methods def __iter__(self): """ Allow iterating over all cells in the row """ - # iterate as a sequence; yield all cells - for cell in self._cells_list: - yield cell + return iter(self.cells) def __contains__(self, item): """ @@ -175,24 +183,22 @@ def __contains__(self, item): Works for both cells in the internal list, and `family` or `(family, qualifier)` pairs associated with the cells """ - if isinstance(item, family_id): - # check if family key is in Row - return item in self._cells_map + if isinstance(item, _family_type): + return item in self._index elif ( isinstance(item, tuple) - and isinstance(item[0], family_id) - and isinstance(item[1], (qualifier, str)) + and isinstance(item[0], _family_type) + and isinstance(item[1], (bytes, str)) ): - # check if (family, qualifier) pair is in Row - qualifer = item[1] if isinstance(item[1], bytes) else item[1].encode() - return item[0] in self._cells_map and qualifer in self._cells_map[item[0]] + q = item[1] if isinstance(item[1], bytes) else item[1].encode("utf-8") + return item[0] in self._index and q in self._index[item[0]] # check if Cell is in Row - return item in self._cells_list + return item in self.cells @overload def __getitem__( self, - index: family_id | tuple[family_id, qualifier | str], + index: str | tuple[str, bytes | str], ) -> list[Cell]: # overload signature for type checking pass @@ -214,17 +220,17 @@ def __getitem__(self, index): Supports indexing by family, (family, qualifier) pair, numerical index, and index slicing """ - if isinstance(index, family_id): + if isinstance(index, _family_type): return self.get_cells(family=index) elif ( isinstance(index, tuple) - and isinstance(index[0], family_id) - and isinstance(index[1], (qualifier, str)) + and isinstance(index[0], _family_type) + and isinstance(index[1], (bytes, str)) ): return self.get_cells(family=index[0], qualifier=index[1]) elif isinstance(index, int) or isinstance(index, slice): # index is int or slice - return self._cells_list[index] + return self.cells[index] else: raise TypeError( "Index must be family_id, (family_id, qualifier), int, or slice" @@ -234,19 +240,15 @@ def __len__(self): """ Implements `len()` operator """ - return len(self._cells_list) + return len(self.cells) - def get_column_components(self): + def get_column_components(self) -> list[tuple[str, bytes]]: """ Returns a list of (family, qualifier) pairs associated with the cells Pairs can be used for indexing """ - key_list = [] - for family in self._cells_map: - for qualifier in self._cells_map[family]: - key_list.append((family, qualifier)) - return key_list + return [(f, q) for f in self._index for q in self._index[f]] def __eq__(self, other): """ @@ -258,7 +260,7 @@ def __eq__(self, other): return False if self.row_key != other.row_key: return False - if len(self._cells_list) != len(other._cells_list): + if len(self.cells) != len(other.cells): return False components = self.get_column_components() other_components = other.get_column_components() @@ -270,7 +272,7 @@ def __eq__(self, other): if len(self[family, qualifier]) != len(other[family, qualifier]): return False # compare individual cell lists - if self._cells_list != other._cells_list: + if self.cells != other.cells: return False return True @@ -281,6 +283,21 @@ def __ne__(self, other) -> bool: return not self == other +class _LastScannedRow(Row): + """A value used to indicate a scanned row that is not returned as part of + a query. + + This is used internally to indicate progress in a scan, and improve retry + performance. It is not intended to be used directly by users. + """ + + def __init__(self, row_key): + super().__init__(row_key, []) + + def __eq__(self, other): + return isinstance(other, _LastScannedRow) + + @total_ordering class Cell: """ @@ -291,12 +308,21 @@ class Cell: Expected to be read-only to users, and written by backend """ + __slots__ = ( + "value", + "row_key", + "family", + "qualifier", + "timestamp_micros", + "labels", + ) + def __init__( self, - value: row_value, - row: row_key, - family: family_id, - column_qualifier: qualifier | str, + value: bytes, + row_key: bytes, + family: str, + qualifier: bytes | str, timestamp_micros: int, labels: list[str] | None = None, ): @@ -307,11 +333,11 @@ def __init__( They are returned by the Bigtable backend. """ self.value = value - self.row_key = row + self.row_key = row_key self.family = family - if isinstance(column_qualifier, str): - column_qualifier = column_qualifier.encode() - self.column_qualifier = column_qualifier + if isinstance(qualifier, str): + qualifier = qualifier.encode() + self.qualifier = qualifier self.timestamp_micros = timestamp_micros self.labels = labels if labels is not None else [] @@ -349,7 +375,7 @@ def __repr__(self): """ Returns a string representation of the cell """ - return f"Cell(value={self.value!r}, row={self.row_key!r}, family='{self.family}', column_qualifier={self.column_qualifier!r}, timestamp_micros={self.timestamp_micros}, labels={self.labels})" + return f"Cell(value={self.value!r}, row_key={self.row_key!r}, family='{self.family}', qualifier={self.qualifier!r}, timestamp_micros={self.timestamp_micros}, labels={self.labels})" """For Bigtable native ordering""" @@ -361,14 +387,14 @@ def __lt__(self, other) -> bool: return NotImplemented this_ordering = ( self.family, - self.column_qualifier, + self.qualifier, -self.timestamp_micros, self.value, self.labels, ) other_ordering = ( other.family, - other.column_qualifier, + other.qualifier, -other.timestamp_micros, other.value, other.labels, @@ -384,7 +410,7 @@ def __eq__(self, other) -> bool: return ( self.row_key == other.row_key and self.family == other.family - and self.column_qualifier == other.column_qualifier + and self.qualifier == other.qualifier and self.value == other.value and self.timestamp_micros == other.timestamp_micros and len(self.labels) == len(other.labels) @@ -405,7 +431,7 @@ def __hash__(self): ( self.row_key, self.family, - self.column_qualifier, + self.qualifier, self.value, self.timestamp_micros, tuple(self.labels), diff --git a/noxfile.py b/noxfile.py index 035599844..164d138bd 100644 --- a/noxfile.py +++ b/noxfile.py @@ -39,7 +39,9 @@ "pytest-cov", "pytest-asyncio", ] -UNIT_TEST_EXTERNAL_DEPENDENCIES = [] +UNIT_TEST_EXTERNAL_DEPENDENCIES = [ + "git+https://github.com/googleapis/python-api-core.git@retry_generators" +] UNIT_TEST_LOCAL_DEPENDENCIES = [] UNIT_TEST_DEPENDENCIES = [] UNIT_TEST_EXTRAS = [] @@ -52,8 +54,11 @@ "pytest-asyncio", "google-cloud-testutils", ] -SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [] +SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [ + "git+https://github.com/googleapis/python-api-core.git@retry_generators" +] SYSTEM_TEST_LOCAL_DEPENDENCIES = [] +UNIT_TEST_DEPENDENCIES = [] SYSTEM_TEST_DEPENDENCIES = [] SYSTEM_TEST_EXTRAS = [] SYSTEM_TEST_EXTRAS_BY_PYTHON = {} @@ -157,16 +162,8 @@ def install_unittest_dependencies(session, *constraints): standard_deps = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_DEPENDENCIES session.install(*standard_deps, *constraints) - if UNIT_TEST_EXTERNAL_DEPENDENCIES: - warnings.warn( - "'unit_test_external_dependencies' is deprecated. Instead, please " - "use 'unit_test_dependencies' or 'unit_test_local_dependencies'.", - DeprecationWarning, - ) - session.install(*UNIT_TEST_EXTERNAL_DEPENDENCIES, *constraints) - if UNIT_TEST_LOCAL_DEPENDENCIES: - session.install(*UNIT_TEST_LOCAL_DEPENDENCIES, *constraints) + session.install("-e", *UNIT_TEST_LOCAL_DEPENDENCIES, *constraints) if UNIT_TEST_EXTRAS_BY_PYTHON: extras = UNIT_TEST_EXTRAS_BY_PYTHON.get(session.python, []) @@ -180,6 +177,20 @@ def install_unittest_dependencies(session, *constraints): else: session.install("-e", ".", *constraints) + if UNIT_TEST_EXTERNAL_DEPENDENCIES: + warnings.warn( + "'unit_test_external_dependencies' is deprecated. Instead, please " + "use 'unit_test_dependencies' or 'unit_test_local_dependencies'.", + DeprecationWarning, + ) + session.install( + "--upgrade", + "--no-deps", + "--force-reinstall", + *UNIT_TEST_EXTERNAL_DEPENDENCIES, + *constraints, + ) + def default(session): # Install all test dependencies, then install this package in-place. @@ -455,6 +466,11 @@ def prerelease_deps(session): ) session.run("python", "-c", "import grpc; print(grpc.__version__)") + # TODO: remove adter merging api-core + session.install( + "--upgrade", "--no-deps", "--force-reinstall", *UNIT_TEST_EXTERNAL_DEPENDENCIES + ) + session.run("py.test", "tests/unit") system_test_path = os.path.join("tests", "system.py") diff --git a/python-api-core b/python-api-core new file mode 160000 index 000000000..9ba76760f --- /dev/null +++ b/python-api-core @@ -0,0 +1 @@ +Subproject commit 9ba76760f5b7ba8128be85ca780811a0b9ec9087 diff --git a/testing/constraints-3.7.txt b/testing/constraints-3.7.txt index d14da7c0c..7bf769c9b 100644 --- a/testing/constraints-3.7.txt +++ b/testing/constraints-3.7.txt @@ -5,7 +5,8 @@ # # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 -google-api-core==1.34.0 +# TODO: reset after merging api-core submodule +# google-api-core==2.11.0 google-cloud-core==1.4.1 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 05633ac91..48caceccd 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -16,6 +16,8 @@ import pytest_asyncio import os import asyncio +from google.api_core import retry +from google.api_core.exceptions import ClientError TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" @@ -146,6 +148,53 @@ async def table(client, table_id, instance_id): yield table +class TempRowBuilder: + """ + Used to add rows to a table for testing purposes. + """ + + def __init__(self, table): + self.rows = [] + self.table = table + + async def add_row( + self, row_key, family=TEST_FAMILY, qualifier=b"q", value=b"test-value" + ): + request = { + "table_name": self.table.table_name, + "row_key": row_key, + "mutations": [ + { + "set_cell": { + "family_name": family, + "column_qualifier": qualifier, + "value": value, + } + } + ], + } + await self.table.client._gapic_client.mutate_row(request) + self.rows.append(row_key) + + async def delete_rows(self): + request = { + "table_name": self.table.table_name, + "entries": [ + {"row_key": row, "mutations": [{"delete_from_row": {}}]} + for row in self.rows + ], + } + await self.table.client._gapic_client.mutate_rows(request) + + +@pytest_asyncio.fixture(scope="function") +async def temp_rows(table): + builder = TempRowBuilder(table) + yield builder + await builder.delete_rows() + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_ping_and_warm_gapic(client, table): """ @@ -154,3 +203,118 @@ async def test_ping_and_warm_gapic(client, table): """ request = {"name": table.instance_name} await client._gapic_client.ping_and_warm(request) + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows_stream(table, temp_rows): + """ + Ensure that the read_rows_stream method works + """ + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + + # full table scan + generator = await table.read_rows_stream({}) + first_row = await generator.__anext__() + second_row = await generator.__anext__() + assert first_row.row_key == b"row_key_1" + assert second_row.row_key == b"row_key_2" + with pytest.raises(StopAsyncIteration): + await generator.__anext__() + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows(table, temp_rows): + """ + Ensure that the read_rows method works + """ + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + # full table scan + row_list = await table.read_rows({}) + assert len(row_list) == 2 + assert row_list[0].row_key == b"row_key_1" + assert row_list[1].row_key == b"row_key_2" + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows_range_query(table, temp_rows): + """ + Ensure that the read_rows method works + """ + from google.cloud.bigtable import ReadRowsQuery + from google.cloud.bigtable import RowRange + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + # full table scan + query = ReadRowsQuery(row_ranges=RowRange(start_key=b"b", end_key=b"d")) + row_list = await table.read_rows(query) + assert len(row_list) == 2 + assert row_list[0].row_key == b"b" + assert row_list[1].row_key == b"c" + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows_key_query(table, temp_rows): + """ + Ensure that the read_rows method works + """ + from google.cloud.bigtable import ReadRowsQuery + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + # full table scan + query = ReadRowsQuery(row_keys=[b"a", b"c"]) + row_list = await table.read_rows(query) + assert len(row_list) == 2 + assert row_list[0].row_key == b"a" + assert row_list[1].row_key == b"c" + + +@pytest.mark.asyncio +async def test_read_rows_stream_close(table, temp_rows): + """ + Ensure that the read_rows_stream can be closed + """ + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + + # full table scan + generator = await table.read_rows_stream({}) + first_row = await generator.__anext__() + assert first_row.row_key == b"row_key_1" + await generator.aclose() + assert generator.active is False + with pytest.raises(StopAsyncIteration) as e: + await generator.__anext__() + assert "closed" in str(e) + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows_stream_inactive_timer(table, temp_rows): + """ + Ensure that the read_rows_stream method works + """ + from google.cloud.bigtable.exceptions import IdleTimeout + + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + + generator = await table.read_rows_stream({}) + await generator._start_idle_timer(0.05) + await asyncio.sleep(0.2) + assert generator.active is False + with pytest.raises(IdleTimeout) as e: + await generator.__anext__() + assert "inactivity" in str(e) + assert "idle_timeout=0.1" in str(e) diff --git a/tests/unit/read-rows-acceptance-test.json b/tests/unit/read-rows-acceptance-test.json new file mode 100644 index 000000000..011ace2b9 --- /dev/null +++ b/tests/unit/read-rows-acceptance-test.json @@ -0,0 +1,1665 @@ +{ + "readRowsTests": [ + { + "description": "invalid - no commit", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "invalid - no cell key before commit", + "chunks": [ + { + "commitRow": true + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "invalid - no cell key before value", + "chunks": [ + { + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "invalid - new col family must specify qualifier", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "familyName": "B", + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "bare commit implies ts=0", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + }, + { + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C" + } + ] + }, + { + "description": "simple row with timestamp", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + } + ] + }, + { + "description": "missing timestamp, implied ts=0", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "value": "value-VAL" + } + ] + }, + { + "description": "empty cell value", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C" + } + ] + }, + { + "description": "two unsplit cells", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "98", + "value": "value-VAL_2" + } + ] + }, + { + "description": "two qualifiers", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "qualifier": "RA==", + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "D", + "timestampMicros": "98", + "value": "value-VAL_2" + } + ] + }, + { + "description": "two families", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "familyName": "B", + "qualifier": "RQ==", + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK", + "familyName": "B", + "qualifier": "E", + "timestampMicros": "98", + "value": "value-VAL_2" + } + ] + }, + { + "description": "with labels", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "labels": [ + "L_1" + ], + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "timestampMicros": "98", + "labels": [ + "L_2" + ], + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1", + "label": "L_1" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "98", + "value": "value-VAL_2", + "label": "L_2" + } + ] + }, + { + "description": "split cell, bare commit", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dg==", + "valueSize": 9, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUw=", + "commitRow": false + }, + { + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C" + } + ] + }, + { + "description": "split cell", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dg==", + "valueSize": 9, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUw=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + } + ] + }, + { + "description": "split four ways", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "labels": [ + "L" + ], + "value": "dg==", + "valueSize": 9, + "commitRow": false + }, + { + "value": "YQ==", + "valueSize": 9, + "commitRow": false + }, + { + "value": "bA==", + "valueSize": 9, + "commitRow": false + }, + { + "value": "dWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL", + "label": "L" + } + ] + }, + { + "description": "two split cells", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUxfMQ==", + "commitRow": false + }, + { + "timestampMicros": "98", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUxfMg==", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "98", + "value": "value-VAL_2" + } + ] + }, + { + "description": "multi-qualifier splits", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUxfMQ==", + "commitRow": false + }, + { + "qualifier": "RA==", + "timestampMicros": "98", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUxfMg==", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "D", + "timestampMicros": "98", + "value": "value-VAL_2" + } + ] + }, + { + "description": "multi-qualifier multi-split", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YQ==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "bHVlLVZBTF8x", + "commitRow": false + }, + { + "qualifier": "RA==", + "timestampMicros": "98", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YQ==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "bHVlLVZBTF8y", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "D", + "timestampMicros": "98", + "value": "value-VAL_2" + } + ] + }, + { + "description": "multi-family split", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUxfMQ==", + "commitRow": false + }, + { + "familyName": "B", + "qualifier": "RQ==", + "timestampMicros": "98", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUxfMg==", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK", + "familyName": "B", + "qualifier": "E", + "timestampMicros": "98", + "value": "value-VAL_2" + } + ] + }, + { + "description": "invalid - no commit between rows", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + }, + { + "rowKey": "UktfMg==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "invalid - no commit after first row", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + }, + { + "rowKey": "UktfMg==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "invalid - last row missing commit", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + }, + { + "error": true + } + ] + }, + { + "description": "invalid - duplicate row key", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + }, + { + "rowKey": "UktfMQ==", + "familyName": "B", + "qualifier": "RA==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + }, + { + "error": true + } + ] + }, + { + "description": "invalid - new row missing row key", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + }, + { + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + }, + { + "error": true + } + ] + }, + { + "description": "two rows", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + }, + { + "rowKey": "RK_2", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + } + ] + }, + { + "description": "two rows implicit timestamp", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "value": "dmFsdWUtVkFM", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "value": "value-VAL" + }, + { + "rowKey": "RK_2", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + } + ] + }, + { + "description": "two rows empty value", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C" + }, + { + "rowKey": "RK_2", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + } + ] + }, + { + "description": "two rows, one with multiple cells", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "B", + "qualifier": "RA==", + "timestampMicros": "97", + "value": "dmFsdWUtVkFMXzM=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "98", + "value": "value-VAL_2" + }, + { + "rowKey": "RK_2", + "familyName": "B", + "qualifier": "D", + "timestampMicros": "97", + "value": "value-VAL_3" + } + ] + }, + { + "description": "two rows, multiple cells", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "qualifier": "RA==", + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "B", + "qualifier": "RQ==", + "timestampMicros": "97", + "value": "dmFsdWUtVkFMXzM=", + "commitRow": false + }, + { + "qualifier": "Rg==", + "timestampMicros": "96", + "value": "dmFsdWUtVkFMXzQ=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "D", + "timestampMicros": "98", + "value": "value-VAL_2" + }, + { + "rowKey": "RK_2", + "familyName": "B", + "qualifier": "E", + "timestampMicros": "97", + "value": "value-VAL_3" + }, + { + "rowKey": "RK_2", + "familyName": "B", + "qualifier": "F", + "timestampMicros": "96", + "value": "value-VAL_4" + } + ] + }, + { + "description": "two rows, multiple cells, multiple families", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "familyName": "B", + "qualifier": "RQ==", + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "M", + "qualifier": "Tw==", + "timestampMicros": "97", + "value": "dmFsdWUtVkFMXzM=", + "commitRow": false + }, + { + "familyName": "N", + "qualifier": "UA==", + "timestampMicros": "96", + "value": "dmFsdWUtVkFMXzQ=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1" + }, + { + "rowKey": "RK_1", + "familyName": "B", + "qualifier": "E", + "timestampMicros": "98", + "value": "value-VAL_2" + }, + { + "rowKey": "RK_2", + "familyName": "M", + "qualifier": "O", + "timestampMicros": "97", + "value": "value-VAL_3" + }, + { + "rowKey": "RK_2", + "familyName": "N", + "qualifier": "P", + "timestampMicros": "96", + "value": "value-VAL_4" + } + ] + }, + { + "description": "two rows, four cells, 2 labels", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "99", + "labels": [ + "L_1" + ], + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "B", + "qualifier": "RA==", + "timestampMicros": "97", + "labels": [ + "L_3" + ], + "value": "dmFsdWUtVkFMXzM=", + "commitRow": false + }, + { + "timestampMicros": "96", + "value": "dmFsdWUtVkFMXzQ=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "99", + "value": "value-VAL_1", + "label": "L_1" + }, + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "98", + "value": "value-VAL_2" + }, + { + "rowKey": "RK_2", + "familyName": "B", + "qualifier": "D", + "timestampMicros": "97", + "value": "value-VAL_3", + "label": "L_3" + }, + { + "rowKey": "RK_2", + "familyName": "B", + "qualifier": "D", + "timestampMicros": "96", + "value": "value-VAL_4" + } + ] + }, + { + "description": "two rows with splits, same timestamp", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUxfMQ==", + "commitRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dg==", + "valueSize": 11, + "commitRow": false + }, + { + "value": "YWx1ZS1WQUxfMg==", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_1" + }, + { + "rowKey": "RK_2", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_2" + } + ] + }, + { + "description": "invalid - bare reset", + "chunks": [ + { + "resetRow": true + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "invalid - bad reset, no commit", + "chunks": [ + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "invalid - missing key after reset", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + }, + { + "resetRow": true + }, + { + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "no data after reset", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + }, + { + "resetRow": true + } + ] + }, + { + "description": "simple reset", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + } + ] + }, + { + "description": "reset to new val", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_2" + } + ] + }, + { + "description": "reset to new qual", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "RA==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "D", + "timestampMicros": "100", + "value": "value-VAL_1" + } + ] + }, + { + "description": "reset with splits", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "timestampMicros": "98", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_2" + } + ] + }, + { + "description": "reset two cells", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": false + }, + { + "timestampMicros": "97", + "value": "dmFsdWUtVkFMXzM=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_2" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "97", + "value": "value-VAL_3" + } + ] + }, + { + "description": "two resets", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzM=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_3" + } + ] + }, + { + "description": "reset then two cells", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "Uks=", + "familyName": "B", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": false + }, + { + "qualifier": "RA==", + "timestampMicros": "97", + "value": "dmFsdWUtVkFMXzM=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "B", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_2" + }, + { + "rowKey": "RK", + "familyName": "B", + "qualifier": "D", + "timestampMicros": "97", + "value": "value-VAL_3" + } + ] + }, + { + "description": "reset to new row", + "chunks": [ + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "UktfMg==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzI=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_2", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_2" + } + ] + }, + { + "description": "reset in between chunks", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "labels": [ + "L" + ], + "value": "dg==", + "valueSize": 10, + "commitRow": false + }, + { + "value": "YQ==", + "valueSize": 10, + "commitRow": false + }, + { + "resetRow": true + }, + { + "rowKey": "UktfMQ==", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFMXzE=", + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK_1", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL_1" + } + ] + }, + { + "description": "invalid - reset with chunk", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "labels": [ + "L" + ], + "value": "dg==", + "valueSize": 10, + "commitRow": false + }, + { + "value": "YQ==", + "valueSize": 10, + "resetRow": true + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "invalid - commit with chunk", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "labels": [ + "L" + ], + "value": "dg==", + "valueSize": 10, + "commitRow": false + }, + { + "value": "YQ==", + "valueSize": 10, + "commitRow": true + } + ], + "results": [ + { + "error": true + } + ] + }, + { + "description": "empty cell chunk", + "chunks": [ + { + "rowKey": "Uks=", + "familyName": "A", + "qualifier": "Qw==", + "timestampMicros": "100", + "value": "dmFsdWUtVkFM", + "commitRow": false + }, + { + "commitRow": false + }, + { + "commitRow": true + } + ], + "results": [ + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C", + "timestampMicros": "100", + "value": "value-VAL" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C" + }, + { + "rowKey": "RK", + "familyName": "A", + "qualifier": "C" + } + ] + } + ] +} diff --git a/tests/unit/test__read_rows.py b/tests/unit/test__read_rows.py new file mode 100644 index 000000000..e57b5d992 --- /dev/null +++ b/tests/unit/test__read_rows.py @@ -0,0 +1,1001 @@ +import unittest +import pytest + +from google.cloud.bigtable.exceptions import InvalidChunk +from google.cloud.bigtable._read_rows import AWAITING_NEW_ROW +from google.cloud.bigtable._read_rows import AWAITING_NEW_CELL +from google.cloud.bigtable._read_rows import AWAITING_CELL_VALUE + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock + from unittest.mock import AsyncMock # type: ignore +except ImportError: # pragma: NO COVER + import mock # type: ignore + from mock import AsyncMock # type: ignore # noqa F401 + +TEST_FAMILY = "family_name" +TEST_QUALIFIER = b"qualifier" +TEST_TIMESTAMP = 123456789 +TEST_LABELS = ["label1", "label2"] + + +class TestReadRowsOperation: + """ + Tests helper functions in the ReadRowsOperation class + in-depth merging logic in merge_row_response_stream and _read_rows_retryable_attempt + is tested in test_read_rows_acceptance test_client_read_rows, and conformance tests + """ + + @staticmethod + def _get_target_class(): + from google.cloud.bigtable._read_rows import _ReadRowsOperation + + return _ReadRowsOperation + + def _make_one(self, *args, **kwargs): + return self._get_target_class()(*args, **kwargs) + + def test_ctor_defaults(self): + request = {} + client = mock.Mock() + client.read_rows = mock.Mock() + client.read_rows.return_value = None + start_time = 123 + default_operation_timeout = 600 + with mock.patch("time.monotonic", return_value=start_time): + instance = self._make_one(request, client) + assert instance.transient_errors == [] + assert instance._last_emitted_row_key is None + assert instance._emit_count == 0 + assert instance.operation_timeout == default_operation_timeout + retryable_fn = instance._partial_retryable + assert retryable_fn.func == instance._read_rows_retryable_attempt + assert retryable_fn.args[0] == client.read_rows + assert retryable_fn.args[1] == default_operation_timeout + assert retryable_fn.args[2] == default_operation_timeout + start_time + assert retryable_fn.args[3] == 0 + assert client.read_rows.call_count == 0 + + def test_ctor(self): + row_limit = 91 + request = {"rows_limit": row_limit} + client = mock.Mock() + client.read_rows = mock.Mock() + client.read_rows.return_value = None + expected_operation_timeout = 42 + expected_request_timeout = 44 + start_time = 123 + with mock.patch("time.monotonic", return_value=start_time): + instance = self._make_one( + request, + client, + operation_timeout=expected_operation_timeout, + per_request_timeout=expected_request_timeout, + ) + assert instance.transient_errors == [] + assert instance._last_emitted_row_key is None + assert instance._emit_count == 0 + assert instance.operation_timeout == expected_operation_timeout + retryable_fn = instance._partial_retryable + assert retryable_fn.func == instance._read_rows_retryable_attempt + assert retryable_fn.args[0] == client.read_rows + assert retryable_fn.args[1] == expected_request_timeout + assert retryable_fn.args[2] == start_time + expected_operation_timeout + assert retryable_fn.args[3] == row_limit + assert client.read_rows.call_count == 0 + + def test___aiter__(self): + request = {} + client = mock.Mock() + client.read_rows = mock.Mock() + instance = self._make_one(request, client) + assert instance.__aiter__() is instance + + @pytest.mark.asyncio + async def test_transient_error_capture(self): + from google.api_core import exceptions as core_exceptions + + client = mock.Mock() + client.read_rows = mock.Mock() + test_exc = core_exceptions.Aborted("test") + test_exc2 = core_exceptions.DeadlineExceeded("test") + client.read_rows.side_effect = [test_exc, test_exc2] + instance = self._make_one({}, client) + with pytest.raises(RuntimeError): + await instance.__anext__() + assert len(instance.transient_errors) == 2 + assert instance.transient_errors[0] == test_exc + assert instance.transient_errors[1] == test_exc2 + + @pytest.mark.parametrize( + "in_keys,last_key,expected", + [ + (["b", "c", "d"], "a", ["b", "c", "d"]), + (["a", "b", "c"], "b", ["c"]), + (["a", "b", "c"], "c", []), + (["a", "b", "c"], "d", []), + (["d", "c", "b", "a"], "b", ["d", "c"]), + ], + ) + def test_revise_request_rowset_keys(self, in_keys, last_key, expected): + sample_range = {"start_key_open": last_key} + row_set = {"row_keys": in_keys, "row_ranges": [sample_range]} + revised = self._get_target_class()._revise_request_rowset(row_set, last_key) + assert revised["row_keys"] == expected + assert revised["row_ranges"] == [sample_range] + + @pytest.mark.parametrize( + "in_ranges,last_key,expected", + [ + ( + [{"start_key_open": "b", "end_key_closed": "d"}], + "a", + [{"start_key_open": "b", "end_key_closed": "d"}], + ), + ( + [{"start_key_closed": "b", "end_key_closed": "d"}], + "a", + [{"start_key_closed": "b", "end_key_closed": "d"}], + ), + ( + [{"start_key_open": "a", "end_key_closed": "d"}], + "b", + [{"start_key_open": "b", "end_key_closed": "d"}], + ), + ( + [{"start_key_closed": "a", "end_key_open": "d"}], + "b", + [{"start_key_open": "b", "end_key_open": "d"}], + ), + ( + [{"start_key_closed": "b", "end_key_closed": "d"}], + "b", + [{"start_key_open": "b", "end_key_closed": "d"}], + ), + ([{"start_key_closed": "b", "end_key_closed": "d"}], "d", []), + ([{"start_key_closed": "b", "end_key_open": "d"}], "d", []), + ([{"start_key_closed": "b", "end_key_closed": "d"}], "e", []), + ([{"start_key_closed": "b"}], "z", [{"start_key_open": "z"}]), + ([{"start_key_closed": "b"}], "a", [{"start_key_closed": "b"}]), + ( + [{"end_key_closed": "z"}], + "a", + [{"start_key_open": "a", "end_key_closed": "z"}], + ), + ( + [{"end_key_open": "z"}], + "a", + [{"start_key_open": "a", "end_key_open": "z"}], + ), + ], + ) + def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected): + next_key = last_key + "a" + row_set = {"row_keys": [next_key], "row_ranges": in_ranges} + revised = self._get_target_class()._revise_request_rowset(row_set, last_key) + assert revised["row_keys"] == [next_key] + assert revised["row_ranges"] == expected + + @pytest.mark.parametrize("last_key", ["a", "b", "c"]) + def test_revise_request_full_table(self, last_key): + row_set = {"row_keys": [], "row_ranges": []} + for selected_set in [row_set, None]: + revised = self._get_target_class()._revise_request_rowset( + selected_set, last_key + ) + assert revised["row_keys"] == [] + assert len(revised["row_ranges"]) == 1 + assert revised["row_ranges"][0]["start_key_open"] == last_key + + def test_revise_to_empty_rowset(self): + """revising to an empty rowset should raise error""" + from google.cloud.bigtable.exceptions import _RowSetComplete + + row_keys = ["a", "b", "c"] + row_set = {"row_keys": row_keys, "row_ranges": [{"end_key_open": "c"}]} + with pytest.raises(_RowSetComplete): + self._get_target_class()._revise_request_rowset(row_set, "d") + + @pytest.mark.parametrize( + "start_limit,emit_num,expected_limit", + [ + (10, 0, 10), + (10, 1, 9), + (10, 10, 0), + (0, 10, 0), + (0, 0, 0), + (4, 2, 2), + ], + ) + @pytest.mark.asyncio + async def test_revise_limit(self, start_limit, emit_num, expected_limit): + """ + revise_limit should revise the request's limit field + - if limit is 0 (unlimited), it should never be revised + - if start_limit-emit_num == 0, the request should end early + - if the number emitted exceeds the new limit, an exception should + should be raised (tested in test_revise_limit_over_limit) + """ + request = {"rows_limit": start_limit} + instance = self._make_one(request, mock.Mock()) + instance._emit_count = emit_num + instance._last_emitted_row_key = "a" + gapic_mock = mock.Mock() + gapic_mock.side_effect = [GeneratorExit("stop_fn")] + attempt = instance._read_rows_retryable_attempt( + gapic_mock, 100, 100, start_limit + ) + if start_limit != 0 and expected_limit == 0: + # if we emitted the expected number of rows, we should receive a StopAsyncIteration + with pytest.raises(StopAsyncIteration): + await attempt.__anext__() + else: + with pytest.raises(GeneratorExit): + await attempt.__anext__() + assert request["rows_limit"] == expected_limit + + @pytest.mark.parametrize("start_limit,emit_num", [(5, 10), (3, 9), (1, 10)]) + @pytest.mark.asyncio + async def test_revise_limit_over_limit(self, start_limit, emit_num): + """ + Should raise runtime error if we get in state where emit_num > start_num + (unless start_num == 0, which represents unlimited) + """ + request = {"rows_limit": start_limit} + instance = self._make_one(request, mock.Mock()) + instance._emit_count = emit_num + instance._last_emitted_row_key = "a" + attempt = instance._read_rows_retryable_attempt( + mock.Mock(), 100, 100, start_limit + ) + with pytest.raises(RuntimeError) as e: + await attempt.__anext__() + assert "emit count exceeds row limit" in str(e.value) + + @pytest.mark.asyncio + async def test_aclose(self): + import asyncio + + instance = self._make_one({}, mock.Mock()) + await instance.aclose() + assert instance._stream is None + assert instance._last_emitted_row_key is None + with pytest.raises(asyncio.InvalidStateError): + await instance.__anext__() + # try calling a second time + await instance.aclose() + + @pytest.mark.parametrize("limit", [1, 3, 10]) + @pytest.mark.asyncio + async def test_retryable_attempt_hit_limit(self, limit): + """ + Stream should end after hitting the limit + """ + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + + instance = self._make_one({}, mock.Mock()) + + async def mock_gapic(*args, **kwargs): + # continuously return a single row + async def gen(): + for i in range(limit * 2): + chunk = ReadRowsResponse.CellChunk( + row_key=str(i).encode(), + family_name="family_name", + qualifier=b"qualifier", + commit_row=True, + ) + yield ReadRowsResponse(chunks=[chunk]) + + return gen() + + gen = instance._read_rows_retryable_attempt(mock_gapic, 100, 100, limit) + # should yield values up to the limit + for i in range(limit): + await gen.__anext__() + # next value should be StopAsyncIteration + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + + @pytest.mark.asyncio + async def test_retryable_ignore_repeated_rows(self): + """ + Duplicate rows should cause an invalid chunk error + """ + from google.cloud.bigtable._read_rows import _ReadRowsOperation + from google.cloud.bigtable.row import Row + from google.cloud.bigtable.exceptions import InvalidChunk + + async def mock_stream(): + while True: + yield Row(b"dup_key", cells=[]) + yield Row(b"dup_key", cells=[]) + + with mock.patch.object( + _ReadRowsOperation, "merge_row_response_stream" + ) as mock_stream_fn: + mock_stream_fn.return_value = mock_stream() + instance = self._make_one({}, mock.AsyncMock()) + first_row = await instance.__anext__() + assert first_row.row_key == b"dup_key" + with pytest.raises(InvalidChunk) as exc: + await instance.__anext__() + assert "Last emitted row key out of order" in str(exc.value) + + @pytest.mark.asyncio + async def test_retryable_ignore_last_scanned_rows(self): + """ + Last scanned rows should not be emitted + """ + from google.cloud.bigtable._read_rows import _ReadRowsOperation + from google.cloud.bigtable.row import Row, _LastScannedRow + + async def mock_stream(): + while True: + yield Row(b"key1", cells=[]) + yield _LastScannedRow(b"key2_ignored") + yield Row(b"key3", cells=[]) + + with mock.patch.object( + _ReadRowsOperation, "merge_row_response_stream" + ) as mock_stream_fn: + mock_stream_fn.return_value = mock_stream() + instance = self._make_one({}, mock.AsyncMock()) + first_row = await instance.__anext__() + assert first_row.row_key == b"key1" + second_row = await instance.__anext__() + assert second_row.row_key == b"key3" + + @pytest.mark.asyncio + async def test_retryable_cancel_on_close(self): + """Underlying gapic call should be cancelled when stream is closed""" + from google.cloud.bigtable._read_rows import _ReadRowsOperation + from google.cloud.bigtable.row import Row + + async def mock_stream(): + while True: + yield Row(b"key1", cells=[]) + + with mock.patch.object( + _ReadRowsOperation, "merge_row_response_stream" + ) as mock_stream_fn: + mock_stream_fn.return_value = mock_stream() + mock_gapic = mock.AsyncMock() + mock_call = await mock_gapic.read_rows() + instance = self._make_one({}, mock_gapic) + await instance.__anext__() + assert mock_call.cancel.call_count == 0 + await instance.aclose() + assert mock_call.cancel.call_count == 1 + + +class TestStateMachine(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigtable._read_rows import _StateMachine + + return _StateMachine + + def _make_one(self, *args, **kwargs): + return self._get_target_class()(*args, **kwargs) + + def test_ctor(self): + from google.cloud.bigtable._read_rows import _RowBuilder + + instance = self._make_one() + assert instance.last_seen_row_key is None + assert instance.current_state == AWAITING_NEW_ROW + assert instance.current_family is None + assert instance.current_qualifier is None + assert isinstance(instance.adapter, _RowBuilder) + assert instance.adapter.current_key is None + assert instance.adapter.working_cell is None + assert instance.adapter.working_value is None + assert instance.adapter.completed_cells == [] + + def test_is_terminal_state(self): + + instance = self._make_one() + assert instance.is_terminal_state() is True + instance.current_state = AWAITING_NEW_ROW + assert instance.is_terminal_state() is True + instance.current_state = AWAITING_NEW_CELL + assert instance.is_terminal_state() is False + instance.current_state = AWAITING_CELL_VALUE + assert instance.is_terminal_state() is False + + def test__reset_row(self): + instance = self._make_one() + instance.current_state = mock.Mock() + instance.current_family = "family" + instance.current_qualifier = "qualifier" + instance.adapter = mock.Mock() + instance._reset_row() + assert instance.current_state == AWAITING_NEW_ROW + assert instance.current_family is None + assert instance.current_qualifier is None + assert instance.adapter.reset.call_count == 1 + + def test_handle_last_scanned_row_wrong_state(self): + from google.cloud.bigtable.exceptions import InvalidChunk + + instance = self._make_one() + instance.current_state = AWAITING_NEW_CELL + with pytest.raises(InvalidChunk) as e: + instance.handle_last_scanned_row("row_key") + assert e.value.args[0] == "Last scanned row key received in invalid state" + instance.current_state = AWAITING_CELL_VALUE + with pytest.raises(InvalidChunk) as e: + instance.handle_last_scanned_row("row_key") + assert e.value.args[0] == "Last scanned row key received in invalid state" + + def test_handle_last_scanned_row_out_of_order(self): + from google.cloud.bigtable.exceptions import InvalidChunk + + instance = self._make_one() + instance.last_seen_row_key = b"b" + with pytest.raises(InvalidChunk) as e: + instance.handle_last_scanned_row(b"a") + assert e.value.args[0] == "Last scanned row key is out of order" + with pytest.raises(InvalidChunk) as e: + instance.handle_last_scanned_row(b"b") + assert e.value.args[0] == "Last scanned row key is out of order" + + def test_handle_last_scanned_row(self): + from google.cloud.bigtable.row import _LastScannedRow + + instance = self._make_one() + instance.adapter = mock.Mock() + instance.last_seen_row_key = b"a" + output_row = instance.handle_last_scanned_row(b"b") + assert instance.last_seen_row_key == b"b" + assert isinstance(output_row, _LastScannedRow) + assert output_row.row_key == b"b" + assert instance.current_state == AWAITING_NEW_ROW + assert instance.current_family is None + assert instance.current_qualifier is None + assert instance.adapter.reset.call_count == 1 + + def test__handle_complete_row(self): + from google.cloud.bigtable.row import Row + + instance = self._make_one() + instance.current_state = mock.Mock() + instance.current_family = "family" + instance.current_qualifier = "qualifier" + instance.adapter = mock.Mock() + instance._handle_complete_row(Row(b"row_key", {})) + assert instance.last_seen_row_key == b"row_key" + assert instance.current_state == AWAITING_NEW_ROW + assert instance.current_family is None + assert instance.current_qualifier is None + assert instance.adapter.reset.call_count == 1 + + def test__handle_reset_chunk_errors(self): + from google.cloud.bigtable.exceptions import InvalidChunk + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + + instance = self._make_one() + with pytest.raises(InvalidChunk) as e: + instance._handle_reset_chunk(mock.Mock()) + instance.current_state = mock.Mock() + assert e.value.args[0] == "Reset chunk received when not processing row" + with pytest.raises(InvalidChunk) as e: + instance._handle_reset_chunk( + ReadRowsResponse.CellChunk(row_key=b"row_key")._pb + ) + assert e.value.args[0] == "Reset chunk has a row key" + with pytest.raises(InvalidChunk) as e: + instance._handle_reset_chunk( + ReadRowsResponse.CellChunk(family_name="family")._pb + ) + assert e.value.args[0] == "Reset chunk has a family name" + with pytest.raises(InvalidChunk) as e: + instance._handle_reset_chunk( + ReadRowsResponse.CellChunk(qualifier=b"qualifier")._pb + ) + assert e.value.args[0] == "Reset chunk has a qualifier" + with pytest.raises(InvalidChunk) as e: + instance._handle_reset_chunk( + ReadRowsResponse.CellChunk(timestamp_micros=1)._pb + ) + assert e.value.args[0] == "Reset chunk has a timestamp" + with pytest.raises(InvalidChunk) as e: + instance._handle_reset_chunk(ReadRowsResponse.CellChunk(value=b"value")._pb) + assert e.value.args[0] == "Reset chunk has a value" + with pytest.raises(InvalidChunk) as e: + instance._handle_reset_chunk( + ReadRowsResponse.CellChunk(labels=["label"])._pb + ) + assert e.value.args[0] == "Reset chunk has labels" + + def test_handle_chunk_out_of_order(self): + from google.cloud.bigtable.exceptions import InvalidChunk + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + + instance = self._make_one() + instance.last_seen_row_key = b"b" + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(row_key=b"a")._pb + instance.handle_chunk(chunk) + assert "increasing" in e.value.args[0] + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(row_key=b"b")._pb + instance.handle_chunk(chunk) + assert "increasing" in e.value.args[0] + + def test_handle_chunk_reset(self): + """Should call _handle_reset_chunk when a chunk with reset_row is encountered""" + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + + instance = self._make_one() + with mock.patch.object(type(instance), "_handle_reset_chunk") as mock_reset: + chunk = ReadRowsResponse.CellChunk(reset_row=True)._pb + output = instance.handle_chunk(chunk) + assert output is None + assert mock_reset.call_count == 1 + + @pytest.mark.parametrize("state", [AWAITING_NEW_ROW, AWAITING_CELL_VALUE]) + def handle_chunk_with_commit_wrong_state(self, state): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + + instance = self._make_one() + with mock.patch.object( + type(instance.current_state), "handle_chunk" + ) as mock_state_handle: + mock_state_handle.return_value = state(mock.Mock()) + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(commit_row=True)._pb + instance.handle_chunk(mock.Mock(), chunk) + assert instance.current_state == state + assert e.value.args[0] == "Commit chunk received with in invalid state" + + def test_handle_chunk_with_commit(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable.row import Row + + instance = self._make_one() + with mock.patch.object(type(instance), "_reset_row") as mock_reset: + chunk = ReadRowsResponse.CellChunk( + row_key=b"row_key", family_name="f", qualifier=b"q", commit_row=True + )._pb + output = instance.handle_chunk(chunk) + assert isinstance(output, Row) + assert output.row_key == b"row_key" + assert output[0].family == "f" + assert output[0].qualifier == b"q" + assert instance.last_seen_row_key == b"row_key" + assert mock_reset.call_count == 1 + + def test_handle_chunk_with_commit_empty_strings(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable.row import Row + + instance = self._make_one() + with mock.patch.object(type(instance), "_reset_row") as mock_reset: + chunk = ReadRowsResponse.CellChunk( + row_key=b"row_key", family_name="", qualifier=b"", commit_row=True + )._pb + output = instance.handle_chunk(chunk) + assert isinstance(output, Row) + assert output.row_key == b"row_key" + assert output[0].family == "" + assert output[0].qualifier == b"" + assert instance.last_seen_row_key == b"row_key" + assert mock_reset.call_count == 1 + + def handle_chunk_incomplete(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + + instance = self._make_one() + chunk = ReadRowsResponse.CellChunk( + row_key=b"row_key", family_name="f", qualifier=b"q", commit_row=False + )._pb + output = instance.handle_chunk(chunk) + assert output is None + assert isinstance(instance.current_state, AWAITING_CELL_VALUE) + assert instance.current_family == "f" + assert instance.current_qualifier == b"q" + + +class TestState(unittest.TestCase): + def test_AWAITING_NEW_ROW_empty_key(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + + instance = AWAITING_NEW_ROW + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(row_key=b"")._pb + instance.handle_chunk(mock.Mock(), chunk) + assert "missing a row key" in e.value.args[0] + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk()._pb + instance.handle_chunk(mock.Mock(), chunk) + assert "missing a row key" in e.value.args[0] + + def test_AWAITING_NEW_ROW(self): + """ + AWAITING_NEW_ROW should start a RowBuilder row, then + delegate the call to AWAITING_NEW_CELL + """ + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + + instance = AWAITING_NEW_ROW + state_machine = mock.Mock() + with mock.patch.object(AWAITING_NEW_CELL, "handle_chunk") as mock_delegate: + chunk = ReadRowsResponse.CellChunk(row_key=b"row_key")._pb + instance.handle_chunk(state_machine, chunk) + assert state_machine.adapter.start_row.call_count == 1 + assert state_machine.adapter.start_row.call_args[0][0] == b"row_key" + mock_delegate.assert_called_once_with(state_machine, chunk) + + def test_AWAITING_NEW_CELL_family_without_qualifier(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + state_machine.current_qualifier = b"q" + instance = AWAITING_NEW_CELL + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(family_name="fam")._pb + instance.handle_chunk(state_machine, chunk) + assert "New family must specify qualifier" in e.value.args[0] + + def test_AWAITING_NEW_CELL_qualifier_without_family(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + instance = AWAITING_NEW_CELL + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(qualifier=b"q")._pb + instance.handle_chunk(state_machine, chunk) + assert "Family not found" in e.value.args[0] + + def test_AWAITING_NEW_CELL_no_row_state(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + instance = AWAITING_NEW_CELL + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk()._pb + instance.handle_chunk(state_machine, chunk) + assert "Missing family for new cell" in e.value.args[0] + state_machine.current_family = "fam" + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk()._pb + instance.handle_chunk(state_machine, chunk) + assert "Missing qualifier for new cell" in e.value.args[0] + + def test_AWAITING_NEW_CELL_invalid_row_key(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + instance = AWAITING_NEW_CELL + state_machine.adapter.current_key = b"abc" + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(row_key=b"123")._pb + instance.handle_chunk(state_machine, chunk) + assert "Row key changed mid row" in e.value.args[0] + + def test_AWAITING_NEW_CELL_success_no_split(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + state_machine.adapter = mock.Mock() + instance = AWAITING_NEW_CELL + row_key = b"row_key" + family = "fam" + qualifier = b"q" + labels = ["label"] + timestamp = 123 + value = b"value" + chunk = ReadRowsResponse.CellChunk( + row_key=row_key, + family_name=family, + qualifier=qualifier, + timestamp_micros=timestamp, + value=value, + labels=labels, + )._pb + state_machine.adapter.current_key = row_key + new_state = instance.handle_chunk(state_machine, chunk) + assert state_machine.adapter.start_cell.call_count == 1 + kwargs = state_machine.adapter.start_cell.call_args[1] + assert kwargs["family"] == family + assert kwargs["qualifier"] == qualifier + assert kwargs["timestamp_micros"] == timestamp + assert kwargs["labels"] == labels + assert state_machine.adapter.cell_value.call_count == 1 + assert state_machine.adapter.cell_value.call_args[0][0] == value + assert state_machine.adapter.finish_cell.call_count == 1 + assert new_state == AWAITING_NEW_CELL + + def test_AWAITING_NEW_CELL_success_with_split(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + state_machine.adapter = mock.Mock() + instance = AWAITING_NEW_CELL + row_key = b"row_key" + family = "fam" + qualifier = b"q" + labels = ["label"] + timestamp = 123 + value = b"value" + chunk = ReadRowsResponse.CellChunk( + value_size=1, + row_key=row_key, + family_name=family, + qualifier=qualifier, + timestamp_micros=timestamp, + value=value, + labels=labels, + )._pb + state_machine.adapter.current_key = row_key + new_state = instance.handle_chunk(state_machine, chunk) + assert state_machine.adapter.start_cell.call_count == 1 + kwargs = state_machine.adapter.start_cell.call_args[1] + assert kwargs["family"] == family + assert kwargs["qualifier"] == qualifier + assert kwargs["timestamp_micros"] == timestamp + assert kwargs["labels"] == labels + assert state_machine.adapter.cell_value.call_count == 1 + assert state_machine.adapter.cell_value.call_args[0][0] == value + assert state_machine.adapter.finish_cell.call_count == 0 + assert new_state == AWAITING_CELL_VALUE + + def test_AWAITING_CELL_VALUE_w_row_key(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + instance = AWAITING_CELL_VALUE + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(row_key=b"123")._pb + instance.handle_chunk(state_machine, chunk) + assert "In progress cell had a row key" in e.value.args[0] + + def test_AWAITING_CELL_VALUE_w_family(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + instance = AWAITING_CELL_VALUE + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(family_name="")._pb + instance.handle_chunk(state_machine, chunk) + assert "In progress cell had a family name" in e.value.args[0] + + def test_AWAITING_CELL_VALUE_w_qualifier(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + instance = AWAITING_CELL_VALUE + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(qualifier=b"")._pb + instance.handle_chunk(state_machine, chunk) + assert "In progress cell had a qualifier" in e.value.args[0] + + def test_AWAITING_CELL_VALUE_w_timestamp(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + instance = AWAITING_CELL_VALUE + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(timestamp_micros=123)._pb + instance.handle_chunk(state_machine, chunk) + assert "In progress cell had a timestamp" in e.value.args[0] + + def test_AWAITING_CELL_VALUE_w_labels(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + instance = AWAITING_CELL_VALUE + with pytest.raises(InvalidChunk) as e: + chunk = ReadRowsResponse.CellChunk(labels=[""])._pb + instance.handle_chunk(state_machine, chunk) + assert "In progress cell had labels" in e.value.args[0] + + def test_AWAITING_CELL_VALUE_continuation(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + state_machine.adapter = mock.Mock() + instance = AWAITING_CELL_VALUE + value = b"value" + chunk = ReadRowsResponse.CellChunk(value=value, value_size=1)._pb + new_state = instance.handle_chunk(state_machine, chunk) + assert state_machine.adapter.cell_value.call_count == 1 + assert state_machine.adapter.cell_value.call_args[0][0] == value + assert state_machine.adapter.finish_cell.call_count == 0 + assert new_state == AWAITING_CELL_VALUE + + def test_AWAITING_CELL_VALUE_final_chunk(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _StateMachine + + state_machine = _StateMachine() + state_machine.adapter = mock.Mock() + instance = AWAITING_CELL_VALUE + value = b"value" + chunk = ReadRowsResponse.CellChunk(value=value, value_size=0)._pb + new_state = instance.handle_chunk(state_machine, chunk) + assert state_machine.adapter.cell_value.call_count == 1 + assert state_machine.adapter.cell_value.call_args[0][0] == value + assert state_machine.adapter.finish_cell.call_count == 1 + assert new_state == AWAITING_NEW_CELL + + +class TestRowBuilder(unittest.TestCase): + @staticmethod + def _get_target_class(): + from google.cloud.bigtable._read_rows import _RowBuilder + + return _RowBuilder + + def _make_one(self, *args, **kwargs): + return self._get_target_class()(*args, **kwargs) + + def test_ctor(self): + with mock.patch.object(self._get_target_class(), "reset") as reset: + self._make_one() + reset.assert_called_once() + row_builder = self._make_one() + self.assertIsNone(row_builder.current_key) + self.assertIsNone(row_builder.working_cell) + self.assertIsNone(row_builder.working_value) + self.assertEqual(row_builder.completed_cells, []) + + def test_start_row(self): + row_builder = self._make_one() + row_builder.start_row(b"row_key") + self.assertEqual(row_builder.current_key, b"row_key") + row_builder.start_row(b"row_key2") + self.assertEqual(row_builder.current_key, b"row_key2") + + def test_start_cell(self): + # test with no family + with self.assertRaises(InvalidChunk) as e: + self._make_one().start_cell("", TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) + self.assertEqual(str(e.exception), "Missing family for a new cell") + # test with no row + with self.assertRaises(InvalidChunk) as e: + row_builder = self._make_one() + row_builder.start_cell( + TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS + ) + self.assertEqual(str(e.exception), "start_cell called without a row") + # test with valid row + row_builder = self._make_one() + row_builder.start_row(b"row_key") + row_builder.start_cell(TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) + self.assertEqual(row_builder.working_cell.family, TEST_FAMILY) + self.assertEqual(row_builder.working_cell.qualifier, TEST_QUALIFIER) + self.assertEqual(row_builder.working_cell.timestamp_micros, TEST_TIMESTAMP) + self.assertEqual(row_builder.working_cell.labels, TEST_LABELS) + self.assertEqual(row_builder.working_value, b"") + + def test_cell_value(self): + row_builder = self._make_one() + row_builder.start_row(b"row_key") + with self.assertRaises(InvalidChunk): + # start_cell must be called before cell_value + row_builder.cell_value(b"cell_value") + row_builder.start_cell(TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) + row_builder.cell_value(b"cell_value") + self.assertEqual(row_builder.working_value, b"cell_value") + # should be able to continuously append to the working value + row_builder.cell_value(b"appended") + self.assertEqual(row_builder.working_value, b"cell_valueappended") + + def test_finish_cell(self): + row_builder = self._make_one() + row_builder.start_row(b"row_key") + row_builder.start_cell(TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) + row_builder.finish_cell() + self.assertEqual(len(row_builder.completed_cells), 1) + self.assertEqual(row_builder.completed_cells[0].family, TEST_FAMILY) + self.assertEqual(row_builder.completed_cells[0].qualifier, TEST_QUALIFIER) + self.assertEqual( + row_builder.completed_cells[0].timestamp_micros, TEST_TIMESTAMP + ) + self.assertEqual(row_builder.completed_cells[0].labels, TEST_LABELS) + self.assertEqual(row_builder.completed_cells[0].value, b"") + self.assertEqual(row_builder.working_cell, None) + self.assertEqual(row_builder.working_value, None) + # add additional cell with value + row_builder.start_cell(TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) + row_builder.cell_value(b"cell_value") + row_builder.cell_value(b"appended") + row_builder.finish_cell() + self.assertEqual(len(row_builder.completed_cells), 2) + self.assertEqual(row_builder.completed_cells[1].family, TEST_FAMILY) + self.assertEqual(row_builder.completed_cells[1].qualifier, TEST_QUALIFIER) + self.assertEqual( + row_builder.completed_cells[1].timestamp_micros, TEST_TIMESTAMP + ) + self.assertEqual(row_builder.completed_cells[1].labels, TEST_LABELS) + self.assertEqual(row_builder.completed_cells[1].value, b"cell_valueappended") + self.assertEqual(row_builder.working_cell, None) + self.assertEqual(row_builder.working_value, None) + + def test_finish_cell_no_cell(self): + with self.assertRaises(InvalidChunk) as e: + self._make_one().finish_cell() + self.assertEqual(str(e.exception), "finish_cell called before start_cell") + with self.assertRaises(InvalidChunk) as e: + row_builder = self._make_one() + row_builder.start_row(b"row_key") + row_builder.finish_cell() + self.assertEqual(str(e.exception), "finish_cell called before start_cell") + + def test_finish_row(self): + row_builder = self._make_one() + row_builder.start_row(b"row_key") + for i in range(3): + row_builder.start_cell(str(i), TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) + row_builder.cell_value(b"cell_value: ") + row_builder.cell_value(str(i).encode("utf-8")) + row_builder.finish_cell() + self.assertEqual(len(row_builder.completed_cells), i + 1) + output = row_builder.finish_row() + self.assertEqual(row_builder.current_key, None) + self.assertEqual(row_builder.working_cell, None) + self.assertEqual(row_builder.working_value, None) + self.assertEqual(len(row_builder.completed_cells), 0) + + self.assertEqual(output.row_key, b"row_key") + self.assertEqual(len(output), 3) + for i in range(3): + self.assertEqual(output[i].family, str(i)) + self.assertEqual(output[i].qualifier, TEST_QUALIFIER) + self.assertEqual(output[i].timestamp_micros, TEST_TIMESTAMP) + self.assertEqual(output[i].labels, TEST_LABELS) + self.assertEqual(output[i].value, b"cell_value: " + str(i).encode("utf-8")) + + def test_finish_row_no_row(self): + with self.assertRaises(InvalidChunk) as e: + self._make_one().finish_row() + self.assertEqual(str(e.exception), "No row in progress") + + def test_reset(self): + row_builder = self._make_one() + row_builder.start_row(b"row_key") + for i in range(3): + row_builder.start_cell(str(i), TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) + row_builder.cell_value(b"cell_value: ") + row_builder.cell_value(str(i).encode("utf-8")) + row_builder.finish_cell() + self.assertEqual(len(row_builder.completed_cells), i + 1) + row_builder.reset() + self.assertEqual(row_builder.current_key, None) + self.assertEqual(row_builder.working_cell, None) + self.assertEqual(row_builder.working_value, None) + self.assertEqual(len(row_builder.completed_cells), 0) + + +class TestChunkHasField: + def test__chunk_has_field_empty(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _chunk_has_field + + chunk = ReadRowsResponse.CellChunk()._pb + assert not _chunk_has_field(chunk, "family_name") + assert not _chunk_has_field(chunk, "qualifier") + + def test__chunk_has_field_populated_empty_strings(self): + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + from google.cloud.bigtable._read_rows import _chunk_has_field + + chunk = ReadRowsResponse.CellChunk(qualifier=b"", family_name="")._pb + assert _chunk_has_field(chunk, "family_name") + assert _chunk_has_field(chunk, "qualifier") diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index ca7220800..8a3402a65 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -11,16 +11,21 @@ # 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. - +from __future__ import annotations import grpc import asyncio import re import sys -from google.auth.credentials import AnonymousCredentials import pytest +from google.auth.credentials import AnonymousCredentials +from google.cloud.bigtable_v2.types import ReadRowsResponse +from google.cloud.bigtable.read_rows_query import ReadRowsQuery +from google.api_core import exceptions as core_exceptions +from google.cloud.bigtable.exceptions import InvalidChunk + # try/except added for compatibility with python < 3.8 try: from unittest import mock @@ -782,6 +787,8 @@ async def test_table_ctor(self): expected_table_id = "table-id" expected_instance_id = "instance-id" expected_app_profile_id = "app-profile-id" + expected_operation_timeout = 123 + expected_per_request_timeout = 12 client = BigtableDataClient() assert not client._active_instances @@ -790,6 +797,8 @@ async def test_table_ctor(self): expected_instance_id, expected_table_id, expected_app_profile_id, + default_operation_timeout=expected_operation_timeout, + default_per_request_timeout=expected_per_request_timeout, ) await asyncio.sleep(0) assert table.table_id == expected_table_id @@ -797,6 +806,8 @@ async def test_table_ctor(self): assert table.app_profile_id == expected_app_profile_id assert table.client is client assert table.instance_name in client._active_instances + assert table.default_operation_timeout == expected_operation_timeout + assert table.default_per_request_timeout == expected_per_request_timeout # ensure task reaches completion await table._register_instance_task assert table._register_instance_task.done() @@ -804,6 +815,33 @@ async def test_table_ctor(self): assert table._register_instance_task.exception() is None await client.close() + @pytest.mark.asyncio + async def test_table_ctor_bad_timeout_values(self): + from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.client import Table + + client = BigtableDataClient() + + with pytest.raises(ValueError) as e: + Table(client, "", "", default_per_request_timeout=-1) + assert "default_per_request_timeout must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + Table(client, "", "", default_operation_timeout=-1) + assert "default_operation_timeout must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + Table( + client, + "", + "", + default_operation_timeout=1, + default_per_request_timeout=2, + ) + assert ( + "default_per_request_timeout must be less than default_operation_timeout" + in str(e.value) + ) + await client.close() + def test_table_ctor_sync(self): # initializing client in a sync context should raise RuntimeError from google.cloud.bigtable.client import Table @@ -812,3 +850,447 @@ def test_table_ctor_sync(self): with pytest.raises(RuntimeError) as e: Table(client, "instance-id", "table-id") assert e.match("Table must be created within an async event loop context.") + + +class TestReadRows: + """ + Tests for table.read_rows and related methods. + """ + + def _make_client(self, *args, **kwargs): + from google.cloud.bigtable.client import BigtableDataClient + + return BigtableDataClient(*args, **kwargs) + + def _make_stats(self): + from google.cloud.bigtable_v2.types import RequestStats + from google.cloud.bigtable_v2.types import FullReadStatsView + from google.cloud.bigtable_v2.types import ReadIterationStats + + return RequestStats( + full_read_stats_view=FullReadStatsView( + read_iteration_stats=ReadIterationStats( + rows_seen_count=1, + rows_returned_count=2, + cells_seen_count=3, + cells_returned_count=4, + ) + ) + ) + + def _make_chunk(self, *args, **kwargs): + from google.cloud.bigtable_v2 import ReadRowsResponse + + kwargs["row_key"] = kwargs.get("row_key", b"row_key") + kwargs["family_name"] = kwargs.get("family_name", "family_name") + kwargs["qualifier"] = kwargs.get("qualifier", b"qualifier") + kwargs["value"] = kwargs.get("value", b"value") + kwargs["commit_row"] = kwargs.get("commit_row", True) + + return ReadRowsResponse.CellChunk(*args, **kwargs) + + async def _make_gapic_stream( + self, + chunk_list: list[ReadRowsResponse.CellChunk | Exception], + sleep_time=0, + ): + from google.cloud.bigtable_v2 import ReadRowsResponse + + class mock_stream: + def __init__(self, chunk_list, sleep_time): + self.chunk_list = chunk_list + self.idx = -1 + self.sleep_time = sleep_time + + def __aiter__(self): + return self + + async def __anext__(self): + self.idx += 1 + if len(self.chunk_list) > self.idx: + if sleep_time: + await asyncio.sleep(self.sleep_time) + chunk = self.chunk_list[self.idx] + if isinstance(chunk, Exception): + raise chunk + else: + return ReadRowsResponse(chunks=[chunk]) + raise StopAsyncIteration + + def cancel(self): + pass + + return mock_stream(chunk_list, sleep_time) + + @pytest.mark.asyncio + async def test_read_rows(self): + client = self._make_client() + table = client.get_table("instance", "table") + query = ReadRowsQuery() + chunks = [ + self._make_chunk(row_key=b"test_1"), + self._make_chunk(row_key=b"test_2"), + ] + with mock.patch.object(table.client._gapic_client, "read_rows") as read_rows: + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + chunks + ) + results = await table.read_rows(query, operation_timeout=3) + assert len(results) == 2 + assert results[0].row_key == b"test_1" + assert results[1].row_key == b"test_2" + await client.close() + + @pytest.mark.asyncio + async def test_read_rows_stream(self): + client = self._make_client() + table = client.get_table("instance", "table") + query = ReadRowsQuery() + chunks = [ + self._make_chunk(row_key=b"test_1"), + self._make_chunk(row_key=b"test_2"), + ] + with mock.patch.object(table.client._gapic_client, "read_rows") as read_rows: + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + chunks + ) + gen = await table.read_rows_stream(query, operation_timeout=3) + results = [row async for row in gen] + assert len(results) == 2 + assert results[0].row_key == b"test_1" + assert results[1].row_key == b"test_2" + await client.close() + + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_read_rows_query_matches_request(self, include_app_profile): + from google.cloud.bigtable import RowRange + + async with self._make_client() as client: + app_profile_id = "app_profile_id" if include_app_profile else None + table = client.get_table("instance", "table", app_profile_id=app_profile_id) + row_keys = [b"test_1", "test_2"] + row_ranges = RowRange("start", "end") + filter_ = {"test": "filter"} + limit = 99 + query = ReadRowsQuery( + row_keys=row_keys, + row_ranges=row_ranges, + row_filter=filter_, + limit=limit, + ) + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + [] + ) + results = await table.read_rows(query, operation_timeout=3) + assert len(results) == 0 + call_request = read_rows.call_args_list[0][0][0] + query_dict = query._to_dict() + if include_app_profile: + assert set(call_request.keys()) == set(query_dict.keys()) | { + "table_name", + "app_profile_id", + } + else: + assert set(call_request.keys()) == set(query_dict.keys()) | { + "table_name" + } + assert call_request["rows"] == query_dict["rows"] + assert call_request["filter"] == filter_ + assert call_request["rows_limit"] == limit + assert call_request["table_name"] == table.table_name + if include_app_profile: + assert call_request["app_profile_id"] == app_profile_id + + @pytest.mark.parametrize("operation_timeout", [0.001, 0.023, 0.1]) + @pytest.mark.asyncio + async def test_read_rows_timeout(self, operation_timeout): + async with self._make_client() as client: + table = client.get_table("instance", "table") + query = ReadRowsQuery() + chunks = [self._make_chunk(row_key=b"test_1")] + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + chunks, sleep_time=1 + ) + try: + await table.read_rows(query, operation_timeout=operation_timeout) + except core_exceptions.DeadlineExceeded as e: + assert ( + e.message + == f"operation_timeout of {operation_timeout:0.1f}s exceeded" + ) + + @pytest.mark.parametrize( + "per_request_t, operation_t, expected_num", + [ + (0.05, 0.08, 2), + (0.05, 0.54, 11), + (0.05, 0.14, 3), + (0.05, 0.24, 5), + ], + ) + @pytest.mark.asyncio + async def test_read_rows_per_request_timeout( + self, per_request_t, operation_t, expected_num + ): + """ + Ensures that the per_request_timeout is respected and that the number of + requests is as expected. + + operation_timeout does not cancel the request, so we expect the number of + requests to be the ceiling of operation_timeout / per_request_timeout. + """ + from google.cloud.bigtable.exceptions import RetryExceptionGroup + + expected_last_timeout = operation_t - (expected_num - 1) * per_request_t + + # mocking uniform ensures there are no sleeps between retries + with mock.patch("random.uniform", side_effect=lambda a, b: 0): + async with self._make_client() as client: + table = client.get_table("instance", "table") + query = ReadRowsQuery() + chunks = [core_exceptions.DeadlineExceeded("mock deadline")] + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = ( + lambda *args, **kwargs: self._make_gapic_stream( + chunks, sleep_time=per_request_t + ) + ) + try: + await table.read_rows( + query, + operation_timeout=operation_t, + per_request_timeout=per_request_t, + ) + except core_exceptions.DeadlineExceeded as e: + retry_exc = e.__cause__ + if expected_num == 0: + assert retry_exc is None + else: + assert type(retry_exc) == RetryExceptionGroup + assert f"{expected_num} failed attempts" in str(retry_exc) + assert len(retry_exc.exceptions) == expected_num + for sub_exc in retry_exc.exceptions: + assert sub_exc.message == "mock deadline" + assert read_rows.call_count == expected_num + # check timeouts + for _, call_kwargs in read_rows.call_args_list[:-1]: + assert call_kwargs["timeout"] == per_request_t + # last timeout should be adjusted to account for the time spent + assert ( + abs( + read_rows.call_args_list[-1][1]["timeout"] + - expected_last_timeout + ) + < 0.05 + ) + + @pytest.mark.asyncio + async def test_read_rows_idle_timeout(self): + from google.cloud.bigtable.client import ReadRowsIterator + from google.cloud.bigtable_v2.services.bigtable.async_client import ( + BigtableAsyncClient, + ) + from google.cloud.bigtable.exceptions import IdleTimeout + from google.cloud.bigtable._read_rows import _ReadRowsOperation + + chunks = [ + self._make_chunk(row_key=b"test_1"), + self._make_chunk(row_key=b"test_2"), + ] + with mock.patch.object(BigtableAsyncClient, "read_rows") as read_rows: + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + chunks + ) + with mock.patch.object( + ReadRowsIterator, "_start_idle_timer" + ) as start_idle_timer: + client = self._make_client() + table = client.get_table("instance", "table") + query = ReadRowsQuery() + gen = await table.read_rows_stream(query) + # should start idle timer on creation + start_idle_timer.assert_called_once() + with mock.patch.object(_ReadRowsOperation, "aclose", AsyncMock()) as aclose: + # start idle timer with our own value + await gen._start_idle_timer(0.1) + # should timeout after being abandoned + await gen.__anext__() + await asyncio.sleep(0.2) + # generator should be expired + assert not gen.active + assert type(gen._error) == IdleTimeout + assert gen._idle_timeout_task is None + await client.close() + with pytest.raises(IdleTimeout) as e: + await gen.__anext__() + + expected_msg = ( + "Timed out waiting for next Row to be consumed. (idle_timeout=0.1s)" + ) + assert e.value.message == expected_msg + aclose.assert_called_once() + aclose.assert_awaited() + + @pytest.mark.parametrize( + "exc_type", + [ + core_exceptions.Aborted, + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ], + ) + @pytest.mark.asyncio + async def test_read_rows_retryable_error(self, exc_type): + async with self._make_client() as client: + table = client.get_table("instance", "table") + query = ReadRowsQuery() + expected_error = exc_type("mock error") + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + [expected_error] + ) + try: + await table.read_rows(query, operation_timeout=0.1) + except core_exceptions.DeadlineExceeded as e: + retry_exc = e.__cause__ + root_cause = retry_exc.exceptions[0] + assert type(root_cause) == exc_type + assert root_cause == expected_error + + @pytest.mark.parametrize( + "exc_type", + [ + core_exceptions.Cancelled, + core_exceptions.PreconditionFailed, + core_exceptions.NotFound, + core_exceptions.PermissionDenied, + core_exceptions.Conflict, + core_exceptions.InternalServerError, + core_exceptions.TooManyRequests, + core_exceptions.ResourceExhausted, + InvalidChunk, + ], + ) + @pytest.mark.asyncio + async def test_read_rows_non_retryable_error(self, exc_type): + async with self._make_client() as client: + table = client.get_table("instance", "table") + query = ReadRowsQuery() + expected_error = exc_type("mock error") + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + [expected_error] + ) + try: + await table.read_rows(query, operation_timeout=0.1) + except exc_type as e: + assert e == expected_error + + @pytest.mark.asyncio + async def test_read_rows_revise_request(self): + """ + Ensure that _revise_request is called between retries + """ + from google.cloud.bigtable._read_rows import _ReadRowsOperation + from google.cloud.bigtable.exceptions import InvalidChunk + + with mock.patch.object( + _ReadRowsOperation, "_revise_request_rowset" + ) as revise_rowset: + with mock.patch.object(_ReadRowsOperation, "aclose"): + revise_rowset.return_value = "modified" + async with self._make_client() as client: + table = client.get_table("instance", "table") + row_keys = [b"test_1", b"test_2", b"test_3"] + query = ReadRowsQuery(row_keys=row_keys) + chunks = [ + self._make_chunk(row_key=b"test_1"), + core_exceptions.Aborted("mock retryable error"), + ] + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = ( + lambda *args, **kwargs: self._make_gapic_stream(chunks) + ) + try: + await table.read_rows(query) + except InvalidChunk: + revise_rowset.assert_called() + revise_call_kwargs = revise_rowset.call_args_list[0].kwargs + assert ( + revise_call_kwargs["row_set"] + == query._to_dict()["rows"] + ) + assert revise_call_kwargs["last_seen_row_key"] == b"test_1" + read_rows_request = read_rows.call_args_list[1].args[0] + assert read_rows_request["rows"] == "modified" + + @pytest.mark.asyncio + async def test_read_rows_default_timeouts(self): + """ + Ensure that the default timeouts are set on the read rows operation when not overridden + """ + from google.cloud.bigtable._read_rows import _ReadRowsOperation + + operation_timeout = 8 + per_request_timeout = 4 + with mock.patch.object(_ReadRowsOperation, "__init__") as mock_op: + mock_op.side_effect = RuntimeError("mock error") + async with self._make_client() as client: + async with client.get_table( + "instance", + "table", + default_operation_timeout=operation_timeout, + default_per_request_timeout=per_request_timeout, + ) as table: + try: + await table.read_rows(ReadRowsQuery()) + except RuntimeError: + pass + kwargs = mock_op.call_args_list[0].kwargs + assert kwargs["operation_timeout"] == operation_timeout + assert kwargs["per_request_timeout"] == per_request_timeout + + @pytest.mark.asyncio + async def test_read_rows_default_timeout_override(self): + """ + When timeouts are passed, they overwrite default values + """ + from google.cloud.bigtable._read_rows import _ReadRowsOperation + + operation_timeout = 8 + per_request_timeout = 4 + with mock.patch.object(_ReadRowsOperation, "__init__") as mock_op: + mock_op.side_effect = RuntimeError("mock error") + async with self._make_client() as client: + async with client.get_table( + "instance", + "table", + default_operation_timeout=99, + default_per_request_timeout=97, + ) as table: + try: + await table.read_rows( + ReadRowsQuery(), + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + except RuntimeError: + pass + kwargs = mock_op.call_args_list[0].kwargs + assert kwargs["operation_timeout"] == operation_timeout + assert kwargs["per_request_timeout"] == per_request_timeout diff --git a/tests/unit/test_iterators.py b/tests/unit/test_iterators.py new file mode 100644 index 000000000..f7aee2822 --- /dev/null +++ b/tests/unit/test_iterators.py @@ -0,0 +1,251 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +from __future__ import annotations + +import sys +import asyncio +import pytest + +from google.cloud.bigtable._read_rows import _ReadRowsOperation + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock +except ImportError: # pragma: NO COVER + import mock # type: ignore + + +class MockStream(_ReadRowsOperation): + """ + Mock a _ReadRowsOperation stream for testing + """ + + def __init__(self, items=None, errors=None, operation_timeout=None): + self.transient_errors = errors + self.operation_timeout = operation_timeout + self.next_idx = 0 + if items is None: + items = list(range(10)) + self.items = items + + def __aiter__(self): + return self + + async def __anext__(self): + if self.next_idx >= len(self.items): + raise StopAsyncIteration + item = self.items[self.next_idx] + self.next_idx += 1 + if isinstance(item, Exception): + raise item + return item + + async def aclose(self): + pass + + +class TestReadRowsIterator: + async def mock_stream(self, size=10): + for i in range(size): + yield i + + def _make_one(self, *args, **kwargs): + from google.cloud.bigtable.iterators import ReadRowsIterator + + stream = MockStream(*args, **kwargs) + return ReadRowsIterator(stream) + + def test_ctor(self): + with mock.patch("time.time", return_value=0): + iterator = self._make_one() + assert iterator.last_interaction_time == 0 + assert iterator._idle_timeout_task is None + assert iterator.active is True + + def test___aiter__(self): + iterator = self._make_one() + assert iterator.__aiter__() is iterator + + @pytest.mark.skipif( + sys.version_info < (3, 8), reason="mock coroutine requires python3.8 or higher" + ) + @pytest.mark.asyncio + async def test__start_idle_timer(self): + """Should start timer coroutine""" + iterator = self._make_one() + expected_timeout = 10 + with mock.patch("time.time", return_value=1): + with mock.patch.object(iterator, "_idle_timeout_coroutine") as mock_coro: + await iterator._start_idle_timer(expected_timeout) + assert mock_coro.call_count == 1 + assert mock_coro.call_args[0] == (expected_timeout,) + assert iterator.last_interaction_time == 1 + assert iterator._idle_timeout_task is not None + + @pytest.mark.skipif( + sys.version_info < (3, 8), reason="mock coroutine requires python3.8 or higher" + ) + @pytest.mark.asyncio + async def test__start_idle_timer_duplicate(self): + """Multiple calls should replace task""" + iterator = self._make_one() + with mock.patch.object(iterator, "_idle_timeout_coroutine") as mock_coro: + await iterator._start_idle_timer(1) + first_task = iterator._idle_timeout_task + await iterator._start_idle_timer(2) + second_task = iterator._idle_timeout_task + assert mock_coro.call_count == 2 + + assert first_task is not None + assert first_task != second_task + # old tasks hould be cancelled + with pytest.raises(asyncio.CancelledError): + await first_task + # new task should not be cancelled + await second_task + + @pytest.mark.asyncio + async def test__idle_timeout_coroutine(self): + from google.cloud.bigtable.exceptions import IdleTimeout + + iterator = self._make_one() + await iterator._idle_timeout_coroutine(0.05) + await asyncio.sleep(0.1) + assert iterator.active is False + with pytest.raises(IdleTimeout): + await iterator.__anext__() + + @pytest.mark.asyncio + async def test__idle_timeout_coroutine_extensions(self): + """touching the generator should reset the idle timer""" + iterator = self._make_one(items=list(range(100))) + await iterator._start_idle_timer(0.05) + for i in range(10): + # will not expire as long as it is in use + assert iterator.active is True + await iterator.__anext__() + await asyncio.sleep(0.03) + # now let it expire + await asyncio.sleep(0.5) + assert iterator.active is False + + @pytest.mark.asyncio + async def test___anext__(self): + num_rows = 10 + iterator = self._make_one(items=list(range(num_rows))) + for i in range(num_rows): + assert await iterator.__anext__() == i + with pytest.raises(StopAsyncIteration): + await iterator.__anext__() + + @pytest.mark.asyncio + async def test___anext__with_deadline_error(self): + """ + RetryErrors mean a deadline has been hit. + Should be wrapped in a DeadlineExceeded exception + """ + from google.api_core import exceptions as core_exceptions + + items = [1, core_exceptions.RetryError("retry error", None)] + expected_timeout = 99 + iterator = self._make_one(items=items, operation_timeout=expected_timeout) + assert await iterator.__anext__() == 1 + with pytest.raises(core_exceptions.DeadlineExceeded) as exc: + await iterator.__anext__() + assert f"operation_timeout of {expected_timeout:0.1f}s exceeded" in str( + exc.value + ) + assert exc.value.__cause__ is None + + @pytest.mark.asyncio + async def test___anext__with_deadline_error_with_cause(self): + """ + Transient errors should be exposed as an error group + """ + from google.api_core import exceptions as core_exceptions + from google.cloud.bigtable.exceptions import RetryExceptionGroup + + items = [1, core_exceptions.RetryError("retry error", None)] + expected_timeout = 99 + errors = [RuntimeError("error1"), ValueError("error2")] + iterator = self._make_one( + items=items, operation_timeout=expected_timeout, errors=errors + ) + assert await iterator.__anext__() == 1 + with pytest.raises(core_exceptions.DeadlineExceeded) as exc: + await iterator.__anext__() + assert f"operation_timeout of {expected_timeout:0.1f}s exceeded" in str( + exc.value + ) + error_group = exc.value.__cause__ + assert isinstance(error_group, RetryExceptionGroup) + assert len(error_group.exceptions) == 2 + assert error_group.exceptions[0] is errors[0] + assert error_group.exceptions[1] is errors[1] + assert "2 failed attempts" in str(error_group) + + @pytest.mark.asyncio + async def test___anext__with_error(self): + """ + Other errors should be raised as-is + """ + from google.api_core import exceptions as core_exceptions + + items = [1, core_exceptions.InternalServerError("mock error")] + iterator = self._make_one(items=items) + assert await iterator.__anext__() == 1 + with pytest.raises(core_exceptions.InternalServerError) as exc: + await iterator.__anext__() + assert exc.value is items[1] + assert iterator.active is False + # next call should raise same error + with pytest.raises(core_exceptions.InternalServerError) as exc: + await iterator.__anext__() + + @pytest.mark.asyncio + async def test__finish_with_error(self): + iterator = self._make_one() + await iterator._start_idle_timer(10) + timeout_task = iterator._idle_timeout_task + assert await iterator.__anext__() == 0 + assert iterator.active is True + err = ZeroDivisionError("mock error") + await iterator._finish_with_error(err) + assert iterator.active is False + assert iterator._error is err + assert iterator._idle_timeout_task is None + with pytest.raises(ZeroDivisionError) as exc: + await iterator.__anext__() + assert exc.value is err + # timeout task should be cancelled + with pytest.raises(asyncio.CancelledError): + await timeout_task + + @pytest.mark.asyncio + async def test_aclose(self): + iterator = self._make_one() + await iterator._start_idle_timer(10) + timeout_task = iterator._idle_timeout_task + assert await iterator.__anext__() == 0 + assert iterator.active is True + await iterator.aclose() + assert iterator.active is False + assert isinstance(iterator._error, StopAsyncIteration) + assert iterator._idle_timeout_task is None + with pytest.raises(StopAsyncIteration) as e: + await iterator.__anext__() + assert "closed" in str(e.value) + # timeout task should be cancelled + with pytest.raises(asyncio.CancelledError): + await timeout_task diff --git a/tests/unit/test_read_rows_acceptance.py b/tests/unit/test_read_rows_acceptance.py new file mode 100644 index 000000000..2349d25c6 --- /dev/null +++ b/tests/unit/test_read_rows_acceptance.py @@ -0,0 +1,314 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +from __future__ import annotations + +import os +from itertools import zip_longest + +import pytest +import mock + +from google.cloud.bigtable_v2 import ReadRowsResponse + +from google.cloud.bigtable.client import BigtableDataClient +from google.cloud.bigtable.exceptions import InvalidChunk +from google.cloud.bigtable._read_rows import _ReadRowsOperation, _StateMachine +from google.cloud.bigtable.row import Row + +from .v2_client.test_row_merger import ReadRowsTest, TestFile + + +def parse_readrows_acceptance_tests(): + dirname = os.path.dirname(__file__) + filename = os.path.join(dirname, "./read-rows-acceptance-test.json") + + with open(filename) as json_file: + test_json = TestFile.from_json(json_file.read()) + return test_json.read_rows_tests + + +def extract_results_from_row(row: Row): + results = [] + for family, col, cells in row.items(): + for cell in cells: + results.append( + ReadRowsTest.Result( + row_key=row.row_key, + family_name=family, + qualifier=col, + timestamp_micros=cell.timestamp_ns // 1000, + value=cell.value, + label=(cell.labels[0] if cell.labels else ""), + ) + ) + return results + + +@pytest.mark.parametrize( + "test_case", parse_readrows_acceptance_tests(), ids=lambda t: t.description +) +@pytest.mark.asyncio +async def test_row_merger_scenario(test_case: ReadRowsTest): + async def _scenerio_stream(): + for chunk in test_case.chunks: + yield ReadRowsResponse(chunks=[chunk]) + + try: + state = _StateMachine() + results = [] + async for row in _ReadRowsOperation.merge_row_response_stream( + _scenerio_stream(), state + ): + for cell in row: + cell_result = ReadRowsTest.Result( + row_key=cell.row_key, + family_name=cell.family, + qualifier=cell.qualifier, + timestamp_micros=cell.timestamp_micros, + value=cell.value, + label=cell.labels[0] if cell.labels else "", + ) + results.append(cell_result) + if not state.is_terminal_state(): + raise InvalidChunk("state machine has partial frame after reading") + except InvalidChunk: + results.append(ReadRowsTest.Result(error=True)) + for expected, actual in zip_longest(test_case.results, results): + assert actual == expected + + +@pytest.mark.parametrize( + "test_case", parse_readrows_acceptance_tests(), ids=lambda t: t.description +) +@pytest.mark.asyncio +async def test_read_rows_scenario(test_case: ReadRowsTest): + async def _make_gapic_stream(chunk_list: list[ReadRowsResponse]): + from google.cloud.bigtable_v2 import ReadRowsResponse + + class mock_stream: + def __init__(self, chunk_list): + self.chunk_list = chunk_list + self.idx = -1 + + def __aiter__(self): + return self + + async def __anext__(self): + self.idx += 1 + if len(self.chunk_list) > self.idx: + chunk = self.chunk_list[self.idx] + return ReadRowsResponse(chunks=[chunk]) + raise StopAsyncIteration + + def cancel(self): + pass + + return mock_stream(chunk_list) + + try: + client = BigtableDataClient() + table = client.get_table("instance", "table") + results = [] + with mock.patch.object(table.client._gapic_client, "read_rows") as read_rows: + # run once, then return error on retry + read_rows.return_value = _make_gapic_stream(test_case.chunks) + async for row in await table.read_rows_stream(query={}): + for cell in row: + cell_result = ReadRowsTest.Result( + row_key=cell.row_key, + family_name=cell.family, + qualifier=cell.qualifier, + timestamp_micros=cell.timestamp_micros, + value=cell.value, + label=cell.labels[0] if cell.labels else "", + ) + results.append(cell_result) + except InvalidChunk: + results.append(ReadRowsTest.Result(error=True)) + finally: + await client.close() + for expected, actual in zip_longest(test_case.results, results): + assert actual == expected + + +@pytest.mark.asyncio +async def test_out_of_order_rows(): + async def _row_stream(): + yield ReadRowsResponse(last_scanned_row_key=b"a") + + state = _StateMachine() + state.last_seen_row_key = b"a" + with pytest.raises(InvalidChunk): + async for _ in _ReadRowsOperation.merge_row_response_stream( + _row_stream(), state + ): + pass + + +@pytest.mark.asyncio +async def test_bare_reset(): + first_chunk = ReadRowsResponse.CellChunk( + ReadRowsResponse.CellChunk( + row_key=b"a", family_name="f", qualifier=b"q", value=b"v" + ) + ) + with pytest.raises(InvalidChunk): + await _process_chunks( + first_chunk, + ReadRowsResponse.CellChunk( + ReadRowsResponse.CellChunk(reset_row=True, row_key=b"a") + ), + ) + with pytest.raises(InvalidChunk): + await _process_chunks( + first_chunk, + ReadRowsResponse.CellChunk( + ReadRowsResponse.CellChunk(reset_row=True, family_name="f") + ), + ) + with pytest.raises(InvalidChunk): + await _process_chunks( + first_chunk, + ReadRowsResponse.CellChunk( + ReadRowsResponse.CellChunk(reset_row=True, qualifier=b"q") + ), + ) + with pytest.raises(InvalidChunk): + await _process_chunks( + first_chunk, + ReadRowsResponse.CellChunk( + ReadRowsResponse.CellChunk(reset_row=True, timestamp_micros=1000) + ), + ) + with pytest.raises(InvalidChunk): + await _process_chunks( + first_chunk, + ReadRowsResponse.CellChunk( + ReadRowsResponse.CellChunk(reset_row=True, labels=["a"]) + ), + ) + with pytest.raises(InvalidChunk): + await _process_chunks( + first_chunk, + ReadRowsResponse.CellChunk( + ReadRowsResponse.CellChunk(reset_row=True, value=b"v") + ), + ) + + +@pytest.mark.asyncio +async def test_missing_family(): + with pytest.raises(InvalidChunk): + await _process_chunks( + ReadRowsResponse.CellChunk( + row_key=b"a", + qualifier=b"q", + timestamp_micros=1000, + value=b"v", + commit_row=True, + ) + ) + + +@pytest.mark.asyncio +async def test_mid_cell_row_key_change(): + with pytest.raises(InvalidChunk): + await _process_chunks( + ReadRowsResponse.CellChunk( + row_key=b"a", + family_name="f", + qualifier=b"q", + timestamp_micros=1000, + value_size=2, + value=b"v", + ), + ReadRowsResponse.CellChunk(row_key=b"b", value=b"v", commit_row=True), + ) + + +@pytest.mark.asyncio +async def test_mid_cell_family_change(): + with pytest.raises(InvalidChunk): + await _process_chunks( + ReadRowsResponse.CellChunk( + row_key=b"a", + family_name="f", + qualifier=b"q", + timestamp_micros=1000, + value_size=2, + value=b"v", + ), + ReadRowsResponse.CellChunk(family_name="f2", value=b"v", commit_row=True), + ) + + +@pytest.mark.asyncio +async def test_mid_cell_qualifier_change(): + with pytest.raises(InvalidChunk): + await _process_chunks( + ReadRowsResponse.CellChunk( + row_key=b"a", + family_name="f", + qualifier=b"q", + timestamp_micros=1000, + value_size=2, + value=b"v", + ), + ReadRowsResponse.CellChunk(qualifier=b"q2", value=b"v", commit_row=True), + ) + + +@pytest.mark.asyncio +async def test_mid_cell_timestamp_change(): + with pytest.raises(InvalidChunk): + await _process_chunks( + ReadRowsResponse.CellChunk( + row_key=b"a", + family_name="f", + qualifier=b"q", + timestamp_micros=1000, + value_size=2, + value=b"v", + ), + ReadRowsResponse.CellChunk( + timestamp_micros=2000, value=b"v", commit_row=True + ), + ) + + +@pytest.mark.asyncio +async def test_mid_cell_labels_change(): + with pytest.raises(InvalidChunk): + await _process_chunks( + ReadRowsResponse.CellChunk( + row_key=b"a", + family_name="f", + qualifier=b"q", + timestamp_micros=1000, + value_size=2, + value=b"v", + ), + ReadRowsResponse.CellChunk(labels=["b"], value=b"v", commit_row=True), + ) + + +async def _process_chunks(*chunks): + async def _row_stream(): + yield ReadRowsResponse(chunks=chunks) + + state = _StateMachine() + results = [] + async for row in _ReadRowsOperation.merge_row_response_stream(_row_stream(), state): + results.append(row) + return results diff --git a/tests/unit/test_row.py b/tests/unit/test_row.py index 32afe92f5..1af09aad9 100644 --- a/tests/unit/test_row.py +++ b/tests/unit/test_row.py @@ -76,7 +76,7 @@ def test_get_cells(self): output = row_response.get_cells(family="1", qualifier=q) self.assertEqual(len(output), 1) self.assertEqual(output[0].family, "1") - self.assertEqual(output[0].column_qualifier, b"a") + self.assertEqual(output[0].qualifier, b"a") self.assertEqual(output[0], cell_list[0]) # calling with just qualifier should raise an error with self.assertRaises(ValueError): @@ -90,7 +90,6 @@ def test_get_cells(self): row_response.get_cells(family="1", qualifier=b"c") def test___repr__(self): - cell_str = ( "{'value': b'1234', 'timestamp_micros': %d, 'labels': ['label1', 'label2']}" % (TEST_TIMESTAMP) @@ -180,7 +179,6 @@ def test_to_dict(self): self.assertEqual(column.cells[1].labels, TEST_LABELS) def test_iteration(self): - from types import GeneratorType from google.cloud.bigtable.row import Cell # should be able to iterate over the Row as a list @@ -189,8 +187,6 @@ def test_iteration(self): cell3 = self._make_cell(value=b"3") row_response = self._make_one(TEST_ROW_KEY, [cell1, cell2, cell3]) self.assertEqual(len(row_response), 3) - # should create generator object - self.assertIsInstance(iter(row_response), GeneratorType) result_list = list(row_response) self.assertEqual(len(result_list), 3) # should be able to iterate over all cells @@ -487,7 +483,7 @@ def test_ctor(self): self.assertEqual(cell.value, TEST_VALUE) self.assertEqual(cell.row_key, TEST_ROW_KEY) self.assertEqual(cell.family, TEST_FAMILY_ID) - self.assertEqual(cell.column_qualifier, TEST_QUALIFIER) + self.assertEqual(cell.qualifier, TEST_QUALIFIER) self.assertEqual(cell.timestamp_micros, TEST_TIMESTAMP) self.assertEqual(cell.labels, TEST_LABELS) @@ -587,8 +583,8 @@ def test___repr__(self): cell = self._make_one() expected = ( - "Cell(value=b'1234', row=b'row', " - + "family='cf1', column_qualifier=b'col', " + "Cell(value=b'1234', row_key=b'row', " + + "family='cf1', qualifier=b'col', " + f"timestamp_micros={TEST_TIMESTAMP}, labels=['label1', 'label2'])" ) self.assertEqual(repr(cell), expected) @@ -608,8 +604,8 @@ def test___repr___no_labels(self): None, ) expected = ( - "Cell(value=b'1234', row=b'row', " - + "family='cf1', column_qualifier=b'col', " + "Cell(value=b'1234', row_key=b'row', " + + "family='cf1', qualifier=b'col', " + f"timestamp_micros={TEST_TIMESTAMP}, labels=[])" ) self.assertEqual(repr(cell_no_labels), expected) From 9b81289a503743084d81b47efbd1801f6d7c9b85 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 6 Jun 2023 12:09:25 -0700 Subject: [PATCH 07/56] feat: implement mutate rows (#769) --- google/cloud/bigtable/__init__.py | 4 +- google/cloud/bigtable/_helpers.py | 111 +++++ google/cloud/bigtable/_mutate_rows.py | 210 ++++++++++ google/cloud/bigtable/_read_rows.py | 31 +- google/cloud/bigtable/client.py | 102 ++++- google/cloud/bigtable/exceptions.py | 123 +++--- google/cloud/bigtable/iterators.py | 2 +- google/cloud/bigtable/mutations.py | 191 ++++++++- tests/system/test_system.py | 73 ++++ tests/unit/test__helpers.py | 145 +++++++ tests/unit/test__mutate_rows.py | 290 +++++++++++++ tests/unit/test__read_rows.py | 43 +- tests/unit/test_client.py | 577 ++++++++++++++++++++++++++ tests/unit/test_exceptions.py | 239 +++++++++++ tests/unit/test_mutations.py | 569 +++++++++++++++++++++++++ 15 files changed, 2596 insertions(+), 114 deletions(-) create mode 100644 google/cloud/bigtable/_helpers.py create mode 100644 google/cloud/bigtable/_mutate_rows.py create mode 100644 tests/unit/test__helpers.py create mode 100644 tests/unit/test__mutate_rows.py create mode 100644 tests/unit/test_exceptions.py create mode 100644 tests/unit/test_mutations.py diff --git a/google/cloud/bigtable/__init__.py b/google/cloud/bigtable/__init__.py index de46f8a75..70c87ade1 100644 --- a/google/cloud/bigtable/__init__.py +++ b/google/cloud/bigtable/__init__.py @@ -28,7 +28,7 @@ from google.cloud.bigtable.mutations_batcher import MutationsBatcher from google.cloud.bigtable.mutations import Mutation -from google.cloud.bigtable.mutations import BulkMutationsEntry +from google.cloud.bigtable.mutations import RowMutationEntry from google.cloud.bigtable.mutations import SetCell from google.cloud.bigtable.mutations import DeleteRangeFromColumn from google.cloud.bigtable.mutations import DeleteAllFromFamily @@ -47,7 +47,7 @@ "RowRange", "MutationsBatcher", "Mutation", - "BulkMutationsEntry", + "RowMutationEntry", "SetCell", "DeleteRangeFromColumn", "DeleteAllFromFamily", diff --git a/google/cloud/bigtable/_helpers.py b/google/cloud/bigtable/_helpers.py new file mode 100644 index 000000000..dec4c2014 --- /dev/null +++ b/google/cloud/bigtable/_helpers.py @@ -0,0 +1,111 @@ +# +# Licensed 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. +# +from __future__ import annotations + +from typing import Callable, Any +from inspect import iscoroutinefunction +import time + +from google.api_core import exceptions as core_exceptions +from google.cloud.bigtable.exceptions import RetryExceptionGroup + +""" +Helper functions used in various places in the library. +""" + + +def _make_metadata( + table_name: str, app_profile_id: str | None +) -> list[tuple[str, str]]: + """ + Create properly formatted gRPC metadata for requests. + """ + params = [] + params.append(f"table_name={table_name}") + if app_profile_id is not None: + params.append(f"app_profile_id={app_profile_id}") + params_str = ",".join(params) + return [("x-goog-request-params", params_str)] + + +def _attempt_timeout_generator( + per_request_timeout: float | None, operation_timeout: float +): + """ + Generator that yields the timeout value for each attempt of a retry loop. + + Will return per_request_timeout until the operation_timeout is approached, + at which point it will return the remaining time in the operation_timeout. + + Args: + - per_request_timeout: The timeout value to use for each request, in seconds. + If None, the operation_timeout will be used for each request. + - operation_timeout: The timeout value to use for the entire operationm in seconds. + Yields: + - The timeout value to use for the next request, in seonds + """ + per_request_timeout = ( + per_request_timeout if per_request_timeout is not None else operation_timeout + ) + deadline = operation_timeout + time.monotonic() + while True: + yield max(0, min(per_request_timeout, deadline - time.monotonic())) + + +def _convert_retry_deadline( + func: Callable[..., Any], + timeout_value: float | None = None, + retry_errors: list[Exception] | None = None, +): + """ + Decorator to convert RetryErrors raised by api_core.retry into + DeadlineExceeded exceptions, indicating that the underlying retries have + exhaused the timeout value. + Optionally attaches a RetryExceptionGroup to the DeadlineExceeded.__cause__, + detailing the failed exceptions associated with each retry. + + Supports both sync and async function wrapping. + + Args: + - func: The function to decorate + - timeout_value: The timeout value to display in the DeadlineExceeded error message + - retry_errors: An optional list of exceptions to attach as a RetryExceptionGroup to the DeadlineExceeded.__cause__ + """ + timeout_str = f" of {timeout_value:.1f}s" if timeout_value is not None else "" + error_str = f"operation_timeout{timeout_str} exceeded" + + def handle_error(): + new_exc = core_exceptions.DeadlineExceeded( + error_str, + ) + source_exc = None + if retry_errors: + source_exc = RetryExceptionGroup(retry_errors) + new_exc.__cause__ = source_exc + raise new_exc from source_exc + + # separate wrappers for async and sync functions + async def wrapper_async(*args, **kwargs): + try: + return await func(*args, **kwargs) + except core_exceptions.RetryError: + handle_error() + + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except core_exceptions.RetryError: + handle_error() + + return wrapper_async if iscoroutinefunction(func) else wrapper diff --git a/google/cloud/bigtable/_mutate_rows.py b/google/cloud/bigtable/_mutate_rows.py new file mode 100644 index 000000000..a422c99b2 --- /dev/null +++ b/google/cloud/bigtable/_mutate_rows.py @@ -0,0 +1,210 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from typing import TYPE_CHECKING +import functools + +from google.api_core import exceptions as core_exceptions +from google.api_core import retry_async as retries +import google.cloud.bigtable.exceptions as bt_exceptions +from google.cloud.bigtable._helpers import _make_metadata +from google.cloud.bigtable._helpers import _convert_retry_deadline +from google.cloud.bigtable._helpers import _attempt_timeout_generator + +if TYPE_CHECKING: + from google.cloud.bigtable_v2.services.bigtable.async_client import ( + BigtableAsyncClient, + ) + from google.cloud.bigtable.client import Table + from google.cloud.bigtable.mutations import RowMutationEntry + + +class _MutateRowsIncomplete(RuntimeError): + """ + Exception raised when a mutate_rows call has unfinished work. + """ + + pass + + +class _MutateRowsOperation: + """ + MutateRowsOperation manages the logic of sending a set of row mutations, + and retrying on failed entries. It manages this using the _run_attempt + function, which attempts to mutate all outstanding entries, and raises + _MutateRowsIncomplete if any retryable errors are encountered. + + Errors are exposed as a MutationsExceptionGroup, which contains a list of + exceptions organized by the related failed mutation entries. + """ + + def __init__( + self, + gapic_client: "BigtableAsyncClient", + table: "Table", + mutation_entries: list["RowMutationEntry"], + operation_timeout: float, + per_request_timeout: float | None, + ): + """ + Args: + - gapic_client: the client to use for the mutate_rows call + - table: the table associated with the request + - mutation_entries: a list of RowMutationEntry objects to send to the server + - operation_timeout: the timeout t o use for the entire operation, in seconds. + - per_request_timeout: the timeoutto use for each mutate_rows attempt, in seconds. + If not specified, the request will run until operation_timeout is reached. + """ + # create partial function to pass to trigger rpc call + metadata = _make_metadata(table.table_name, table.app_profile_id) + self._gapic_fn = functools.partial( + gapic_client.mutate_rows, + table_name=table.table_name, + app_profile_id=table.app_profile_id, + metadata=metadata, + ) + # create predicate for determining which errors are retryable + self.is_retryable = retries.if_exception_type( + # RPC level errors + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + # Entry level errors + _MutateRowsIncomplete, + ) + # build retryable operation + retry = retries.AsyncRetry( + predicate=self.is_retryable, + timeout=operation_timeout, + initial=0.01, + multiplier=2, + maximum=60, + ) + retry_wrapped = retry(self._run_attempt) + self._operation = _convert_retry_deadline(retry_wrapped, operation_timeout) + # initialize state + self.timeout_generator = _attempt_timeout_generator( + per_request_timeout, operation_timeout + ) + self.mutations = mutation_entries + self.remaining_indices = list(range(len(self.mutations))) + self.errors: dict[int, list[Exception]] = {} + + async def start(self): + """ + Start the operation, and run until completion + + Raises: + - MutationsExceptionGroup: if any mutations failed + """ + try: + # trigger mutate_rows + await self._operation() + except Exception as exc: + # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations + incomplete_indices = self.remaining_indices.copy() + for idx in incomplete_indices: + self._handle_entry_error(idx, exc) + finally: + # raise exception detailing incomplete mutations + all_errors = [] + for idx, exc_list in self.errors.items(): + if len(exc_list) == 0: + raise core_exceptions.ClientError( + f"Mutation {idx} failed with no associated errors" + ) + elif len(exc_list) == 1: + cause_exc = exc_list[0] + else: + cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) + entry = self.mutations[idx] + all_errors.append( + bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) + ) + if all_errors: + raise bt_exceptions.MutationsExceptionGroup( + all_errors, len(self.mutations) + ) + + async def _run_attempt(self): + """ + Run a single attempt of the mutate_rows rpc. + + Raises: + - _MutateRowsIncomplete: if there are failed mutations eligible for + retry after the attempt is complete + - GoogleAPICallError: if the gapic rpc fails + """ + request_entries = [ + self.mutations[idx]._to_dict() for idx in self.remaining_indices + ] + # track mutations in this request that have not been finalized yet + active_request_indices = { + req_idx: orig_idx for req_idx, orig_idx in enumerate(self.remaining_indices) + } + self.remaining_indices = [] + if not request_entries: + # no more mutations. return early + return + # make gapic request + try: + result_generator = await self._gapic_fn( + timeout=next(self.timeout_generator), + entries=request_entries, + ) + async for result_list in result_generator: + for result in result_list.entries: + # convert sub-request index to global index + orig_idx = active_request_indices[result.index] + entry_error = core_exceptions.from_grpc_status( + result.status.code, + result.status.message, + details=result.status.details, + ) + if result.status.code != 0: + # mutation failed; update error list (and remaining_indices if retryable) + self._handle_entry_error(orig_idx, entry_error) + # remove processed entry from active list + del active_request_indices[result.index] + except Exception as exc: + # add this exception to list for each mutation that wasn't + # already handled, and update remaining_indices if mutation is retryable + for idx in active_request_indices.values(): + self._handle_entry_error(idx, exc) + # bubble up exception to be handled by retry wrapper + raise + # check if attempt succeeded, or needs to be retried + if self.remaining_indices: + # unfinished work; raise exception to trigger retry + raise _MutateRowsIncomplete + + def _handle_entry_error(self, idx: int, exc: Exception): + """ + Add an exception to the list of exceptions for a given mutation index, + and add the index to the list of remaining indices if the exception is + retryable. + + Args: + - idx: the index of the mutation that failed + - exc: the exception to add to the list + """ + entry = self.mutations[idx] + self.errors.setdefault(idx, []).append(exc) + if ( + entry.is_idempotent() + and self.is_retryable(exc) + and idx not in self.remaining_indices + ): + self.remaining_indices.append(idx) diff --git a/google/cloud/bigtable/_read_rows.py b/google/cloud/bigtable/_read_rows.py index 1c9e02d5a..ee094f1a7 100644 --- a/google/cloud/bigtable/_read_rows.py +++ b/google/cloud/bigtable/_read_rows.py @@ -20,13 +20,13 @@ AsyncIterable, AsyncIterator, AsyncGenerator, + Iterator, Callable, Awaitable, Type, ) import asyncio -import time from functools import partial from grpc.aio import RpcContext @@ -37,6 +37,8 @@ from google.cloud.bigtable.exceptions import _RowSetComplete from google.api_core import retry_async as retries from google.api_core import exceptions as core_exceptions +from google.cloud.bigtable._helpers import _make_metadata +from google.cloud.bigtable._helpers import _attempt_timeout_generator """ This module provides a set of classes for merging ReadRowsResponse chunks @@ -87,16 +89,16 @@ def __init__( self._emit_count = 0 self._request = request self.operation_timeout = operation_timeout - deadline = operation_timeout + time.monotonic() + # use generator to lower per-attempt timeout as we approach operation_timeout deadline + attempt_timeout_gen = _attempt_timeout_generator( + per_request_timeout, operation_timeout + ) row_limit = request.get("rows_limit", 0) - if per_request_timeout is None: - per_request_timeout = operation_timeout # lock in paramters for retryable wrapper self._partial_retryable = partial( self._read_rows_retryable_attempt, client.read_rows, - per_request_timeout, - deadline, + attempt_timeout_gen, row_limit, ) predicate = retries.if_exception_type( @@ -145,8 +147,7 @@ async def aclose(self): async def _read_rows_retryable_attempt( self, gapic_fn: Callable[..., Awaitable[AsyncIterable[ReadRowsResponse]]], - per_request_timeout: float, - operation_deadline: float, + timeout_generator: Iterator[float], total_row_limit: int, ) -> AsyncGenerator[Row, None]: """ @@ -183,16 +184,14 @@ async def _read_rows_retryable_attempt( raise RuntimeError("unexpected state: emit count exceeds row limit") else: self._request["rows_limit"] = new_limit - params_str = f'table_name={self._request.get("table_name", "")}' - app_profile_id = self._request.get("app_profile_id", None) - if app_profile_id: - params_str = f"{params_str},app_profile_id={app_profile_id}" - time_to_deadline = operation_deadline - time.monotonic() - gapic_timeout = max(0, min(time_to_deadline, per_request_timeout)) + metadata = _make_metadata( + self._request.get("table_name", None), + self._request.get("app_profile_id", None), + ) new_gapic_stream: RpcContext = await gapic_fn( self._request, - timeout=gapic_timeout, - metadata=[("x-goog-request-params", params_str)], + timeout=next(timeout_generator), + metadata=metadata, ) try: state_machine = _StateMachine() diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index e8c7b5e04..3921d6640 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -20,6 +20,8 @@ Any, Optional, Set, + Callable, + Coroutine, TYPE_CHECKING, ) @@ -38,6 +40,8 @@ ) from google.cloud.client import ClientWithProject from google.api_core.exceptions import GoogleAPICallError +from google.api_core import retry_async as retries +from google.api_core import exceptions as core_exceptions from google.cloud.bigtable._read_rows import _ReadRowsOperation import google.auth.credentials @@ -46,9 +50,12 @@ from google.cloud.bigtable.row import Row from google.cloud.bigtable.read_rows_query import ReadRowsQuery from google.cloud.bigtable.iterators import ReadRowsIterator +from google.cloud.bigtable.mutations import Mutation, RowMutationEntry +from google.cloud.bigtable._mutate_rows import _MutateRowsOperation +from google.cloud.bigtable._helpers import _make_metadata +from google.cloud.bigtable._helpers import _convert_retry_deadline if TYPE_CHECKING: - from google.cloud.bigtable.mutations import Mutation, BulkMutationsEntry from google.cloud.bigtable.mutations_batcher import MutationsBatcher from google.cloud.bigtable import RowKeySamples from google.cloud.bigtable.row_filters import RowFilter @@ -586,8 +593,8 @@ async def mutate_row( row_key: str | bytes, mutations: list[Mutation] | Mutation, *, - operation_timeout: int | float | None = 60, - per_request_timeout: int | float | None = None, + operation_timeout: float | None = 60, + per_request_timeout: float | None = None, ): """ Mutates a row atomically. @@ -617,19 +624,75 @@ async def mutate_row( - GoogleAPIError: raised on non-idempotent operations that cannot be safely retried. """ - raise NotImplementedError + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = per_request_timeout or self.default_per_request_timeout + + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout <= 0: + raise ValueError("per_request_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout > operation_timeout: + raise ValueError("per_request_timeout must be less than operation_timeout") + + if isinstance(row_key, str): + row_key = row_key.encode("utf-8") + request = {"table_name": self.table_name, "row_key": row_key} + if self.app_profile_id: + request["app_profile_id"] = self.app_profile_id + + if isinstance(mutations, Mutation): + mutations = [mutations] + request["mutations"] = [mutation._to_dict() for mutation in mutations] + + if all(mutation.is_idempotent() for mutation in mutations): + # mutations are all idempotent and safe to retry + predicate = retries.if_exception_type( + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ) + else: + # mutations should not be retried + predicate = retries.if_exception_type() + + transient_errors = [] + + def on_error_fn(exc): + if predicate(exc): + transient_errors.append(exc) + + retry = retries.AsyncRetry( + predicate=predicate, + on_error=on_error_fn, + timeout=operation_timeout, + initial=0.01, + multiplier=2, + maximum=60, + ) + # wrap rpc in retry logic + retry_wrapped = retry(self.client._gapic_client.mutate_row) + # convert RetryErrors from retry wrapper into DeadlineExceeded errors + deadline_wrapped = _convert_retry_deadline( + retry_wrapped, operation_timeout, transient_errors + ) + metadata = _make_metadata(self.table_name, self.app_profile_id) + # trigger rpc + await deadline_wrapped(request, timeout=per_request_timeout, metadata=metadata) async def bulk_mutate_rows( self, - mutation_entries: list[BulkMutationsEntry], + mutation_entries: list[RowMutationEntry], *, - operation_timeout: int | float | None = 60, - per_request_timeout: int | float | None = None, + operation_timeout: float | None = 60, + per_request_timeout: float | None = None, + on_success: Callable[ + [int, RowMutationEntry], None | Coroutine[None, None, None] + ] + | None = None, ): """ Applies mutations for multiple rows in a single batched request. - Each individual BulkMutationsEntry is applied atomically, but separate entries + Each individual RowMutationEntry is applied atomically, but separate entries may be applied in arbitrary order (even for entries targetting the same row) In total, the row_mutations can contain at most 100000 individual mutations across all entries @@ -650,12 +713,31 @@ async def bulk_mutate_rows( in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted if within operation_timeout budget - + - on_success: a callback function that will be called when each mutation + entry is confirmed to be applied successfully. Will be passed the + index and the entry itself. Raises: - MutationsExceptionGroup if one or more mutations fails Contains details about any failed entries in .exceptions """ - raise NotImplementedError + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = per_request_timeout or self.default_per_request_timeout + + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout <= 0: + raise ValueError("per_request_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout > operation_timeout: + raise ValueError("per_request_timeout must be less than operation_timeout") + + operation = _MutateRowsOperation( + self.client._gapic_client, + self, + mutation_entries, + operation_timeout, + per_request_timeout, + ) + await operation.start() async def check_and_mutate_row( self, diff --git a/google/cloud/bigtable/exceptions.py b/google/cloud/bigtable/exceptions.py index 975feb101..fe3bec7e9 100644 --- a/google/cloud/bigtable/exceptions.py +++ b/google/cloud/bigtable/exceptions.py @@ -15,63 +15,15 @@ from __future__ import annotations import sys -from inspect import iscoroutinefunction -from typing import Callable, Any +from typing import TYPE_CHECKING from google.api_core import exceptions as core_exceptions is_311_plus = sys.version_info >= (3, 11) - -def _convert_retry_deadline( - func: Callable[..., Any], - timeout_value: float | None = None, - retry_errors: list[Exception] | None = None, -): - """ - Decorator to convert RetryErrors raised by api_core.retry into - DeadlineExceeded exceptions, indicating that the underlying retries have - exhaused the timeout value. - Optionally attaches a RetryExceptionGroup to the DeadlineExceeded.__cause__, - detailing the failed exceptions associated with each retry. - - Supports both sync and async function wrapping. - - Args: - - func: The function to decorate - - timeout_value: The timeout value to display in the DeadlineExceeded error message - - retry_errors: An optional list of exceptions to attach as a RetryExceptionGroup to the DeadlineExceeded.__cause__ - """ - timeout_str = f" of {timeout_value:.1f}s" if timeout_value is not None else "" - error_str = f"operation_timeout{timeout_str} exceeded" - - def handle_error(): - new_exc = core_exceptions.DeadlineExceeded( - error_str, - ) - source_exc = None - if retry_errors: - source_exc = RetryExceptionGroup( - f"{len(retry_errors)} failed attempts", retry_errors - ) - new_exc.__cause__ = source_exc - raise new_exc from source_exc - - # separate wrappers for async and sync functions - async def wrapper_async(*args, **kwargs): - try: - return await func(*args, **kwargs) - except core_exceptions.RetryError: - handle_error() - - def wrapper(*args, **kwargs): - try: - return func(*args, **kwargs) - except core_exceptions.RetryError: - handle_error() - - return wrapper_async if iscoroutinefunction(func) else wrapper +if TYPE_CHECKING: + from google.cloud.bigtable.mutations import RowMutationEntry class IdleTimeout(core_exceptions.DeadlineExceeded): @@ -109,9 +61,23 @@ def __init__(self, message, excs): if is_311_plus: super().__init__(message, excs) else: - self.exceptions = excs - revised_message = f"{message} ({len(excs)} sub-exceptions)" - super().__init__(revised_message) + if len(excs) == 0: + raise ValueError("exceptions must be a non-empty sequence") + self.exceptions = tuple(excs) + super().__init__(message) + + def __new__(cls, message, excs): + if is_311_plus: + return super().__new__(cls, message, excs) + else: + return super().__new__(cls) + + def __str__(self): + """ + String representation doesn't display sub-exceptions. Subexceptions are + described in message + """ + return self.args[0] class MutationsExceptionGroup(BigtableExceptionGroup): @@ -119,10 +85,55 @@ class MutationsExceptionGroup(BigtableExceptionGroup): Represents one or more exceptions that occur during a bulk mutation operation """ - pass + @staticmethod + def _format_message(excs: list[FailedMutationEntryError], total_entries: int): + entry_str = "entry" if total_entries == 1 else "entries" + plural_str = "" if len(excs) == 1 else "s" + return f"{len(excs)} sub-exception{plural_str} (from {total_entries} {entry_str} attempted)" + + def __init__(self, excs: list[FailedMutationEntryError], total_entries: int): + super().__init__(self._format_message(excs, total_entries), excs) + + def __new__(cls, excs: list[FailedMutationEntryError], total_entries: int): + return super().__new__(cls, cls._format_message(excs, total_entries), excs) + + +class FailedMutationEntryError(Exception): + """ + Represents a single failed RowMutationEntry in a bulk_mutate_rows request. + A collection of FailedMutationEntryErrors will be raised in a MutationsExceptionGroup + """ + + def __init__( + self, + failed_idx: int, + failed_mutation_entry: "RowMutationEntry", + cause: Exception, + ): + idempotent_msg = ( + "idempotent" if failed_mutation_entry.is_idempotent() else "non-idempotent" + ) + message = f"Failed {idempotent_msg} mutation entry at index {failed_idx} with cause: {cause!r}" + super().__init__(message) + self.index = failed_idx + self.entry = failed_mutation_entry + self.__cause__ = cause class RetryExceptionGroup(BigtableExceptionGroup): """Represents one or more exceptions that occur during a retryable operation""" - pass + @staticmethod + def _format_message(excs: list[Exception]): + if len(excs) == 0: + return "No exceptions" + if len(excs) == 1: + return f"1 failed attempt: {type(excs[0]).__name__}" + else: + return f"{len(excs)} failed attempts. Latest: {type(excs[-1]).__name__}" + + def __init__(self, excs: list[Exception]): + super().__init__(self._format_message(excs), excs) + + def __new__(cls, excs: list[Exception]): + return super().__new__(cls, cls._format_message(excs), excs) diff --git a/google/cloud/bigtable/iterators.py b/google/cloud/bigtable/iterators.py index 169bbc3f3..b20932fb2 100644 --- a/google/cloud/bigtable/iterators.py +++ b/google/cloud/bigtable/iterators.py @@ -22,7 +22,7 @@ from google.cloud.bigtable._read_rows import _ReadRowsOperation from google.cloud.bigtable.exceptions import IdleTimeout -from google.cloud.bigtable.exceptions import _convert_retry_deadline +from google.cloud.bigtable._helpers import _convert_retry_deadline from google.cloud.bigtable.row import Row diff --git a/google/cloud/bigtable/mutations.py b/google/cloud/bigtable/mutations.py index 3bb5b2ed6..c72f132c8 100644 --- a/google/cloud/bigtable/mutations.py +++ b/google/cloud/bigtable/mutations.py @@ -13,41 +13,200 @@ # limitations under the License. # from __future__ import annotations - +from typing import Any +import time from dataclasses import dataclass +from abc import ABC, abstractmethod +# special value for SetCell mutation timestamps. If set, server will assign a timestamp +SERVER_SIDE_TIMESTAMP = -1 -class Mutation: - pass + +class Mutation(ABC): + """Model class for mutations""" + + @abstractmethod + def _to_dict(self) -> dict[str, Any]: + raise NotImplementedError + + def is_idempotent(self) -> bool: + """ + Check if the mutation is idempotent + If false, the mutation will not be retried + """ + return True + + def __str__(self) -> str: + return str(self._to_dict()) + + @classmethod + def _from_dict(cls, input_dict: dict[str, Any]) -> Mutation: + instance: Mutation | None = None + try: + if "set_cell" in input_dict: + details = input_dict["set_cell"] + instance = SetCell( + details["family_name"], + details["column_qualifier"], + details["value"], + details["timestamp_micros"], + ) + elif "delete_from_column" in input_dict: + details = input_dict["delete_from_column"] + time_range = details.get("time_range", {}) + start = time_range.get("start_timestamp_micros", None) + end = time_range.get("end_timestamp_micros", None) + instance = DeleteRangeFromColumn( + details["family_name"], details["column_qualifier"], start, end + ) + elif "delete_from_family" in input_dict: + details = input_dict["delete_from_family"] + instance = DeleteAllFromFamily(details["family_name"]) + elif "delete_from_row" in input_dict: + instance = DeleteAllFromRow() + except KeyError as e: + raise ValueError("Invalid mutation dictionary") from e + if instance is None: + raise ValueError("No valid mutation found") + if not issubclass(instance.__class__, cls): + raise ValueError("Mutation type mismatch") + return instance -@dataclass class SetCell(Mutation): - family: str - column_qualifier: bytes - new_value: bytes | str | int - timestamp_ms: int | None = None + def __init__( + self, + family: str, + qualifier: bytes | str, + new_value: bytes | str | int, + timestamp_micros: int | None = None, + ): + """ + Mutation to set the value of a cell + + Args: + - family: The name of the column family to which the new cell belongs. + - qualifier: The column qualifier of the new cell. + - new_value: The value of the new cell. str or int input will be converted to bytes + - timestamp_micros: The timestamp of the new cell. If None, the current timestamp will be used. + Timestamps will be sent with milisecond-percision. Extra precision will be truncated. + If -1, the server will assign a timestamp. Note that SetCell mutations with server-side + timestamps are non-idempotent operations and will not be retried. + """ + qualifier = qualifier.encode() if isinstance(qualifier, str) else qualifier + if not isinstance(qualifier, bytes): + raise TypeError("qualifier must be bytes or str") + if isinstance(new_value, str): + new_value = new_value.encode() + elif isinstance(new_value, int): + new_value = new_value.to_bytes(8, "big", signed=True) + if not isinstance(new_value, bytes): + raise TypeError("new_value must be bytes, str, or int") + if timestamp_micros is None: + # use current timestamp, with milisecond precision + timestamp_micros = time.time_ns() // 1000 + timestamp_micros = timestamp_micros - (timestamp_micros % 1000) + if timestamp_micros < SERVER_SIDE_TIMESTAMP: + raise ValueError( + "timestamp_micros must be positive (or -1 for server-side timestamp)" + ) + self.family = family + self.qualifier = qualifier + self.new_value = new_value + self.timestamp_micros = timestamp_micros + + def _to_dict(self) -> dict[str, Any]: + """Convert the mutation to a dictionary representation""" + return { + "set_cell": { + "family_name": self.family, + "column_qualifier": self.qualifier, + "timestamp_micros": self.timestamp_micros, + "value": self.new_value, + } + } + + def is_idempotent(self) -> bool: + """Check if the mutation is idempotent""" + return self.timestamp_micros != SERVER_SIDE_TIMESTAMP @dataclass class DeleteRangeFromColumn(Mutation): family: str - column_qualifier: bytes - start_timestamp_ms: int - end_timestamp_ms: int + qualifier: bytes + # None represents 0 + start_timestamp_micros: int | None = None + # None represents infinity + end_timestamp_micros: int | None = None + + def __post_init__(self): + if ( + self.start_timestamp_micros is not None + and self.end_timestamp_micros is not None + and self.start_timestamp_micros > self.end_timestamp_micros + ): + raise ValueError("start_timestamp_micros must be <= end_timestamp_micros") + + def _to_dict(self) -> dict[str, Any]: + timestamp_range = {} + if self.start_timestamp_micros is not None: + timestamp_range["start_timestamp_micros"] = self.start_timestamp_micros + if self.end_timestamp_micros is not None: + timestamp_range["end_timestamp_micros"] = self.end_timestamp_micros + return { + "delete_from_column": { + "family_name": self.family, + "column_qualifier": self.qualifier, + "time_range": timestamp_range, + } + } @dataclass class DeleteAllFromFamily(Mutation): family_to_delete: str + def _to_dict(self) -> dict[str, Any]: + return { + "delete_from_family": { + "family_name": self.family_to_delete, + } + } + @dataclass class DeleteAllFromRow(Mutation): - pass + def _to_dict(self) -> dict[str, Any]: + return { + "delete_from_row": {}, + } -@dataclass -class BulkMutationsEntry: - row: bytes - mutations: list[Mutation] | Mutation +class RowMutationEntry: + def __init__(self, row_key: bytes | str, mutations: Mutation | list[Mutation]): + if isinstance(row_key, str): + row_key = row_key.encode("utf-8") + if isinstance(mutations, Mutation): + mutations = [mutations] + self.row_key = row_key + self.mutations = tuple(mutations) + + def _to_dict(self) -> dict[str, Any]: + return { + "row_key": self.row_key, + "mutations": [mutation._to_dict() for mutation in self.mutations], + } + + def is_idempotent(self) -> bool: + """Check if the mutation is idempotent""" + return all(mutation.is_idempotent() for mutation in self.mutations) + + @classmethod + def _from_dict(cls, input_dict: dict[str, Any]) -> RowMutationEntry: + return RowMutationEntry( + row_key=input_dict["row_key"], + mutations=[ + Mutation._from_dict(mutation) for mutation in input_dict["mutations"] + ], + ) diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 48caceccd..7d015224c 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -187,6 +187,19 @@ async def delete_rows(self): await self.table.client._gapic_client.mutate_rows(request) +async def _retrieve_cell_value(table, row_key): + """ + Helper to read an individual row + """ + from google.cloud.bigtable import ReadRowsQuery + + row_list = await table.read_rows(ReadRowsQuery(row_keys=row_key)) + assert len(row_list) == 1 + row = row_list[0] + cell = row.cells[0] + return cell.value + + @pytest_asyncio.fixture(scope="function") async def temp_rows(table): builder = TempRowBuilder(table) @@ -205,6 +218,66 @@ async def test_ping_and_warm_gapic(client, table): await client._gapic_client.ping_and_warm(request) +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_mutation_set_cell(table, temp_rows): + """ + Ensure cells can be set properly + """ + from google.cloud.bigtable.mutations import SetCell + + row_key = b"mutate" + family = TEST_FAMILY + qualifier = b"test-qualifier" + start_value = b"start" + await temp_rows.add_row( + row_key, family=family, qualifier=qualifier, value=start_value + ) + + # ensure cell is initialized + assert (await _retrieve_cell_value(table, row_key)) == start_value + + expected_value = b"new-value" + mutation = SetCell( + family=TEST_FAMILY, qualifier=b"test-qualifier", new_value=expected_value + ) + + await table.mutate_row(row_key, mutation) + + # ensure cell is updated + assert (await _retrieve_cell_value(table, row_key)) == expected_value + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_bulk_mutations_set_cell(client, table, temp_rows): + """ + Ensure cells can be set properly + """ + from google.cloud.bigtable.mutations import SetCell, RowMutationEntry + + row_key = b"bulk_mutate" + family = TEST_FAMILY + qualifier = b"test-qualifier" + start_value = b"start" + await temp_rows.add_row( + row_key, family=family, qualifier=qualifier, value=start_value + ) + + # ensure cell is initialized + assert (await _retrieve_cell_value(table, row_key)) == start_value + + expected_value = b"new-value" + mutation = SetCell( + family=TEST_FAMILY, qualifier=b"test-qualifier", new_value=expected_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + await table.bulk_mutate_rows([bulk_mutation]) + + # ensure cell is updated + assert (await _retrieve_cell_value(table, row_key)) == expected_value + + @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_stream(table, temp_rows): diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py new file mode 100644 index 000000000..2765afe24 --- /dev/null +++ b/tests/unit/test__helpers.py @@ -0,0 +1,145 @@ +# +# Licensed 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. +# + +import pytest +import google.cloud.bigtable._helpers as _helpers +import google.cloud.bigtable.exceptions as bigtable_exceptions + +import mock + + +class TestMakeMetadata: + @pytest.mark.parametrize( + "table,profile,expected", + [ + ("table", "profile", "table_name=table,app_profile_id=profile"), + ("table", None, "table_name=table"), + ], + ) + def test__make_metadata(self, table, profile, expected): + metadata = _helpers._make_metadata(table, profile) + assert metadata == [("x-goog-request-params", expected)] + + +class TestAttemptTimeoutGenerator: + @pytest.mark.parametrize( + "request_t,operation_t,expected_list", + [ + (1, 3.5, [1, 1, 1, 0.5, 0, 0]), + (None, 3.5, [3.5, 2.5, 1.5, 0.5, 0, 0]), + (10, 5, [5, 4, 3, 2, 1, 0, 0]), + (3, 3, [3, 2, 1, 0, 0, 0, 0]), + (0, 3, [0, 0, 0]), + (3, 0, [0, 0, 0]), + (-1, 3, [0, 0, 0]), + (3, -1, [0, 0, 0]), + ], + ) + def test_attempt_timeout_generator(self, request_t, operation_t, expected_list): + """ + test different values for timeouts. Clock is incremented by 1 second for each item in expected_list + """ + timestamp_start = 123 + with mock.patch("time.monotonic") as mock_monotonic: + mock_monotonic.return_value = timestamp_start + generator = _helpers._attempt_timeout_generator(request_t, operation_t) + for val in expected_list: + mock_monotonic.return_value += 1 + assert next(generator) == val + + @pytest.mark.parametrize( + "request_t,operation_t,expected", + [ + (1, 3.5, 1), + (None, 3.5, 3.5), + (10, 5, 5), + (5, 10, 5), + (3, 3, 3), + (0, 3, 0), + (3, 0, 0), + (-1, 3, 0), + (3, -1, 0), + ], + ) + def test_attempt_timeout_frozen_time(self, request_t, operation_t, expected): + """test with time.monotonic frozen""" + timestamp_start = 123 + with mock.patch("time.monotonic") as mock_monotonic: + mock_monotonic.return_value = timestamp_start + generator = _helpers._attempt_timeout_generator(request_t, operation_t) + assert next(generator) == expected + # value should not change without time.monotonic changing + assert next(generator) == expected + + def test_attempt_timeout_w_sleeps(self): + """use real sleep values to make sure it matches expectations""" + from time import sleep + + operation_timeout = 1 + generator = _helpers._attempt_timeout_generator(None, operation_timeout) + expected_value = operation_timeout + sleep_time = 0.1 + for i in range(3): + found_value = next(generator) + assert abs(found_value - expected_value) < 0.001 + sleep(sleep_time) + expected_value -= sleep_time + + +class TestConvertRetryDeadline: + """ + Test _convert_retry_deadline wrapper + """ + + @pytest.mark.asyncio + async def test_no_error(self): + async def test_func(): + return 1 + + wrapped = _helpers._convert_retry_deadline(test_func, 0.1) + assert await wrapped() == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("timeout", [0.1, 2.0, 30.0]) + async def test_retry_error(self, timeout): + from google.api_core.exceptions import RetryError, DeadlineExceeded + + async def test_func(): + raise RetryError("retry error", None) + + wrapped = _helpers._convert_retry_deadline(test_func, timeout) + with pytest.raises(DeadlineExceeded) as e: + await wrapped() + assert e.value.__cause__ is None + assert f"operation_timeout of {timeout}s exceeded" in str(e.value) + + @pytest.mark.asyncio + async def test_with_retry_errors(self): + from google.api_core.exceptions import RetryError, DeadlineExceeded + + timeout = 10.0 + + async def test_func(): + raise RetryError("retry error", None) + + associated_errors = [RuntimeError("error1"), ZeroDivisionError("other")] + wrapped = _helpers._convert_retry_deadline( + test_func, timeout, associated_errors + ) + with pytest.raises(DeadlineExceeded) as e: + await wrapped() + cause = e.value.__cause__ + assert isinstance(cause, bigtable_exceptions.RetryExceptionGroup) + assert cause.exceptions == tuple(associated_errors) + assert f"operation_timeout of {timeout}s exceeded" in str(e.value) diff --git a/tests/unit/test__mutate_rows.py b/tests/unit/test__mutate_rows.py new file mode 100644 index 000000000..4fba16f23 --- /dev/null +++ b/tests/unit/test__mutate_rows.py @@ -0,0 +1,290 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +import pytest + +from google.cloud.bigtable_v2.types import MutateRowsResponse +from google.rpc import status_pb2 +import google.api_core.exceptions as core_exceptions + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock + from unittest.mock import AsyncMock # type: ignore +except ImportError: # pragma: NO COVER + import mock # type: ignore + from mock import AsyncMock # type: ignore + + +class TestMutateRowsOperation: + def _target_class(self): + from google.cloud.bigtable._mutate_rows import _MutateRowsOperation + + return _MutateRowsOperation + + def _make_one(self, *args, **kwargs): + if not args: + kwargs["gapic_client"] = kwargs.pop("gapic_client", mock.Mock()) + kwargs["table"] = kwargs.pop("table", AsyncMock()) + kwargs["mutation_entries"] = kwargs.pop("mutation_entries", []) + kwargs["operation_timeout"] = kwargs.pop("operation_timeout", 5) + kwargs["per_request_timeout"] = kwargs.pop("per_request_timeout", 0.1) + return self._target_class()(*args, **kwargs) + + async def _mock_stream(self, mutation_list, error_dict): + for idx, entry in enumerate(mutation_list): + code = error_dict.get(idx, 0) + yield MutateRowsResponse( + entries=[ + MutateRowsResponse.Entry( + index=idx, status=status_pb2.Status(code=code) + ) + ] + ) + + def _make_mock_gapic(self, mutation_list, error_dict=None): + mock_fn = AsyncMock() + if error_dict is None: + error_dict = {} + mock_fn.side_effect = lambda *args, **kwargs: self._mock_stream( + mutation_list, error_dict + ) + return mock_fn + + def test_ctor(self): + """ + test that constructor sets all the attributes correctly + """ + from google.cloud.bigtable._mutate_rows import _MutateRowsIncomplete + from google.api_core.exceptions import DeadlineExceeded + from google.api_core.exceptions import ServiceUnavailable + + client = mock.Mock() + table = mock.Mock() + entries = [mock.Mock(), mock.Mock()] + operation_timeout = 0.05 + attempt_timeout = 0.01 + instance = self._make_one( + client, table, entries, operation_timeout, attempt_timeout + ) + # running gapic_fn should trigger a client call + assert client.mutate_rows.call_count == 0 + instance._gapic_fn() + assert client.mutate_rows.call_count == 1 + # gapic_fn should call with table details + inner_kwargs = client.mutate_rows.call_args[1] + assert len(inner_kwargs) == 3 + assert inner_kwargs["table_name"] == table.table_name + assert inner_kwargs["app_profile_id"] == table.app_profile_id + metadata = inner_kwargs["metadata"] + assert len(metadata) == 1 + assert metadata[0][0] == "x-goog-request-params" + assert str(table.table_name) in metadata[0][1] + assert str(table.app_profile_id) in metadata[0][1] + # entries should be passed down + assert instance.mutations == entries + # timeout_gen should generate per-attempt timeout + assert next(instance.timeout_generator) == attempt_timeout + # ensure predicate is set + assert instance.is_retryable is not None + assert instance.is_retryable(DeadlineExceeded("")) is True + assert instance.is_retryable(ServiceUnavailable("")) is True + assert instance.is_retryable(_MutateRowsIncomplete("")) is True + assert instance.is_retryable(RuntimeError("")) is False + assert instance.remaining_indices == list(range(len(entries))) + assert instance.errors == {} + + @pytest.mark.asyncio + async def test_mutate_rows_operation(self): + """ + Test successful case of mutate_rows_operation + """ + client = mock.Mock() + table = mock.Mock() + entries = [mock.Mock(), mock.Mock()] + operation_timeout = 0.05 + instance = self._make_one( + client, table, entries, operation_timeout, operation_timeout + ) + with mock.patch.object(instance, "_operation", AsyncMock()) as attempt_mock: + attempt_mock.return_value = None + await instance.start() + assert attempt_mock.call_count == 1 + + @pytest.mark.parametrize( + "exc_type", [RuntimeError, ZeroDivisionError, core_exceptions.Forbidden] + ) + @pytest.mark.asyncio + async def test_mutate_rows_exception(self, exc_type): + """ + exceptions raised from retryable should be raised in MutationsExceptionGroup + """ + from google.cloud.bigtable.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.exceptions import FailedMutationEntryError + + client = mock.Mock() + table = mock.Mock() + entries = [mock.Mock()] + operation_timeout = 0.05 + expected_cause = exc_type("abort") + with mock.patch.object( + self._target_class(), + "_run_attempt", + AsyncMock(), + ) as attempt_mock: + attempt_mock.side_effect = expected_cause + found_exc = None + try: + instance = self._make_one( + client, table, entries, operation_timeout, operation_timeout + ) + await instance.start() + except MutationsExceptionGroup as e: + found_exc = e + assert attempt_mock.call_count == 1 + assert len(found_exc.exceptions) == 1 + assert isinstance(found_exc.exceptions[0], FailedMutationEntryError) + assert found_exc.exceptions[0].__cause__ == expected_cause + + @pytest.mark.parametrize( + "exc_type", + [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable], + ) + @pytest.mark.asyncio + async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): + """ + If an exception fails but eventually passes, it should not raise an exception + """ + from google.cloud.bigtable._mutate_rows import _MutateRowsOperation + + client = mock.Mock() + table = mock.Mock() + entries = [mock.Mock()] + operation_timeout = 1 + expected_cause = exc_type("retry") + num_retries = 2 + with mock.patch.object( + _MutateRowsOperation, + "_run_attempt", + AsyncMock(), + ) as attempt_mock: + attempt_mock.side_effect = [expected_cause] * num_retries + [None] + instance = self._make_one( + client, table, entries, operation_timeout, operation_timeout + ) + await instance.start() + assert attempt_mock.call_count == num_retries + 1 + + @pytest.mark.asyncio + async def test_mutate_rows_incomplete_ignored(self): + """ + MutateRowsIncomplete exceptions should not be added to error list + """ + from google.cloud.bigtable._mutate_rows import _MutateRowsIncomplete + from google.cloud.bigtable.exceptions import MutationsExceptionGroup + from google.api_core.exceptions import DeadlineExceeded + + client = mock.Mock() + table = mock.Mock() + entries = [mock.Mock()] + operation_timeout = 0.05 + with mock.patch.object( + self._target_class(), + "_run_attempt", + AsyncMock(), + ) as attempt_mock: + attempt_mock.side_effect = _MutateRowsIncomplete("ignored") + found_exc = None + try: + instance = self._make_one( + client, table, entries, operation_timeout, operation_timeout + ) + await instance.start() + except MutationsExceptionGroup as e: + found_exc = e + assert attempt_mock.call_count > 0 + assert len(found_exc.exceptions) == 1 + assert isinstance(found_exc.exceptions[0].__cause__, DeadlineExceeded) + + @pytest.mark.asyncio + async def test_run_attempt_single_entry_success(self): + """Test mutating a single entry""" + mutation = mock.Mock() + mutations = {0: mutation} + expected_timeout = 1.3 + mock_gapic_fn = self._make_mock_gapic(mutations) + instance = self._make_one( + mutation_entries=mutations, + per_request_timeout=expected_timeout, + ) + with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): + await instance._run_attempt() + assert len(instance.remaining_indices) == 0 + assert mock_gapic_fn.call_count == 1 + _, kwargs = mock_gapic_fn.call_args + assert kwargs["timeout"] == expected_timeout + assert kwargs["entries"] == [mutation._to_dict()] + + @pytest.mark.asyncio + async def test_run_attempt_empty_request(self): + """Calling with no mutations should result in no API calls""" + mock_gapic_fn = self._make_mock_gapic([]) + instance = self._make_one( + mutation_entries=[], + ) + await instance._run_attempt() + assert mock_gapic_fn.call_count == 0 + + @pytest.mark.asyncio + async def test_run_attempt_partial_success_retryable(self): + """Some entries succeed, but one fails. Should report the proper index, and raise incomplete exception""" + from google.cloud.bigtable._mutate_rows import _MutateRowsIncomplete + + success_mutation = mock.Mock() + success_mutation_2 = mock.Mock() + failure_mutation = mock.Mock() + mutations = [success_mutation, failure_mutation, success_mutation_2] + mock_gapic_fn = self._make_mock_gapic(mutations, error_dict={1: 300}) + instance = self._make_one( + mutation_entries=mutations, + ) + instance.is_retryable = lambda x: True + with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): + with pytest.raises(_MutateRowsIncomplete): + await instance._run_attempt() + assert instance.remaining_indices == [1] + assert 0 not in instance.errors + assert len(instance.errors[1]) == 1 + assert instance.errors[1][0].grpc_status_code == 300 + assert 2 not in instance.errors + + @pytest.mark.asyncio + async def test_run_attempt_partial_success_non_retryable(self): + """Some entries succeed, but one fails. Exception marked as non-retryable. Do not raise incomplete error""" + success_mutation = mock.Mock() + success_mutation_2 = mock.Mock() + failure_mutation = mock.Mock() + mutations = [success_mutation, failure_mutation, success_mutation_2] + mock_gapic_fn = self._make_mock_gapic(mutations, error_dict={1: 300}) + instance = self._make_one( + mutation_entries=mutations, + ) + instance.is_retryable = lambda x: False + with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): + await instance._run_attempt() + assert instance.remaining_indices == [] + assert 0 not in instance.errors + assert len(instance.errors[1]) == 1 + assert instance.errors[1][0].grpc_status_code == 300 + assert 2 not in instance.errors diff --git a/tests/unit/test__read_rows.py b/tests/unit/test__read_rows.py index e57b5d992..c893c56cd 100644 --- a/tests/unit/test__read_rows.py +++ b/tests/unit/test__read_rows.py @@ -41,10 +41,14 @@ def test_ctor_defaults(self): client = mock.Mock() client.read_rows = mock.Mock() client.read_rows.return_value = None - start_time = 123 default_operation_timeout = 600 - with mock.patch("time.monotonic", return_value=start_time): + time_gen_mock = mock.Mock() + with mock.patch( + "google.cloud.bigtable._read_rows._attempt_timeout_generator", time_gen_mock + ): instance = self._make_one(request, client) + assert time_gen_mock.call_count == 1 + time_gen_mock.assert_called_once_with(None, default_operation_timeout) assert instance.transient_errors == [] assert instance._last_emitted_row_key is None assert instance._emit_count == 0 @@ -52,9 +56,8 @@ def test_ctor_defaults(self): retryable_fn = instance._partial_retryable assert retryable_fn.func == instance._read_rows_retryable_attempt assert retryable_fn.args[0] == client.read_rows - assert retryable_fn.args[1] == default_operation_timeout - assert retryable_fn.args[2] == default_operation_timeout + start_time - assert retryable_fn.args[3] == 0 + assert retryable_fn.args[1] == time_gen_mock.return_value + assert retryable_fn.args[2] == 0 assert client.read_rows.call_count == 0 def test_ctor(self): @@ -65,14 +68,20 @@ def test_ctor(self): client.read_rows.return_value = None expected_operation_timeout = 42 expected_request_timeout = 44 - start_time = 123 - with mock.patch("time.monotonic", return_value=start_time): + time_gen_mock = mock.Mock() + with mock.patch( + "google.cloud.bigtable._read_rows._attempt_timeout_generator", time_gen_mock + ): instance = self._make_one( request, client, operation_timeout=expected_operation_timeout, per_request_timeout=expected_request_timeout, ) + assert time_gen_mock.call_count == 1 + time_gen_mock.assert_called_once_with( + expected_request_timeout, expected_operation_timeout + ) assert instance.transient_errors == [] assert instance._last_emitted_row_key is None assert instance._emit_count == 0 @@ -80,9 +89,8 @@ def test_ctor(self): retryable_fn = instance._partial_retryable assert retryable_fn.func == instance._read_rows_retryable_attempt assert retryable_fn.args[0] == client.read_rows - assert retryable_fn.args[1] == expected_request_timeout - assert retryable_fn.args[2] == start_time + expected_operation_timeout - assert retryable_fn.args[3] == row_limit + assert retryable_fn.args[1] == time_gen_mock.return_value + assert retryable_fn.args[2] == row_limit assert client.read_rows.call_count == 0 def test___aiter__(self): @@ -217,14 +225,18 @@ async def test_revise_limit(self, start_limit, emit_num, expected_limit): - if the number emitted exceeds the new limit, an exception should should be raised (tested in test_revise_limit_over_limit) """ + import itertools + request = {"rows_limit": start_limit} instance = self._make_one(request, mock.Mock()) instance._emit_count = emit_num instance._last_emitted_row_key = "a" gapic_mock = mock.Mock() gapic_mock.side_effect = [GeneratorExit("stop_fn")] + mock_timeout_gen = itertools.repeat(5) + attempt = instance._read_rows_retryable_attempt( - gapic_mock, 100, 100, start_limit + gapic_mock, mock_timeout_gen, start_limit ) if start_limit != 0 and expected_limit == 0: # if we emitted the expected number of rows, we should receive a StopAsyncIteration @@ -242,12 +254,15 @@ async def test_revise_limit_over_limit(self, start_limit, emit_num): Should raise runtime error if we get in state where emit_num > start_num (unless start_num == 0, which represents unlimited) """ + import itertools + request = {"rows_limit": start_limit} instance = self._make_one(request, mock.Mock()) instance._emit_count = emit_num instance._last_emitted_row_key = "a" + mock_timeout_gen = itertools.repeat(5) attempt = instance._read_rows_retryable_attempt( - mock.Mock(), 100, 100, start_limit + mock.Mock(), mock_timeout_gen, start_limit ) with pytest.raises(RuntimeError) as e: await attempt.__anext__() @@ -273,6 +288,7 @@ async def test_retryable_attempt_hit_limit(self, limit): Stream should end after hitting the limit """ from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + import itertools instance = self._make_one({}, mock.Mock()) @@ -290,7 +306,8 @@ async def gen(): return gen() - gen = instance._read_rows_retryable_attempt(mock_gapic, 100, 100, limit) + mock_timeout_gen = itertools.repeat(5) + gen = instance._read_rows_retryable_attempt(mock_gapic, mock_timeout_gen, limit) # should yield values up to the limit for i in range(limit): await gen.__anext__() diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 8a3402a65..be3703a23 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -20,6 +20,7 @@ import pytest +from google.cloud.bigtable import mutations from google.auth.credentials import AnonymousCredentials from google.cloud.bigtable_v2.types import ReadRowsResponse from google.cloud.bigtable.read_rows_query import ReadRowsQuery @@ -1294,3 +1295,579 @@ async def test_read_rows_default_timeout_override(self): kwargs = mock_op.call_args_list[0].kwargs assert kwargs["operation_timeout"] == operation_timeout assert kwargs["per_request_timeout"] == per_request_timeout + + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_read_rows_metadata(self, include_app_profile): + """request should attach metadata headers""" + profile = "profile" if include_app_profile else None + async with self._make_client() as client: + async with client.get_table("i", "t", app_profile_id=profile) as table: + with mock.patch.object( + client._gapic_client, "read_rows", AsyncMock() + ) as read_rows: + await table.read_rows(ReadRowsQuery()) + kwargs = read_rows.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata + + +class TestMutateRow: + def _make_client(self, *args, **kwargs): + from google.cloud.bigtable.client import BigtableDataClient + + return BigtableDataClient(*args, **kwargs) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "mutation_arg", + [ + mutations.SetCell("family", b"qualifier", b"value"), + mutations.SetCell( + "family", b"qualifier", b"value", timestamp_micros=1234567890 + ), + mutations.DeleteRangeFromColumn("family", b"qualifier"), + mutations.DeleteAllFromFamily("family"), + mutations.DeleteAllFromRow(), + [mutations.SetCell("family", b"qualifier", b"value")], + [ + mutations.DeleteRangeFromColumn("family", b"qualifier"), + mutations.DeleteAllFromRow(), + ], + ], + ) + async def test_mutate_row(self, mutation_arg): + """Test mutations with no errors""" + expected_per_request_timeout = 19 + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_row" + ) as mock_gapic: + mock_gapic.return_value = None + await table.mutate_row( + "row_key", + mutation_arg, + per_request_timeout=expected_per_request_timeout, + ) + assert mock_gapic.call_count == 1 + request = mock_gapic.call_args[0][0] + assert ( + request["table_name"] + == "projects/project/instances/instance/tables/table" + ) + assert request["row_key"] == b"row_key" + formatted_mutations = ( + [mutation._to_dict() for mutation in mutation_arg] + if isinstance(mutation_arg, list) + else [mutation_arg._to_dict()] + ) + assert request["mutations"] == formatted_mutations + found_per_request_timeout = mock_gapic.call_args[1]["timeout"] + assert found_per_request_timeout == expected_per_request_timeout + + @pytest.mark.parametrize( + "retryable_exception", + [ + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ], + ) + @pytest.mark.asyncio + async def test_mutate_row_retryable_errors(self, retryable_exception): + from google.api_core.exceptions import DeadlineExceeded + from google.cloud.bigtable.exceptions import RetryExceptionGroup + + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_row" + ) as mock_gapic: + mock_gapic.side_effect = retryable_exception("mock") + with pytest.raises(DeadlineExceeded) as e: + mutation = mutations.DeleteAllFromRow() + assert mutation.is_idempotent() is True + await table.mutate_row( + "row_key", mutation, operation_timeout=0.05 + ) + cause = e.value.__cause__ + assert isinstance(cause, RetryExceptionGroup) + assert isinstance(cause.exceptions[0], retryable_exception) + + @pytest.mark.parametrize( + "retryable_exception", + [ + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ], + ) + @pytest.mark.asyncio + async def test_mutate_row_non_idempotent_retryable_errors( + self, retryable_exception + ): + """ + Non-idempotent mutations should not be retried + """ + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_row" + ) as mock_gapic: + mock_gapic.side_effect = retryable_exception("mock") + with pytest.raises(retryable_exception): + mutation = mutations.SetCell( + "family", b"qualifier", b"value", -1 + ) + assert mutation.is_idempotent() is False + await table.mutate_row( + "row_key", mutation, operation_timeout=0.2 + ) + + @pytest.mark.parametrize( + "non_retryable_exception", + [ + core_exceptions.OutOfRange, + core_exceptions.NotFound, + core_exceptions.FailedPrecondition, + RuntimeError, + ValueError, + core_exceptions.Aborted, + ], + ) + @pytest.mark.asyncio + async def test_mutate_row_non_retryable_errors(self, non_retryable_exception): + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_row" + ) as mock_gapic: + mock_gapic.side_effect = non_retryable_exception("mock") + with pytest.raises(non_retryable_exception): + mutation = mutations.SetCell( + "family", + b"qualifier", + b"value", + timestamp_micros=1234567890, + ) + assert mutation.is_idempotent() is True + await table.mutate_row( + "row_key", mutation, operation_timeout=0.2 + ) + + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_mutate_row_metadata(self, include_app_profile): + """request should attach metadata headers""" + profile = "profile" if include_app_profile else None + async with self._make_client() as client: + async with client.get_table("i", "t", app_profile_id=profile) as table: + with mock.patch.object( + client._gapic_client, "mutate_row", AsyncMock() + ) as read_rows: + await table.mutate_row("rk", {}) + kwargs = read_rows.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata + + +class TestBulkMutateRows: + def _make_client(self, *args, **kwargs): + from google.cloud.bigtable.client import BigtableDataClient + + return BigtableDataClient(*args, **kwargs) + + async def _mock_response(self, response_list): + from google.cloud.bigtable_v2.types import MutateRowsResponse + from google.rpc import status_pb2 + + statuses = [] + for response in response_list: + if isinstance(response, core_exceptions.GoogleAPICallError): + statuses.append( + status_pb2.Status( + message=str(response), code=response.grpc_status_code.value[0] + ) + ) + else: + statuses.append(status_pb2.Status(code=0)) + entries = [ + MutateRowsResponse.Entry(index=i, status=statuses[i]) + for i in range(len(response_list)) + ] + + async def generator(): + yield MutateRowsResponse(entries=entries) + + return generator() + + @pytest.mark.asyncio + @pytest.mark.asyncio + @pytest.mark.parametrize( + "mutation_arg", + [ + [mutations.SetCell("family", b"qualifier", b"value")], + [ + mutations.SetCell( + "family", b"qualifier", b"value", timestamp_micros=1234567890 + ) + ], + [mutations.DeleteRangeFromColumn("family", b"qualifier")], + [mutations.DeleteAllFromFamily("family")], + [mutations.DeleteAllFromRow()], + [mutations.SetCell("family", b"qualifier", b"value")], + [ + mutations.DeleteRangeFromColumn("family", b"qualifier"), + mutations.DeleteAllFromRow(), + ], + ], + ) + async def test_bulk_mutate_rows(self, mutation_arg): + """Test mutations with no errors""" + expected_per_request_timeout = 19 + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_rows" + ) as mock_gapic: + mock_gapic.return_value = self._mock_response([None]) + bulk_mutation = mutations.RowMutationEntry(b"row_key", mutation_arg) + await table.bulk_mutate_rows( + [bulk_mutation], + per_request_timeout=expected_per_request_timeout, + ) + assert mock_gapic.call_count == 1 + kwargs = mock_gapic.call_args[1] + assert ( + kwargs["table_name"] + == "projects/project/instances/instance/tables/table" + ) + assert kwargs["entries"] == [bulk_mutation._to_dict()] + assert kwargs["timeout"] == expected_per_request_timeout + + @pytest.mark.asyncio + async def test_bulk_mutate_rows_multiple_entries(self): + """Test mutations with no errors""" + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_rows" + ) as mock_gapic: + mock_gapic.return_value = self._mock_response([None, None]) + mutation_list = [mutations.DeleteAllFromRow()] + entry_1 = mutations.RowMutationEntry(b"row_key_1", mutation_list) + entry_2 = mutations.RowMutationEntry(b"row_key_2", mutation_list) + await table.bulk_mutate_rows( + [entry_1, entry_2], + ) + assert mock_gapic.call_count == 1 + kwargs = mock_gapic.call_args[1] + assert ( + kwargs["table_name"] + == "projects/project/instances/instance/tables/table" + ) + assert kwargs["entries"][0] == entry_1._to_dict() + assert kwargs["entries"][1] == entry_2._to_dict() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "exception", + [ + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ], + ) + async def test_bulk_mutate_rows_idempotent_mutation_error_retryable( + self, exception + ): + """ + Individual idempotent mutations should be retried if they fail with a retryable error + """ + from google.cloud.bigtable.exceptions import ( + RetryExceptionGroup, + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_rows" + ) as mock_gapic: + mock_gapic.side_effect = lambda *a, **k: self._mock_response( + [exception("mock")] + ) + with pytest.raises(MutationsExceptionGroup) as e: + mutation = mutations.DeleteAllFromRow() + entry = mutations.RowMutationEntry(b"row_key", [mutation]) + assert mutation.is_idempotent() is True + await table.bulk_mutate_rows([entry], operation_timeout=0.05) + assert len(e.value.exceptions) == 1 + failed_exception = e.value.exceptions[0] + assert "non-idempotent" not in str(failed_exception) + assert isinstance(failed_exception, FailedMutationEntryError) + cause = failed_exception.__cause__ + assert isinstance(cause, RetryExceptionGroup) + assert isinstance(cause.exceptions[0], exception) + # last exception should be due to retry timeout + assert isinstance( + cause.exceptions[-1], core_exceptions.DeadlineExceeded + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "exception", + [ + core_exceptions.OutOfRange, + core_exceptions.NotFound, + core_exceptions.FailedPrecondition, + core_exceptions.Aborted, + ], + ) + async def test_bulk_mutate_rows_idempotent_mutation_error_non_retryable( + self, exception + ): + """ + Individual idempotent mutations should not be retried if they fail with a non-retryable error + """ + from google.cloud.bigtable.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_rows" + ) as mock_gapic: + mock_gapic.side_effect = lambda *a, **k: self._mock_response( + [exception("mock")] + ) + with pytest.raises(MutationsExceptionGroup) as e: + mutation = mutations.DeleteAllFromRow() + entry = mutations.RowMutationEntry(b"row_key", [mutation]) + assert mutation.is_idempotent() is True + await table.bulk_mutate_rows([entry], operation_timeout=0.05) + assert len(e.value.exceptions) == 1 + failed_exception = e.value.exceptions[0] + assert "non-idempotent" not in str(failed_exception) + assert isinstance(failed_exception, FailedMutationEntryError) + cause = failed_exception.__cause__ + assert isinstance(cause, exception) + + @pytest.mark.parametrize( + "retryable_exception", + [ + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ], + ) + @pytest.mark.asyncio + async def test_bulk_mutate_idempotent_retryable_request_errors( + self, retryable_exception + ): + """ + Individual idempotent mutations should be retried if the request fails with a retryable error + """ + from google.cloud.bigtable.exceptions import ( + RetryExceptionGroup, + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_rows" + ) as mock_gapic: + mock_gapic.side_effect = retryable_exception("mock") + with pytest.raises(MutationsExceptionGroup) as e: + mutation = mutations.SetCell( + "family", b"qualifier", b"value", timestamp_micros=123 + ) + entry = mutations.RowMutationEntry(b"row_key", [mutation]) + assert mutation.is_idempotent() is True + await table.bulk_mutate_rows([entry], operation_timeout=0.05) + assert len(e.value.exceptions) == 1 + failed_exception = e.value.exceptions[0] + assert isinstance(failed_exception, FailedMutationEntryError) + assert "non-idempotent" not in str(failed_exception) + cause = failed_exception.__cause__ + assert isinstance(cause, RetryExceptionGroup) + assert isinstance(cause.exceptions[0], retryable_exception) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "retryable_exception", + [ + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ], + ) + async def test_bulk_mutate_rows_non_idempotent_retryable_errors( + self, retryable_exception + ): + """Non-Idempotent mutations should never be retried""" + from google.cloud.bigtable.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_rows" + ) as mock_gapic: + mock_gapic.side_effect = lambda *a, **k: self._mock_response( + [retryable_exception("mock")] + ) + with pytest.raises(MutationsExceptionGroup) as e: + mutation = mutations.SetCell( + "family", b"qualifier", b"value", -1 + ) + entry = mutations.RowMutationEntry(b"row_key", [mutation]) + assert mutation.is_idempotent() is False + await table.bulk_mutate_rows([entry], operation_timeout=0.2) + assert len(e.value.exceptions) == 1 + failed_exception = e.value.exceptions[0] + assert isinstance(failed_exception, FailedMutationEntryError) + assert "non-idempotent" in str(failed_exception) + cause = failed_exception.__cause__ + assert isinstance(cause, retryable_exception) + + @pytest.mark.parametrize( + "non_retryable_exception", + [ + core_exceptions.OutOfRange, + core_exceptions.NotFound, + core_exceptions.FailedPrecondition, + RuntimeError, + ValueError, + ], + ) + @pytest.mark.asyncio + async def test_bulk_mutate_rows_non_retryable_errors(self, non_retryable_exception): + """ + If the request fails with a non-retryable error, mutations should not be retried + """ + from google.cloud.bigtable.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_rows" + ) as mock_gapic: + mock_gapic.side_effect = non_retryable_exception("mock") + with pytest.raises(MutationsExceptionGroup) as e: + mutation = mutations.SetCell( + "family", b"qualifier", b"value", timestamp_micros=123 + ) + entry = mutations.RowMutationEntry(b"row_key", [mutation]) + assert mutation.is_idempotent() is True + await table.bulk_mutate_rows([entry], operation_timeout=0.2) + assert len(e.value.exceptions) == 1 + failed_exception = e.value.exceptions[0] + assert isinstance(failed_exception, FailedMutationEntryError) + assert "non-idempotent" not in str(failed_exception) + cause = failed_exception.__cause__ + assert isinstance(cause, non_retryable_exception) + + @pytest.mark.asyncio + async def test_bulk_mutate_error_index(self): + """ + Test partial failure, partial success. Errors should be associated with the correct index + """ + from google.api_core.exceptions import ( + DeadlineExceeded, + ServiceUnavailable, + FailedPrecondition, + ) + from google.cloud.bigtable.exceptions import ( + RetryExceptionGroup, + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + async with self._make_client(project="project") as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "mutate_rows" + ) as mock_gapic: + # fail with retryable errors, then a non-retryable one + mock_gapic.side_effect = [ + self._mock_response([None, ServiceUnavailable("mock"), None]), + self._mock_response([DeadlineExceeded("mock")]), + self._mock_response([FailedPrecondition("final")]), + ] + with pytest.raises(MutationsExceptionGroup) as e: + mutation = mutations.SetCell( + "family", b"qualifier", b"value", timestamp_micros=123 + ) + entries = [ + mutations.RowMutationEntry( + (f"row_key_{i}").encode(), [mutation] + ) + for i in range(3) + ] + assert mutation.is_idempotent() is True + await table.bulk_mutate_rows(entries, operation_timeout=1000) + assert len(e.value.exceptions) == 1 + failed = e.value.exceptions[0] + assert isinstance(failed, FailedMutationEntryError) + assert failed.index == 1 + assert failed.entry == entries[1] + cause = failed.__cause__ + assert isinstance(cause, RetryExceptionGroup) + assert len(cause.exceptions) == 3 + assert isinstance(cause.exceptions[0], ServiceUnavailable) + assert isinstance(cause.exceptions[1], DeadlineExceeded) + assert isinstance(cause.exceptions[2], FailedPrecondition) + + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_bulk_mutate_row_metadata(self, include_app_profile): + """request should attach metadata headers""" + profile = "profile" if include_app_profile else None + async with self._make_client() as client: + async with client.get_table("i", "t", app_profile_id=profile) as table: + with mock.patch.object( + client._gapic_client, "mutate_rows", AsyncMock() + ) as read_rows: + read_rows.side_effect = core_exceptions.Aborted("mock") + try: + await table.bulk_mutate_rows([mock.Mock()]) + except Exception: + # exception used to end early + pass + kwargs = read_rows.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata diff --git a/tests/unit/test_exceptions.py b/tests/unit/test_exceptions.py new file mode 100644 index 000000000..49b90a8a9 --- /dev/null +++ b/tests/unit/test_exceptions.py @@ -0,0 +1,239 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +import unittest +import pytest +import sys + +import google.cloud.bigtable.exceptions as bigtable_exceptions + + +class TestBigtableExceptionGroup: + """ + Subclass for MutationsExceptionGroup and RetryExceptionGroup + """ + + def _get_class(self): + from google.cloud.bigtable.exceptions import BigtableExceptionGroup + + return BigtableExceptionGroup + + def _make_one(self, message="test_message", excs=None): + if excs is None: + excs = [RuntimeError("mock")] + + return self._get_class()(message, excs=excs) + + def test_raise(self): + """ + Create exception in raise statement, which calls __new__ and __init__ + """ + test_msg = "test message" + test_excs = [Exception(test_msg)] + with pytest.raises(self._get_class()) as e: + raise self._get_class()(test_msg, test_excs) + assert str(e.value) == test_msg + assert list(e.value.exceptions) == test_excs + + def test_raise_empty_list(self): + """ + Empty exception lists are not supported + """ + with pytest.raises(ValueError) as e: + raise self._make_one(excs=[]) + assert "non-empty sequence" in str(e.value) + + @pytest.mark.skipif( + sys.version_info < (3, 11), reason="requires python3.11 or higher" + ) + def test_311_traceback(self): + """ + Exception customizations should not break rich exception group traceback in python 3.11 + """ + import traceback + + sub_exc1 = RuntimeError("first sub exception") + sub_exc2 = ZeroDivisionError("second sub exception") + exc_group = self._make_one(excs=[sub_exc1, sub_exc2]) + + expected_traceback = ( + f" | google.cloud.bigtable.exceptions.{type(exc_group).__name__}: {str(exc_group)}", + " +-+---------------- 1 ----------------", + " | RuntimeError: first sub exception", + " +---------------- 2 ----------------", + " | ZeroDivisionError: second sub exception", + " +------------------------------------", + ) + exception_caught = False + try: + raise exc_group + except self._get_class(): + exception_caught = True + tb = traceback.format_exc() + tb_relevant_lines = tuple(tb.splitlines()[3:]) + assert expected_traceback == tb_relevant_lines + assert exception_caught + + @pytest.mark.skipif( + sys.version_info < (3, 11), reason="requires python3.11 or higher" + ) + def test_311_exception_group(self): + """ + Python 3.11+ should handle exepctions as native exception groups + """ + exceptions = [RuntimeError("mock"), ValueError("mock")] + instance = self._make_one(excs=exceptions) + # ensure split works as expected + runtime_error, others = instance.split(lambda e: isinstance(e, RuntimeError)) + assert runtime_error.exceptions[0] == exceptions[0] + assert others.exceptions[0] == exceptions[1] + + def test_exception_handling(self): + """ + All versions should inherit from exception + and support tranditional exception handling + """ + instance = self._make_one() + assert isinstance(instance, Exception) + try: + raise instance + except Exception as e: + assert isinstance(e, Exception) + assert e == instance + was_raised = True + assert was_raised + + +class TestMutationsExceptionGroup(TestBigtableExceptionGroup): + def _get_class(self): + from google.cloud.bigtable.exceptions import MutationsExceptionGroup + + return MutationsExceptionGroup + + def _make_one(self, excs=None, num_entries=3): + if excs is None: + excs = [RuntimeError("mock")] + + return self._get_class()(excs, num_entries) + + @pytest.mark.parametrize( + "exception_list,total_entries,expected_message", + [ + ([Exception()], 1, "1 sub-exception (from 1 entry attempted)"), + ([Exception()], 2, "1 sub-exception (from 2 entries attempted)"), + ( + [Exception(), RuntimeError()], + 2, + "2 sub-exceptions (from 2 entries attempted)", + ), + ], + ) + def test_raise(self, exception_list, total_entries, expected_message): + """ + Create exception in raise statement, which calls __new__ and __init__ + """ + with pytest.raises(self._get_class()) as e: + raise self._get_class()(exception_list, total_entries) + assert str(e.value) == expected_message + assert list(e.value.exceptions) == exception_list + + +class TestRetryExceptionGroup(TestBigtableExceptionGroup): + def _get_class(self): + from google.cloud.bigtable.exceptions import RetryExceptionGroup + + return RetryExceptionGroup + + def _make_one(self, excs=None): + if excs is None: + excs = [RuntimeError("mock")] + + return self._get_class()(excs=excs) + + @pytest.mark.parametrize( + "exception_list,expected_message", + [ + ([Exception()], "1 failed attempt: Exception"), + ([Exception(), RuntimeError()], "2 failed attempts. Latest: RuntimeError"), + ( + [Exception(), ValueError("test")], + "2 failed attempts. Latest: ValueError", + ), + ( + [ + bigtable_exceptions.RetryExceptionGroup( + [Exception(), ValueError("test")] + ) + ], + "1 failed attempt: RetryExceptionGroup", + ), + ], + ) + def test_raise(self, exception_list, expected_message): + """ + Create exception in raise statement, which calls __new__ and __init__ + """ + with pytest.raises(self._get_class()) as e: + raise self._get_class()(exception_list) + assert str(e.value) == expected_message + assert list(e.value.exceptions) == exception_list + + +class TestFailedMutationEntryError: + def _get_class(self): + from google.cloud.bigtable.exceptions import FailedMutationEntryError + + return FailedMutationEntryError + + def _make_one(self, idx=9, entry=unittest.mock.Mock(), cause=RuntimeError("mock")): + + return self._get_class()(idx, entry, cause) + + def test_raise(self): + """ + Create exception in raise statement, which calls __new__ and __init__ + """ + test_idx = 2 + test_entry = unittest.mock.Mock() + test_exc = ValueError("test") + with pytest.raises(self._get_class()) as e: + raise self._get_class()(test_idx, test_entry, test_exc) + assert ( + str(e.value) + == "Failed idempotent mutation entry at index 2 with cause: ValueError('test')" + ) + assert e.value.index == test_idx + assert e.value.entry == test_entry + assert e.value.__cause__ == test_exc + assert isinstance(e.value, Exception) + assert test_entry.is_idempotent.call_count == 1 + + def test_raise_idempotent(self): + """ + Test raise with non idempotent entry + """ + test_idx = 2 + test_entry = unittest.mock.Mock() + test_entry.is_idempotent.return_value = False + test_exc = ValueError("test") + with pytest.raises(self._get_class()) as e: + raise self._get_class()(test_idx, test_entry, test_exc) + assert ( + str(e.value) + == "Failed non-idempotent mutation entry at index 2 with cause: ValueError('test')" + ) + assert e.value.index == test_idx + assert e.value.entry == test_entry + assert e.value.__cause__ == test_exc + assert test_entry.is_idempotent.call_count == 1 diff --git a/tests/unit/test_mutations.py b/tests/unit/test_mutations.py new file mode 100644 index 000000000..2a376609e --- /dev/null +++ b/tests/unit/test_mutations.py @@ -0,0 +1,569 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +import pytest + +import google.cloud.bigtable.mutations as mutations + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock +except ImportError: # pragma: NO COVER + import mock # type: ignore + + +class TestBaseMutation: + def _target_class(self): + from google.cloud.bigtable.mutations import Mutation + + return Mutation + + def test__to_dict(self): + """Should be unimplemented in the base class""" + with pytest.raises(NotImplementedError): + self._target_class()._to_dict(mock.Mock()) + + def test_is_idempotent(self): + """is_idempotent should assume True""" + assert self._target_class().is_idempotent(mock.Mock()) + + def test___str__(self): + """Str representation of mutations should be to_dict""" + self_mock = mock.Mock() + str_value = self._target_class().__str__(self_mock) + assert self_mock._to_dict.called + assert str_value == str(self_mock._to_dict.return_value) + + @pytest.mark.parametrize( + "expected_class,input_dict", + [ + ( + mutations.SetCell, + { + "set_cell": { + "family_name": "foo", + "column_qualifier": b"bar", + "value": b"test", + "timestamp_micros": 12345, + } + }, + ), + ( + mutations.DeleteRangeFromColumn, + { + "delete_from_column": { + "family_name": "foo", + "column_qualifier": b"bar", + "time_range": {}, + } + }, + ), + ( + mutations.DeleteRangeFromColumn, + { + "delete_from_column": { + "family_name": "foo", + "column_qualifier": b"bar", + "time_range": {"start_timestamp_micros": 123456789}, + } + }, + ), + ( + mutations.DeleteRangeFromColumn, + { + "delete_from_column": { + "family_name": "foo", + "column_qualifier": b"bar", + "time_range": {"end_timestamp_micros": 123456789}, + } + }, + ), + ( + mutations.DeleteRangeFromColumn, + { + "delete_from_column": { + "family_name": "foo", + "column_qualifier": b"bar", + "time_range": { + "start_timestamp_micros": 123, + "end_timestamp_micros": 123456789, + }, + } + }, + ), + ( + mutations.DeleteAllFromFamily, + {"delete_from_family": {"family_name": "foo"}}, + ), + (mutations.DeleteAllFromRow, {"delete_from_row": {}}), + ], + ) + def test__from_dict(self, expected_class, input_dict): + """Should be able to create instance from dict""" + instance = self._target_class()._from_dict(input_dict) + assert isinstance(instance, expected_class) + found_dict = instance._to_dict() + assert found_dict == input_dict + + @pytest.mark.parametrize( + "input_dict", + [ + {"set_cell": {}}, + { + "set_cell": { + "column_qualifier": b"bar", + "value": b"test", + "timestamp_micros": 12345, + } + }, + { + "set_cell": { + "family_name": "f", + "column_qualifier": b"bar", + "value": b"test", + } + }, + {"delete_from_family": {}}, + {"delete_from_column": {}}, + {"fake-type"}, + {}, + ], + ) + def test__from_dict_missing_fields(self, input_dict): + """If dict is malformed or fields are missing, should raise ValueError""" + with pytest.raises(ValueError): + self._target_class()._from_dict(input_dict) + + def test__from_dict_wrong_subclass(self): + """You shouldn't be able to instantiate one mutation type using the dict of another""" + subclasses = [ + mutations.SetCell("foo", b"bar", b"test"), + mutations.DeleteRangeFromColumn("foo", b"bar"), + mutations.DeleteAllFromFamily("foo"), + mutations.DeleteAllFromRow(), + ] + for instance in subclasses: + others = [other for other in subclasses if other != instance] + for other in others: + with pytest.raises(ValueError) as e: + type(other)._from_dict(instance._to_dict()) + assert "Mutation type mismatch" in str(e.value) + + +class TestSetCell: + def _target_class(self): + from google.cloud.bigtable.mutations import SetCell + + return SetCell + + def _make_one(self, *args, **kwargs): + return self._target_class()(*args, **kwargs) + + def test_ctor(self): + """Ensure constructor sets expected values""" + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + expected_value = b"test-value" + expected_timestamp = 1234567890 + instance = self._make_one( + expected_family, expected_qualifier, expected_value, expected_timestamp + ) + assert instance.family == expected_family + assert instance.qualifier == expected_qualifier + assert instance.new_value == expected_value + assert instance.timestamp_micros == expected_timestamp + + def test_ctor_str_inputs(self): + """Test with string qualifier and value""" + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + expected_value = b"test-value" + instance = self._make_one(expected_family, "test-qualifier", "test-value") + assert instance.family == expected_family + assert instance.qualifier == expected_qualifier + assert instance.new_value == expected_value + + @pytest.mark.parametrize( + "int_value,expected_bytes", + [ + (-42, b"\xff\xff\xff\xff\xff\xff\xff\xd6"), + (-2, b"\xff\xff\xff\xff\xff\xff\xff\xfe"), + (-1, b"\xff\xff\xff\xff\xff\xff\xff\xff"), + (0, b"\x00\x00\x00\x00\x00\x00\x00\x00"), + (1, b"\x00\x00\x00\x00\x00\x00\x00\x01"), + (2, b"\x00\x00\x00\x00\x00\x00\x00\x02"), + (100, b"\x00\x00\x00\x00\x00\x00\x00d"), + ], + ) + def test_ctor_int_value(self, int_value, expected_bytes): + """Test with int value""" + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + instance = self._make_one(expected_family, expected_qualifier, int_value) + assert instance.family == expected_family + assert instance.qualifier == expected_qualifier + assert instance.new_value == expected_bytes + + def test_ctor_negative_timestamp(self): + """Only positive or -1 timestamps are valid""" + with pytest.raises(ValueError) as e: + self._make_one("test-family", b"test-qualifier", b"test-value", -2) + assert ( + "timestamp_micros must be positive (or -1 for server-side timestamp)" + in str(e.value) + ) + + @pytest.mark.parametrize( + "timestamp_ns,expected_timestamp_micros", + [ + (0, 0), + (1, 0), + (123, 0), + (999, 0), + (999_999, 0), + (1_000_000, 1000), + (1_234_567, 1000), + (1_999_999, 1000), + (2_000_000, 2000), + (1_234_567_890_123, 1_234_567_000), + ], + ) + def test_ctor_no_timestamp(self, timestamp_ns, expected_timestamp_micros): + """If no timestamp is given, should use current time with millisecond precision""" + with mock.patch("time.time_ns", return_value=timestamp_ns): + instance = self._make_one("test-family", b"test-qualifier", b"test-value") + assert instance.timestamp_micros == expected_timestamp_micros + + def test__to_dict(self): + """ensure dict representation is as expected""" + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + expected_value = b"test-value" + expected_timestamp = 123456789 + instance = self._make_one( + expected_family, expected_qualifier, expected_value, expected_timestamp + ) + got_dict = instance._to_dict() + assert list(got_dict.keys()) == ["set_cell"] + got_inner_dict = got_dict["set_cell"] + assert got_inner_dict["family_name"] == expected_family + assert got_inner_dict["column_qualifier"] == expected_qualifier + assert got_inner_dict["timestamp_micros"] == expected_timestamp + assert got_inner_dict["value"] == expected_value + assert len(got_inner_dict.keys()) == 4 + + def test__to_dict_server_timestamp(self): + """test with server side timestamp -1 value""" + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + expected_value = b"test-value" + expected_timestamp = -1 + instance = self._make_one( + expected_family, expected_qualifier, expected_value, expected_timestamp + ) + got_dict = instance._to_dict() + assert list(got_dict.keys()) == ["set_cell"] + got_inner_dict = got_dict["set_cell"] + assert got_inner_dict["family_name"] == expected_family + assert got_inner_dict["column_qualifier"] == expected_qualifier + assert got_inner_dict["timestamp_micros"] == expected_timestamp + assert got_inner_dict["value"] == expected_value + assert len(got_inner_dict.keys()) == 4 + + @pytest.mark.parametrize( + "timestamp,expected_value", + [ + (1234567890, True), + (1, True), + (0, True), + (-1, False), + (None, True), + ], + ) + def test_is_idempotent(self, timestamp, expected_value): + """is_idempotent is based on whether an explicit timestamp is set""" + instance = self._make_one( + "test-family", b"test-qualifier", b"test-value", timestamp + ) + assert instance.is_idempotent() is expected_value + + def test___str__(self): + """Str representation of mutations should be to_dict""" + instance = self._make_one( + "test-family", b"test-qualifier", b"test-value", 1234567890 + ) + str_value = instance.__str__() + dict_value = instance._to_dict() + assert str_value == str(dict_value) + + +class TestDeleteRangeFromColumn: + def _target_class(self): + from google.cloud.bigtable.mutations import DeleteRangeFromColumn + + return DeleteRangeFromColumn + + def _make_one(self, *args, **kwargs): + return self._target_class()(*args, **kwargs) + + def test_ctor(self): + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + expected_start = 1234567890 + expected_end = 1234567891 + instance = self._make_one( + expected_family, expected_qualifier, expected_start, expected_end + ) + assert instance.family == expected_family + assert instance.qualifier == expected_qualifier + assert instance.start_timestamp_micros == expected_start + assert instance.end_timestamp_micros == expected_end + + def test_ctor_no_timestamps(self): + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + instance = self._make_one(expected_family, expected_qualifier) + assert instance.family == expected_family + assert instance.qualifier == expected_qualifier + assert instance.start_timestamp_micros is None + assert instance.end_timestamp_micros is None + + def test_ctor_timestamps_out_of_order(self): + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + expected_start = 10 + expected_end = 1 + with pytest.raises(ValueError) as excinfo: + self._make_one( + expected_family, expected_qualifier, expected_start, expected_end + ) + assert "start_timestamp_micros must be <= end_timestamp_micros" in str( + excinfo.value + ) + + @pytest.mark.parametrize( + "start,end", + [ + (0, 1), + (None, 1), + (0, None), + ], + ) + def test__to_dict(self, start, end): + """Should be unimplemented in the base class""" + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + + instance = self._make_one(expected_family, expected_qualifier, start, end) + got_dict = instance._to_dict() + assert list(got_dict.keys()) == ["delete_from_column"] + got_inner_dict = got_dict["delete_from_column"] + assert len(got_inner_dict.keys()) == 3 + assert got_inner_dict["family_name"] == expected_family + assert got_inner_dict["column_qualifier"] == expected_qualifier + time_range_dict = got_inner_dict["time_range"] + expected_len = int(isinstance(start, int)) + int(isinstance(end, int)) + assert len(time_range_dict.keys()) == expected_len + if start is not None: + assert time_range_dict["start_timestamp_micros"] == start + if end is not None: + assert time_range_dict["end_timestamp_micros"] == end + + def test_is_idempotent(self): + """is_idempotent is always true""" + instance = self._make_one( + "test-family", b"test-qualifier", 1234567890, 1234567891 + ) + assert instance.is_idempotent() is True + + def test___str__(self): + """Str representation of mutations should be to_dict""" + instance = self._make_one("test-family", b"test-qualifier") + str_value = instance.__str__() + dict_value = instance._to_dict() + assert str_value == str(dict_value) + + +class TestDeleteAllFromFamily: + def _target_class(self): + from google.cloud.bigtable.mutations import DeleteAllFromFamily + + return DeleteAllFromFamily + + def _make_one(self, *args, **kwargs): + return self._target_class()(*args, **kwargs) + + def test_ctor(self): + expected_family = "test-family" + instance = self._make_one(expected_family) + assert instance.family_to_delete == expected_family + + def test__to_dict(self): + """Should be unimplemented in the base class""" + expected_family = "test-family" + instance = self._make_one(expected_family) + got_dict = instance._to_dict() + assert list(got_dict.keys()) == ["delete_from_family"] + got_inner_dict = got_dict["delete_from_family"] + assert len(got_inner_dict.keys()) == 1 + assert got_inner_dict["family_name"] == expected_family + + def test_is_idempotent(self): + """is_idempotent is always true""" + instance = self._make_one("test-family") + assert instance.is_idempotent() is True + + def test___str__(self): + """Str representation of mutations should be to_dict""" + instance = self._make_one("test-family") + str_value = instance.__str__() + dict_value = instance._to_dict() + assert str_value == str(dict_value) + + +class TestDeleteFromRow: + def _target_class(self): + from google.cloud.bigtable.mutations import DeleteAllFromRow + + return DeleteAllFromRow + + def _make_one(self, *args, **kwargs): + return self._target_class()(*args, **kwargs) + + def test_ctor(self): + self._make_one() + + def test__to_dict(self): + """Should be unimplemented in the base class""" + instance = self._make_one() + got_dict = instance._to_dict() + assert list(got_dict.keys()) == ["delete_from_row"] + assert len(got_dict["delete_from_row"].keys()) == 0 + + def test_is_idempotent(self): + """is_idempotent is always true""" + instance = self._make_one() + assert instance.is_idempotent() is True + + def test___str__(self): + """Str representation of mutations should be to_dict""" + instance = self._make_one() + assert instance.__str__() == "{'delete_from_row': {}}" + + +class TestRowMutationEntry: + def _target_class(self): + from google.cloud.bigtable.mutations import RowMutationEntry + + return RowMutationEntry + + def _make_one(self, row_key, mutations): + return self._target_class()(row_key, mutations) + + def test_ctor(self): + expected_key = b"row_key" + expected_mutations = [mock.Mock()] + instance = self._make_one(expected_key, expected_mutations) + assert instance.row_key == expected_key + assert list(instance.mutations) == expected_mutations + + def test_ctor_str_key(self): + expected_key = "row_key" + expected_mutations = [mock.Mock(), mock.Mock()] + instance = self._make_one(expected_key, expected_mutations) + assert instance.row_key == b"row_key" + assert list(instance.mutations) == expected_mutations + + def test_ctor_single_mutation(self): + from google.cloud.bigtable.mutations import DeleteAllFromRow + + expected_key = b"row_key" + expected_mutations = DeleteAllFromRow() + instance = self._make_one(expected_key, expected_mutations) + assert instance.row_key == expected_key + assert instance.mutations == (expected_mutations,) + + def test__to_dict(self): + expected_key = "row_key" + mutation_mock = mock.Mock() + n_mutations = 3 + expected_mutations = [mutation_mock for i in range(n_mutations)] + for mock_mutations in expected_mutations: + mock_mutations._to_dict.return_value = {"test": "data"} + instance = self._make_one(expected_key, expected_mutations) + expected_result = { + "row_key": b"row_key", + "mutations": [{"test": "data"}] * n_mutations, + } + assert instance._to_dict() == expected_result + assert mutation_mock._to_dict.call_count == n_mutations + + @pytest.mark.parametrize( + "mutations,result", + [ + ([], True), + ([mock.Mock(is_idempotent=lambda: True)], True), + ([mock.Mock(is_idempotent=lambda: False)], False), + ( + [ + mock.Mock(is_idempotent=lambda: True), + mock.Mock(is_idempotent=lambda: False), + ], + False, + ), + ( + [ + mock.Mock(is_idempotent=lambda: True), + mock.Mock(is_idempotent=lambda: True), + ], + True, + ), + ], + ) + def test_is_idempotent(self, mutations, result): + instance = self._make_one("row_key", mutations) + assert instance.is_idempotent() == result + + def test__from_dict_mock(self): + """ + test creating instance from entry dict, with mocked mutation._from_dict + """ + expected_key = b"row_key" + expected_mutations = [mock.Mock(), mock.Mock()] + input_dict = { + "row_key": expected_key, + "mutations": [{"test": "data"}, {"another": "data"}], + } + with mock.patch.object(mutations.Mutation, "_from_dict") as inner_from_dict: + inner_from_dict.side_effect = expected_mutations + instance = self._target_class()._from_dict(input_dict) + assert instance.row_key == b"row_key" + assert inner_from_dict.call_count == 2 + assert len(instance.mutations) == 2 + assert instance.mutations[0] == expected_mutations[0] + assert instance.mutations[1] == expected_mutations[1] + + def test__from_dict(self): + """ + test creating end-to-end with a real mutation instance + """ + input_dict = { + "row_key": b"row_key", + "mutations": [{"delete_from_family": {"family_name": "test_family"}}], + } + instance = self._target_class()._from_dict(input_dict) + assert instance.row_key == b"row_key" + assert len(instance.mutations) == 1 + assert isinstance(instance.mutations[0], mutations.DeleteAllFromFamily) + assert instance.mutations[0].family_to_delete == "test_family" From ec3fd01732a72204edf105382116712fd672d80d Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 6 Jun 2023 12:46:09 -0700 Subject: [PATCH 08/56] feat: literal value filter (#767) --- google/cloud/bigtable/row_filters.py | 44 +++++++++- tests/system/test_system.py | 50 ++++++++++++ tests/unit/test_row_filters.py | 116 +++++++++++++++++---------- 3 files changed, 163 insertions(+), 47 deletions(-) diff --git a/google/cloud/bigtable/row_filters.py b/google/cloud/bigtable/row_filters.py index 48a1d4d8a..b2fae6971 100644 --- a/google/cloud/bigtable/row_filters.py +++ b/google/cloud/bigtable/row_filters.py @@ -481,20 +481,56 @@ def to_dict(self) -> dict[str, bytes]: return {"value_regex_filter": self.regex} -class ExactValueFilter(ValueRegexFilter): +class LiteralValueFilter(ValueRegexFilter): """Row filter for an exact value. :type value: bytes or str or int :param value: - a literal string encodable as ASCII, or the - equivalent bytes, or an integer (which will be packed into 8-bytes). + a literal string, integer, or the equivalent bytes. + Integer values will be packed into signed 8-bytes. """ def __init__(self, value: bytes | str | int): if isinstance(value, int): value = _PACK_I64(value) - super(ExactValueFilter, self).__init__(value) + elif isinstance(value, str): + value = value.encode("utf-8") + value = self._write_literal_regex(value) + super(LiteralValueFilter, self).__init__(value) + + @staticmethod + def _write_literal_regex(input_bytes: bytes) -> bytes: + """ + Escape re2 special characters from literal bytes. + + Extracted from: re2 QuoteMeta: + https://github.com/google/re2/blob/70f66454c255080a54a8da806c52d1f618707f8a/re2/re2.cc#L456 + """ + result = bytearray() + for byte in input_bytes: + # If this is the part of a UTF8 or Latin1 character, we need \ + # to copy this byte without escaping. Experimentally this is \ + # what works correctly with the regexp library. \ + utf8_latin1_check = (byte & 128) == 0 + if ( + (byte < ord("a") or byte > ord("z")) + and (byte < ord("A") or byte > ord("Z")) + and (byte < ord("0") or byte > ord("9")) + and byte != ord("_") + and utf8_latin1_check + ): + if byte == 0: + # Special handling for null chars. + # Note that this special handling is not strictly required for RE2, + # but this quoting is required for other regexp libraries such as + # PCRE. + # Can't use "\\0" since the next character might be a digit. + result.extend([ord("\\"), ord("x"), ord("0"), ord("0")]) + continue + result.append(ord(b"\\")) + result.append(byte) + return bytes(result) def __repr__(self) -> str: return f"{self.__class__.__name__}(value={self.regex!r})" diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 7d015224c..f0fab7d45 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -160,6 +160,10 @@ def __init__(self, table): async def add_row( self, row_key, family=TEST_FAMILY, qualifier=b"q", value=b"test-value" ): + if isinstance(value, str): + value = value.encode("utf-8") + elif isinstance(value, int): + value = value.to_bytes(8, byteorder="big", signed=True) request = { "table_name": self.table.table_name, "row_key": row_key, @@ -391,3 +395,49 @@ async def test_read_rows_stream_inactive_timer(table, temp_rows): await generator.__anext__() assert "inactivity" in str(e) assert "idle_timeout=0.1" in str(e) + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.parametrize( + "cell_value,filter_input,expect_match", + [ + (b"abc", b"abc", True), + (b"abc", "abc", True), + (b".", ".", True), + (".*", ".*", True), + (".*", b".*", True), + ("a", ".*", False), + (b".*", b".*", True), + (r"\a", r"\a", True), + (b"\xe2\x98\x83", "☃", True), + ("☃", "☃", True), + (r"\C☃", r"\C☃", True), + (1, 1, True), + (2, 1, False), + (68, 68, True), + ("D", 68, False), + (68, "D", False), + (-1, -1, True), + (2852126720, 2852126720, True), + (-1431655766, -1431655766, True), + (-1431655766, -1, False), + ], +) +@pytest.mark.asyncio +async def test_literal_value_filter( + table, temp_rows, cell_value, filter_input, expect_match +): + """ + Literal value filter does complex escaping on re2 strings. + Make sure inputs are properly interpreted by the server + """ + from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable import ReadRowsQuery + + f = LiteralValueFilter(filter_input) + await temp_rows.add_row(b"row_key_1", value=cell_value) + query = ReadRowsQuery(row_filter=f) + row_list = await table.read_rows(query) + assert len(row_list) == bool( + expect_match + ), f"row {type(cell_value)}({cell_value}) not found with {type(filter_input)}({filter_input}) filter" diff --git a/tests/unit/test_row_filters.py b/tests/unit/test_row_filters.py index d0fbad42f..11ff9f2f1 100644 --- a/tests/unit/test_row_filters.py +++ b/tests/unit/test_row_filters.py @@ -822,84 +822,90 @@ def test_value_regex_filter___repr__(): assert eval(repr(row_filter)) == row_filter -def test_exact_value_filter_to_pb_w_bytes(): - from google.cloud.bigtable.row_filters import ExactValueFilter +def test_literal_value_filter_to_pb_w_bytes(): + from google.cloud.bigtable.row_filters import LiteralValueFilter - value = regex = b"value-regex" - row_filter = ExactValueFilter(value) + value = regex = b"value_regex" + row_filter = LiteralValueFilter(value) pb_val = row_filter._to_pb() expected_pb = _RowFilterPB(value_regex_filter=regex) assert pb_val == expected_pb -def test_exact_value_filter_to_dict_w_bytes(): - from google.cloud.bigtable.row_filters import ExactValueFilter +def test_literal_value_filter_to_dict_w_bytes(): + from google.cloud.bigtable.row_filters import LiteralValueFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 - value = regex = b"value-regex" - row_filter = ExactValueFilter(value) + value = regex = b"value_regex" + row_filter = LiteralValueFilter(value) expected_dict = {"value_regex_filter": regex} assert row_filter.to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value -def test_exact_value_filter_to_pb_w_str(): - from google.cloud.bigtable.row_filters import ExactValueFilter +def test_literal_value_filter_to_pb_w_str(): + from google.cloud.bigtable.row_filters import LiteralValueFilter - value = "value-regex" + value = "value_regex" regex = value.encode("ascii") - row_filter = ExactValueFilter(value) + row_filter = LiteralValueFilter(value) pb_val = row_filter._to_pb() expected_pb = _RowFilterPB(value_regex_filter=regex) assert pb_val == expected_pb -def test_exact_value_filter_to_dict_w_str(): - from google.cloud.bigtable.row_filters import ExactValueFilter +def test_literal_value_filter_to_dict_w_str(): + from google.cloud.bigtable.row_filters import LiteralValueFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 - value = "value-regex" + value = "value_regex" regex = value.encode("ascii") - row_filter = ExactValueFilter(value) + row_filter = LiteralValueFilter(value) expected_dict = {"value_regex_filter": regex} assert row_filter.to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value -def test_exact_value_filter_to_pb_w_int(): - import struct - from google.cloud.bigtable.row_filters import ExactValueFilter +@pytest.mark.parametrize( + "value,expected_byte_string", + [ + # null bytes are encoded as "\x00" in ascii characters + # others are just prefixed with "\" + (0, b"\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00"), + (1, b"\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\\x01"), + ( + 68, + b"\\x00\\x00\\x00\\x00\\x00\\x00\\x00D", + ), # bytes that encode to alphanum are not escaped + (570, b"\\x00\\x00\\x00\\x00\\x00\\x00\\\x02\\\x3a"), + (2852126720, b"\\x00\\x00\\x00\\x00\xaa\\x00\\x00\\x00"), + (-1, b"\xff\xff\xff\xff\xff\xff\xff\xff"), + (-1096642724096, b"\xff\xff\xff\\x00\xaa\xff\xff\\x00"), + ], +) +def test_literal_value_filter_w_int(value, expected_byte_string): + from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable_v2.types import data as data_v2_pb2 - value = 1 - regex = struct.Struct(">q").pack(value) - row_filter = ExactValueFilter(value) + row_filter = LiteralValueFilter(value) + # test pb pb_val = row_filter._to_pb() - expected_pb = _RowFilterPB(value_regex_filter=regex) + expected_pb = _RowFilterPB(value_regex_filter=expected_byte_string) assert pb_val == expected_pb - - -def test_exact_value_filter_to_dict_w_int(): - import struct - from google.cloud.bigtable.row_filters import ExactValueFilter - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - value = 1 - regex = struct.Struct(">q").pack(value) - row_filter = ExactValueFilter(value) - expected_dict = {"value_regex_filter": regex} + # test dict + expected_dict = {"value_regex_filter": expected_byte_string} assert row_filter.to_dict() == expected_dict - expected_pb_value = row_filter._to_pb() - assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value + assert data_v2_pb2.RowFilter(**expected_dict) == pb_val -def test_exact_value_filter___repr__(): - from google.cloud.bigtable.row_filters import ExactValueFilter +def test_literal_value_filter___repr__(): + from google.cloud.bigtable.row_filters import LiteralValueFilter - value = "value-regex" - row_filter = ExactValueFilter(value) - expected = "ExactValueFilter(value=b'value-regex')" + value = "value_regex" + row_filter = LiteralValueFilter(value) + expected = "LiteralValueFilter(value=b'value_regex')" assert repr(row_filter) == expected assert repr(row_filter) == str(row_filter) assert eval(repr(row_filter)) == row_filter @@ -1907,6 +1913,30 @@ def test_conditional_row_filter___str__(): assert str(row_filter4) == expected +@pytest.mark.parametrize( + "input_arg, expected_bytes", + [ + (b"abc", b"abc"), + ("abc", b"abc"), + (1, b"\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\\x01"), # null bytes are ascii + (b"*", b"\\*"), + (".", b"\\."), + (b"\\", b"\\\\"), + (b"h.*i", b"h\\.\\*i"), + (b'""', b'\\"\\"'), + (b"[xyz]", b"\\[xyz\\]"), + (b"\xe2\x98\xba\xef\xb8\x8f", b"\xe2\x98\xba\xef\xb8\x8f"), + ("☃", b"\xe2\x98\x83"), + (r"\C☃", b"\\\\C\xe2\x98\x83"), + ], +) +def test_literal_value__write_literal_regex(input_arg, expected_bytes): + from google.cloud.bigtable.row_filters import LiteralValueFilter + + filter_ = LiteralValueFilter(input_arg) + assert filter_.regex == expected_bytes + + def _ColumnRangePB(*args, **kw): from google.cloud.bigtable_v2.types import data as data_v2_pb2 @@ -1955,7 +1985,7 @@ def _get_regex_filters(): FamilyNameRegexFilter, ColumnQualifierRegexFilter, ValueRegexFilter, - ExactValueFilter, + LiteralValueFilter, ) return [ @@ -1963,7 +1993,7 @@ def _get_regex_filters(): FamilyNameRegexFilter, ColumnQualifierRegexFilter, ValueRegexFilter, - ExactValueFilter, + LiteralValueFilter, ] From 5d65703783eadf3b82fb0d4cb0ef36393add97a7 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 16 Jun 2023 12:18:19 -0700 Subject: [PATCH 09/56] feat: row_exists and read_row (#778) --- google/cloud/bigtable/client.py | 36 +++++- google/cloud/bigtable/read_rows_query.py | 10 +- tests/system/test_system.py | 97 ++++++++++++++- tests/unit/test_client.py | 152 +++++++++++++++++++++++ tests/unit/test_read_rows_query.py | 7 +- 5 files changed, 290 insertions(+), 12 deletions(-) diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index 3921d6640..0544bcb78 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -55,6 +55,10 @@ from google.cloud.bigtable._helpers import _make_metadata from google.cloud.bigtable._helpers import _convert_retry_deadline +from google.cloud.bigtable.row_filters import StripValueTransformerFilter +from google.cloud.bigtable.row_filters import CellsRowLimitFilter +from google.cloud.bigtable.row_filters import RowFilterChain + if TYPE_CHECKING: from google.cloud.bigtable.mutations_batcher import MutationsBatcher from google.cloud.bigtable import RowKeySamples @@ -500,18 +504,31 @@ async def read_row( self, row_key: str | bytes, *, + row_filter: RowFilter | None = None, operation_timeout: int | float | None = 60, per_request_timeout: int | float | None = None, - ) -> Row: + ) -> Row | None: """ Helper function to return a single row See read_rows_stream + Raises: + - google.cloud.bigtable.exceptions.RowNotFound: if the row does not exist Returns: - - the individual row requested + - the individual row requested, or None if it does not exist """ - raise NotImplementedError + if row_key is None: + raise ValueError("row_key must be string or bytes") + query = ReadRowsQuery(row_keys=row_key, row_filter=row_filter, limit=1) + results = await self.read_rows( + query, + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + if len(results) == 0: + return None + return results[0] async def read_rows_sharded( self, @@ -547,7 +564,18 @@ async def row_exists( Returns: - a bool indicating whether the row exists """ - raise NotImplementedError + if row_key is None: + raise ValueError("row_key must be string or bytes") + strip_filter = StripValueTransformerFilter(flag=True) + limit_filter = CellsRowLimitFilter(1) + chain_filter = RowFilterChain(filters=[limit_filter, strip_filter]) + query = ReadRowsQuery(row_keys=row_key, limit=1, row_filter=chain_filter) + results = await self.read_rows( + query, + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + return len(results) > 0 async def sample_keys( self, diff --git a/google/cloud/bigtable/read_rows_query.py b/google/cloud/bigtable/read_rows_query.py index e26f99d34..6de84e918 100644 --- a/google/cloud/bigtable/read_rows_query.py +++ b/google/cloud/bigtable/read_rows_query.py @@ -106,12 +106,12 @@ def __init__( """ self.row_keys: set[bytes] = set() self.row_ranges: list[RowRange | dict[str, bytes]] = [] - if row_ranges: + if row_ranges is not None: if isinstance(row_ranges, RowRange): row_ranges = [row_ranges] for r in row_ranges: self.add_range(r) - if row_keys: + if row_keys is not None: if not isinstance(row_keys, list): row_keys = [row_keys] for k in row_keys: @@ -221,7 +221,11 @@ def _to_dict(self) -> dict[str, Any]: row_ranges.append(dict_range) row_keys = list(self.row_keys) row_keys.sort() - row_set = {"row_keys": row_keys, "row_ranges": row_ranges} + row_set: dict[str, Any] = {} + if row_keys: + row_set["row_keys"] = row_keys + if row_ranges: + row_set["row_ranges"] = row_ranges final_dict: dict[str, Any] = { "rows": row_set, } diff --git a/tests/system/test_system.py b/tests/system/test_system.py index f0fab7d45..f6730576d 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -158,7 +158,7 @@ def __init__(self, table): self.table = table async def add_row( - self, row_key, family=TEST_FAMILY, qualifier=b"q", value=b"test-value" + self, row_key, *, family=TEST_FAMILY, qualifier=b"q", value=b"test-value" ): if isinstance(value, str): value = value.encode("utf-8") @@ -339,9 +339,9 @@ async def test_read_rows_range_query(table, temp_rows): @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio -async def test_read_rows_key_query(table, temp_rows): +async def test_read_rows_single_key_query(table, temp_rows): """ - Ensure that the read_rows method works + Ensure that the read_rows method works with specified query """ from google.cloud.bigtable import ReadRowsQuery @@ -349,7 +349,7 @@ async def test_read_rows_key_query(table, temp_rows): await temp_rows.add_row(b"b") await temp_rows.add_row(b"c") await temp_rows.add_row(b"d") - # full table scan + # retrieve specific keys query = ReadRowsQuery(row_keys=[b"a", b"c"]) row_list = await table.read_rows(query) assert len(row_list) == 2 @@ -357,6 +357,29 @@ async def test_read_rows_key_query(table, temp_rows): assert row_list[1].row_key == b"c" +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows_with_filter(table, temp_rows): + """ + ensure filters are applied + """ + from google.cloud.bigtable import ReadRowsQuery + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + # retrieve keys with filter + expected_label = "test-label" + row_filter = ApplyLabelFilter(expected_label) + query = ReadRowsQuery(row_filter=row_filter) + row_list = await table.read_rows(query) + assert len(row_list) == 4 + for row in row_list: + assert row[0].labels == [expected_label] + + @pytest.mark.asyncio async def test_read_rows_stream_close(table, temp_rows): """ @@ -397,6 +420,72 @@ async def test_read_rows_stream_inactive_timer(table, temp_rows): assert "idle_timeout=0.1" in str(e) +@pytest.mark.asyncio +async def test_read_row(table, temp_rows): + """ + Test read_row (single row helper) + """ + from google.cloud.bigtable import Row + + await temp_rows.add_row(b"row_key_1", value=b"value") + row = await table.read_row(b"row_key_1") + assert isinstance(row, Row) + assert row.row_key == b"row_key_1" + assert row.cells[0].value == b"value" + + +@pytest.mark.asyncio +async def test_read_row_missing(table): + """ + Test read_row when row does not exist + """ + from google.api_core import exceptions + + row_key = "row_key_not_exist" + result = await table.read_row(row_key) + assert result is None + with pytest.raises(exceptions.InvalidArgument) as e: + await table.read_row("") + assert "Row key must be non-empty" in str(e) + + +@pytest.mark.asyncio +async def test_read_row_w_filter(table, temp_rows): + """ + Test read_row (single row helper) + """ + from google.cloud.bigtable import Row + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + await temp_rows.add_row(b"row_key_1", value=b"value") + expected_label = "test-label" + label_filter = ApplyLabelFilter(expected_label) + row = await table.read_row(b"row_key_1", row_filter=label_filter) + assert isinstance(row, Row) + assert row.row_key == b"row_key_1" + assert row.cells[0].value == b"value" + assert row.cells[0].labels == [expected_label] + + +@pytest.mark.asyncio +async def test_row_exists(table, temp_rows): + from google.api_core import exceptions + + """Test row_exists with rows that exist and don't exist""" + assert await table.row_exists(b"row_key_1") is False + await temp_rows.add_row(b"row_key_1") + assert await table.row_exists(b"row_key_1") is True + assert await table.row_exists("row_key_1") is True + assert await table.row_exists(b"row_key_2") is False + assert await table.row_exists("row_key_2") is False + assert await table.row_exists("3") is False + await temp_rows.add_row(b"3") + assert await table.row_exists(b"3") is True + with pytest.raises(exceptions.InvalidArgument) as e: + await table.row_exists("") + assert "Row kest must be non-empty" in str(e) + + @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.parametrize( "cell_value,filter_input,expect_match", diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index be3703a23..14da80dae 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1296,6 +1296,158 @@ async def test_read_rows_default_timeout_override(self): assert kwargs["operation_timeout"] == operation_timeout assert kwargs["per_request_timeout"] == per_request_timeout + @pytest.mark.asyncio + async def test_read_row(self): + """Test reading a single row""" + async with self._make_client() as client: + table = client.get_table("instance", "table") + row_key = b"test_1" + with mock.patch.object(table, "read_rows") as read_rows: + expected_result = object() + read_rows.side_effect = lambda *args, **kwargs: [expected_result] + expected_op_timeout = 8 + expected_req_timeout = 4 + row = await table.read_row( + row_key, + operation_timeout=expected_op_timeout, + per_request_timeout=expected_req_timeout, + ) + assert row == expected_result + assert read_rows.call_count == 1 + args, kwargs = read_rows.call_args_list[0] + assert kwargs["operation_timeout"] == expected_op_timeout + assert kwargs["per_request_timeout"] == expected_req_timeout + assert len(args) == 1 + assert isinstance(args[0], ReadRowsQuery) + assert args[0]._to_dict() == { + "rows": {"row_keys": [row_key]}, + "rows_limit": 1, + } + + @pytest.mark.asyncio + async def test_read_row_w_filter(self): + """Test reading a single row with an added filter""" + async with self._make_client() as client: + table = client.get_table("instance", "table") + row_key = b"test_1" + with mock.patch.object(table, "read_rows") as read_rows: + expected_result = object() + read_rows.side_effect = lambda *args, **kwargs: [expected_result] + expected_op_timeout = 8 + expected_req_timeout = 4 + mock_filter = mock.Mock() + expected_filter = {"filter": "mock filter"} + mock_filter._to_dict.return_value = expected_filter + row = await table.read_row( + row_key, + operation_timeout=expected_op_timeout, + per_request_timeout=expected_req_timeout, + row_filter=expected_filter, + ) + assert row == expected_result + assert read_rows.call_count == 1 + args, kwargs = read_rows.call_args_list[0] + assert kwargs["operation_timeout"] == expected_op_timeout + assert kwargs["per_request_timeout"] == expected_req_timeout + assert len(args) == 1 + assert isinstance(args[0], ReadRowsQuery) + assert args[0]._to_dict() == { + "rows": {"row_keys": [row_key]}, + "rows_limit": 1, + "filter": expected_filter, + } + + @pytest.mark.asyncio + async def test_read_row_no_response(self): + """should return None if row does not exist""" + async with self._make_client() as client: + table = client.get_table("instance", "table") + row_key = b"test_1" + with mock.patch.object(table, "read_rows") as read_rows: + # return no rows + read_rows.side_effect = lambda *args, **kwargs: [] + expected_op_timeout = 8 + expected_req_timeout = 4 + result = await table.read_row( + row_key, + operation_timeout=expected_op_timeout, + per_request_timeout=expected_req_timeout, + ) + assert result is None + assert read_rows.call_count == 1 + args, kwargs = read_rows.call_args_list[0] + assert kwargs["operation_timeout"] == expected_op_timeout + assert kwargs["per_request_timeout"] == expected_req_timeout + assert isinstance(args[0], ReadRowsQuery) + assert args[0]._to_dict() == { + "rows": {"row_keys": [row_key]}, + "rows_limit": 1, + } + + @pytest.mark.parametrize("input_row", [None, 5, object()]) + @pytest.mark.asyncio + async def test_read_row_w_invalid_input(self, input_row): + """Should raise error when passed None""" + async with self._make_client() as client: + table = client.get_table("instance", "table") + with pytest.raises(ValueError) as e: + await table.read_row(input_row) + assert "must be string or bytes" in e + + @pytest.mark.parametrize( + "return_value,expected_result", + [ + ([], False), + ([object()], True), + ([object(), object()], True), + ], + ) + @pytest.mark.asyncio + async def test_row_exists(self, return_value, expected_result): + """Test checking for row existence""" + async with self._make_client() as client: + table = client.get_table("instance", "table") + row_key = b"test_1" + with mock.patch.object(table, "read_rows") as read_rows: + # return no rows + read_rows.side_effect = lambda *args, **kwargs: return_value + expected_op_timeout = 1 + expected_req_timeout = 2 + result = await table.row_exists( + row_key, + operation_timeout=expected_op_timeout, + per_request_timeout=expected_req_timeout, + ) + assert expected_result == result + assert read_rows.call_count == 1 + args, kwargs = read_rows.call_args_list[0] + assert kwargs["operation_timeout"] == expected_op_timeout + assert kwargs["per_request_timeout"] == expected_req_timeout + assert isinstance(args[0], ReadRowsQuery) + expected_filter = { + "chain": { + "filters": [ + {"cells_per_row_limit_filter": 1}, + {"strip_value_transformer": True}, + ] + } + } + assert args[0]._to_dict() == { + "rows": {"row_keys": [row_key]}, + "rows_limit": 1, + "filter": expected_filter, + } + + @pytest.mark.parametrize("input_row", [None, 5, object()]) + @pytest.mark.asyncio + async def test_row_exists_w_invalid_input(self, input_row): + """Should raise error when passed None""" + async with self._make_client() as client: + table = client.get_table("instance", "table") + with pytest.raises(ValueError) as e: + await table.row_exists(input_row) + assert "must be string or bytes" in e + @pytest.mark.parametrize("include_app_profile", [True, False]) @pytest.mark.asyncio async def test_read_rows_metadata(self, include_app_profile): diff --git a/tests/unit/test_read_rows_query.py b/tests/unit/test_read_rows_query.py index aa690bc86..f630f2eab 100644 --- a/tests/unit/test_read_rows_query.py +++ b/tests/unit/test_read_rows_query.py @@ -300,7 +300,7 @@ def test_to_dict_rows_default(self): output = query._to_dict() self.assertTrue(isinstance(output, dict)) self.assertEqual(len(output.keys()), 1) - expected = {"rows": {"row_keys": [], "row_ranges": []}} + expected = {"rows": {}} self.assertEqual(output, expected) request_proto = ReadRowsRequest(**output) @@ -355,5 +355,10 @@ def test_to_dict_rows_populated(self): filter_proto = request_proto.filter self.assertEqual(filter_proto, row_filter._to_pb()) + def test_empty_row_set(self): + """Empty strings should be treated as keys inputs""" + query = self._make_one(row_keys="") + self.assertEqual(query.row_keys, {b""}) + def test_shard(self): pass From 432d159925233125212eafa490dae241060bc50d Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 16 Jun 2023 13:29:53 -0700 Subject: [PATCH 10/56] feat: read_modify_write and check_and_mutate_row (#780) --- google/cloud/bigtable/client.py | 71 +++- google/cloud/bigtable/mutations.py | 6 + .../cloud/bigtable/read_modify_write_rules.py | 59 ++- google/cloud/bigtable/row.py | 26 ++ tests/system/test_system.py | 162 +++++++- tests/unit/test_client.py | 353 ++++++++++++++++++ tests/unit/test_mutations.py | 18 +- tests/unit/test_read_modify_write_rules.py | 142 +++++++ tests/unit/test_row.py | 44 +++ 9 files changed, 857 insertions(+), 24 deletions(-) create mode 100644 tests/unit/test_read_modify_write_rules.py diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index 0544bcb78..f75613098 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -55,6 +55,8 @@ from google.cloud.bigtable._helpers import _make_metadata from google.cloud.bigtable._helpers import _convert_retry_deadline +from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule +from google.cloud.bigtable.row_filters import RowFilter from google.cloud.bigtable.row_filters import StripValueTransformerFilter from google.cloud.bigtable.row_filters import CellsRowLimitFilter from google.cloud.bigtable.row_filters import RowFilterChain @@ -62,8 +64,6 @@ if TYPE_CHECKING: from google.cloud.bigtable.mutations_batcher import MutationsBatcher from google.cloud.bigtable import RowKeySamples - from google.cloud.bigtable.row_filters import RowFilter - from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule class BigtableDataClient(ClientWithProject): @@ -770,10 +770,11 @@ async def bulk_mutate_rows( async def check_and_mutate_row( self, row_key: str | bytes, - predicate: RowFilter | None, + predicate: RowFilter | dict[str, Any] | None, + *, true_case_mutations: Mutation | list[Mutation] | None = None, false_case_mutations: Mutation | list[Mutation] | None = None, - operation_timeout: int | float | None = 60, + operation_timeout: int | float | None = 20, ) -> bool: """ Mutates a row atomically based on the output of a predicate filter @@ -807,17 +808,43 @@ async def check_and_mutate_row( Raises: - GoogleAPIError exceptions from grpc call """ - raise NotImplementedError + operation_timeout = operation_timeout or self.default_operation_timeout + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key + if true_case_mutations is not None and not isinstance( + true_case_mutations, list + ): + true_case_mutations = [true_case_mutations] + true_case_dict = [m._to_dict() for m in true_case_mutations or []] + if false_case_mutations is not None and not isinstance( + false_case_mutations, list + ): + false_case_mutations = [false_case_mutations] + false_case_dict = [m._to_dict() for m in false_case_mutations or []] + if predicate is not None and not isinstance(predicate, dict): + predicate = predicate.to_dict() + metadata = _make_metadata(self.table_name, self.app_profile_id) + result = await self.client._gapic_client.check_and_mutate_row( + request={ + "predicate_filter": predicate, + "true_mutations": true_case_dict, + "false_mutations": false_case_dict, + "table_name": self.table_name, + "row_key": row_key, + "app_profile_id": self.app_profile_id, + }, + metadata=metadata, + timeout=operation_timeout, + ) + return result.predicate_matched async def read_modify_write_row( self, row_key: str | bytes, - rules: ReadModifyWriteRule - | list[ReadModifyWriteRule] - | dict[str, Any] - | list[dict[str, Any]], + rules: ReadModifyWriteRule | list[ReadModifyWriteRule], *, - operation_timeout: int | float | None = 60, + operation_timeout: int | float | None = 20, ) -> Row: """ Reads and modifies a row atomically according to input ReadModifyWriteRules, @@ -841,7 +868,29 @@ async def read_modify_write_row( Raises: - GoogleAPIError exceptions from grpc call """ - raise NotImplementedError + operation_timeout = operation_timeout or self.default_operation_timeout + row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if rules is not None and not isinstance(rules, list): + rules = [rules] + if not rules: + raise ValueError("rules must contain at least one item") + # concert to dict representation + rules_dict = [rule._to_dict() for rule in rules] + metadata = _make_metadata(self.table_name, self.app_profile_id) + result = await self.client._gapic_client.read_modify_write_row( + request={ + "rules": rules_dict, + "table_name": self.table_name, + "row_key": row_key, + "app_profile_id": self.app_profile_id, + }, + metadata=metadata, + timeout=operation_timeout, + ) + # construct Row from result + return Row._from_pb(result.row) async def close(self): """ diff --git a/google/cloud/bigtable/mutations.py b/google/cloud/bigtable/mutations.py index c72f132c8..fe136f8d9 100644 --- a/google/cloud/bigtable/mutations.py +++ b/google/cloud/bigtable/mutations.py @@ -18,6 +18,8 @@ from dataclasses import dataclass from abc import ABC, abstractmethod +from google.cloud.bigtable.read_modify_write_rules import MAX_INCREMENT_VALUE + # special value for SetCell mutation timestamps. If set, server will assign a timestamp SERVER_SIDE_TIMESTAMP = -1 @@ -99,6 +101,10 @@ def __init__( if isinstance(new_value, str): new_value = new_value.encode() elif isinstance(new_value, int): + if abs(new_value) > MAX_INCREMENT_VALUE: + raise ValueError( + "int values must be between -2**63 and 2**63 (64-bit signed int)" + ) new_value = new_value.to_bytes(8, "big", signed=True) if not isinstance(new_value, bytes): raise TypeError("new_value must be bytes, str, or int") diff --git a/google/cloud/bigtable/read_modify_write_rules.py b/google/cloud/bigtable/read_modify_write_rules.py index cd6b370df..aa282b1a6 100644 --- a/google/cloud/bigtable/read_modify_write_rules.py +++ b/google/cloud/bigtable/read_modify_write_rules.py @@ -14,22 +14,59 @@ # from __future__ import annotations -from dataclasses import dataclass +import abc +# value must fit in 64-bit signed integer +MAX_INCREMENT_VALUE = (1 << 63) - 1 -class ReadModifyWriteRule: - pass + +class ReadModifyWriteRule(abc.ABC): + def __init__(self, family: str, qualifier: bytes | str): + qualifier = ( + qualifier if isinstance(qualifier, bytes) else qualifier.encode("utf-8") + ) + self.family = family + self.qualifier = qualifier + + @abc.abstractmethod + def _to_dict(self): + raise NotImplementedError -@dataclass class IncrementRule(ReadModifyWriteRule): - increment_amount: int - family: str - qualifier: bytes + def __init__(self, family: str, qualifier: bytes | str, increment_amount: int = 1): + if not isinstance(increment_amount, int): + raise TypeError("increment_amount must be an integer") + if abs(increment_amount) > MAX_INCREMENT_VALUE: + raise ValueError( + "increment_amount must be between -2**63 and 2**63 (64-bit signed int)" + ) + super().__init__(family, qualifier) + self.increment_amount = increment_amount + + def _to_dict(self): + return { + "family_name": self.family, + "column_qualifier": self.qualifier, + "increment_amount": self.increment_amount, + } -@dataclass class AppendValueRule(ReadModifyWriteRule): - append_value: bytes - family: str - qualifier: bytes + def __init__(self, family: str, qualifier: bytes | str, append_value: bytes | str): + append_value = ( + append_value.encode("utf-8") + if isinstance(append_value, str) + else append_value + ) + if not isinstance(append_value, bytes): + raise TypeError("append_value must be bytes or str") + super().__init__(family, qualifier) + self.append_value = append_value + + def _to_dict(self): + return { + "family_name": self.family, + "column_qualifier": self.qualifier, + "append_value": self.append_value, + } diff --git a/google/cloud/bigtable/row.py b/google/cloud/bigtable/row.py index a5fb033e6..5fdc1b365 100644 --- a/google/cloud/bigtable/row.py +++ b/google/cloud/bigtable/row.py @@ -18,6 +18,8 @@ from typing import Sequence, Generator, overload, Any from functools import total_ordering +from google.cloud.bigtable_v2.types import Row as RowPB + # Type aliases used internally for readability. _family_type = str _qualifier_type = bytes @@ -72,6 +74,30 @@ def _index( ).append(cell) return self._index_data + @classmethod + def _from_pb(cls, row_pb: RowPB) -> Row: + """ + Creates a row from a protobuf representation + + Row objects are not intended to be created by users. + They are returned by the Bigtable backend. + """ + row_key: bytes = row_pb.key + cell_list: list[Cell] = [] + for family in row_pb.families: + for column in family.columns: + for cell in column.cells: + new_cell = Cell( + value=cell.value, + row_key=row_key, + family=family.name, + qualifier=column.qualifier, + timestamp_micros=cell.timestamp_micros, + labels=list(cell.labels) if cell.labels else None, + ) + cell_list.append(new_cell) + return cls(row_key, cells=cell_list) + def get_cells( self, family: str | None = None, qualifier: str | bytes | None = None ) -> list[Cell]: diff --git a/tests/system/test_system.py b/tests/system/test_system.py index f6730576d..692911b10 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -19,6 +19,8 @@ from google.api_core import retry from google.api_core.exceptions import ClientError +from google.cloud.bigtable.read_modify_write_rules import MAX_INCREMENT_VALUE + TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" @@ -245,7 +247,6 @@ async def test_mutation_set_cell(table, temp_rows): mutation = SetCell( family=TEST_FAMILY, qualifier=b"test-qualifier", new_value=expected_value ) - await table.mutate_row(row_key, mutation) # ensure cell is updated @@ -282,6 +283,165 @@ async def test_bulk_mutations_set_cell(client, table, temp_rows): assert (await _retrieve_cell_value(table, row_key)) == expected_value +@pytest.mark.parametrize( + "start,increment,expected", + [ + (0, 0, 0), + (0, 1, 1), + (0, -1, -1), + (1, 0, 1), + (0, -100, -100), + (0, 3000, 3000), + (10, 4, 14), + (MAX_INCREMENT_VALUE, -MAX_INCREMENT_VALUE, 0), + (MAX_INCREMENT_VALUE, 2, -MAX_INCREMENT_VALUE), + (-MAX_INCREMENT_VALUE, -2, MAX_INCREMENT_VALUE), + ], +) +@pytest.mark.asyncio +async def test_read_modify_write_row_increment( + client, table, temp_rows, start, increment, expected +): + """ + test read_modify_write_row + """ + from google.cloud.bigtable.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=start, family=family, qualifier=qualifier) + + rule = IncrementRule(family, qualifier, increment) + result = await table.read_modify_write_row(row_key, rule) + assert result.row_key == row_key + assert len(result) == 1 + assert result[0].family == family + assert result[0].qualifier == qualifier + assert int(result[0]) == expected + # ensure that reading from server gives same value + assert (await _retrieve_cell_value(table, row_key)) == result[0].value + + +@pytest.mark.parametrize( + "start,append,expected", + [ + (b"", b"", b""), + ("", "", b""), + (b"abc", b"123", b"abc123"), + (b"abc", "123", b"abc123"), + ("", b"1", b"1"), + (b"abc", "", b"abc"), + (b"hello", b"world", b"helloworld"), + ], +) +@pytest.mark.asyncio +async def test_read_modify_write_row_append( + client, table, temp_rows, start, append, expected +): + """ + test read_modify_write_row + """ + from google.cloud.bigtable.read_modify_write_rules import AppendValueRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=start, family=family, qualifier=qualifier) + + rule = AppendValueRule(family, qualifier, append) + result = await table.read_modify_write_row(row_key, rule) + assert result.row_key == row_key + assert len(result) == 1 + assert result[0].family == family + assert result[0].qualifier == qualifier + assert result[0].value == expected + # ensure that reading from server gives same value + assert (await _retrieve_cell_value(table, row_key)) == result[0].value + + +@pytest.mark.asyncio +async def test_read_modify_write_row_chained(client, table, temp_rows): + """ + test read_modify_write_row with multiple rules + """ + from google.cloud.bigtable.read_modify_write_rules import AppendValueRule + from google.cloud.bigtable.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + start_amount = 1 + increment_amount = 10 + await temp_rows.add_row( + row_key, value=start_amount, family=family, qualifier=qualifier + ) + rule = [ + IncrementRule(family, qualifier, increment_amount), + AppendValueRule(family, qualifier, "hello"), + AppendValueRule(family, qualifier, "world"), + AppendValueRule(family, qualifier, "!"), + ] + result = await table.read_modify_write_row(row_key, rule) + assert result.row_key == row_key + assert result[0].family == family + assert result[0].qualifier == qualifier + # result should be a bytes number string for the IncrementRules, followed by the AppendValueRule values + assert ( + result[0].value + == (start_amount + increment_amount).to_bytes(8, "big", signed=True) + + b"helloworld!" + ) + # ensure that reading from server gives same value + assert (await _retrieve_cell_value(table, row_key)) == result[0].value + + +@pytest.mark.parametrize( + "start_val,predicate_range,expected_result", + [ + (1, (0, 2), True), + (-1, (0, 2), False), + ], +) +@pytest.mark.asyncio +async def test_check_and_mutate( + client, table, temp_rows, start_val, predicate_range, expected_result +): + """ + test that check_and_mutate_row works applies the right mutations, and returns the right result + """ + from google.cloud.bigtable.mutations import SetCell + from google.cloud.bigtable.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + + await temp_rows.add_row( + row_key, value=start_val, family=family, qualifier=qualifier + ) + + false_mutation_value = b"false-mutation-value" + false_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=false_mutation_value + ) + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + predicate = ValueRangeFilter(predicate_range[0], predicate_range[1]) + result = await table.check_and_mutate_row( + row_key, + predicate, + true_case_mutations=true_mutation, + false_case_mutations=false_mutation, + ) + assert result == expected_result + # ensure cell is updated + expected_value = true_mutation_value if expected_result else false_mutation_value + assert (await _retrieve_cell_value(table, row_key)) == expected_value + + @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_stream(table, temp_rows): diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 14da80dae..7009069d1 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -27,6 +27,9 @@ from google.api_core import exceptions as core_exceptions from google.cloud.bigtable.exceptions import InvalidChunk +from google.cloud.bigtable.read_modify_write_rules import IncrementRule +from google.cloud.bigtable.read_modify_write_rules import AppendValueRule + # try/except added for compatibility with python < 3.8 try: from unittest import mock @@ -2023,3 +2026,353 @@ async def test_bulk_mutate_row_metadata(self, include_app_profile): assert "app_profile_id=profile" in goog_metadata else: assert "app_profile_id=" not in goog_metadata + + +class TestCheckAndMutateRow: + def _make_client(self, *args, **kwargs): + from google.cloud.bigtable.client import BigtableDataClient + + return BigtableDataClient(*args, **kwargs) + + @pytest.mark.parametrize("gapic_result", [True, False]) + @pytest.mark.asyncio + async def test_check_and_mutate(self, gapic_result): + from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse + + app_profile = "app_profile_id" + async with self._make_client() as client: + async with client.get_table( + "instance", "table", app_profile_id=app_profile + ) as table: + with mock.patch.object( + client._gapic_client, "check_and_mutate_row" + ) as mock_gapic: + mock_gapic.return_value = CheckAndMutateRowResponse( + predicate_matched=gapic_result + ) + row_key = b"row_key" + predicate = None + true_mutations = [mock.Mock()] + false_mutations = [mock.Mock(), mock.Mock()] + operation_timeout = 0.2 + found = await table.check_and_mutate_row( + row_key, + predicate, + true_case_mutations=true_mutations, + false_case_mutations=false_mutations, + operation_timeout=operation_timeout, + ) + assert found == gapic_result + kwargs = mock_gapic.call_args[1] + request = kwargs["request"] + assert request["table_name"] == table.table_name + assert request["row_key"] == row_key + assert request["predicate_filter"] == predicate + assert request["true_mutations"] == [ + m._to_dict() for m in true_mutations + ] + assert request["false_mutations"] == [ + m._to_dict() for m in false_mutations + ] + assert request["app_profile_id"] == app_profile + assert kwargs["timeout"] == operation_timeout + + @pytest.mark.asyncio + async def test_check_and_mutate_bad_timeout(self): + """Should raise error if operation_timeout < 0""" + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with pytest.raises(ValueError) as e: + await table.check_and_mutate_row( + b"row_key", + None, + true_case_mutations=[mock.Mock()], + false_case_mutations=[], + operation_timeout=-1, + ) + assert str(e.value) == "operation_timeout must be greater than 0" + + @pytest.mark.asyncio + async def test_check_and_mutate_no_mutations(self): + """Requests require either true_case_mutations or false_case_mutations""" + from google.api_core.exceptions import InvalidArgument + + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with pytest.raises(InvalidArgument) as e: + await table.check_and_mutate_row( + b"row_key", + None, + true_case_mutations=None, + false_case_mutations=None, + ) + assert "No mutations provided" in str(e.value) + + @pytest.mark.asyncio + async def test_check_and_mutate_single_mutations(self): + """if single mutations are passed, they should be internally wrapped in a list""" + from google.cloud.bigtable.mutations import SetCell + from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse + + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "check_and_mutate_row" + ) as mock_gapic: + mock_gapic.return_value = CheckAndMutateRowResponse( + predicate_matched=True + ) + true_mutation = SetCell("family", b"qualifier", b"value") + false_mutation = SetCell("family", b"qualifier", b"value") + await table.check_and_mutate_row( + b"row_key", + None, + true_case_mutations=true_mutation, + false_case_mutations=false_mutation, + ) + kwargs = mock_gapic.call_args[1] + request = kwargs["request"] + assert request["true_mutations"] == [true_mutation._to_dict()] + assert request["false_mutations"] == [false_mutation._to_dict()] + + @pytest.mark.asyncio + async def test_check_and_mutate_predicate_object(self): + """predicate object should be converted to dict""" + from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse + + mock_predicate = mock.Mock() + fake_dict = {"fake": "dict"} + mock_predicate.to_dict.return_value = fake_dict + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "check_and_mutate_row" + ) as mock_gapic: + mock_gapic.return_value = CheckAndMutateRowResponse( + predicate_matched=True + ) + await table.check_and_mutate_row( + b"row_key", + mock_predicate, + false_case_mutations=[mock.Mock()], + ) + kwargs = mock_gapic.call_args[1] + assert kwargs["request"]["predicate_filter"] == fake_dict + assert mock_predicate.to_dict.call_count == 1 + + @pytest.mark.asyncio + async def test_check_and_mutate_mutations_parsing(self): + """mutations objects should be converted to dicts""" + from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse + from google.cloud.bigtable.mutations import DeleteAllFromRow + + mutations = [mock.Mock() for _ in range(5)] + for idx, mutation in enumerate(mutations): + mutation._to_dict.return_value = {"fake": idx} + mutations.append(DeleteAllFromRow()) + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "check_and_mutate_row" + ) as mock_gapic: + mock_gapic.return_value = CheckAndMutateRowResponse( + predicate_matched=True + ) + await table.check_and_mutate_row( + b"row_key", + None, + true_case_mutations=mutations[0:2], + false_case_mutations=mutations[2:], + ) + kwargs = mock_gapic.call_args[1]["request"] + assert kwargs["true_mutations"] == [{"fake": 0}, {"fake": 1}] + assert kwargs["false_mutations"] == [ + {"fake": 2}, + {"fake": 3}, + {"fake": 4}, + {"delete_from_row": {}}, + ] + assert all( + mutation._to_dict.call_count == 1 for mutation in mutations[:5] + ) + + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_check_and_mutate_metadata(self, include_app_profile): + """request should attach metadata headers""" + profile = "profile" if include_app_profile else None + async with self._make_client() as client: + async with client.get_table("i", "t", app_profile_id=profile) as table: + with mock.patch.object( + client._gapic_client, "check_and_mutate_row", AsyncMock() + ) as mock_gapic: + await table.check_and_mutate_row(b"key", mock.Mock()) + kwargs = mock_gapic.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata + + +class TestReadModifyWriteRow: + def _make_client(self, *args, **kwargs): + from google.cloud.bigtable.client import BigtableDataClient + + return BigtableDataClient(*args, **kwargs) + + @pytest.mark.parametrize( + "call_rules,expected_rules", + [ + ( + AppendValueRule("f", "c", b"1"), + [AppendValueRule("f", "c", b"1")._to_dict()], + ), + ( + [AppendValueRule("f", "c", b"1")], + [AppendValueRule("f", "c", b"1")._to_dict()], + ), + (IncrementRule("f", "c", 1), [IncrementRule("f", "c", 1)._to_dict()]), + ( + [AppendValueRule("f", "c", b"1"), IncrementRule("f", "c", 1)], + [ + AppendValueRule("f", "c", b"1")._to_dict(), + IncrementRule("f", "c", 1)._to_dict(), + ], + ), + ], + ) + @pytest.mark.asyncio + async def test_read_modify_write_call_rule_args(self, call_rules, expected_rules): + """ + Test that the gapic call is called with given rules + """ + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + client._gapic_client, "read_modify_write_row" + ) as mock_gapic: + await table.read_modify_write_row("key", call_rules) + assert mock_gapic.call_count == 1 + found_kwargs = mock_gapic.call_args_list[0][1] + assert found_kwargs["request"]["rules"] == expected_rules + + @pytest.mark.parametrize("rules", [[], None]) + @pytest.mark.asyncio + async def test_read_modify_write_no_rules(self, rules): + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with pytest.raises(ValueError) as e: + await table.read_modify_write_row("key", rules=rules) + assert e.value.args[0] == "rules must contain at least one item" + + @pytest.mark.asyncio + async def test_read_modify_write_call_defaults(self): + instance = "instance1" + table_id = "table1" + project = "project1" + row_key = "row_key1" + async with self._make_client(project=project) as client: + async with client.get_table(instance, table_id) as table: + with mock.patch.object( + client._gapic_client, "read_modify_write_row" + ) as mock_gapic: + await table.read_modify_write_row(row_key, mock.Mock()) + assert mock_gapic.call_count == 1 + found_kwargs = mock_gapic.call_args_list[0][1] + request = found_kwargs["request"] + assert ( + request["table_name"] + == f"projects/{project}/instances/{instance}/tables/{table_id}" + ) + assert request["app_profile_id"] is None + assert request["row_key"] == row_key.encode() + assert found_kwargs["timeout"] > 1 + + @pytest.mark.asyncio + async def test_read_modify_write_call_overrides(self): + row_key = b"row_key1" + expected_timeout = 12345 + profile_id = "profile1" + async with self._make_client() as client: + async with client.get_table( + "instance", "table_id", app_profile_id=profile_id + ) as table: + with mock.patch.object( + client._gapic_client, "read_modify_write_row" + ) as mock_gapic: + await table.read_modify_write_row( + row_key, + mock.Mock(), + operation_timeout=expected_timeout, + ) + assert mock_gapic.call_count == 1 + found_kwargs = mock_gapic.call_args_list[0][1] + request = found_kwargs["request"] + assert request["app_profile_id"] is profile_id + assert request["row_key"] == row_key + assert found_kwargs["timeout"] == expected_timeout + + @pytest.mark.asyncio + async def test_read_modify_write_string_key(self): + row_key = "string_row_key1" + async with self._make_client() as client: + async with client.get_table("instance", "table_id") as table: + with mock.patch.object( + client._gapic_client, "read_modify_write_row" + ) as mock_gapic: + await table.read_modify_write_row(row_key, mock.Mock()) + assert mock_gapic.call_count == 1 + found_kwargs = mock_gapic.call_args_list[0][1] + assert found_kwargs["request"]["row_key"] == row_key.encode() + + @pytest.mark.asyncio + async def test_read_modify_write_row_building(self): + """ + results from gapic call should be used to construct row + """ + from google.cloud.bigtable.row import Row + from google.cloud.bigtable_v2.types import ReadModifyWriteRowResponse + from google.cloud.bigtable_v2.types import Row as RowPB + + mock_response = ReadModifyWriteRowResponse(row=RowPB()) + async with self._make_client() as client: + async with client.get_table("instance", "table_id") as table: + with mock.patch.object( + client._gapic_client, "read_modify_write_row" + ) as mock_gapic: + with mock.patch.object(Row, "_from_pb") as constructor_mock: + mock_gapic.return_value = mock_response + await table.read_modify_write_row("key", mock.Mock()) + assert constructor_mock.call_count == 1 + constructor_mock.assert_called_once_with(mock_response.row) + + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_read_modify_write_metadata(self, include_app_profile): + """request should attach metadata headers""" + profile = "profile" if include_app_profile else None + async with self._make_client() as client: + async with client.get_table("i", "t", app_profile_id=profile) as table: + with mock.patch.object( + client._gapic_client, "read_modify_write_row", AsyncMock() + ) as mock_gapic: + await table.read_modify_write_row("key", mock.Mock()) + kwargs = mock_gapic.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata diff --git a/tests/unit/test_mutations.py b/tests/unit/test_mutations.py index 2a376609e..5730c53c9 100644 --- a/tests/unit/test_mutations.py +++ b/tests/unit/test_mutations.py @@ -170,6 +170,17 @@ def _target_class(self): def _make_one(self, *args, **kwargs): return self._target_class()(*args, **kwargs) + @pytest.mark.parametrize("input_val", [2**64, -(2**64)]) + def test_ctor_large_int(self, input_val): + with pytest.raises(ValueError) as e: + self._make_one(family="f", qualifier=b"b", new_value=input_val) + assert "int values must be between" in str(e.value) + + @pytest.mark.parametrize("input_val", ["", "a", "abc", "hello world!"]) + def test_ctor_str_value(self, input_val): + found = self._make_one(family="f", qualifier=b"b", new_value=input_val) + assert found.new_value == input_val.encode("utf-8") + def test_ctor(self): """Ensure constructor sets expected values""" expected_family = "test-family" @@ -194,6 +205,11 @@ def test_ctor_str_inputs(self): assert instance.qualifier == expected_qualifier assert instance.new_value == expected_value + @pytest.mark.parametrize("input_val", [-20, -1, 0, 1, 100, int(2**60)]) + def test_ctor_int_value(self, input_val): + found = self._make_one(family="f", qualifier=b"b", new_value=input_val) + assert found.new_value == input_val.to_bytes(8, "big", signed=True) + @pytest.mark.parametrize( "int_value,expected_bytes", [ @@ -206,7 +222,7 @@ def test_ctor_str_inputs(self): (100, b"\x00\x00\x00\x00\x00\x00\x00d"), ], ) - def test_ctor_int_value(self, int_value, expected_bytes): + def test_ctor_int_value_bytes(self, int_value, expected_bytes): """Test with int value""" expected_family = "test-family" expected_qualifier = b"test-qualifier" diff --git a/tests/unit/test_read_modify_write_rules.py b/tests/unit/test_read_modify_write_rules.py new file mode 100644 index 000000000..02240df6d --- /dev/null +++ b/tests/unit/test_read_modify_write_rules.py @@ -0,0 +1,142 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + + +import pytest + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock +except ImportError: # pragma: NO COVER + import mock # type: ignore + + +class TestBaseReadModifyWriteRule: + def _target_class(self): + from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule + + return ReadModifyWriteRule + + def test_abstract(self): + """should not be able to instantiate""" + with pytest.raises(TypeError): + self._target_class()(family="foo", qualifier=b"bar") + + def test__to_dict(self): + with pytest.raises(NotImplementedError): + self._target_class()._to_dict(mock.Mock()) + + +class TestIncrementRule: + def _target_class(self): + from google.cloud.bigtable.read_modify_write_rules import IncrementRule + + return IncrementRule + + @pytest.mark.parametrize( + "args,expected", + [ + (("fam", b"qual", 1), ("fam", b"qual", 1)), + (("fam", b"qual", -12), ("fam", b"qual", -12)), + (("fam", "qual", 1), ("fam", b"qual", 1)), + (("fam", "qual", 0), ("fam", b"qual", 0)), + (("", "", 0), ("", b"", 0)), + (("f", b"q"), ("f", b"q", 1)), + ], + ) + def test_ctor(self, args, expected): + instance = self._target_class()(*args) + assert instance.family == expected[0] + assert instance.qualifier == expected[1] + assert instance.increment_amount == expected[2] + + @pytest.mark.parametrize("input_amount", [1.1, None, "1", object(), "", b"", b"1"]) + def test_ctor_bad_input(self, input_amount): + with pytest.raises(TypeError) as e: + self._target_class()("fam", b"qual", input_amount) + assert "increment_amount must be an integer" in str(e.value) + + @pytest.mark.parametrize( + "large_value", [2**64, 2**64 + 1, -(2**64), -(2**64) - 1] + ) + def test_ctor_large_values(self, large_value): + with pytest.raises(ValueError) as e: + self._target_class()("fam", b"qual", large_value) + assert "too large" in str(e.value) + + @pytest.mark.parametrize( + "args,expected", + [ + (("fam", b"qual", 1), ("fam", b"qual", 1)), + (("fam", b"qual", -12), ("fam", b"qual", -12)), + (("fam", "qual", 1), ("fam", b"qual", 1)), + (("fam", "qual", 0), ("fam", b"qual", 0)), + (("", "", 0), ("", b"", 0)), + (("f", b"q"), ("f", b"q", 1)), + ], + ) + def test__to_dict(self, args, expected): + instance = self._target_class()(*args) + expected = { + "family_name": expected[0], + "column_qualifier": expected[1], + "increment_amount": expected[2], + } + assert instance._to_dict() == expected + + +class TestAppendValueRule: + def _target_class(self): + from google.cloud.bigtable.read_modify_write_rules import AppendValueRule + + return AppendValueRule + + @pytest.mark.parametrize( + "args,expected", + [ + (("fam", b"qual", b"val"), ("fam", b"qual", b"val")), + (("fam", "qual", b"val"), ("fam", b"qual", b"val")), + (("", "", b""), ("", b"", b"")), + (("f", "q", "str_val"), ("f", b"q", b"str_val")), + (("f", "q", ""), ("f", b"q", b"")), + ], + ) + def test_ctor(self, args, expected): + instance = self._target_class()(*args) + assert instance.family == expected[0] + assert instance.qualifier == expected[1] + assert instance.append_value == expected[2] + + @pytest.mark.parametrize("input_val", [5, 1.1, None, object()]) + def test_ctor_bad_input(self, input_val): + with pytest.raises(TypeError) as e: + self._target_class()("fam", b"qual", input_val) + assert "append_value must be bytes or str" in str(e.value) + + @pytest.mark.parametrize( + "args,expected", + [ + (("fam", b"qual", b"val"), ("fam", b"qual", b"val")), + (("fam", "qual", b"val"), ("fam", b"qual", b"val")), + (("", "", b""), ("", b"", b"")), + ], + ) + def test__to_dict(self, args, expected): + instance = self._target_class()(*args) + expected = { + "family_name": expected[0], + "column_qualifier": expected[1], + "append_value": expected[2], + } + assert instance._to_dict() == expected diff --git a/tests/unit/test_row.py b/tests/unit/test_row.py index 1af09aad9..0413b2889 100644 --- a/tests/unit/test_row.py +++ b/tests/unit/test_row.py @@ -55,6 +55,50 @@ def test_ctor(self): self.assertEqual(list(row_response), cells) self.assertEqual(row_response.row_key, TEST_ROW_KEY) + def test__from_pb(self): + """ + Construct from protobuf. + """ + from google.cloud.bigtable_v2.types import Row as RowPB + from google.cloud.bigtable_v2.types import Family as FamilyPB + from google.cloud.bigtable_v2.types import Column as ColumnPB + from google.cloud.bigtable_v2.types import Cell as CellPB + + row_key = b"row_key" + cells = [ + CellPB( + value=str(i).encode(), + timestamp_micros=TEST_TIMESTAMP, + labels=TEST_LABELS, + ) + for i in range(2) + ] + column = ColumnPB(qualifier=TEST_QUALIFIER, cells=cells) + families_pb = [FamilyPB(name=TEST_FAMILY_ID, columns=[column])] + row_pb = RowPB(key=row_key, families=families_pb) + output = self._get_target_class()._from_pb(row_pb) + self.assertEqual(output.row_key, row_key) + self.assertEqual(len(output), 2) + self.assertEqual(output[0].value, b"0") + self.assertEqual(output[1].value, b"1") + self.assertEqual(output[0].timestamp_micros, TEST_TIMESTAMP) + self.assertEqual(output[0].labels, TEST_LABELS) + assert output[0].row_key == row_key + assert output[0].family == TEST_FAMILY_ID + assert output[0].qualifier == TEST_QUALIFIER + + def test__from_pb_sparse(self): + """ + Construct from minimal protobuf. + """ + from google.cloud.bigtable_v2.types import Row as RowPB + + row_key = b"row_key" + row_pb = RowPB(key=row_key) + output = self._get_target_class()._from_pb(row_pb) + self.assertEqual(output.row_key, row_key) + self.assertEqual(len(output), 0) + def test_get_cells(self): cell_list = [] for family_id in ["1", "2"]: From ec2b9835d3fd050414755b30a769a769d74f802a Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 23 Jun 2023 11:29:15 -0700 Subject: [PATCH 11/56] feat: sharded read rows (#766) --- google/cloud/bigtable/__init__.py | 2 + google/cloud/bigtable/client.py | 163 +++++- google/cloud/bigtable/exceptions.py | 50 +- google/cloud/bigtable/read_rows_query.py | 210 ++++++- tests/system/test_system.py | 104 ++++ tests/unit/test_client.py | 620 ++++++++++++++++---- tests/unit/test_exceptions.py | 75 ++- tests/unit/test_read_rows_query.py | 686 ++++++++++++++++++----- 8 files changed, 1605 insertions(+), 305 deletions(-) diff --git a/google/cloud/bigtable/__init__.py b/google/cloud/bigtable/__init__.py index 70c87ade1..06b45bc4d 100644 --- a/google/cloud/bigtable/__init__.py +++ b/google/cloud/bigtable/__init__.py @@ -36,6 +36,8 @@ # Type alias for the output of sample_keys RowKeySamples = List[Tuple[bytes, int]] +# type alias for the output of query.shard() +ShardedQuery = List[ReadRowsQuery] __version__: str = package_version.__version__ diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index f75613098..055263cfa 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -38,6 +38,7 @@ from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( PooledBigtableGrpcAsyncIOTransport, ) +from google.cloud.bigtable_v2.types.bigtable import PingAndWarmRequest from google.cloud.client import ClientWithProject from google.api_core.exceptions import GoogleAPICallError from google.api_core import retry_async as retries @@ -50,10 +51,14 @@ from google.cloud.bigtable.row import Row from google.cloud.bigtable.read_rows_query import ReadRowsQuery from google.cloud.bigtable.iterators import ReadRowsIterator +from google.cloud.bigtable.exceptions import FailedQueryShardError +from google.cloud.bigtable.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.mutations import Mutation, RowMutationEntry from google.cloud.bigtable._mutate_rows import _MutateRowsOperation from google.cloud.bigtable._helpers import _make_metadata from google.cloud.bigtable._helpers import _convert_retry_deadline +from google.cloud.bigtable._helpers import _attempt_timeout_generator from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule from google.cloud.bigtable.row_filters import RowFilter @@ -64,6 +69,10 @@ if TYPE_CHECKING: from google.cloud.bigtable.mutations_batcher import MutationsBatcher from google.cloud.bigtable import RowKeySamples + from google.cloud.bigtable import ShardedQuery + +# used by read_rows_sharded to limit how many requests are attempted in parallel +CONCURRENCY_LIMIT = 10 class BigtableDataClient(ClientWithProject): @@ -190,10 +199,13 @@ async def _ping_and_warm_instances( - sequence of results or exceptions from the ping requests """ ping_rpc = channel.unary_unary( - "/google.bigtable.v2.Bigtable/PingAndWarmChannel" + "/google.bigtable.v2.Bigtable/PingAndWarm", + request_serializer=PingAndWarmRequest.serialize, ) tasks = [ping_rpc({"name": n}) for n in self._active_instances] - return await asyncio.gather(*tasks, return_exceptions=True) + result = await asyncio.gather(*tasks, return_exceptions=True) + # return None in place of empty successful responses + return [r or None for r in result] async def _manage_channel( self, @@ -532,22 +544,79 @@ async def read_row( async def read_rows_sharded( self, - query_list: list[ReadRowsQuery] | list[dict[str, Any]], + sharded_query: ShardedQuery, *, - limit: int | None, - operation_timeout: int | float | None = 60, + operation_timeout: int | float | None = None, per_request_timeout: int | float | None = None, - ) -> ReadRowsIterator: + ) -> list[Row]: """ - Runs a sharded query in parallel + Runs a sharded query in parallel, then return the results in a single list. + Results will be returned in the order of the input queries. + + This function is intended to be run on the results on a query.shard() call: - Each query in query list will be run concurrently, with results yielded as they are ready - yielded results may be out of order + ``` + table_shard_keys = await table.sample_row_keys() + query = ReadRowsQuery(...) + shard_queries = query.shard(table_shard_keys) + results = await table.read_rows_sharded(shard_queries) + ``` Args: - - query_list: a list of queries to run in parallel + - sharded_query: a sharded query to execute + Raises: + - ShardedReadRowsExceptionGroup: if any of the queries failed + - ValueError: if the query_list is empty """ - raise NotImplementedError + if not sharded_query: + raise ValueError("empty sharded_query") + # reduce operation_timeout between batches + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = ( + per_request_timeout or self.default_per_request_timeout or operation_timeout + ) + timeout_generator = _attempt_timeout_generator( + operation_timeout, operation_timeout + ) + # submit shards in batches if the number of shards goes over CONCURRENCY_LIMIT + batched_queries = [ + sharded_query[i : i + CONCURRENCY_LIMIT] + for i in range(0, len(sharded_query), CONCURRENCY_LIMIT) + ] + # run batches and collect results + results_list = [] + error_dict = {} + shard_idx = 0 + for batch in batched_queries: + batch_operation_timeout = next(timeout_generator) + routine_list = [ + self.read_rows( + query, + operation_timeout=batch_operation_timeout, + per_request_timeout=min( + per_request_timeout, batch_operation_timeout + ), + ) + for query in batch + ] + batch_result = await asyncio.gather(*routine_list, return_exceptions=True) + for result in batch_result: + if isinstance(result, Exception): + error_dict[shard_idx] = result + else: + results_list.extend(result) + shard_idx += 1 + if error_dict: + # if any sub-request failed, raise an exception instead of returning results + raise ShardedReadRowsExceptionGroup( + [ + FailedQueryShardError(idx, sharded_query[idx], e) + for idx, e in error_dict.items() + ], + results_list, + len(sharded_query), + ) + return results_list async def row_exists( self, @@ -577,12 +646,11 @@ async def row_exists( ) return len(results) > 0 - async def sample_keys( + async def sample_row_keys( self, *, - operation_timeout: int | float | None = 60, - per_sample_timeout: int | float | None = 10, - per_request_timeout: int | float | None = None, + operation_timeout: float | None = None, + per_request_timeout: float | None = None, ) -> RowKeySamples: """ Return a set of RowKeySamples that delimit contiguous sections of the table of @@ -590,7 +658,7 @@ async def sample_keys( RowKeySamples output can be used with ReadRowsQuery.shard() to create a sharded query that can be parallelized across multiple backend nodes read_rows and read_rows_stream - requests will call sample_keys internally for this purpose when sharding is enabled + requests will call sample_row_keys internally for this purpose when sharding is enabled RowKeySamples is simply a type alias for list[tuple[bytes, int]]; a list of row_keys, along with offset positions in the table @@ -598,11 +666,61 @@ async def sample_keys( Returns: - a set of RowKeySamples the delimit contiguous sections of the table Raises: - - DeadlineExceeded: raised after operation timeout - will be chained with a RetryExceptionGroup containing all GoogleAPIError - exceptions from any retries that failed + - GoogleAPICallError: if the sample_row_keys request fails """ - raise NotImplementedError + # prepare timeouts + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = per_request_timeout or self.default_per_request_timeout + + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout <= 0: + raise ValueError("per_request_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout > operation_timeout: + raise ValueError( + "per_request_timeout must not be greater than operation_timeout" + ) + attempt_timeout_gen = _attempt_timeout_generator( + per_request_timeout, operation_timeout + ) + # prepare retryable + predicate = retries.if_exception_type( + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ) + transient_errors = [] + + def on_error_fn(exc): + # add errors to list if retryable + if predicate(exc): + transient_errors.append(exc) + + retry = retries.AsyncRetry( + predicate=predicate, + timeout=operation_timeout, + initial=0.01, + multiplier=2, + maximum=60, + on_error=on_error_fn, + is_stream=False, + ) + + # prepare request + metadata = _make_metadata(self.table_name, self.app_profile_id) + + async def execute_rpc(): + results = await self.client._gapic_client.sample_row_keys( + table_name=self.table_name, + app_profile_id=self.app_profile_id, + timeout=next(attempt_timeout_gen), + metadata=metadata, + ) + return [(s.row_key, s.offset_bytes) async for s in results] + + wrapped_fn = _convert_retry_deadline( + retry(execute_rpc), operation_timeout, transient_errors + ) + return await wrapped_fn() def mutations_batcher(self, **kwargs) -> MutationsBatcher: """ @@ -896,16 +1014,17 @@ async def close(self): """ Called to close the Table instance and release any resources held by it. """ + self._register_instance_task.cancel() await self.client._remove_instance_registration(self.instance_id, self) async def __aenter__(self): """ Implement async context manager protocol - Register this instance with the client, so that + Ensure registration task has time to run, so that grpc channels will be warmed for the specified instance """ - await self.client._register_instance(self.instance_id, self) + await self._register_instance_task return self async def __aexit__(self, exc_type, exc_val, exc_tb): diff --git a/google/cloud/bigtable/exceptions.py b/google/cloud/bigtable/exceptions.py index fe3bec7e9..b2cf0ce6b 100644 --- a/google/cloud/bigtable/exceptions.py +++ b/google/cloud/bigtable/exceptions.py @@ -16,14 +16,16 @@ import sys -from typing import TYPE_CHECKING +from typing import Any, TYPE_CHECKING from google.api_core import exceptions as core_exceptions +from google.cloud.bigtable.row import Row is_311_plus = sys.version_info >= (3, 11) if TYPE_CHECKING: from google.cloud.bigtable.mutations import RowMutationEntry + from google.cloud.bigtable.read_rows_query import ReadRowsQuery class IdleTimeout(core_exceptions.DeadlineExceeded): @@ -137,3 +139,49 @@ def __init__(self, excs: list[Exception]): def __new__(cls, excs: list[Exception]): return super().__new__(cls, cls._format_message(excs), excs) + + +class ShardedReadRowsExceptionGroup(BigtableExceptionGroup): + """ + Represents one or more exceptions that occur during a sharded read rows operation + """ + + @staticmethod + def _format_message(excs: list[FailedQueryShardError], total_queries: int): + query_str = "query" if total_queries == 1 else "queries" + plural_str = "" if len(excs) == 1 else "s" + return f"{len(excs)} sub-exception{plural_str} (from {total_queries} {query_str} attempted)" + + def __init__( + self, + excs: list[FailedQueryShardError], + succeeded: list[Row], + total_queries: int, + ): + super().__init__(self._format_message(excs, total_queries), excs) + self.successful_rows = succeeded + + def __new__( + cls, excs: list[FailedQueryShardError], succeeded: list[Row], total_queries: int + ): + instance = super().__new__(cls, cls._format_message(excs, total_queries), excs) + instance.successful_rows = succeeded + return instance + + +class FailedQueryShardError(Exception): + """ + Represents an individual failed query in a sharded read rows operation + """ + + def __init__( + self, + failed_index: int, + failed_query: "ReadRowsQuery" | dict[str, Any], + cause: Exception, + ): + message = f"Failed query at index {failed_index} with cause: {cause!r}" + super().__init__(message) + self.index = failed_index + self.query = failed_query + self.__cause__ = cause diff --git a/google/cloud/bigtable/read_rows_query.py b/google/cloud/bigtable/read_rows_query.py index 6de84e918..eb28eeda3 100644 --- a/google/cloud/bigtable/read_rows_query.py +++ b/google/cloud/bigtable/read_rows_query.py @@ -14,11 +14,15 @@ # from __future__ import annotations from typing import TYPE_CHECKING, Any +from bisect import bisect_left +from bisect import bisect_right +from collections import defaultdict from dataclasses import dataclass from google.cloud.bigtable.row_filters import RowFilter if TYPE_CHECKING: from google.cloud.bigtable import RowKeySamples + from google.cloud.bigtable import ShardedQuery @dataclass @@ -28,6 +32,9 @@ class _RangePoint: key: bytes is_inclusive: bool + def __hash__(self) -> int: + return hash((self.key, self.is_inclusive)) + @dataclass class RowRange: @@ -80,6 +87,44 @@ def _to_dict(self) -> dict[str, bytes]: output[key] = self.end.key return output + def __hash__(self) -> int: + return hash((self.start, self.end)) + + @classmethod + def _from_dict(cls, data: dict[str, bytes]) -> RowRange: + """Creates a RowRange from a dictionary""" + start_key = data.get("start_key_closed", data.get("start_key_open")) + end_key = data.get("end_key_closed", data.get("end_key_open")) + start_is_inclusive = "start_key_closed" in data if start_key else None + end_is_inclusive = "end_key_closed" in data if end_key else None + return cls( + start_key, + end_key, + start_is_inclusive, + end_is_inclusive, + ) + + @classmethod + def _from_points( + cls, start: _RangePoint | None, end: _RangePoint | None + ) -> RowRange: + """Creates a RowRange from two RangePoints""" + kwargs: dict[str, Any] = {} + if start is not None: + kwargs["start_key"] = start.key + kwargs["start_is_inclusive"] = start.is_inclusive + if end is not None: + kwargs["end_key"] = end.key + kwargs["end_is_inclusive"] = end.is_inclusive + return cls(**kwargs) + + def __bool__(self) -> bool: + """ + Empty RowRanges (representing a full table scan) are falsy, because + they can be substituted with None. Non-empty RowRanges are truthy. + """ + return self.start is not None or self.end is not None + class ReadRowsQuery: """ @@ -91,7 +136,7 @@ def __init__( row_keys: list[str | bytes] | str | bytes | None = None, row_ranges: list[RowRange] | RowRange | None = None, limit: int | None = None, - row_filter: RowFilter | None = None, + row_filter: RowFilter | dict[str, Any] | None = None, ): """ Create a new ReadRowsQuery @@ -105,7 +150,7 @@ def __init__( - row_filter: a RowFilter to apply to the query """ self.row_keys: set[bytes] = set() - self.row_ranges: list[RowRange | dict[str, bytes]] = [] + self.row_ranges: set[RowRange] = set() if row_ranges is not None: if isinstance(row_ranges, RowRange): row_ranges = [row_ranges] @@ -197,18 +242,133 @@ def add_range( """ if not (isinstance(row_range, dict) or isinstance(row_range, RowRange)): raise ValueError("row_range must be a RowRange or dict") - self.row_ranges.append(row_range) + if isinstance(row_range, dict): + row_range = RowRange._from_dict(row_range) + self.row_ranges.add(row_range) - def shard(self, shard_keys: "RowKeySamples" | None = None) -> list[ReadRowsQuery]: + def shard(self, shard_keys: RowKeySamples) -> ShardedQuery: """ Split this query into multiple queries that can be evenly distributed - across nodes and be run in parallel + across nodes and run in parallel + + Returns: + - a ShardedQuery that can be used in sharded_read_rows calls + Raises: + - AttributeError if the query contains a limit + """ + if self.limit is not None: + raise AttributeError("Cannot shard query with a limit") + if len(self.row_keys) == 0 and len(self.row_ranges) == 0: + # empty query represents full scan + # ensure that we have at least one key or range + full_scan_query = ReadRowsQuery( + row_ranges=RowRange(), row_filter=self.filter + ) + return full_scan_query.shard(shard_keys) + sharded_queries: dict[int, ReadRowsQuery] = defaultdict( + lambda: ReadRowsQuery(row_filter=self.filter) + ) + # the split_points divde our key space into segments + # each split_point defines last key that belongs to a segment + # our goal is to break up the query into subqueries that each operate in a single segment + split_points = [sample[0] for sample in shard_keys if sample[0]] + + # handle row_keys + # use binary search to find the segment that each key belongs to + for this_key in list(self.row_keys): + # bisect_left: in case of exact match, pick left side (keys are inclusive ends) + segment_index = bisect_left(split_points, this_key) + sharded_queries[segment_index].add_key(this_key) + + # handle row_ranges + for this_range in self.row_ranges: + # defer to _shard_range helper + for segment_index, added_range in self._shard_range( + this_range, split_points + ): + sharded_queries[segment_index].add_range(added_range) + # return list of queries ordered by segment index + # pull populated segments out of sharded_queries dict + keys = sorted(list(sharded_queries.keys())) + # return list of queries + return [sharded_queries[k] for k in keys] + + @staticmethod + def _shard_range( + orig_range: RowRange, split_points: list[bytes] + ) -> list[tuple[int, RowRange]]: + """ + Helper function for sharding row_range into subranges that fit into + segments of the key-space, determined by split_points + + Args: + - orig_range: a row range to split + - split_points: a list of row keys that define the boundaries of segments. + each point represents the inclusive end of a segment Returns: - - a list of queries that represent a sharded version of the original - query (if possible) + - a list of tuples, containing a segment index and a new sub-range. """ - raise NotImplementedError + # 1. find the index of the segment the start key belongs to + if orig_range.start is None: + # if range is open on the left, include first segment + start_segment = 0 + else: + # use binary search to find the segment the start key belongs to + # bisect method determines how we break ties when the start key matches a split point + # if inclusive, bisect_left to the left segment, otherwise bisect_right + bisect = bisect_left if orig_range.start.is_inclusive else bisect_right + start_segment = bisect(split_points, orig_range.start.key) + + # 2. find the index of the segment the end key belongs to + if orig_range.end is None: + # if range is open on the right, include final segment + end_segment = len(split_points) + else: + # use binary search to find the segment the end key belongs to. + end_segment = bisect_left( + split_points, orig_range.end.key, lo=start_segment + ) + # note: end_segment will always bisect_left, because split points represent inclusive ends + # whether the end_key is includes the split point or not, the result is the same segment + # 3. create new range definitions for each segment this_range spans + if start_segment == end_segment: + # this_range is contained in a single segment. + # Add this_range to that segment's query only + return [(start_segment, orig_range)] + else: + results: list[tuple[int, RowRange]] = [] + # this_range spans multiple segments. Create a new range for each segment's query + # 3a. add new range for first segment this_range spans + # first range spans from start_key to the split_point representing the last key in the segment + last_key_in_first_segment = split_points[start_segment] + start_range = RowRange._from_points( + start=orig_range.start, + end=_RangePoint(last_key_in_first_segment, is_inclusive=True), + ) + results.append((start_segment, start_range)) + # 3b. add new range for last segment this_range spans + # we start the final range using the end key from of the previous segment, with is_inclusive=False + previous_segment = end_segment - 1 + last_key_before_segment = split_points[previous_segment] + end_range = RowRange._from_points( + start=_RangePoint(last_key_before_segment, is_inclusive=False), + end=orig_range.end, + ) + results.append((end_segment, end_range)) + # 3c. add new spanning range to all segments other than the first and last + for this_segment in range(start_segment + 1, end_segment): + prev_segment = this_segment - 1 + prev_end_key = split_points[prev_segment] + this_end_key = split_points[prev_segment + 1] + new_range = RowRange( + start_key=prev_end_key, + start_is_inclusive=False, + end_key=this_end_key, + end_is_inclusive=True, + ) + results.append((this_segment, new_range)) + return results def _to_dict(self) -> dict[str, Any]: """ @@ -221,11 +381,7 @@ def _to_dict(self) -> dict[str, Any]: row_ranges.append(dict_range) row_keys = list(self.row_keys) row_keys.sort() - row_set: dict[str, Any] = {} - if row_keys: - row_set["row_keys"] = row_keys - if row_ranges: - row_set["row_ranges"] = row_ranges + row_set = {"row_keys": row_keys, "row_ranges": row_ranges} final_dict: dict[str, Any] = { "rows": row_set, } @@ -237,3 +393,31 @@ def _to_dict(self) -> dict[str, Any]: if self.limit is not None: final_dict["rows_limit"] = self.limit return final_dict + + def __eq__(self, other): + """ + RowRanges are equal if they have the same row keys, row ranges, + filter and limit, or if they both represent a full scan with the + same filter and limit + """ + if not isinstance(other, ReadRowsQuery): + return False + # empty queries are equal + if len(self.row_keys) == 0 and len(other.row_keys) == 0: + this_range_empty = len(self.row_ranges) == 0 or all( + [bool(r) is False for r in self.row_ranges] + ) + other_range_empty = len(other.row_ranges) == 0 or all( + [bool(r) is False for r in other.row_ranges] + ) + if this_range_empty and other_range_empty: + return self.filter == other.filter and self.limit == other.limit + return ( + self.row_keys == other.row_keys + and self.row_ranges == other.row_ranges + and self.filter == other.filter + and self.limit == other.limit + ) + + def __repr__(self): + return f"ReadRowsQuery(row_keys={list(self.row_keys)}, row_ranges={list(self.row_ranges)}, row_filter={self.filter}, limit={self.limit})" diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 692911b10..1ba022ae0 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -225,6 +225,21 @@ async def test_ping_and_warm_gapic(client, table): @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_ping_and_warm(client, table): + """ + Test ping and warm from handwritten client + """ + try: + channel = client.transport._grpc_channel.pool[0] + except Exception: + # for sync client + channel = client.transport._grpc_channel + results = await client._ping_and_warm_instances(channel) + assert len(results) == 1 + assert results[0] is None + + @pytest.mark.asyncio async def test_mutation_set_cell(table, temp_rows): """ @@ -254,6 +269,21 @@ async def test_mutation_set_cell(table, temp_rows): @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_sample_row_keys(client, table, temp_rows): + """ + Sample keys should return a single sample in small test tables + """ + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + + results = await table.sample_row_keys() + assert len(results) == 1 + sample = results[0] + assert isinstance(sample[0], bytes) + assert isinstance(sample[1], int) + + @pytest.mark.asyncio async def test_bulk_mutations_set_cell(client, table, temp_rows): """ @@ -476,6 +506,80 @@ async def test_read_rows(table, temp_rows): assert row_list[1].row_key == b"row_key_2" +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows_sharded_simple(table, temp_rows): + """ + Test read rows sharded with two queries + """ + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) + query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) + row_list = await table.read_rows_sharded([query1, query2]) + assert len(row_list) == 4 + assert row_list[0].row_key == b"a" + assert row_list[1].row_key == b"c" + assert row_list[2].row_key == b"b" + assert row_list[3].row_key == b"d" + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows_sharded_from_sample(table, temp_rows): + """ + Test end-to-end sharding + """ + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.read_rows_query import RowRange + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + + table_shard_keys = await table.sample_row_keys() + query = ReadRowsQuery(row_ranges=[RowRange(start_key=b"b", end_key=b"z")]) + shard_queries = query.shard(table_shard_keys) + row_list = await table.read_rows_sharded(shard_queries) + assert len(row_list) == 3 + assert row_list[0].row_key == b"b" + assert row_list[1].row_key == b"c" + assert row_list[2].row_key == b"d" + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_read_rows_sharded_filters_limits(table, temp_rows): + """ + Test read rows sharded with filters and limits + """ + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.row_filters import ApplyLabelFilter + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + + label_filter1 = ApplyLabelFilter("first") + label_filter2 = ApplyLabelFilter("second") + query1 = ReadRowsQuery(row_keys=[b"a", b"c"], limit=1, row_filter=label_filter1) + query2 = ReadRowsQuery(row_keys=[b"b", b"d"], row_filter=label_filter2) + row_list = await table.read_rows_sharded([query1, query2]) + assert len(row_list) == 3 + assert row_list[0].row_key == b"a" + assert row_list[1].row_key == b"b" + assert row_list[2].row_key == b"d" + assert row_list[0][0].labels == ["first"] + assert row_list[1][0].labels == ["second"] + assert row_list[2][0].labels == ["second"] + + @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_range_query(table, temp_rows): diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 7009069d1..706ab973d 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -882,7 +882,8 @@ def _make_stats(self): ) ) - def _make_chunk(self, *args, **kwargs): + @staticmethod + def _make_chunk(*args, **kwargs): from google.cloud.bigtable_v2 import ReadRowsResponse kwargs["row_key"] = kwargs.get("row_key", b"row_key") @@ -893,8 +894,8 @@ def _make_chunk(self, *args, **kwargs): return ReadRowsResponse.CellChunk(*args, **kwargs) + @staticmethod async def _make_gapic_stream( - self, chunk_list: list[ReadRowsResponse.CellChunk | Exception], sleep_time=0, ): @@ -926,6 +927,9 @@ def cancel(self): return mock_stream(chunk_list, sleep_time) + async def execute_fn(self, table, *args, **kwargs): + return await table.read_rows(*args, **kwargs) + @pytest.mark.asyncio async def test_read_rows(self): client = self._make_client() @@ -939,7 +943,7 @@ async def test_read_rows(self): read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( chunks ) - results = await table.read_rows(query, operation_timeout=3) + results = await self.execute_fn(table, query, operation_timeout=3) assert len(results) == 2 assert results[0].row_key == b"test_1" assert results[1].row_key == b"test_2" @@ -989,7 +993,7 @@ async def test_read_rows_query_matches_request(self, include_app_profile): read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( [] ) - results = await table.read_rows(query, operation_timeout=3) + results = await self.execute_fn(table, query, operation_timeout=3) assert len(results) == 0 call_request = read_rows.call_args_list[0][0][0] query_dict = query._to_dict() @@ -1013,22 +1017,26 @@ async def test_read_rows_query_matches_request(self, include_app_profile): @pytest.mark.asyncio async def test_read_rows_timeout(self, operation_timeout): async with self._make_client() as client: - table = client.get_table("instance", "table") - query = ReadRowsQuery() - chunks = [self._make_chunk(row_key=b"test_1")] - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( - chunks, sleep_time=1 - ) - try: - await table.read_rows(query, operation_timeout=operation_timeout) - except core_exceptions.DeadlineExceeded as e: - assert ( - e.message - == f"operation_timeout of {operation_timeout:0.1f}s exceeded" + async with client.get_table("instance", "table") as table: + query = ReadRowsQuery() + chunks = [self._make_chunk(row_key=b"test_1")] + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = ( + lambda *args, **kwargs: self._make_gapic_stream( + chunks, sleep_time=1 + ) ) + try: + await self.execute_fn( + table, query, operation_timeout=operation_timeout + ) + except core_exceptions.DeadlineExceeded as e: + assert ( + e.message + == f"operation_timeout of {operation_timeout:0.1f}s exceeded" + ) @pytest.mark.parametrize( "per_request_t, operation_t, expected_num", @@ -1057,45 +1065,48 @@ async def test_read_rows_per_request_timeout( # mocking uniform ensures there are no sleeps between retries with mock.patch("random.uniform", side_effect=lambda a, b: 0): async with self._make_client() as client: - table = client.get_table("instance", "table") - query = ReadRowsQuery() - chunks = [core_exceptions.DeadlineExceeded("mock deadline")] - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = ( - lambda *args, **kwargs: self._make_gapic_stream( - chunks, sleep_time=per_request_t - ) - ) - try: - await table.read_rows( - query, - operation_timeout=operation_t, - per_request_timeout=per_request_t, + async with client.get_table("instance", "table") as table: + query = ReadRowsQuery() + chunks = [core_exceptions.DeadlineExceeded("mock deadline")] + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = ( + lambda *args, **kwargs: self._make_gapic_stream( + chunks, sleep_time=per_request_t + ) ) - except core_exceptions.DeadlineExceeded as e: - retry_exc = e.__cause__ - if expected_num == 0: - assert retry_exc is None - else: - assert type(retry_exc) == RetryExceptionGroup - assert f"{expected_num} failed attempts" in str(retry_exc) - assert len(retry_exc.exceptions) == expected_num - for sub_exc in retry_exc.exceptions: - assert sub_exc.message == "mock deadline" - assert read_rows.call_count == expected_num - # check timeouts - for _, call_kwargs in read_rows.call_args_list[:-1]: - assert call_kwargs["timeout"] == per_request_t - # last timeout should be adjusted to account for the time spent - assert ( - abs( - read_rows.call_args_list[-1][1]["timeout"] - - expected_last_timeout + try: + await self.execute_fn( + table, + query, + operation_timeout=operation_t, + per_request_timeout=per_request_t, + ) + except core_exceptions.DeadlineExceeded as e: + retry_exc = e.__cause__ + if expected_num == 0: + assert retry_exc is None + else: + assert type(retry_exc) == RetryExceptionGroup + assert f"{expected_num} failed attempts" in str( + retry_exc + ) + assert len(retry_exc.exceptions) == expected_num + for sub_exc in retry_exc.exceptions: + assert sub_exc.message == "mock deadline" + assert read_rows.call_count == expected_num + # check timeouts + for _, call_kwargs in read_rows.call_args_list[:-1]: + assert call_kwargs["timeout"] == per_request_t + # last timeout should be adjusted to account for the time spent + assert ( + abs( + read_rows.call_args_list[-1][1]["timeout"] + - expected_last_timeout + ) + < 0.05 ) - < 0.05 - ) @pytest.mark.asyncio async def test_read_rows_idle_timeout(self): @@ -1155,22 +1166,24 @@ async def test_read_rows_idle_timeout(self): @pytest.mark.asyncio async def test_read_rows_retryable_error(self, exc_type): async with self._make_client() as client: - table = client.get_table("instance", "table") - query = ReadRowsQuery() - expected_error = exc_type("mock error") - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( - [expected_error] - ) - try: - await table.read_rows(query, operation_timeout=0.1) - except core_exceptions.DeadlineExceeded as e: - retry_exc = e.__cause__ - root_cause = retry_exc.exceptions[0] - assert type(root_cause) == exc_type - assert root_cause == expected_error + async with client.get_table("instance", "table") as table: + query = ReadRowsQuery() + expected_error = exc_type("mock error") + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = ( + lambda *args, **kwargs: self._make_gapic_stream( + [expected_error] + ) + ) + try: + await self.execute_fn(table, query, operation_timeout=0.1) + except core_exceptions.DeadlineExceeded as e: + retry_exc = e.__cause__ + root_cause = retry_exc.exceptions[0] + assert type(root_cause) == exc_type + assert root_cause == expected_error @pytest.mark.parametrize( "exc_type", @@ -1189,19 +1202,21 @@ async def test_read_rows_retryable_error(self, exc_type): @pytest.mark.asyncio async def test_read_rows_non_retryable_error(self, exc_type): async with self._make_client() as client: - table = client.get_table("instance", "table") - query = ReadRowsQuery() - expected_error = exc_type("mock error") - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( - [expected_error] - ) - try: - await table.read_rows(query, operation_timeout=0.1) - except exc_type as e: - assert e == expected_error + async with client.get_table("instance", "table") as table: + query = ReadRowsQuery() + expected_error = exc_type("mock error") + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = ( + lambda *args, **kwargs: self._make_gapic_stream( + [expected_error] + ) + ) + try: + await self.execute_fn(table, query, operation_timeout=0.1) + except exc_type as e: + assert e == expected_error @pytest.mark.asyncio async def test_read_rows_revise_request(self): @@ -1217,31 +1232,35 @@ async def test_read_rows_revise_request(self): with mock.patch.object(_ReadRowsOperation, "aclose"): revise_rowset.return_value = "modified" async with self._make_client() as client: - table = client.get_table("instance", "table") - row_keys = [b"test_1", b"test_2", b"test_3"] - query = ReadRowsQuery(row_keys=row_keys) - chunks = [ - self._make_chunk(row_key=b"test_1"), - core_exceptions.Aborted("mock retryable error"), - ] - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = ( - lambda *args, **kwargs: self._make_gapic_stream(chunks) - ) - try: - await table.read_rows(query) - except InvalidChunk: - revise_rowset.assert_called() - revise_call_kwargs = revise_rowset.call_args_list[0].kwargs - assert ( - revise_call_kwargs["row_set"] - == query._to_dict()["rows"] + async with client.get_table("instance", "table") as table: + row_keys = [b"test_1", b"test_2", b"test_3"] + query = ReadRowsQuery(row_keys=row_keys) + chunks = [ + self._make_chunk(row_key=b"test_1"), + core_exceptions.Aborted("mock retryable error"), + ] + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = ( + lambda *args, **kwargs: self._make_gapic_stream(chunks) ) - assert revise_call_kwargs["last_seen_row_key"] == b"test_1" - read_rows_request = read_rows.call_args_list[1].args[0] - assert read_rows_request["rows"] == "modified" + try: + await self.execute_fn(table, query) + except InvalidChunk: + revise_rowset.assert_called() + revise_call_kwargs = revise_rowset.call_args_list[ + 0 + ].kwargs + assert ( + revise_call_kwargs["row_set"] + == query._to_dict()["rows"] + ) + assert ( + revise_call_kwargs["last_seen_row_key"] == b"test_1" + ) + read_rows_request = read_rows.call_args_list[1].args[0] + assert read_rows_request["rows"] == "modified" @pytest.mark.asyncio async def test_read_rows_default_timeouts(self): @@ -1262,7 +1281,7 @@ async def test_read_rows_default_timeouts(self): default_per_request_timeout=per_request_timeout, ) as table: try: - await table.read_rows(ReadRowsQuery()) + await self.execute_fn(table, ReadRowsQuery()) except RuntimeError: pass kwargs = mock_op.call_args_list[0].kwargs @@ -1288,7 +1307,8 @@ async def test_read_rows_default_timeout_override(self): default_per_request_timeout=97, ) as table: try: - await table.read_rows( + await self.execute_fn( + table, ReadRowsQuery(), operation_timeout=operation_timeout, per_request_timeout=per_request_timeout, @@ -1323,7 +1343,7 @@ async def test_read_row(self): assert len(args) == 1 assert isinstance(args[0], ReadRowsQuery) assert args[0]._to_dict() == { - "rows": {"row_keys": [row_key]}, + "rows": {"row_keys": [row_key], "row_ranges": []}, "rows_limit": 1, } @@ -1355,7 +1375,7 @@ async def test_read_row_w_filter(self): assert len(args) == 1 assert isinstance(args[0], ReadRowsQuery) assert args[0]._to_dict() == { - "rows": {"row_keys": [row_key]}, + "rows": {"row_keys": [row_key], "row_ranges": []}, "rows_limit": 1, "filter": expected_filter, } @@ -1383,7 +1403,7 @@ async def test_read_row_no_response(self): assert kwargs["per_request_timeout"] == expected_req_timeout assert isinstance(args[0], ReadRowsQuery) assert args[0]._to_dict() == { - "rows": {"row_keys": [row_key]}, + "rows": {"row_keys": [row_key], "row_ranges": []}, "rows_limit": 1, } @@ -1436,7 +1456,7 @@ async def test_row_exists(self, return_value, expected_result): } } assert args[0]._to_dict() == { - "rows": {"row_keys": [row_key]}, + "rows": {"row_keys": [row_key], "row_ranges": []}, "rows_limit": 1, "filter": expected_filter, } @@ -1476,6 +1496,376 @@ async def test_read_rows_metadata(self, include_app_profile): assert "app_profile_id=" not in goog_metadata +class TestReadRowsSharded: + def _make_client(self, *args, **kwargs): + from google.cloud.bigtable.client import BigtableDataClient + + return BigtableDataClient(*args, **kwargs) + + @pytest.mark.asyncio + async def test_read_rows_sharded_empty_query(self): + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with pytest.raises(ValueError) as exc: + await table.read_rows_sharded([]) + assert "empty sharded_query" in str(exc.value) + + @pytest.mark.asyncio + async def test_read_rows_sharded_multiple_queries(self): + """ + Test with multiple queries. Should return results from both + """ + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + table.client._gapic_client, "read_rows" + ) as read_rows: + read_rows.side_effect = ( + lambda *args, **kwargs: TestReadRows._make_gapic_stream( + [ + TestReadRows._make_chunk(row_key=k) + for k in args[0]["rows"]["row_keys"] + ] + ) + ) + query_1 = ReadRowsQuery(b"test_1") + query_2 = ReadRowsQuery(b"test_2") + result = await table.read_rows_sharded([query_1, query_2]) + assert len(result) == 2 + assert result[0].row_key == b"test_1" + assert result[1].row_key == b"test_2" + + @pytest.mark.parametrize("n_queries", [1, 2, 5, 11, 24]) + @pytest.mark.asyncio + async def test_read_rows_sharded_multiple_queries_calls(self, n_queries): + """ + Each query should trigger a separate read_rows call + """ + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object(table, "read_rows") as read_rows: + query_list = [ReadRowsQuery() for _ in range(n_queries)] + await table.read_rows_sharded(query_list) + assert read_rows.call_count == n_queries + + @pytest.mark.asyncio + async def test_read_rows_sharded_errors(self): + """ + Errors should be exposed as ShardedReadRowsExceptionGroups + """ + from google.cloud.bigtable.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.exceptions import FailedQueryShardError + + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object(table, "read_rows") as read_rows: + read_rows.side_effect = RuntimeError("mock error") + query_1 = ReadRowsQuery(b"test_1") + query_2 = ReadRowsQuery(b"test_2") + with pytest.raises(ShardedReadRowsExceptionGroup) as exc: + await table.read_rows_sharded([query_1, query_2]) + exc_group = exc.value + assert isinstance(exc_group, ShardedReadRowsExceptionGroup) + assert len(exc.value.exceptions) == 2 + assert isinstance(exc.value.exceptions[0], FailedQueryShardError) + assert isinstance(exc.value.exceptions[0].__cause__, RuntimeError) + assert exc.value.exceptions[0].index == 0 + assert exc.value.exceptions[0].query == query_1 + assert isinstance(exc.value.exceptions[1], FailedQueryShardError) + assert isinstance(exc.value.exceptions[1].__cause__, RuntimeError) + assert exc.value.exceptions[1].index == 1 + assert exc.value.exceptions[1].query == query_2 + + @pytest.mark.asyncio + async def test_read_rows_sharded_concurrent(self): + """ + Ensure sharded requests are concurrent + """ + import time + + async def mock_call(*args, **kwargs): + await asyncio.sleep(0.1) + return [mock.Mock()] + + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object(table, "read_rows") as read_rows: + read_rows.side_effect = mock_call + queries = [ReadRowsQuery() for _ in range(10)] + start_time = time.monotonic() + result = await table.read_rows_sharded(queries) + call_time = time.monotonic() - start_time + assert read_rows.call_count == 10 + assert len(result) == 10 + # if run in sequence, we would expect this to take 1 second + assert call_time < 0.2 + + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_read_rows_sharded_metadata(self, include_app_profile): + """request should attach metadata headers""" + profile = "profile" if include_app_profile else None + async with self._make_client() as client: + async with client.get_table("i", "t", app_profile_id=profile) as table: + with mock.patch.object( + client._gapic_client, "read_rows", AsyncMock() + ) as read_rows: + await table.read_rows_sharded([ReadRowsQuery()]) + kwargs = read_rows.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata + + @pytest.mark.asyncio + async def test_read_rows_sharded_batching(self): + """ + Large queries should be processed in batches to limit concurrency + operation timeout should change between batches + """ + from google.cloud.bigtable.client import Table + from google.cloud.bigtable.client import CONCURRENCY_LIMIT + + assert CONCURRENCY_LIMIT == 10 # change this test if this changes + + n_queries = 90 + expected_num_batches = n_queries // CONCURRENCY_LIMIT + query_list = [ReadRowsQuery() for _ in range(n_queries)] + + table_mock = AsyncMock() + start_operation_timeout = 10 + start_per_request_timeout = 3 + table_mock.default_operation_timeout = start_operation_timeout + table_mock.default_per_request_timeout = start_per_request_timeout + # clock ticks one second on each check + with mock.patch("time.monotonic", side_effect=range(0, 100000)): + with mock.patch("asyncio.gather", AsyncMock()) as gather_mock: + await Table.read_rows_sharded(table_mock, query_list) + # should have individual calls for each query + assert table_mock.read_rows.call_count == n_queries + # should have single gather call for each batch + assert gather_mock.call_count == expected_num_batches + # ensure that timeouts decrease over time + kwargs = [ + table_mock.read_rows.call_args_list[idx][1] + for idx in range(n_queries) + ] + for batch_idx in range(expected_num_batches): + batch_kwargs = kwargs[ + batch_idx + * CONCURRENCY_LIMIT : (batch_idx + 1) + * CONCURRENCY_LIMIT + ] + for req_kwargs in batch_kwargs: + # each batch should have the same operation_timeout, and it should decrease in each batch + expected_operation_timeout = start_operation_timeout - ( + batch_idx + 1 + ) + assert ( + req_kwargs["operation_timeout"] + == expected_operation_timeout + ) + # each per_request_timeout should start with default value, but decrease when operation_timeout reaches it + expected_per_request_timeout = min( + start_per_request_timeout, expected_operation_timeout + ) + assert ( + req_kwargs["per_request_timeout"] + == expected_per_request_timeout + ) + # await all created coroutines to avoid warnings + for i in range(len(gather_mock.call_args_list)): + for j in range(len(gather_mock.call_args_list[i][0])): + await gather_mock.call_args_list[i][0][j] + + +class TestSampleRowKeys: + def _make_client(self, *args, **kwargs): + from google.cloud.bigtable.client import BigtableDataClient + + return BigtableDataClient(*args, **kwargs) + + async def _make_gapic_stream(self, sample_list: list[tuple[bytes, int]]): + from google.cloud.bigtable_v2.types import SampleRowKeysResponse + + for value in sample_list: + yield SampleRowKeysResponse(row_key=value[0], offset_bytes=value[1]) + + @pytest.mark.asyncio + async def test_sample_row_keys(self): + """ + Test that method returns the expected key samples + """ + samples = [ + (b"test_1", 0), + (b"test_2", 100), + (b"test_3", 200), + ] + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + table.client._gapic_client, "sample_row_keys", AsyncMock() + ) as sample_row_keys: + sample_row_keys.return_value = self._make_gapic_stream(samples) + result = await table.sample_row_keys() + assert len(result) == 3 + assert all(isinstance(r, tuple) for r in result) + assert all(isinstance(r[0], bytes) for r in result) + assert all(isinstance(r[1], int) for r in result) + assert result[0] == samples[0] + assert result[1] == samples[1] + assert result[2] == samples[2] + + @pytest.mark.asyncio + async def test_sample_row_keys_bad_timeout(self): + """ + should raise error if timeout is negative + """ + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with pytest.raises(ValueError) as e: + await table.sample_row_keys(operation_timeout=-1) + assert "operation_timeout must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + await table.sample_row_keys(per_request_timeout=-1) + assert "per_request_timeout must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + await table.sample_row_keys( + operation_timeout=10, per_request_timeout=20 + ) + assert ( + "per_request_timeout must not be greater than operation_timeout" + in str(e.value) + ) + + @pytest.mark.asyncio + async def test_sample_row_keys_default_timeout(self): + """Should fallback to using table default operation_timeout""" + expected_timeout = 99 + async with self._make_client() as client: + async with client.get_table( + "i", "t", default_operation_timeout=expected_timeout + ) as table: + with mock.patch.object( + table.client._gapic_client, "sample_row_keys", AsyncMock() + ) as sample_row_keys: + sample_row_keys.return_value = self._make_gapic_stream([]) + result = await table.sample_row_keys() + _, kwargs = sample_row_keys.call_args + assert abs(kwargs["timeout"] - expected_timeout) < 0.1 + assert result == [] + + @pytest.mark.asyncio + async def test_sample_row_keys_gapic_params(self): + """ + make sure arguments are propagated to gapic call as expected + """ + expected_timeout = 10 + expected_profile = "test1" + instance = "instance_name" + table_id = "my_table" + async with self._make_client() as client: + async with client.get_table( + instance, table_id, app_profile_id=expected_profile + ) as table: + with mock.patch.object( + table.client._gapic_client, "sample_row_keys", AsyncMock() + ) as sample_row_keys: + sample_row_keys.return_value = self._make_gapic_stream([]) + await table.sample_row_keys(per_request_timeout=expected_timeout) + args, kwargs = sample_row_keys.call_args + assert len(args) == 0 + assert len(kwargs) == 4 + assert kwargs["timeout"] == expected_timeout + assert kwargs["app_profile_id"] == expected_profile + assert kwargs["table_name"] == table.table_name + assert kwargs["metadata"] is not None + + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_sample_row_keys_metadata(self, include_app_profile): + """request should attach metadata headers""" + profile = "profile" if include_app_profile else None + async with self._make_client() as client: + async with client.get_table("i", "t", app_profile_id=profile) as table: + with mock.patch.object( + client._gapic_client, "sample_row_keys", AsyncMock() + ) as read_rows: + await table.sample_row_keys() + kwargs = read_rows.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata + + @pytest.mark.parametrize( + "retryable_exception", + [ + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ], + ) + @pytest.mark.asyncio + async def test_sample_row_keys_retryable_errors(self, retryable_exception): + """ + retryable errors should be retried until timeout + """ + from google.api_core.exceptions import DeadlineExceeded + from google.cloud.bigtable.exceptions import RetryExceptionGroup + + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + table.client._gapic_client, "sample_row_keys", AsyncMock() + ) as sample_row_keys: + sample_row_keys.side_effect = retryable_exception("mock") + with pytest.raises(DeadlineExceeded) as e: + await table.sample_row_keys(operation_timeout=0.05) + cause = e.value.__cause__ + assert isinstance(cause, RetryExceptionGroup) + assert len(cause.exceptions) > 0 + assert isinstance(cause.exceptions[0], retryable_exception) + + @pytest.mark.parametrize( + "non_retryable_exception", + [ + core_exceptions.OutOfRange, + core_exceptions.NotFound, + core_exceptions.FailedPrecondition, + RuntimeError, + ValueError, + core_exceptions.Aborted, + ], + ) + @pytest.mark.asyncio + async def test_sample_row_keys_non_retryable_errors(self, non_retryable_exception): + """ + non-retryable errors should cause a raise + """ + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with mock.patch.object( + table.client._gapic_client, "sample_row_keys", AsyncMock() + ) as sample_row_keys: + sample_row_keys.side_effect = non_retryable_exception("mock") + with pytest.raises(non_retryable_exception): + await table.sample_row_keys() + + class TestMutateRow: def _make_client(self, *args, **kwargs): from google.cloud.bigtable.client import BigtableDataClient diff --git a/tests/unit/test_exceptions.py b/tests/unit/test_exceptions.py index 49b90a8a9..e68ccf5e8 100644 --- a/tests/unit/test_exceptions.py +++ b/tests/unit/test_exceptions.py @@ -18,10 +18,16 @@ import google.cloud.bigtable.exceptions as bigtable_exceptions +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock +except ImportError: # pragma: NO COVER + import mock # type: ignore + class TestBigtableExceptionGroup: """ - Subclass for MutationsExceptionGroup and RetryExceptionGroup + Subclass for MutationsExceptionGroup, RetryExceptionGroup, and ShardedReadRowsExceptionGroup """ def _get_class(self): @@ -190,13 +196,50 @@ def test_raise(self, exception_list, expected_message): assert list(e.value.exceptions) == exception_list +class TestShardedReadRowsExceptionGroup(TestBigtableExceptionGroup): + def _get_class(self): + from google.cloud.bigtable.exceptions import ShardedReadRowsExceptionGroup + + return ShardedReadRowsExceptionGroup + + def _make_one(self, excs=None, succeeded=None, num_entries=3): + if excs is None: + excs = [RuntimeError("mock")] + succeeded = succeeded or [] + + return self._get_class()(excs, succeeded, num_entries) + + @pytest.mark.parametrize( + "exception_list,succeeded,total_entries,expected_message", + [ + ([Exception()], [], 1, "1 sub-exception (from 1 query attempted)"), + ([Exception()], [1], 2, "1 sub-exception (from 2 queries attempted)"), + ( + [Exception(), RuntimeError()], + [0, 1], + 2, + "2 sub-exceptions (from 2 queries attempted)", + ), + ], + ) + def test_raise(self, exception_list, succeeded, total_entries, expected_message): + """ + Create exception in raise statement, which calls __new__ and __init__ + """ + with pytest.raises(self._get_class()) as e: + raise self._get_class()(exception_list, succeeded, total_entries) + assert str(e.value) == expected_message + assert list(e.value.exceptions) == exception_list + assert e.value.successful_rows == succeeded + + class TestFailedMutationEntryError: def _get_class(self): from google.cloud.bigtable.exceptions import FailedMutationEntryError return FailedMutationEntryError - def _make_one(self, idx=9, entry=unittest.mock.Mock(), cause=RuntimeError("mock")): + def _make_one(self, idx=9, entry=mock.Mock(), cause=RuntimeError("mock")): return self._get_class()(idx, entry, cause) @@ -205,7 +248,7 @@ def test_raise(self): Create exception in raise statement, which calls __new__ and __init__ """ test_idx = 2 - test_entry = unittest.mock.Mock() + test_entry = mock.Mock() test_exc = ValueError("test") with pytest.raises(self._get_class()) as e: raise self._get_class()(test_idx, test_entry, test_exc) @@ -237,3 +280,29 @@ def test_raise_idempotent(self): assert e.value.entry == test_entry assert e.value.__cause__ == test_exc assert test_entry.is_idempotent.call_count == 1 + + +class TestFailedQueryShardError: + def _get_class(self): + from google.cloud.bigtable.exceptions import FailedQueryShardError + + return FailedQueryShardError + + def _make_one(self, idx=9, query=mock.Mock(), cause=RuntimeError("mock")): + + return self._get_class()(idx, query, cause) + + def test_raise(self): + """ + Create exception in raise statement, which calls __new__ and __init__ + """ + test_idx = 2 + test_query = mock.Mock() + test_exc = ValueError("test") + with pytest.raises(self._get_class()) as e: + raise self._get_class()(test_idx, test_query, test_exc) + assert str(e.value) == "Failed query at index 2 with cause: ValueError('test')" + assert e.value.index == test_idx + assert e.value.query == test_query + assert e.value.__cause__ == test_exc + assert isinstance(e.value, Exception) diff --git a/tests/unit/test_read_rows_query.py b/tests/unit/test_read_rows_query.py index f630f2eab..7ecd91f8c 100644 --- a/tests/unit/test_read_rows_query.py +++ b/tests/unit/test_read_rows_query.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest +import pytest TEST_ROWS = [ "row_key_1", @@ -20,7 +20,7 @@ ] -class TestRowRange(unittest.TestCase): +class TestRowRange: @staticmethod def _get_target_class(): from google.cloud.bigtable.read_rows_query import RowRange @@ -32,68 +32,57 @@ def _make_one(self, *args, **kwargs): def test_ctor_start_end(self): row_range = self._make_one("test_row", "test_row2") - self.assertEqual(row_range.start.key, "test_row".encode()) - self.assertEqual(row_range.end.key, "test_row2".encode()) - self.assertEqual(row_range.start.is_inclusive, True) - self.assertEqual(row_range.end.is_inclusive, False) + assert row_range.start.key == "test_row".encode() + assert row_range.end.key == "test_row2".encode() + assert row_range.start.is_inclusive is True + assert row_range.end.is_inclusive is False def test_ctor_start_only(self): row_range = self._make_one("test_row3") - self.assertEqual(row_range.start.key, "test_row3".encode()) - self.assertEqual(row_range.start.is_inclusive, True) - self.assertEqual(row_range.end, None) + assert row_range.start.key == "test_row3".encode() + assert row_range.start.is_inclusive is True + assert row_range.end is None def test_ctor_end_only(self): row_range = self._make_one(end_key="test_row4") - self.assertEqual(row_range.end.key, "test_row4".encode()) - self.assertEqual(row_range.end.is_inclusive, False) - self.assertEqual(row_range.start, None) + assert row_range.end.key == "test_row4".encode() + assert row_range.end.is_inclusive is False + assert row_range.start is None def test_ctor_inclusive_flags(self): row_range = self._make_one("test_row5", "test_row6", False, True) - self.assertEqual(row_range.start.key, "test_row5".encode()) - self.assertEqual(row_range.end.key, "test_row6".encode()) - self.assertEqual(row_range.start.is_inclusive, False) - self.assertEqual(row_range.end.is_inclusive, True) + assert row_range.start.key == "test_row5".encode() + assert row_range.end.key == "test_row6".encode() + assert row_range.start.is_inclusive is False + assert row_range.end.is_inclusive is True def test_ctor_defaults(self): row_range = self._make_one() - self.assertEqual(row_range.start, None) - self.assertEqual(row_range.end, None) + assert row_range.start is None + assert row_range.end is None def test_ctor_flags_only(self): - with self.assertRaises(ValueError) as exc: + with pytest.raises(ValueError) as exc: self._make_one(start_is_inclusive=True, end_is_inclusive=True) - self.assertEqual( - exc.exception.args, - ("start_is_inclusive must be set with start_key",), - ) - with self.assertRaises(ValueError) as exc: + assert str(exc.value) == "start_is_inclusive must be set with start_key" + with pytest.raises(ValueError) as exc: self._make_one(start_is_inclusive=False, end_is_inclusive=False) - self.assertEqual( - exc.exception.args, - ("start_is_inclusive must be set with start_key",), - ) - with self.assertRaises(ValueError) as exc: + assert str(exc.value) == "start_is_inclusive must be set with start_key" + with pytest.raises(ValueError) as exc: self._make_one(start_is_inclusive=False) - self.assertEqual( - exc.exception.args, - ("start_is_inclusive must be set with start_key",), - ) - with self.assertRaises(ValueError) as exc: + assert str(exc.value) == "start_is_inclusive must be set with start_key" + with pytest.raises(ValueError) as exc: self._make_one(end_is_inclusive=True) - self.assertEqual( - exc.exception.args, ("end_is_inclusive must be set with end_key",) - ) + assert str(exc.value) == "end_is_inclusive must be set with end_key" def test_ctor_invalid_keys(self): # test with invalid keys - with self.assertRaises(ValueError) as exc: + with pytest.raises(ValueError) as exc: self._make_one(1, "2") - self.assertEqual(exc.exception.args, ("start_key must be a string or bytes",)) - with self.assertRaises(ValueError) as exc: + assert str(exc.value) == "start_key must be a string or bytes" + with pytest.raises(ValueError) as exc: self._make_one("1", 2) - self.assertEqual(exc.exception.args, ("end_key must be a string or bytes",)) + assert str(exc.value) == "end_key must be a string or bytes" def test__to_dict_defaults(self): row_range = self._make_one("test_row", "test_row2") @@ -101,7 +90,7 @@ def test__to_dict_defaults(self): "start_key_closed": b"test_row", "end_key_open": b"test_row2", } - self.assertEqual(row_range._to_dict(), expected) + assert row_range._to_dict() == expected def test__to_dict_inclusive_flags(self): row_range = self._make_one("test_row", "test_row2", False, True) @@ -109,10 +98,148 @@ def test__to_dict_inclusive_flags(self): "start_key_open": b"test_row", "end_key_closed": b"test_row2", } - self.assertEqual(row_range._to_dict(), expected) + assert row_range._to_dict() == expected + + @pytest.mark.parametrize( + "input_dict,expected_start,expected_end,start_is_inclusive,end_is_inclusive", + [ + ( + {"start_key_closed": "test_row", "end_key_open": "test_row2"}, + b"test_row", + b"test_row2", + True, + False, + ), + ( + {"start_key_closed": b"test_row", "end_key_open": b"test_row2"}, + b"test_row", + b"test_row2", + True, + False, + ), + ( + {"start_key_open": "test_row", "end_key_closed": "test_row2"}, + b"test_row", + b"test_row2", + False, + True, + ), + ({"start_key_open": b"a"}, b"a", None, False, None), + ({"end_key_closed": b"b"}, None, b"b", None, True), + ({"start_key_closed": "a"}, b"a", None, True, None), + ({"end_key_open": b"b"}, None, b"b", None, False), + ({}, None, None, None, None), + ], + ) + def test__from_dict( + self, + input_dict, + expected_start, + expected_end, + start_is_inclusive, + end_is_inclusive, + ): + from google.cloud.bigtable.read_rows_query import RowRange + + row_range = RowRange._from_dict(input_dict) + assert row_range._to_dict().keys() == input_dict.keys() + found_start = row_range.start + found_end = row_range.end + if expected_start is None: + assert found_start is None + assert start_is_inclusive is None + else: + assert found_start.key == expected_start + assert found_start.is_inclusive == start_is_inclusive + if expected_end is None: + assert found_end is None + assert end_is_inclusive is None + else: + assert found_end.key == expected_end + assert found_end.is_inclusive == end_is_inclusive + + @pytest.mark.parametrize( + "dict_repr", + [ + {"start_key_closed": "test_row", "end_key_open": "test_row2"}, + {"start_key_closed": b"test_row", "end_key_open": b"test_row2"}, + {"start_key_open": "test_row", "end_key_closed": "test_row2"}, + {"start_key_open": b"a"}, + {"end_key_closed": b"b"}, + {"start_key_closed": "a"}, + {"end_key_open": b"b"}, + {}, + ], + ) + def test__from_points(self, dict_repr): + from google.cloud.bigtable.read_rows_query import RowRange + + row_range_from_dict = RowRange._from_dict(dict_repr) + row_range_from_points = RowRange._from_points( + row_range_from_dict.start, row_range_from_dict.end + ) + assert row_range_from_points._to_dict() == row_range_from_dict._to_dict() + + @pytest.mark.parametrize( + "first_dict,second_dict,should_match", + [ + ( + {"start_key_closed": "a", "end_key_open": "b"}, + {"start_key_closed": "a", "end_key_open": "b"}, + True, + ), + ( + {"start_key_closed": "a", "end_key_open": "b"}, + {"start_key_closed": "a", "end_key_open": "c"}, + False, + ), + ( + {"start_key_closed": "a", "end_key_open": "b"}, + {"start_key_closed": "a", "end_key_closed": "b"}, + False, + ), + ( + {"start_key_closed": b"a", "end_key_open": b"b"}, + {"start_key_closed": "a", "end_key_open": "b"}, + True, + ), + ({}, {}, True), + ({"start_key_closed": "a"}, {}, False), + ({"start_key_closed": "a"}, {"start_key_closed": "a"}, True), + ({"start_key_closed": "a"}, {"start_key_open": "a"}, False), + ], + ) + def test___hash__(self, first_dict, second_dict, should_match): + from google.cloud.bigtable.read_rows_query import RowRange + + row_range1 = RowRange._from_dict(first_dict) + row_range2 = RowRange._from_dict(second_dict) + assert (hash(row_range1) == hash(row_range2)) == should_match + + @pytest.mark.parametrize( + "dict_repr,expected", + [ + ({"start_key_closed": "test_row", "end_key_open": "test_row2"}, True), + ({"start_key_closed": b"test_row", "end_key_open": b"test_row2"}, True), + ({"start_key_open": "test_row", "end_key_closed": "test_row2"}, True), + ({"start_key_open": b"a"}, True), + ({"end_key_closed": b"b"}, True), + ({"start_key_closed": "a"}, True), + ({"end_key_open": b"b"}, True), + ({}, False), + ], + ) + def test___bool__(self, dict_repr, expected): + """ + Only row range with both points empty should be falsy + """ + from google.cloud.bigtable.read_rows_query import RowRange + + row_range = RowRange._from_dict(dict_repr) + assert bool(row_range) is expected -class TestReadRowsQuery(unittest.TestCase): +class TestReadRowsQuery: @staticmethod def _get_target_class(): from google.cloud.bigtable.read_rows_query import ReadRowsQuery @@ -124,48 +251,53 @@ def _make_one(self, *args, **kwargs): def test_ctor_defaults(self): query = self._make_one() - self.assertEqual(query.row_keys, set()) - self.assertEqual(query.row_ranges, []) - self.assertEqual(query.filter, None) - self.assertEqual(query.limit, None) + assert query.row_keys == set() + assert query.row_ranges == set() + assert query.filter is None + assert query.limit is None def test_ctor_explicit(self): from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.read_rows_query import RowRange filter_ = RowFilterChain() - query = self._make_one(["row_key_1", "row_key_2"], limit=10, row_filter=filter_) - self.assertEqual(len(query.row_keys), 2) - self.assertIn("row_key_1".encode(), query.row_keys) - self.assertIn("row_key_2".encode(), query.row_keys) - self.assertEqual(query.row_ranges, []) - self.assertEqual(query.filter, filter_) - self.assertEqual(query.limit, 10) + query = self._make_one( + ["row_key_1", "row_key_2"], + row_ranges=[RowRange("row_key_3", "row_key_4")], + limit=10, + row_filter=filter_, + ) + assert len(query.row_keys) == 2 + assert "row_key_1".encode() in query.row_keys + assert "row_key_2".encode() in query.row_keys + assert len(query.row_ranges) == 1 + assert RowRange("row_key_3", "row_key_4") in query.row_ranges + assert query.filter == filter_ + assert query.limit == 10 def test_ctor_invalid_limit(self): - with self.assertRaises(ValueError) as exc: + with pytest.raises(ValueError) as exc: self._make_one(limit=-1) - self.assertEqual(exc.exception.args, ("limit must be >= 0",)) + assert str(exc.value) == "limit must be >= 0" def test_set_filter(self): from google.cloud.bigtable.row_filters import RowFilterChain filter1 = RowFilterChain() query = self._make_one() - self.assertEqual(query.filter, None) + assert query.filter is None query.filter = filter1 - self.assertEqual(query.filter, filter1) + assert query.filter == filter1 filter2 = RowFilterChain() query.filter = filter2 - self.assertEqual(query.filter, filter2) + assert query.filter == filter2 query.filter = None - self.assertEqual(query.filter, None) + assert query.filter is None query.filter = RowFilterChain() - self.assertEqual(query.filter, RowFilterChain()) - with self.assertRaises(ValueError) as exc: + assert query.filter == RowFilterChain() + with pytest.raises(ValueError) as exc: query.filter = 1 - self.assertEqual( - exc.exception.args, ("row_filter must be a RowFilter or dict",) - ) + assert str(exc.value) == "row_filter must be a RowFilter or dict" def test_set_filter_dict(self): from google.cloud.bigtable.row_filters import RowSampleFilter @@ -174,123 +306,128 @@ def test_set_filter_dict(self): filter1 = RowSampleFilter(0.5) filter1_dict = filter1.to_dict() query = self._make_one() - self.assertEqual(query.filter, None) + assert query.filter is None query.filter = filter1_dict - self.assertEqual(query.filter, filter1_dict) + assert query.filter == filter1_dict output = query._to_dict() - self.assertEqual(output["filter"], filter1_dict) + assert output["filter"] == filter1_dict proto_output = ReadRowsRequest(**output) - self.assertEqual(proto_output.filter, filter1._to_pb()) + assert proto_output.filter == filter1._to_pb() query.filter = None - self.assertEqual(query.filter, None) + assert query.filter is None def test_set_limit(self): query = self._make_one() - self.assertEqual(query.limit, None) + assert query.limit is None query.limit = 10 - self.assertEqual(query.limit, 10) + assert query.limit == 10 query.limit = 9 - self.assertEqual(query.limit, 9) + assert query.limit == 9 query.limit = 0 - self.assertEqual(query.limit, 0) - with self.assertRaises(ValueError) as exc: + assert query.limit == 0 + with pytest.raises(ValueError) as exc: query.limit = -1 - self.assertEqual(exc.exception.args, ("limit must be >= 0",)) - with self.assertRaises(ValueError) as exc: + assert str(exc.value) == "limit must be >= 0" + with pytest.raises(ValueError) as exc: query.limit = -100 - self.assertEqual(exc.exception.args, ("limit must be >= 0",)) + assert str(exc.value) == "limit must be >= 0" def test_add_key_str(self): query = self._make_one() - self.assertEqual(query.row_keys, set()) + assert query.row_keys == set() input_str = "test_row" query.add_key(input_str) - self.assertEqual(len(query.row_keys), 1) - self.assertIn(input_str.encode(), query.row_keys) + assert len(query.row_keys) == 1 + assert input_str.encode() in query.row_keys input_str2 = "test_row2" query.add_key(input_str2) - self.assertEqual(len(query.row_keys), 2) - self.assertIn(input_str.encode(), query.row_keys) - self.assertIn(input_str2.encode(), query.row_keys) + assert len(query.row_keys) == 2 + assert input_str.encode() in query.row_keys + assert input_str2.encode() in query.row_keys def test_add_key_bytes(self): query = self._make_one() - self.assertEqual(query.row_keys, set()) + assert query.row_keys == set() input_bytes = b"test_row" query.add_key(input_bytes) - self.assertEqual(len(query.row_keys), 1) - self.assertIn(input_bytes, query.row_keys) + assert len(query.row_keys) == 1 + assert input_bytes in query.row_keys input_bytes2 = b"test_row2" query.add_key(input_bytes2) - self.assertEqual(len(query.row_keys), 2) - self.assertIn(input_bytes, query.row_keys) - self.assertIn(input_bytes2, query.row_keys) + assert len(query.row_keys) == 2 + assert input_bytes in query.row_keys + assert input_bytes2 in query.row_keys def test_add_rows_batch(self): query = self._make_one() - self.assertEqual(query.row_keys, set()) + assert query.row_keys == set() input_batch = ["test_row", b"test_row2", "test_row3"] for k in input_batch: query.add_key(k) - self.assertEqual(len(query.row_keys), 3) - self.assertIn(b"test_row", query.row_keys) - self.assertIn(b"test_row2", query.row_keys) - self.assertIn(b"test_row3", query.row_keys) + assert len(query.row_keys) == 3 + assert b"test_row" in query.row_keys + assert b"test_row2" in query.row_keys + assert b"test_row3" in query.row_keys # test adding another batch for k in ["test_row4", b"test_row5"]: query.add_key(k) - self.assertEqual(len(query.row_keys), 5) - self.assertIn(input_batch[0].encode(), query.row_keys) - self.assertIn(input_batch[1], query.row_keys) - self.assertIn(input_batch[2].encode(), query.row_keys) - self.assertIn(b"test_row4", query.row_keys) - self.assertIn(b"test_row5", query.row_keys) + assert len(query.row_keys) == 5 + assert input_batch[0].encode() in query.row_keys + assert input_batch[1] in query.row_keys + assert input_batch[2].encode() in query.row_keys + assert b"test_row4" in query.row_keys + assert b"test_row5" in query.row_keys def test_add_key_invalid(self): query = self._make_one() - with self.assertRaises(ValueError) as exc: + with pytest.raises(ValueError) as exc: query.add_key(1) - self.assertEqual(exc.exception.args, ("row_key must be string or bytes",)) - with self.assertRaises(ValueError) as exc: + assert str(exc.value) == "row_key must be string or bytes" + with pytest.raises(ValueError) as exc: query.add_key(["s"]) - self.assertEqual(exc.exception.args, ("row_key must be string or bytes",)) + assert str(exc.value) == "row_key must be string or bytes" def test_duplicate_rows(self): # should only hold one of each input key key_1 = b"test_row" key_2 = b"test_row2" query = self._make_one(row_keys=[key_1, key_1, key_2]) - self.assertEqual(len(query.row_keys), 2) - self.assertIn(key_1, query.row_keys) - self.assertIn(key_2, query.row_keys) + assert len(query.row_keys) == 2 + assert key_1 in query.row_keys + assert key_2 in query.row_keys key_3 = "test_row3" for i in range(10): query.add_key(key_3) - self.assertEqual(len(query.row_keys), 3) + assert len(query.row_keys) == 3 def test_add_range(self): from google.cloud.bigtable.read_rows_query import RowRange query = self._make_one() - self.assertEqual(query.row_ranges, []) + assert query.row_ranges == set() input_range = RowRange(start_key=b"test_row") query.add_range(input_range) - self.assertEqual(len(query.row_ranges), 1) - self.assertEqual(query.row_ranges[0], input_range) + assert len(query.row_ranges) == 1 + assert input_range in query.row_ranges input_range2 = RowRange(start_key=b"test_row2") query.add_range(input_range2) - self.assertEqual(len(query.row_ranges), 2) - self.assertEqual(query.row_ranges[0], input_range) - self.assertEqual(query.row_ranges[1], input_range2) + assert len(query.row_ranges) == 2 + assert input_range in query.row_ranges + assert input_range2 in query.row_ranges + query.add_range(input_range2) + assert len(query.row_ranges) == 2 def test_add_range_dict(self): + from google.cloud.bigtable.read_rows_query import RowRange + query = self._make_one() - self.assertEqual(query.row_ranges, []) + assert query.row_ranges == set() input_range = {"start_key_closed": b"test_row"} query.add_range(input_range) - self.assertEqual(len(query.row_ranges), 1) - self.assertEqual(query.row_ranges[0], input_range) + assert len(query.row_ranges) == 1 + range_obj = RowRange._from_dict(input_range) + assert range_obj in query.row_ranges def test_to_dict_rows_default(self): # dictionary should be in rowset proto format @@ -298,16 +435,16 @@ def test_to_dict_rows_default(self): query = self._make_one() output = query._to_dict() - self.assertTrue(isinstance(output, dict)) - self.assertEqual(len(output.keys()), 1) - expected = {"rows": {}} - self.assertEqual(output, expected) + assert isinstance(output, dict) + assert len(output.keys()) == 1 + expected = {"rows": {"row_keys": [], "row_ranges": []}} + assert output == expected request_proto = ReadRowsRequest(**output) - self.assertEqual(request_proto.rows.row_keys, []) - self.assertEqual(request_proto.rows.row_ranges, []) - self.assertFalse(request_proto.filter) - self.assertEqual(request_proto.rows_limit, 0) + assert request_proto.rows.row_keys == [] + assert request_proto.rows.row_ranges == [] + assert not request_proto.filter + assert request_proto.rows_limit == 0 def test_to_dict_rows_populated(self): # dictionary should be in rowset proto format @@ -321,44 +458,291 @@ def test_to_dict_rows_populated(self): query.add_range(RowRange("test_row3")) query.add_range(RowRange(start_key=None, end_key="test_row5")) query.add_range(RowRange(b"test_row6", b"test_row7", False, True)) - query.add_range(RowRange()) + query.add_range({}) query.add_key("test_row") query.add_key(b"test_row2") query.add_key("test_row3") query.add_key(b"test_row3") query.add_key(b"test_row4") output = query._to_dict() - self.assertTrue(isinstance(output, dict)) + assert isinstance(output, dict) request_proto = ReadRowsRequest(**output) rowset_proto = request_proto.rows # check rows - self.assertEqual(len(rowset_proto.row_keys), 4) - self.assertEqual(rowset_proto.row_keys[0], b"test_row") - self.assertEqual(rowset_proto.row_keys[1], b"test_row2") - self.assertEqual(rowset_proto.row_keys[2], b"test_row3") - self.assertEqual(rowset_proto.row_keys[3], b"test_row4") + assert len(rowset_proto.row_keys) == 4 + assert rowset_proto.row_keys[0] == b"test_row" + assert rowset_proto.row_keys[1] == b"test_row2" + assert rowset_proto.row_keys[2] == b"test_row3" + assert rowset_proto.row_keys[3] == b"test_row4" # check ranges - self.assertEqual(len(rowset_proto.row_ranges), 5) - self.assertEqual(rowset_proto.row_ranges[0].start_key_closed, b"test_row") - self.assertEqual(rowset_proto.row_ranges[0].end_key_open, b"test_row2") - self.assertEqual(rowset_proto.row_ranges[1].start_key_closed, b"test_row3") - self.assertEqual(rowset_proto.row_ranges[1].end_key_open, b"") - self.assertEqual(rowset_proto.row_ranges[2].start_key_closed, b"") - self.assertEqual(rowset_proto.row_ranges[2].end_key_open, b"test_row5") - self.assertEqual(rowset_proto.row_ranges[3].start_key_open, b"test_row6") - self.assertEqual(rowset_proto.row_ranges[3].end_key_closed, b"test_row7") - self.assertEqual(rowset_proto.row_ranges[4].start_key_closed, b"") - self.assertEqual(rowset_proto.row_ranges[4].end_key_open, b"") + assert len(rowset_proto.row_ranges) == 5 + assert { + "start_key_closed": b"test_row", + "end_key_open": b"test_row2", + } in output["rows"]["row_ranges"] + assert {"start_key_closed": b"test_row3"} in output["rows"]["row_ranges"] + assert {"end_key_open": b"test_row5"} in output["rows"]["row_ranges"] + assert { + "start_key_open": b"test_row6", + "end_key_closed": b"test_row7", + } in output["rows"]["row_ranges"] + assert {} in output["rows"]["row_ranges"] # check limit - self.assertEqual(request_proto.rows_limit, 100) + assert request_proto.rows_limit == 100 # check filter filter_proto = request_proto.filter - self.assertEqual(filter_proto, row_filter._to_pb()) + assert filter_proto == row_filter._to_pb() + + def _parse_query_string(self, query_string): + from google.cloud.bigtable.read_rows_query import ReadRowsQuery, RowRange + + query = ReadRowsQuery() + segments = query_string.split(",") + for segment in segments: + if "-" in segment: + start, end = segment.split("-") + s_open, e_open = True, True + if start == "": + start = None + s_open = None + else: + if start[0] == "(": + s_open = False + start = start[1:] + if end == "": + end = None + e_open = None + else: + if end[-1] == ")": + e_open = False + end = end[:-1] + query.add_range(RowRange(start, end, s_open, e_open)) + else: + query.add_key(segment) + return query + + @pytest.mark.parametrize( + "query_string,shard_points", + [ + ("a,[p-q)", []), + ("0_key,[1_range_start-2_range_end)", ["3_split"]), + ("0_key,[1_range_start-2_range_end)", ["2_range_end"]), + ("0_key,[1_range_start-2_range_end]", ["2_range_end"]), + ("-1_range_end)", ["5_split"]), + ("8_key,(1_range_start-2_range_end]", ["1_range_start"]), + ("9_row_key,(5_range_start-7_range_end)", ["3_split"]), + ("3_row_key,(5_range_start-7_range_end)", ["2_row_key"]), + ("4_split,4_split,(3_split-5_split]", ["3_split", "5_split"]), + ("(3_split-", ["3_split"]), + ], + ) + def test_shard_no_split(self, query_string, shard_points): + """ + Test sharding with a set of queries that should not result in any splits. + """ + initial_query = self._parse_query_string(query_string) + row_samples = [(point.encode(), None) for point in shard_points] + sharded_queries = initial_query.shard(row_samples) + assert len(sharded_queries) == 1 + assert initial_query == sharded_queries[0] + + def test_shard_full_table_scan_empty_split(self): + """ + Sharding a full table scan with no split should return another full table scan. + """ + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + + full_scan_query = ReadRowsQuery() + split_points = [] + sharded_queries = full_scan_query.shard(split_points) + assert len(sharded_queries) == 1 + result_query = sharded_queries[0] + assert result_query == full_scan_query + + def test_shard_full_table_scan_with_split(self): + """ + Test splitting a full table scan into two queries + """ + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + + full_scan_query = ReadRowsQuery() + split_points = [(b"a", None)] + sharded_queries = full_scan_query.shard(split_points) + assert len(sharded_queries) == 2 + assert sharded_queries[0] == self._parse_query_string("-a]") + assert sharded_queries[1] == self._parse_query_string("(a-") + + def test_shard_full_table_scan_with_multiple_split(self): + """ + Test splitting a full table scan into three queries + """ + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + + full_scan_query = ReadRowsQuery() + split_points = [(b"a", None), (b"z", None)] + sharded_queries = full_scan_query.shard(split_points) + assert len(sharded_queries) == 3 + assert sharded_queries[0] == self._parse_query_string("-a]") + assert sharded_queries[1] == self._parse_query_string("(a-z]") + assert sharded_queries[2] == self._parse_query_string("(z-") + + def test_shard_multiple_keys(self): + """ + Test splitting multiple individual keys into separate queries + """ + initial_query = self._parse_query_string("1_beforeSplit,2_onSplit,3_afterSplit") + split_points = [(b"2_onSplit", None)] + sharded_queries = initial_query.shard(split_points) + assert len(sharded_queries) == 2 + assert sharded_queries[0] == self._parse_query_string("1_beforeSplit,2_onSplit") + assert sharded_queries[1] == self._parse_query_string("3_afterSplit") + + def test_shard_keys_empty_left(self): + """ + Test with the left-most split point empty + """ + initial_query = self._parse_query_string("5_test,8_test") + split_points = [(b"0_split", None), (b"6_split", None)] + sharded_queries = initial_query.shard(split_points) + assert len(sharded_queries) == 2 + assert sharded_queries[0] == self._parse_query_string("5_test") + assert sharded_queries[1] == self._parse_query_string("8_test") + + def test_shard_keys_empty_right(self): + """ + Test with the right-most split point empty + """ + initial_query = self._parse_query_string("0_test,2_test") + split_points = [(b"1_split", None), (b"5_split", None)] + sharded_queries = initial_query.shard(split_points) + assert len(sharded_queries) == 2 + assert sharded_queries[0] == self._parse_query_string("0_test") + assert sharded_queries[1] == self._parse_query_string("2_test") + + def test_shard_mixed_split(self): + """ + Test splitting a complex query with multiple split points + """ + initial_query = self._parse_query_string("0,a,c,-a],-b],(c-e],(d-f],(m-") + split_points = [(s.encode(), None) for s in ["a", "d", "j", "o"]] + sharded_queries = initial_query.shard(split_points) + assert len(sharded_queries) == 5 + assert sharded_queries[0] == self._parse_query_string("0,a,-a]") + assert sharded_queries[1] == self._parse_query_string("c,(a-b],(c-d]") + assert sharded_queries[2] == self._parse_query_string("(d-e],(d-f]") + assert sharded_queries[3] == self._parse_query_string("(m-o]") + assert sharded_queries[4] == self._parse_query_string("(o-") + + def test_shard_unsorted_request(self): + """ + Test with a query that contains rows and queries in a random order + """ + initial_query = self._parse_query_string( + "7_row_key_1,2_row_key_2,[8_range_1_start-9_range_1_end),[3_range_2_start-4_range_2_end)" + ) + split_points = [(b"5-split", None)] + sharded_queries = initial_query.shard(split_points) + assert len(sharded_queries) == 2 + assert sharded_queries[0] == self._parse_query_string( + "2_row_key_2,[3_range_2_start-4_range_2_end)" + ) + assert sharded_queries[1] == self._parse_query_string( + "7_row_key_1,[8_range_1_start-9_range_1_end)" + ) + + @pytest.mark.parametrize( + "query_string,shard_points", + [ + ("a,[p-q)", []), + ("0_key,[1_range_start-2_range_end)", ["3_split"]), + ("-1_range_end)", ["5_split"]), + ("0_key,[1_range_start-2_range_end)", ["2_range_end"]), + ("9_row_key,(5_range_start-7_range_end)", ["3_split"]), + ("(5_range_start-", ["3_split"]), + ("3_split,[3_split-5_split)", ["3_split", "5_split"]), + ("[3_split-", ["3_split"]), + ("", []), + ("", ["3_split"]), + ("", ["3_split", "5_split"]), + ("1,2,3,4,5,6,7,8,9", ["3_split"]), + ], + ) + def test_shard_keeps_filter(self, query_string, shard_points): + """ + sharded queries should keep the filter from the original query + """ + initial_query = self._parse_query_string(query_string) + expected_filter = {"test": "filter"} + initial_query.filter = expected_filter + row_samples = [(point.encode(), None) for point in shard_points] + sharded_queries = initial_query.shard(row_samples) + assert len(sharded_queries) > 0 + for query in sharded_queries: + assert query.filter == expected_filter + + def test_shard_limit_exception(self): + """ + queries with a limit should raise an exception when a shard is attempted + """ + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + + query = ReadRowsQuery(limit=10) + with pytest.raises(AttributeError) as e: + query.shard([]) + assert "Cannot shard query with a limit" in str(e.value) + + @pytest.mark.parametrize( + "first_args,second_args,expected", + [ + ((), (), True), + ((), ("a",), False), + (("a",), (), False), + (("a",), ("a",), True), + (("a",), (b"a",), True), + (("a",), ("b",), False), + (("a",), ("a", ["b"]), False), + (("a", ["b"]), ("a", ["b"]), True), + (("a", ["b"]), ("a", ["b", "c"]), False), + (("a", ["b", "c"]), ("a", [b"b", "c"]), True), + (("a", ["b", "c"], 1), ("a", ["b", b"c"], 1), True), + (("a", ["b"], 1), ("a", ["b"], 2), False), + (("a", ["b"], 1, {"a": "b"}), ("a", ["b"], 1, {"a": "b"}), True), + (("a", ["b"], 1, {"a": "b"}), ("a", ["b"], 1), False), + ( + (), + (None, [None], None, None), + True, + ), # empty query is equal to empty row range + ((), (None, [None], 1, None), False), + ((), (None, [None], None, {"a": "b"}), False), + ], + ) + def test___eq__(self, first_args, second_args, expected): + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.read_rows_query import RowRange + + # replace row_range placeholders with a RowRange object + if len(first_args) > 1: + first_args = list(first_args) + first_args[1] = [RowRange(c) for c in first_args[1]] + if len(second_args) > 1: + second_args = list(second_args) + second_args[1] = [RowRange(c) for c in second_args[1]] + first = ReadRowsQuery(*first_args) + second = ReadRowsQuery(*second_args) + assert (first == second) == expected + + def test___repr__(self): + from google.cloud.bigtable.read_rows_query import ReadRowsQuery + + instance = self._make_one(row_keys=["a", "b"], row_filter={}, limit=10) + # should be able to recreate the instance from the repr + repr_str = repr(instance) + recreated = eval(repr_str) + assert isinstance(recreated, ReadRowsQuery) + assert recreated == instance def test_empty_row_set(self): """Empty strings should be treated as keys inputs""" query = self._make_one(row_keys="") - self.assertEqual(query.row_keys, {b""}) - - def test_shard(self): - pass + assert query.row_keys == {b""} From ceaf598e1743eb145dafca984cff2630b42431f3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 26 Jun 2023 11:10:50 -0700 Subject: [PATCH 12/56] feat: ping and warm with metadata (#810) --- google/cloud/bigtable/_helpers.py | 2 +- google/cloud/bigtable/client.py | 59 +- tests/system/test_system.py | 1 + tests/unit/test__helpers.py | 2 +- tests/unit/test_client.py | 857 ++++++++++++++++++------------ 5 files changed, 554 insertions(+), 367 deletions(-) diff --git a/google/cloud/bigtable/_helpers.py b/google/cloud/bigtable/_helpers.py index dec4c2014..722fac9f4 100644 --- a/google/cloud/bigtable/_helpers.py +++ b/google/cloud/bigtable/_helpers.py @@ -35,7 +35,7 @@ def _make_metadata( params.append(f"table_name={table_name}") if app_profile_id is not None: params.append(f"app_profile_id={app_profile_id}") - params_str = ",".join(params) + params_str = "&".join(params) return [("x-goog-request-params", params_str)] diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index 055263cfa..3d33eebf9 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -32,6 +32,8 @@ import sys import random +from collections import namedtuple + from google.cloud.bigtable_v2.services.bigtable.client import BigtableClientMeta from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient from google.cloud.bigtable_v2.services.bigtable.async_client import DEFAULT_CLIENT_INFO @@ -74,6 +76,11 @@ # used by read_rows_sharded to limit how many requests are attempted in parallel CONCURRENCY_LIMIT = 10 +# used to register instance data with the client for channel warming +_WarmedInstanceKey = namedtuple( + "_WarmedInstanceKey", ["instance_name", "table_name", "app_profile_id"] +) + class BigtableDataClient(ClientWithProject): def __init__( @@ -139,12 +146,12 @@ def __init__( PooledBigtableGrpcAsyncIOTransport, self._gapic_client.transport ) # keep track of active instances to for warmup on channel refresh - self._active_instances: Set[str] = set() + self._active_instances: Set[_WarmedInstanceKey] = set() # keep track of table objects associated with each instance # only remove instance from _active_instances when all associated tables remove it - self._instance_owners: dict[str, Set[int]] = {} + self._instance_owners: dict[_WarmedInstanceKey, Set[int]] = {} # attempt to start background tasks - self._channel_init_time = time.time() + self._channel_init_time = time.monotonic() self._channel_refresh_tasks: list[asyncio.Task[None]] = [] try: self.start_background_channel_refresh() @@ -186,7 +193,7 @@ async def close(self, timeout: float = 2.0): self._channel_refresh_tasks = [] async def _ping_and_warm_instances( - self, channel: grpc.aio.Channel + self, channel: grpc.aio.Channel, instance_key: _WarmedInstanceKey | None = None ) -> list[GoogleAPICallError | None]: """ Prepares the backend for requests on a channel @@ -194,18 +201,36 @@ async def _ping_and_warm_instances( Pings each Bigtable instance registered in `_active_instances` on the client Args: - channel: grpc channel to ping + - channel: grpc channel to warm + - instance_key: if provided, only warm the instance associated with the key Returns: - sequence of results or exceptions from the ping requests """ + instance_list = ( + [instance_key] if instance_key is not None else self._active_instances + ) ping_rpc = channel.unary_unary( "/google.bigtable.v2.Bigtable/PingAndWarm", request_serializer=PingAndWarmRequest.serialize, ) - tasks = [ping_rpc({"name": n}) for n in self._active_instances] - result = await asyncio.gather(*tasks, return_exceptions=True) + # prepare list of coroutines to run + tasks = [ + ping_rpc( + request={"name": instance_name, "app_profile_id": app_profile_id}, + metadata=[ + ( + "x-goog-request-params", + f"name={instance_name}&app_profile_id={app_profile_id}", + ) + ], + wait_for_ready=True, + ) + for (instance_name, table_name, app_profile_id) in instance_list + ] + # execute coroutines in parallel + result_list = await asyncio.gather(*tasks, return_exceptions=True) # return None in place of empty successful responses - return [r or None for r in result] + return [r or None for r in result_list] async def _manage_channel( self, @@ -236,7 +261,7 @@ async def _manage_channel( first_refresh = self._channel_init_time + random.uniform( refresh_interval_min, refresh_interval_max ) - next_sleep = max(first_refresh - time.time(), 0) + next_sleep = max(first_refresh - time.monotonic(), 0) if next_sleep > 0: # warm the current channel immediately channel = self.transport.channels[channel_idx] @@ -271,14 +296,17 @@ async def _register_instance(self, instance_id: str, owner: Table) -> None: owners call _remove_instance_registration """ instance_name = self._gapic_client.instance_path(self.project, instance_id) - self._instance_owners.setdefault(instance_name, set()).add(id(owner)) + instance_key = _WarmedInstanceKey( + instance_name, owner.table_name, owner.app_profile_id + ) + self._instance_owners.setdefault(instance_key, set()).add(id(owner)) if instance_name not in self._active_instances: - self._active_instances.add(instance_name) + self._active_instances.add(instance_key) if self._channel_refresh_tasks: # refresh tasks already running # call ping and warm on all existing channels for channel in self.transport.channels: - await self._ping_and_warm_instances(channel) + await self._ping_and_warm_instances(channel, instance_key) else: # refresh tasks aren't active. start them as background tasks self.start_background_channel_refresh() @@ -301,11 +329,14 @@ async def _remove_instance_registration( - True if instance was removed """ instance_name = self._gapic_client.instance_path(self.project, instance_id) - owner_list = self._instance_owners.get(instance_name, set()) + instance_key = _WarmedInstanceKey( + instance_name, owner.table_name, owner.app_profile_id + ) + owner_list = self._instance_owners.get(instance_key, set()) try: owner_list.remove(id(owner)) if len(owner_list) == 0: - self._active_instances.remove(instance_name) + self._active_instances.remove(instance_key) return True except KeyError: return False diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 1ba022ae0..45a3e17d2 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -240,6 +240,7 @@ async def test_ping_and_warm(client, table): assert results[0] is None +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutation_set_cell(table, temp_rows): """ diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py index 2765afe24..9aa1a7bb4 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/test__helpers.py @@ -23,7 +23,7 @@ class TestMakeMetadata: @pytest.mark.parametrize( "table,profile,expected", [ - ("table", "profile", "table_name=table,app_profile_id=profile"), + ("table", "profile", "table_name=table&app_profile_id=profile"), ("table", None, "table_name=table"), ], ) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 706ab973d..805a6340d 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -280,32 +280,87 @@ async def test_start_background_channel_refresh_tasks_names(self): @pytest.mark.asyncio async def test__ping_and_warm_instances(self): - # test with no instances + """ + test ping and warm with mocked asyncio.gather + """ + client_mock = mock.Mock() with mock.patch.object(asyncio, "gather", AsyncMock()) as gather: - client = self._make_one(project="project-id", pool_size=1) - channel = client.transport._grpc_channel._pool[0] - await client._ping_and_warm_instances(channel) + # simulate gather by returning the same number of items as passed in + gather.side_effect = lambda *args, **kwargs: [None for _ in args] + channel = mock.Mock() + # test with no instances + client_mock._active_instances = [] + result = await self._get_target_class()._ping_and_warm_instances( + client_mock, channel + ) + assert len(result) == 0 gather.assert_called_once() gather.assert_awaited_once() assert not gather.call_args.args assert gather.call_args.kwargs == {"return_exceptions": True} # test with instances - client._active_instances = [ - "instance-1", - "instance-2", - "instance-3", - "instance-4", - ] - with mock.patch.object(asyncio, "gather", AsyncMock()) as gather: - await client._ping_and_warm_instances(channel) + client_mock._active_instances = [ + (mock.Mock(), mock.Mock(), mock.Mock()) + ] * 4 + gather.reset_mock() + channel.reset_mock() + result = await self._get_target_class()._ping_and_warm_instances( + client_mock, channel + ) + assert len(result) == 4 gather.assert_called_once() gather.assert_awaited_once() assert len(gather.call_args.args) == 4 - assert gather.call_args.kwargs == {"return_exceptions": True} - for idx, call in enumerate(gather.call_args.args): - assert isinstance(call, grpc.aio.UnaryUnaryCall) - call._request["name"] = client._active_instances[idx] - await client.close() + # check grpc call arguments + grpc_call_args = channel.unary_unary().call_args_list + for idx, (_, kwargs) in enumerate(grpc_call_args): + ( + expected_instance, + expected_table, + expected_app_profile, + ) = client_mock._active_instances[idx] + request = kwargs["request"] + assert request["name"] == expected_instance + assert request["app_profile_id"] == expected_app_profile + metadata = kwargs["metadata"] + assert len(metadata) == 1 + assert metadata[0][0] == "x-goog-request-params" + assert ( + metadata[0][1] + == f"name={expected_instance}&app_profile_id={expected_app_profile}" + ) + + @pytest.mark.asyncio + async def test__ping_and_warm_single_instance(self): + """ + should be able to call ping and warm with single instance + """ + client_mock = mock.Mock() + with mock.patch.object(asyncio, "gather", AsyncMock()) as gather: + # simulate gather by returning the same number of items as passed in + gather.side_effect = lambda *args, **kwargs: [None for _ in args] + channel = mock.Mock() + # test with large set of instances + client_mock._active_instances = [mock.Mock()] * 100 + test_key = ("test-instance", "test-table", "test-app-profile") + result = await self._get_target_class()._ping_and_warm_instances( + client_mock, channel, test_key + ) + # should only have been called with test instance + assert len(result) == 1 + # check grpc call arguments + grpc_call_args = channel.unary_unary().call_args_list + assert len(grpc_call_args) == 1 + kwargs = grpc_call_args[0][1] + request = kwargs["request"] + assert request["name"] == "test-instance" + assert request["app_profile_id"] == "test-app-profile" + metadata = kwargs["metadata"] + assert len(metadata) == 1 + assert metadata[0][0] == "x-goog-request-params" + assert ( + metadata[0][1] == "name=test-instance&app_profile_id=test-app-profile" + ) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -325,7 +380,7 @@ async def test__manage_channel_first_sleep( # first sleep time should be `refresh_interval` seconds after client init import time - with mock.patch.object(time, "time") as time: + with mock.patch.object(time, "monotonic") as time: time.return_value = 0 with mock.patch.object(asyncio, "sleep") as sleep: sleep.side_effect = asyncio.CancelledError @@ -344,46 +399,47 @@ async def test__manage_channel_first_sleep( @pytest.mark.asyncio async def test__manage_channel_ping_and_warm(self): - from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( - PooledBigtableGrpcAsyncIOTransport, - ) + """ + _manage channel should call ping and warm internally + """ + import time + client_mock = mock.Mock() + client_mock._channel_init_time = time.monotonic() + channel_list = [mock.Mock(), mock.Mock()] + client_mock.transport.channels = channel_list + new_channel = mock.Mock() + client_mock.transport.grpc_channel._create_channel.return_value = new_channel # should ping an warm all new channels, and old channels if sleeping - client = self._make_one(project="project-id") - new_channel = grpc.aio.insecure_channel("localhost:8080") with mock.patch.object(asyncio, "sleep"): - create_channel = mock.Mock() - create_channel.return_value = new_channel - client.transport.grpc_channel._create_channel = create_channel - with mock.patch.object( - PooledBigtableGrpcAsyncIOTransport, "replace_channel" - ) as replace_channel: - replace_channel.side_effect = asyncio.CancelledError - # should ping and warm old channel then new if sleep > 0 - with mock.patch.object( - type(self._make_one()), "_ping_and_warm_instances" - ) as ping_and_warm: - try: - channel_idx = 2 - old_channel = client.transport._grpc_channel._pool[channel_idx] - await client._manage_channel(channel_idx, 10) - except asyncio.CancelledError: - pass - assert ping_and_warm.call_count == 2 - assert old_channel != new_channel - called_with = [call[0][0] for call in ping_and_warm.call_args_list] - assert old_channel in called_with - assert new_channel in called_with - # should ping and warm instantly new channel only if not sleeping - with mock.patch.object( - type(self._make_one()), "_ping_and_warm_instances" - ) as ping_and_warm: - try: - await client._manage_channel(0, 0, 0) - except asyncio.CancelledError: - pass - ping_and_warm.assert_called_once_with(new_channel) - await client.close() + # stop process after replace_channel is called + client_mock.transport.replace_channel.side_effect = asyncio.CancelledError + ping_and_warm = client_mock._ping_and_warm_instances = AsyncMock() + # should ping and warm old channel then new if sleep > 0 + try: + channel_idx = 1 + await self._get_target_class()._manage_channel( + client_mock, channel_idx, 10 + ) + except asyncio.CancelledError: + pass + # should have called at loop start, and after replacement + assert ping_and_warm.call_count == 2 + # should have replaced channel once + assert client_mock.transport.replace_channel.call_count == 1 + # make sure new and old channels were warmed + old_channel = channel_list[channel_idx] + assert old_channel != new_channel + called_with = [call[0][0] for call in ping_and_warm.call_args_list] + assert old_channel in called_with + assert new_channel in called_with + # should ping and warm instantly new channel only if not sleeping + ping_and_warm.reset_mock() + try: + await self._get_target_class()._manage_channel(client_mock, 0, 0, 0) + except asyncio.CancelledError: + pass + ping_and_warm.assert_called_once_with(new_channel) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -502,58 +558,134 @@ async def test__manage_channel_refresh(self, num_cycles): await client.close() @pytest.mark.asyncio - @pytest.mark.filterwarnings("ignore::RuntimeWarning") async def test__register_instance(self): - # create the client without calling start_background_channel_refresh - with mock.patch.object(asyncio, "get_running_loop") as get_event_loop: - get_event_loop.side_effect = RuntimeError("no event loop") - client = self._make_one(project="project-id") - assert not client._channel_refresh_tasks + """ + test instance registration + """ + # set up mock client + client_mock = mock.Mock() + client_mock._gapic_client.instance_path.side_effect = lambda a, b: f"prefix/{b}" + active_instances = set() + instance_owners = {} + client_mock._active_instances = active_instances + client_mock._instance_owners = instance_owners + client_mock._channel_refresh_tasks = [] + client_mock.start_background_channel_refresh.side_effect = ( + lambda: client_mock._channel_refresh_tasks.append(mock.Mock) + ) + mock_channels = [mock.Mock() for i in range(5)] + client_mock.transport.channels = mock_channels + client_mock._ping_and_warm_instances = AsyncMock() + table_mock = mock.Mock() + await self._get_target_class()._register_instance( + client_mock, "instance-1", table_mock + ) # first call should start background refresh - assert client._active_instances == set() - await client._register_instance("instance-1", mock.Mock()) - assert len(client._active_instances) == 1 - assert client._active_instances == {"projects/project-id/instances/instance-1"} - assert client._channel_refresh_tasks - # next call should not - with mock.patch.object( - type(self._make_one()), "start_background_channel_refresh" - ) as refresh_mock: - await client._register_instance("instance-2", mock.Mock()) - assert len(client._active_instances) == 2 - assert client._active_instances == { - "projects/project-id/instances/instance-1", - "projects/project-id/instances/instance-2", - } - refresh_mock.assert_not_called() + assert client_mock.start_background_channel_refresh.call_count == 1 + # ensure active_instances and instance_owners were updated properly + expected_key = ( + "prefix/instance-1", + table_mock.table_name, + table_mock.app_profile_id, + ) + assert len(active_instances) == 1 + assert expected_key == tuple(list(active_instances)[0]) + assert len(instance_owners) == 1 + assert expected_key == tuple(list(instance_owners)[0]) + # should be a new task set + assert client_mock._channel_refresh_tasks + # # next call should not call start_background_channel_refresh again + table_mock2 = mock.Mock() + await self._get_target_class()._register_instance( + client_mock, "instance-2", table_mock2 + ) + assert client_mock.start_background_channel_refresh.call_count == 1 + # but it should call ping and warm with new instance key + assert client_mock._ping_and_warm_instances.call_count == len(mock_channels) + for channel in mock_channels: + assert channel in [ + call[0][0] + for call in client_mock._ping_and_warm_instances.call_args_list + ] + # check for updated lists + assert len(active_instances) == 2 + assert len(instance_owners) == 2 + expected_key2 = ( + "prefix/instance-2", + table_mock2.table_name, + table_mock2.app_profile_id, + ) + assert any( + [ + expected_key2 == tuple(list(active_instances)[i]) + for i in range(len(active_instances)) + ] + ) + assert any( + [ + expected_key2 == tuple(list(instance_owners)[i]) + for i in range(len(instance_owners)) + ] + ) @pytest.mark.asyncio - @pytest.mark.filterwarnings("ignore::RuntimeWarning") - async def test__register_instance_ping_and_warm(self): - # should ping and warm each new instance - pool_size = 7 - with mock.patch.object(asyncio, "get_running_loop") as get_event_loop: - get_event_loop.side_effect = RuntimeError("no event loop") - client = self._make_one(project="project-id", pool_size=pool_size) - # first call should start background refresh - assert not client._channel_refresh_tasks - await client._register_instance("instance-1", mock.Mock()) - client = self._make_one(project="project-id", pool_size=pool_size) - assert len(client._channel_refresh_tasks) == pool_size - assert not client._active_instances - # next calls should trigger ping and warm - with mock.patch.object( - type(self._make_one()), "_ping_and_warm_instances" - ) as ping_mock: - # new instance should trigger ping and warm - await client._register_instance("instance-2", mock.Mock()) - assert ping_mock.call_count == pool_size - await client._register_instance("instance-3", mock.Mock()) - assert ping_mock.call_count == pool_size * 2 - # duplcate instances should not trigger ping and warm - await client._register_instance("instance-3", mock.Mock()) - assert ping_mock.call_count == pool_size * 2 - await client.close() + @pytest.mark.parametrize( + "insert_instances,expected_active,expected_owner_keys", + [ + ([("i", "t", None)], [("i", "t", None)], [("i", "t", None)]), + ([("i", "t", "p")], [("i", "t", "p")], [("i", "t", "p")]), + ([("1", "t", "p"), ("1", "t", "p")], [("1", "t", "p")], [("1", "t", "p")]), + ( + [("1", "t", "p"), ("2", "t", "p")], + [("1", "t", "p"), ("2", "t", "p")], + [("1", "t", "p"), ("2", "t", "p")], + ), + ], + ) + async def test__register_instance_state( + self, insert_instances, expected_active, expected_owner_keys + ): + """ + test that active_instances and instance_owners are updated as expected + """ + # set up mock client + client_mock = mock.Mock() + client_mock._gapic_client.instance_path.side_effect = lambda a, b: b + active_instances = set() + instance_owners = {} + client_mock._active_instances = active_instances + client_mock._instance_owners = instance_owners + client_mock._channel_refresh_tasks = [] + client_mock.start_background_channel_refresh.side_effect = ( + lambda: client_mock._channel_refresh_tasks.append(mock.Mock) + ) + mock_channels = [mock.Mock() for i in range(5)] + client_mock.transport.channels = mock_channels + client_mock._ping_and_warm_instances = AsyncMock() + table_mock = mock.Mock() + # register instances + for instance, table, profile in insert_instances: + table_mock.table_name = table + table_mock.app_profile_id = profile + await self._get_target_class()._register_instance( + client_mock, instance, table_mock + ) + assert len(active_instances) == len(expected_active) + assert len(instance_owners) == len(expected_owner_keys) + for expected in expected_active: + assert any( + [ + expected == tuple(list(active_instances)[i]) + for i in range(len(active_instances)) + ] + ) + for expected in expected_owner_keys: + assert any( + [ + expected == tuple(list(instance_owners)[i]) + for i in range(len(instance_owners)) + ] + ) @pytest.mark.asyncio async def test__remove_instance_registration(self): @@ -566,78 +698,118 @@ async def test__remove_instance_registration(self): instance_1_path = client._gapic_client.instance_path( client.project, "instance-1" ) + instance_1_key = (instance_1_path, table.table_name, table.app_profile_id) instance_2_path = client._gapic_client.instance_path( client.project, "instance-2" ) - assert len(client._instance_owners[instance_1_path]) == 1 - assert list(client._instance_owners[instance_1_path])[0] == id(table) - assert len(client._instance_owners[instance_2_path]) == 1 - assert list(client._instance_owners[instance_2_path])[0] == id(table) + instance_2_key = (instance_2_path, table.table_name, table.app_profile_id) + assert len(client._instance_owners[instance_1_key]) == 1 + assert list(client._instance_owners[instance_1_key])[0] == id(table) + assert len(client._instance_owners[instance_2_key]) == 1 + assert list(client._instance_owners[instance_2_key])[0] == id(table) success = await client._remove_instance_registration("instance-1", table) assert success assert len(client._active_instances) == 1 - assert len(client._instance_owners[instance_1_path]) == 0 - assert len(client._instance_owners[instance_2_path]) == 1 - assert client._active_instances == {"projects/project-id/instances/instance-2"} - success = await client._remove_instance_registration("nonexistant", table) + assert len(client._instance_owners[instance_1_key]) == 0 + assert len(client._instance_owners[instance_2_key]) == 1 + assert client._active_instances == {instance_2_key} + success = await client._remove_instance_registration("fake-key", table) assert not success assert len(client._active_instances) == 1 await client.close() @pytest.mark.asyncio async def test__multiple_table_registration(self): + """ + registering with multiple tables with the same key should + add multiple owners to instance_owners, but only keep one copy + of shared key in active_instances + """ + from google.cloud.bigtable.client import _WarmedInstanceKey + async with self._make_one(project="project-id") as client: async with client.get_table("instance_1", "table_1") as table_1: instance_1_path = client._gapic_client.instance_path( client.project, "instance_1" ) - assert len(client._instance_owners[instance_1_path]) == 1 + instance_1_key = _WarmedInstanceKey( + instance_1_path, table_1.table_name, table_1.app_profile_id + ) + assert len(client._instance_owners[instance_1_key]) == 1 assert len(client._active_instances) == 1 - assert id(table_1) in client._instance_owners[instance_1_path] - async with client.get_table("instance_1", "table_2") as table_2: - assert len(client._instance_owners[instance_1_path]) == 2 + assert id(table_1) in client._instance_owners[instance_1_key] + # duplicate table should register in instance_owners under same key + async with client.get_table("instance_1", "table_1") as table_2: + assert len(client._instance_owners[instance_1_key]) == 2 assert len(client._active_instances) == 1 - assert id(table_1) in client._instance_owners[instance_1_path] - assert id(table_2) in client._instance_owners[instance_1_path] - # table_2 should be unregistered, but instance should still be active + assert id(table_1) in client._instance_owners[instance_1_key] + assert id(table_2) in client._instance_owners[instance_1_key] + # unique table should register in instance_owners and active_instances + async with client.get_table("instance_1", "table_3") as table_3: + instance_3_path = client._gapic_client.instance_path( + client.project, "instance_1" + ) + instance_3_key = _WarmedInstanceKey( + instance_3_path, table_3.table_name, table_3.app_profile_id + ) + assert len(client._instance_owners[instance_1_key]) == 2 + assert len(client._instance_owners[instance_3_key]) == 1 + assert len(client._active_instances) == 2 + assert id(table_1) in client._instance_owners[instance_1_key] + assert id(table_2) in client._instance_owners[instance_1_key] + assert id(table_3) in client._instance_owners[instance_3_key] + # sub-tables should be unregistered, but instance should still be active assert len(client._active_instances) == 1 - assert instance_1_path in client._active_instances - assert id(table_2) not in client._instance_owners[instance_1_path] + assert instance_1_key in client._active_instances + assert id(table_2) not in client._instance_owners[instance_1_key] # both tables are gone. instance should be unregistered assert len(client._active_instances) == 0 - assert instance_1_path not in client._active_instances - assert len(client._instance_owners[instance_1_path]) == 0 + assert instance_1_key not in client._active_instances + assert len(client._instance_owners[instance_1_key]) == 0 @pytest.mark.asyncio async def test__multiple_instance_registration(self): + """ + registering with multiple instance keys should update the key + in instance_owners and active_instances + """ + from google.cloud.bigtable.client import _WarmedInstanceKey + async with self._make_one(project="project-id") as client: async with client.get_table("instance_1", "table_1") as table_1: async with client.get_table("instance_2", "table_2") as table_2: instance_1_path = client._gapic_client.instance_path( client.project, "instance_1" ) + instance_1_key = _WarmedInstanceKey( + instance_1_path, table_1.table_name, table_1.app_profile_id + ) instance_2_path = client._gapic_client.instance_path( client.project, "instance_2" ) - assert len(client._instance_owners[instance_1_path]) == 1 - assert len(client._instance_owners[instance_2_path]) == 1 + instance_2_key = _WarmedInstanceKey( + instance_2_path, table_2.table_name, table_2.app_profile_id + ) + assert len(client._instance_owners[instance_1_key]) == 1 + assert len(client._instance_owners[instance_2_key]) == 1 assert len(client._active_instances) == 2 - assert id(table_1) in client._instance_owners[instance_1_path] - assert id(table_2) in client._instance_owners[instance_2_path] + assert id(table_1) in client._instance_owners[instance_1_key] + assert id(table_2) in client._instance_owners[instance_2_key] # instance2 should be unregistered, but instance1 should still be active assert len(client._active_instances) == 1 - assert instance_1_path in client._active_instances - assert len(client._instance_owners[instance_2_path]) == 0 - assert len(client._instance_owners[instance_1_path]) == 1 - assert id(table_1) in client._instance_owners[instance_1_path] + assert instance_1_key in client._active_instances + assert len(client._instance_owners[instance_2_key]) == 0 + assert len(client._instance_owners[instance_1_key]) == 1 + assert id(table_1) in client._instance_owners[instance_1_key] # both tables are gone. instances should both be unregistered assert len(client._active_instances) == 0 - assert len(client._instance_owners[instance_1_path]) == 0 - assert len(client._instance_owners[instance_2_path]) == 0 + assert len(client._instance_owners[instance_1_key]) == 0 + assert len(client._instance_owners[instance_2_key]) == 0 @pytest.mark.asyncio async def test_get_table(self): from google.cloud.bigtable.client import Table + from google.cloud.bigtable.client import _WarmedInstanceKey client = self._make_one(project="project-id") assert not client._active_instances @@ -663,12 +835,17 @@ async def test_get_table(self): ) assert table.app_profile_id == expected_app_profile_id assert table.client is client - assert table.instance_name in client._active_instances + instance_key = _WarmedInstanceKey( + table.instance_name, table.table_name, table.app_profile_id + ) + assert instance_key in client._active_instances + assert client._instance_owners[instance_key] == {id(table)} await client.close() @pytest.mark.asyncio async def test_get_table_context_manager(self): from google.cloud.bigtable.client import Table + from google.cloud.bigtable.client import _WarmedInstanceKey expected_table_id = "table-id" expected_instance_id = "instance-id" @@ -696,7 +873,11 @@ async def test_get_table_context_manager(self): ) assert table.app_profile_id == expected_app_profile_id assert table.client is client - assert table.instance_name in client._active_instances + instance_key = _WarmedInstanceKey( + table.instance_name, table.table_name, table.app_profile_id + ) + assert instance_key in client._active_instances + assert client._instance_owners[instance_key] == {id(table)} assert close_mock.call_count == 1 @pytest.mark.asyncio @@ -787,6 +968,7 @@ class TestTable: async def test_table_ctor(self): from google.cloud.bigtable.client import BigtableDataClient from google.cloud.bigtable.client import Table + from google.cloud.bigtable.client import _WarmedInstanceKey expected_table_id = "table-id" expected_instance_id = "instance-id" @@ -809,7 +991,11 @@ async def test_table_ctor(self): assert table.instance_id == expected_instance_id assert table.app_profile_id == expected_app_profile_id assert table.client is client - assert table.instance_name in client._active_instances + instance_key = _WarmedInstanceKey( + table.instance_name, table.table_name, table.app_profile_id + ) + assert instance_key in client._active_instances + assert client._instance_owners[instance_key] == {id(table)} assert table.default_operation_timeout == expected_operation_timeout assert table.default_per_request_timeout == expected_per_request_timeout # ensure task reaches completion @@ -866,6 +1052,26 @@ def _make_client(self, *args, **kwargs): return BigtableDataClient(*args, **kwargs) + def _make_table(self, *args, **kwargs): + from google.cloud.bigtable.client import Table + + client_mock = mock.Mock() + client_mock._register_instance.side_effect = ( + lambda *args, **kwargs: asyncio.sleep(0) + ) + client_mock._remove_instance_registration.side_effect = ( + lambda *args, **kwargs: asyncio.sleep(0) + ) + kwargs["instance_id"] = kwargs.get( + "instance_id", args[0] if args else "instance" + ) + kwargs["table_id"] = kwargs.get( + "table_id", args[1] if len(args) > 1 else "table" + ) + client_mock._gapic_client.table_path.return_value = kwargs["table_id"] + client_mock._gapic_client.instance_path.return_value = kwargs["instance_id"] + return Table(client_mock, *args, **kwargs) + def _make_stats(self): from google.cloud.bigtable_v2.types import RequestStats from google.cloud.bigtable_v2.types import FullReadStatsView @@ -932,14 +1138,13 @@ async def execute_fn(self, table, *args, **kwargs): @pytest.mark.asyncio async def test_read_rows(self): - client = self._make_client() - table = client.get_table("instance", "table") query = ReadRowsQuery() chunks = [ self._make_chunk(row_key=b"test_1"), self._make_chunk(row_key=b"test_2"), ] - with mock.patch.object(table.client._gapic_client, "read_rows") as read_rows: + async with self._make_table() as table: + read_rows = table.client._gapic_client.read_rows read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( chunks ) @@ -947,18 +1152,16 @@ async def test_read_rows(self): assert len(results) == 2 assert results[0].row_key == b"test_1" assert results[1].row_key == b"test_2" - await client.close() @pytest.mark.asyncio async def test_read_rows_stream(self): - client = self._make_client() - table = client.get_table("instance", "table") query = ReadRowsQuery() chunks = [ self._make_chunk(row_key=b"test_1"), self._make_chunk(row_key=b"test_2"), ] - with mock.patch.object(table.client._gapic_client, "read_rows") as read_rows: + async with self._make_table() as table: + read_rows = table.client._gapic_client.read_rows read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( chunks ) @@ -967,16 +1170,16 @@ async def test_read_rows_stream(self): assert len(results) == 2 assert results[0].row_key == b"test_1" assert results[1].row_key == b"test_2" - await client.close() @pytest.mark.parametrize("include_app_profile", [True, False]) @pytest.mark.asyncio async def test_read_rows_query_matches_request(self, include_app_profile): from google.cloud.bigtable import RowRange - async with self._make_client() as client: - app_profile_id = "app_profile_id" if include_app_profile else None - table = client.get_table("instance", "table", app_profile_id=app_profile_id) + app_profile_id = "app_profile_id" if include_app_profile else None + async with self._make_table(app_profile_id=app_profile_id) as table: + read_rows = table.client._gapic_client.read_rows + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream([]) row_keys = [b"test_1", "test_2"] row_ranges = RowRange("start", "end") filter_ = {"test": "filter"} @@ -987,56 +1190,45 @@ async def test_read_rows_query_matches_request(self, include_app_profile): row_filter=filter_, limit=limit, ) - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( - [] - ) - results = await self.execute_fn(table, query, operation_timeout=3) - assert len(results) == 0 - call_request = read_rows.call_args_list[0][0][0] - query_dict = query._to_dict() - if include_app_profile: - assert set(call_request.keys()) == set(query_dict.keys()) | { - "table_name", - "app_profile_id", - } - else: - assert set(call_request.keys()) == set(query_dict.keys()) | { - "table_name" - } - assert call_request["rows"] == query_dict["rows"] - assert call_request["filter"] == filter_ - assert call_request["rows_limit"] == limit - assert call_request["table_name"] == table.table_name - if include_app_profile: - assert call_request["app_profile_id"] == app_profile_id + + results = await table.read_rows(query, operation_timeout=3) + assert len(results) == 0 + call_request = read_rows.call_args_list[0][0][0] + query_dict = query._to_dict() + if include_app_profile: + assert set(call_request.keys()) == set(query_dict.keys()) | { + "table_name", + "app_profile_id", + } + else: + assert set(call_request.keys()) == set(query_dict.keys()) | { + "table_name" + } + assert call_request["rows"] == query_dict["rows"] + assert call_request["filter"] == filter_ + assert call_request["rows_limit"] == limit + assert call_request["table_name"] == table.table_name + if include_app_profile: + assert call_request["app_profile_id"] == app_profile_id @pytest.mark.parametrize("operation_timeout", [0.001, 0.023, 0.1]) @pytest.mark.asyncio async def test_read_rows_timeout(self, operation_timeout): - async with self._make_client() as client: - async with client.get_table("instance", "table") as table: - query = ReadRowsQuery() - chunks = [self._make_chunk(row_key=b"test_1")] - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = ( - lambda *args, **kwargs: self._make_gapic_stream( - chunks, sleep_time=1 - ) - ) - try: - await self.execute_fn( - table, query, operation_timeout=operation_timeout - ) - except core_exceptions.DeadlineExceeded as e: - assert ( - e.message - == f"operation_timeout of {operation_timeout:0.1f}s exceeded" - ) + + async with self._make_table() as table: + read_rows = table.client._gapic_client.read_rows + query = ReadRowsQuery() + chunks = [self._make_chunk(row_key=b"test_1")] + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + chunks, sleep_time=1 + ) + try: + await table.read_rows(query, operation_timeout=operation_timeout) + except core_exceptions.DeadlineExceeded as e: + assert ( + e.message + == f"operation_timeout of {operation_timeout:0.1f}s exceeded" + ) @pytest.mark.parametrize( "per_request_t, operation_t, expected_num", @@ -1064,49 +1256,42 @@ async def test_read_rows_per_request_timeout( # mocking uniform ensures there are no sleeps between retries with mock.patch("random.uniform", side_effect=lambda a, b: 0): - async with self._make_client() as client: - async with client.get_table("instance", "table") as table: - query = ReadRowsQuery() - chunks = [core_exceptions.DeadlineExceeded("mock deadline")] - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = ( - lambda *args, **kwargs: self._make_gapic_stream( - chunks, sleep_time=per_request_t - ) - ) - try: - await self.execute_fn( - table, - query, - operation_timeout=operation_t, - per_request_timeout=per_request_t, - ) - except core_exceptions.DeadlineExceeded as e: - retry_exc = e.__cause__ - if expected_num == 0: - assert retry_exc is None - else: - assert type(retry_exc) == RetryExceptionGroup - assert f"{expected_num} failed attempts" in str( - retry_exc - ) - assert len(retry_exc.exceptions) == expected_num - for sub_exc in retry_exc.exceptions: - assert sub_exc.message == "mock deadline" - assert read_rows.call_count == expected_num - # check timeouts - for _, call_kwargs in read_rows.call_args_list[:-1]: - assert call_kwargs["timeout"] == per_request_t - # last timeout should be adjusted to account for the time spent - assert ( - abs( - read_rows.call_args_list[-1][1]["timeout"] - - expected_last_timeout - ) - < 0.05 - ) + async with self._make_table() as table: + read_rows = table.client._gapic_client.read_rows + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + chunks, sleep_time=per_request_t + ) + query = ReadRowsQuery() + chunks = [core_exceptions.DeadlineExceeded("mock deadline")] + + try: + await table.read_rows( + query, + operation_timeout=operation_t, + per_request_timeout=per_request_t, + ) + except core_exceptions.DeadlineExceeded as e: + retry_exc = e.__cause__ + if expected_num == 0: + assert retry_exc is None + else: + assert type(retry_exc) == RetryExceptionGroup + assert f"{expected_num} failed attempts" in str(retry_exc) + assert len(retry_exc.exceptions) == expected_num + for sub_exc in retry_exc.exceptions: + assert sub_exc.message == "mock deadline" + assert read_rows.call_count == expected_num + # check timeouts + for _, call_kwargs in read_rows.call_args_list[:-1]: + assert call_kwargs["timeout"] == per_request_t + # last timeout should be adjusted to account for the time spent + assert ( + abs( + read_rows.call_args_list[-1][1]["timeout"] + - expected_last_timeout + ) + < 0.05 + ) @pytest.mark.asyncio async def test_read_rows_idle_timeout(self): @@ -1165,25 +1350,20 @@ async def test_read_rows_idle_timeout(self): ) @pytest.mark.asyncio async def test_read_rows_retryable_error(self, exc_type): - async with self._make_client() as client: - async with client.get_table("instance", "table") as table: - query = ReadRowsQuery() - expected_error = exc_type("mock error") - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = ( - lambda *args, **kwargs: self._make_gapic_stream( - [expected_error] - ) - ) - try: - await self.execute_fn(table, query, operation_timeout=0.1) - except core_exceptions.DeadlineExceeded as e: - retry_exc = e.__cause__ - root_cause = retry_exc.exceptions[0] - assert type(root_cause) == exc_type - assert root_cause == expected_error + async with self._make_table() as table: + read_rows = table.client._gapic_client.read_rows + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + [expected_error] + ) + query = ReadRowsQuery() + expected_error = exc_type("mock error") + try: + await table.read_rows(query, operation_timeout=0.1) + except core_exceptions.DeadlineExceeded as e: + retry_exc = e.__cause__ + root_cause = retry_exc.exceptions[0] + assert type(root_cause) == exc_type + assert root_cause == expected_error @pytest.mark.parametrize( "exc_type", @@ -1201,22 +1381,17 @@ async def test_read_rows_retryable_error(self, exc_type): ) @pytest.mark.asyncio async def test_read_rows_non_retryable_error(self, exc_type): - async with self._make_client() as client: - async with client.get_table("instance", "table") as table: - query = ReadRowsQuery() - expected_error = exc_type("mock error") - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = ( - lambda *args, **kwargs: self._make_gapic_stream( - [expected_error] - ) - ) - try: - await self.execute_fn(table, query, operation_timeout=0.1) - except exc_type as e: - assert e == expected_error + async with self._make_table() as table: + read_rows = table.client._gapic_client.read_rows + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + [expected_error] + ) + query = ReadRowsQuery() + expected_error = exc_type("mock error") + try: + await table.read_rows(query, operation_timeout=0.1) + except exc_type as e: + assert e == expected_error @pytest.mark.asyncio async def test_read_rows_revise_request(self): @@ -1231,36 +1406,26 @@ async def test_read_rows_revise_request(self): ) as revise_rowset: with mock.patch.object(_ReadRowsOperation, "aclose"): revise_rowset.return_value = "modified" - async with self._make_client() as client: - async with client.get_table("instance", "table") as table: - row_keys = [b"test_1", b"test_2", b"test_3"] - query = ReadRowsQuery(row_keys=row_keys) - chunks = [ - self._make_chunk(row_key=b"test_1"), - core_exceptions.Aborted("mock retryable error"), - ] - with mock.patch.object( - table.client._gapic_client, "read_rows" - ) as read_rows: - read_rows.side_effect = ( - lambda *args, **kwargs: self._make_gapic_stream(chunks) - ) - try: - await self.execute_fn(table, query) - except InvalidChunk: - revise_rowset.assert_called() - revise_call_kwargs = revise_rowset.call_args_list[ - 0 - ].kwargs - assert ( - revise_call_kwargs["row_set"] - == query._to_dict()["rows"] - ) - assert ( - revise_call_kwargs["last_seen_row_key"] == b"test_1" - ) - read_rows_request = read_rows.call_args_list[1].args[0] - assert read_rows_request["rows"] == "modified" + async with self._make_table() as table: + read_rows = table.client._gapic_client.read_rows + read_rows.side_effect = ( + lambda *args, **kwargs: self._make_gapic_stream(chunks) + ) + row_keys = [b"test_1", b"test_2", b"test_3"] + query = ReadRowsQuery(row_keys=row_keys) + chunks = [ + self._make_chunk(row_key=b"test_1"), + core_exceptions.Aborted("mock retryable error"), + ] + try: + await table.read_rows(query) + except InvalidChunk: + revise_rowset.assert_called() + revise_call_kwargs = revise_rowset.call_args_list[0].kwargs + assert revise_call_kwargs["row_set"] == query._to_dict()["rows"] + assert revise_call_kwargs["last_seen_row_key"] == b"test_1" + read_rows_request = read_rows.call_args_list[1].args[0] + assert read_rows_request["rows"] == "modified" @pytest.mark.asyncio async def test_read_rows_default_timeouts(self): @@ -1273,20 +1438,17 @@ async def test_read_rows_default_timeouts(self): per_request_timeout = 4 with mock.patch.object(_ReadRowsOperation, "__init__") as mock_op: mock_op.side_effect = RuntimeError("mock error") - async with self._make_client() as client: - async with client.get_table( - "instance", - "table", - default_operation_timeout=operation_timeout, - default_per_request_timeout=per_request_timeout, - ) as table: - try: - await self.execute_fn(table, ReadRowsQuery()) - except RuntimeError: - pass - kwargs = mock_op.call_args_list[0].kwargs - assert kwargs["operation_timeout"] == operation_timeout - assert kwargs["per_request_timeout"] == per_request_timeout + async with self._make_table( + default_operation_timeout=operation_timeout, + default_per_request_timeout=per_request_timeout, + ) as table: + try: + await table.read_rows(ReadRowsQuery()) + except RuntimeError: + pass + kwargs = mock_op.call_args_list[0].kwargs + assert kwargs["operation_timeout"] == operation_timeout + assert kwargs["per_request_timeout"] == per_request_timeout @pytest.mark.asyncio async def test_read_rows_default_timeout_override(self): @@ -1299,25 +1461,20 @@ async def test_read_rows_default_timeout_override(self): per_request_timeout = 4 with mock.patch.object(_ReadRowsOperation, "__init__") as mock_op: mock_op.side_effect = RuntimeError("mock error") - async with self._make_client() as client: - async with client.get_table( - "instance", - "table", - default_operation_timeout=99, - default_per_request_timeout=97, - ) as table: - try: - await self.execute_fn( - table, - ReadRowsQuery(), - operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, - ) - except RuntimeError: - pass - kwargs = mock_op.call_args_list[0].kwargs - assert kwargs["operation_timeout"] == operation_timeout - assert kwargs["per_request_timeout"] == per_request_timeout + async with self._make_table( + default_operation_timeout=99, default_per_request_timeout=97 + ) as table: + try: + await table.read_rows( + ReadRowsQuery(), + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + except RuntimeError: + pass + kwargs = mock_op.call_args_list[0].kwargs + assert kwargs["operation_timeout"] == operation_timeout + assert kwargs["per_request_timeout"] == per_request_timeout @pytest.mark.asyncio async def test_read_row(self): @@ -1476,24 +1633,22 @@ async def test_row_exists_w_invalid_input(self, input_row): async def test_read_rows_metadata(self, include_app_profile): """request should attach metadata headers""" profile = "profile" if include_app_profile else None - async with self._make_client() as client: - async with client.get_table("i", "t", app_profile_id=profile) as table: - with mock.patch.object( - client._gapic_client, "read_rows", AsyncMock() - ) as read_rows: - await table.read_rows(ReadRowsQuery()) - kwargs = read_rows.call_args_list[0].kwargs - metadata = kwargs["metadata"] - goog_metadata = None - for key, value in metadata: - if key == "x-goog-request-params": - goog_metadata = value - assert goog_metadata is not None, "x-goog-request-params not found" - assert "table_name=" + table.table_name in goog_metadata - if include_app_profile: - assert "app_profile_id=profile" in goog_metadata - else: - assert "app_profile_id=" not in goog_metadata + async with self._make_table(app_profile_id=profile) as table: + read_rows = table.client._gapic_client.read_rows + read_rows.return_value = self._make_gapic_stream([]) + await table.read_rows(ReadRowsQuery()) + kwargs = read_rows.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata class TestReadRowsSharded: From 1ecf65fcd1bf843449a70e25686a7a8492e0828b Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 26 Jun 2023 11:56:04 -0700 Subject: [PATCH 13/56] feat: mutate rows batching (#770) --- google/cloud/bigtable/_mutate_rows.py | 13 +- google/cloud/bigtable/client.py | 47 +- google/cloud/bigtable/exceptions.py | 100 +- google/cloud/bigtable/mutations.py | 23 + google/cloud/bigtable/mutations_batcher.py | 465 ++++++++- tests/system/test_system.py | 224 +++- tests/unit/test__mutate_rows.py | 67 +- tests/unit/test_client.py | 12 +- tests/unit/test_exceptions.py | 96 +- tests/unit/test_mutations.py | 41 +- tests/unit/test_mutations_batcher.py | 1097 ++++++++++++++++++++ 11 files changed, 2070 insertions(+), 115 deletions(-) create mode 100644 tests/unit/test_mutations_batcher.py diff --git a/google/cloud/bigtable/_mutate_rows.py b/google/cloud/bigtable/_mutate_rows.py index a422c99b2..e34ebaeb6 100644 --- a/google/cloud/bigtable/_mutate_rows.py +++ b/google/cloud/bigtable/_mutate_rows.py @@ -31,6 +31,9 @@ from google.cloud.bigtable.client import Table from google.cloud.bigtable.mutations import RowMutationEntry +# mutate_rows requests are limited to this value +MUTATE_ROWS_REQUEST_MUTATION_LIMIT = 100_000 + class _MutateRowsIncomplete(RuntimeError): """ @@ -68,6 +71,14 @@ def __init__( - per_request_timeout: the timeoutto use for each mutate_rows attempt, in seconds. If not specified, the request will run until operation_timeout is reached. """ + # check that mutations are within limits + total_mutations = sum(len(entry.mutations) for entry in mutation_entries) + if total_mutations > MUTATE_ROWS_REQUEST_MUTATION_LIMIT: + raise ValueError( + "mutate_rows requests can contain at most " + f"{MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations across " + f"all entries. Found {total_mutations}." + ) # create partial function to pass to trigger rpc call metadata = _make_metadata(table.table_name, table.app_profile_id) self._gapic_fn = functools.partial( @@ -119,7 +130,7 @@ async def start(self): self._handle_entry_error(idx, exc) finally: # raise exception detailing incomplete mutations - all_errors = [] + all_errors: list[Exception] = [] for idx, exc_list in self.errors.items(): if len(exc_list) == 0: raise core_exceptions.ClientError( diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index 3d33eebf9..4ec3cea27 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -20,8 +20,6 @@ Any, Optional, Set, - Callable, - Coroutine, TYPE_CHECKING, ) @@ -60,6 +58,8 @@ from google.cloud.bigtable._mutate_rows import _MutateRowsOperation from google.cloud.bigtable._helpers import _make_metadata from google.cloud.bigtable._helpers import _convert_retry_deadline +from google.cloud.bigtable.mutations_batcher import MutationsBatcher +from google.cloud.bigtable.mutations_batcher import _MB_SIZE from google.cloud.bigtable._helpers import _attempt_timeout_generator from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule @@ -69,7 +69,6 @@ from google.cloud.bigtable.row_filters import RowFilterChain if TYPE_CHECKING: - from google.cloud.bigtable.mutations_batcher import MutationsBatcher from google.cloud.bigtable import RowKeySamples from google.cloud.bigtable import ShardedQuery @@ -753,17 +752,48 @@ async def execute_rpc(): ) return await wrapped_fn() - def mutations_batcher(self, **kwargs) -> MutationsBatcher: + def mutations_batcher( + self, + *, + flush_interval: float | None = 5, + flush_limit_mutation_count: int | None = 1000, + flush_limit_bytes: int = 20 * _MB_SIZE, + flow_control_max_mutation_count: int = 100_000, + flow_control_max_bytes: int = 100 * _MB_SIZE, + batch_operation_timeout: float | None = None, + batch_per_request_timeout: float | None = None, + ) -> MutationsBatcher: """ Returns a new mutations batcher instance. Can be used to iteratively add mutations that are flushed as a group, to avoid excess network calls + Args: + - flush_interval: Automatically flush every flush_interval seconds. If None, + a table default will be used + - flush_limit_mutation_count: Flush immediately after flush_limit_mutation_count + mutations are added across all entries. If None, this limit is ignored. + - flush_limit_bytes: Flush immediately after flush_limit_bytes bytes are added. + - flow_control_max_mutation_count: Maximum number of inflight mutations. + - flow_control_max_bytes: Maximum number of inflight bytes. + - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None, + table default_operation_timeout will be used + - batch_per_request_timeout: timeout for each individual request, in seconds. If None, + table default_per_request_timeout will be used Returns: - a MutationsBatcher context manager that can batch requests """ - return MutationsBatcher(self, **kwargs) + return MutationsBatcher( + self, + flush_interval=flush_interval, + flush_limit_mutation_count=flush_limit_mutation_count, + flush_limit_bytes=flush_limit_bytes, + flow_control_max_mutation_count=flow_control_max_mutation_count, + flow_control_max_bytes=flow_control_max_bytes, + batch_operation_timeout=batch_operation_timeout, + batch_per_request_timeout=batch_per_request_timeout, + ) async def mutate_row( self, @@ -861,10 +891,6 @@ async def bulk_mutate_rows( *, operation_timeout: float | None = 60, per_request_timeout: float | None = None, - on_success: Callable[ - [int, RowMutationEntry], None | Coroutine[None, None, None] - ] - | None = None, ): """ Applies mutations for multiple rows in a single batched request. @@ -890,9 +916,6 @@ async def bulk_mutate_rows( in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted if within operation_timeout budget - - on_success: a callback function that will be called when each mutation - entry is confirmed to be applied successfully. Will be passed the - index and the entry itself. Raises: - MutationsExceptionGroup if one or more mutations fails Contains details about any failed entries in .exceptions diff --git a/google/cloud/bigtable/exceptions.py b/google/cloud/bigtable/exceptions.py index b2cf0ce6b..fc4e368b9 100644 --- a/google/cloud/bigtable/exceptions.py +++ b/google/cloud/bigtable/exceptions.py @@ -85,19 +85,96 @@ def __str__(self): class MutationsExceptionGroup(BigtableExceptionGroup): """ Represents one or more exceptions that occur during a bulk mutation operation + + Exceptions will typically be of type FailedMutationEntryError, but other exceptions may + be included if they are raised during the mutation operation """ @staticmethod - def _format_message(excs: list[FailedMutationEntryError], total_entries: int): - entry_str = "entry" if total_entries == 1 else "entries" - plural_str = "" if len(excs) == 1 else "s" - return f"{len(excs)} sub-exception{plural_str} (from {total_entries} {entry_str} attempted)" + def _format_message( + excs: list[Exception], total_entries: int, exc_count: int | None = None + ) -> str: + """ + Format a message for the exception group + + Args: + - excs: the exceptions in the group + - total_entries: the total number of entries attempted, successful or not + - exc_count: the number of exceptions associated with the request + if None, this will be len(excs) + """ + exc_count = exc_count if exc_count is not None else len(excs) + entry_str = "entry" if exc_count == 1 else "entries" + return f"{exc_count} failed {entry_str} from {total_entries} attempted." + + def __init__( + self, excs: list[Exception], total_entries: int, message: str | None = None + ): + """ + Args: + - excs: the exceptions in the group + - total_entries: the total number of entries attempted, successful or not + - message: the message for the exception group. If None, a default message + will be generated + """ + message = ( + message + if message is not None + else self._format_message(excs, total_entries) + ) + super().__init__(message, excs) + self.total_entries_attempted = total_entries - def __init__(self, excs: list[FailedMutationEntryError], total_entries: int): - super().__init__(self._format_message(excs, total_entries), excs) + def __new__( + cls, excs: list[Exception], total_entries: int, message: str | None = None + ): + """ + Args: + - excs: the exceptions in the group + - total_entries: the total number of entries attempted, successful or not + - message: the message for the exception group. If None, a default message + """ + message = ( + message if message is not None else cls._format_message(excs, total_entries) + ) + instance = super().__new__(cls, message, excs) + instance.total_entries_attempted = total_entries + return instance - def __new__(cls, excs: list[FailedMutationEntryError], total_entries: int): - return super().__new__(cls, cls._format_message(excs, total_entries), excs) + @classmethod + def from_truncated_lists( + cls, + first_list: list[Exception], + last_list: list[Exception], + total_excs: int, + entry_count: int, + ) -> MutationsExceptionGroup: + """ + Create a MutationsExceptionGroup from two lists of exceptions, representing + a larger set that has been truncated. The MutationsExceptionGroup will + contain the union of the two lists as sub-exceptions, and the error message + describe the number of exceptions that were truncated. + + Args: + - first_list: the set of oldest exceptions to add to the ExceptionGroup + - last_list: the set of newest exceptions to add to the ExceptionGroup + - total_excs: the total number of exceptions associated with the request + Should be len(first_list) + len(last_list) + number of dropped exceptions + in the middle + - entry_count: the total number of entries attempted, successful or not + """ + first_count, last_count = len(first_list), len(last_list) + if first_count + last_count >= total_excs: + # no exceptions were dropped + return cls(first_list + last_list, entry_count) + excs = first_list + last_list + truncation_count = total_excs - (first_count + last_count) + base_message = cls._format_message(excs, entry_count, total_excs) + first_message = f"first {first_count}" if first_count else "" + last_message = f"last {last_count}" if last_count else "" + conjunction = " and " if first_message and last_message else "" + message = f"{base_message} ({first_message}{conjunction}{last_message} attached as sub-exceptions; {truncation_count} truncated)" + return cls(excs, entry_count, message) class FailedMutationEntryError(Exception): @@ -108,14 +185,17 @@ class FailedMutationEntryError(Exception): def __init__( self, - failed_idx: int, + failed_idx: int | None, failed_mutation_entry: "RowMutationEntry", cause: Exception, ): idempotent_msg = ( "idempotent" if failed_mutation_entry.is_idempotent() else "non-idempotent" ) - message = f"Failed {idempotent_msg} mutation entry at index {failed_idx} with cause: {cause!r}" + index_msg = f" at index {failed_idx} " if failed_idx is not None else " " + message = ( + f"Failed {idempotent_msg} mutation entry{index_msg}with cause: {cause!r}" + ) super().__init__(message) self.index = failed_idx self.entry = failed_mutation_entry diff --git a/google/cloud/bigtable/mutations.py b/google/cloud/bigtable/mutations.py index fe136f8d9..a4c02cd74 100644 --- a/google/cloud/bigtable/mutations.py +++ b/google/cloud/bigtable/mutations.py @@ -17,6 +17,11 @@ import time from dataclasses import dataclass from abc import ABC, abstractmethod +from sys import getsizeof + +# mutation entries above this should be rejected +from google.cloud.bigtable._mutate_rows import MUTATE_ROWS_REQUEST_MUTATION_LIMIT + from google.cloud.bigtable.read_modify_write_rules import MAX_INCREMENT_VALUE @@ -41,6 +46,12 @@ def is_idempotent(self) -> bool: def __str__(self) -> str: return str(self._to_dict()) + def size(self) -> int: + """ + Get the size of the mutation in bytes + """ + return getsizeof(self._to_dict()) + @classmethod def _from_dict(cls, input_dict: dict[str, Any]) -> Mutation: instance: Mutation | None = None @@ -195,6 +206,12 @@ def __init__(self, row_key: bytes | str, mutations: Mutation | list[Mutation]): row_key = row_key.encode("utf-8") if isinstance(mutations, Mutation): mutations = [mutations] + if len(mutations) == 0: + raise ValueError("mutations must not be empty") + elif len(mutations) > MUTATE_ROWS_REQUEST_MUTATION_LIMIT: + raise ValueError( + f"entries must have <= {MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations" + ) self.row_key = row_key self.mutations = tuple(mutations) @@ -208,6 +225,12 @@ def is_idempotent(self) -> bool: """Check if the mutation is idempotent""" return all(mutation.is_idempotent() for mutation in self.mutations) + def size(self) -> int: + """ + Get the size of the mutation in bytes + """ + return getsizeof(self._to_dict()) + @classmethod def _from_dict(cls, input_dict: dict[str, Any]) -> RowMutationEntry: return RowMutationEntry( diff --git a/google/cloud/bigtable/mutations_batcher.py b/google/cloud/bigtable/mutations_batcher.py index 9681f4382..68c3f9fbe 100644 --- a/google/cloud/bigtable/mutations_batcher.py +++ b/google/cloud/bigtable/mutations_batcher.py @@ -14,17 +14,149 @@ # from __future__ import annotations +from typing import Any, TYPE_CHECKING import asyncio -from typing import TYPE_CHECKING +import atexit +import warnings +from collections import deque +from google.cloud.bigtable.mutations import RowMutationEntry +from google.cloud.bigtable.exceptions import MutationsExceptionGroup +from google.cloud.bigtable.exceptions import FailedMutationEntryError + +from google.cloud.bigtable._mutate_rows import _MutateRowsOperation +from google.cloud.bigtable._mutate_rows import MUTATE_ROWS_REQUEST_MUTATION_LIMIT from google.cloud.bigtable.mutations import Mutation -from google.cloud.bigtable.row_filters import RowFilter if TYPE_CHECKING: from google.cloud.bigtable.client import Table # pragma: no cover -# Type alias used internally for readability. -_row_key_type = bytes +# used to make more readable default values +_MB_SIZE = 1024 * 1024 + + +class _FlowControl: + """ + Manages flow control for batched mutations. Mutations are registered against + the FlowControl object before being sent, which will block if size or count + limits have reached capacity. As mutations completed, they are removed from + the FlowControl object, which will notify any blocked requests that there + is additional capacity. + + Flow limits are not hard limits. If a single mutation exceeds the configured + limits, it will be allowed as a single batch when the capacity is available. + """ + + def __init__( + self, + max_mutation_count: int, + max_mutation_bytes: int, + ): + """ + Args: + - max_mutation_count: maximum number of mutations to send in a single rpc. + This corresponds to individual mutations in a single RowMutationEntry. + - max_mutation_bytes: maximum number of bytes to send in a single rpc. + """ + self._max_mutation_count = max_mutation_count + self._max_mutation_bytes = max_mutation_bytes + if self._max_mutation_count < 1: + raise ValueError("max_mutation_count must be greater than 0") + if self._max_mutation_bytes < 1: + raise ValueError("max_mutation_bytes must be greater than 0") + self._capacity_condition = asyncio.Condition() + self._in_flight_mutation_count = 0 + self._in_flight_mutation_bytes = 0 + + def _has_capacity(self, additional_count: int, additional_size: int) -> bool: + """ + Checks if there is capacity to send a new entry with the given size and count + + FlowControl limits are not hard limits. If a single mutation exceeds + the configured flow limits, it will be sent in a single batch when + previous batches have completed. + + Args: + - additional_count: number of mutations in the pending entry + - additional_size: size of the pending entry + Returns: + - True if there is capacity to send the pending entry, False otherwise + """ + # adjust limits to allow overly large mutations + acceptable_size = max(self._max_mutation_bytes, additional_size) + acceptable_count = max(self._max_mutation_count, additional_count) + # check if we have capacity for new mutation + new_size = self._in_flight_mutation_bytes + additional_size + new_count = self._in_flight_mutation_count + additional_count + return new_size <= acceptable_size and new_count <= acceptable_count + + async def remove_from_flow( + self, mutations: RowMutationEntry | list[RowMutationEntry] + ) -> None: + """ + Removes mutations from flow control. This method should be called once + for each mutation that was sent to add_to_flow, after the corresponding + operation is complete. + + Args: + - mutations: mutation or list of mutations to remove from flow control + """ + if not isinstance(mutations, list): + mutations = [mutations] + total_count = sum(len(entry.mutations) for entry in mutations) + total_size = sum(entry.size() for entry in mutations) + self._in_flight_mutation_count -= total_count + self._in_flight_mutation_bytes -= total_size + # notify any blocked requests that there is additional capacity + async with self._capacity_condition: + self._capacity_condition.notify_all() + + async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]): + """ + Generator function that registers mutations with flow control. As mutations + are accepted into the flow control, they are yielded back to the caller, + to be sent in a batch. If the flow control is at capacity, the generator + will block until there is capacity available. + + Args: + - mutations: list mutations to break up into batches + Yields: + - list of mutations that have reserved space in the flow control. + Each batch contains at least one mutation. + """ + if not isinstance(mutations, list): + mutations = [mutations] + start_idx = 0 + end_idx = 0 + while end_idx < len(mutations): + start_idx = end_idx + batch_mutation_count = 0 + # fill up batch until we hit capacity + async with self._capacity_condition: + while end_idx < len(mutations): + next_entry = mutations[end_idx] + next_size = next_entry.size() + next_count = len(next_entry.mutations) + if ( + self._has_capacity(next_count, next_size) + # make sure not to exceed per-request mutation count limits + and (batch_mutation_count + next_count) + <= MUTATE_ROWS_REQUEST_MUTATION_LIMIT + ): + # room for new mutation; add to batch + end_idx += 1 + batch_mutation_count += next_count + self._in_flight_mutation_bytes += next_size + self._in_flight_mutation_count += next_count + elif start_idx != end_idx: + # we have at least one mutation in the batch, so send it + break + else: + # batch is empty. Block until we have capacity + await self._capacity_condition.wait_for( + lambda: self._has_capacity(next_count, next_size) + ) + yield mutations[start_idx:end_idx] class MutationsBatcher: @@ -35,7 +167,6 @@ class MutationsBatcher: to use as few network requests as required Flushes: - - manually - every flush_interval seconds - after queue reaches flush_count in quantity - after queue reaches flush_size_bytes in storage size @@ -46,61 +177,323 @@ class MutationsBatcher: batcher.add(row, mut) """ - queue: asyncio.Queue[tuple[_row_key_type, list[Mutation]]] - conditional_queues: dict[RowFilter, tuple[list[Mutation], list[Mutation]]] - - MB_SIZE = 1024 * 1024 - def __init__( self, table: "Table", - flush_count: int = 100, - flush_size_bytes: int = 100 * MB_SIZE, - max_mutation_bytes: int = 20 * MB_SIZE, - flush_interval: int = 5, - metadata: list[tuple[str, str]] | None = None, + *, + flush_interval: float | None = 5, + flush_limit_mutation_count: int | None = 1000, + flush_limit_bytes: int = 20 * _MB_SIZE, + flow_control_max_mutation_count: int = 100_000, + flow_control_max_bytes: int = 100 * _MB_SIZE, + batch_operation_timeout: float | None = None, + batch_per_request_timeout: float | None = None, ): - raise NotImplementedError + """ + Args: + - table: Table to preform rpc calls + - flush_interval: Automatically flush every flush_interval seconds. + If None, no time-based flushing is performed. + - flush_limit_mutation_count: Flush immediately after flush_limit_mutation_count + mutations are added across all entries. If None, this limit is ignored. + - flush_limit_bytes: Flush immediately after flush_limit_bytes bytes are added. + - flow_control_max_mutation_count: Maximum number of inflight mutations. + - flow_control_max_bytes: Maximum number of inflight bytes. + - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None, + table default_operation_timeout will be used + - batch_per_request_timeout: timeout for each individual request, in seconds. If None, + table default_per_request_timeout will be used + """ + self._operation_timeout: float = ( + batch_operation_timeout or table.default_operation_timeout + ) + self._per_request_timeout: float = ( + batch_per_request_timeout + or table.default_per_request_timeout + or self._operation_timeout + ) + if self._operation_timeout <= 0: + raise ValueError("batch_operation_timeout must be greater than 0") + if self._per_request_timeout <= 0: + raise ValueError("batch_per_request_timeout must be greater than 0") + if self._per_request_timeout > self._operation_timeout: + raise ValueError( + "batch_per_request_timeout must be less than batch_operation_timeout" + ) + self.closed: bool = False + self._table = table + self._staged_entries: list[RowMutationEntry] = [] + self._staged_count, self._staged_bytes = 0, 0 + self._flow_control = _FlowControl( + flow_control_max_mutation_count, flow_control_max_bytes + ) + self._flush_limit_bytes = flush_limit_bytes + self._flush_limit_count = ( + flush_limit_mutation_count + if flush_limit_mutation_count is not None + else float("inf") + ) + self._flush_timer = self._start_flush_timer(flush_interval) + self._flush_jobs: set[asyncio.Future[None]] = set() + # MutationExceptionGroup reports number of successful entries along with failures + self._entries_processed_since_last_raise: int = 0 + self._exceptions_since_last_raise: int = 0 + # keep track of the first and last _exception_list_limit exceptions + self._exception_list_limit: int = 10 + self._oldest_exceptions: list[Exception] = [] + self._newest_exceptions: deque[Exception] = deque( + maxlen=self._exception_list_limit + ) + # clean up on program exit + atexit.register(self._on_exit) - async def append(self, row_key: str | bytes, mutation: Mutation | list[Mutation]): + def _start_flush_timer(self, interval: float | None) -> asyncio.Future[None]: """ - Add a new mutation to the internal queue + Set up a background task to flush the batcher every interval seconds + + If interval is None, an empty future is returned + + Args: + - flush_interval: Automatically flush every flush_interval seconds. + If None, no time-based flushing is performed. + Returns: + - asyncio.Future that represents the background task """ - raise NotImplementedError + if interval is None or self.closed: + empty_future: asyncio.Future[None] = asyncio.Future() + empty_future.set_result(None) + return empty_future - async def append_conditional( - self, - predicate_filter: RowFilter, - row_key: str | bytes, - if_true_mutations: Mutation | list[Mutation] | None = None, - if_false_mutations: Mutation | list[Mutation] | None = None, - ): + async def timer_routine(self, interval: float): + """ + Triggers new flush tasks every `interval` seconds + """ + while not self.closed: + await asyncio.sleep(interval) + # add new flush task to list + if not self.closed and self._staged_entries: + self._schedule_flush() + + timer_task = asyncio.create_task(timer_routine(self, interval)) + return timer_task + + async def append(self, mutation_entry: RowMutationEntry): + """ + Add a new set of mutations to the internal queue + + TODO: return a future to track completion of this entry + + Args: + - mutation_entry: new entry to add to flush queue + Raises: + - RuntimeError if batcher is closed + - ValueError if an invalid mutation type is added + """ + if self.closed: + raise RuntimeError("Cannot append to closed MutationsBatcher") + if isinstance(mutation_entry, Mutation): # type: ignore + raise ValueError( + f"invalid mutation type: {type(mutation_entry).__name__}. Only RowMutationEntry objects are supported by batcher" + ) + self._staged_entries.append(mutation_entry) + # start a new flush task if limits exceeded + self._staged_count += len(mutation_entry.mutations) + self._staged_bytes += mutation_entry.size() + if ( + self._staged_count >= self._flush_limit_count + or self._staged_bytes >= self._flush_limit_bytes + ): + self._schedule_flush() + # yield to the event loop to allow flush to run + await asyncio.sleep(0) + + def _schedule_flush(self) -> asyncio.Future[None] | None: + """Update the flush task to include the latest staged entries""" + if self._staged_entries: + entries, self._staged_entries = self._staged_entries, [] + self._staged_count, self._staged_bytes = 0, 0 + new_task = self._create_bg_task(self._flush_internal, entries) + new_task.add_done_callback(self._flush_jobs.remove) + self._flush_jobs.add(new_task) + return new_task + return None + + async def _flush_internal(self, new_entries: list[RowMutationEntry]): + """ + Flushes a set of mutations to the server, and updates internal state + + Args: + - new_entries: list of RowMutationEntry objects to flush + """ + # flush new entries + in_process_requests: list[asyncio.Future[list[FailedMutationEntryError]]] = [] + async for batch in self._flow_control.add_to_flow(new_entries): + batch_task = self._create_bg_task(self._execute_mutate_rows, batch) + in_process_requests.append(batch_task) + # wait for all inflight requests to complete + found_exceptions = await self._wait_for_batch_results(*in_process_requests) + # update exception data to reflect any new errors + self._entries_processed_since_last_raise += len(new_entries) + self._add_exceptions(found_exceptions) + + async def _execute_mutate_rows( + self, batch: list[RowMutationEntry] + ) -> list[FailedMutationEntryError]: """ - Apply a different set of mutations based on the outcome of predicate_filter + Helper to execute mutation operation on a batch - Calls check_and_mutate_row internally on flush + Args: + - batch: list of RowMutationEntry objects to send to server + - timeout: timeout in seconds. Used as operation_timeout and per_request_timeout. + If not given, will use table defaults + Returns: + - list of FailedMutationEntryError objects for mutations that failed. + FailedMutationEntryError objects will not contain index information """ - raise NotImplementedError + request = {"table_name": self._table.table_name} + if self._table.app_profile_id: + request["app_profile_id"] = self._table.app_profile_id + try: + operation = _MutateRowsOperation( + self._table.client._gapic_client, + self._table, + batch, + operation_timeout=self._operation_timeout, + per_request_timeout=self._per_request_timeout, + ) + await operation.start() + except MutationsExceptionGroup as e: + # strip index information from exceptions, since it is not useful in a batch context + for subexc in e.exceptions: + subexc.index = None + return list(e.exceptions) + finally: + # mark batch as complete in flow control + await self._flow_control.remove_from_flow(batch) + return [] - async def flush(self): + def _add_exceptions(self, excs: list[Exception]): """ - Send queue over network in as few calls as possible + Add new list of exceptions to internal store. To avoid unbounded memory, + the batcher will store the first and last _exception_list_limit exceptions, + and discard any in between. + """ + self._exceptions_since_last_raise += len(excs) + if excs and len(self._oldest_exceptions) < self._exception_list_limit: + # populate oldest_exceptions with found_exceptions + addition_count = self._exception_list_limit - len(self._oldest_exceptions) + self._oldest_exceptions.extend(excs[:addition_count]) + excs = excs[addition_count:] + if excs: + # populate newest_exceptions with remaining found_exceptions + self._newest_exceptions.extend(excs[-self._exception_list_limit :]) + + def _raise_exceptions(self): + """ + Raise any unreported exceptions from background flush operations Raises: - - MutationsExceptionGroup if any mutation in the batch fails + - MutationsExceptionGroup with all unreported exceptions """ - raise NotImplementedError + if self._oldest_exceptions or self._newest_exceptions: + oldest, self._oldest_exceptions = self._oldest_exceptions, [] + newest = list(self._newest_exceptions) + self._newest_exceptions.clear() + entry_count, self._entries_processed_since_last_raise = ( + self._entries_processed_since_last_raise, + 0, + ) + exc_count, self._exceptions_since_last_raise = ( + self._exceptions_since_last_raise, + 0, + ) + raise MutationsExceptionGroup.from_truncated_lists( + first_list=oldest, + last_list=newest, + total_excs=exc_count, + entry_count=entry_count, + ) async def __aenter__(self): """For context manager API""" - raise NotImplementedError + return self async def __aexit__(self, exc_type, exc, tb): """For context manager API""" - raise NotImplementedError + await self.close() async def close(self): """ Flush queue and clean up resources """ - raise NotImplementedError + self.closed = True + self._flush_timer.cancel() + self._schedule_flush() + if self._flush_jobs: + await asyncio.gather(*self._flush_jobs, return_exceptions=True) + try: + await self._flush_timer + except asyncio.CancelledError: + pass + atexit.unregister(self._on_exit) + # raise unreported exceptions + self._raise_exceptions() + + def _on_exit(self): + """ + Called when program is exited. Raises warning if unflushed mutations remain + """ + if not self.closed and self._staged_entries: + warnings.warn( + f"MutationsBatcher for table {self._table.table_name} was not closed. " + f"{len(self._staged_entries)} Unflushed mutations will not be sent to the server." + ) + + @staticmethod + def _create_bg_task(func, *args, **kwargs) -> asyncio.Future[Any]: + """ + Create a new background task, and return a future + + This method wraps asyncio to make it easier to maintain subclasses + with different concurrency models. + + Args: + - func: function to execute in background task + - *args: positional arguments to pass to func + - **kwargs: keyword arguments to pass to func + Returns: + - Future object representing the background task + """ + return asyncio.create_task(func(*args, **kwargs)) + + @staticmethod + async def _wait_for_batch_results( + *tasks: asyncio.Future[list[FailedMutationEntryError]] | asyncio.Future[None], + ) -> list[Exception]: + """ + Takes in a list of futures representing _execute_mutate_rows tasks, + waits for them to complete, and returns a list of errors encountered. + + Args: + - *tasks: futures representing _execute_mutate_rows or _flush_internal tasks + Returns: + - list of Exceptions encountered by any of the tasks. Errors are expected + to be FailedMutationEntryError, representing a failed mutation operation. + If a task fails with a different exception, it will be included in the + output list. Successful tasks will not be represented in the output list. + """ + if not tasks: + return [] + all_results = await asyncio.gather(*tasks, return_exceptions=True) + found_errors = [] + for result in all_results: + if isinstance(result, Exception): + # will receive direct Exception objects if request task fails + found_errors.append(result) + elif result: + # completed requests will return a list of FailedMutationEntryError + for e in result: + # strip index information + e.index = None + found_errors.extend(result) + return found_errors diff --git a/tests/system/test_system.py b/tests/system/test_system.py index 45a3e17d2..e1771202a 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -16,6 +16,7 @@ import pytest_asyncio import os import asyncio +import uuid from google.api_core import retry from google.api_core.exceptions import ClientError @@ -27,7 +28,10 @@ @pytest.fixture(scope="session") def event_loop(): - return asyncio.get_event_loop() + loop = asyncio.get_event_loop() + yield loop + loop.stop() + loop.close() @pytest.fixture(scope="session") @@ -206,6 +210,27 @@ async def _retrieve_cell_value(table, row_key): return cell.value +async def _create_row_and_mutation( + table, temp_rows, *, start_value=b"start", new_value=b"new_value" +): + """ + Helper to create a new row, and a sample set_cell mutation to change its value + """ + from google.cloud.bigtable.mutations import SetCell + + row_key = uuid.uuid4().hex.encode() + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row( + row_key, family=family, qualifier=qualifier, value=start_value + ) + # ensure cell is initialized + assert (await _retrieve_cell_value(table, row_key)) == start_value + + mutation = SetCell(family=TEST_FAMILY, qualifier=qualifier, new_value=new_value) + return row_key, mutation + + @pytest_asyncio.fixture(scope="function") async def temp_rows(table): builder = TempRowBuilder(table) @@ -213,7 +238,7 @@ async def temp_rows(table): await builder.delete_rows() -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=10) @pytest.mark.asyncio async def test_ping_and_warm_gapic(client, table): """ @@ -246,27 +271,15 @@ async def test_mutation_set_cell(table, temp_rows): """ Ensure cells can be set properly """ - from google.cloud.bigtable.mutations import SetCell - - row_key = b"mutate" - family = TEST_FAMILY - qualifier = b"test-qualifier" - start_value = b"start" - await temp_rows.add_row( - row_key, family=family, qualifier=qualifier, value=start_value - ) - - # ensure cell is initialized - assert (await _retrieve_cell_value(table, row_key)) == start_value - - expected_value = b"new-value" - mutation = SetCell( - family=TEST_FAMILY, qualifier=b"test-qualifier", new_value=expected_value + row_key = b"bulk_mutate" + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await _create_row_and_mutation( + table, temp_rows, new_value=new_value ) await table.mutate_row(row_key, mutation) # ensure cell is updated - assert (await _retrieve_cell_value(table, row_key)) == expected_value + assert (await _retrieve_cell_value(table, row_key)) == new_value @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @@ -290,28 +303,173 @@ async def test_bulk_mutations_set_cell(client, table, temp_rows): """ Ensure cells can be set properly """ - from google.cloud.bigtable.mutations import SetCell, RowMutationEntry + from google.cloud.bigtable.mutations import RowMutationEntry - row_key = b"bulk_mutate" - family = TEST_FAMILY - qualifier = b"test-qualifier" - start_value = b"start" - await temp_rows.add_row( - row_key, family=family, qualifier=qualifier, value=start_value + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await _create_row_and_mutation( + table, temp_rows, new_value=new_value ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) - # ensure cell is initialized - assert (await _retrieve_cell_value(table, row_key)) == start_value + await table.bulk_mutate_rows([bulk_mutation]) + + # ensure cell is updated + assert (await _retrieve_cell_value(table, row_key)) == new_value + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_mutations_batcher_context_manager(client, table, temp_rows): + """ + test batcher with context manager. Should flush on exit + """ + from google.cloud.bigtable.mutations import RowMutationEntry - expected_value = b"new-value" - mutation = SetCell( - family=TEST_FAMILY, qualifier=b"test-qualifier", new_value=expected_value + new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] + row_key, mutation = await _create_row_and_mutation( + table, temp_rows, new_value=new_value + ) + row_key2, mutation2 = await _create_row_and_mutation( + table, temp_rows, new_value=new_value2 ) bulk_mutation = RowMutationEntry(row_key, [mutation]) - await table.bulk_mutate_rows([bulk_mutation]) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + async with table.mutations_batcher() as batcher: + await batcher.append(bulk_mutation) + await batcher.append(bulk_mutation2) # ensure cell is updated - assert (await _retrieve_cell_value(table, row_key)) == expected_value + assert (await _retrieve_cell_value(table, row_key)) == new_value + assert len(batcher._staged_entries) == 0 + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_mutations_batcher_timer_flush(client, table, temp_rows): + """ + batch should occur after flush_interval seconds + """ + from google.cloud.bigtable.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await _create_row_and_mutation( + table, temp_rows, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + flush_interval = 0.1 + async with table.mutations_batcher(flush_interval=flush_interval) as batcher: + await batcher.append(bulk_mutation) + await asyncio.sleep(0) + assert len(batcher._staged_entries) == 1 + await asyncio.sleep(flush_interval + 0.1) + assert len(batcher._staged_entries) == 0 + # ensure cell is updated + assert (await _retrieve_cell_value(table, row_key)) == new_value + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_mutations_batcher_count_flush(client, table, temp_rows): + """ + batch should flush after flush_limit_mutation_count mutations + """ + from google.cloud.bigtable.mutations import RowMutationEntry + + new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] + row_key, mutation = await _create_row_and_mutation( + table, temp_rows, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + row_key2, mutation2 = await _create_row_and_mutation( + table, temp_rows, new_value=new_value2 + ) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + + async with table.mutations_batcher(flush_limit_mutation_count=2) as batcher: + await batcher.append(bulk_mutation) + assert len(batcher._flush_jobs) == 0 + # should be noop; flush not scheduled + assert len(batcher._staged_entries) == 1 + await batcher.append(bulk_mutation2) + # task should now be scheduled + assert len(batcher._flush_jobs) == 1 + await asyncio.gather(*batcher._flush_jobs) + assert len(batcher._staged_entries) == 0 + assert len(batcher._flush_jobs) == 0 + # ensure cells were updated + assert (await _retrieve_cell_value(table, row_key)) == new_value + assert (await _retrieve_cell_value(table, row_key2)) == new_value2 + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_mutations_batcher_bytes_flush(client, table, temp_rows): + """ + batch should flush after flush_limit_bytes bytes + """ + from google.cloud.bigtable.mutations import RowMutationEntry + + new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] + row_key, mutation = await _create_row_and_mutation( + table, temp_rows, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + row_key2, mutation2 = await _create_row_and_mutation( + table, temp_rows, new_value=new_value2 + ) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + + flush_limit = bulk_mutation.size() + bulk_mutation2.size() - 1 + + async with table.mutations_batcher(flush_limit_bytes=flush_limit) as batcher: + await batcher.append(bulk_mutation) + assert len(batcher._flush_jobs) == 0 + assert len(batcher._staged_entries) == 1 + await batcher.append(bulk_mutation2) + # task should now be scheduled + assert len(batcher._flush_jobs) == 1 + assert len(batcher._staged_entries) == 0 + # let flush complete + await asyncio.gather(*batcher._flush_jobs) + # ensure cells were updated + assert (await _retrieve_cell_value(table, row_key)) == new_value + assert (await _retrieve_cell_value(table, row_key2)) == new_value2 + + +@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@pytest.mark.asyncio +async def test_mutations_batcher_no_flush(client, table, temp_rows): + """ + test with no flush requirements met + """ + from google.cloud.bigtable.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + start_value = b"unchanged" + row_key, mutation = await _create_row_and_mutation( + table, temp_rows, start_value=start_value, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + row_key2, mutation2 = await _create_row_and_mutation( + table, temp_rows, start_value=start_value, new_value=new_value + ) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + + size_limit = bulk_mutation.size() + bulk_mutation2.size() + 1 + async with table.mutations_batcher( + flush_limit_bytes=size_limit, flush_limit_mutation_count=3, flush_interval=1 + ) as batcher: + await batcher.append(bulk_mutation) + assert len(batcher._staged_entries) == 1 + await batcher.append(bulk_mutation2) + # flush not scheduled + assert len(batcher._flush_jobs) == 0 + await asyncio.sleep(0.01) + assert len(batcher._staged_entries) == 2 + assert len(batcher._flush_jobs) == 0 + # ensure cells were not updated + assert (await _retrieve_cell_value(table, row_key)) == start_value + assert (await _retrieve_cell_value(table, row_key2)) == start_value @pytest.mark.parametrize( diff --git a/tests/unit/test__mutate_rows.py b/tests/unit/test__mutate_rows.py index 4fba16f23..18b2beede 100644 --- a/tests/unit/test__mutate_rows.py +++ b/tests/unit/test__mutate_rows.py @@ -27,6 +27,13 @@ from mock import AsyncMock # type: ignore +def _make_mutation(count=1, size=1): + mutation = mock.Mock() + mutation.size.return_value = size + mutation.mutations = [mock.Mock()] * count + return mutation + + class TestMutateRowsOperation: def _target_class(self): from google.cloud.bigtable._mutate_rows import _MutateRowsOperation @@ -72,7 +79,7 @@ def test_ctor(self): client = mock.Mock() table = mock.Mock() - entries = [mock.Mock(), mock.Mock()] + entries = [_make_mutation(), _make_mutation()] operation_timeout = 0.05 attempt_timeout = 0.01 instance = self._make_one( @@ -105,6 +112,37 @@ def test_ctor(self): assert instance.remaining_indices == list(range(len(entries))) assert instance.errors == {} + def test_ctor_too_many_entries(self): + """ + should raise an error if an operation is created with more than 100,000 entries + """ + from google.cloud.bigtable._mutate_rows import ( + MUTATE_ROWS_REQUEST_MUTATION_LIMIT, + ) + + assert MUTATE_ROWS_REQUEST_MUTATION_LIMIT == 100_000 + + client = mock.Mock() + table = mock.Mock() + entries = [_make_mutation()] * MUTATE_ROWS_REQUEST_MUTATION_LIMIT + operation_timeout = 0.05 + attempt_timeout = 0.01 + # no errors if at limit + self._make_one(client, table, entries, operation_timeout, attempt_timeout) + # raise error after crossing + with pytest.raises(ValueError) as e: + self._make_one( + client, + table, + entries + [_make_mutation()], + operation_timeout, + attempt_timeout, + ) + assert "mutate_rows requests can contain at most 100000 mutations" in str( + e.value + ) + assert "Found 100001" in str(e.value) + @pytest.mark.asyncio async def test_mutate_rows_operation(self): """ @@ -112,7 +150,7 @@ async def test_mutate_rows_operation(self): """ client = mock.Mock() table = mock.Mock() - entries = [mock.Mock(), mock.Mock()] + entries = [_make_mutation(), _make_mutation()] operation_timeout = 0.05 instance = self._make_one( client, table, entries, operation_timeout, operation_timeout @@ -135,7 +173,7 @@ async def test_mutate_rows_exception(self, exc_type): client = mock.Mock() table = mock.Mock() - entries = [mock.Mock()] + entries = [_make_mutation()] operation_timeout = 0.05 expected_cause = exc_type("abort") with mock.patch.object( @@ -170,7 +208,7 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): client = mock.Mock() table = mock.Mock() - entries = [mock.Mock()] + entries = [_make_mutation()] operation_timeout = 1 expected_cause = exc_type("retry") num_retries = 2 @@ -197,7 +235,7 @@ async def test_mutate_rows_incomplete_ignored(self): client = mock.Mock() table = mock.Mock() - entries = [mock.Mock()] + entries = [_make_mutation()] operation_timeout = 0.05 with mock.patch.object( self._target_class(), @@ -220,12 +258,11 @@ async def test_mutate_rows_incomplete_ignored(self): @pytest.mark.asyncio async def test_run_attempt_single_entry_success(self): """Test mutating a single entry""" - mutation = mock.Mock() - mutations = {0: mutation} + mutation = _make_mutation() expected_timeout = 1.3 - mock_gapic_fn = self._make_mock_gapic(mutations) + mock_gapic_fn = self._make_mock_gapic({0: mutation}) instance = self._make_one( - mutation_entries=mutations, + mutation_entries=[mutation], per_request_timeout=expected_timeout, ) with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): @@ -251,9 +288,9 @@ async def test_run_attempt_partial_success_retryable(self): """Some entries succeed, but one fails. Should report the proper index, and raise incomplete exception""" from google.cloud.bigtable._mutate_rows import _MutateRowsIncomplete - success_mutation = mock.Mock() - success_mutation_2 = mock.Mock() - failure_mutation = mock.Mock() + success_mutation = _make_mutation() + success_mutation_2 = _make_mutation() + failure_mutation = _make_mutation() mutations = [success_mutation, failure_mutation, success_mutation_2] mock_gapic_fn = self._make_mock_gapic(mutations, error_dict={1: 300}) instance = self._make_one( @@ -272,9 +309,9 @@ async def test_run_attempt_partial_success_retryable(self): @pytest.mark.asyncio async def test_run_attempt_partial_success_non_retryable(self): """Some entries succeed, but one fails. Exception marked as non-retryable. Do not raise incomplete error""" - success_mutation = mock.Mock() - success_mutation_2 = mock.Mock() - failure_mutation = mock.Mock() + success_mutation = _make_mutation() + success_mutation_2 = _make_mutation() + failure_mutation = _make_mutation() mutations = [success_mutation, failure_mutation, success_mutation_2] mock_gapic_fn = self._make_mock_gapic(mutations, error_dict={1: 300}) instance = self._make_one( diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 805a6340d..3557c1c16 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -2552,14 +2552,18 @@ async def test_bulk_mutate_row_metadata(self, include_app_profile): async with client.get_table("i", "t", app_profile_id=profile) as table: with mock.patch.object( client._gapic_client, "mutate_rows", AsyncMock() - ) as read_rows: - read_rows.side_effect = core_exceptions.Aborted("mock") + ) as mutate_rows: + mutate_rows.side_effect = core_exceptions.Aborted("mock") + mutation = mock.Mock() + mutation.size.return_value = 1 + entry = mock.Mock() + entry.mutations = [mutation] try: - await table.bulk_mutate_rows([mock.Mock()]) + await table.bulk_mutate_rows([entry]) except Exception: # exception used to end early pass - kwargs = read_rows.call_args_list[0].kwargs + kwargs = mutate_rows.call_args_list[0].kwargs metadata = kwargs["metadata"] goog_metadata = None for key, value in metadata: diff --git a/tests/unit/test_exceptions.py b/tests/unit/test_exceptions.py index e68ccf5e8..ef186a47c 100644 --- a/tests/unit/test_exceptions.py +++ b/tests/unit/test_exceptions.py @@ -136,12 +136,12 @@ def _make_one(self, excs=None, num_entries=3): @pytest.mark.parametrize( "exception_list,total_entries,expected_message", [ - ([Exception()], 1, "1 sub-exception (from 1 entry attempted)"), - ([Exception()], 2, "1 sub-exception (from 2 entries attempted)"), + ([Exception()], 1, "1 failed entry from 1 attempted."), + ([Exception()], 2, "1 failed entry from 2 attempted."), ( [Exception(), RuntimeError()], 2, - "2 sub-exceptions (from 2 entries attempted)", + "2 failed entries from 2 attempted.", ), ], ) @@ -154,6 +154,77 @@ def test_raise(self, exception_list, total_entries, expected_message): assert str(e.value) == expected_message assert list(e.value.exceptions) == exception_list + def test_raise_custom_message(self): + """ + should be able to set a custom error message + """ + custom_message = "custom message" + exception_list = [Exception()] + with pytest.raises(self._get_class()) as e: + raise self._get_class()(exception_list, 5, message=custom_message) + assert str(e.value) == custom_message + assert list(e.value.exceptions) == exception_list + + @pytest.mark.parametrize( + "first_list_len,second_list_len,total_excs,entry_count,expected_message", + [ + (3, 0, 3, 4, "3 failed entries from 4 attempted."), + (1, 0, 1, 2, "1 failed entry from 2 attempted."), + (0, 1, 1, 2, "1 failed entry from 2 attempted."), + (2, 2, 4, 4, "4 failed entries from 4 attempted."), + ( + 1, + 1, + 3, + 2, + "3 failed entries from 2 attempted. (first 1 and last 1 attached as sub-exceptions; 1 truncated)", + ), + ( + 1, + 2, + 100, + 2, + "100 failed entries from 2 attempted. (first 1 and last 2 attached as sub-exceptions; 97 truncated)", + ), + ( + 2, + 1, + 4, + 9, + "4 failed entries from 9 attempted. (first 2 and last 1 attached as sub-exceptions; 1 truncated)", + ), + ( + 3, + 0, + 10, + 10, + "10 failed entries from 10 attempted. (first 3 attached as sub-exceptions; 7 truncated)", + ), + ( + 0, + 3, + 10, + 10, + "10 failed entries from 10 attempted. (last 3 attached as sub-exceptions; 7 truncated)", + ), + ], + ) + def test_from_truncated_lists( + self, first_list_len, second_list_len, total_excs, entry_count, expected_message + ): + """ + Should be able to make MutationsExceptionGroup using a pair of + lists representing a larger truncated list of exceptions + """ + first_list = [Exception()] * first_list_len + second_list = [Exception()] * second_list_len + with pytest.raises(self._get_class()) as e: + raise self._get_class().from_truncated_lists( + first_list, second_list, total_excs, entry_count + ) + assert str(e.value) == expected_message + assert list(e.value.exceptions) == first_list + second_list + class TestRetryExceptionGroup(TestBigtableExceptionGroup): def _get_class(self): @@ -281,6 +352,25 @@ def test_raise_idempotent(self): assert e.value.__cause__ == test_exc assert test_entry.is_idempotent.call_count == 1 + def test_no_index(self): + """ + Instances without an index should display different error string + """ + test_idx = None + test_entry = unittest.mock.Mock() + test_exc = ValueError("test") + with pytest.raises(self._get_class()) as e: + raise self._get_class()(test_idx, test_entry, test_exc) + assert ( + str(e.value) + == "Failed idempotent mutation entry with cause: ValueError('test')" + ) + assert e.value.index == test_idx + assert e.value.entry == test_entry + assert e.value.__cause__ == test_exc + assert isinstance(e.value, Exception) + assert test_entry.is_idempotent.call_count == 1 + class TestFailedQueryShardError: def _get_class(self): diff --git a/tests/unit/test_mutations.py b/tests/unit/test_mutations.py index 5730c53c9..c8c6788b1 100644 --- a/tests/unit/test_mutations.py +++ b/tests/unit/test_mutations.py @@ -45,6 +45,16 @@ def test___str__(self): assert self_mock._to_dict.called assert str_value == str(self_mock._to_dict.return_value) + @pytest.mark.parametrize("test_dict", [{}, {"key": "value"}]) + def test_size(self, test_dict): + from sys import getsizeof + + """Size should return size of dict representation""" + self_mock = mock.Mock() + self_mock._to_dict.return_value = test_dict + size_value = self._target_class().size(self_mock) + assert size_value == getsizeof(test_dict) + @pytest.mark.parametrize( "expected_class,input_dict", [ @@ -494,6 +504,21 @@ def test_ctor(self): assert instance.row_key == expected_key assert list(instance.mutations) == expected_mutations + def test_ctor_over_limit(self): + """Should raise error if mutations exceed MAX_MUTATIONS_PER_ENTRY""" + from google.cloud.bigtable._mutate_rows import ( + MUTATE_ROWS_REQUEST_MUTATION_LIMIT, + ) + + assert MUTATE_ROWS_REQUEST_MUTATION_LIMIT == 100_000 + # no errors at limit + expected_mutations = [None for _ in range(MUTATE_ROWS_REQUEST_MUTATION_LIMIT)] + self._make_one(b"row_key", expected_mutations) + # error if over limit + with pytest.raises(ValueError) as e: + self._make_one("key", expected_mutations + [mock.Mock()]) + assert "entries must have <= 100000 mutations" in str(e.value) + def test_ctor_str_key(self): expected_key = "row_key" expected_mutations = [mock.Mock(), mock.Mock()] @@ -528,7 +553,6 @@ def test__to_dict(self): @pytest.mark.parametrize( "mutations,result", [ - ([], True), ([mock.Mock(is_idempotent=lambda: True)], True), ([mock.Mock(is_idempotent=lambda: False)], False), ( @@ -551,6 +575,21 @@ def test_is_idempotent(self, mutations, result): instance = self._make_one("row_key", mutations) assert instance.is_idempotent() == result + def test_empty_mutations(self): + with pytest.raises(ValueError) as e: + self._make_one("row_key", []) + assert "must not be empty" in str(e.value) + + @pytest.mark.parametrize("test_dict", [{}, {"key": "value"}]) + def test_size(self, test_dict): + from sys import getsizeof + + """Size should return size of dict representation""" + self_mock = mock.Mock() + self_mock._to_dict.return_value = test_dict + size_value = self._target_class().size(self_mock) + assert size_value == getsizeof(test_dict) + def test__from_dict_mock(self): """ test creating instance from entry dict, with mocked mutation._from_dict diff --git a/tests/unit/test_mutations_batcher.py b/tests/unit/test_mutations_batcher.py new file mode 100644 index 000000000..a900468d5 --- /dev/null +++ b/tests/unit/test_mutations_batcher.py @@ -0,0 +1,1097 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +import pytest +import asyncio + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock + from unittest.mock import AsyncMock +except ImportError: # pragma: NO COVER + import mock # type: ignore + from mock import AsyncMock # type: ignore + + +def _make_mutation(count=1, size=1): + mutation = mock.Mock() + mutation.size.return_value = size + mutation.mutations = [mock.Mock()] * count + return mutation + + +class Test_FlowControl: + def _make_one(self, max_mutation_count=10, max_mutation_bytes=100): + from google.cloud.bigtable.mutations_batcher import _FlowControl + + return _FlowControl(max_mutation_count, max_mutation_bytes) + + def test_ctor(self): + max_mutation_count = 9 + max_mutation_bytes = 19 + instance = self._make_one(max_mutation_count, max_mutation_bytes) + assert instance._max_mutation_count == max_mutation_count + assert instance._max_mutation_bytes == max_mutation_bytes + assert instance._in_flight_mutation_count == 0 + assert instance._in_flight_mutation_bytes == 0 + assert isinstance(instance._capacity_condition, asyncio.Condition) + + def test_ctor_invalid_values(self): + """Test that values are positive, and fit within expected limits""" + with pytest.raises(ValueError) as e: + self._make_one(0, 1) + assert "max_mutation_count must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + self._make_one(1, 0) + assert "max_mutation_bytes must be greater than 0" in str(e.value) + + @pytest.mark.parametrize( + "max_count,max_size,existing_count,existing_size,new_count,new_size,expected", + [ + (1, 1, 0, 0, 0, 0, True), + (1, 1, 1, 1, 1, 1, False), + (10, 10, 0, 0, 0, 0, True), + (10, 10, 0, 0, 9, 9, True), + (10, 10, 0, 0, 11, 9, True), + (10, 10, 0, 1, 11, 9, True), + (10, 10, 1, 0, 11, 9, False), + (10, 10, 0, 0, 9, 11, True), + (10, 10, 1, 0, 9, 11, True), + (10, 10, 0, 1, 9, 11, False), + (10, 1, 0, 0, 1, 0, True), + (1, 10, 0, 0, 0, 8, True), + (float("inf"), float("inf"), 0, 0, 1e10, 1e10, True), + (8, 8, 0, 0, 1e10, 1e10, True), + (12, 12, 6, 6, 5, 5, True), + (12, 12, 5, 5, 6, 6, True), + (12, 12, 6, 6, 6, 6, True), + (12, 12, 6, 6, 7, 7, False), + # allow capacity check if new_count or new_size exceeds limits + (12, 12, 0, 0, 13, 13, True), + (12, 12, 12, 0, 0, 13, True), + (12, 12, 0, 12, 13, 0, True), + # but not if there's already values in flight + (12, 12, 1, 1, 13, 13, False), + (12, 12, 1, 1, 0, 13, False), + (12, 12, 1, 1, 13, 0, False), + ], + ) + def test__has_capacity( + self, + max_count, + max_size, + existing_count, + existing_size, + new_count, + new_size, + expected, + ): + """ + _has_capacity should return True if the new mutation will will not exceed the max count or size + """ + instance = self._make_one(max_count, max_size) + instance._in_flight_mutation_count = existing_count + instance._in_flight_mutation_bytes = existing_size + assert instance._has_capacity(new_count, new_size) == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "existing_count,existing_size,added_count,added_size,new_count,new_size", + [ + (0, 0, 0, 0, 0, 0), + (2, 2, 1, 1, 1, 1), + (2, 0, 1, 0, 1, 0), + (0, 2, 0, 1, 0, 1), + (10, 10, 0, 0, 10, 10), + (10, 10, 5, 5, 5, 5), + (0, 0, 1, 1, -1, -1), + ], + ) + async def test_remove_from_flow_value_update( + self, + existing_count, + existing_size, + added_count, + added_size, + new_count, + new_size, + ): + """ + completed mutations should lower the inflight values + """ + instance = self._make_one() + instance._in_flight_mutation_count = existing_count + instance._in_flight_mutation_bytes = existing_size + mutation = _make_mutation(added_count, added_size) + await instance.remove_from_flow(mutation) + assert instance._in_flight_mutation_count == new_count + assert instance._in_flight_mutation_bytes == new_size + + @pytest.mark.asyncio + async def test__remove_from_flow_unlock(self): + """capacity condition should notify after mutation is complete""" + instance = self._make_one(10, 10) + instance._in_flight_mutation_count = 10 + instance._in_flight_mutation_bytes = 10 + + async def task_routine(): + async with instance._capacity_condition: + await instance._capacity_condition.wait_for( + lambda: instance._has_capacity(1, 1) + ) + + task = asyncio.create_task(task_routine()) + await asyncio.sleep(0.05) + # should be blocked due to capacity + assert task.done() is False + # try changing size + mutation = _make_mutation(count=0, size=5) + await instance.remove_from_flow([mutation]) + await asyncio.sleep(0.05) + assert instance._in_flight_mutation_count == 10 + assert instance._in_flight_mutation_bytes == 5 + assert task.done() is False + # try changing count + instance._in_flight_mutation_bytes = 10 + mutation = _make_mutation(count=5, size=0) + await instance.remove_from_flow([mutation]) + await asyncio.sleep(0.05) + assert instance._in_flight_mutation_count == 5 + assert instance._in_flight_mutation_bytes == 10 + assert task.done() is False + # try changing both + instance._in_flight_mutation_count = 10 + mutation = _make_mutation(count=5, size=5) + await instance.remove_from_flow([mutation]) + await asyncio.sleep(0.05) + assert instance._in_flight_mutation_count == 5 + assert instance._in_flight_mutation_bytes == 5 + # task should be complete + assert task.done() is True + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "mutations,count_cap,size_cap,expected_results", + [ + # high capacity results in no batching + ([(5, 5), (1, 1), (1, 1)], 10, 10, [[(5, 5), (1, 1), (1, 1)]]), + # low capacity splits up into batches + ([(1, 1), (1, 1), (1, 1)], 1, 1, [[(1, 1)], [(1, 1)], [(1, 1)]]), + # test count as limiting factor + ([(1, 1), (1, 1), (1, 1)], 2, 10, [[(1, 1), (1, 1)], [(1, 1)]]), + # test size as limiting factor + ([(1, 1), (1, 1), (1, 1)], 10, 2, [[(1, 1), (1, 1)], [(1, 1)]]), + # test with some bloackages and some flows + ( + [(1, 1), (5, 5), (4, 1), (1, 4), (1, 1)], + 5, + 5, + [[(1, 1)], [(5, 5)], [(4, 1), (1, 4)], [(1, 1)]], + ), + ], + ) + async def test_add_to_flow(self, mutations, count_cap, size_cap, expected_results): + """ + Test batching with various flow control settings + """ + mutation_objs = [_make_mutation(count=m[0], size=m[1]) for m in mutations] + instance = self._make_one(count_cap, size_cap) + i = 0 + async for batch in instance.add_to_flow(mutation_objs): + expected_batch = expected_results[i] + assert len(batch) == len(expected_batch) + for j in range(len(expected_batch)): + # check counts + assert len(batch[j].mutations) == expected_batch[j][0] + # check sizes + assert batch[j].size() == expected_batch[j][1] + # update lock + await instance.remove_from_flow(batch) + i += 1 + assert i == len(expected_results) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "mutations,max_limit,expected_results", + [ + ([(1, 1)] * 11, 10, [[(1, 1)] * 10, [(1, 1)]]), + ([(1, 1)] * 10, 1, [[(1, 1)] for _ in range(10)]), + ([(1, 1)] * 10, 2, [[(1, 1), (1, 1)] for _ in range(5)]), + ], + ) + async def test_add_to_flow_max_mutation_limits( + self, mutations, max_limit, expected_results + ): + """ + Test flow control running up against the max API limit + Should submit request early, even if the flow control has room for more + """ + with mock.patch( + "google.cloud.bigtable.mutations_batcher.MUTATE_ROWS_REQUEST_MUTATION_LIMIT", + max_limit, + ): + mutation_objs = [_make_mutation(count=m[0], size=m[1]) for m in mutations] + # flow control has no limits except API restrictions + instance = self._make_one(float("inf"), float("inf")) + i = 0 + async for batch in instance.add_to_flow(mutation_objs): + expected_batch = expected_results[i] + assert len(batch) == len(expected_batch) + for j in range(len(expected_batch)): + # check counts + assert len(batch[j].mutations) == expected_batch[j][0] + # check sizes + assert batch[j].size() == expected_batch[j][1] + # update lock + await instance.remove_from_flow(batch) + i += 1 + assert i == len(expected_results) + + @pytest.mark.asyncio + async def test_add_to_flow_oversize(self): + """ + mutations over the flow control limits should still be accepted + """ + instance = self._make_one(2, 3) + large_size_mutation = _make_mutation(count=1, size=10) + large_count_mutation = _make_mutation(count=10, size=1) + results = [out async for out in instance.add_to_flow([large_size_mutation])] + assert len(results) == 1 + await instance.remove_from_flow(results[0]) + count_results = [ + out async for out in instance.add_to_flow(large_count_mutation) + ] + assert len(count_results) == 1 + + +class TestMutationsBatcher: + def _get_target_class(self): + from google.cloud.bigtable.mutations_batcher import MutationsBatcher + + return MutationsBatcher + + def _make_one(self, table=None, **kwargs): + if table is None: + table = mock.Mock() + table.default_operation_timeout = 10 + table.default_per_request_timeout = 10 + + return self._get_target_class()(table, **kwargs) + + @mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._start_flush_timer" + ) + @pytest.mark.asyncio + async def test_ctor_defaults(self, flush_timer_mock): + flush_timer_mock.return_value = asyncio.create_task(asyncio.sleep(0)) + table = mock.Mock() + table.default_operation_timeout = 10 + table.default_per_request_timeout = 8 + async with self._make_one(table) as instance: + assert instance._table == table + assert instance.closed is False + assert instance._flush_jobs == set() + assert len(instance._staged_entries) == 0 + assert len(instance._oldest_exceptions) == 0 + assert len(instance._newest_exceptions) == 0 + assert instance._exception_list_limit == 10 + assert instance._exceptions_since_last_raise == 0 + assert instance._flow_control._max_mutation_count == 100000 + assert instance._flow_control._max_mutation_bytes == 104857600 + assert instance._flow_control._in_flight_mutation_count == 0 + assert instance._flow_control._in_flight_mutation_bytes == 0 + assert instance._entries_processed_since_last_raise == 0 + assert instance._operation_timeout == table.default_operation_timeout + assert instance._per_request_timeout == table.default_per_request_timeout + await asyncio.sleep(0) + assert flush_timer_mock.call_count == 1 + assert flush_timer_mock.call_args[0][0] == 5 + assert isinstance(instance._flush_timer, asyncio.Future) + + @mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._start_flush_timer", + ) + @pytest.mark.asyncio + async def test_ctor_explicit(self, flush_timer_mock): + """Test with explicit parameters""" + flush_timer_mock.return_value = asyncio.create_task(asyncio.sleep(0)) + table = mock.Mock() + flush_interval = 20 + flush_limit_count = 17 + flush_limit_bytes = 19 + flow_control_max_mutation_count = 1001 + flow_control_max_bytes = 12 + operation_timeout = 11 + per_request_timeout = 2 + async with self._make_one( + table, + flush_interval=flush_interval, + flush_limit_mutation_count=flush_limit_count, + flush_limit_bytes=flush_limit_bytes, + flow_control_max_mutation_count=flow_control_max_mutation_count, + flow_control_max_bytes=flow_control_max_bytes, + batch_operation_timeout=operation_timeout, + batch_per_request_timeout=per_request_timeout, + ) as instance: + assert instance._table == table + assert instance.closed is False + assert instance._flush_jobs == set() + assert len(instance._staged_entries) == 0 + assert len(instance._oldest_exceptions) == 0 + assert len(instance._newest_exceptions) == 0 + assert instance._exception_list_limit == 10 + assert instance._exceptions_since_last_raise == 0 + assert ( + instance._flow_control._max_mutation_count + == flow_control_max_mutation_count + ) + assert instance._flow_control._max_mutation_bytes == flow_control_max_bytes + assert instance._flow_control._in_flight_mutation_count == 0 + assert instance._flow_control._in_flight_mutation_bytes == 0 + assert instance._entries_processed_since_last_raise == 0 + assert instance._operation_timeout == operation_timeout + assert instance._per_request_timeout == per_request_timeout + await asyncio.sleep(0) + assert flush_timer_mock.call_count == 1 + assert flush_timer_mock.call_args[0][0] == flush_interval + assert isinstance(instance._flush_timer, asyncio.Future) + + @mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._start_flush_timer" + ) + @pytest.mark.asyncio + async def test_ctor_no_flush_limits(self, flush_timer_mock): + """Test with None for flush limits""" + flush_timer_mock.return_value = asyncio.create_task(asyncio.sleep(0)) + table = mock.Mock() + table.default_operation_timeout = 10 + table.default_per_request_timeout = 8 + flush_interval = None + flush_limit_count = None + flush_limit_bytes = None + async with self._make_one( + table, + flush_interval=flush_interval, + flush_limit_mutation_count=flush_limit_count, + flush_limit_bytes=flush_limit_bytes, + ) as instance: + assert instance._table == table + assert instance.closed is False + assert instance._staged_entries == [] + assert len(instance._oldest_exceptions) == 0 + assert len(instance._newest_exceptions) == 0 + assert instance._exception_list_limit == 10 + assert instance._exceptions_since_last_raise == 0 + assert instance._flow_control._in_flight_mutation_count == 0 + assert instance._flow_control._in_flight_mutation_bytes == 0 + assert instance._entries_processed_since_last_raise == 0 + await asyncio.sleep(0) + assert flush_timer_mock.call_count == 1 + assert flush_timer_mock.call_args[0][0] is None + assert isinstance(instance._flush_timer, asyncio.Future) + + @pytest.mark.asyncio + async def test_ctor_invalid_values(self): + """Test that timeout values are positive, and fit within expected limits""" + with pytest.raises(ValueError) as e: + self._make_one(batch_operation_timeout=-1) + assert "batch_operation_timeout must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + self._make_one(batch_per_request_timeout=-1) + assert "batch_per_request_timeout must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + self._make_one(batch_operation_timeout=1, batch_per_request_timeout=2) + assert ( + "batch_per_request_timeout must be less than batch_operation_timeout" + in str(e.value) + ) + + def test_default_argument_consistency(self): + """ + We supply default arguments in MutationsBatcher.__init__, and in + table.mutations_batcher. Make sure any changes to defaults are applied to + both places + """ + from google.cloud.bigtable.client import Table + from google.cloud.bigtable.mutations_batcher import MutationsBatcher + import inspect + + get_batcher_signature = dict( + inspect.signature(Table.mutations_batcher).parameters + ) + get_batcher_signature.pop("self") + batcher_init_signature = dict(inspect.signature(MutationsBatcher).parameters) + batcher_init_signature.pop("table") + # both should have same number of arguments + assert len(get_batcher_signature.keys()) == len(batcher_init_signature.keys()) + assert len(get_batcher_signature) == 7 # update if expected params change + # both should have same argument names + assert set(get_batcher_signature.keys()) == set(batcher_init_signature.keys()) + # both should have same default values + for arg_name in get_batcher_signature.keys(): + assert ( + get_batcher_signature[arg_name].default + == batcher_init_signature[arg_name].default + ) + + @mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + ) + @pytest.mark.asyncio + async def test__start_flush_timer_w_None(self, flush_mock): + """Empty timer should return immediately""" + async with self._make_one() as instance: + with mock.patch("asyncio.sleep") as sleep_mock: + await instance._start_flush_timer(None) + assert sleep_mock.call_count == 0 + assert flush_mock.call_count == 0 + + @mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + ) + @pytest.mark.asyncio + async def test__start_flush_timer_call_when_closed(self, flush_mock): + """closed batcher's timer should return immediately""" + async with self._make_one() as instance: + await instance.close() + flush_mock.reset_mock() + with mock.patch("asyncio.sleep") as sleep_mock: + await instance._start_flush_timer(1) + assert sleep_mock.call_count == 0 + assert flush_mock.call_count == 0 + + @mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + ) + @pytest.mark.asyncio + async def test__flush_timer(self, flush_mock): + """Timer should continue to call _schedule_flush in a loop""" + expected_sleep = 12 + async with self._make_one(flush_interval=expected_sleep) as instance: + instance._staged_entries = [mock.Mock()] + loop_num = 3 + with mock.patch("asyncio.sleep") as sleep_mock: + sleep_mock.side_effect = [None] * loop_num + [asyncio.CancelledError()] + try: + await instance._flush_timer + except asyncio.CancelledError: + pass + assert sleep_mock.call_count == loop_num + 1 + sleep_mock.assert_called_with(expected_sleep) + assert flush_mock.call_count == loop_num + + @mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + ) + @pytest.mark.asyncio + async def test__flush_timer_no_mutations(self, flush_mock): + """Timer should not flush if no new mutations have been staged""" + expected_sleep = 12 + async with self._make_one(flush_interval=expected_sleep) as instance: + loop_num = 3 + with mock.patch("asyncio.sleep") as sleep_mock: + sleep_mock.side_effect = [None] * loop_num + [asyncio.CancelledError()] + try: + await instance._flush_timer + except asyncio.CancelledError: + pass + assert sleep_mock.call_count == loop_num + 1 + sleep_mock.assert_called_with(expected_sleep) + assert flush_mock.call_count == 0 + + @mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + ) + @pytest.mark.asyncio + async def test__flush_timer_close(self, flush_mock): + """Timer should continue terminate after close""" + async with self._make_one() as instance: + with mock.patch("asyncio.sleep"): + # let task run in background + await asyncio.sleep(0.5) + assert instance._flush_timer.done() is False + # close the batcher + await instance.close() + await asyncio.sleep(0.1) + # task should be complete + assert instance._flush_timer.done() is True + + @pytest.mark.asyncio + async def test_append_closed(self): + """Should raise exception""" + with pytest.raises(RuntimeError): + instance = self._make_one() + await instance.close() + await instance.append(mock.Mock()) + + @pytest.mark.asyncio + async def test_append_wrong_mutation(self): + """ + Mutation objects should raise an exception. + Only support RowMutationEntry + """ + from google.cloud.bigtable.mutations import DeleteAllFromRow + + async with self._make_one() as instance: + expected_error = "invalid mutation type: DeleteAllFromRow. Only RowMutationEntry objects are supported by batcher" + with pytest.raises(ValueError) as e: + await instance.append(DeleteAllFromRow()) + assert str(e.value) == expected_error + + @pytest.mark.asyncio + async def test_append_outside_flow_limits(self): + """entries larger than mutation limits are still processed""" + async with self._make_one( + flow_control_max_mutation_count=1, flow_control_max_bytes=1 + ) as instance: + oversized_entry = _make_mutation(count=0, size=2) + await instance.append(oversized_entry) + assert instance._staged_entries == [oversized_entry] + assert instance._staged_count == 0 + assert instance._staged_bytes == 2 + instance._staged_entries = [] + async with self._make_one( + flow_control_max_mutation_count=1, flow_control_max_bytes=1 + ) as instance: + overcount_entry = _make_mutation(count=2, size=0) + await instance.append(overcount_entry) + assert instance._staged_entries == [overcount_entry] + assert instance._staged_count == 2 + assert instance._staged_bytes == 0 + instance._staged_entries = [] + + @pytest.mark.asyncio + async def test_append_flush_runs_after_limit_hit(self): + """ + If the user appends a bunch of entries above the flush limits back-to-back, + it should still flush in a single task + """ + from google.cloud.bigtable.mutations_batcher import MutationsBatcher + + with mock.patch.object(MutationsBatcher, "_execute_mutate_rows") as op_mock: + async with self._make_one(flush_limit_bytes=100) as instance: + # mock network calls + async def mock_call(*args, **kwargs): + return [] + + op_mock.side_effect = mock_call + # append a mutation just under the size limit + await instance.append(_make_mutation(size=99)) + # append a bunch of entries back-to-back in a loop + num_entries = 10 + for _ in range(num_entries): + await instance.append(_make_mutation(size=1)) + # let any flush jobs finish + await asyncio.gather(*instance._flush_jobs) + # should have only flushed once, with large mutation and first mutation in loop + assert op_mock.call_count == 1 + sent_batch = op_mock.call_args[0][0] + assert len(sent_batch) == 2 + # others should still be pending + assert len(instance._staged_entries) == num_entries - 1 + + @pytest.mark.parametrize( + "flush_count,flush_bytes,mutation_count,mutation_bytes,expect_flush", + [ + (10, 10, 1, 1, False), + (10, 10, 9, 9, False), + (10, 10, 10, 1, True), + (10, 10, 1, 10, True), + (10, 10, 10, 10, True), + (1, 1, 10, 10, True), + (1, 1, 0, 0, False), + ], + ) + @pytest.mark.asyncio + async def test_append( + self, flush_count, flush_bytes, mutation_count, mutation_bytes, expect_flush + ): + """test appending different mutations, and checking if it causes a flush""" + async with self._make_one( + flush_limit_mutation_count=flush_count, flush_limit_bytes=flush_bytes + ) as instance: + assert instance._staged_count == 0 + assert instance._staged_bytes == 0 + assert instance._staged_entries == [] + mutation = _make_mutation(count=mutation_count, size=mutation_bytes) + with mock.patch.object(instance, "_schedule_flush") as flush_mock: + await instance.append(mutation) + assert flush_mock.call_count == bool(expect_flush) + assert instance._staged_count == mutation_count + assert instance._staged_bytes == mutation_bytes + assert instance._staged_entries == [mutation] + instance._staged_entries = [] + + @pytest.mark.asyncio + async def test_append_multiple_sequentially(self): + """Append multiple mutations""" + async with self._make_one( + flush_limit_mutation_count=8, flush_limit_bytes=8 + ) as instance: + assert instance._staged_count == 0 + assert instance._staged_bytes == 0 + assert instance._staged_entries == [] + mutation = _make_mutation(count=2, size=3) + with mock.patch.object(instance, "_schedule_flush") as flush_mock: + await instance.append(mutation) + assert flush_mock.call_count == 0 + assert instance._staged_count == 2 + assert instance._staged_bytes == 3 + assert len(instance._staged_entries) == 1 + await instance.append(mutation) + assert flush_mock.call_count == 0 + assert instance._staged_count == 4 + assert instance._staged_bytes == 6 + assert len(instance._staged_entries) == 2 + await instance.append(mutation) + assert flush_mock.call_count == 1 + assert instance._staged_count == 6 + assert instance._staged_bytes == 9 + assert len(instance._staged_entries) == 3 + instance._staged_entries = [] + + @pytest.mark.asyncio + async def test_flush_flow_control_concurrent_requests(self): + """ + requests should happen in parallel if flow control breaks up single flush into batches + """ + import time + + num_calls = 10 + fake_mutations = [_make_mutation(count=1) for _ in range(num_calls)] + async with self._make_one(flow_control_max_mutation_count=1) as instance: + with mock.patch.object( + instance, "_execute_mutate_rows", AsyncMock() + ) as op_mock: + # mock network calls + async def mock_call(*args, **kwargs): + await asyncio.sleep(0.1) + return [] + + op_mock.side_effect = mock_call + start_time = time.monotonic() + # flush one large batch, that will be broken up into smaller batches + instance._staged_entries = fake_mutations + instance._schedule_flush() + await asyncio.sleep(0.01) + # make room for new mutations + for i in range(num_calls): + await instance._flow_control.remove_from_flow( + [_make_mutation(count=1)] + ) + await asyncio.sleep(0.01) + # allow flushes to complete + await asyncio.gather(*instance._flush_jobs) + duration = time.monotonic() - start_time + assert len(instance._oldest_exceptions) == 0 + assert len(instance._newest_exceptions) == 0 + # if flushes were sequential, total duration would be 1s + assert duration < 0.25 + assert op_mock.call_count == num_calls + + @pytest.mark.asyncio + async def test_schedule_flush_no_mutations(self): + """schedule flush should return None if no staged mutations""" + async with self._make_one() as instance: + with mock.patch.object(instance, "_flush_internal") as flush_mock: + for i in range(3): + assert instance._schedule_flush() is None + assert flush_mock.call_count == 0 + + @pytest.mark.asyncio + async def test_schedule_flush_with_mutations(self): + """if new mutations exist, should add a new flush task to _flush_jobs""" + async with self._make_one() as instance: + with mock.patch.object(instance, "_flush_internal") as flush_mock: + for i in range(1, 4): + mutation = mock.Mock() + instance._staged_entries = [mutation] + instance._schedule_flush() + assert instance._staged_entries == [] + # let flush task run + await asyncio.sleep(0) + assert instance._staged_entries == [] + assert instance._staged_count == 0 + assert instance._staged_bytes == 0 + assert flush_mock.call_count == i + + @pytest.mark.asyncio + async def test__flush_internal(self): + """ + _flush_internal should: + - await previous flush call + - delegate batching to _flow_control + - call _execute_mutate_rows on each batch + - update self.exceptions and self._entries_processed_since_last_raise + """ + num_entries = 10 + async with self._make_one() as instance: + with mock.patch.object(instance, "_execute_mutate_rows") as execute_mock: + with mock.patch.object( + instance._flow_control, "add_to_flow" + ) as flow_mock: + # mock flow control to always return a single batch + async def gen(x): + yield x + + flow_mock.side_effect = lambda x: gen(x) + mutations = [_make_mutation(count=1, size=1)] * num_entries + await instance._flush_internal(mutations) + assert instance._entries_processed_since_last_raise == num_entries + assert execute_mock.call_count == 1 + assert flow_mock.call_count == 1 + instance._oldest_exceptions.clear() + instance._newest_exceptions.clear() + + @pytest.mark.asyncio + async def test_flush_clears_job_list(self): + """ + a job should be added to _flush_jobs when _schedule_flush is called, + and removed when it completes + """ + async with self._make_one() as instance: + with mock.patch.object(instance, "_flush_internal", AsyncMock()): + mutations = [_make_mutation(count=1, size=1)] + instance._staged_entries = mutations + assert instance._flush_jobs == set() + new_job = instance._schedule_flush() + assert instance._flush_jobs == {new_job} + await new_job + assert instance._flush_jobs == set() + + @pytest.mark.parametrize( + "num_starting,num_new_errors,expected_total_errors", + [ + (0, 0, 0), + (0, 1, 1), + (0, 2, 2), + (1, 0, 1), + (1, 1, 2), + (10, 2, 12), + (10, 20, 20), # should cap at 20 + ], + ) + @pytest.mark.asyncio + async def test__flush_internal_with_errors( + self, num_starting, num_new_errors, expected_total_errors + ): + """ + errors returned from _execute_mutate_rows should be added to internal exceptions + """ + from google.cloud.bigtable import exceptions + + num_entries = 10 + expected_errors = [ + exceptions.FailedMutationEntryError(mock.Mock(), mock.Mock(), ValueError()) + ] * num_new_errors + async with self._make_one() as instance: + instance._oldest_exceptions = [mock.Mock()] * num_starting + with mock.patch.object(instance, "_execute_mutate_rows") as execute_mock: + execute_mock.return_value = expected_errors + with mock.patch.object( + instance._flow_control, "add_to_flow" + ) as flow_mock: + # mock flow control to always return a single batch + async def gen(x): + yield x + + flow_mock.side_effect = lambda x: gen(x) + mutations = [_make_mutation(count=1, size=1)] * num_entries + await instance._flush_internal(mutations) + assert instance._entries_processed_since_last_raise == num_entries + assert execute_mock.call_count == 1 + assert flow_mock.call_count == 1 + found_exceptions = instance._oldest_exceptions + list( + instance._newest_exceptions + ) + assert len(found_exceptions) == expected_total_errors + for i in range(num_starting, expected_total_errors): + assert found_exceptions[i] == expected_errors[i - num_starting] + # errors should have index stripped + assert found_exceptions[i].index is None + # clear out exceptions + instance._oldest_exceptions.clear() + instance._newest_exceptions.clear() + + async def _mock_gapic_return(self, num=5): + from google.cloud.bigtable_v2.types import MutateRowsResponse + from google.rpc import status_pb2 + + async def gen(num): + for i in range(num): + entry = MutateRowsResponse.Entry( + index=i, status=status_pb2.Status(code=0) + ) + yield MutateRowsResponse(entries=[entry]) + + return gen(num) + + @pytest.mark.asyncio + async def test_timer_flush_end_to_end(self): + """Flush should automatically trigger after flush_interval""" + num_nutations = 10 + mutations = [_make_mutation(count=2, size=2)] * num_nutations + + async with self._make_one(flush_interval=0.05) as instance: + instance._table.default_operation_timeout = 10 + instance._table.default_per_request_timeout = 9 + with mock.patch.object( + instance._table.client._gapic_client, "mutate_rows" + ) as gapic_mock: + gapic_mock.side_effect = ( + lambda *args, **kwargs: self._mock_gapic_return(num_nutations) + ) + for m in mutations: + await instance.append(m) + assert instance._entries_processed_since_last_raise == 0 + # let flush trigger due to timer + await asyncio.sleep(0.1) + assert instance._entries_processed_since_last_raise == num_nutations + + @pytest.mark.asyncio + @mock.patch( + "google.cloud.bigtable.mutations_batcher._MutateRowsOperation", + ) + async def test__execute_mutate_rows(self, mutate_rows): + mutate_rows.return_value = AsyncMock() + start_operation = mutate_rows().start + table = mock.Mock() + table.table_name = "test-table" + table.app_profile_id = "test-app-profile" + table.default_operation_timeout = 17 + table.default_per_request_timeout = 13 + async with self._make_one(table) as instance: + batch = [_make_mutation()] + result = await instance._execute_mutate_rows(batch) + assert start_operation.call_count == 1 + args, kwargs = mutate_rows.call_args + assert args[0] == table.client._gapic_client + assert args[1] == table + assert args[2] == batch + kwargs["operation_timeout"] == 17 + kwargs["per_request_timeout"] == 13 + assert result == [] + + @pytest.mark.asyncio + @mock.patch("google.cloud.bigtable.mutations_batcher._MutateRowsOperation.start") + async def test__execute_mutate_rows_returns_errors(self, mutate_rows): + """Errors from operation should be retruned as list""" + from google.cloud.bigtable.exceptions import ( + MutationsExceptionGroup, + FailedMutationEntryError, + ) + + err1 = FailedMutationEntryError(0, mock.Mock(), RuntimeError("test error")) + err2 = FailedMutationEntryError(1, mock.Mock(), RuntimeError("test error")) + mutate_rows.side_effect = MutationsExceptionGroup([err1, err2], 10) + table = mock.Mock() + table.default_operation_timeout = 17 + table.default_per_request_timeout = 13 + async with self._make_one(table) as instance: + batch = [_make_mutation()] + result = await instance._execute_mutate_rows(batch) + assert len(result) == 2 + assert result[0] == err1 + assert result[1] == err2 + # indices should be set to None + assert result[0].index is None + assert result[1].index is None + + @pytest.mark.asyncio + async def test__raise_exceptions(self): + """Raise exceptions and reset error state""" + from google.cloud.bigtable import exceptions + + expected_total = 1201 + expected_exceptions = [RuntimeError("mock")] * 3 + async with self._make_one() as instance: + instance._oldest_exceptions = expected_exceptions + instance._entries_processed_since_last_raise = expected_total + try: + instance._raise_exceptions() + except exceptions.MutationsExceptionGroup as exc: + assert list(exc.exceptions) == expected_exceptions + assert str(expected_total) in str(exc) + assert instance._entries_processed_since_last_raise == 0 + instance._oldest_exceptions, instance._newest_exceptions = ([], []) + # try calling again + instance._raise_exceptions() + + @pytest.mark.asyncio + async def test___aenter__(self): + """Should return self""" + async with self._make_one() as instance: + assert await instance.__aenter__() == instance + + @pytest.mark.asyncio + async def test___aexit__(self): + """aexit should call close""" + async with self._make_one() as instance: + with mock.patch.object(instance, "close") as close_mock: + await instance.__aexit__(None, None, None) + assert close_mock.call_count == 1 + + @pytest.mark.asyncio + async def test_close(self): + """Should clean up all resources""" + async with self._make_one() as instance: + with mock.patch.object(instance, "_schedule_flush") as flush_mock: + with mock.patch.object(instance, "_raise_exceptions") as raise_mock: + await instance.close() + assert instance.closed is True + assert instance._flush_timer.done() is True + assert instance._flush_jobs == set() + assert flush_mock.call_count == 1 + assert raise_mock.call_count == 1 + + @pytest.mark.asyncio + async def test_close_w_exceptions(self): + """Raise exceptions on close""" + from google.cloud.bigtable import exceptions + + expected_total = 10 + expected_exceptions = [RuntimeError("mock")] + async with self._make_one() as instance: + instance._oldest_exceptions = expected_exceptions + instance._entries_processed_since_last_raise = expected_total + try: + await instance.close() + except exceptions.MutationsExceptionGroup as exc: + assert list(exc.exceptions) == expected_exceptions + assert str(expected_total) in str(exc) + assert instance._entries_processed_since_last_raise == 0 + # clear out exceptions + instance._oldest_exceptions, instance._newest_exceptions = ([], []) + + @pytest.mark.asyncio + async def test__on_exit(self, recwarn): + """Should raise warnings if unflushed mutations exist""" + async with self._make_one() as instance: + # calling without mutations is noop + instance._on_exit() + assert len(recwarn) == 0 + # calling with existing mutations should raise warning + num_left = 4 + instance._staged_entries = [mock.Mock()] * num_left + with pytest.warns(UserWarning) as w: + instance._on_exit() + assert len(w) == 1 + assert "unflushed mutations" in str(w[0].message).lower() + assert str(num_left) in str(w[0].message) + # calling while closed is noop + instance.closed = True + instance._on_exit() + assert len(recwarn) == 0 + # reset staged mutations for cleanup + instance._staged_entries = [] + + @pytest.mark.asyncio + async def test_atexit_registration(self): + """Should run _on_exit on program termination""" + import atexit + + with mock.patch( + "google.cloud.bigtable.mutations_batcher.MutationsBatcher._on_exit" + ) as on_exit_mock: + async with self._make_one(): + assert on_exit_mock.call_count == 0 + atexit._run_exitfuncs() + assert on_exit_mock.call_count == 1 + # should not call after close + atexit._run_exitfuncs() + assert on_exit_mock.call_count == 1 + + @pytest.mark.asyncio + @mock.patch( + "google.cloud.bigtable.mutations_batcher._MutateRowsOperation", + ) + async def test_timeout_args_passed(self, mutate_rows): + """ + batch_operation_timeout and batch_per_request_timeout should be used + in api calls + """ + mutate_rows.return_value = AsyncMock() + expected_operation_timeout = 17 + expected_per_request_timeout = 13 + async with self._make_one( + batch_operation_timeout=expected_operation_timeout, + batch_per_request_timeout=expected_per_request_timeout, + ) as instance: + assert instance._operation_timeout == expected_operation_timeout + assert instance._per_request_timeout == expected_per_request_timeout + # make simulated gapic call + await instance._execute_mutate_rows([_make_mutation()]) + assert mutate_rows.call_count == 1 + kwargs = mutate_rows.call_args[1] + assert kwargs["operation_timeout"] == expected_operation_timeout + assert kwargs["per_request_timeout"] == expected_per_request_timeout + + @pytest.mark.parametrize( + "limit,in_e,start_e,end_e", + [ + (10, 0, (10, 0), (10, 0)), + (1, 10, (0, 0), (1, 1)), + (10, 1, (0, 0), (1, 0)), + (10, 10, (0, 0), (10, 0)), + (10, 11, (0, 0), (10, 1)), + (3, 20, (0, 0), (3, 3)), + (10, 20, (0, 0), (10, 10)), + (10, 21, (0, 0), (10, 10)), + (2, 1, (2, 0), (2, 1)), + (2, 1, (1, 0), (2, 0)), + (2, 2, (1, 0), (2, 1)), + (3, 1, (3, 1), (3, 2)), + (3, 3, (3, 1), (3, 3)), + (1000, 5, (999, 0), (1000, 4)), + (1000, 5, (0, 0), (5, 0)), + (1000, 5, (1000, 0), (1000, 5)), + ], + ) + def test__add_exceptions(self, limit, in_e, start_e, end_e): + """ + Test that the _add_exceptions function properly updates the + _oldest_exceptions and _newest_exceptions lists + Args: + - limit: the _exception_list_limit representing the max size of either list + - in_e: size of list of exceptions to send to _add_exceptions + - start_e: a tuple of ints representing the initial sizes of _oldest_exceptions and _newest_exceptions + - end_e: a tuple of ints representing the expected sizes of _oldest_exceptions and _newest_exceptions + """ + from collections import deque + + input_list = [RuntimeError(f"mock {i}") for i in range(in_e)] + mock_batcher = mock.Mock() + mock_batcher._oldest_exceptions = [ + RuntimeError(f"starting mock {i}") for i in range(start_e[0]) + ] + mock_batcher._newest_exceptions = deque( + [RuntimeError(f"starting mock {i}") for i in range(start_e[1])], + maxlen=limit, + ) + mock_batcher._exception_list_limit = limit + mock_batcher._exceptions_since_last_raise = 0 + self._get_target_class()._add_exceptions(mock_batcher, input_list) + assert len(mock_batcher._oldest_exceptions) == end_e[0] + assert len(mock_batcher._newest_exceptions) == end_e[1] + assert mock_batcher._exceptions_since_last_raise == in_e + # make sure that the right items ended up in the right spots + # should fill the oldest slots first + oldest_list_diff = end_e[0] - start_e[0] + # new items should by added on top of the starting list + newest_list_diff = min(max(in_e - oldest_list_diff, 0), limit) + for i in range(oldest_list_diff): + assert mock_batcher._oldest_exceptions[i + start_e[0]] == input_list[i] + # then, the newest slots should be filled with the last items of the input list + for i in range(1, newest_list_diff + 1): + assert mock_batcher._newest_exceptions[-i] == input_list[-i] From eedde1e22f25678a5b43db162245e93c32904ae0 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 28 Jun 2023 10:22:24 -0700 Subject: [PATCH 14/56] chore: restructure module paths (#816) --- .github/.OwlBot.lock.yaml | 3 +- .kokoro/requirements.in | 2 +- .kokoro/requirements.txt | 62 +- .release-please-manifest.json | 2 +- CHANGELOG.md | 43 + docs/app-profile.rst | 2 +- docs/backup.rst | 2 +- docs/batcher.rst | 6 + docs/client-intro.rst | 18 +- docs/client.rst | 2 +- docs/cluster.rst | 2 +- docs/column-family.rst | 22 +- docs/data-api.rst | 82 +- docs/encryption-info.rst | 2 +- docs/instance-api.rst | 32 +- docs/instance.rst | 2 +- docs/row-data.rst | 2 +- docs/row-filters.rst | 12 +- docs/row-set.rst | 2 +- docs/row.rst | 2 +- docs/snippets.py | 118 +- docs/snippets_table.py | 154 +- docs/table-api.rst | 40 +- docs/table.rst | 2 +- docs/usage.rst | 17 +- google/cloud/bigtable/__init__.py | 48 +- .../bigtable/{deprecated => }/app_profile.py | 8 +- .../cloud/bigtable/{deprecated => }/backup.py | 20 +- google/cloud/bigtable/batcher.py | 395 +++++ google/cloud/bigtable/client.py | 1412 +++++---------- .../bigtable/{deprecated => }/cluster.py | 22 +- .../{deprecated => }/column_family.py | 2 +- google/cloud/bigtable/data/__init__.py | 78 + google/cloud/bigtable/data/_async/__init__.py | 26 + .../{ => data/_async}/_mutate_rows.py | 34 +- .../cloud/bigtable/data/_async/_read_rows.py | 403 +++++ google/cloud/bigtable/data/_async/client.py | 1091 ++++++++++++ .../{ => data/_async}/mutations_batcher.py | 26 +- google/cloud/bigtable/{ => data}/_helpers.py | 2 +- .../_read_rows_state_machine.py} | 280 +-- .../cloud/bigtable/{ => data}/exceptions.py | 22 +- google/cloud/bigtable/{ => data}/mutations.py | 10 +- .../{ => data}/read_modify_write_rules.py | 0 .../bigtable/{ => data}/read_rows_query.py | 6 +- google/cloud/bigtable/data/row.py | 465 +++++ .../{deprecated => data}/row_filters.py | 526 +++--- google/cloud/bigtable/deprecated/batcher.py | 146 -- google/cloud/bigtable/deprecated/client.py | 521 ------ google/cloud/bigtable/deprecated/py.typed | 2 - google/cloud/bigtable/deprecated/row.py | 1267 -------------- .../{deprecated => }/encryption_info.py | 4 +- .../cloud/bigtable/{deprecated => }/enums.py | 0 .../cloud/bigtable/{deprecated => }/error.py | 0 google/cloud/bigtable/gapic_version.py | 2 +- .../bigtable/{deprecated => }/instance.py | 68 +- google/cloud/bigtable/iterators.py | 129 -- .../cloud/bigtable/{deprecated => }/policy.py | 0 google/cloud/bigtable/row.py | 1558 +++++++++++++---- .../bigtable/{deprecated => }/row_data.py | 10 +- google/cloud/bigtable/row_filters.py | 526 +++--- .../bigtable/{deprecated => }/row_merger.py | 2 +- .../bigtable/{deprecated => }/row_set.py | 0 .../cloud/bigtable/{deprecated => }/table.py | 68 +- google/cloud/bigtable_admin/__init__.py | 2 + google/cloud/bigtable_admin/gapic_version.py | 2 +- google/cloud/bigtable_admin_v2/__init__.py | 2 + .../cloud/bigtable_admin_v2/gapic_version.py | 2 +- .../bigtable_instance_admin/async_client.py | 13 +- .../bigtable_instance_admin/client.py | 13 +- .../transports/rest.py | 18 - .../bigtable_table_admin/async_client.py | 27 +- .../services/bigtable_table_admin/client.py | 27 +- .../bigtable_table_admin/transports/rest.py | 21 - .../cloud/bigtable_admin_v2/types/__init__.py | 2 + .../types/bigtable_table_admin.py | 20 +- google/cloud/bigtable_admin_v2/types/table.py | 31 + google/cloud/bigtable_v2/__init__.py | 4 + google/cloud/bigtable_v2/gapic_version.py | 2 +- .../services/bigtable/async_client.py | 21 +- .../bigtable_v2/services/bigtable/client.py | 21 +- .../services/bigtable/transports/rest.py | 9 - google/cloud/bigtable_v2/types/__init__.py | 6 + google/cloud/bigtable_v2/types/bigtable.py | 59 +- .../cloud/bigtable_v2/types/feature_flags.py | 54 + noxfile.py | 14 +- python-api-core | 2 +- samples/beam/requirements-test.txt | 2 +- samples/beam/requirements.txt | 2 +- samples/hello/README.md | 2 +- samples/hello/main.py | 6 +- samples/hello/requirements-test.txt | 2 +- samples/hello_happybase/requirements-test.txt | 2 +- samples/instanceadmin/requirements-test.txt | 2 +- samples/metricscaler/requirements-test.txt | 4 +- samples/metricscaler/requirements.txt | 2 +- samples/quickstart/requirements-test.txt | 2 +- .../requirements-test.txt | 2 +- samples/snippets/deletes/deletes_snippets.py | 2 +- .../snippets/deletes/requirements-test.txt | 2 +- .../snippets/filters/requirements-test.txt | 2 +- samples/snippets/reads/requirements-test.txt | 2 +- samples/snippets/writes/requirements-test.txt | 2 +- samples/tableadmin/requirements-test.txt | 2 +- setup.py | 2 +- testing/constraints-3.7.txt | 5 +- .../system/data}/__init__.py | 16 +- tests/system/{ => data}/test_system.py | 66 +- tests/system/v2_client/conftest.py | 2 +- tests/system/v2_client/test_data_api.py | 20 +- tests/system/v2_client/test_instance_admin.py | 6 +- tests/system/v2_client/test_table_admin.py | 12 +- tests/unit/data/__init__.py | 15 + .../{ => data/_async}/test__mutate_rows.py | 26 +- tests/unit/data/_async/test__read_rows.py | 625 +++++++ tests/unit/{ => data/_async}/test_client.py | 174 +- .../_async}/test_mutations_batcher.py | 84 +- .../{ => data}/read-rows-acceptance-test.json | 0 tests/unit/{ => data}/test__helpers.py | 4 +- .../test__read_rows_state_machine.py} | 439 +---- tests/unit/{ => data}/test_exceptions.py | 18 +- tests/unit/{ => data}/test_mutations.py | 18 +- .../test_read_modify_write_rules.py | 8 +- .../{ => data}/test_read_rows_acceptance.py | 21 +- tests/unit/{ => data}/test_read_rows_query.py | 44 +- tests/unit/{ => data}/test_row.py | 12 +- tests/unit/{ => data}/test_row_filters.py | 360 ++-- .../test_bigtable_table_admin.py | 2 + tests/unit/gapic/bigtable_v2/test_bigtable.py | 505 +----- tests/unit/test_iterators.py | 251 --- tests/unit/v2_client/test_app_profile.py | 44 +- tests/unit/v2_client/test_backup.py | 36 +- tests/unit/v2_client/test_batcher.py | 236 ++- tests/unit/v2_client/test_client.py | 59 +- tests/unit/v2_client/test_cluster.py | 68 +- tests/unit/v2_client/test_column_family.py | 62 +- tests/unit/v2_client/test_encryption_info.py | 8 +- tests/unit/v2_client/test_error.py | 2 +- tests/unit/v2_client/test_instance.py | 50 +- tests/unit/v2_client/test_policy.py | 28 +- tests/unit/v2_client/test_row.py | 30 +- tests/unit/v2_client/test_row_data.py | 62 +- tests/unit/v2_client/test_row_filters.py | 214 +-- tests/unit/v2_client/test_row_merger.py | 8 +- tests/unit/v2_client/test_row_set.py | 60 +- tests/unit/v2_client/test_table.py | 188 +- 145 files changed, 7125 insertions(+), 6989 deletions(-) create mode 100644 docs/batcher.rst rename google/cloud/bigtable/{deprecated => }/app_profile.py (97%) rename google/cloud/bigtable/{deprecated => }/backup.py (96%) create mode 100644 google/cloud/bigtable/batcher.py rename google/cloud/bigtable/{deprecated => }/cluster.py (95%) rename google/cloud/bigtable/{deprecated => }/column_family.py (99%) create mode 100644 google/cloud/bigtable/data/__init__.py create mode 100644 google/cloud/bigtable/data/_async/__init__.py rename google/cloud/bigtable/{ => data/_async}/_mutate_rows.py (91%) create mode 100644 google/cloud/bigtable/data/_async/_read_rows.py create mode 100644 google/cloud/bigtable/data/_async/client.py rename google/cloud/bigtable/{ => data/_async}/mutations_batcher.py (96%) rename google/cloud/bigtable/{ => data}/_helpers.py (98%) rename google/cloud/bigtable/{_read_rows.py => data/_read_rows_state_machine.py} (54%) rename google/cloud/bigtable/{ => data}/exceptions.py (93%) rename google/cloud/bigtable/{ => data}/mutations.py (98%) rename google/cloud/bigtable/{ => data}/read_modify_write_rules.py (100%) rename google/cloud/bigtable/{ => data}/read_rows_query.py (98%) create mode 100644 google/cloud/bigtable/data/row.py rename google/cloud/bigtable/{deprecated => data}/row_filters.py (58%) delete mode 100644 google/cloud/bigtable/deprecated/batcher.py delete mode 100644 google/cloud/bigtable/deprecated/client.py delete mode 100644 google/cloud/bigtable/deprecated/py.typed delete mode 100644 google/cloud/bigtable/deprecated/row.py rename google/cloud/bigtable/{deprecated => }/encryption_info.py (93%) rename google/cloud/bigtable/{deprecated => }/enums.py (100%) rename google/cloud/bigtable/{deprecated => }/error.py (100%) rename google/cloud/bigtable/{deprecated => }/instance.py (91%) delete mode 100644 google/cloud/bigtable/iterators.py rename google/cloud/bigtable/{deprecated => }/policy.py (100%) rename google/cloud/bigtable/{deprecated => }/row_data.py (97%) rename google/cloud/bigtable/{deprecated => }/row_merger.py (99%) rename google/cloud/bigtable/{deprecated => }/row_set.py (100%) rename google/cloud/bigtable/{deprecated => }/table.py (95%) create mode 100644 google/cloud/bigtable_v2/types/feature_flags.py rename {google/cloud/bigtable/deprecated => tests/system/data}/__init__.py (64%) rename tests/system/{ => data}/test_system.py (93%) create mode 100644 tests/unit/data/__init__.py rename tests/unit/{ => data/_async}/test__mutate_rows.py (93%) create mode 100644 tests/unit/data/_async/test__read_rows.py rename tests/unit/{ => data/_async}/test_client.py (95%) rename tests/unit/{ => data/_async}/test_mutations_batcher.py (94%) rename tests/unit/{ => data}/read-rows-acceptance-test.json (100%) rename tests/unit/{ => data}/test__helpers.py (97%) rename tests/unit/{test__read_rows.py => data/test__read_rows_state_machine.py} (62%) rename tests/unit/{ => data}/test_exceptions.py (95%) rename tests/unit/{ => data}/test_mutations.py (97%) rename tests/unit/{ => data}/test_read_modify_write_rules.py (94%) rename tests/unit/{ => data}/test_read_rows_acceptance.py (93%) rename tests/unit/{ => data}/test_read_rows_query.py (94%) rename tests/unit/{ => data}/test_row.py (98%) rename tests/unit/{ => data}/test_row_filters.py (81%) delete mode 100644 tests/unit/test_iterators.py diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index 5fc5daa31..02a4dedce 100644 --- a/.github/.OwlBot.lock.yaml +++ b/.github/.OwlBot.lock.yaml @@ -13,4 +13,5 @@ # limitations under the License. docker: image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest - digest: sha256:8555f0e37e6261408f792bfd6635102d2da5ad73f8f09bcb24f25e6afb5fac97 + digest: sha256:240b5bcc2bafd450912d2da2be15e62bc6de2cf839823ae4bf94d4f392b451dc +# created: 2023-06-03T21:25:37.968717478Z diff --git a/.kokoro/requirements.in b/.kokoro/requirements.in index 882178ce6..ec867d9fd 100644 --- a/.kokoro/requirements.in +++ b/.kokoro/requirements.in @@ -5,6 +5,6 @@ typing-extensions twine wheel setuptools -nox +nox>=2022.11.21 # required to remove dependency on py charset-normalizer<3 click<8.1.0 diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index fa99c1290..c7929db6d 100644 --- a/.kokoro/requirements.txt +++ b/.kokoro/requirements.txt @@ -1,6 +1,6 @@ # -# This file is autogenerated by pip-compile with python 3.10 -# To update, run: +# This file is autogenerated by pip-compile with Python 3.9 +# by the following command: # # pip-compile --allow-unsafe --generate-hashes requirements.in # @@ -113,28 +113,26 @@ commonmark==0.9.1 \ --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \ --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9 # via rich -cryptography==39.0.1 \ - --hash=sha256:0f8da300b5c8af9f98111ffd512910bc792b4c77392a9523624680f7956a99d4 \ - --hash=sha256:35f7c7d015d474f4011e859e93e789c87d21f6f4880ebdc29896a60403328f1f \ - --hash=sha256:5aa67414fcdfa22cf052e640cb5ddc461924a045cacf325cd164e65312d99502 \ - --hash=sha256:5d2d8b87a490bfcd407ed9d49093793d0f75198a35e6eb1a923ce1ee86c62b41 \ - --hash=sha256:6687ef6d0a6497e2b58e7c5b852b53f62142cfa7cd1555795758934da363a965 \ - --hash=sha256:6f8ba7f0328b79f08bdacc3e4e66fb4d7aab0c3584e0bd41328dce5262e26b2e \ - --hash=sha256:706843b48f9a3f9b9911979761c91541e3d90db1ca905fd63fee540a217698bc \ - --hash=sha256:807ce09d4434881ca3a7594733669bd834f5b2c6d5c7e36f8c00f691887042ad \ - --hash=sha256:83e17b26de248c33f3acffb922748151d71827d6021d98c70e6c1a25ddd78505 \ - --hash=sha256:96f1157a7c08b5b189b16b47bc9db2332269d6680a196341bf30046330d15388 \ - --hash=sha256:aec5a6c9864be7df2240c382740fcf3b96928c46604eaa7f3091f58b878c0bb6 \ - --hash=sha256:b0afd054cd42f3d213bf82c629efb1ee5f22eba35bf0eec88ea9ea7304f511a2 \ - --hash=sha256:ced4e447ae29ca194449a3f1ce132ded8fcab06971ef5f618605aacaa612beac \ - --hash=sha256:d1f6198ee6d9148405e49887803907fe8962a23e6c6f83ea7d98f1c0de375695 \ - --hash=sha256:e124352fd3db36a9d4a21c1aa27fd5d051e621845cb87fb851c08f4f75ce8be6 \ - --hash=sha256:e422abdec8b5fa8462aa016786680720d78bdce7a30c652b7fadf83a4ba35336 \ - --hash=sha256:ef8b72fa70b348724ff1218267e7f7375b8de4e8194d1636ee60510aae104cd0 \ - --hash=sha256:f0c64d1bd842ca2633e74a1a28033d139368ad959872533b1bab8c80e8240a0c \ - --hash=sha256:f24077a3b5298a5a06a8e0536e3ea9ec60e4c7ac486755e5fb6e6ea9b3500106 \ - --hash=sha256:fdd188c8a6ef8769f148f88f859884507b954cc64db6b52f66ef199bb9ad660a \ - --hash=sha256:fe913f20024eb2cb2f323e42a64bdf2911bb9738a15dba7d3cce48151034e3a8 +cryptography==41.0.0 \ + --hash=sha256:0ddaee209d1cf1f180f1efa338a68c4621154de0afaef92b89486f5f96047c55 \ + --hash=sha256:14754bcdae909d66ff24b7b5f166d69340ccc6cb15731670435efd5719294895 \ + --hash=sha256:344c6de9f8bda3c425b3a41b319522ba3208551b70c2ae00099c205f0d9fd3be \ + --hash=sha256:34d405ea69a8b34566ba3dfb0521379b210ea5d560fafedf9f800a9a94a41928 \ + --hash=sha256:3680248309d340fda9611498a5319b0193a8dbdb73586a1acf8109d06f25b92d \ + --hash=sha256:3c5ef25d060c80d6d9f7f9892e1d41bb1c79b78ce74805b8cb4aa373cb7d5ec8 \ + --hash=sha256:4ab14d567f7bbe7f1cdff1c53d5324ed4d3fc8bd17c481b395db224fb405c237 \ + --hash=sha256:5c1f7293c31ebc72163a9a0df246f890d65f66b4a40d9ec80081969ba8c78cc9 \ + --hash=sha256:6b71f64beeea341c9b4f963b48ee3b62d62d57ba93eb120e1196b31dc1025e78 \ + --hash=sha256:7d92f0248d38faa411d17f4107fc0bce0c42cae0b0ba5415505df72d751bf62d \ + --hash=sha256:8362565b3835ceacf4dc8f3b56471a2289cf51ac80946f9087e66dc283a810e0 \ + --hash=sha256:84a165379cb9d411d58ed739e4af3396e544eac190805a54ba2e0322feb55c46 \ + --hash=sha256:88ff107f211ea696455ea8d911389f6d2b276aabf3231bf72c8853d22db755c5 \ + --hash=sha256:9f65e842cb02550fac96536edb1d17f24c0a338fd84eaf582be25926e993dde4 \ + --hash=sha256:a4fc68d1c5b951cfb72dfd54702afdbbf0fb7acdc9b7dc4301bbf2225a27714d \ + --hash=sha256:b7f2f5c525a642cecad24ee8670443ba27ac1fab81bba4cc24c7b6b41f2d0c75 \ + --hash=sha256:b846d59a8d5a9ba87e2c3d757ca019fa576793e8758174d3868aecb88d6fc8eb \ + --hash=sha256:bf8fc66012ca857d62f6a347007e166ed59c0bc150cefa49f28376ebe7d992a2 \ + --hash=sha256:f5d0bf9b252f30a31664b6f64432b4730bb7038339bd18b1fafe129cfc2be9be # via # gcp-releasetool # secretstorage @@ -335,9 +333,9 @@ more-itertools==9.0.0 \ --hash=sha256:250e83d7e81d0c87ca6bd942e6aeab8cc9daa6096d12c5308f3f92fa5e5c1f41 \ --hash=sha256:5a6257e40878ef0520b1803990e3e22303a41b5714006c32a3fd8304b26ea1ab # via jaraco-classes -nox==2022.8.7 \ - --hash=sha256:1b894940551dc5c389f9271d197ca5d655d40bdc6ccf93ed6880e4042760a34b \ - --hash=sha256:96cca88779e08282a699d672258ec01eb7c792d35bbbf538c723172bce23212c +nox==2022.11.21 \ + --hash=sha256:0e41a990e290e274cb205a976c4c97ee3c5234441a8132c8c3fd9ea3c22149eb \ + --hash=sha256:e21c31de0711d1274ca585a2c5fde36b1aa962005ba8e9322bf5eeed16dcd684 # via -r requirements.in packaging==21.3 \ --hash=sha256:dd47c42927d89ab911e606518907cc2d3a1f38bbd026385970643f9c5b8ecfeb \ @@ -380,10 +378,6 @@ protobuf==3.20.3 \ # gcp-docuploader # gcp-releasetool # google-api-core -py==1.11.0 \ - --hash=sha256:51c75c4126074b472f746a24399ad32f6053d1b34b68d2fa41e558e6f4a98719 \ - --hash=sha256:607c53218732647dff4acdfcd50cb62615cedf612e72d1724fb1a0cc6405b378 - # via nox pyasn1==0.4.8 \ --hash=sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d \ --hash=sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba @@ -423,9 +417,9 @@ readme-renderer==37.3 \ --hash=sha256:cd653186dfc73055656f090f227f5cb22a046d7f71a841dfa305f55c9a513273 \ --hash=sha256:f67a16caedfa71eef48a31b39708637a6f4664c4394801a7b0d6432d13907343 # via twine -requests==2.28.1 \ - --hash=sha256:7c5599b102feddaa661c826c56ab4fee28bfd17f5abca1ebbe3e7f19d7c97983 \ - --hash=sha256:8fefa2a1a1365bf5520aac41836fbee479da67864514bdb821f31ce07ce65349 +requests==2.31.0 \ + --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f \ + --hash=sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1 # via # gcp-releasetool # google-api-core diff --git a/.release-please-manifest.json b/.release-please-manifest.json index 882f663e6..b7f666a68 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,3 +1,3 @@ { - ".": "2.17.0" + ".": "2.19.0" } \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 78b4d1b29..dc80386a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,49 @@ [1]: https://pypi.org/project/google-cloud-bigtable/#history +## [2.19.0](https://github.com/googleapis/python-bigtable/compare/v2.18.1...v2.19.0) (2023-06-08) + + +### Features + +* Add ChangeStreamConfig to CreateTable and UpdateTable ([#786](https://github.com/googleapis/python-bigtable/issues/786)) ([cef70f2](https://github.com/googleapis/python-bigtable/commit/cef70f243541820225f86a520e0b2abd3a7354f7)) + + +### Bug Fixes + +* Add a callback function on flush_rows ([#796](https://github.com/googleapis/python-bigtable/issues/796)) ([589aa5d](https://github.com/googleapis/python-bigtable/commit/589aa5d04f6b5a2bd310d0bf06aeb7058fb6fcd2)) + + +### Documentation + +* **samples:** Add region tags ([#788](https://github.com/googleapis/python-bigtable/issues/788)) ([ecf539c](https://github.com/googleapis/python-bigtable/commit/ecf539c4c976fd9e5505b8abf0b697b218f09fef)) + +## [2.18.1](https://github.com/googleapis/python-bigtable/compare/v2.18.0...v2.18.1) (2023-05-11) + + +### Bug Fixes + +* Revert "Feat: Threaded MutationsBatcher" ([#773](https://github.com/googleapis/python-bigtable/issues/773)) ([a767cff](https://github.com/googleapis/python-bigtable/commit/a767cff95d990994f85f5fd05cc10f952087b49d)) + +## [2.18.0](https://github.com/googleapis/python-bigtable/compare/v2.17.0...v2.18.0) (2023-05-10) + + +### Features + +* Publish RateLimitInfo and FeatureFlag protos ([#768](https://github.com/googleapis/python-bigtable/issues/768)) ([171fea6](https://github.com/googleapis/python-bigtable/commit/171fea6de57a47f92a2a56050f8bfe7518144df7)) +* Threaded MutationsBatcher ([#722](https://github.com/googleapis/python-bigtable/issues/722)) ([7521a61](https://github.com/googleapis/python-bigtable/commit/7521a617c121ead96a21ca47959a53b2db2da090)) + + +### Bug Fixes + +* Pass the "retry" when calling read_rows. ([#759](https://github.com/googleapis/python-bigtable/issues/759)) ([505273b](https://github.com/googleapis/python-bigtable/commit/505273b72bf83d8f92d0e0a92d62f22bce96cc3d)) + + +### Documentation + +* Fix delete from column family example ([#764](https://github.com/googleapis/python-bigtable/issues/764)) ([128b4e1](https://github.com/googleapis/python-bigtable/commit/128b4e1f3eea2dad903d84c8f2933b17a5f0d226)) +* Fix formatting of request arg in docstring ([#756](https://github.com/googleapis/python-bigtable/issues/756)) ([45d3e43](https://github.com/googleapis/python-bigtable/commit/45d3e4308c4f494228c2e6e18a36285c557cb0c3)) + ## [2.17.0](https://github.com/googleapis/python-bigtable/compare/v2.16.0...v2.17.0) (2023-03-01) diff --git a/docs/app-profile.rst b/docs/app-profile.rst index 50e57c179..5c9d426c2 100644 --- a/docs/app-profile.rst +++ b/docs/app-profile.rst @@ -1,6 +1,6 @@ App Profile ~~~~~~~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.app_profile +.. automodule:: google.cloud.bigtable.app_profile :members: :show-inheritance: diff --git a/docs/backup.rst b/docs/backup.rst index 46c32c91b..e75abd431 100644 --- a/docs/backup.rst +++ b/docs/backup.rst @@ -1,6 +1,6 @@ Backup ~~~~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.backup +.. automodule:: google.cloud.bigtable.backup :members: :show-inheritance: diff --git a/docs/batcher.rst b/docs/batcher.rst new file mode 100644 index 000000000..9ac335be1 --- /dev/null +++ b/docs/batcher.rst @@ -0,0 +1,6 @@ +Mutations Batching +~~~~~~~~~~~~~~~~~~ + +.. automodule:: google.cloud.bigtable.batcher + :members: + :show-inheritance: diff --git a/docs/client-intro.rst b/docs/client-intro.rst index d75cf5f96..242068499 100644 --- a/docs/client-intro.rst +++ b/docs/client-intro.rst @@ -1,21 +1,21 @@ Base for Everything =================== -To use the API, the :class:`Client ` +To use the API, the :class:`Client ` class defines a high-level interface which handles authorization and creating other objects: .. code:: python - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client client = Client() Long-lived Defaults ------------------- -When creating a :class:`Client `, the +When creating a :class:`Client `, the ``user_agent`` argument has sensible a default -(:data:`DEFAULT_USER_AGENT `). +(:data:`DEFAULT_USER_AGENT `). However, you may over-ride it and the value will be used throughout all API requests made with the ``client`` you create. @@ -38,14 +38,14 @@ Configuration .. code:: - >>> import google.cloud.deprecated as bigtable + >>> from google.cloud import bigtable >>> client = bigtable.Client() or pass in ``credentials`` and ``project`` explicitly .. code:: - >>> import google.cloud.deprecated as bigtable + >>> from google.cloud import bigtable >>> client = bigtable.Client(project='my-project', credentials=creds) .. tip:: @@ -73,15 +73,15 @@ you can pass the ``read_only`` argument: client = bigtable.Client(read_only=True) This will ensure that the -:data:`READ_ONLY_SCOPE ` is used +:data:`READ_ONLY_SCOPE ` is used for API requests (so any accidental requests that would modify data will fail). Next Step --------- -After a :class:`Client `, the next highest-level -object is an :class:`Instance `. You'll need +After a :class:`Client `, the next highest-level +object is an :class:`Instance `. You'll need one before you can interact with tables or data. Head next to learn about the :doc:`instance-api`. diff --git a/docs/client.rst b/docs/client.rst index df92a9861..c48595c8a 100644 --- a/docs/client.rst +++ b/docs/client.rst @@ -1,6 +1,6 @@ Client ~~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.client +.. automodule:: google.cloud.bigtable.client :members: :show-inheritance: diff --git a/docs/cluster.rst b/docs/cluster.rst index 9747b226f..ad33aae5e 100644 --- a/docs/cluster.rst +++ b/docs/cluster.rst @@ -1,6 +1,6 @@ Cluster ~~~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.cluster +.. automodule:: google.cloud.bigtable.cluster :members: :show-inheritance: diff --git a/docs/column-family.rst b/docs/column-family.rst index 39095000d..de6c1eb1f 100644 --- a/docs/column-family.rst +++ b/docs/column-family.rst @@ -2,7 +2,7 @@ Column Families =============== When creating a -:class:`ColumnFamily `, it is +:class:`ColumnFamily `, it is possible to set garbage collection rules for expired data. By setting a rule, cells in the table matching the rule will be deleted @@ -10,19 +10,19 @@ during periodic garbage collection (which executes opportunistically in the background). The types -:class:`MaxAgeGCRule `, -:class:`MaxVersionsGCRule `, -:class:`GarbageCollectionRuleUnion ` and -:class:`GarbageCollectionRuleIntersection ` +:class:`MaxAgeGCRule `, +:class:`MaxVersionsGCRule `, +:class:`GarbageCollectionRuleUnion ` and +:class:`GarbageCollectionRuleIntersection ` can all be used as the optional ``gc_rule`` argument in the -:class:`ColumnFamily ` +:class:`ColumnFamily ` constructor. This value is then used in the -:meth:`create() ` and -:meth:`update() ` methods. +:meth:`create() ` and +:meth:`update() ` methods. These rules can be nested arbitrarily, with a -:class:`MaxAgeGCRule ` or -:class:`MaxVersionsGCRule ` +:class:`MaxAgeGCRule ` or +:class:`MaxVersionsGCRule ` at the lowest level of the nesting: .. code:: python @@ -44,6 +44,6 @@ at the lowest level of the nesting: ---- -.. automodule:: google.cloud.bigtable.deprecated.column_family +.. automodule:: google.cloud.bigtable.column_family :members: :show-inheritance: diff --git a/docs/data-api.rst b/docs/data-api.rst index e68835d1a..01a49178f 100644 --- a/docs/data-api.rst +++ b/docs/data-api.rst @@ -1,7 +1,7 @@ Data API ======== -After creating a :class:`Table ` and some +After creating a :class:`Table ` and some column families, you are ready to store and retrieve data. Cells vs. Columns vs. Column Families @@ -27,7 +27,7 @@ Modifying Data Since data is stored in cells, which are stored in rows, we use the metaphor of a **row** in classes that are used to modify (write, update, delete) data in a -:class:`Table `. +:class:`Table `. Direct vs. Conditional vs. Append --------------------------------- @@ -38,26 +38,26 @@ methods. * The **direct** way is via `MutateRow`_ which involves simply adding, overwriting or deleting cells. The - :class:`DirectRow ` class + :class:`DirectRow ` class handles direct mutations. * The **conditional** way is via `CheckAndMutateRow`_. This method first checks if some filter is matched in a given row, then applies one of two sets of mutations, depending on if a match occurred or not. (These mutation sets are called the "true mutations" and "false mutations".) The - :class:`ConditionalRow ` class + :class:`ConditionalRow ` class handles conditional mutations. * The **append** way is via `ReadModifyWriteRow`_. This simply appends (as bytes) or increments (as an integer) data in a presumed existing cell in a row. The - :class:`AppendRow ` class + :class:`AppendRow ` class handles append mutations. Row Factory ----------- A single factory can be used to create any of the three row types. -To create a :class:`DirectRow `: +To create a :class:`DirectRow `: .. code:: python @@ -66,15 +66,15 @@ To create a :class:`DirectRow `: Unlike the previous string values we've used before, the row key must be ``bytes``. -To create a :class:`ConditionalRow `, -first create a :class:`RowFilter ` and +To create a :class:`ConditionalRow `, +first create a :class:`RowFilter ` and then .. code:: python cond_row = table.row(row_key, filter_=filter_) -To create an :class:`AppendRow ` +To create an :class:`AppendRow ` .. code:: python @@ -95,7 +95,7 @@ Direct Mutations Direct mutations can be added via one of four methods -* :meth:`set_cell() ` allows a +* :meth:`set_cell() ` allows a single value to be written to a column .. code:: python @@ -109,7 +109,7 @@ Direct mutations can be added via one of four methods The value can either be bytes or an integer, which will be converted to bytes as a signed 64-bit integer. -* :meth:`delete_cell() ` deletes +* :meth:`delete_cell() ` deletes all cells (i.e. for all timestamps) in a given column .. code:: python @@ -119,7 +119,7 @@ Direct mutations can be added via one of four methods Remember, this only happens in the ``row`` we are using. If we only want to delete cells from a limited range of time, a - :class:`TimestampRange ` can + :class:`TimestampRange ` can be used .. code:: python @@ -127,9 +127,9 @@ Direct mutations can be added via one of four methods row.delete_cell(column_family_id, column, time_range=time_range) -* :meth:`delete_cells() ` does +* :meth:`delete_cells() ` does the same thing as - :meth:`delete_cell() `, + :meth:`delete_cell() `, but accepts a list of columns in a column family rather than a single one. .. code:: python @@ -138,7 +138,7 @@ Direct mutations can be added via one of four methods time_range=time_range) In addition, if we want to delete cells from every column in a column family, - the special :attr:`ALL_COLUMNS ` + the special :attr:`ALL_COLUMNS ` value can be used .. code:: python @@ -146,7 +146,7 @@ Direct mutations can be added via one of four methods row.delete_cells(column_family_id, row.ALL_COLUMNS, time_range=time_range) -* :meth:`delete() ` will delete the +* :meth:`delete() ` will delete the entire row .. code:: python @@ -177,14 +177,14 @@ Append Mutations Append mutations can be added via one of two methods -* :meth:`append_cell_value() ` +* :meth:`append_cell_value() ` appends a bytes value to an existing cell: .. code:: python append_row.append_cell_value(column_family_id, column, bytes_value) -* :meth:`increment_cell_value() ` +* :meth:`increment_cell_value() ` increments an integer value in an existing cell: .. code:: python @@ -217,7 +217,7 @@ Read Single Row from a Table ---------------------------- To make a `ReadRows`_ API request for a single row key, use -:meth:`Table.read_row() `: +:meth:`Table.read_row() `: .. code:: python @@ -226,34 +226,34 @@ To make a `ReadRows`_ API request for a single row key, use { u'fam1': { b'col1': [ - , - , + , + , ], b'col2': [ - , + , ], }, u'fam2': { b'col3': [ - , - , - , + , + , + , ], }, } >>> cell = row_data.cells[u'fam1'][b'col1'][0] >>> cell - + >>> cell.value b'val1' >>> cell.timestamp datetime.datetime(2016, 2, 27, 3, 41, 18, 122823, tzinfo=) -Rather than returning a :class:`DirectRow ` +Rather than returning a :class:`DirectRow ` or similar class, this method returns a -:class:`PartialRowData ` +:class:`PartialRowData ` instance. This class is used for reading and parsing data rather than for -modifying data (as :class:`DirectRow ` is). +modifying data (as :class:`DirectRow ` is). A filter can also be applied to the results: @@ -262,15 +262,15 @@ A filter can also be applied to the results: row_data = table.read_row(row_key, filter_=filter_val) The allowable ``filter_`` values are the same as those used for a -:class:`ConditionalRow `. For +:class:`ConditionalRow `. For more information, see the -:meth:`Table.read_row() ` documentation. +:meth:`Table.read_row() ` documentation. Stream Many Rows from a Table ----------------------------- To make a `ReadRows`_ API request for a stream of rows, use -:meth:`Table.read_rows() `: +:meth:`Table.read_rows() `: .. code:: python @@ -279,32 +279,32 @@ To make a `ReadRows`_ API request for a stream of rows, use Using gRPC over HTTP/2, a continual stream of responses will be delivered. In particular -* :meth:`consume_next() ` +* :meth:`consume_next() ` pulls the next result from the stream, parses it and stores it on the - :class:`PartialRowsData ` instance -* :meth:`consume_all() ` + :class:`PartialRowsData ` instance +* :meth:`consume_all() ` pulls results from the stream until there are no more -* :meth:`cancel() ` closes +* :meth:`cancel() ` closes the stream -See the :class:`PartialRowsData ` +See the :class:`PartialRowsData ` documentation for more information. As with -:meth:`Table.read_row() `, an optional +:meth:`Table.read_row() `, an optional ``filter_`` can be applied. In addition a ``start_key`` and / or ``end_key`` can be supplied for the stream, a ``limit`` can be set and a boolean ``allow_row_interleaving`` can be specified to allow faster streamed results at the potential cost of non-sequential reads. -See the :meth:`Table.read_rows() ` +See the :meth:`Table.read_rows() ` documentation for more information on the optional arguments. Sample Keys in a Table ---------------------- Make a `SampleRowKeys`_ API request with -:meth:`Table.sample_row_keys() `: +:meth:`Table.sample_row_keys() `: .. code:: python @@ -315,7 +315,7 @@ approximately equal size, which can be used to break up the data for distributed tasks like mapreduces. As with -:meth:`Table.read_rows() `, the +:meth:`Table.read_rows() `, the returned ``keys_iterator`` is connected to a cancellable HTTP/2 stream. The next key in the result can be accessed via diff --git a/docs/encryption-info.rst b/docs/encryption-info.rst index 62b77ea0c..46f19880f 100644 --- a/docs/encryption-info.rst +++ b/docs/encryption-info.rst @@ -1,6 +1,6 @@ Encryption Info ~~~~~~~~~~~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.encryption_info +.. automodule:: google.cloud.bigtable.encryption_info :members: :show-inheritance: diff --git a/docs/instance-api.rst b/docs/instance-api.rst index 78123e8ca..88b4eb4dc 100644 --- a/docs/instance-api.rst +++ b/docs/instance-api.rst @@ -1,7 +1,7 @@ Instance Admin API ================== -After creating a :class:`Client `, you can +After creating a :class:`Client `, you can interact with individual instances for a project. List Instances @@ -9,7 +9,7 @@ List Instances If you want a comprehensive list of all existing instances, make a `ListInstances`_ API request with -:meth:`Client.list_instances() `: +:meth:`Client.list_instances() `: .. code:: python @@ -18,7 +18,7 @@ If you want a comprehensive list of all existing instances, make a Instance Factory ---------------- -To create an :class:`Instance ` object: +To create an :class:`Instance ` object: .. code:: python @@ -40,7 +40,7 @@ Create a new Instance --------------------- After creating the instance object, make a `CreateInstance`_ API request -with :meth:`create() `: +with :meth:`create() `: .. code:: python @@ -54,14 +54,14 @@ Check on Current Operation When modifying an instance (via a `CreateInstance`_ request), the Bigtable API will return a `long-running operation`_ and a corresponding - :class:`Operation ` object + :class:`Operation ` object will be returned by - :meth:`create() `. + :meth:`create() `. You can check if a long-running operation (for a -:meth:`create() ` has finished +:meth:`create() ` has finished by making a `GetOperation`_ request with -:meth:`Operation.finished() `: +:meth:`Operation.finished() `: .. code:: python @@ -71,18 +71,18 @@ by making a `GetOperation`_ request with .. note:: - Once an :class:`Operation ` object + Once an :class:`Operation ` object has returned :data:`True` from - :meth:`finished() `, the + :meth:`finished() `, the object should not be re-used. Subsequent calls to - :meth:`finished() ` + :meth:`finished() ` will result in a :class:`ValueError `. Get metadata for an existing Instance ------------------------------------- After creating the instance object, make a `GetInstance`_ API request -with :meth:`reload() `: +with :meth:`reload() `: .. code:: python @@ -94,7 +94,7 @@ Update an existing Instance --------------------------- After creating the instance object, make an `UpdateInstance`_ API request -with :meth:`update() `: +with :meth:`update() `: .. code:: python @@ -105,7 +105,7 @@ Delete an existing Instance --------------------------- Make a `DeleteInstance`_ API request with -:meth:`delete() `: +:meth:`delete() `: .. code:: python @@ -115,8 +115,8 @@ Next Step --------- Now we go down the hierarchy from -:class:`Instance ` to a -:class:`Table `. +:class:`Instance ` to a +:class:`Table `. Head next to learn about the :doc:`table-api`. diff --git a/docs/instance.rst b/docs/instance.rst index 3a61faf1c..f9be9672f 100644 --- a/docs/instance.rst +++ b/docs/instance.rst @@ -1,6 +1,6 @@ Instance ~~~~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.instance +.. automodule:: google.cloud.bigtable.instance :members: :show-inheritance: diff --git a/docs/row-data.rst b/docs/row-data.rst index b9013ebf5..503f9b1cb 100644 --- a/docs/row-data.rst +++ b/docs/row-data.rst @@ -1,6 +1,6 @@ Row Data ~~~~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.row_data +.. automodule:: google.cloud.bigtable.row_data :members: :show-inheritance: diff --git a/docs/row-filters.rst b/docs/row-filters.rst index 8d1fac46b..9884ce400 100644 --- a/docs/row-filters.rst +++ b/docs/row-filters.rst @@ -2,11 +2,11 @@ Bigtable Row Filters ==================== It is possible to use a -:class:`RowFilter ` +:class:`RowFilter ` when adding mutations to a -:class:`ConditionalRow ` and when -reading row data with :meth:`read_row() ` -or :meth:`read_rows() `. +:class:`ConditionalRow ` and when +reading row data with :meth:`read_row() ` +or :meth:`read_rows() `. As laid out in the `RowFilter definition`_, the following basic filters are provided: @@ -60,8 +60,8 @@ level. For example: ---- -.. automodule:: google.cloud.bigtable.deprecated.row_filters +.. automodule:: google.cloud.bigtable.row_filters :members: :show-inheritance: -.. _RowFilter definition: https://googleapis.dev/python/bigtable/latest/row-filters.html?highlight=rowfilter#google.cloud.bigtable.deprecated.row_filters.RowFilter +.. _RowFilter definition: https://googleapis.dev/python/bigtable/latest/row-filters.html?highlight=rowfilter#google.cloud.bigtable.row_filters.RowFilter diff --git a/docs/row-set.rst b/docs/row-set.rst index 92cd107e8..5f7a16a02 100644 --- a/docs/row-set.rst +++ b/docs/row-set.rst @@ -1,6 +1,6 @@ Row Set ~~~~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.row_set +.. automodule:: google.cloud.bigtable.row_set :members: :show-inheritance: diff --git a/docs/row.rst b/docs/row.rst index e8fa48cdd..33686608b 100644 --- a/docs/row.rst +++ b/docs/row.rst @@ -1,7 +1,7 @@ Bigtable Row ============ -.. automodule:: google.cloud.bigtable.deprecated.row +.. automodule:: google.cloud.bigtable.row :members: :show-inheritance: :inherited-members: diff --git a/docs/snippets.py b/docs/snippets.py index 084f10270..1d93fdf12 100644 --- a/docs/snippets.py +++ b/docs/snippets.py @@ -16,7 +16,7 @@ """Testable usage examples for Google Cloud Bigtable API wrapper Each example function takes a ``client`` argument (which must be an instance -of :class:`google.cloud.bigtable.deprecated.client.Client`) and uses it to perform a task +of :class:`google.cloud.bigtable.client.Client`) and uses it to perform a task with the API. To facilitate running the examples as system tests, each example is also passed @@ -40,8 +40,8 @@ from test_utils.retry import RetryErrors from google.cloud._helpers import UTC -from google.cloud.bigtable.deprecated import Client -from google.cloud.bigtable.deprecated import enums +from google.cloud.bigtable import Client +from google.cloud.bigtable import enums UNIQUE_SUFFIX = unique_resource_id("-") @@ -110,8 +110,8 @@ def teardown_module(): def test_bigtable_create_instance(): # [START bigtable_api_create_prod_instance] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import Client + from google.cloud.bigtable import enums my_instance_id = "inst-my-" + UNIQUE_SUFFIX my_cluster_id = "clus-my-" + UNIQUE_SUFFIX @@ -144,8 +144,8 @@ def test_bigtable_create_instance(): def test_bigtable_create_additional_cluster(): # [START bigtable_api_create_cluster] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import Client + from google.cloud.bigtable import enums # Assuming that there is an existing instance with `INSTANCE_ID` # on the server already. @@ -181,8 +181,8 @@ def test_bigtable_create_reload_delete_app_profile(): import re # [START bigtable_api_create_app_profile] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import Client + from google.cloud.bigtable import enums routing_policy_type = enums.RoutingPolicyType.ANY @@ -202,7 +202,7 @@ def test_bigtable_create_reload_delete_app_profile(): # [END bigtable_api_create_app_profile] # [START bigtable_api_app_profile_name] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -219,7 +219,7 @@ def test_bigtable_create_reload_delete_app_profile(): assert _profile_name_re.match(app_profile_name) # [START bigtable_api_app_profile_exists] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -230,7 +230,7 @@ def test_bigtable_create_reload_delete_app_profile(): assert app_profile_exists # [START bigtable_api_reload_app_profile] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -241,7 +241,7 @@ def test_bigtable_create_reload_delete_app_profile(): assert app_profile.routing_policy_type == ROUTING_POLICY_TYPE # [START bigtable_api_update_app_profile] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -255,7 +255,7 @@ def test_bigtable_create_reload_delete_app_profile(): assert app_profile.description == description # [START bigtable_api_delete_app_profile] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -269,7 +269,7 @@ def test_bigtable_create_reload_delete_app_profile(): def test_bigtable_list_instances(): # [START bigtable_api_list_instances] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) (instances_list, failed_locations_list) = client.list_instances() @@ -280,7 +280,7 @@ def test_bigtable_list_instances(): def test_bigtable_list_clusters_on_instance(): # [START bigtable_api_list_clusters_on_instance] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -292,7 +292,7 @@ def test_bigtable_list_clusters_on_instance(): def test_bigtable_list_clusters_in_project(): # [START bigtable_api_list_clusters_in_project] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) (clusters_list, failed_locations_list) = client.list_clusters() @@ -309,7 +309,7 @@ def test_bigtable_list_app_profiles(): app_profile = app_profile.create(ignore_warnings=True) # [START bigtable_api_list_app_profiles] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -325,7 +325,7 @@ def test_bigtable_list_app_profiles(): def test_bigtable_instance_exists(): # [START bigtable_api_check_instance_exists] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -337,7 +337,7 @@ def test_bigtable_instance_exists(): def test_bigtable_cluster_exists(): # [START bigtable_api_check_cluster_exists] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -350,7 +350,7 @@ def test_bigtable_cluster_exists(): def test_bigtable_reload_instance(): # [START bigtable_api_reload_instance] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -362,7 +362,7 @@ def test_bigtable_reload_instance(): def test_bigtable_reload_cluster(): # [START bigtable_api_reload_cluster] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -375,7 +375,7 @@ def test_bigtable_reload_cluster(): def test_bigtable_update_instance(): # [START bigtable_api_update_instance] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -389,7 +389,7 @@ def test_bigtable_update_instance(): def test_bigtable_update_cluster(): # [START bigtable_api_update_cluster] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -403,7 +403,7 @@ def test_bigtable_update_cluster(): def test_bigtable_cluster_disable_autoscaling(): # [START bigtable_api_cluster_disable_autoscaling] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -424,8 +424,8 @@ def test_bigtable_create_table(): # [START bigtable_api_create_table] from google.api_core import exceptions from google.api_core import retry - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -450,7 +450,7 @@ def test_bigtable_create_table(): def test_bigtable_list_tables(): # [START bigtable_api_list_tables] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -463,7 +463,7 @@ def test_bigtable_list_tables(): def test_bigtable_delete_cluster(): - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -480,7 +480,7 @@ def test_bigtable_delete_cluster(): operation.result(timeout=1000) # [START bigtable_api_delete_cluster] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -493,7 +493,7 @@ def test_bigtable_delete_cluster(): def test_bigtable_delete_instance(): - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) @@ -515,7 +515,7 @@ def test_bigtable_delete_instance(): INSTANCES_TO_DELETE.append(instance) # [START bigtable_api_delete_instance] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) @@ -531,7 +531,7 @@ def test_bigtable_delete_instance(): def test_bigtable_test_iam_permissions(): # [START bigtable_api_test_iam_permissions] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -547,9 +547,9 @@ def test_bigtable_set_iam_policy_then_get_iam_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_set_iam_policy] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable import Client + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -563,7 +563,7 @@ def test_bigtable_set_iam_policy_then_get_iam_policy(): assert len(policy_latest.bigtable_admins) > 0 # [START bigtable_api_get_iam_policy] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -577,7 +577,7 @@ def test_bigtable_project_path(): import re # [START bigtable_api_project_path] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) project_path = client.project_path @@ -586,7 +586,7 @@ def test_bigtable_project_path(): def test_bigtable_table_data_client(): # [START bigtable_api_table_data_client] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) table_data_client = client.table_data_client @@ -595,7 +595,7 @@ def test_bigtable_table_data_client(): def test_bigtable_table_admin_client(): # [START bigtable_api_table_admin_client] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) table_admin_client = client.table_admin_client @@ -604,7 +604,7 @@ def test_bigtable_table_admin_client(): def test_bigtable_instance_admin_client(): # [START bigtable_api_instance_admin_client] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance_admin_client = client.instance_admin_client @@ -615,9 +615,9 @@ def test_bigtable_admins_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_admins_policy] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable import Client + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -636,9 +636,9 @@ def test_bigtable_readers_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_readers_policy] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_READER_ROLE + from google.cloud.bigtable import Client + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_READER_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -657,9 +657,9 @@ def test_bigtable_users_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_users_policy] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_USER_ROLE + from google.cloud.bigtable import Client + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_USER_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -678,9 +678,9 @@ def test_bigtable_viewers_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_viewers_policy] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_VIEWER_ROLE + from google.cloud.bigtable import Client + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_VIEWER_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -699,7 +699,7 @@ def test_bigtable_instance_name(): import re # [START bigtable_api_instance_name] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -711,7 +711,7 @@ def test_bigtable_cluster_name(): import re # [START bigtable_api_cluster_name] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -722,7 +722,7 @@ def test_bigtable_cluster_name(): def test_bigtable_instance_from_pb(): # [START bigtable_api_instance_from_pb] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 client = Client(admin=True) @@ -741,7 +741,7 @@ def test_bigtable_instance_from_pb(): def test_bigtable_cluster_from_pb(): # [START bigtable_api_cluster_from_pb] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 client = Client(admin=True) @@ -767,7 +767,7 @@ def test_bigtable_cluster_from_pb(): def test_bigtable_instance_state(): # [START bigtable_api_instance_state] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -779,7 +779,7 @@ def test_bigtable_instance_state(): def test_bigtable_cluster_state(): # [START bigtable_api_cluster_state] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) diff --git a/docs/snippets_table.py b/docs/snippets_table.py index 72c342907..f27260425 100644 --- a/docs/snippets_table.py +++ b/docs/snippets_table.py @@ -16,7 +16,7 @@ """Testable usage examples for Google Cloud Bigtable API wrapper Each example function takes a ``client`` argument (which must be an instance -of :class:`google.cloud.bigtable.deprecated.client.Client`) and uses it to perform a task +of :class:`google.cloud.bigtable.client.Client`) and uses it to perform a task with the API. To facilitate running the examples as system tests, each example is also passed @@ -38,9 +38,9 @@ from test_utils.retry import RetryErrors from google.cloud._helpers import UTC -from google.cloud.bigtable.deprecated import Client -from google.cloud.bigtable.deprecated import enums -from google.cloud.bigtable.deprecated import column_family +from google.cloud.bigtable import Client +from google.cloud.bigtable import enums +from google.cloud.bigtable import column_family INSTANCE_ID = "snippet" + unique_resource_id("-") @@ -113,8 +113,8 @@ def teardown_module(): def test_bigtable_create_table(): # [START bigtable_api_create_table] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -143,7 +143,7 @@ def test_bigtable_sample_row_keys(): assert table_sample.exists() # [START bigtable_api_sample_row_keys] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -159,7 +159,7 @@ def test_bigtable_sample_row_keys(): def test_bigtable_write_read_drop_truncate(): # [START bigtable_api_mutate_rows] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -190,7 +190,7 @@ def test_bigtable_write_read_drop_truncate(): # [END bigtable_api_mutate_rows] assert len(response) == len(rows) # [START bigtable_api_read_row] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -200,7 +200,7 @@ def test_bigtable_write_read_drop_truncate(): # [END bigtable_api_read_row] assert row.row_key.decode("utf-8") == row_key # [START bigtable_api_read_rows] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -218,7 +218,7 @@ def test_bigtable_write_read_drop_truncate(): # [END bigtable_api_read_rows] assert len(total_rows) == len(rows) # [START bigtable_api_drop_by_prefix] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -231,7 +231,7 @@ def test_bigtable_write_read_drop_truncate(): assert row.row_key.decode("utf-8") not in dropped_row_keys # [START bigtable_api_truncate_table] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -246,7 +246,7 @@ def test_bigtable_write_read_drop_truncate(): def test_bigtable_mutations_batcher(): # [START bigtable_api_mutations_batcher] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -297,7 +297,7 @@ def test_bigtable_mutations_batcher(): def test_bigtable_table_column_family(): # [START bigtable_api_table_column_family] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -311,7 +311,7 @@ def test_bigtable_table_column_family(): def test_bigtable_list_tables(): # [START bigtable_api_list_tables] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -324,7 +324,7 @@ def test_bigtable_table_name(): import re # [START bigtable_api_table_name] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -342,7 +342,7 @@ def test_bigtable_table_name(): def test_bigtable_list_column_families(): # [START bigtable_api_list_column_families] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -356,7 +356,7 @@ def test_bigtable_list_column_families(): def test_bigtable_get_cluster_states(): # [START bigtable_api_get_cluster_states] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -374,7 +374,7 @@ def test_bigtable_table_test_iam_permissions(): assert table_policy.exists # [START bigtable_api_table_test_iam_permissions] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -392,9 +392,9 @@ def test_bigtable_table_set_iam_policy_then_get_iam_policy(): service_account_email = Config.CLIENT._credentials.service_account_email # [START bigtable_api_table_set_iam_policy] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable import Client + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -407,7 +407,7 @@ def test_bigtable_table_set_iam_policy_then_get_iam_policy(): assert len(policy_latest.bigtable_admins) > 0 # [START bigtable_api_table_get_iam_policy] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -419,7 +419,7 @@ def test_bigtable_table_set_iam_policy_then_get_iam_policy(): def test_bigtable_table_exists(): # [START bigtable_api_check_table_exists] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -435,7 +435,7 @@ def test_bigtable_delete_table(): assert table_del.exists() # [START bigtable_api_delete_table] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -448,7 +448,7 @@ def test_bigtable_delete_table(): def test_bigtable_table_row(): # [START bigtable_api_table_row] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -475,7 +475,7 @@ def test_bigtable_table_row(): def test_bigtable_table_append_row(): # [START bigtable_api_table_append_row] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -502,7 +502,7 @@ def test_bigtable_table_append_row(): def test_bigtable_table_direct_row(): # [START bigtable_api_table_direct_row] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -529,8 +529,8 @@ def test_bigtable_table_direct_row(): def test_bigtable_table_conditional_row(): # [START bigtable_api_table_conditional_row] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.row_filters import PassAllFilter + from google.cloud.bigtable import Client + from google.cloud.bigtable.row_filters import PassAllFilter client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -558,7 +558,7 @@ def test_bigtable_table_conditional_row(): def test_bigtable_column_family_name(): # [START bigtable_api_column_family_name] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -581,8 +581,8 @@ def test_bigtable_column_family_name(): def test_bigtable_create_update_delete_column_family(): # [START bigtable_api_create_column_family] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -598,8 +598,8 @@ def test_bigtable_create_update_delete_column_family(): assert column_families[column_family_id].gc_rule == gc_rule # [START bigtable_api_update_column_family] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -617,8 +617,8 @@ def test_bigtable_create_update_delete_column_family(): assert updated_families[column_family_id].gc_rule == max_age_rule # [START bigtable_api_delete_column_family] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -653,8 +653,8 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): Config.TABLE.mutate_rows(rows) # [START bigtable_api_add_row_key] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable import Client + from google.cloud.bigtable.row_set import RowSet client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -670,9 +670,9 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): assert found_row_keys == expected_row_keys # [START bigtable_api_add_row_range] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.row_set import RowSet - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable import Client + from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -688,8 +688,8 @@ def test_bigtable_add_row_add_row_range_add_row_range_from_keys(): assert found_row_keys == expected_row_keys # [START bigtable_api_row_range_from_keys] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable import Client + from google.cloud.bigtable.row_set import RowSet client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -723,8 +723,8 @@ def test_bigtable_add_row_range_with_prefix(): Config.TABLE.mutate_rows(rows) # [START bigtable_api_add_row_range_with_prefix] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable import Client + from google.cloud.bigtable.row_set import RowSet client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -747,7 +747,7 @@ def test_bigtable_add_row_range_with_prefix(): def test_bigtable_batcher_mutate_flush_mutate_rows(): # [START bigtable_api_batcher_mutate] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -769,7 +769,7 @@ def test_bigtable_batcher_mutate_flush_mutate_rows(): # [END bigtable_api_batcher_mutate] # [START bigtable_api_batcher_flush] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -795,7 +795,7 @@ def test_bigtable_batcher_mutate_flush_mutate_rows(): table.truncate(timeout=200) # [START bigtable_api_batcher_mutate_rows] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -829,8 +829,8 @@ def test_bigtable_batcher_mutate_flush_mutate_rows(): def test_bigtable_create_family_gc_max_age(): # [START bigtable_api_create_family_gc_max_age] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -851,8 +851,8 @@ def test_bigtable_create_family_gc_max_age(): def test_bigtable_create_family_gc_max_versions(): # [START bigtable_api_create_family_gc_max_versions] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -872,8 +872,8 @@ def test_bigtable_create_family_gc_max_versions(): def test_bigtable_create_family_gc_union(): # [START bigtable_api_create_family_gc_union] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -898,8 +898,8 @@ def test_bigtable_create_family_gc_union(): def test_bigtable_create_family_gc_intersection(): # [START bigtable_api_create_family_gc_intersection] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -927,8 +927,8 @@ def test_bigtable_create_family_gc_intersection(): def test_bigtable_create_family_gc_nested(): # [START bigtable_api_create_family_gc_nested] - from google.cloud.bigtable.deprecated import Client - from google.cloud.bigtable.deprecated import column_family + from google.cloud.bigtable import Client + from google.cloud.bigtable import column_family client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -978,7 +978,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): row.commit() # [START bigtable_api_row_data_cells] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -993,7 +993,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): assert actual_cell_value == value # [START bigtable_api_row_cell_value] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1006,7 +1006,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): assert cell_value == value # [START bigtable_api_row_cell_values] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1025,7 +1025,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): row.commit() # [START bigtable_api_row_find_cells] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1042,7 +1042,7 @@ def test_bigtable_row_data_cells_cell_value_cell_values(): def test_bigtable_row_setcell_rowkey(): # [START bigtable_api_row_set_cell] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1061,7 +1061,7 @@ def test_bigtable_row_setcell_rowkey(): assert status.code == 0 # [START bigtable_api_row_row_key] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1073,7 +1073,7 @@ def test_bigtable_row_setcell_rowkey(): assert row_key == ROW_KEY1 # [START bigtable_api_row_table] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1098,7 +1098,7 @@ def test_bigtable_row_delete(): assert written_row_keys == [b"row_key_1"] # [START bigtable_api_row_delete] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1130,7 +1130,7 @@ def test_bigtable_row_delete_cell(): assert written_row_keys == [row_key1] # [START bigtable_api_row_delete_cell] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1163,7 +1163,7 @@ def test_bigtable_row_delete_cells(): assert written_row_keys == [row_key1] # [START bigtable_api_row_delete_cells] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1189,7 +1189,7 @@ def test_bigtable_row_clear(): assert mutation_size > 0 # [START bigtable_api_row_clear] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1208,7 +1208,7 @@ def test_bigtable_row_clear(): def test_bigtable_row_clear_get_mutations_size(): # [START bigtable_api_row_get_mutations_size] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1230,7 +1230,7 @@ def test_bigtable_row_clear_get_mutations_size(): def test_bigtable_row_setcell_commit_rowkey(): # [START bigtable_api_row_set_cell] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1244,7 +1244,7 @@ def test_bigtable_row_setcell_commit_rowkey(): row_obj.commit() # [START bigtable_api_row_commit] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1264,7 +1264,7 @@ def test_bigtable_row_setcell_commit_rowkey(): assert written_row_keys == [b"row_key_1", b"row_key_2"] # [START bigtable_api_row_row_key] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1286,7 +1286,7 @@ def test_bigtable_row_append_cell_value(): row.commit() # [START bigtable_api_row_append_cell_value] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1303,7 +1303,7 @@ def test_bigtable_row_append_cell_value(): assert actual_value == cell_val1 + cell_val2 # [START bigtable_api_row_commit] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) @@ -1315,7 +1315,7 @@ def test_bigtable_row_append_cell_value(): # [END bigtable_api_row_commit] # [START bigtable_api_row_increment_cell_value] - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client client = Client(admin=True) instance = client.instance(INSTANCE_ID) diff --git a/docs/table-api.rst b/docs/table-api.rst index ce05a3419..1bbf85146 100644 --- a/docs/table-api.rst +++ b/docs/table-api.rst @@ -1,7 +1,7 @@ Table Admin API =============== -After creating an :class:`Instance `, you can +After creating an :class:`Instance `, you can interact with individual tables, groups of tables or column families within a table. @@ -10,33 +10,33 @@ List Tables If you want a comprehensive list of all existing tables in a instance, make a `ListTables`_ API request with -:meth:`Instance.list_tables() `: +:meth:`Instance.list_tables() `: .. code:: python >>> instance.list_tables() - [, - ] + [, + ] Table Factory ------------- -To create a :class:`Table ` object: +To create a :class:`Table ` object: .. code:: python table = instance.table(table_id) -Even if this :class:`Table ` already +Even if this :class:`Table ` already has been created with the API, you'll want this object to use as a -parent of a :class:`ColumnFamily ` -or :class:`Row `. +parent of a :class:`ColumnFamily ` +or :class:`Row `. Create a new Table ------------------ After creating the table object, make a `CreateTable`_ API request -with :meth:`create() `: +with :meth:`create() `: .. code:: python @@ -53,7 +53,7 @@ Delete an existing Table ------------------------ Make a `DeleteTable`_ API request with -:meth:`delete() `: +:meth:`delete() `: .. code:: python @@ -67,7 +67,7 @@ associated with a table, the `GetTable`_ API method returns a table object with the names of the column families. To retrieve the list of column families use -:meth:`list_column_families() `: +:meth:`list_column_families() `: .. code:: python @@ -77,7 +77,7 @@ Column Family Factory --------------------- To create a -:class:`ColumnFamily ` object: +:class:`ColumnFamily ` object: .. code:: python @@ -87,7 +87,7 @@ There is no real reason to use this factory unless you intend to create or delete a column family. In addition, you can specify an optional ``gc_rule`` (a -:class:`GarbageCollectionRule ` +:class:`GarbageCollectionRule ` or similar): .. code:: python @@ -99,7 +99,7 @@ This rule helps the backend determine when and how to clean up old cells in the column family. See :doc:`column-family` for more information about -:class:`GarbageCollectionRule ` +:class:`GarbageCollectionRule ` and related classes. Create a new Column Family @@ -107,7 +107,7 @@ Create a new Column Family After creating the column family object, make a `CreateColumnFamily`_ API request with -:meth:`ColumnFamily.create() ` +:meth:`ColumnFamily.create() ` .. code:: python @@ -117,7 +117,7 @@ Delete an existing Column Family -------------------------------- Make a `DeleteColumnFamily`_ API request with -:meth:`ColumnFamily.delete() ` +:meth:`ColumnFamily.delete() ` .. code:: python @@ -127,7 +127,7 @@ Update an existing Column Family -------------------------------- Make an `UpdateColumnFamily`_ API request with -:meth:`ColumnFamily.delete() ` +:meth:`ColumnFamily.delete() ` .. code:: python @@ -137,9 +137,9 @@ Next Step --------- Now we go down the final step of the hierarchy from -:class:`Table ` to -:class:`Row ` as well as streaming -data directly via a :class:`Table `. +:class:`Table ` to +:class:`Row ` as well as streaming +data directly via a :class:`Table `. Head next to learn about the :doc:`data-api`. diff --git a/docs/table.rst b/docs/table.rst index 0d938e0af..c230725d1 100644 --- a/docs/table.rst +++ b/docs/table.rst @@ -1,6 +1,6 @@ Table ~~~~~ -.. automodule:: google.cloud.bigtable.deprecated.table +.. automodule:: google.cloud.bigtable.table :members: :show-inheritance: diff --git a/docs/usage.rst b/docs/usage.rst index 80fb65898..73a32b039 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -17,16 +17,17 @@ Using the API row-data row-filters row-set + batcher In the hierarchy of API concepts -* a :class:`Client ` owns an - :class:`Instance ` -* an :class:`Instance ` owns a - :class:`Table ` -* a :class:`Table ` owns a - :class:`ColumnFamily ` -* a :class:`Table ` owns a - :class:`Row ` +* a :class:`Client ` owns an + :class:`Instance ` +* an :class:`Instance ` owns a + :class:`Table ` +* a :class:`Table ` owns a + :class:`ColumnFamily ` +* a :class:`Table ` owns a + :class:`Row ` (and all the cells in the row) diff --git a/google/cloud/bigtable/__init__.py b/google/cloud/bigtable/__init__.py index 06b45bc4d..7331ff241 100644 --- a/google/cloud/bigtable/__init__.py +++ b/google/cloud/bigtable/__init__.py @@ -1,5 +1,4 @@ -# -*- coding: utf-8 -*- -# Copyright 2023 Google LLC +# Copyright 2015 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,48 +11,15 @@ # 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. -# - -from typing import List, Tuple -from google.cloud.bigtable import gapic_version as package_version +"""Google Cloud Bigtable API package.""" -from google.cloud.bigtable.client import BigtableDataClient -from google.cloud.bigtable.client import Table +from google.cloud.bigtable.client import Client -from google.cloud.bigtable.read_rows_query import ReadRowsQuery -from google.cloud.bigtable.read_rows_query import RowRange -from google.cloud.bigtable.row import Row -from google.cloud.bigtable.row import Cell - -from google.cloud.bigtable.mutations_batcher import MutationsBatcher -from google.cloud.bigtable.mutations import Mutation -from google.cloud.bigtable.mutations import RowMutationEntry -from google.cloud.bigtable.mutations import SetCell -from google.cloud.bigtable.mutations import DeleteRangeFromColumn -from google.cloud.bigtable.mutations import DeleteAllFromFamily -from google.cloud.bigtable.mutations import DeleteAllFromRow +from google.cloud.bigtable import gapic_version as package_version -# Type alias for the output of sample_keys -RowKeySamples = List[Tuple[bytes, int]] -# type alias for the output of query.shard() -ShardedQuery = List[ReadRowsQuery] +__version__: str -__version__: str = package_version.__version__ +__version__ = package_version.__version__ -__all__ = ( - "BigtableDataClient", - "Table", - "RowKeySamples", - "ReadRowsQuery", - "RowRange", - "MutationsBatcher", - "Mutation", - "RowMutationEntry", - "SetCell", - "DeleteRangeFromColumn", - "DeleteAllFromFamily", - "DeleteAllFromRow", - "Row", - "Cell", -) +__all__ = ["__version__", "Client"] diff --git a/google/cloud/bigtable/deprecated/app_profile.py b/google/cloud/bigtable/app_profile.py similarity index 97% rename from google/cloud/bigtable/deprecated/app_profile.py rename to google/cloud/bigtable/app_profile.py index a5c3df356..8cde66146 100644 --- a/google/cloud/bigtable/deprecated/app_profile.py +++ b/google/cloud/bigtable/app_profile.py @@ -17,7 +17,7 @@ import re -from google.cloud.bigtable.deprecated.enums import RoutingPolicyType +from google.cloud.bigtable.enums import RoutingPolicyType from google.cloud.bigtable_admin_v2.types import instance from google.protobuf import field_mask_pb2 from google.api_core.exceptions import NotFound @@ -47,8 +47,8 @@ class AppProfile(object): :param: routing_policy_type: (Optional) The type of the routing policy. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.deprecated.enums.RoutingPolicyType.ANY` - :data:`google.cloud.bigtable.deprecated.enums.RoutingPolicyType.SINGLE` + :data:`google.cloud.bigtable.enums.RoutingPolicyType.ANY` + :data:`google.cloud.bigtable.enums.RoutingPolicyType.SINGLE` :type: description: str :param: description: (Optional) Long form description of the use @@ -148,7 +148,7 @@ def from_pb(cls, app_profile_pb, instance): :type app_profile_pb: :class:`instance.app_profile_pb` :param app_profile_pb: An instance protobuf object. - :type instance: :class:`google.cloud.bigtable.deprecated.instance.Instance` + :type instance: :class:`google.cloud.bigtable.instance.Instance` :param instance: The instance that owns the cluster. :rtype: :class:`AppProfile` diff --git a/google/cloud/bigtable/deprecated/backup.py b/google/cloud/bigtable/backup.py similarity index 96% rename from google/cloud/bigtable/deprecated/backup.py rename to google/cloud/bigtable/backup.py index fc15318bc..6986d730a 100644 --- a/google/cloud/bigtable/deprecated/backup.py +++ b/google/cloud/bigtable/backup.py @@ -19,8 +19,8 @@ from google.cloud._helpers import _datetime_to_pb_timestamp # type: ignore from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient from google.cloud.bigtable_admin_v2.types import table -from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo -from google.cloud.bigtable.deprecated.policy import Policy +from google.cloud.bigtable.encryption_info import EncryptionInfo +from google.cloud.bigtable.policy import Policy from google.cloud.exceptions import NotFound # type: ignore from google.protobuf import field_mask_pb2 @@ -50,7 +50,7 @@ class Backup(object): :type backup_id: str :param backup_id: The ID of the backup. - :type instance: :class:`~google.cloud.bigtable.deprecated.instance.Instance` + :type instance: :class:`~google.cloud.bigtable.instance.Instance` :param instance: The Instance that owns this Backup. :type cluster_id: str @@ -188,7 +188,7 @@ def expire_time(self, new_expire_time): def encryption_info(self): """Encryption info for this Backup. - :rtype: :class:`google.cloud.bigtable.deprecated.encryption.EncryptionInfo` + :rtype: :class:`google.cloud.bigtable.encryption.EncryptionInfo` :returns: The encryption information for this backup. """ return self._encryption_info @@ -238,10 +238,10 @@ def from_pb(cls, backup_pb, instance): :type backup_pb: :class:`table.Backup` :param backup_pb: A Backup protobuf object. - :type instance: :class:`Instance ` + :type instance: :class:`Instance ` :param instance: The Instance that owns the Backup. - :rtype: :class:`~google.cloud.bigtable.deprecated.backup.Backup` + :rtype: :class:`~google.cloud.bigtable.backup.Backup` :returns: The backup parsed from the protobuf response. :raises: ValueError: If the backup name does not match the expected format or the parsed project ID does not match the @@ -440,7 +440,7 @@ def restore(self, table_id, instance_id=None): def get_iam_policy(self): """Gets the IAM access control policy for this backup. - :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :rtype: :class:`google.cloud.bigtable.policy.Policy` :returns: The current IAM policy of this backup. """ table_api = self._instance._client.table_admin_client @@ -452,13 +452,13 @@ def set_iam_policy(self, policy): existing policy. For more information about policy, please see documentation of - class `google.cloud.bigtable.deprecated.policy.Policy` + class `google.cloud.bigtable.policy.Policy` - :type policy: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :type policy: :class:`google.cloud.bigtable.policy.Policy` :param policy: A new IAM policy to replace the current IAM policy of this backup. - :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :rtype: :class:`google.cloud.bigtable.policy.Policy` :returns: The current IAM policy of this backup. """ table_api = self._instance._client.table_admin_client diff --git a/google/cloud/bigtable/batcher.py b/google/cloud/bigtable/batcher.py new file mode 100644 index 000000000..a6eb806e9 --- /dev/null +++ b/google/cloud/bigtable/batcher.py @@ -0,0 +1,395 @@ +# Copyright 2018 Google LLC +# +# Licensed 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. + +"""User friendly container for Google Cloud Bigtable MutationBatcher.""" +import threading +import queue +import concurrent.futures +import atexit + + +from google.api_core.exceptions import from_grpc_status +from dataclasses import dataclass + + +FLUSH_COUNT = 100 # after this many elements, send out the batch + +MAX_MUTATION_SIZE = 20 * 1024 * 1024 # 20MB # after this many bytes, send out the batch + +MAX_OUTSTANDING_BYTES = 100 * 1024 * 1024 # 100MB # max inflight byte size. + +MAX_OUTSTANDING_ELEMENTS = 100000 # max inflight mutations. + + +class MutationsBatchError(Exception): + """Error in the batch request""" + + def __init__(self, message, exc): + self.exc = exc + self.message = message + super().__init__(self.message) + + +class _MutationsBatchQueue(object): + """Private Threadsafe Queue to hold rows for batching.""" + + def __init__(self, max_mutation_bytes=MAX_MUTATION_SIZE, flush_count=FLUSH_COUNT): + """Specify the queue constraints""" + self._queue = queue.Queue() + self.total_mutation_count = 0 + self.total_size = 0 + self.max_mutation_bytes = max_mutation_bytes + self.flush_count = flush_count + + def get(self): + """Retrieve an item from the queue. Recalculate queue size.""" + row = self._queue.get() + mutation_size = row.get_mutations_size() + self.total_mutation_count -= len(row._get_mutations()) + self.total_size -= mutation_size + return row + + def put(self, item): + """Insert an item to the queue. Recalculate queue size.""" + + mutation_count = len(item._get_mutations()) + + self._queue.put(item) + + self.total_size += item.get_mutations_size() + self.total_mutation_count += mutation_count + + def full(self): + """Check if the queue is full.""" + if ( + self.total_mutation_count >= self.flush_count + or self.total_size >= self.max_mutation_bytes + ): + return True + return False + + def empty(self): + return self._queue.empty() + + +@dataclass +class _BatchInfo: + """Keeping track of size of a batch""" + + mutations_count: int = 0 + rows_count: int = 0 + mutations_size: int = 0 + + +class _FlowControl(object): + def __init__( + self, + max_mutations=MAX_OUTSTANDING_ELEMENTS, + max_mutation_bytes=MAX_OUTSTANDING_BYTES, + ): + """Control the inflight requests. Keep track of the mutations, row bytes and row counts. + As requests to backend are being made, adjust the number of mutations being processed. + + If threshold is reached, block the flow. + Reopen the flow as requests are finished. + """ + self.max_mutations = max_mutations + self.max_mutation_bytes = max_mutation_bytes + self.inflight_mutations = 0 + self.inflight_size = 0 + self.event = threading.Event() + self.event.set() + + def is_blocked(self): + """Returns True if: + + - inflight mutations >= max_mutations, or + - inflight bytes size >= max_mutation_bytes, or + """ + + return ( + self.inflight_mutations >= self.max_mutations + or self.inflight_size >= self.max_mutation_bytes + ) + + def control_flow(self, batch_info): + """ + Calculate the resources used by this batch + """ + + self.inflight_mutations += batch_info.mutations_count + self.inflight_size += batch_info.mutations_size + self.set_flow_control_status() + + def wait(self): + """ + Wait until flow control pushback has been released. + It awakens as soon as `event` is set. + """ + self.event.wait() + + def set_flow_control_status(self): + """Check the inflight mutations and size. + + If values exceed the allowed threshold, block the event. + """ + if self.is_blocked(): + self.event.clear() # sleep + else: + self.event.set() # awaken the threads + + def release(self, batch_info): + """ + Release the resources. + Decrement the row size to allow enqueued mutations to be run. + """ + self.inflight_mutations -= batch_info.mutations_count + self.inflight_size -= batch_info.mutations_size + self.set_flow_control_status() + + +class MutationsBatcher(object): + """A MutationsBatcher is used in batch cases where the number of mutations + is large or unknown. It will store :class:`DirectRow` in memory until one of the + size limits is reached, or an explicit call to :func:`flush()` is performed. When + a flush event occurs, the :class:`DirectRow` in memory will be sent to Cloud + Bigtable. Batching mutations is more efficient than sending individual + request. + + This class is not suited for usage in systems where each mutation + must be guaranteed to be sent, since calling mutate may only result in an + in-memory change. In a case of a system crash, any :class:`DirectRow` remaining in + memory will not necessarily be sent to the service, even after the + completion of the :func:`mutate()` method. + + Note on thread safety: The same :class:`MutationBatcher` cannot be shared by multiple end-user threads. + + :type table: class + :param table: class:`~google.cloud.bigtable.table.Table`. + + :type flush_count: int + :param flush_count: (Optional) Max number of rows to flush. If it + reaches the max number of rows it calls finish_batch() to mutate the + current row batch. Default is FLUSH_COUNT (1000 rows). + + :type max_row_bytes: int + :param max_row_bytes: (Optional) Max number of row mutations size to + flush. If it reaches the max number of row mutations size it calls + finish_batch() to mutate the current row batch. Default is MAX_ROW_BYTES + (5 MB). + + :type flush_interval: float + :param flush_interval: (Optional) The interval (in seconds) between asynchronous flush. + Default is 1 second. + + :type batch_completed_callback: Callable[list:[`~google.rpc.status_pb2.Status`]] = None + :param batch_completed_callback: (Optional) A callable for handling responses + after the current batch is sent. The callable function expect a list of grpc + Status. + """ + + def __init__( + self, + table, + flush_count=FLUSH_COUNT, + max_row_bytes=MAX_MUTATION_SIZE, + flush_interval=1, + batch_completed_callback=None, + ): + self._rows = _MutationsBatchQueue( + max_mutation_bytes=max_row_bytes, flush_count=flush_count + ) + self.table = table + self._executor = concurrent.futures.ThreadPoolExecutor() + atexit.register(self.close) + self._timer = threading.Timer(flush_interval, self.flush) + self._timer.start() + self.flow_control = _FlowControl( + max_mutations=MAX_OUTSTANDING_ELEMENTS, + max_mutation_bytes=MAX_OUTSTANDING_BYTES, + ) + self.futures_mapping = {} + self.exceptions = queue.Queue() + self._user_batch_completed_callback = batch_completed_callback + + @property + def flush_count(self): + return self._rows.flush_count + + @property + def max_row_bytes(self): + return self._rows.max_mutation_bytes + + def __enter__(self): + """Starting the MutationsBatcher as a context manager""" + return self + + def mutate(self, row): + """Add a row to the batch. If the current batch meets one of the size + limits, the batch is sent asynchronously. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_batcher_mutate] + :end-before: [END bigtable_api_batcher_mutate] + :dedent: 4 + + :type row: class + :param row: :class:`~google.cloud.bigtable.row.DirectRow`. + + :raises: One of the following: + * :exc:`~.table._BigtableRetryableError` if any row returned a transient error. + * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried + """ + self._rows.put(row) + + if self._rows.full(): + self._flush_async() + + def mutate_rows(self, rows): + """Add multiple rows to the batch. If the current batch meets one of the size + limits, the batch is sent asynchronously. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_batcher_mutate_rows] + :end-before: [END bigtable_api_batcher_mutate_rows] + :dedent: 4 + + :type rows: list:[`~google.cloud.bigtable.row.DirectRow`] + :param rows: list:[`~google.cloud.bigtable.row.DirectRow`]. + + :raises: One of the following: + * :exc:`~.table._BigtableRetryableError` if any row returned a transient error. + * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried + """ + for row in rows: + self.mutate(row) + + def flush(self): + """Sends the current batch to Cloud Bigtable synchronously. + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_batcher_flush] + :end-before: [END bigtable_api_batcher_flush] + :dedent: 4 + + :raises: + * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. + """ + rows_to_flush = [] + while not self._rows.empty(): + rows_to_flush.append(self._rows.get()) + response = self._flush_rows(rows_to_flush) + return response + + def _flush_async(self): + """Sends the current batch to Cloud Bigtable asynchronously. + + :raises: + * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. + """ + + rows_to_flush = [] + mutations_count = 0 + mutations_size = 0 + rows_count = 0 + batch_info = _BatchInfo() + + while not self._rows.empty(): + row = self._rows.get() + mutations_count += len(row._get_mutations()) + mutations_size += row.get_mutations_size() + rows_count += 1 + rows_to_flush.append(row) + batch_info.mutations_count = mutations_count + batch_info.rows_count = rows_count + batch_info.mutations_size = mutations_size + + if ( + rows_count >= self.flush_count + or mutations_size >= self.max_row_bytes + or mutations_count >= self.flow_control.max_mutations + or mutations_size >= self.flow_control.max_mutation_bytes + or self._rows.empty() # submit when it reached the end of the queue + ): + # wait for resources to become available, before submitting any new batch + self.flow_control.wait() + # once unblocked, submit a batch + # event flag will be set by control_flow to block subsequent thread, but not blocking this one + self.flow_control.control_flow(batch_info) + future = self._executor.submit(self._flush_rows, rows_to_flush) + self.futures_mapping[future] = batch_info + future.add_done_callback(self._batch_completed_callback) + + # reset and start a new batch + rows_to_flush = [] + mutations_size = 0 + rows_count = 0 + mutations_count = 0 + batch_info = _BatchInfo() + + def _batch_completed_callback(self, future): + """Callback for when the mutation has finished to clean up the current batch + and release items from the flow controller. + + Raise exceptions if there's any. + Release the resources locked by the flow control and allow enqueued tasks to be run. + """ + + processed_rows = self.futures_mapping[future] + self.flow_control.release(processed_rows) + del self.futures_mapping[future] + + def _flush_rows(self, rows_to_flush): + """Mutate the specified rows. + + :raises: + * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. + """ + responses = [] + if len(rows_to_flush) > 0: + response = self.table.mutate_rows(rows_to_flush) + + if self._user_batch_completed_callback: + self._user_batch_completed_callback(response) + + for result in response: + if result.code != 0: + exc = from_grpc_status(result.code, result.message) + self.exceptions.put(exc) + responses.append(result) + + return responses + + def __exit__(self, exc_type, exc_value, exc_traceback): + """Clean up resources. Flush and shutdown the ThreadPoolExecutor.""" + self.close() + + def close(self): + """Clean up resources. Flush and shutdown the ThreadPoolExecutor. + Any errors will be raised. + + :raises: + * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. + """ + self.flush() + self._executor.shutdown(wait=True) + atexit.unregister(self.close) + if self.exceptions.qsize() > 0: + exc = list(self.exceptions.queue) + raise MutationsBatchError("Errors in batch mutations.", exc=exc) diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index 4ec3cea27..c82a268c6 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -1,4 +1,4 @@ -# Copyright 2023 Google LLC +# Copyright 2015 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,1081 +11,503 @@ # 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. -# -from __future__ import annotations +"""Parent client for calling the Google Cloud Bigtable API. -from typing import ( - cast, - Any, - Optional, - Set, - TYPE_CHECKING, -) +This is the base from which all interactions with the API occur. + +In the hierarchy of API concepts -import asyncio -import grpc -import time +* a :class:`~google.cloud.bigtable.client.Client` owns an + :class:`~google.cloud.bigtable.instance.Instance` +* an :class:`~google.cloud.bigtable.instance.Instance` owns a + :class:`~google.cloud.bigtable.table.Table` +* a :class:`~google.cloud.bigtable.table.Table` owns a + :class:`~.column_family.ColumnFamily` +* a :class:`~google.cloud.bigtable.table.Table` owns a + :class:`~google.cloud.bigtable.row.Row` (and all the cells in the row) +""" +import os import warnings -import sys -import random +import grpc # type: ignore -from collections import namedtuple +from google.api_core.gapic_v1 import client_info as client_info_lib +import google.auth # type: ignore +from google.auth.credentials import AnonymousCredentials # type: ignore -from google.cloud.bigtable_v2.services.bigtable.client import BigtableClientMeta -from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient -from google.cloud.bigtable_v2.services.bigtable.async_client import DEFAULT_CLIENT_INFO -from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( - PooledBigtableGrpcAsyncIOTransport, +from google.cloud import bigtable_v2 +from google.cloud import bigtable_admin_v2 +from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport +from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin.transports import ( + BigtableInstanceAdminGrpcTransport, ) -from google.cloud.bigtable_v2.types.bigtable import PingAndWarmRequest -from google.cloud.client import ClientWithProject -from google.api_core.exceptions import GoogleAPICallError -from google.api_core import retry_async as retries -from google.api_core import exceptions as core_exceptions -from google.cloud.bigtable._read_rows import _ReadRowsOperation - -import google.auth.credentials -import google.auth._default -from google.api_core import client_options as client_options_lib -from google.cloud.bigtable.row import Row -from google.cloud.bigtable.read_rows_query import ReadRowsQuery -from google.cloud.bigtable.iterators import ReadRowsIterator -from google.cloud.bigtable.exceptions import FailedQueryShardError -from google.cloud.bigtable.exceptions import ShardedReadRowsExceptionGroup - -from google.cloud.bigtable.mutations import Mutation, RowMutationEntry -from google.cloud.bigtable._mutate_rows import _MutateRowsOperation -from google.cloud.bigtable._helpers import _make_metadata -from google.cloud.bigtable._helpers import _convert_retry_deadline -from google.cloud.bigtable.mutations_batcher import MutationsBatcher -from google.cloud.bigtable.mutations_batcher import _MB_SIZE -from google.cloud.bigtable._helpers import _attempt_timeout_generator - -from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule -from google.cloud.bigtable.row_filters import RowFilter -from google.cloud.bigtable.row_filters import StripValueTransformerFilter -from google.cloud.bigtable.row_filters import CellsRowLimitFilter -from google.cloud.bigtable.row_filters import RowFilterChain - -if TYPE_CHECKING: - from google.cloud.bigtable import RowKeySamples - from google.cloud.bigtable import ShardedQuery - -# used by read_rows_sharded to limit how many requests are attempted in parallel -CONCURRENCY_LIMIT = 10 - -# used to register instance data with the client for channel warming -_WarmedInstanceKey = namedtuple( - "_WarmedInstanceKey", ["instance_name", "table_name", "app_profile_id"] +from google.cloud.bigtable_admin_v2.services.bigtable_table_admin.transports import ( + BigtableTableAdminGrpcTransport, ) +from google.cloud import bigtable +from google.cloud.bigtable.instance import Instance +from google.cloud.bigtable.cluster import Cluster + +from google.cloud.client import ClientWithProject # type: ignore + +from google.cloud.bigtable_admin_v2.types import instance +from google.cloud.bigtable.cluster import _CLUSTER_NAME_RE +from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore + + +INSTANCE_TYPE_PRODUCTION = instance.Instance.Type.PRODUCTION +INSTANCE_TYPE_DEVELOPMENT = instance.Instance.Type.DEVELOPMENT +INSTANCE_TYPE_UNSPECIFIED = instance.Instance.Type.TYPE_UNSPECIFIED +SPANNER_ADMIN_SCOPE = "https://www.googleapis.com/auth/spanner.admin" +ADMIN_SCOPE = "https://www.googleapis.com/auth/bigtable.admin" +"""Scope for interacting with the Cluster Admin and Table Admin APIs.""" +DATA_SCOPE = "https://www.googleapis.com/auth/bigtable.data" +"""Scope for reading and writing table data.""" +READ_ONLY_SCOPE = "https://www.googleapis.com/auth/bigtable.data.readonly" +"""Scope for reading table data.""" + +_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator" +_GRPC_CHANNEL_OPTIONS = ( + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ("grpc.keepalive_time_ms", 30000), + ("grpc.keepalive_timeout_ms", 10000), +) -class BigtableDataClient(ClientWithProject): - def __init__( - self, - *, - project: str | None = None, - pool_size: int = 3, - credentials: google.auth.credentials.Credentials | None = None, - client_options: dict[str, Any] - | "google.api_core.client_options.ClientOptions" - | None = None, - ): - """ - Create a client instance for the Bigtable Data API - - Client should be created within an async context (running event loop) - - Args: - project: the project which the client acts on behalf of. - If not passed, falls back to the default inferred - from the environment. - pool_size: The number of grpc channels to maintain - in the internal channel pool. - credentials: - Thehe OAuth2 Credentials to use for this - client. If not passed (and if no ``_http`` object is - passed), falls back to the default inferred from the - environment. - client_options (Optional[Union[dict, google.api_core.client_options.ClientOptions]]): - Client options used to set user options - on the client. API Endpoint should be set through client_options. - Raises: - - RuntimeError if called outside of an async context (no running event loop) - - ValueError if pool_size is less than 1 - """ - # set up transport in registry - transport_str = f"pooled_grpc_asyncio_{pool_size}" - transport = PooledBigtableGrpcAsyncIOTransport.with_fixed_size(pool_size) - BigtableClientMeta._transport_registry[transport_str] = transport - # set up client info headers for veneer library - client_info = DEFAULT_CLIENT_INFO - client_info.client_library_version = client_info.gapic_version - # parse client options - if type(client_options) is dict: - client_options = client_options_lib.from_dict(client_options) - client_options = cast( - Optional[client_options_lib.ClientOptions], client_options - ) - # initialize client - ClientWithProject.__init__( - self, - credentials=credentials, - project=project, - client_options=client_options, - ) - self._gapic_client = BigtableAsyncClient( - transport=transport_str, - credentials=credentials, + +def _create_gapic_client(client_class, client_options=None, transport=None): + def inner(self): + return client_class( + credentials=None, + client_info=self._client_info, client_options=client_options, - client_info=client_info, - ) - self.transport = cast( - PooledBigtableGrpcAsyncIOTransport, self._gapic_client.transport + transport=transport, ) - # keep track of active instances to for warmup on channel refresh - self._active_instances: Set[_WarmedInstanceKey] = set() - # keep track of table objects associated with each instance - # only remove instance from _active_instances when all associated tables remove it - self._instance_owners: dict[_WarmedInstanceKey, Set[int]] = {} - # attempt to start background tasks - self._channel_init_time = time.monotonic() - self._channel_refresh_tasks: list[asyncio.Task[None]] = [] - try: - self.start_background_channel_refresh() - except RuntimeError: - warnings.warn( - f"{self.__class__.__name__} should be started in an " - "asyncio event loop. Channel refresh will not be started", - RuntimeWarning, - stacklevel=2, - ) - def start_background_channel_refresh(self) -> None: - """ - Starts a background task to ping and warm each channel in the pool - Raises: - - RuntimeError if not called in an asyncio event loop - """ - if not self._channel_refresh_tasks: - # raise RuntimeError if there is no event loop - asyncio.get_running_loop() - for channel_idx in range(self.transport.pool_size): - refresh_task = asyncio.create_task(self._manage_channel(channel_idx)) - if sys.version_info >= (3, 8): - # task names supported in Python 3.8+ - refresh_task.set_name( - f"{self.__class__.__name__} channel refresh {channel_idx}" - ) - self._channel_refresh_tasks.append(refresh_task) - - async def close(self, timeout: float = 2.0): - """ - Cancel all background tasks - """ - for task in self._channel_refresh_tasks: - task.cancel() - group = asyncio.gather(*self._channel_refresh_tasks, return_exceptions=True) - await asyncio.wait_for(group, timeout=timeout) - await self.transport.close() - self._channel_refresh_tasks = [] - - async def _ping_and_warm_instances( - self, channel: grpc.aio.Channel, instance_key: _WarmedInstanceKey | None = None - ) -> list[GoogleAPICallError | None]: - """ - Prepares the backend for requests on a channel + return inner - Pings each Bigtable instance registered in `_active_instances` on the client - Args: - - channel: grpc channel to warm - - instance_key: if provided, only warm the instance associated with the key - Returns: - - sequence of results or exceptions from the ping requests - """ - instance_list = ( - [instance_key] if instance_key is not None else self._active_instances - ) - ping_rpc = channel.unary_unary( - "/google.bigtable.v2.Bigtable/PingAndWarm", - request_serializer=PingAndWarmRequest.serialize, - ) - # prepare list of coroutines to run - tasks = [ - ping_rpc( - request={"name": instance_name, "app_profile_id": app_profile_id}, - metadata=[ - ( - "x-goog-request-params", - f"name={instance_name}&app_profile_id={app_profile_id}", - ) - ], - wait_for_ready=True, - ) - for (instance_name, table_name, app_profile_id) in instance_list - ] - # execute coroutines in parallel - result_list = await asyncio.gather(*tasks, return_exceptions=True) - # return None in place of empty successful responses - return [r or None for r in result_list] - - async def _manage_channel( - self, - channel_idx: int, - refresh_interval_min: float = 60 * 35, - refresh_interval_max: float = 60 * 45, - grace_period: float = 60 * 10, - ) -> None: - """ - Background coroutine that periodically refreshes and warms a grpc channel - - The backend will automatically close channels after 60 minutes, so - `refresh_interval` + `grace_period` should be < 60 minutes - - Runs continuously until the client is closed - - Args: - channel_idx: index of the channel in the transport's channel pool - refresh_interval_min: minimum interval before initiating refresh - process in seconds. Actual interval will be a random value - between `refresh_interval_min` and `refresh_interval_max` - refresh_interval_max: maximum interval before initiating refresh - process in seconds. Actual interval will be a random value - between `refresh_interval_min` and `refresh_interval_max` - grace_period: time to allow previous channel to serve existing - requests before closing, in seconds - """ - first_refresh = self._channel_init_time + random.uniform( - refresh_interval_min, refresh_interval_max - ) - next_sleep = max(first_refresh - time.monotonic(), 0) - if next_sleep > 0: - # warm the current channel immediately - channel = self.transport.channels[channel_idx] - await self._ping_and_warm_instances(channel) - # continuously refresh the channel every `refresh_interval` seconds - while True: - await asyncio.sleep(next_sleep) - # prepare new channel for use - new_channel = self.transport.grpc_channel._create_channel() - await self._ping_and_warm_instances(new_channel) - # cycle channel out of use, with long grace window before closure - start_timestamp = time.time() - await self.transport.replace_channel( - channel_idx, grace=grace_period, swap_sleep=10, new_channel=new_channel - ) - # subtract the time spent waiting for the channel to be replaced - next_refresh = random.uniform(refresh_interval_min, refresh_interval_max) - next_sleep = next_refresh - (time.time() - start_timestamp) +class Client(ClientWithProject): + """Client for interacting with Google Cloud Bigtable API. - async def _register_instance(self, instance_id: str, owner: Table) -> None: - """ - Registers an instance with the client, and warms the channel pool - for the instance - The client will periodically refresh grpc channel pool used to make - requests, and new channels will be warmed for each registered instance - Channels will not be refreshed unless at least one instance is registered - - Args: - - instance_id: id of the instance to register. - - owner: table that owns the instance. Owners will be tracked in - _instance_owners, and instances will only be unregistered when all - owners call _remove_instance_registration - """ - instance_name = self._gapic_client.instance_path(self.project, instance_id) - instance_key = _WarmedInstanceKey( - instance_name, owner.table_name, owner.app_profile_id - ) - self._instance_owners.setdefault(instance_key, set()).add(id(owner)) - if instance_name not in self._active_instances: - self._active_instances.add(instance_key) - if self._channel_refresh_tasks: - # refresh tasks already running - # call ping and warm on all existing channels - for channel in self.transport.channels: - await self._ping_and_warm_instances(channel, instance_key) - else: - # refresh tasks aren't active. start them as background tasks - self.start_background_channel_refresh() - - async def _remove_instance_registration( - self, instance_id: str, owner: Table - ) -> bool: - """ - Removes an instance from the client's registered instances, to prevent - warming new channels for the instance + .. note:: - If instance_id is not registered, or is still in use by other tables, returns False + Since the Cloud Bigtable API requires the gRPC transport, no + ``_http`` argument is accepted by this class. - Args: - - instance_id: id of the instance to remove - - owner: table that owns the instance. Owners will be tracked in - _instance_owners, and instances will only be unregistered when all - owners call _remove_instance_registration - Returns: - - True if instance was removed - """ - instance_name = self._gapic_client.instance_path(self.project, instance_id) - instance_key = _WarmedInstanceKey( - instance_name, owner.table_name, owner.app_profile_id - ) - owner_list = self._instance_owners.get(instance_key, set()) - try: - owner_list.remove(id(owner)) - if len(owner_list) == 0: - self._active_instances.remove(instance_key) - return True - except KeyError: - return False - - # TODO: revisit timeouts https://github.com/googleapis/python-bigtable/issues/782 - def get_table( - self, - instance_id: str, - table_id: str, - app_profile_id: str | None = None, - default_operation_timeout: float = 600, - default_per_request_timeout: float | None = None, - ) -> Table: - """ - Returns a table instance for making data API requests - - Args: - instance_id: The Bigtable instance ID to associate with this client. - instance_id is combined with the client's project to fully - specify the instance - table_id: The ID of the table. - app_profile_id: (Optional) The app profile to associate with requests. - https://cloud.google.com/bigtable/docs/app-profiles - """ - return Table( - self, - instance_id, - table_id, - app_profile_id, - default_operation_timeout=default_operation_timeout, - default_per_request_timeout=default_per_request_timeout, - ) + :type project: :class:`str` or :func:`unicode ` + :param project: (Optional) The ID of the project which owns the + instances, tables and data. If not provided, will + attempt to determine from the environment. - async def __aenter__(self): - self.start_background_channel_refresh() - return self + :type credentials: :class:`~google.auth.credentials.Credentials` + :param credentials: (Optional) The OAuth2 Credentials to use for this + client. If not passed, falls back to the default + inferred from the environment. - async def __aexit__(self, exc_type, exc_val, exc_tb): - await self.close() - await self._gapic_client.__aexit__(exc_type, exc_val, exc_tb) + :type read_only: bool + :param read_only: (Optional) Boolean indicating if the data scope should be + for reading only (or for writing as well). Defaults to + :data:`False`. + :type admin: bool + :param admin: (Optional) Boolean indicating if the client will be used to + interact with the Instance Admin or Table Admin APIs. This + requires the :const:`ADMIN_SCOPE`. Defaults to :data:`False`. -class Table: - """ - Main Data API surface + :type: client_info: :class:`google.api_core.gapic_v1.client_info.ClientInfo` + :param client_info: + The client info used to send a user-agent string along with API + requests. If ``None``, then default info will be used. Generally, + you only need to set this if you're developing your own library + or partner tool. + + :type client_options: :class:`~google.api_core.client_options.ClientOptions` + or :class:`dict` + :param client_options: (Optional) Client options used to set user options + on the client. API Endpoint should be set through client_options. + + :type admin_client_options: + :class:`~google.api_core.client_options.ClientOptions` or :class:`dict` + :param admin_client_options: (Optional) Client options used to set user + options on the client. API Endpoint for admin operations should be set + through admin_client_options. - Table object maintains table_id, and app_profile_id context, and passes them with - each call + :type channel: :instance: grpc.Channel + :param channel (grpc.Channel): (Optional) DEPRECATED: + A ``Channel`` instance through which to make calls. + This argument is mutually exclusive with ``credentials``; + providing both will raise an exception. No longer used. + + :raises: :class:`ValueError ` if both ``read_only`` + and ``admin`` are :data:`True` """ + _table_data_client = None + _table_admin_client = None + _instance_admin_client = None + def __init__( self, - client: BigtableDataClient, - instance_id: str, - table_id: str, - app_profile_id: str | None = None, - *, - default_operation_timeout: float = 600, - default_per_request_timeout: float | None = None, + project=None, + credentials=None, + read_only=False, + admin=False, + client_info=None, + client_options=None, + admin_client_options=None, + channel=None, ): - """ - Initialize a Table instance - - Must be created within an async context (running event loop) - - Args: - instance_id: The Bigtable instance ID to associate with this client. - instance_id is combined with the client's project to fully - specify the instance - table_id: The ID of the table. table_id is combined with the - instance_id and the client's project to fully specify the table - app_profile_id: (Optional) The app profile to associate with requests. - https://cloud.google.com/bigtable/docs/app-profiles - default_operation_timeout: (Optional) The default timeout, in seconds - default_per_request_timeout: (Optional) The default timeout for individual - rpc requests, in seconds - Raises: - - RuntimeError if called outside of an async context (no running event loop) - """ - # validate timeouts - if default_operation_timeout <= 0: - raise ValueError("default_operation_timeout must be greater than 0") - if default_per_request_timeout is not None and default_per_request_timeout <= 0: - raise ValueError("default_per_request_timeout must be greater than 0") - if ( - default_per_request_timeout is not None - and default_per_request_timeout > default_operation_timeout - ): + if client_info is None: + client_info = client_info_lib.ClientInfo( + client_library_version=bigtable.__version__, + ) + if read_only and admin: raise ValueError( - "default_per_request_timeout must be less than default_operation_timeout" + "A read-only client cannot also perform" "administrative actions." ) - self.client = client - self.instance_id = instance_id - self.instance_name = self.client._gapic_client.instance_path( - self.client.project, instance_id - ) - self.table_id = table_id - self.table_name = self.client._gapic_client.table_path( - self.client.project, instance_id, table_id - ) - self.app_profile_id = app_profile_id - self.default_operation_timeout = default_operation_timeout - self.default_per_request_timeout = default_per_request_timeout + # NOTE: We set the scopes **before** calling the parent constructor. + # It **may** use those scopes in ``with_scopes_if_required``. + self._read_only = bool(read_only) + self._admin = bool(admin) + self._client_info = client_info + self._emulator_host = os.getenv(BIGTABLE_EMULATOR) - # raises RuntimeError if called outside of an async context (no running event loop) - try: - self._register_instance_task = asyncio.create_task( - self.client._register_instance(instance_id, self) - ) - except RuntimeError as e: - raise RuntimeError( - f"{self.__class__.__name__} must be created within an async event loop context." - ) from e + if self._emulator_host is not None: + if credentials is None: + credentials = AnonymousCredentials() + if project is None: + project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT - async def read_rows_stream( - self, - query: ReadRowsQuery | dict[str, Any], - *, - operation_timeout: float | None = None, - per_request_timeout: float | None = None, - ) -> ReadRowsIterator: - """ - Returns an iterator to asynchronously stream back row data. + if channel is not None: + warnings.warn( + "'channel' is deprecated and no longer used.", + DeprecationWarning, + stacklevel=2, + ) - Failed requests within operation_timeout and operation_deadline policies will be retried. + self._client_options = client_options + self._admin_client_options = admin_client_options + self._channel = channel + self.SCOPE = self._get_scopes() + super(Client, self).__init__( + project=project, + credentials=credentials, + client_options=client_options, + ) - Args: - - query: contains details about which rows to return - - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will be retried within the budget. - time is only counted while actively waiting on the network. - If None, defaults to the Table's default_operation_timeout - - per_request_timeout: the time budget for an individual network request, in seconds. - If it takes longer than this time to complete, the request will be cancelled with - a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_per_request_timeout + def _get_scopes(self): + """Get the scopes corresponding to admin / read-only state. Returns: - - an asynchronous iterator that yields rows returned by the query - Raises: - - DeadlineExceeded: raised after operation timeout - will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions - from any retries that failed - - GoogleAPIError: raised if the request encounters an unrecoverable error - - IdleTimeout: if iterator was abandoned + Tuple[str, ...]: The tuple of scopes. """ + if self._read_only: + scopes = (READ_ONLY_SCOPE,) + else: + scopes = (DATA_SCOPE,) - operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = per_request_timeout or self.default_per_request_timeout + if self._admin: + scopes += (ADMIN_SCOPE,) - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout <= 0: - raise ValueError("per_request_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout > operation_timeout: - raise ValueError( - "per_request_timeout must not be greater than operation_timeout" - ) - if per_request_timeout is None: - per_request_timeout = operation_timeout - request = query._to_dict() if isinstance(query, ReadRowsQuery) else query - request["table_name"] = self.table_name - if self.app_profile_id: - request["app_profile_id"] = self.app_profile_id - - # read_rows smart retries is implemented using a series of iterators: - # - client.read_rows: outputs raw ReadRowsResponse objects from backend. Has per_request_timeout - # - ReadRowsOperation.merge_row_response_stream: parses chunks into rows - # - ReadRowsOperation.retryable_merge_rows: adds retries, caching, revised requests, per_request_timeout - # - ReadRowsIterator: adds idle_timeout, moves stats out of stream and into attribute - row_merger = _ReadRowsOperation( - request, - self.client._gapic_client, - operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, - ) - output_generator = ReadRowsIterator(row_merger) - # add idle timeout to clear resources if generator is abandoned - idle_timeout_seconds = 300 - await output_generator._start_idle_timer(idle_timeout_seconds) - return output_generator + return scopes - async def read_rows( - self, - query: ReadRowsQuery | dict[str, Any], - *, - operation_timeout: float | None = None, - per_request_timeout: float | None = None, - ) -> list[Row]: - """ - Helper function that returns a full list instead of a generator + def _emulator_channel(self, transport, options): + """Create a channel using self._credentials - See read_rows_stream + Works in a similar way to ``grpc.secure_channel`` but using + ``grpc.local_channel_credentials`` rather than + ``grpc.ssh_channel_credentials`` to allow easy connection to a + local emulator. Returns: - - a list of the rows returned by the query - """ - row_generator = await self.read_rows_stream( - query, - operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, - ) - results = [row async for row in row_generator] - return results - - async def read_row( - self, - row_key: str | bytes, - *, - row_filter: RowFilter | None = None, - operation_timeout: int | float | None = 60, - per_request_timeout: int | float | None = None, - ) -> Row | None: - """ - Helper function to return a single row + grpc.Channel or grpc.aio.Channel + """ + # TODO: Implement a special credentials type for emulator and use + # "transport.create_channel" to create gRPC channels once google-auth + # extends it's allowed credentials types. + # Note: this code also exists in the firestore client. + if "GrpcAsyncIOTransport" in str(transport.__name__): + return grpc.aio.secure_channel( + self._emulator_host, + self._local_composite_credentials(), + options=options, + ) + else: + return grpc.secure_channel( + self._emulator_host, + self._local_composite_credentials(), + options=options, + ) - See read_rows_stream + def _local_composite_credentials(self): + """Create credentials for the local emulator channel. - Raises: - - google.cloud.bigtable.exceptions.RowNotFound: if the row does not exist - Returns: - - the individual row requested, or None if it does not exist + :return: grpc.ChannelCredentials """ - if row_key is None: - raise ValueError("row_key must be string or bytes") - query = ReadRowsQuery(row_keys=row_key, row_filter=row_filter, limit=1) - results = await self.read_rows( - query, - operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, + credentials = google.auth.credentials.with_scopes_if_required( + self._credentials, None ) - if len(results) == 0: - return None - return results[0] + request = google.auth.transport.requests.Request() - async def read_rows_sharded( - self, - sharded_query: ShardedQuery, - *, - operation_timeout: int | float | None = None, - per_request_timeout: int | float | None = None, - ) -> list[Row]: - """ - Runs a sharded query in parallel, then return the results in a single list. - Results will be returned in the order of the input queries. - - This function is intended to be run on the results on a query.shard() call: - - ``` - table_shard_keys = await table.sample_row_keys() - query = ReadRowsQuery(...) - shard_queries = query.shard(table_shard_keys) - results = await table.read_rows_sharded(shard_queries) - ``` - - Args: - - sharded_query: a sharded query to execute - Raises: - - ShardedReadRowsExceptionGroup: if any of the queries failed - - ValueError: if the query_list is empty - """ - if not sharded_query: - raise ValueError("empty sharded_query") - # reduce operation_timeout between batches - operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = ( - per_request_timeout or self.default_per_request_timeout or operation_timeout + # Create the metadata plugin for inserting the authorization header. + metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin( + credentials, request ) - timeout_generator = _attempt_timeout_generator( - operation_timeout, operation_timeout + + # Create a set of grpc.CallCredentials using the metadata plugin. + google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin) + + # Using the local_credentials to allow connection to emulator + local_credentials = grpc.local_channel_credentials() + + # Combine the local credentials and the authorization credentials. + return grpc.composite_channel_credentials( + local_credentials, google_auth_credentials ) - # submit shards in batches if the number of shards goes over CONCURRENCY_LIMIT - batched_queries = [ - sharded_query[i : i + CONCURRENCY_LIMIT] - for i in range(0, len(sharded_query), CONCURRENCY_LIMIT) - ] - # run batches and collect results - results_list = [] - error_dict = {} - shard_idx = 0 - for batch in batched_queries: - batch_operation_timeout = next(timeout_generator) - routine_list = [ - self.read_rows( - query, - operation_timeout=batch_operation_timeout, - per_request_timeout=min( - per_request_timeout, batch_operation_timeout - ), - ) - for query in batch - ] - batch_result = await asyncio.gather(*routine_list, return_exceptions=True) - for result in batch_result: - if isinstance(result, Exception): - error_dict[shard_idx] = result - else: - results_list.extend(result) - shard_idx += 1 - if error_dict: - # if any sub-request failed, raise an exception instead of returning results - raise ShardedReadRowsExceptionGroup( - [ - FailedQueryShardError(idx, sharded_query[idx], e) - for idx, e in error_dict.items() - ], - results_list, - len(sharded_query), + + def _create_gapic_client_channel(self, client_class, grpc_transport): + if self._emulator_host is not None: + api_endpoint = self._emulator_host + elif self._client_options and self._client_options.api_endpoint: + api_endpoint = self._client_options.api_endpoint + else: + api_endpoint = client_class.DEFAULT_ENDPOINT + + if self._emulator_host is not None: + channel = self._emulator_channel( + transport=grpc_transport, + options=_GRPC_CHANNEL_OPTIONS, ) - return results_list + else: + channel = grpc_transport.create_channel( + host=api_endpoint, + credentials=self._credentials, + options=_GRPC_CHANNEL_OPTIONS, + ) + return grpc_transport(channel=channel, host=api_endpoint) - async def row_exists( - self, - row_key: str | bytes, - *, - operation_timeout: int | float | None = 60, - per_request_timeout: int | float | None = None, - ) -> bool: - """ - Helper function to determine if a row exists + @property + def project_path(self): + """Project name to be used with Instance Admin API. - uses the filters: chain(limit cells per row = 1, strip value) + .. note:: - Returns: - - a bool indicating whether the row exists - """ - if row_key is None: - raise ValueError("row_key must be string or bytes") - strip_filter = StripValueTransformerFilter(flag=True) - limit_filter = CellsRowLimitFilter(1) - chain_filter = RowFilterChain(filters=[limit_filter, strip_filter]) - query = ReadRowsQuery(row_keys=row_key, limit=1, row_filter=chain_filter) - results = await self.read_rows( - query, - operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, - ) - return len(results) > 0 + This property will not change if ``project`` does not, but the + return value is not cached. - async def sample_row_keys( - self, - *, - operation_timeout: float | None = None, - per_request_timeout: float | None = None, - ) -> RowKeySamples: + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_project_path] + :end-before: [END bigtable_api_project_path] + :dedent: 4 + + The project name is of the form + + ``"projects/{project}"`` + + :rtype: str + :returns: Return a fully-qualified project string. """ - Return a set of RowKeySamples that delimit contiguous sections of the table of - approximately equal size + return self.instance_admin_client.common_project_path(self.project) - RowKeySamples output can be used with ReadRowsQuery.shard() to create a sharded query that - can be parallelized across multiple backend nodes read_rows and read_rows_stream - requests will call sample_row_keys internally for this purpose when sharding is enabled + @property + def table_data_client(self): + """Getter for the gRPC stub used for the Table Admin API. - RowKeySamples is simply a type alias for list[tuple[bytes, int]]; a list of - row_keys, along with offset positions in the table + For example: - Returns: - - a set of RowKeySamples the delimit contiguous sections of the table - Raises: - - GoogleAPICallError: if the sample_row_keys request fails + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_table_data_client] + :end-before: [END bigtable_api_table_data_client] + :dedent: 4 + + :rtype: :class:`.bigtable_v2.BigtableClient` + :returns: A BigtableClient object. """ - # prepare timeouts - operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = per_request_timeout or self.default_per_request_timeout - - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout <= 0: - raise ValueError("per_request_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout > operation_timeout: - raise ValueError( - "per_request_timeout must not be greater than operation_timeout" + if self._table_data_client is None: + transport = self._create_gapic_client_channel( + bigtable_v2.BigtableClient, + BigtableGrpcTransport, ) - attempt_timeout_gen = _attempt_timeout_generator( - per_request_timeout, operation_timeout - ) - # prepare retryable - predicate = retries.if_exception_type( - core_exceptions.DeadlineExceeded, - core_exceptions.ServiceUnavailable, - ) - transient_errors = [] - - def on_error_fn(exc): - # add errors to list if retryable - if predicate(exc): - transient_errors.append(exc) - - retry = retries.AsyncRetry( - predicate=predicate, - timeout=operation_timeout, - initial=0.01, - multiplier=2, - maximum=60, - on_error=on_error_fn, - is_stream=False, - ) + klass = _create_gapic_client( + bigtable_v2.BigtableClient, + client_options=self._client_options, + transport=transport, + ) + self._table_data_client = klass(self) + return self._table_data_client - # prepare request - metadata = _make_metadata(self.table_name, self.app_profile_id) + @property + def table_admin_client(self): + """Getter for the gRPC stub used for the Table Admin API. - async def execute_rpc(): - results = await self.client._gapic_client.sample_row_keys( - table_name=self.table_name, - app_profile_id=self.app_profile_id, - timeout=next(attempt_timeout_gen), - metadata=metadata, - ) - return [(s.row_key, s.offset_bytes) async for s in results] + For example: - wrapped_fn = _convert_retry_deadline( - retry(execute_rpc), operation_timeout, transient_errors - ) - return await wrapped_fn() + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_table_admin_client] + :end-before: [END bigtable_api_table_admin_client] + :dedent: 4 - def mutations_batcher( - self, - *, - flush_interval: float | None = 5, - flush_limit_mutation_count: int | None = 1000, - flush_limit_bytes: int = 20 * _MB_SIZE, - flow_control_max_mutation_count: int = 100_000, - flow_control_max_bytes: int = 100 * _MB_SIZE, - batch_operation_timeout: float | None = None, - batch_per_request_timeout: float | None = None, - ) -> MutationsBatcher: - """ - Returns a new mutations batcher instance. - - Can be used to iteratively add mutations that are flushed as a group, - to avoid excess network calls - - Args: - - flush_interval: Automatically flush every flush_interval seconds. If None, - a table default will be used - - flush_limit_mutation_count: Flush immediately after flush_limit_mutation_count - mutations are added across all entries. If None, this limit is ignored. - - flush_limit_bytes: Flush immediately after flush_limit_bytes bytes are added. - - flow_control_max_mutation_count: Maximum number of inflight mutations. - - flow_control_max_bytes: Maximum number of inflight bytes. - - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None, - table default_operation_timeout will be used - - batch_per_request_timeout: timeout for each individual request, in seconds. If None, - table default_per_request_timeout will be used - Returns: - - a MutationsBatcher context manager that can batch requests + :rtype: :class:`.bigtable_admin_pb2.BigtableTableAdmin` + :returns: A BigtableTableAdmin instance. + :raises: :class:`ValueError ` if the current + client is not an admin client or if it has not been + :meth:`start`-ed. """ - return MutationsBatcher( - self, - flush_interval=flush_interval, - flush_limit_mutation_count=flush_limit_mutation_count, - flush_limit_bytes=flush_limit_bytes, - flow_control_max_mutation_count=flow_control_max_mutation_count, - flow_control_max_bytes=flow_control_max_bytes, - batch_operation_timeout=batch_operation_timeout, - batch_per_request_timeout=batch_per_request_timeout, - ) + if self._table_admin_client is None: + if not self._admin: + raise ValueError("Client is not an admin client.") - async def mutate_row( - self, - row_key: str | bytes, - mutations: list[Mutation] | Mutation, - *, - operation_timeout: float | None = 60, - per_request_timeout: float | None = None, - ): - """ - Mutates a row atomically. - - Cells already present in the row are left unchanged unless explicitly changed - by ``mutation``. - - Idempotent operations (i.e, all mutations have an explicit timestamp) will be - retried on server failure. Non-idempotent operations will not. - - Args: - - row_key: the row to apply mutations to - - mutations: the set of mutations to apply to the row - - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will be retried within the budget. - time is only counted while actively waiting on the network. - DeadlineExceeded exception raised after timeout - - per_request_timeout: the time budget for an individual network request, - in seconds. If it takes longer than this time to complete, the request - will be cancelled with a DeadlineExceeded exception, and a retry will be - attempted if within operation_timeout budget - - Raises: - - DeadlineExceeded: raised after operation timeout - will be chained with a RetryExceptionGroup containing all - GoogleAPIError exceptions from any retries that failed - - GoogleAPIError: raised on non-idempotent operations that cannot be - safely retried. - """ - operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = per_request_timeout or self.default_per_request_timeout - - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout <= 0: - raise ValueError("per_request_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout > operation_timeout: - raise ValueError("per_request_timeout must be less than operation_timeout") - - if isinstance(row_key, str): - row_key = row_key.encode("utf-8") - request = {"table_name": self.table_name, "row_key": row_key} - if self.app_profile_id: - request["app_profile_id"] = self.app_profile_id - - if isinstance(mutations, Mutation): - mutations = [mutations] - request["mutations"] = [mutation._to_dict() for mutation in mutations] - - if all(mutation.is_idempotent() for mutation in mutations): - # mutations are all idempotent and safe to retry - predicate = retries.if_exception_type( - core_exceptions.DeadlineExceeded, - core_exceptions.ServiceUnavailable, + transport = self._create_gapic_client_channel( + bigtable_admin_v2.BigtableTableAdminClient, + BigtableTableAdminGrpcTransport, ) - else: - # mutations should not be retried - predicate = retries.if_exception_type() - - transient_errors = [] - - def on_error_fn(exc): - if predicate(exc): - transient_errors.append(exc) - - retry = retries.AsyncRetry( - predicate=predicate, - on_error=on_error_fn, - timeout=operation_timeout, - initial=0.01, - multiplier=2, - maximum=60, - ) - # wrap rpc in retry logic - retry_wrapped = retry(self.client._gapic_client.mutate_row) - # convert RetryErrors from retry wrapper into DeadlineExceeded errors - deadline_wrapped = _convert_retry_deadline( - retry_wrapped, operation_timeout, transient_errors - ) - metadata = _make_metadata(self.table_name, self.app_profile_id) - # trigger rpc - await deadline_wrapped(request, timeout=per_request_timeout, metadata=metadata) + klass = _create_gapic_client( + bigtable_admin_v2.BigtableTableAdminClient, + client_options=self._admin_client_options, + transport=transport, + ) + self._table_admin_client = klass(self) + return self._table_admin_client - async def bulk_mutate_rows( - self, - mutation_entries: list[RowMutationEntry], - *, - operation_timeout: float | None = 60, - per_request_timeout: float | None = None, - ): - """ - Applies mutations for multiple rows in a single batched request. - - Each individual RowMutationEntry is applied atomically, but separate entries - may be applied in arbitrary order (even for entries targetting the same row) - In total, the row_mutations can contain at most 100000 individual mutations - across all entries - - Idempotent entries (i.e., entries with mutations with explicit timestamps) - will be retried on failure. Non-idempotent will not, and will reported in a - raised exception group - - Args: - - mutation_entries: the batches of mutations to apply - Each entry will be applied atomically, but entries will be applied - in arbitrary order - - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will be retried within the budget. - time is only counted while actively waiting on the network. - DeadlineExceeded exception raised after timeout - - per_request_timeout: the time budget for an individual network request, - in seconds. If it takes longer than this time to complete, the request - will be cancelled with a DeadlineExceeded exception, and a retry will - be attempted if within operation_timeout budget - Raises: - - MutationsExceptionGroup if one or more mutations fails - Contains details about any failed entries in .exceptions - """ - operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = per_request_timeout or self.default_per_request_timeout - - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout <= 0: - raise ValueError("per_request_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout > operation_timeout: - raise ValueError("per_request_timeout must be less than operation_timeout") - - operation = _MutateRowsOperation( - self.client._gapic_client, - self, - mutation_entries, - operation_timeout, - per_request_timeout, - ) - await operation.start() + @property + def instance_admin_client(self): + """Getter for the gRPC stub used for the Table Admin API. - async def check_and_mutate_row( - self, - row_key: str | bytes, - predicate: RowFilter | dict[str, Any] | None, - *, - true_case_mutations: Mutation | list[Mutation] | None = None, - false_case_mutations: Mutation | list[Mutation] | None = None, - operation_timeout: int | float | None = 20, - ) -> bool: - """ - Mutates a row atomically based on the output of a predicate filter - - Non-idempotent operation: will not be retried - - Args: - - row_key: the key of the row to mutate - - predicate: the filter to be applied to the contents of the specified row. - Depending on whether or not any results are yielded, - either true_case_mutations or false_case_mutations will be executed. - If None, checks that the row contains any values at all. - - true_case_mutations: - Changes to be atomically applied to the specified row if - predicate yields at least one cell when - applied to row_key. Entries are applied in order, - meaning that earlier mutations can be masked by later - ones. Must contain at least one entry if - false_case_mutations is empty, and at most 100000. - - false_case_mutations: - Changes to be atomically applied to the specified row if - predicate_filter does not yield any cells when - applied to row_key. Entries are applied in order, - meaning that earlier mutations can be masked by later - ones. Must contain at least one entry if - `true_case_mutations is empty, and at most 100000. - - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will not be retried. - Returns: - - bool indicating whether the predicate was true or false - Raises: - - GoogleAPIError exceptions from grpc call + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_instance_admin_client] + :end-before: [END bigtable_api_instance_admin_client] + :dedent: 4 + + :rtype: :class:`.bigtable_admin_pb2.BigtableInstanceAdmin` + :returns: A BigtableInstanceAdmin instance. + :raises: :class:`ValueError ` if the current + client is not an admin client or if it has not been + :meth:`start`-ed. """ - operation_timeout = operation_timeout or self.default_operation_timeout - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key - if true_case_mutations is not None and not isinstance( - true_case_mutations, list - ): - true_case_mutations = [true_case_mutations] - true_case_dict = [m._to_dict() for m in true_case_mutations or []] - if false_case_mutations is not None and not isinstance( - false_case_mutations, list - ): - false_case_mutations = [false_case_mutations] - false_case_dict = [m._to_dict() for m in false_case_mutations or []] - if predicate is not None and not isinstance(predicate, dict): - predicate = predicate.to_dict() - metadata = _make_metadata(self.table_name, self.app_profile_id) - result = await self.client._gapic_client.check_and_mutate_row( - request={ - "predicate_filter": predicate, - "true_mutations": true_case_dict, - "false_mutations": false_case_dict, - "table_name": self.table_name, - "row_key": row_key, - "app_profile_id": self.app_profile_id, - }, - metadata=metadata, - timeout=operation_timeout, + if self._instance_admin_client is None: + if not self._admin: + raise ValueError("Client is not an admin client.") + + transport = self._create_gapic_client_channel( + bigtable_admin_v2.BigtableInstanceAdminClient, + BigtableInstanceAdminGrpcTransport, + ) + klass = _create_gapic_client( + bigtable_admin_v2.BigtableInstanceAdminClient, + client_options=self._admin_client_options, + transport=transport, + ) + self._instance_admin_client = klass(self) + return self._instance_admin_client + + def instance(self, instance_id, display_name=None, instance_type=None, labels=None): + """Factory to create a instance associated with this client. + + For example: + + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_create_prod_instance] + :end-before: [END bigtable_api_create_prod_instance] + :dedent: 4 + + :type instance_id: str + :param instance_id: The ID of the instance. + + :type display_name: str + :param display_name: (Optional) The display name for the instance in + the Cloud Console UI. (Must be between 4 and 30 + characters.) If this value is not set in the + constructor, will fall back to the instance ID. + + :type instance_type: int + :param instance_type: (Optional) The type of the instance. + Possible values are represented + by the following constants: + :data:`google.cloud.bigtable.instance.InstanceType.PRODUCTION`. + :data:`google.cloud.bigtable.instance.InstanceType.DEVELOPMENT`, + Defaults to + :data:`google.cloud.bigtable.instance.InstanceType.UNSPECIFIED`. + + :type labels: dict + :param labels: (Optional) Labels are a flexible and lightweight + mechanism for organizing cloud resources into groups + that reflect a customer's organizational needs and + deployment strategies. They can be used to filter + resources and aggregate metrics. Label keys must be + between 1 and 63 characters long. Maximum 64 labels can + be associated with a given resource. Label values must + be between 0 and 63 characters long. Keys and values + must both be under 128 bytes. + + :rtype: :class:`~google.cloud.bigtable.instance.Instance` + :returns: an instance owned by this client. + """ + return Instance( + instance_id, + self, + display_name=display_name, + instance_type=instance_type, + labels=labels, ) - return result.predicate_matched - async def read_modify_write_row( - self, - row_key: str | bytes, - rules: ReadModifyWriteRule | list[ReadModifyWriteRule], - *, - operation_timeout: int | float | None = 20, - ) -> Row: - """ - Reads and modifies a row atomically according to input ReadModifyWriteRules, - and returns the contents of all modified cells + def list_instances(self): + """List instances owned by the project. - The new value for the timestamp is the greater of the existing timestamp or - the current server time. + For example: - Non-idempotent operation: will not be retried + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_list_instances] + :end-before: [END bigtable_api_list_instances] + :dedent: 4 - Args: - - row_key: the key of the row to apply read/modify/write rules to - - rules: A rule or set of rules to apply to the row. - Rules are applied in order, meaning that earlier rules will affect the - results of later ones. - - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will not be retried. - Returns: - - Row: containing cell data that was modified as part of the - operation - Raises: - - GoogleAPIError exceptions from grpc call + :rtype: tuple + :returns: + (instances, failed_locations), where 'instances' is list of + :class:`google.cloud.bigtable.instance.Instance`, and + 'failed_locations' is a list of locations which could not + be resolved. """ - operation_timeout = operation_timeout or self.default_operation_timeout - row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if rules is not None and not isinstance(rules, list): - rules = [rules] - if not rules: - raise ValueError("rules must contain at least one item") - # concert to dict representation - rules_dict = [rule._to_dict() for rule in rules] - metadata = _make_metadata(self.table_name, self.app_profile_id) - result = await self.client._gapic_client.read_modify_write_row( - request={ - "rules": rules_dict, - "table_name": self.table_name, - "row_key": row_key, - "app_profile_id": self.app_profile_id, - }, - metadata=metadata, - timeout=operation_timeout, + resp = self.instance_admin_client.list_instances( + request={"parent": self.project_path} ) - # construct Row from result - return Row._from_pb(result.row) + instances = [Instance.from_pb(instance, self) for instance in resp.instances] + return instances, resp.failed_locations - async def close(self): - """ - Called to close the Table instance and release any resources held by it. - """ - self._register_instance_task.cancel() - await self.client._remove_instance_registration(self.instance_id, self) + def list_clusters(self): + """List the clusters in the project. - async def __aenter__(self): - """ - Implement async context manager protocol + For example: - Ensure registration task has time to run, so that - grpc channels will be warmed for the specified instance - """ - await self._register_instance_task - return self + .. literalinclude:: snippets.py + :start-after: [START bigtable_api_list_clusters_in_project] + :end-before: [END bigtable_api_list_clusters_in_project] + :dedent: 4 - async def __aexit__(self, exc_type, exc_val, exc_tb): + :rtype: tuple + :returns: + (clusters, failed_locations), where 'clusters' is list of + :class:`google.cloud.bigtable.instance.Cluster`, and + 'failed_locations' is a list of strings representing + locations which could not be resolved. """ - Implement async context manager protocol - - Unregister this instance with the client, so that - grpc channels will no longer be warmed - """ - await self.close() + resp = self.instance_admin_client.list_clusters( + request={ + "parent": self.instance_admin_client.instance_path(self.project, "-") + } + ) + clusters = [] + instances = {} + for cluster in resp.clusters: + match_cluster_name = _CLUSTER_NAME_RE.match(cluster.name) + instance_id = match_cluster_name.group("instance") + if instance_id not in instances: + instances[instance_id] = self.instance(instance_id) + clusters.append(Cluster.from_pb(cluster, instances[instance_id])) + return clusters, resp.failed_locations diff --git a/google/cloud/bigtable/deprecated/cluster.py b/google/cloud/bigtable/cluster.py similarity index 95% rename from google/cloud/bigtable/deprecated/cluster.py rename to google/cloud/bigtable/cluster.py index b60d3503c..11fb5492d 100644 --- a/google/cloud/bigtable/deprecated/cluster.py +++ b/google/cloud/bigtable/cluster.py @@ -42,7 +42,7 @@ class Cluster(object): :type cluster_id: str :param cluster_id: The ID of the cluster. - :type instance: :class:`~google.cloud.bigtable.deprecated.instance.Instance` + :type instance: :class:`~google.cloud.bigtable.instance.Instance` :param instance: The instance where the cluster resides. :type location_id: str @@ -62,10 +62,10 @@ class Cluster(object): :param default_storage_type: (Optional) The type of storage Possible values are represented by the following constants: - :data:`google.cloud.bigtable.deprecated.enums.StorageType.SSD`. - :data:`google.cloud.bigtable.deprecated.enums.StorageType.HDD`, + :data:`google.cloud.bigtable.enums.StorageType.SSD`. + :data:`google.cloud.bigtable.enums.StorageType.HDD`, Defaults to - :data:`google.cloud.bigtable.deprecated.enums.StorageType.UNSPECIFIED`. + :data:`google.cloud.bigtable.enums.StorageType.UNSPECIFIED`. :type kms_key_name: str :param kms_key_name: (Optional, Creation Only) The name of the KMS customer managed @@ -84,11 +84,11 @@ class Cluster(object): :param _state: (`OutputOnly`) The current state of the cluster. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.NOT_KNOWN`. - :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.READY`. - :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.CREATING`. - :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.RESIZING`. - :data:`google.cloud.bigtable.deprecated.enums.Cluster.State.DISABLED`. + :data:`google.cloud.bigtable.enums.Cluster.State.NOT_KNOWN`. + :data:`google.cloud.bigtable.enums.Cluster.State.READY`. + :data:`google.cloud.bigtable.enums.Cluster.State.CREATING`. + :data:`google.cloud.bigtable.enums.Cluster.State.RESIZING`. + :data:`google.cloud.bigtable.enums.Cluster.State.DISABLED`. :type min_serve_nodes: int :param min_serve_nodes: (Optional) The minimum number of nodes to be set in the cluster for autoscaling. @@ -150,7 +150,7 @@ def from_pb(cls, cluster_pb, instance): :type cluster_pb: :class:`instance.Cluster` :param cluster_pb: An instance protobuf object. - :type instance: :class:`google.cloud.bigtable.deprecated.instance.Instance` + :type instance: :class:`google.cloud.bigtable.instance.Instance` :param instance: The instance that owns the cluster. :rtype: :class:`Cluster` @@ -236,7 +236,7 @@ def name(self): @property def state(self): - """google.cloud.bigtable.deprecated.enums.Cluster.State: state of cluster. + """google.cloud.bigtable.enums.Cluster.State: state of cluster. For example: diff --git a/google/cloud/bigtable/deprecated/column_family.py b/google/cloud/bigtable/column_family.py similarity index 99% rename from google/cloud/bigtable/deprecated/column_family.py rename to google/cloud/bigtable/column_family.py index 3d4c1a642..80232958d 100644 --- a/google/cloud/bigtable/deprecated/column_family.py +++ b/google/cloud/bigtable/column_family.py @@ -195,7 +195,7 @@ class ColumnFamily(object): :param column_family_id: The ID of the column family. Must be of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - :type table: :class:`Table ` + :type table: :class:`Table ` :param table: The table that owns the column family. :type gc_rule: :class:`GarbageCollectionRule` diff --git a/google/cloud/bigtable/data/__init__.py b/google/cloud/bigtable/data/__init__.py new file mode 100644 index 000000000..c68e78c6f --- /dev/null +++ b/google/cloud/bigtable/data/__init__.py @@ -0,0 +1,78 @@ +# -*- coding: utf-8 -*- +# Copyright 2023 Google LLC +# +# Licensed 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. +# + +from typing import List, Tuple + +from google.cloud.bigtable import gapic_version as package_version + +from google.cloud.bigtable.data._async.client import BigtableDataClientAsync +from google.cloud.bigtable.data._async.client import TableAsync +from google.cloud.bigtable.data._async._read_rows import ReadRowsAsyncIterator +from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync + +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery +from google.cloud.bigtable.data.read_rows_query import RowRange +from google.cloud.bigtable.data.row import Row +from google.cloud.bigtable.data.row import Cell + +from google.cloud.bigtable.data.mutations import Mutation +from google.cloud.bigtable.data.mutations import RowMutationEntry +from google.cloud.bigtable.data.mutations import SetCell +from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn +from google.cloud.bigtable.data.mutations import DeleteAllFromFamily +from google.cloud.bigtable.data.mutations import DeleteAllFromRow + +from google.cloud.bigtable.data.exceptions import IdleTimeout +from google.cloud.bigtable.data.exceptions import InvalidChunk +from google.cloud.bigtable.data.exceptions import FailedMutationEntryError +from google.cloud.bigtable.data.exceptions import FailedQueryShardError + +from google.cloud.bigtable.data.exceptions import RetryExceptionGroup +from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup +from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + +# Type alias for the output of sample_keys +RowKeySamples = List[Tuple[bytes, int]] +# type alias for the output of query.shard() +ShardedQuery = List[ReadRowsQuery] + +__version__: str = package_version.__version__ + +__all__ = ( + "BigtableDataClientAsync", + "TableAsync", + "RowKeySamples", + "ReadRowsQuery", + "RowRange", + "MutationsBatcherAsync", + "Mutation", + "RowMutationEntry", + "SetCell", + "DeleteRangeFromColumn", + "DeleteAllFromFamily", + "DeleteAllFromRow", + "Row", + "Cell", + "ReadRowsAsyncIterator", + "IdleTimeout", + "InvalidChunk", + "FailedMutationEntryError", + "FailedQueryShardError", + "RetryExceptionGroup", + "MutationsExceptionGroup", + "ShardedReadRowsExceptionGroup", + "ShardedQuery", +) diff --git a/google/cloud/bigtable/data/_async/__init__.py b/google/cloud/bigtable/data/_async/__init__.py new file mode 100644 index 000000000..1e92e58dc --- /dev/null +++ b/google/cloud/bigtable/data/_async/__init__.py @@ -0,0 +1,26 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +from google.cloud.bigtable.data._async.client import BigtableDataClientAsync +from google.cloud.bigtable.data._async.client import TableAsync +from google.cloud.bigtable.data._async._read_rows import ReadRowsAsyncIterator +from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync + + +__all__ = [ + "BigtableDataClientAsync", + "TableAsync", + "ReadRowsAsyncIterator", + "MutationsBatcherAsync", +] diff --git a/google/cloud/bigtable/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py similarity index 91% rename from google/cloud/bigtable/_mutate_rows.py rename to google/cloud/bigtable/data/_async/_mutate_rows.py index e34ebaeb6..ac491adaf 100644 --- a/google/cloud/bigtable/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -19,31 +19,23 @@ from google.api_core import exceptions as core_exceptions from google.api_core import retry_async as retries -import google.cloud.bigtable.exceptions as bt_exceptions -from google.cloud.bigtable._helpers import _make_metadata -from google.cloud.bigtable._helpers import _convert_retry_deadline -from google.cloud.bigtable._helpers import _attempt_timeout_generator +import google.cloud.bigtable.data.exceptions as bt_exceptions +from google.cloud.bigtable.data._helpers import _make_metadata +from google.cloud.bigtable.data._helpers import _convert_retry_deadline +from google.cloud.bigtable.data._helpers import _attempt_timeout_generator + +# mutate_rows requests are limited to this number of mutations +from google.cloud.bigtable.data.mutations import MUTATE_ROWS_REQUEST_MUTATION_LIMIT if TYPE_CHECKING: from google.cloud.bigtable_v2.services.bigtable.async_client import ( BigtableAsyncClient, ) - from google.cloud.bigtable.client import Table - from google.cloud.bigtable.mutations import RowMutationEntry - -# mutate_rows requests are limited to this value -MUTATE_ROWS_REQUEST_MUTATION_LIMIT = 100_000 - - -class _MutateRowsIncomplete(RuntimeError): - """ - Exception raised when a mutate_rows call has unfinished work. - """ - - pass + from google.cloud.bigtable.data.mutations import RowMutationEntry + from google.cloud.bigtable.data._async.client import TableAsync -class _MutateRowsOperation: +class _MutateRowsOperationAsync: """ MutateRowsOperation manages the logic of sending a set of row mutations, and retrying on failed entries. It manages this using the _run_attempt @@ -57,7 +49,7 @@ class _MutateRowsOperation: def __init__( self, gapic_client: "BigtableAsyncClient", - table: "Table", + table: "TableAsync", mutation_entries: list["RowMutationEntry"], operation_timeout: float, per_request_timeout: float | None, @@ -93,7 +85,7 @@ def __init__( core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable, # Entry level errors - _MutateRowsIncomplete, + bt_exceptions._MutateRowsIncomplete, ) # build retryable operation retry = retries.AsyncRetry( @@ -199,7 +191,7 @@ async def _run_attempt(self): # check if attempt succeeded, or needs to be retried if self.remaining_indices: # unfinished work; raise exception to trigger retry - raise _MutateRowsIncomplete + raise bt_exceptions._MutateRowsIncomplete def _handle_entry_error(self, idx: int, exc: Exception): """ diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py new file mode 100644 index 000000000..910a01c4c --- /dev/null +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -0,0 +1,403 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from typing import ( + List, + Any, + AsyncIterable, + AsyncIterator, + AsyncGenerator, + Iterator, + Callable, + Awaitable, +) +import sys +import time +import asyncio +from functools import partial +from grpc.aio import RpcContext + +from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse +from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient +from google.cloud.bigtable.data.row import Row, _LastScannedRow +from google.cloud.bigtable.data.exceptions import InvalidChunk +from google.cloud.bigtable.data.exceptions import _RowSetComplete +from google.cloud.bigtable.data.exceptions import IdleTimeout +from google.cloud.bigtable.data._read_rows_state_machine import _StateMachine +from google.api_core import retry_async as retries +from google.api_core import exceptions as core_exceptions +from google.cloud.bigtable.data._helpers import _make_metadata +from google.cloud.bigtable.data._helpers import _attempt_timeout_generator +from google.cloud.bigtable.data._helpers import _convert_retry_deadline + + +class _ReadRowsOperationAsync(AsyncIterable[Row]): + """ + ReadRowsOperation handles the logic of merging chunks from a ReadRowsResponse stream + into a stream of Row objects. + + ReadRowsOperation.merge_row_response_stream takes in a stream of ReadRowsResponse + and turns them into a stream of Row objects using an internal + StateMachine. + + ReadRowsOperation(request, client) handles row merging logic end-to-end, including + performing retries on stream errors. + """ + + def __init__( + self, + request: dict[str, Any], + client: BigtableAsyncClient, + *, + operation_timeout: float = 600.0, + per_request_timeout: float | None = None, + ): + """ + Args: + - request: the request dict to send to the Bigtable API + - client: the Bigtable client to use to make the request + - operation_timeout: the timeout to use for the entire operation, in seconds + - per_request_timeout: the timeout to use when waiting for each individual grpc request, in seconds + If not specified, defaults to operation_timeout + """ + self._last_emitted_row_key: bytes | None = None + self._emit_count = 0 + self._request = request + self.operation_timeout = operation_timeout + # use generator to lower per-attempt timeout as we approach operation_timeout deadline + attempt_timeout_gen = _attempt_timeout_generator( + per_request_timeout, operation_timeout + ) + row_limit = request.get("rows_limit", 0) + # lock in paramters for retryable wrapper + self._partial_retryable = partial( + self._read_rows_retryable_attempt, + client.read_rows, + attempt_timeout_gen, + row_limit, + ) + predicate = retries.if_exception_type( + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + core_exceptions.Aborted, + ) + + def on_error_fn(exc): + if predicate(exc): + self.transient_errors.append(exc) + + retry = retries.AsyncRetry( + predicate=predicate, + timeout=self.operation_timeout, + initial=0.01, + multiplier=2, + maximum=60, + on_error=on_error_fn, + is_stream=True, + ) + self._stream: AsyncGenerator[Row, None] | None = retry( + self._partial_retryable + )() + # contains the list of errors that were retried + self.transient_errors: List[Exception] = [] + + def __aiter__(self) -> AsyncIterator[Row]: + """Implements the AsyncIterable interface""" + return self + + async def __anext__(self) -> Row: + """Implements the AsyncIterator interface""" + if self._stream is not None: + return await self._stream.__anext__() + else: + raise asyncio.InvalidStateError("stream is closed") + + async def aclose(self): + """Close the stream and release resources""" + if self._stream is not None: + await self._stream.aclose() + self._stream = None + self._emitted_seen_row_key = None + + async def _read_rows_retryable_attempt( + self, + gapic_fn: Callable[..., Awaitable[AsyncIterable[ReadRowsResponse]]], + timeout_generator: Iterator[float], + total_row_limit: int, + ) -> AsyncGenerator[Row, None]: + """ + Retryable wrapper for merge_rows. This function is called each time + a retry is attempted. + + Some fresh state is created on each retry: + - grpc network stream + - state machine to hold merge chunks received from stream + Some state is shared between retries: + - _last_emitted_row_key is used to ensure that + duplicate rows are not emitted + - request is stored and (potentially) modified on each retry + """ + if self._last_emitted_row_key is not None: + # if this is a retry, try to trim down the request to avoid ones we've already processed + try: + self._request["rows"] = _ReadRowsOperationAsync._revise_request_rowset( + row_set=self._request.get("rows", None), + last_seen_row_key=self._last_emitted_row_key, + ) + except _RowSetComplete: + # if there are no rows left to process, we're done + # This is not expected to happen often, but could occur if + # a retry is triggered quickly after the last row is emitted + return + # revise next request's row limit based on number emitted + if total_row_limit: + new_limit = total_row_limit - self._emit_count + if new_limit == 0: + # we have hit the row limit, so we're done + return + elif new_limit < 0: + raise RuntimeError("unexpected state: emit count exceeds row limit") + else: + self._request["rows_limit"] = new_limit + metadata = _make_metadata( + self._request.get("table_name", None), + self._request.get("app_profile_id", None), + ) + new_gapic_stream: RpcContext = await gapic_fn( + self._request, + timeout=next(timeout_generator), + metadata=metadata, + ) + try: + state_machine = _StateMachine() + stream = _ReadRowsOperationAsync.merge_row_response_stream( + new_gapic_stream, state_machine + ) + # run until we get a timeout or the stream is exhausted + async for new_item in stream: + if ( + self._last_emitted_row_key is not None + and new_item.row_key <= self._last_emitted_row_key + ): + raise InvalidChunk("Last emitted row key out of order") + # don't yeild _LastScannedRow markers; they + # should only update last_seen_row_key + if not isinstance(new_item, _LastScannedRow): + yield new_item + self._emit_count += 1 + self._last_emitted_row_key = new_item.row_key + if total_row_limit and self._emit_count >= total_row_limit: + return + except (Exception, GeneratorExit) as exc: + # ensure grpc stream is closed + new_gapic_stream.cancel() + raise exc + + @staticmethod + def _revise_request_rowset( + row_set: dict[str, Any] | None, + last_seen_row_key: bytes, + ) -> dict[str, Any]: + """ + Revise the rows in the request to avoid ones we've already processed. + + Args: + - row_set: the row set from the request + - last_seen_row_key: the last row key encountered + Raises: + - _RowSetComplete: if there are no rows left to process after the revision + """ + # if user is doing a whole table scan, start a new one with the last seen key + if row_set is None or ( + len(row_set.get("row_ranges", [])) == 0 + and len(row_set.get("row_keys", [])) == 0 + ): + last_seen = last_seen_row_key + return { + "row_keys": [], + "row_ranges": [{"start_key_open": last_seen}], + } + # remove seen keys from user-specific key list + row_keys: list[bytes] = row_set.get("row_keys", []) + adjusted_keys = [k for k in row_keys if k > last_seen_row_key] + # adjust ranges to ignore keys before last seen + row_ranges: list[dict[str, Any]] = row_set.get("row_ranges", []) + adjusted_ranges = [] + for row_range in row_ranges: + end_key = row_range.get("end_key_closed", None) or row_range.get( + "end_key_open", None + ) + if end_key is None or end_key > last_seen_row_key: + # end range is after last seen key + new_range = row_range.copy() + start_key = row_range.get("start_key_closed", None) or row_range.get( + "start_key_open", None + ) + if start_key is None or start_key <= last_seen_row_key: + # replace start key with last seen + new_range["start_key_open"] = last_seen_row_key + new_range.pop("start_key_closed", None) + adjusted_ranges.append(new_range) + if len(adjusted_keys) == 0 and len(adjusted_ranges) == 0: + # if the query is empty after revision, raise an exception + # this will avoid an unwanted full table scan + raise _RowSetComplete() + return {"row_keys": adjusted_keys, "row_ranges": adjusted_ranges} + + @staticmethod + async def merge_row_response_stream( + response_generator: AsyncIterable[ReadRowsResponse], + state_machine: _StateMachine, + ) -> AsyncGenerator[Row, None]: + """ + Consume chunks from a ReadRowsResponse stream into a set of Rows + + Args: + - response_generator: AsyncIterable of ReadRowsResponse objects. Typically + this is a stream of chunks from the Bigtable API + Returns: + - AsyncGenerator of Rows + Raises: + - InvalidChunk: if the chunk stream is invalid + """ + async for row_response in response_generator: + # unwrap protoplus object for increased performance + response_pb = row_response._pb + last_scanned = response_pb.last_scanned_row_key + # if the server sends a scan heartbeat, notify the state machine. + if last_scanned: + yield state_machine.handle_last_scanned_row(last_scanned) + # process new chunks through the state machine. + for chunk in response_pb.chunks: + complete_row = state_machine.handle_chunk(chunk) + if complete_row is not None: + yield complete_row + # TODO: handle request stats + if not state_machine.is_terminal_state(): + # read rows is complete, but there's still data in the merger + raise InvalidChunk("read_rows completed with partial state remaining") + + +class ReadRowsAsyncIterator(AsyncIterable[Row]): + """ + Async iterator for ReadRows responses. + + Supports the AsyncIterator protocol for use in async for loops, + along with: + - `aclose` for closing the underlying stream + - `active` for checking if the iterator is still active + - an internal idle timer for closing the stream after a period of inactivity + """ + + def __init__(self, merger: _ReadRowsOperationAsync): + self._merger: _ReadRowsOperationAsync = merger + self._error: Exception | None = None + self._last_interaction_time = time.monotonic() + self._idle_timeout_task: asyncio.Task[None] | None = None + # wrap merger with a wrapper that properly formats exceptions + self._next_fn = _convert_retry_deadline( + self._merger.__anext__, + self._merger.operation_timeout, + self._merger.transient_errors, + ) + + async def _start_idle_timer(self, idle_timeout: float): + """ + Start a coroutine that will cancel a stream if no interaction + with the iterator occurs for the specified number of seconds. + + Subsequent access to the iterator will raise an IdleTimeout exception. + + Args: + - idle_timeout: number of seconds of inactivity before cancelling the stream + """ + self._last_interaction_time = time.monotonic() + if self._idle_timeout_task is not None: + self._idle_timeout_task.cancel() + self._idle_timeout_task = asyncio.create_task( + self._idle_timeout_coroutine(idle_timeout) + ) + if sys.version_info >= (3, 8): + self._idle_timeout_task.name = f"{self.__class__.__name__}.idle_timeout" + + @property + def active(self): + """ + Returns True if the iterator is still active and has not been closed + """ + return self._error is None + + async def _idle_timeout_coroutine(self, idle_timeout: float): + """ + Coroutine that will cancel a stream if no interaction with the iterator + in the last `idle_timeout` seconds. + """ + while self.active: + next_timeout = self._last_interaction_time + idle_timeout + await asyncio.sleep(next_timeout - time.monotonic()) + if ( + self._last_interaction_time + idle_timeout < time.monotonic() + and self.active + ): + # idle timeout has expired + await self._finish_with_error( + IdleTimeout( + ( + "Timed out waiting for next Row to be consumed. " + f"(idle_timeout={idle_timeout:0.1f}s)" + ) + ) + ) + + def __aiter__(self): + """Implement the async iterator protocol.""" + return self + + async def __anext__(self) -> Row: + """ + Implement the async iterator potocol. + + Return the next item in the stream if active, or + raise an exception if the stream has been closed. + """ + if self._error is not None: + raise self._error + try: + self._last_interaction_time = time.monotonic() + return await self._next_fn() + except Exception as e: + await self._finish_with_error(e) + raise e + + async def _finish_with_error(self, e: Exception): + """ + Helper function to close the stream and clean up resources + after an error has occurred. + """ + if self.active: + await self._merger.aclose() + self._error = e + if self._idle_timeout_task is not None: + self._idle_timeout_task.cancel() + self._idle_timeout_task = None + + async def aclose(self): + """ + Support closing the stream with an explicit call to aclose() + """ + await self._finish_with_error( + StopAsyncIteration(f"{self.__class__.__name__} closed") + ) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py new file mode 100644 index 000000000..3a5831799 --- /dev/null +++ b/google/cloud/bigtable/data/_async/client.py @@ -0,0 +1,1091 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# + +from __future__ import annotations + +from typing import ( + cast, + Any, + Optional, + Set, + TYPE_CHECKING, +) + +import asyncio +import grpc +import time +import warnings +import sys +import random + +from collections import namedtuple + +from google.cloud.bigtable_v2.services.bigtable.client import BigtableClientMeta +from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient +from google.cloud.bigtable_v2.services.bigtable.async_client import DEFAULT_CLIENT_INFO +from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( + PooledBigtableGrpcAsyncIOTransport, +) +from google.cloud.bigtable_v2.types.bigtable import PingAndWarmRequest +from google.cloud.client import ClientWithProject +from google.api_core.exceptions import GoogleAPICallError +from google.api_core import retry_async as retries +from google.api_core import exceptions as core_exceptions +from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync +from google.cloud.bigtable.data._async._read_rows import ReadRowsAsyncIterator + +import google.auth.credentials +import google.auth._default +from google.api_core import client_options as client_options_lib +from google.cloud.bigtable.data.row import Row +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery +from google.cloud.bigtable.data.exceptions import FailedQueryShardError +from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + +from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry +from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync +from google.cloud.bigtable.data._helpers import _make_metadata +from google.cloud.bigtable.data._helpers import _convert_retry_deadline +from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync +from google.cloud.bigtable.data._async.mutations_batcher import _MB_SIZE +from google.cloud.bigtable.data._helpers import _attempt_timeout_generator + +from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule +from google.cloud.bigtable.data.row_filters import RowFilter +from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter +from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter +from google.cloud.bigtable.data.row_filters import RowFilterChain + +if TYPE_CHECKING: + from google.cloud.bigtable.data import RowKeySamples + from google.cloud.bigtable.data import ShardedQuery + +# used by read_rows_sharded to limit how many requests are attempted in parallel +CONCURRENCY_LIMIT = 10 + +# used to register instance data with the client for channel warming +_WarmedInstanceKey = namedtuple( + "_WarmedInstanceKey", ["instance_name", "table_name", "app_profile_id"] +) + + +class BigtableDataClientAsync(ClientWithProject): + def __init__( + self, + *, + project: str | None = None, + pool_size: int = 3, + credentials: google.auth.credentials.Credentials | None = None, + client_options: dict[str, Any] + | "google.api_core.client_options.ClientOptions" + | None = None, + ): + """ + Create a client instance for the Bigtable Data API + + Client should be created within an async context (running event loop) + + Args: + project: the project which the client acts on behalf of. + If not passed, falls back to the default inferred + from the environment. + pool_size: The number of grpc channels to maintain + in the internal channel pool. + credentials: + Thehe OAuth2 Credentials to use for this + client. If not passed (and if no ``_http`` object is + passed), falls back to the default inferred from the + environment. + client_options (Optional[Union[dict, google.api_core.client_options.ClientOptions]]): + Client options used to set user options + on the client. API Endpoint should be set through client_options. + Raises: + - RuntimeError if called outside of an async context (no running event loop) + - ValueError if pool_size is less than 1 + """ + # set up transport in registry + transport_str = f"pooled_grpc_asyncio_{pool_size}" + transport = PooledBigtableGrpcAsyncIOTransport.with_fixed_size(pool_size) + BigtableClientMeta._transport_registry[transport_str] = transport + # set up client info headers for veneer library + client_info = DEFAULT_CLIENT_INFO + client_info.client_library_version = client_info.gapic_version + # parse client options + if type(client_options) is dict: + client_options = client_options_lib.from_dict(client_options) + client_options = cast( + Optional[client_options_lib.ClientOptions], client_options + ) + # initialize client + ClientWithProject.__init__( + self, + credentials=credentials, + project=project, + client_options=client_options, + ) + self._gapic_client = BigtableAsyncClient( + transport=transport_str, + credentials=credentials, + client_options=client_options, + client_info=client_info, + ) + self.transport = cast( + PooledBigtableGrpcAsyncIOTransport, self._gapic_client.transport + ) + # keep track of active instances to for warmup on channel refresh + self._active_instances: Set[_WarmedInstanceKey] = set() + # keep track of table objects associated with each instance + # only remove instance from _active_instances when all associated tables remove it + self._instance_owners: dict[_WarmedInstanceKey, Set[int]] = {} + # attempt to start background tasks + self._channel_init_time = time.monotonic() + self._channel_refresh_tasks: list[asyncio.Task[None]] = [] + try: + self.start_background_channel_refresh() + except RuntimeError: + warnings.warn( + f"{self.__class__.__name__} should be started in an " + "asyncio event loop. Channel refresh will not be started", + RuntimeWarning, + stacklevel=2, + ) + + def start_background_channel_refresh(self) -> None: + """ + Starts a background task to ping and warm each channel in the pool + Raises: + - RuntimeError if not called in an asyncio event loop + """ + if not self._channel_refresh_tasks: + # raise RuntimeError if there is no event loop + asyncio.get_running_loop() + for channel_idx in range(self.transport.pool_size): + refresh_task = asyncio.create_task(self._manage_channel(channel_idx)) + if sys.version_info >= (3, 8): + # task names supported in Python 3.8+ + refresh_task.set_name( + f"{self.__class__.__name__} channel refresh {channel_idx}" + ) + self._channel_refresh_tasks.append(refresh_task) + + async def close(self, timeout: float = 2.0): + """ + Cancel all background tasks + """ + for task in self._channel_refresh_tasks: + task.cancel() + group = asyncio.gather(*self._channel_refresh_tasks, return_exceptions=True) + await asyncio.wait_for(group, timeout=timeout) + await self.transport.close() + self._channel_refresh_tasks = [] + + async def _ping_and_warm_instances( + self, channel: grpc.aio.Channel, instance_key: _WarmedInstanceKey | None = None + ) -> list[GoogleAPICallError | None]: + """ + Prepares the backend for requests on a channel + + Pings each Bigtable instance registered in `_active_instances` on the client + + Args: + - channel: grpc channel to warm + - instance_key: if provided, only warm the instance associated with the key + Returns: + - sequence of results or exceptions from the ping requests + """ + instance_list = ( + [instance_key] if instance_key is not None else self._active_instances + ) + ping_rpc = channel.unary_unary( + "/google.bigtable.v2.Bigtable/PingAndWarm", + request_serializer=PingAndWarmRequest.serialize, + ) + # prepare list of coroutines to run + tasks = [ + ping_rpc( + request={"name": instance_name, "app_profile_id": app_profile_id}, + metadata=[ + ( + "x-goog-request-params", + f"name={instance_name}&app_profile_id={app_profile_id}", + ) + ], + wait_for_ready=True, + ) + for (instance_name, table_name, app_profile_id) in instance_list + ] + # execute coroutines in parallel + result_list = await asyncio.gather(*tasks, return_exceptions=True) + # return None in place of empty successful responses + return [r or None for r in result_list] + + async def _manage_channel( + self, + channel_idx: int, + refresh_interval_min: float = 60 * 35, + refresh_interval_max: float = 60 * 45, + grace_period: float = 60 * 10, + ) -> None: + """ + Background coroutine that periodically refreshes and warms a grpc channel + + The backend will automatically close channels after 60 minutes, so + `refresh_interval` + `grace_period` should be < 60 minutes + + Runs continuously until the client is closed + + Args: + channel_idx: index of the channel in the transport's channel pool + refresh_interval_min: minimum interval before initiating refresh + process in seconds. Actual interval will be a random value + between `refresh_interval_min` and `refresh_interval_max` + refresh_interval_max: maximum interval before initiating refresh + process in seconds. Actual interval will be a random value + between `refresh_interval_min` and `refresh_interval_max` + grace_period: time to allow previous channel to serve existing + requests before closing, in seconds + """ + first_refresh = self._channel_init_time + random.uniform( + refresh_interval_min, refresh_interval_max + ) + next_sleep = max(first_refresh - time.monotonic(), 0) + if next_sleep > 0: + # warm the current channel immediately + channel = self.transport.channels[channel_idx] + await self._ping_and_warm_instances(channel) + # continuously refresh the channel every `refresh_interval` seconds + while True: + await asyncio.sleep(next_sleep) + # prepare new channel for use + new_channel = self.transport.grpc_channel._create_channel() + await self._ping_and_warm_instances(new_channel) + # cycle channel out of use, with long grace window before closure + start_timestamp = time.time() + await self.transport.replace_channel( + channel_idx, grace=grace_period, swap_sleep=10, new_channel=new_channel + ) + # subtract the time spent waiting for the channel to be replaced + next_refresh = random.uniform(refresh_interval_min, refresh_interval_max) + next_sleep = next_refresh - (time.time() - start_timestamp) + + async def _register_instance(self, instance_id: str, owner: TableAsync) -> None: + """ + Registers an instance with the client, and warms the channel pool + for the instance + The client will periodically refresh grpc channel pool used to make + requests, and new channels will be warmed for each registered instance + Channels will not be refreshed unless at least one instance is registered + + Args: + - instance_id: id of the instance to register. + - owner: table that owns the instance. Owners will be tracked in + _instance_owners, and instances will only be unregistered when all + owners call _remove_instance_registration + """ + instance_name = self._gapic_client.instance_path(self.project, instance_id) + instance_key = _WarmedInstanceKey( + instance_name, owner.table_name, owner.app_profile_id + ) + self._instance_owners.setdefault(instance_key, set()).add(id(owner)) + if instance_name not in self._active_instances: + self._active_instances.add(instance_key) + if self._channel_refresh_tasks: + # refresh tasks already running + # call ping and warm on all existing channels + for channel in self.transport.channels: + await self._ping_and_warm_instances(channel, instance_key) + else: + # refresh tasks aren't active. start them as background tasks + self.start_background_channel_refresh() + + async def _remove_instance_registration( + self, instance_id: str, owner: TableAsync + ) -> bool: + """ + Removes an instance from the client's registered instances, to prevent + warming new channels for the instance + + If instance_id is not registered, or is still in use by other tables, returns False + + Args: + - instance_id: id of the instance to remove + - owner: table that owns the instance. Owners will be tracked in + _instance_owners, and instances will only be unregistered when all + owners call _remove_instance_registration + Returns: + - True if instance was removed + """ + instance_name = self._gapic_client.instance_path(self.project, instance_id) + instance_key = _WarmedInstanceKey( + instance_name, owner.table_name, owner.app_profile_id + ) + owner_list = self._instance_owners.get(instance_key, set()) + try: + owner_list.remove(id(owner)) + if len(owner_list) == 0: + self._active_instances.remove(instance_key) + return True + except KeyError: + return False + + # TODO: revisit timeouts https://github.com/googleapis/python-bigtable/issues/782 + def get_table( + self, + instance_id: str, + table_id: str, + app_profile_id: str | None = None, + default_operation_timeout: float = 600, + default_per_request_timeout: float | None = None, + ) -> TableAsync: + """ + Returns a table instance for making data API requests + + Args: + instance_id: The Bigtable instance ID to associate with this client. + instance_id is combined with the client's project to fully + specify the instance + table_id: The ID of the table. + app_profile_id: (Optional) The app profile to associate with requests. + https://cloud.google.com/bigtable/docs/app-profiles + """ + return TableAsync( + self, + instance_id, + table_id, + app_profile_id, + default_operation_timeout=default_operation_timeout, + default_per_request_timeout=default_per_request_timeout, + ) + + async def __aenter__(self): + self.start_background_channel_refresh() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.close() + await self._gapic_client.__aexit__(exc_type, exc_val, exc_tb) + + +class TableAsync: + """ + Main Data API surface + + Table object maintains table_id, and app_profile_id context, and passes them with + each call + """ + + def __init__( + self, + client: BigtableDataClientAsync, + instance_id: str, + table_id: str, + app_profile_id: str | None = None, + *, + default_operation_timeout: float = 600, + default_per_request_timeout: float | None = None, + ): + """ + Initialize a Table instance + + Must be created within an async context (running event loop) + + Args: + instance_id: The Bigtable instance ID to associate with this client. + instance_id is combined with the client's project to fully + specify the instance + table_id: The ID of the table. table_id is combined with the + instance_id and the client's project to fully specify the table + app_profile_id: (Optional) The app profile to associate with requests. + https://cloud.google.com/bigtable/docs/app-profiles + default_operation_timeout: (Optional) The default timeout, in seconds + default_per_request_timeout: (Optional) The default timeout for individual + rpc requests, in seconds + Raises: + - RuntimeError if called outside of an async context (no running event loop) + """ + # validate timeouts + if default_operation_timeout <= 0: + raise ValueError("default_operation_timeout must be greater than 0") + if default_per_request_timeout is not None and default_per_request_timeout <= 0: + raise ValueError("default_per_request_timeout must be greater than 0") + if ( + default_per_request_timeout is not None + and default_per_request_timeout > default_operation_timeout + ): + raise ValueError( + "default_per_request_timeout must be less than default_operation_timeout" + ) + self.client = client + self.instance_id = instance_id + self.instance_name = self.client._gapic_client.instance_path( + self.client.project, instance_id + ) + self.table_id = table_id + self.table_name = self.client._gapic_client.table_path( + self.client.project, instance_id, table_id + ) + self.app_profile_id = app_profile_id + + self.default_operation_timeout = default_operation_timeout + self.default_per_request_timeout = default_per_request_timeout + + # raises RuntimeError if called outside of an async context (no running event loop) + try: + self._register_instance_task = asyncio.create_task( + self.client._register_instance(instance_id, self) + ) + except RuntimeError as e: + raise RuntimeError( + f"{self.__class__.__name__} must be created within an async event loop context." + ) from e + + async def read_rows_stream( + self, + query: ReadRowsQuery | dict[str, Any], + *, + operation_timeout: float | None = None, + per_request_timeout: float | None = None, + ) -> ReadRowsAsyncIterator: + """ + Returns an iterator to asynchronously stream back row data. + + Failed requests within operation_timeout and operation_deadline policies will be retried. + + Args: + - query: contains details about which rows to return + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + time is only counted while actively waiting on the network. + If None, defaults to the Table's default_operation_timeout + - per_request_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_per_request_timeout + + Returns: + - an asynchronous iterator that yields rows returned by the query + Raises: + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions + from any retries that failed + - GoogleAPIError: raised if the request encounters an unrecoverable error + - IdleTimeout: if iterator was abandoned + """ + + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = per_request_timeout or self.default_per_request_timeout + + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout <= 0: + raise ValueError("per_request_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout > operation_timeout: + raise ValueError( + "per_request_timeout must not be greater than operation_timeout" + ) + if per_request_timeout is None: + per_request_timeout = operation_timeout + request = query._to_dict() if isinstance(query, ReadRowsQuery) else query + request["table_name"] = self.table_name + if self.app_profile_id: + request["app_profile_id"] = self.app_profile_id + + # read_rows smart retries is implemented using a series of iterators: + # - client.read_rows: outputs raw ReadRowsResponse objects from backend. Has per_request_timeout + # - ReadRowsOperation.merge_row_response_stream: parses chunks into rows + # - ReadRowsOperation.retryable_merge_rows: adds retries, caching, revised requests, per_request_timeout + # - ReadRowsAsyncIterator: adds idle_timeout, moves stats out of stream and into attribute + row_merger = _ReadRowsOperationAsync( + request, + self.client._gapic_client, + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + output_generator = ReadRowsAsyncIterator(row_merger) + # add idle timeout to clear resources if generator is abandoned + idle_timeout_seconds = 300 + await output_generator._start_idle_timer(idle_timeout_seconds) + return output_generator + + async def read_rows( + self, + query: ReadRowsQuery | dict[str, Any], + *, + operation_timeout: float | None = None, + per_request_timeout: float | None = None, + ) -> list[Row]: + """ + Helper function that returns a full list instead of a generator + + See read_rows_stream + + Returns: + - a list of the rows returned by the query + """ + row_generator = await self.read_rows_stream( + query, + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + results = [row async for row in row_generator] + return results + + async def read_row( + self, + row_key: str | bytes, + *, + row_filter: RowFilter | None = None, + operation_timeout: int | float | None = 60, + per_request_timeout: int | float | None = None, + ) -> Row | None: + """ + Helper function to return a single row + + See read_rows_stream + + Raises: + - google.cloud.bigtable.data.exceptions.RowNotFound: if the row does not exist + Returns: + - the individual row requested, or None if it does not exist + """ + if row_key is None: + raise ValueError("row_key must be string or bytes") + query = ReadRowsQuery(row_keys=row_key, row_filter=row_filter, limit=1) + results = await self.read_rows( + query, + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + if len(results) == 0: + return None + return results[0] + + async def read_rows_sharded( + self, + sharded_query: ShardedQuery, + *, + operation_timeout: int | float | None = None, + per_request_timeout: int | float | None = None, + ) -> list[Row]: + """ + Runs a sharded query in parallel, then return the results in a single list. + Results will be returned in the order of the input queries. + + This function is intended to be run on the results on a query.shard() call: + + ``` + table_shard_keys = await table.sample_row_keys() + query = ReadRowsQuery(...) + shard_queries = query.shard(table_shard_keys) + results = await table.read_rows_sharded(shard_queries) + ``` + + Args: + - sharded_query: a sharded query to execute + Raises: + - ShardedReadRowsExceptionGroup: if any of the queries failed + - ValueError: if the query_list is empty + """ + if not sharded_query: + raise ValueError("empty sharded_query") + # reduce operation_timeout between batches + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = ( + per_request_timeout or self.default_per_request_timeout or operation_timeout + ) + timeout_generator = _attempt_timeout_generator( + operation_timeout, operation_timeout + ) + # submit shards in batches if the number of shards goes over CONCURRENCY_LIMIT + batched_queries = [ + sharded_query[i : i + CONCURRENCY_LIMIT] + for i in range(0, len(sharded_query), CONCURRENCY_LIMIT) + ] + # run batches and collect results + results_list = [] + error_dict = {} + shard_idx = 0 + for batch in batched_queries: + batch_operation_timeout = next(timeout_generator) + routine_list = [ + self.read_rows( + query, + operation_timeout=batch_operation_timeout, + per_request_timeout=min( + per_request_timeout, batch_operation_timeout + ), + ) + for query in batch + ] + batch_result = await asyncio.gather(*routine_list, return_exceptions=True) + for result in batch_result: + if isinstance(result, Exception): + error_dict[shard_idx] = result + else: + results_list.extend(result) + shard_idx += 1 + if error_dict: + # if any sub-request failed, raise an exception instead of returning results + raise ShardedReadRowsExceptionGroup( + [ + FailedQueryShardError(idx, sharded_query[idx], e) + for idx, e in error_dict.items() + ], + results_list, + len(sharded_query), + ) + return results_list + + async def row_exists( + self, + row_key: str | bytes, + *, + operation_timeout: int | float | None = 60, + per_request_timeout: int | float | None = None, + ) -> bool: + """ + Helper function to determine if a row exists + + uses the filters: chain(limit cells per row = 1, strip value) + + Returns: + - a bool indicating whether the row exists + """ + if row_key is None: + raise ValueError("row_key must be string or bytes") + strip_filter = StripValueTransformerFilter(flag=True) + limit_filter = CellsRowLimitFilter(1) + chain_filter = RowFilterChain(filters=[limit_filter, strip_filter]) + query = ReadRowsQuery(row_keys=row_key, limit=1, row_filter=chain_filter) + results = await self.read_rows( + query, + operation_timeout=operation_timeout, + per_request_timeout=per_request_timeout, + ) + return len(results) > 0 + + async def sample_row_keys( + self, + *, + operation_timeout: float | None = None, + per_request_timeout: float | None = None, + ) -> RowKeySamples: + """ + Return a set of RowKeySamples that delimit contiguous sections of the table of + approximately equal size + + RowKeySamples output can be used with ReadRowsQuery.shard() to create a sharded query that + can be parallelized across multiple backend nodes read_rows and read_rows_stream + requests will call sample_row_keys internally for this purpose when sharding is enabled + + RowKeySamples is simply a type alias for list[tuple[bytes, int]]; a list of + row_keys, along with offset positions in the table + + Returns: + - a set of RowKeySamples the delimit contiguous sections of the table + Raises: + - GoogleAPICallError: if the sample_row_keys request fails + """ + # prepare timeouts + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = per_request_timeout or self.default_per_request_timeout + + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout <= 0: + raise ValueError("per_request_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout > operation_timeout: + raise ValueError( + "per_request_timeout must not be greater than operation_timeout" + ) + attempt_timeout_gen = _attempt_timeout_generator( + per_request_timeout, operation_timeout + ) + # prepare retryable + predicate = retries.if_exception_type( + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ) + transient_errors = [] + + def on_error_fn(exc): + # add errors to list if retryable + if predicate(exc): + transient_errors.append(exc) + + retry = retries.AsyncRetry( + predicate=predicate, + timeout=operation_timeout, + initial=0.01, + multiplier=2, + maximum=60, + on_error=on_error_fn, + is_stream=False, + ) + + # prepare request + metadata = _make_metadata(self.table_name, self.app_profile_id) + + async def execute_rpc(): + results = await self.client._gapic_client.sample_row_keys( + table_name=self.table_name, + app_profile_id=self.app_profile_id, + timeout=next(attempt_timeout_gen), + metadata=metadata, + ) + return [(s.row_key, s.offset_bytes) async for s in results] + + wrapped_fn = _convert_retry_deadline( + retry(execute_rpc), operation_timeout, transient_errors + ) + return await wrapped_fn() + + def mutations_batcher( + self, + *, + flush_interval: float | None = 5, + flush_limit_mutation_count: int | None = 1000, + flush_limit_bytes: int = 20 * _MB_SIZE, + flow_control_max_mutation_count: int = 100_000, + flow_control_max_bytes: int = 100 * _MB_SIZE, + batch_operation_timeout: float | None = None, + batch_per_request_timeout: float | None = None, + ) -> MutationsBatcherAsync: + """ + Returns a new mutations batcher instance. + + Can be used to iteratively add mutations that are flushed as a group, + to avoid excess network calls + + Args: + - flush_interval: Automatically flush every flush_interval seconds. If None, + a table default will be used + - flush_limit_mutation_count: Flush immediately after flush_limit_mutation_count + mutations are added across all entries. If None, this limit is ignored. + - flush_limit_bytes: Flush immediately after flush_limit_bytes bytes are added. + - flow_control_max_mutation_count: Maximum number of inflight mutations. + - flow_control_max_bytes: Maximum number of inflight bytes. + - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None, + table default_operation_timeout will be used + - batch_per_request_timeout: timeout for each individual request, in seconds. If None, + table default_per_request_timeout will be used + Returns: + - a MutationsBatcherAsync context manager that can batch requests + """ + return MutationsBatcherAsync( + self, + flush_interval=flush_interval, + flush_limit_mutation_count=flush_limit_mutation_count, + flush_limit_bytes=flush_limit_bytes, + flow_control_max_mutation_count=flow_control_max_mutation_count, + flow_control_max_bytes=flow_control_max_bytes, + batch_operation_timeout=batch_operation_timeout, + batch_per_request_timeout=batch_per_request_timeout, + ) + + async def mutate_row( + self, + row_key: str | bytes, + mutations: list[Mutation] | Mutation, + *, + operation_timeout: float | None = 60, + per_request_timeout: float | None = None, + ): + """ + Mutates a row atomically. + + Cells already present in the row are left unchanged unless explicitly changed + by ``mutation``. + + Idempotent operations (i.e, all mutations have an explicit timestamp) will be + retried on server failure. Non-idempotent operations will not. + + Args: + - row_key: the row to apply mutations to + - mutations: the set of mutations to apply to the row + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + time is only counted while actively waiting on the network. + DeadlineExceeded exception raised after timeout + - per_request_timeout: the time budget for an individual network request, + in seconds. If it takes longer than this time to complete, the request + will be cancelled with a DeadlineExceeded exception, and a retry will be + attempted if within operation_timeout budget + + Raises: + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing all + GoogleAPIError exceptions from any retries that failed + - GoogleAPIError: raised on non-idempotent operations that cannot be + safely retried. + """ + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = per_request_timeout or self.default_per_request_timeout + + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout <= 0: + raise ValueError("per_request_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout > operation_timeout: + raise ValueError("per_request_timeout must be less than operation_timeout") + + if isinstance(row_key, str): + row_key = row_key.encode("utf-8") + request = {"table_name": self.table_name, "row_key": row_key} + if self.app_profile_id: + request["app_profile_id"] = self.app_profile_id + + if isinstance(mutations, Mutation): + mutations = [mutations] + request["mutations"] = [mutation._to_dict() for mutation in mutations] + + if all(mutation.is_idempotent() for mutation in mutations): + # mutations are all idempotent and safe to retry + predicate = retries.if_exception_type( + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ) + else: + # mutations should not be retried + predicate = retries.if_exception_type() + + transient_errors = [] + + def on_error_fn(exc): + if predicate(exc): + transient_errors.append(exc) + + retry = retries.AsyncRetry( + predicate=predicate, + on_error=on_error_fn, + timeout=operation_timeout, + initial=0.01, + multiplier=2, + maximum=60, + ) + # wrap rpc in retry logic + retry_wrapped = retry(self.client._gapic_client.mutate_row) + # convert RetryErrors from retry wrapper into DeadlineExceeded errors + deadline_wrapped = _convert_retry_deadline( + retry_wrapped, operation_timeout, transient_errors + ) + metadata = _make_metadata(self.table_name, self.app_profile_id) + # trigger rpc + await deadline_wrapped(request, timeout=per_request_timeout, metadata=metadata) + + async def bulk_mutate_rows( + self, + mutation_entries: list[RowMutationEntry], + *, + operation_timeout: float | None = 60, + per_request_timeout: float | None = None, + ): + """ + Applies mutations for multiple rows in a single batched request. + + Each individual RowMutationEntry is applied atomically, but separate entries + may be applied in arbitrary order (even for entries targetting the same row) + In total, the row_mutations can contain at most 100000 individual mutations + across all entries + + Idempotent entries (i.e., entries with mutations with explicit timestamps) + will be retried on failure. Non-idempotent will not, and will reported in a + raised exception group + + Args: + - mutation_entries: the batches of mutations to apply + Each entry will be applied atomically, but entries will be applied + in arbitrary order + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + time is only counted while actively waiting on the network. + DeadlineExceeded exception raised after timeout + - per_request_timeout: the time budget for an individual network request, + in seconds. If it takes longer than this time to complete, the request + will be cancelled with a DeadlineExceeded exception, and a retry will + be attempted if within operation_timeout budget + Raises: + - MutationsExceptionGroup if one or more mutations fails + Contains details about any failed entries in .exceptions + """ + operation_timeout = operation_timeout or self.default_operation_timeout + per_request_timeout = per_request_timeout or self.default_per_request_timeout + + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout <= 0: + raise ValueError("per_request_timeout must be greater than 0") + if per_request_timeout is not None and per_request_timeout > operation_timeout: + raise ValueError("per_request_timeout must be less than operation_timeout") + + operation = _MutateRowsOperationAsync( + self.client._gapic_client, + self, + mutation_entries, + operation_timeout, + per_request_timeout, + ) + await operation.start() + + async def check_and_mutate_row( + self, + row_key: str | bytes, + predicate: RowFilter | dict[str, Any] | None, + *, + true_case_mutations: Mutation | list[Mutation] | None = None, + false_case_mutations: Mutation | list[Mutation] | None = None, + operation_timeout: int | float | None = 20, + ) -> bool: + """ + Mutates a row atomically based on the output of a predicate filter + + Non-idempotent operation: will not be retried + + Args: + - row_key: the key of the row to mutate + - predicate: the filter to be applied to the contents of the specified row. + Depending on whether or not any results are yielded, + either true_case_mutations or false_case_mutations will be executed. + If None, checks that the row contains any values at all. + - true_case_mutations: + Changes to be atomically applied to the specified row if + predicate yields at least one cell when + applied to row_key. Entries are applied in order, + meaning that earlier mutations can be masked by later + ones. Must contain at least one entry if + false_case_mutations is empty, and at most 100000. + - false_case_mutations: + Changes to be atomically applied to the specified row if + predicate_filter does not yield any cells when + applied to row_key. Entries are applied in order, + meaning that earlier mutations can be masked by later + ones. Must contain at least one entry if + `true_case_mutations is empty, and at most 100000. + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will not be retried. + Returns: + - bool indicating whether the predicate was true or false + Raises: + - GoogleAPIError exceptions from grpc call + """ + operation_timeout = operation_timeout or self.default_operation_timeout + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key + if true_case_mutations is not None and not isinstance( + true_case_mutations, list + ): + true_case_mutations = [true_case_mutations] + true_case_dict = [m._to_dict() for m in true_case_mutations or []] + if false_case_mutations is not None and not isinstance( + false_case_mutations, list + ): + false_case_mutations = [false_case_mutations] + false_case_dict = [m._to_dict() for m in false_case_mutations or []] + if predicate is not None and not isinstance(predicate, dict): + predicate = predicate.to_dict() + metadata = _make_metadata(self.table_name, self.app_profile_id) + result = await self.client._gapic_client.check_and_mutate_row( + request={ + "predicate_filter": predicate, + "true_mutations": true_case_dict, + "false_mutations": false_case_dict, + "table_name": self.table_name, + "row_key": row_key, + "app_profile_id": self.app_profile_id, + }, + metadata=metadata, + timeout=operation_timeout, + ) + return result.predicate_matched + + async def read_modify_write_row( + self, + row_key: str | bytes, + rules: ReadModifyWriteRule | list[ReadModifyWriteRule], + *, + operation_timeout: int | float | None = 20, + ) -> Row: + """ + Reads and modifies a row atomically according to input ReadModifyWriteRules, + and returns the contents of all modified cells + + The new value for the timestamp is the greater of the existing timestamp or + the current server time. + + Non-idempotent operation: will not be retried + + Args: + - row_key: the key of the row to apply read/modify/write rules to + - rules: A rule or set of rules to apply to the row. + Rules are applied in order, meaning that earlier rules will affect the + results of later ones. + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will not be retried. + Returns: + - Row: containing cell data that was modified as part of the + operation + Raises: + - GoogleAPIError exceptions from grpc call + """ + operation_timeout = operation_timeout or self.default_operation_timeout + row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if rules is not None and not isinstance(rules, list): + rules = [rules] + if not rules: + raise ValueError("rules must contain at least one item") + # concert to dict representation + rules_dict = [rule._to_dict() for rule in rules] + metadata = _make_metadata(self.table_name, self.app_profile_id) + result = await self.client._gapic_client.read_modify_write_row( + request={ + "rules": rules_dict, + "table_name": self.table_name, + "row_key": row_key, + "app_profile_id": self.app_profile_id, + }, + metadata=metadata, + timeout=operation_timeout, + ) + # construct Row from result + return Row._from_pb(result.row) + + async def close(self): + """ + Called to close the Table instance and release any resources held by it. + """ + self._register_instance_task.cancel() + await self.client._remove_instance_registration(self.instance_id, self) + + async def __aenter__(self): + """ + Implement async context manager protocol + + Ensure registration task has time to run, so that + grpc channels will be warmed for the specified instance + """ + await self._register_instance_task + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """ + Implement async context manager protocol + + Unregister this instance with the client, so that + grpc channels will no longer be warmed + """ + await self.close() diff --git a/google/cloud/bigtable/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py similarity index 96% rename from google/cloud/bigtable/mutations_batcher.py rename to google/cloud/bigtable/data/_async/mutations_batcher.py index 68c3f9fbe..25aafc2a1 100644 --- a/google/cloud/bigtable/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -20,22 +20,24 @@ import warnings from collections import deque -from google.cloud.bigtable.mutations import RowMutationEntry -from google.cloud.bigtable.exceptions import MutationsExceptionGroup -from google.cloud.bigtable.exceptions import FailedMutationEntryError +from google.cloud.bigtable.data.mutations import RowMutationEntry +from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup +from google.cloud.bigtable.data.exceptions import FailedMutationEntryError -from google.cloud.bigtable._mutate_rows import _MutateRowsOperation -from google.cloud.bigtable._mutate_rows import MUTATE_ROWS_REQUEST_MUTATION_LIMIT -from google.cloud.bigtable.mutations import Mutation +from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync +from google.cloud.bigtable.data._async._mutate_rows import ( + MUTATE_ROWS_REQUEST_MUTATION_LIMIT, +) +from google.cloud.bigtable.data.mutations import Mutation if TYPE_CHECKING: - from google.cloud.bigtable.client import Table # pragma: no cover + from google.cloud.bigtable.data._async.client import TableAsync # used to make more readable default values _MB_SIZE = 1024 * 1024 -class _FlowControl: +class _FlowControlAsync: """ Manages flow control for batched mutations. Mutations are registered against the FlowControl object before being sent, which will block if size or count @@ -159,7 +161,7 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry] yield mutations[start_idx:end_idx] -class MutationsBatcher: +class MutationsBatcherAsync: """ Allows users to send batches using context manager API: @@ -179,7 +181,7 @@ class MutationsBatcher: def __init__( self, - table: "Table", + table: "TableAsync", *, flush_interval: float | None = 5, flush_limit_mutation_count: int | None = 1000, @@ -224,7 +226,7 @@ def __init__( self._table = table self._staged_entries: list[RowMutationEntry] = [] self._staged_count, self._staged_bytes = 0, 0 - self._flow_control = _FlowControl( + self._flow_control = _FlowControlAsync( flow_control_max_mutation_count, flow_control_max_bytes ) self._flush_limit_bytes = flush_limit_bytes @@ -354,7 +356,7 @@ async def _execute_mutate_rows( if self._table.app_profile_id: request["app_profile_id"] = self._table.app_profile_id try: - operation = _MutateRowsOperation( + operation = _MutateRowsOperationAsync( self._table.client._gapic_client, self._table, batch, diff --git a/google/cloud/bigtable/_helpers.py b/google/cloud/bigtable/data/_helpers.py similarity index 98% rename from google/cloud/bigtable/_helpers.py rename to google/cloud/bigtable/data/_helpers.py index 722fac9f4..64d91e108 100644 --- a/google/cloud/bigtable/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -18,7 +18,7 @@ import time from google.api_core import exceptions as core_exceptions -from google.cloud.bigtable.exceptions import RetryExceptionGroup +from google.cloud.bigtable.data.exceptions import RetryExceptionGroup """ Helper functions used in various places in the library. diff --git a/google/cloud/bigtable/_read_rows.py b/google/cloud/bigtable/data/_read_rows_state_machine.py similarity index 54% rename from google/cloud/bigtable/_read_rows.py rename to google/cloud/bigtable/data/_read_rows_state_machine.py index ee094f1a7..7c0d05fb9 100644 --- a/google/cloud/bigtable/_read_rows.py +++ b/google/cloud/bigtable/data/_read_rows_state_machine.py @@ -14,35 +14,14 @@ # from __future__ import annotations -from typing import ( - List, - Any, - AsyncIterable, - AsyncIterator, - AsyncGenerator, - Iterator, - Callable, - Awaitable, - Type, -) - -import asyncio -from functools import partial -from grpc.aio import RpcContext +from typing import Type from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse -from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient -from google.cloud.bigtable.row import Row, Cell, _LastScannedRow -from google.cloud.bigtable.exceptions import InvalidChunk -from google.cloud.bigtable.exceptions import _RowSetComplete -from google.api_core import retry_async as retries -from google.api_core import exceptions as core_exceptions -from google.cloud.bigtable._helpers import _make_metadata -from google.cloud.bigtable._helpers import _attempt_timeout_generator +from google.cloud.bigtable.data.row import Row, Cell, _LastScannedRow +from google.cloud.bigtable.data.exceptions import InvalidChunk """ -This module provides a set of classes for merging ReadRowsResponse chunks -into Row objects. +This module provides classes for the read_rows state machine: - ReadRowsOperation is the highest level class, providing an interface for asynchronous merging end-to-end @@ -56,253 +35,6 @@ """ -class _ReadRowsOperation(AsyncIterable[Row]): - """ - ReadRowsOperation handles the logic of merging chunks from a ReadRowsResponse stream - into a stream of Row objects. - - ReadRowsOperation.merge_row_response_stream takes in a stream of ReadRowsResponse - and turns them into a stream of Row objects using an internal - StateMachine. - - ReadRowsOperation(request, client) handles row merging logic end-to-end, including - performing retries on stream errors. - """ - - def __init__( - self, - request: dict[str, Any], - client: BigtableAsyncClient, - *, - operation_timeout: float = 600.0, - per_request_timeout: float | None = None, - ): - """ - Args: - - request: the request dict to send to the Bigtable API - - client: the Bigtable client to use to make the request - - operation_timeout: the timeout to use for the entire operation, in seconds - - per_request_timeout: the timeout to use when waiting for each individual grpc request, in seconds - If not specified, defaults to operation_timeout - """ - self._last_emitted_row_key: bytes | None = None - self._emit_count = 0 - self._request = request - self.operation_timeout = operation_timeout - # use generator to lower per-attempt timeout as we approach operation_timeout deadline - attempt_timeout_gen = _attempt_timeout_generator( - per_request_timeout, operation_timeout - ) - row_limit = request.get("rows_limit", 0) - # lock in paramters for retryable wrapper - self._partial_retryable = partial( - self._read_rows_retryable_attempt, - client.read_rows, - attempt_timeout_gen, - row_limit, - ) - predicate = retries.if_exception_type( - core_exceptions.DeadlineExceeded, - core_exceptions.ServiceUnavailable, - core_exceptions.Aborted, - ) - - def on_error_fn(exc): - if predicate(exc): - self.transient_errors.append(exc) - - retry = retries.AsyncRetry( - predicate=predicate, - timeout=self.operation_timeout, - initial=0.01, - multiplier=2, - maximum=60, - on_error=on_error_fn, - is_stream=True, - ) - self._stream: AsyncGenerator[Row, None] | None = retry( - self._partial_retryable - )() - # contains the list of errors that were retried - self.transient_errors: List[Exception] = [] - - def __aiter__(self) -> AsyncIterator[Row]: - """Implements the AsyncIterable interface""" - return self - - async def __anext__(self) -> Row: - """Implements the AsyncIterator interface""" - if self._stream is not None: - return await self._stream.__anext__() - else: - raise asyncio.InvalidStateError("stream is closed") - - async def aclose(self): - """Close the stream and release resources""" - if self._stream is not None: - await self._stream.aclose() - self._stream = None - self._emitted_seen_row_key = None - - async def _read_rows_retryable_attempt( - self, - gapic_fn: Callable[..., Awaitable[AsyncIterable[ReadRowsResponse]]], - timeout_generator: Iterator[float], - total_row_limit: int, - ) -> AsyncGenerator[Row, None]: - """ - Retryable wrapper for merge_rows. This function is called each time - a retry is attempted. - - Some fresh state is created on each retry: - - grpc network stream - - state machine to hold merge chunks received from stream - Some state is shared between retries: - - _last_emitted_row_key is used to ensure that - duplicate rows are not emitted - - request is stored and (potentially) modified on each retry - """ - if self._last_emitted_row_key is not None: - # if this is a retry, try to trim down the request to avoid ones we've already processed - try: - self._request["rows"] = _ReadRowsOperation._revise_request_rowset( - row_set=self._request.get("rows", None), - last_seen_row_key=self._last_emitted_row_key, - ) - except _RowSetComplete: - # if there are no rows left to process, we're done - # This is not expected to happen often, but could occur if - # a retry is triggered quickly after the last row is emitted - return - # revise next request's row limit based on number emitted - if total_row_limit: - new_limit = total_row_limit - self._emit_count - if new_limit == 0: - # we have hit the row limit, so we're done - return - elif new_limit < 0: - raise RuntimeError("unexpected state: emit count exceeds row limit") - else: - self._request["rows_limit"] = new_limit - metadata = _make_metadata( - self._request.get("table_name", None), - self._request.get("app_profile_id", None), - ) - new_gapic_stream: RpcContext = await gapic_fn( - self._request, - timeout=next(timeout_generator), - metadata=metadata, - ) - try: - state_machine = _StateMachine() - stream = _ReadRowsOperation.merge_row_response_stream( - new_gapic_stream, state_machine - ) - # run until we get a timeout or the stream is exhausted - async for new_item in stream: - if ( - self._last_emitted_row_key is not None - and new_item.row_key <= self._last_emitted_row_key - ): - raise InvalidChunk("Last emitted row key out of order") - # don't yeild _LastScannedRow markers; they - # should only update last_seen_row_key - if not isinstance(new_item, _LastScannedRow): - yield new_item - self._emit_count += 1 - self._last_emitted_row_key = new_item.row_key - if total_row_limit and self._emit_count >= total_row_limit: - return - except (Exception, GeneratorExit) as exc: - # ensure grpc stream is closed - new_gapic_stream.cancel() - raise exc - - @staticmethod - def _revise_request_rowset( - row_set: dict[str, Any] | None, - last_seen_row_key: bytes, - ) -> dict[str, Any]: - """ - Revise the rows in the request to avoid ones we've already processed. - - Args: - - row_set: the row set from the request - - last_seen_row_key: the last row key encountered - Raises: - - _RowSetComplete: if there are no rows left to process after the revision - """ - # if user is doing a whole table scan, start a new one with the last seen key - if row_set is None or ( - len(row_set.get("row_ranges", [])) == 0 - and len(row_set.get("row_keys", [])) == 0 - ): - last_seen = last_seen_row_key - return { - "row_keys": [], - "row_ranges": [{"start_key_open": last_seen}], - } - # remove seen keys from user-specific key list - row_keys: list[bytes] = row_set.get("row_keys", []) - adjusted_keys = [k for k in row_keys if k > last_seen_row_key] - # adjust ranges to ignore keys before last seen - row_ranges: list[dict[str, Any]] = row_set.get("row_ranges", []) - adjusted_ranges = [] - for row_range in row_ranges: - end_key = row_range.get("end_key_closed", None) or row_range.get( - "end_key_open", None - ) - if end_key is None or end_key > last_seen_row_key: - # end range is after last seen key - new_range = row_range.copy() - start_key = row_range.get("start_key_closed", None) or row_range.get( - "start_key_open", None - ) - if start_key is None or start_key <= last_seen_row_key: - # replace start key with last seen - new_range["start_key_open"] = last_seen_row_key - new_range.pop("start_key_closed", None) - adjusted_ranges.append(new_range) - if len(adjusted_keys) == 0 and len(adjusted_ranges) == 0: - # if the query is empty after revision, raise an exception - # this will avoid an unwanted full table scan - raise _RowSetComplete() - return {"row_keys": adjusted_keys, "row_ranges": adjusted_ranges} - - @staticmethod - async def merge_row_response_stream( - response_generator: AsyncIterable[ReadRowsResponse], - state_machine: _StateMachine, - ) -> AsyncGenerator[Row, None]: - """ - Consume chunks from a ReadRowsResponse stream into a set of Rows - - Args: - - response_generator: AsyncIterable of ReadRowsResponse objects. Typically - this is a stream of chunks from the Bigtable API - Returns: - - AsyncGenerator of Rows - Raises: - - InvalidChunk: if the chunk stream is invalid - """ - async for row_response in response_generator: - # unwrap protoplus object for increased performance - response_pb = row_response._pb - last_scanned = response_pb.last_scanned_row_key - # if the server sends a scan heartbeat, notify the state machine. - if last_scanned: - yield state_machine.handle_last_scanned_row(last_scanned) - # process new chunks through the state machine. - for chunk in response_pb.chunks: - complete_row = state_machine.handle_chunk(chunk) - if complete_row is not None: - yield complete_row - # TODO: handle request stats - if not state_machine.is_terminal_state(): - # read rows is complete, but there's still data in the merger - raise InvalidChunk("read_rows completed with partial state remaining") - - class _StateMachine: """ State Machine converts chunks into Rows @@ -579,7 +311,7 @@ def reset(self) -> None: self.current_key: bytes | None = None self.working_cell: Cell | None = None self.working_value: bytearray | None = None - self.completed_cells: List[Cell] = [] + self.completed_cells: list[Cell] = [] def start_row(self, key: bytes) -> None: """Called to start a new row. This will be called once per row""" @@ -590,7 +322,7 @@ def start_cell( family: str, qualifier: bytes, timestamp_micros: int, - labels: List[str], + labels: list[str], ) -> None: """called to start a new cell in a row.""" if self.current_key is None: diff --git a/google/cloud/bigtable/exceptions.py b/google/cloud/bigtable/data/exceptions.py similarity index 93% rename from google/cloud/bigtable/exceptions.py rename to google/cloud/bigtable/data/exceptions.py index fc4e368b9..9b6b4fe3f 100644 --- a/google/cloud/bigtable/exceptions.py +++ b/google/cloud/bigtable/data/exceptions.py @@ -19,13 +19,13 @@ from typing import Any, TYPE_CHECKING from google.api_core import exceptions as core_exceptions -from google.cloud.bigtable.row import Row +from google.cloud.bigtable.data.row import Row is_311_plus = sys.version_info >= (3, 11) if TYPE_CHECKING: - from google.cloud.bigtable.mutations import RowMutationEntry - from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.mutations import RowMutationEntry + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery class IdleTimeout(core_exceptions.DeadlineExceeded): @@ -50,7 +50,15 @@ class _RowSetComplete(Exception): pass -class BigtableExceptionGroup(ExceptionGroup if is_311_plus else Exception): # type: ignore # noqa: F821 +class _MutateRowsIncomplete(RuntimeError): + """ + Exception raised when a mutate_rows call has unfinished work. + """ + + pass + + +class _BigtableExceptionGroup(ExceptionGroup if is_311_plus else Exception): # type: ignore # noqa: F821 """ Represents one or more exceptions that occur during a bulk Bigtable operation @@ -82,7 +90,7 @@ def __str__(self): return self.args[0] -class MutationsExceptionGroup(BigtableExceptionGroup): +class MutationsExceptionGroup(_BigtableExceptionGroup): """ Represents one or more exceptions that occur during a bulk mutation operation @@ -202,7 +210,7 @@ def __init__( self.__cause__ = cause -class RetryExceptionGroup(BigtableExceptionGroup): +class RetryExceptionGroup(_BigtableExceptionGroup): """Represents one or more exceptions that occur during a retryable operation""" @staticmethod @@ -221,7 +229,7 @@ def __new__(cls, excs: list[Exception]): return super().__new__(cls, cls._format_message(excs), excs) -class ShardedReadRowsExceptionGroup(BigtableExceptionGroup): +class ShardedReadRowsExceptionGroup(_BigtableExceptionGroup): """ Represents one or more exceptions that occur during a sharded read rows operation """ diff --git a/google/cloud/bigtable/mutations.py b/google/cloud/bigtable/data/mutations.py similarity index 98% rename from google/cloud/bigtable/mutations.py rename to google/cloud/bigtable/data/mutations.py index a4c02cd74..de1b3b137 100644 --- a/google/cloud/bigtable/mutations.py +++ b/google/cloud/bigtable/data/mutations.py @@ -19,16 +19,16 @@ from abc import ABC, abstractmethod from sys import getsizeof -# mutation entries above this should be rejected -from google.cloud.bigtable._mutate_rows import MUTATE_ROWS_REQUEST_MUTATION_LIMIT - - -from google.cloud.bigtable.read_modify_write_rules import MAX_INCREMENT_VALUE +from google.cloud.bigtable.data.read_modify_write_rules import MAX_INCREMENT_VALUE # special value for SetCell mutation timestamps. If set, server will assign a timestamp SERVER_SIDE_TIMESTAMP = -1 +# mutation entries above this should be rejected +MUTATE_ROWS_REQUEST_MUTATION_LIMIT = 100_000 + + class Mutation(ABC): """Model class for mutations""" diff --git a/google/cloud/bigtable/read_modify_write_rules.py b/google/cloud/bigtable/data/read_modify_write_rules.py similarity index 100% rename from google/cloud/bigtable/read_modify_write_rules.py rename to google/cloud/bigtable/data/read_modify_write_rules.py diff --git a/google/cloud/bigtable/read_rows_query.py b/google/cloud/bigtable/data/read_rows_query.py similarity index 98% rename from google/cloud/bigtable/read_rows_query.py rename to google/cloud/bigtable/data/read_rows_query.py index eb28eeda3..7d7e1f99f 100644 --- a/google/cloud/bigtable/read_rows_query.py +++ b/google/cloud/bigtable/data/read_rows_query.py @@ -18,11 +18,11 @@ from bisect import bisect_right from collections import defaultdict from dataclasses import dataclass -from google.cloud.bigtable.row_filters import RowFilter +from google.cloud.bigtable.data.row_filters import RowFilter if TYPE_CHECKING: - from google.cloud.bigtable import RowKeySamples - from google.cloud.bigtable import ShardedQuery + from google.cloud.bigtable.data import RowKeySamples + from google.cloud.bigtable.data import ShardedQuery @dataclass diff --git a/google/cloud/bigtable/data/row.py b/google/cloud/bigtable/data/row.py new file mode 100644 index 000000000..5fdc1b365 --- /dev/null +++ b/google/cloud/bigtable/data/row.py @@ -0,0 +1,465 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from collections import OrderedDict +from typing import Sequence, Generator, overload, Any +from functools import total_ordering + +from google.cloud.bigtable_v2.types import Row as RowPB + +# Type aliases used internally for readability. +_family_type = str +_qualifier_type = bytes + + +class Row(Sequence["Cell"]): + """ + Model class for row data returned from server + + Does not represent all data contained in the row, only data returned by a + query. + Expected to be read-only to users, and written by backend + + Can be indexed: + cells = row["family", "qualifier"] + """ + + __slots__ = ("row_key", "cells", "_index_data") + + def __init__( + self, + key: bytes, + cells: list[Cell], + ): + """ + Initializes a Row object + + Row objects are not intended to be created by users. + They are returned by the Bigtable backend. + """ + self.row_key = key + self.cells: list[Cell] = cells + # index is lazily created when needed + self._index_data: OrderedDict[ + _family_type, OrderedDict[_qualifier_type, list[Cell]] + ] | None = None + + @property + def _index( + self, + ) -> OrderedDict[_family_type, OrderedDict[_qualifier_type, list[Cell]]]: + """ + Returns an index of cells associated with each family and qualifier. + + The index is lazily created when needed + """ + if self._index_data is None: + self._index_data = OrderedDict() + for cell in self.cells: + self._index_data.setdefault(cell.family, OrderedDict()).setdefault( + cell.qualifier, [] + ).append(cell) + return self._index_data + + @classmethod + def _from_pb(cls, row_pb: RowPB) -> Row: + """ + Creates a row from a protobuf representation + + Row objects are not intended to be created by users. + They are returned by the Bigtable backend. + """ + row_key: bytes = row_pb.key + cell_list: list[Cell] = [] + for family in row_pb.families: + for column in family.columns: + for cell in column.cells: + new_cell = Cell( + value=cell.value, + row_key=row_key, + family=family.name, + qualifier=column.qualifier, + timestamp_micros=cell.timestamp_micros, + labels=list(cell.labels) if cell.labels else None, + ) + cell_list.append(new_cell) + return cls(row_key, cells=cell_list) + + def get_cells( + self, family: str | None = None, qualifier: str | bytes | None = None + ) -> list[Cell]: + """ + Returns cells sorted in Bigtable native order: + - Family lexicographically ascending + - Qualifier ascending + - Timestamp in reverse chronological order + + If family or qualifier not passed, will include all + + Can also be accessed through indexing: + cells = row["family", "qualifier"] + cells = row["family"] + """ + if family is None: + if qualifier is not None: + # get_cells(None, "qualifier") is not allowed + raise ValueError("Qualifier passed without family") + else: + # return all cells on get_cells() + return self.cells + if qualifier is None: + # return all cells in family on get_cells(family) + return list(self._get_all_from_family(family)) + if isinstance(qualifier, str): + qualifier = qualifier.encode("utf-8") + # return cells in family and qualifier on get_cells(family, qualifier) + if family not in self._index: + raise ValueError(f"Family '{family}' not found in row '{self.row_key!r}'") + if qualifier not in self._index[family]: + raise ValueError( + f"Qualifier '{qualifier!r}' not found in family '{family}' in row '{self.row_key!r}'" + ) + return self._index[family][qualifier] + + def _get_all_from_family(self, family: str) -> Generator[Cell, None, None]: + """ + Returns all cells in the row for the family_id + """ + if family not in self._index: + raise ValueError(f"Family '{family}' not found in row '{self.row_key!r}'") + for qualifier in self._index[family]: + yield from self._index[family][qualifier] + + def __str__(self) -> str: + """ + Human-readable string representation + + { + (family='fam', qualifier=b'col'): [b'value', (+1 more),], + (family='fam', qualifier=b'col2'): [b'other'], + } + """ + output = ["{"] + for family, qualifier in self.get_column_components(): + cell_list = self[family, qualifier] + line = [f" (family={family!r}, qualifier={qualifier!r}): "] + if len(cell_list) == 0: + line.append("[],") + elif len(cell_list) == 1: + line.append(f"[{cell_list[0]}],") + else: + line.append(f"[{cell_list[0]}, (+{len(cell_list)-1} more)],") + output.append("".join(line)) + output.append("}") + return "\n".join(output) + + def __repr__(self): + cell_str_buffer = ["{"] + for family, qualifier in self.get_column_components(): + cell_list = self[family, qualifier] + repr_list = [cell.to_dict() for cell in cell_list] + cell_str_buffer.append(f" ('{family}', {qualifier!r}): {repr_list},") + cell_str_buffer.append("}") + cell_str = "\n".join(cell_str_buffer) + output = f"Row(key={self.row_key!r}, cells={cell_str})" + return output + + def to_dict(self) -> dict[str, Any]: + """ + Returns a dictionary representation of the cell in the Bigtable Row + proto format + + https://cloud.google.com/bigtable/docs/reference/data/rpc/google.bigtable.v2#row + """ + family_list = [] + for family_name, qualifier_dict in self._index.items(): + qualifier_list = [] + for qualifier_name, cell_list in qualifier_dict.items(): + cell_dicts = [cell.to_dict() for cell in cell_list] + qualifier_list.append( + {"qualifier": qualifier_name, "cells": cell_dicts} + ) + family_list.append({"name": family_name, "columns": qualifier_list}) + return {"key": self.row_key, "families": family_list} + + # Sequence and Mapping methods + def __iter__(self): + """ + Allow iterating over all cells in the row + """ + return iter(self.cells) + + def __contains__(self, item): + """ + Implements `in` operator + + Works for both cells in the internal list, and `family` or + `(family, qualifier)` pairs associated with the cells + """ + if isinstance(item, _family_type): + return item in self._index + elif ( + isinstance(item, tuple) + and isinstance(item[0], _family_type) + and isinstance(item[1], (bytes, str)) + ): + q = item[1] if isinstance(item[1], bytes) else item[1].encode("utf-8") + return item[0] in self._index and q in self._index[item[0]] + # check if Cell is in Row + return item in self.cells + + @overload + def __getitem__( + self, + index: str | tuple[str, bytes | str], + ) -> list[Cell]: + # overload signature for type checking + pass + + @overload + def __getitem__(self, index: int) -> Cell: + # overload signature for type checking + pass + + @overload + def __getitem__(self, index: slice) -> list[Cell]: + # overload signature for type checking + pass + + def __getitem__(self, index): + """ + Implements [] indexing + + Supports indexing by family, (family, qualifier) pair, + numerical index, and index slicing + """ + if isinstance(index, _family_type): + return self.get_cells(family=index) + elif ( + isinstance(index, tuple) + and isinstance(index[0], _family_type) + and isinstance(index[1], (bytes, str)) + ): + return self.get_cells(family=index[0], qualifier=index[1]) + elif isinstance(index, int) or isinstance(index, slice): + # index is int or slice + return self.cells[index] + else: + raise TypeError( + "Index must be family_id, (family_id, qualifier), int, or slice" + ) + + def __len__(self): + """ + Implements `len()` operator + """ + return len(self.cells) + + def get_column_components(self) -> list[tuple[str, bytes]]: + """ + Returns a list of (family, qualifier) pairs associated with the cells + + Pairs can be used for indexing + """ + return [(f, q) for f in self._index for q in self._index[f]] + + def __eq__(self, other): + """ + Implements `==` operator + """ + # for performance reasons, check row metadata + # before checking individual cells + if not isinstance(other, Row): + return False + if self.row_key != other.row_key: + return False + if len(self.cells) != len(other.cells): + return False + components = self.get_column_components() + other_components = other.get_column_components() + if len(components) != len(other_components): + return False + if components != other_components: + return False + for family, qualifier in components: + if len(self[family, qualifier]) != len(other[family, qualifier]): + return False + # compare individual cell lists + if self.cells != other.cells: + return False + return True + + def __ne__(self, other) -> bool: + """ + Implements `!=` operator + """ + return not self == other + + +class _LastScannedRow(Row): + """A value used to indicate a scanned row that is not returned as part of + a query. + + This is used internally to indicate progress in a scan, and improve retry + performance. It is not intended to be used directly by users. + """ + + def __init__(self, row_key): + super().__init__(row_key, []) + + def __eq__(self, other): + return isinstance(other, _LastScannedRow) + + +@total_ordering +class Cell: + """ + Model class for cell data + + Does not represent all data contained in the cell, only data returned by a + query. + Expected to be read-only to users, and written by backend + """ + + __slots__ = ( + "value", + "row_key", + "family", + "qualifier", + "timestamp_micros", + "labels", + ) + + def __init__( + self, + value: bytes, + row_key: bytes, + family: str, + qualifier: bytes | str, + timestamp_micros: int, + labels: list[str] | None = None, + ): + """ + Cell constructor + + Cell objects are not intended to be constructed by users. + They are returned by the Bigtable backend. + """ + self.value = value + self.row_key = row_key + self.family = family + if isinstance(qualifier, str): + qualifier = qualifier.encode() + self.qualifier = qualifier + self.timestamp_micros = timestamp_micros + self.labels = labels if labels is not None else [] + + def __int__(self) -> int: + """ + Allows casting cell to int + Interprets value as a 64-bit big-endian signed integer, as expected by + ReadModifyWrite increment rule + """ + return int.from_bytes(self.value, byteorder="big", signed=True) + + def to_dict(self) -> dict[str, Any]: + """ + Returns a dictionary representation of the cell in the Bigtable Cell + proto format + + https://cloud.google.com/bigtable/docs/reference/data/rpc/google.bigtable.v2#cell + """ + cell_dict: dict[str, Any] = { + "value": self.value, + } + cell_dict["timestamp_micros"] = self.timestamp_micros + if self.labels: + cell_dict["labels"] = self.labels + return cell_dict + + def __str__(self) -> str: + """ + Allows casting cell to str + Prints encoded byte string, same as printing value directly. + """ + return str(self.value) + + def __repr__(self): + """ + Returns a string representation of the cell + """ + return f"Cell(value={self.value!r}, row_key={self.row_key!r}, family='{self.family}', qualifier={self.qualifier!r}, timestamp_micros={self.timestamp_micros}, labels={self.labels})" + + """For Bigtable native ordering""" + + def __lt__(self, other) -> bool: + """ + Implements `<` operator + """ + if not isinstance(other, Cell): + return NotImplemented + this_ordering = ( + self.family, + self.qualifier, + -self.timestamp_micros, + self.value, + self.labels, + ) + other_ordering = ( + other.family, + other.qualifier, + -other.timestamp_micros, + other.value, + other.labels, + ) + return this_ordering < other_ordering + + def __eq__(self, other) -> bool: + """ + Implements `==` operator + """ + if not isinstance(other, Cell): + return NotImplemented + return ( + self.row_key == other.row_key + and self.family == other.family + and self.qualifier == other.qualifier + and self.value == other.value + and self.timestamp_micros == other.timestamp_micros + and len(self.labels) == len(other.labels) + and all([label in other.labels for label in self.labels]) + ) + + def __ne__(self, other) -> bool: + """ + Implements `!=` operator + """ + return not self == other + + def __hash__(self): + """ + Implements `hash()` function to fingerprint cell + """ + return hash( + ( + self.row_key, + self.family, + self.qualifier, + self.value, + self.timestamp_micros, + tuple(self.labels), + ) + ) diff --git a/google/cloud/bigtable/deprecated/row_filters.py b/google/cloud/bigtable/data/row_filters.py similarity index 58% rename from google/cloud/bigtable/deprecated/row_filters.py rename to google/cloud/bigtable/data/row_filters.py index 53192acc8..b2fae6971 100644 --- a/google/cloud/bigtable/deprecated/row_filters.py +++ b/google/cloud/bigtable/data/row_filters.py @@ -13,18 +13,25 @@ # limitations under the License. """Filters for Google Cloud Bigtable Row classes.""" +from __future__ import annotations import struct +from typing import Any, Sequence, TYPE_CHECKING, overload +from abc import ABC, abstractmethod from google.cloud._helpers import _microseconds_from_datetime # type: ignore from google.cloud._helpers import _to_bytes # type: ignore from google.cloud.bigtable_v2.types import data as data_v2_pb2 +if TYPE_CHECKING: + # import dependencies when type checking + from datetime import datetime + _PACK_I64 = struct.Struct(">q").pack -class RowFilter(object): +class RowFilter(ABC): """Basic filter to apply to cells in a row. These values can be combined via :class:`RowFilterChain`, @@ -35,15 +42,30 @@ class RowFilter(object): This class is a do-nothing base class for all row filters. """ + def _to_pb(self) -> data_v2_pb2.RowFilter: + """Converts the row filter to a protobuf. + + Returns: The converted current object. + """ + return data_v2_pb2.RowFilter(**self.to_dict()) + + @abstractmethod + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + pass + + def __repr__(self) -> str: + return f"{self.__class__.__name__}()" + -class _BoolFilter(RowFilter): +class _BoolFilter(RowFilter, ABC): """Row filter that uses a boolean flag. :type flag: bool :param flag: An indicator if a setting is turned on or off. """ - def __init__(self, flag): + def __init__(self, flag: bool): self.flag = flag def __eq__(self, other): @@ -54,6 +76,9 @@ def __eq__(self, other): def __ne__(self, other): return not self == other + def __repr__(self) -> str: + return f"{self.__class__.__name__}(flag={self.flag})" + class SinkFilter(_BoolFilter): """Advanced row filter to skip parent filters. @@ -66,13 +91,9 @@ class SinkFilter(_BoolFilter): of a :class:`ConditionalRowFilter`. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(sink=self.flag) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"sink": self.flag} class PassAllFilter(_BoolFilter): @@ -84,13 +105,9 @@ class PassAllFilter(_BoolFilter): completeness. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(pass_all_filter=self.flag) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"pass_all_filter": self.flag} class BlockAllFilter(_BoolFilter): @@ -101,16 +118,12 @@ class BlockAllFilter(_BoolFilter): temporarily disabling just part of a filter. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(block_all_filter=self.flag) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"block_all_filter": self.flag} -class _RegexFilter(RowFilter): +class _RegexFilter(RowFilter, ABC): """Row filter that uses a regular expression. The ``regex`` must be valid RE2 patterns. See Google's @@ -124,8 +137,8 @@ class _RegexFilter(RowFilter): will be encoded as ASCII. """ - def __init__(self, regex): - self.regex = _to_bytes(regex) + def __init__(self, regex: str | bytes): + self.regex: bytes = _to_bytes(regex) def __eq__(self, other): if not isinstance(other, self.__class__): @@ -135,6 +148,9 @@ def __eq__(self, other): def __ne__(self, other): return not self == other + def __repr__(self) -> str: + return f"{self.__class__.__name__}(regex={self.regex!r})" + class RowKeyRegexFilter(_RegexFilter): """Row filter for a row key regular expression. @@ -159,13 +175,9 @@ class RowKeyRegexFilter(_RegexFilter): since the row key is already specified. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(row_key_regex_filter=self.regex) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"row_key_regex_filter": self.regex} class RowSampleFilter(RowFilter): @@ -176,8 +188,8 @@ class RowSampleFilter(RowFilter): interval ``(0, 1)`` The end points are excluded). """ - def __init__(self, sample): - self.sample = sample + def __init__(self, sample: float): + self.sample: float = sample def __eq__(self, other): if not isinstance(other, self.__class__): @@ -187,13 +199,12 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): - """Converts the row filter to a protobuf. + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"row_sample_filter": self.sample} - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(row_sample_filter=self.sample) + def __repr__(self) -> str: + return f"{self.__class__.__name__}(sample={self.sample})" class FamilyNameRegexFilter(_RegexFilter): @@ -211,13 +222,9 @@ class FamilyNameRegexFilter(_RegexFilter): used as a literal. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(family_name_regex_filter=self.regex) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"family_name_regex_filter": self.regex} class ColumnQualifierRegexFilter(_RegexFilter): @@ -241,13 +248,9 @@ class ColumnQualifierRegexFilter(_RegexFilter): match this regex (irrespective of column family). """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(column_qualifier_regex_filter=self.regex) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"column_qualifier_regex_filter": self.regex} class TimestampRange(object): @@ -262,9 +265,9 @@ class TimestampRange(object): range. If omitted, no upper bound is used. """ - def __init__(self, start=None, end=None): - self.start = start - self.end = end + def __init__(self, start: "datetime" | None = None, end: "datetime" | None = None): + self.start: "datetime" | None = start + self.end: "datetime" | None = end def __eq__(self, other): if not isinstance(other, self.__class__): @@ -274,23 +277,29 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.TimestampRange: """Converts the :class:`TimestampRange` to a protobuf. - :rtype: :class:`.data_v2_pb2.TimestampRange` - :returns: The converted current object. + Returns: The converted current object. """ + return data_v2_pb2.TimestampRange(**self.to_dict()) + + def to_dict(self) -> dict[str, int]: + """Converts the timestamp range to a dict representation.""" timestamp_range_kwargs = {} if self.start is not None: - timestamp_range_kwargs["start_timestamp_micros"] = ( - _microseconds_from_datetime(self.start) // 1000 * 1000 - ) + start_time = _microseconds_from_datetime(self.start) // 1000 * 1000 + timestamp_range_kwargs["start_timestamp_micros"] = start_time if self.end is not None: end_time = _microseconds_from_datetime(self.end) if end_time % 1000 != 0: + # if not a whole milisecond value, round up end_time = end_time // 1000 * 1000 + 1000 timestamp_range_kwargs["end_timestamp_micros"] = end_time - return data_v2_pb2.TimestampRange(**timestamp_range_kwargs) + return timestamp_range_kwargs + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(start={self.start}, end={self.end})" class TimestampRangeFilter(RowFilter): @@ -300,8 +309,8 @@ class TimestampRangeFilter(RowFilter): :param range_: Range of time that cells should match against. """ - def __init__(self, range_): - self.range_ = range_ + def __init__(self, start: "datetime" | None = None, end: "datetime" | None = None): + self.range_: TimestampRange = TimestampRange(start, end) def __eq__(self, other): if not isinstance(other, self.__class__): @@ -311,16 +320,22 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. First converts the ``range_`` on the current object to a protobuf and then uses it in the ``timestamp_range_filter`` field. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ - return data_v2_pb2.RowFilter(timestamp_range_filter=self.range_.to_pb()) + return data_v2_pb2.RowFilter(timestamp_range_filter=self.range_._to_pb()) + + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"timestamp_range_filter": self.range_.to_dict()} + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(start={self.range_.start!r}, end={self.range_.end!r})" class ColumnRangeFilter(RowFilter): @@ -330,71 +345,72 @@ class ColumnRangeFilter(RowFilter): By default, we include them both, but this can be changed with optional flags. - :type column_family_id: str - :param column_family_id: The column family that contains the columns. Must + :type family_id: str + :param family_id: The column family that contains the columns. Must be of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - :type start_column: bytes - :param start_column: The start of the range of columns. If no value is + :type start_qualifier: bytes + :param start_qualifier: The start of the range of columns. If no value is used, the backend applies no upper bound to the values. - :type end_column: bytes - :param end_column: The end of the range of columns. If no value is used, + :type end_qualifier: bytes + :param end_qualifier: The end of the range of columns. If no value is used, the backend applies no upper bound to the values. :type inclusive_start: bool :param inclusive_start: Boolean indicating if the start column should be included in the range (or excluded). Defaults - to :data:`True` if ``start_column`` is passed and + to :data:`True` if ``start_qualifier`` is passed and no ``inclusive_start`` was given. :type inclusive_end: bool :param inclusive_end: Boolean indicating if the end column should be included in the range (or excluded). Defaults - to :data:`True` if ``end_column`` is passed and + to :data:`True` if ``end_qualifier`` is passed and no ``inclusive_end`` was given. :raises: :class:`ValueError ` if ``inclusive_start`` - is set but no ``start_column`` is given or if ``inclusive_end`` - is set but no ``end_column`` is given + is set but no ``start_qualifier`` is given or if ``inclusive_end`` + is set but no ``end_qualifier`` is given """ def __init__( self, - column_family_id, - start_column=None, - end_column=None, - inclusive_start=None, - inclusive_end=None, + family_id: str, + start_qualifier: bytes | None = None, + end_qualifier: bytes | None = None, + inclusive_start: bool | None = None, + inclusive_end: bool | None = None, ): - self.column_family_id = column_family_id - if inclusive_start is None: inclusive_start = True - elif start_column is None: + elif start_qualifier is None: raise ValueError( - "Inclusive start was specified but no " "start column was given." + "inclusive_start was specified but no start_qualifier was given." ) - self.start_column = start_column - self.inclusive_start = inclusive_start - if inclusive_end is None: inclusive_end = True - elif end_column is None: + elif end_qualifier is None: raise ValueError( - "Inclusive end was specified but no " "end column was given." + "inclusive_end was specified but no end_qualifier was given." ) - self.end_column = end_column + + self.family_id = family_id + + self.start_qualifier = start_qualifier + self.inclusive_start = inclusive_start + + self.end_qualifier = end_qualifier self.inclusive_end = inclusive_end def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented return ( - other.column_family_id == self.column_family_id - and other.start_column == self.start_column - and other.end_column == self.end_column + other.family_id == self.family_id + and other.start_qualifier == self.start_qualifier + and other.end_qualifier == self.end_qualifier and other.inclusive_start == self.inclusive_start and other.inclusive_end == self.inclusive_end ) @@ -402,31 +418,41 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. First converts to a :class:`.data_v2_pb2.ColumnRange` and then uses it in the ``column_range_filter`` field. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ - column_range_kwargs = {"family_name": self.column_family_id} - if self.start_column is not None: + column_range = data_v2_pb2.ColumnRange(**self.range_to_dict()) + return data_v2_pb2.RowFilter(column_range_filter=column_range) + + def range_to_dict(self) -> dict[str, str | bytes]: + """Converts the column range range to a dict representation.""" + column_range_kwargs: dict[str, str | bytes] = {} + column_range_kwargs["family_name"] = self.family_id + if self.start_qualifier is not None: if self.inclusive_start: key = "start_qualifier_closed" else: key = "start_qualifier_open" - column_range_kwargs[key] = _to_bytes(self.start_column) - if self.end_column is not None: + column_range_kwargs[key] = _to_bytes(self.start_qualifier) + if self.end_qualifier is not None: if self.inclusive_end: key = "end_qualifier_closed" else: key = "end_qualifier_open" - column_range_kwargs[key] = _to_bytes(self.end_column) + column_range_kwargs[key] = _to_bytes(self.end_qualifier) + return column_range_kwargs - column_range = data_v2_pb2.ColumnRange(**column_range_kwargs) - return data_v2_pb2.RowFilter(column_range_filter=column_range) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"column_range_filter": self.range_to_dict()} + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(family_id='{self.family_id}', start_qualifier={self.start_qualifier!r}, end_qualifier={self.end_qualifier!r}, inclusive_start={self.inclusive_start}, inclusive_end={self.inclusive_end})" class ValueRegexFilter(_RegexFilter): @@ -450,29 +476,64 @@ class ValueRegexFilter(_RegexFilter): match this regex. String values will be encoded as ASCII. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(value_regex_filter=self.regex) + def to_dict(self) -> dict[str, bytes]: + """Converts the row filter to a dict representation.""" + return {"value_regex_filter": self.regex} -class ExactValueFilter(ValueRegexFilter): +class LiteralValueFilter(ValueRegexFilter): """Row filter for an exact value. :type value: bytes or str or int :param value: - a literal string encodable as ASCII, or the - equivalent bytes, or an integer (which will be packed into 8-bytes). + a literal string, integer, or the equivalent bytes. + Integer values will be packed into signed 8-bytes. """ - def __init__(self, value): + def __init__(self, value: bytes | str | int): if isinstance(value, int): value = _PACK_I64(value) - super(ExactValueFilter, self).__init__(value) + elif isinstance(value, str): + value = value.encode("utf-8") + value = self._write_literal_regex(value) + super(LiteralValueFilter, self).__init__(value) + + @staticmethod + def _write_literal_regex(input_bytes: bytes) -> bytes: + """ + Escape re2 special characters from literal bytes. + + Extracted from: re2 QuoteMeta: + https://github.com/google/re2/blob/70f66454c255080a54a8da806c52d1f618707f8a/re2/re2.cc#L456 + """ + result = bytearray() + for byte in input_bytes: + # If this is the part of a UTF8 or Latin1 character, we need \ + # to copy this byte without escaping. Experimentally this is \ + # what works correctly with the regexp library. \ + utf8_latin1_check = (byte & 128) == 0 + if ( + (byte < ord("a") or byte > ord("z")) + and (byte < ord("A") or byte > ord("Z")) + and (byte < ord("0") or byte > ord("9")) + and byte != ord("_") + and utf8_latin1_check + ): + if byte == 0: + # Special handling for null chars. + # Note that this special handling is not strictly required for RE2, + # but this quoting is required for other regexp libraries such as + # PCRE. + # Can't use "\\0" since the next character might be a digit. + result.extend([ord("\\"), ord("x"), ord("0"), ord("0")]) + continue + result.append(ord(b"\\")) + result.append(byte) + return bytes(result) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(value={self.regex!r})" class ValueRangeFilter(RowFilter): @@ -510,25 +571,29 @@ class ValueRangeFilter(RowFilter): """ def __init__( - self, start_value=None, end_value=None, inclusive_start=None, inclusive_end=None + self, + start_value: bytes | int | None = None, + end_value: bytes | int | None = None, + inclusive_start: bool | None = None, + inclusive_end: bool | None = None, ): if inclusive_start is None: inclusive_start = True elif start_value is None: raise ValueError( - "Inclusive start was specified but no " "start value was given." + "inclusive_start was specified but no start_value was given." ) - if isinstance(start_value, int): - start_value = _PACK_I64(start_value) - self.start_value = start_value - self.inclusive_start = inclusive_start - if inclusive_end is None: inclusive_end = True elif end_value is None: raise ValueError( - "Inclusive end was specified but no " "end value was given." + "inclusive_end was specified but no end_qualifier was given." ) + if isinstance(start_value, int): + start_value = _PACK_I64(start_value) + self.start_value = start_value + self.inclusive_start = inclusive_start + if isinstance(end_value, int): end_value = _PACK_I64(end_value) self.end_value = end_value @@ -547,15 +612,19 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. First converts to a :class:`.data_v2_pb2.ValueRange` and then uses it to create a row filter protobuf. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ + value_range = data_v2_pb2.ValueRange(**self.range_to_dict()) + return data_v2_pb2.RowFilter(value_range_filter=value_range) + + def range_to_dict(self) -> dict[str, bytes]: + """Converts the value range range to a dict representation.""" value_range_kwargs = {} if self.start_value is not None: if self.inclusive_start: @@ -569,12 +638,17 @@ def to_pb(self): else: key = "end_value_open" value_range_kwargs[key] = _to_bytes(self.end_value) + return value_range_kwargs - value_range = data_v2_pb2.ValueRange(**value_range_kwargs) - return data_v2_pb2.RowFilter(value_range_filter=value_range) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"value_range_filter": self.range_to_dict()} + def __repr__(self) -> str: + return f"{self.__class__.__name__}(start_value={self.start_value!r}, end_value={self.end_value!r}, inclusive_start={self.inclusive_start}, inclusive_end={self.inclusive_end})" -class _CellCountFilter(RowFilter): + +class _CellCountFilter(RowFilter, ABC): """Row filter that uses an integer count of cells. The cell count is used as an offset or a limit for the number @@ -584,7 +658,7 @@ class _CellCountFilter(RowFilter): :param num_cells: An integer count / offset / limit. """ - def __init__(self, num_cells): + def __init__(self, num_cells: int): self.num_cells = num_cells def __eq__(self, other): @@ -595,6 +669,9 @@ def __eq__(self, other): def __ne__(self, other): return not self == other + def __repr__(self) -> str: + return f"{self.__class__.__name__}(num_cells={self.num_cells})" + class CellsRowOffsetFilter(_CellCountFilter): """Row filter to skip cells in a row. @@ -603,13 +680,9 @@ class CellsRowOffsetFilter(_CellCountFilter): :param num_cells: Skips the first N cells of the row. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(cells_per_row_offset_filter=self.num_cells) + def to_dict(self) -> dict[str, int]: + """Converts the row filter to a dict representation.""" + return {"cells_per_row_offset_filter": self.num_cells} class CellsRowLimitFilter(_CellCountFilter): @@ -619,13 +692,9 @@ class CellsRowLimitFilter(_CellCountFilter): :param num_cells: Matches only the first N cells of the row. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(cells_per_row_limit_filter=self.num_cells) + def to_dict(self) -> dict[str, int]: + """Converts the row filter to a dict representation.""" + return {"cells_per_row_limit_filter": self.num_cells} class CellsColumnLimitFilter(_CellCountFilter): @@ -637,13 +706,9 @@ class CellsColumnLimitFilter(_CellCountFilter): timestamps of each cell. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(cells_per_column_limit_filter=self.num_cells) + def to_dict(self) -> dict[str, int]: + """Converts the row filter to a dict representation.""" + return {"cells_per_column_limit_filter": self.num_cells} class StripValueTransformerFilter(_BoolFilter): @@ -655,13 +720,9 @@ class StripValueTransformerFilter(_BoolFilter): transformer than a generic query / filter. """ - def to_pb(self): - """Converts the row filter to a protobuf. - - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(strip_value_transformer=self.flag) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"strip_value_transformer": self.flag} class ApplyLabelFilter(RowFilter): @@ -683,7 +744,7 @@ class ApplyLabelFilter(RowFilter): ``[a-z0-9\\-]+``. """ - def __init__(self, label): + def __init__(self, label: str): self.label = label def __eq__(self, other): @@ -694,16 +755,15 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): - """Converts the row filter to a protobuf. + def to_dict(self) -> dict[str, str]: + """Converts the row filter to a dict representation.""" + return {"apply_label_transformer": self.label} - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. - """ - return data_v2_pb2.RowFilter(apply_label_transformer=self.label) + def __repr__(self) -> str: + return f"{self.__class__.__name__}(label={self.label})" -class _FilterCombination(RowFilter): +class _FilterCombination(RowFilter, Sequence[RowFilter], ABC): """Chain of row filters. Sends rows through several filters in sequence. The filters are "chained" @@ -714,10 +774,10 @@ class _FilterCombination(RowFilter): :param filters: List of :class:`RowFilter` """ - def __init__(self, filters=None): + def __init__(self, filters: list[RowFilter] | None = None): if filters is None: filters = [] - self.filters = filters + self.filters: list[RowFilter] = filters def __eq__(self, other): if not isinstance(other, self.__class__): @@ -727,6 +787,38 @@ def __eq__(self, other): def __ne__(self, other): return not self == other + def __len__(self) -> int: + return len(self.filters) + + @overload + def __getitem__(self, index: int) -> RowFilter: + # overload signature for type checking + pass + + @overload + def __getitem__(self, index: slice) -> list[RowFilter]: + # overload signature for type checking + pass + + def __getitem__(self, index): + return self.filters[index] + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(filters={self.filters})" + + def __str__(self) -> str: + """ + Returns a string representation of the filter chain. + + Adds line breaks between each sub-filter for readability. + """ + output = [f"{self.__class__.__name__}(["] + for filter_ in self.filters: + filter_lines = f"{filter_},".splitlines() + output.extend([f" {line}" for line in filter_lines]) + output.append("])") + return "\n".join(output) + class RowFilterChain(_FilterCombination): """Chain of row filters. @@ -739,17 +831,20 @@ class RowFilterChain(_FilterCombination): :param filters: List of :class:`RowFilter` """ - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ chain = data_v2_pb2.RowFilter.Chain( - filters=[row_filter.to_pb() for row_filter in self.filters] + filters=[row_filter._to_pb() for row_filter in self.filters] ) return data_v2_pb2.RowFilter(chain=chain) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"chain": {"filters": [f.to_dict() for f in self.filters]}} + class RowFilterUnion(_FilterCombination): """Union of row filters. @@ -764,50 +859,58 @@ class RowFilterUnion(_FilterCombination): :param filters: List of :class:`RowFilter` """ - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ interleave = data_v2_pb2.RowFilter.Interleave( - filters=[row_filter.to_pb() for row_filter in self.filters] + filters=[row_filter._to_pb() for row_filter in self.filters] ) return data_v2_pb2.RowFilter(interleave=interleave) + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"interleave": {"filters": [f.to_dict() for f in self.filters]}} + class ConditionalRowFilter(RowFilter): """Conditional row filter which exhibits ternary behavior. - Executes one of two filters based on another filter. If the ``base_filter`` + Executes one of two filters based on another filter. If the ``predicate_filter`` returns any cells in the row, then ``true_filter`` is executed. If not, then ``false_filter`` is executed. .. note:: - The ``base_filter`` does not execute atomically with the true and false + The ``predicate_filter`` does not execute atomically with the true and false filters, which may lead to inconsistent or unexpected results. Additionally, executing a :class:`ConditionalRowFilter` has poor performance on the server, especially when ``false_filter`` is set. - :type base_filter: :class:`RowFilter` - :param base_filter: The filter to condition on before executing the + :type predicate_filter: :class:`RowFilter` + :param predicate_filter: The filter to condition on before executing the true/false filters. :type true_filter: :class:`RowFilter` :param true_filter: (Optional) The filter to execute if there are any cells - matching ``base_filter``. If not provided, no results + matching ``predicate_filter``. If not provided, no results will be returned in the true case. :type false_filter: :class:`RowFilter` :param false_filter: (Optional) The filter to execute if there are no cells - matching ``base_filter``. If not provided, no results + matching ``predicate_filter``. If not provided, no results will be returned in the false case. """ - def __init__(self, base_filter, true_filter=None, false_filter=None): - self.base_filter = base_filter + def __init__( + self, + predicate_filter: RowFilter, + true_filter: RowFilter | None = None, + false_filter: RowFilter | None = None, + ): + self.predicate_filter = predicate_filter self.true_filter = true_filter self.false_filter = false_filter @@ -815,7 +918,7 @@ def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented return ( - other.base_filter == self.base_filter + other.predicate_filter == self.predicate_filter and other.true_filter == self.true_filter and other.false_filter == self.false_filter ) @@ -823,16 +926,43 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_pb(self): + def _to_pb(self) -> data_v2_pb2.RowFilter: """Converts the row filter to a protobuf. - :rtype: :class:`.data_v2_pb2.RowFilter` - :returns: The converted current object. + Returns: The converted current object. """ - condition_kwargs = {"predicate_filter": self.base_filter.to_pb()} + condition_kwargs = {"predicate_filter": self.predicate_filter._to_pb()} if self.true_filter is not None: - condition_kwargs["true_filter"] = self.true_filter.to_pb() + condition_kwargs["true_filter"] = self.true_filter._to_pb() if self.false_filter is not None: - condition_kwargs["false_filter"] = self.false_filter.to_pb() + condition_kwargs["false_filter"] = self.false_filter._to_pb() condition = data_v2_pb2.RowFilter.Condition(**condition_kwargs) return data_v2_pb2.RowFilter(condition=condition) + + def condition_to_dict(self) -> dict[str, Any]: + """Converts the condition to a dict representation.""" + condition_kwargs = {"predicate_filter": self.predicate_filter.to_dict()} + if self.true_filter is not None: + condition_kwargs["true_filter"] = self.true_filter.to_dict() + if self.false_filter is not None: + condition_kwargs["false_filter"] = self.false_filter.to_dict() + return condition_kwargs + + def to_dict(self) -> dict[str, Any]: + """Converts the row filter to a dict representation.""" + return {"condition": self.condition_to_dict()} + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(predicate_filter={self.predicate_filter!r}, true_filter={self.true_filter!r}, false_filter={self.false_filter!r})" + + def __str__(self) -> str: + output = [f"{self.__class__.__name__}("] + for filter_type in ("predicate_filter", "true_filter", "false_filter"): + filter_ = getattr(self, filter_type) + if filter_ is None: + continue + # add the new filter set, adding indentations for readability + filter_lines = f"{filter_type}={filter_},".splitlines() + output.extend(f" {line}" for line in filter_lines) + output.append(")") + return "\n".join(output) diff --git a/google/cloud/bigtable/deprecated/batcher.py b/google/cloud/bigtable/deprecated/batcher.py deleted file mode 100644 index 58cf6b6e3..000000000 --- a/google/cloud/bigtable/deprecated/batcher.py +++ /dev/null @@ -1,146 +0,0 @@ -# Copyright 2018 Google LLC -# -# Licensed 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. - -"""User friendly container for Google Cloud Bigtable MutationBatcher.""" - - -FLUSH_COUNT = 1000 -MAX_MUTATIONS = 100000 -MAX_ROW_BYTES = 5242880 # 5MB - - -class MaxMutationsError(ValueError): - """The number of mutations for bulk request is too big.""" - - -class MutationsBatcher(object): - """A MutationsBatcher is used in batch cases where the number of mutations - is large or unknown. It will store DirectRows in memory until one of the - size limits is reached, or an explicit call to flush() is performed. When - a flush event occurs, the DirectRows in memory will be sent to Cloud - Bigtable. Batching mutations is more efficient than sending individual - request. - - This class is not suited for usage in systems where each mutation - must be guaranteed to be sent, since calling mutate may only result in an - in-memory change. In a case of a system crash, any DirectRows remaining in - memory will not necessarily be sent to the service, even after the - completion of the mutate() method. - - TODO: Performance would dramatically improve if this class had the - capability of asynchronous, parallel RPCs. - - :type table: class - :param table: class:`~google.cloud.bigtable.deprecated.table.Table`. - - :type flush_count: int - :param flush_count: (Optional) Max number of rows to flush. If it - reaches the max number of rows it calls finish_batch() to mutate the - current row batch. Default is FLUSH_COUNT (1000 rows). - - :type max_row_bytes: int - :param max_row_bytes: (Optional) Max number of row mutations size to - flush. If it reaches the max number of row mutations size it calls - finish_batch() to mutate the current row batch. Default is MAX_ROW_BYTES - (5 MB). - """ - - def __init__(self, table, flush_count=FLUSH_COUNT, max_row_bytes=MAX_ROW_BYTES): - self.rows = [] - self.total_mutation_count = 0 - self.total_size = 0 - self.table = table - self.flush_count = flush_count - self.max_row_bytes = max_row_bytes - - def mutate(self, row): - """Add a row to the batch. If the current batch meets one of the size - limits, the batch is sent synchronously. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_batcher_mutate] - :end-before: [END bigtable_api_batcher_mutate] - :dedent: 4 - - :type row: class - :param row: class:`~google.cloud.bigtable.deprecated.row.DirectRow`. - - :raises: One of the following: - * :exc:`~.table._BigtableRetryableError` if any - row returned a transient error. - * :exc:`RuntimeError` if the number of responses doesn't - match the number of rows that were retried - * :exc:`.batcher.MaxMutationsError` if any row exceeds max - mutations count. - """ - mutation_count = len(row._get_mutations()) - if mutation_count > MAX_MUTATIONS: - raise MaxMutationsError( - "The row key {} exceeds the number of mutations {}.".format( - row.row_key, mutation_count - ) - ) - - if (self.total_mutation_count + mutation_count) >= MAX_MUTATIONS: - self.flush() - - self.rows.append(row) - self.total_mutation_count += mutation_count - self.total_size += row.get_mutations_size() - - if self.total_size >= self.max_row_bytes or len(self.rows) >= self.flush_count: - self.flush() - - def mutate_rows(self, rows): - """Add multiple rows to the batch. If the current batch meets one of the size - limits, the batch is sent synchronously. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_batcher_mutate_rows] - :end-before: [END bigtable_api_batcher_mutate_rows] - :dedent: 4 - - :type rows: list:[`~google.cloud.bigtable.deprecated.row.DirectRow`] - :param rows: list:[`~google.cloud.bigtable.deprecated.row.DirectRow`]. - - :raises: One of the following: - * :exc:`~.table._BigtableRetryableError` if any - row returned a transient error. - * :exc:`RuntimeError` if the number of responses doesn't - match the number of rows that were retried - * :exc:`.batcher.MaxMutationsError` if any row exceeds max - mutations count. - """ - for row in rows: - self.mutate(row) - - def flush(self): - """Sends the current. batch to Cloud Bigtable. - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_batcher_flush] - :end-before: [END bigtable_api_batcher_flush] - :dedent: 4 - - """ - if len(self.rows) != 0: - self.table.mutate_rows(self.rows) - self.total_mutation_count = 0 - self.total_size = 0 - self.rows = [] diff --git a/google/cloud/bigtable/deprecated/client.py b/google/cloud/bigtable/deprecated/client.py deleted file mode 100644 index c13e5f0da..000000000 --- a/google/cloud/bigtable/deprecated/client.py +++ /dev/null @@ -1,521 +0,0 @@ -# Copyright 2015 Google LLC -# -# Licensed 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. - -"""Parent client for calling the Google Cloud Bigtable API. - -This is the base from which all interactions with the API occur. - -In the hierarchy of API concepts - -* a :class:`~google.cloud.bigtable.deprecated.client.Client` owns an - :class:`~google.cloud.bigtable.deprecated.instance.Instance` -* an :class:`~google.cloud.bigtable.deprecated.instance.Instance` owns a - :class:`~google.cloud.bigtable.deprecated.table.Table` -* a :class:`~google.cloud.bigtable.deprecated.table.Table` owns a - :class:`~.column_family.ColumnFamily` -* a :class:`~google.cloud.bigtable.deprecated.table.Table` owns a - :class:`~google.cloud.bigtable.deprecated.row.Row` (and all the cells in the row) -""" -import os -import warnings -import grpc # type: ignore - -from google.api_core.gapic_v1 import client_info as client_info_lib -import google.auth # type: ignore -from google.auth.credentials import AnonymousCredentials # type: ignore - -from google.cloud import bigtable_v2 -from google.cloud import bigtable_admin_v2 -from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport -from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin.transports import ( - BigtableInstanceAdminGrpcTransport, -) -from google.cloud.bigtable_admin_v2.services.bigtable_table_admin.transports import ( - BigtableTableAdminGrpcTransport, -) - -from google.cloud import bigtable -from google.cloud.bigtable.deprecated.instance import Instance -from google.cloud.bigtable.deprecated.cluster import Cluster - -from google.cloud.client import ClientWithProject # type: ignore - -from google.cloud.bigtable_admin_v2.types import instance -from google.cloud.bigtable.deprecated.cluster import _CLUSTER_NAME_RE -from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore - - -INSTANCE_TYPE_PRODUCTION = instance.Instance.Type.PRODUCTION -INSTANCE_TYPE_DEVELOPMENT = instance.Instance.Type.DEVELOPMENT -INSTANCE_TYPE_UNSPECIFIED = instance.Instance.Type.TYPE_UNSPECIFIED -SPANNER_ADMIN_SCOPE = "https://www.googleapis.com/auth/spanner.admin" -ADMIN_SCOPE = "https://www.googleapis.com/auth/bigtable.admin" -"""Scope for interacting with the Cluster Admin and Table Admin APIs.""" -DATA_SCOPE = "https://www.googleapis.com/auth/bigtable.data" -"""Scope for reading and writing table data.""" -READ_ONLY_SCOPE = "https://www.googleapis.com/auth/bigtable.data.readonly" -"""Scope for reading table data.""" - -_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator" -_GRPC_CHANNEL_OPTIONS = ( - ("grpc.max_send_message_length", -1), - ("grpc.max_receive_message_length", -1), - ("grpc.keepalive_time_ms", 30000), - ("grpc.keepalive_timeout_ms", 10000), -) - - -def _create_gapic_client(client_class, client_options=None, transport=None): - def inner(self): - return client_class( - credentials=None, - client_info=self._client_info, - client_options=client_options, - transport=transport, - ) - - return inner - - -class Client(ClientWithProject): - """Client for interacting with Google Cloud Bigtable API. - - DEPRECATED: This class is deprecated and may be removed in a future version - Please use `google.cloud.bigtable.BigtableDataClient` instead. - - .. note:: - - Since the Cloud Bigtable API requires the gRPC transport, no - ``_http`` argument is accepted by this class. - - :type project: :class:`str` or :func:`unicode ` - :param project: (Optional) The ID of the project which owns the - instances, tables and data. If not provided, will - attempt to determine from the environment. - - :type credentials: :class:`~google.auth.credentials.Credentials` - :param credentials: (Optional) The OAuth2 Credentials to use for this - client. If not passed, falls back to the default - inferred from the environment. - - :type read_only: bool - :param read_only: (Optional) Boolean indicating if the data scope should be - for reading only (or for writing as well). Defaults to - :data:`False`. - - :type admin: bool - :param admin: (Optional) Boolean indicating if the client will be used to - interact with the Instance Admin or Table Admin APIs. This - requires the :const:`ADMIN_SCOPE`. Defaults to :data:`False`. - - :type: client_info: :class:`google.api_core.gapic_v1.client_info.ClientInfo` - :param client_info: - The client info used to send a user-agent string along with API - requests. If ``None``, then default info will be used. Generally, - you only need to set this if you're developing your own library - or partner tool. - - :type client_options: :class:`~google.api_core.client_options.ClientOptions` - or :class:`dict` - :param client_options: (Optional) Client options used to set user options - on the client. API Endpoint should be set through client_options. - - :type admin_client_options: - :class:`~google.api_core.client_options.ClientOptions` or :class:`dict` - :param admin_client_options: (Optional) Client options used to set user - options on the client. API Endpoint for admin operations should be set - through admin_client_options. - - :type channel: :instance: grpc.Channel - :param channel (grpc.Channel): (Optional) DEPRECATED: - A ``Channel`` instance through which to make calls. - This argument is mutually exclusive with ``credentials``; - providing both will raise an exception. No longer used. - - :raises: :class:`ValueError ` if both ``read_only`` - and ``admin`` are :data:`True` - """ - - _table_data_client = None - _table_admin_client = None - _instance_admin_client = None - - def __init__( - self, - project=None, - credentials=None, - read_only=False, - admin=False, - client_info=None, - client_options=None, - admin_client_options=None, - channel=None, - ): - warnings.warn( - "'Client' is deprecated. Please use 'google.cloud.bigtable.BigtableDataClient' instead.", - DeprecationWarning, - stacklevel=2, - ) - if client_info is None: - client_info = client_info_lib.ClientInfo( - client_library_version=bigtable.__version__, - ) - if read_only and admin: - raise ValueError( - "A read-only client cannot also perform" "administrative actions." - ) - - # NOTE: We set the scopes **before** calling the parent constructor. - # It **may** use those scopes in ``with_scopes_if_required``. - self._read_only = bool(read_only) - self._admin = bool(admin) - self._client_info = client_info - self._emulator_host = os.getenv(BIGTABLE_EMULATOR) - - if self._emulator_host is not None: - if credentials is None: - credentials = AnonymousCredentials() - if project is None: - project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT - - if channel is not None: - warnings.warn( - "'channel' is deprecated and no longer used.", - DeprecationWarning, - stacklevel=2, - ) - - self._client_options = client_options - self._admin_client_options = admin_client_options - self._channel = channel - self.SCOPE = self._get_scopes() - super(Client, self).__init__( - project=project, - credentials=credentials, - client_options=client_options, - ) - - def _get_scopes(self): - """Get the scopes corresponding to admin / read-only state. - - Returns: - Tuple[str, ...]: The tuple of scopes. - """ - if self._read_only: - scopes = (READ_ONLY_SCOPE,) - else: - scopes = (DATA_SCOPE,) - - if self._admin: - scopes += (ADMIN_SCOPE,) - - return scopes - - def _emulator_channel(self, transport, options): - """Create a channel using self._credentials - - Works in a similar way to ``grpc.secure_channel`` but using - ``grpc.local_channel_credentials`` rather than - ``grpc.ssh_channel_credentials`` to allow easy connection to a - local emulator. - - Returns: - grpc.Channel or grpc.aio.Channel - """ - # TODO: Implement a special credentials type for emulator and use - # "transport.create_channel" to create gRPC channels once google-auth - # extends it's allowed credentials types. - # Note: this code also exists in the firestore client. - if "GrpcAsyncIOTransport" in str(transport.__name__): - return grpc.aio.secure_channel( - self._emulator_host, - self._local_composite_credentials(), - options=options, - ) - else: - return grpc.secure_channel( - self._emulator_host, - self._local_composite_credentials(), - options=options, - ) - - def _local_composite_credentials(self): - """Create credentials for the local emulator channel. - - :return: grpc.ChannelCredentials - """ - credentials = google.auth.credentials.with_scopes_if_required( - self._credentials, None - ) - request = google.auth.transport.requests.Request() - - # Create the metadata plugin for inserting the authorization header. - metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin( - credentials, request - ) - - # Create a set of grpc.CallCredentials using the metadata plugin. - google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin) - - # Using the local_credentials to allow connection to emulator - local_credentials = grpc.local_channel_credentials() - - # Combine the local credentials and the authorization credentials. - return grpc.composite_channel_credentials( - local_credentials, google_auth_credentials - ) - - def _create_gapic_client_channel(self, client_class, grpc_transport): - if self._emulator_host is not None: - api_endpoint = self._emulator_host - elif self._client_options and self._client_options.api_endpoint: - api_endpoint = self._client_options.api_endpoint - else: - api_endpoint = client_class.DEFAULT_ENDPOINT - - if self._emulator_host is not None: - channel = self._emulator_channel( - transport=grpc_transport, - options=_GRPC_CHANNEL_OPTIONS, - ) - else: - channel = grpc_transport.create_channel( - host=api_endpoint, - credentials=self._credentials, - options=_GRPC_CHANNEL_OPTIONS, - ) - return grpc_transport(channel=channel, host=api_endpoint) - - @property - def project_path(self): - """Project name to be used with Instance Admin API. - - .. note:: - - This property will not change if ``project`` does not, but the - return value is not cached. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_project_path] - :end-before: [END bigtable_api_project_path] - :dedent: 4 - - The project name is of the form - - ``"projects/{project}"`` - - :rtype: str - :returns: Return a fully-qualified project string. - """ - return self.instance_admin_client.common_project_path(self.project) - - @property - def table_data_client(self): - """Getter for the gRPC stub used for the Table Admin API. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_table_data_client] - :end-before: [END bigtable_api_table_data_client] - :dedent: 4 - - :rtype: :class:`.bigtable_v2.BigtableClient` - :returns: A BigtableClient object. - """ - if self._table_data_client is None: - transport = self._create_gapic_client_channel( - bigtable_v2.BigtableClient, - BigtableGrpcTransport, - ) - klass = _create_gapic_client( - bigtable_v2.BigtableClient, - client_options=self._client_options, - transport=transport, - ) - self._table_data_client = klass(self) - return self._table_data_client - - @property - def table_admin_client(self): - """Getter for the gRPC stub used for the Table Admin API. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_table_admin_client] - :end-before: [END bigtable_api_table_admin_client] - :dedent: 4 - - :rtype: :class:`.bigtable_admin_pb2.BigtableTableAdmin` - :returns: A BigtableTableAdmin instance. - :raises: :class:`ValueError ` if the current - client is not an admin client or if it has not been - :meth:`start`-ed. - """ - if self._table_admin_client is None: - if not self._admin: - raise ValueError("Client is not an admin client.") - - transport = self._create_gapic_client_channel( - bigtable_admin_v2.BigtableTableAdminClient, - BigtableTableAdminGrpcTransport, - ) - klass = _create_gapic_client( - bigtable_admin_v2.BigtableTableAdminClient, - client_options=self._admin_client_options, - transport=transport, - ) - self._table_admin_client = klass(self) - return self._table_admin_client - - @property - def instance_admin_client(self): - """Getter for the gRPC stub used for the Table Admin API. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_instance_admin_client] - :end-before: [END bigtable_api_instance_admin_client] - :dedent: 4 - - :rtype: :class:`.bigtable_admin_pb2.BigtableInstanceAdmin` - :returns: A BigtableInstanceAdmin instance. - :raises: :class:`ValueError ` if the current - client is not an admin client or if it has not been - :meth:`start`-ed. - """ - if self._instance_admin_client is None: - if not self._admin: - raise ValueError("Client is not an admin client.") - - transport = self._create_gapic_client_channel( - bigtable_admin_v2.BigtableInstanceAdminClient, - BigtableInstanceAdminGrpcTransport, - ) - klass = _create_gapic_client( - bigtable_admin_v2.BigtableInstanceAdminClient, - client_options=self._admin_client_options, - transport=transport, - ) - self._instance_admin_client = klass(self) - return self._instance_admin_client - - def instance(self, instance_id, display_name=None, instance_type=None, labels=None): - """Factory to create a instance associated with this client. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_create_prod_instance] - :end-before: [END bigtable_api_create_prod_instance] - :dedent: 4 - - :type instance_id: str - :param instance_id: The ID of the instance. - - :type display_name: str - :param display_name: (Optional) The display name for the instance in - the Cloud Console UI. (Must be between 4 and 30 - characters.) If this value is not set in the - constructor, will fall back to the instance ID. - - :type instance_type: int - :param instance_type: (Optional) The type of the instance. - Possible values are represented - by the following constants: - :data:`google.cloud.bigtable.deprecated.instance.InstanceType.PRODUCTION`. - :data:`google.cloud.bigtable.deprecated.instance.InstanceType.DEVELOPMENT`, - Defaults to - :data:`google.cloud.bigtable.deprecated.instance.InstanceType.UNSPECIFIED`. - - :type labels: dict - :param labels: (Optional) Labels are a flexible and lightweight - mechanism for organizing cloud resources into groups - that reflect a customer's organizational needs and - deployment strategies. They can be used to filter - resources and aggregate metrics. Label keys must be - between 1 and 63 characters long. Maximum 64 labels can - be associated with a given resource. Label values must - be between 0 and 63 characters long. Keys and values - must both be under 128 bytes. - - :rtype: :class:`~google.cloud.bigtable.deprecated.instance.Instance` - :returns: an instance owned by this client. - """ - return Instance( - instance_id, - self, - display_name=display_name, - instance_type=instance_type, - labels=labels, - ) - - def list_instances(self): - """List instances owned by the project. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_list_instances] - :end-before: [END bigtable_api_list_instances] - :dedent: 4 - - :rtype: tuple - :returns: - (instances, failed_locations), where 'instances' is list of - :class:`google.cloud.bigtable.deprecated.instance.Instance`, and - 'failed_locations' is a list of locations which could not - be resolved. - """ - resp = self.instance_admin_client.list_instances( - request={"parent": self.project_path} - ) - instances = [Instance.from_pb(instance, self) for instance in resp.instances] - return instances, resp.failed_locations - - def list_clusters(self): - """List the clusters in the project. - - For example: - - .. literalinclude:: snippets.py - :start-after: [START bigtable_api_list_clusters_in_project] - :end-before: [END bigtable_api_list_clusters_in_project] - :dedent: 4 - - :rtype: tuple - :returns: - (clusters, failed_locations), where 'clusters' is list of - :class:`google.cloud.bigtable.deprecated.instance.Cluster`, and - 'failed_locations' is a list of strings representing - locations which could not be resolved. - """ - resp = self.instance_admin_client.list_clusters( - request={ - "parent": self.instance_admin_client.instance_path(self.project, "-") - } - ) - clusters = [] - instances = {} - for cluster in resp.clusters: - match_cluster_name = _CLUSTER_NAME_RE.match(cluster.name) - instance_id = match_cluster_name.group("instance") - if instance_id not in instances: - instances[instance_id] = self.instance(instance_id) - clusters.append(Cluster.from_pb(cluster, instances[instance_id])) - return clusters, resp.failed_locations diff --git a/google/cloud/bigtable/deprecated/py.typed b/google/cloud/bigtable/deprecated/py.typed deleted file mode 100644 index 7bd4705d4..000000000 --- a/google/cloud/bigtable/deprecated/py.typed +++ /dev/null @@ -1,2 +0,0 @@ -# Marker file for PEP 561. -# The google-cloud-bigtable package uses inline types. diff --git a/google/cloud/bigtable/deprecated/row.py b/google/cloud/bigtable/deprecated/row.py deleted file mode 100644 index 3b114a74a..000000000 --- a/google/cloud/bigtable/deprecated/row.py +++ /dev/null @@ -1,1267 +0,0 @@ -# Copyright 2015 Google LLC -# -# Licensed 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. - -"""User-friendly container for Google Cloud Bigtable Row.""" - - -import struct - -from google.cloud._helpers import _datetime_from_microseconds # type: ignore -from google.cloud._helpers import _microseconds_from_datetime # type: ignore -from google.cloud._helpers import _to_bytes # type: ignore -from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - -_PACK_I64 = struct.Struct(">q").pack - -MAX_MUTATIONS = 100000 -"""The maximum number of mutations that a row can accumulate.""" - -_MISSING_COLUMN_FAMILY = "Column family {} is not among the cells stored in this row." -_MISSING_COLUMN = ( - "Column {} is not among the cells stored in this row in the column family {}." -) -_MISSING_INDEX = ( - "Index {!r} is not valid for the cells stored in this row for column {} " - "in the column family {}. There are {} such cells." -) - - -class Row(object): - """Base representation of a Google Cloud Bigtable Row. - - This class has three subclasses corresponding to the three - RPC methods for sending row mutations: - - * :class:`DirectRow` for ``MutateRow`` - * :class:`ConditionalRow` for ``CheckAndMutateRow`` - * :class:`AppendRow` for ``ReadModifyWriteRow`` - - :type row_key: bytes - :param row_key: The key for the current row. - - :type table: :class:`Table ` - :param table: (Optional) The table that owns the row. - """ - - def __init__(self, row_key, table=None): - self._row_key = _to_bytes(row_key) - self._table = table - - @property - def row_key(self): - """Row key. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_row_key] - :end-before: [END bigtable_api_row_row_key] - :dedent: 4 - - :rtype: bytes - :returns: The key for the current row. - """ - return self._row_key - - @property - def table(self): - """Row table. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_table] - :end-before: [END bigtable_api_row_table] - :dedent: 4 - - :rtype: table: :class:`Table ` - :returns: table: The table that owns the row. - """ - return self._table - - -class _SetDeleteRow(Row): - """Row helper for setting or deleting cell values. - - Implements helper methods to add mutations to set or delete cell contents: - - * :meth:`set_cell` - * :meth:`delete` - * :meth:`delete_cell` - * :meth:`delete_cells` - - :type row_key: bytes - :param row_key: The key for the current row. - - :type table: :class:`Table ` - :param table: The table that owns the row. - """ - - ALL_COLUMNS = object() - """Sentinel value used to indicate all columns in a column family.""" - - def _get_mutations(self, state=None): - """Gets the list of mutations for a given state. - - This method intended to be implemented by subclasses. - - ``state`` may not need to be used by all subclasses. - - :type state: bool - :param state: The state that the mutation should be - applied in. - - :raises: :class:`NotImplementedError ` - always. - """ - raise NotImplementedError - - def _set_cell(self, column_family_id, column, value, timestamp=None, state=None): - """Helper for :meth:`set_cell` - - Adds a mutation to set the value in a specific cell. - - ``state`` is unused by :class:`DirectRow` but is used by - subclasses. - - :type column_family_id: str - :param column_family_id: The column family that contains the column. - Must be of the form - ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type column: bytes - :param column: The column within the column family where the cell - is located. - - :type value: bytes or :class:`int` - :param value: The value to set in the cell. If an integer is used, - will be interpreted as a 64-bit big-endian signed - integer (8 bytes). - - :type timestamp: :class:`datetime.datetime` - :param timestamp: (Optional) The timestamp of the operation. - - :type state: bool - :param state: (Optional) The state that is passed along to - :meth:`_get_mutations`. - """ - column = _to_bytes(column) - if isinstance(value, int): - value = _PACK_I64(value) - value = _to_bytes(value) - if timestamp is None: - # Use -1 for current Bigtable server time. - timestamp_micros = -1 - else: - timestamp_micros = _microseconds_from_datetime(timestamp) - # Truncate to millisecond granularity. - timestamp_micros -= timestamp_micros % 1000 - - mutation_val = data_v2_pb2.Mutation.SetCell( - family_name=column_family_id, - column_qualifier=column, - timestamp_micros=timestamp_micros, - value=value, - ) - mutation_pb = data_v2_pb2.Mutation(set_cell=mutation_val) - self._get_mutations(state).append(mutation_pb) - - def _delete(self, state=None): - """Helper for :meth:`delete` - - Adds a delete mutation (for the entire row) to the accumulated - mutations. - - ``state`` is unused by :class:`DirectRow` but is used by - subclasses. - - :type state: bool - :param state: (Optional) The state that is passed along to - :meth:`_get_mutations`. - """ - mutation_val = data_v2_pb2.Mutation.DeleteFromRow() - mutation_pb = data_v2_pb2.Mutation(delete_from_row=mutation_val) - self._get_mutations(state).append(mutation_pb) - - def _delete_cells(self, column_family_id, columns, time_range=None, state=None): - """Helper for :meth:`delete_cell` and :meth:`delete_cells`. - - ``state`` is unused by :class:`DirectRow` but is used by - subclasses. - - :type column_family_id: str - :param column_family_id: The column family that contains the column - or columns with cells being deleted. Must be - of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type columns: :class:`list` of :class:`str` / - :func:`unicode `, or :class:`object` - :param columns: The columns within the column family that will have - cells deleted. If :attr:`ALL_COLUMNS` is used then - the entire column family will be deleted from the row. - - :type time_range: :class:`TimestampRange` - :param time_range: (Optional) The range of time within which cells - should be deleted. - - :type state: bool - :param state: (Optional) The state that is passed along to - :meth:`_get_mutations`. - """ - mutations_list = self._get_mutations(state) - if columns is self.ALL_COLUMNS: - mutation_val = data_v2_pb2.Mutation.DeleteFromFamily( - family_name=column_family_id - ) - mutation_pb = data_v2_pb2.Mutation(delete_from_family=mutation_val) - mutations_list.append(mutation_pb) - else: - delete_kwargs = {} - if time_range is not None: - delete_kwargs["time_range"] = time_range.to_pb() - - to_append = [] - for column in columns: - column = _to_bytes(column) - # time_range will never change if present, but the rest of - # delete_kwargs will - delete_kwargs.update( - family_name=column_family_id, column_qualifier=column - ) - mutation_val = data_v2_pb2.Mutation.DeleteFromColumn(**delete_kwargs) - mutation_pb = data_v2_pb2.Mutation(delete_from_column=mutation_val) - to_append.append(mutation_pb) - - # We don't add the mutations until all columns have been - # processed without error. - mutations_list.extend(to_append) - - -class DirectRow(_SetDeleteRow): - """Google Cloud Bigtable Row for sending "direct" mutations. - - These mutations directly set or delete cell contents: - - * :meth:`set_cell` - * :meth:`delete` - * :meth:`delete_cell` - * :meth:`delete_cells` - - These methods can be used directly:: - - >>> row = table.row(b'row-key1') - >>> row.set_cell(u'fam', b'col1', b'cell-val') - >>> row.delete_cell(u'fam', b'col2') - - .. note:: - - A :class:`DirectRow` accumulates mutations locally via the - :meth:`set_cell`, :meth:`delete`, :meth:`delete_cell` and - :meth:`delete_cells` methods. To actually send these mutations to the - Google Cloud Bigtable API, you must call :meth:`commit`. - - :type row_key: bytes - :param row_key: The key for the current row. - - :type table: :class:`Table ` - :param table: (Optional) The table that owns the row. This is - used for the :meth: `commit` only. Alternatively, - DirectRows can be persisted via - :meth:`~google.cloud.bigtable.deprecated.table.Table.mutate_rows`. - """ - - def __init__(self, row_key, table=None): - super(DirectRow, self).__init__(row_key, table) - self._pb_mutations = [] - - def _get_mutations(self, state=None): # pylint: disable=unused-argument - """Gets the list of mutations for a given state. - - ``state`` is unused by :class:`DirectRow` but is used by - subclasses. - - :type state: bool - :param state: The state that the mutation should be - applied in. - - :rtype: list - :returns: The list to add new mutations to (for the current state). - """ - return self._pb_mutations - - def get_mutations_size(self): - """Gets the total mutations size for current row - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_get_mutations_size] - :end-before: [END bigtable_api_row_get_mutations_size] - :dedent: 4 - """ - - mutation_size = 0 - for mutation in self._get_mutations(): - mutation_size += mutation._pb.ByteSize() - - return mutation_size - - def set_cell(self, column_family_id, column, value, timestamp=None): - """Sets a value in this row. - - The cell is determined by the ``row_key`` of this :class:`DirectRow` - and the ``column``. The ``column`` must be in an existing - :class:`.ColumnFamily` (as determined by ``column_family_id``). - - .. note:: - - This method adds a mutation to the accumulated mutations on this - row, but does not make an API request. To actually - send an API request (with the mutations) to the Google Cloud - Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_set_cell] - :end-before: [END bigtable_api_row_set_cell] - :dedent: 4 - - :type column_family_id: str - :param column_family_id: The column family that contains the column. - Must be of the form - ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type column: bytes - :param column: The column within the column family where the cell - is located. - - :type value: bytes or :class:`int` - :param value: The value to set in the cell. If an integer is used, - will be interpreted as a 64-bit big-endian signed - integer (8 bytes). - - :type timestamp: :class:`datetime.datetime` - :param timestamp: (Optional) The timestamp of the operation. - """ - self._set_cell(column_family_id, column, value, timestamp=timestamp, state=None) - - def delete(self): - """Deletes this row from the table. - - .. note:: - - This method adds a mutation to the accumulated mutations on this - row, but does not make an API request. To actually - send an API request (with the mutations) to the Google Cloud - Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_delete] - :end-before: [END bigtable_api_row_delete] - :dedent: 4 - """ - self._delete(state=None) - - def delete_cell(self, column_family_id, column, time_range=None): - """Deletes cell in this row. - - .. note:: - - This method adds a mutation to the accumulated mutations on this - row, but does not make an API request. To actually - send an API request (with the mutations) to the Google Cloud - Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_delete_cell] - :end-before: [END bigtable_api_row_delete_cell] - :dedent: 4 - - :type column_family_id: str - :param column_family_id: The column family that contains the column - or columns with cells being deleted. Must be - of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type column: bytes - :param column: The column within the column family that will have a - cell deleted. - - :type time_range: :class:`TimestampRange` - :param time_range: (Optional) The range of time within which cells - should be deleted. - """ - self._delete_cells( - column_family_id, [column], time_range=time_range, state=None - ) - - def delete_cells(self, column_family_id, columns, time_range=None): - """Deletes cells in this row. - - .. note:: - - This method adds a mutation to the accumulated mutations on this - row, but does not make an API request. To actually - send an API request (with the mutations) to the Google Cloud - Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_delete_cells] - :end-before: [END bigtable_api_row_delete_cells] - :dedent: 4 - - :type column_family_id: str - :param column_family_id: The column family that contains the column - or columns with cells being deleted. Must be - of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type columns: :class:`list` of :class:`str` / - :func:`unicode `, or :class:`object` - :param columns: The columns within the column family that will have - cells deleted. If :attr:`ALL_COLUMNS` is used then - the entire column family will be deleted from the row. - - :type time_range: :class:`TimestampRange` - :param time_range: (Optional) The range of time within which cells - should be deleted. - """ - self._delete_cells(column_family_id, columns, time_range=time_range, state=None) - - def commit(self): - """Makes a ``MutateRow`` API request. - - If no mutations have been created in the row, no request is made. - - Mutations are applied atomically and in order, meaning that earlier - mutations can be masked / negated by later ones. Cells already present - in the row are left unchanged unless explicitly changed by a mutation. - - After committing the accumulated mutations, resets the local - mutations to an empty list. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_commit] - :end-before: [END bigtable_api_row_commit] - :dedent: 4 - - :rtype: :class:`~google.rpc.status_pb2.Status` - :returns: A response status (`google.rpc.status_pb2.Status`) - representing success or failure of the row committed. - :raises: :exc:`~.table.TooManyMutationsError` if the number of - mutations is greater than 100,000. - """ - response = self._table.mutate_rows([self]) - - self.clear() - - return response[0] - - def clear(self): - """Removes all currently accumulated mutations on the current row. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_clear] - :end-before: [END bigtable_api_row_clear] - :dedent: 4 - """ - del self._pb_mutations[:] - - -class ConditionalRow(_SetDeleteRow): - """Google Cloud Bigtable Row for sending mutations conditionally. - - Each mutation has an associated state: :data:`True` or :data:`False`. - When :meth:`commit`-ed, the mutations for the :data:`True` - state will be applied if the filter matches any cells in - the row, otherwise the :data:`False` state will be applied. - - A :class:`ConditionalRow` accumulates mutations in the same way a - :class:`DirectRow` does: - - * :meth:`set_cell` - * :meth:`delete` - * :meth:`delete_cell` - * :meth:`delete_cells` - - with the only change the extra ``state`` parameter:: - - >>> row_cond = table.row(b'row-key2', filter_=row_filter) - >>> row_cond.set_cell(u'fam', b'col', b'cell-val', state=True) - >>> row_cond.delete_cell(u'fam', b'col', state=False) - - .. note:: - - As with :class:`DirectRow`, to actually send these mutations to the - Google Cloud Bigtable API, you must call :meth:`commit`. - - :type row_key: bytes - :param row_key: The key for the current row. - - :type table: :class:`Table ` - :param table: The table that owns the row. - - :type filter_: :class:`.RowFilter` - :param filter_: Filter to be used for conditional mutations. - """ - - def __init__(self, row_key, table, filter_): - super(ConditionalRow, self).__init__(row_key, table) - self._filter = filter_ - self._true_pb_mutations = [] - self._false_pb_mutations = [] - - def _get_mutations(self, state=None): - """Gets the list of mutations for a given state. - - Over-ridden so that the state can be used in: - - * :meth:`set_cell` - * :meth:`delete` - * :meth:`delete_cell` - * :meth:`delete_cells` - - :type state: bool - :param state: The state that the mutation should be - applied in. - - :rtype: list - :returns: The list to add new mutations to (for the current state). - """ - if state: - return self._true_pb_mutations - else: - return self._false_pb_mutations - - def commit(self): - """Makes a ``CheckAndMutateRow`` API request. - - If no mutations have been created in the row, no request is made. - - The mutations will be applied conditionally, based on whether the - filter matches any cells in the :class:`ConditionalRow` or not. (Each - method which adds a mutation has a ``state`` parameter for this - purpose.) - - Mutations are applied atomically and in order, meaning that earlier - mutations can be masked / negated by later ones. Cells already present - in the row are left unchanged unless explicitly changed by a mutation. - - After committing the accumulated mutations, resets the local - mutations. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_commit] - :end-before: [END bigtable_api_row_commit] - :dedent: 4 - - :rtype: bool - :returns: Flag indicating if the filter was matched (which also - indicates which set of mutations were applied by the server). - :raises: :class:`ValueError ` if the number of - mutations exceeds the :data:`MAX_MUTATIONS`. - """ - true_mutations = self._get_mutations(state=True) - false_mutations = self._get_mutations(state=False) - num_true_mutations = len(true_mutations) - num_false_mutations = len(false_mutations) - if num_true_mutations == 0 and num_false_mutations == 0: - return - if num_true_mutations > MAX_MUTATIONS or num_false_mutations > MAX_MUTATIONS: - raise ValueError( - "Exceed the maximum allowable mutations (%d). Had %s true " - "mutations and %d false mutations." - % (MAX_MUTATIONS, num_true_mutations, num_false_mutations) - ) - - data_client = self._table._instance._client.table_data_client - resp = data_client.check_and_mutate_row( - table_name=self._table.name, - row_key=self._row_key, - predicate_filter=self._filter.to_pb(), - app_profile_id=self._table._app_profile_id, - true_mutations=true_mutations, - false_mutations=false_mutations, - ) - self.clear() - return resp.predicate_matched - - # pylint: disable=arguments-differ - def set_cell(self, column_family_id, column, value, timestamp=None, state=True): - """Sets a value in this row. - - The cell is determined by the ``row_key`` of this - :class:`ConditionalRow` and the ``column``. The ``column`` must be in - an existing :class:`.ColumnFamily` (as determined by - ``column_family_id``). - - .. note:: - - This method adds a mutation to the accumulated mutations on this - row, but does not make an API request. To actually - send an API request (with the mutations) to the Google Cloud - Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_set_cell] - :end-before: [END bigtable_api_row_set_cell] - :dedent: 4 - - :type column_family_id: str - :param column_family_id: The column family that contains the column. - Must be of the form - ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type column: bytes - :param column: The column within the column family where the cell - is located. - - :type value: bytes or :class:`int` - :param value: The value to set in the cell. If an integer is used, - will be interpreted as a 64-bit big-endian signed - integer (8 bytes). - - :type timestamp: :class:`datetime.datetime` - :param timestamp: (Optional) The timestamp of the operation. - - :type state: bool - :param state: (Optional) The state that the mutation should be - applied in. Defaults to :data:`True`. - """ - self._set_cell( - column_family_id, column, value, timestamp=timestamp, state=state - ) - - def delete(self, state=True): - """Deletes this row from the table. - - .. note:: - - This method adds a mutation to the accumulated mutations on this - row, but does not make an API request. To actually - send an API request (with the mutations) to the Google Cloud - Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_delete] - :end-before: [END bigtable_api_row_delete] - :dedent: 4 - - :type state: bool - :param state: (Optional) The state that the mutation should be - applied in. Defaults to :data:`True`. - """ - self._delete(state=state) - - def delete_cell(self, column_family_id, column, time_range=None, state=True): - """Deletes cell in this row. - - .. note:: - - This method adds a mutation to the accumulated mutations on this - row, but does not make an API request. To actually - send an API request (with the mutations) to the Google Cloud - Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_delete_cell] - :end-before: [END bigtable_api_row_delete_cell] - :dedent: 4 - - :type column_family_id: str - :param column_family_id: The column family that contains the column - or columns with cells being deleted. Must be - of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type column: bytes - :param column: The column within the column family that will have a - cell deleted. - - :type time_range: :class:`TimestampRange` - :param time_range: (Optional) The range of time within which cells - should be deleted. - - :type state: bool - :param state: (Optional) The state that the mutation should be - applied in. Defaults to :data:`True`. - """ - self._delete_cells( - column_family_id, [column], time_range=time_range, state=state - ) - - def delete_cells(self, column_family_id, columns, time_range=None, state=True): - """Deletes cells in this row. - - .. note:: - - This method adds a mutation to the accumulated mutations on this - row, but does not make an API request. To actually - send an API request (with the mutations) to the Google Cloud - Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_delete_cells] - :end-before: [END bigtable_api_row_delete_cells] - :dedent: 4 - - :type column_family_id: str - :param column_family_id: The column family that contains the column - or columns with cells being deleted. Must be - of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type columns: :class:`list` of :class:`str` / - :func:`unicode `, or :class:`object` - :param columns: The columns within the column family that will have - cells deleted. If :attr:`ALL_COLUMNS` is used then the - entire column family will be deleted from the row. - - :type time_range: :class:`TimestampRange` - :param time_range: (Optional) The range of time within which cells - should be deleted. - - :type state: bool - :param state: (Optional) The state that the mutation should be - applied in. Defaults to :data:`True`. - """ - self._delete_cells( - column_family_id, columns, time_range=time_range, state=state - ) - - # pylint: enable=arguments-differ - - def clear(self): - """Removes all currently accumulated mutations on the current row. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_clear] - :end-before: [END bigtable_api_row_clear] - :dedent: 4 - """ - del self._true_pb_mutations[:] - del self._false_pb_mutations[:] - - -class AppendRow(Row): - """Google Cloud Bigtable Row for sending append mutations. - - These mutations are intended to augment the value of an existing cell - and uses the methods: - - * :meth:`append_cell_value` - * :meth:`increment_cell_value` - - The first works by appending bytes and the second by incrementing an - integer (stored in the cell as 8 bytes). In either case, if the - cell is empty, assumes the default empty value (empty string for - bytes or 0 for integer). - - :type row_key: bytes - :param row_key: The key for the current row. - - :type table: :class:`Table ` - :param table: The table that owns the row. - """ - - def __init__(self, row_key, table): - super(AppendRow, self).__init__(row_key, table) - self._rule_pb_list = [] - - def clear(self): - """Removes all currently accumulated modifications on current row. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_clear] - :end-before: [END bigtable_api_row_clear] - :dedent: 4 - """ - del self._rule_pb_list[:] - - def append_cell_value(self, column_family_id, column, value): - """Appends a value to an existing cell. - - .. note:: - - This method adds a read-modify rule protobuf to the accumulated - read-modify rules on this row, but does not make an API - request. To actually send an API request (with the rules) to the - Google Cloud Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_append_cell_value] - :end-before: [END bigtable_api_row_append_cell_value] - :dedent: 4 - - :type column_family_id: str - :param column_family_id: The column family that contains the column. - Must be of the form - ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type column: bytes - :param column: The column within the column family where the cell - is located. - - :type value: bytes - :param value: The value to append to the existing value in the cell. If - the targeted cell is unset, it will be treated as - containing the empty string. - """ - column = _to_bytes(column) - value = _to_bytes(value) - rule_pb = data_v2_pb2.ReadModifyWriteRule( - family_name=column_family_id, column_qualifier=column, append_value=value - ) - self._rule_pb_list.append(rule_pb) - - def increment_cell_value(self, column_family_id, column, int_value): - """Increments a value in an existing cell. - - Assumes the value in the cell is stored as a 64 bit integer - serialized to bytes. - - .. note:: - - This method adds a read-modify rule protobuf to the accumulated - read-modify rules on this row, but does not make an API - request. To actually send an API request (with the rules) to the - Google Cloud Bigtable API, call :meth:`commit`. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_increment_cell_value] - :end-before: [END bigtable_api_row_increment_cell_value] - :dedent: 4 - - :type column_family_id: str - :param column_family_id: The column family that contains the column. - Must be of the form - ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - - :type column: bytes - :param column: The column within the column family where the cell - is located. - - :type int_value: int - :param int_value: The value to increment the existing value in the cell - by. If the targeted cell is unset, it will be treated - as containing a zero. Otherwise, the targeted cell - must contain an 8-byte value (interpreted as a 64-bit - big-endian signed integer), or the entire request - will fail. - """ - column = _to_bytes(column) - rule_pb = data_v2_pb2.ReadModifyWriteRule( - family_name=column_family_id, - column_qualifier=column, - increment_amount=int_value, - ) - self._rule_pb_list.append(rule_pb) - - def commit(self): - """Makes a ``ReadModifyWriteRow`` API request. - - This commits modifications made by :meth:`append_cell_value` and - :meth:`increment_cell_value`. If no modifications were made, makes - no API request and just returns ``{}``. - - Modifies a row atomically, reading the latest existing - timestamp / value from the specified columns and writing a new value by - appending / incrementing. The new cell created uses either the current - server time or the highest timestamp of a cell in that column (if it - exceeds the server time). - - After committing the accumulated mutations, resets the local mutations. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_commit] - :end-before: [END bigtable_api_row_commit] - :dedent: 4 - - :rtype: dict - :returns: The new contents of all modified cells. Returned as a - dictionary of column families, each of which holds a - dictionary of columns. Each column contains a list of cells - modified. Each cell is represented with a two-tuple with the - value (in bytes) and the timestamp for the cell. - :raises: :class:`ValueError ` if the number of - mutations exceeds the :data:`MAX_MUTATIONS`. - """ - num_mutations = len(self._rule_pb_list) - if num_mutations == 0: - return {} - if num_mutations > MAX_MUTATIONS: - raise ValueError( - "%d total append mutations exceed the maximum " - "allowable %d." % (num_mutations, MAX_MUTATIONS) - ) - - data_client = self._table._instance._client.table_data_client - row_response = data_client.read_modify_write_row( - table_name=self._table.name, - row_key=self._row_key, - rules=self._rule_pb_list, - app_profile_id=self._table._app_profile_id, - ) - - # Reset modifications after commit-ing request. - self.clear() - - # NOTE: We expect row_response.key == self._row_key but don't check. - return _parse_rmw_row_response(row_response) - - -def _parse_rmw_row_response(row_response): - """Parses the response to a ``ReadModifyWriteRow`` request. - - :type row_response: :class:`.data_v2_pb2.Row` - :param row_response: The response row (with only modified cells) from a - ``ReadModifyWriteRow`` request. - - :rtype: dict - :returns: The new contents of all modified cells. Returned as a - dictionary of column families, each of which holds a - dictionary of columns. Each column contains a list of cells - modified. Each cell is represented with a two-tuple with the - value (in bytes) and the timestamp for the cell. For example: - - .. code:: python - - { - u'col-fam-id': { - b'col-name1': [ - (b'cell-val', datetime.datetime(...)), - (b'cell-val-newer', datetime.datetime(...)), - ], - b'col-name2': [ - (b'altcol-cell-val', datetime.datetime(...)), - ], - }, - u'col-fam-id2': { - b'col-name3-but-other-fam': [ - (b'foo', datetime.datetime(...)), - ], - }, - } - """ - result = {} - for column_family in row_response.row.families: - column_family_id, curr_family = _parse_family_pb(column_family) - result[column_family_id] = curr_family - return result - - -def _parse_family_pb(family_pb): - """Parses a Family protobuf into a dictionary. - - :type family_pb: :class:`._generated.data_pb2.Family` - :param family_pb: A protobuf - - :rtype: tuple - :returns: A string and dictionary. The string is the name of the - column family and the dictionary has column names (within the - family) as keys and cell lists as values. Each cell is - represented with a two-tuple with the value (in bytes) and the - timestamp for the cell. For example: - - .. code:: python - - { - b'col-name1': [ - (b'cell-val', datetime.datetime(...)), - (b'cell-val-newer', datetime.datetime(...)), - ], - b'col-name2': [ - (b'altcol-cell-val', datetime.datetime(...)), - ], - } - """ - result = {} - for column in family_pb.columns: - result[column.qualifier] = cells = [] - for cell in column.cells: - val_pair = (cell.value, _datetime_from_microseconds(cell.timestamp_micros)) - cells.append(val_pair) - - return family_pb.name, result - - -class PartialRowData(object): - """Representation of partial row in a Google Cloud Bigtable Table. - - These are expected to be updated directly from a - :class:`._generated.bigtable_service_messages_pb2.ReadRowsResponse` - - :type row_key: bytes - :param row_key: The key for the row holding the (partial) data. - """ - - def __init__(self, row_key): - self._row_key = row_key - self._cells = {} - - def __eq__(self, other): - if not isinstance(other, self.__class__): - return NotImplemented - return other._row_key == self._row_key and other._cells == self._cells - - def __ne__(self, other): - return not self == other - - def to_dict(self): - """Convert the cells to a dictionary. - - This is intended to be used with HappyBase, so the column family and - column qualiers are combined (with ``:``). - - :rtype: dict - :returns: Dictionary containing all the data in the cells of this row. - """ - result = {} - for column_family_id, columns in self._cells.items(): - for column_qual, cells in columns.items(): - key = _to_bytes(column_family_id) + b":" + _to_bytes(column_qual) - result[key] = cells - return result - - @property - def cells(self): - """Property returning all the cells accumulated on this partial row. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_data_cells] - :end-before: [END bigtable_api_row_data_cells] - :dedent: 4 - - :rtype: dict - :returns: Dictionary of the :class:`Cell` objects accumulated. This - dictionary has two-levels of keys (first for column families - and second for column names/qualifiers within a family). For - a given column, a list of :class:`Cell` objects is stored. - """ - return self._cells - - @property - def row_key(self): - """Getter for the current (partial) row's key. - - :rtype: bytes - :returns: The current (partial) row's key. - """ - return self._row_key - - def find_cells(self, column_family_id, column): - """Get a time series of cells stored on this instance. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_find_cells] - :end-before: [END bigtable_api_row_find_cells] - :dedent: 4 - - Args: - column_family_id (str): The ID of the column family. Must be of the - form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - column (bytes): The column within the column family where the cells - are located. - - Returns: - List[~google.cloud.bigtable.deprecated.row_data.Cell]: The cells stored in the - specified column. - - Raises: - KeyError: If ``column_family_id`` is not among the cells stored - in this row. - KeyError: If ``column`` is not among the cells stored in this row - for the given ``column_family_id``. - """ - try: - column_family = self._cells[column_family_id] - except KeyError: - raise KeyError(_MISSING_COLUMN_FAMILY.format(column_family_id)) - - try: - cells = column_family[column] - except KeyError: - raise KeyError(_MISSING_COLUMN.format(column, column_family_id)) - - return cells - - def cell_value(self, column_family_id, column, index=0): - """Get a single cell value stored on this instance. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_cell_value] - :end-before: [END bigtable_api_row_cell_value] - :dedent: 4 - - Args: - column_family_id (str): The ID of the column family. Must be of the - form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - column (bytes): The column within the column family where the cell - is located. - index (Optional[int]): The offset within the series of values. If - not specified, will return the first cell. - - Returns: - ~google.cloud.bigtable.deprecated.row_data.Cell value: The cell value stored - in the specified column and specified index. - - Raises: - KeyError: If ``column_family_id`` is not among the cells stored - in this row. - KeyError: If ``column`` is not among the cells stored in this row - for the given ``column_family_id``. - IndexError: If ``index`` cannot be found within the cells stored - in this row for the given ``column_family_id``, ``column`` - pair. - """ - cells = self.find_cells(column_family_id, column) - - try: - cell = cells[index] - except (TypeError, IndexError): - num_cells = len(cells) - msg = _MISSING_INDEX.format(index, column, column_family_id, num_cells) - raise IndexError(msg) - - return cell.value - - def cell_values(self, column_family_id, column, max_count=None): - """Get a time series of cells stored on this instance. - - For example: - - .. literalinclude:: snippets_table.py - :start-after: [START bigtable_api_row_cell_values] - :end-before: [END bigtable_api_row_cell_values] - :dedent: 4 - - Args: - column_family_id (str): The ID of the column family. Must be of the - form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - column (bytes): The column within the column family where the cells - are located. - max_count (int): The maximum number of cells to use. - - Returns: - A generator which provides: cell.value, cell.timestamp_micros - for each cell in the list of cells - - Raises: - KeyError: If ``column_family_id`` is not among the cells stored - in this row. - KeyError: If ``column`` is not among the cells stored in this row - for the given ``column_family_id``. - """ - cells = self.find_cells(column_family_id, column) - if max_count is None: - max_count = len(cells) - - for index, cell in enumerate(cells): - if index == max_count: - break - - yield cell.value, cell.timestamp_micros - - -class Cell(object): - """Representation of a Google Cloud Bigtable Cell. - - :type value: bytes - :param value: The value stored in the cell. - - :type timestamp_micros: int - :param timestamp_micros: The timestamp_micros when the cell was stored. - - :type labels: list - :param labels: (Optional) List of strings. Labels applied to the cell. - """ - - def __init__(self, value, timestamp_micros, labels=None): - self.value = value - self.timestamp_micros = timestamp_micros - self.labels = list(labels) if labels is not None else [] - - @classmethod - def from_pb(cls, cell_pb): - """Create a new cell from a Cell protobuf. - - :type cell_pb: :class:`._generated.data_pb2.Cell` - :param cell_pb: The protobuf to convert. - - :rtype: :class:`Cell` - :returns: The cell corresponding to the protobuf. - """ - if cell_pb.labels: - return cls(cell_pb.value, cell_pb.timestamp_micros, labels=cell_pb.labels) - else: - return cls(cell_pb.value, cell_pb.timestamp_micros) - - @property - def timestamp(self): - return _datetime_from_microseconds(self.timestamp_micros) - - def __eq__(self, other): - if not isinstance(other, self.__class__): - return NotImplemented - return ( - other.value == self.value - and other.timestamp_micros == self.timestamp_micros - and other.labels == self.labels - ) - - def __ne__(self, other): - return not self == other - - def __repr__(self): - return "<{name} value={value!r} timestamp={timestamp}>".format( - name=self.__class__.__name__, value=self.value, timestamp=self.timestamp - ) - - -class InvalidChunk(RuntimeError): - """Exception raised to invalid chunk data from back-end.""" diff --git a/google/cloud/bigtable/deprecated/encryption_info.py b/google/cloud/bigtable/encryption_info.py similarity index 93% rename from google/cloud/bigtable/deprecated/encryption_info.py rename to google/cloud/bigtable/encryption_info.py index daa0d9232..1757297bc 100644 --- a/google/cloud/bigtable/deprecated/encryption_info.py +++ b/google/cloud/bigtable/encryption_info.py @@ -14,7 +14,7 @@ """Class for encryption info for tables and backups.""" -from google.cloud.bigtable.deprecated.error import Status +from google.cloud.bigtable.error import Status class EncryptionInfo: @@ -27,7 +27,7 @@ class EncryptionInfo: :type encryption_type: int :param encryption_type: See :class:`enums.EncryptionInfo.EncryptionType` - :type encryption_status: google.cloud.bigtable.deprecated.encryption.Status + :type encryption_status: google.cloud.bigtable.encryption.Status :param encryption_status: The encryption status. :type kms_key_version: str diff --git a/google/cloud/bigtable/deprecated/enums.py b/google/cloud/bigtable/enums.py similarity index 100% rename from google/cloud/bigtable/deprecated/enums.py rename to google/cloud/bigtable/enums.py diff --git a/google/cloud/bigtable/deprecated/error.py b/google/cloud/bigtable/error.py similarity index 100% rename from google/cloud/bigtable/deprecated/error.py rename to google/cloud/bigtable/error.py diff --git a/google/cloud/bigtable/gapic_version.py b/google/cloud/bigtable/gapic_version.py index 8d4f4cfb6..0f1a446f3 100644 --- a/google/cloud/bigtable/gapic_version.py +++ b/google/cloud/bigtable/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.17.0" # {x-release-please-version} +__version__ = "2.19.0" # {x-release-please-version} diff --git a/google/cloud/bigtable/deprecated/instance.py b/google/cloud/bigtable/instance.py similarity index 91% rename from google/cloud/bigtable/deprecated/instance.py rename to google/cloud/bigtable/instance.py index 33475d261..6d092cefd 100644 --- a/google/cloud/bigtable/deprecated/instance.py +++ b/google/cloud/bigtable/instance.py @@ -16,9 +16,9 @@ import re -from google.cloud.bigtable.deprecated.app_profile import AppProfile -from google.cloud.bigtable.deprecated.cluster import Cluster -from google.cloud.bigtable.deprecated.table import Table +from google.cloud.bigtable.app_profile import AppProfile +from google.cloud.bigtable.cluster import Cluster +from google.cloud.bigtable.table import Table from google.protobuf import field_mask_pb2 @@ -28,7 +28,7 @@ from google.api_core.exceptions import NotFound -from google.cloud.bigtable.deprecated.policy import Policy +from google.cloud.bigtable.policy import Policy import warnings @@ -61,7 +61,7 @@ class Instance(object): :type instance_id: str :param instance_id: The ID of the instance. - :type client: :class:`Client ` + :type client: :class:`Client ` :param client: The client that owns the instance. Provides authorization and a project ID. @@ -75,10 +75,10 @@ class Instance(object): :param instance_type: (Optional) The type of the instance. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.deprecated.enums.Instance.Type.PRODUCTION`. - :data:`google.cloud.bigtable.deprecated.enums.Instance.Type.DEVELOPMENT`, + :data:`google.cloud.bigtable.enums.Instance.Type.PRODUCTION`. + :data:`google.cloud.bigtable.enums.Instance.Type.DEVELOPMENT`, Defaults to - :data:`google.cloud.bigtable.deprecated.enums.Instance.Type.UNSPECIFIED`. + :data:`google.cloud.bigtable.enums.Instance.Type.UNSPECIFIED`. :type labels: dict :param labels: (Optional) Labels are a flexible and lightweight @@ -95,9 +95,9 @@ class Instance(object): :param _state: (`OutputOnly`) The current state of the instance. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.deprecated.enums.Instance.State.STATE_NOT_KNOWN`. - :data:`google.cloud.bigtable.deprecated.enums.Instance.State.READY`. - :data:`google.cloud.bigtable.deprecated.enums.Instance.State.CREATING`. + :data:`google.cloud.bigtable.enums.Instance.State.STATE_NOT_KNOWN`. + :data:`google.cloud.bigtable.enums.Instance.State.READY`. + :data:`google.cloud.bigtable.enums.Instance.State.CREATING`. """ def __init__( @@ -141,7 +141,7 @@ def from_pb(cls, instance_pb, client): :type instance_pb: :class:`instance.Instance` :param instance_pb: An instance protobuf object. - :type client: :class:`Client ` + :type client: :class:`Client ` :param client: The client that owns the instance. :rtype: :class:`Instance` @@ -196,7 +196,7 @@ def name(self): @property def state(self): - """google.cloud.bigtable.deprecated.enums.Instance.State: state of Instance. + """google.cloud.bigtable.enums.Instance.State: state of Instance. For example: @@ -272,12 +272,12 @@ def create( persisting Bigtable data. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.deprecated.enums.StorageType.SSD`. - :data:`google.cloud.bigtable.deprecated.enums.StorageType.HDD`, + :data:`google.cloud.bigtable.enums.StorageType.SSD`. + :data:`google.cloud.bigtable.enums.StorageType.HDD`, Defaults to - :data:`google.cloud.bigtable.deprecated.enums.StorageType.UNSPECIFIED`. + :data:`google.cloud.bigtable.enums.StorageType.UNSPECIFIED`. - :type clusters: class:`~[~google.cloud.bigtable.deprecated.cluster.Cluster]` + :type clusters: class:`~[~google.cloud.bigtable.cluster.Cluster]` :param clusters: List of clusters to be created. :rtype: :class:`~google.api_core.operation.Operation` @@ -478,7 +478,7 @@ def get_iam_policy(self, requested_policy_version=None): than the one that was requested, based on the feature syntax in the policy fetched. - :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :rtype: :class:`google.cloud.bigtable.policy.Policy` :returns: The current IAM policy of this instance """ args = {"resource": self.name} @@ -497,7 +497,7 @@ def set_iam_policy(self, policy): existing policy. For more information about policy, please see documentation of - class `google.cloud.bigtable.deprecated.policy.Policy` + class `google.cloud.bigtable.policy.Policy` For example: @@ -506,11 +506,11 @@ class `google.cloud.bigtable.deprecated.policy.Policy` :end-before: [END bigtable_api_set_iam_policy] :dedent: 4 - :type policy: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :type policy: :class:`google.cloud.bigtable.policy.Policy` :param policy: A new IAM policy to replace the current IAM policy of this instance - :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :rtype: :class:`google.cloud.bigtable.policy.Policy` :returns: The current IAM policy of this instance. """ instance_admin_client = self._client.instance_admin_client @@ -586,12 +586,12 @@ def cluster( :param default_storage_type: (Optional) The type of storage Possible values are represented by the following constants: - :data:`google.cloud.bigtable.deprecated.enums.StorageType.SSD`. - :data:`google.cloud.bigtable.deprecated.enums.StorageType.HDD`, + :data:`google.cloud.bigtable.enums.StorageType.SSD`. + :data:`google.cloud.bigtable.enums.StorageType.HDD`, Defaults to - :data:`google.cloud.bigtable.deprecated.enums.StorageType.UNSPECIFIED`. + :data:`google.cloud.bigtable.enums.StorageType.UNSPECIFIED`. - :rtype: :class:`~google.cloud.bigtable.deprecated.instance.Cluster` + :rtype: :class:`~google.cloud.bigtable.instance.Cluster` :returns: a cluster owned by this instance. :type kms_key_name: str @@ -635,7 +635,7 @@ def list_clusters(self): :rtype: tuple :returns: (clusters, failed_locations), where 'clusters' is list of - :class:`google.cloud.bigtable.deprecated.instance.Cluster`, and + :class:`google.cloud.bigtable.instance.Cluster`, and 'failed_locations' is a list of locations which could not be resolved. """ @@ -664,7 +664,7 @@ def table(self, table_id, mutation_timeout=None, app_profile_id=None): :type app_profile_id: str :param app_profile_id: (Optional) The unique name of the AppProfile. - :rtype: :class:`Table ` + :rtype: :class:`Table ` :returns: The table owned by this instance. """ return Table( @@ -684,7 +684,7 @@ def list_tables(self): :end-before: [END bigtable_api_list_tables] :dedent: 4 - :rtype: list of :class:`Table ` + :rtype: list of :class:`Table ` :returns: The list of tables owned by the instance. :raises: :class:`ValueError ` if one of the returned tables has a name that is not of the expected format. @@ -731,8 +731,8 @@ def app_profile( :param: routing_policy_type: The type of the routing policy. Possible values are represented by the following constants: - :data:`google.cloud.bigtable.deprecated.enums.RoutingPolicyType.ANY` - :data:`google.cloud.bigtable.deprecated.enums.RoutingPolicyType.SINGLE` + :data:`google.cloud.bigtable.enums.RoutingPolicyType.ANY` + :data:`google.cloud.bigtable.enums.RoutingPolicyType.SINGLE` :type: description: str :param: description: (Optional) Long form description of the use @@ -753,7 +753,7 @@ def app_profile( transactional writes for ROUTING_POLICY_TYPE_SINGLE. - :rtype: :class:`~google.cloud.bigtable.deprecated.app_profile.AppProfile>` + :rtype: :class:`~google.cloud.bigtable.app_profile.AppProfile>` :returns: AppProfile for this instance. """ return AppProfile( @@ -776,10 +776,10 @@ def list_app_profiles(self): :end-before: [END bigtable_api_list_app_profiles] :dedent: 4 - :rtype: :list:[`~google.cloud.bigtable.deprecated.app_profile.AppProfile`] - :returns: A :list:[`~google.cloud.bigtable.deprecated.app_profile.AppProfile`]. + :rtype: :list:[`~google.cloud.bigtable.app_profile.AppProfile`] + :returns: A :list:[`~google.cloud.bigtable.app_profile.AppProfile`]. By default, this is a list of - :class:`~google.cloud.bigtable.deprecated.app_profile.AppProfile` + :class:`~google.cloud.bigtable.app_profile.AppProfile` instances. """ resp = self._client.instance_admin_client.list_app_profiles( diff --git a/google/cloud/bigtable/iterators.py b/google/cloud/bigtable/iterators.py deleted file mode 100644 index b20932fb2..000000000 --- a/google/cloud/bigtable/iterators.py +++ /dev/null @@ -1,129 +0,0 @@ -# Copyright 2023 Google LLC -# -# Licensed 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. -# -from __future__ import annotations - -from typing import AsyncIterable - -import asyncio -import time -import sys - -from google.cloud.bigtable._read_rows import _ReadRowsOperation -from google.cloud.bigtable.exceptions import IdleTimeout -from google.cloud.bigtable._helpers import _convert_retry_deadline -from google.cloud.bigtable.row import Row - - -class ReadRowsIterator(AsyncIterable[Row]): - """ - Async iterator for ReadRows responses. - """ - - def __init__(self, merger: _ReadRowsOperation): - self._merger: _ReadRowsOperation = merger - self._error: Exception | None = None - self.last_interaction_time = time.time() - self._idle_timeout_task: asyncio.Task[None] | None = None - # wrap merger with a wrapper that properly formats exceptions - self._next_fn = _convert_retry_deadline( - self._merger.__anext__, - self._merger.operation_timeout, - self._merger.transient_errors, - ) - - async def _start_idle_timer(self, idle_timeout: float): - """ - Start a coroutine that will cancel a stream if no interaction - with the iterator occurs for the specified number of seconds. - - Subsequent access to the iterator will raise an IdleTimeout exception. - - Args: - - idle_timeout: number of seconds of inactivity before cancelling the stream - """ - self.last_interaction_time = time.time() - if self._idle_timeout_task is not None: - self._idle_timeout_task.cancel() - self._idle_timeout_task = asyncio.create_task( - self._idle_timeout_coroutine(idle_timeout) - ) - if sys.version_info >= (3, 8): - self._idle_timeout_task.name = "ReadRowsIterator._idle_timeout" - - @property - def active(self): - """ - Returns True if the iterator is still active and has not been closed - """ - return self._error is None - - async def _idle_timeout_coroutine(self, idle_timeout: float): - """ - Coroutine that will cancel a stream if no interaction with the iterator - in the last `idle_timeout` seconds. - """ - while self.active: - next_timeout = self.last_interaction_time + idle_timeout - await asyncio.sleep(next_timeout - time.time()) - if self.last_interaction_time + idle_timeout < time.time() and self.active: - # idle timeout has expired - await self._finish_with_error( - IdleTimeout( - ( - "Timed out waiting for next Row to be consumed. " - f"(idle_timeout={idle_timeout:0.1f}s)" - ) - ) - ) - - def __aiter__(self): - """Implement the async iterator protocol.""" - return self - - async def __anext__(self) -> Row: - """ - Implement the async iterator potocol. - - Return the next item in the stream if active, or - raise an exception if the stream has been closed. - """ - if self._error is not None: - raise self._error - try: - self.last_interaction_time = time.time() - return await self._next_fn() - except Exception as e: - await self._finish_with_error(e) - raise e - - async def _finish_with_error(self, e: Exception): - """ - Helper function to close the stream and clean up resources - after an error has occurred. - """ - if self.active: - await self._merger.aclose() - self._error = e - if self._idle_timeout_task is not None: - self._idle_timeout_task.cancel() - self._idle_timeout_task = None - - async def aclose(self): - """ - Support closing the stream with an explicit call to aclose() - """ - await self._finish_with_error( - StopAsyncIteration(f"{self.__class__.__name__} closed") - ) diff --git a/google/cloud/bigtable/deprecated/policy.py b/google/cloud/bigtable/policy.py similarity index 100% rename from google/cloud/bigtable/deprecated/policy.py rename to google/cloud/bigtable/policy.py diff --git a/google/cloud/bigtable/row.py b/google/cloud/bigtable/row.py index 5fdc1b365..752458a08 100644 --- a/google/cloud/bigtable/row.py +++ b/google/cloud/bigtable/row.py @@ -1,4 +1,4 @@ -# Copyright 2023 Google LLC +# Copyright 2015 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,455 +11,1257 @@ # 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. -# -from __future__ import annotations -from collections import OrderedDict -from typing import Sequence, Generator, overload, Any -from functools import total_ordering +"""User-friendly container for Google Cloud Bigtable Row.""" -from google.cloud.bigtable_v2.types import Row as RowPB -# Type aliases used internally for readability. -_family_type = str -_qualifier_type = bytes +import struct +from google.cloud._helpers import _datetime_from_microseconds # type: ignore +from google.cloud._helpers import _microseconds_from_datetime # type: ignore +from google.cloud._helpers import _to_bytes # type: ignore +from google.cloud.bigtable_v2.types import data as data_v2_pb2 + + +_PACK_I64 = struct.Struct(">q").pack + +MAX_MUTATIONS = 100000 +"""The maximum number of mutations that a row can accumulate.""" + +_MISSING_COLUMN_FAMILY = "Column family {} is not among the cells stored in this row." +_MISSING_COLUMN = ( + "Column {} is not among the cells stored in this row in the column family {}." +) +_MISSING_INDEX = ( + "Index {!r} is not valid for the cells stored in this row for column {} " + "in the column family {}. There are {} such cells." +) -class Row(Sequence["Cell"]): - """ - Model class for row data returned from server - Does not represent all data contained in the row, only data returned by a - query. - Expected to be read-only to users, and written by backend +class Row(object): + """Base representation of a Google Cloud Bigtable Row. - Can be indexed: - cells = row["family", "qualifier"] + This class has three subclasses corresponding to the three + RPC methods for sending row mutations: + + * :class:`DirectRow` for ``MutateRow`` + * :class:`ConditionalRow` for ``CheckAndMutateRow`` + * :class:`AppendRow` for ``ReadModifyWriteRow`` + + :type row_key: bytes + :param row_key: The key for the current row. + + :type table: :class:`Table ` + :param table: (Optional) The table that owns the row. """ - __slots__ = ("row_key", "cells", "_index_data") + def __init__(self, row_key, table=None): + self._row_key = _to_bytes(row_key) + self._table = table - def __init__( - self, - key: bytes, - cells: list[Cell], - ): - """ - Initializes a Row object + @property + def row_key(self): + """Row key. - Row objects are not intended to be created by users. - They are returned by the Bigtable backend. + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_row_key] + :end-before: [END bigtable_api_row_row_key] + :dedent: 4 + + :rtype: bytes + :returns: The key for the current row. """ - self.row_key = key - self.cells: list[Cell] = cells - # index is lazily created when needed - self._index_data: OrderedDict[ - _family_type, OrderedDict[_qualifier_type, list[Cell]] - ] | None = None + return self._row_key @property - def _index( - self, - ) -> OrderedDict[_family_type, OrderedDict[_qualifier_type, list[Cell]]]: - """ - Returns an index of cells associated with each family and qualifier. + def table(self): + """Row table. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_table] + :end-before: [END bigtable_api_row_table] + :dedent: 4 - The index is lazily created when needed + :rtype: table: :class:`Table ` + :returns: table: The table that owns the row. """ - if self._index_data is None: - self._index_data = OrderedDict() - for cell in self.cells: - self._index_data.setdefault(cell.family, OrderedDict()).setdefault( - cell.qualifier, [] - ).append(cell) - return self._index_data + return self._table - @classmethod - def _from_pb(cls, row_pb: RowPB) -> Row: - """ - Creates a row from a protobuf representation - - Row objects are not intended to be created by users. - They are returned by the Bigtable backend. - """ - row_key: bytes = row_pb.key - cell_list: list[Cell] = [] - for family in row_pb.families: - for column in family.columns: - for cell in column.cells: - new_cell = Cell( - value=cell.value, - row_key=row_key, - family=family.name, - qualifier=column.qualifier, - timestamp_micros=cell.timestamp_micros, - labels=list(cell.labels) if cell.labels else None, - ) - cell_list.append(new_cell) - return cls(row_key, cells=cell_list) - - def get_cells( - self, family: str | None = None, qualifier: str | bytes | None = None - ) -> list[Cell]: - """ - Returns cells sorted in Bigtable native order: - - Family lexicographically ascending - - Qualifier ascending - - Timestamp in reverse chronological order - - If family or qualifier not passed, will include all - - Can also be accessed through indexing: - cells = row["family", "qualifier"] - cells = row["family"] - """ - if family is None: - if qualifier is not None: - # get_cells(None, "qualifier") is not allowed - raise ValueError("Qualifier passed without family") - else: - # return all cells on get_cells() - return self.cells - if qualifier is None: - # return all cells in family on get_cells(family) - return list(self._get_all_from_family(family)) - if isinstance(qualifier, str): - qualifier = qualifier.encode("utf-8") - # return cells in family and qualifier on get_cells(family, qualifier) - if family not in self._index: - raise ValueError(f"Family '{family}' not found in row '{self.row_key!r}'") - if qualifier not in self._index[family]: - raise ValueError( - f"Qualifier '{qualifier!r}' not found in family '{family}' in row '{self.row_key!r}'" - ) - return self._index[family][qualifier] - def _get_all_from_family(self, family: str) -> Generator[Cell, None, None]: +class _SetDeleteRow(Row): + """Row helper for setting or deleting cell values. + + Implements helper methods to add mutations to set or delete cell contents: + + * :meth:`set_cell` + * :meth:`delete` + * :meth:`delete_cell` + * :meth:`delete_cells` + + :type row_key: bytes + :param row_key: The key for the current row. + + :type table: :class:`Table ` + :param table: The table that owns the row. + """ + + ALL_COLUMNS = object() + """Sentinel value used to indicate all columns in a column family.""" + + def _get_mutations(self, state=None): + """Gets the list of mutations for a given state. + + This method intended to be implemented by subclasses. + + ``state`` may not need to be used by all subclasses. + + :type state: bool + :param state: The state that the mutation should be + applied in. + + :raises: :class:`NotImplementedError ` + always. """ - Returns all cells in the row for the family_id + raise NotImplementedError + + def _set_cell(self, column_family_id, column, value, timestamp=None, state=None): + """Helper for :meth:`set_cell` + + Adds a mutation to set the value in a specific cell. + + ``state`` is unused by :class:`DirectRow` but is used by + subclasses. + + :type column_family_id: str + :param column_family_id: The column family that contains the column. + Must be of the form + ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type column: bytes + :param column: The column within the column family where the cell + is located. + + :type value: bytes or :class:`int` + :param value: The value to set in the cell. If an integer is used, + will be interpreted as a 64-bit big-endian signed + integer (8 bytes). + + :type timestamp: :class:`datetime.datetime` + :param timestamp: (Optional) The timestamp of the operation. + + :type state: bool + :param state: (Optional) The state that is passed along to + :meth:`_get_mutations`. """ - if family not in self._index: - raise ValueError(f"Family '{family}' not found in row '{self.row_key!r}'") - for qualifier in self._index[family]: - yield from self._index[family][qualifier] + column = _to_bytes(column) + if isinstance(value, int): + value = _PACK_I64(value) + value = _to_bytes(value) + if timestamp is None: + # Use -1 for current Bigtable server time. + timestamp_micros = -1 + else: + timestamp_micros = _microseconds_from_datetime(timestamp) + # Truncate to millisecond granularity. + timestamp_micros -= timestamp_micros % 1000 + + mutation_val = data_v2_pb2.Mutation.SetCell( + family_name=column_family_id, + column_qualifier=column, + timestamp_micros=timestamp_micros, + value=value, + ) + mutation_pb = data_v2_pb2.Mutation(set_cell=mutation_val) + self._get_mutations(state).append(mutation_pb) - def __str__(self) -> str: + def _delete(self, state=None): + """Helper for :meth:`delete` + + Adds a delete mutation (for the entire row) to the accumulated + mutations. + + ``state`` is unused by :class:`DirectRow` but is used by + subclasses. + + :type state: bool + :param state: (Optional) The state that is passed along to + :meth:`_get_mutations`. """ - Human-readable string representation + mutation_val = data_v2_pb2.Mutation.DeleteFromRow() + mutation_pb = data_v2_pb2.Mutation(delete_from_row=mutation_val) + self._get_mutations(state).append(mutation_pb) + + def _delete_cells(self, column_family_id, columns, time_range=None, state=None): + """Helper for :meth:`delete_cell` and :meth:`delete_cells`. + + ``state`` is unused by :class:`DirectRow` but is used by + subclasses. - { - (family='fam', qualifier=b'col'): [b'value', (+1 more),], - (family='fam', qualifier=b'col2'): [b'other'], - } + :type column_family_id: str + :param column_family_id: The column family that contains the column + or columns with cells being deleted. Must be + of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type columns: :class:`list` of :class:`str` / + :func:`unicode `, or :class:`object` + :param columns: The columns within the column family that will have + cells deleted. If :attr:`ALL_COLUMNS` is used then + the entire column family will be deleted from the row. + + :type time_range: :class:`TimestampRange` + :param time_range: (Optional) The range of time within which cells + should be deleted. + + :type state: bool + :param state: (Optional) The state that is passed along to + :meth:`_get_mutations`. """ - output = ["{"] - for family, qualifier in self.get_column_components(): - cell_list = self[family, qualifier] - line = [f" (family={family!r}, qualifier={qualifier!r}): "] - if len(cell_list) == 0: - line.append("[],") - elif len(cell_list) == 1: - line.append(f"[{cell_list[0]}],") - else: - line.append(f"[{cell_list[0]}, (+{len(cell_list)-1} more)],") - output.append("".join(line)) - output.append("}") - return "\n".join(output) + mutations_list = self._get_mutations(state) + if columns is self.ALL_COLUMNS: + mutation_val = data_v2_pb2.Mutation.DeleteFromFamily( + family_name=column_family_id + ) + mutation_pb = data_v2_pb2.Mutation(delete_from_family=mutation_val) + mutations_list.append(mutation_pb) + else: + delete_kwargs = {} + if time_range is not None: + delete_kwargs["time_range"] = time_range.to_pb() - def __repr__(self): - cell_str_buffer = ["{"] - for family, qualifier in self.get_column_components(): - cell_list = self[family, qualifier] - repr_list = [cell.to_dict() for cell in cell_list] - cell_str_buffer.append(f" ('{family}', {qualifier!r}): {repr_list},") - cell_str_buffer.append("}") - cell_str = "\n".join(cell_str_buffer) - output = f"Row(key={self.row_key!r}, cells={cell_str})" - return output - - def to_dict(self) -> dict[str, Any]: - """ - Returns a dictionary representation of the cell in the Bigtable Row - proto format - - https://cloud.google.com/bigtable/docs/reference/data/rpc/google.bigtable.v2#row - """ - family_list = [] - for family_name, qualifier_dict in self._index.items(): - qualifier_list = [] - for qualifier_name, cell_list in qualifier_dict.items(): - cell_dicts = [cell.to_dict() for cell in cell_list] - qualifier_list.append( - {"qualifier": qualifier_name, "cells": cell_dicts} + to_append = [] + for column in columns: + column = _to_bytes(column) + # time_range will never change if present, but the rest of + # delete_kwargs will + delete_kwargs.update( + family_name=column_family_id, column_qualifier=column ) - family_list.append({"name": family_name, "columns": qualifier_list}) - return {"key": self.row_key, "families": family_list} - - # Sequence and Mapping methods - def __iter__(self): - """ - Allow iterating over all cells in the row - """ - return iter(self.cells) - - def __contains__(self, item): - """ - Implements `in` operator - - Works for both cells in the internal list, and `family` or - `(family, qualifier)` pairs associated with the cells - """ - if isinstance(item, _family_type): - return item in self._index - elif ( - isinstance(item, tuple) - and isinstance(item[0], _family_type) - and isinstance(item[1], (bytes, str)) - ): - q = item[1] if isinstance(item[1], bytes) else item[1].encode("utf-8") - return item[0] in self._index and q in self._index[item[0]] - # check if Cell is in Row - return item in self.cells - - @overload - def __getitem__( - self, - index: str | tuple[str, bytes | str], - ) -> list[Cell]: - # overload signature for type checking - pass - - @overload - def __getitem__(self, index: int) -> Cell: - # overload signature for type checking - pass - - @overload - def __getitem__(self, index: slice) -> list[Cell]: - # overload signature for type checking - pass - - def __getitem__(self, index): - """ - Implements [] indexing - - Supports indexing by family, (family, qualifier) pair, - numerical index, and index slicing - """ - if isinstance(index, _family_type): - return self.get_cells(family=index) - elif ( - isinstance(index, tuple) - and isinstance(index[0], _family_type) - and isinstance(index[1], (bytes, str)) - ): - return self.get_cells(family=index[0], qualifier=index[1]) - elif isinstance(index, int) or isinstance(index, slice): - # index is int or slice - return self.cells[index] - else: - raise TypeError( - "Index must be family_id, (family_id, qualifier), int, or slice" - ) + mutation_val = data_v2_pb2.Mutation.DeleteFromColumn(**delete_kwargs) + mutation_pb = data_v2_pb2.Mutation(delete_from_column=mutation_val) + to_append.append(mutation_pb) + + # We don't add the mutations until all columns have been + # processed without error. + mutations_list.extend(to_append) + + +class DirectRow(_SetDeleteRow): + """Google Cloud Bigtable Row for sending "direct" mutations. + + These mutations directly set or delete cell contents: + + * :meth:`set_cell` + * :meth:`delete` + * :meth:`delete_cell` + * :meth:`delete_cells` + + These methods can be used directly:: + + >>> row = table.row(b'row-key1') + >>> row.set_cell(u'fam', b'col1', b'cell-val') + >>> row.delete_cell(u'fam', b'col2') - def __len__(self): + .. note:: + + A :class:`DirectRow` accumulates mutations locally via the + :meth:`set_cell`, :meth:`delete`, :meth:`delete_cell` and + :meth:`delete_cells` methods. To actually send these mutations to the + Google Cloud Bigtable API, you must call :meth:`commit`. + + :type row_key: bytes + :param row_key: The key for the current row. + + :type table: :class:`Table ` + :param table: (Optional) The table that owns the row. This is + used for the :meth: `commit` only. Alternatively, + DirectRows can be persisted via + :meth:`~google.cloud.bigtable.table.Table.mutate_rows`. + """ + + def __init__(self, row_key, table=None): + super(DirectRow, self).__init__(row_key, table) + self._pb_mutations = [] + + def _get_mutations(self, state=None): # pylint: disable=unused-argument + """Gets the list of mutations for a given state. + + ``state`` is unused by :class:`DirectRow` but is used by + subclasses. + + :type state: bool + :param state: The state that the mutation should be + applied in. + + :rtype: list + :returns: The list to add new mutations to (for the current state). """ - Implements `len()` operator + return self._pb_mutations + + def get_mutations_size(self): + """Gets the total mutations size for current row + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_get_mutations_size] + :end-before: [END bigtable_api_row_get_mutations_size] + :dedent: 4 """ - return len(self.cells) - def get_column_components(self) -> list[tuple[str, bytes]]: + mutation_size = 0 + for mutation in self._get_mutations(): + mutation_size += mutation._pb.ByteSize() + + return mutation_size + + def set_cell(self, column_family_id, column, value, timestamp=None): + """Sets a value in this row. + + The cell is determined by the ``row_key`` of this :class:`DirectRow` + and the ``column``. The ``column`` must be in an existing + :class:`.ColumnFamily` (as determined by ``column_family_id``). + + .. note:: + + This method adds a mutation to the accumulated mutations on this + row, but does not make an API request. To actually + send an API request (with the mutations) to the Google Cloud + Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_set_cell] + :end-before: [END bigtable_api_row_set_cell] + :dedent: 4 + + :type column_family_id: str + :param column_family_id: The column family that contains the column. + Must be of the form + ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type column: bytes + :param column: The column within the column family where the cell + is located. + + :type value: bytes or :class:`int` + :param value: The value to set in the cell. If an integer is used, + will be interpreted as a 64-bit big-endian signed + integer (8 bytes). + + :type timestamp: :class:`datetime.datetime` + :param timestamp: (Optional) The timestamp of the operation. """ - Returns a list of (family, qualifier) pairs associated with the cells + self._set_cell(column_family_id, column, value, timestamp=timestamp, state=None) + + def delete(self): + """Deletes this row from the table. + + .. note:: - Pairs can be used for indexing + This method adds a mutation to the accumulated mutations on this + row, but does not make an API request. To actually + send an API request (with the mutations) to the Google Cloud + Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_delete] + :end-before: [END bigtable_api_row_delete] + :dedent: 4 """ - return [(f, q) for f in self._index for q in self._index[f]] + self._delete(state=None) - def __eq__(self, other): + def delete_cell(self, column_family_id, column, time_range=None): + """Deletes cell in this row. + + .. note:: + + This method adds a mutation to the accumulated mutations on this + row, but does not make an API request. To actually + send an API request (with the mutations) to the Google Cloud + Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_delete_cell] + :end-before: [END bigtable_api_row_delete_cell] + :dedent: 4 + + :type column_family_id: str + :param column_family_id: The column family that contains the column + or columns with cells being deleted. Must be + of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type column: bytes + :param column: The column within the column family that will have a + cell deleted. + + :type time_range: :class:`TimestampRange` + :param time_range: (Optional) The range of time within which cells + should be deleted. """ - Implements `==` operator - """ - # for performance reasons, check row metadata - # before checking individual cells - if not isinstance(other, Row): - return False - if self.row_key != other.row_key: - return False - if len(self.cells) != len(other.cells): - return False - components = self.get_column_components() - other_components = other.get_column_components() - if len(components) != len(other_components): - return False - if components != other_components: - return False - for family, qualifier in components: - if len(self[family, qualifier]) != len(other[family, qualifier]): - return False - # compare individual cell lists - if self.cells != other.cells: - return False - return True - - def __ne__(self, other) -> bool: - """ - Implements `!=` operator + self._delete_cells( + column_family_id, [column], time_range=time_range, state=None + ) + + def delete_cells(self, column_family_id, columns, time_range=None): + """Deletes cells in this row. + + .. note:: + + This method adds a mutation to the accumulated mutations on this + row, but does not make an API request. To actually + send an API request (with the mutations) to the Google Cloud + Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_delete_cells] + :end-before: [END bigtable_api_row_delete_cells] + :dedent: 4 + + :type column_family_id: str + :param column_family_id: The column family that contains the column + or columns with cells being deleted. Must be + of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type columns: :class:`list` of :class:`str` / + :func:`unicode `, or :class:`object` + :param columns: The columns within the column family that will have + cells deleted. If :attr:`ALL_COLUMNS` is used then + the entire column family will be deleted from the row. + + :type time_range: :class:`TimestampRange` + :param time_range: (Optional) The range of time within which cells + should be deleted. """ - return not self == other + self._delete_cells(column_family_id, columns, time_range=time_range, state=None) + def commit(self): + """Makes a ``MutateRow`` API request. -class _LastScannedRow(Row): - """A value used to indicate a scanned row that is not returned as part of - a query. + If no mutations have been created in the row, no request is made. - This is used internally to indicate progress in a scan, and improve retry - performance. It is not intended to be used directly by users. - """ + Mutations are applied atomically and in order, meaning that earlier + mutations can be masked / negated by later ones. Cells already present + in the row are left unchanged unless explicitly changed by a mutation. - def __init__(self, row_key): - super().__init__(row_key, []) + After committing the accumulated mutations, resets the local + mutations to an empty list. - def __eq__(self, other): - return isinstance(other, _LastScannedRow) + For example: + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_commit] + :end-before: [END bigtable_api_row_commit] + :dedent: 4 -@total_ordering -class Cell: - """ - Model class for cell data + :rtype: :class:`~google.rpc.status_pb2.Status` + :returns: A response status (`google.rpc.status_pb2.Status`) + representing success or failure of the row committed. + :raises: :exc:`~.table.TooManyMutationsError` if the number of + mutations is greater than 100,000. + """ + response = self._table.mutate_rows([self]) - Does not represent all data contained in the cell, only data returned by a - query. - Expected to be read-only to users, and written by backend - """ + self.clear() + + return response[0] + + def clear(self): + """Removes all currently accumulated mutations on the current row. - __slots__ = ( - "value", - "row_key", - "family", - "qualifier", - "timestamp_micros", - "labels", - ) - - def __init__( - self, - value: bytes, - row_key: bytes, - family: str, - qualifier: bytes | str, - timestamp_micros: int, - labels: list[str] | None = None, - ): - """ - Cell constructor - - Cell objects are not intended to be constructed by users. - They are returned by the Bigtable backend. + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_clear] + :end-before: [END bigtable_api_row_clear] + :dedent: 4 """ - self.value = value - self.row_key = row_key - self.family = family - if isinstance(qualifier, str): - qualifier = qualifier.encode() - self.qualifier = qualifier - self.timestamp_micros = timestamp_micros - self.labels = labels if labels is not None else [] + del self._pb_mutations[:] + + +class ConditionalRow(_SetDeleteRow): + """Google Cloud Bigtable Row for sending mutations conditionally. + + Each mutation has an associated state: :data:`True` or :data:`False`. + When :meth:`commit`-ed, the mutations for the :data:`True` + state will be applied if the filter matches any cells in + the row, otherwise the :data:`False` state will be applied. + + A :class:`ConditionalRow` accumulates mutations in the same way a + :class:`DirectRow` does: - def __int__(self) -> int: + * :meth:`set_cell` + * :meth:`delete` + * :meth:`delete_cell` + * :meth:`delete_cells` + + with the only change the extra ``state`` parameter:: + + >>> row_cond = table.row(b'row-key2', filter_=row_filter) + >>> row_cond.set_cell(u'fam', b'col', b'cell-val', state=True) + >>> row_cond.delete_cell(u'fam', b'col', state=False) + + .. note:: + + As with :class:`DirectRow`, to actually send these mutations to the + Google Cloud Bigtable API, you must call :meth:`commit`. + + :type row_key: bytes + :param row_key: The key for the current row. + + :type table: :class:`Table ` + :param table: The table that owns the row. + + :type filter_: :class:`.RowFilter` + :param filter_: Filter to be used for conditional mutations. + """ + + def __init__(self, row_key, table, filter_): + super(ConditionalRow, self).__init__(row_key, table) + self._filter = filter_ + self._true_pb_mutations = [] + self._false_pb_mutations = [] + + def _get_mutations(self, state=None): + """Gets the list of mutations for a given state. + + Over-ridden so that the state can be used in: + + * :meth:`set_cell` + * :meth:`delete` + * :meth:`delete_cell` + * :meth:`delete_cells` + + :type state: bool + :param state: The state that the mutation should be + applied in. + + :rtype: list + :returns: The list to add new mutations to (for the current state). """ - Allows casting cell to int - Interprets value as a 64-bit big-endian signed integer, as expected by - ReadModifyWrite increment rule + if state: + return self._true_pb_mutations + else: + return self._false_pb_mutations + + def commit(self): + """Makes a ``CheckAndMutateRow`` API request. + + If no mutations have been created in the row, no request is made. + + The mutations will be applied conditionally, based on whether the + filter matches any cells in the :class:`ConditionalRow` or not. (Each + method which adds a mutation has a ``state`` parameter for this + purpose.) + + Mutations are applied atomically and in order, meaning that earlier + mutations can be masked / negated by later ones. Cells already present + in the row are left unchanged unless explicitly changed by a mutation. + + After committing the accumulated mutations, resets the local + mutations. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_commit] + :end-before: [END bigtable_api_row_commit] + :dedent: 4 + + :rtype: bool + :returns: Flag indicating if the filter was matched (which also + indicates which set of mutations were applied by the server). + :raises: :class:`ValueError ` if the number of + mutations exceeds the :data:`MAX_MUTATIONS`. """ - return int.from_bytes(self.value, byteorder="big", signed=True) + true_mutations = self._get_mutations(state=True) + false_mutations = self._get_mutations(state=False) + num_true_mutations = len(true_mutations) + num_false_mutations = len(false_mutations) + if num_true_mutations == 0 and num_false_mutations == 0: + return + if num_true_mutations > MAX_MUTATIONS or num_false_mutations > MAX_MUTATIONS: + raise ValueError( + "Exceed the maximum allowable mutations (%d). Had %s true " + "mutations and %d false mutations." + % (MAX_MUTATIONS, num_true_mutations, num_false_mutations) + ) + + data_client = self._table._instance._client.table_data_client + resp = data_client.check_and_mutate_row( + table_name=self._table.name, + row_key=self._row_key, + predicate_filter=self._filter.to_pb(), + app_profile_id=self._table._app_profile_id, + true_mutations=true_mutations, + false_mutations=false_mutations, + ) + self.clear() + return resp.predicate_matched + + # pylint: disable=arguments-differ + def set_cell(self, column_family_id, column, value, timestamp=None, state=True): + """Sets a value in this row. + + The cell is determined by the ``row_key`` of this + :class:`ConditionalRow` and the ``column``. The ``column`` must be in + an existing :class:`.ColumnFamily` (as determined by + ``column_family_id``). + + .. note:: - def to_dict(self) -> dict[str, Any]: + This method adds a mutation to the accumulated mutations on this + row, but does not make an API request. To actually + send an API request (with the mutations) to the Google Cloud + Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_set_cell] + :end-before: [END bigtable_api_row_set_cell] + :dedent: 4 + + :type column_family_id: str + :param column_family_id: The column family that contains the column. + Must be of the form + ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type column: bytes + :param column: The column within the column family where the cell + is located. + + :type value: bytes or :class:`int` + :param value: The value to set in the cell. If an integer is used, + will be interpreted as a 64-bit big-endian signed + integer (8 bytes). + + :type timestamp: :class:`datetime.datetime` + :param timestamp: (Optional) The timestamp of the operation. + + :type state: bool + :param state: (Optional) The state that the mutation should be + applied in. Defaults to :data:`True`. """ - Returns a dictionary representation of the cell in the Bigtable Cell - proto format + self._set_cell( + column_family_id, column, value, timestamp=timestamp, state=state + ) + + def delete(self, state=True): + """Deletes this row from the table. + + .. note:: - https://cloud.google.com/bigtable/docs/reference/data/rpc/google.bigtable.v2#cell + This method adds a mutation to the accumulated mutations on this + row, but does not make an API request. To actually + send an API request (with the mutations) to the Google Cloud + Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_delete] + :end-before: [END bigtable_api_row_delete] + :dedent: 4 + + :type state: bool + :param state: (Optional) The state that the mutation should be + applied in. Defaults to :data:`True`. """ - cell_dict: dict[str, Any] = { - "value": self.value, - } - cell_dict["timestamp_micros"] = self.timestamp_micros - if self.labels: - cell_dict["labels"] = self.labels - return cell_dict + self._delete(state=state) + + def delete_cell(self, column_family_id, column, time_range=None, state=True): + """Deletes cell in this row. + + .. note:: + + This method adds a mutation to the accumulated mutations on this + row, but does not make an API request. To actually + send an API request (with the mutations) to the Google Cloud + Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_delete_cell] + :end-before: [END bigtable_api_row_delete_cell] + :dedent: 4 + + :type column_family_id: str + :param column_family_id: The column family that contains the column + or columns with cells being deleted. Must be + of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type column: bytes + :param column: The column within the column family that will have a + cell deleted. - def __str__(self) -> str: + :type time_range: :class:`TimestampRange` + :param time_range: (Optional) The range of time within which cells + should be deleted. + + :type state: bool + :param state: (Optional) The state that the mutation should be + applied in. Defaults to :data:`True`. """ - Allows casting cell to str - Prints encoded byte string, same as printing value directly. + self._delete_cells( + column_family_id, [column], time_range=time_range, state=state + ) + + def delete_cells(self, column_family_id, columns, time_range=None, state=True): + """Deletes cells in this row. + + .. note:: + + This method adds a mutation to the accumulated mutations on this + row, but does not make an API request. To actually + send an API request (with the mutations) to the Google Cloud + Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_delete_cells] + :end-before: [END bigtable_api_row_delete_cells] + :dedent: 4 + + :type column_family_id: str + :param column_family_id: The column family that contains the column + or columns with cells being deleted. Must be + of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type columns: :class:`list` of :class:`str` / + :func:`unicode `, or :class:`object` + :param columns: The columns within the column family that will have + cells deleted. If :attr:`ALL_COLUMNS` is used then the + entire column family will be deleted from the row. + + :type time_range: :class:`TimestampRange` + :param time_range: (Optional) The range of time within which cells + should be deleted. + + :type state: bool + :param state: (Optional) The state that the mutation should be + applied in. Defaults to :data:`True`. """ - return str(self.value) + self._delete_cells( + column_family_id, columns, time_range=time_range, state=state + ) - def __repr__(self): + # pylint: enable=arguments-differ + + def clear(self): + """Removes all currently accumulated mutations on the current row. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_clear] + :end-before: [END bigtable_api_row_clear] + :dedent: 4 """ - Returns a string representation of the cell + del self._true_pb_mutations[:] + del self._false_pb_mutations[:] + + +class AppendRow(Row): + """Google Cloud Bigtable Row for sending append mutations. + + These mutations are intended to augment the value of an existing cell + and uses the methods: + + * :meth:`append_cell_value` + * :meth:`increment_cell_value` + + The first works by appending bytes and the second by incrementing an + integer (stored in the cell as 8 bytes). In either case, if the + cell is empty, assumes the default empty value (empty string for + bytes or 0 for integer). + + :type row_key: bytes + :param row_key: The key for the current row. + + :type table: :class:`Table ` + :param table: The table that owns the row. + """ + + def __init__(self, row_key, table): + super(AppendRow, self).__init__(row_key, table) + self._rule_pb_list = [] + + def clear(self): + """Removes all currently accumulated modifications on current row. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_clear] + :end-before: [END bigtable_api_row_clear] + :dedent: 4 """ - return f"Cell(value={self.value!r}, row_key={self.row_key!r}, family='{self.family}', qualifier={self.qualifier!r}, timestamp_micros={self.timestamp_micros}, labels={self.labels})" + del self._rule_pb_list[:] + + def append_cell_value(self, column_family_id, column, value): + """Appends a value to an existing cell. + + .. note:: - """For Bigtable native ordering""" + This method adds a read-modify rule protobuf to the accumulated + read-modify rules on this row, but does not make an API + request. To actually send an API request (with the rules) to the + Google Cloud Bigtable API, call :meth:`commit`. - def __lt__(self, other) -> bool: + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_append_cell_value] + :end-before: [END bigtable_api_row_append_cell_value] + :dedent: 4 + + :type column_family_id: str + :param column_family_id: The column family that contains the column. + Must be of the form + ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type column: bytes + :param column: The column within the column family where the cell + is located. + + :type value: bytes + :param value: The value to append to the existing value in the cell. If + the targeted cell is unset, it will be treated as + containing the empty string. """ - Implements `<` operator + column = _to_bytes(column) + value = _to_bytes(value) + rule_pb = data_v2_pb2.ReadModifyWriteRule( + family_name=column_family_id, column_qualifier=column, append_value=value + ) + self._rule_pb_list.append(rule_pb) + + def increment_cell_value(self, column_family_id, column, int_value): + """Increments a value in an existing cell. + + Assumes the value in the cell is stored as a 64 bit integer + serialized to bytes. + + .. note:: + + This method adds a read-modify rule protobuf to the accumulated + read-modify rules on this row, but does not make an API + request. To actually send an API request (with the rules) to the + Google Cloud Bigtable API, call :meth:`commit`. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_increment_cell_value] + :end-before: [END bigtable_api_row_increment_cell_value] + :dedent: 4 + + :type column_family_id: str + :param column_family_id: The column family that contains the column. + Must be of the form + ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + + :type column: bytes + :param column: The column within the column family where the cell + is located. + + :type int_value: int + :param int_value: The value to increment the existing value in the cell + by. If the targeted cell is unset, it will be treated + as containing a zero. Otherwise, the targeted cell + must contain an 8-byte value (interpreted as a 64-bit + big-endian signed integer), or the entire request + will fail. """ - if not isinstance(other, Cell): - return NotImplemented - this_ordering = ( - self.family, - self.qualifier, - -self.timestamp_micros, - self.value, - self.labels, + column = _to_bytes(column) + rule_pb = data_v2_pb2.ReadModifyWriteRule( + family_name=column_family_id, + column_qualifier=column, + increment_amount=int_value, ) - other_ordering = ( - other.family, - other.qualifier, - -other.timestamp_micros, - other.value, - other.labels, + self._rule_pb_list.append(rule_pb) + + def commit(self): + """Makes a ``ReadModifyWriteRow`` API request. + + This commits modifications made by :meth:`append_cell_value` and + :meth:`increment_cell_value`. If no modifications were made, makes + no API request and just returns ``{}``. + + Modifies a row atomically, reading the latest existing + timestamp / value from the specified columns and writing a new value by + appending / incrementing. The new cell created uses either the current + server time or the highest timestamp of a cell in that column (if it + exceeds the server time). + + After committing the accumulated mutations, resets the local mutations. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_commit] + :end-before: [END bigtable_api_row_commit] + :dedent: 4 + + :rtype: dict + :returns: The new contents of all modified cells. Returned as a + dictionary of column families, each of which holds a + dictionary of columns. Each column contains a list of cells + modified. Each cell is represented with a two-tuple with the + value (in bytes) and the timestamp for the cell. + :raises: :class:`ValueError ` if the number of + mutations exceeds the :data:`MAX_MUTATIONS`. + """ + num_mutations = len(self._rule_pb_list) + if num_mutations == 0: + return {} + if num_mutations > MAX_MUTATIONS: + raise ValueError( + "%d total append mutations exceed the maximum " + "allowable %d." % (num_mutations, MAX_MUTATIONS) + ) + + data_client = self._table._instance._client.table_data_client + row_response = data_client.read_modify_write_row( + table_name=self._table.name, + row_key=self._row_key, + rules=self._rule_pb_list, + app_profile_id=self._table._app_profile_id, ) - return this_ordering < other_ordering - def __eq__(self, other) -> bool: + # Reset modifications after commit-ing request. + self.clear() + + # NOTE: We expect row_response.key == self._row_key but don't check. + return _parse_rmw_row_response(row_response) + + +def _parse_rmw_row_response(row_response): + """Parses the response to a ``ReadModifyWriteRow`` request. + + :type row_response: :class:`.data_v2_pb2.Row` + :param row_response: The response row (with only modified cells) from a + ``ReadModifyWriteRow`` request. + + :rtype: dict + :returns: The new contents of all modified cells. Returned as a + dictionary of column families, each of which holds a + dictionary of columns. Each column contains a list of cells + modified. Each cell is represented with a two-tuple with the + value (in bytes) and the timestamp for the cell. For example: + + .. code:: python + + { + u'col-fam-id': { + b'col-name1': [ + (b'cell-val', datetime.datetime(...)), + (b'cell-val-newer', datetime.datetime(...)), + ], + b'col-name2': [ + (b'altcol-cell-val', datetime.datetime(...)), + ], + }, + u'col-fam-id2': { + b'col-name3-but-other-fam': [ + (b'foo', datetime.datetime(...)), + ], + }, + } + """ + result = {} + for column_family in row_response.row.families: + column_family_id, curr_family = _parse_family_pb(column_family) + result[column_family_id] = curr_family + return result + + +def _parse_family_pb(family_pb): + """Parses a Family protobuf into a dictionary. + + :type family_pb: :class:`._generated.data_pb2.Family` + :param family_pb: A protobuf + + :rtype: tuple + :returns: A string and dictionary. The string is the name of the + column family and the dictionary has column names (within the + family) as keys and cell lists as values. Each cell is + represented with a two-tuple with the value (in bytes) and the + timestamp for the cell. For example: + + .. code:: python + + { + b'col-name1': [ + (b'cell-val', datetime.datetime(...)), + (b'cell-val-newer', datetime.datetime(...)), + ], + b'col-name2': [ + (b'altcol-cell-val', datetime.datetime(...)), + ], + } + """ + result = {} + for column in family_pb.columns: + result[column.qualifier] = cells = [] + for cell in column.cells: + val_pair = (cell.value, _datetime_from_microseconds(cell.timestamp_micros)) + cells.append(val_pair) + + return family_pb.name, result + + +class PartialRowData(object): + """Representation of partial row in a Google Cloud Bigtable Table. + + These are expected to be updated directly from a + :class:`._generated.bigtable_service_messages_pb2.ReadRowsResponse` + + :type row_key: bytes + :param row_key: The key for the row holding the (partial) data. + """ + + def __init__(self, row_key): + self._row_key = row_key + self._cells = {} + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return other._row_key == self._row_key and other._cells == self._cells + + def __ne__(self, other): + return not self == other + + def to_dict(self): + """Convert the cells to a dictionary. + + This is intended to be used with HappyBase, so the column family and + column qualiers are combined (with ``:``). + + :rtype: dict + :returns: Dictionary containing all the data in the cells of this row. """ - Implements `==` operator + result = {} + for column_family_id, columns in self._cells.items(): + for column_qual, cells in columns.items(): + key = _to_bytes(column_family_id) + b":" + _to_bytes(column_qual) + result[key] = cells + return result + + @property + def cells(self): + """Property returning all the cells accumulated on this partial row. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_data_cells] + :end-before: [END bigtable_api_row_data_cells] + :dedent: 4 + + :rtype: dict + :returns: Dictionary of the :class:`Cell` objects accumulated. This + dictionary has two-levels of keys (first for column families + and second for column names/qualifiers within a family). For + a given column, a list of :class:`Cell` objects is stored. """ - if not isinstance(other, Cell): - return NotImplemented - return ( - self.row_key == other.row_key - and self.family == other.family - and self.qualifier == other.qualifier - and self.value == other.value - and self.timestamp_micros == other.timestamp_micros - and len(self.labels) == len(other.labels) - and all([label in other.labels for label in self.labels]) - ) + return self._cells - def __ne__(self, other) -> bool: + @property + def row_key(self): + """Getter for the current (partial) row's key. + + :rtype: bytes + :returns: The current (partial) row's key. """ - Implements `!=` operator + return self._row_key + + def find_cells(self, column_family_id, column): + """Get a time series of cells stored on this instance. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_find_cells] + :end-before: [END bigtable_api_row_find_cells] + :dedent: 4 + + Args: + column_family_id (str): The ID of the column family. Must be of the + form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + column (bytes): The column within the column family where the cells + are located. + + Returns: + List[~google.cloud.bigtable.row_data.Cell]: The cells stored in the + specified column. + + Raises: + KeyError: If ``column_family_id`` is not among the cells stored + in this row. + KeyError: If ``column`` is not among the cells stored in this row + for the given ``column_family_id``. """ - return not self == other + try: + column_family = self._cells[column_family_id] + except KeyError: + raise KeyError(_MISSING_COLUMN_FAMILY.format(column_family_id)) + + try: + cells = column_family[column] + except KeyError: + raise KeyError(_MISSING_COLUMN.format(column, column_family_id)) + + return cells + + def cell_value(self, column_family_id, column, index=0): + """Get a single cell value stored on this instance. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_cell_value] + :end-before: [END bigtable_api_row_cell_value] + :dedent: 4 + + Args: + column_family_id (str): The ID of the column family. Must be of the + form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + column (bytes): The column within the column family where the cell + is located. + index (Optional[int]): The offset within the series of values. If + not specified, will return the first cell. - def __hash__(self): + Returns: + ~google.cloud.bigtable.row_data.Cell value: The cell value stored + in the specified column and specified index. + + Raises: + KeyError: If ``column_family_id`` is not among the cells stored + in this row. + KeyError: If ``column`` is not among the cells stored in this row + for the given ``column_family_id``. + IndexError: If ``index`` cannot be found within the cells stored + in this row for the given ``column_family_id``, ``column`` + pair. """ - Implements `hash()` function to fingerprint cell + cells = self.find_cells(column_family_id, column) + + try: + cell = cells[index] + except (TypeError, IndexError): + num_cells = len(cells) + msg = _MISSING_INDEX.format(index, column, column_family_id, num_cells) + raise IndexError(msg) + + return cell.value + + def cell_values(self, column_family_id, column, max_count=None): + """Get a time series of cells stored on this instance. + + For example: + + .. literalinclude:: snippets_table.py + :start-after: [START bigtable_api_row_cell_values] + :end-before: [END bigtable_api_row_cell_values] + :dedent: 4 + + Args: + column_family_id (str): The ID of the column family. Must be of the + form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. + column (bytes): The column within the column family where the cells + are located. + max_count (int): The maximum number of cells to use. + + Returns: + A generator which provides: cell.value, cell.timestamp_micros + for each cell in the list of cells + + Raises: + KeyError: If ``column_family_id`` is not among the cells stored + in this row. + KeyError: If ``column`` is not among the cells stored in this row + for the given ``column_family_id``. """ - return hash( - ( - self.row_key, - self.family, - self.qualifier, - self.value, - self.timestamp_micros, - tuple(self.labels), - ) + cells = self.find_cells(column_family_id, column) + if max_count is None: + max_count = len(cells) + + for index, cell in enumerate(cells): + if index == max_count: + break + + yield cell.value, cell.timestamp_micros + + +class Cell(object): + """Representation of a Google Cloud Bigtable Cell. + + :type value: bytes + :param value: The value stored in the cell. + + :type timestamp_micros: int + :param timestamp_micros: The timestamp_micros when the cell was stored. + + :type labels: list + :param labels: (Optional) List of strings. Labels applied to the cell. + """ + + def __init__(self, value, timestamp_micros, labels=None): + self.value = value + self.timestamp_micros = timestamp_micros + self.labels = list(labels) if labels is not None else [] + + @classmethod + def from_pb(cls, cell_pb): + """Create a new cell from a Cell protobuf. + + :type cell_pb: :class:`._generated.data_pb2.Cell` + :param cell_pb: The protobuf to convert. + + :rtype: :class:`Cell` + :returns: The cell corresponding to the protobuf. + """ + if cell_pb.labels: + return cls(cell_pb.value, cell_pb.timestamp_micros, labels=cell_pb.labels) + else: + return cls(cell_pb.value, cell_pb.timestamp_micros) + + @property + def timestamp(self): + return _datetime_from_microseconds(self.timestamp_micros) + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return NotImplemented + return ( + other.value == self.value + and other.timestamp_micros == self.timestamp_micros + and other.labels == self.labels + ) + + def __ne__(self, other): + return not self == other + + def __repr__(self): + return "<{name} value={value!r} timestamp={timestamp}>".format( + name=self.__class__.__name__, value=self.value, timestamp=self.timestamp ) + + +class InvalidChunk(RuntimeError): + """Exception raised to invalid chunk data from back-end.""" diff --git a/google/cloud/bigtable/deprecated/row_data.py b/google/cloud/bigtable/row_data.py similarity index 97% rename from google/cloud/bigtable/deprecated/row_data.py rename to google/cloud/bigtable/row_data.py index 9daa1ed8f..e11379108 100644 --- a/google/cloud/bigtable/deprecated/row_data.py +++ b/google/cloud/bigtable/row_data.py @@ -23,10 +23,10 @@ from google.api_core import retry from google.cloud._helpers import _to_bytes # type: ignore -from google.cloud.bigtable.deprecated.row_merger import _RowMerger, _State +from google.cloud.bigtable.row_merger import _RowMerger, _State from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 from google.cloud.bigtable_v2.types import data as data_v2_pb2 -from google.cloud.bigtable.deprecated.row import Cell, InvalidChunk, PartialRowData +from google.cloud.bigtable.row import Cell, InvalidChunk, PartialRowData # Some classes need to be re-exported here to keep backwards @@ -98,7 +98,7 @@ def _retry_read_rows_exception(exc): """The default retry strategy to be used on retry-able errors. Used by -:meth:`~google.cloud.bigtable.deprecated.row_data.PartialRowsData._read_next_response`. +:meth:`~google.cloud.bigtable.row_data.PartialRowsData._read_next_response`. """ @@ -157,7 +157,9 @@ def __init__(self, read_method, request, retry=DEFAULT_RETRY_READ_ROWS): # Otherwise there is a risk of entering an infinite loop that resets # the timeout counter just before it being triggered. The increment # by 1 second here is customary but should not be much less than that. - self.response_iterator = read_method(request, timeout=self.retry._deadline + 1) + self.response_iterator = read_method( + request, timeout=self.retry._deadline + 1, retry=self.retry + ) self.rows = {} diff --git a/google/cloud/bigtable/row_filters.py b/google/cloud/bigtable/row_filters.py index b2fae6971..53192acc8 100644 --- a/google/cloud/bigtable/row_filters.py +++ b/google/cloud/bigtable/row_filters.py @@ -13,25 +13,18 @@ # limitations under the License. """Filters for Google Cloud Bigtable Row classes.""" -from __future__ import annotations import struct -from typing import Any, Sequence, TYPE_CHECKING, overload -from abc import ABC, abstractmethod from google.cloud._helpers import _microseconds_from_datetime # type: ignore from google.cloud._helpers import _to_bytes # type: ignore from google.cloud.bigtable_v2.types import data as data_v2_pb2 -if TYPE_CHECKING: - # import dependencies when type checking - from datetime import datetime - _PACK_I64 = struct.Struct(">q").pack -class RowFilter(ABC): +class RowFilter(object): """Basic filter to apply to cells in a row. These values can be combined via :class:`RowFilterChain`, @@ -42,30 +35,15 @@ class RowFilter(ABC): This class is a do-nothing base class for all row filters. """ - def _to_pb(self) -> data_v2_pb2.RowFilter: - """Converts the row filter to a protobuf. - - Returns: The converted current object. - """ - return data_v2_pb2.RowFilter(**self.to_dict()) - - @abstractmethod - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - pass - - def __repr__(self) -> str: - return f"{self.__class__.__name__}()" - -class _BoolFilter(RowFilter, ABC): +class _BoolFilter(RowFilter): """Row filter that uses a boolean flag. :type flag: bool :param flag: An indicator if a setting is turned on or off. """ - def __init__(self, flag: bool): + def __init__(self, flag): self.flag = flag def __eq__(self, other): @@ -76,9 +54,6 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def __repr__(self) -> str: - return f"{self.__class__.__name__}(flag={self.flag})" - class SinkFilter(_BoolFilter): """Advanced row filter to skip parent filters. @@ -91,9 +66,13 @@ class SinkFilter(_BoolFilter): of a :class:`ConditionalRowFilter`. """ - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"sink": self.flag} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(sink=self.flag) class PassAllFilter(_BoolFilter): @@ -105,9 +84,13 @@ class PassAllFilter(_BoolFilter): completeness. """ - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"pass_all_filter": self.flag} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(pass_all_filter=self.flag) class BlockAllFilter(_BoolFilter): @@ -118,12 +101,16 @@ class BlockAllFilter(_BoolFilter): temporarily disabling just part of a filter. """ - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"block_all_filter": self.flag} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(block_all_filter=self.flag) -class _RegexFilter(RowFilter, ABC): +class _RegexFilter(RowFilter): """Row filter that uses a regular expression. The ``regex`` must be valid RE2 patterns. See Google's @@ -137,8 +124,8 @@ class _RegexFilter(RowFilter, ABC): will be encoded as ASCII. """ - def __init__(self, regex: str | bytes): - self.regex: bytes = _to_bytes(regex) + def __init__(self, regex): + self.regex = _to_bytes(regex) def __eq__(self, other): if not isinstance(other, self.__class__): @@ -148,9 +135,6 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def __repr__(self) -> str: - return f"{self.__class__.__name__}(regex={self.regex!r})" - class RowKeyRegexFilter(_RegexFilter): """Row filter for a row key regular expression. @@ -175,9 +159,13 @@ class RowKeyRegexFilter(_RegexFilter): since the row key is already specified. """ - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"row_key_regex_filter": self.regex} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(row_key_regex_filter=self.regex) class RowSampleFilter(RowFilter): @@ -188,8 +176,8 @@ class RowSampleFilter(RowFilter): interval ``(0, 1)`` The end points are excluded). """ - def __init__(self, sample: float): - self.sample: float = sample + def __init__(self, sample): + self.sample = sample def __eq__(self, other): if not isinstance(other, self.__class__): @@ -199,12 +187,13 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"row_sample_filter": self.sample} + def to_pb(self): + """Converts the row filter to a protobuf. - def __repr__(self) -> str: - return f"{self.__class__.__name__}(sample={self.sample})" + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(row_sample_filter=self.sample) class FamilyNameRegexFilter(_RegexFilter): @@ -222,9 +211,13 @@ class FamilyNameRegexFilter(_RegexFilter): used as a literal. """ - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"family_name_regex_filter": self.regex} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(family_name_regex_filter=self.regex) class ColumnQualifierRegexFilter(_RegexFilter): @@ -248,9 +241,13 @@ class ColumnQualifierRegexFilter(_RegexFilter): match this regex (irrespective of column family). """ - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"column_qualifier_regex_filter": self.regex} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(column_qualifier_regex_filter=self.regex) class TimestampRange(object): @@ -265,9 +262,9 @@ class TimestampRange(object): range. If omitted, no upper bound is used. """ - def __init__(self, start: "datetime" | None = None, end: "datetime" | None = None): - self.start: "datetime" | None = start - self.end: "datetime" | None = end + def __init__(self, start=None, end=None): + self.start = start + self.end = end def __eq__(self, other): if not isinstance(other, self.__class__): @@ -277,29 +274,23 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def _to_pb(self) -> data_v2_pb2.TimestampRange: + def to_pb(self): """Converts the :class:`TimestampRange` to a protobuf. - Returns: The converted current object. + :rtype: :class:`.data_v2_pb2.TimestampRange` + :returns: The converted current object. """ - return data_v2_pb2.TimestampRange(**self.to_dict()) - - def to_dict(self) -> dict[str, int]: - """Converts the timestamp range to a dict representation.""" timestamp_range_kwargs = {} if self.start is not None: - start_time = _microseconds_from_datetime(self.start) // 1000 * 1000 - timestamp_range_kwargs["start_timestamp_micros"] = start_time + timestamp_range_kwargs["start_timestamp_micros"] = ( + _microseconds_from_datetime(self.start) // 1000 * 1000 + ) if self.end is not None: end_time = _microseconds_from_datetime(self.end) if end_time % 1000 != 0: - # if not a whole milisecond value, round up end_time = end_time // 1000 * 1000 + 1000 timestamp_range_kwargs["end_timestamp_micros"] = end_time - return timestamp_range_kwargs - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(start={self.start}, end={self.end})" + return data_v2_pb2.TimestampRange(**timestamp_range_kwargs) class TimestampRangeFilter(RowFilter): @@ -309,8 +300,8 @@ class TimestampRangeFilter(RowFilter): :param range_: Range of time that cells should match against. """ - def __init__(self, start: "datetime" | None = None, end: "datetime" | None = None): - self.range_: TimestampRange = TimestampRange(start, end) + def __init__(self, range_): + self.range_ = range_ def __eq__(self, other): if not isinstance(other, self.__class__): @@ -320,22 +311,16 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def _to_pb(self) -> data_v2_pb2.RowFilter: + def to_pb(self): """Converts the row filter to a protobuf. First converts the ``range_`` on the current object to a protobuf and then uses it in the ``timestamp_range_filter`` field. - Returns: The converted current object. + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. """ - return data_v2_pb2.RowFilter(timestamp_range_filter=self.range_._to_pb()) - - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"timestamp_range_filter": self.range_.to_dict()} - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(start={self.range_.start!r}, end={self.range_.end!r})" + return data_v2_pb2.RowFilter(timestamp_range_filter=self.range_.to_pb()) class ColumnRangeFilter(RowFilter): @@ -345,72 +330,71 @@ class ColumnRangeFilter(RowFilter): By default, we include them both, but this can be changed with optional flags. - :type family_id: str - :param family_id: The column family that contains the columns. Must + :type column_family_id: str + :param column_family_id: The column family that contains the columns. Must be of the form ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``. - :type start_qualifier: bytes - :param start_qualifier: The start of the range of columns. If no value is + :type start_column: bytes + :param start_column: The start of the range of columns. If no value is used, the backend applies no upper bound to the values. - :type end_qualifier: bytes - :param end_qualifier: The end of the range of columns. If no value is used, + :type end_column: bytes + :param end_column: The end of the range of columns. If no value is used, the backend applies no upper bound to the values. :type inclusive_start: bool :param inclusive_start: Boolean indicating if the start column should be included in the range (or excluded). Defaults - to :data:`True` if ``start_qualifier`` is passed and + to :data:`True` if ``start_column`` is passed and no ``inclusive_start`` was given. :type inclusive_end: bool :param inclusive_end: Boolean indicating if the end column should be included in the range (or excluded). Defaults - to :data:`True` if ``end_qualifier`` is passed and + to :data:`True` if ``end_column`` is passed and no ``inclusive_end`` was given. :raises: :class:`ValueError ` if ``inclusive_start`` - is set but no ``start_qualifier`` is given or if ``inclusive_end`` - is set but no ``end_qualifier`` is given + is set but no ``start_column`` is given or if ``inclusive_end`` + is set but no ``end_column`` is given """ def __init__( self, - family_id: str, - start_qualifier: bytes | None = None, - end_qualifier: bytes | None = None, - inclusive_start: bool | None = None, - inclusive_end: bool | None = None, + column_family_id, + start_column=None, + end_column=None, + inclusive_start=None, + inclusive_end=None, ): + self.column_family_id = column_family_id + if inclusive_start is None: inclusive_start = True - elif start_qualifier is None: + elif start_column is None: raise ValueError( - "inclusive_start was specified but no start_qualifier was given." + "Inclusive start was specified but no " "start column was given." ) + self.start_column = start_column + self.inclusive_start = inclusive_start + if inclusive_end is None: inclusive_end = True - elif end_qualifier is None: + elif end_column is None: raise ValueError( - "inclusive_end was specified but no end_qualifier was given." + "Inclusive end was specified but no " "end column was given." ) - - self.family_id = family_id - - self.start_qualifier = start_qualifier - self.inclusive_start = inclusive_start - - self.end_qualifier = end_qualifier + self.end_column = end_column self.inclusive_end = inclusive_end def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented return ( - other.family_id == self.family_id - and other.start_qualifier == self.start_qualifier - and other.end_qualifier == self.end_qualifier + other.column_family_id == self.column_family_id + and other.start_column == self.start_column + and other.end_column == self.end_column and other.inclusive_start == self.inclusive_start and other.inclusive_end == self.inclusive_end ) @@ -418,41 +402,31 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def _to_pb(self) -> data_v2_pb2.RowFilter: + def to_pb(self): """Converts the row filter to a protobuf. First converts to a :class:`.data_v2_pb2.ColumnRange` and then uses it in the ``column_range_filter`` field. - Returns: The converted current object. + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. """ - column_range = data_v2_pb2.ColumnRange(**self.range_to_dict()) - return data_v2_pb2.RowFilter(column_range_filter=column_range) - - def range_to_dict(self) -> dict[str, str | bytes]: - """Converts the column range range to a dict representation.""" - column_range_kwargs: dict[str, str | bytes] = {} - column_range_kwargs["family_name"] = self.family_id - if self.start_qualifier is not None: + column_range_kwargs = {"family_name": self.column_family_id} + if self.start_column is not None: if self.inclusive_start: key = "start_qualifier_closed" else: key = "start_qualifier_open" - column_range_kwargs[key] = _to_bytes(self.start_qualifier) - if self.end_qualifier is not None: + column_range_kwargs[key] = _to_bytes(self.start_column) + if self.end_column is not None: if self.inclusive_end: key = "end_qualifier_closed" else: key = "end_qualifier_open" - column_range_kwargs[key] = _to_bytes(self.end_qualifier) - return column_range_kwargs - - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"column_range_filter": self.range_to_dict()} + column_range_kwargs[key] = _to_bytes(self.end_column) - def __repr__(self) -> str: - return f"{self.__class__.__name__}(family_id='{self.family_id}', start_qualifier={self.start_qualifier!r}, end_qualifier={self.end_qualifier!r}, inclusive_start={self.inclusive_start}, inclusive_end={self.inclusive_end})" + column_range = data_v2_pb2.ColumnRange(**column_range_kwargs) + return data_v2_pb2.RowFilter(column_range_filter=column_range) class ValueRegexFilter(_RegexFilter): @@ -476,64 +450,29 @@ class ValueRegexFilter(_RegexFilter): match this regex. String values will be encoded as ASCII. """ - def to_dict(self) -> dict[str, bytes]: - """Converts the row filter to a dict representation.""" - return {"value_regex_filter": self.regex} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(value_regex_filter=self.regex) -class LiteralValueFilter(ValueRegexFilter): +class ExactValueFilter(ValueRegexFilter): """Row filter for an exact value. :type value: bytes or str or int :param value: - a literal string, integer, or the equivalent bytes. - Integer values will be packed into signed 8-bytes. + a literal string encodable as ASCII, or the + equivalent bytes, or an integer (which will be packed into 8-bytes). """ - def __init__(self, value: bytes | str | int): + def __init__(self, value): if isinstance(value, int): value = _PACK_I64(value) - elif isinstance(value, str): - value = value.encode("utf-8") - value = self._write_literal_regex(value) - super(LiteralValueFilter, self).__init__(value) - - @staticmethod - def _write_literal_regex(input_bytes: bytes) -> bytes: - """ - Escape re2 special characters from literal bytes. - - Extracted from: re2 QuoteMeta: - https://github.com/google/re2/blob/70f66454c255080a54a8da806c52d1f618707f8a/re2/re2.cc#L456 - """ - result = bytearray() - for byte in input_bytes: - # If this is the part of a UTF8 or Latin1 character, we need \ - # to copy this byte without escaping. Experimentally this is \ - # what works correctly with the regexp library. \ - utf8_latin1_check = (byte & 128) == 0 - if ( - (byte < ord("a") or byte > ord("z")) - and (byte < ord("A") or byte > ord("Z")) - and (byte < ord("0") or byte > ord("9")) - and byte != ord("_") - and utf8_latin1_check - ): - if byte == 0: - # Special handling for null chars. - # Note that this special handling is not strictly required for RE2, - # but this quoting is required for other regexp libraries such as - # PCRE. - # Can't use "\\0" since the next character might be a digit. - result.extend([ord("\\"), ord("x"), ord("0"), ord("0")]) - continue - result.append(ord(b"\\")) - result.append(byte) - return bytes(result) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(value={self.regex!r})" + super(ExactValueFilter, self).__init__(value) class ValueRangeFilter(RowFilter): @@ -571,29 +510,25 @@ class ValueRangeFilter(RowFilter): """ def __init__( - self, - start_value: bytes | int | None = None, - end_value: bytes | int | None = None, - inclusive_start: bool | None = None, - inclusive_end: bool | None = None, + self, start_value=None, end_value=None, inclusive_start=None, inclusive_end=None ): if inclusive_start is None: inclusive_start = True elif start_value is None: raise ValueError( - "inclusive_start was specified but no start_value was given." - ) - if inclusive_end is None: - inclusive_end = True - elif end_value is None: - raise ValueError( - "inclusive_end was specified but no end_qualifier was given." + "Inclusive start was specified but no " "start value was given." ) if isinstance(start_value, int): start_value = _PACK_I64(start_value) self.start_value = start_value self.inclusive_start = inclusive_start + if inclusive_end is None: + inclusive_end = True + elif end_value is None: + raise ValueError( + "Inclusive end was specified but no " "end value was given." + ) if isinstance(end_value, int): end_value = _PACK_I64(end_value) self.end_value = end_value @@ -612,19 +547,15 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def _to_pb(self) -> data_v2_pb2.RowFilter: + def to_pb(self): """Converts the row filter to a protobuf. First converts to a :class:`.data_v2_pb2.ValueRange` and then uses it to create a row filter protobuf. - Returns: The converted current object. + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. """ - value_range = data_v2_pb2.ValueRange(**self.range_to_dict()) - return data_v2_pb2.RowFilter(value_range_filter=value_range) - - def range_to_dict(self) -> dict[str, bytes]: - """Converts the value range range to a dict representation.""" value_range_kwargs = {} if self.start_value is not None: if self.inclusive_start: @@ -638,17 +569,12 @@ def range_to_dict(self) -> dict[str, bytes]: else: key = "end_value_open" value_range_kwargs[key] = _to_bytes(self.end_value) - return value_range_kwargs - - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"value_range_filter": self.range_to_dict()} - def __repr__(self) -> str: - return f"{self.__class__.__name__}(start_value={self.start_value!r}, end_value={self.end_value!r}, inclusive_start={self.inclusive_start}, inclusive_end={self.inclusive_end})" + value_range = data_v2_pb2.ValueRange(**value_range_kwargs) + return data_v2_pb2.RowFilter(value_range_filter=value_range) -class _CellCountFilter(RowFilter, ABC): +class _CellCountFilter(RowFilter): """Row filter that uses an integer count of cells. The cell count is used as an offset or a limit for the number @@ -658,7 +584,7 @@ class _CellCountFilter(RowFilter, ABC): :param num_cells: An integer count / offset / limit. """ - def __init__(self, num_cells: int): + def __init__(self, num_cells): self.num_cells = num_cells def __eq__(self, other): @@ -669,9 +595,6 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def __repr__(self) -> str: - return f"{self.__class__.__name__}(num_cells={self.num_cells})" - class CellsRowOffsetFilter(_CellCountFilter): """Row filter to skip cells in a row. @@ -680,9 +603,13 @@ class CellsRowOffsetFilter(_CellCountFilter): :param num_cells: Skips the first N cells of the row. """ - def to_dict(self) -> dict[str, int]: - """Converts the row filter to a dict representation.""" - return {"cells_per_row_offset_filter": self.num_cells} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(cells_per_row_offset_filter=self.num_cells) class CellsRowLimitFilter(_CellCountFilter): @@ -692,9 +619,13 @@ class CellsRowLimitFilter(_CellCountFilter): :param num_cells: Matches only the first N cells of the row. """ - def to_dict(self) -> dict[str, int]: - """Converts the row filter to a dict representation.""" - return {"cells_per_row_limit_filter": self.num_cells} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(cells_per_row_limit_filter=self.num_cells) class CellsColumnLimitFilter(_CellCountFilter): @@ -706,9 +637,13 @@ class CellsColumnLimitFilter(_CellCountFilter): timestamps of each cell. """ - def to_dict(self) -> dict[str, int]: - """Converts the row filter to a dict representation.""" - return {"cells_per_column_limit_filter": self.num_cells} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(cells_per_column_limit_filter=self.num_cells) class StripValueTransformerFilter(_BoolFilter): @@ -720,9 +655,13 @@ class StripValueTransformerFilter(_BoolFilter): transformer than a generic query / filter. """ - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"strip_value_transformer": self.flag} + def to_pb(self): + """Converts the row filter to a protobuf. + + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(strip_value_transformer=self.flag) class ApplyLabelFilter(RowFilter): @@ -744,7 +683,7 @@ class ApplyLabelFilter(RowFilter): ``[a-z0-9\\-]+``. """ - def __init__(self, label: str): + def __init__(self, label): self.label = label def __eq__(self, other): @@ -755,15 +694,16 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_dict(self) -> dict[str, str]: - """Converts the row filter to a dict representation.""" - return {"apply_label_transformer": self.label} + def to_pb(self): + """Converts the row filter to a protobuf. - def __repr__(self) -> str: - return f"{self.__class__.__name__}(label={self.label})" + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. + """ + return data_v2_pb2.RowFilter(apply_label_transformer=self.label) -class _FilterCombination(RowFilter, Sequence[RowFilter], ABC): +class _FilterCombination(RowFilter): """Chain of row filters. Sends rows through several filters in sequence. The filters are "chained" @@ -774,10 +714,10 @@ class _FilterCombination(RowFilter, Sequence[RowFilter], ABC): :param filters: List of :class:`RowFilter` """ - def __init__(self, filters: list[RowFilter] | None = None): + def __init__(self, filters=None): if filters is None: filters = [] - self.filters: list[RowFilter] = filters + self.filters = filters def __eq__(self, other): if not isinstance(other, self.__class__): @@ -787,38 +727,6 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def __len__(self) -> int: - return len(self.filters) - - @overload - def __getitem__(self, index: int) -> RowFilter: - # overload signature for type checking - pass - - @overload - def __getitem__(self, index: slice) -> list[RowFilter]: - # overload signature for type checking - pass - - def __getitem__(self, index): - return self.filters[index] - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(filters={self.filters})" - - def __str__(self) -> str: - """ - Returns a string representation of the filter chain. - - Adds line breaks between each sub-filter for readability. - """ - output = [f"{self.__class__.__name__}(["] - for filter_ in self.filters: - filter_lines = f"{filter_},".splitlines() - output.extend([f" {line}" for line in filter_lines]) - output.append("])") - return "\n".join(output) - class RowFilterChain(_FilterCombination): """Chain of row filters. @@ -831,20 +739,17 @@ class RowFilterChain(_FilterCombination): :param filters: List of :class:`RowFilter` """ - def _to_pb(self) -> data_v2_pb2.RowFilter: + def to_pb(self): """Converts the row filter to a protobuf. - Returns: The converted current object. + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. """ chain = data_v2_pb2.RowFilter.Chain( - filters=[row_filter._to_pb() for row_filter in self.filters] + filters=[row_filter.to_pb() for row_filter in self.filters] ) return data_v2_pb2.RowFilter(chain=chain) - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"chain": {"filters": [f.to_dict() for f in self.filters]}} - class RowFilterUnion(_FilterCombination): """Union of row filters. @@ -859,58 +764,50 @@ class RowFilterUnion(_FilterCombination): :param filters: List of :class:`RowFilter` """ - def _to_pb(self) -> data_v2_pb2.RowFilter: + def to_pb(self): """Converts the row filter to a protobuf. - Returns: The converted current object. + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. """ interleave = data_v2_pb2.RowFilter.Interleave( - filters=[row_filter._to_pb() for row_filter in self.filters] + filters=[row_filter.to_pb() for row_filter in self.filters] ) return data_v2_pb2.RowFilter(interleave=interleave) - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"interleave": {"filters": [f.to_dict() for f in self.filters]}} - class ConditionalRowFilter(RowFilter): """Conditional row filter which exhibits ternary behavior. - Executes one of two filters based on another filter. If the ``predicate_filter`` + Executes one of two filters based on another filter. If the ``base_filter`` returns any cells in the row, then ``true_filter`` is executed. If not, then ``false_filter`` is executed. .. note:: - The ``predicate_filter`` does not execute atomically with the true and false + The ``base_filter`` does not execute atomically with the true and false filters, which may lead to inconsistent or unexpected results. Additionally, executing a :class:`ConditionalRowFilter` has poor performance on the server, especially when ``false_filter`` is set. - :type predicate_filter: :class:`RowFilter` - :param predicate_filter: The filter to condition on before executing the + :type base_filter: :class:`RowFilter` + :param base_filter: The filter to condition on before executing the true/false filters. :type true_filter: :class:`RowFilter` :param true_filter: (Optional) The filter to execute if there are any cells - matching ``predicate_filter``. If not provided, no results + matching ``base_filter``. If not provided, no results will be returned in the true case. :type false_filter: :class:`RowFilter` :param false_filter: (Optional) The filter to execute if there are no cells - matching ``predicate_filter``. If not provided, no results + matching ``base_filter``. If not provided, no results will be returned in the false case. """ - def __init__( - self, - predicate_filter: RowFilter, - true_filter: RowFilter | None = None, - false_filter: RowFilter | None = None, - ): - self.predicate_filter = predicate_filter + def __init__(self, base_filter, true_filter=None, false_filter=None): + self.base_filter = base_filter self.true_filter = true_filter self.false_filter = false_filter @@ -918,7 +815,7 @@ def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented return ( - other.predicate_filter == self.predicate_filter + other.base_filter == self.base_filter and other.true_filter == self.true_filter and other.false_filter == self.false_filter ) @@ -926,43 +823,16 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def _to_pb(self) -> data_v2_pb2.RowFilter: + def to_pb(self): """Converts the row filter to a protobuf. - Returns: The converted current object. + :rtype: :class:`.data_v2_pb2.RowFilter` + :returns: The converted current object. """ - condition_kwargs = {"predicate_filter": self.predicate_filter._to_pb()} + condition_kwargs = {"predicate_filter": self.base_filter.to_pb()} if self.true_filter is not None: - condition_kwargs["true_filter"] = self.true_filter._to_pb() + condition_kwargs["true_filter"] = self.true_filter.to_pb() if self.false_filter is not None: - condition_kwargs["false_filter"] = self.false_filter._to_pb() + condition_kwargs["false_filter"] = self.false_filter.to_pb() condition = data_v2_pb2.RowFilter.Condition(**condition_kwargs) return data_v2_pb2.RowFilter(condition=condition) - - def condition_to_dict(self) -> dict[str, Any]: - """Converts the condition to a dict representation.""" - condition_kwargs = {"predicate_filter": self.predicate_filter.to_dict()} - if self.true_filter is not None: - condition_kwargs["true_filter"] = self.true_filter.to_dict() - if self.false_filter is not None: - condition_kwargs["false_filter"] = self.false_filter.to_dict() - return condition_kwargs - - def to_dict(self) -> dict[str, Any]: - """Converts the row filter to a dict representation.""" - return {"condition": self.condition_to_dict()} - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(predicate_filter={self.predicate_filter!r}, true_filter={self.true_filter!r}, false_filter={self.false_filter!r})" - - def __str__(self) -> str: - output = [f"{self.__class__.__name__}("] - for filter_type in ("predicate_filter", "true_filter", "false_filter"): - filter_ = getattr(self, filter_type) - if filter_ is None: - continue - # add the new filter set, adding indentations for readability - filter_lines = f"{filter_type}={filter_},".splitlines() - output.extend(f" {line}" for line in filter_lines) - output.append(")") - return "\n".join(output) diff --git a/google/cloud/bigtable/deprecated/row_merger.py b/google/cloud/bigtable/row_merger.py similarity index 99% rename from google/cloud/bigtable/deprecated/row_merger.py rename to google/cloud/bigtable/row_merger.py index d29d64eb2..515b91df7 100644 --- a/google/cloud/bigtable/deprecated/row_merger.py +++ b/google/cloud/bigtable/row_merger.py @@ -1,6 +1,6 @@ from enum import Enum from collections import OrderedDict -from google.cloud.bigtable.deprecated.row import Cell, PartialRowData, InvalidChunk +from google.cloud.bigtable.row import Cell, PartialRowData, InvalidChunk _MISSING_COLUMN_FAMILY = "Column family {} is not among the cells stored in this row." _MISSING_COLUMN = ( diff --git a/google/cloud/bigtable/deprecated/row_set.py b/google/cloud/bigtable/row_set.py similarity index 100% rename from google/cloud/bigtable/deprecated/row_set.py rename to google/cloud/bigtable/row_set.py diff --git a/google/cloud/bigtable/deprecated/table.py b/google/cloud/bigtable/table.py similarity index 95% rename from google/cloud/bigtable/deprecated/table.py rename to google/cloud/bigtable/table.py index cf60b066e..e3191a729 100644 --- a/google/cloud/bigtable/deprecated/table.py +++ b/google/cloud/bigtable/table.py @@ -28,24 +28,24 @@ from google.api_core.retry import if_exception_type from google.api_core.retry import Retry from google.cloud._helpers import _to_bytes # type: ignore -from google.cloud.bigtable.deprecated.backup import Backup -from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb -from google.cloud.bigtable.deprecated.column_family import ColumnFamily -from google.cloud.bigtable.deprecated.batcher import MutationsBatcher -from google.cloud.bigtable.deprecated.batcher import FLUSH_COUNT, MAX_ROW_BYTES -from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo -from google.cloud.bigtable.deprecated.policy import Policy -from google.cloud.bigtable.deprecated.row import AppendRow -from google.cloud.bigtable.deprecated.row import ConditionalRow -from google.cloud.bigtable.deprecated.row import DirectRow -from google.cloud.bigtable.deprecated.row_data import ( +from google.cloud.bigtable.backup import Backup +from google.cloud.bigtable.column_family import _gc_rule_from_pb +from google.cloud.bigtable.column_family import ColumnFamily +from google.cloud.bigtable.batcher import MutationsBatcher +from google.cloud.bigtable.batcher import FLUSH_COUNT, MAX_MUTATION_SIZE +from google.cloud.bigtable.encryption_info import EncryptionInfo +from google.cloud.bigtable.policy import Policy +from google.cloud.bigtable.row import AppendRow +from google.cloud.bigtable.row import ConditionalRow +from google.cloud.bigtable.row import DirectRow +from google.cloud.bigtable.row_data import ( PartialRowsData, _retriable_internal_server_error, ) -from google.cloud.bigtable.deprecated.row_data import DEFAULT_RETRY_READ_ROWS -from google.cloud.bigtable.deprecated.row_set import RowSet -from google.cloud.bigtable.deprecated.row_set import RowRange -from google.cloud.bigtable.deprecated import enums +from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS +from google.cloud.bigtable.row_set import RowSet +from google.cloud.bigtable.row_set import RowRange +from google.cloud.bigtable import enums from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient from google.cloud.bigtable_admin_v2.types import table as admin_messages_v2_pb2 @@ -88,7 +88,7 @@ class _BigtableRetryableError(Exception): ) """The default retry strategy to be used on retry-able errors. -Used by :meth:`~google.cloud.bigtable.deprecated.table.Table.mutate_rows`. +Used by :meth:`~google.cloud.bigtable.table.Table.mutate_rows`. """ @@ -119,7 +119,7 @@ class Table(object): :type table_id: str :param table_id: The ID of the table. - :type instance: :class:`~google.cloud.bigtable.deprecated.instance.Instance` + :type instance: :class:`~google.cloud.bigtable.instance.Instance` :param instance: The instance that owns the table. :type app_profile_id: str @@ -172,7 +172,7 @@ def get_iam_policy(self): :end-before: [END bigtable_api_table_get_iam_policy] :dedent: 4 - :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :rtype: :class:`google.cloud.bigtable.policy.Policy` :returns: The current IAM policy of this table. """ table_client = self._instance._client.table_admin_client @@ -184,7 +184,7 @@ def set_iam_policy(self, policy): existing policy. For more information about policy, please see documentation of - class `google.cloud.bigtable.deprecated.policy.Policy` + class `google.cloud.bigtable.policy.Policy` For example: @@ -193,11 +193,11 @@ class `google.cloud.bigtable.deprecated.policy.Policy` :end-before: [END bigtable_api_table_set_iam_policy] :dedent: 4 - :type policy: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :type policy: :class:`google.cloud.bigtable.policy.Policy` :param policy: A new IAM policy to replace the current IAM policy of this table. - :rtype: :class:`google.cloud.bigtable.deprecated.policy.Policy` + :rtype: :class:`google.cloud.bigtable.policy.Policy` :returns: The current IAM policy of this table. """ table_client = self._instance._client.table_admin_client @@ -271,7 +271,7 @@ def row(self, row_key, filter_=None, append=False): .. warning:: At most one of ``filter_`` and ``append`` can be used in a - :class:`~google.cloud.bigtable.deprecated.row.Row`. + :class:`~google.cloud.bigtable.row.Row`. :type row_key: bytes :param row_key: The key for the row being created. @@ -284,7 +284,7 @@ def row(self, row_key, filter_=None, append=False): :param append: (Optional) Flag to determine if the row should be used for append mutations. - :rtype: :class:`~google.cloud.bigtable.deprecated.row.Row` + :rtype: :class:`~google.cloud.bigtable.row.Row` :returns: A row owned by this table. :raises: :class:`ValueError ` if both ``filter_`` and ``append`` are used. @@ -307,7 +307,7 @@ def row(self, row_key, filter_=None, append=False): return DirectRow(row_key, self) def append_row(self, row_key): - """Create a :class:`~google.cloud.bigtable.deprecated.row.AppendRow` associated with this table. + """Create a :class:`~google.cloud.bigtable.row.AppendRow` associated with this table. For example: @@ -325,7 +325,7 @@ def append_row(self, row_key): return AppendRow(row_key, self) def direct_row(self, row_key): - """Create a :class:`~google.cloud.bigtable.deprecated.row.DirectRow` associated with this table. + """Create a :class:`~google.cloud.bigtable.row.DirectRow` associated with this table. For example: @@ -343,7 +343,7 @@ def direct_row(self, row_key): return DirectRow(row_key, self) def conditional_row(self, row_key, filter_): - """Create a :class:`~google.cloud.bigtable.deprecated.row.ConditionalRow` associated with this table. + """Create a :class:`~google.cloud.bigtable.row.ConditionalRow` associated with this table. For example: @@ -515,7 +515,7 @@ def get_encryption_info(self): :rtype: dict :returns: Dictionary of encryption info for this table. Keys are cluster ids and - values are tuples of :class:`google.cloud.bigtable.deprecated.encryption.EncryptionInfo` instances. + values are tuples of :class:`google.cloud.bigtable.encryption.EncryptionInfo` instances. """ ENCRYPTION_VIEW = enums.Table.View.ENCRYPTION_VIEW table_client = self._instance._client.table_admin_client @@ -844,7 +844,9 @@ def drop_by_prefix(self, row_key_prefix, timeout=None): request={"name": self.name, "row_key_prefix": _to_bytes(row_key_prefix)} ) - def mutations_batcher(self, flush_count=FLUSH_COUNT, max_row_bytes=MAX_ROW_BYTES): + def mutations_batcher( + self, flush_count=FLUSH_COUNT, max_row_bytes=MAX_MUTATION_SIZE + ): """Factory to create a mutation batcher associated with this instance. For example: @@ -967,7 +969,7 @@ def list_backups(self, cluster_id=None, filter_=None, order_by=None, page_size=0 number of resources in a page. :rtype: :class:`~google.api_core.page_iterator.Iterator` - :returns: Iterator of :class:`~google.cloud.bigtable.deprecated.backup.Backup` + :returns: Iterator of :class:`~google.cloud.bigtable.backup.Backup` resources within the current Instance. :raises: :class:`ValueError ` if one of the returned Backups' name is not of the expected format. @@ -1367,8 +1369,8 @@ def _check_row_table_name(table_name, row): :type table_name: str :param table_name: The name of the table. - :type row: :class:`~google.cloud.bigtable.deprecated.row.Row` - :param row: An instance of :class:`~google.cloud.bigtable.deprecated.row.Row` + :type row: :class:`~google.cloud.bigtable.row.Row` + :param row: An instance of :class:`~google.cloud.bigtable.row.Row` subclasses. :raises: :exc:`~.table.TableMismatchError` if the row does not belong to @@ -1384,8 +1386,8 @@ def _check_row_table_name(table_name, row): def _check_row_type(row): """Checks that a row is an instance of :class:`.DirectRow`. - :type row: :class:`~google.cloud.bigtable.deprecated.row.Row` - :param row: An instance of :class:`~google.cloud.bigtable.deprecated.row.Row` + :type row: :class:`~google.cloud.bigtable.row.Row` + :param row: An instance of :class:`~google.cloud.bigtable.row.Row` subclasses. :raises: :class:`TypeError ` if the row is not an diff --git a/google/cloud/bigtable_admin/__init__.py b/google/cloud/bigtable_admin/__init__.py index 6ddc6acb2..0ba93ec63 100644 --- a/google/cloud/bigtable_admin/__init__.py +++ b/google/cloud/bigtable_admin/__init__.py @@ -200,6 +200,7 @@ from google.cloud.bigtable_admin_v2.types.instance import Instance from google.cloud.bigtable_admin_v2.types.table import Backup from google.cloud.bigtable_admin_v2.types.table import BackupInfo +from google.cloud.bigtable_admin_v2.types.table import ChangeStreamConfig from google.cloud.bigtable_admin_v2.types.table import ColumnFamily from google.cloud.bigtable_admin_v2.types.table import EncryptionInfo from google.cloud.bigtable_admin_v2.types.table import GcRule @@ -282,6 +283,7 @@ "Instance", "Backup", "BackupInfo", + "ChangeStreamConfig", "ColumnFamily", "EncryptionInfo", "GcRule", diff --git a/google/cloud/bigtable_admin/gapic_version.py b/google/cloud/bigtable_admin/gapic_version.py index 8d4f4cfb6..0f1a446f3 100644 --- a/google/cloud/bigtable_admin/gapic_version.py +++ b/google/cloud/bigtable_admin/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.17.0" # {x-release-please-version} +__version__ = "2.19.0" # {x-release-please-version} diff --git a/google/cloud/bigtable_admin_v2/__init__.py b/google/cloud/bigtable_admin_v2/__init__.py index 282834fe7..c030ec1bd 100644 --- a/google/cloud/bigtable_admin_v2/__init__.py +++ b/google/cloud/bigtable_admin_v2/__init__.py @@ -92,6 +92,7 @@ from .types.instance import Instance from .types.table import Backup from .types.table import BackupInfo +from .types.table import ChangeStreamConfig from .types.table import ColumnFamily from .types.table import EncryptionInfo from .types.table import GcRule @@ -110,6 +111,7 @@ "BackupInfo", "BigtableInstanceAdminClient", "BigtableTableAdminClient", + "ChangeStreamConfig", "CheckConsistencyRequest", "CheckConsistencyResponse", "Cluster", diff --git a/google/cloud/bigtable_admin_v2/gapic_version.py b/google/cloud/bigtable_admin_v2/gapic_version.py index 8d4f4cfb6..0f1a446f3 100644 --- a/google/cloud/bigtable_admin_v2/gapic_version.py +++ b/google/cloud/bigtable_admin_v2/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.17.0" # {x-release-please-version} +__version__ = "2.19.0" # {x-release-please-version} diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/async_client.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/async_client.py index ddeaf979a..12811bcea 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/async_client.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/async_client.py @@ -1137,8 +1137,8 @@ async def update_cluster( Args: request (Optional[Union[google.cloud.bigtable_admin_v2.types.Cluster, dict]]): - The request object. A resizable group of nodes in a - particular cloud location, capable of serving all + The request object. A resizable group of nodes in a particular cloud + location, capable of serving all [Tables][google.bigtable.admin.v2.Table] in the parent [Instance][google.bigtable.admin.v2.Instance]. retry (google.api_core.retry.Retry): Designation of what errors, if any, @@ -1880,8 +1880,7 @@ async def get_iam_policy( Args: request (Optional[Union[google.iam.v1.iam_policy_pb2.GetIamPolicyRequest, dict]]): - The request object. Request message for `GetIamPolicy` - method. + The request object. Request message for ``GetIamPolicy`` method. resource (:class:`str`): REQUIRED: The resource for which the policy is being requested. See the @@ -2030,8 +2029,7 @@ async def set_iam_policy( Args: request (Optional[Union[google.iam.v1.iam_policy_pb2.SetIamPolicyRequest, dict]]): - The request object. Request message for `SetIamPolicy` - method. + The request object. Request message for ``SetIamPolicy`` method. resource (:class:`str`): REQUIRED: The resource for which the policy is being specified. See the @@ -2171,8 +2169,7 @@ async def test_iam_permissions( Args: request (Optional[Union[google.iam.v1.iam_policy_pb2.TestIamPermissionsRequest, dict]]): - The request object. Request message for - `TestIamPermissions` method. + The request object. Request message for ``TestIamPermissions`` method. resource (:class:`str`): REQUIRED: The resource for which the policy detail is being requested. See diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/client.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/client.py index fcb767a3d..ecc9bf1e2 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/client.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/client.py @@ -1400,8 +1400,8 @@ def update_cluster( Args: request (Union[google.cloud.bigtable_admin_v2.types.Cluster, dict]): - The request object. A resizable group of nodes in a - particular cloud location, capable of serving all + The request object. A resizable group of nodes in a particular cloud + location, capable of serving all [Tables][google.bigtable.admin.v2.Table] in the parent [Instance][google.bigtable.admin.v2.Instance]. retry (google.api_core.retry.Retry): Designation of what errors, if any, @@ -2104,8 +2104,7 @@ def get_iam_policy( Args: request (Union[google.iam.v1.iam_policy_pb2.GetIamPolicyRequest, dict]): - The request object. Request message for `GetIamPolicy` - method. + The request object. Request message for ``GetIamPolicy`` method. resource (str): REQUIRED: The resource for which the policy is being requested. See the @@ -2241,8 +2240,7 @@ def set_iam_policy( Args: request (Union[google.iam.v1.iam_policy_pb2.SetIamPolicyRequest, dict]): - The request object. Request message for `SetIamPolicy` - method. + The request object. Request message for ``SetIamPolicy`` method. resource (str): REQUIRED: The resource for which the policy is being specified. See the @@ -2379,8 +2377,7 @@ def test_iam_permissions( Args: request (Union[google.iam.v1.iam_policy_pb2.TestIamPermissionsRequest, dict]): - The request object. Request message for - `TestIamPermissions` method. + The request object. Request message for ``TestIamPermissions`` method. resource (str): REQUIRED: The resource for which the policy detail is being requested. See diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/rest.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/rest.py index 5ae9600a9..e9b94cf78 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/rest.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/rest.py @@ -874,7 +874,6 @@ def __call__( request (~.bigtable_instance_admin.CreateAppProfileRequest): The request object. Request message for BigtableInstanceAdmin.CreateAppProfile. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -978,7 +977,6 @@ def __call__( request (~.bigtable_instance_admin.CreateClusterRequest): The request object. Request message for BigtableInstanceAdmin.CreateCluster. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1076,7 +1074,6 @@ def __call__( request (~.bigtable_instance_admin.CreateInstanceRequest): The request object. Request message for BigtableInstanceAdmin.CreateInstance. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1176,7 +1173,6 @@ def __call__( request (~.bigtable_instance_admin.DeleteAppProfileRequest): The request object. Request message for BigtableInstanceAdmin.DeleteAppProfile. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1254,7 +1250,6 @@ def __call__( request (~.bigtable_instance_admin.DeleteClusterRequest): The request object. Request message for BigtableInstanceAdmin.DeleteCluster. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1330,7 +1325,6 @@ def __call__( request (~.bigtable_instance_admin.DeleteInstanceRequest): The request object. Request message for BigtableInstanceAdmin.DeleteInstance. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1406,7 +1400,6 @@ def __call__( request (~.bigtable_instance_admin.GetAppProfileRequest): The request object. Request message for BigtableInstanceAdmin.GetAppProfile. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1497,7 +1490,6 @@ def __call__( request (~.bigtable_instance_admin.GetClusterRequest): The request object. Request message for BigtableInstanceAdmin.GetCluster. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1759,7 +1751,6 @@ def __call__( request (~.bigtable_instance_admin.GetInstanceRequest): The request object. Request message for BigtableInstanceAdmin.GetInstance. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1853,7 +1844,6 @@ def __call__( request (~.bigtable_instance_admin.ListAppProfilesRequest): The request object. Request message for BigtableInstanceAdmin.ListAppProfiles. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1945,7 +1935,6 @@ def __call__( request (~.bigtable_instance_admin.ListClustersRequest): The request object. Request message for BigtableInstanceAdmin.ListClusters. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2035,7 +2024,6 @@ def __call__( request (~.bigtable_instance_admin.ListHotTabletsRequest): The request object. Request message for BigtableInstanceAdmin.ListHotTablets. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2127,7 +2115,6 @@ def __call__( request (~.bigtable_instance_admin.ListInstancesRequest): The request object. Request message for BigtableInstanceAdmin.ListInstances. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2219,7 +2206,6 @@ def __call__( request (~.bigtable_instance_admin.PartialUpdateClusterRequest): The request object. Request message for BigtableInstanceAdmin.PartialUpdateCluster. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2321,7 +2307,6 @@ def __call__( request (~.bigtable_instance_admin.PartialUpdateInstanceRequest): The request object. Request message for BigtableInstanceAdmin.PartialUpdateInstance. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2692,7 +2677,6 @@ def __call__( request (~.bigtable_instance_admin.UpdateAppProfileRequest): The request object. Request message for BigtableInstanceAdmin.UpdateAppProfile. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2784,7 +2768,6 @@ def __call__( location, capable of serving all [Tables][google.bigtable.admin.v2.Table] in the parent [Instance][google.bigtable.admin.v2.Instance]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2885,7 +2868,6 @@ def __call__( served from all [Clusters][google.bigtable.admin.v2.Cluster] in the instance. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/async_client.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/async_client.py index bc85e5c5d..1663c16eb 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/async_client.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/async_client.py @@ -369,6 +369,7 @@ async def create_table_from_snapshot( request (Optional[Union[google.cloud.bigtable_admin_v2.types.CreateTableFromSnapshotRequest, dict]]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.CreateTableFromSnapshot][google.bigtable.admin.v2.BigtableTableAdmin.CreateTableFromSnapshot] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -682,16 +683,19 @@ async def update_table( should not be set. update_mask (:class:`google.protobuf.field_mask_pb2.FieldMask`): Required. The list of fields to update. A mask - specifying which fields (e.g. ``deletion_protection``) + specifying which fields (e.g. ``change_stream_config``) in the ``table`` field should be updated. This mask is relative to the ``table`` field, not to the request message. The wildcard (*) path is currently not supported. Currently UpdateTable is only supported for - the following field: + the following fields: + + - ``change_stream_config`` + - ``change_stream_config.retention_period`` + - ``deletion_protection`` - - ``deletion_protection`` If ``column_families`` is set - in ``update_mask``, it will return an UNIMPLEMENTED - error. + If ``column_families`` is set in ``update_mask``, it + will return an UNIMPLEMENTED error. This corresponds to the ``update_mask`` field on the ``request`` instance; if ``request`` is provided, this @@ -1300,6 +1304,7 @@ async def snapshot_table( request (Optional[Union[google.cloud.bigtable_admin_v2.types.SnapshotTableRequest, dict]]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.SnapshotTable][google.bigtable.admin.v2.BigtableTableAdmin.SnapshotTable] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1437,6 +1442,7 @@ async def get_snapshot( request (Optional[Union[google.cloud.bigtable_admin_v2.types.GetSnapshotRequest, dict]]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.GetSnapshot][google.bigtable.admin.v2.BigtableTableAdmin.GetSnapshot] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1549,6 +1555,7 @@ async def list_snapshots( request (Optional[Union[google.cloud.bigtable_admin_v2.types.ListSnapshotsRequest, dict]]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.ListSnapshots][google.bigtable.admin.v2.BigtableTableAdmin.ListSnapshots] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1672,6 +1679,7 @@ async def delete_snapshot( request (Optional[Union[google.cloud.bigtable_admin_v2.types.DeleteSnapshotRequest, dict]]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.DeleteSnapshot][google.bigtable.admin.v2.BigtableTableAdmin.DeleteSnapshot] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -2290,8 +2298,7 @@ async def get_iam_policy( Args: request (Optional[Union[google.iam.v1.iam_policy_pb2.GetIamPolicyRequest, dict]]): - The request object. Request message for `GetIamPolicy` - method. + The request object. Request message for ``GetIamPolicy`` method. resource (:class:`str`): REQUIRED: The resource for which the policy is being requested. See the @@ -2440,8 +2447,7 @@ async def set_iam_policy( Args: request (Optional[Union[google.iam.v1.iam_policy_pb2.SetIamPolicyRequest, dict]]): - The request object. Request message for `SetIamPolicy` - method. + The request object. Request message for ``SetIamPolicy`` method. resource (:class:`str`): REQUIRED: The resource for which the policy is being specified. See the @@ -2581,8 +2587,7 @@ async def test_iam_permissions( Args: request (Optional[Union[google.iam.v1.iam_policy_pb2.TestIamPermissionsRequest, dict]]): - The request object. Request message for - `TestIamPermissions` method. + The request object. Request message for ``TestIamPermissions`` method. resource (:class:`str`): REQUIRED: The resource for which the policy detail is being requested. See diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/client.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/client.py index aa7eaa197..e043aa224 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/client.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/client.py @@ -696,6 +696,7 @@ def create_table_from_snapshot( request (Union[google.cloud.bigtable_admin_v2.types.CreateTableFromSnapshotRequest, dict]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.CreateTableFromSnapshot][google.bigtable.admin.v2.BigtableTableAdmin.CreateTableFromSnapshot] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -991,16 +992,19 @@ def update_table( should not be set. update_mask (google.protobuf.field_mask_pb2.FieldMask): Required. The list of fields to update. A mask - specifying which fields (e.g. ``deletion_protection``) + specifying which fields (e.g. ``change_stream_config``) in the ``table`` field should be updated. This mask is relative to the ``table`` field, not to the request message. The wildcard (*) path is currently not supported. Currently UpdateTable is only supported for - the following field: + the following fields: + + - ``change_stream_config`` + - ``change_stream_config.retention_period`` + - ``deletion_protection`` - - ``deletion_protection`` If ``column_families`` is set - in ``update_mask``, it will return an UNIMPLEMENTED - error. + If ``column_families`` is set in ``update_mask``, it + will return an UNIMPLEMENTED error. This corresponds to the ``update_mask`` field on the ``request`` instance; if ``request`` is provided, this @@ -1594,6 +1598,7 @@ def snapshot_table( request (Union[google.cloud.bigtable_admin_v2.types.SnapshotTableRequest, dict]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.SnapshotTable][google.bigtable.admin.v2.BigtableTableAdmin.SnapshotTable] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1731,6 +1736,7 @@ def get_snapshot( request (Union[google.cloud.bigtable_admin_v2.types.GetSnapshotRequest, dict]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.GetSnapshot][google.bigtable.admin.v2.BigtableTableAdmin.GetSnapshot] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1833,6 +1839,7 @@ def list_snapshots( request (Union[google.cloud.bigtable_admin_v2.types.ListSnapshotsRequest, dict]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.ListSnapshots][google.bigtable.admin.v2.BigtableTableAdmin.ListSnapshots] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1946,6 +1953,7 @@ def delete_snapshot( request (Union[google.cloud.bigtable_admin_v2.types.DeleteSnapshotRequest, dict]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.DeleteSnapshot][google.bigtable.admin.v2.BigtableTableAdmin.DeleteSnapshot] + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -2545,8 +2553,7 @@ def get_iam_policy( Args: request (Union[google.iam.v1.iam_policy_pb2.GetIamPolicyRequest, dict]): - The request object. Request message for `GetIamPolicy` - method. + The request object. Request message for ``GetIamPolicy`` method. resource (str): REQUIRED: The resource for which the policy is being requested. See the @@ -2682,8 +2689,7 @@ def set_iam_policy( Args: request (Union[google.iam.v1.iam_policy_pb2.SetIamPolicyRequest, dict]): - The request object. Request message for `SetIamPolicy` - method. + The request object. Request message for ``SetIamPolicy`` method. resource (str): REQUIRED: The resource for which the policy is being specified. See the @@ -2820,8 +2826,7 @@ def test_iam_permissions( Args: request (Union[google.iam.v1.iam_policy_pb2.TestIamPermissionsRequest, dict]): - The request object. Request message for - `TestIamPermissions` method. + The request object. Request message for ``TestIamPermissions`` method. resource (str): REQUIRED: The resource for which the policy detail is being requested. See diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/rest.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/rest.py index 5c25ac556..4d5b2ed1c 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/rest.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/rest.py @@ -938,7 +938,6 @@ def __call__( request (~.bigtable_table_admin.CheckConsistencyRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.CheckConsistency][google.bigtable.admin.v2.BigtableTableAdmin.CheckConsistency] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1041,7 +1040,6 @@ def __call__( request (~.bigtable_table_admin.CreateBackupRequest): The request object. The request for [CreateBackup][google.bigtable.admin.v2.BigtableTableAdmin.CreateBackup]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1139,7 +1137,6 @@ def __call__( request (~.bigtable_table_admin.CreateTableRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.CreateTable][google.bigtable.admin.v2.BigtableTableAdmin.CreateTable] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1248,7 +1245,6 @@ def __call__( changed in backward-incompatible ways and is not recommended for production use. It is not subject to any SLA or deprecation policy. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1348,7 +1344,6 @@ def __call__( request (~.bigtable_table_admin.DeleteBackupRequest): The request object. The request for [DeleteBackup][google.bigtable.admin.v2.BigtableTableAdmin.DeleteBackup]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1431,7 +1426,6 @@ def __call__( changed in backward-incompatible ways and is not recommended for production use. It is not subject to any SLA or deprecation policy. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1507,7 +1501,6 @@ def __call__( request (~.bigtable_table_admin.DeleteTableRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.DeleteTable][google.bigtable.admin.v2.BigtableTableAdmin.DeleteTable] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1583,7 +1576,6 @@ def __call__( request (~.bigtable_table_admin.DropRowRangeRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.DropRowRange][google.bigtable.admin.v2.BigtableTableAdmin.DropRowRange] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1669,7 +1661,6 @@ def __call__( request (~.bigtable_table_admin.GenerateConsistencyTokenRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.GenerateConsistencyToken][google.bigtable.admin.v2.BigtableTableAdmin.GenerateConsistencyToken] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1772,7 +1763,6 @@ def __call__( request (~.bigtable_table_admin.GetBackupRequest): The request object. The request for [GetBackup][google.bigtable.admin.v2.BigtableTableAdmin.GetBackup]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2042,7 +2032,6 @@ def __call__( changed in backward-incompatible ways and is not recommended for production use. It is not subject to any SLA or deprecation policy. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2142,7 +2131,6 @@ def __call__( request (~.bigtable_table_admin.GetTableRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.GetTable][google.bigtable.admin.v2.BigtableTableAdmin.GetTable] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2234,7 +2222,6 @@ def __call__( request (~.bigtable_table_admin.ListBackupsRequest): The request object. The request for [ListBackups][google.bigtable.admin.v2.BigtableTableAdmin.ListBackups]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2331,7 +2318,6 @@ def __call__( changed in backward-incompatible ways and is not recommended for production use. It is not subject to any SLA or deprecation policy. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2428,7 +2414,6 @@ def __call__( request (~.bigtable_table_admin.ListTablesRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.ListTables][google.bigtable.admin.v2.BigtableTableAdmin.ListTables] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2518,7 +2503,6 @@ def __call__( request (~.bigtable_table_admin.ModifyColumnFamiliesRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.ModifyColumnFamilies][google.bigtable.admin.v2.BigtableTableAdmin.ModifyColumnFamilies] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2621,7 +2605,6 @@ def __call__( request (~.bigtable_table_admin.RestoreTableRequest): The request object. The request for [RestoreTable][google.bigtable.admin.v2.BigtableTableAdmin.RestoreTable]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -2901,7 +2884,6 @@ def __call__( changed in backward-incompatible ways and is not recommended for production use. It is not subject to any SLA or deprecation policy. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -3101,7 +3083,6 @@ def __call__( request (~.bigtable_table_admin.UndeleteTableRequest): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.UndeleteTable][google.bigtable.admin.v2.BigtableTableAdmin.UndeleteTable] - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -3201,7 +3182,6 @@ def __call__( request (~.bigtable_table_admin.UpdateBackupRequest): The request object. The request for [UpdateBackup][google.bigtable.admin.v2.BigtableTableAdmin.UpdateBackup]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -3300,7 +3280,6 @@ def __call__( request (~.bigtable_table_admin.UpdateTableRequest): The request object. The request for [UpdateTable][google.bigtable.admin.v2.BigtableTableAdmin.UpdateTable]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. diff --git a/google/cloud/bigtable_admin_v2/types/__init__.py b/google/cloud/bigtable_admin_v2/types/__init__.py index 5a66ddf09..69153c9fc 100644 --- a/google/cloud/bigtable_admin_v2/types/__init__.py +++ b/google/cloud/bigtable_admin_v2/types/__init__.py @@ -91,6 +91,7 @@ from .table import ( Backup, BackupInfo, + ChangeStreamConfig, ColumnFamily, EncryptionInfo, GcRule, @@ -170,6 +171,7 @@ "Instance", "Backup", "BackupInfo", + "ChangeStreamConfig", "ColumnFamily", "EncryptionInfo", "GcRule", diff --git a/google/cloud/bigtable_admin_v2/types/bigtable_table_admin.py b/google/cloud/bigtable_admin_v2/types/bigtable_table_admin.py index 9b236fea9..4c4b9e9e2 100644 --- a/google/cloud/bigtable_admin_v2/types/bigtable_table_admin.py +++ b/google/cloud/bigtable_admin_v2/types/bigtable_table_admin.py @@ -460,14 +460,18 @@ class UpdateTableRequest(proto.Message): used to identify the table to update. update_mask (google.protobuf.field_mask_pb2.FieldMask): Required. The list of fields to update. A mask specifying - which fields (e.g. ``deletion_protection``) in the ``table`` - field should be updated. This mask is relative to the - ``table`` field, not to the request message. The wildcard - (*) path is currently not supported. Currently UpdateTable - is only supported for the following field: - - - ``deletion_protection`` If ``column_families`` is set in - ``update_mask``, it will return an UNIMPLEMENTED error. + which fields (e.g. ``change_stream_config``) in the + ``table`` field should be updated. This mask is relative to + the ``table`` field, not to the request message. The + wildcard (*) path is currently not supported. Currently + UpdateTable is only supported for the following fields: + + - ``change_stream_config`` + - ``change_stream_config.retention_period`` + - ``deletion_protection`` + + If ``column_families`` is set in ``update_mask``, it will + return an UNIMPLEMENTED error. """ table: gba_table.Table = proto.Field( diff --git a/google/cloud/bigtable_admin_v2/types/table.py b/google/cloud/bigtable_admin_v2/types/table.py index fd936df63..16d136e16 100644 --- a/google/cloud/bigtable_admin_v2/types/table.py +++ b/google/cloud/bigtable_admin_v2/types/table.py @@ -29,6 +29,7 @@ manifest={ "RestoreSourceType", "RestoreInfo", + "ChangeStreamConfig", "Table", "ColumnFamily", "GcRule", @@ -82,6 +83,27 @@ class RestoreInfo(proto.Message): ) +class ChangeStreamConfig(proto.Message): + r"""Change stream configuration. + + Attributes: + retention_period (google.protobuf.duration_pb2.Duration): + How long the change stream should be + retained. Change stream data older than the + retention period will not be returned when + reading the change stream from the table. + Values must be at least 1 day and at most 7 + days, and will be truncated to microsecond + granularity. + """ + + retention_period: duration_pb2.Duration = proto.Field( + proto.MESSAGE, + number=1, + message=duration_pb2.Duration, + ) + + class Table(proto.Message): r"""A collection of user data indexed by row, column, and timestamp. Each table is served using the resources of its @@ -114,6 +136,10 @@ class Table(proto.Message): another data source (e.g. a backup), this field will be populated with information about the restore. + change_stream_config (google.cloud.bigtable_admin_v2.types.ChangeStreamConfig): + If specified, enable the change stream on + this table. Otherwise, the change stream is + disabled and the change stream is not retained. deletion_protection (bool): Set to true to make the table protected against data loss. i.e. deleting the following @@ -263,6 +289,11 @@ class ReplicationState(proto.Enum): number=6, message="RestoreInfo", ) + change_stream_config: "ChangeStreamConfig" = proto.Field( + proto.MESSAGE, + number=8, + message="ChangeStreamConfig", + ) deletion_protection: bool = proto.Field( proto.BOOL, number=9, diff --git a/google/cloud/bigtable_v2/__init__.py b/google/cloud/bigtable_v2/__init__.py index 342718dea..ee3bd8c0c 100644 --- a/google/cloud/bigtable_v2/__init__.py +++ b/google/cloud/bigtable_v2/__init__.py @@ -31,6 +31,7 @@ from .types.bigtable import MutateRowsResponse from .types.bigtable import PingAndWarmRequest from .types.bigtable import PingAndWarmResponse +from .types.bigtable import RateLimitInfo from .types.bigtable import ReadChangeStreamRequest from .types.bigtable import ReadChangeStreamResponse from .types.bigtable import ReadModifyWriteRowRequest @@ -54,6 +55,7 @@ from .types.data import StreamPartition from .types.data import TimestampRange from .types.data import ValueRange +from .types.feature_flags import FeatureFlags from .types.request_stats import FullReadStatsView from .types.request_stats import ReadIterationStats from .types.request_stats import RequestLatencyStats @@ -69,6 +71,7 @@ "Column", "ColumnRange", "Family", + "FeatureFlags", "FullReadStatsView", "GenerateInitialChangeStreamPartitionsRequest", "GenerateInitialChangeStreamPartitionsResponse", @@ -79,6 +82,7 @@ "Mutation", "PingAndWarmRequest", "PingAndWarmResponse", + "RateLimitInfo", "ReadChangeStreamRequest", "ReadChangeStreamResponse", "ReadIterationStats", diff --git a/google/cloud/bigtable_v2/gapic_version.py b/google/cloud/bigtable_v2/gapic_version.py index 8d4f4cfb6..0f1a446f3 100644 --- a/google/cloud/bigtable_v2/gapic_version.py +++ b/google/cloud/bigtable_v2/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.17.0" # {x-release-please-version} +__version__ = "2.19.0" # {x-release-please-version} diff --git a/google/cloud/bigtable_v2/services/bigtable/async_client.py b/google/cloud/bigtable_v2/services/bigtable/async_client.py index 3465569b3..abd82d4d8 100644 --- a/google/cloud/bigtable_v2/services/bigtable/async_client.py +++ b/google/cloud/bigtable_v2/services/bigtable/async_client.py @@ -242,8 +242,10 @@ def read_rows( on the ``request`` instance; if ``request`` is provided, this should not be set. app_profile_id (:class:`str`): - This value specifies routing for replication. This API - only accepts the empty value of app_profile_id. + This value specifies routing for + replication. If not specified, the + "default" application profile will be + used. This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this @@ -807,8 +809,8 @@ async def ping_and_warm( Args: request (Optional[Union[google.cloud.bigtable_v2.types.PingAndWarmRequest, dict]]): - The request object. Request message for client - connection keep-alive and warming. + The request object. Request message for client connection + keep-alive and warming. name (:class:`str`): Required. The unique name of the instance to check permissions for as well as respond. Values are of the @@ -1027,8 +1029,9 @@ def generate_initial_change_stream_partitions( Args: request (Optional[Union[google.cloud.bigtable_v2.types.GenerateInitialChangeStreamPartitionsRequest, dict]]): - The request object. NOTE: This API is intended to be - used by Apache Beam BigtableIO. Request message for + The request object. NOTE: This API is intended to be used + by Apache Beam BigtableIO. Request + message for Bigtable.GenerateInitialChangeStreamPartitions. table_name (:class:`str`): Required. The unique name of the table from which to get @@ -1126,9 +1129,9 @@ def read_change_stream( Args: request (Optional[Union[google.cloud.bigtable_v2.types.ReadChangeStreamRequest, dict]]): - The request object. NOTE: This API is intended to be - used by Apache Beam BigtableIO. Request message for - Bigtable.ReadChangeStream. + The request object. NOTE: This API is intended to be used + by Apache Beam BigtableIO. Request + message for Bigtable.ReadChangeStream. table_name (:class:`str`): Required. The unique name of the table from which to read a change stream. Values are of the form diff --git a/google/cloud/bigtable_v2/services/bigtable/client.py b/google/cloud/bigtable_v2/services/bigtable/client.py index 60622509a..b0efc8a0b 100644 --- a/google/cloud/bigtable_v2/services/bigtable/client.py +++ b/google/cloud/bigtable_v2/services/bigtable/client.py @@ -493,8 +493,10 @@ def read_rows( on the ``request`` instance; if ``request`` is provided, this should not be set. app_profile_id (str): - This value specifies routing for replication. This API - only accepts the empty value of app_profile_id. + This value specifies routing for + replication. If not specified, the + "default" application profile will be + used. This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this @@ -1093,8 +1095,8 @@ def ping_and_warm( Args: request (Union[google.cloud.bigtable_v2.types.PingAndWarmRequest, dict]): - The request object. Request message for client - connection keep-alive and warming. + The request object. Request message for client connection + keep-alive and warming. name (str): Required. The unique name of the instance to check permissions for as well as respond. Values are of the @@ -1329,8 +1331,9 @@ def generate_initial_change_stream_partitions( Args: request (Union[google.cloud.bigtable_v2.types.GenerateInitialChangeStreamPartitionsRequest, dict]): - The request object. NOTE: This API is intended to be - used by Apache Beam BigtableIO. Request message for + The request object. NOTE: This API is intended to be used + by Apache Beam BigtableIO. Request + message for Bigtable.GenerateInitialChangeStreamPartitions. table_name (str): Required. The unique name of the table from which to get @@ -1432,9 +1435,9 @@ def read_change_stream( Args: request (Union[google.cloud.bigtable_v2.types.ReadChangeStreamRequest, dict]): - The request object. NOTE: This API is intended to be - used by Apache Beam BigtableIO. Request message for - Bigtable.ReadChangeStream. + The request object. NOTE: This API is intended to be used + by Apache Beam BigtableIO. Request + message for Bigtable.ReadChangeStream. table_name (str): Required. The unique name of the table from which to read a change stream. Values are of the form diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/rest.py b/google/cloud/bigtable_v2/services/bigtable/transports/rest.py index ee9cb046f..4343fbb90 100644 --- a/google/cloud/bigtable_v2/services/bigtable/transports/rest.py +++ b/google/cloud/bigtable_v2/services/bigtable/transports/rest.py @@ -471,7 +471,6 @@ def __call__( request (~.bigtable.CheckAndMutateRowRequest): The request object. Request message for Bigtable.CheckAndMutateRow. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -575,7 +574,6 @@ def __call__( by Apache Beam BigtableIO. Request message for Bigtable.GenerateInitialChangeStreamPartitions. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -684,7 +682,6 @@ def __call__( request (~.bigtable.MutateRowRequest): The request object. Request message for Bigtable.MutateRow. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -783,7 +780,6 @@ def __call__( request (~.bigtable.MutateRowsRequest): The request object. Request message for BigtableService.MutateRows. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -881,7 +877,6 @@ def __call__( request (~.bigtable.PingAndWarmRequest): The request object. Request message for client connection keep-alive and warming. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -982,7 +977,6 @@ def __call__( The request object. NOTE: This API is intended to be used by Apache Beam BigtableIO. Request message for Bigtable.ReadChangeStream. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1083,7 +1077,6 @@ def __call__( request (~.bigtable.ReadModifyWriteRowRequest): The request object. Request message for Bigtable.ReadModifyWriteRow. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1184,7 +1177,6 @@ def __call__( request (~.bigtable.ReadRowsRequest): The request object. Request message for Bigtable.ReadRows. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. @@ -1280,7 +1272,6 @@ def __call__( request (~.bigtable.SampleRowKeysRequest): The request object. Request message for Bigtable.SampleRowKeys. - retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. diff --git a/google/cloud/bigtable_v2/types/__init__.py b/google/cloud/bigtable_v2/types/__init__.py index bb2533e33..9f15efaf5 100644 --- a/google/cloud/bigtable_v2/types/__init__.py +++ b/google/cloud/bigtable_v2/types/__init__.py @@ -24,6 +24,7 @@ MutateRowsResponse, PingAndWarmRequest, PingAndWarmResponse, + RateLimitInfo, ReadChangeStreamRequest, ReadChangeStreamResponse, ReadModifyWriteRowRequest, @@ -50,6 +51,9 @@ TimestampRange, ValueRange, ) +from .feature_flags import ( + FeatureFlags, +) from .request_stats import ( FullReadStatsView, ReadIterationStats, @@ -71,6 +75,7 @@ "MutateRowsResponse", "PingAndWarmRequest", "PingAndWarmResponse", + "RateLimitInfo", "ReadChangeStreamRequest", "ReadChangeStreamResponse", "ReadModifyWriteRowRequest", @@ -94,6 +99,7 @@ "StreamPartition", "TimestampRange", "ValueRange", + "FeatureFlags", "FullReadStatsView", "ReadIterationStats", "RequestLatencyStats", diff --git a/google/cloud/bigtable_v2/types/bigtable.py b/google/cloud/bigtable_v2/types/bigtable.py index ea97588c2..13f6ac0db 100644 --- a/google/cloud/bigtable_v2/types/bigtable.py +++ b/google/cloud/bigtable_v2/types/bigtable.py @@ -38,6 +38,7 @@ "MutateRowResponse", "MutateRowsRequest", "MutateRowsResponse", + "RateLimitInfo", "CheckAndMutateRowRequest", "CheckAndMutateRowResponse", "PingAndWarmRequest", @@ -61,8 +62,9 @@ class ReadRowsRequest(proto.Message): Values are of the form ``projects//instances//tables/``. app_profile_id (str): - This value specifies routing for replication. This API only - accepts the empty value of app_profile_id. + This value specifies routing for replication. + If not specified, the "default" application + profile will be used. rows (google.cloud.bigtable_v2.types.RowSet): The row keys and/or ranges to read sequentially. If not specified, reads from all @@ -469,10 +471,19 @@ class Entry(proto.Message): class MutateRowsResponse(proto.Message): r"""Response message for BigtableService.MutateRows. + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + Attributes: entries (MutableSequence[google.cloud.bigtable_v2.types.MutateRowsResponse.Entry]): One or more results for Entries from the batch request. + rate_limit_info (google.cloud.bigtable_v2.types.RateLimitInfo): + Information about how client should limit the + rate (QPS). Primirily used by supported official + Cloud Bigtable clients. If unset, the rate limit + info is not provided by the server. + + This field is a member of `oneof`_ ``_rate_limit_info``. """ class Entry(proto.Message): @@ -506,6 +517,50 @@ class Entry(proto.Message): number=1, message=Entry, ) + rate_limit_info: "RateLimitInfo" = proto.Field( + proto.MESSAGE, + number=3, + optional=True, + message="RateLimitInfo", + ) + + +class RateLimitInfo(proto.Message): + r"""Information about how client should adjust the load to + Bigtable. + + Attributes: + period (google.protobuf.duration_pb2.Duration): + Time that clients should wait before + adjusting the target rate again. If clients + adjust rate too frequently, the impact of the + previous adjustment may not have been taken into + account and may over-throttle or under-throttle. + If clients adjust rate too slowly, they will not + be responsive to load changes on server side, + and may over-throttle or under-throttle. + factor (float): + If it has been at least one ``period`` since the last load + adjustment, the client should multiply the current load by + this value to get the new target load. For example, if the + current load is 100 and ``factor`` is 0.8, the new target + load should be 80. After adjusting, the client should ignore + ``factor`` until another ``period`` has passed. + + The client can measure its load using any unit that's + comparable over time For example, QPS can be used as long as + each request involves a similar amount of work. + """ + + period: duration_pb2.Duration = proto.Field( + proto.MESSAGE, + number=1, + message=duration_pb2.Duration, + ) + factor: float = proto.Field( + proto.DOUBLE, + number=2, + ) class CheckAndMutateRowRequest(proto.Message): diff --git a/google/cloud/bigtable_v2/types/feature_flags.py b/google/cloud/bigtable_v2/types/feature_flags.py new file mode 100644 index 000000000..1b5f76e24 --- /dev/null +++ b/google/cloud/bigtable_v2/types/feature_flags.py @@ -0,0 +1,54 @@ +# -*- coding: utf-8 -*- +# Copyright 2022 Google LLC +# +# Licensed 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. +# +from __future__ import annotations + +from typing import MutableMapping, MutableSequence + +import proto # type: ignore + + +__protobuf__ = proto.module( + package="google.bigtable.v2", + manifest={ + "FeatureFlags", + }, +) + + +class FeatureFlags(proto.Message): + r"""Feature flags supported by a client. This is intended to be sent as + part of request metadata to assure the server that certain behaviors + are safe to enable. This proto is meant to be serialized and + websafe-base64 encoded under the ``bigtable-features`` metadata key. + The value will remain constant for the lifetime of a client and due + to HTTP2's HPACK compression, the request overhead will be tiny. + This is an internal implementation detail and should not be used by + endusers directly. + + Attributes: + mutate_rows_rate_limit (bool): + Notify the server that the client enables + batch write flow control by requesting + RateLimitInfo from MutateRowsResponse. + """ + + mutate_rows_rate_limit: bool = proto.Field( + proto.BOOL, + number=3, + ) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/noxfile.py b/noxfile.py index 164d138bd..8499a610f 100644 --- a/noxfile.py +++ b/noxfile.py @@ -40,7 +40,7 @@ "pytest-asyncio", ] UNIT_TEST_EXTERNAL_DEPENDENCIES = [ - "git+https://github.com/googleapis/python-api-core.git@retry_generators" + # "git+https://github.com/googleapis/python-api-core.git@retry_generators" ] UNIT_TEST_LOCAL_DEPENDENCIES = [] UNIT_TEST_DEPENDENCIES = [] @@ -55,7 +55,7 @@ "google-cloud-testutils", ] SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [ - "git+https://github.com/googleapis/python-api-core.git@retry_generators" + # "git+https://github.com/googleapis/python-api-core.git@retry_generators" ] SYSTEM_TEST_LOCAL_DEPENDENCIES = [] UNIT_TEST_DEPENDENCIES = [] @@ -138,13 +138,11 @@ def mypy(session): session.install("google-cloud-testutils") session.run( "mypy", - "google/cloud/bigtable", + "google/cloud/bigtable/data", "--check-untyped-defs", "--warn-unreachable", "--disallow-any-generics", "--exclude", - "google/cloud/bigtable/deprecated", - "--exclude", "tests/system/v2_client", "--exclude", "tests/unit/v2_client", @@ -318,7 +316,6 @@ def system(session): "py.test", "--quiet", f"--junitxml=system_{session.python}_sponge_log.xml", - "--ignore=tests/system/v2_client", system_test_folder_path, *session.posargs, ) @@ -466,11 +463,6 @@ def prerelease_deps(session): ) session.run("python", "-c", "import grpc; print(grpc.__version__)") - # TODO: remove adter merging api-core - session.install( - "--upgrade", "--no-deps", "--force-reinstall", *UNIT_TEST_EXTERNAL_DEPENDENCIES - ) - session.run("py.test", "tests/unit") system_test_path = os.path.join("tests", "system.py") diff --git a/python-api-core b/python-api-core index 9ba76760f..a526d6593 160000 --- a/python-api-core +++ b/python-api-core @@ -1 +1 @@ -Subproject commit 9ba76760f5b7ba8128be85ca780811a0b9ec9087 +Subproject commit a526d659320939cd7f47ee775b250e8a3e3ab16b diff --git a/samples/beam/requirements-test.txt b/samples/beam/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/beam/requirements-test.txt +++ b/samples/beam/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/beam/requirements.txt b/samples/beam/requirements.txt index bcb270e72..8be9b98e0 100644 --- a/samples/beam/requirements.txt +++ b/samples/beam/requirements.txt @@ -1,3 +1,3 @@ -apache-beam==2.45.0 +apache-beam==2.46.0 google-cloud-bigtable==2.17.0 google-cloud-core==2.3.2 diff --git a/samples/hello/README.md b/samples/hello/README.md index 0e1fc92f9..b3779fb43 100644 --- a/samples/hello/README.md +++ b/samples/hello/README.md @@ -17,7 +17,7 @@ Demonstrates how to connect to Cloud Bigtable and run some basic operations. Mor To run this sample: -1. If this is your first time working with GCP products, you will need to set up [the Cloud SDK][cloud_sdk] or utilize [Google Cloud Shell][gcloud_shell]. This sample may [require authetication][authentication] and you will need to [enable billing][enable_billing]. +1. If this is your first time working with GCP products, you will need to set up [the Cloud SDK][cloud_sdk] or utilize [Google Cloud Shell][gcloud_shell]. This sample may [require authentication][authentication] and you will need to [enable billing][enable_billing]. 1. Make a fork of this repo and clone the branch locally, then navigate to the sample directory you want to use. diff --git a/samples/hello/main.py b/samples/hello/main.py index 7b2b1764a..5e47b4a38 100644 --- a/samples/hello/main.py +++ b/samples/hello/main.py @@ -87,26 +87,30 @@ def main(project_id, instance_id, table_id): # [START bigtable_hw_create_filter] # Create a filter to only retrieve the most recent version of the cell - # for each column accross entire row. + # for each column across entire row. row_filter = row_filters.CellsColumnLimitFilter(1) # [END bigtable_hw_create_filter] # [START bigtable_hw_get_with_filter] + # [START bigtable_hw_get_by_key] print("Getting a single greeting by row key.") key = "greeting0".encode() row = table.read_row(key, row_filter) cell = row.cells[column_family_id][column][0] print(cell.value.decode("utf-8")) + # [END bigtable_hw_get_by_key] # [END bigtable_hw_get_with_filter] # [START bigtable_hw_scan_with_filter] + # [START bigtable_hw_scan_all] print("Scanning for all greetings:") partial_rows = table.read_rows(filter_=row_filter) for row in partial_rows: cell = row.cells[column_family_id][column][0] print(cell.value.decode("utf-8")) + # [END bigtable_hw_scan_all] # [END bigtable_hw_scan_with_filter] # [START bigtable_hw_delete_table] diff --git a/samples/hello/requirements-test.txt b/samples/hello/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/hello/requirements-test.txt +++ b/samples/hello/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/hello_happybase/requirements-test.txt b/samples/hello_happybase/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/hello_happybase/requirements-test.txt +++ b/samples/hello_happybase/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/instanceadmin/requirements-test.txt b/samples/instanceadmin/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/instanceadmin/requirements-test.txt +++ b/samples/instanceadmin/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/metricscaler/requirements-test.txt b/samples/metricscaler/requirements-test.txt index 82f315c7f..761227068 100644 --- a/samples/metricscaler/requirements-test.txt +++ b/samples/metricscaler/requirements-test.txt @@ -1,3 +1,3 @@ -pytest==7.2.2 -mock==5.0.1 +pytest==7.3.1 +mock==5.0.2 google-cloud-testutils diff --git a/samples/metricscaler/requirements.txt b/samples/metricscaler/requirements.txt index e9647809f..02e08b4c8 100644 --- a/samples/metricscaler/requirements.txt +++ b/samples/metricscaler/requirements.txt @@ -1,2 +1,2 @@ google-cloud-bigtable==2.17.0 -google-cloud-monitoring==2.14.1 +google-cloud-monitoring==2.14.2 diff --git a/samples/quickstart/requirements-test.txt b/samples/quickstart/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/quickstart/requirements-test.txt +++ b/samples/quickstart/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/quickstart_happybase/requirements-test.txt b/samples/quickstart_happybase/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/quickstart_happybase/requirements-test.txt +++ b/samples/quickstart_happybase/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/snippets/deletes/deletes_snippets.py b/samples/snippets/deletes/deletes_snippets.py index 4e89189db..8e78083bf 100644 --- a/samples/snippets/deletes/deletes_snippets.py +++ b/samples/snippets/deletes/deletes_snippets.py @@ -38,7 +38,7 @@ def delete_from_column_family(project_id, instance_id, table_id): table = instance.table(table_id) row = table.row("phone#4c410523#20190501") row.delete_cells( - column_family_id="cell_plan", columns=["data_plan_01gb", "data_plan_05gb"] + column_family_id="cell_plan", columns=row.ALL_COLUMNS ) row.commit() diff --git a/samples/snippets/deletes/requirements-test.txt b/samples/snippets/deletes/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/snippets/deletes/requirements-test.txt +++ b/samples/snippets/deletes/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/snippets/filters/requirements-test.txt b/samples/snippets/filters/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/snippets/filters/requirements-test.txt +++ b/samples/snippets/filters/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/snippets/reads/requirements-test.txt b/samples/snippets/reads/requirements-test.txt index c021c5b5b..c4d04a08d 100644 --- a/samples/snippets/reads/requirements-test.txt +++ b/samples/snippets/reads/requirements-test.txt @@ -1 +1 @@ -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/snippets/writes/requirements-test.txt b/samples/snippets/writes/requirements-test.txt index 8d6117f16..96aa71dab 100644 --- a/samples/snippets/writes/requirements-test.txt +++ b/samples/snippets/writes/requirements-test.txt @@ -1,2 +1,2 @@ backoff==2.2.1 -pytest==7.2.2 +pytest==7.3.1 diff --git a/samples/tableadmin/requirements-test.txt b/samples/tableadmin/requirements-test.txt index d3ddc990f..ca1f33bd3 100644 --- a/samples/tableadmin/requirements-test.txt +++ b/samples/tableadmin/requirements-test.txt @@ -1,2 +1,2 @@ -pytest==7.2.2 +pytest==7.3.1 google-cloud-testutils==1.3.3 diff --git a/setup.py b/setup.py index 49bb10adc..e05b37c79 100644 --- a/setup.py +++ b/setup.py @@ -37,7 +37,7 @@ # 'Development Status :: 5 - Production/Stable' release_status = "Development Status :: 5 - Production/Stable" dependencies = [ - "google-api-core[grpc] >= 1.34.0, <3.0.0dev,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,!=2.8.*,!=2.9.*,!=2.10.*", + "google-api-core[grpc] == 2.12.0.dev0", # TODO: change to >= after streaming retries is merged "google-cloud-core >= 1.4.1, <3.0.0dev", "grpc-google-iam-v1 >= 0.12.4, <1.0.0dev", "proto-plus >= 1.22.0, <2.0.0dev", diff --git a/testing/constraints-3.7.txt b/testing/constraints-3.7.txt index 7bf769c9b..92b616563 100644 --- a/testing/constraints-3.7.txt +++ b/testing/constraints-3.7.txt @@ -5,9 +5,8 @@ # # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 -# TODO: reset after merging api-core submodule -# google-api-core==2.11.0 -google-cloud-core==1.4.1 +google-api-core==2.12.0.dev0 +google-cloud-core==2.3.2 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 libcst==0.2.5 diff --git a/google/cloud/bigtable/deprecated/__init__.py b/tests/system/data/__init__.py similarity index 64% rename from google/cloud/bigtable/deprecated/__init__.py rename to tests/system/data/__init__.py index a54fffdf1..89a37dc92 100644 --- a/google/cloud/bigtable/deprecated/__init__.py +++ b/tests/system/data/__init__.py @@ -1,4 +1,5 @@ -# Copyright 2015 Google LLC +# -*- coding: utf-8 -*- +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,15 +12,4 @@ # 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. - -"""Google Cloud Bigtable API package.""" - -from google.cloud.bigtable.deprecated.client import Client - -from google.cloud.bigtable import gapic_version as package_version - -__version__: str - -__version__ = package_version.__version__ - -__all__ = ["__version__", "Client"] +# diff --git a/tests/system/test_system.py b/tests/system/data/test_system.py similarity index 93% rename from tests/system/test_system.py rename to tests/system/data/test_system.py index e1771202a..548433444 100644 --- a/tests/system/test_system.py +++ b/tests/system/data/test_system.py @@ -20,7 +20,7 @@ from google.api_core import retry from google.api_core.exceptions import ClientError -from google.cloud.bigtable.read_modify_write_rules import MAX_INCREMENT_VALUE +from google.cloud.bigtable.data.read_modify_write_rules import MAX_INCREMENT_VALUE TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" @@ -135,10 +135,10 @@ def table_id(table_admin_client, project_id, instance_id): @pytest_asyncio.fixture(scope="session") async def client(): - from google.cloud.bigtable import BigtableDataClient + from google.cloud.bigtable.data import BigtableDataClientAsync project = os.getenv("GOOGLE_CLOUD_PROJECT") or None - async with BigtableDataClient(project=project) as client: + async with BigtableDataClientAsync(project=project) as client: yield client @@ -201,7 +201,7 @@ async def _retrieve_cell_value(table, row_key): """ Helper to read an individual row """ - from google.cloud.bigtable import ReadRowsQuery + from google.cloud.bigtable.data import ReadRowsQuery row_list = await table.read_rows(ReadRowsQuery(row_keys=row_key)) assert len(row_list) == 1 @@ -216,7 +216,7 @@ async def _create_row_and_mutation( """ Helper to create a new row, and a sample set_cell mutation to change its value """ - from google.cloud.bigtable.mutations import SetCell + from google.cloud.bigtable.data.mutations import SetCell row_key = uuid.uuid4().hex.encode() family = TEST_FAMILY @@ -303,7 +303,7 @@ async def test_bulk_mutations_set_cell(client, table, temp_rows): """ Ensure cells can be set properly """ - from google.cloud.bigtable.mutations import RowMutationEntry + from google.cloud.bigtable.data.mutations import RowMutationEntry new_value = uuid.uuid4().hex.encode() row_key, mutation = await _create_row_and_mutation( @@ -323,7 +323,7 @@ async def test_mutations_batcher_context_manager(client, table, temp_rows): """ test batcher with context manager. Should flush on exit """ - from google.cloud.bigtable.mutations import RowMutationEntry + from google.cloud.bigtable.data.mutations import RowMutationEntry new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] row_key, mutation = await _create_row_and_mutation( @@ -349,7 +349,7 @@ async def test_mutations_batcher_timer_flush(client, table, temp_rows): """ batch should occur after flush_interval seconds """ - from google.cloud.bigtable.mutations import RowMutationEntry + from google.cloud.bigtable.data.mutations import RowMutationEntry new_value = uuid.uuid4().hex.encode() row_key, mutation = await _create_row_and_mutation( @@ -373,7 +373,7 @@ async def test_mutations_batcher_count_flush(client, table, temp_rows): """ batch should flush after flush_limit_mutation_count mutations """ - from google.cloud.bigtable.mutations import RowMutationEntry + from google.cloud.bigtable.data.mutations import RowMutationEntry new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] row_key, mutation = await _create_row_and_mutation( @@ -407,7 +407,7 @@ async def test_mutations_batcher_bytes_flush(client, table, temp_rows): """ batch should flush after flush_limit_bytes bytes """ - from google.cloud.bigtable.mutations import RowMutationEntry + from google.cloud.bigtable.data.mutations import RowMutationEntry new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] row_key, mutation = await _create_row_and_mutation( @@ -442,7 +442,7 @@ async def test_mutations_batcher_no_flush(client, table, temp_rows): """ test with no flush requirements met """ - from google.cloud.bigtable.mutations import RowMutationEntry + from google.cloud.bigtable.data.mutations import RowMutationEntry new_value = uuid.uuid4().hex.encode() start_value = b"unchanged" @@ -494,7 +494,7 @@ async def test_read_modify_write_row_increment( """ test read_modify_write_row """ - from google.cloud.bigtable.read_modify_write_rules import IncrementRule + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule row_key = b"test-row-key" family = TEST_FAMILY @@ -531,7 +531,7 @@ async def test_read_modify_write_row_append( """ test read_modify_write_row """ - from google.cloud.bigtable.read_modify_write_rules import AppendValueRule + from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule row_key = b"test-row-key" family = TEST_FAMILY @@ -554,8 +554,8 @@ async def test_read_modify_write_row_chained(client, table, temp_rows): """ test read_modify_write_row with multiple rules """ - from google.cloud.bigtable.read_modify_write_rules import AppendValueRule - from google.cloud.bigtable.read_modify_write_rules import IncrementRule + from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule row_key = b"test-row-key" family = TEST_FAMILY @@ -599,8 +599,8 @@ async def test_check_and_mutate( """ test that check_and_mutate_row works applies the right mutations, and returns the right result """ - from google.cloud.bigtable.mutations import SetCell - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter row_key = b"test-row-key" family = TEST_FAMILY @@ -671,7 +671,7 @@ async def test_read_rows_sharded_simple(table, temp_rows): """ Test read rows sharded with two queries """ - from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery await temp_rows.add_row(b"a") await temp_rows.add_row(b"b") @@ -693,8 +693,8 @@ async def test_read_rows_sharded_from_sample(table, temp_rows): """ Test end-to-end sharding """ - from google.cloud.bigtable.read_rows_query import ReadRowsQuery - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import RowRange await temp_rows.add_row(b"a") await temp_rows.add_row(b"b") @@ -717,8 +717,8 @@ async def test_read_rows_sharded_filters_limits(table, temp_rows): """ Test read rows sharded with filters and limits """ - from google.cloud.bigtable.read_rows_query import ReadRowsQuery - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter await temp_rows.add_row(b"a") await temp_rows.add_row(b"b") @@ -745,8 +745,8 @@ async def test_read_rows_range_query(table, temp_rows): """ Ensure that the read_rows method works """ - from google.cloud.bigtable import ReadRowsQuery - from google.cloud.bigtable import RowRange + from google.cloud.bigtable.data import ReadRowsQuery + from google.cloud.bigtable.data import RowRange await temp_rows.add_row(b"a") await temp_rows.add_row(b"b") @@ -766,7 +766,7 @@ async def test_read_rows_single_key_query(table, temp_rows): """ Ensure that the read_rows method works with specified query """ - from google.cloud.bigtable import ReadRowsQuery + from google.cloud.bigtable.data import ReadRowsQuery await temp_rows.add_row(b"a") await temp_rows.add_row(b"b") @@ -786,8 +786,8 @@ async def test_read_rows_with_filter(table, temp_rows): """ ensure filters are applied """ - from google.cloud.bigtable import ReadRowsQuery - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter await temp_rows.add_row(b"a") await temp_rows.add_row(b"b") @@ -828,7 +828,7 @@ async def test_read_rows_stream_inactive_timer(table, temp_rows): """ Ensure that the read_rows_stream method works """ - from google.cloud.bigtable.exceptions import IdleTimeout + from google.cloud.bigtable.data.exceptions import IdleTimeout await temp_rows.add_row(b"row_key_1") await temp_rows.add_row(b"row_key_2") @@ -848,7 +848,7 @@ async def test_read_row(table, temp_rows): """ Test read_row (single row helper) """ - from google.cloud.bigtable import Row + from google.cloud.bigtable.data import Row await temp_rows.add_row(b"row_key_1", value=b"value") row = await table.read_row(b"row_key_1") @@ -877,8 +877,8 @@ async def test_read_row_w_filter(table, temp_rows): """ Test read_row (single row helper) """ - from google.cloud.bigtable import Row - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data import Row + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter await temp_rows.add_row(b"row_key_1", value=b"value") expected_label = "test-label" @@ -943,8 +943,8 @@ async def test_literal_value_filter( Literal value filter does complex escaping on re2 strings. Make sure inputs are properly interpreted by the server """ - from google.cloud.bigtable.row_filters import LiteralValueFilter - from google.cloud.bigtable import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import LiteralValueFilter + from google.cloud.bigtable.data import ReadRowsQuery f = LiteralValueFilter(filter_input) await temp_rows.add_row(b"row_key_1", value=cell_value) diff --git a/tests/system/v2_client/conftest.py b/tests/system/v2_client/conftest.py index bb4f54b41..f39fcba88 100644 --- a/tests/system/v2_client/conftest.py +++ b/tests/system/v2_client/conftest.py @@ -17,7 +17,7 @@ import pytest from test_utils.system import unique_resource_id -from google.cloud.bigtable.deprecated.client import Client +from google.cloud.bigtable.client import Client from google.cloud.environment_vars import BIGTABLE_EMULATOR from . import _helpers diff --git a/tests/system/v2_client/test_data_api.py b/tests/system/v2_client/test_data_api.py index 551a221ee..2ca7e1504 100644 --- a/tests/system/v2_client/test_data_api.py +++ b/tests/system/v2_client/test_data_api.py @@ -60,7 +60,7 @@ def rows_to_delete(): def test_table_read_rows_filter_millis(data_table): - from google.cloud.bigtable.deprecated import row_filters + from google.cloud.bigtable import row_filters end = datetime.datetime.now() start = end - datetime.timedelta(minutes=60) @@ -158,8 +158,8 @@ def test_table_drop_by_prefix(data_table, rows_to_delete): def test_table_read_rows_w_row_set(data_table, rows_to_delete): - from google.cloud.bigtable.deprecated.row_set import RowSet - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange row_keys = [ b"row_key_1", @@ -189,7 +189,7 @@ def test_table_read_rows_w_row_set(data_table, rows_to_delete): def test_rowset_add_row_range_w_pfx(data_table, rows_to_delete): - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowSet row_keys = [ b"row_key_1", @@ -234,7 +234,7 @@ def _write_to_row(row1, row2, row3, row4): from google.cloud._helpers import _datetime_from_microseconds from google.cloud._helpers import _microseconds_from_datetime from google.cloud._helpers import UTC - from google.cloud.bigtable.deprecated.row_data import Cell + from google.cloud.bigtable.row_data import Cell timestamp1 = datetime.datetime.utcnow().replace(tzinfo=UTC) timestamp1_micros = _microseconds_from_datetime(timestamp1) @@ -290,7 +290,7 @@ def test_table_read_row(data_table, rows_to_delete): def test_table_read_rows(data_table, rows_to_delete): - from google.cloud.bigtable.deprecated.row_data import PartialRowData + from google.cloud.bigtable.row_data import PartialRowData row = data_table.direct_row(ROW_KEY) rows_to_delete.append(row) @@ -326,10 +326,10 @@ def test_table_read_rows(data_table, rows_to_delete): def test_read_with_label_applied(data_table, rows_to_delete, skip_on_emulator): - from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter - from google.cloud.bigtable.deprecated.row_filters import ColumnQualifierRegexFilter - from google.cloud.bigtable.deprecated.row_filters import RowFilterChain - from google.cloud.bigtable.deprecated.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowFilterUnion row = data_table.direct_row(ROW_KEY) rows_to_delete.append(row) diff --git a/tests/system/v2_client/test_instance_admin.py b/tests/system/v2_client/test_instance_admin.py index debe1ab56..e5e311213 100644 --- a/tests/system/v2_client/test_instance_admin.py +++ b/tests/system/v2_client/test_instance_admin.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from google.cloud.bigtable.deprecated import enums -from google.cloud.bigtable.deprecated.table import ClusterState +from google.cloud.bigtable import enums +from google.cloud.bigtable.table import ClusterState from . import _helpers @@ -149,7 +149,7 @@ def test_instance_create_prod( instances_to_delete, skip_on_emulator, ): - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums alt_instance_id = f"ndef{unique_suffix}" instance = admin_client.instance(alt_instance_id, labels=instance_labels) diff --git a/tests/system/v2_client/test_table_admin.py b/tests/system/v2_client/test_table_admin.py index 107ed41bf..c50189013 100644 --- a/tests/system/v2_client/test_table_admin.py +++ b/tests/system/v2_client/test_table_admin.py @@ -97,7 +97,7 @@ def test_table_create_w_families( data_instance_populated, tables_to_delete, ): - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule temp_table_id = "test-create-table-with-failies" column_family_id = "col-fam-id1" @@ -134,7 +134,7 @@ def test_table_create_w_split_keys( def test_column_family_create(data_instance_populated, tables_to_delete): - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule temp_table_id = "test-create-column-family" temp_table = data_instance_populated.table(temp_table_id) @@ -158,7 +158,7 @@ def test_column_family_create(data_instance_populated, tables_to_delete): def test_column_family_update(data_instance_populated, tables_to_delete): - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule temp_table_id = "test-update-column-family" temp_table = data_instance_populated.table(temp_table_id) @@ -219,8 +219,8 @@ def test_table_get_iam_policy( def test_table_set_iam_policy( service_account, data_instance_populated, tables_to_delete, skip_on_emulator ): - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE - from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import Policy temp_table_id = "test-set-iam-policy-table" temp_table = data_instance_populated.table(temp_table_id) @@ -264,7 +264,7 @@ def test_table_backup( skip_on_emulator, ): from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums temp_table_id = "test-backup-table" temp_table = data_instance_populated.table(temp_table_id) diff --git a/tests/unit/data/__init__.py b/tests/unit/data/__init__.py new file mode 100644 index 000000000..89a37dc92 --- /dev/null +++ b/tests/unit/data/__init__.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +# Copyright 2023 Google LLC +# +# Licensed 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. +# diff --git a/tests/unit/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py similarity index 93% rename from tests/unit/test__mutate_rows.py rename to tests/unit/data/_async/test__mutate_rows.py index 18b2beede..f77455d60 100644 --- a/tests/unit/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -36,9 +36,11 @@ def _make_mutation(count=1, size=1): class TestMutateRowsOperation: def _target_class(self): - from google.cloud.bigtable._mutate_rows import _MutateRowsOperation + from google.cloud.bigtable.data._async._mutate_rows import ( + _MutateRowsOperationAsync, + ) - return _MutateRowsOperation + return _MutateRowsOperationAsync def _make_one(self, *args, **kwargs): if not args: @@ -73,7 +75,7 @@ def test_ctor(self): """ test that constructor sets all the attributes correctly """ - from google.cloud.bigtable._mutate_rows import _MutateRowsIncomplete + from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import ServiceUnavailable @@ -116,7 +118,7 @@ def test_ctor_too_many_entries(self): """ should raise an error if an operation is created with more than 100,000 entries """ - from google.cloud.bigtable._mutate_rows import ( + from google.cloud.bigtable.data._async._mutate_rows import ( MUTATE_ROWS_REQUEST_MUTATION_LIMIT, ) @@ -168,8 +170,8 @@ async def test_mutate_rows_exception(self, exc_type): """ exceptions raised from retryable should be raised in MutationsExceptionGroup """ - from google.cloud.bigtable.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.exceptions import FailedMutationEntryError + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.exceptions import FailedMutationEntryError client = mock.Mock() table = mock.Mock() @@ -204,7 +206,9 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): """ If an exception fails but eventually passes, it should not raise an exception """ - from google.cloud.bigtable._mutate_rows import _MutateRowsOperation + from google.cloud.bigtable.data._async._mutate_rows import ( + _MutateRowsOperationAsync, + ) client = mock.Mock() table = mock.Mock() @@ -213,7 +217,7 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): expected_cause = exc_type("retry") num_retries = 2 with mock.patch.object( - _MutateRowsOperation, + _MutateRowsOperationAsync, "_run_attempt", AsyncMock(), ) as attempt_mock: @@ -229,8 +233,8 @@ async def test_mutate_rows_incomplete_ignored(self): """ MutateRowsIncomplete exceptions should not be added to error list """ - from google.cloud.bigtable._mutate_rows import _MutateRowsIncomplete - from google.cloud.bigtable.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup from google.api_core.exceptions import DeadlineExceeded client = mock.Mock() @@ -286,7 +290,7 @@ async def test_run_attempt_empty_request(self): @pytest.mark.asyncio async def test_run_attempt_partial_success_retryable(self): """Some entries succeed, but one fails. Should report the proper index, and raise incomplete exception""" - from google.cloud.bigtable._mutate_rows import _MutateRowsIncomplete + from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete success_mutation = _make_mutation() success_mutation_2 = _make_mutation() diff --git a/tests/unit/data/_async/test__read_rows.py b/tests/unit/data/_async/test__read_rows.py new file mode 100644 index 000000000..c7d52280c --- /dev/null +++ b/tests/unit/data/_async/test__read_rows.py @@ -0,0 +1,625 @@ +# +# Licensed 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. + +import pytest +import sys +import asyncio + +from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock + from unittest.mock import AsyncMock # type: ignore +except ImportError: # pragma: NO COVER + import mock # type: ignore + from mock import AsyncMock # type: ignore # noqa F401 + +TEST_FAMILY = "family_name" +TEST_QUALIFIER = b"qualifier" +TEST_TIMESTAMP = 123456789 +TEST_LABELS = ["label1", "label2"] + + +class TestReadRowsOperation: + """ + Tests helper functions in the ReadRowsOperation class + in-depth merging logic in merge_row_response_stream and _read_rows_retryable_attempt + is tested in test_read_rows_acceptance test_client_read_rows, and conformance tests + """ + + @staticmethod + def _get_target_class(): + from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync + + return _ReadRowsOperationAsync + + def _make_one(self, *args, **kwargs): + return self._get_target_class()(*args, **kwargs) + + def test_ctor_defaults(self): + request = {} + client = mock.Mock() + client.read_rows = mock.Mock() + client.read_rows.return_value = None + default_operation_timeout = 600 + time_gen_mock = mock.Mock() + with mock.patch( + "google.cloud.bigtable.data._async._read_rows._attempt_timeout_generator", + time_gen_mock, + ): + instance = self._make_one(request, client) + assert time_gen_mock.call_count == 1 + time_gen_mock.assert_called_once_with(None, default_operation_timeout) + assert instance.transient_errors == [] + assert instance._last_emitted_row_key is None + assert instance._emit_count == 0 + assert instance.operation_timeout == default_operation_timeout + retryable_fn = instance._partial_retryable + assert retryable_fn.func == instance._read_rows_retryable_attempt + assert retryable_fn.args[0] == client.read_rows + assert retryable_fn.args[1] == time_gen_mock.return_value + assert retryable_fn.args[2] == 0 + assert client.read_rows.call_count == 0 + + def test_ctor(self): + row_limit = 91 + request = {"rows_limit": row_limit} + client = mock.Mock() + client.read_rows = mock.Mock() + client.read_rows.return_value = None + expected_operation_timeout = 42 + expected_request_timeout = 44 + time_gen_mock = mock.Mock() + with mock.patch( + "google.cloud.bigtable.data._async._read_rows._attempt_timeout_generator", + time_gen_mock, + ): + instance = self._make_one( + request, + client, + operation_timeout=expected_operation_timeout, + per_request_timeout=expected_request_timeout, + ) + assert time_gen_mock.call_count == 1 + time_gen_mock.assert_called_once_with( + expected_request_timeout, expected_operation_timeout + ) + assert instance.transient_errors == [] + assert instance._last_emitted_row_key is None + assert instance._emit_count == 0 + assert instance.operation_timeout == expected_operation_timeout + retryable_fn = instance._partial_retryable + assert retryable_fn.func == instance._read_rows_retryable_attempt + assert retryable_fn.args[0] == client.read_rows + assert retryable_fn.args[1] == time_gen_mock.return_value + assert retryable_fn.args[2] == row_limit + assert client.read_rows.call_count == 0 + + def test___aiter__(self): + request = {} + client = mock.Mock() + client.read_rows = mock.Mock() + instance = self._make_one(request, client) + assert instance.__aiter__() is instance + + @pytest.mark.asyncio + async def test_transient_error_capture(self): + from google.api_core import exceptions as core_exceptions + + client = mock.Mock() + client.read_rows = mock.Mock() + test_exc = core_exceptions.Aborted("test") + test_exc2 = core_exceptions.DeadlineExceeded("test") + client.read_rows.side_effect = [test_exc, test_exc2] + instance = self._make_one({}, client) + with pytest.raises(RuntimeError): + await instance.__anext__() + assert len(instance.transient_errors) == 2 + assert instance.transient_errors[0] == test_exc + assert instance.transient_errors[1] == test_exc2 + + @pytest.mark.parametrize( + "in_keys,last_key,expected", + [ + (["b", "c", "d"], "a", ["b", "c", "d"]), + (["a", "b", "c"], "b", ["c"]), + (["a", "b", "c"], "c", []), + (["a", "b", "c"], "d", []), + (["d", "c", "b", "a"], "b", ["d", "c"]), + ], + ) + def test_revise_request_rowset_keys(self, in_keys, last_key, expected): + sample_range = {"start_key_open": last_key} + row_set = {"row_keys": in_keys, "row_ranges": [sample_range]} + revised = self._get_target_class()._revise_request_rowset(row_set, last_key) + assert revised["row_keys"] == expected + assert revised["row_ranges"] == [sample_range] + + @pytest.mark.parametrize( + "in_ranges,last_key,expected", + [ + ( + [{"start_key_open": "b", "end_key_closed": "d"}], + "a", + [{"start_key_open": "b", "end_key_closed": "d"}], + ), + ( + [{"start_key_closed": "b", "end_key_closed": "d"}], + "a", + [{"start_key_closed": "b", "end_key_closed": "d"}], + ), + ( + [{"start_key_open": "a", "end_key_closed": "d"}], + "b", + [{"start_key_open": "b", "end_key_closed": "d"}], + ), + ( + [{"start_key_closed": "a", "end_key_open": "d"}], + "b", + [{"start_key_open": "b", "end_key_open": "d"}], + ), + ( + [{"start_key_closed": "b", "end_key_closed": "d"}], + "b", + [{"start_key_open": "b", "end_key_closed": "d"}], + ), + ([{"start_key_closed": "b", "end_key_closed": "d"}], "d", []), + ([{"start_key_closed": "b", "end_key_open": "d"}], "d", []), + ([{"start_key_closed": "b", "end_key_closed": "d"}], "e", []), + ([{"start_key_closed": "b"}], "z", [{"start_key_open": "z"}]), + ([{"start_key_closed": "b"}], "a", [{"start_key_closed": "b"}]), + ( + [{"end_key_closed": "z"}], + "a", + [{"start_key_open": "a", "end_key_closed": "z"}], + ), + ( + [{"end_key_open": "z"}], + "a", + [{"start_key_open": "a", "end_key_open": "z"}], + ), + ], + ) + def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected): + next_key = last_key + "a" + row_set = {"row_keys": [next_key], "row_ranges": in_ranges} + revised = self._get_target_class()._revise_request_rowset(row_set, last_key) + assert revised["row_keys"] == [next_key] + assert revised["row_ranges"] == expected + + @pytest.mark.parametrize("last_key", ["a", "b", "c"]) + def test_revise_request_full_table(self, last_key): + row_set = {"row_keys": [], "row_ranges": []} + for selected_set in [row_set, None]: + revised = self._get_target_class()._revise_request_rowset( + selected_set, last_key + ) + assert revised["row_keys"] == [] + assert len(revised["row_ranges"]) == 1 + assert revised["row_ranges"][0]["start_key_open"] == last_key + + def test_revise_to_empty_rowset(self): + """revising to an empty rowset should raise error""" + from google.cloud.bigtable.data.exceptions import _RowSetComplete + + row_keys = ["a", "b", "c"] + row_set = {"row_keys": row_keys, "row_ranges": [{"end_key_open": "c"}]} + with pytest.raises(_RowSetComplete): + self._get_target_class()._revise_request_rowset(row_set, "d") + + @pytest.mark.parametrize( + "start_limit,emit_num,expected_limit", + [ + (10, 0, 10), + (10, 1, 9), + (10, 10, 0), + (0, 10, 0), + (0, 0, 0), + (4, 2, 2), + ], + ) + @pytest.mark.asyncio + async def test_revise_limit(self, start_limit, emit_num, expected_limit): + """ + revise_limit should revise the request's limit field + - if limit is 0 (unlimited), it should never be revised + - if start_limit-emit_num == 0, the request should end early + - if the number emitted exceeds the new limit, an exception should + should be raised (tested in test_revise_limit_over_limit) + """ + import itertools + + request = {"rows_limit": start_limit} + instance = self._make_one(request, mock.Mock()) + instance._emit_count = emit_num + instance._last_emitted_row_key = "a" + gapic_mock = mock.Mock() + gapic_mock.side_effect = [GeneratorExit("stop_fn")] + mock_timeout_gen = itertools.repeat(5) + + attempt = instance._read_rows_retryable_attempt( + gapic_mock, mock_timeout_gen, start_limit + ) + if start_limit != 0 and expected_limit == 0: + # if we emitted the expected number of rows, we should receive a StopAsyncIteration + with pytest.raises(StopAsyncIteration): + await attempt.__anext__() + else: + with pytest.raises(GeneratorExit): + await attempt.__anext__() + assert request["rows_limit"] == expected_limit + + @pytest.mark.parametrize("start_limit,emit_num", [(5, 10), (3, 9), (1, 10)]) + @pytest.mark.asyncio + async def test_revise_limit_over_limit(self, start_limit, emit_num): + """ + Should raise runtime error if we get in state where emit_num > start_num + (unless start_num == 0, which represents unlimited) + """ + import itertools + + request = {"rows_limit": start_limit} + instance = self._make_one(request, mock.Mock()) + instance._emit_count = emit_num + instance._last_emitted_row_key = "a" + mock_timeout_gen = itertools.repeat(5) + attempt = instance._read_rows_retryable_attempt( + mock.Mock(), mock_timeout_gen, start_limit + ) + with pytest.raises(RuntimeError) as e: + await attempt.__anext__() + assert "emit count exceeds row limit" in str(e.value) + + @pytest.mark.asyncio + async def test_aclose(self): + import asyncio + + instance = self._make_one({}, mock.Mock()) + await instance.aclose() + assert instance._stream is None + assert instance._last_emitted_row_key is None + with pytest.raises(asyncio.InvalidStateError): + await instance.__anext__() + # try calling a second time + await instance.aclose() + + @pytest.mark.parametrize("limit", [1, 3, 10]) + @pytest.mark.asyncio + async def test_retryable_attempt_hit_limit(self, limit): + """ + Stream should end after hitting the limit + """ + from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse + import itertools + + instance = self._make_one({}, mock.Mock()) + + async def mock_gapic(*args, **kwargs): + # continuously return a single row + async def gen(): + for i in range(limit * 2): + chunk = ReadRowsResponse.CellChunk( + row_key=str(i).encode(), + family_name="family_name", + qualifier=b"qualifier", + commit_row=True, + ) + yield ReadRowsResponse(chunks=[chunk]) + + return gen() + + mock_timeout_gen = itertools.repeat(5) + gen = instance._read_rows_retryable_attempt(mock_gapic, mock_timeout_gen, limit) + # should yield values up to the limit + for i in range(limit): + await gen.__anext__() + # next value should be StopAsyncIteration + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + + @pytest.mark.asyncio + async def test_retryable_ignore_repeated_rows(self): + """ + Duplicate rows should cause an invalid chunk error + """ + from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync + from google.cloud.bigtable.data.row import Row + from google.cloud.bigtable.data.exceptions import InvalidChunk + + async def mock_stream(): + while True: + yield Row(b"dup_key", cells=[]) + yield Row(b"dup_key", cells=[]) + + with mock.patch.object( + _ReadRowsOperationAsync, "merge_row_response_stream" + ) as mock_stream_fn: + mock_stream_fn.return_value = mock_stream() + instance = self._make_one({}, mock.AsyncMock()) + first_row = await instance.__anext__() + assert first_row.row_key == b"dup_key" + with pytest.raises(InvalidChunk) as exc: + await instance.__anext__() + assert "Last emitted row key out of order" in str(exc.value) + + @pytest.mark.asyncio + async def test_retryable_ignore_last_scanned_rows(self): + """ + Last scanned rows should not be emitted + """ + from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync + from google.cloud.bigtable.data.row import Row, _LastScannedRow + + async def mock_stream(): + while True: + yield Row(b"key1", cells=[]) + yield _LastScannedRow(b"key2_ignored") + yield Row(b"key3", cells=[]) + + with mock.patch.object( + _ReadRowsOperationAsync, "merge_row_response_stream" + ) as mock_stream_fn: + mock_stream_fn.return_value = mock_stream() + instance = self._make_one({}, mock.AsyncMock()) + first_row = await instance.__anext__() + assert first_row.row_key == b"key1" + second_row = await instance.__anext__() + assert second_row.row_key == b"key3" + + @pytest.mark.asyncio + async def test_retryable_cancel_on_close(self): + """Underlying gapic call should be cancelled when stream is closed""" + from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync + from google.cloud.bigtable.data.row import Row + + async def mock_stream(): + while True: + yield Row(b"key1", cells=[]) + + with mock.patch.object( + _ReadRowsOperationAsync, "merge_row_response_stream" + ) as mock_stream_fn: + mock_stream_fn.return_value = mock_stream() + mock_gapic = mock.AsyncMock() + mock_call = await mock_gapic.read_rows() + instance = self._make_one({}, mock_gapic) + await instance.__anext__() + assert mock_call.cancel.call_count == 0 + await instance.aclose() + assert mock_call.cancel.call_count == 1 + + +class MockStream(_ReadRowsOperationAsync): + """ + Mock a _ReadRowsOperationAsync stream for testing + """ + + def __init__(self, items=None, errors=None, operation_timeout=None): + self.transient_errors = errors + self.operation_timeout = operation_timeout + self.next_idx = 0 + if items is None: + items = list(range(10)) + self.items = items + + def __aiter__(self): + return self + + async def __anext__(self): + if self.next_idx >= len(self.items): + raise StopAsyncIteration + item = self.items[self.next_idx] + self.next_idx += 1 + if isinstance(item, Exception): + raise item + return item + + async def aclose(self): + pass + + +class TestReadRowsAsyncIterator: + async def mock_stream(self, size=10): + for i in range(size): + yield i + + def _make_one(self, *args, **kwargs): + from google.cloud.bigtable.data._async._read_rows import ReadRowsAsyncIterator + + stream = MockStream(*args, **kwargs) + return ReadRowsAsyncIterator(stream) + + def test_ctor(self): + with mock.patch("time.monotonic", return_value=0): + iterator = self._make_one() + assert iterator._last_interaction_time == 0 + assert iterator._idle_timeout_task is None + assert iterator.active is True + + def test___aiter__(self): + iterator = self._make_one() + assert iterator.__aiter__() is iterator + + @pytest.mark.skipif( + sys.version_info < (3, 8), reason="mock coroutine requires python3.8 or higher" + ) + @pytest.mark.asyncio + async def test__start_idle_timer(self): + """Should start timer coroutine""" + iterator = self._make_one() + expected_timeout = 10 + with mock.patch("time.monotonic", return_value=1): + with mock.patch.object(iterator, "_idle_timeout_coroutine") as mock_coro: + await iterator._start_idle_timer(expected_timeout) + assert mock_coro.call_count == 1 + assert mock_coro.call_args[0] == (expected_timeout,) + assert iterator._last_interaction_time == 1 + assert iterator._idle_timeout_task is not None + + @pytest.mark.skipif( + sys.version_info < (3, 8), reason="mock coroutine requires python3.8 or higher" + ) + @pytest.mark.asyncio + async def test__start_idle_timer_duplicate(self): + """Multiple calls should replace task""" + iterator = self._make_one() + with mock.patch.object(iterator, "_idle_timeout_coroutine") as mock_coro: + await iterator._start_idle_timer(1) + first_task = iterator._idle_timeout_task + await iterator._start_idle_timer(2) + second_task = iterator._idle_timeout_task + assert mock_coro.call_count == 2 + + assert first_task is not None + assert first_task != second_task + # old tasks hould be cancelled + with pytest.raises(asyncio.CancelledError): + await first_task + # new task should not be cancelled + await second_task + + @pytest.mark.asyncio + async def test__idle_timeout_coroutine(self): + from google.cloud.bigtable.data.exceptions import IdleTimeout + + iterator = self._make_one() + await iterator._idle_timeout_coroutine(0.05) + await asyncio.sleep(0.1) + assert iterator.active is False + with pytest.raises(IdleTimeout): + await iterator.__anext__() + + @pytest.mark.asyncio + async def test__idle_timeout_coroutine_extensions(self): + """touching the generator should reset the idle timer""" + iterator = self._make_one(items=list(range(100))) + await iterator._start_idle_timer(0.05) + for i in range(10): + # will not expire as long as it is in use + assert iterator.active is True + await iterator.__anext__() + await asyncio.sleep(0.03) + # now let it expire + await asyncio.sleep(0.5) + assert iterator.active is False + + @pytest.mark.asyncio + async def test___anext__(self): + num_rows = 10 + iterator = self._make_one(items=list(range(num_rows))) + for i in range(num_rows): + assert await iterator.__anext__() == i + with pytest.raises(StopAsyncIteration): + await iterator.__anext__() + + @pytest.mark.asyncio + async def test___anext__with_deadline_error(self): + """ + RetryErrors mean a deadline has been hit. + Should be wrapped in a DeadlineExceeded exception + """ + from google.api_core import exceptions as core_exceptions + + items = [1, core_exceptions.RetryError("retry error", None)] + expected_timeout = 99 + iterator = self._make_one(items=items, operation_timeout=expected_timeout) + assert await iterator.__anext__() == 1 + with pytest.raises(core_exceptions.DeadlineExceeded) as exc: + await iterator.__anext__() + assert f"operation_timeout of {expected_timeout:0.1f}s exceeded" in str( + exc.value + ) + assert exc.value.__cause__ is None + + @pytest.mark.asyncio + async def test___anext__with_deadline_error_with_cause(self): + """ + Transient errors should be exposed as an error group + """ + from google.api_core import exceptions as core_exceptions + from google.cloud.bigtable.data.exceptions import RetryExceptionGroup + + items = [1, core_exceptions.RetryError("retry error", None)] + expected_timeout = 99 + errors = [RuntimeError("error1"), ValueError("error2")] + iterator = self._make_one( + items=items, operation_timeout=expected_timeout, errors=errors + ) + assert await iterator.__anext__() == 1 + with pytest.raises(core_exceptions.DeadlineExceeded) as exc: + await iterator.__anext__() + assert f"operation_timeout of {expected_timeout:0.1f}s exceeded" in str( + exc.value + ) + error_group = exc.value.__cause__ + assert isinstance(error_group, RetryExceptionGroup) + assert len(error_group.exceptions) == 2 + assert error_group.exceptions[0] is errors[0] + assert error_group.exceptions[1] is errors[1] + assert "2 failed attempts" in str(error_group) + + @pytest.mark.asyncio + async def test___anext__with_error(self): + """ + Other errors should be raised as-is + """ + from google.api_core import exceptions as core_exceptions + + items = [1, core_exceptions.InternalServerError("mock error")] + iterator = self._make_one(items=items) + assert await iterator.__anext__() == 1 + with pytest.raises(core_exceptions.InternalServerError) as exc: + await iterator.__anext__() + assert exc.value is items[1] + assert iterator.active is False + # next call should raise same error + with pytest.raises(core_exceptions.InternalServerError) as exc: + await iterator.__anext__() + + @pytest.mark.asyncio + async def test__finish_with_error(self): + iterator = self._make_one() + await iterator._start_idle_timer(10) + timeout_task = iterator._idle_timeout_task + assert await iterator.__anext__() == 0 + assert iterator.active is True + err = ZeroDivisionError("mock error") + await iterator._finish_with_error(err) + assert iterator.active is False + assert iterator._error is err + assert iterator._idle_timeout_task is None + with pytest.raises(ZeroDivisionError) as exc: + await iterator.__anext__() + assert exc.value is err + # timeout task should be cancelled + with pytest.raises(asyncio.CancelledError): + await timeout_task + + @pytest.mark.asyncio + async def test_aclose(self): + iterator = self._make_one() + await iterator._start_idle_timer(10) + timeout_task = iterator._idle_timeout_task + assert await iterator.__anext__() == 0 + assert iterator.active is True + await iterator.aclose() + assert iterator.active is False + assert isinstance(iterator._error, StopAsyncIteration) + assert iterator._idle_timeout_task is None + with pytest.raises(StopAsyncIteration) as e: + await iterator.__anext__() + assert "closed" in str(e.value) + # timeout task should be cancelled + with pytest.raises(asyncio.CancelledError): + await timeout_task diff --git a/tests/unit/test_client.py b/tests/unit/data/_async/test_client.py similarity index 95% rename from tests/unit/test_client.py rename to tests/unit/data/_async/test_client.py index 3557c1c16..25006d725 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -20,15 +20,15 @@ import pytest -from google.cloud.bigtable import mutations +from google.cloud.bigtable.data import mutations from google.auth.credentials import AnonymousCredentials from google.cloud.bigtable_v2.types import ReadRowsResponse -from google.cloud.bigtable.read_rows_query import ReadRowsQuery +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.api_core import exceptions as core_exceptions -from google.cloud.bigtable.exceptions import InvalidChunk +from google.cloud.bigtable.data.exceptions import InvalidChunk -from google.cloud.bigtable.read_modify_write_rules import IncrementRule -from google.cloud.bigtable.read_modify_write_rules import AppendValueRule +from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule +from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule # try/except added for compatibility with python < 3.8 try: @@ -43,11 +43,11 @@ ) -class TestBigtableDataClient: +class TestBigtableDataClientAsync: def _get_target_class(self): - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - return BigtableDataClient + return BigtableDataClientAsync def _make_one(self, *args, **kwargs): return self._get_target_class()(*args, **kwargs) @@ -118,7 +118,6 @@ async def test_ctor_dict_options(self): BigtableAsyncClient, ) from google.api_core.client_options import ClientOptions - from google.cloud.bigtable.client import BigtableDataClient client_options = {"api_endpoint": "foo.bar:1234"} with mock.patch.object(BigtableAsyncClient, "__init__") as bigtable_client_init: @@ -132,7 +131,7 @@ async def test_ctor_dict_options(self): assert called_options.api_endpoint == "foo.bar:1234" assert isinstance(called_options, ClientOptions) with mock.patch.object( - BigtableDataClient, "start_background_channel_refresh" + self._get_target_class(), "start_background_channel_refresh" ) as start_background_refresh: client = self._make_one(client_options=client_options) start_background_refresh.assert_called_once() @@ -275,7 +274,7 @@ async def test_start_background_channel_refresh_tasks_names(self): for i in range(pool_size): name = client._channel_refresh_tasks[i].get_name() assert str(i) in name - assert "BigtableDataClient channel refresh " in name + assert "BigtableDataClientAsync channel refresh " in name await client.close() @pytest.mark.asyncio @@ -725,7 +724,7 @@ async def test__multiple_table_registration(self): add multiple owners to instance_owners, but only keep one copy of shared key in active_instances """ - from google.cloud.bigtable.client import _WarmedInstanceKey + from google.cloud.bigtable.data._async.client import _WarmedInstanceKey async with self._make_one(project="project-id") as client: async with client.get_table("instance_1", "table_1") as table_1: @@ -773,7 +772,7 @@ async def test__multiple_instance_registration(self): registering with multiple instance keys should update the key in instance_owners and active_instances """ - from google.cloud.bigtable.client import _WarmedInstanceKey + from google.cloud.bigtable.data._async.client import _WarmedInstanceKey async with self._make_one(project="project-id") as client: async with client.get_table("instance_1", "table_1") as table_1: @@ -808,8 +807,8 @@ async def test__multiple_instance_registration(self): @pytest.mark.asyncio async def test_get_table(self): - from google.cloud.bigtable.client import Table - from google.cloud.bigtable.client import _WarmedInstanceKey + from google.cloud.bigtable.data._async.client import TableAsync + from google.cloud.bigtable.data._async.client import _WarmedInstanceKey client = self._make_one(project="project-id") assert not client._active_instances @@ -822,7 +821,7 @@ async def test_get_table(self): expected_app_profile_id, ) await asyncio.sleep(0) - assert isinstance(table, Table) + assert isinstance(table, TableAsync) assert table.table_id == expected_table_id assert ( table.table_name @@ -844,15 +843,15 @@ async def test_get_table(self): @pytest.mark.asyncio async def test_get_table_context_manager(self): - from google.cloud.bigtable.client import Table - from google.cloud.bigtable.client import _WarmedInstanceKey + from google.cloud.bigtable.data._async.client import TableAsync + from google.cloud.bigtable.data._async.client import _WarmedInstanceKey expected_table_id = "table-id" expected_instance_id = "instance-id" expected_app_profile_id = "app-profile-id" expected_project_id = "project-id" - with mock.patch.object(Table, "close") as close_mock: + with mock.patch.object(TableAsync, "close") as close_mock: async with self._make_one(project=expected_project_id) as client: async with client.get_table( expected_instance_id, @@ -860,7 +859,7 @@ async def test_get_table_context_manager(self): expected_app_profile_id, ) as table: await asyncio.sleep(0) - assert isinstance(table, Table) + assert isinstance(table, TableAsync) assert table.table_id == expected_table_id assert ( table.table_name @@ -950,35 +949,36 @@ async def test_context_manager(self): def test_client_ctor_sync(self): # initializing client in a sync context should raise RuntimeError - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync with pytest.warns(RuntimeWarning) as warnings: - client = BigtableDataClient(project="project-id") + client = BigtableDataClientAsync(project="project-id") expected_warning = [w for w in warnings if "client.py" in w.filename] assert len(expected_warning) == 1 - assert "BigtableDataClient should be started in an asyncio event loop." in str( - expected_warning[0].message + assert ( + "BigtableDataClientAsync should be started in an asyncio event loop." + in str(expected_warning[0].message) ) assert client.project == "project-id" assert client._channel_refresh_tasks == [] -class TestTable: +class TestTableAsync: @pytest.mark.asyncio async def test_table_ctor(self): - from google.cloud.bigtable.client import BigtableDataClient - from google.cloud.bigtable.client import Table - from google.cloud.bigtable.client import _WarmedInstanceKey + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync + from google.cloud.bigtable.data._async.client import TableAsync + from google.cloud.bigtable.data._async.client import _WarmedInstanceKey expected_table_id = "table-id" expected_instance_id = "instance-id" expected_app_profile_id = "app-profile-id" expected_operation_timeout = 123 expected_per_request_timeout = 12 - client = BigtableDataClient() + client = BigtableDataClientAsync() assert not client._active_instances - table = Table( + table = TableAsync( client, expected_instance_id, expected_table_id, @@ -1007,19 +1007,19 @@ async def test_table_ctor(self): @pytest.mark.asyncio async def test_table_ctor_bad_timeout_values(self): - from google.cloud.bigtable.client import BigtableDataClient - from google.cloud.bigtable.client import Table + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync + from google.cloud.bigtable.data._async.client import TableAsync - client = BigtableDataClient() + client = BigtableDataClientAsync() with pytest.raises(ValueError) as e: - Table(client, "", "", default_per_request_timeout=-1) + TableAsync(client, "", "", default_per_request_timeout=-1) assert "default_per_request_timeout must be greater than 0" in str(e.value) with pytest.raises(ValueError) as e: - Table(client, "", "", default_operation_timeout=-1) + TableAsync(client, "", "", default_operation_timeout=-1) assert "default_operation_timeout must be greater than 0" in str(e.value) with pytest.raises(ValueError) as e: - Table( + TableAsync( client, "", "", @@ -1034,12 +1034,12 @@ async def test_table_ctor_bad_timeout_values(self): def test_table_ctor_sync(self): # initializing client in a sync context should raise RuntimeError - from google.cloud.bigtable.client import Table + from google.cloud.bigtable.data._async.client import TableAsync client = mock.Mock() with pytest.raises(RuntimeError) as e: - Table(client, "instance-id", "table-id") - assert e.match("Table must be created within an async event loop context.") + TableAsync(client, "instance-id", "table-id") + assert e.match("TableAsync must be created within an async event loop context.") class TestReadRows: @@ -1048,12 +1048,12 @@ class TestReadRows: """ def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - return BigtableDataClient(*args, **kwargs) + return BigtableDataClientAsync(*args, **kwargs) def _make_table(self, *args, **kwargs): - from google.cloud.bigtable.client import Table + from google.cloud.bigtable.data._async.client import TableAsync client_mock = mock.Mock() client_mock._register_instance.side_effect = ( @@ -1070,7 +1070,7 @@ def _make_table(self, *args, **kwargs): ) client_mock._gapic_client.table_path.return_value = kwargs["table_id"] client_mock._gapic_client.instance_path.return_value = kwargs["instance_id"] - return Table(client_mock, *args, **kwargs) + return TableAsync(client_mock, *args, **kwargs) def _make_stats(self): from google.cloud.bigtable_v2.types import RequestStats @@ -1174,7 +1174,7 @@ async def test_read_rows_stream(self): @pytest.mark.parametrize("include_app_profile", [True, False]) @pytest.mark.asyncio async def test_read_rows_query_matches_request(self, include_app_profile): - from google.cloud.bigtable import RowRange + from google.cloud.bigtable.data import RowRange app_profile_id = "app_profile_id" if include_app_profile else None async with self._make_table(app_profile_id=app_profile_id) as table: @@ -1250,7 +1250,7 @@ async def test_read_rows_per_request_timeout( operation_timeout does not cancel the request, so we expect the number of requests to be the ceiling of operation_timeout / per_request_timeout. """ - from google.cloud.bigtable.exceptions import RetryExceptionGroup + from google.cloud.bigtable.data.exceptions import RetryExceptionGroup expected_last_timeout = operation_t - (expected_num - 1) * per_request_t @@ -1295,12 +1295,12 @@ async def test_read_rows_per_request_timeout( @pytest.mark.asyncio async def test_read_rows_idle_timeout(self): - from google.cloud.bigtable.client import ReadRowsIterator + from google.cloud.bigtable.data._async.client import ReadRowsAsyncIterator from google.cloud.bigtable_v2.services.bigtable.async_client import ( BigtableAsyncClient, ) - from google.cloud.bigtable.exceptions import IdleTimeout - from google.cloud.bigtable._read_rows import _ReadRowsOperation + from google.cloud.bigtable.data.exceptions import IdleTimeout + from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync chunks = [ self._make_chunk(row_key=b"test_1"), @@ -1311,7 +1311,7 @@ async def test_read_rows_idle_timeout(self): chunks ) with mock.patch.object( - ReadRowsIterator, "_start_idle_timer" + ReadRowsAsyncIterator, "_start_idle_timer" ) as start_idle_timer: client = self._make_client() table = client.get_table("instance", "table") @@ -1319,7 +1319,9 @@ async def test_read_rows_idle_timeout(self): gen = await table.read_rows_stream(query) # should start idle timer on creation start_idle_timer.assert_called_once() - with mock.patch.object(_ReadRowsOperation, "aclose", AsyncMock()) as aclose: + with mock.patch.object( + _ReadRowsOperationAsync, "aclose", AsyncMock() + ) as aclose: # start idle timer with our own value await gen._start_idle_timer(0.1) # should timeout after being abandoned @@ -1398,13 +1400,13 @@ async def test_read_rows_revise_request(self): """ Ensure that _revise_request is called between retries """ - from google.cloud.bigtable._read_rows import _ReadRowsOperation - from google.cloud.bigtable.exceptions import InvalidChunk + from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync + from google.cloud.bigtable.data.exceptions import InvalidChunk with mock.patch.object( - _ReadRowsOperation, "_revise_request_rowset" + _ReadRowsOperationAsync, "_revise_request_rowset" ) as revise_rowset: - with mock.patch.object(_ReadRowsOperation, "aclose"): + with mock.patch.object(_ReadRowsOperationAsync, "aclose"): revise_rowset.return_value = "modified" async with self._make_table() as table: read_rows = table.client._gapic_client.read_rows @@ -1432,11 +1434,11 @@ async def test_read_rows_default_timeouts(self): """ Ensure that the default timeouts are set on the read rows operation when not overridden """ - from google.cloud.bigtable._read_rows import _ReadRowsOperation + from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync operation_timeout = 8 per_request_timeout = 4 - with mock.patch.object(_ReadRowsOperation, "__init__") as mock_op: + with mock.patch.object(_ReadRowsOperationAsync, "__init__") as mock_op: mock_op.side_effect = RuntimeError("mock error") async with self._make_table( default_operation_timeout=operation_timeout, @@ -1455,11 +1457,11 @@ async def test_read_rows_default_timeout_override(self): """ When timeouts are passed, they overwrite default values """ - from google.cloud.bigtable._read_rows import _ReadRowsOperation + from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync operation_timeout = 8 per_request_timeout = 4 - with mock.patch.object(_ReadRowsOperation, "__init__") as mock_op: + with mock.patch.object(_ReadRowsOperationAsync, "__init__") as mock_op: mock_op.side_effect = RuntimeError("mock error") async with self._make_table( default_operation_timeout=99, default_per_request_timeout=97 @@ -1653,9 +1655,9 @@ async def test_read_rows_metadata(self, include_app_profile): class TestReadRowsSharded: def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - return BigtableDataClient(*args, **kwargs) + return BigtableDataClientAsync(*args, **kwargs) @pytest.mark.asyncio async def test_read_rows_sharded_empty_query(self): @@ -1708,8 +1710,8 @@ async def test_read_rows_sharded_errors(self): """ Errors should be exposed as ShardedReadRowsExceptionGroups """ - from google.cloud.bigtable.exceptions import ShardedReadRowsExceptionGroup - from google.cloud.bigtable.exceptions import FailedQueryShardError + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.data.exceptions import FailedQueryShardError async with self._make_client() as client: async with client.get_table("instance", "table") as table: @@ -1785,8 +1787,8 @@ async def test_read_rows_sharded_batching(self): Large queries should be processed in batches to limit concurrency operation timeout should change between batches """ - from google.cloud.bigtable.client import Table - from google.cloud.bigtable.client import CONCURRENCY_LIMIT + from google.cloud.bigtable.data._async.client import TableAsync + from google.cloud.bigtable.data._async.client import CONCURRENCY_LIMIT assert CONCURRENCY_LIMIT == 10 # change this test if this changes @@ -1802,7 +1804,7 @@ async def test_read_rows_sharded_batching(self): # clock ticks one second on each check with mock.patch("time.monotonic", side_effect=range(0, 100000)): with mock.patch("asyncio.gather", AsyncMock()) as gather_mock: - await Table.read_rows_sharded(table_mock, query_list) + await TableAsync.read_rows_sharded(table_mock, query_list) # should have individual calls for each query assert table_mock.read_rows.call_count == n_queries # should have single gather call for each batch @@ -1843,9 +1845,9 @@ async def test_read_rows_sharded_batching(self): class TestSampleRowKeys: def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - return BigtableDataClient(*args, **kwargs) + return BigtableDataClientAsync(*args, **kwargs) async def _make_gapic_stream(self, sample_list: list[tuple[bytes, int]]): from google.cloud.bigtable_v2.types import SampleRowKeysResponse @@ -1980,7 +1982,7 @@ async def test_sample_row_keys_retryable_errors(self, retryable_exception): retryable errors should be retried until timeout """ from google.api_core.exceptions import DeadlineExceeded - from google.cloud.bigtable.exceptions import RetryExceptionGroup + from google.cloud.bigtable.data.exceptions import RetryExceptionGroup async with self._make_client() as client: async with client.get_table("instance", "table") as table: @@ -2023,9 +2025,9 @@ async def test_sample_row_keys_non_retryable_errors(self, non_retryable_exceptio class TestMutateRow: def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - return BigtableDataClient(*args, **kwargs) + return BigtableDataClientAsync(*args, **kwargs) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -2085,7 +2087,7 @@ async def test_mutate_row(self, mutation_arg): @pytest.mark.asyncio async def test_mutate_row_retryable_errors(self, retryable_exception): from google.api_core.exceptions import DeadlineExceeded - from google.cloud.bigtable.exceptions import RetryExceptionGroup + from google.cloud.bigtable.data.exceptions import RetryExceptionGroup async with self._make_client(project="project") as client: async with client.get_table("instance", "table") as table: @@ -2190,9 +2192,9 @@ async def test_mutate_row_metadata(self, include_app_profile): class TestBulkMutateRows: def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - return BigtableDataClient(*args, **kwargs) + return BigtableDataClientAsync(*args, **kwargs) async def _mock_response(self, response_list): from google.cloud.bigtable_v2.types import MutateRowsResponse @@ -2300,7 +2302,7 @@ async def test_bulk_mutate_rows_idempotent_mutation_error_retryable( """ Individual idempotent mutations should be retried if they fail with a retryable error """ - from google.cloud.bigtable.exceptions import ( + from google.cloud.bigtable.data.exceptions import ( RetryExceptionGroup, FailedMutationEntryError, MutationsExceptionGroup, @@ -2347,7 +2349,7 @@ async def test_bulk_mutate_rows_idempotent_mutation_error_non_retryable( """ Individual idempotent mutations should not be retried if they fail with a non-retryable error """ - from google.cloud.bigtable.exceptions import ( + from google.cloud.bigtable.data.exceptions import ( FailedMutationEntryError, MutationsExceptionGroup, ) @@ -2386,7 +2388,7 @@ async def test_bulk_mutate_idempotent_retryable_request_errors( """ Individual idempotent mutations should be retried if the request fails with a retryable error """ - from google.cloud.bigtable.exceptions import ( + from google.cloud.bigtable.data.exceptions import ( RetryExceptionGroup, FailedMutationEntryError, MutationsExceptionGroup, @@ -2425,7 +2427,7 @@ async def test_bulk_mutate_rows_non_idempotent_retryable_errors( self, retryable_exception ): """Non-Idempotent mutations should never be retried""" - from google.cloud.bigtable.exceptions import ( + from google.cloud.bigtable.data.exceptions import ( FailedMutationEntryError, MutationsExceptionGroup, ) @@ -2467,7 +2469,7 @@ async def test_bulk_mutate_rows_non_retryable_errors(self, non_retryable_excepti """ If the request fails with a non-retryable error, mutations should not be retried """ - from google.cloud.bigtable.exceptions import ( + from google.cloud.bigtable.data.exceptions import ( FailedMutationEntryError, MutationsExceptionGroup, ) @@ -2502,7 +2504,7 @@ async def test_bulk_mutate_error_index(self): ServiceUnavailable, FailedPrecondition, ) - from google.cloud.bigtable.exceptions import ( + from google.cloud.bigtable.data.exceptions import ( RetryExceptionGroup, FailedMutationEntryError, MutationsExceptionGroup, @@ -2579,9 +2581,9 @@ async def test_bulk_mutate_row_metadata(self, include_app_profile): class TestCheckAndMutateRow: def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - return BigtableDataClient(*args, **kwargs) + return BigtableDataClientAsync(*args, **kwargs) @pytest.mark.parametrize("gapic_result", [True, False]) @pytest.mark.asyncio @@ -2660,7 +2662,7 @@ async def test_check_and_mutate_no_mutations(self): @pytest.mark.asyncio async def test_check_and_mutate_single_mutations(self): """if single mutations are passed, they should be internally wrapped in a list""" - from google.cloud.bigtable.mutations import SetCell + from google.cloud.bigtable.data.mutations import SetCell from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse async with self._make_client() as client: @@ -2713,7 +2715,7 @@ async def test_check_and_mutate_predicate_object(self): async def test_check_and_mutate_mutations_parsing(self): """mutations objects should be converted to dicts""" from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse - from google.cloud.bigtable.mutations import DeleteAllFromRow + from google.cloud.bigtable.data.mutations import DeleteAllFromRow mutations = [mock.Mock() for _ in range(5)] for idx, mutation in enumerate(mutations): @@ -2772,9 +2774,9 @@ async def test_check_and_mutate_metadata(self, include_app_profile): class TestReadModifyWriteRow: def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.client import BigtableDataClient + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - return BigtableDataClient(*args, **kwargs) + return BigtableDataClientAsync(*args, **kwargs) @pytest.mark.parametrize( "call_rules,expected_rules", @@ -2886,7 +2888,7 @@ async def test_read_modify_write_row_building(self): """ results from gapic call should be used to construct row """ - from google.cloud.bigtable.row import Row + from google.cloud.bigtable.data.row import Row from google.cloud.bigtable_v2.types import ReadModifyWriteRowResponse from google.cloud.bigtable_v2.types import Row as RowPB diff --git a/tests/unit/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py similarity index 94% rename from tests/unit/test_mutations_batcher.py rename to tests/unit/data/_async/test_mutations_batcher.py index a900468d5..1b14cc128 100644 --- a/tests/unit/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -33,9 +33,11 @@ def _make_mutation(count=1, size=1): class Test_FlowControl: def _make_one(self, max_mutation_count=10, max_mutation_bytes=100): - from google.cloud.bigtable.mutations_batcher import _FlowControl + from google.cloud.bigtable.data._async.mutations_batcher import ( + _FlowControlAsync, + ) - return _FlowControl(max_mutation_count, max_mutation_bytes) + return _FlowControlAsync(max_mutation_count, max_mutation_bytes) def test_ctor(self): max_mutation_count = 9 @@ -238,7 +240,7 @@ async def test_add_to_flow_max_mutation_limits( Should submit request early, even if the flow control has room for more """ with mock.patch( - "google.cloud.bigtable.mutations_batcher.MUTATE_ROWS_REQUEST_MUTATION_LIMIT", + "google.cloud.bigtable.data._async.mutations_batcher.MUTATE_ROWS_REQUEST_MUTATION_LIMIT", max_limit, ): mutation_objs = [_make_mutation(count=m[0], size=m[1]) for m in mutations] @@ -275,11 +277,13 @@ async def test_add_to_flow_oversize(self): assert len(count_results) == 1 -class TestMutationsBatcher: +class TestMutationsBatcherAsync: def _get_target_class(self): - from google.cloud.bigtable.mutations_batcher import MutationsBatcher + from google.cloud.bigtable.data._async.mutations_batcher import ( + MutationsBatcherAsync, + ) - return MutationsBatcher + return MutationsBatcherAsync def _make_one(self, table=None, **kwargs): if table is None: @@ -290,7 +294,7 @@ def _make_one(self, table=None, **kwargs): return self._get_target_class()(table, **kwargs) @mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._start_flush_timer" + "google.cloud.bigtable.data._async.mutations_batcher.MutationsBatcherAsync._start_flush_timer" ) @pytest.mark.asyncio async def test_ctor_defaults(self, flush_timer_mock): @@ -320,7 +324,7 @@ async def test_ctor_defaults(self, flush_timer_mock): assert isinstance(instance._flush_timer, asyncio.Future) @mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._start_flush_timer", + "google.cloud.bigtable.data._async.mutations_batcher.MutationsBatcherAsync._start_flush_timer", ) @pytest.mark.asyncio async def test_ctor_explicit(self, flush_timer_mock): @@ -368,7 +372,7 @@ async def test_ctor_explicit(self, flush_timer_mock): assert isinstance(instance._flush_timer, asyncio.Future) @mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._start_flush_timer" + "google.cloud.bigtable.data._async.mutations_batcher.MutationsBatcherAsync._start_flush_timer" ) @pytest.mark.asyncio async def test_ctor_no_flush_limits(self, flush_timer_mock): @@ -419,19 +423,23 @@ async def test_ctor_invalid_values(self): def test_default_argument_consistency(self): """ - We supply default arguments in MutationsBatcher.__init__, and in + We supply default arguments in MutationsBatcherAsync.__init__, and in table.mutations_batcher. Make sure any changes to defaults are applied to both places """ - from google.cloud.bigtable.client import Table - from google.cloud.bigtable.mutations_batcher import MutationsBatcher + from google.cloud.bigtable.data._async.client import TableAsync + from google.cloud.bigtable.data._async.mutations_batcher import ( + MutationsBatcherAsync, + ) import inspect get_batcher_signature = dict( - inspect.signature(Table.mutations_batcher).parameters + inspect.signature(TableAsync.mutations_batcher).parameters ) get_batcher_signature.pop("self") - batcher_init_signature = dict(inspect.signature(MutationsBatcher).parameters) + batcher_init_signature = dict( + inspect.signature(MutationsBatcherAsync).parameters + ) batcher_init_signature.pop("table") # both should have same number of arguments assert len(get_batcher_signature.keys()) == len(batcher_init_signature.keys()) @@ -446,7 +454,7 @@ def test_default_argument_consistency(self): ) @mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + "google.cloud.bigtable.data._async.mutations_batcher.MutationsBatcherAsync._schedule_flush" ) @pytest.mark.asyncio async def test__start_flush_timer_w_None(self, flush_mock): @@ -458,7 +466,7 @@ async def test__start_flush_timer_w_None(self, flush_mock): assert flush_mock.call_count == 0 @mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + "google.cloud.bigtable.data._async.mutations_batcher.MutationsBatcherAsync._schedule_flush" ) @pytest.mark.asyncio async def test__start_flush_timer_call_when_closed(self, flush_mock): @@ -472,7 +480,7 @@ async def test__start_flush_timer_call_when_closed(self, flush_mock): assert flush_mock.call_count == 0 @mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + "google.cloud.bigtable.data._async.mutations_batcher.MutationsBatcherAsync._schedule_flush" ) @pytest.mark.asyncio async def test__flush_timer(self, flush_mock): @@ -492,7 +500,7 @@ async def test__flush_timer(self, flush_mock): assert flush_mock.call_count == loop_num @mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + "google.cloud.bigtable.data._async.mutations_batcher.MutationsBatcherAsync._schedule_flush" ) @pytest.mark.asyncio async def test__flush_timer_no_mutations(self, flush_mock): @@ -511,7 +519,7 @@ async def test__flush_timer_no_mutations(self, flush_mock): assert flush_mock.call_count == 0 @mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._schedule_flush" + "google.cloud.bigtable.data._async.mutations_batcher.MutationsBatcherAsync._schedule_flush" ) @pytest.mark.asyncio async def test__flush_timer_close(self, flush_mock): @@ -541,7 +549,7 @@ async def test_append_wrong_mutation(self): Mutation objects should raise an exception. Only support RowMutationEntry """ - from google.cloud.bigtable.mutations import DeleteAllFromRow + from google.cloud.bigtable.data.mutations import DeleteAllFromRow async with self._make_one() as instance: expected_error = "invalid mutation type: DeleteAllFromRow. Only RowMutationEntry objects are supported by batcher" @@ -577,9 +585,13 @@ async def test_append_flush_runs_after_limit_hit(self): If the user appends a bunch of entries above the flush limits back-to-back, it should still flush in a single task """ - from google.cloud.bigtable.mutations_batcher import MutationsBatcher + from google.cloud.bigtable.data._async.mutations_batcher import ( + MutationsBatcherAsync, + ) - with mock.patch.object(MutationsBatcher, "_execute_mutate_rows") as op_mock: + with mock.patch.object( + MutationsBatcherAsync, "_execute_mutate_rows" + ) as op_mock: async with self._make_one(flush_limit_bytes=100) as instance: # mock network calls async def mock_call(*args, **kwargs): @@ -789,7 +801,7 @@ async def test__flush_internal_with_errors( """ errors returned from _execute_mutate_rows should be added to internal exceptions """ - from google.cloud.bigtable import exceptions + from google.cloud.bigtable.data import exceptions num_entries = 10 expected_errors = [ @@ -861,7 +873,7 @@ async def test_timer_flush_end_to_end(self): @pytest.mark.asyncio @mock.patch( - "google.cloud.bigtable.mutations_batcher._MutateRowsOperation", + "google.cloud.bigtable.data._async.mutations_batcher._MutateRowsOperationAsync", ) async def test__execute_mutate_rows(self, mutate_rows): mutate_rows.return_value = AsyncMock() @@ -884,10 +896,12 @@ async def test__execute_mutate_rows(self, mutate_rows): assert result == [] @pytest.mark.asyncio - @mock.patch("google.cloud.bigtable.mutations_batcher._MutateRowsOperation.start") + @mock.patch( + "google.cloud.bigtable.data._async.mutations_batcher._MutateRowsOperationAsync.start" + ) async def test__execute_mutate_rows_returns_errors(self, mutate_rows): """Errors from operation should be retruned as list""" - from google.cloud.bigtable.exceptions import ( + from google.cloud.bigtable.data.exceptions import ( MutationsExceptionGroup, FailedMutationEntryError, ) @@ -911,7 +925,7 @@ async def test__execute_mutate_rows_returns_errors(self, mutate_rows): @pytest.mark.asyncio async def test__raise_exceptions(self): """Raise exceptions and reset error state""" - from google.cloud.bigtable import exceptions + from google.cloud.bigtable.data import exceptions expected_total = 1201 expected_exceptions = [RuntimeError("mock")] * 3 @@ -958,7 +972,7 @@ async def test_close(self): @pytest.mark.asyncio async def test_close_w_exceptions(self): """Raise exceptions on close""" - from google.cloud.bigtable import exceptions + from google.cloud.bigtable.data import exceptions expected_total = 10 expected_exceptions = [RuntimeError("mock")] @@ -1001,20 +1015,14 @@ async def test_atexit_registration(self): """Should run _on_exit on program termination""" import atexit - with mock.patch( - "google.cloud.bigtable.mutations_batcher.MutationsBatcher._on_exit" - ) as on_exit_mock: + with mock.patch.object(atexit, "register") as register_mock: + assert register_mock.call_count == 0 async with self._make_one(): - assert on_exit_mock.call_count == 0 - atexit._run_exitfuncs() - assert on_exit_mock.call_count == 1 - # should not call after close - atexit._run_exitfuncs() - assert on_exit_mock.call_count == 1 + assert register_mock.call_count == 1 @pytest.mark.asyncio @mock.patch( - "google.cloud.bigtable.mutations_batcher._MutateRowsOperation", + "google.cloud.bigtable.data._async.mutations_batcher._MutateRowsOperationAsync", ) async def test_timeout_args_passed(self, mutate_rows): """ diff --git a/tests/unit/read-rows-acceptance-test.json b/tests/unit/data/read-rows-acceptance-test.json similarity index 100% rename from tests/unit/read-rows-acceptance-test.json rename to tests/unit/data/read-rows-acceptance-test.json diff --git a/tests/unit/test__helpers.py b/tests/unit/data/test__helpers.py similarity index 97% rename from tests/unit/test__helpers.py rename to tests/unit/data/test__helpers.py index 9aa1a7bb4..dc688bb0c 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/data/test__helpers.py @@ -13,8 +13,8 @@ # import pytest -import google.cloud.bigtable._helpers as _helpers -import google.cloud.bigtable.exceptions as bigtable_exceptions +import google.cloud.bigtable.data._helpers as _helpers +import google.cloud.bigtable.data.exceptions as bigtable_exceptions import mock diff --git a/tests/unit/test__read_rows.py b/tests/unit/data/test__read_rows_state_machine.py similarity index 62% rename from tests/unit/test__read_rows.py rename to tests/unit/data/test__read_rows_state_machine.py index c893c56cd..0d1ee6b06 100644 --- a/tests/unit/test__read_rows.py +++ b/tests/unit/data/test__read_rows_state_machine.py @@ -1,10 +1,22 @@ +# Licensed 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. + import unittest import pytest -from google.cloud.bigtable.exceptions import InvalidChunk -from google.cloud.bigtable._read_rows import AWAITING_NEW_ROW -from google.cloud.bigtable._read_rows import AWAITING_NEW_CELL -from google.cloud.bigtable._read_rows import AWAITING_CELL_VALUE +from google.cloud.bigtable.data.exceptions import InvalidChunk +from google.cloud.bigtable.data._read_rows_state_machine import AWAITING_NEW_ROW +from google.cloud.bigtable.data._read_rows_state_machine import AWAITING_NEW_CELL +from google.cloud.bigtable.data._read_rows_state_machine import AWAITING_CELL_VALUE # try/except added for compatibility with python < 3.8 try: @@ -20,377 +32,10 @@ TEST_LABELS = ["label1", "label2"] -class TestReadRowsOperation: - """ - Tests helper functions in the ReadRowsOperation class - in-depth merging logic in merge_row_response_stream and _read_rows_retryable_attempt - is tested in test_read_rows_acceptance test_client_read_rows, and conformance tests - """ - - @staticmethod - def _get_target_class(): - from google.cloud.bigtable._read_rows import _ReadRowsOperation - - return _ReadRowsOperation - - def _make_one(self, *args, **kwargs): - return self._get_target_class()(*args, **kwargs) - - def test_ctor_defaults(self): - request = {} - client = mock.Mock() - client.read_rows = mock.Mock() - client.read_rows.return_value = None - default_operation_timeout = 600 - time_gen_mock = mock.Mock() - with mock.patch( - "google.cloud.bigtable._read_rows._attempt_timeout_generator", time_gen_mock - ): - instance = self._make_one(request, client) - assert time_gen_mock.call_count == 1 - time_gen_mock.assert_called_once_with(None, default_operation_timeout) - assert instance.transient_errors == [] - assert instance._last_emitted_row_key is None - assert instance._emit_count == 0 - assert instance.operation_timeout == default_operation_timeout - retryable_fn = instance._partial_retryable - assert retryable_fn.func == instance._read_rows_retryable_attempt - assert retryable_fn.args[0] == client.read_rows - assert retryable_fn.args[1] == time_gen_mock.return_value - assert retryable_fn.args[2] == 0 - assert client.read_rows.call_count == 0 - - def test_ctor(self): - row_limit = 91 - request = {"rows_limit": row_limit} - client = mock.Mock() - client.read_rows = mock.Mock() - client.read_rows.return_value = None - expected_operation_timeout = 42 - expected_request_timeout = 44 - time_gen_mock = mock.Mock() - with mock.patch( - "google.cloud.bigtable._read_rows._attempt_timeout_generator", time_gen_mock - ): - instance = self._make_one( - request, - client, - operation_timeout=expected_operation_timeout, - per_request_timeout=expected_request_timeout, - ) - assert time_gen_mock.call_count == 1 - time_gen_mock.assert_called_once_with( - expected_request_timeout, expected_operation_timeout - ) - assert instance.transient_errors == [] - assert instance._last_emitted_row_key is None - assert instance._emit_count == 0 - assert instance.operation_timeout == expected_operation_timeout - retryable_fn = instance._partial_retryable - assert retryable_fn.func == instance._read_rows_retryable_attempt - assert retryable_fn.args[0] == client.read_rows - assert retryable_fn.args[1] == time_gen_mock.return_value - assert retryable_fn.args[2] == row_limit - assert client.read_rows.call_count == 0 - - def test___aiter__(self): - request = {} - client = mock.Mock() - client.read_rows = mock.Mock() - instance = self._make_one(request, client) - assert instance.__aiter__() is instance - - @pytest.mark.asyncio - async def test_transient_error_capture(self): - from google.api_core import exceptions as core_exceptions - - client = mock.Mock() - client.read_rows = mock.Mock() - test_exc = core_exceptions.Aborted("test") - test_exc2 = core_exceptions.DeadlineExceeded("test") - client.read_rows.side_effect = [test_exc, test_exc2] - instance = self._make_one({}, client) - with pytest.raises(RuntimeError): - await instance.__anext__() - assert len(instance.transient_errors) == 2 - assert instance.transient_errors[0] == test_exc - assert instance.transient_errors[1] == test_exc2 - - @pytest.mark.parametrize( - "in_keys,last_key,expected", - [ - (["b", "c", "d"], "a", ["b", "c", "d"]), - (["a", "b", "c"], "b", ["c"]), - (["a", "b", "c"], "c", []), - (["a", "b", "c"], "d", []), - (["d", "c", "b", "a"], "b", ["d", "c"]), - ], - ) - def test_revise_request_rowset_keys(self, in_keys, last_key, expected): - sample_range = {"start_key_open": last_key} - row_set = {"row_keys": in_keys, "row_ranges": [sample_range]} - revised = self._get_target_class()._revise_request_rowset(row_set, last_key) - assert revised["row_keys"] == expected - assert revised["row_ranges"] == [sample_range] - - @pytest.mark.parametrize( - "in_ranges,last_key,expected", - [ - ( - [{"start_key_open": "b", "end_key_closed": "d"}], - "a", - [{"start_key_open": "b", "end_key_closed": "d"}], - ), - ( - [{"start_key_closed": "b", "end_key_closed": "d"}], - "a", - [{"start_key_closed": "b", "end_key_closed": "d"}], - ), - ( - [{"start_key_open": "a", "end_key_closed": "d"}], - "b", - [{"start_key_open": "b", "end_key_closed": "d"}], - ), - ( - [{"start_key_closed": "a", "end_key_open": "d"}], - "b", - [{"start_key_open": "b", "end_key_open": "d"}], - ), - ( - [{"start_key_closed": "b", "end_key_closed": "d"}], - "b", - [{"start_key_open": "b", "end_key_closed": "d"}], - ), - ([{"start_key_closed": "b", "end_key_closed": "d"}], "d", []), - ([{"start_key_closed": "b", "end_key_open": "d"}], "d", []), - ([{"start_key_closed": "b", "end_key_closed": "d"}], "e", []), - ([{"start_key_closed": "b"}], "z", [{"start_key_open": "z"}]), - ([{"start_key_closed": "b"}], "a", [{"start_key_closed": "b"}]), - ( - [{"end_key_closed": "z"}], - "a", - [{"start_key_open": "a", "end_key_closed": "z"}], - ), - ( - [{"end_key_open": "z"}], - "a", - [{"start_key_open": "a", "end_key_open": "z"}], - ), - ], - ) - def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected): - next_key = last_key + "a" - row_set = {"row_keys": [next_key], "row_ranges": in_ranges} - revised = self._get_target_class()._revise_request_rowset(row_set, last_key) - assert revised["row_keys"] == [next_key] - assert revised["row_ranges"] == expected - - @pytest.mark.parametrize("last_key", ["a", "b", "c"]) - def test_revise_request_full_table(self, last_key): - row_set = {"row_keys": [], "row_ranges": []} - for selected_set in [row_set, None]: - revised = self._get_target_class()._revise_request_rowset( - selected_set, last_key - ) - assert revised["row_keys"] == [] - assert len(revised["row_ranges"]) == 1 - assert revised["row_ranges"][0]["start_key_open"] == last_key - - def test_revise_to_empty_rowset(self): - """revising to an empty rowset should raise error""" - from google.cloud.bigtable.exceptions import _RowSetComplete - - row_keys = ["a", "b", "c"] - row_set = {"row_keys": row_keys, "row_ranges": [{"end_key_open": "c"}]} - with pytest.raises(_RowSetComplete): - self._get_target_class()._revise_request_rowset(row_set, "d") - - @pytest.mark.parametrize( - "start_limit,emit_num,expected_limit", - [ - (10, 0, 10), - (10, 1, 9), - (10, 10, 0), - (0, 10, 0), - (0, 0, 0), - (4, 2, 2), - ], - ) - @pytest.mark.asyncio - async def test_revise_limit(self, start_limit, emit_num, expected_limit): - """ - revise_limit should revise the request's limit field - - if limit is 0 (unlimited), it should never be revised - - if start_limit-emit_num == 0, the request should end early - - if the number emitted exceeds the new limit, an exception should - should be raised (tested in test_revise_limit_over_limit) - """ - import itertools - - request = {"rows_limit": start_limit} - instance = self._make_one(request, mock.Mock()) - instance._emit_count = emit_num - instance._last_emitted_row_key = "a" - gapic_mock = mock.Mock() - gapic_mock.side_effect = [GeneratorExit("stop_fn")] - mock_timeout_gen = itertools.repeat(5) - - attempt = instance._read_rows_retryable_attempt( - gapic_mock, mock_timeout_gen, start_limit - ) - if start_limit != 0 and expected_limit == 0: - # if we emitted the expected number of rows, we should receive a StopAsyncIteration - with pytest.raises(StopAsyncIteration): - await attempt.__anext__() - else: - with pytest.raises(GeneratorExit): - await attempt.__anext__() - assert request["rows_limit"] == expected_limit - - @pytest.mark.parametrize("start_limit,emit_num", [(5, 10), (3, 9), (1, 10)]) - @pytest.mark.asyncio - async def test_revise_limit_over_limit(self, start_limit, emit_num): - """ - Should raise runtime error if we get in state where emit_num > start_num - (unless start_num == 0, which represents unlimited) - """ - import itertools - - request = {"rows_limit": start_limit} - instance = self._make_one(request, mock.Mock()) - instance._emit_count = emit_num - instance._last_emitted_row_key = "a" - mock_timeout_gen = itertools.repeat(5) - attempt = instance._read_rows_retryable_attempt( - mock.Mock(), mock_timeout_gen, start_limit - ) - with pytest.raises(RuntimeError) as e: - await attempt.__anext__() - assert "emit count exceeds row limit" in str(e.value) - - @pytest.mark.asyncio - async def test_aclose(self): - import asyncio - - instance = self._make_one({}, mock.Mock()) - await instance.aclose() - assert instance._stream is None - assert instance._last_emitted_row_key is None - with pytest.raises(asyncio.InvalidStateError): - await instance.__anext__() - # try calling a second time - await instance.aclose() - - @pytest.mark.parametrize("limit", [1, 3, 10]) - @pytest.mark.asyncio - async def test_retryable_attempt_hit_limit(self, limit): - """ - Stream should end after hitting the limit - """ - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - import itertools - - instance = self._make_one({}, mock.Mock()) - - async def mock_gapic(*args, **kwargs): - # continuously return a single row - async def gen(): - for i in range(limit * 2): - chunk = ReadRowsResponse.CellChunk( - row_key=str(i).encode(), - family_name="family_name", - qualifier=b"qualifier", - commit_row=True, - ) - yield ReadRowsResponse(chunks=[chunk]) - - return gen() - - mock_timeout_gen = itertools.repeat(5) - gen = instance._read_rows_retryable_attempt(mock_gapic, mock_timeout_gen, limit) - # should yield values up to the limit - for i in range(limit): - await gen.__anext__() - # next value should be StopAsyncIteration - with pytest.raises(StopAsyncIteration): - await gen.__anext__() - - @pytest.mark.asyncio - async def test_retryable_ignore_repeated_rows(self): - """ - Duplicate rows should cause an invalid chunk error - """ - from google.cloud.bigtable._read_rows import _ReadRowsOperation - from google.cloud.bigtable.row import Row - from google.cloud.bigtable.exceptions import InvalidChunk - - async def mock_stream(): - while True: - yield Row(b"dup_key", cells=[]) - yield Row(b"dup_key", cells=[]) - - with mock.patch.object( - _ReadRowsOperation, "merge_row_response_stream" - ) as mock_stream_fn: - mock_stream_fn.return_value = mock_stream() - instance = self._make_one({}, mock.AsyncMock()) - first_row = await instance.__anext__() - assert first_row.row_key == b"dup_key" - with pytest.raises(InvalidChunk) as exc: - await instance.__anext__() - assert "Last emitted row key out of order" in str(exc.value) - - @pytest.mark.asyncio - async def test_retryable_ignore_last_scanned_rows(self): - """ - Last scanned rows should not be emitted - """ - from google.cloud.bigtable._read_rows import _ReadRowsOperation - from google.cloud.bigtable.row import Row, _LastScannedRow - - async def mock_stream(): - while True: - yield Row(b"key1", cells=[]) - yield _LastScannedRow(b"key2_ignored") - yield Row(b"key3", cells=[]) - - with mock.patch.object( - _ReadRowsOperation, "merge_row_response_stream" - ) as mock_stream_fn: - mock_stream_fn.return_value = mock_stream() - instance = self._make_one({}, mock.AsyncMock()) - first_row = await instance.__anext__() - assert first_row.row_key == b"key1" - second_row = await instance.__anext__() - assert second_row.row_key == b"key3" - - @pytest.mark.asyncio - async def test_retryable_cancel_on_close(self): - """Underlying gapic call should be cancelled when stream is closed""" - from google.cloud.bigtable._read_rows import _ReadRowsOperation - from google.cloud.bigtable.row import Row - - async def mock_stream(): - while True: - yield Row(b"key1", cells=[]) - - with mock.patch.object( - _ReadRowsOperation, "merge_row_response_stream" - ) as mock_stream_fn: - mock_stream_fn.return_value = mock_stream() - mock_gapic = mock.AsyncMock() - mock_call = await mock_gapic.read_rows() - instance = self._make_one({}, mock_gapic) - await instance.__anext__() - assert mock_call.cancel.call_count == 0 - await instance.aclose() - assert mock_call.cancel.call_count == 1 - - class TestStateMachine(unittest.TestCase): @staticmethod def _get_target_class(): - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._read_rows_state_machine import _StateMachine return _StateMachine @@ -398,7 +43,7 @@ def _make_one(self, *args, **kwargs): return self._get_target_class()(*args, **kwargs) def test_ctor(self): - from google.cloud.bigtable._read_rows import _RowBuilder + from google.cloud.bigtable.data._read_rows_state_machine import _RowBuilder instance = self._make_one() assert instance.last_seen_row_key is None @@ -435,7 +80,7 @@ def test__reset_row(self): assert instance.adapter.reset.call_count == 1 def test_handle_last_scanned_row_wrong_state(self): - from google.cloud.bigtable.exceptions import InvalidChunk + from google.cloud.bigtable.data.exceptions import InvalidChunk instance = self._make_one() instance.current_state = AWAITING_NEW_CELL @@ -448,7 +93,7 @@ def test_handle_last_scanned_row_wrong_state(self): assert e.value.args[0] == "Last scanned row key received in invalid state" def test_handle_last_scanned_row_out_of_order(self): - from google.cloud.bigtable.exceptions import InvalidChunk + from google.cloud.bigtable.data.exceptions import InvalidChunk instance = self._make_one() instance.last_seen_row_key = b"b" @@ -460,7 +105,7 @@ def test_handle_last_scanned_row_out_of_order(self): assert e.value.args[0] == "Last scanned row key is out of order" def test_handle_last_scanned_row(self): - from google.cloud.bigtable.row import _LastScannedRow + from google.cloud.bigtable.data.row import _LastScannedRow instance = self._make_one() instance.adapter = mock.Mock() @@ -475,7 +120,7 @@ def test_handle_last_scanned_row(self): assert instance.adapter.reset.call_count == 1 def test__handle_complete_row(self): - from google.cloud.bigtable.row import Row + from google.cloud.bigtable.data.row import Row instance = self._make_one() instance.current_state = mock.Mock() @@ -490,7 +135,7 @@ def test__handle_complete_row(self): assert instance.adapter.reset.call_count == 1 def test__handle_reset_chunk_errors(self): - from google.cloud.bigtable.exceptions import InvalidChunk + from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse instance = self._make_one() @@ -528,7 +173,7 @@ def test__handle_reset_chunk_errors(self): assert e.value.args[0] == "Reset chunk has labels" def test_handle_chunk_out_of_order(self): - from google.cloud.bigtable.exceptions import InvalidChunk + from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse instance = self._make_one() @@ -570,7 +215,7 @@ def handle_chunk_with_commit_wrong_state(self, state): def test_handle_chunk_with_commit(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.row import Row + from google.cloud.bigtable.data.row import Row instance = self._make_one() with mock.patch.object(type(instance), "_reset_row") as mock_reset: @@ -587,7 +232,7 @@ def test_handle_chunk_with_commit(self): def test_handle_chunk_with_commit_empty_strings(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.row import Row + from google.cloud.bigtable.data.row import Row instance = self._make_one() with mock.patch.object(type(instance), "_reset_row") as mock_reset: @@ -648,7 +293,7 @@ def test_AWAITING_NEW_ROW(self): def test_AWAITING_NEW_CELL_family_without_qualifier(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() state_machine.current_qualifier = b"q" @@ -660,7 +305,7 @@ def test_AWAITING_NEW_CELL_family_without_qualifier(self): def test_AWAITING_NEW_CELL_qualifier_without_family(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() instance = AWAITING_NEW_CELL @@ -671,7 +316,7 @@ def test_AWAITING_NEW_CELL_qualifier_without_family(self): def test_AWAITING_NEW_CELL_no_row_state(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() instance = AWAITING_NEW_CELL @@ -687,7 +332,7 @@ def test_AWAITING_NEW_CELL_no_row_state(self): def test_AWAITING_NEW_CELL_invalid_row_key(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() instance = AWAITING_NEW_CELL @@ -699,7 +344,7 @@ def test_AWAITING_NEW_CELL_invalid_row_key(self): def test_AWAITING_NEW_CELL_success_no_split(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() state_machine.adapter = mock.Mock() @@ -733,7 +378,7 @@ def test_AWAITING_NEW_CELL_success_no_split(self): def test_AWAITING_NEW_CELL_success_with_split(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() state_machine.adapter = mock.Mock() @@ -768,7 +413,7 @@ def test_AWAITING_NEW_CELL_success_with_split(self): def test_AWAITING_CELL_VALUE_w_row_key(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() instance = AWAITING_CELL_VALUE @@ -779,7 +424,7 @@ def test_AWAITING_CELL_VALUE_w_row_key(self): def test_AWAITING_CELL_VALUE_w_family(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() instance = AWAITING_CELL_VALUE @@ -790,7 +435,7 @@ def test_AWAITING_CELL_VALUE_w_family(self): def test_AWAITING_CELL_VALUE_w_qualifier(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() instance = AWAITING_CELL_VALUE @@ -801,7 +446,7 @@ def test_AWAITING_CELL_VALUE_w_qualifier(self): def test_AWAITING_CELL_VALUE_w_timestamp(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() instance = AWAITING_CELL_VALUE @@ -812,7 +457,7 @@ def test_AWAITING_CELL_VALUE_w_timestamp(self): def test_AWAITING_CELL_VALUE_w_labels(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() instance = AWAITING_CELL_VALUE @@ -823,7 +468,7 @@ def test_AWAITING_CELL_VALUE_w_labels(self): def test_AWAITING_CELL_VALUE_continuation(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() state_machine.adapter = mock.Mock() @@ -838,7 +483,7 @@ def test_AWAITING_CELL_VALUE_continuation(self): def test_AWAITING_CELL_VALUE_final_chunk(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _StateMachine + from google.cloud.bigtable.data._async._read_rows import _StateMachine state_machine = _StateMachine() state_machine.adapter = mock.Mock() @@ -855,7 +500,7 @@ def test_AWAITING_CELL_VALUE_final_chunk(self): class TestRowBuilder(unittest.TestCase): @staticmethod def _get_target_class(): - from google.cloud.bigtable._read_rows import _RowBuilder + from google.cloud.bigtable.data._read_rows_state_machine import _RowBuilder return _RowBuilder @@ -1003,7 +648,7 @@ def test_reset(self): class TestChunkHasField: def test__chunk_has_field_empty(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _chunk_has_field + from google.cloud.bigtable.data._read_rows_state_machine import _chunk_has_field chunk = ReadRowsResponse.CellChunk()._pb assert not _chunk_has_field(chunk, "family_name") @@ -1011,7 +656,7 @@ def test__chunk_has_field_empty(self): def test__chunk_has_field_populated_empty_strings(self): from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable._read_rows import _chunk_has_field + from google.cloud.bigtable.data._read_rows_state_machine import _chunk_has_field chunk = ReadRowsResponse.CellChunk(qualifier=b"", family_name="")._pb assert _chunk_has_field(chunk, "family_name") diff --git a/tests/unit/test_exceptions.py b/tests/unit/data/test_exceptions.py similarity index 95% rename from tests/unit/test_exceptions.py rename to tests/unit/data/test_exceptions.py index ef186a47c..9d1145e36 100644 --- a/tests/unit/test_exceptions.py +++ b/tests/unit/data/test_exceptions.py @@ -16,7 +16,7 @@ import pytest import sys -import google.cloud.bigtable.exceptions as bigtable_exceptions +import google.cloud.bigtable.data.exceptions as bigtable_exceptions # try/except added for compatibility with python < 3.8 try: @@ -31,9 +31,9 @@ class TestBigtableExceptionGroup: """ def _get_class(self): - from google.cloud.bigtable.exceptions import BigtableExceptionGroup + from google.cloud.bigtable.data.exceptions import _BigtableExceptionGroup - return BigtableExceptionGroup + return _BigtableExceptionGroup def _make_one(self, message="test_message", excs=None): if excs is None: @@ -74,7 +74,7 @@ def test_311_traceback(self): exc_group = self._make_one(excs=[sub_exc1, sub_exc2]) expected_traceback = ( - f" | google.cloud.bigtable.exceptions.{type(exc_group).__name__}: {str(exc_group)}", + f" | google.cloud.bigtable.data.exceptions.{type(exc_group).__name__}: {str(exc_group)}", " +-+---------------- 1 ----------------", " | RuntimeError: first sub exception", " +---------------- 2 ----------------", @@ -123,7 +123,7 @@ def test_exception_handling(self): class TestMutationsExceptionGroup(TestBigtableExceptionGroup): def _get_class(self): - from google.cloud.bigtable.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup return MutationsExceptionGroup @@ -228,7 +228,7 @@ def test_from_truncated_lists( class TestRetryExceptionGroup(TestBigtableExceptionGroup): def _get_class(self): - from google.cloud.bigtable.exceptions import RetryExceptionGroup + from google.cloud.bigtable.data.exceptions import RetryExceptionGroup return RetryExceptionGroup @@ -269,7 +269,7 @@ def test_raise(self, exception_list, expected_message): class TestShardedReadRowsExceptionGroup(TestBigtableExceptionGroup): def _get_class(self): - from google.cloud.bigtable.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup return ShardedReadRowsExceptionGroup @@ -306,7 +306,7 @@ def test_raise(self, exception_list, succeeded, total_entries, expected_message) class TestFailedMutationEntryError: def _get_class(self): - from google.cloud.bigtable.exceptions import FailedMutationEntryError + from google.cloud.bigtable.data.exceptions import FailedMutationEntryError return FailedMutationEntryError @@ -374,7 +374,7 @@ def test_no_index(self): class TestFailedQueryShardError: def _get_class(self): - from google.cloud.bigtable.exceptions import FailedQueryShardError + from google.cloud.bigtable.data.exceptions import FailedQueryShardError return FailedQueryShardError diff --git a/tests/unit/test_mutations.py b/tests/unit/data/test_mutations.py similarity index 97% rename from tests/unit/test_mutations.py rename to tests/unit/data/test_mutations.py index c8c6788b1..8365dbd02 100644 --- a/tests/unit/test_mutations.py +++ b/tests/unit/data/test_mutations.py @@ -14,7 +14,7 @@ import pytest -import google.cloud.bigtable.mutations as mutations +import google.cloud.bigtable.data.mutations as mutations # try/except added for compatibility with python < 3.8 try: @@ -25,7 +25,7 @@ class TestBaseMutation: def _target_class(self): - from google.cloud.bigtable.mutations import Mutation + from google.cloud.bigtable.data.mutations import Mutation return Mutation @@ -173,7 +173,7 @@ def test__from_dict_wrong_subclass(self): class TestSetCell: def _target_class(self): - from google.cloud.bigtable.mutations import SetCell + from google.cloud.bigtable.data.mutations import SetCell return SetCell @@ -336,7 +336,7 @@ def test___str__(self): class TestDeleteRangeFromColumn: def _target_class(self): - from google.cloud.bigtable.mutations import DeleteRangeFromColumn + from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn return DeleteRangeFromColumn @@ -423,7 +423,7 @@ def test___str__(self): class TestDeleteAllFromFamily: def _target_class(self): - from google.cloud.bigtable.mutations import DeleteAllFromFamily + from google.cloud.bigtable.data.mutations import DeleteAllFromFamily return DeleteAllFromFamily @@ -460,7 +460,7 @@ def test___str__(self): class TestDeleteFromRow: def _target_class(self): - from google.cloud.bigtable.mutations import DeleteAllFromRow + from google.cloud.bigtable.data.mutations import DeleteAllFromRow return DeleteAllFromRow @@ -490,7 +490,7 @@ def test___str__(self): class TestRowMutationEntry: def _target_class(self): - from google.cloud.bigtable.mutations import RowMutationEntry + from google.cloud.bigtable.data.mutations import RowMutationEntry return RowMutationEntry @@ -506,7 +506,7 @@ def test_ctor(self): def test_ctor_over_limit(self): """Should raise error if mutations exceed MAX_MUTATIONS_PER_ENTRY""" - from google.cloud.bigtable._mutate_rows import ( + from google.cloud.bigtable.data.mutations import ( MUTATE_ROWS_REQUEST_MUTATION_LIMIT, ) @@ -527,7 +527,7 @@ def test_ctor_str_key(self): assert list(instance.mutations) == expected_mutations def test_ctor_single_mutation(self): - from google.cloud.bigtable.mutations import DeleteAllFromRow + from google.cloud.bigtable.data.mutations import DeleteAllFromRow expected_key = b"row_key" expected_mutations = DeleteAllFromRow() diff --git a/tests/unit/test_read_modify_write_rules.py b/tests/unit/data/test_read_modify_write_rules.py similarity index 94% rename from tests/unit/test_read_modify_write_rules.py rename to tests/unit/data/test_read_modify_write_rules.py index 02240df6d..aeb41f19c 100644 --- a/tests/unit/test_read_modify_write_rules.py +++ b/tests/unit/data/test_read_modify_write_rules.py @@ -24,7 +24,9 @@ class TestBaseReadModifyWriteRule: def _target_class(self): - from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule + from google.cloud.bigtable.data.read_modify_write_rules import ( + ReadModifyWriteRule, + ) return ReadModifyWriteRule @@ -40,7 +42,7 @@ def test__to_dict(self): class TestIncrementRule: def _target_class(self): - from google.cloud.bigtable.read_modify_write_rules import IncrementRule + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule return IncrementRule @@ -98,7 +100,7 @@ def test__to_dict(self, args, expected): class TestAppendValueRule: def _target_class(self): - from google.cloud.bigtable.read_modify_write_rules import AppendValueRule + from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule return AppendValueRule diff --git a/tests/unit/test_read_rows_acceptance.py b/tests/unit/data/test_read_rows_acceptance.py similarity index 93% rename from tests/unit/test_read_rows_acceptance.py rename to tests/unit/data/test_read_rows_acceptance.py index 2349d25c6..804e4e0fb 100644 --- a/tests/unit/test_read_rows_acceptance.py +++ b/tests/unit/data/test_read_rows_acceptance.py @@ -21,12 +21,13 @@ from google.cloud.bigtable_v2 import ReadRowsResponse -from google.cloud.bigtable.client import BigtableDataClient -from google.cloud.bigtable.exceptions import InvalidChunk -from google.cloud.bigtable._read_rows import _ReadRowsOperation, _StateMachine -from google.cloud.bigtable.row import Row +from google.cloud.bigtable.data._async.client import BigtableDataClientAsync +from google.cloud.bigtable.data.exceptions import InvalidChunk +from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync +from google.cloud.bigtable.data._read_rows_state_machine import _StateMachine +from google.cloud.bigtable.data.row import Row -from .v2_client.test_row_merger import ReadRowsTest, TestFile +from ..v2_client.test_row_merger import ReadRowsTest, TestFile def parse_readrows_acceptance_tests(): @@ -67,7 +68,7 @@ async def _scenerio_stream(): try: state = _StateMachine() results = [] - async for row in _ReadRowsOperation.merge_row_response_stream( + async for row in _ReadRowsOperationAsync.merge_row_response_stream( _scenerio_stream(), state ): for cell in row: @@ -117,7 +118,7 @@ def cancel(self): return mock_stream(chunk_list) try: - client = BigtableDataClient() + client = BigtableDataClientAsync() table = client.get_table("instance", "table") results = [] with mock.patch.object(table.client._gapic_client, "read_rows") as read_rows: @@ -150,7 +151,7 @@ async def _row_stream(): state = _StateMachine() state.last_seen_row_key = b"a" with pytest.raises(InvalidChunk): - async for _ in _ReadRowsOperation.merge_row_response_stream( + async for _ in _ReadRowsOperationAsync.merge_row_response_stream( _row_stream(), state ): pass @@ -309,6 +310,8 @@ async def _row_stream(): state = _StateMachine() results = [] - async for row in _ReadRowsOperation.merge_row_response_stream(_row_stream(), state): + async for row in _ReadRowsOperationAsync.merge_row_response_stream( + _row_stream(), state + ): results.append(row) return results diff --git a/tests/unit/test_read_rows_query.py b/tests/unit/data/test_read_rows_query.py similarity index 94% rename from tests/unit/test_read_rows_query.py rename to tests/unit/data/test_read_rows_query.py index 7ecd91f8c..88fde2d24 100644 --- a/tests/unit/test_read_rows_query.py +++ b/tests/unit/data/test_read_rows_query.py @@ -23,7 +23,7 @@ class TestRowRange: @staticmethod def _get_target_class(): - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import RowRange return RowRange @@ -139,7 +139,7 @@ def test__from_dict( start_is_inclusive, end_is_inclusive, ): - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import RowRange row_range = RowRange._from_dict(input_dict) assert row_range._to_dict().keys() == input_dict.keys() @@ -172,7 +172,7 @@ def test__from_dict( ], ) def test__from_points(self, dict_repr): - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import RowRange row_range_from_dict = RowRange._from_dict(dict_repr) row_range_from_points = RowRange._from_points( @@ -210,7 +210,7 @@ def test__from_points(self, dict_repr): ], ) def test___hash__(self, first_dict, second_dict, should_match): - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import RowRange row_range1 = RowRange._from_dict(first_dict) row_range2 = RowRange._from_dict(second_dict) @@ -233,7 +233,7 @@ def test___bool__(self, dict_repr, expected): """ Only row range with both points empty should be falsy """ - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import RowRange row_range = RowRange._from_dict(dict_repr) assert bool(row_range) is expected @@ -242,7 +242,7 @@ def test___bool__(self, dict_repr, expected): class TestReadRowsQuery: @staticmethod def _get_target_class(): - from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery return ReadRowsQuery @@ -257,8 +257,8 @@ def test_ctor_defaults(self): assert query.limit is None def test_ctor_explicit(self): - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.row_filters import RowFilterChain + from google.cloud.bigtable.data.read_rows_query import RowRange filter_ = RowFilterChain() query = self._make_one( @@ -281,7 +281,7 @@ def test_ctor_invalid_limit(self): assert str(exc.value) == "limit must be >= 0" def test_set_filter(self): - from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.data.row_filters import RowFilterChain filter1 = RowFilterChain() query = self._make_one() @@ -300,7 +300,7 @@ def test_set_filter(self): assert str(exc.value) == "row_filter must be a RowFilter or dict" def test_set_filter_dict(self): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter from google.cloud.bigtable_v2.types.bigtable import ReadRowsRequest filter1 = RowSampleFilter(0.5) @@ -402,7 +402,7 @@ def test_duplicate_rows(self): assert len(query.row_keys) == 3 def test_add_range(self): - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import RowRange query = self._make_one() assert query.row_ranges == set() @@ -419,7 +419,7 @@ def test_add_range(self): assert len(query.row_ranges) == 2 def test_add_range_dict(self): - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import RowRange query = self._make_one() assert query.row_ranges == set() @@ -449,8 +449,8 @@ def test_to_dict_rows_default(self): def test_to_dict_rows_populated(self): # dictionary should be in rowset proto format from google.cloud.bigtable_v2.types.bigtable import ReadRowsRequest - from google.cloud.bigtable.row_filters import PassAllFilter - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.row_filters import PassAllFilter + from google.cloud.bigtable.data.read_rows_query import RowRange row_filter = PassAllFilter(False) query = self._make_one(limit=100, row_filter=row_filter) @@ -494,7 +494,7 @@ def test_to_dict_rows_populated(self): assert filter_proto == row_filter._to_pb() def _parse_query_string(self, query_string): - from google.cloud.bigtable.read_rows_query import ReadRowsQuery, RowRange + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery, RowRange query = ReadRowsQuery() segments = query_string.split(",") @@ -550,7 +550,7 @@ def test_shard_full_table_scan_empty_split(self): """ Sharding a full table scan with no split should return another full table scan. """ - from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery full_scan_query = ReadRowsQuery() split_points = [] @@ -563,7 +563,7 @@ def test_shard_full_table_scan_with_split(self): """ Test splitting a full table scan into two queries """ - from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery full_scan_query = ReadRowsQuery() split_points = [(b"a", None)] @@ -576,7 +576,7 @@ def test_shard_full_table_scan_with_multiple_split(self): """ Test splitting a full table scan into three queries """ - from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery full_scan_query = ReadRowsQuery() split_points = [(b"a", None), (b"z", None)] @@ -684,7 +684,7 @@ def test_shard_limit_exception(self): """ queries with a limit should raise an exception when a shard is attempted """ - from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery query = ReadRowsQuery(limit=10) with pytest.raises(AttributeError) as e: @@ -718,8 +718,8 @@ def test_shard_limit_exception(self): ], ) def test___eq__(self, first_args, second_args, expected): - from google.cloud.bigtable.read_rows_query import ReadRowsQuery - from google.cloud.bigtable.read_rows_query import RowRange + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import RowRange # replace row_range placeholders with a RowRange object if len(first_args) > 1: @@ -733,7 +733,7 @@ def test___eq__(self, first_args, second_args, expected): assert (first == second) == expected def test___repr__(self): - from google.cloud.bigtable.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery instance = self._make_one(row_keys=["a", "b"], row_filter={}, limit=10) # should be able to recreate the instance from the repr diff --git a/tests/unit/test_row.py b/tests/unit/data/test_row.py similarity index 98% rename from tests/unit/test_row.py rename to tests/unit/data/test_row.py index 0413b2889..c9c797b61 100644 --- a/tests/unit/test_row.py +++ b/tests/unit/data/test_row.py @@ -27,7 +27,7 @@ class TestRow(unittest.TestCase): @staticmethod def _get_target_class(): - from google.cloud.bigtable.row import Row + from google.cloud.bigtable.data.row import Row return Row @@ -45,7 +45,7 @@ def _make_cell( timestamp=TEST_TIMESTAMP, labels=TEST_LABELS, ): - from google.cloud.bigtable.row import Cell + from google.cloud.bigtable.data.row import Cell return Cell(value, row_key, family_id, qualifier, timestamp, labels) @@ -223,7 +223,7 @@ def test_to_dict(self): self.assertEqual(column.cells[1].labels, TEST_LABELS) def test_iteration(self): - from google.cloud.bigtable.row import Cell + from google.cloud.bigtable.data.row import Cell # should be able to iterate over the Row as a list cell1 = self._make_cell(value=b"1") @@ -499,7 +499,7 @@ def test_index_of(self): class TestCell(unittest.TestCase): @staticmethod def _get_target_class(): - from google.cloud.bigtable.row import Cell + from google.cloud.bigtable.data.row import Cell return Cell @@ -623,7 +623,7 @@ def test___str__(self): self.assertEqual(str(cell), str(test_value)) def test___repr__(self): - from google.cloud.bigtable.row import Cell # type: ignore # noqa: F401 + from google.cloud.bigtable.data.row import Cell # type: ignore # noqa: F401 cell = self._make_one() expected = ( @@ -637,7 +637,7 @@ def test___repr__(self): self.assertEqual(result, cell) def test___repr___no_labels(self): - from google.cloud.bigtable.row import Cell # type: ignore # noqa: F401 + from google.cloud.bigtable.data.row import Cell # type: ignore # noqa: F401 cell_no_labels = self._make_one( TEST_VALUE, diff --git a/tests/unit/test_row_filters.py b/tests/unit/data/test_row_filters.py similarity index 81% rename from tests/unit/test_row_filters.py rename to tests/unit/data/test_row_filters.py index 11ff9f2f1..a3e275e70 100644 --- a/tests/unit/test_row_filters.py +++ b/tests/unit/data/test_row_filters.py @@ -17,10 +17,10 @@ def test_abstract_class_constructors(): - from google.cloud.bigtable.row_filters import RowFilter - from google.cloud.bigtable.row_filters import _BoolFilter - from google.cloud.bigtable.row_filters import _FilterCombination - from google.cloud.bigtable.row_filters import _CellCountFilter + from google.cloud.bigtable.data.row_filters import RowFilter + from google.cloud.bigtable.data.row_filters import _BoolFilter + from google.cloud.bigtable.data.row_filters import _FilterCombination + from google.cloud.bigtable.data.row_filters import _CellCountFilter with pytest.raises(TypeError): RowFilter() @@ -64,7 +64,7 @@ def test_bool_filter___ne__same_value(): def test_sink_filter_to_pb(): - from google.cloud.bigtable.row_filters import SinkFilter + from google.cloud.bigtable.data.row_filters import SinkFilter flag = True row_filter = SinkFilter(flag) @@ -74,7 +74,7 @@ def test_sink_filter_to_pb(): def test_sink_filter_to_dict(): - from google.cloud.bigtable.row_filters import SinkFilter + from google.cloud.bigtable.data.row_filters import SinkFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 flag = True @@ -86,7 +86,7 @@ def test_sink_filter_to_dict(): def test_sink_filter___repr__(): - from google.cloud.bigtable.row_filters import SinkFilter + from google.cloud.bigtable.data.row_filters import SinkFilter flag = True row_filter = SinkFilter(flag) @@ -96,7 +96,7 @@ def test_sink_filter___repr__(): def test_pass_all_filter_to_pb(): - from google.cloud.bigtable.row_filters import PassAllFilter + from google.cloud.bigtable.data.row_filters import PassAllFilter flag = True row_filter = PassAllFilter(flag) @@ -106,7 +106,7 @@ def test_pass_all_filter_to_pb(): def test_pass_all_filter_to_dict(): - from google.cloud.bigtable.row_filters import PassAllFilter + from google.cloud.bigtable.data.row_filters import PassAllFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 flag = True @@ -118,7 +118,7 @@ def test_pass_all_filter_to_dict(): def test_pass_all_filter___repr__(): - from google.cloud.bigtable.row_filters import PassAllFilter + from google.cloud.bigtable.data.row_filters import PassAllFilter flag = True row_filter = PassAllFilter(flag) @@ -128,7 +128,7 @@ def test_pass_all_filter___repr__(): def test_block_all_filter_to_pb(): - from google.cloud.bigtable.row_filters import BlockAllFilter + from google.cloud.bigtable.data.row_filters import BlockAllFilter flag = True row_filter = BlockAllFilter(flag) @@ -138,7 +138,7 @@ def test_block_all_filter_to_pb(): def test_block_all_filter_to_dict(): - from google.cloud.bigtable.row_filters import BlockAllFilter + from google.cloud.bigtable.data.row_filters import BlockAllFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 flag = True @@ -150,7 +150,7 @@ def test_block_all_filter_to_dict(): def test_block_all_filter___repr__(): - from google.cloud.bigtable.row_filters import BlockAllFilter + from google.cloud.bigtable.data.row_filters import BlockAllFilter flag = True row_filter = BlockAllFilter(flag) @@ -198,7 +198,7 @@ def test_regex_filter__ne__same_value(): def test_row_key_regex_filter_to_pb(): - from google.cloud.bigtable.row_filters import RowKeyRegexFilter + from google.cloud.bigtable.data.row_filters import RowKeyRegexFilter regex = b"row-key-regex" row_filter = RowKeyRegexFilter(regex) @@ -208,7 +208,7 @@ def test_row_key_regex_filter_to_pb(): def test_row_key_regex_filter_to_dict(): - from google.cloud.bigtable.row_filters import RowKeyRegexFilter + from google.cloud.bigtable.data.row_filters import RowKeyRegexFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 regex = b"row-key-regex" @@ -220,7 +220,7 @@ def test_row_key_regex_filter_to_dict(): def test_row_key_regex_filter___repr__(): - from google.cloud.bigtable.row_filters import RowKeyRegexFilter + from google.cloud.bigtable.data.row_filters import RowKeyRegexFilter regex = b"row-key-regex" row_filter = RowKeyRegexFilter(regex) @@ -230,7 +230,7 @@ def test_row_key_regex_filter___repr__(): def test_row_sample_filter_constructor(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter sample = object() row_filter = RowSampleFilter(sample) @@ -238,7 +238,7 @@ def test_row_sample_filter_constructor(): def test_row_sample_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter sample = object() row_filter1 = RowSampleFilter(sample) @@ -247,7 +247,7 @@ def test_row_sample_filter___eq__type_differ(): def test_row_sample_filter___eq__same_value(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter sample = object() row_filter1 = RowSampleFilter(sample) @@ -256,7 +256,7 @@ def test_row_sample_filter___eq__same_value(): def test_row_sample_filter___ne__(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter sample = object() other_sample = object() @@ -266,7 +266,7 @@ def test_row_sample_filter___ne__(): def test_row_sample_filter_to_pb(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter sample = 0.25 row_filter = RowSampleFilter(sample) @@ -276,7 +276,7 @@ def test_row_sample_filter_to_pb(): def test_row_sample_filter___repr__(): - from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter sample = 0.25 row_filter = RowSampleFilter(sample) @@ -286,7 +286,7 @@ def test_row_sample_filter___repr__(): def test_family_name_regex_filter_to_pb(): - from google.cloud.bigtable.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter regex = "family-regex" row_filter = FamilyNameRegexFilter(regex) @@ -296,7 +296,7 @@ def test_family_name_regex_filter_to_pb(): def test_family_name_regex_filter_to_dict(): - from google.cloud.bigtable.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 regex = "family-regex" @@ -308,7 +308,7 @@ def test_family_name_regex_filter_to_dict(): def test_family_name_regex_filter___repr__(): - from google.cloud.bigtable.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter regex = "family-regex" row_filter = FamilyNameRegexFilter(regex) @@ -319,7 +319,7 @@ def test_family_name_regex_filter___repr__(): def test_column_qualifier_regex_filter_to_pb(): - from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter + from google.cloud.bigtable.data.row_filters import ColumnQualifierRegexFilter regex = b"column-regex" row_filter = ColumnQualifierRegexFilter(regex) @@ -329,7 +329,7 @@ def test_column_qualifier_regex_filter_to_pb(): def test_column_qualifier_regex_filter_to_dict(): - from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter + from google.cloud.bigtable.data.row_filters import ColumnQualifierRegexFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 regex = b"column-regex" @@ -341,7 +341,7 @@ def test_column_qualifier_regex_filter_to_dict(): def test_column_qualifier_regex_filter___repr__(): - from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter + from google.cloud.bigtable.data.row_filters import ColumnQualifierRegexFilter regex = b"column-regex" row_filter = ColumnQualifierRegexFilter(regex) @@ -351,7 +351,7 @@ def test_column_qualifier_regex_filter___repr__(): def test_timestamp_range_constructor(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange start = object() end = object() @@ -361,7 +361,7 @@ def test_timestamp_range_constructor(): def test_timestamp_range___eq__(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange start = object() end = object() @@ -371,7 +371,7 @@ def test_timestamp_range___eq__(): def test_timestamp_range___eq__type_differ(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange start = object() end = object() @@ -381,7 +381,7 @@ def test_timestamp_range___eq__type_differ(): def test_timestamp_range___ne__same_value(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange start = object() end = object() @@ -393,7 +393,7 @@ def test_timestamp_range___ne__same_value(): def _timestamp_range_to_pb_helper(pb_kwargs, start=None, end=None): import datetime from google.cloud._helpers import _EPOCH - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange if start is not None: start = _EPOCH + datetime.timedelta(microseconds=start) @@ -421,7 +421,7 @@ def test_timestamp_range_to_pb(): def test_timestamp_range_to_dict(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange from google.cloud.bigtable_v2.types import data as data_v2_pb2 import datetime @@ -448,7 +448,7 @@ def test_timestamp_range_to_pb_start_only(): def test_timestamp_range_to_dict_start_only(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange from google.cloud.bigtable_v2.types import data as data_v2_pb2 import datetime @@ -470,7 +470,7 @@ def test_timestamp_range_to_pb_end_only(): def test_timestamp_range_to_dict_end_only(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange from google.cloud.bigtable_v2.types import data as data_v2_pb2 import datetime @@ -482,7 +482,7 @@ def test_timestamp_range_to_dict_end_only(): def timestamp_range___repr__(): - from google.cloud.bigtable.row_filters import TimestampRange + from google.cloud.bigtable.data.row_filters import TimestampRange start = object() end = object() @@ -493,7 +493,7 @@ def timestamp_range___repr__(): def test_timestamp_range_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.data.row_filters import TimestampRangeFilter range_ = object() row_filter1 = TimestampRangeFilter(range_) @@ -502,7 +502,7 @@ def test_timestamp_range_filter___eq__type_differ(): def test_timestamp_range_filter___eq__same_value(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.data.row_filters import TimestampRangeFilter range_ = object() row_filter1 = TimestampRangeFilter(range_) @@ -511,7 +511,7 @@ def test_timestamp_range_filter___eq__same_value(): def test_timestamp_range_filter___ne__(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.data.row_filters import TimestampRangeFilter range_ = object() other_range_ = object() @@ -521,7 +521,7 @@ def test_timestamp_range_filter___ne__(): def test_timestamp_range_filter_to_pb(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.data.row_filters import TimestampRangeFilter row_filter = TimestampRangeFilter() pb_val = row_filter._to_pb() @@ -530,7 +530,7 @@ def test_timestamp_range_filter_to_pb(): def test_timestamp_range_filter_to_dict(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.data.row_filters import TimestampRangeFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 import datetime @@ -549,7 +549,7 @@ def test_timestamp_range_filter_to_dict(): def test_timestamp_range_filter_empty_to_dict(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.data.row_filters import TimestampRangeFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter = TimestampRangeFilter() @@ -560,7 +560,7 @@ def test_timestamp_range_filter_empty_to_dict(): def test_timestamp_range_filter___repr__(): - from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.data.row_filters import TimestampRangeFilter import datetime start = datetime.datetime(2019, 1, 1) @@ -575,7 +575,7 @@ def test_timestamp_range_filter___repr__(): def test_column_range_filter_constructor_defaults(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = object() row_filter = ColumnRangeFilter(family_id) @@ -587,7 +587,7 @@ def test_column_range_filter_constructor_defaults(): def test_column_range_filter_constructor_explicit(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = object() start_qualifier = object() @@ -609,7 +609,7 @@ def test_column_range_filter_constructor_explicit(): def test_column_range_filter_constructor_(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = object() with pytest.raises(ValueError): @@ -617,7 +617,7 @@ def test_column_range_filter_constructor_(): def test_column_range_filter_constructor_bad_end(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = object() with pytest.raises(ValueError): @@ -625,7 +625,7 @@ def test_column_range_filter_constructor_bad_end(): def test_column_range_filter___eq__(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = object() start_qualifier = object() @@ -650,7 +650,7 @@ def test_column_range_filter___eq__(): def test_column_range_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = object() row_filter1 = ColumnRangeFilter(family_id) @@ -659,7 +659,7 @@ def test_column_range_filter___eq__type_differ(): def test_column_range_filter___ne__(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = object() other_family_id = object() @@ -685,7 +685,7 @@ def test_column_range_filter___ne__(): def test_column_range_filter_to_pb(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = "column-family-id" row_filter = ColumnRangeFilter(family_id) @@ -695,7 +695,7 @@ def test_column_range_filter_to_pb(): def test_column_range_filter_to_dict(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 family_id = "column-family-id" @@ -707,7 +707,7 @@ def test_column_range_filter_to_dict(): def test_column_range_filter_to_pb_inclusive_start(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = "column-family-id" column = b"column" @@ -718,7 +718,7 @@ def test_column_range_filter_to_pb_inclusive_start(): def test_column_range_filter_to_pb_exclusive_start(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = "column-family-id" column = b"column" @@ -731,7 +731,7 @@ def test_column_range_filter_to_pb_exclusive_start(): def test_column_range_filter_to_pb_inclusive_end(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = "column-family-id" column = b"column" @@ -742,7 +742,7 @@ def test_column_range_filter_to_pb_inclusive_end(): def test_column_range_filter_to_pb_exclusive_end(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = "column-family-id" column = b"column" @@ -753,7 +753,7 @@ def test_column_range_filter_to_pb_exclusive_end(): def test_column_range_filter___repr__(): - from google.cloud.bigtable.row_filters import ColumnRangeFilter + from google.cloud.bigtable.data.row_filters import ColumnRangeFilter family_id = "column-family-id" start_qualifier = b"column" @@ -766,7 +766,7 @@ def test_column_range_filter___repr__(): def test_value_regex_filter_to_pb_w_bytes(): - from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable.data.row_filters import ValueRegexFilter value = regex = b"value-regex" row_filter = ValueRegexFilter(value) @@ -776,7 +776,7 @@ def test_value_regex_filter_to_pb_w_bytes(): def test_value_regex_filter_to_dict_w_bytes(): - from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable.data.row_filters import ValueRegexFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 value = regex = b"value-regex" @@ -788,7 +788,7 @@ def test_value_regex_filter_to_dict_w_bytes(): def test_value_regex_filter_to_pb_w_str(): - from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable.data.row_filters import ValueRegexFilter value = "value-regex" regex = value.encode("ascii") @@ -799,7 +799,7 @@ def test_value_regex_filter_to_pb_w_str(): def test_value_regex_filter_to_dict_w_str(): - from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable.data.row_filters import ValueRegexFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 value = "value-regex" @@ -812,7 +812,7 @@ def test_value_regex_filter_to_dict_w_str(): def test_value_regex_filter___repr__(): - from google.cloud.bigtable.row_filters import ValueRegexFilter + from google.cloud.bigtable.data.row_filters import ValueRegexFilter value = "value-regex" row_filter = ValueRegexFilter(value) @@ -823,7 +823,7 @@ def test_value_regex_filter___repr__(): def test_literal_value_filter_to_pb_w_bytes(): - from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable.data.row_filters import LiteralValueFilter value = regex = b"value_regex" row_filter = LiteralValueFilter(value) @@ -833,7 +833,7 @@ def test_literal_value_filter_to_pb_w_bytes(): def test_literal_value_filter_to_dict_w_bytes(): - from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable.data.row_filters import LiteralValueFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 value = regex = b"value_regex" @@ -845,7 +845,7 @@ def test_literal_value_filter_to_dict_w_bytes(): def test_literal_value_filter_to_pb_w_str(): - from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable.data.row_filters import LiteralValueFilter value = "value_regex" regex = value.encode("ascii") @@ -856,7 +856,7 @@ def test_literal_value_filter_to_pb_w_str(): def test_literal_value_filter_to_dict_w_str(): - from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable.data.row_filters import LiteralValueFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 value = "value_regex" @@ -886,7 +886,7 @@ def test_literal_value_filter_to_dict_w_str(): ], ) def test_literal_value_filter_w_int(value, expected_byte_string): - from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable.data.row_filters import LiteralValueFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter = LiteralValueFilter(value) @@ -901,7 +901,7 @@ def test_literal_value_filter_w_int(value, expected_byte_string): def test_literal_value_filter___repr__(): - from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable.data.row_filters import LiteralValueFilter value = "value_regex" row_filter = LiteralValueFilter(value) @@ -912,7 +912,7 @@ def test_literal_value_filter___repr__(): def test_value_range_filter_constructor_defaults(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter row_filter = ValueRangeFilter() @@ -923,7 +923,7 @@ def test_value_range_filter_constructor_defaults(): def test_value_range_filter_constructor_explicit(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter start_value = object() end_value = object() @@ -944,7 +944,7 @@ def test_value_range_filter_constructor_explicit(): def test_value_range_filter_constructor_w_int_values(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter import struct start_value = 1 @@ -962,21 +962,21 @@ def test_value_range_filter_constructor_w_int_values(): def test_value_range_filter_constructor_bad_start(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter with pytest.raises(ValueError): ValueRangeFilter(inclusive_start=True) def test_value_range_filter_constructor_bad_end(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter with pytest.raises(ValueError): ValueRangeFilter(inclusive_end=True) def test_value_range_filter___eq__(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter start_value = object() end_value = object() @@ -998,7 +998,7 @@ def test_value_range_filter___eq__(): def test_value_range_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter row_filter1 = ValueRangeFilter() row_filter2 = object() @@ -1006,7 +1006,7 @@ def test_value_range_filter___eq__type_differ(): def test_value_range_filter___ne__(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter start_value = object() other_start_value = object() @@ -1029,7 +1029,7 @@ def test_value_range_filter___ne__(): def test_value_range_filter_to_pb(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter row_filter = ValueRangeFilter() expected_pb = _RowFilterPB(value_range_filter=_ValueRangePB()) @@ -1037,7 +1037,7 @@ def test_value_range_filter_to_pb(): def test_value_range_filter_to_dict(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter = ValueRangeFilter() @@ -1048,7 +1048,7 @@ def test_value_range_filter_to_dict(): def test_value_range_filter_to_pb_inclusive_start(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(start_value=value) @@ -1058,7 +1058,7 @@ def test_value_range_filter_to_pb_inclusive_start(): def test_value_range_filter_to_pb_exclusive_start(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(start_value=value, inclusive_start=False) @@ -1068,7 +1068,7 @@ def test_value_range_filter_to_pb_exclusive_start(): def test_value_range_filter_to_pb_inclusive_end(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(end_value=value) @@ -1078,7 +1078,7 @@ def test_value_range_filter_to_pb_inclusive_end(): def test_value_range_filter_to_pb_exclusive_end(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(end_value=value, inclusive_end=False) @@ -1088,7 +1088,7 @@ def test_value_range_filter_to_pb_exclusive_end(): def test_value_range_filter___repr__(): - from google.cloud.bigtable.row_filters import ValueRangeFilter + from google.cloud.bigtable.data.row_filters import ValueRangeFilter start_value = b"some-value" end_value = b"some-other-value" @@ -1133,7 +1133,7 @@ def test_cell_count___ne__same_value(): def test_cells_row_offset_filter_to_pb(): - from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.data.row_filters import CellsRowOffsetFilter num_cells = 76 row_filter = CellsRowOffsetFilter(num_cells) @@ -1143,7 +1143,7 @@ def test_cells_row_offset_filter_to_pb(): def test_cells_row_offset_filter_to_dict(): - from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.data.row_filters import CellsRowOffsetFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 num_cells = 76 @@ -1155,7 +1155,7 @@ def test_cells_row_offset_filter_to_dict(): def test_cells_row_offset_filter___repr__(): - from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.data.row_filters import CellsRowOffsetFilter num_cells = 76 row_filter = CellsRowOffsetFilter(num_cells) @@ -1166,7 +1166,7 @@ def test_cells_row_offset_filter___repr__(): def test_cells_row_limit_filter_to_pb(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter num_cells = 189 row_filter = CellsRowLimitFilter(num_cells) @@ -1176,7 +1176,7 @@ def test_cells_row_limit_filter_to_pb(): def test_cells_row_limit_filter_to_dict(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 num_cells = 189 @@ -1188,7 +1188,7 @@ def test_cells_row_limit_filter_to_dict(): def test_cells_row_limit_filter___repr__(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter num_cells = 189 row_filter = CellsRowLimitFilter(num_cells) @@ -1199,7 +1199,7 @@ def test_cells_row_limit_filter___repr__(): def test_cells_column_limit_filter_to_pb(): - from google.cloud.bigtable.row_filters import CellsColumnLimitFilter + from google.cloud.bigtable.data.row_filters import CellsColumnLimitFilter num_cells = 10 row_filter = CellsColumnLimitFilter(num_cells) @@ -1209,7 +1209,7 @@ def test_cells_column_limit_filter_to_pb(): def test_cells_column_limit_filter_to_dict(): - from google.cloud.bigtable.row_filters import CellsColumnLimitFilter + from google.cloud.bigtable.data.row_filters import CellsColumnLimitFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 num_cells = 10 @@ -1221,7 +1221,7 @@ def test_cells_column_limit_filter_to_dict(): def test_cells_column_limit_filter___repr__(): - from google.cloud.bigtable.row_filters import CellsColumnLimitFilter + from google.cloud.bigtable.data.row_filters import CellsColumnLimitFilter num_cells = 10 row_filter = CellsColumnLimitFilter(num_cells) @@ -1232,7 +1232,7 @@ def test_cells_column_limit_filter___repr__(): def test_strip_value_transformer_filter_to_pb(): - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter flag = True row_filter = StripValueTransformerFilter(flag) @@ -1242,7 +1242,7 @@ def test_strip_value_transformer_filter_to_pb(): def test_strip_value_transformer_filter_to_dict(): - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 flag = True @@ -1254,7 +1254,7 @@ def test_strip_value_transformer_filter_to_dict(): def test_strip_value_transformer_filter___repr__(): - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter flag = True row_filter = StripValueTransformerFilter(flag) @@ -1265,7 +1265,7 @@ def test_strip_value_transformer_filter___repr__(): def test_apply_label_filter_constructor(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter label = object() row_filter = ApplyLabelFilter(label) @@ -1273,7 +1273,7 @@ def test_apply_label_filter_constructor(): def test_apply_label_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter label = object() row_filter1 = ApplyLabelFilter(label) @@ -1282,7 +1282,7 @@ def test_apply_label_filter___eq__type_differ(): def test_apply_label_filter___eq__same_value(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter label = object() row_filter1 = ApplyLabelFilter(label) @@ -1291,7 +1291,7 @@ def test_apply_label_filter___eq__same_value(): def test_apply_label_filter___ne__(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter label = object() other_label = object() @@ -1301,7 +1301,7 @@ def test_apply_label_filter___ne__(): def test_apply_label_filter_to_pb(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter label = "label" row_filter = ApplyLabelFilter(label) @@ -1311,7 +1311,7 @@ def test_apply_label_filter_to_pb(): def test_apply_label_filter_to_dict(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 label = "label" @@ -1323,7 +1323,7 @@ def test_apply_label_filter_to_dict(): def test_apply_label_filter___repr__(): - from google.cloud.bigtable.row_filters import ApplyLabelFilter + from google.cloud.bigtable.data.row_filters import ApplyLabelFilter label = "label" row_filter = ApplyLabelFilter(label) @@ -1399,7 +1399,7 @@ def test_filter_combination___getitem__(): def test_filter_combination___str__(): - from google.cloud.bigtable.row_filters import PassAllFilter + from google.cloud.bigtable.data.row_filters import PassAllFilter for FilterType in _get_filter_combination_filters(): filters = [PassAllFilter(True), PassAllFilter(False)] @@ -1411,9 +1411,9 @@ def test_filter_combination___str__(): def test_row_filter_chain_to_pb(): - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import RowFilterChain + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1._to_pb() @@ -1431,9 +1431,9 @@ def test_row_filter_chain_to_pb(): def test_row_filter_chain_to_dict(): - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import RowFilterChain + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) @@ -1452,10 +1452,10 @@ def test_row_filter_chain_to_dict(): def test_row_filter_chain_to_pb_nested(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.data.row_filters import RowFilterChain + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -1476,10 +1476,10 @@ def test_row_filter_chain_to_pb_nested(): def test_row_filter_chain_to_dict_nested(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.data.row_filters import RowFilterChain + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) @@ -1502,9 +1502,9 @@ def test_row_filter_chain_to_dict_nested(): def test_row_filter_chain___repr__(): - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import RowFilterChain + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -1516,9 +1516,9 @@ def test_row_filter_chain___repr__(): def test_row_filter_chain___str__(): - from google.cloud.bigtable.row_filters import RowFilterChain - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import RowFilterChain + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -1533,9 +1533,9 @@ def test_row_filter_chain___str__(): def test_row_filter_union_to_pb(): - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import RowFilterUnion + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1._to_pb() @@ -1553,9 +1553,9 @@ def test_row_filter_union_to_pb(): def test_row_filter_union_to_dict(): - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import RowFilterUnion + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) @@ -1574,10 +1574,10 @@ def test_row_filter_union_to_dict(): def test_row_filter_union_to_pb_nested(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.data.row_filters import RowFilterUnion + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -1598,10 +1598,10 @@ def test_row_filter_union_to_pb_nested(): def test_row_filter_union_to_dict_nested(): - from google.cloud.bigtable.row_filters import CellsRowLimitFilter - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.data.row_filters import RowFilterUnion + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) @@ -1624,9 +1624,9 @@ def test_row_filter_union_to_dict_nested(): def test_row_filter_union___repr__(): - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import RowFilterUnion + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -1638,9 +1638,9 @@ def test_row_filter_union___repr__(): def test_row_filter_union___str__(): - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import RowFilterUnion + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -1655,7 +1655,7 @@ def test_row_filter_union___str__(): def test_conditional_row_filter_constructor(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter predicate_filter = object() true_filter = object() @@ -1669,7 +1669,7 @@ def test_conditional_row_filter_constructor(): def test_conditional_row_filter___eq__(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter predicate_filter = object() true_filter = object() @@ -1684,7 +1684,7 @@ def test_conditional_row_filter___eq__(): def test_conditional_row_filter___eq__type_differ(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter predicate_filter = object() true_filter = object() @@ -1697,7 +1697,7 @@ def test_conditional_row_filter___eq__type_differ(): def test_conditional_row_filter___ne__(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter predicate_filter = object() other_predicate_filter = object() @@ -1713,10 +1713,10 @@ def test_conditional_row_filter___ne__(): def test_conditional_row_filter_to_pb(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import CellsRowOffsetFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1._to_pb() @@ -1743,10 +1743,10 @@ def test_conditional_row_filter_to_pb(): def test_conditional_row_filter_to_dict(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import CellsRowOffsetFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) @@ -1776,9 +1776,9 @@ def test_conditional_row_filter_to_dict(): def test_conditional_row_filter_to_pb_true_only(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1._to_pb() @@ -1798,9 +1798,9 @@ def test_conditional_row_filter_to_pb_true_only(): def test_conditional_row_filter_to_dict_true_only(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) @@ -1824,9 +1824,9 @@ def test_conditional_row_filter_to_dict_true_only(): def test_conditional_row_filter_to_pb_false_only(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1._to_pb() @@ -1846,9 +1846,9 @@ def test_conditional_row_filter_to_pb_false_only(): def test_conditional_row_filter_to_dict_false_only(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) @@ -1872,9 +1872,9 @@ def test_conditional_row_filter_to_dict_false_only(): def test_conditional_row_filter___repr__(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -1893,10 +1893,10 @@ def test_conditional_row_filter___repr__(): def test_conditional_row_filter___str__(): - from google.cloud.bigtable.row_filters import ConditionalRowFilter - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_filters import RowFilterUnion - from google.cloud.bigtable.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.data.row_filters import ConditionalRowFilter + from google.cloud.bigtable.data.row_filters import RowSampleFilter + from google.cloud.bigtable.data.row_filters import RowFilterUnion + from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -1931,7 +1931,7 @@ def test_conditional_row_filter___str__(): ], ) def test_literal_value__write_literal_regex(input_arg, expected_bytes): - from google.cloud.bigtable.row_filters import LiteralValueFilter + from google.cloud.bigtable.data.row_filters import LiteralValueFilter filter_ = LiteralValueFilter(input_arg) assert filter_.regex == expected_bytes @@ -1980,7 +1980,7 @@ def _ValueRangePB(*args, **kw): def _get_regex_filters(): - from google.cloud.bigtable.row_filters import ( + from google.cloud.bigtable.data.row_filters import ( RowKeyRegexFilter, FamilyNameRegexFilter, ColumnQualifierRegexFilter, @@ -1998,7 +1998,7 @@ def _get_regex_filters(): def _get_bool_filters(): - from google.cloud.bigtable.row_filters import ( + from google.cloud.bigtable.data.row_filters import ( SinkFilter, PassAllFilter, BlockAllFilter, @@ -2014,7 +2014,7 @@ def _get_bool_filters(): def _get_cell_count_filters(): - from google.cloud.bigtable.row_filters import ( + from google.cloud.bigtable.data.row_filters import ( CellsRowLimitFilter, CellsRowOffsetFilter, CellsColumnLimitFilter, @@ -2028,7 +2028,7 @@ def _get_cell_count_filters(): def _get_filter_combination_filters(): - from google.cloud.bigtable.row_filters import ( + from google.cloud.bigtable.data.row_filters import ( RowFilterChain, RowFilterUnion, ) diff --git a/tests/unit/gapic/bigtable_admin_v2/test_bigtable_table_admin.py b/tests/unit/gapic/bigtable_admin_v2/test_bigtable_table_admin.py index 8e4004ab1..8498e4fa5 100644 --- a/tests/unit/gapic/bigtable_admin_v2/test_bigtable_table_admin.py +++ b/tests/unit/gapic/bigtable_admin_v2/test_bigtable_table_admin.py @@ -8202,6 +8202,7 @@ def test_update_table_rest(request_type): "source_table": "source_table_value", }, }, + "change_stream_config": {"retention_period": {"seconds": 751, "nanos": 543}}, "deletion_protection": True, } request = request_type(**request_init) @@ -8399,6 +8400,7 @@ def test_update_table_rest_bad_request( "source_table": "source_table_value", }, }, + "change_stream_config": {"retention_period": {"seconds": 751, "nanos": 543}}, "deletion_protection": True, } request = request_type(**request_init) diff --git a/tests/unit/gapic/bigtable_v2/test_bigtable.py b/tests/unit/gapic/bigtable_v2/test_bigtable.py index b1500aa48..03ba3044f 100644 --- a/tests/unit/gapic/bigtable_v2/test_bigtable.py +++ b/tests/unit/gapic/bigtable_v2/test_bigtable.py @@ -100,7 +100,6 @@ def test__get_default_mtls_endpoint(): [ (BigtableClient, "grpc"), (BigtableAsyncClient, "grpc_asyncio"), - (BigtableAsyncClient, "pooled_grpc_asyncio"), (BigtableClient, "rest"), ], ) @@ -117,7 +116,7 @@ def test_bigtable_client_from_service_account_info(client_class, transport_name) assert client.transport._host == ( "bigtable.googleapis.com:443" - if transport_name in ["grpc", "grpc_asyncio", "pooled_grpc_asyncio"] + if transport_name in ["grpc", "grpc_asyncio"] else "https://bigtable.googleapis.com" ) @@ -127,7 +126,6 @@ def test_bigtable_client_from_service_account_info(client_class, transport_name) [ (transports.BigtableGrpcTransport, "grpc"), (transports.BigtableGrpcAsyncIOTransport, "grpc_asyncio"), - (transports.PooledBigtableGrpcAsyncIOTransport, "pooled_grpc_asyncio"), (transports.BigtableRestTransport, "rest"), ], ) @@ -154,7 +152,6 @@ def test_bigtable_client_service_account_always_use_jwt( [ (BigtableClient, "grpc"), (BigtableAsyncClient, "grpc_asyncio"), - (BigtableAsyncClient, "pooled_grpc_asyncio"), (BigtableClient, "rest"), ], ) @@ -178,7 +175,7 @@ def test_bigtable_client_from_service_account_file(client_class, transport_name) assert client.transport._host == ( "bigtable.googleapis.com:443" - if transport_name in ["grpc", "grpc_asyncio", "pooled_grpc_asyncio"] + if transport_name in ["grpc", "grpc_asyncio"] else "https://bigtable.googleapis.com" ) @@ -200,11 +197,6 @@ def test_bigtable_client_get_transport_class(): [ (BigtableClient, transports.BigtableGrpcTransport, "grpc"), (BigtableAsyncClient, transports.BigtableGrpcAsyncIOTransport, "grpc_asyncio"), - ( - BigtableAsyncClient, - transports.PooledBigtableGrpcAsyncIOTransport, - "pooled_grpc_asyncio", - ), (BigtableClient, transports.BigtableRestTransport, "rest"), ], ) @@ -340,12 +332,6 @@ def test_bigtable_client_client_options(client_class, transport_class, transport "grpc_asyncio", "true", ), - ( - BigtableAsyncClient, - transports.PooledBigtableGrpcAsyncIOTransport, - "pooled_grpc_asyncio", - "true", - ), (BigtableClient, transports.BigtableGrpcTransport, "grpc", "false"), ( BigtableAsyncClient, @@ -353,12 +339,6 @@ def test_bigtable_client_client_options(client_class, transport_class, transport "grpc_asyncio", "false", ), - ( - BigtableAsyncClient, - transports.PooledBigtableGrpcAsyncIOTransport, - "pooled_grpc_asyncio", - "false", - ), (BigtableClient, transports.BigtableRestTransport, "rest", "true"), (BigtableClient, transports.BigtableRestTransport, "rest", "false"), ], @@ -550,11 +530,6 @@ def test_bigtable_client_get_mtls_endpoint_and_cert_source(client_class): [ (BigtableClient, transports.BigtableGrpcTransport, "grpc"), (BigtableAsyncClient, transports.BigtableGrpcAsyncIOTransport, "grpc_asyncio"), - ( - BigtableAsyncClient, - transports.PooledBigtableGrpcAsyncIOTransport, - "pooled_grpc_asyncio", - ), (BigtableClient, transports.BigtableRestTransport, "rest"), ], ) @@ -591,12 +566,6 @@ def test_bigtable_client_client_options_scopes( "grpc_asyncio", grpc_helpers_async, ), - ( - BigtableAsyncClient, - transports.PooledBigtableGrpcAsyncIOTransport, - "pooled_grpc_asyncio", - grpc_helpers_async, - ), (BigtableClient, transports.BigtableRestTransport, "rest", None), ], ) @@ -743,35 +712,6 @@ def test_read_rows(request_type, transport: str = "grpc"): assert isinstance(message, bigtable.ReadRowsResponse) -def test_read_rows_pooled_rotation(transport: str = "pooled_grpc_asyncio"): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.read_rows(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.read_rows(request) - assert next_channel.call_count == i - - def test_read_rows_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -991,35 +931,6 @@ def test_sample_row_keys(request_type, transport: str = "grpc"): assert isinstance(message, bigtable.SampleRowKeysResponse) -def test_sample_row_keys_pooled_rotation(transport: str = "pooled_grpc_asyncio"): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.sample_row_keys(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.sample_row_keys(request) - assert next_channel.call_count == i - - def test_sample_row_keys_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -1238,35 +1149,6 @@ def test_mutate_row(request_type, transport: str = "grpc"): assert isinstance(response, bigtable.MutateRowResponse) -def test_mutate_row_pooled_rotation(transport: str = "pooled_grpc_asyncio"): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.mutate_row(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.mutate_row(request) - assert next_channel.call_count == i - - def test_mutate_row_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -1530,35 +1412,6 @@ def test_mutate_rows(request_type, transport: str = "grpc"): assert isinstance(message, bigtable.MutateRowsResponse) -def test_mutate_rows_pooled_rotation(transport: str = "pooled_grpc_asyncio"): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.mutate_rows(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.mutate_rows(request) - assert next_channel.call_count == i - - def test_mutate_rows_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -1792,35 +1645,6 @@ def test_check_and_mutate_row(request_type, transport: str = "grpc"): assert response.predicate_matched is True -def test_check_and_mutate_row_pooled_rotation(transport: str = "pooled_grpc_asyncio"): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.check_and_mutate_row(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.check_and_mutate_row(request) - assert next_channel.call_count == i - - def test_check_and_mutate_row_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -2198,35 +2022,6 @@ def test_ping_and_warm(request_type, transport: str = "grpc"): assert isinstance(response, bigtable.PingAndWarmResponse) -def test_ping_and_warm_pooled_rotation(transport: str = "pooled_grpc_asyncio"): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.ping_and_warm(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.ping_and_warm(request) - assert next_channel.call_count == i - - def test_ping_and_warm_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -2447,35 +2242,6 @@ def test_read_modify_write_row(request_type, transport: str = "grpc"): assert isinstance(response, bigtable.ReadModifyWriteRowResponse) -def test_read_modify_write_row_pooled_rotation(transport: str = "pooled_grpc_asyncio"): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.read_modify_write_row(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.read_modify_write_row(request) - assert next_channel.call_count == i - - def test_read_modify_write_row_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -2735,37 +2501,6 @@ def test_generate_initial_change_stream_partitions( ) -def test_generate_initial_change_stream_partitions_pooled_rotation( - transport: str = "pooled_grpc_asyncio", -): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.generate_initial_change_stream_partitions(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.generate_initial_change_stream_partitions(request) - assert next_channel.call_count == i - - def test_generate_initial_change_stream_partitions_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -3025,35 +2760,6 @@ def test_read_change_stream(request_type, transport: str = "grpc"): assert isinstance(message, bigtable.ReadChangeStreamResponse) -def test_read_change_stream_pooled_rotation(transport: str = "pooled_grpc_asyncio"): - with mock.patch.object( - transports.pooled_grpc_asyncio.PooledChannel, "next_channel" - ) as next_channel: - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport=transport, - ) - - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = {} - - channel = client.transport._grpc_channel._pool[ - client.transport._grpc_channel._next_idx - ] - next_channel.return_value = channel - - response = client.read_change_stream(request) - - # Establish that next_channel was called - next_channel.assert_called_once() - # Establish that subsequent calls all call next_channel - starting_idx = client.transport._grpc_channel._next_idx - for i in range(2, 10): - response = client.read_change_stream(request) - assert next_channel.call_count == i - - def test_read_change_stream_empty_call(): # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @@ -5957,7 +5663,6 @@ def test_transport_get_channel(): [ transports.BigtableGrpcTransport, transports.BigtableGrpcAsyncIOTransport, - transports.PooledBigtableGrpcAsyncIOTransport, transports.BigtableRestTransport, ], ) @@ -6105,7 +5810,6 @@ def test_bigtable_auth_adc(): [ transports.BigtableGrpcTransport, transports.BigtableGrpcAsyncIOTransport, - transports.PooledBigtableGrpcAsyncIOTransport, ], ) def test_bigtable_transport_auth_adc(transport_class): @@ -6133,7 +5837,6 @@ def test_bigtable_transport_auth_adc(transport_class): [ transports.BigtableGrpcTransport, transports.BigtableGrpcAsyncIOTransport, - transports.PooledBigtableGrpcAsyncIOTransport, transports.BigtableRestTransport, ], ) @@ -6236,61 +5939,6 @@ def test_bigtable_grpc_transport_client_cert_source_for_mtls(transport_class): ) -@pytest.mark.parametrize( - "transport_class", [transports.PooledBigtableGrpcAsyncIOTransport] -) -def test_bigtable_pooled_grpc_transport_client_cert_source_for_mtls(transport_class): - cred = ga_credentials.AnonymousCredentials() - - # test with invalid pool size - with pytest.raises(ValueError): - transport_class( - host="squid.clam.whelk", - credentials=cred, - pool_size=0, - ) - - # Check ssl_channel_credentials is used if provided. - for pool_num in range(1, 5): - with mock.patch.object( - transport_class, "create_channel" - ) as mock_create_channel: - mock_ssl_channel_creds = mock.Mock() - transport_class( - host="squid.clam.whelk", - credentials=cred, - ssl_channel_credentials=mock_ssl_channel_creds, - pool_size=pool_num, - ) - mock_create_channel.assert_called_with( - pool_num, - "squid.clam.whelk:443", - credentials=cred, - credentials_file=None, - scopes=None, - ssl_credentials=mock_ssl_channel_creds, - quota_project_id=None, - options=[ - ("grpc.max_send_message_length", -1), - ("grpc.max_receive_message_length", -1), - ], - ) - assert mock_create_channel.call_count == 1 - - # Check if ssl_channel_credentials is not provided, then client_cert_source_for_mtls - # is used. - with mock.patch.object(transport_class, "create_channel", return_value=mock.Mock()): - with mock.patch("grpc.ssl_channel_credentials") as mock_ssl_cred: - transport_class( - credentials=cred, - client_cert_source_for_mtls=client_cert_source_callback, - ) - expected_cert, expected_key = client_cert_source_callback() - mock_ssl_cred.assert_called_once_with( - certificate_chain=expected_cert, private_key=expected_key - ) - - def test_bigtable_http_transport_client_cert_source_for_mtls(): cred = ga_credentials.AnonymousCredentials() with mock.patch( @@ -6307,7 +5955,6 @@ def test_bigtable_http_transport_client_cert_source_for_mtls(): [ "grpc", "grpc_asyncio", - "pooled_grpc_asyncio", "rest", ], ) @@ -6321,7 +5968,7 @@ def test_bigtable_host_no_port(transport_name): ) assert client.transport._host == ( "bigtable.googleapis.com:443" - if transport_name in ["grpc", "grpc_asyncio", "pooled_grpc_asyncio"] + if transport_name in ["grpc", "grpc_asyncio"] else "https://bigtable.googleapis.com" ) @@ -6331,7 +5978,6 @@ def test_bigtable_host_no_port(transport_name): [ "grpc", "grpc_asyncio", - "pooled_grpc_asyncio", "rest", ], ) @@ -6345,7 +5991,7 @@ def test_bigtable_host_with_port(transport_name): ) assert client.transport._host == ( "bigtable.googleapis.com:8000" - if transport_name in ["grpc", "grpc_asyncio", "pooled_grpc_asyncio"] + if transport_name in ["grpc", "grpc_asyncio"] else "https://bigtable.googleapis.com:8000" ) @@ -6701,24 +6347,6 @@ async def test_transport_close_async(): async with client: close.assert_not_called() close.assert_called_once() - close.assert_awaited() - - -@pytest.mark.asyncio -async def test_pooled_transport_close_async(): - client = BigtableAsyncClient( - credentials=ga_credentials.AnonymousCredentials(), - transport="pooled_grpc_asyncio", - ) - num_channels = len(client.transport._grpc_channel._pool) - with mock.patch.object( - type(client.transport._grpc_channel._pool[0]), "close" - ) as close: - async with client: - close.assert_not_called() - close.assert_called() - assert close.call_count == num_channels - close.assert_awaited() def test_transport_close(): @@ -6785,128 +6413,3 @@ def test_api_key_credentials(client_class, transport_class): always_use_jwt_access=True, api_audience=None, ) - - -@pytest.mark.asyncio -async def test_pooled_transport_replace_default(): - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport="pooled_grpc_asyncio", - ) - num_channels = len(client.transport._grpc_channel._pool) - for replace_idx in range(num_channels): - prev_pool = [channel for channel in client.transport._grpc_channel._pool] - grace_period = 4 - with mock.patch.object( - type(client.transport._grpc_channel._pool[0]), "close" - ) as close: - await client.transport.replace_channel(replace_idx, grace=grace_period) - close.assert_called_once() - close.assert_awaited() - close.assert_called_with(grace=grace_period) - assert isinstance( - client.transport._grpc_channel._pool[replace_idx], grpc.aio.Channel - ) - # only the specified channel should be replaced - for i in range(num_channels): - if i == replace_idx: - assert client.transport._grpc_channel._pool[i] != prev_pool[i] - else: - assert client.transport._grpc_channel._pool[i] == prev_pool[i] - with pytest.raises(ValueError): - await client.transport.replace_channel(num_channels + 1) - with pytest.raises(ValueError): - await client.transport.replace_channel(-1) - - -@pytest.mark.asyncio -async def test_pooled_transport_replace_explicit(): - client = BigtableClient( - credentials=ga_credentials.AnonymousCredentials(), - transport="pooled_grpc_asyncio", - ) - num_channels = len(client.transport._grpc_channel._pool) - for replace_idx in range(num_channels): - prev_pool = [channel for channel in client.transport._grpc_channel._pool] - grace_period = 0 - with mock.patch.object( - type(client.transport._grpc_channel._pool[0]), "close" - ) as close: - new_channel = grpc.aio.insecure_channel("localhost:8080") - await client.transport.replace_channel( - replace_idx, grace=grace_period, new_channel=new_channel - ) - close.assert_called_once() - close.assert_awaited() - close.assert_called_with(grace=grace_period) - assert client.transport._grpc_channel._pool[replace_idx] == new_channel - # only the specified channel should be replaced - for i in range(num_channels): - if i == replace_idx: - assert client.transport._grpc_channel._pool[i] != prev_pool[i] - else: - assert client.transport._grpc_channel._pool[i] == prev_pool[i] - - -def test_pooled_transport_next_channel(): - num_channels = 10 - transport = transports.PooledBigtableGrpcAsyncIOTransport( - credentials=ga_credentials.AnonymousCredentials(), - pool_size=num_channels, - ) - assert len(transport._grpc_channel._pool) == num_channels - transport._grpc_channel._next_idx = 0 - # rotate through all channels multiple times - num_cycles = 4 - for _ in range(num_cycles): - for i in range(num_channels - 1): - assert transport._grpc_channel._next_idx == i - got_channel = transport._grpc_channel.next_channel() - assert got_channel == transport._grpc_channel._pool[i] - assert transport._grpc_channel._next_idx == (i + 1) - # test wrap around - assert transport._grpc_channel._next_idx == num_channels - 1 - got_channel = transport._grpc_channel.next_channel() - assert got_channel == transport._grpc_channel._pool[num_channels - 1] - assert transport._grpc_channel._next_idx == 0 - - -def test_pooled_transport_pool_unique_channels(): - num_channels = 50 - - transport = transports.PooledBigtableGrpcAsyncIOTransport( - credentials=ga_credentials.AnonymousCredentials(), - pool_size=num_channels, - ) - channel_list = [channel for channel in transport._grpc_channel._pool] - channel_set = set(channel_list) - assert len(channel_list) == num_channels - assert len(channel_set) == num_channels - for channel in channel_list: - assert isinstance(channel, grpc.aio.Channel) - - -def test_pooled_transport_pool_creation(): - # channels should be created with the specified options - num_channels = 50 - creds = ga_credentials.AnonymousCredentials() - scopes = ["test1", "test2"] - quota_project_id = "test3" - host = "testhost:8080" - with mock.patch( - "google.api_core.grpc_helpers_async.create_channel" - ) as create_channel: - transport = transports.PooledBigtableGrpcAsyncIOTransport( - credentials=creds, - pool_size=num_channels, - scopes=scopes, - quota_project_id=quota_project_id, - host=host, - ) - assert create_channel.call_count == num_channels - for i in range(num_channels): - kwargs = create_channel.call_args_list[i][1] - assert kwargs["target"] == host - assert kwargs["credentials"] == creds - assert kwargs["scopes"] == scopes - assert kwargs["quota_project_id"] == quota_project_id diff --git a/tests/unit/test_iterators.py b/tests/unit/test_iterators.py deleted file mode 100644 index f7aee2822..000000000 --- a/tests/unit/test_iterators.py +++ /dev/null @@ -1,251 +0,0 @@ -# Copyright 2023 Google LLC -# -# Licensed 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. -from __future__ import annotations - -import sys -import asyncio -import pytest - -from google.cloud.bigtable._read_rows import _ReadRowsOperation - -# try/except added for compatibility with python < 3.8 -try: - from unittest import mock -except ImportError: # pragma: NO COVER - import mock # type: ignore - - -class MockStream(_ReadRowsOperation): - """ - Mock a _ReadRowsOperation stream for testing - """ - - def __init__(self, items=None, errors=None, operation_timeout=None): - self.transient_errors = errors - self.operation_timeout = operation_timeout - self.next_idx = 0 - if items is None: - items = list(range(10)) - self.items = items - - def __aiter__(self): - return self - - async def __anext__(self): - if self.next_idx >= len(self.items): - raise StopAsyncIteration - item = self.items[self.next_idx] - self.next_idx += 1 - if isinstance(item, Exception): - raise item - return item - - async def aclose(self): - pass - - -class TestReadRowsIterator: - async def mock_stream(self, size=10): - for i in range(size): - yield i - - def _make_one(self, *args, **kwargs): - from google.cloud.bigtable.iterators import ReadRowsIterator - - stream = MockStream(*args, **kwargs) - return ReadRowsIterator(stream) - - def test_ctor(self): - with mock.patch("time.time", return_value=0): - iterator = self._make_one() - assert iterator.last_interaction_time == 0 - assert iterator._idle_timeout_task is None - assert iterator.active is True - - def test___aiter__(self): - iterator = self._make_one() - assert iterator.__aiter__() is iterator - - @pytest.mark.skipif( - sys.version_info < (3, 8), reason="mock coroutine requires python3.8 or higher" - ) - @pytest.mark.asyncio - async def test__start_idle_timer(self): - """Should start timer coroutine""" - iterator = self._make_one() - expected_timeout = 10 - with mock.patch("time.time", return_value=1): - with mock.patch.object(iterator, "_idle_timeout_coroutine") as mock_coro: - await iterator._start_idle_timer(expected_timeout) - assert mock_coro.call_count == 1 - assert mock_coro.call_args[0] == (expected_timeout,) - assert iterator.last_interaction_time == 1 - assert iterator._idle_timeout_task is not None - - @pytest.mark.skipif( - sys.version_info < (3, 8), reason="mock coroutine requires python3.8 or higher" - ) - @pytest.mark.asyncio - async def test__start_idle_timer_duplicate(self): - """Multiple calls should replace task""" - iterator = self._make_one() - with mock.patch.object(iterator, "_idle_timeout_coroutine") as mock_coro: - await iterator._start_idle_timer(1) - first_task = iterator._idle_timeout_task - await iterator._start_idle_timer(2) - second_task = iterator._idle_timeout_task - assert mock_coro.call_count == 2 - - assert first_task is not None - assert first_task != second_task - # old tasks hould be cancelled - with pytest.raises(asyncio.CancelledError): - await first_task - # new task should not be cancelled - await second_task - - @pytest.mark.asyncio - async def test__idle_timeout_coroutine(self): - from google.cloud.bigtable.exceptions import IdleTimeout - - iterator = self._make_one() - await iterator._idle_timeout_coroutine(0.05) - await asyncio.sleep(0.1) - assert iterator.active is False - with pytest.raises(IdleTimeout): - await iterator.__anext__() - - @pytest.mark.asyncio - async def test__idle_timeout_coroutine_extensions(self): - """touching the generator should reset the idle timer""" - iterator = self._make_one(items=list(range(100))) - await iterator._start_idle_timer(0.05) - for i in range(10): - # will not expire as long as it is in use - assert iterator.active is True - await iterator.__anext__() - await asyncio.sleep(0.03) - # now let it expire - await asyncio.sleep(0.5) - assert iterator.active is False - - @pytest.mark.asyncio - async def test___anext__(self): - num_rows = 10 - iterator = self._make_one(items=list(range(num_rows))) - for i in range(num_rows): - assert await iterator.__anext__() == i - with pytest.raises(StopAsyncIteration): - await iterator.__anext__() - - @pytest.mark.asyncio - async def test___anext__with_deadline_error(self): - """ - RetryErrors mean a deadline has been hit. - Should be wrapped in a DeadlineExceeded exception - """ - from google.api_core import exceptions as core_exceptions - - items = [1, core_exceptions.RetryError("retry error", None)] - expected_timeout = 99 - iterator = self._make_one(items=items, operation_timeout=expected_timeout) - assert await iterator.__anext__() == 1 - with pytest.raises(core_exceptions.DeadlineExceeded) as exc: - await iterator.__anext__() - assert f"operation_timeout of {expected_timeout:0.1f}s exceeded" in str( - exc.value - ) - assert exc.value.__cause__ is None - - @pytest.mark.asyncio - async def test___anext__with_deadline_error_with_cause(self): - """ - Transient errors should be exposed as an error group - """ - from google.api_core import exceptions as core_exceptions - from google.cloud.bigtable.exceptions import RetryExceptionGroup - - items = [1, core_exceptions.RetryError("retry error", None)] - expected_timeout = 99 - errors = [RuntimeError("error1"), ValueError("error2")] - iterator = self._make_one( - items=items, operation_timeout=expected_timeout, errors=errors - ) - assert await iterator.__anext__() == 1 - with pytest.raises(core_exceptions.DeadlineExceeded) as exc: - await iterator.__anext__() - assert f"operation_timeout of {expected_timeout:0.1f}s exceeded" in str( - exc.value - ) - error_group = exc.value.__cause__ - assert isinstance(error_group, RetryExceptionGroup) - assert len(error_group.exceptions) == 2 - assert error_group.exceptions[0] is errors[0] - assert error_group.exceptions[1] is errors[1] - assert "2 failed attempts" in str(error_group) - - @pytest.mark.asyncio - async def test___anext__with_error(self): - """ - Other errors should be raised as-is - """ - from google.api_core import exceptions as core_exceptions - - items = [1, core_exceptions.InternalServerError("mock error")] - iterator = self._make_one(items=items) - assert await iterator.__anext__() == 1 - with pytest.raises(core_exceptions.InternalServerError) as exc: - await iterator.__anext__() - assert exc.value is items[1] - assert iterator.active is False - # next call should raise same error - with pytest.raises(core_exceptions.InternalServerError) as exc: - await iterator.__anext__() - - @pytest.mark.asyncio - async def test__finish_with_error(self): - iterator = self._make_one() - await iterator._start_idle_timer(10) - timeout_task = iterator._idle_timeout_task - assert await iterator.__anext__() == 0 - assert iterator.active is True - err = ZeroDivisionError("mock error") - await iterator._finish_with_error(err) - assert iterator.active is False - assert iterator._error is err - assert iterator._idle_timeout_task is None - with pytest.raises(ZeroDivisionError) as exc: - await iterator.__anext__() - assert exc.value is err - # timeout task should be cancelled - with pytest.raises(asyncio.CancelledError): - await timeout_task - - @pytest.mark.asyncio - async def test_aclose(self): - iterator = self._make_one() - await iterator._start_idle_timer(10) - timeout_task = iterator._idle_timeout_task - assert await iterator.__anext__() == 0 - assert iterator.active is True - await iterator.aclose() - assert iterator.active is False - assert isinstance(iterator._error, StopAsyncIteration) - assert iterator._idle_timeout_task is None - with pytest.raises(StopAsyncIteration) as e: - await iterator.__anext__() - assert "closed" in str(e.value) - # timeout task should be cancelled - with pytest.raises(asyncio.CancelledError): - await timeout_task diff --git a/tests/unit/v2_client/test_app_profile.py b/tests/unit/v2_client/test_app_profile.py index 575f25194..660ee7899 100644 --- a/tests/unit/v2_client/test_app_profile.py +++ b/tests/unit/v2_client/test_app_profile.py @@ -32,19 +32,19 @@ def _make_app_profile(*args, **kwargs): - from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.app_profile import AppProfile return AppProfile(*args, **kwargs) def _make_client(*args, **kwargs): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client return Client(*args, **kwargs) def test_app_profile_constructor_defaults(): - from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.app_profile import AppProfile client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -60,7 +60,7 @@ def test_app_profile_constructor_defaults(): def test_app_profile_constructor_explicit(): - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.enums import RoutingPolicyType ANY = RoutingPolicyType.ANY DESCRIPTION_1 = "routing policy any" @@ -99,7 +99,7 @@ def test_app_profile_constructor_explicit(): def test_app_profile_constructor_multi_cluster_ids(): - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.enums import RoutingPolicyType ANY = RoutingPolicyType.ANY DESCRIPTION_1 = "routing policy any" @@ -166,8 +166,8 @@ def test_app_profile___ne__(): def test_app_profile_from_pb_success_w_routing_any(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.app_profile import AppProfile - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.enums import RoutingPolicyType client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -195,8 +195,8 @@ def test_app_profile_from_pb_success_w_routing_any(): def test_app_profile_from_pb_success_w_routing_any_multi_cluster_ids(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.app_profile import AppProfile - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.enums import RoutingPolicyType client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -226,8 +226,8 @@ def test_app_profile_from_pb_success_w_routing_any_multi_cluster_ids(): def test_app_profile_from_pb_success_w_routing_single(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.app_profile import AppProfile - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.enums import RoutingPolicyType client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -259,7 +259,7 @@ def test_app_profile_from_pb_success_w_routing_single(): def test_app_profile_from_pb_w_bad_app_profile_name(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.app_profile import AppProfile bad_app_profile_name = "BAD_NAME" @@ -271,7 +271,7 @@ def test_app_profile_from_pb_w_bad_app_profile_name(): def test_app_profile_from_pb_w_instance_id_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.app_profile import AppProfile ALT_INSTANCE_ID = "ALT_INSTANCE_ID" client = _Client(PROJECT) @@ -286,7 +286,7 @@ def test_app_profile_from_pb_w_instance_id_mistmatch(): def test_app_profile_from_pb_w_project_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.app_profile import AppProfile ALT_PROJECT = "ALT_PROJECT" client = _Client(project=ALT_PROJECT) @@ -304,7 +304,7 @@ def test_app_profile_reload_w_routing_any(): BigtableInstanceAdminClient, ) from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.enums import RoutingPolicyType api = mock.create_autospec(BigtableInstanceAdminClient) credentials = _make_credentials() @@ -400,8 +400,8 @@ def test_app_profile_create_w_routing_any(): from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) - from google.cloud.bigtable.deprecated.app_profile import AppProfile - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.enums import RoutingPolicyType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -461,8 +461,8 @@ def test_app_profile_create_w_routing_single(): from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) - from google.cloud.bigtable.deprecated.app_profile import AppProfile - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.app_profile import AppProfile + from google.cloud.bigtable.enums import RoutingPolicyType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -533,7 +533,7 @@ def test_app_profile_update_w_routing_any(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.enums import RoutingPolicyType from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -608,7 +608,7 @@ def test_app_profile_update_w_routing_any_multi_cluster_ids(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.enums import RoutingPolicyType from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -684,7 +684,7 @@ def test_app_profile_update_w_routing_single(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.enums import RoutingPolicyType from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) diff --git a/tests/unit/v2_client/test_backup.py b/tests/unit/v2_client/test_backup.py index 34cc8823a..9882ca339 100644 --- a/tests/unit/v2_client/test_backup.py +++ b/tests/unit/v2_client/test_backup.py @@ -48,7 +48,7 @@ def _make_table_admin_client(): def _make_backup(*args, **kwargs): - from google.cloud.bigtable.deprecated.backup import Backup + from google.cloud.bigtable.backup import Backup return Backup(*args, **kwargs) @@ -102,7 +102,7 @@ def test_backup_constructor_explicit(): def test_backup_from_pb_w_project_mismatch(): from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.deprecated.backup import Backup + from google.cloud.bigtable.backup import Backup alt_project_id = "alt-project-id" client = _Client(project=alt_project_id) @@ -115,7 +115,7 @@ def test_backup_from_pb_w_project_mismatch(): def test_backup_from_pb_w_instance_mismatch(): from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.deprecated.backup import Backup + from google.cloud.bigtable.backup import Backup alt_instance = "/projects/%s/instances/alt-instance" % PROJECT_ID client = _Client() @@ -128,7 +128,7 @@ def test_backup_from_pb_w_instance_mismatch(): def test_backup_from_pb_w_bad_name(): from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.deprecated.backup import Backup + from google.cloud.bigtable.backup import Backup client = _Client() instance = _Instance(INSTANCE_NAME, client) @@ -139,10 +139,10 @@ def test_backup_from_pb_w_bad_name(): def test_backup_from_pb_success(): - from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo - from google.cloud.bigtable.deprecated.error import Status + from google.cloud.bigtable.encryption_info import EncryptionInfo + from google.cloud.bigtable.error import Status from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.deprecated.backup import Backup + from google.cloud.bigtable.backup import Backup from google.cloud._helpers import _datetime_to_pb_timestamp from google.rpc.code_pb2 import Code @@ -190,7 +190,7 @@ def test_backup_from_pb_success(): def test_backup_name(): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -225,7 +225,7 @@ def test_backup_parent_none(): def test_backup_parent_w_cluster(): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -242,7 +242,7 @@ def test_backup_parent_w_cluster(): def test_backup_source_table_none(): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -258,7 +258,7 @@ def test_backup_source_table_none(): def test_backup_source_table_valid(): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) @@ -473,7 +473,7 @@ def test_backup_create_w_expire_time_not_set(): def test_backup_create_success(): from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.bigtable_admin_v2.types import table - from google.cloud.bigtable.deprecated import Client + from google.cloud.bigtable import Client op_future = object() credentials = _make_credentials() @@ -806,12 +806,12 @@ def test_backup_restore_to_another_instance(): def test_backup_get_iam_policy(): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = Client(project=PROJECT_ID, credentials=credentials, admin=True) @@ -842,13 +842,13 @@ def test_backup_get_iam_policy(): def test_backup_set_iam_policy(): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = Client(project=PROJECT_ID, credentials=credentials, admin=True) @@ -887,7 +887,7 @@ def test_backup_set_iam_policy(): def test_backup_test_iam_permissions(): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client from google.cloud.bigtable_admin_v2.services.bigtable_table_admin import ( BigtableTableAdminClient, ) diff --git a/tests/unit/v2_client/test_batcher.py b/tests/unit/v2_client/test_batcher.py index 0793ed480..ab511e030 100644 --- a/tests/unit/v2_client/test_batcher.py +++ b/tests/unit/v2_client/test_batcher.py @@ -14,122 +14,139 @@ import mock +import time + import pytest -from google.cloud.bigtable.deprecated.row import DirectRow +from google.cloud.bigtable.row import DirectRow +from google.cloud.bigtable.batcher import ( + _FlowControl, + MutationsBatcher, + MutationsBatchError, +) TABLE_ID = "table-id" TABLE_NAME = "/tables/" + TABLE_ID -def _make_mutation_batcher(table, **kw): - from google.cloud.bigtable.deprecated.batcher import MutationsBatcher - - return MutationsBatcher(table, **kw) +def test_mutation_batcher_constructor(): + table = _Table(TABLE_NAME) + with MutationsBatcher(table) as mutation_batcher: + assert table is mutation_batcher.table -def test_mutation_batcher_constructor(): +def test_mutation_batcher_w_user_callback(): table = _Table(TABLE_NAME) - mutation_batcher = _make_mutation_batcher(table) - assert table is mutation_batcher.table + def callback_fn(response): + callback_fn.count = len(response) + + with MutationsBatcher( + table, flush_count=1, batch_completed_callback=callback_fn + ) as mutation_batcher: + rows = [ + DirectRow(row_key=b"row_key"), + DirectRow(row_key=b"row_key_2"), + DirectRow(row_key=b"row_key_3"), + DirectRow(row_key=b"row_key_4"), + ] + + mutation_batcher.mutate_rows(rows) + + assert callback_fn.count == 4 def test_mutation_batcher_mutate_row(): table = _Table(TABLE_NAME) - mutation_batcher = _make_mutation_batcher(table=table) + with MutationsBatcher(table=table) as mutation_batcher: - rows = [ - DirectRow(row_key=b"row_key"), - DirectRow(row_key=b"row_key_2"), - DirectRow(row_key=b"row_key_3"), - DirectRow(row_key=b"row_key_4"), - ] + rows = [ + DirectRow(row_key=b"row_key"), + DirectRow(row_key=b"row_key_2"), + DirectRow(row_key=b"row_key_3"), + DirectRow(row_key=b"row_key_4"), + ] - mutation_batcher.mutate_rows(rows) - mutation_batcher.flush() + mutation_batcher.mutate_rows(rows) assert table.mutation_calls == 1 def test_mutation_batcher_mutate(): table = _Table(TABLE_NAME) - mutation_batcher = _make_mutation_batcher(table=table) - - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", 1) - row.set_cell("cf1", b"c2", 2) - row.set_cell("cf1", b"c3", 3) - row.set_cell("cf1", b"c4", 4) + with MutationsBatcher(table=table) as mutation_batcher: - mutation_batcher.mutate(row) + row = DirectRow(row_key=b"row_key") + row.set_cell("cf1", b"c1", 1) + row.set_cell("cf1", b"c2", 2) + row.set_cell("cf1", b"c3", 3) + row.set_cell("cf1", b"c4", 4) - mutation_batcher.flush() + mutation_batcher.mutate(row) assert table.mutation_calls == 1 def test_mutation_batcher_flush_w_no_rows(): table = _Table(TABLE_NAME) - mutation_batcher = _make_mutation_batcher(table=table) - mutation_batcher.flush() + with MutationsBatcher(table=table) as mutation_batcher: + mutation_batcher.flush() assert table.mutation_calls == 0 def test_mutation_batcher_mutate_w_max_flush_count(): table = _Table(TABLE_NAME) - mutation_batcher = _make_mutation_batcher(table=table, flush_count=3) + with MutationsBatcher(table=table, flush_count=3) as mutation_batcher: - row_1 = DirectRow(row_key=b"row_key_1") - row_2 = DirectRow(row_key=b"row_key_2") - row_3 = DirectRow(row_key=b"row_key_3") + row_1 = DirectRow(row_key=b"row_key_1") + row_2 = DirectRow(row_key=b"row_key_2") + row_3 = DirectRow(row_key=b"row_key_3") - mutation_batcher.mutate(row_1) - mutation_batcher.mutate(row_2) - mutation_batcher.mutate(row_3) + mutation_batcher.mutate(row_1) + mutation_batcher.mutate(row_2) + mutation_batcher.mutate(row_3) assert table.mutation_calls == 1 -@mock.patch("google.cloud.bigtable.deprecated.batcher.MAX_MUTATIONS", new=3) -def test_mutation_batcher_mutate_with_max_mutations_failure(): - from google.cloud.bigtable.deprecated.batcher import MaxMutationsError - +@mock.patch("google.cloud.bigtable.batcher.MAX_OUTSTANDING_ELEMENTS", new=3) +def test_mutation_batcher_mutate_w_max_mutations(): table = _Table(TABLE_NAME) - mutation_batcher = _make_mutation_batcher(table=table) + with MutationsBatcher(table=table) as mutation_batcher: - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", 1) - row.set_cell("cf1", b"c2", 2) - row.set_cell("cf1", b"c3", 3) - row.set_cell("cf1", b"c4", 4) + row = DirectRow(row_key=b"row_key") + row.set_cell("cf1", b"c1", 1) + row.set_cell("cf1", b"c2", 2) + row.set_cell("cf1", b"c3", 3) - with pytest.raises(MaxMutationsError): mutation_batcher.mutate(row) + assert table.mutation_calls == 1 + -@mock.patch("google.cloud.bigtable.deprecated.batcher.MAX_MUTATIONS", new=3) -def test_mutation_batcher_mutate_w_max_mutations(): +def test_mutation_batcher_mutate_w_max_row_bytes(): table = _Table(TABLE_NAME) - mutation_batcher = _make_mutation_batcher(table=table) + with MutationsBatcher( + table=table, max_row_bytes=3 * 1024 * 1024 + ) as mutation_batcher: - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", 1) - row.set_cell("cf1", b"c2", 2) - row.set_cell("cf1", b"c3", 3) + number_of_bytes = 1 * 1024 * 1024 + max_value = b"1" * number_of_bytes - mutation_batcher.mutate(row) - mutation_batcher.flush() + row = DirectRow(row_key=b"row_key") + row.set_cell("cf1", b"c1", max_value) + row.set_cell("cf1", b"c2", max_value) + row.set_cell("cf1", b"c3", max_value) + + mutation_batcher.mutate(row) assert table.mutation_calls == 1 -def test_mutation_batcher_mutate_w_max_row_bytes(): +def test_mutations_batcher_flushed_when_closed(): table = _Table(TABLE_NAME) - mutation_batcher = _make_mutation_batcher( - table=table, max_row_bytes=3 * 1024 * 1024 - ) + mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024) number_of_bytes = 1 * 1024 * 1024 max_value = b"1" * number_of_bytes @@ -137,13 +154,107 @@ def test_mutation_batcher_mutate_w_max_row_bytes(): row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", max_value) row.set_cell("cf1", b"c2", max_value) - row.set_cell("cf1", b"c3", max_value) mutation_batcher.mutate(row) + assert table.mutation_calls == 0 + + mutation_batcher.close() assert table.mutation_calls == 1 +def test_mutations_batcher_context_manager_flushed_when_closed(): + table = _Table(TABLE_NAME) + with MutationsBatcher( + table=table, max_row_bytes=3 * 1024 * 1024 + ) as mutation_batcher: + + number_of_bytes = 1 * 1024 * 1024 + max_value = b"1" * number_of_bytes + + row = DirectRow(row_key=b"row_key") + row.set_cell("cf1", b"c1", max_value) + row.set_cell("cf1", b"c2", max_value) + + mutation_batcher.mutate(row) + + assert table.mutation_calls == 1 + + +@mock.patch("google.cloud.bigtable.batcher.MutationsBatcher.flush") +def test_mutations_batcher_flush_interval(mocked_flush): + table = _Table(TABLE_NAME) + flush_interval = 0.5 + mutation_batcher = MutationsBatcher(table=table, flush_interval=flush_interval) + + assert mutation_batcher._timer.interval == flush_interval + mocked_flush.assert_not_called() + + time.sleep(0.4) + mocked_flush.assert_not_called() + + time.sleep(0.1) + mocked_flush.assert_called_once_with() + + mutation_batcher.close() + + +def test_mutations_batcher_response_with_error_codes(): + from google.rpc.status_pb2 import Status + + mocked_response = [Status(code=1), Status(code=5)] + + table = mock.Mock() + mutation_batcher = MutationsBatcher(table=table) + + row1 = DirectRow(row_key=b"row_key") + row2 = DirectRow(row_key=b"row_key") + table.mutate_rows.return_value = mocked_response + + mutation_batcher.mutate_rows([row1, row2]) + with pytest.raises(MutationsBatchError) as exc: + mutation_batcher.close() + assert exc.value.message == "Errors in batch mutations." + assert len(exc.value.exc) == 2 + + assert exc.value.exc[0].message == mocked_response[0].message + assert exc.value.exc[1].message == mocked_response[1].message + + +def test_flow_control_event_is_set_when_not_blocked(): + flow_control = _FlowControl() + + flow_control.set_flow_control_status() + assert flow_control.event.is_set() + + +def test_flow_control_event_is_not_set_when_blocked(): + flow_control = _FlowControl() + + flow_control.inflight_mutations = flow_control.max_mutations + flow_control.inflight_size = flow_control.max_mutation_bytes + + flow_control.set_flow_control_status() + assert not flow_control.event.is_set() + + +@mock.patch("concurrent.futures.ThreadPoolExecutor.submit") +def test_flush_async_batch_count(mocked_executor_submit): + table = _Table(TABLE_NAME) + mutation_batcher = MutationsBatcher(table=table, flush_count=2) + + number_of_bytes = 1 * 1024 * 1024 + max_value = b"1" * number_of_bytes + for index in range(5): + row = DirectRow(row_key=f"row_key_{index}") + row.set_cell("cf1", b"c1", max_value) + mutation_batcher.mutate(row) + mutation_batcher._flush_async() + + # 3 batches submitted. 2 batches of 2 items, and the last one a single item batch. + assert mocked_executor_submit.call_count == 3 + + class _Instance(object): def __init__(self, client=None): self._client = client @@ -156,5 +267,8 @@ def __init__(self, name, client=None): self.mutation_calls = 0 def mutate_rows(self, rows): + from google.rpc.status_pb2 import Status + self.mutation_calls += 1 - return rows + + return [Status(code=0) for _ in rows] diff --git a/tests/unit/v2_client/test_client.py b/tests/unit/v2_client/test_client.py index 9deac6a25..5944c58a3 100644 --- a/tests/unit/v2_client/test_client.py +++ b/tests/unit/v2_client/test_client.py @@ -25,7 +25,7 @@ def _invoke_client_factory(client_class, **kw): - from google.cloud.bigtable.deprecated.client import _create_gapic_client + from google.cloud.bigtable.client import _create_gapic_client return _create_gapic_client(client_class, **kw) @@ -101,27 +101,23 @@ def __init__(self, credentials, emulator_host=None, emulator_channel=None): def _make_client(*args, **kwargs): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client return Client(*args, **kwargs) @mock.patch("os.environ", {}) def test_client_constructor_defaults(): - import warnings from google.api_core import client_info - from google.cloud.bigtable.deprecated import __version__ - from google.cloud.bigtable.deprecated.client import DATA_SCOPE + from google.cloud.bigtable import __version__ + from google.cloud.bigtable.client import DATA_SCOPE credentials = _make_credentials() - with warnings.catch_warnings(record=True) as warned: - with mock.patch("google.auth.default") as mocked: - mocked.return_value = credentials, PROJECT - client = _make_client() + with mock.patch("google.auth.default") as mocked: + mocked.return_value = credentials, PROJECT + client = _make_client() - # warn about client deprecation - assert len(warned) == 1 assert client.project == PROJECT assert client._credentials is credentials.with_scopes.return_value assert not client._read_only @@ -135,8 +131,8 @@ def test_client_constructor_defaults(): def test_client_constructor_explicit(): import warnings - from google.cloud.bigtable.deprecated.client import ADMIN_SCOPE - from google.cloud.bigtable.deprecated.client import DATA_SCOPE + from google.cloud.bigtable.client import ADMIN_SCOPE + from google.cloud.bigtable.client import DATA_SCOPE credentials = _make_credentials() client_info = mock.Mock() @@ -151,8 +147,7 @@ def test_client_constructor_explicit(): channel=mock.sentinel.channel, ) - # deprecationw arnning for channel and Client deprecation - assert len(warned) == 2 + assert len(warned) == 1 assert client.project == PROJECT assert client._credentials is credentials.with_scopes.return_value @@ -176,10 +171,8 @@ def test_client_constructor_w_both_admin_and_read_only(): def test_client_constructor_w_emulator_host(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.deprecated.client import ( - _DEFAULT_BIGTABLE_EMULATOR_CLIENT, - ) - from google.cloud.bigtable.deprecated.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT + from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS emulator_host = "localhost:8081" with mock.patch("os.environ", {BIGTABLE_EMULATOR: emulator_host}): @@ -202,7 +195,7 @@ def test_client_constructor_w_emulator_host(): def test_client_constructor_w_emulator_host_w_project(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.deprecated.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS emulator_host = "localhost:8081" with mock.patch("os.environ", {BIGTABLE_EMULATOR: emulator_host}): @@ -223,10 +216,8 @@ def test_client_constructor_w_emulator_host_w_project(): def test_client_constructor_w_emulator_host_w_credentials(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.deprecated.client import ( - _DEFAULT_BIGTABLE_EMULATOR_CLIENT, - ) - from google.cloud.bigtable.deprecated.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT + from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS emulator_host = "localhost:8081" credentials = _make_credentials() @@ -247,15 +238,15 @@ def test_client_constructor_w_emulator_host_w_credentials(): def test_client__get_scopes_default(): - from google.cloud.bigtable.deprecated.client import DATA_SCOPE + from google.cloud.bigtable.client import DATA_SCOPE client = _make_client(project=PROJECT, credentials=_make_credentials()) assert client._get_scopes() == (DATA_SCOPE,) def test_client__get_scopes_w_admin(): - from google.cloud.bigtable.deprecated.client import ADMIN_SCOPE - from google.cloud.bigtable.deprecated.client import DATA_SCOPE + from google.cloud.bigtable.client import ADMIN_SCOPE + from google.cloud.bigtable.client import DATA_SCOPE client = _make_client(project=PROJECT, credentials=_make_credentials(), admin=True) expected_scopes = (DATA_SCOPE, ADMIN_SCOPE) @@ -263,7 +254,7 @@ def test_client__get_scopes_w_admin(): def test_client__get_scopes_w_read_only(): - from google.cloud.bigtable.deprecated.client import READ_ONLY_SCOPE + from google.cloud.bigtable.client import READ_ONLY_SCOPE client = _make_client( project=PROJECT, credentials=_make_credentials(), read_only=True @@ -353,7 +344,7 @@ def test_client__local_composite_credentials(): def _create_gapic_client_channel_helper(endpoint=None, emulator_host=None): - from google.cloud.bigtable.deprecated.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS client_class = mock.Mock(spec=["DEFAULT_ENDPOINT"]) credentials = _make_credentials() @@ -627,7 +618,7 @@ def test_client_instance_admin_client_initialized(): def test_client_instance_factory_defaults(): - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials) @@ -643,8 +634,8 @@ def test_client_instance_factory_defaults(): def test_client_instance_factory_non_defaults(): - from google.cloud.bigtable.deprecated.instance import Instance - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable import enums instance_type = enums.Instance.Type.DEVELOPMENT labels = {"foo": "bar"} @@ -674,7 +665,7 @@ def test_client_list_instances(): from google.cloud.bigtable_admin_v2.services.bigtable_instance_admin import ( BigtableInstanceAdminClient, ) - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance FAILED_LOCATION = "FAILED" INSTANCE_ID1 = "instance-id1" @@ -726,7 +717,7 @@ def test_client_list_clusters(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.instance import Cluster + from google.cloud.bigtable.instance import Cluster instance_api = mock.create_autospec(BigtableInstanceAdminClient) diff --git a/tests/unit/v2_client/test_cluster.py b/tests/unit/v2_client/test_cluster.py index e667c2af4..cb0312b0c 100644 --- a/tests/unit/v2_client/test_cluster.py +++ b/tests/unit/v2_client/test_cluster.py @@ -42,13 +42,13 @@ def _make_cluster(*args, **kwargs): - from google.cloud.bigtable.deprecated.cluster import Cluster + from google.cloud.bigtable.cluster import Cluster return Cluster(*args, **kwargs) def _make_client(*args, **kwargs): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client return Client(*args, **kwargs) @@ -72,8 +72,8 @@ def test_cluster_constructor_defaults(): def test_cluster_constructor_explicit(): - from google.cloud.bigtable.deprecated.enums import StorageType - from google.cloud.bigtable.deprecated.enums import Cluster + from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.enums import Cluster STATE = Cluster.State.READY STORAGE_TYPE_SSD = StorageType.SSD @@ -126,8 +126,8 @@ def test_cluster_kms_key_name_setter(): def test_cluster_from_pb_success(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.cluster import Cluster - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable.cluster import Cluster + from google.cloud.bigtable import enums client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -162,7 +162,7 @@ def test_cluster_from_pb_success(): def test_cluster_from_pb_w_bad_cluster_name(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.cluster import Cluster + from google.cloud.bigtable.cluster import Cluster bad_cluster_name = "BAD_NAME" @@ -174,7 +174,7 @@ def test_cluster_from_pb_w_bad_cluster_name(): def test_cluster_from_pb_w_instance_id_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.cluster import Cluster + from google.cloud.bigtable.cluster import Cluster ALT_INSTANCE_ID = "ALT_INSTANCE_ID" client = _Client(PROJECT) @@ -189,7 +189,7 @@ def test_cluster_from_pb_w_instance_id_mistmatch(): def test_cluster_from_pb_w_project_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.cluster import Cluster + from google.cloud.bigtable.cluster import Cluster ALT_PROJECT = "ALT_PROJECT" client = _Client(project=ALT_PROJECT) @@ -204,8 +204,8 @@ def test_cluster_from_pb_w_project_mistmatch(): def test_cluster_from_pb_w_autoscaling(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.cluster import Cluster - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable.cluster import Cluster + from google.cloud.bigtable import enums client = _Client(PROJECT) instance = _Instance(INSTANCE_ID, client) @@ -292,8 +292,8 @@ def _make_instance_admin_client(): def test_cluster_reload(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.enums import StorageType - from google.cloud.bigtable.deprecated.enums import Cluster + from google.cloud.bigtable.enums import StorageType + from google.cloud.bigtable.enums import Cluster credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -349,7 +349,7 @@ def test_cluster_reload(): def test_cluster_exists_hit(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -371,7 +371,7 @@ def test_cluster_exists_hit(): def test_cluster_exists_miss(): - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance from google.api_core import exceptions credentials = _make_credentials() @@ -390,7 +390,7 @@ def test_cluster_exists_miss(): def test_cluster_exists_w_error(): - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance from google.api_core import exceptions credentials = _make_credentials() @@ -416,9 +416,9 @@ def test_cluster_create(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2 - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -471,9 +471,9 @@ def test_cluster_create_w_cmek(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2 - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -531,9 +531,9 @@ def test_cluster_create_w_autoscaling(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance from google.cloud.bigtable_admin_v2.types import instance as instance_v2_pb2 - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -600,7 +600,7 @@ def test_cluster_update(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -667,7 +667,7 @@ def test_cluster_update_w_autoscaling(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -726,7 +726,7 @@ def test_cluster_update_w_partial_autoscaling_config(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -811,7 +811,7 @@ def test_cluster_update_w_both_manual_and_autoscaling(): from google.cloud.bigtable_admin_v2.types import ( bigtable_instance_admin as messages_v2_pb2, ) - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -871,8 +871,8 @@ def test_cluster_disable_autoscaling(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud._helpers import _datetime_to_pb_timestamp - from google.cloud.bigtable.deprecated.instance import Instance - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.enums import StorageType NOW = datetime.datetime.utcnow() NOW_PB = _datetime_to_pb_timestamp(NOW) @@ -928,8 +928,8 @@ def test_cluster_disable_autoscaling(): def test_create_cluster_with_both_manual_and_autoscaling(): - from google.cloud.bigtable.deprecated.instance import Instance - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.enums import StorageType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -956,8 +956,8 @@ def test_create_cluster_with_both_manual_and_autoscaling(): def test_create_cluster_with_partial_autoscaling_config(): - from google.cloud.bigtable.deprecated.instance import Instance - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.enums import StorageType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -997,8 +997,8 @@ def test_create_cluster_with_partial_autoscaling_config(): def test_create_cluster_with_no_scaling_config(): - from google.cloud.bigtable.deprecated.instance import Instance - from google.cloud.bigtable.deprecated.enums import StorageType + from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.enums import StorageType credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) diff --git a/tests/unit/v2_client/test_column_family.py b/tests/unit/v2_client/test_column_family.py index d16d2b20c..b164b2fc1 100644 --- a/tests/unit/v2_client/test_column_family.py +++ b/tests/unit/v2_client/test_column_family.py @@ -19,7 +19,7 @@ def _make_max_versions_gc_rule(*args, **kwargs): - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule return MaxVersionsGCRule(*args, **kwargs) @@ -51,7 +51,7 @@ def test_max_versions_gc_rule_to_pb(): def _make_max_age_gc_rule(*args, **kwargs): - from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule + from google.cloud.bigtable.column_family import MaxAgeGCRule return MaxAgeGCRule(*args, **kwargs) @@ -89,7 +89,7 @@ def test_max_age_gc_rule_to_pb(): def _make_gc_rule_union(*args, **kwargs): - from google.cloud.bigtable.deprecated.column_family import GCRuleUnion + from google.cloud.bigtable.column_family import GCRuleUnion return GCRuleUnion(*args, **kwargs) @@ -124,8 +124,8 @@ def test_gc_rule_union___ne__same_value(): def test_gc_rule_union_to_pb(): import datetime from google.protobuf import duration_pb2 - from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxAgeGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule max_num_versions = 42 rule1 = MaxVersionsGCRule(max_num_versions) @@ -145,8 +145,8 @@ def test_gc_rule_union_to_pb(): def test_gc_rule_union_to_pb_nested(): import datetime from google.protobuf import duration_pb2 - from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxAgeGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule max_num_versions1 = 42 rule1 = MaxVersionsGCRule(max_num_versions1) @@ -171,7 +171,7 @@ def test_gc_rule_union_to_pb_nested(): def _make_gc_rule_intersection(*args, **kwargs): - from google.cloud.bigtable.deprecated.column_family import GCRuleIntersection + from google.cloud.bigtable.column_family import GCRuleIntersection return GCRuleIntersection(*args, **kwargs) @@ -206,8 +206,8 @@ def test_gc_rule_intersection___ne__same_value(): def test_gc_rule_intersection_to_pb(): import datetime from google.protobuf import duration_pb2 - from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxAgeGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule max_num_versions = 42 rule1 = MaxVersionsGCRule(max_num_versions) @@ -227,8 +227,8 @@ def test_gc_rule_intersection_to_pb(): def test_gc_rule_intersection_to_pb_nested(): import datetime from google.protobuf import duration_pb2 - from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxAgeGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule max_num_versions1 = 42 rule1 = MaxVersionsGCRule(max_num_versions1) @@ -253,13 +253,13 @@ def test_gc_rule_intersection_to_pb_nested(): def _make_column_family(*args, **kwargs): - from google.cloud.bigtable.deprecated.column_family import ColumnFamily + from google.cloud.bigtable.column_family import ColumnFamily return ColumnFamily(*args, **kwargs) def _make_client(*args, **kwargs): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client return Client(*args, **kwargs) @@ -323,7 +323,7 @@ def test_column_family_to_pb_no_rules(): def test_column_family_to_pb_with_rule(): - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule gc_rule = MaxVersionsGCRule(1) column_family = _make_column_family("column_family_id", None, gc_rule=gc_rule) @@ -397,7 +397,7 @@ def test_column_family_create(): def test_column_family_create_with_gc_rule(): - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule gc_rule = MaxVersionsGCRule(1337) _create_test_helper(gc_rule=gc_rule) @@ -467,7 +467,7 @@ def test_column_family_update(): def test_column_family_update_with_gc_rule(): - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule gc_rule = MaxVersionsGCRule(1337) _update_test_helper(gc_rule=gc_rule) @@ -530,15 +530,15 @@ def test_column_family_delete(): def test__gc_rule_from_pb_empty(): - from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb + from google.cloud.bigtable.column_family import _gc_rule_from_pb gc_rule_pb = _GcRulePB() assert _gc_rule_from_pb(gc_rule_pb) is None def test__gc_rule_from_pb_max_num_versions(): - from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import _gc_rule_from_pb + from google.cloud.bigtable.column_family import MaxVersionsGCRule orig_rule = MaxVersionsGCRule(1) gc_rule_pb = orig_rule.to_pb() @@ -549,8 +549,8 @@ def test__gc_rule_from_pb_max_num_versions(): def test__gc_rule_from_pb_max_age(): import datetime - from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb - from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule + from google.cloud.bigtable.column_family import _gc_rule_from_pb + from google.cloud.bigtable.column_family import MaxAgeGCRule orig_rule = MaxAgeGCRule(datetime.timedelta(seconds=1)) gc_rule_pb = orig_rule.to_pb() @@ -561,10 +561,10 @@ def test__gc_rule_from_pb_max_age(): def test__gc_rule_from_pb_union(): import datetime - from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb - from google.cloud.bigtable.deprecated.column_family import GCRuleUnion - from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import _gc_rule_from_pb + from google.cloud.bigtable.column_family import GCRuleUnion + from google.cloud.bigtable.column_family import MaxAgeGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule rule1 = MaxVersionsGCRule(1) rule2 = MaxAgeGCRule(datetime.timedelta(seconds=1)) @@ -577,10 +577,10 @@ def test__gc_rule_from_pb_union(): def test__gc_rule_from_pb_intersection(): import datetime - from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb - from google.cloud.bigtable.deprecated.column_family import GCRuleIntersection - from google.cloud.bigtable.deprecated.column_family import MaxAgeGCRule - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import _gc_rule_from_pb + from google.cloud.bigtable.column_family import GCRuleIntersection + from google.cloud.bigtable.column_family import MaxAgeGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule rule1 = MaxVersionsGCRule(1) rule2 = MaxAgeGCRule(datetime.timedelta(seconds=1)) @@ -592,7 +592,7 @@ def test__gc_rule_from_pb_intersection(): def test__gc_rule_from_pb_unknown_field_name(): - from google.cloud.bigtable.deprecated.column_family import _gc_rule_from_pb + from google.cloud.bigtable.column_family import _gc_rule_from_pb class MockProto(object): diff --git a/tests/unit/v2_client/test_encryption_info.py b/tests/unit/v2_client/test_encryption_info.py index 0b6a93e9e..8b92a83ed 100644 --- a/tests/unit/v2_client/test_encryption_info.py +++ b/tests/unit/v2_client/test_encryption_info.py @@ -14,7 +14,7 @@ import mock -from google.cloud.bigtable.deprecated import enums +from google.cloud.bigtable import enums EncryptionType = enums.EncryptionInfo.EncryptionType @@ -30,7 +30,7 @@ def _make_status_pb(code=_STATUS_CODE, message=_STATUS_MESSAGE): def _make_status(code=_STATUS_CODE, message=_STATUS_MESSAGE): - from google.cloud.bigtable.deprecated.error import Status + from google.cloud.bigtable.error import Status status_pb = _make_status_pb(code=code, message=message) return Status(status_pb) @@ -54,7 +54,7 @@ def _make_info_pb( def _make_encryption_info(*args, **kwargs): - from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo + from google.cloud.bigtable.encryption_info import EncryptionInfo return EncryptionInfo(*args, **kwargs) @@ -70,7 +70,7 @@ def _make_encryption_info_defaults( def test_encryption_info__from_pb(): - from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo + from google.cloud.bigtable.encryption_info import EncryptionInfo info_pb = _make_info_pb() diff --git a/tests/unit/v2_client/test_error.py b/tests/unit/v2_client/test_error.py index 072a3b3c3..8b148473c 100644 --- a/tests/unit/v2_client/test_error.py +++ b/tests/unit/v2_client/test_error.py @@ -20,7 +20,7 @@ def _make_status_pb(**kwargs): def _make_status(status_pb): - from google.cloud.bigtable.deprecated.error import Status + from google.cloud.bigtable.error import Status return Status(status_pb) diff --git a/tests/unit/v2_client/test_instance.py b/tests/unit/v2_client/test_instance.py index b43e8bb38..c577adca5 100644 --- a/tests/unit/v2_client/test_instance.py +++ b/tests/unit/v2_client/test_instance.py @@ -17,7 +17,7 @@ import pytest from ._testing import _make_credentials -from google.cloud.bigtable.deprecated.cluster import Cluster +from google.cloud.bigtable.cluster import Cluster PROJECT = "project" INSTANCE_ID = "instance-id" @@ -47,7 +47,7 @@ def _make_client(*args, **kwargs): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client return Client(*args, **kwargs) @@ -61,7 +61,7 @@ def _make_instance_admin_api(): def _make_instance(*args, **kwargs): - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance return Instance(*args, **kwargs) @@ -79,7 +79,7 @@ def test_instance_constructor_defaults(): def test_instance_constructor_non_default(): - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums instance_type = enums.Instance.Type.DEVELOPMENT state = enums.Instance.State.READY @@ -104,7 +104,7 @@ def test_instance_constructor_non_default(): def test_instance__update_from_pb_success(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums instance_type = data_v2_pb2.Instance.Type.PRODUCTION state = enums.Instance.State.READY @@ -129,7 +129,7 @@ def test_instance__update_from_pb_success(): def test_instance__update_from_pb_success_defaults(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums instance_pb = data_v2_pb2.Instance(display_name=DISPLAY_NAME) @@ -156,8 +156,8 @@ def test_instance__update_from_pb_wo_display_name(): def test_instance_from_pb_success(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated import enums - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable import enums + from google.cloud.bigtable.instance import Instance credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -184,7 +184,7 @@ def test_instance_from_pb_success(): def test_instance_from_pb_bad_instance_name(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance instance_name = "INCORRECT_FORMAT" instance_pb = data_v2_pb2.Instance(name=instance_name) @@ -195,7 +195,7 @@ def test_instance_from_pb_bad_instance_name(): def test_instance_from_pb_project_mistmatch(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance ALT_PROJECT = "ALT_PROJECT" credentials = _make_credentials() @@ -304,7 +304,7 @@ def _instance_api_response_for_create(): def test_instance_create(): - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums from google.cloud.bigtable_admin_v2.types import Instance from google.cloud.bigtable_admin_v2.types import Cluster import warnings @@ -353,8 +353,8 @@ def test_instance_create(): def test_instance_create_w_clusters(): - from google.cloud.bigtable.deprecated import enums - from google.cloud.bigtable.deprecated.cluster import Cluster + from google.cloud.bigtable import enums + from google.cloud.bigtable.cluster import Cluster from google.cloud.bigtable_admin_v2.types import Cluster as cluster_pb from google.cloud.bigtable_admin_v2.types import Instance as instance_pb @@ -473,7 +473,7 @@ def test_instance_exists_w_error(): def test_instance_reload(): from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums DISPLAY_NAME = "hey-hi-hello" credentials = _make_credentials() @@ -527,7 +527,7 @@ def _instance_api_response_for_update(): def test_instance_update(): - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums from google.protobuf import field_mask_pb2 from google.cloud.bigtable_admin_v2.types import Instance @@ -603,7 +603,7 @@ def test_instance_delete(): def test_instance_get_iam_policy(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -631,7 +631,7 @@ def test_instance_get_iam_policy(): def test_instance_get_iam_policy_w_requested_policy_version(): from google.iam.v1 import policy_pb2, options_pb2 - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -665,8 +665,8 @@ def test_instance_get_iam_policy_w_requested_policy_version(): def test_instance_set_iam_policy(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -721,7 +721,7 @@ def test_instance_test_iam_permissions(): def test_instance_cluster_factory(): - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums CLUSTER_ID = "{}-cluster".format(INSTANCE_ID) LOCATION_ID = "us-central1-c" @@ -749,8 +749,8 @@ def test_instance_list_clusters(): bigtable_instance_admin as messages_v2_pb2, ) from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.instance import Instance - from google.cloud.bigtable.deprecated.instance import Cluster + from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.instance import Cluster credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) @@ -788,7 +788,7 @@ def test_instance_list_clusters(): def test_instance_table_factory(): - from google.cloud.bigtable.deprecated.table import Table + from google.cloud.bigtable.table import Table app_profile_id = "appProfileId1262094415" instance = _make_instance(INSTANCE_ID, None) @@ -857,7 +857,7 @@ def test_instance_list_tables_failure_name_bad_before(): def test_instance_app_profile_factory(): - from google.cloud.bigtable.deprecated.enums import RoutingPolicyType + from google.cloud.bigtable.enums import RoutingPolicyType instance = _make_instance(INSTANCE_ID, None) @@ -890,7 +890,7 @@ def test_instance_list_app_profiles(): from google.api_core.page_iterator import Iterator from google.api_core.page_iterator import Page from google.cloud.bigtable_admin_v2.types import instance as data_v2_pb2 - from google.cloud.bigtable.deprecated.app_profile import AppProfile + from google.cloud.bigtable.app_profile import AppProfile class _Iterator(Iterator): def __init__(self, pages): diff --git a/tests/unit/v2_client/test_policy.py b/tests/unit/v2_client/test_policy.py index ef3df2d2b..77674517e 100644 --- a/tests/unit/v2_client/test_policy.py +++ b/tests/unit/v2_client/test_policy.py @@ -14,7 +14,7 @@ def _make_policy(*args, **kw): - from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.policy import Policy return Policy(*args, **kw) @@ -48,7 +48,7 @@ def test_policy_ctor_explicit(): def test_policy_bigtable_admins(): - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE MEMBER = "user:phred@example.com" expected = frozenset([MEMBER]) @@ -58,7 +58,7 @@ def test_policy_bigtable_admins(): def test_policy_bigtable_readers(): - from google.cloud.bigtable.deprecated.policy import BIGTABLE_READER_ROLE + from google.cloud.bigtable.policy import BIGTABLE_READER_ROLE MEMBER = "user:phred@example.com" expected = frozenset([MEMBER]) @@ -68,7 +68,7 @@ def test_policy_bigtable_readers(): def test_policy_bigtable_users(): - from google.cloud.bigtable.deprecated.policy import BIGTABLE_USER_ROLE + from google.cloud.bigtable.policy import BIGTABLE_USER_ROLE MEMBER = "user:phred@example.com" expected = frozenset([MEMBER]) @@ -78,7 +78,7 @@ def test_policy_bigtable_users(): def test_policy_bigtable_viewers(): - from google.cloud.bigtable.deprecated.policy import BIGTABLE_VIEWER_ROLE + from google.cloud.bigtable.policy import BIGTABLE_VIEWER_ROLE MEMBER = "user:phred@example.com" expected = frozenset([MEMBER]) @@ -89,7 +89,7 @@ def test_policy_bigtable_viewers(): def test_policy_from_pb_w_empty(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.policy import Policy empty = frozenset() message = policy_pb2.Policy() @@ -106,8 +106,8 @@ def test_policy_from_pb_w_empty(): def test_policy_from_pb_w_non_empty(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE - from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import Policy ETAG = b"ETAG" VERSION = 1 @@ -133,8 +133,8 @@ def test_policy_from_pb_w_condition(): import pytest from google.iam.v1 import policy_pb2 from google.api_core.iam import InvalidOperationException, _DICT_ACCESS_MSG - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE - from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import Policy ETAG = b"ETAG" VERSION = 3 @@ -184,7 +184,7 @@ def test_policy_to_pb_empty(): def test_policy_to_pb_explicit(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE VERSION = 1 ETAG = b"ETAG" @@ -204,7 +204,7 @@ def test_policy_to_pb_explicit(): def test_policy_to_pb_w_condition(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE VERSION = 3 ETAG = b"ETAG" @@ -234,7 +234,7 @@ def test_policy_to_pb_w_condition(): def test_policy_from_api_repr_wo_etag(): - from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.policy import Policy VERSION = 1 empty = frozenset() @@ -252,7 +252,7 @@ def test_policy_from_api_repr_wo_etag(): def test_policy_from_api_repr_w_etag(): import base64 - from google.cloud.bigtable.deprecated.policy import Policy + from google.cloud.bigtable.policy import Policy ETAG = b"ETAG" empty = frozenset() diff --git a/tests/unit/v2_client/test_row.py b/tests/unit/v2_client/test_row.py index 4850b18c3..f04802f5c 100644 --- a/tests/unit/v2_client/test_row.py +++ b/tests/unit/v2_client/test_row.py @@ -20,13 +20,13 @@ def _make_client(*args, **kwargs): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client return Client(*args, **kwargs) def _make_row(*args, **kwargs): - from google.cloud.bigtable.deprecated.row import Row + from google.cloud.bigtable.row import Row return Row(*args, **kwargs) @@ -42,7 +42,7 @@ def test_row_table_getter(): def _make__set_delete_row(*args, **kwargs): - from google.cloud.bigtable.deprecated.row import _SetDeleteRow + from google.cloud.bigtable.row import _SetDeleteRow return _SetDeleteRow(*args, **kwargs) @@ -54,7 +54,7 @@ def test__set_detlete_row__get_mutations_virtual(): def _make_direct_row(*args, **kwargs): - from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.row import DirectRow return DirectRow(*args, **kwargs) @@ -193,7 +193,7 @@ def test_direct_row_delete(): def test_direct_row_delete_cell(): - from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.row import DirectRow class MockRow(DirectRow): def __init__(self, *args, **kwargs): @@ -237,7 +237,7 @@ def test_direct_row_delete_cells_non_iterable(): def test_direct_row_delete_cells_all_columns(): - from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.row import DirectRow row_key = b"row_key" column_family_id = "column_family_id" @@ -293,7 +293,7 @@ def test_direct_row_delete_cells_no_time_range(): def test_direct_row_delete_cells_with_time_range(): import datetime from google.cloud._helpers import _EPOCH - from google.cloud.bigtable.deprecated.row_filters import TimestampRange + from google.cloud.bigtable.row_filters import TimestampRange microseconds = 30871000 # Makes sure already milliseconds granularity start = _EPOCH + datetime.timedelta(microseconds=microseconds) @@ -386,7 +386,7 @@ def test_direct_row_commit_with_exception(): def _make_conditional_row(*args, **kwargs): - from google.cloud.bigtable.deprecated.row import ConditionalRow + from google.cloud.bigtable.row import ConditionalRow return ConditionalRow(*args, **kwargs) @@ -417,7 +417,7 @@ def test_conditional_row__get_mutations(): def test_conditional_row_commit(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter from google.cloud.bigtable_v2.services.bigtable import BigtableClient project_id = "project-id" @@ -466,7 +466,7 @@ def test_conditional_row_commit(): def test_conditional_row_commit_too_many_mutations(): from google.cloud._testing import _Monkey - from google.cloud.bigtable.deprecated import row as MUT + from google.cloud.bigtable import row as MUT row_key = b"row_key" table = object() @@ -504,7 +504,7 @@ def test_conditional_row_commit_no_mutations(): def _make_append_row(*args, **kwargs): - from google.cloud.bigtable.deprecated.row import AppendRow + from google.cloud.bigtable.row import AppendRow return AppendRow(*args, **kwargs) @@ -564,7 +564,7 @@ def test_append_row_increment_cell_value(): def test_append_row_commit(): from google.cloud._testing import _Monkey - from google.cloud.bigtable.deprecated import row as MUT + from google.cloud.bigtable import row as MUT from google.cloud.bigtable_v2.services.bigtable import BigtableClient project_id = "project-id" @@ -630,7 +630,7 @@ def test_append_row_commit_no_rules(): def test_append_row_commit_too_many_mutations(): from google.cloud._testing import _Monkey - from google.cloud.bigtable.deprecated import row as MUT + from google.cloud.bigtable import row as MUT row_key = b"row_key" table = object() @@ -644,7 +644,7 @@ def test_append_row_commit_too_many_mutations(): def test__parse_rmw_row_response(): from google.cloud._helpers import _datetime_from_microseconds - from google.cloud.bigtable.deprecated.row import _parse_rmw_row_response + from google.cloud.bigtable.row import _parse_rmw_row_response col_fam1 = "col-fam-id" col_fam2 = "col-fam-id2" @@ -700,7 +700,7 @@ def test__parse_rmw_row_response(): def test__parse_family_pb(): from google.cloud._helpers import _datetime_from_microseconds - from google.cloud.bigtable.deprecated.row import _parse_family_pb + from google.cloud.bigtable.row import _parse_family_pb col_fam1 = "col-fam-id" col_name1 = b"col-name1" diff --git a/tests/unit/v2_client/test_row_data.py b/tests/unit/v2_client/test_row_data.py index ee9b065c8..fba69ceba 100644 --- a/tests/unit/v2_client/test_row_data.py +++ b/tests/unit/v2_client/test_row_data.py @@ -27,7 +27,7 @@ def _make_cell(*args, **kwargs): - from google.cloud.bigtable.deprecated.row_data import Cell + from google.cloud.bigtable.row_data import Cell return Cell(*args, **kwargs) @@ -36,7 +36,7 @@ def _cell_from_pb_test_helper(labels=None): import datetime from google.cloud._helpers import _EPOCH from google.cloud.bigtable_v2.types import data as data_v2_pb2 - from google.cloud.bigtable.deprecated.row_data import Cell + from google.cloud.bigtable.row_data import Cell timestamp = _EPOCH + datetime.timedelta(microseconds=TIMESTAMP_MICROS) value = b"value-bytes" @@ -100,7 +100,7 @@ def test_cell___ne__(): def _make_partial_row_data(*args, **kwargs): - from google.cloud.bigtable.deprecated.row_data import PartialRowData + from google.cloud.bigtable.row_data import PartialRowData return PartialRowData(*args, **kwargs) @@ -288,7 +288,7 @@ def trailing_metadata(self): def test__retry_read_rows_exception_miss(): from google.api_core.exceptions import Conflict - from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception + from google.cloud.bigtable.row_data import _retry_read_rows_exception exception = Conflict("testing") assert not _retry_read_rows_exception(exception) @@ -296,7 +296,7 @@ def test__retry_read_rows_exception_miss(): def test__retry_read_rows_exception_service_unavailable(): from google.api_core.exceptions import ServiceUnavailable - from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception + from google.cloud.bigtable.row_data import _retry_read_rows_exception exception = ServiceUnavailable("testing") assert _retry_read_rows_exception(exception) @@ -304,7 +304,7 @@ def test__retry_read_rows_exception_service_unavailable(): def test__retry_read_rows_exception_deadline_exceeded(): from google.api_core.exceptions import DeadlineExceeded - from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception + from google.cloud.bigtable.row_data import _retry_read_rows_exception exception = DeadlineExceeded("testing") assert _retry_read_rows_exception(exception) @@ -312,7 +312,7 @@ def test__retry_read_rows_exception_deadline_exceeded(): def test__retry_read_rows_exception_internal_server_not_retriable(): from google.api_core.exceptions import InternalServerError - from google.cloud.bigtable.deprecated.row_data import ( + from google.cloud.bigtable.row_data import ( _retry_read_rows_exception, RETRYABLE_INTERNAL_ERROR_MESSAGES, ) @@ -325,7 +325,7 @@ def test__retry_read_rows_exception_internal_server_not_retriable(): def test__retry_read_rows_exception_internal_server_retriable(): from google.api_core.exceptions import InternalServerError - from google.cloud.bigtable.deprecated.row_data import ( + from google.cloud.bigtable.row_data import ( _retry_read_rows_exception, RETRYABLE_INTERNAL_ERROR_MESSAGES, ) @@ -337,7 +337,7 @@ def test__retry_read_rows_exception_internal_server_retriable(): def test__retry_read_rows_exception_miss_wrapped_in_grpc(): from google.api_core.exceptions import Conflict - from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception + from google.cloud.bigtable.row_data import _retry_read_rows_exception wrapped = Conflict("testing") exception = _make_grpc_call_error(wrapped) @@ -346,7 +346,7 @@ def test__retry_read_rows_exception_miss_wrapped_in_grpc(): def test__retry_read_rows_exception_service_unavailable_wrapped_in_grpc(): from google.api_core.exceptions import ServiceUnavailable - from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception + from google.cloud.bigtable.row_data import _retry_read_rows_exception wrapped = ServiceUnavailable("testing") exception = _make_grpc_call_error(wrapped) @@ -355,7 +355,7 @@ def test__retry_read_rows_exception_service_unavailable_wrapped_in_grpc(): def test__retry_read_rows_exception_deadline_exceeded_wrapped_in_grpc(): from google.api_core.exceptions import DeadlineExceeded - from google.cloud.bigtable.deprecated.row_data import _retry_read_rows_exception + from google.cloud.bigtable.row_data import _retry_read_rows_exception wrapped = DeadlineExceeded("testing") exception = _make_grpc_call_error(wrapped) @@ -363,7 +363,7 @@ def test__retry_read_rows_exception_deadline_exceeded_wrapped_in_grpc(): def _make_partial_rows_data(*args, **kwargs): - from google.cloud.bigtable.deprecated.row_data import PartialRowsData + from google.cloud.bigtable.row_data import PartialRowsData return PartialRowsData(*args, **kwargs) @@ -373,13 +373,13 @@ def _partial_rows_data_consume_all(yrd): def _make_client(*args, **kwargs): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client return Client(*args, **kwargs) def test_partial_rows_data_constructor(): - from google.cloud.bigtable.deprecated.row_data import DEFAULT_RETRY_READ_ROWS + from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS client = _Client() client._data_stub = mock.MagicMock() @@ -436,7 +436,7 @@ def fake_read(*args, **kwargs): def test_partial_rows_data_constructor_with_retry(): - from google.cloud.bigtable.deprecated.row_data import DEFAULT_RETRY_READ_ROWS + from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS client = _Client() client._data_stub = mock.MagicMock() @@ -446,7 +446,9 @@ def test_partial_rows_data_constructor_with_retry(): client._data_stub.ReadRows, request, retry ) partial_rows_data.read_method.assert_called_once_with( - request, timeout=DEFAULT_RETRY_READ_ROWS.deadline + 1 + request, + timeout=DEFAULT_RETRY_READ_ROWS.deadline + 1, + retry=DEFAULT_RETRY_READ_ROWS, ) assert partial_rows_data.request is request assert partial_rows_data.rows == {} @@ -644,7 +646,7 @@ def test_partial_rows_data_valid_last_scanned_row_key_on_start(): def test_partial_rows_data_invalid_empty_chunk(): - from google.cloud.bigtable.deprecated.row_data import InvalidChunk + from google.cloud.bigtable.row_data import InvalidChunk from google.cloud.bigtable_v2.services.bigtable import BigtableClient client = _Client() @@ -755,14 +757,14 @@ def test_partial_rows_data_yield_retry_rows_data(): def _make_read_rows_request_manager(*args, **kwargs): - from google.cloud.bigtable.deprecated.row_data import _ReadRowsRequestManager + from google.cloud.bigtable.row_data import _ReadRowsRequestManager return _ReadRowsRequestManager(*args, **kwargs) @pytest.fixture(scope="session") def rrrm_data(): - from google.cloud.bigtable.deprecated import row_set + from google.cloud.bigtable import row_set row_range1 = row_set.RowRange(b"row_key21", b"row_key29") row_range2 = row_set.RowRange(b"row_key31", b"row_key39") @@ -851,7 +853,7 @@ def test_RRRM__filter_row_ranges_all_ranges_already_read(rrrm_data): def test_RRRM__filter_row_ranges_all_ranges_already_read_open_closed(): - from google.cloud.bigtable.deprecated import row_set + from google.cloud.bigtable import row_set last_scanned_key = b"row_key54" @@ -895,7 +897,7 @@ def test_RRRM__filter_row_ranges_some_ranges_already_read(rrrm_data): def test_RRRM_build_updated_request(rrrm_data): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter from google.cloud.bigtable_v2 import types row_range1 = rrrm_data["row_range1"] @@ -944,7 +946,7 @@ def test_RRRM_build_updated_request_full_table(): def test_RRRM_build_updated_request_no_start_key(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter from google.cloud.bigtable_v2 import types row_filter = RowSampleFilter(0.33) @@ -972,7 +974,7 @@ def test_RRRM_build_updated_request_no_start_key(): def test_RRRM_build_updated_request_no_end_key(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter from google.cloud.bigtable_v2 import types row_filter = RowSampleFilter(0.33) @@ -998,7 +1000,7 @@ def test_RRRM_build_updated_request_no_end_key(): def test_RRRM_build_updated_request_rows(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter row_filter = RowSampleFilter(0.33) last_scanned_key = b"row_key4" @@ -1046,7 +1048,7 @@ def test_RRRM__key_already_read(): def test_RRRM__rows_limit_reached(): - from google.cloud.bigtable.deprecated.row_data import InvalidRetryRequest + from google.cloud.bigtable.row_data import InvalidRetryRequest last_scanned_key = b"row_key14" request = _ReadRowsRequestPB(table_name=TABLE_NAME) @@ -1059,7 +1061,7 @@ def test_RRRM__rows_limit_reached(): def test_RRRM_build_updated_request_last_row_read_raises_invalid_retry_request(): - from google.cloud.bigtable.deprecated.row_data import InvalidRetryRequest + from google.cloud.bigtable.row_data import InvalidRetryRequest last_scanned_key = b"row_key4" request = _ReadRowsRequestPB(table_name=TABLE_NAME) @@ -1073,8 +1075,8 @@ def test_RRRM_build_updated_request_last_row_read_raises_invalid_retry_request() def test_RRRM_build_updated_request_row_ranges_read_raises_invalid_retry_request(): - from google.cloud.bigtable.deprecated.row_data import InvalidRetryRequest - from google.cloud.bigtable.deprecated import row_set + from google.cloud.bigtable.row_data import InvalidRetryRequest + from google.cloud.bigtable import row_set row_range1 = row_set.RowRange(b"row_key21", b"row_key29") @@ -1095,7 +1097,7 @@ def test_RRRM_build_updated_request_row_ranges_read_raises_invalid_retry_request def test_RRRM_build_updated_request_row_ranges_valid(): - from google.cloud.bigtable.deprecated import row_set + from google.cloud.bigtable import row_set row_range1 = row_set.RowRange(b"row_key21", b"row_key29") @@ -1179,7 +1181,7 @@ def _ReadRowsResponseCellChunkPB(*args, **kw): def _make_cell_pb(value): - from google.cloud.bigtable.deprecated import row_data + from google.cloud.bigtable import row_data return row_data.Cell(value, TIMESTAMP_MICROS) diff --git a/tests/unit/v2_client/test_row_filters.py b/tests/unit/v2_client/test_row_filters.py index dfb16ba16..b312cb942 100644 --- a/tests/unit/v2_client/test_row_filters.py +++ b/tests/unit/v2_client/test_row_filters.py @@ -17,7 +17,7 @@ def test_bool_filter_constructor(): - from google.cloud.bigtable.deprecated.row_filters import _BoolFilter + from google.cloud.bigtable.row_filters import _BoolFilter flag = object() row_filter = _BoolFilter(flag) @@ -25,7 +25,7 @@ def test_bool_filter_constructor(): def test_bool_filter___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import _BoolFilter + from google.cloud.bigtable.row_filters import _BoolFilter flag = object() row_filter1 = _BoolFilter(flag) @@ -34,7 +34,7 @@ def test_bool_filter___eq__type_differ(): def test_bool_filter___eq__same_value(): - from google.cloud.bigtable.deprecated.row_filters import _BoolFilter + from google.cloud.bigtable.row_filters import _BoolFilter flag = object() row_filter1 = _BoolFilter(flag) @@ -43,7 +43,7 @@ def test_bool_filter___eq__same_value(): def test_bool_filter___ne__same_value(): - from google.cloud.bigtable.deprecated.row_filters import _BoolFilter + from google.cloud.bigtable.row_filters import _BoolFilter flag = object() row_filter1 = _BoolFilter(flag) @@ -52,7 +52,7 @@ def test_bool_filter___ne__same_value(): def test_sink_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import SinkFilter + from google.cloud.bigtable.row_filters import SinkFilter flag = True row_filter = SinkFilter(flag) @@ -62,7 +62,7 @@ def test_sink_filter_to_pb(): def test_pass_all_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import PassAllFilter + from google.cloud.bigtable.row_filters import PassAllFilter flag = True row_filter = PassAllFilter(flag) @@ -72,7 +72,7 @@ def test_pass_all_filter_to_pb(): def test_block_all_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import BlockAllFilter + from google.cloud.bigtable.row_filters import BlockAllFilter flag = True row_filter = BlockAllFilter(flag) @@ -82,7 +82,7 @@ def test_block_all_filter_to_pb(): def test_regex_filterconstructor(): - from google.cloud.bigtable.deprecated.row_filters import _RegexFilter + from google.cloud.bigtable.row_filters import _RegexFilter regex = b"abc" row_filter = _RegexFilter(regex) @@ -90,7 +90,7 @@ def test_regex_filterconstructor(): def test_regex_filterconstructor_non_bytes(): - from google.cloud.bigtable.deprecated.row_filters import _RegexFilter + from google.cloud.bigtable.row_filters import _RegexFilter regex = "abc" row_filter = _RegexFilter(regex) @@ -98,7 +98,7 @@ def test_regex_filterconstructor_non_bytes(): def test_regex_filter__eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import _RegexFilter + from google.cloud.bigtable.row_filters import _RegexFilter regex = b"def-rgx" row_filter1 = _RegexFilter(regex) @@ -107,7 +107,7 @@ def test_regex_filter__eq__type_differ(): def test_regex_filter__eq__same_value(): - from google.cloud.bigtable.deprecated.row_filters import _RegexFilter + from google.cloud.bigtable.row_filters import _RegexFilter regex = b"trex-regex" row_filter1 = _RegexFilter(regex) @@ -116,7 +116,7 @@ def test_regex_filter__eq__same_value(): def test_regex_filter__ne__same_value(): - from google.cloud.bigtable.deprecated.row_filters import _RegexFilter + from google.cloud.bigtable.row_filters import _RegexFilter regex = b"abc" row_filter1 = _RegexFilter(regex) @@ -125,7 +125,7 @@ def test_regex_filter__ne__same_value(): def test_row_key_regex_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import RowKeyRegexFilter + from google.cloud.bigtable.row_filters import RowKeyRegexFilter regex = b"row-key-regex" row_filter = RowKeyRegexFilter(regex) @@ -135,7 +135,7 @@ def test_row_key_regex_filter_to_pb(): def test_row_sample_filter_constructor(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter sample = object() row_filter = RowSampleFilter(sample) @@ -143,7 +143,7 @@ def test_row_sample_filter_constructor(): def test_row_sample_filter___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter sample = object() row_filter1 = RowSampleFilter(sample) @@ -152,7 +152,7 @@ def test_row_sample_filter___eq__type_differ(): def test_row_sample_filter___eq__same_value(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter sample = object() row_filter1 = RowSampleFilter(sample) @@ -161,7 +161,7 @@ def test_row_sample_filter___eq__same_value(): def test_row_sample_filter___ne__(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter sample = object() other_sample = object() @@ -171,7 +171,7 @@ def test_row_sample_filter___ne__(): def test_row_sample_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import RowSampleFilter sample = 0.25 row_filter = RowSampleFilter(sample) @@ -181,7 +181,7 @@ def test_row_sample_filter_to_pb(): def test_family_name_regex_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.row_filters import FamilyNameRegexFilter regex = "family-regex" row_filter = FamilyNameRegexFilter(regex) @@ -191,7 +191,7 @@ def test_family_name_regex_filter_to_pb(): def test_column_qualifier_regext_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import ColumnQualifierRegexFilter + from google.cloud.bigtable.row_filters import ColumnQualifierRegexFilter regex = b"column-regex" row_filter = ColumnQualifierRegexFilter(regex) @@ -201,7 +201,7 @@ def test_column_qualifier_regext_filter_to_pb(): def test_timestamp_range_constructor(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRange + from google.cloud.bigtable.row_filters import TimestampRange start = object() end = object() @@ -211,7 +211,7 @@ def test_timestamp_range_constructor(): def test_timestamp_range___eq__(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRange + from google.cloud.bigtable.row_filters import TimestampRange start = object() end = object() @@ -221,7 +221,7 @@ def test_timestamp_range___eq__(): def test_timestamp_range___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRange + from google.cloud.bigtable.row_filters import TimestampRange start = object() end = object() @@ -231,7 +231,7 @@ def test_timestamp_range___eq__type_differ(): def test_timestamp_range___ne__same_value(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRange + from google.cloud.bigtable.row_filters import TimestampRange start = object() end = object() @@ -243,7 +243,7 @@ def test_timestamp_range___ne__same_value(): def _timestamp_range_to_pb_helper(pb_kwargs, start=None, end=None): import datetime from google.cloud._helpers import _EPOCH - from google.cloud.bigtable.deprecated.row_filters import TimestampRange + from google.cloud.bigtable.row_filters import TimestampRange if start is not None: start = _EPOCH + datetime.timedelta(microseconds=start) @@ -291,7 +291,7 @@ def test_timestamp_range_to_pb_end_only(): def test_timestamp_range_filter_constructor(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter + from google.cloud.bigtable.row_filters import TimestampRangeFilter range_ = object() row_filter = TimestampRangeFilter(range_) @@ -299,7 +299,7 @@ def test_timestamp_range_filter_constructor(): def test_timestamp_range_filter___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter + from google.cloud.bigtable.row_filters import TimestampRangeFilter range_ = object() row_filter1 = TimestampRangeFilter(range_) @@ -308,7 +308,7 @@ def test_timestamp_range_filter___eq__type_differ(): def test_timestamp_range_filter___eq__same_value(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter + from google.cloud.bigtable.row_filters import TimestampRangeFilter range_ = object() row_filter1 = TimestampRangeFilter(range_) @@ -317,7 +317,7 @@ def test_timestamp_range_filter___eq__same_value(): def test_timestamp_range_filter___ne__(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter + from google.cloud.bigtable.row_filters import TimestampRangeFilter range_ = object() other_range_ = object() @@ -327,8 +327,8 @@ def test_timestamp_range_filter___ne__(): def test_timestamp_range_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import TimestampRangeFilter - from google.cloud.bigtable.deprecated.row_filters import TimestampRange + from google.cloud.bigtable.row_filters import TimestampRangeFilter + from google.cloud.bigtable.row_filters import TimestampRange range_ = TimestampRange() row_filter = TimestampRangeFilter(range_) @@ -338,7 +338,7 @@ def test_timestamp_range_filter_to_pb(): def test_column_range_filter_constructor_defaults(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = object() row_filter = ColumnRangeFilter(column_family_id) @@ -350,7 +350,7 @@ def test_column_range_filter_constructor_defaults(): def test_column_range_filter_constructor_explicit(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = object() start_column = object() @@ -372,7 +372,7 @@ def test_column_range_filter_constructor_explicit(): def test_column_range_filter_constructor_bad_start(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = object() with pytest.raises(ValueError): @@ -380,7 +380,7 @@ def test_column_range_filter_constructor_bad_start(): def test_column_range_filter_constructor_bad_end(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = object() with pytest.raises(ValueError): @@ -388,7 +388,7 @@ def test_column_range_filter_constructor_bad_end(): def test_column_range_filter___eq__(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = object() start_column = object() @@ -413,7 +413,7 @@ def test_column_range_filter___eq__(): def test_column_range_filter___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = object() row_filter1 = ColumnRangeFilter(column_family_id) @@ -422,7 +422,7 @@ def test_column_range_filter___eq__type_differ(): def test_column_range_filter___ne__(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = object() other_column_family_id = object() @@ -448,7 +448,7 @@ def test_column_range_filter___ne__(): def test_column_range_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = "column-family-id" row_filter = ColumnRangeFilter(column_family_id) @@ -458,7 +458,7 @@ def test_column_range_filter_to_pb(): def test_column_range_filter_to_pb_inclusive_start(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = "column-family-id" column = b"column" @@ -471,7 +471,7 @@ def test_column_range_filter_to_pb_inclusive_start(): def test_column_range_filter_to_pb_exclusive_start(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = "column-family-id" column = b"column" @@ -486,7 +486,7 @@ def test_column_range_filter_to_pb_exclusive_start(): def test_column_range_filter_to_pb_inclusive_end(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = "column-family-id" column = b"column" @@ -499,7 +499,7 @@ def test_column_range_filter_to_pb_inclusive_end(): def test_column_range_filter_to_pb_exclusive_end(): - from google.cloud.bigtable.deprecated.row_filters import ColumnRangeFilter + from google.cloud.bigtable.row_filters import ColumnRangeFilter column_family_id = "column-family-id" column = b"column" @@ -514,7 +514,7 @@ def test_column_range_filter_to_pb_exclusive_end(): def test_value_regex_filter_to_pb_w_bytes(): - from google.cloud.bigtable.deprecated.row_filters import ValueRegexFilter + from google.cloud.bigtable.row_filters import ValueRegexFilter value = regex = b"value-regex" row_filter = ValueRegexFilter(value) @@ -524,7 +524,7 @@ def test_value_regex_filter_to_pb_w_bytes(): def test_value_regex_filter_to_pb_w_str(): - from google.cloud.bigtable.deprecated.row_filters import ValueRegexFilter + from google.cloud.bigtable.row_filters import ValueRegexFilter value = "value-regex" regex = value.encode("ascii") @@ -535,7 +535,7 @@ def test_value_regex_filter_to_pb_w_str(): def test_exact_value_filter_to_pb_w_bytes(): - from google.cloud.bigtable.deprecated.row_filters import ExactValueFilter + from google.cloud.bigtable.row_filters import ExactValueFilter value = regex = b"value-regex" row_filter = ExactValueFilter(value) @@ -545,7 +545,7 @@ def test_exact_value_filter_to_pb_w_bytes(): def test_exact_value_filter_to_pb_w_str(): - from google.cloud.bigtable.deprecated.row_filters import ExactValueFilter + from google.cloud.bigtable.row_filters import ExactValueFilter value = "value-regex" regex = value.encode("ascii") @@ -557,7 +557,7 @@ def test_exact_value_filter_to_pb_w_str(): def test_exact_value_filter_to_pb_w_int(): import struct - from google.cloud.bigtable.deprecated.row_filters import ExactValueFilter + from google.cloud.bigtable.row_filters import ExactValueFilter value = 1 regex = struct.Struct(">q").pack(value) @@ -568,7 +568,7 @@ def test_exact_value_filter_to_pb_w_int(): def test_value_range_filter_constructor_defaults(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter row_filter = ValueRangeFilter() @@ -579,7 +579,7 @@ def test_value_range_filter_constructor_defaults(): def test_value_range_filter_constructor_explicit(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter start_value = object() end_value = object() @@ -600,7 +600,7 @@ def test_value_range_filter_constructor_explicit(): def test_value_range_filter_constructor_w_int_values(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter import struct start_value = 1 @@ -618,21 +618,21 @@ def test_value_range_filter_constructor_w_int_values(): def test_value_range_filter_constructor_bad_start(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter with pytest.raises(ValueError): ValueRangeFilter(inclusive_start=True) def test_value_range_filter_constructor_bad_end(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter with pytest.raises(ValueError): ValueRangeFilter(inclusive_end=True) def test_value_range_filter___eq__(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter start_value = object() end_value = object() @@ -654,7 +654,7 @@ def test_value_range_filter___eq__(): def test_value_range_filter___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter row_filter1 = ValueRangeFilter() row_filter2 = object() @@ -662,7 +662,7 @@ def test_value_range_filter___eq__type_differ(): def test_value_range_filter___ne__(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter start_value = object() other_start_value = object() @@ -685,7 +685,7 @@ def test_value_range_filter___ne__(): def test_value_range_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter row_filter = ValueRangeFilter() expected_pb = _RowFilterPB(value_range_filter=_ValueRangePB()) @@ -693,7 +693,7 @@ def test_value_range_filter_to_pb(): def test_value_range_filter_to_pb_inclusive_start(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(start_value=value) @@ -703,7 +703,7 @@ def test_value_range_filter_to_pb_inclusive_start(): def test_value_range_filter_to_pb_exclusive_start(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(start_value=value, inclusive_start=False) @@ -713,7 +713,7 @@ def test_value_range_filter_to_pb_exclusive_start(): def test_value_range_filter_to_pb_inclusive_end(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(end_value=value) @@ -723,7 +723,7 @@ def test_value_range_filter_to_pb_inclusive_end(): def test_value_range_filter_to_pb_exclusive_end(): - from google.cloud.bigtable.deprecated.row_filters import ValueRangeFilter + from google.cloud.bigtable.row_filters import ValueRangeFilter value = b"some-value" row_filter = ValueRangeFilter(end_value=value, inclusive_end=False) @@ -733,7 +733,7 @@ def test_value_range_filter_to_pb_exclusive_end(): def test_cell_count_constructor(): - from google.cloud.bigtable.deprecated.row_filters import _CellCountFilter + from google.cloud.bigtable.row_filters import _CellCountFilter num_cells = object() row_filter = _CellCountFilter(num_cells) @@ -741,7 +741,7 @@ def test_cell_count_constructor(): def test_cell_count___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import _CellCountFilter + from google.cloud.bigtable.row_filters import _CellCountFilter num_cells = object() row_filter1 = _CellCountFilter(num_cells) @@ -750,7 +750,7 @@ def test_cell_count___eq__type_differ(): def test_cell_count___eq__same_value(): - from google.cloud.bigtable.deprecated.row_filters import _CellCountFilter + from google.cloud.bigtable.row_filters import _CellCountFilter num_cells = object() row_filter1 = _CellCountFilter(num_cells) @@ -759,7 +759,7 @@ def test_cell_count___eq__same_value(): def test_cell_count___ne__same_value(): - from google.cloud.bigtable.deprecated.row_filters import _CellCountFilter + from google.cloud.bigtable.row_filters import _CellCountFilter num_cells = object() row_filter1 = _CellCountFilter(num_cells) @@ -768,7 +768,7 @@ def test_cell_count___ne__same_value(): def test_cells_row_offset_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.row_filters import CellsRowOffsetFilter num_cells = 76 row_filter = CellsRowOffsetFilter(num_cells) @@ -778,7 +778,7 @@ def test_cells_row_offset_filter_to_pb(): def test_cells_row_limit_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.row_filters import CellsRowLimitFilter num_cells = 189 row_filter = CellsRowLimitFilter(num_cells) @@ -788,7 +788,7 @@ def test_cells_row_limit_filter_to_pb(): def test_cells_column_limit_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import CellsColumnLimitFilter + from google.cloud.bigtable.row_filters import CellsColumnLimitFilter num_cells = 10 row_filter = CellsColumnLimitFilter(num_cells) @@ -798,7 +798,7 @@ def test_cells_column_limit_filter_to_pb(): def test_strip_value_transformer_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter flag = True row_filter = StripValueTransformerFilter(flag) @@ -808,7 +808,7 @@ def test_strip_value_transformer_filter_to_pb(): def test_apply_label_filter_constructor(): - from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter + from google.cloud.bigtable.row_filters import ApplyLabelFilter label = object() row_filter = ApplyLabelFilter(label) @@ -816,7 +816,7 @@ def test_apply_label_filter_constructor(): def test_apply_label_filter___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter + from google.cloud.bigtable.row_filters import ApplyLabelFilter label = object() row_filter1 = ApplyLabelFilter(label) @@ -825,7 +825,7 @@ def test_apply_label_filter___eq__type_differ(): def test_apply_label_filter___eq__same_value(): - from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter + from google.cloud.bigtable.row_filters import ApplyLabelFilter label = object() row_filter1 = ApplyLabelFilter(label) @@ -834,7 +834,7 @@ def test_apply_label_filter___eq__same_value(): def test_apply_label_filter___ne__(): - from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter + from google.cloud.bigtable.row_filters import ApplyLabelFilter label = object() other_label = object() @@ -844,7 +844,7 @@ def test_apply_label_filter___ne__(): def test_apply_label_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import ApplyLabelFilter + from google.cloud.bigtable.row_filters import ApplyLabelFilter label = "label" row_filter = ApplyLabelFilter(label) @@ -854,14 +854,14 @@ def test_apply_label_filter_to_pb(): def test_filter_combination_constructor_defaults(): - from google.cloud.bigtable.deprecated.row_filters import _FilterCombination + from google.cloud.bigtable.row_filters import _FilterCombination row_filter = _FilterCombination() assert row_filter.filters == [] def test_filter_combination_constructor_explicit(): - from google.cloud.bigtable.deprecated.row_filters import _FilterCombination + from google.cloud.bigtable.row_filters import _FilterCombination filters = object() row_filter = _FilterCombination(filters=filters) @@ -869,7 +869,7 @@ def test_filter_combination_constructor_explicit(): def test_filter_combination___eq__(): - from google.cloud.bigtable.deprecated.row_filters import _FilterCombination + from google.cloud.bigtable.row_filters import _FilterCombination filters = object() row_filter1 = _FilterCombination(filters=filters) @@ -878,7 +878,7 @@ def test_filter_combination___eq__(): def test_filter_combination___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import _FilterCombination + from google.cloud.bigtable.row_filters import _FilterCombination filters = object() row_filter1 = _FilterCombination(filters=filters) @@ -887,7 +887,7 @@ def test_filter_combination___eq__type_differ(): def test_filter_combination___ne__(): - from google.cloud.bigtable.deprecated.row_filters import _FilterCombination + from google.cloud.bigtable.row_filters import _FilterCombination filters = object() other_filters = object() @@ -897,9 +897,9 @@ def test_filter_combination___ne__(): def test_row_filter_chain_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import RowFilterChain - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter - from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() @@ -917,10 +917,10 @@ def test_row_filter_chain_to_pb(): def test_row_filter_chain_to_pb_nested(): - from google.cloud.bigtable.deprecated.row_filters import CellsRowLimitFilter - from google.cloud.bigtable.deprecated.row_filters import RowFilterChain - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter - from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.row_filters import RowFilterChain + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -941,9 +941,9 @@ def test_row_filter_chain_to_pb_nested(): def test_row_filter_union_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import RowFilterUnion - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter - from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() @@ -961,10 +961,10 @@ def test_row_filter_union_to_pb(): def test_row_filter_union_to_pb_nested(): - from google.cloud.bigtable.deprecated.row_filters import CellsRowLimitFilter - from google.cloud.bigtable.deprecated.row_filters import RowFilterUnion - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter - from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.row_filters import CellsRowLimitFilter + from google.cloud.bigtable.row_filters import RowFilterUnion + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter2 = RowSampleFilter(0.25) @@ -985,7 +985,7 @@ def test_row_filter_union_to_pb_nested(): def test_conditional_row_filter_constructor(): - from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import ConditionalRowFilter base_filter = object() true_filter = object() @@ -999,7 +999,7 @@ def test_conditional_row_filter_constructor(): def test_conditional_row_filter___eq__(): - from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import ConditionalRowFilter base_filter = object() true_filter = object() @@ -1014,7 +1014,7 @@ def test_conditional_row_filter___eq__(): def test_conditional_row_filter___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import ConditionalRowFilter base_filter = object() true_filter = object() @@ -1027,7 +1027,7 @@ def test_conditional_row_filter___eq__type_differ(): def test_conditional_row_filter___ne__(): - from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import ConditionalRowFilter base_filter = object() other_base_filter = object() @@ -1043,10 +1043,10 @@ def test_conditional_row_filter___ne__(): def test_conditional_row_filter_to_pb(): - from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter - from google.cloud.bigtable.deprecated.row_filters import CellsRowOffsetFilter - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter - from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import CellsRowOffsetFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() @@ -1073,9 +1073,9 @@ def test_conditional_row_filter_to_pb(): def test_conditional_row_filter_to_pb_true_only(): - from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter - from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() @@ -1095,9 +1095,9 @@ def test_conditional_row_filter_to_pb_true_only(): def test_conditional_row_filter_to_pb_false_only(): - from google.cloud.bigtable.deprecated.row_filters import ConditionalRowFilter - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter - from google.cloud.bigtable.deprecated.row_filters import StripValueTransformerFilter + from google.cloud.bigtable.row_filters import ConditionalRowFilter + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_filters import StripValueTransformerFilter row_filter1 = StripValueTransformerFilter(True) row_filter1_pb = row_filter1.to_pb() diff --git a/tests/unit/v2_client/test_row_merger.py b/tests/unit/v2_client/test_row_merger.py index 26cedb34d..483c04536 100644 --- a/tests/unit/v2_client/test_row_merger.py +++ b/tests/unit/v2_client/test_row_merger.py @@ -5,13 +5,9 @@ import proto import pytest -from google.cloud.bigtable.deprecated.row_data import ( - PartialRowsData, - PartialRowData, - InvalidChunk, -) +from google.cloud.bigtable.row_data import PartialRowsData, PartialRowData, InvalidChunk from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse -from google.cloud.bigtable.deprecated.row_merger import _RowMerger +from google.cloud.bigtable.row_merger import _RowMerger # TODO: autogenerate protos from diff --git a/tests/unit/v2_client/test_row_set.py b/tests/unit/v2_client/test_row_set.py index ce0e9bfea..1a33be720 100644 --- a/tests/unit/v2_client/test_row_set.py +++ b/tests/unit/v2_client/test_row_set.py @@ -14,7 +14,7 @@ def test_row_set_constructor(): - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowSet row_set = RowSet() assert [] == row_set.row_keys @@ -22,8 +22,8 @@ def test_row_set_constructor(): def test_row_set__eq__(): - from google.cloud.bigtable.deprecated.row_set import RowRange - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.row_set import RowSet row_key1 = b"row_key1" row_key2 = b"row_key1" @@ -42,7 +42,7 @@ def test_row_set__eq__(): def test_row_set__eq__type_differ(): - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowSet row_set1 = RowSet() row_set2 = object() @@ -50,7 +50,7 @@ def test_row_set__eq__type_differ(): def test_row_set__eq__len_row_keys_differ(): - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowSet row_key1 = b"row_key1" row_key2 = b"row_key1" @@ -66,8 +66,8 @@ def test_row_set__eq__len_row_keys_differ(): def test_row_set__eq__len_row_ranges_differ(): - from google.cloud.bigtable.deprecated.row_set import RowRange - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.row_set import RowSet row_range1 = RowRange(b"row_key4", b"row_key9") row_range2 = RowRange(b"row_key4", b"row_key9") @@ -83,7 +83,7 @@ def test_row_set__eq__len_row_ranges_differ(): def test_row_set__eq__row_keys_differ(): - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowSet row_set1 = RowSet() row_set2 = RowSet() @@ -99,8 +99,8 @@ def test_row_set__eq__row_keys_differ(): def test_row_set__eq__row_ranges_differ(): - from google.cloud.bigtable.deprecated.row_set import RowRange - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.row_set import RowSet row_range1 = RowRange(b"row_key4", b"row_key9") row_range2 = RowRange(b"row_key14", b"row_key19") @@ -119,8 +119,8 @@ def test_row_set__eq__row_ranges_differ(): def test_row_set__ne__(): - from google.cloud.bigtable.deprecated.row_set import RowRange - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.row_set import RowSet row_key1 = b"row_key1" row_key2 = b"row_key1" @@ -139,8 +139,8 @@ def test_row_set__ne__(): def test_row_set__ne__same_value(): - from google.cloud.bigtable.deprecated.row_set import RowRange - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.row_set import RowSet row_key1 = b"row_key1" row_key2 = b"row_key1" @@ -159,7 +159,7 @@ def test_row_set__ne__same_value(): def test_row_set_add_row_key(): - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowSet row_set = RowSet() row_set.add_row_key("row_key1") @@ -168,8 +168,8 @@ def test_row_set_add_row_key(): def test_row_set_add_row_range(): - from google.cloud.bigtable.deprecated.row_set import RowRange - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.row_set import RowSet row_set = RowSet() row_range1 = RowRange(b"row_key1", b"row_key9") @@ -181,7 +181,7 @@ def test_row_set_add_row_range(): def test_row_set_add_row_range_from_keys(): - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowSet row_set = RowSet() row_set.add_row_range_from_keys( @@ -194,7 +194,7 @@ def test_row_set_add_row_range_from_keys(): def test_row_set_add_row_range_with_prefix(): - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowSet row_set = RowSet() row_set.add_row_range_with_prefix("row") @@ -203,8 +203,8 @@ def test_row_set_add_row_range_with_prefix(): def test_row_set__update_message_request(): from google.cloud._helpers import _to_bytes - from google.cloud.bigtable.deprecated.row_set import RowRange - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.row_set import RowSet row_set = RowSet() table_name = "table_name" @@ -224,7 +224,7 @@ def test_row_set__update_message_request(): def test_row_range_constructor(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange start_key = "row_key1" end_key = "row_key9" @@ -236,7 +236,7 @@ def test_row_range_constructor(): def test_row_range___hash__set_equality(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange row_range1 = RowRange("row_key1", "row_key9") row_range2 = RowRange("row_key1", "row_key9") @@ -246,7 +246,7 @@ def test_row_range___hash__set_equality(): def test_row_range___hash__not_equals(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange row_range1 = RowRange("row_key1", "row_key9") row_range2 = RowRange("row_key1", "row_key19") @@ -256,7 +256,7 @@ def test_row_range___hash__not_equals(): def test_row_range__eq__(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -266,7 +266,7 @@ def test_row_range__eq__(): def test_row_range___eq__type_differ(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -276,7 +276,7 @@ def test_row_range___eq__type_differ(): def test_row_range__ne__(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -286,7 +286,7 @@ def test_row_range__ne__(): def test_row_range__ne__same_value(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -296,7 +296,7 @@ def test_row_range__ne__same_value(): def test_row_range_get_range_kwargs_closed_open(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" @@ -307,7 +307,7 @@ def test_row_range_get_range_kwargs_closed_open(): def test_row_range_get_range_kwargs_open_closed(): - from google.cloud.bigtable.deprecated.row_set import RowRange + from google.cloud.bigtable.row_set import RowRange start_key = b"row_key1" end_key = b"row_key9" diff --git a/tests/unit/v2_client/test_table.py b/tests/unit/v2_client/test_table.py index ad31e8bc9..3d7d2e8ee 100644 --- a/tests/unit/v2_client/test_table.py +++ b/tests/unit/v2_client/test_table.py @@ -50,11 +50,11 @@ STATUS_INTERNAL = StatusCode.INTERNAL.value[0] -@mock.patch("google.cloud.bigtable.deprecated.table._MAX_BULK_MUTATIONS", new=3) +@mock.patch("google.cloud.bigtable.table._MAX_BULK_MUTATIONS", new=3) def test__compile_mutation_entries_w_too_many_mutations(): - from google.cloud.bigtable.deprecated.row import DirectRow - from google.cloud.bigtable.deprecated.table import TooManyMutationsError - from google.cloud.bigtable.deprecated.table import _compile_mutation_entries + from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.table import TooManyMutationsError + from google.cloud.bigtable.table import _compile_mutation_entries table = mock.Mock(name="table", spec=["name"]) table.name = "table" @@ -72,8 +72,8 @@ def test__compile_mutation_entries_w_too_many_mutations(): def test__compile_mutation_entries_normal(): - from google.cloud.bigtable.deprecated.row import DirectRow - from google.cloud.bigtable.deprecated.table import _compile_mutation_entries + from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.table import _compile_mutation_entries from google.cloud.bigtable_v2.types import MutateRowsRequest from google.cloud.bigtable_v2.types import data @@ -109,9 +109,9 @@ def test__compile_mutation_entries_normal(): def test__check_row_table_name_w_wrong_table_name(): - from google.cloud.bigtable.deprecated.table import _check_row_table_name - from google.cloud.bigtable.deprecated.table import TableMismatchError - from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.table import _check_row_table_name + from google.cloud.bigtable.table import TableMismatchError + from google.cloud.bigtable.row import DirectRow table = mock.Mock(name="table", spec=["name"]) table.name = "table" @@ -122,8 +122,8 @@ def test__check_row_table_name_w_wrong_table_name(): def test__check_row_table_name_w_right_table_name(): - from google.cloud.bigtable.deprecated.row import DirectRow - from google.cloud.bigtable.deprecated.table import _check_row_table_name + from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.table import _check_row_table_name table = mock.Mock(name="table", spec=["name"]) table.name = "table" @@ -133,8 +133,8 @@ def test__check_row_table_name_w_right_table_name(): def test__check_row_type_w_wrong_row_type(): - from google.cloud.bigtable.deprecated.row import ConditionalRow - from google.cloud.bigtable.deprecated.table import _check_row_type + from google.cloud.bigtable.row import ConditionalRow + from google.cloud.bigtable.table import _check_row_type row = ConditionalRow(row_key=b"row_key", table="table", filter_=None) with pytest.raises(TypeError): @@ -142,21 +142,21 @@ def test__check_row_type_w_wrong_row_type(): def test__check_row_type_w_right_row_type(): - from google.cloud.bigtable.deprecated.row import DirectRow - from google.cloud.bigtable.deprecated.table import _check_row_type + from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.table import _check_row_type row = DirectRow(row_key=b"row_key", table="table") assert not _check_row_type(row) def _make_client(*args, **kwargs): - from google.cloud.bigtable.deprecated.client import Client + from google.cloud.bigtable.client import Client return Client(*args, **kwargs) def _make_table(*args, **kwargs): - from google.cloud.bigtable.deprecated.table import Table + from google.cloud.bigtable.table import Table return Table(*args, **kwargs) @@ -219,7 +219,7 @@ def _table_row_methods_helper(): def test_table_row_factory_direct(): - from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.row import DirectRow table, row_key = _table_row_methods_helper() with warnings.catch_warnings(record=True) as warned: @@ -234,7 +234,7 @@ def test_table_row_factory_direct(): def test_table_row_factory_conditional(): - from google.cloud.bigtable.deprecated.row import ConditionalRow + from google.cloud.bigtable.row import ConditionalRow table, row_key = _table_row_methods_helper() filter_ = object() @@ -251,7 +251,7 @@ def test_table_row_factory_conditional(): def test_table_row_factory_append(): - from google.cloud.bigtable.deprecated.row import AppendRow + from google.cloud.bigtable.row import AppendRow table, row_key = _table_row_methods_helper() @@ -278,7 +278,7 @@ def test_table_row_factory_failure(): def test_table_direct_row(): - from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.row import DirectRow table, row_key = _table_row_methods_helper() row = table.direct_row(row_key) @@ -289,7 +289,7 @@ def test_table_direct_row(): def test_table_conditional_row(): - from google.cloud.bigtable.deprecated.row import ConditionalRow + from google.cloud.bigtable.row import ConditionalRow table, row_key = _table_row_methods_helper() filter_ = object() @@ -301,7 +301,7 @@ def test_table_conditional_row(): def test_table_append_row(): - from google.cloud.bigtable.deprecated.row import AppendRow + from google.cloud.bigtable.row import AppendRow table, row_key = _table_row_methods_helper() row = table.append_row(row_key) @@ -357,7 +357,7 @@ def _create_table_helper(split_keys=[], column_families={}): from google.cloud.bigtable_admin_v2.types import ( bigtable_table_admin as table_admin_messages_v2_pb2, ) - from google.cloud.bigtable.deprecated.column_family import ColumnFamily + from google.cloud.bigtable.column_family import ColumnFamily credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -391,7 +391,7 @@ def test_table_create(): def test_table_create_with_families(): - from google.cloud.bigtable.deprecated.column_family import MaxVersionsGCRule + from google.cloud.bigtable.column_family import MaxVersionsGCRule families = {"family": MaxVersionsGCRule(5)} _create_table_helper(column_families=families) @@ -404,7 +404,7 @@ def test_table_create_with_split_keys(): def test_table_exists_hit(): from google.cloud.bigtable_admin_v2.types import ListTablesResponse from google.cloud.bigtable_admin_v2.types import Table - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -426,7 +426,7 @@ def test_table_exists_hit(): def test_table_exists_miss(): from google.api_core.exceptions import NotFound - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -447,7 +447,7 @@ def test_table_exists_miss(): def test_table_exists_error(): from google.api_core.exceptions import BadRequest - from google.cloud.bigtable.deprecated import enums + from google.cloud.bigtable import enums credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -512,8 +512,8 @@ def test_table_list_column_families(): def test_table_get_cluster_states(): - from google.cloud.bigtable.deprecated.enums import Table as enum_table - from google.cloud.bigtable.deprecated.table import ClusterState + from google.cloud.bigtable.enums import Table as enum_table + from google.cloud.bigtable.table import ClusterState INITIALIZING = enum_table.ReplicationState.INITIALIZING PLANNED_MAINTENANCE = enum_table.ReplicationState.PLANNED_MAINTENANCE @@ -557,10 +557,10 @@ def test_table_get_cluster_states(): def test_table_get_encryption_info(): from google.rpc.code_pb2 import Code - from google.cloud.bigtable.deprecated.encryption_info import EncryptionInfo - from google.cloud.bigtable.deprecated.enums import EncryptionInfo as enum_crypto - from google.cloud.bigtable.deprecated.enums import Table as enum_table - from google.cloud.bigtable.deprecated.error import Status + from google.cloud.bigtable.encryption_info import EncryptionInfo + from google.cloud.bigtable.enums import EncryptionInfo as enum_crypto + from google.cloud.bigtable.enums import Table as enum_table + from google.cloud.bigtable.error import Status ENCRYPTION_TYPE_UNSPECIFIED = enum_crypto.EncryptionType.ENCRYPTION_TYPE_UNSPECIFIED GOOGLE_DEFAULT_ENCRYPTION = enum_crypto.EncryptionType.GOOGLE_DEFAULT_ENCRYPTION @@ -640,9 +640,10 @@ def _make_data_api(): def _table_read_row_helper(chunks, expected_result, app_profile_id=None): from google.cloud._testing import _Monkey - from google.cloud.bigtable.deprecated import table as MUT - from google.cloud.bigtable.deprecated.row_set import RowSet - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable import table as MUT + from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.row_filters import RowSampleFilter + from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -691,7 +692,9 @@ def mock_create_row_request(table_name, **kwargs): assert result == expected_result assert mock_created == expected_request - data_api.read_rows.assert_called_once_with(request_pb, timeout=61.0) + data_api.read_rows.assert_called_once_with( + request_pb, timeout=61.0, retry=DEFAULT_RETRY_READ_ROWS + ) def test_table_read_row_miss_no__responses(): @@ -704,8 +707,8 @@ def test_table_read_row_miss_no_chunks_in_response(): def test_table_read_row_complete(): - from google.cloud.bigtable.deprecated.row_data import Cell - from google.cloud.bigtable.deprecated.row_data import PartialRowData + from google.cloud.bigtable.row_data import Cell + from google.cloud.bigtable.row_data import PartialRowData app_profile_id = "app-profile-id" chunk = _ReadRowsResponseCellChunkPB( @@ -768,7 +771,7 @@ def _table_mutate_rows_helper( mutation_timeout=None, app_profile_id=None, retry=None, timeout=None ): from google.rpc.status_pb2 import Status - from google.cloud.bigtable.deprecated.table import DEFAULT_RETRY + from google.cloud.bigtable.table import DEFAULT_RETRY credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -787,7 +790,7 @@ def _table_mutate_rows_helper( response = [Status(code=0), Status(code=1)] instance_mock = mock.Mock(return_value=response) klass_mock = mock.patch( - "google.cloud.bigtable.deprecated.table._RetryableMutateRowsWorker", + "google.cloud.bigtable.table._RetryableMutateRowsWorker", new=mock.MagicMock(return_value=instance_mock), ) @@ -854,9 +857,9 @@ def test_table_mutate_rows_w_mutation_timeout_and_timeout_arg(): def test_table_read_rows(): from google.cloud._testing import _Monkey - from google.cloud.bigtable.deprecated.row_data import PartialRowsData - from google.cloud.bigtable.deprecated import table as MUT - from google.cloud.bigtable.deprecated.row_data import DEFAULT_RETRY_READ_ROWS + from google.cloud.bigtable.row_data import PartialRowsData + from google.cloud.bigtable import table as MUT + from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -906,7 +909,7 @@ def mock_create_row_request(table_name, **kwargs): } assert mock_created == [(table.name, created_kwargs)] - data_api.read_rows.assert_called_once_with(request_pb, timeout=61.0) + data_api.read_rows.assert_called_once_with(request_pb, timeout=61.0, retry=retry) def test_table_read_retry_rows(): @@ -1017,7 +1020,7 @@ def test_table_read_retry_rows_no_full_table_scan(): def test_table_yield_retry_rows(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1079,9 +1082,10 @@ def test_table_yield_retry_rows(): def test_table_yield_rows_with_row_set(): - from google.cloud.bigtable.deprecated.row_set import RowSet - from google.cloud.bigtable.deprecated.row_set import RowRange - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.row_set import RowSet + from google.cloud.bigtable.row_set import RowRange + from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1149,7 +1153,9 @@ def test_table_yield_rows_with_row_set(): end_key=ROW_KEY_2, ) expected_request.rows.row_keys.append(ROW_KEY_3) - data_api.read_rows.assert_called_once_with(expected_request, timeout=61.0) + data_api.read_rows.assert_called_once_with( + expected_request, timeout=61.0, retry=DEFAULT_RETRY_READ_ROWS + ) def test_table_sample_row_keys(): @@ -1174,9 +1180,7 @@ def test_table_truncate(): table = _make_table(TABLE_ID, instance) table_api = client._table_admin_client = _make_table_api() - with mock.patch( - "google.cloud.bigtable.deprecated.table.Table.name", new=TABLE_NAME - ): + with mock.patch("google.cloud.bigtable.table.Table.name", new=TABLE_NAME): result = table.truncate() assert result is None @@ -1257,7 +1261,7 @@ def test_table_mutations_batcher_factory(): def test_table_get_iam_policy(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1288,8 +1292,8 @@ def test_table_get_iam_policy(): def test_table_set_iam_policy(): from google.iam.v1 import policy_pb2 - from google.cloud.bigtable.deprecated.policy import Policy - from google.cloud.bigtable.deprecated.policy import BIGTABLE_ADMIN_ROLE + from google.cloud.bigtable.policy import Policy + from google.cloud.bigtable.policy import BIGTABLE_ADMIN_ROLE credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1351,7 +1355,7 @@ def test_table_test_iam_permissions(): def test_table_backup_factory_defaults(): - from google.cloud.bigtable.deprecated.backup import Backup + from google.cloud.bigtable.backup import Backup instance = _make_table(INSTANCE_ID, None) table = _make_table(TABLE_ID, instance) @@ -1375,8 +1379,8 @@ def test_table_backup_factory_defaults(): def test_table_backup_factory_non_defaults(): import datetime from google.cloud._helpers import UTC - from google.cloud.bigtable.deprecated.backup import Backup - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.instance import Instance instance = Instance(INSTANCE_ID, None) table = _make_table(TABLE_ID, instance) @@ -1406,7 +1410,7 @@ def _table_list_backups_helper(cluster_id=None, filter_=None, **kwargs): Backup as backup_pb, bigtable_table_admin, ) - from google.cloud.bigtable.deprecated.backup import Backup + from google.cloud.bigtable.backup import Backup client = _make_client( project=PROJECT_ID, credentials=_make_credentials(), admin=True @@ -1468,7 +1472,7 @@ def test_table_list_backups_w_options(): def _table_restore_helper(backup_name=None): - from google.cloud.bigtable.deprecated.instance import Instance + from google.cloud.bigtable.instance import Instance op_future = object() credentials = _make_credentials() @@ -1504,7 +1508,7 @@ def test_table_restore_table_w_backup_name(): def _make_worker(*args, **kwargs): - from google.cloud.bigtable.deprecated.table import _RetryableMutateRowsWorker + from google.cloud.bigtable.table import _RetryableMutateRowsWorker return _RetryableMutateRowsWorker(*args, **kwargs) @@ -1545,7 +1549,7 @@ def test_rmrw_callable_empty_rows(): def test_rmrw_callable_no_retry_strategy(): - from google.cloud.bigtable.deprecated.row import DirectRow + from google.cloud.bigtable.row import DirectRow # Setup: # - Mutate 3 rows. @@ -1587,8 +1591,8 @@ def test_rmrw_callable_no_retry_strategy(): def test_rmrw_callable_retry(): - from google.cloud.bigtable.deprecated.row import DirectRow - from google.cloud.bigtable.deprecated.table import DEFAULT_RETRY + from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.table import DEFAULT_RETRY # Setup: # - Mutate 3 rows. @@ -1642,8 +1646,8 @@ def _do_mutate_retryable_rows_helper( mutate_rows_side_effect=None, ): from google.api_core.exceptions import ServiceUnavailable - from google.cloud.bigtable.deprecated.row import DirectRow - from google.cloud.bigtable.deprecated.table import _BigtableRetryableError + from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.table import _BigtableRetryableError from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 # Setup: @@ -1799,9 +1803,7 @@ def test_rmrw_do_mutate_retryable_rows_w_retryable_error_internal_rst_stream_err # Raise internal server error with RST STREAM error messages # There should be no error raised and that the request is retried from google.api_core.exceptions import InternalServerError - from google.cloud.bigtable.deprecated.row_data import ( - RETRYABLE_INTERNAL_ERROR_MESSAGES, - ) + from google.cloud.bigtable.row_data import RETRYABLE_INTERNAL_ERROR_MESSAGES row_cells = [ (b"row_key_1", ("cf", b"col", b"value1")), @@ -2007,7 +2009,7 @@ def test_rmrw_do_mutate_retryable_rows_mismatch_num_responses(): def test__create_row_request_table_name_only(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request table_name = "table_name" result = _create_row_request(table_name) @@ -2016,14 +2018,14 @@ def test__create_row_request_table_name_only(): def test__create_row_request_row_range_row_set_conflict(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request with pytest.raises(ValueError): _create_row_request(None, end_key=object(), row_set=object()) def test__create_row_request_row_range_start_key(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange table_name = "table_name" @@ -2036,7 +2038,7 @@ def test__create_row_request_row_range_start_key(): def test__create_row_request_row_range_end_key(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange table_name = "table_name" @@ -2049,7 +2051,7 @@ def test__create_row_request_row_range_end_key(): def test__create_row_request_row_range_both_keys(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange table_name = "table_name" @@ -2063,7 +2065,7 @@ def test__create_row_request_row_range_both_keys(): def test__create_row_request_row_range_both_keys_inclusive(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange table_name = "table_name" @@ -2079,8 +2081,8 @@ def test__create_row_request_row_range_both_keys_inclusive(): def test__create_row_request_with_filter(): - from google.cloud.bigtable.deprecated.table import _create_row_request - from google.cloud.bigtable.deprecated.row_filters import RowSampleFilter + from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.row_filters import RowSampleFilter table_name = "table_name" row_filter = RowSampleFilter(0.33) @@ -2092,7 +2094,7 @@ def test__create_row_request_with_filter(): def test__create_row_request_with_limit(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request table_name = "table_name" limit = 1337 @@ -2102,8 +2104,8 @@ def test__create_row_request_with_limit(): def test__create_row_request_with_row_set(): - from google.cloud.bigtable.deprecated.table import _create_row_request - from google.cloud.bigtable.deprecated.row_set import RowSet + from google.cloud.bigtable.table import _create_row_request + from google.cloud.bigtable.row_set import RowSet table_name = "table_name" row_set = RowSet() @@ -2113,7 +2115,7 @@ def test__create_row_request_with_row_set(): def test__create_row_request_with_app_profile_id(): - from google.cloud.bigtable.deprecated.table import _create_row_request + from google.cloud.bigtable.table import _create_row_request table_name = "table_name" limit = 1337 @@ -2132,8 +2134,8 @@ def _ReadRowsRequestPB(*args, **kw): def test_cluster_state___eq__(): - from google.cloud.bigtable.deprecated.enums import Table as enum_table - from google.cloud.bigtable.deprecated.table import ClusterState + from google.cloud.bigtable.enums import Table as enum_table + from google.cloud.bigtable.table import ClusterState READY = enum_table.ReplicationState.READY state1 = ClusterState(READY) @@ -2142,8 +2144,8 @@ def test_cluster_state___eq__(): def test_cluster_state___eq__type_differ(): - from google.cloud.bigtable.deprecated.enums import Table as enum_table - from google.cloud.bigtable.deprecated.table import ClusterState + from google.cloud.bigtable.enums import Table as enum_table + from google.cloud.bigtable.table import ClusterState READY = enum_table.ReplicationState.READY state1 = ClusterState(READY) @@ -2152,8 +2154,8 @@ def test_cluster_state___eq__type_differ(): def test_cluster_state___ne__same_value(): - from google.cloud.bigtable.deprecated.enums import Table as enum_table - from google.cloud.bigtable.deprecated.table import ClusterState + from google.cloud.bigtable.enums import Table as enum_table + from google.cloud.bigtable.table import ClusterState READY = enum_table.ReplicationState.READY state1 = ClusterState(READY) @@ -2162,8 +2164,8 @@ def test_cluster_state___ne__same_value(): def test_cluster_state___ne__(): - from google.cloud.bigtable.deprecated.enums import Table as enum_table - from google.cloud.bigtable.deprecated.table import ClusterState + from google.cloud.bigtable.enums import Table as enum_table + from google.cloud.bigtable.table import ClusterState READY = enum_table.ReplicationState.READY INITIALIZING = enum_table.ReplicationState.INITIALIZING @@ -2173,8 +2175,8 @@ def test_cluster_state___ne__(): def test_cluster_state__repr__(): - from google.cloud.bigtable.deprecated.enums import Table as enum_table - from google.cloud.bigtable.deprecated.table import ClusterState + from google.cloud.bigtable.enums import Table as enum_table + from google.cloud.bigtable.table import ClusterState STATE_NOT_KNOWN = enum_table.ReplicationState.STATE_NOT_KNOWN INITIALIZING = enum_table.ReplicationState.INITIALIZING From 07438ca301f2bdbb3945b6330c5bb4c7c687747d Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 29 Jun 2023 16:04:10 -0700 Subject: [PATCH 15/56] feat: improve timeout structure (#819) --- .github/.OwlBot.lock.yaml | 4 +- .../bigtable/data/_async/_mutate_rows.py | 8 +- .../cloud/bigtable/data/_async/_read_rows.py | 6 +- google/cloud/bigtable/data/_async/client.py | 368 +++++++++++------- .../bigtable/data/_async/mutations_batcher.py | 31 +- google/cloud/bigtable/data/_helpers.py | 23 ++ noxfile.py | 3 +- tests/unit/data/_async/test__mutate_rows.py | 4 +- tests/unit/data/_async/test__read_rows.py | 2 +- tests/unit/data/_async/test_client.py | 187 +++++---- .../data/_async/test_mutations_batcher.py | 61 ++- tests/unit/data/test__helpers.py | 39 ++ 12 files changed, 469 insertions(+), 267 deletions(-) diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index 02a4dedce..1b3cb6c52 100644 --- a/.github/.OwlBot.lock.yaml +++ b/.github/.OwlBot.lock.yaml @@ -13,5 +13,5 @@ # limitations under the License. docker: image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest - digest: sha256:240b5bcc2bafd450912d2da2be15e62bc6de2cf839823ae4bf94d4f392b451dc -# created: 2023-06-03T21:25:37.968717478Z + digest: sha256:ddf4551385d566771dc713090feb7b4c1164fb8a698fe52bbe7670b24236565b +# created: 2023-06-27T13:04:21.96690344Z diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index ac491adaf..519043a92 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -52,15 +52,15 @@ def __init__( table: "TableAsync", mutation_entries: list["RowMutationEntry"], operation_timeout: float, - per_request_timeout: float | None, + attempt_timeout: float | None, ): """ Args: - gapic_client: the client to use for the mutate_rows call - table: the table associated with the request - mutation_entries: a list of RowMutationEntry objects to send to the server - - operation_timeout: the timeout t o use for the entire operation, in seconds. - - per_request_timeout: the timeoutto use for each mutate_rows attempt, in seconds. + - operation_timeout: the timeout to use for the entire operation, in seconds. + - attempt_timeout: the timeout to use for each mutate_rows attempt, in seconds. If not specified, the request will run until operation_timeout is reached. """ # check that mutations are within limits @@ -99,7 +99,7 @@ def __init__( self._operation = _convert_retry_deadline(retry_wrapped, operation_timeout) # initialize state self.timeout_generator = _attempt_timeout_generator( - per_request_timeout, operation_timeout + attempt_timeout, operation_timeout ) self.mutations = mutation_entries self.remaining_indices = list(range(len(self.mutations))) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 910a01c4c..ec1e488c6 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -63,14 +63,14 @@ def __init__( client: BigtableAsyncClient, *, operation_timeout: float = 600.0, - per_request_timeout: float | None = None, + attempt_timeout: float | None = None, ): """ Args: - request: the request dict to send to the Bigtable API - client: the Bigtable client to use to make the request - operation_timeout: the timeout to use for the entire operation, in seconds - - per_request_timeout: the timeout to use when waiting for each individual grpc request, in seconds + - attempt_timeout: the timeout to use when waiting for each individual grpc request, in seconds If not specified, defaults to operation_timeout """ self._last_emitted_row_key: bytes | None = None @@ -79,7 +79,7 @@ def __init__( self.operation_timeout = operation_timeout # use generator to lower per-attempt timeout as we approach operation_timeout deadline attempt_timeout_gen = _attempt_timeout_generator( - per_request_timeout, operation_timeout + attempt_timeout, operation_timeout ) row_limit = request.get("rows_limit", 0) # lock in paramters for retryable wrapper diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 3a5831799..6c321fe62 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -58,6 +58,7 @@ from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync from google.cloud.bigtable.data._helpers import _make_metadata from google.cloud.bigtable.data._helpers import _convert_retry_deadline +from google.cloud.bigtable.data._helpers import _validate_timeouts from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync from google.cloud.bigtable.data._async.mutations_batcher import _MB_SIZE from google.cloud.bigtable.data._helpers import _attempt_timeout_generator @@ -340,14 +341,18 @@ async def _remove_instance_registration( except KeyError: return False - # TODO: revisit timeouts https://github.com/googleapis/python-bigtable/issues/782 def get_table( self, instance_id: str, table_id: str, app_profile_id: str | None = None, - default_operation_timeout: float = 600, - default_per_request_timeout: float | None = None, + *, + default_read_rows_operation_timeout: float = 600, + default_read_rows_attempt_timeout: float | None = None, + default_mutate_rows_operation_timeout: float = 600, + default_mutate_rows_attempt_timeout: float | None = None, + default_operation_timeout: float = 60, + default_attempt_timeout: float | None = None, ) -> TableAsync: """ Returns a table instance for making data API requests @@ -356,9 +361,22 @@ def get_table( instance_id: The Bigtable instance ID to associate with this client. instance_id is combined with the client's project to fully specify the instance - table_id: The ID of the table. - app_profile_id: (Optional) The app profile to associate with requests. + table_id: The ID of the table. table_id is combined with the + instance_id and the client's project to fully specify the table + app_profile_id: The app profile to associate with requests. https://cloud.google.com/bigtable/docs/app-profiles + default_read_rows_operation_timeout: The default timeout for read rows + operations, in seconds. If not set, defaults to 600 seconds (10 minutes) + default_read_rows_attempt_timeout: The default timeout for individual + read rows rpc requests, in seconds. If not set, defaults to 20 seconds + default_mutate_rows_operation_timeout: The default timeout for mutate rows + operations, in seconds. If not set, defaults to 600 seconds (10 minutes) + default_mutate_rows_attempt_timeout: The default timeout for individual + mutate rows rpc requests, in seconds. If not set, defaults to 60 seconds + default_operation_timeout: The default timeout for all other operations, in + seconds. If not set, defaults to 60 seconds + default_attempt_timeout: The default timeout for all other individual rpc + requests, in seconds. If not set, defaults to 20 seconds """ return TableAsync( self, @@ -366,7 +384,7 @@ def get_table( table_id, app_profile_id, default_operation_timeout=default_operation_timeout, - default_per_request_timeout=default_per_request_timeout, + default_attempt_timeout=default_attempt_timeout, ) async def __aenter__(self): @@ -393,8 +411,12 @@ def __init__( table_id: str, app_profile_id: str | None = None, *, - default_operation_timeout: float = 600, - default_per_request_timeout: float | None = None, + default_read_rows_operation_timeout: float = 600, + default_read_rows_attempt_timeout: float | None = 20, + default_mutate_rows_operation_timeout: float = 600, + default_mutate_rows_attempt_timeout: float | None = 60, + default_operation_timeout: float = 60, + default_attempt_timeout: float | None = 20, ): """ Initialize a Table instance @@ -407,26 +429,38 @@ def __init__( specify the instance table_id: The ID of the table. table_id is combined with the instance_id and the client's project to fully specify the table - app_profile_id: (Optional) The app profile to associate with requests. + app_profile_id: The app profile to associate with requests. https://cloud.google.com/bigtable/docs/app-profiles - default_operation_timeout: (Optional) The default timeout, in seconds - default_per_request_timeout: (Optional) The default timeout for individual - rpc requests, in seconds + default_read_rows_operation_timeout: The default timeout for read rows + operations, in seconds. If not set, defaults to 600 seconds (10 minutes) + default_read_rows_attempt_timeout: The default timeout for individual + read rows rpc requests, in seconds. If not set, defaults to 20 seconds + default_mutate_rows_operation_timeout: The default timeout for mutate rows + operations, in seconds. If not set, defaults to 600 seconds (10 minutes) + default_mutate_rows_attempt_timeout: The default timeout for individual + mutate rows rpc requests, in seconds. If not set, defaults to 60 seconds + default_operation_timeout: The default timeout for all other operations, in + seconds. If not set, defaults to 60 seconds + default_attempt_timeout: The default timeout for all other individual rpc + requests, in seconds. If not set, defaults to 20 seconds Raises: - RuntimeError if called outside of an async context (no running event loop) """ # validate timeouts - if default_operation_timeout <= 0: - raise ValueError("default_operation_timeout must be greater than 0") - if default_per_request_timeout is not None and default_per_request_timeout <= 0: - raise ValueError("default_per_request_timeout must be greater than 0") - if ( - default_per_request_timeout is not None - and default_per_request_timeout > default_operation_timeout - ): - raise ValueError( - "default_per_request_timeout must be less than default_operation_timeout" - ) + _validate_timeouts( + default_operation_timeout, default_attempt_timeout, allow_none=True + ) + _validate_timeouts( + default_read_rows_operation_timeout, + default_read_rows_attempt_timeout, + allow_none=True, + ) + _validate_timeouts( + default_mutate_rows_operation_timeout, + default_mutate_rows_attempt_timeout, + allow_none=True, + ) + self.client = client self.instance_id = instance_id self.instance_name = self.client._gapic_client.instance_path( @@ -439,7 +473,13 @@ def __init__( self.app_profile_id = app_profile_id self.default_operation_timeout = default_operation_timeout - self.default_per_request_timeout = default_per_request_timeout + self.default_attempt_timeout = default_attempt_timeout + self.default_read_rows_operation_timeout = default_read_rows_operation_timeout + self.default_read_rows_attempt_timeout = default_read_rows_attempt_timeout + self.default_mutate_rows_operation_timeout = ( + default_mutate_rows_operation_timeout + ) + self.default_mutate_rows_attempt_timeout = default_mutate_rows_attempt_timeout # raises RuntimeError if called outside of an async context (no running event loop) try: @@ -456,24 +496,24 @@ async def read_rows_stream( query: ReadRowsQuery | dict[str, Any], *, operation_timeout: float | None = None, - per_request_timeout: float | None = None, + attempt_timeout: float | None = None, ) -> ReadRowsAsyncIterator: """ + Read a set of rows from the table, based on the specified query. Returns an iterator to asynchronously stream back row data. - Failed requests within operation_timeout and operation_deadline policies will be retried. + Failed requests within operation_timeout will be retried. Args: - query: contains details about which rows to return - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - time is only counted while actively waiting on the network. - If None, defaults to the Table's default_operation_timeout - - per_request_timeout: the time budget for an individual network request, in seconds. + If None, defaults to the Table's default_read_rows_operation_timeout + - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_per_request_timeout - + If None, defaults to the Table's default_read_rows_attempt_timeout, + or the operation_timeout if that is also None. Returns: - an asynchronous iterator that yields rows returned by the query Raises: @@ -483,35 +523,31 @@ async def read_rows_stream( - GoogleAPIError: raised if the request encounters an unrecoverable error - IdleTimeout: if iterator was abandoned """ + operation_timeout = ( + operation_timeout or self.default_read_rows_operation_timeout + ) + attempt_timeout = ( + attempt_timeout + or self.default_read_rows_attempt_timeout + or operation_timeout + ) + _validate_timeouts(operation_timeout, attempt_timeout) - operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = per_request_timeout or self.default_per_request_timeout - - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout <= 0: - raise ValueError("per_request_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout > operation_timeout: - raise ValueError( - "per_request_timeout must not be greater than operation_timeout" - ) - if per_request_timeout is None: - per_request_timeout = operation_timeout request = query._to_dict() if isinstance(query, ReadRowsQuery) else query request["table_name"] = self.table_name if self.app_profile_id: request["app_profile_id"] = self.app_profile_id # read_rows smart retries is implemented using a series of iterators: - # - client.read_rows: outputs raw ReadRowsResponse objects from backend. Has per_request_timeout + # - client.read_rows: outputs raw ReadRowsResponse objects from backend. Has attempt_timeout # - ReadRowsOperation.merge_row_response_stream: parses chunks into rows - # - ReadRowsOperation.retryable_merge_rows: adds retries, caching, revised requests, per_request_timeout + # - ReadRowsOperation.retryable_merge_rows: adds retries, caching, revised requests, operation_timeout # - ReadRowsAsyncIterator: adds idle_timeout, moves stats out of stream and into attribute row_merger = _ReadRowsOperationAsync( request, self.client._gapic_client, operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, + attempt_timeout=attempt_timeout, ) output_generator = ReadRowsAsyncIterator(row_merger) # add idle timeout to clear resources if generator is abandoned @@ -524,20 +560,37 @@ async def read_rows( query: ReadRowsQuery | dict[str, Any], *, operation_timeout: float | None = None, - per_request_timeout: float | None = None, + attempt_timeout: float | None = None, ) -> list[Row]: """ - Helper function that returns a full list instead of a generator + Read a set of rows from the table, based on the specified query. + Retruns results as a list of Row objects when the request is complete. + For streamed results, use read_rows_stream. - See read_rows_stream + Failed requests within operation_timeout will be retried. + Args: + - query: contains details about which rows to return + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + If None, defaults to the Table's default_read_rows_operation_timeout + - attempt_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_read_rows_attempt_timeout, + or the operation_timeout if that is also None. Returns: - - a list of the rows returned by the query + - a list of Rows returned by the query + Raises: + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions + from any retries that failed + - GoogleAPIError: raised if the request encounters an unrecoverable error """ row_generator = await self.read_rows_stream( query, operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, + attempt_timeout=attempt_timeout, ) results = [row async for row in row_generator] return results @@ -547,18 +600,31 @@ async def read_row( row_key: str | bytes, *, row_filter: RowFilter | None = None, - operation_timeout: int | float | None = 60, - per_request_timeout: int | float | None = None, + operation_timeout: int | float | None = None, + attempt_timeout: int | float | None = None, ) -> Row | None: """ - Helper function to return a single row + Read a single row from the table, based on the specified key. - See read_rows_stream + Failed requests within operation_timeout will be retried. - Raises: - - google.cloud.bigtable.data.exceptions.RowNotFound: if the row does not exist + Args: + - query: contains details about which rows to return + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + If None, defaults to the Table's default_read_rows_operation_timeout + - attempt_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_read_rows_attempt_timeout, or the operation_timeout + if that is also None. Returns: - - the individual row requested, or None if it does not exist + - a Row object if the row exists, otherwise None + Raises: + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions + from any retries that failed + - GoogleAPIError: raised if the request encounters an unrecoverable error """ if row_key is None: raise ValueError("row_key must be string or bytes") @@ -566,7 +632,7 @@ async def read_row( results = await self.read_rows( query, operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, + attempt_timeout=attempt_timeout, ) if len(results) == 0: return None @@ -577,7 +643,7 @@ async def read_rows_sharded( sharded_query: ShardedQuery, *, operation_timeout: int | float | None = None, - per_request_timeout: int | float | None = None, + attempt_timeout: int | float | None = None, ) -> list[Row]: """ Runs a sharded query in parallel, then return the results in a single list. @@ -594,6 +660,14 @@ async def read_rows_sharded( Args: - sharded_query: a sharded query to execute + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + If None, defaults to the Table's default_read_rows_operation_timeout + - attempt_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_read_rows_attempt_timeout, or the operation_timeout + if that is also None. Raises: - ShardedReadRowsExceptionGroup: if any of the queries failed - ValueError: if the query_list is empty @@ -601,10 +675,15 @@ async def read_rows_sharded( if not sharded_query: raise ValueError("empty sharded_query") # reduce operation_timeout between batches - operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = ( - per_request_timeout or self.default_per_request_timeout or operation_timeout + operation_timeout = ( + operation_timeout or self.default_read_rows_operation_timeout ) + attempt_timeout = ( + attempt_timeout + or self.default_read_rows_attempt_timeout + or operation_timeout + ) + _validate_timeouts(operation_timeout, attempt_timeout) timeout_generator = _attempt_timeout_generator( operation_timeout, operation_timeout ) @@ -623,9 +702,7 @@ async def read_rows_sharded( self.read_rows( query, operation_timeout=batch_operation_timeout, - per_request_timeout=min( - per_request_timeout, batch_operation_timeout - ), + attempt_timeout=min(attempt_timeout, batch_operation_timeout), ) for query in batch ] @@ -652,19 +729,33 @@ async def row_exists( self, row_key: str | bytes, *, - operation_timeout: int | float | None = 60, - per_request_timeout: int | float | None = None, + operation_timeout: int | float | None = None, + attempt_timeout: int | float | None = None, ) -> bool: """ - Helper function to determine if a row exists - + Return a boolean indicating whether the specified row exists in the table. uses the filters: chain(limit cells per row = 1, strip value) - + Args: + - row_key: the key of the row to check + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + If None, defaults to the Table's default_read_rows_operation_timeout + - attempt_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_read_rows_attempt_timeout, or the operation_timeout + if that is also None. Returns: - a bool indicating whether the row exists + Raises: + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions + from any retries that failed + - GoogleAPIError: raised if the request encounters an unrecoverable error """ if row_key is None: raise ValueError("row_key must be string or bytes") + strip_filter = StripValueTransformerFilter(flag=True) limit_filter = CellsRowLimitFilter(1) chain_filter = RowFilterChain(filters=[limit_filter, strip_filter]) @@ -672,7 +763,7 @@ async def row_exists( results = await self.read_rows( query, operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, + attempt_timeout=attempt_timeout, ) return len(results) > 0 @@ -680,7 +771,7 @@ async def sample_row_keys( self, *, operation_timeout: float | None = None, - per_request_timeout: float | None = None, + attempt_timeout: float | None = None, ) -> RowKeySamples: """ Return a set of RowKeySamples that delimit contiguous sections of the table of @@ -693,25 +784,32 @@ async def sample_row_keys( RowKeySamples is simply a type alias for list[tuple[bytes, int]]; a list of row_keys, along with offset positions in the table + Args: + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + If None, defaults to the Table's default_operation_timeout + - attempt_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_attempt_timeout, or the operation_timeout + if that is also None. Returns: - a set of RowKeySamples the delimit contiguous sections of the table Raises: - - GoogleAPICallError: if the sample_row_keys request fails + - DeadlineExceeded: raised after operation timeout + will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions + from any retries that failed + - GoogleAPIError: raised if the request encounters an unrecoverable error """ # prepare timeouts operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = per_request_timeout or self.default_per_request_timeout + attempt_timeout = ( + attempt_timeout or self.default_attempt_timeout or operation_timeout + ) + _validate_timeouts(operation_timeout, attempt_timeout) - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout <= 0: - raise ValueError("per_request_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout > operation_timeout: - raise ValueError( - "per_request_timeout must not be greater than operation_timeout" - ) attempt_timeout_gen = _attempt_timeout_generator( - per_request_timeout, operation_timeout + attempt_timeout, operation_timeout ) # prepare retryable predicate = retries.if_exception_type( @@ -761,7 +859,7 @@ def mutations_batcher( flow_control_max_mutation_count: int = 100_000, flow_control_max_bytes: int = 100 * _MB_SIZE, batch_operation_timeout: float | None = None, - batch_per_request_timeout: float | None = None, + batch_attempt_timeout: float | None = None, ) -> MutationsBatcherAsync: """ Returns a new mutations batcher instance. @@ -778,9 +876,10 @@ def mutations_batcher( - flow_control_max_mutation_count: Maximum number of inflight mutations. - flow_control_max_bytes: Maximum number of inflight bytes. - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None, - table default_operation_timeout will be used - - batch_per_request_timeout: timeout for each individual request, in seconds. If None, - table default_per_request_timeout will be used + table default_mutate_rows_operation_timeout will be used + - batch_attempt_timeout: timeout for each individual request, in seconds. If None, + table default_mutate_rows_attempt_timeout will be used, or batch_operation_timeout + if that is also None. Returns: - a MutationsBatcherAsync context manager that can batch requests """ @@ -792,7 +891,7 @@ def mutations_batcher( flow_control_max_mutation_count=flow_control_max_mutation_count, flow_control_max_bytes=flow_control_max_bytes, batch_operation_timeout=batch_operation_timeout, - batch_per_request_timeout=batch_per_request_timeout, + batch_attempt_timeout=batch_attempt_timeout, ) async def mutate_row( @@ -800,8 +899,8 @@ async def mutate_row( row_key: str | bytes, mutations: list[Mutation] | Mutation, *, - operation_timeout: float | None = 60, - per_request_timeout: float | None = None, + operation_timeout: float | None = None, + attempt_timeout: float | None = None, ): """ Mutates a row atomically. @@ -813,17 +912,16 @@ async def mutate_row( retried on server failure. Non-idempotent operations will not. Args: - - row_key: the row to apply mutations to - - mutations: the set of mutations to apply to the row - - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will be retried within the budget. - time is only counted while actively waiting on the network. - DeadlineExceeded exception raised after timeout - - per_request_timeout: the time budget for an individual network request, - in seconds. If it takes longer than this time to complete, the request - will be cancelled with a DeadlineExceeded exception, and a retry will be - attempted if within operation_timeout budget - + - row_key: the row to apply mutations to + - mutations: the set of mutations to apply to the row + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will be retried within the budget. + If None, defaults to the Table's default_operation_timeout + - attempt_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_attempt_timeout, or the operation_timeout + if that is also None. Raises: - DeadlineExceeded: raised after operation timeout will be chained with a RetryExceptionGroup containing all @@ -832,14 +930,10 @@ async def mutate_row( safely retried. """ operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = per_request_timeout or self.default_per_request_timeout - - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout <= 0: - raise ValueError("per_request_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout > operation_timeout: - raise ValueError("per_request_timeout must be less than operation_timeout") + attempt_timeout = ( + attempt_timeout or self.default_attempt_timeout or operation_timeout + ) + _validate_timeouts(operation_timeout, attempt_timeout) if isinstance(row_key, str): row_key = row_key.encode("utf-8") @@ -883,14 +977,16 @@ def on_error_fn(exc): ) metadata = _make_metadata(self.table_name, self.app_profile_id) # trigger rpc - await deadline_wrapped(request, timeout=per_request_timeout, metadata=metadata) + await deadline_wrapped( + request, timeout=attempt_timeout, metadata=metadata, retry=None + ) async def bulk_mutate_rows( self, mutation_entries: list[RowMutationEntry], *, - operation_timeout: float | None = 60, - per_request_timeout: float | None = None, + operation_timeout: float | None = None, + attempt_timeout: float | None = None, ): """ Applies mutations for multiple rows in a single batched request. @@ -910,32 +1006,32 @@ async def bulk_mutate_rows( in arbitrary order - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - time is only counted while actively waiting on the network. - DeadlineExceeded exception raised after timeout - - per_request_timeout: the time budget for an individual network request, - in seconds. If it takes longer than this time to complete, the request - will be cancelled with a DeadlineExceeded exception, and a retry will - be attempted if within operation_timeout budget + If None, defaults to the Table's default_mutate_rows_operation_timeout + - attempt_timeout: the time budget for an individual network request, in seconds. + If it takes longer than this time to complete, the request will be cancelled with + a DeadlineExceeded exception, and a retry will be attempted. + If None, defaults to the Table's default_mutate_rows_attempt_timeout, + or the operation_timeout if that is also None. Raises: - MutationsExceptionGroup if one or more mutations fails Contains details about any failed entries in .exceptions """ - operation_timeout = operation_timeout or self.default_operation_timeout - per_request_timeout = per_request_timeout or self.default_per_request_timeout - - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout <= 0: - raise ValueError("per_request_timeout must be greater than 0") - if per_request_timeout is not None and per_request_timeout > operation_timeout: - raise ValueError("per_request_timeout must be less than operation_timeout") + operation_timeout = ( + operation_timeout or self.default_mutate_rows_operation_timeout + ) + attempt_timeout = ( + attempt_timeout + or self.default_mutate_rows_attempt_timeout + or operation_timeout + ) + _validate_timeouts(operation_timeout, attempt_timeout) operation = _MutateRowsOperationAsync( self.client._gapic_client, self, mutation_entries, operation_timeout, - per_request_timeout, + attempt_timeout, ) await operation.start() @@ -946,7 +1042,7 @@ async def check_and_mutate_row( *, true_case_mutations: Mutation | list[Mutation] | None = None, false_case_mutations: Mutation | list[Mutation] | None = None, - operation_timeout: int | float | None = 20, + operation_timeout: int | float | None = None, ) -> bool: """ Mutates a row atomically based on the output of a predicate filter @@ -974,7 +1070,8 @@ async def check_and_mutate_row( ones. Must contain at least one entry if `true_case_mutations is empty, and at most 100000. - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will not be retried. + Failed requests will not be retried. Defaults to the Table's default_operation_timeout + if None. Returns: - bool indicating whether the predicate was true or false Raises: @@ -1016,7 +1113,7 @@ async def read_modify_write_row( row_key: str | bytes, rules: ReadModifyWriteRule | list[ReadModifyWriteRule], *, - operation_timeout: int | float | None = 20, + operation_timeout: int | float | None = None, ) -> Row: """ Reads and modifies a row atomically according to input ReadModifyWriteRules, @@ -1032,8 +1129,9 @@ async def read_modify_write_row( - rules: A rule or set of rules to apply to the row. Rules are applied in order, meaning that earlier rules will affect the results of later ones. - - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will not be retried. + - operation_timeout: the time budget for the entire operation, in seconds. + Failed requests will not be retried. Defaults to the Table's default_operation_timeout + if None. Returns: - Row: containing cell data that was modified as part of the operation diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index 25aafc2a1..2e3ad52dc 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -23,6 +23,7 @@ from google.cloud.bigtable.data.mutations import RowMutationEntry from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup from google.cloud.bigtable.data.exceptions import FailedMutationEntryError +from google.cloud.bigtable.data._helpers import _validate_timeouts from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync from google.cloud.bigtable.data._async._mutate_rows import ( @@ -189,7 +190,7 @@ def __init__( flow_control_max_mutation_count: int = 100_000, flow_control_max_bytes: int = 100 * _MB_SIZE, batch_operation_timeout: float | None = None, - batch_per_request_timeout: float | None = None, + batch_attempt_timeout: float | None = None, ): """ Args: @@ -202,26 +203,20 @@ def __init__( - flow_control_max_mutation_count: Maximum number of inflight mutations. - flow_control_max_bytes: Maximum number of inflight bytes. - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None, - table default_operation_timeout will be used - - batch_per_request_timeout: timeout for each individual request, in seconds. If None, - table default_per_request_timeout will be used + table default_mutate_rows_operation_timeout will be used + - batch_attempt_timeout: timeout for each individual request, in seconds. If None, + table default_mutate_rows_attempt_timeout will be used, or batch_operation_timeout + if that is also None. """ self._operation_timeout: float = ( - batch_operation_timeout or table.default_operation_timeout + batch_operation_timeout or table.default_mutate_rows_operation_timeout ) - self._per_request_timeout: float = ( - batch_per_request_timeout - or table.default_per_request_timeout + self._attempt_timeout: float = ( + batch_attempt_timeout + or table.default_mutate_rows_attempt_timeout or self._operation_timeout ) - if self._operation_timeout <= 0: - raise ValueError("batch_operation_timeout must be greater than 0") - if self._per_request_timeout <= 0: - raise ValueError("batch_per_request_timeout must be greater than 0") - if self._per_request_timeout > self._operation_timeout: - raise ValueError( - "batch_per_request_timeout must be less than batch_operation_timeout" - ) + _validate_timeouts(self._operation_timeout, self._attempt_timeout) self.closed: bool = False self._table = table self._staged_entries: list[RowMutationEntry] = [] @@ -346,7 +341,7 @@ async def _execute_mutate_rows( Args: - batch: list of RowMutationEntry objects to send to server - - timeout: timeout in seconds. Used as operation_timeout and per_request_timeout. + - timeout: timeout in seconds. Used as operation_timeout and attempt_timeout. If not given, will use table defaults Returns: - list of FailedMutationEntryError objects for mutations that failed. @@ -361,7 +356,7 @@ async def _execute_mutate_rows( self._table, batch, operation_timeout=self._operation_timeout, - per_request_timeout=self._per_request_timeout, + attempt_timeout=self._attempt_timeout, ) await operation.start() except MutationsExceptionGroup as e: diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 64d91e108..68f310b49 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -109,3 +109,26 @@ def wrapper(*args, **kwargs): handle_error() return wrapper_async if iscoroutinefunction(func) else wrapper + + +def _validate_timeouts( + operation_timeout: float, attempt_timeout: float | None, allow_none: bool = False +): + """ + Helper function that will verify that timeout values are valid, and raise + an exception if they are not. + + Args: + - operation_timeout: The timeout value to use for the entire operation, in seconds. + - attempt_timeout: The timeout value to use for each attempt, in seconds. + - allow_none: If True, attempt_timeout can be None. If False, None values will raise an exception. + Raises: + - ValueError if operation_timeout or attempt_timeout are invalid. + """ + if operation_timeout <= 0: + raise ValueError("operation_timeout must be greater than 0") + if not allow_none and attempt_timeout is None: + raise ValueError("attempt_timeout must not be None") + elif attempt_timeout is not None: + if attempt_timeout <= 0: + raise ValueError("attempt_timeout must be greater than 0") diff --git a/noxfile.py b/noxfile.py index 8499a610f..16447778e 100644 --- a/noxfile.py +++ b/noxfile.py @@ -366,10 +366,9 @@ def docfx(session): session.install("-e", ".") session.install( - "sphinx==4.0.1", + "gcp-sphinx-docfx-yaml", "alabaster", "recommonmark", - "gcp-sphinx-docfx-yaml", ) shutil.rmtree(os.path.join("docs", "_build"), ignore_errors=True) diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index f77455d60..9bebd35e6 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -48,7 +48,7 @@ def _make_one(self, *args, **kwargs): kwargs["table"] = kwargs.pop("table", AsyncMock()) kwargs["mutation_entries"] = kwargs.pop("mutation_entries", []) kwargs["operation_timeout"] = kwargs.pop("operation_timeout", 5) - kwargs["per_request_timeout"] = kwargs.pop("per_request_timeout", 0.1) + kwargs["attempt_timeout"] = kwargs.pop("attempt_timeout", 0.1) return self._target_class()(*args, **kwargs) async def _mock_stream(self, mutation_list, error_dict): @@ -267,7 +267,7 @@ async def test_run_attempt_single_entry_success(self): mock_gapic_fn = self._make_mock_gapic({0: mutation}) instance = self._make_one( mutation_entries=[mutation], - per_request_timeout=expected_timeout, + attempt_timeout=expected_timeout, ) with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): await instance._run_attempt() diff --git a/tests/unit/data/_async/test__read_rows.py b/tests/unit/data/_async/test__read_rows.py index c7d52280c..76e1148de 100644 --- a/tests/unit/data/_async/test__read_rows.py +++ b/tests/unit/data/_async/test__read_rows.py @@ -89,7 +89,7 @@ def test_ctor(self): request, client, operation_timeout=expected_operation_timeout, - per_request_timeout=expected_request_timeout, + attempt_timeout=expected_request_timeout, ) assert time_gen_mock.call_count == 1 time_gen_mock.assert_called_once_with( diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 25006d725..b7a7b3ae7 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -974,7 +974,11 @@ async def test_table_ctor(self): expected_instance_id = "instance-id" expected_app_profile_id = "app-profile-id" expected_operation_timeout = 123 - expected_per_request_timeout = 12 + expected_attempt_timeout = 12 + expected_read_rows_operation_timeout = 1.5 + expected_read_rows_attempt_timeout = 0.5 + expected_mutate_rows_operation_timeout = 2.5 + expected_mutate_rows_attempt_timeout = 0.75 client = BigtableDataClientAsync() assert not client._active_instances @@ -984,7 +988,11 @@ async def test_table_ctor(self): expected_table_id, expected_app_profile_id, default_operation_timeout=expected_operation_timeout, - default_per_request_timeout=expected_per_request_timeout, + default_attempt_timeout=expected_attempt_timeout, + default_read_rows_operation_timeout=expected_read_rows_operation_timeout, + default_read_rows_attempt_timeout=expected_read_rows_attempt_timeout, + default_mutate_rows_operation_timeout=expected_mutate_rows_operation_timeout, + default_mutate_rows_attempt_timeout=expected_mutate_rows_attempt_timeout, ) await asyncio.sleep(0) assert table.table_id == expected_table_id @@ -997,7 +1005,23 @@ async def test_table_ctor(self): assert instance_key in client._active_instances assert client._instance_owners[instance_key] == {id(table)} assert table.default_operation_timeout == expected_operation_timeout - assert table.default_per_request_timeout == expected_per_request_timeout + assert table.default_attempt_timeout == expected_attempt_timeout + assert ( + table.default_read_rows_operation_timeout + == expected_read_rows_operation_timeout + ) + assert ( + table.default_read_rows_attempt_timeout + == expected_read_rows_attempt_timeout + ) + assert ( + table.default_mutate_rows_operation_timeout + == expected_mutate_rows_operation_timeout + ) + assert ( + table.default_mutate_rows_attempt_timeout + == expected_mutate_rows_attempt_timeout + ) # ensure task reaches completion await table._register_instance_task assert table._register_instance_task.done() @@ -1006,30 +1030,64 @@ async def test_table_ctor(self): await client.close() @pytest.mark.asyncio - async def test_table_ctor_bad_timeout_values(self): + async def test_table_ctor_defaults(self): + """ + should provide default timeout values and app_profile_id + """ from google.cloud.bigtable.data._async.client import BigtableDataClientAsync from google.cloud.bigtable.data._async.client import TableAsync + expected_table_id = "table-id" + expected_instance_id = "instance-id" client = BigtableDataClientAsync() + assert not client._active_instances - with pytest.raises(ValueError) as e: - TableAsync(client, "", "", default_per_request_timeout=-1) - assert "default_per_request_timeout must be greater than 0" in str(e.value) - with pytest.raises(ValueError) as e: - TableAsync(client, "", "", default_operation_timeout=-1) - assert "default_operation_timeout must be greater than 0" in str(e.value) - with pytest.raises(ValueError) as e: - TableAsync( - client, - "", - "", - default_operation_timeout=1, - default_per_request_timeout=2, - ) - assert ( - "default_per_request_timeout must be less than default_operation_timeout" - in str(e.value) + table = TableAsync( + client, + expected_instance_id, + expected_table_id, ) + await asyncio.sleep(0) + assert table.table_id == expected_table_id + assert table.instance_id == expected_instance_id + assert table.app_profile_id is None + assert table.client is client + assert table.default_operation_timeout == 60 + assert table.default_read_rows_operation_timeout == 600 + assert table.default_mutate_rows_operation_timeout == 600 + assert table.default_attempt_timeout == 20 + assert table.default_read_rows_attempt_timeout == 20 + assert table.default_mutate_rows_attempt_timeout == 60 + await client.close() + + @pytest.mark.asyncio + async def test_table_ctor_invalid_timeout_values(self): + """ + bad timeout values should raise ValueError + """ + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync + from google.cloud.bigtable.data._async.client import TableAsync + + client = BigtableDataClientAsync() + + timeout_pairs = [ + ("default_operation_timeout", "default_attempt_timeout"), + ( + "default_read_rows_operation_timeout", + "default_read_rows_attempt_timeout", + ), + ( + "default_mutate_rows_operation_timeout", + "default_mutate_rows_attempt_timeout", + ), + ] + for operation_timeout, attempt_timeout in timeout_pairs: + with pytest.raises(ValueError) as e: + TableAsync(client, "", "", **{attempt_timeout: -1}) + assert "attempt_timeout must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + TableAsync(client, "", "", **{operation_timeout: -1}) + assert "operation_timeout must be greater than 0" in str(e.value) await client.close() def test_table_ctor_sync(self): @@ -1240,15 +1298,15 @@ async def test_read_rows_timeout(self, operation_timeout): ], ) @pytest.mark.asyncio - async def test_read_rows_per_request_timeout( + async def test_read_rows_attempt_timeout( self, per_request_t, operation_t, expected_num ): """ - Ensures that the per_request_timeout is respected and that the number of + Ensures that the attempt_timeout is respected and that the number of requests is as expected. operation_timeout does not cancel the request, so we expect the number of - requests to be the ceiling of operation_timeout / per_request_timeout. + requests to be the ceiling of operation_timeout / attempt_timeout. """ from google.cloud.bigtable.data.exceptions import RetryExceptionGroup @@ -1268,7 +1326,7 @@ async def test_read_rows_per_request_timeout( await table.read_rows( query, operation_timeout=operation_t, - per_request_timeout=per_request_t, + attempt_timeout=per_request_t, ) except core_exceptions.DeadlineExceeded as e: retry_exc = e.__cause__ @@ -1437,12 +1495,12 @@ async def test_read_rows_default_timeouts(self): from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync operation_timeout = 8 - per_request_timeout = 4 + attempt_timeout = 4 with mock.patch.object(_ReadRowsOperationAsync, "__init__") as mock_op: mock_op.side_effect = RuntimeError("mock error") async with self._make_table( - default_operation_timeout=operation_timeout, - default_per_request_timeout=per_request_timeout, + default_read_rows_operation_timeout=operation_timeout, + default_read_rows_attempt_timeout=attempt_timeout, ) as table: try: await table.read_rows(ReadRowsQuery()) @@ -1450,7 +1508,7 @@ async def test_read_rows_default_timeouts(self): pass kwargs = mock_op.call_args_list[0].kwargs assert kwargs["operation_timeout"] == operation_timeout - assert kwargs["per_request_timeout"] == per_request_timeout + assert kwargs["attempt_timeout"] == attempt_timeout @pytest.mark.asyncio async def test_read_rows_default_timeout_override(self): @@ -1460,23 +1518,23 @@ async def test_read_rows_default_timeout_override(self): from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync operation_timeout = 8 - per_request_timeout = 4 + attempt_timeout = 4 with mock.patch.object(_ReadRowsOperationAsync, "__init__") as mock_op: mock_op.side_effect = RuntimeError("mock error") async with self._make_table( - default_operation_timeout=99, default_per_request_timeout=97 + default_operation_timeout=99, default_attempt_timeout=97 ) as table: try: await table.read_rows( ReadRowsQuery(), operation_timeout=operation_timeout, - per_request_timeout=per_request_timeout, + attempt_timeout=attempt_timeout, ) except RuntimeError: pass kwargs = mock_op.call_args_list[0].kwargs assert kwargs["operation_timeout"] == operation_timeout - assert kwargs["per_request_timeout"] == per_request_timeout + assert kwargs["attempt_timeout"] == attempt_timeout @pytest.mark.asyncio async def test_read_row(self): @@ -1492,13 +1550,13 @@ async def test_read_row(self): row = await table.read_row( row_key, operation_timeout=expected_op_timeout, - per_request_timeout=expected_req_timeout, + attempt_timeout=expected_req_timeout, ) assert row == expected_result assert read_rows.call_count == 1 args, kwargs = read_rows.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout - assert kwargs["per_request_timeout"] == expected_req_timeout + assert kwargs["attempt_timeout"] == expected_req_timeout assert len(args) == 1 assert isinstance(args[0], ReadRowsQuery) assert args[0]._to_dict() == { @@ -1523,14 +1581,14 @@ async def test_read_row_w_filter(self): row = await table.read_row( row_key, operation_timeout=expected_op_timeout, - per_request_timeout=expected_req_timeout, + attempt_timeout=expected_req_timeout, row_filter=expected_filter, ) assert row == expected_result assert read_rows.call_count == 1 args, kwargs = read_rows.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout - assert kwargs["per_request_timeout"] == expected_req_timeout + assert kwargs["attempt_timeout"] == expected_req_timeout assert len(args) == 1 assert isinstance(args[0], ReadRowsQuery) assert args[0]._to_dict() == { @@ -1553,13 +1611,13 @@ async def test_read_row_no_response(self): result = await table.read_row( row_key, operation_timeout=expected_op_timeout, - per_request_timeout=expected_req_timeout, + attempt_timeout=expected_req_timeout, ) assert result is None assert read_rows.call_count == 1 args, kwargs = read_rows.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout - assert kwargs["per_request_timeout"] == expected_req_timeout + assert kwargs["attempt_timeout"] == expected_req_timeout assert isinstance(args[0], ReadRowsQuery) assert args[0]._to_dict() == { "rows": {"row_keys": [row_key], "row_ranges": []}, @@ -1598,13 +1656,13 @@ async def test_row_exists(self, return_value, expected_result): result = await table.row_exists( row_key, operation_timeout=expected_op_timeout, - per_request_timeout=expected_req_timeout, + attempt_timeout=expected_req_timeout, ) assert expected_result == result assert read_rows.call_count == 1 args, kwargs = read_rows.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout - assert kwargs["per_request_timeout"] == expected_req_timeout + assert kwargs["attempt_timeout"] == expected_req_timeout assert isinstance(args[0], ReadRowsQuery) expected_filter = { "chain": { @@ -1798,9 +1856,9 @@ async def test_read_rows_sharded_batching(self): table_mock = AsyncMock() start_operation_timeout = 10 - start_per_request_timeout = 3 - table_mock.default_operation_timeout = start_operation_timeout - table_mock.default_per_request_timeout = start_per_request_timeout + start_attempt_timeout = 3 + table_mock.default_read_rows_operation_timeout = start_operation_timeout + table_mock.default_read_rows_attempt_timeout = start_attempt_timeout # clock ticks one second on each check with mock.patch("time.monotonic", side_effect=range(0, 100000)): with mock.patch("asyncio.gather", AsyncMock()) as gather_mock: @@ -1829,14 +1887,11 @@ async def test_read_rows_sharded_batching(self): req_kwargs["operation_timeout"] == expected_operation_timeout ) - # each per_request_timeout should start with default value, but decrease when operation_timeout reaches it - expected_per_request_timeout = min( - start_per_request_timeout, expected_operation_timeout - ) - assert ( - req_kwargs["per_request_timeout"] - == expected_per_request_timeout + # each attempt_timeout should start with default value, but decrease when operation_timeout reaches it + expected_attempt_timeout = min( + start_attempt_timeout, expected_operation_timeout ) + assert req_kwargs["attempt_timeout"] == expected_attempt_timeout # await all created coroutines to avoid warnings for i in range(len(gather_mock.call_args_list)): for j in range(len(gather_mock.call_args_list[i][0])): @@ -1891,16 +1946,8 @@ async def test_sample_row_keys_bad_timeout(self): await table.sample_row_keys(operation_timeout=-1) assert "operation_timeout must be greater than 0" in str(e.value) with pytest.raises(ValueError) as e: - await table.sample_row_keys(per_request_timeout=-1) - assert "per_request_timeout must be greater than 0" in str(e.value) - with pytest.raises(ValueError) as e: - await table.sample_row_keys( - operation_timeout=10, per_request_timeout=20 - ) - assert ( - "per_request_timeout must not be greater than operation_timeout" - in str(e.value) - ) + await table.sample_row_keys(attempt_timeout=-1) + assert "attempt_timeout must be greater than 0" in str(e.value) @pytest.mark.asyncio async def test_sample_row_keys_default_timeout(self): @@ -1936,7 +1983,7 @@ async def test_sample_row_keys_gapic_params(self): table.client._gapic_client, "sample_row_keys", AsyncMock() ) as sample_row_keys: sample_row_keys.return_value = self._make_gapic_stream([]) - await table.sample_row_keys(per_request_timeout=expected_timeout) + await table.sample_row_keys(attempt_timeout=expected_timeout) args, kwargs = sample_row_keys.call_args assert len(args) == 0 assert len(kwargs) == 4 @@ -2049,7 +2096,7 @@ def _make_client(self, *args, **kwargs): ) async def test_mutate_row(self, mutation_arg): """Test mutations with no errors""" - expected_per_request_timeout = 19 + expected_attempt_timeout = 19 async with self._make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( @@ -2059,9 +2106,10 @@ async def test_mutate_row(self, mutation_arg): await table.mutate_row( "row_key", mutation_arg, - per_request_timeout=expected_per_request_timeout, + attempt_timeout=expected_attempt_timeout, ) assert mock_gapic.call_count == 1 + kwargs = mock_gapic.call_args_list[0].kwargs request = mock_gapic.call_args[0][0] assert ( request["table_name"] @@ -2074,8 +2122,9 @@ async def test_mutate_row(self, mutation_arg): else [mutation_arg._to_dict()] ) assert request["mutations"] == formatted_mutations - found_per_request_timeout = mock_gapic.call_args[1]["timeout"] - assert found_per_request_timeout == expected_per_request_timeout + assert kwargs["timeout"] == expected_attempt_timeout + # make sure gapic layer is not retrying + assert kwargs["retry"] is None @pytest.mark.parametrize( "retryable_exception", @@ -2243,7 +2292,7 @@ async def generator(): ) async def test_bulk_mutate_rows(self, mutation_arg): """Test mutations with no errors""" - expected_per_request_timeout = 19 + expected_attempt_timeout = 19 async with self._make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( @@ -2253,7 +2302,7 @@ async def test_bulk_mutate_rows(self, mutation_arg): bulk_mutation = mutations.RowMutationEntry(b"row_key", mutation_arg) await table.bulk_mutate_rows( [bulk_mutation], - per_request_timeout=expected_per_request_timeout, + attempt_timeout=expected_attempt_timeout, ) assert mock_gapic.call_count == 1 kwargs = mock_gapic.call_args[1] @@ -2262,7 +2311,7 @@ async def test_bulk_mutate_rows(self, mutation_arg): == "projects/project/instances/instance/tables/table" ) assert kwargs["entries"] == [bulk_mutation._to_dict()] - assert kwargs["timeout"] == expected_per_request_timeout + assert kwargs["timeout"] == expected_attempt_timeout @pytest.mark.asyncio async def test_bulk_mutate_rows_multiple_entries(self): diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py index 1b14cc128..a55c4f6ef 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -288,8 +288,8 @@ def _get_target_class(self): def _make_one(self, table=None, **kwargs): if table is None: table = mock.Mock() - table.default_operation_timeout = 10 - table.default_per_request_timeout = 10 + table.default_mutate_rows_operation_timeout = 10 + table.default_mutate_rows_attempt_timeout = 10 return self._get_target_class()(table, **kwargs) @@ -300,8 +300,8 @@ def _make_one(self, table=None, **kwargs): async def test_ctor_defaults(self, flush_timer_mock): flush_timer_mock.return_value = asyncio.create_task(asyncio.sleep(0)) table = mock.Mock() - table.default_operation_timeout = 10 - table.default_per_request_timeout = 8 + table.default_mutate_rows_operation_timeout = 10 + table.default_mutate_rows_attempt_timeout = 8 async with self._make_one(table) as instance: assert instance._table == table assert instance.closed is False @@ -316,8 +316,13 @@ async def test_ctor_defaults(self, flush_timer_mock): assert instance._flow_control._in_flight_mutation_count == 0 assert instance._flow_control._in_flight_mutation_bytes == 0 assert instance._entries_processed_since_last_raise == 0 - assert instance._operation_timeout == table.default_operation_timeout - assert instance._per_request_timeout == table.default_per_request_timeout + assert ( + instance._operation_timeout + == table.default_mutate_rows_operation_timeout + ) + assert ( + instance._attempt_timeout == table.default_mutate_rows_attempt_timeout + ) await asyncio.sleep(0) assert flush_timer_mock.call_count == 1 assert flush_timer_mock.call_args[0][0] == 5 @@ -337,7 +342,7 @@ async def test_ctor_explicit(self, flush_timer_mock): flow_control_max_mutation_count = 1001 flow_control_max_bytes = 12 operation_timeout = 11 - per_request_timeout = 2 + attempt_timeout = 2 async with self._make_one( table, flush_interval=flush_interval, @@ -346,7 +351,7 @@ async def test_ctor_explicit(self, flush_timer_mock): flow_control_max_mutation_count=flow_control_max_mutation_count, flow_control_max_bytes=flow_control_max_bytes, batch_operation_timeout=operation_timeout, - batch_per_request_timeout=per_request_timeout, + batch_attempt_timeout=attempt_timeout, ) as instance: assert instance._table == table assert instance.closed is False @@ -365,7 +370,7 @@ async def test_ctor_explicit(self, flush_timer_mock): assert instance._flow_control._in_flight_mutation_bytes == 0 assert instance._entries_processed_since_last_raise == 0 assert instance._operation_timeout == operation_timeout - assert instance._per_request_timeout == per_request_timeout + assert instance._attempt_timeout == attempt_timeout await asyncio.sleep(0) assert flush_timer_mock.call_count == 1 assert flush_timer_mock.call_args[0][0] == flush_interval @@ -379,8 +384,8 @@ async def test_ctor_no_flush_limits(self, flush_timer_mock): """Test with None for flush limits""" flush_timer_mock.return_value = asyncio.create_task(asyncio.sleep(0)) table = mock.Mock() - table.default_operation_timeout = 10 - table.default_per_request_timeout = 8 + table.default_mutate_rows_operation_timeout = 10 + table.default_mutate_rows_attempt_timeout = 8 flush_interval = None flush_limit_count = None flush_limit_bytes = None @@ -410,16 +415,10 @@ async def test_ctor_invalid_values(self): """Test that timeout values are positive, and fit within expected limits""" with pytest.raises(ValueError) as e: self._make_one(batch_operation_timeout=-1) - assert "batch_operation_timeout must be greater than 0" in str(e.value) - with pytest.raises(ValueError) as e: - self._make_one(batch_per_request_timeout=-1) - assert "batch_per_request_timeout must be greater than 0" in str(e.value) + assert "operation_timeout must be greater than 0" in str(e.value) with pytest.raises(ValueError) as e: - self._make_one(batch_operation_timeout=1, batch_per_request_timeout=2) - assert ( - "batch_per_request_timeout must be less than batch_operation_timeout" - in str(e.value) - ) + self._make_one(batch_attempt_timeout=-1) + assert "attempt_timeout must be greater than 0" in str(e.value) def test_default_argument_consistency(self): """ @@ -857,7 +856,7 @@ async def test_timer_flush_end_to_end(self): async with self._make_one(flush_interval=0.05) as instance: instance._table.default_operation_timeout = 10 - instance._table.default_per_request_timeout = 9 + instance._table.default_attempt_timeout = 9 with mock.patch.object( instance._table.client._gapic_client, "mutate_rows" ) as gapic_mock: @@ -881,8 +880,8 @@ async def test__execute_mutate_rows(self, mutate_rows): table = mock.Mock() table.table_name = "test-table" table.app_profile_id = "test-app-profile" - table.default_operation_timeout = 17 - table.default_per_request_timeout = 13 + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 async with self._make_one(table) as instance: batch = [_make_mutation()] result = await instance._execute_mutate_rows(batch) @@ -892,7 +891,7 @@ async def test__execute_mutate_rows(self, mutate_rows): assert args[1] == table assert args[2] == batch kwargs["operation_timeout"] == 17 - kwargs["per_request_timeout"] == 13 + kwargs["attempt_timeout"] == 13 assert result == [] @pytest.mark.asyncio @@ -910,8 +909,8 @@ async def test__execute_mutate_rows_returns_errors(self, mutate_rows): err2 = FailedMutationEntryError(1, mock.Mock(), RuntimeError("test error")) mutate_rows.side_effect = MutationsExceptionGroup([err1, err2], 10) table = mock.Mock() - table.default_operation_timeout = 17 - table.default_per_request_timeout = 13 + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 async with self._make_one(table) as instance: batch = [_make_mutation()] result = await instance._execute_mutate_rows(batch) @@ -1026,24 +1025,24 @@ async def test_atexit_registration(self): ) async def test_timeout_args_passed(self, mutate_rows): """ - batch_operation_timeout and batch_per_request_timeout should be used + batch_operation_timeout and batch_attempt_timeout should be used in api calls """ mutate_rows.return_value = AsyncMock() expected_operation_timeout = 17 - expected_per_request_timeout = 13 + expected_attempt_timeout = 13 async with self._make_one( batch_operation_timeout=expected_operation_timeout, - batch_per_request_timeout=expected_per_request_timeout, + batch_attempt_timeout=expected_attempt_timeout, ) as instance: assert instance._operation_timeout == expected_operation_timeout - assert instance._per_request_timeout == expected_per_request_timeout + assert instance._attempt_timeout == expected_attempt_timeout # make simulated gapic call await instance._execute_mutate_rows([_make_mutation()]) assert mutate_rows.call_count == 1 kwargs = mutate_rows.call_args[1] assert kwargs["operation_timeout"] == expected_operation_timeout - assert kwargs["per_request_timeout"] == expected_per_request_timeout + assert kwargs["attempt_timeout"] == expected_attempt_timeout @pytest.mark.parametrize( "limit,in_e,start_e,end_e", diff --git a/tests/unit/data/test__helpers.py b/tests/unit/data/test__helpers.py index dc688bb0c..e3bfd4750 100644 --- a/tests/unit/data/test__helpers.py +++ b/tests/unit/data/test__helpers.py @@ -143,3 +143,42 @@ async def test_func(): assert isinstance(cause, bigtable_exceptions.RetryExceptionGroup) assert cause.exceptions == tuple(associated_errors) assert f"operation_timeout of {timeout}s exceeded" in str(e.value) + + +class TestValidateTimeouts: + def test_validate_timeouts_error_messages(self): + with pytest.raises(ValueError) as e: + _helpers._validate_timeouts(operation_timeout=1, attempt_timeout=-1) + assert "attempt_timeout must be greater than 0" in str(e.value) + with pytest.raises(ValueError) as e: + _helpers._validate_timeouts(operation_timeout=-1, attempt_timeout=1) + assert "operation_timeout must be greater than 0" in str(e.value) + + @pytest.mark.parametrize( + "args,expected", + [ + ([1, None, False], False), + ([1, None, True], True), + ([1, 1, False], True), + ([1, 1, True], True), + ([1, 1], True), + ([1, None], False), + ([2, 1], True), + ([0, 1], False), + ([1, 0], False), + ([60, None], False), + ([600, None], False), + ([600, 600], True), + ], + ) + def test_validate_with_inputs(self, args, expected): + """ + test whether an exception is thrown with different inputs + """ + success = False + try: + _helpers._validate_timeouts(*args) + success = True + except ValueError: + pass + assert success == expected From 0d92a84ef498ea4a9e7dd0c63bfe0a92e6c17824 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 24 Jul 2023 16:14:20 -0700 Subject: [PATCH 16/56] fix: api errors apply to all bulk mutations --- .../bigtable/data/_async/_mutate_rows.py | 4 +-- tests/unit/data/_async/test__mutate_rows.py | 34 +++++++++++++++++-- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 519043a92..0f197e051 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -186,8 +186,8 @@ async def _run_attempt(self): # already handled, and update remaining_indices if mutation is retryable for idx in active_request_indices.values(): self._handle_entry_error(idx, exc) - # bubble up exception to be handled by retry wrapper - raise + # bubble up exception to be handled by retry wrapper + raise # check if attempt succeeded, or needs to be retried if self.remaining_indices: # unfinished work; raise exception to trigger retry diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index 9bebd35e6..82cf135ac 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -162,6 +162,34 @@ async def test_mutate_rows_operation(self): await instance.start() assert attempt_mock.call_count == 1 + @pytest.mark.parametrize( + "exc_type", [RuntimeError, ZeroDivisionError, core_exceptions.Forbidden] + ) + @pytest.mark.asyncio + async def test_mutate_rows_attempt_exception(self, exc_type): + """ + exceptions raised from attempt should be raised in MutationsExceptionGroup + """ + client = AsyncMock() + table = mock.Mock() + entries = [_make_mutation(), _make_mutation()] + operation_timeout = 0.05 + expected_exception = exc_type("test") + client.mutate_rows.side_effect = expected_exception + found_exc = None + try: + instance = self._make_one( + client, table, entries, operation_timeout, operation_timeout + ) + await instance._run_attempt() + except Exception as e: + found_exc = e + assert client.mutate_rows.call_count == 1 + assert type(found_exc) == exc_type + assert found_exc == expected_exception + assert len(instance.errors) == 2 + assert len(instance.remaining_indices) == 0 + @pytest.mark.parametrize( "exc_type", [RuntimeError, ZeroDivisionError, core_exceptions.Forbidden] ) @@ -175,7 +203,7 @@ async def test_mutate_rows_exception(self, exc_type): client = mock.Mock() table = mock.Mock() - entries = [_make_mutation()] + entries = [_make_mutation(), _make_mutation()] operation_timeout = 0.05 expected_cause = exc_type("abort") with mock.patch.object( @@ -193,9 +221,11 @@ async def test_mutate_rows_exception(self, exc_type): except MutationsExceptionGroup as e: found_exc = e assert attempt_mock.call_count == 1 - assert len(found_exc.exceptions) == 1 + assert len(found_exc.exceptions) == 2 assert isinstance(found_exc.exceptions[0], FailedMutationEntryError) + assert isinstance(found_exc.exceptions[1], FailedMutationEntryError) assert found_exc.exceptions[0].__cause__ == expected_cause + assert found_exc.exceptions[1].__cause__ == expected_cause @pytest.mark.parametrize( "exc_type", From a8cdf7cac86b5d2207bcc14d295d86ee4eac8308 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 1 Aug 2023 22:45:26 +0000 Subject: [PATCH 17/56] chore: reduce public api surface (#820) --- .../bigtable/data/_async/_mutate_rows.py | 6 +- google/cloud/bigtable/data/_async/client.py | 28 +-- .../bigtable/data/_async/mutations_batcher.py | 4 +- google/cloud/bigtable/data/mutations.py | 20 +-- .../bigtable/data/read_modify_write_rules.py | 4 +- google/cloud/bigtable/data/read_rows_query.py | 142 ++++++++++++--- google/cloud/bigtable/data/row.py | 18 +- google/cloud/bigtable/data/row_filters.py | 74 ++++---- tests/system/data/test_system.py | 8 +- tests/unit/data/_async/test__mutate_rows.py | 6 +- tests/unit/data/_async/test_client.py | 48 ++--- .../data/_async/test_mutations_batcher.py | 2 +- tests/unit/data/test_mutations.py | 6 +- tests/unit/data/test_read_rows_query.py | 164 ++++++++++++++---- tests/unit/data/test_row.py | 18 +- tests/unit/data/test_row_filters.py | 90 +++++----- 16 files changed, 408 insertions(+), 230 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 0f197e051..d41f1aeea 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -25,7 +25,7 @@ from google.cloud.bigtable.data._helpers import _attempt_timeout_generator # mutate_rows requests are limited to this number of mutations -from google.cloud.bigtable.data.mutations import MUTATE_ROWS_REQUEST_MUTATION_LIMIT +from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT if TYPE_CHECKING: from google.cloud.bigtable_v2.services.bigtable.async_client import ( @@ -65,10 +65,10 @@ def __init__( """ # check that mutations are within limits total_mutations = sum(len(entry.mutations) for entry in mutation_entries) - if total_mutations > MUTATE_ROWS_REQUEST_MUTATION_LIMIT: + if total_mutations > _MUTATE_ROWS_REQUEST_MUTATION_LIMIT: raise ValueError( "mutate_rows requests can contain at most " - f"{MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations across " + f"{_MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations across " f"all entries. Found {total_mutations}." ) # create partial function to pass to trigger rpc call diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 6c321fe62..0f8748bab 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -74,7 +74,7 @@ from google.cloud.bigtable.data import ShardedQuery # used by read_rows_sharded to limit how many requests are attempted in parallel -CONCURRENCY_LIMIT = 10 +_CONCURRENCY_LIMIT = 10 # used to register instance data with the client for channel warming _WarmedInstanceKey = namedtuple( @@ -154,7 +154,7 @@ def __init__( self._channel_init_time = time.monotonic() self._channel_refresh_tasks: list[asyncio.Task[None]] = [] try: - self.start_background_channel_refresh() + self._start_background_channel_refresh() except RuntimeError: warnings.warn( f"{self.__class__.__name__} should be started in an " @@ -163,7 +163,7 @@ def __init__( stacklevel=2, ) - def start_background_channel_refresh(self) -> None: + def _start_background_channel_refresh(self) -> None: """ Starts a background task to ping and warm each channel in the pool Raises: @@ -309,7 +309,7 @@ async def _register_instance(self, instance_id: str, owner: TableAsync) -> None: await self._ping_and_warm_instances(channel, instance_key) else: # refresh tasks aren't active. start them as background tasks - self.start_background_channel_refresh() + self._start_background_channel_refresh() async def _remove_instance_registration( self, instance_id: str, owner: TableAsync @@ -388,7 +388,7 @@ def get_table( ) async def __aenter__(self): - self.start_background_channel_refresh() + self._start_background_channel_refresh() return self async def __aexit__(self, exc_type, exc_val, exc_tb): @@ -493,7 +493,7 @@ def __init__( async def read_rows_stream( self, - query: ReadRowsQuery | dict[str, Any], + query: ReadRowsQuery, *, operation_timeout: float | None = None, attempt_timeout: float | None = None, @@ -557,7 +557,7 @@ async def read_rows_stream( async def read_rows( self, - query: ReadRowsQuery | dict[str, Any], + query: ReadRowsQuery, *, operation_timeout: float | None = None, attempt_timeout: float | None = None, @@ -687,10 +687,10 @@ async def read_rows_sharded( timeout_generator = _attempt_timeout_generator( operation_timeout, operation_timeout ) - # submit shards in batches if the number of shards goes over CONCURRENCY_LIMIT + # submit shards in batches if the number of shards goes over _CONCURRENCY_LIMIT batched_queries = [ - sharded_query[i : i + CONCURRENCY_LIMIT] - for i in range(0, len(sharded_query), CONCURRENCY_LIMIT) + sharded_query[i : i + _CONCURRENCY_LIMIT] + for i in range(0, len(sharded_query), _CONCURRENCY_LIMIT) ] # run batches and collect results results_list = [] @@ -1038,7 +1038,7 @@ async def bulk_mutate_rows( async def check_and_mutate_row( self, row_key: str | bytes, - predicate: RowFilter | dict[str, Any] | None, + predicate: RowFilter | None, *, true_case_mutations: Mutation | list[Mutation] | None = None, false_case_mutations: Mutation | list[Mutation] | None = None, @@ -1091,12 +1091,12 @@ async def check_and_mutate_row( ): false_case_mutations = [false_case_mutations] false_case_dict = [m._to_dict() for m in false_case_mutations or []] - if predicate is not None and not isinstance(predicate, dict): - predicate = predicate.to_dict() metadata = _make_metadata(self.table_name, self.app_profile_id) result = await self.client._gapic_client.check_and_mutate_row( request={ - "predicate_filter": predicate, + "predicate_filter": predicate._to_dict() + if predicate is not None + else None, "true_mutations": true_case_dict, "false_mutations": false_case_dict, "table_name": self.table_name, diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index 2e3ad52dc..34e1bfb5d 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -27,7 +27,7 @@ from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync from google.cloud.bigtable.data._async._mutate_rows import ( - MUTATE_ROWS_REQUEST_MUTATION_LIMIT, + _MUTATE_ROWS_REQUEST_MUTATION_LIMIT, ) from google.cloud.bigtable.data.mutations import Mutation @@ -144,7 +144,7 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry] self._has_capacity(next_count, next_size) # make sure not to exceed per-request mutation count limits and (batch_mutation_count + next_count) - <= MUTATE_ROWS_REQUEST_MUTATION_LIMIT + <= _MUTATE_ROWS_REQUEST_MUTATION_LIMIT ): # room for new mutation; add to batch end_idx += 1 diff --git a/google/cloud/bigtable/data/mutations.py b/google/cloud/bigtable/data/mutations.py index de1b3b137..06db21879 100644 --- a/google/cloud/bigtable/data/mutations.py +++ b/google/cloud/bigtable/data/mutations.py @@ -19,14 +19,14 @@ from abc import ABC, abstractmethod from sys import getsizeof -from google.cloud.bigtable.data.read_modify_write_rules import MAX_INCREMENT_VALUE -# special value for SetCell mutation timestamps. If set, server will assign a timestamp -SERVER_SIDE_TIMESTAMP = -1 +from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE +# special value for SetCell mutation timestamps. If set, server will assign a timestamp +_SERVER_SIDE_TIMESTAMP = -1 # mutation entries above this should be rejected -MUTATE_ROWS_REQUEST_MUTATION_LIMIT = 100_000 +_MUTATE_ROWS_REQUEST_MUTATION_LIMIT = 100_000 class Mutation(ABC): @@ -112,7 +112,7 @@ def __init__( if isinstance(new_value, str): new_value = new_value.encode() elif isinstance(new_value, int): - if abs(new_value) > MAX_INCREMENT_VALUE: + if abs(new_value) > _MAX_INCREMENT_VALUE: raise ValueError( "int values must be between -2**63 and 2**63 (64-bit signed int)" ) @@ -123,9 +123,9 @@ def __init__( # use current timestamp, with milisecond precision timestamp_micros = time.time_ns() // 1000 timestamp_micros = timestamp_micros - (timestamp_micros % 1000) - if timestamp_micros < SERVER_SIDE_TIMESTAMP: + if timestamp_micros < _SERVER_SIDE_TIMESTAMP: raise ValueError( - "timestamp_micros must be positive (or -1 for server-side timestamp)" + f"timestamp_micros must be positive (or {_SERVER_SIDE_TIMESTAMP} for server-side timestamp)" ) self.family = family self.qualifier = qualifier @@ -145,7 +145,7 @@ def _to_dict(self) -> dict[str, Any]: def is_idempotent(self) -> bool: """Check if the mutation is idempotent""" - return self.timestamp_micros != SERVER_SIDE_TIMESTAMP + return self.timestamp_micros != _SERVER_SIDE_TIMESTAMP @dataclass @@ -208,9 +208,9 @@ def __init__(self, row_key: bytes | str, mutations: Mutation | list[Mutation]): mutations = [mutations] if len(mutations) == 0: raise ValueError("mutations must not be empty") - elif len(mutations) > MUTATE_ROWS_REQUEST_MUTATION_LIMIT: + elif len(mutations) > _MUTATE_ROWS_REQUEST_MUTATION_LIMIT: raise ValueError( - f"entries must have <= {MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations" + f"entries must have <= {_MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations" ) self.row_key = row_key self.mutations = tuple(mutations) diff --git a/google/cloud/bigtable/data/read_modify_write_rules.py b/google/cloud/bigtable/data/read_modify_write_rules.py index aa282b1a6..3a3eb3752 100644 --- a/google/cloud/bigtable/data/read_modify_write_rules.py +++ b/google/cloud/bigtable/data/read_modify_write_rules.py @@ -17,7 +17,7 @@ import abc # value must fit in 64-bit signed integer -MAX_INCREMENT_VALUE = (1 << 63) - 1 +_MAX_INCREMENT_VALUE = (1 << 63) - 1 class ReadModifyWriteRule(abc.ABC): @@ -37,7 +37,7 @@ class IncrementRule(ReadModifyWriteRule): def __init__(self, family: str, qualifier: bytes | str, increment_amount: int = 1): if not isinstance(increment_amount, int): raise TypeError("increment_amount must be an integer") - if abs(increment_amount) > MAX_INCREMENT_VALUE: + if abs(increment_amount) > _MAX_INCREMENT_VALUE: raise ValueError( "increment_amount must be between -2**63 and 2**63 (64-bit signed int)" ) diff --git a/google/cloud/bigtable/data/read_rows_query.py b/google/cloud/bigtable/data/read_rows_query.py index 7d7e1f99f..cf3cd316c 100644 --- a/google/cloud/bigtable/data/read_rows_query.py +++ b/google/cloud/bigtable/data/read_rows_query.py @@ -35,11 +35,16 @@ class _RangePoint: def __hash__(self) -> int: return hash((self.key, self.is_inclusive)) + def __eq__(self, other: Any) -> bool: + if not isinstance(other, _RangePoint): + return NotImplemented + return self.key == other.key and self.is_inclusive == other.is_inclusive + -@dataclass class RowRange: - start: _RangePoint | None - end: _RangePoint | None + """ + Represents a range of keys in a ReadRowsQuery + """ def __init__( self, @@ -48,11 +53,27 @@ def __init__( start_is_inclusive: bool | None = None, end_is_inclusive: bool | None = None, ): + """ + Args: + - start_key: The start key of the range. If empty, the range is unbounded on the left. + - end_key: The end key of the range. If empty, the range is unbounded on the right. + - start_is_inclusive: Whether the start key is inclusive. If None, the start key is + inclusive. + - end_is_inclusive: Whether the end key is inclusive. If None, the end key is not inclusive. + Raises: + - ValueError: if start_key is greater than end_key, or start_is_inclusive, + or end_is_inclusive is set when the corresponding key is None, + or start_key or end_key is not a string or bytes. + """ + # convert empty key inputs to None for consistency + start_key = None if not start_key else start_key + end_key = None if not end_key else end_key # check for invalid combinations of arguments if start_is_inclusive is None: start_is_inclusive = True elif start_key is None: raise ValueError("start_is_inclusive must be set with start_key") + if end_is_inclusive is None: end_is_inclusive = False elif end_key is None: @@ -66,29 +87,62 @@ def __init__( end_key = end_key.encode() elif end_key is not None and not isinstance(end_key, bytes): raise ValueError("end_key must be a string or bytes") + # ensure that start_key is less than or equal to end_key + if start_key is not None and end_key is not None and start_key > end_key: + raise ValueError("start_key must be less than or equal to end_key") - self.start = ( + self._start: _RangePoint | None = ( _RangePoint(start_key, start_is_inclusive) if start_key is not None else None ) - self.end = ( + self._end: _RangePoint | None = ( _RangePoint(end_key, end_is_inclusive) if end_key is not None else None ) + @property + def start_key(self) -> bytes | None: + """ + Returns the start key of the range. If None, the range is unbounded on the left. + """ + return self._start.key if self._start is not None else None + + @property + def end_key(self) -> bytes | None: + """ + Returns the end key of the range. If None, the range is unbounded on the right. + """ + return self._end.key if self._end is not None else None + + @property + def start_is_inclusive(self) -> bool: + """ + Returns whether the range is inclusive of the start key. + Returns True if the range is unbounded on the left. + """ + return self._start.is_inclusive if self._start is not None else True + + @property + def end_is_inclusive(self) -> bool: + """ + Returns whether the range is inclusive of the end key. + Returns True if the range is unbounded on the right. + """ + return self._end.is_inclusive if self._end is not None else True + def _to_dict(self) -> dict[str, bytes]: """Converts this object to a dictionary""" output = {} - if self.start is not None: - key = "start_key_closed" if self.start.is_inclusive else "start_key_open" - output[key] = self.start.key - if self.end is not None: - key = "end_key_closed" if self.end.is_inclusive else "end_key_open" - output[key] = self.end.key + if self._start is not None: + key = "start_key_closed" if self.start_is_inclusive else "start_key_open" + output[key] = self._start.key + if self._end is not None: + key = "end_key_closed" if self.end_is_inclusive else "end_key_open" + output[key] = self._end.key return output def __hash__(self) -> int: - return hash((self.start, self.end)) + return hash((self._start, self._end)) @classmethod def _from_dict(cls, data: dict[str, bytes]) -> RowRange: @@ -123,7 +177,35 @@ def __bool__(self) -> bool: Empty RowRanges (representing a full table scan) are falsy, because they can be substituted with None. Non-empty RowRanges are truthy. """ - return self.start is not None or self.end is not None + return self._start is not None or self._end is not None + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, RowRange): + return NotImplemented + return self._start == other._start and self._end == other._end + + def __str__(self) -> str: + """ + Represent range as a string, e.g. "[b'a', b'z)" + Unbounded start or end keys are represented as "-inf" or "+inf" + """ + left = "[" if self.start_is_inclusive else "(" + right = "]" if self.end_is_inclusive else ")" + start = repr(self.start_key) if self.start_key is not None else "-inf" + end = repr(self.end_key) if self.end_key is not None else "+inf" + return f"{left}{start}, {end}{right}" + + def __repr__(self) -> str: + args_list = [] + args_list.append(f"start_key={self.start_key!r}") + args_list.append(f"end_key={self.end_key!r}") + if self.start_is_inclusive is False: + # only show start_is_inclusive if it is different from the default + args_list.append(f"start_is_inclusive={self.start_is_inclusive}") + if self.end_is_inclusive is True and self._end is not None: + # only show end_is_inclusive if it is different from the default + args_list.append(f"end_is_inclusive={self.end_is_inclusive}") + return f"RowRange({', '.join(args_list)})" class ReadRowsQuery: @@ -136,7 +218,7 @@ def __init__( row_keys: list[str | bytes] | str | bytes | None = None, row_ranges: list[RowRange] | RowRange | None = None, limit: int | None = None, - row_filter: RowFilter | dict[str, Any] | None = None, + row_filter: RowFilter | None = None, ): """ Create a new ReadRowsQuery @@ -162,7 +244,7 @@ def __init__( for k in row_keys: self.add_key(k) self.limit: int | None = limit - self.filter: RowFilter | dict[str, Any] | None = row_filter + self.filter: RowFilter | None = row_filter @property def limit(self) -> int | None: @@ -187,11 +269,11 @@ def limit(self, new_limit: int | None): self._limit = new_limit @property - def filter(self) -> RowFilter | dict[str, Any] | None: + def filter(self) -> RowFilter | None: return self._filter @filter.setter - def filter(self, row_filter: RowFilter | dict[str, Any] | None): + def filter(self, row_filter: RowFilter | None): """ Set a RowFilter to apply to this query @@ -310,24 +392,24 @@ def _shard_range( - a list of tuples, containing a segment index and a new sub-range. """ # 1. find the index of the segment the start key belongs to - if orig_range.start is None: + if orig_range._start is None: # if range is open on the left, include first segment start_segment = 0 else: # use binary search to find the segment the start key belongs to # bisect method determines how we break ties when the start key matches a split point # if inclusive, bisect_left to the left segment, otherwise bisect_right - bisect = bisect_left if orig_range.start.is_inclusive else bisect_right - start_segment = bisect(split_points, orig_range.start.key) + bisect = bisect_left if orig_range._start.is_inclusive else bisect_right + start_segment = bisect(split_points, orig_range._start.key) # 2. find the index of the segment the end key belongs to - if orig_range.end is None: + if orig_range._end is None: # if range is open on the right, include final segment end_segment = len(split_points) else: # use binary search to find the segment the end key belongs to. end_segment = bisect_left( - split_points, orig_range.end.key, lo=start_segment + split_points, orig_range._end.key, lo=start_segment ) # note: end_segment will always bisect_left, because split points represent inclusive ends # whether the end_key is includes the split point or not, the result is the same segment @@ -343,7 +425,7 @@ def _shard_range( # first range spans from start_key to the split_point representing the last key in the segment last_key_in_first_segment = split_points[start_segment] start_range = RowRange._from_points( - start=orig_range.start, + start=orig_range._start, end=_RangePoint(last_key_in_first_segment, is_inclusive=True), ) results.append((start_segment, start_range)) @@ -353,7 +435,7 @@ def _shard_range( last_key_before_segment = split_points[previous_segment] end_range = RowRange._from_points( start=_RangePoint(last_key_before_segment, is_inclusive=False), - end=orig_range.end, + end=orig_range._end, ) results.append((end_segment, end_range)) # 3c. add new spanning range to all segments other than the first and last @@ -386,7 +468,9 @@ def _to_dict(self) -> dict[str, Any]: "rows": row_set, } dict_filter = ( - self.filter.to_dict() if isinstance(self.filter, RowFilter) else self.filter + self.filter._to_dict() + if isinstance(self.filter, RowFilter) + else self.filter ) if dict_filter: final_dict["filter"] = dict_filter @@ -412,9 +496,15 @@ def __eq__(self, other): ) if this_range_empty and other_range_empty: return self.filter == other.filter and self.limit == other.limit + # otherwise, sets should have same sizes + if len(self.row_keys) != len(other.row_keys): + return False + if len(self.row_ranges) != len(other.row_ranges): + return False + ranges_match = all([row in other.row_ranges for row in self.row_ranges]) return ( self.row_keys == other.row_keys - and self.row_ranges == other.row_ranges + and ranges_match and self.filter == other.filter and self.limit == other.limit ) diff --git a/google/cloud/bigtable/data/row.py b/google/cloud/bigtable/data/row.py index 5fdc1b365..f562e96d6 100644 --- a/google/cloud/bigtable/data/row.py +++ b/google/cloud/bigtable/data/row.py @@ -153,7 +153,7 @@ def __str__(self) -> str: } """ output = ["{"] - for family, qualifier in self.get_column_components(): + for family, qualifier in self._get_column_components(): cell_list = self[family, qualifier] line = [f" (family={family!r}, qualifier={qualifier!r}): "] if len(cell_list) == 0: @@ -168,16 +168,16 @@ def __str__(self) -> str: def __repr__(self): cell_str_buffer = ["{"] - for family, qualifier in self.get_column_components(): + for family, qualifier in self._get_column_components(): cell_list = self[family, qualifier] - repr_list = [cell.to_dict() for cell in cell_list] + repr_list = [cell._to_dict() for cell in cell_list] cell_str_buffer.append(f" ('{family}', {qualifier!r}): {repr_list},") cell_str_buffer.append("}") cell_str = "\n".join(cell_str_buffer) output = f"Row(key={self.row_key!r}, cells={cell_str})" return output - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """ Returns a dictionary representation of the cell in the Bigtable Row proto format @@ -188,7 +188,7 @@ def to_dict(self) -> dict[str, Any]: for family_name, qualifier_dict in self._index.items(): qualifier_list = [] for qualifier_name, cell_list in qualifier_dict.items(): - cell_dicts = [cell.to_dict() for cell in cell_list] + cell_dicts = [cell._to_dict() for cell in cell_list] qualifier_list.append( {"qualifier": qualifier_name, "cells": cell_dicts} ) @@ -268,7 +268,7 @@ def __len__(self): """ return len(self.cells) - def get_column_components(self) -> list[tuple[str, bytes]]: + def _get_column_components(self) -> list[tuple[str, bytes]]: """ Returns a list of (family, qualifier) pairs associated with the cells @@ -288,8 +288,8 @@ def __eq__(self, other): return False if len(self.cells) != len(other.cells): return False - components = self.get_column_components() - other_components = other.get_column_components() + components = self._get_column_components() + other_components = other._get_column_components() if len(components) != len(other_components): return False if components != other_components: @@ -375,7 +375,7 @@ def __int__(self) -> int: """ return int.from_bytes(self.value, byteorder="big", signed=True) - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """ Returns a dictionary representation of the cell in the Bigtable Cell proto format diff --git a/google/cloud/bigtable/data/row_filters.py b/google/cloud/bigtable/data/row_filters.py index b2fae6971..9f09133d5 100644 --- a/google/cloud/bigtable/data/row_filters.py +++ b/google/cloud/bigtable/data/row_filters.py @@ -47,10 +47,10 @@ def _to_pb(self) -> data_v2_pb2.RowFilter: Returns: The converted current object. """ - return data_v2_pb2.RowFilter(**self.to_dict()) + return data_v2_pb2.RowFilter(**self._to_dict()) @abstractmethod - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" pass @@ -91,7 +91,7 @@ class SinkFilter(_BoolFilter): of a :class:`ConditionalRowFilter`. """ - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" return {"sink": self.flag} @@ -105,7 +105,7 @@ class PassAllFilter(_BoolFilter): completeness. """ - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" return {"pass_all_filter": self.flag} @@ -118,7 +118,7 @@ class BlockAllFilter(_BoolFilter): temporarily disabling just part of a filter. """ - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" return {"block_all_filter": self.flag} @@ -175,7 +175,7 @@ class RowKeyRegexFilter(_RegexFilter): since the row key is already specified. """ - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" return {"row_key_regex_filter": self.regex} @@ -199,7 +199,7 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" return {"row_sample_filter": self.sample} @@ -222,7 +222,7 @@ class FamilyNameRegexFilter(_RegexFilter): used as a literal. """ - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" return {"family_name_regex_filter": self.regex} @@ -248,7 +248,7 @@ class ColumnQualifierRegexFilter(_RegexFilter): match this regex (irrespective of column family). """ - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" return {"column_qualifier_regex_filter": self.regex} @@ -282,9 +282,9 @@ def _to_pb(self) -> data_v2_pb2.TimestampRange: Returns: The converted current object. """ - return data_v2_pb2.TimestampRange(**self.to_dict()) + return data_v2_pb2.TimestampRange(**self._to_dict()) - def to_dict(self) -> dict[str, int]: + def _to_dict(self) -> dict[str, int]: """Converts the timestamp range to a dict representation.""" timestamp_range_kwargs = {} if self.start is not None: @@ -330,9 +330,9 @@ def _to_pb(self) -> data_v2_pb2.RowFilter: """ return data_v2_pb2.RowFilter(timestamp_range_filter=self.range_._to_pb()) - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" - return {"timestamp_range_filter": self.range_.to_dict()} + return {"timestamp_range_filter": self.range_._to_dict()} def __repr__(self) -> str: return f"{self.__class__.__name__}(start={self.range_.start!r}, end={self.range_.end!r})" @@ -426,10 +426,10 @@ def _to_pb(self) -> data_v2_pb2.RowFilter: Returns: The converted current object. """ - column_range = data_v2_pb2.ColumnRange(**self.range_to_dict()) + column_range = data_v2_pb2.ColumnRange(**self._range_to_dict()) return data_v2_pb2.RowFilter(column_range_filter=column_range) - def range_to_dict(self) -> dict[str, str | bytes]: + def _range_to_dict(self) -> dict[str, str | bytes]: """Converts the column range range to a dict representation.""" column_range_kwargs: dict[str, str | bytes] = {} column_range_kwargs["family_name"] = self.family_id @@ -447,9 +447,9 @@ def range_to_dict(self) -> dict[str, str | bytes]: column_range_kwargs[key] = _to_bytes(self.end_qualifier) return column_range_kwargs - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" - return {"column_range_filter": self.range_to_dict()} + return {"column_range_filter": self._range_to_dict()} def __repr__(self) -> str: return f"{self.__class__.__name__}(family_id='{self.family_id}', start_qualifier={self.start_qualifier!r}, end_qualifier={self.end_qualifier!r}, inclusive_start={self.inclusive_start}, inclusive_end={self.inclusive_end})" @@ -476,7 +476,7 @@ class ValueRegexFilter(_RegexFilter): match this regex. String values will be encoded as ASCII. """ - def to_dict(self) -> dict[str, bytes]: + def _to_dict(self) -> dict[str, bytes]: """Converts the row filter to a dict representation.""" return {"value_regex_filter": self.regex} @@ -620,10 +620,10 @@ def _to_pb(self) -> data_v2_pb2.RowFilter: Returns: The converted current object. """ - value_range = data_v2_pb2.ValueRange(**self.range_to_dict()) + value_range = data_v2_pb2.ValueRange(**self._range_to_dict()) return data_v2_pb2.RowFilter(value_range_filter=value_range) - def range_to_dict(self) -> dict[str, bytes]: + def _range_to_dict(self) -> dict[str, bytes]: """Converts the value range range to a dict representation.""" value_range_kwargs = {} if self.start_value is not None: @@ -640,9 +640,9 @@ def range_to_dict(self) -> dict[str, bytes]: value_range_kwargs[key] = _to_bytes(self.end_value) return value_range_kwargs - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" - return {"value_range_filter": self.range_to_dict()} + return {"value_range_filter": self._range_to_dict()} def __repr__(self) -> str: return f"{self.__class__.__name__}(start_value={self.start_value!r}, end_value={self.end_value!r}, inclusive_start={self.inclusive_start}, inclusive_end={self.inclusive_end})" @@ -680,7 +680,7 @@ class CellsRowOffsetFilter(_CellCountFilter): :param num_cells: Skips the first N cells of the row. """ - def to_dict(self) -> dict[str, int]: + def _to_dict(self) -> dict[str, int]: """Converts the row filter to a dict representation.""" return {"cells_per_row_offset_filter": self.num_cells} @@ -692,7 +692,7 @@ class CellsRowLimitFilter(_CellCountFilter): :param num_cells: Matches only the first N cells of the row. """ - def to_dict(self) -> dict[str, int]: + def _to_dict(self) -> dict[str, int]: """Converts the row filter to a dict representation.""" return {"cells_per_row_limit_filter": self.num_cells} @@ -706,7 +706,7 @@ class CellsColumnLimitFilter(_CellCountFilter): timestamps of each cell. """ - def to_dict(self) -> dict[str, int]: + def _to_dict(self) -> dict[str, int]: """Converts the row filter to a dict representation.""" return {"cells_per_column_limit_filter": self.num_cells} @@ -720,7 +720,7 @@ class StripValueTransformerFilter(_BoolFilter): transformer than a generic query / filter. """ - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" return {"strip_value_transformer": self.flag} @@ -755,7 +755,7 @@ def __eq__(self, other): def __ne__(self, other): return not self == other - def to_dict(self) -> dict[str, str]: + def _to_dict(self) -> dict[str, str]: """Converts the row filter to a dict representation.""" return {"apply_label_transformer": self.label} @@ -841,9 +841,9 @@ def _to_pb(self) -> data_v2_pb2.RowFilter: ) return data_v2_pb2.RowFilter(chain=chain) - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" - return {"chain": {"filters": [f.to_dict() for f in self.filters]}} + return {"chain": {"filters": [f._to_dict() for f in self.filters]}} class RowFilterUnion(_FilterCombination): @@ -869,9 +869,9 @@ def _to_pb(self) -> data_v2_pb2.RowFilter: ) return data_v2_pb2.RowFilter(interleave=interleave) - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" - return {"interleave": {"filters": [f.to_dict() for f in self.filters]}} + return {"interleave": {"filters": [f._to_dict() for f in self.filters]}} class ConditionalRowFilter(RowFilter): @@ -939,18 +939,18 @@ def _to_pb(self) -> data_v2_pb2.RowFilter: condition = data_v2_pb2.RowFilter.Condition(**condition_kwargs) return data_v2_pb2.RowFilter(condition=condition) - def condition_to_dict(self) -> dict[str, Any]: + def _condition_to_dict(self) -> dict[str, Any]: """Converts the condition to a dict representation.""" - condition_kwargs = {"predicate_filter": self.predicate_filter.to_dict()} + condition_kwargs = {"predicate_filter": self.predicate_filter._to_dict()} if self.true_filter is not None: - condition_kwargs["true_filter"] = self.true_filter.to_dict() + condition_kwargs["true_filter"] = self.true_filter._to_dict() if self.false_filter is not None: - condition_kwargs["false_filter"] = self.false_filter.to_dict() + condition_kwargs["false_filter"] = self.false_filter._to_dict() return condition_kwargs - def to_dict(self) -> dict[str, Any]: + def _to_dict(self) -> dict[str, Any]: """Converts the row filter to a dict representation.""" - return {"condition": self.condition_to_dict()} + return {"condition": self._condition_to_dict()} def __repr__(self) -> str: return f"{self.__class__.__name__}(predicate_filter={self.predicate_filter!r}, true_filter={self.true_filter!r}, false_filter={self.false_filter!r})" diff --git a/tests/system/data/test_system.py b/tests/system/data/test_system.py index 548433444..fe341e4a8 100644 --- a/tests/system/data/test_system.py +++ b/tests/system/data/test_system.py @@ -20,7 +20,7 @@ from google.api_core import retry from google.api_core.exceptions import ClientError -from google.cloud.bigtable.data.read_modify_write_rules import MAX_INCREMENT_VALUE +from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" @@ -482,9 +482,9 @@ async def test_mutations_batcher_no_flush(client, table, temp_rows): (0, -100, -100), (0, 3000, 3000), (10, 4, 14), - (MAX_INCREMENT_VALUE, -MAX_INCREMENT_VALUE, 0), - (MAX_INCREMENT_VALUE, 2, -MAX_INCREMENT_VALUE), - (-MAX_INCREMENT_VALUE, -2, MAX_INCREMENT_VALUE), + (_MAX_INCREMENT_VALUE, -_MAX_INCREMENT_VALUE, 0), + (_MAX_INCREMENT_VALUE, 2, -_MAX_INCREMENT_VALUE), + (-_MAX_INCREMENT_VALUE, -2, _MAX_INCREMENT_VALUE), ], ) @pytest.mark.asyncio diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index 82cf135ac..0388d62e9 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -119,14 +119,14 @@ def test_ctor_too_many_entries(self): should raise an error if an operation is created with more than 100,000 entries """ from google.cloud.bigtable.data._async._mutate_rows import ( - MUTATE_ROWS_REQUEST_MUTATION_LIMIT, + _MUTATE_ROWS_REQUEST_MUTATION_LIMIT, ) - assert MUTATE_ROWS_REQUEST_MUTATION_LIMIT == 100_000 + assert _MUTATE_ROWS_REQUEST_MUTATION_LIMIT == 100_000 client = mock.Mock() table = mock.Mock() - entries = [_make_mutation()] * MUTATE_ROWS_REQUEST_MUTATION_LIMIT + entries = [_make_mutation()] * _MUTATE_ROWS_REQUEST_MUTATION_LIMIT operation_timeout = 0.05 attempt_timeout = 0.01 # no errors if at limit diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index b7a7b3ae7..5a96c6c16 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -131,7 +131,7 @@ async def test_ctor_dict_options(self): assert called_options.api_endpoint == "foo.bar:1234" assert isinstance(called_options, ClientOptions) with mock.patch.object( - self._get_target_class(), "start_background_channel_refresh" + self._get_target_class(), "_start_background_channel_refresh" ) as start_background_refresh: client = self._make_one(client_options=client_options) start_background_refresh.assert_called_once() @@ -231,29 +231,29 @@ async def test_channel_pool_replace(self): await client.close() @pytest.mark.filterwarnings("ignore::RuntimeWarning") - def test_start_background_channel_refresh_sync(self): + def test__start_background_channel_refresh_sync(self): # should raise RuntimeError if called in a sync context client = self._make_one(project="project-id") with pytest.raises(RuntimeError): - client.start_background_channel_refresh() + client._start_background_channel_refresh() @pytest.mark.asyncio - async def test_start_background_channel_refresh_tasks_exist(self): + async def test__start_background_channel_refresh_tasks_exist(self): # if tasks exist, should do nothing client = self._make_one(project="project-id") with mock.patch.object(asyncio, "create_task") as create_task: - client.start_background_channel_refresh() + client._start_background_channel_refresh() create_task.assert_not_called() await client.close() @pytest.mark.asyncio @pytest.mark.parametrize("pool_size", [1, 3, 7]) - async def test_start_background_channel_refresh(self, pool_size): + async def test__start_background_channel_refresh(self, pool_size): # should create background tasks for each channel client = self._make_one(project="project-id", pool_size=pool_size) ping_and_warm = AsyncMock() client._ping_and_warm_instances = ping_and_warm - client.start_background_channel_refresh() + client._start_background_channel_refresh() assert len(client._channel_refresh_tasks) == pool_size for task in client._channel_refresh_tasks: assert isinstance(task, asyncio.Task) @@ -267,7 +267,7 @@ async def test_start_background_channel_refresh(self, pool_size): @pytest.mark.skipif( sys.version_info < (3, 8), reason="Task.name requires python3.8 or higher" ) - async def test_start_background_channel_refresh_tasks_names(self): + async def test__start_background_channel_refresh_tasks_names(self): # if tasks exist, should do nothing pool_size = 3 client = self._make_one(project="project-id", pool_size=pool_size) @@ -569,7 +569,7 @@ async def test__register_instance(self): client_mock._active_instances = active_instances client_mock._instance_owners = instance_owners client_mock._channel_refresh_tasks = [] - client_mock.start_background_channel_refresh.side_effect = ( + client_mock._start_background_channel_refresh.side_effect = ( lambda: client_mock._channel_refresh_tasks.append(mock.Mock) ) mock_channels = [mock.Mock() for i in range(5)] @@ -580,7 +580,7 @@ async def test__register_instance(self): client_mock, "instance-1", table_mock ) # first call should start background refresh - assert client_mock.start_background_channel_refresh.call_count == 1 + assert client_mock._start_background_channel_refresh.call_count == 1 # ensure active_instances and instance_owners were updated properly expected_key = ( "prefix/instance-1", @@ -593,12 +593,12 @@ async def test__register_instance(self): assert expected_key == tuple(list(instance_owners)[0]) # should be a new task set assert client_mock._channel_refresh_tasks - # # next call should not call start_background_channel_refresh again + # next call should not call _start_background_channel_refresh again table_mock2 = mock.Mock() await self._get_target_class()._register_instance( client_mock, "instance-2", table_mock2 ) - assert client_mock.start_background_channel_refresh.call_count == 1 + assert client_mock._start_background_channel_refresh.call_count == 1 # but it should call ping and warm with new instance key assert client_mock._ping_and_warm_instances.call_count == len(mock_channels) for channel in mock_channels: @@ -655,7 +655,7 @@ async def test__register_instance_state( client_mock._active_instances = active_instances client_mock._instance_owners = instance_owners client_mock._channel_refresh_tasks = [] - client_mock.start_background_channel_refresh.side_effect = ( + client_mock._start_background_channel_refresh.side_effect = ( lambda: client_mock._channel_refresh_tasks.append(mock.Mock) ) mock_channels = [mock.Mock() for i in range(5)] @@ -1239,7 +1239,7 @@ async def test_read_rows_query_matches_request(self, include_app_profile): read_rows = table.client._gapic_client.read_rows read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream([]) row_keys = [b"test_1", "test_2"] - row_ranges = RowRange("start", "end") + row_ranges = RowRange("1start", "2end") filter_ = {"test": "filter"} limit = 99 query = ReadRowsQuery( @@ -1846,12 +1846,12 @@ async def test_read_rows_sharded_batching(self): operation timeout should change between batches """ from google.cloud.bigtable.data._async.client import TableAsync - from google.cloud.bigtable.data._async.client import CONCURRENCY_LIMIT + from google.cloud.bigtable.data._async.client import _CONCURRENCY_LIMIT - assert CONCURRENCY_LIMIT == 10 # change this test if this changes + assert _CONCURRENCY_LIMIT == 10 # change this test if this changes n_queries = 90 - expected_num_batches = n_queries // CONCURRENCY_LIMIT + expected_num_batches = n_queries // _CONCURRENCY_LIMIT query_list = [ReadRowsQuery() for _ in range(n_queries)] table_mock = AsyncMock() @@ -1875,8 +1875,8 @@ async def test_read_rows_sharded_batching(self): for batch_idx in range(expected_num_batches): batch_kwargs = kwargs[ batch_idx - * CONCURRENCY_LIMIT : (batch_idx + 1) - * CONCURRENCY_LIMIT + * _CONCURRENCY_LIMIT : (batch_idx + 1) + * _CONCURRENCY_LIMIT ] for req_kwargs in batch_kwargs: # each batch should have the same operation_timeout, and it should decrease in each batch @@ -2737,12 +2737,12 @@ async def test_check_and_mutate_single_mutations(self): @pytest.mark.asyncio async def test_check_and_mutate_predicate_object(self): - """predicate object should be converted to dict""" + """predicate filter should be passed to gapic request""" from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse mock_predicate = mock.Mock() - fake_dict = {"fake": "dict"} - mock_predicate.to_dict.return_value = fake_dict + predicate_dict = {"predicate": "dict"} + mock_predicate._to_dict.return_value = predicate_dict async with self._make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( @@ -2757,8 +2757,8 @@ async def test_check_and_mutate_predicate_object(self): false_case_mutations=[mock.Mock()], ) kwargs = mock_gapic.call_args[1] - assert kwargs["request"]["predicate_filter"] == fake_dict - assert mock_predicate.to_dict.call_count == 1 + assert kwargs["request"]["predicate_filter"] == predicate_dict + assert mock_predicate._to_dict.call_count == 1 @pytest.mark.asyncio async def test_check_and_mutate_mutations_parsing(self): diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py index a55c4f6ef..f95b53271 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -240,7 +240,7 @@ async def test_add_to_flow_max_mutation_limits( Should submit request early, even if the flow control has room for more """ with mock.patch( - "google.cloud.bigtable.data._async.mutations_batcher.MUTATE_ROWS_REQUEST_MUTATION_LIMIT", + "google.cloud.bigtable.data._async.mutations_batcher._MUTATE_ROWS_REQUEST_MUTATION_LIMIT", max_limit, ): mutation_objs = [_make_mutation(count=m[0], size=m[1]) for m in mutations] diff --git a/tests/unit/data/test_mutations.py b/tests/unit/data/test_mutations.py index 8365dbd02..8680a8da9 100644 --- a/tests/unit/data/test_mutations.py +++ b/tests/unit/data/test_mutations.py @@ -507,12 +507,12 @@ def test_ctor(self): def test_ctor_over_limit(self): """Should raise error if mutations exceed MAX_MUTATIONS_PER_ENTRY""" from google.cloud.bigtable.data.mutations import ( - MUTATE_ROWS_REQUEST_MUTATION_LIMIT, + _MUTATE_ROWS_REQUEST_MUTATION_LIMIT, ) - assert MUTATE_ROWS_REQUEST_MUTATION_LIMIT == 100_000 + assert _MUTATE_ROWS_REQUEST_MUTATION_LIMIT == 100_000 # no errors at limit - expected_mutations = [None for _ in range(MUTATE_ROWS_REQUEST_MUTATION_LIMIT)] + expected_mutations = [None for _ in range(_MUTATE_ROWS_REQUEST_MUTATION_LIMIT)] self._make_one(b"row_key", expected_mutations) # error if over limit with pytest.raises(ValueError) as e: diff --git a/tests/unit/data/test_read_rows_query.py b/tests/unit/data/test_read_rows_query.py index 88fde2d24..1e4e27d36 100644 --- a/tests/unit/data/test_read_rows_query.py +++ b/tests/unit/data/test_read_rows_query.py @@ -32,34 +32,52 @@ def _make_one(self, *args, **kwargs): def test_ctor_start_end(self): row_range = self._make_one("test_row", "test_row2") - assert row_range.start.key == "test_row".encode() - assert row_range.end.key == "test_row2".encode() - assert row_range.start.is_inclusive is True - assert row_range.end.is_inclusive is False + assert row_range._start.key == "test_row".encode() + assert row_range._end.key == "test_row2".encode() + assert row_range._start.is_inclusive is True + assert row_range._end.is_inclusive is False + assert row_range.start_key == "test_row".encode() + assert row_range.end_key == "test_row2".encode() + assert row_range.start_is_inclusive is True + assert row_range.end_is_inclusive is False def test_ctor_start_only(self): row_range = self._make_one("test_row3") - assert row_range.start.key == "test_row3".encode() - assert row_range.start.is_inclusive is True - assert row_range.end is None + assert row_range.start_key == "test_row3".encode() + assert row_range.start_is_inclusive is True + assert row_range.end_key is None + assert row_range.end_is_inclusive is True def test_ctor_end_only(self): row_range = self._make_one(end_key="test_row4") - assert row_range.end.key == "test_row4".encode() - assert row_range.end.is_inclusive is False - assert row_range.start is None + assert row_range.end_key == "test_row4".encode() + assert row_range.end_is_inclusive is False + assert row_range.start_key is None + assert row_range.start_is_inclusive is True + + def test_ctor_empty_strings(self): + """ + empty strings should be treated as None + """ + row_range = self._make_one("", "") + assert row_range._start is None + assert row_range._end is None + assert row_range.start_key is None + assert row_range.end_key is None + assert row_range.start_is_inclusive is True + assert row_range.end_is_inclusive is True def test_ctor_inclusive_flags(self): row_range = self._make_one("test_row5", "test_row6", False, True) - assert row_range.start.key == "test_row5".encode() - assert row_range.end.key == "test_row6".encode() - assert row_range.start.is_inclusive is False - assert row_range.end.is_inclusive is True + assert row_range.start_key == "test_row5".encode() + assert row_range.end_key == "test_row6".encode() + assert row_range.start_is_inclusive is False + assert row_range.end_is_inclusive is True def test_ctor_defaults(self): row_range = self._make_one() - assert row_range.start is None - assert row_range.end is None + assert row_range.start_key is None + assert row_range.end_key is None def test_ctor_flags_only(self): with pytest.raises(ValueError) as exc: @@ -83,6 +101,9 @@ def test_ctor_invalid_keys(self): with pytest.raises(ValueError) as exc: self._make_one("1", 2) assert str(exc.value) == "end_key must be a string or bytes" + with pytest.raises(ValueError) as exc: + self._make_one("2", "1") + assert str(exc.value) == "start_key must be less than or equal to end_key" def test__to_dict_defaults(self): row_range = self._make_one("test_row", "test_row2") @@ -143,8 +164,8 @@ def test__from_dict( row_range = RowRange._from_dict(input_dict) assert row_range._to_dict().keys() == input_dict.keys() - found_start = row_range.start - found_end = row_range.end + found_start = row_range._start + found_end = row_range._end if expected_start is None: assert found_start is None assert start_is_inclusive is None @@ -176,7 +197,7 @@ def test__from_points(self, dict_repr): row_range_from_dict = RowRange._from_dict(dict_repr) row_range_from_points = RowRange._from_points( - row_range_from_dict.start, row_range_from_dict.end + row_range_from_dict._start, row_range_from_dict._end ) assert row_range_from_points._to_dict() == row_range_from_dict._to_dict() @@ -238,6 +259,86 @@ def test___bool__(self, dict_repr, expected): row_range = RowRange._from_dict(dict_repr) assert bool(row_range) is expected + def test__eq__(self): + """ + test that row ranges can be compared for equality + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + range1 = RowRange("1", "2") + range1_dup = RowRange("1", "2") + range2 = RowRange("1", "3") + range_w_empty = RowRange(None, "2") + assert range1 == range1_dup + assert range1 != range2 + assert range1 != range_w_empty + range_1_w_inclusive_start = RowRange("1", "2", start_is_inclusive=True) + range_1_w_exclusive_start = RowRange("1", "2", start_is_inclusive=False) + range_1_w_inclusive_end = RowRange("1", "2", end_is_inclusive=True) + range_1_w_exclusive_end = RowRange("1", "2", end_is_inclusive=False) + assert range1 == range_1_w_inclusive_start + assert range1 == range_1_w_exclusive_end + assert range1 != range_1_w_exclusive_start + assert range1 != range_1_w_inclusive_end + + @pytest.mark.parametrize( + "dict_repr,expected", + [ + ( + {"start_key_closed": "test_row", "end_key_open": "test_row2"}, + "[b'test_row', b'test_row2')", + ), + ( + {"start_key_open": "test_row", "end_key_closed": "test_row2"}, + "(b'test_row', b'test_row2']", + ), + ({"start_key_open": b"a"}, "(b'a', +inf]"), + ({"end_key_closed": b"b"}, "[-inf, b'b']"), + ({"end_key_open": b"b"}, "[-inf, b'b')"), + ({}, "[-inf, +inf]"), + ], + ) + def test___str__(self, dict_repr, expected): + """ + test string representations of row ranges + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + row_range = RowRange._from_dict(dict_repr) + assert str(row_range) == expected + + @pytest.mark.parametrize( + "dict_repr,expected", + [ + ( + {"start_key_closed": "test_row", "end_key_open": "test_row2"}, + "RowRange(start_key=b'test_row', end_key=b'test_row2')", + ), + ( + {"start_key_open": "test_row", "end_key_closed": "test_row2"}, + "RowRange(start_key=b'test_row', end_key=b'test_row2', start_is_inclusive=False, end_is_inclusive=True)", + ), + ( + {"start_key_open": b"a"}, + "RowRange(start_key=b'a', end_key=None, start_is_inclusive=False)", + ), + ( + {"end_key_closed": b"b"}, + "RowRange(start_key=None, end_key=b'b', end_is_inclusive=True)", + ), + ({"end_key_open": b"b"}, "RowRange(start_key=None, end_key=b'b')"), + ({}, "RowRange(start_key=None, end_key=None)"), + ], + ) + def test___repr__(self, dict_repr, expected): + """ + test repr representations of row ranges + """ + from google.cloud.bigtable.data.read_rows_query import RowRange + + row_range = RowRange._from_dict(dict_repr) + assert repr(row_range) == expected + class TestReadRowsQuery: @staticmethod @@ -299,24 +400,6 @@ def test_set_filter(self): query.filter = 1 assert str(exc.value) == "row_filter must be a RowFilter or dict" - def test_set_filter_dict(self): - from google.cloud.bigtable.data.row_filters import RowSampleFilter - from google.cloud.bigtable_v2.types.bigtable import ReadRowsRequest - - filter1 = RowSampleFilter(0.5) - filter1_dict = filter1.to_dict() - query = self._make_one() - assert query.filter is None - query.filter = filter1_dict - assert query.filter == filter1_dict - output = query._to_dict() - assert output["filter"] == filter1_dict - proto_output = ReadRowsRequest(**output) - assert proto_output.filter == filter1._to_pb() - - query.filter = None - assert query.filter is None - def test_set_limit(self): query = self._make_one() assert query.limit is None @@ -698,13 +781,18 @@ def test_shard_limit_exception(self): ((), ("a",), False), (("a",), (), False), (("a",), ("a",), True), + (("a",), (["a", b"a"],), True), # duplicate keys + ((["a"],), (["a", "b"],), False), + ((["a", "b"],), (["a", "b"],), True), + ((["a", b"b"],), ([b"a", "b"],), True), (("a",), (b"a",), True), (("a",), ("b",), False), (("a",), ("a", ["b"]), False), - (("a", ["b"]), ("a", ["b"]), True), + (("a", "b"), ("a", ["b"]), True), (("a", ["b"]), ("a", ["b", "c"]), False), (("a", ["b", "c"]), ("a", [b"b", "c"]), True), (("a", ["b", "c"], 1), ("a", ["b", b"c"], 1), True), + (("a", ["b"], 1), ("a", ["b", b"b", "b"], 1), True), # duplicate ranges (("a", ["b"], 1), ("a", ["b"], 2), False), (("a", ["b"], 1, {"a": "b"}), ("a", ["b"], 1, {"a": "b"}), True), (("a", ["b"], 1, {"a": "b"}), ("a", ["b"], 1), False), diff --git a/tests/unit/data/test_row.py b/tests/unit/data/test_row.py index c9c797b61..df2fc72c0 100644 --- a/tests/unit/data/test_row.py +++ b/tests/unit/data/test_row.py @@ -176,7 +176,7 @@ def test_to_dict(self): cell2 = self._make_cell() cell2.value = b"other" row = self._make_one(TEST_ROW_KEY, [cell1, cell2]) - row_dict = row.to_dict() + row_dict = row._to_dict() expected_dict = { "key": TEST_ROW_KEY, "families": [ @@ -465,20 +465,20 @@ def test_get_column_components(self): ) row_response = self._make_one(TEST_ROW_KEY, [cell, cell2, cell3]) - self.assertEqual(len(row_response.get_column_components()), 2) + self.assertEqual(len(row_response._get_column_components()), 2) self.assertEqual( - row_response.get_column_components(), + row_response._get_column_components(), [(TEST_FAMILY_ID, TEST_QUALIFIER), (new_family_id, new_qualifier)], ) row_response = self._make_one(TEST_ROW_KEY, []) - self.assertEqual(len(row_response.get_column_components()), 0) - self.assertEqual(row_response.get_column_components(), []) + self.assertEqual(len(row_response._get_column_components()), 0) + self.assertEqual(row_response._get_column_components(), []) row_response = self._make_one(TEST_ROW_KEY, [cell]) - self.assertEqual(len(row_response.get_column_components()), 1) + self.assertEqual(len(row_response._get_column_components()), 1) self.assertEqual( - row_response.get_column_components(), [(TEST_FAMILY_ID, TEST_QUALIFIER)] + row_response._get_column_components(), [(TEST_FAMILY_ID, TEST_QUALIFIER)] ) def test_index_of(self): @@ -535,7 +535,7 @@ def test_to_dict(self): from google.cloud.bigtable_v2.types import Cell cell = self._make_one() - cell_dict = cell.to_dict() + cell_dict = cell._to_dict() expected_dict = { "value": TEST_VALUE, "timestamp_micros": TEST_TIMESTAMP, @@ -561,7 +561,7 @@ def test_to_dict_no_labels(self): TEST_TIMESTAMP, None, ) - cell_dict = cell_no_labels.to_dict() + cell_dict = cell_no_labels._to_dict() expected_dict = { "value": TEST_VALUE, "timestamp_micros": TEST_TIMESTAMP, diff --git a/tests/unit/data/test_row_filters.py b/tests/unit/data/test_row_filters.py index a3e275e70..e90b6f270 100644 --- a/tests/unit/data/test_row_filters.py +++ b/tests/unit/data/test_row_filters.py @@ -80,7 +80,7 @@ def test_sink_filter_to_dict(): flag = True row_filter = SinkFilter(flag) expected_dict = {"sink": flag} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -112,7 +112,7 @@ def test_pass_all_filter_to_dict(): flag = True row_filter = PassAllFilter(flag) expected_dict = {"pass_all_filter": flag} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -144,7 +144,7 @@ def test_block_all_filter_to_dict(): flag = True row_filter = BlockAllFilter(flag) expected_dict = {"block_all_filter": flag} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -214,7 +214,7 @@ def test_row_key_regex_filter_to_dict(): regex = b"row-key-regex" row_filter = RowKeyRegexFilter(regex) expected_dict = {"row_key_regex_filter": regex} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -302,7 +302,7 @@ def test_family_name_regex_filter_to_dict(): regex = "family-regex" row_filter = FamilyNameRegexFilter(regex) expected_dict = {"family_name_regex_filter": regex.encode()} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -335,7 +335,7 @@ def test_column_qualifier_regex_filter_to_dict(): regex = b"column-regex" row_filter = ColumnQualifierRegexFilter(regex) expected_dict = {"column_qualifier_regex_filter": regex} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -432,7 +432,7 @@ def test_timestamp_range_to_dict(): "start_timestamp_micros": 1546300800000000, "end_timestamp_micros": 1546387200000000, } - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.TimestampRange(**expected_dict) == expected_pb_value @@ -454,7 +454,7 @@ def test_timestamp_range_to_dict_start_only(): row_filter = TimestampRange(start=datetime.datetime(2019, 1, 1)) expected_dict = {"start_timestamp_micros": 1546300800000000} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.TimestampRange(**expected_dict) == expected_pb_value @@ -476,7 +476,7 @@ def test_timestamp_range_to_dict_end_only(): row_filter = TimestampRange(end=datetime.datetime(2019, 1, 2)) expected_dict = {"end_timestamp_micros": 1546387200000000} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.TimestampRange(**expected_dict) == expected_pb_value @@ -543,7 +543,7 @@ def test_timestamp_range_filter_to_dict(): "end_timestamp_micros": 1546387200000000, } } - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -554,7 +554,7 @@ def test_timestamp_range_filter_empty_to_dict(): row_filter = TimestampRangeFilter() expected_dict = {"timestamp_range_filter": {}} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -701,7 +701,7 @@ def test_column_range_filter_to_dict(): family_id = "column-family-id" row_filter = ColumnRangeFilter(family_id) expected_dict = {"column_range_filter": {"family_name": family_id}} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -782,7 +782,7 @@ def test_value_regex_filter_to_dict_w_bytes(): value = regex = b"value-regex" row_filter = ValueRegexFilter(value) expected_dict = {"value_regex_filter": regex} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -806,7 +806,7 @@ def test_value_regex_filter_to_dict_w_str(): regex = value.encode("ascii") row_filter = ValueRegexFilter(value) expected_dict = {"value_regex_filter": regex} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -839,7 +839,7 @@ def test_literal_value_filter_to_dict_w_bytes(): value = regex = b"value_regex" row_filter = LiteralValueFilter(value) expected_dict = {"value_regex_filter": regex} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -863,7 +863,7 @@ def test_literal_value_filter_to_dict_w_str(): regex = value.encode("ascii") row_filter = LiteralValueFilter(value) expected_dict = {"value_regex_filter": regex} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -896,7 +896,7 @@ def test_literal_value_filter_w_int(value, expected_byte_string): assert pb_val == expected_pb # test dict expected_dict = {"value_regex_filter": expected_byte_string} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict assert data_v2_pb2.RowFilter(**expected_dict) == pb_val @@ -1042,7 +1042,7 @@ def test_value_range_filter_to_dict(): row_filter = ValueRangeFilter() expected_dict = {"value_range_filter": {}} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -1149,7 +1149,7 @@ def test_cells_row_offset_filter_to_dict(): num_cells = 76 row_filter = CellsRowOffsetFilter(num_cells) expected_dict = {"cells_per_row_offset_filter": num_cells} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -1182,7 +1182,7 @@ def test_cells_row_limit_filter_to_dict(): num_cells = 189 row_filter = CellsRowLimitFilter(num_cells) expected_dict = {"cells_per_row_limit_filter": num_cells} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -1215,7 +1215,7 @@ def test_cells_column_limit_filter_to_dict(): num_cells = 10 row_filter = CellsColumnLimitFilter(num_cells) expected_dict = {"cells_per_column_limit_filter": num_cells} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -1248,7 +1248,7 @@ def test_strip_value_transformer_filter_to_dict(): flag = True row_filter = StripValueTransformerFilter(flag) expected_dict = {"strip_value_transformer": flag} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -1317,7 +1317,7 @@ def test_apply_label_filter_to_dict(): label = "label" row_filter = ApplyLabelFilter(label) expected_dict = {"apply_label_transformer": label} - assert row_filter.to_dict() == expected_dict + assert row_filter._to_dict() == expected_dict expected_pb_value = row_filter._to_pb() assert data_v2_pb2.RowFilter(**expected_dict) == expected_pb_value @@ -1437,13 +1437,13 @@ def test_row_filter_chain_to_dict(): from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) - row_filter1_dict = row_filter1.to_dict() + row_filter1_dict = row_filter1._to_dict() row_filter2 = RowSampleFilter(0.25) - row_filter2_dict = row_filter2.to_dict() + row_filter2_dict = row_filter2._to_dict() row_filter3 = RowFilterChain(filters=[row_filter1, row_filter2]) - filter_dict = row_filter3.to_dict() + filter_dict = row_filter3._to_dict() expected_dict = {"chain": {"filters": [row_filter1_dict, row_filter2_dict]}} assert filter_dict == expected_dict @@ -1487,13 +1487,13 @@ def test_row_filter_chain_to_dict_nested(): row_filter2 = RowSampleFilter(0.25) row_filter3 = RowFilterChain(filters=[row_filter1, row_filter2]) - row_filter3_dict = row_filter3.to_dict() + row_filter3_dict = row_filter3._to_dict() row_filter4 = CellsRowLimitFilter(11) - row_filter4_dict = row_filter4.to_dict() + row_filter4_dict = row_filter4._to_dict() row_filter5 = RowFilterChain(filters=[row_filter3, row_filter4]) - filter_dict = row_filter5.to_dict() + filter_dict = row_filter5._to_dict() expected_dict = {"chain": {"filters": [row_filter3_dict, row_filter4_dict]}} assert filter_dict == expected_dict @@ -1559,13 +1559,13 @@ def test_row_filter_union_to_dict(): from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) - row_filter1_dict = row_filter1.to_dict() + row_filter1_dict = row_filter1._to_dict() row_filter2 = RowSampleFilter(0.25) - row_filter2_dict = row_filter2.to_dict() + row_filter2_dict = row_filter2._to_dict() row_filter3 = RowFilterUnion(filters=[row_filter1, row_filter2]) - filter_dict = row_filter3.to_dict() + filter_dict = row_filter3._to_dict() expected_dict = {"interleave": {"filters": [row_filter1_dict, row_filter2_dict]}} assert filter_dict == expected_dict @@ -1609,13 +1609,13 @@ def test_row_filter_union_to_dict_nested(): row_filter2 = RowSampleFilter(0.25) row_filter3 = RowFilterUnion(filters=[row_filter1, row_filter2]) - row_filter3_dict = row_filter3.to_dict() + row_filter3_dict = row_filter3._to_dict() row_filter4 = CellsRowLimitFilter(11) - row_filter4_dict = row_filter4.to_dict() + row_filter4_dict = row_filter4._to_dict() row_filter5 = RowFilterUnion(filters=[row_filter3, row_filter4]) - filter_dict = row_filter5.to_dict() + filter_dict = row_filter5._to_dict() expected_dict = {"interleave": {"filters": [row_filter3_dict, row_filter4_dict]}} assert filter_dict == expected_dict @@ -1750,18 +1750,18 @@ def test_conditional_row_filter_to_dict(): from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) - row_filter1_dict = row_filter1.to_dict() + row_filter1_dict = row_filter1._to_dict() row_filter2 = RowSampleFilter(0.25) - row_filter2_dict = row_filter2.to_dict() + row_filter2_dict = row_filter2._to_dict() row_filter3 = CellsRowOffsetFilter(11) - row_filter3_dict = row_filter3.to_dict() + row_filter3_dict = row_filter3._to_dict() row_filter4 = ConditionalRowFilter( row_filter1, true_filter=row_filter2, false_filter=row_filter3 ) - filter_dict = row_filter4.to_dict() + filter_dict = row_filter4._to_dict() expected_dict = { "condition": { @@ -1804,13 +1804,13 @@ def test_conditional_row_filter_to_dict_true_only(): from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) - row_filter1_dict = row_filter1.to_dict() + row_filter1_dict = row_filter1._to_dict() row_filter2 = RowSampleFilter(0.25) - row_filter2_dict = row_filter2.to_dict() + row_filter2_dict = row_filter2._to_dict() row_filter3 = ConditionalRowFilter(row_filter1, true_filter=row_filter2) - filter_dict = row_filter3.to_dict() + filter_dict = row_filter3._to_dict() expected_dict = { "condition": { @@ -1852,13 +1852,13 @@ def test_conditional_row_filter_to_dict_false_only(): from google.cloud.bigtable_v2.types import data as data_v2_pb2 row_filter1 = StripValueTransformerFilter(True) - row_filter1_dict = row_filter1.to_dict() + row_filter1_dict = row_filter1._to_dict() row_filter2 = RowSampleFilter(0.25) - row_filter2_dict = row_filter2.to_dict() + row_filter2_dict = row_filter2._to_dict() row_filter3 = ConditionalRowFilter(row_filter1, false_filter=row_filter2) - filter_dict = row_filter3.to_dict() + filter_dict = row_filter3._to_dict() expected_dict = { "condition": { From aa760b2ce418eb67a94eff36267b2d75417bad8c Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 16 Aug 2023 11:42:22 -0600 Subject: [PATCH 18/56] feat: improve error group tracebacks on < py11 (#825) --- google/cloud/bigtable/data/exceptions.py | 71 +++++-- tests/unit/data/test_exceptions.py | 241 ++++++++++++++++++----- 2 files changed, 245 insertions(+), 67 deletions(-) diff --git a/google/cloud/bigtable/data/exceptions.py b/google/cloud/bigtable/data/exceptions.py index 9b6b4fe3f..7344874df 100644 --- a/google/cloud/bigtable/data/exceptions.py +++ b/google/cloud/bigtable/data/exceptions.py @@ -74,7 +74,44 @@ def __init__(self, message, excs): if len(excs) == 0: raise ValueError("exceptions must be a non-empty sequence") self.exceptions = tuple(excs) - super().__init__(message) + # simulate an exception group in Python < 3.11 by adding exception info + # to the message + first_line = "--+---------------- 1 ----------------" + last_line = "+------------------------------------" + message_parts = [message + "\n" + first_line] + # print error info for each exception in the group + for idx, e in enumerate(excs[:15]): + # apply index header + if idx != 0: + message_parts.append( + f"+---------------- {str(idx+1).rjust(2)} ----------------" + ) + cause = e.__cause__ + # if this exception was had a cause, print the cause first + # used to display root causes of FailedMutationEntryError and FailedQueryShardError + # format matches the error output of Python 3.11+ + if cause is not None: + message_parts.extend( + f"| {type(cause).__name__}: {cause}".splitlines() + ) + message_parts.append("| ") + message_parts.append( + "| The above exception was the direct cause of the following exception:" + ) + message_parts.append("| ") + # attach error message for this sub-exception + # if the subexception is also a _BigtableExceptionGroup, + # error messages will be nested + message_parts.extend(f"| {type(e).__name__}: {e}".splitlines()) + # truncate the message if there are more than 15 exceptions + if len(excs) > 15: + message_parts.append("+---------------- ... ---------------") + message_parts.append(f"| and {len(excs) - 15} more") + if last_line not in message_parts[-1]: + # in the case of nested _BigtableExceptionGroups, the last line + # does not need to be added, since one was added by the final sub-exception + message_parts.append(last_line) + super().__init__("\n ".join(message_parts)) def __new__(cls, message, excs): if is_311_plus: @@ -83,11 +120,19 @@ def __new__(cls, message, excs): return super().__new__(cls) def __str__(self): + if is_311_plus: + # don't return built-in sub-exception message + return self.args[0] + return super().__str__() + + def __repr__(self): """ - String representation doesn't display sub-exceptions. Subexceptions are - described in message + repr representation should strip out sub-exception details """ - return self.args[0] + if is_311_plus: + return super().__repr__() + message = self.args[0].split("\n")[0] + return f"{self.__class__.__name__}({message!r}, {self.exceptions!r})" class MutationsExceptionGroup(_BigtableExceptionGroup): @@ -200,14 +245,12 @@ def __init__( idempotent_msg = ( "idempotent" if failed_mutation_entry.is_idempotent() else "non-idempotent" ) - index_msg = f" at index {failed_idx} " if failed_idx is not None else " " - message = ( - f"Failed {idempotent_msg} mutation entry{index_msg}with cause: {cause!r}" - ) + index_msg = f" at index {failed_idx}" if failed_idx is not None else "" + message = f"Failed {idempotent_msg} mutation entry{index_msg}" super().__init__(message) + self.__cause__ = cause self.index = failed_idx self.entry = failed_mutation_entry - self.__cause__ = cause class RetryExceptionGroup(_BigtableExceptionGroup): @@ -217,10 +260,8 @@ class RetryExceptionGroup(_BigtableExceptionGroup): def _format_message(excs: list[Exception]): if len(excs) == 0: return "No exceptions" - if len(excs) == 1: - return f"1 failed attempt: {type(excs[0]).__name__}" - else: - return f"{len(excs)} failed attempts. Latest: {type(excs[-1]).__name__}" + plural = "s" if len(excs) > 1 else "" + return f"{len(excs)} failed attempt{plural}" def __init__(self, excs: list[Exception]): super().__init__(self._format_message(excs), excs) @@ -268,8 +309,8 @@ def __init__( failed_query: "ReadRowsQuery" | dict[str, Any], cause: Exception, ): - message = f"Failed query at index {failed_index} with cause: {cause!r}" + message = f"Failed query at index {failed_index}" super().__init__(message) + self.__cause__ = cause self.index = failed_index self.query = failed_query - self.__cause__ = cause diff --git a/tests/unit/data/test_exceptions.py b/tests/unit/data/test_exceptions.py index 9d1145e36..2defffc86 100644 --- a/tests/unit/data/test_exceptions.py +++ b/tests/unit/data/test_exceptions.py @@ -25,57 +25,68 @@ import mock # type: ignore -class TestBigtableExceptionGroup: +class TracebackTests311: """ - Subclass for MutationsExceptionGroup, RetryExceptionGroup, and ShardedReadRowsExceptionGroup + Provides a set of tests that should be run on python 3.11 and above, + to verify that the exception traceback looks as expected """ - def _get_class(self): - from google.cloud.bigtable.data.exceptions import _BigtableExceptionGroup - - return _BigtableExceptionGroup - - def _make_one(self, message="test_message", excs=None): - if excs is None: - excs = [RuntimeError("mock")] - - return self._get_class()(message, excs=excs) - - def test_raise(self): + @pytest.mark.skipif( + sys.version_info < (3, 11), reason="requires python3.11 or higher" + ) + def test_311_traceback(self): """ - Create exception in raise statement, which calls __new__ and __init__ + Exception customizations should not break rich exception group traceback in python 3.11 """ - test_msg = "test message" - test_excs = [Exception(test_msg)] - with pytest.raises(self._get_class()) as e: - raise self._get_class()(test_msg, test_excs) - assert str(e.value) == test_msg - assert list(e.value.exceptions) == test_excs + import traceback - def test_raise_empty_list(self): - """ - Empty exception lists are not supported - """ - with pytest.raises(ValueError) as e: - raise self._make_one(excs=[]) - assert "non-empty sequence" in str(e.value) + sub_exc1 = RuntimeError("first sub exception") + sub_exc2 = ZeroDivisionError("second sub exception") + sub_group = self._make_one(excs=[sub_exc2]) + exc_group = self._make_one(excs=[sub_exc1, sub_group]) + + expected_traceback = ( + f" | google.cloud.bigtable.data.exceptions.{type(exc_group).__name__}: {str(exc_group)}", + " +-+---------------- 1 ----------------", + " | RuntimeError: first sub exception", + " +---------------- 2 ----------------", + f" | google.cloud.bigtable.data.exceptions.{type(sub_group).__name__}: {str(sub_group)}", + " +-+---------------- 1 ----------------", + " | ZeroDivisionError: second sub exception", + " +------------------------------------", + ) + exception_caught = False + try: + raise exc_group + except self._get_class(): + exception_caught = True + tb = traceback.format_exc() + tb_relevant_lines = tuple(tb.splitlines()[3:]) + assert expected_traceback == tb_relevant_lines + assert exception_caught @pytest.mark.skipif( sys.version_info < (3, 11), reason="requires python3.11 or higher" ) - def test_311_traceback(self): + def test_311_traceback_with_cause(self): """ - Exception customizations should not break rich exception group traceback in python 3.11 + traceback should display nicely with sub-exceptions with __cause__ set """ import traceback sub_exc1 = RuntimeError("first sub exception") + cause_exc = ImportError("cause exception") + sub_exc1.__cause__ = cause_exc sub_exc2 = ZeroDivisionError("second sub exception") exc_group = self._make_one(excs=[sub_exc1, sub_exc2]) expected_traceback = ( f" | google.cloud.bigtable.data.exceptions.{type(exc_group).__name__}: {str(exc_group)}", " +-+---------------- 1 ----------------", + " | ImportError: cause exception", + " | ", + " | The above exception was the direct cause of the following exception:", + " | ", " | RuntimeError: first sub exception", " +---------------- 2 ----------------", " | ZeroDivisionError: second sub exception", @@ -105,6 +116,126 @@ def test_311_exception_group(self): assert runtime_error.exceptions[0] == exceptions[0] assert others.exceptions[0] == exceptions[1] + +class TracebackTests310: + """ + Provides a set of tests that should be run on python 3.10 and under, + to verify that the exception traceback looks as expected + """ + + @pytest.mark.skipif( + sys.version_info >= (3, 11), reason="requires python3.10 or lower" + ) + def test_310_traceback(self): + """ + Exception customizations should not break rich exception group traceback in python 3.10 + """ + import traceback + + sub_exc1 = RuntimeError("first sub exception") + sub_exc2 = ZeroDivisionError("second sub exception") + sub_group = self._make_one(excs=[sub_exc2]) + exc_group = self._make_one(excs=[sub_exc1, sub_group]) + found_message = str(exc_group).splitlines()[0] + found_sub_message = str(sub_group).splitlines()[0] + + expected_traceback = ( + f"google.cloud.bigtable.data.exceptions.{type(exc_group).__name__}: {found_message}", + "--+---------------- 1 ----------------", + " | RuntimeError: first sub exception", + " +---------------- 2 ----------------", + f" | {type(sub_group).__name__}: {found_sub_message}", + " --+---------------- 1 ----------------", + " | ZeroDivisionError: second sub exception", + " +------------------------------------", + ) + exception_caught = False + try: + raise exc_group + except self._get_class(): + exception_caught = True + tb = traceback.format_exc() + tb_relevant_lines = tuple(tb.splitlines()[3:]) + assert expected_traceback == tb_relevant_lines + assert exception_caught + + @pytest.mark.skipif( + sys.version_info >= (3, 11), reason="requires python3.10 or lower" + ) + def test_310_traceback_with_cause(self): + """ + traceback should display nicely with sub-exceptions with __cause__ set + """ + import traceback + + sub_exc1 = RuntimeError("first sub exception") + cause_exc = ImportError("cause exception") + sub_exc1.__cause__ = cause_exc + sub_exc2 = ZeroDivisionError("second sub exception") + exc_group = self._make_one(excs=[sub_exc1, sub_exc2]) + found_message = str(exc_group).splitlines()[0] + + expected_traceback = ( + f"google.cloud.bigtable.data.exceptions.{type(exc_group).__name__}: {found_message}", + "--+---------------- 1 ----------------", + " | ImportError: cause exception", + " | ", + " | The above exception was the direct cause of the following exception:", + " | ", + " | RuntimeError: first sub exception", + " +---------------- 2 ----------------", + " | ZeroDivisionError: second sub exception", + " +------------------------------------", + ) + exception_caught = False + try: + raise exc_group + except self._get_class(): + exception_caught = True + tb = traceback.format_exc() + tb_relevant_lines = tuple(tb.splitlines()[3:]) + assert expected_traceback == tb_relevant_lines + assert exception_caught + + +class TestBigtableExceptionGroup(TracebackTests311, TracebackTests310): + """ + Subclass for MutationsExceptionGroup, RetryExceptionGroup, and ShardedReadRowsExceptionGroup + """ + + def _get_class(self): + from google.cloud.bigtable.data.exceptions import _BigtableExceptionGroup + + return _BigtableExceptionGroup + + def _make_one(self, message="test_message", excs=None): + if excs is None: + excs = [RuntimeError("mock")] + + return self._get_class()(message, excs=excs) + + def test_raise(self): + """ + Create exception in raise statement, which calls __new__ and __init__ + """ + test_msg = "test message" + test_excs = [Exception(test_msg)] + with pytest.raises(self._get_class()) as e: + raise self._get_class()(test_msg, test_excs) + found_message = str(e.value).splitlines()[ + 0 + ] # added to prase out subexceptions in <3.11 + assert found_message == test_msg + assert list(e.value.exceptions) == test_excs + + def test_raise_empty_list(self): + """ + Empty exception lists are not supported + """ + with pytest.raises(ValueError) as e: + raise self._make_one(excs=[]) + assert "non-empty sequence" in str(e.value) + def test_exception_handling(self): """ All versions should inherit from exception @@ -151,7 +282,10 @@ def test_raise(self, exception_list, total_entries, expected_message): """ with pytest.raises(self._get_class()) as e: raise self._get_class()(exception_list, total_entries) - assert str(e.value) == expected_message + found_message = str(e.value).splitlines()[ + 0 + ] # added to prase out subexceptions in <3.11 + assert found_message == expected_message assert list(e.value.exceptions) == exception_list def test_raise_custom_message(self): @@ -162,7 +296,10 @@ def test_raise_custom_message(self): exception_list = [Exception()] with pytest.raises(self._get_class()) as e: raise self._get_class()(exception_list, 5, message=custom_message) - assert str(e.value) == custom_message + found_message = str(e.value).splitlines()[ + 0 + ] # added to prase out subexceptions in <3.11 + assert found_message == custom_message assert list(e.value.exceptions) == exception_list @pytest.mark.parametrize( @@ -222,7 +359,10 @@ def test_from_truncated_lists( raise self._get_class().from_truncated_lists( first_list, second_list, total_excs, entry_count ) - assert str(e.value) == expected_message + found_message = str(e.value).splitlines()[ + 0 + ] # added to prase out subexceptions in <3.11 + assert found_message == expected_message assert list(e.value.exceptions) == first_list + second_list @@ -241,11 +381,11 @@ def _make_one(self, excs=None): @pytest.mark.parametrize( "exception_list,expected_message", [ - ([Exception()], "1 failed attempt: Exception"), - ([Exception(), RuntimeError()], "2 failed attempts. Latest: RuntimeError"), + ([Exception()], "1 failed attempt"), + ([Exception(), RuntimeError()], "2 failed attempts"), ( [Exception(), ValueError("test")], - "2 failed attempts. Latest: ValueError", + "2 failed attempts", ), ( [ @@ -253,7 +393,7 @@ def _make_one(self, excs=None): [Exception(), ValueError("test")] ) ], - "1 failed attempt: RetryExceptionGroup", + "1 failed attempt", ), ], ) @@ -263,7 +403,10 @@ def test_raise(self, exception_list, expected_message): """ with pytest.raises(self._get_class()) as e: raise self._get_class()(exception_list) - assert str(e.value) == expected_message + found_message = str(e.value).splitlines()[ + 0 + ] # added to prase out subexceptions in <3.11 + assert found_message == expected_message assert list(e.value.exceptions) == exception_list @@ -299,7 +442,10 @@ def test_raise(self, exception_list, succeeded, total_entries, expected_message) """ with pytest.raises(self._get_class()) as e: raise self._get_class()(exception_list, succeeded, total_entries) - assert str(e.value) == expected_message + found_message = str(e.value).splitlines()[ + 0 + ] # added to prase out subexceptions in <3.11 + assert found_message == expected_message assert list(e.value.exceptions) == exception_list assert e.value.successful_rows == succeeded @@ -323,10 +469,7 @@ def test_raise(self): test_exc = ValueError("test") with pytest.raises(self._get_class()) as e: raise self._get_class()(test_idx, test_entry, test_exc) - assert ( - str(e.value) - == "Failed idempotent mutation entry at index 2 with cause: ValueError('test')" - ) + assert str(e.value) == "Failed idempotent mutation entry at index 2" assert e.value.index == test_idx assert e.value.entry == test_entry assert e.value.__cause__ == test_exc @@ -343,10 +486,7 @@ def test_raise_idempotent(self): test_exc = ValueError("test") with pytest.raises(self._get_class()) as e: raise self._get_class()(test_idx, test_entry, test_exc) - assert ( - str(e.value) - == "Failed non-idempotent mutation entry at index 2 with cause: ValueError('test')" - ) + assert str(e.value) == "Failed non-idempotent mutation entry at index 2" assert e.value.index == test_idx assert e.value.entry == test_entry assert e.value.__cause__ == test_exc @@ -361,10 +501,7 @@ def test_no_index(self): test_exc = ValueError("test") with pytest.raises(self._get_class()) as e: raise self._get_class()(test_idx, test_entry, test_exc) - assert ( - str(e.value) - == "Failed idempotent mutation entry with cause: ValueError('test')" - ) + assert str(e.value) == "Failed idempotent mutation entry" assert e.value.index == test_idx assert e.value.entry == test_entry assert e.value.__cause__ == test_exc @@ -391,7 +528,7 @@ def test_raise(self): test_exc = ValueError("test") with pytest.raises(self._get_class()) as e: raise self._get_class()(test_idx, test_query, test_exc) - assert str(e.value) == "Failed query at index 2 with cause: ValueError('test')" + assert str(e.value) == "Failed query at index 2" assert e.value.index == test_idx assert e.value.query == test_query assert e.value.__cause__ == test_exc From 0323ddec4d49eda87afffb0dccd50f2e328f0817 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 17 Aug 2023 10:19:17 -0600 Subject: [PATCH 19/56] feat: optimize read_rows (#852) --- google/cloud/bigtable/data/__init__.py | 3 +- google/cloud/bigtable/data/_async/__init__.py | 3 +- .../bigtable/data/_async/_mutate_rows.py | 4 +- .../cloud/bigtable/data/_async/_read_rows.py | 605 ++++++++-------- google/cloud/bigtable/data/_async/client.py | 31 +- google/cloud/bigtable/data/_helpers.py | 4 +- .../bigtable/data/_read_rows_state_machine.py | 368 ---------- google/cloud/bigtable/data/read_rows_query.py | 233 +++--- google/cloud/bigtable/data/row.py | 19 +- tests/system/conftest.py | 25 + tests/system/data/setup_fixtures.py | 171 +++++ tests/system/data/test_system.py | 225 +++--- tests/unit/data/_async/test__mutate_rows.py | 2 +- tests/unit/data/_async/test__read_rows.py | 556 ++++----------- tests/unit/data/_async/test_client.py | 177 ++--- tests/unit/data/test__helpers.py | 39 +- .../data/test__read_rows_state_machine.py | 663 ------------------ tests/unit/data/test_read_rows_acceptance.py | 44 +- tests/unit/data/test_read_rows_query.py | 263 +------ tests/unit/data/test_row.py | 14 - 20 files changed, 940 insertions(+), 2509 deletions(-) delete mode 100644 google/cloud/bigtable/data/_read_rows_state_machine.py create mode 100644 tests/system/conftest.py create mode 100644 tests/system/data/setup_fixtures.py delete mode 100644 tests/unit/data/test__read_rows_state_machine.py diff --git a/google/cloud/bigtable/data/__init__.py b/google/cloud/bigtable/data/__init__.py index c68e78c6f..4b01d0e6b 100644 --- a/google/cloud/bigtable/data/__init__.py +++ b/google/cloud/bigtable/data/__init__.py @@ -20,7 +20,7 @@ from google.cloud.bigtable.data._async.client import BigtableDataClientAsync from google.cloud.bigtable.data._async.client import TableAsync -from google.cloud.bigtable.data._async._read_rows import ReadRowsAsyncIterator + from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery @@ -66,7 +66,6 @@ "DeleteAllFromRow", "Row", "Cell", - "ReadRowsAsyncIterator", "IdleTimeout", "InvalidChunk", "FailedMutationEntryError", diff --git a/google/cloud/bigtable/data/_async/__init__.py b/google/cloud/bigtable/data/_async/__init__.py index 1e92e58dc..e13c9acb7 100644 --- a/google/cloud/bigtable/data/_async/__init__.py +++ b/google/cloud/bigtable/data/_async/__init__.py @@ -14,13 +14,12 @@ from google.cloud.bigtable.data._async.client import BigtableDataClientAsync from google.cloud.bigtable.data._async.client import TableAsync -from google.cloud.bigtable.data._async._read_rows import ReadRowsAsyncIterator + from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync __all__ = [ "BigtableDataClientAsync", "TableAsync", - "ReadRowsAsyncIterator", "MutationsBatcherAsync", ] diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index d41f1aeea..d4f5bc215 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -96,7 +96,9 @@ def __init__( maximum=60, ) retry_wrapped = retry(self._run_attempt) - self._operation = _convert_retry_deadline(retry_wrapped, operation_timeout) + self._operation = _convert_retry_deadline( + retry_wrapped, operation_timeout, is_async=True + ) # initialize state self.timeout_generator = _attempt_timeout_generator( attempt_timeout, operation_timeout diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index ec1e488c6..6edb72858 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -12,39 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. # + from __future__ import annotations -from typing import ( - List, - Any, - AsyncIterable, - AsyncIterator, - AsyncGenerator, - Iterator, - Callable, - Awaitable, -) -import sys -import time -import asyncio -from functools import partial -from grpc.aio import RpcContext - -from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse -from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient -from google.cloud.bigtable.data.row import Row, _LastScannedRow +from typing import TYPE_CHECKING, AsyncGenerator, AsyncIterable, Awaitable + +from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB +from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB +from google.cloud.bigtable_v2.types import RowSet as RowSetPB +from google.cloud.bigtable_v2.types import RowRange as RowRangePB + +from google.cloud.bigtable.data.row import Row, Cell +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data.exceptions import InvalidChunk +from google.cloud.bigtable.data.exceptions import RetryExceptionGroup from google.cloud.bigtable.data.exceptions import _RowSetComplete -from google.cloud.bigtable.data.exceptions import IdleTimeout -from google.cloud.bigtable.data._read_rows_state_machine import _StateMachine +from google.cloud.bigtable.data._helpers import _attempt_timeout_generator +from google.cloud.bigtable.data._helpers import _make_metadata + from google.api_core import retry_async as retries +from google.api_core.retry_streaming_async import AsyncRetryableGenerator +from google.api_core.retry import exponential_sleep_generator from google.api_core import exceptions as core_exceptions -from google.cloud.bigtable.data._helpers import _make_metadata -from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _convert_retry_deadline + +if TYPE_CHECKING: + from google.cloud.bigtable.data._async.client import TableAsync -class _ReadRowsOperationAsync(AsyncIterable[Row]): +class _ResetRow(Exception): + def __init__(self, chunk): + self.chunk = chunk + + +class _ReadRowsOperationAsync: """ ReadRowsOperation handles the logic of merging chunks from a ReadRowsResponse stream into a stream of Row objects. @@ -57,160 +57,268 @@ class _ReadRowsOperationAsync(AsyncIterable[Row]): performing retries on stream errors. """ + __slots__ = ( + "attempt_timeout_gen", + "operation_timeout", + "request", + "table", + "_predicate", + "_metadata", + "_last_yielded_row_key", + "_remaining_count", + ) + def __init__( self, - request: dict[str, Any], - client: BigtableAsyncClient, - *, - operation_timeout: float = 600.0, - attempt_timeout: float | None = None, + query: ReadRowsQuery, + table: "TableAsync", + operation_timeout: float, + attempt_timeout: float, ): - """ - Args: - - request: the request dict to send to the Bigtable API - - client: the Bigtable client to use to make the request - - operation_timeout: the timeout to use for the entire operation, in seconds - - attempt_timeout: the timeout to use when waiting for each individual grpc request, in seconds - If not specified, defaults to operation_timeout - """ - self._last_emitted_row_key: bytes | None = None - self._emit_count = 0 - self._request = request - self.operation_timeout = operation_timeout - # use generator to lower per-attempt timeout as we approach operation_timeout deadline - attempt_timeout_gen = _attempt_timeout_generator( + self.attempt_timeout_gen = _attempt_timeout_generator( attempt_timeout, operation_timeout ) - row_limit = request.get("rows_limit", 0) - # lock in paramters for retryable wrapper - self._partial_retryable = partial( - self._read_rows_retryable_attempt, - client.read_rows, - attempt_timeout_gen, - row_limit, - ) - predicate = retries.if_exception_type( + self.operation_timeout = operation_timeout + if isinstance(query, dict): + self.request = ReadRowsRequestPB( + **query, + table_name=table.table_name, + app_profile_id=table.app_profile_id, + ) + else: + self.request = query._to_pb(table) + self.table = table + self._predicate = retries.if_exception_type( core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable, core_exceptions.Aborted, ) - - def on_error_fn(exc): - if predicate(exc): - self.transient_errors.append(exc) - - retry = retries.AsyncRetry( - predicate=predicate, - timeout=self.operation_timeout, - initial=0.01, - multiplier=2, - maximum=60, - on_error=on_error_fn, - is_stream=True, + self._metadata = _make_metadata( + table.table_name, + table.app_profile_id, ) - self._stream: AsyncGenerator[Row, None] | None = retry( - self._partial_retryable - )() - # contains the list of errors that were retried - self.transient_errors: List[Exception] = [] - - def __aiter__(self) -> AsyncIterator[Row]: - """Implements the AsyncIterable interface""" - return self - - async def __anext__(self) -> Row: - """Implements the AsyncIterator interface""" - if self._stream is not None: - return await self._stream.__anext__() - else: - raise asyncio.InvalidStateError("stream is closed") + self._last_yielded_row_key: bytes | None = None + self._remaining_count: int | None = self.request.rows_limit or None - async def aclose(self): - """Close the stream and release resources""" - if self._stream is not None: - await self._stream.aclose() - self._stream = None - self._emitted_seen_row_key = None + async def start_operation(self) -> AsyncGenerator[Row, None]: + """ + Start the read_rows operation, retrying on retryable errors. + """ + transient_errors = [] - async def _read_rows_retryable_attempt( - self, - gapic_fn: Callable[..., Awaitable[AsyncIterable[ReadRowsResponse]]], - timeout_generator: Iterator[float], - total_row_limit: int, - ) -> AsyncGenerator[Row, None]: + def on_error_fn(exc): + if self._predicate(exc): + transient_errors.append(exc) + + retry_gen = AsyncRetryableGenerator( + self._read_rows_attempt, + self._predicate, + exponential_sleep_generator(0.01, 60, multiplier=2), + self.operation_timeout, + on_error_fn, + ) + try: + async for row in retry_gen: + yield row + if self._remaining_count is not None: + self._remaining_count -= 1 + if self._remaining_count < 0: + raise RuntimeError("emit count exceeds row limit") + except core_exceptions.RetryError: + self._raise_retry_error(transient_errors) + except GeneratorExit: + # propagate close to wrapped generator + await retry_gen.aclose() + + def _read_rows_attempt(self) -> AsyncGenerator[Row, None]: """ - Retryable wrapper for merge_rows. This function is called each time - a retry is attempted. - - Some fresh state is created on each retry: - - grpc network stream - - state machine to hold merge chunks received from stream - Some state is shared between retries: - - _last_emitted_row_key is used to ensure that - duplicate rows are not emitted - - request is stored and (potentially) modified on each retry + Attempt a single read_rows rpc call. + This function is intended to be wrapped by retry logic, + which will call this function until it succeeds or + a non-retryable error is raised. """ - if self._last_emitted_row_key is not None: + # revise request keys and ranges between attempts + if self._last_yielded_row_key is not None: # if this is a retry, try to trim down the request to avoid ones we've already processed try: - self._request["rows"] = _ReadRowsOperationAsync._revise_request_rowset( - row_set=self._request.get("rows", None), - last_seen_row_key=self._last_emitted_row_key, + self.request.rows = self._revise_request_rowset( + row_set=self.request.rows, + last_seen_row_key=self._last_yielded_row_key, ) except _RowSetComplete: - # if there are no rows left to process, we're done - # This is not expected to happen often, but could occur if - # a retry is triggered quickly after the last row is emitted - return - # revise next request's row limit based on number emitted - if total_row_limit: - new_limit = total_row_limit - self._emit_count - if new_limit == 0: - # we have hit the row limit, so we're done - return - elif new_limit < 0: - raise RuntimeError("unexpected state: emit count exceeds row limit") - else: - self._request["rows_limit"] = new_limit - metadata = _make_metadata( - self._request.get("table_name", None), - self._request.get("app_profile_id", None), - ) - new_gapic_stream: RpcContext = await gapic_fn( - self._request, - timeout=next(timeout_generator), - metadata=metadata, + # if we've already seen all the rows, we're done + return self.merge_rows(None) + # revise the limit based on number of rows already yielded + if self._remaining_count is not None: + self.request.rows_limit = self._remaining_count + if self._remaining_count == 0: + return self.merge_rows(None) + # create and return a new row merger + gapic_stream = self.table.client._gapic_client.read_rows( + self.request, + timeout=next(self.attempt_timeout_gen), + metadata=self._metadata, ) - try: - state_machine = _StateMachine() - stream = _ReadRowsOperationAsync.merge_row_response_stream( - new_gapic_stream, state_machine - ) - # run until we get a timeout or the stream is exhausted - async for new_item in stream: + chunked_stream = self.chunk_stream(gapic_stream) + return self.merge_rows(chunked_stream) + + async def chunk_stream( + self, stream: Awaitable[AsyncIterable[ReadRowsResponsePB]] + ) -> AsyncGenerator[ReadRowsResponsePB.CellChunk, None]: + """ + process chunks out of raw read_rows stream + """ + async for resp in await stream: + # extract proto from proto-plus wrapper + resp = resp._pb + + # handle last_scanned_row_key packets, sent when server + # has scanned past the end of the row range + if resp.last_scanned_row_key: if ( - self._last_emitted_row_key is not None - and new_item.row_key <= self._last_emitted_row_key + self._last_yielded_row_key is not None + and resp.last_scanned_row_key <= self._last_yielded_row_key ): - raise InvalidChunk("Last emitted row key out of order") - # don't yeild _LastScannedRow markers; they - # should only update last_seen_row_key - if not isinstance(new_item, _LastScannedRow): - yield new_item - self._emit_count += 1 - self._last_emitted_row_key = new_item.row_key - if total_row_limit and self._emit_count >= total_row_limit: - return - except (Exception, GeneratorExit) as exc: - # ensure grpc stream is closed - new_gapic_stream.cancel() - raise exc + raise InvalidChunk("last scanned out of order") + self._last_yielded_row_key = resp.last_scanned_row_key + + current_key = None + # process each chunk in the response + for c in resp.chunks: + if current_key is None: + current_key = c.row_key + if current_key is None: + raise InvalidChunk("first chunk is missing a row key") + elif ( + self._last_yielded_row_key + and current_key <= self._last_yielded_row_key + ): + raise InvalidChunk("row keys should be strictly increasing") + + yield c + + if c.reset_row: + current_key = None + elif c.commit_row: + # update row state after each commit + self._last_yielded_row_key = current_key + current_key = None + + @staticmethod + async def merge_rows( + chunks: AsyncGenerator[ReadRowsResponsePB.CellChunk, None] | None + ): + """ + Merge chunks into rows + """ + if chunks is None: + return + it = chunks.__aiter__() + # For each row + while True: + try: + c = await it.__anext__() + except StopAsyncIteration: + # stream complete + return + row_key = c.row_key + + if not row_key: + raise InvalidChunk("first row chunk is missing key") + + cells = [] + + # shared per cell storage + family: str | None = None + qualifier: bytes | None = None + + try: + # for each cell + while True: + if c.reset_row: + raise _ResetRow(c) + k = c.row_key + f = c.family_name.value + q = c.qualifier.value if c.HasField("qualifier") else None + if k and k != row_key: + raise InvalidChunk("unexpected new row key") + if f: + family = f + if q is not None: + qualifier = q + else: + raise InvalidChunk("new family without qualifier") + elif family is None: + raise InvalidChunk("missing family") + elif q is not None: + if family is None: + raise InvalidChunk("new qualifier without family") + qualifier = q + elif qualifier is None: + raise InvalidChunk("missing qualifier") + + ts = c.timestamp_micros + labels = c.labels if c.labels else [] + value = c.value + + # merge split cells + if c.value_size > 0: + buffer = [value] + while c.value_size > 0: + # throws when premature end + c = await it.__anext__() + + t = c.timestamp_micros + cl = c.labels + k = c.row_key + if ( + c.HasField("family_name") + and c.family_name.value != family + ): + raise InvalidChunk("family changed mid cell") + if ( + c.HasField("qualifier") + and c.qualifier.value != qualifier + ): + raise InvalidChunk("qualifier changed mid cell") + if t and t != ts: + raise InvalidChunk("timestamp changed mid cell") + if cl and cl != labels: + raise InvalidChunk("labels changed mid cell") + if k and k != row_key: + raise InvalidChunk("row key changed mid cell") + + if c.reset_row: + raise _ResetRow(c) + buffer.append(c.value) + value = b"".join(buffer) + cells.append( + Cell(value, row_key, family, qualifier, ts, list(labels)) + ) + if c.commit_row: + yield Row(row_key, cells) + break + c = await it.__anext__() + except _ResetRow as e: + c = e.chunk + if ( + c.row_key + or c.HasField("family_name") + or c.HasField("qualifier") + or c.timestamp_micros + or c.labels + or c.value + ): + raise InvalidChunk("reset row with data") + continue + except StopAsyncIteration: + raise InvalidChunk("premature end of stream") @staticmethod def _revise_request_rowset( - row_set: dict[str, Any] | None, + row_set: RowSetPB, last_seen_row_key: bytes, - ) -> dict[str, Any]: + ) -> RowSetPB: """ Revise the rows in the request to avoid ones we've already processed. @@ -221,183 +329,44 @@ def _revise_request_rowset( - _RowSetComplete: if there are no rows left to process after the revision """ # if user is doing a whole table scan, start a new one with the last seen key - if row_set is None or ( - len(row_set.get("row_ranges", [])) == 0 - and len(row_set.get("row_keys", [])) == 0 - ): + if row_set is None or (not row_set.row_ranges and row_set.row_keys is not None): last_seen = last_seen_row_key - return { - "row_keys": [], - "row_ranges": [{"start_key_open": last_seen}], - } + return RowSetPB(row_ranges=[RowRangePB(start_key_open=last_seen)]) # remove seen keys from user-specific key list - row_keys: list[bytes] = row_set.get("row_keys", []) - adjusted_keys = [k for k in row_keys if k > last_seen_row_key] + adjusted_keys: list[bytes] = [ + k for k in row_set.row_keys if k > last_seen_row_key + ] # adjust ranges to ignore keys before last seen - row_ranges: list[dict[str, Any]] = row_set.get("row_ranges", []) - adjusted_ranges = [] - for row_range in row_ranges: - end_key = row_range.get("end_key_closed", None) or row_range.get( - "end_key_open", None - ) + adjusted_ranges: list[RowRangePB] = [] + for row_range in row_set.row_ranges: + end_key = row_range.end_key_closed or row_range.end_key_open or None if end_key is None or end_key > last_seen_row_key: # end range is after last seen key - new_range = row_range.copy() - start_key = row_range.get("start_key_closed", None) or row_range.get( - "start_key_open", None - ) + new_range = RowRangePB(row_range) + start_key = row_range.start_key_closed or row_range.start_key_open if start_key is None or start_key <= last_seen_row_key: # replace start key with last seen - new_range["start_key_open"] = last_seen_row_key - new_range.pop("start_key_closed", None) + new_range.start_key_open = last_seen_row_key adjusted_ranges.append(new_range) if len(adjusted_keys) == 0 and len(adjusted_ranges) == 0: # if the query is empty after revision, raise an exception # this will avoid an unwanted full table scan raise _RowSetComplete() - return {"row_keys": adjusted_keys, "row_ranges": adjusted_ranges} - - @staticmethod - async def merge_row_response_stream( - response_generator: AsyncIterable[ReadRowsResponse], - state_machine: _StateMachine, - ) -> AsyncGenerator[Row, None]: - """ - Consume chunks from a ReadRowsResponse stream into a set of Rows - - Args: - - response_generator: AsyncIterable of ReadRowsResponse objects. Typically - this is a stream of chunks from the Bigtable API - Returns: - - AsyncGenerator of Rows - Raises: - - InvalidChunk: if the chunk stream is invalid - """ - async for row_response in response_generator: - # unwrap protoplus object for increased performance - response_pb = row_response._pb - last_scanned = response_pb.last_scanned_row_key - # if the server sends a scan heartbeat, notify the state machine. - if last_scanned: - yield state_machine.handle_last_scanned_row(last_scanned) - # process new chunks through the state machine. - for chunk in response_pb.chunks: - complete_row = state_machine.handle_chunk(chunk) - if complete_row is not None: - yield complete_row - # TODO: handle request stats - if not state_machine.is_terminal_state(): - # read rows is complete, but there's still data in the merger - raise InvalidChunk("read_rows completed with partial state remaining") - - -class ReadRowsAsyncIterator(AsyncIterable[Row]): - """ - Async iterator for ReadRows responses. - - Supports the AsyncIterator protocol for use in async for loops, - along with: - - `aclose` for closing the underlying stream - - `active` for checking if the iterator is still active - - an internal idle timer for closing the stream after a period of inactivity - """ - - def __init__(self, merger: _ReadRowsOperationAsync): - self._merger: _ReadRowsOperationAsync = merger - self._error: Exception | None = None - self._last_interaction_time = time.monotonic() - self._idle_timeout_task: asyncio.Task[None] | None = None - # wrap merger with a wrapper that properly formats exceptions - self._next_fn = _convert_retry_deadline( - self._merger.__anext__, - self._merger.operation_timeout, - self._merger.transient_errors, - ) - - async def _start_idle_timer(self, idle_timeout: float): - """ - Start a coroutine that will cancel a stream if no interaction - with the iterator occurs for the specified number of seconds. - - Subsequent access to the iterator will raise an IdleTimeout exception. - - Args: - - idle_timeout: number of seconds of inactivity before cancelling the stream - """ - self._last_interaction_time = time.monotonic() - if self._idle_timeout_task is not None: - self._idle_timeout_task.cancel() - self._idle_timeout_task = asyncio.create_task( - self._idle_timeout_coroutine(idle_timeout) - ) - if sys.version_info >= (3, 8): - self._idle_timeout_task.name = f"{self.__class__.__name__}.idle_timeout" - - @property - def active(self): - """ - Returns True if the iterator is still active and has not been closed - """ - return self._error is None - - async def _idle_timeout_coroutine(self, idle_timeout: float): - """ - Coroutine that will cancel a stream if no interaction with the iterator - in the last `idle_timeout` seconds. - """ - while self.active: - next_timeout = self._last_interaction_time + idle_timeout - await asyncio.sleep(next_timeout - time.monotonic()) - if ( - self._last_interaction_time + idle_timeout < time.monotonic() - and self.active - ): - # idle timeout has expired - await self._finish_with_error( - IdleTimeout( - ( - "Timed out waiting for next Row to be consumed. " - f"(idle_timeout={idle_timeout:0.1f}s)" - ) - ) - ) - - def __aiter__(self): - """Implement the async iterator protocol.""" - return self - - async def __anext__(self) -> Row: - """ - Implement the async iterator potocol. + return RowSetPB(row_keys=adjusted_keys, row_ranges=adjusted_ranges) - Return the next item in the stream if active, or - raise an exception if the stream has been closed. - """ - if self._error is not None: - raise self._error - try: - self._last_interaction_time = time.monotonic() - return await self._next_fn() - except Exception as e: - await self._finish_with_error(e) - raise e - - async def _finish_with_error(self, e: Exception): - """ - Helper function to close the stream and clean up resources - after an error has occurred. - """ - if self.active: - await self._merger.aclose() - self._error = e - if self._idle_timeout_task is not None: - self._idle_timeout_task.cancel() - self._idle_timeout_task = None - - async def aclose(self): + def _raise_retry_error(self, transient_errors: list[Exception]) -> None: """ - Support closing the stream with an explicit call to aclose() + If the retryable deadline is hit, wrap the raised exception + in a RetryExceptionGroup """ - await self._finish_with_error( - StopAsyncIteration(f"{self.__class__.__name__} closed") + timeout_value = self.operation_timeout + timeout_str = f" of {timeout_value:.1f}s" if timeout_value is not None else "" + error_str = f"operation_timeout{timeout_str} exceeded" + new_exc = core_exceptions.DeadlineExceeded( + error_str, ) + source_exc = None + if transient_errors: + source_exc = RetryExceptionGroup(transient_errors) + new_exc.__cause__ = source_exc + raise new_exc from source_exc diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 0f8748bab..ff323b9bc 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -18,6 +18,7 @@ from typing import ( cast, Any, + AsyncIterable, Optional, Set, TYPE_CHECKING, @@ -44,7 +45,6 @@ from google.api_core import retry_async as retries from google.api_core import exceptions as core_exceptions from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync -from google.cloud.bigtable.data._async._read_rows import ReadRowsAsyncIterator import google.auth.credentials import google.auth._default @@ -497,7 +497,7 @@ async def read_rows_stream( *, operation_timeout: float | None = None, attempt_timeout: float | None = None, - ) -> ReadRowsAsyncIterator: + ) -> AsyncIterable[Row]: """ Read a set of rows from the table, based on the specified query. Returns an iterator to asynchronously stream back row data. @@ -533,27 +533,13 @@ async def read_rows_stream( ) _validate_timeouts(operation_timeout, attempt_timeout) - request = query._to_dict() if isinstance(query, ReadRowsQuery) else query - request["table_name"] = self.table_name - if self.app_profile_id: - request["app_profile_id"] = self.app_profile_id - - # read_rows smart retries is implemented using a series of iterators: - # - client.read_rows: outputs raw ReadRowsResponse objects from backend. Has attempt_timeout - # - ReadRowsOperation.merge_row_response_stream: parses chunks into rows - # - ReadRowsOperation.retryable_merge_rows: adds retries, caching, revised requests, operation_timeout - # - ReadRowsAsyncIterator: adds idle_timeout, moves stats out of stream and into attribute row_merger = _ReadRowsOperationAsync( - request, - self.client._gapic_client, + query, + self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, ) - output_generator = ReadRowsAsyncIterator(row_merger) - # add idle timeout to clear resources if generator is abandoned - idle_timeout_seconds = 300 - await output_generator._start_idle_timer(idle_timeout_seconds) - return output_generator + return row_merger.start_operation() async def read_rows( self, @@ -592,8 +578,7 @@ async def read_rows( operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, ) - results = [row async for row in row_generator] - return results + return [row async for row in row_generator] async def read_row( self, @@ -846,7 +831,7 @@ async def execute_rpc(): return [(s.row_key, s.offset_bytes) async for s in results] wrapped_fn = _convert_retry_deadline( - retry(execute_rpc), operation_timeout, transient_errors + retry(execute_rpc), operation_timeout, transient_errors, is_async=True ) return await wrapped_fn() @@ -973,7 +958,7 @@ def on_error_fn(exc): retry_wrapped = retry(self.client._gapic_client.mutate_row) # convert RetryErrors from retry wrapper into DeadlineExceeded errors deadline_wrapped = _convert_retry_deadline( - retry_wrapped, operation_timeout, transient_errors + retry_wrapped, operation_timeout, transient_errors, is_async=True ) metadata = _make_metadata(self.table_name, self.app_profile_id) # trigger rpc diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 68f310b49..b13b670d4 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -14,7 +14,6 @@ from __future__ import annotations from typing import Callable, Any -from inspect import iscoroutinefunction import time from google.api_core import exceptions as core_exceptions @@ -67,6 +66,7 @@ def _convert_retry_deadline( func: Callable[..., Any], timeout_value: float | None = None, retry_errors: list[Exception] | None = None, + is_async: bool = False, ): """ Decorator to convert RetryErrors raised by api_core.retry into @@ -108,7 +108,7 @@ def wrapper(*args, **kwargs): except core_exceptions.RetryError: handle_error() - return wrapper_async if iscoroutinefunction(func) else wrapper + return wrapper_async if is_async else wrapper def _validate_timeouts( diff --git a/google/cloud/bigtable/data/_read_rows_state_machine.py b/google/cloud/bigtable/data/_read_rows_state_machine.py deleted file mode 100644 index 7c0d05fb9..000000000 --- a/google/cloud/bigtable/data/_read_rows_state_machine.py +++ /dev/null @@ -1,368 +0,0 @@ -# Copyright 2023 Google LLC -# -# Licensed 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. -# -from __future__ import annotations - -from typing import Type - -from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse -from google.cloud.bigtable.data.row import Row, Cell, _LastScannedRow -from google.cloud.bigtable.data.exceptions import InvalidChunk - -""" -This module provides classes for the read_rows state machine: - -- ReadRowsOperation is the highest level class, providing an interface for asynchronous - merging end-to-end -- StateMachine is used internally to track the state of the merge, including - the current row key and the keys of the rows that have been processed. - It processes a stream of chunks, and will raise InvalidChunk if it reaches - an invalid state. -- State classes track the current state of the StateMachine, and define what - to do on the next chunk. -- RowBuilder is used by the StateMachine to build a Row object. -""" - - -class _StateMachine: - """ - State Machine converts chunks into Rows - - Chunks are added to the state machine via handle_chunk, which - transitions the state machine through the various states. - - When a row is complete, it will be returned from handle_chunk, - and the state machine will reset to AWAITING_NEW_ROW - - If an unexpected chunk is received for the current state, - the state machine will raise an InvalidChunk exception - - The server may send a heartbeat message indicating that it has - processed a particular row, to facilitate retries. This will be passed - to the state machine via handle_last_scanned_row, which emit a - _LastScannedRow marker to the stream. - """ - - __slots__ = ( - "current_state", - "current_family", - "current_qualifier", - "last_seen_row_key", - "adapter", - ) - - def __init__(self): - # represents either the last row emitted, or the last_scanned_key sent from backend - # all future rows should have keys > last_seen_row_key - self.last_seen_row_key: bytes | None = None - self.adapter = _RowBuilder() - self._reset_row() - - def _reset_row(self) -> None: - """ - Drops the current row and transitions to AWAITING_NEW_ROW to start a fresh one - """ - self.current_state: Type[_State] = AWAITING_NEW_ROW - self.current_family: str | None = None - self.current_qualifier: bytes | None = None - self.adapter.reset() - - def is_terminal_state(self) -> bool: - """ - Returns true if the state machine is in a terminal state (AWAITING_NEW_ROW) - - At the end of the read_rows stream, if the state machine is not in a terminal - state, an exception should be raised - """ - return self.current_state == AWAITING_NEW_ROW - - def handle_last_scanned_row(self, last_scanned_row_key: bytes) -> Row: - """ - Called by ReadRowsOperation to notify the state machine of a scan heartbeat - - Returns an empty row with the last_scanned_row_key - """ - if self.last_seen_row_key and self.last_seen_row_key >= last_scanned_row_key: - raise InvalidChunk("Last scanned row key is out of order") - if not self.current_state == AWAITING_NEW_ROW: - raise InvalidChunk("Last scanned row key received in invalid state") - scan_marker = _LastScannedRow(last_scanned_row_key) - self._handle_complete_row(scan_marker) - return scan_marker - - def handle_chunk(self, chunk: ReadRowsResponse.CellChunk) -> Row | None: - """ - Called by ReadRowsOperation to process a new chunk - - Returns a Row if the chunk completes a row, otherwise returns None - """ - if ( - self.last_seen_row_key - and chunk.row_key - and self.last_seen_row_key >= chunk.row_key - ): - raise InvalidChunk("row keys should be strictly increasing") - if chunk.reset_row: - # reset row if requested - self._handle_reset_chunk(chunk) - return None - - # process the chunk and update the state - self.current_state = self.current_state.handle_chunk(self, chunk) - if chunk.commit_row: - # check if row is complete, and return it if so - if not self.current_state == AWAITING_NEW_CELL: - raise InvalidChunk("Commit chunk received in invalid state") - complete_row = self.adapter.finish_row() - self._handle_complete_row(complete_row) - return complete_row - else: - # row is not complete, return None - return None - - def _handle_complete_row(self, complete_row: Row) -> None: - """ - Complete row, update seen keys, and move back to AWAITING_NEW_ROW - - Called by StateMachine when a commit_row flag is set on a chunk, - or when a scan heartbeat is received - """ - self.last_seen_row_key = complete_row.row_key - self._reset_row() - - def _handle_reset_chunk(self, chunk: ReadRowsResponse.CellChunk): - """ - Drop all buffers and reset the row in progress - - Called by StateMachine when a reset_row flag is set on a chunk - """ - # ensure reset chunk matches expectations - if self.current_state == AWAITING_NEW_ROW: - raise InvalidChunk("Reset chunk received when not processing row") - if chunk.row_key: - raise InvalidChunk("Reset chunk has a row key") - if _chunk_has_field(chunk, "family_name"): - raise InvalidChunk("Reset chunk has a family name") - if _chunk_has_field(chunk, "qualifier"): - raise InvalidChunk("Reset chunk has a qualifier") - if chunk.timestamp_micros: - raise InvalidChunk("Reset chunk has a timestamp") - if chunk.labels: - raise InvalidChunk("Reset chunk has labels") - if chunk.value: - raise InvalidChunk("Reset chunk has a value") - self._reset_row() - - -class _State: - """ - Represents a state the state machine can be in - - Each state is responsible for handling the next chunk, and then - transitioning to the next state - """ - - @staticmethod - def handle_chunk( - owner: _StateMachine, chunk: ReadRowsResponse.CellChunk - ) -> Type["_State"]: - raise NotImplementedError - - -class AWAITING_NEW_ROW(_State): - """ - Default state - Awaiting a chunk to start a new row - Exit states: - - AWAITING_NEW_CELL: when a chunk with a row_key is received - """ - - @staticmethod - def handle_chunk( - owner: _StateMachine, chunk: ReadRowsResponse.CellChunk - ) -> Type["_State"]: - if not chunk.row_key: - raise InvalidChunk("New row is missing a row key") - owner.adapter.start_row(chunk.row_key) - # the first chunk signals both the start of a new row and the start of a new cell, so - # force the chunk processing in the AWAITING_CELL_VALUE. - return AWAITING_NEW_CELL.handle_chunk(owner, chunk) - - -class AWAITING_NEW_CELL(_State): - """ - Represents a cell boundary witin a row - - Exit states: - - AWAITING_NEW_CELL: when the incoming cell is complete and ready for another - - AWAITING_CELL_VALUE: when the value is split across multiple chunks - """ - - @staticmethod - def handle_chunk( - owner: _StateMachine, chunk: ReadRowsResponse.CellChunk - ) -> Type["_State"]: - is_split = chunk.value_size > 0 - # track latest cell data. New chunks won't send repeated data - has_family = _chunk_has_field(chunk, "family_name") - has_qualifier = _chunk_has_field(chunk, "qualifier") - if has_family: - owner.current_family = chunk.family_name.value - if not has_qualifier: - raise InvalidChunk("New family must specify qualifier") - if has_qualifier: - owner.current_qualifier = chunk.qualifier.value - if owner.current_family is None: - raise InvalidChunk("Family not found") - - # ensure that all chunks after the first one are either missing a row - # key or the row is the same - if chunk.row_key and chunk.row_key != owner.adapter.current_key: - raise InvalidChunk("Row key changed mid row") - - if owner.current_family is None: - raise InvalidChunk("Missing family for new cell") - if owner.current_qualifier is None: - raise InvalidChunk("Missing qualifier for new cell") - - owner.adapter.start_cell( - family=owner.current_family, - qualifier=owner.current_qualifier, - labels=list(chunk.labels), - timestamp_micros=chunk.timestamp_micros, - ) - owner.adapter.cell_value(chunk.value) - # transition to new state - if is_split: - return AWAITING_CELL_VALUE - else: - # cell is complete - owner.adapter.finish_cell() - return AWAITING_NEW_CELL - - -class AWAITING_CELL_VALUE(_State): - """ - State that represents a split cell's continuation - - Exit states: - - AWAITING_NEW_CELL: when the cell is complete - - AWAITING_CELL_VALUE: when additional value chunks are required - """ - - @staticmethod - def handle_chunk( - owner: _StateMachine, chunk: ReadRowsResponse.CellChunk - ) -> Type["_State"]: - # ensure reset chunk matches expectations - if chunk.row_key: - raise InvalidChunk("In progress cell had a row key") - if _chunk_has_field(chunk, "family_name"): - raise InvalidChunk("In progress cell had a family name") - if _chunk_has_field(chunk, "qualifier"): - raise InvalidChunk("In progress cell had a qualifier") - if chunk.timestamp_micros: - raise InvalidChunk("In progress cell had a timestamp") - if chunk.labels: - raise InvalidChunk("In progress cell had labels") - is_last = chunk.value_size == 0 - owner.adapter.cell_value(chunk.value) - # transition to new state - if not is_last: - return AWAITING_CELL_VALUE - else: - # cell is complete - owner.adapter.finish_cell() - return AWAITING_NEW_CELL - - -class _RowBuilder: - """ - called by state machine to build rows - State machine makes the following guarantees: - Exactly 1 `start_row` for each row. - Exactly 1 `start_cell` for each cell. - At least 1 `cell_value` for each cell. - Exactly 1 `finish_cell` for each cell. - Exactly 1 `finish_row` for each row. - `reset` can be called at any point and can be invoked multiple times in - a row. - """ - - __slots__ = "current_key", "working_cell", "working_value", "completed_cells" - - def __init__(self): - # initialize state - self.reset() - - def reset(self) -> None: - """called when the current in progress row should be dropped""" - self.current_key: bytes | None = None - self.working_cell: Cell | None = None - self.working_value: bytearray | None = None - self.completed_cells: list[Cell] = [] - - def start_row(self, key: bytes) -> None: - """Called to start a new row. This will be called once per row""" - self.current_key = key - - def start_cell( - self, - family: str, - qualifier: bytes, - timestamp_micros: int, - labels: list[str], - ) -> None: - """called to start a new cell in a row.""" - if self.current_key is None: - raise InvalidChunk("start_cell called without a row") - self.working_value = bytearray() - self.working_cell = Cell( - b"", self.current_key, family, qualifier, timestamp_micros, labels - ) - - def cell_value(self, value: bytes) -> None: - """called multiple times per cell to concatenate the cell value""" - if self.working_value is None: - raise InvalidChunk("Cell value received before start_cell") - self.working_value.extend(value) - - def finish_cell(self) -> None: - """called once per cell to signal the end of the value (unless reset)""" - if self.working_cell is None or self.working_value is None: - raise InvalidChunk("finish_cell called before start_cell") - self.working_cell.value = bytes(self.working_value) - self.completed_cells.append(self.working_cell) - self.working_cell = None - self.working_value = None - - def finish_row(self) -> Row: - """called once per row to signal that all cells have been processed (unless reset)""" - if self.current_key is None: - raise InvalidChunk("No row in progress") - new_row = Row(self.current_key, self.completed_cells) - self.reset() - return new_row - - -def _chunk_has_field(chunk: ReadRowsResponse.CellChunk, field: str) -> bool: - """ - Returns true if the field is set on the chunk - - Required to disambiguate between empty strings and unset values - """ - try: - return chunk.HasField(field) - except ValueError: - return False diff --git a/google/cloud/bigtable/data/read_rows_query.py b/google/cloud/bigtable/data/read_rows_query.py index cf3cd316c..362f54c3e 100644 --- a/google/cloud/bigtable/data/read_rows_query.py +++ b/google/cloud/bigtable/data/read_rows_query.py @@ -17,35 +17,24 @@ from bisect import bisect_left from bisect import bisect_right from collections import defaultdict -from dataclasses import dataclass from google.cloud.bigtable.data.row_filters import RowFilter +from google.cloud.bigtable_v2.types import RowRange as RowRangePB +from google.cloud.bigtable_v2.types import RowSet as RowSetPB +from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB + if TYPE_CHECKING: from google.cloud.bigtable.data import RowKeySamples from google.cloud.bigtable.data import ShardedQuery -@dataclass -class _RangePoint: - """Model class for a point in a row range""" - - key: bytes - is_inclusive: bool - - def __hash__(self) -> int: - return hash((self.key, self.is_inclusive)) - - def __eq__(self, other: Any) -> bool: - if not isinstance(other, _RangePoint): - return NotImplemented - return self.key == other.key and self.is_inclusive == other.is_inclusive - - class RowRange: """ Represents a range of keys in a ReadRowsQuery """ + __slots__ = ("_pb",) + def __init__( self, start_key: str | bytes | None = None, @@ -71,13 +60,9 @@ def __init__( # check for invalid combinations of arguments if start_is_inclusive is None: start_is_inclusive = True - elif start_key is None: - raise ValueError("start_is_inclusive must be set with start_key") if end_is_inclusive is None: end_is_inclusive = False - elif end_key is None: - raise ValueError("end_is_inclusive must be set with end_key") # ensure that start_key and end_key are bytes if isinstance(start_key, str): start_key = start_key.encode() @@ -91,28 +76,32 @@ def __init__( if start_key is not None and end_key is not None and start_key > end_key: raise ValueError("start_key must be less than or equal to end_key") - self._start: _RangePoint | None = ( - _RangePoint(start_key, start_is_inclusive) - if start_key is not None - else None - ) - self._end: _RangePoint | None = ( - _RangePoint(end_key, end_is_inclusive) if end_key is not None else None - ) + init_dict = {} + if start_key is not None: + if start_is_inclusive: + init_dict["start_key_closed"] = start_key + else: + init_dict["start_key_open"] = start_key + if end_key is not None: + if end_is_inclusive: + init_dict["end_key_closed"] = end_key + else: + init_dict["end_key_open"] = end_key + self._pb = RowRangePB(**init_dict) @property def start_key(self) -> bytes | None: """ Returns the start key of the range. If None, the range is unbounded on the left. """ - return self._start.key if self._start is not None else None + return self._pb.start_key_closed or self._pb.start_key_open or None @property def end_key(self) -> bytes | None: """ Returns the end key of the range. If None, the range is unbounded on the right. """ - return self._end.key if self._end is not None else None + return self._pb.end_key_closed or self._pb.end_key_open or None @property def start_is_inclusive(self) -> bool: @@ -120,7 +109,7 @@ def start_is_inclusive(self) -> bool: Returns whether the range is inclusive of the start key. Returns True if the range is unbounded on the left. """ - return self._start.is_inclusive if self._start is not None else True + return not bool(self._pb.start_key_open) @property def end_is_inclusive(self) -> bool: @@ -128,61 +117,45 @@ def end_is_inclusive(self) -> bool: Returns whether the range is inclusive of the end key. Returns True if the range is unbounded on the right. """ - return self._end.is_inclusive if self._end is not None else True - - def _to_dict(self) -> dict[str, bytes]: - """Converts this object to a dictionary""" - output = {} - if self._start is not None: - key = "start_key_closed" if self.start_is_inclusive else "start_key_open" - output[key] = self._start.key - if self._end is not None: - key = "end_key_closed" if self.end_is_inclusive else "end_key_open" - output[key] = self._end.key - return output + return not bool(self._pb.end_key_open) - def __hash__(self) -> int: - return hash((self._start, self._end)) + def _to_pb(self) -> RowRangePB: + """Converts this object to a protobuf""" + return self._pb @classmethod - def _from_dict(cls, data: dict[str, bytes]) -> RowRange: - """Creates a RowRange from a dictionary""" - start_key = data.get("start_key_closed", data.get("start_key_open")) - end_key = data.get("end_key_closed", data.get("end_key_open")) - start_is_inclusive = "start_key_closed" in data if start_key else None - end_is_inclusive = "end_key_closed" in data if end_key else None - return cls( - start_key, - end_key, - start_is_inclusive, - end_is_inclusive, - ) + def _from_pb(cls, data: RowRangePB) -> RowRange: + """Creates a RowRange from a protobuf""" + instance = cls() + instance._pb = data + return instance @classmethod - def _from_points( - cls, start: _RangePoint | None, end: _RangePoint | None - ) -> RowRange: - """Creates a RowRange from two RangePoints""" - kwargs: dict[str, Any] = {} - if start is not None: - kwargs["start_key"] = start.key - kwargs["start_is_inclusive"] = start.is_inclusive - if end is not None: - kwargs["end_key"] = end.key - kwargs["end_is_inclusive"] = end.is_inclusive - return cls(**kwargs) + def _from_dict(cls, data: dict[str, bytes | str]) -> RowRange: + """Creates a RowRange from a protobuf""" + formatted_data = { + k: v.encode() if isinstance(v, str) else v for k, v in data.items() + } + instance = cls() + instance._pb = RowRangePB(**formatted_data) + return instance def __bool__(self) -> bool: """ Empty RowRanges (representing a full table scan) are falsy, because they can be substituted with None. Non-empty RowRanges are truthy. """ - return self._start is not None or self._end is not None + return bool( + self._pb.start_key_closed + or self._pb.start_key_open + or self._pb.end_key_closed + or self._pb.end_key_open + ) def __eq__(self, other: Any) -> bool: if not isinstance(other, RowRange): return NotImplemented - return self._start == other._start and self._end == other._end + return self._pb == other._pb def __str__(self) -> str: """ @@ -202,7 +175,7 @@ def __repr__(self) -> str: if self.start_is_inclusive is False: # only show start_is_inclusive if it is different from the default args_list.append(f"start_is_inclusive={self.start_is_inclusive}") - if self.end_is_inclusive is True and self._end is not None: + if self.end_is_inclusive is True and self.end_key is not None: # only show end_is_inclusive if it is different from the default args_list.append(f"end_is_inclusive={self.end_is_inclusive}") return f"RowRange({', '.join(args_list)})" @@ -213,6 +186,8 @@ class ReadRowsQuery: Class to encapsulate details of a read row request """ + slots = ("_limit", "_filter", "_row_set") + def __init__( self, row_keys: list[str | bytes] | str | bytes | None = None, @@ -231,24 +206,32 @@ def __init__( default: None (no limit) - row_filter: a RowFilter to apply to the query """ - self.row_keys: set[bytes] = set() - self.row_ranges: set[RowRange] = set() - if row_ranges is not None: - if isinstance(row_ranges, RowRange): - row_ranges = [row_ranges] - for r in row_ranges: - self.add_range(r) - if row_keys is not None: - if not isinstance(row_keys, list): - row_keys = [row_keys] - for k in row_keys: - self.add_key(k) - self.limit: int | None = limit - self.filter: RowFilter | None = row_filter + if row_keys is None: + row_keys = [] + if row_ranges is None: + row_ranges = [] + if not isinstance(row_ranges, list): + row_ranges = [row_ranges] + if not isinstance(row_keys, list): + row_keys = [row_keys] + row_keys = [key.encode() if isinstance(key, str) else key for key in row_keys] + self._row_set = RowSetPB( + row_keys=row_keys, row_ranges=[r._pb for r in row_ranges] + ) + self.limit = limit or None + self.filter = row_filter + + @property + def row_keys(self) -> list[bytes]: + return list(self._row_set.row_keys) + + @property + def row_ranges(self) -> list[RowRange]: + return [RowRange._from_pb(r) for r in self._row_set.row_ranges] @property def limit(self) -> int | None: - return self._limit + return self._limit or None @limit.setter def limit(self, new_limit: int | None): @@ -279,16 +262,9 @@ def filter(self, row_filter: RowFilter | None): Args: - row_filter: a RowFilter to apply to this query - Can be a RowFilter object or a dict representation Returns: - a reference to this query for chaining """ - if not ( - isinstance(row_filter, dict) - or isinstance(row_filter, RowFilter) - or row_filter is None - ): - raise ValueError("row_filter must be a RowFilter or dict") self._filter = row_filter def add_key(self, row_key: str | bytes): @@ -308,25 +284,21 @@ def add_key(self, row_key: str | bytes): row_key = row_key.encode() elif not isinstance(row_key, bytes): raise ValueError("row_key must be string or bytes") - self.row_keys.add(row_key) + if row_key not in self._row_set.row_keys: + self._row_set.row_keys.append(row_key) def add_range( self, - row_range: RowRange | dict[str, bytes], + row_range: RowRange, ): """ Add a range of row keys to this query. Args: - row_range: a range of row keys to add to this query - Can be a RowRange object or a dict representation in - RowRange proto format """ - if not (isinstance(row_range, dict) or isinstance(row_range, RowRange)): - raise ValueError("row_range must be a RowRange or dict") - if isinstance(row_range, dict): - row_range = RowRange._from_dict(row_range) - self.row_ranges.add(row_range) + if row_range not in self.row_ranges: + self._row_set.row_ranges.append(row_range._pb) def shard(self, shard_keys: RowKeySamples) -> ShardedQuery: """ @@ -392,24 +364,24 @@ def _shard_range( - a list of tuples, containing a segment index and a new sub-range. """ # 1. find the index of the segment the start key belongs to - if orig_range._start is None: + if orig_range.start_key is None: # if range is open on the left, include first segment start_segment = 0 else: # use binary search to find the segment the start key belongs to # bisect method determines how we break ties when the start key matches a split point # if inclusive, bisect_left to the left segment, otherwise bisect_right - bisect = bisect_left if orig_range._start.is_inclusive else bisect_right - start_segment = bisect(split_points, orig_range._start.key) + bisect = bisect_left if orig_range.start_is_inclusive else bisect_right + start_segment = bisect(split_points, orig_range.start_key) # 2. find the index of the segment the end key belongs to - if orig_range._end is None: + if orig_range.end_key is None: # if range is open on the right, include final segment end_segment = len(split_points) else: # use binary search to find the segment the end key belongs to. end_segment = bisect_left( - split_points, orig_range._end.key, lo=start_segment + split_points, orig_range.end_key, lo=start_segment ) # note: end_segment will always bisect_left, because split points represent inclusive ends # whether the end_key is includes the split point or not, the result is the same segment @@ -424,18 +396,22 @@ def _shard_range( # 3a. add new range for first segment this_range spans # first range spans from start_key to the split_point representing the last key in the segment last_key_in_first_segment = split_points[start_segment] - start_range = RowRange._from_points( - start=orig_range._start, - end=_RangePoint(last_key_in_first_segment, is_inclusive=True), + start_range = RowRange( + start_key=orig_range.start_key, + start_is_inclusive=orig_range.start_is_inclusive, + end_key=last_key_in_first_segment, + end_is_inclusive=True, ) results.append((start_segment, start_range)) # 3b. add new range for last segment this_range spans # we start the final range using the end key from of the previous segment, with is_inclusive=False previous_segment = end_segment - 1 last_key_before_segment = split_points[previous_segment] - end_range = RowRange._from_points( - start=_RangePoint(last_key_before_segment, is_inclusive=False), - end=orig_range._end, + end_range = RowRange( + start_key=last_key_before_segment, + start_is_inclusive=False, + end_key=orig_range.end_key, + end_is_inclusive=orig_range.end_is_inclusive, ) results.append((end_segment, end_range)) # 3c. add new spanning range to all segments other than the first and last @@ -452,31 +428,18 @@ def _shard_range( results.append((this_segment, new_range)) return results - def _to_dict(self) -> dict[str, Any]: + def _to_pb(self, table) -> ReadRowsRequestPB: """ Convert this query into a dictionary that can be used to construct a ReadRowsRequest protobuf """ - row_ranges = [] - for r in self.row_ranges: - dict_range = r._to_dict() if isinstance(r, RowRange) else r - row_ranges.append(dict_range) - row_keys = list(self.row_keys) - row_keys.sort() - row_set = {"row_keys": row_keys, "row_ranges": row_ranges} - final_dict: dict[str, Any] = { - "rows": row_set, - } - dict_filter = ( - self.filter._to_dict() - if isinstance(self.filter, RowFilter) - else self.filter + return ReadRowsRequestPB( + table_name=table.table_name, + app_profile_id=table.app_profile_id, + filter=self.filter._to_pb() if self.filter else None, + rows_limit=self.limit or 0, + rows=self._row_set, ) - if dict_filter: - final_dict["filter"] = dict_filter - if self.limit is not None: - final_dict["rows_limit"] = self.limit - return final_dict def __eq__(self, other): """ diff --git a/google/cloud/bigtable/data/row.py b/google/cloud/bigtable/data/row.py index f562e96d6..ecf9cea66 100644 --- a/google/cloud/bigtable/data/row.py +++ b/google/cloud/bigtable/data/row.py @@ -15,7 +15,7 @@ from __future__ import annotations from collections import OrderedDict -from typing import Sequence, Generator, overload, Any +from typing import Generator, overload, Any from functools import total_ordering from google.cloud.bigtable_v2.types import Row as RowPB @@ -25,7 +25,7 @@ _qualifier_type = bytes -class Row(Sequence["Cell"]): +class Row: """ Model class for row data returned from server @@ -309,21 +309,6 @@ def __ne__(self, other) -> bool: return not self == other -class _LastScannedRow(Row): - """A value used to indicate a scanned row that is not returned as part of - a query. - - This is used internally to indicate progress in a scan, and improve retry - performance. It is not intended to be used directly by users. - """ - - def __init__(self, row_key): - super().__init__(row_key, []) - - def __eq__(self, other): - return isinstance(other, _LastScannedRow) - - @total_ordering class Cell: """ diff --git a/tests/system/conftest.py b/tests/system/conftest.py new file mode 100644 index 000000000..b8862ea4b --- /dev/null +++ b/tests/system/conftest.py @@ -0,0 +1,25 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +""" +Import pytest fixtures for setting up table for data client system tests +""" +import sys +import os + +script_path = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(script_path) + +pytest_plugins = [ + "data.setup_fixtures", +] diff --git a/tests/system/data/setup_fixtures.py b/tests/system/data/setup_fixtures.py new file mode 100644 index 000000000..7a7faa5f5 --- /dev/null +++ b/tests/system/data/setup_fixtures.py @@ -0,0 +1,171 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +""" +Contains a set of pytest fixtures for setting up and populating a +Bigtable database for testing purposes. +""" + +import pytest +import pytest_asyncio +import os +import asyncio +import uuid + + +@pytest.fixture(scope="session") +def event_loop(): + loop = asyncio.get_event_loop() + yield loop + loop.stop() + loop.close() + + +@pytest.fixture(scope="session") +def instance_admin_client(): + """Client for interacting with the Instance Admin API.""" + from google.cloud.bigtable_admin_v2 import BigtableInstanceAdminClient + + with BigtableInstanceAdminClient() as client: + yield client + + +@pytest.fixture(scope="session") +def table_admin_client(): + """Client for interacting with the Table Admin API.""" + from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient + + with BigtableTableAdminClient() as client: + yield client + + +@pytest.fixture(scope="session") +def instance_id(instance_admin_client, project_id, cluster_config): + """ + Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session + """ + from google.cloud.bigtable_admin_v2 import types + from google.api_core import exceptions + + # use user-specified instance if available + user_specified_instance = os.getenv("BIGTABLE_TEST_INSTANCE") + if user_specified_instance: + print("Using user-specified instance: {}".format(user_specified_instance)) + yield user_specified_instance + return + + # create a new temporary test instance + instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" + try: + operation = instance_admin_client.create_instance( + parent=f"projects/{project_id}", + instance_id=instance_id, + instance=types.Instance( + display_name="Test Instance", + labels={"python-system-test": "true"}, + ), + clusters=cluster_config, + ) + operation.result(timeout=240) + except exceptions.AlreadyExists: + pass + yield instance_id + instance_admin_client.delete_instance( + name=f"projects/{project_id}/instances/{instance_id}" + ) + + +@pytest.fixture(scope="session") +def column_split_config(): + """ + specify initial splits to create when creating a new test table + """ + return [(num * 1000).to_bytes(8, "big") for num in range(1, 10)] + + +@pytest.fixture(scope="session") +def table_id( + table_admin_client, + project_id, + instance_id, + column_family_config, + init_table_id, + column_split_config, +): + """ + Returns BIGTABLE_TEST_TABLE if set, otherwise creates a new temporary table for the test session + + Args: + - table_admin_client: Client for interacting with the Table Admin API. Supplied by the table_admin_client fixture. + - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture. + - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture. + - init_column_families: A list of column families to initialize the table with, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. + Supplied by the init_column_families fixture. + - init_table_id: The table ID to give to the test table, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. + Supplied by the init_table_id fixture. + - column_split_config: A list of row keys to use as initial splits when creating the test table. + """ + from google.api_core import exceptions + from google.api_core import retry + + # use user-specified instance if available + user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") + if user_specified_table: + print("Using user-specified table: {}".format(user_specified_table)) + yield user_specified_table + return + + retry = retry.Retry( + predicate=retry.if_exception_type(exceptions.FailedPrecondition) + ) + try: + parent_path = f"projects/{project_id}/instances/{instance_id}" + print(f"Creating table: {parent_path}/tables/{init_table_id}") + table_admin_client.create_table( + request={ + "parent": parent_path, + "table_id": init_table_id, + "table": {"column_families": column_family_config}, + "initial_splits": [{"key": key} for key in column_split_config], + }, + retry=retry, + ) + except exceptions.AlreadyExists: + pass + yield init_table_id + print(f"Deleting table: {parent_path}/tables/{init_table_id}") + try: + table_admin_client.delete_table(name=f"{parent_path}/tables/{init_table_id}") + except exceptions.NotFound: + print(f"Table {init_table_id} not found, skipping deletion") + + +@pytest_asyncio.fixture(scope="session") +async def client(): + from google.cloud.bigtable.data import BigtableDataClientAsync + + project = os.getenv("GOOGLE_CLOUD_PROJECT") or None + async with BigtableDataClientAsync(project=project, pool_size=4) as client: + yield client + + +@pytest.fixture(scope="session") +def project_id(client): + """Returns the project ID from the client.""" + yield client.project + + +@pytest_asyncio.fixture(scope="session") +async def table(client, table_id, instance_id): + async with client.get_table(instance_id, table_id) as table: + yield table diff --git a/tests/system/data/test_system.py b/tests/system/data/test_system.py index fe341e4a8..6bd21f386 100644 --- a/tests/system/data/test_system.py +++ b/tests/system/data/test_system.py @@ -14,7 +14,6 @@ import pytest import pytest_asyncio -import os import asyncio import uuid from google.api_core import retry @@ -27,131 +26,37 @@ @pytest.fixture(scope="session") -def event_loop(): - loop = asyncio.get_event_loop() - yield loop - loop.stop() - loop.close() - - -@pytest.fixture(scope="session") -def instance_admin_client(): - """Client for interacting with the Instance Admin API.""" - from google.cloud.bigtable_admin_v2 import BigtableInstanceAdminClient - - with BigtableInstanceAdminClient() as client: - yield client - - -@pytest.fixture(scope="session") -def table_admin_client(): - """Client for interacting with the Table Admin API.""" - from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient - - with BigtableTableAdminClient() as client: - yield client - - -@pytest.fixture(scope="session") -def instance_id(instance_admin_client, project_id): +def column_family_config(): """ - Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session + specify column families to create when creating a new test table """ from google.cloud.bigtable_admin_v2 import types - from google.api_core import exceptions - - # use user-specified instance if available - user_specified_instance = os.getenv("BIGTABLE_TEST_INSTANCE") - if user_specified_instance: - print("Using user-specified instance: {}".format(user_specified_instance)) - yield user_specified_instance - return - # create a new temporary test instance - instance_id = "test-instance" - try: - operation = instance_admin_client.create_instance( - parent=f"projects/{project_id}", - instance_id=instance_id, - instance=types.Instance( - display_name="Test Instance", - labels={"python-system-test": "true"}, - ), - clusters={ - "test-cluster": types.Cluster( - location=f"projects/{project_id}/locations/us-central1-b", - serve_nodes=3, - ) - }, - ) - operation.result(timeout=240) - except exceptions.AlreadyExists: - pass - yield instance_id - instance_admin_client.delete_instance( - name=f"projects/{project_id}/instances/{instance_id}" - ) + return {TEST_FAMILY: types.ColumnFamily(), TEST_FAMILY_2: types.ColumnFamily()} @pytest.fixture(scope="session") -def table_id(table_admin_client, project_id, instance_id): +def init_table_id(): """ - Returns BIGTABLE_TEST_TABLE if set, otherwise creates a new temporary table for the test session + The table_id to use when creating a new test table """ - from google.cloud.bigtable_admin_v2 import types - from google.api_core import exceptions - from google.api_core import retry - - # use user-specified instance if available - user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") - if user_specified_table: - print("Using user-specified table: {}".format(user_specified_table)) - yield user_specified_table - return - - table_id = "test-table" - retry = retry.Retry( - predicate=retry.if_exception_type(exceptions.FailedPrecondition) - ) - try: - table_admin_client.create_table( - parent=f"projects/{project_id}/instances/{instance_id}", - table_id=table_id, - table=types.Table( - column_families={ - TEST_FAMILY: types.ColumnFamily(), - TEST_FAMILY_2: types.ColumnFamily(), - }, - ), - retry=retry, - ) - except exceptions.AlreadyExists: - pass - yield table_id - table_admin_client.delete_table( - name=f"projects/{project_id}/instances/{instance_id}/tables/{table_id}" - ) - - -@pytest_asyncio.fixture(scope="session") -async def client(): - from google.cloud.bigtable.data import BigtableDataClientAsync - - project = os.getenv("GOOGLE_CLOUD_PROJECT") or None - async with BigtableDataClientAsync(project=project) as client: - yield client + return f"test-table-{uuid.uuid4().hex}" @pytest.fixture(scope="session") -def project_id(client): - """Returns the project ID from the client.""" - yield client.project - +def cluster_config(project_id): + """ + Configuration for the clusters to use when creating a new instance + """ + from google.cloud.bigtable_admin_v2 import types -@pytest_asyncio.fixture(scope="session") -async def table(client, table_id, instance_id): - async with client.get_table(instance_id, table_id) as table: - yield table + cluster = { + "test-cluster": types.Cluster( + location=f"projects/{project_id}/locations/us-central1-b", + serve_nodes=1, + ) + } + return cluster class TempRowBuilder: @@ -197,6 +102,7 @@ async def delete_rows(self): await self.table.client._gapic_client.mutate_rows(request) +@pytest.mark.usefixtures("table") async def _retrieve_cell_value(table, row_key): """ Helper to read an individual row @@ -231,6 +137,7 @@ async def _create_row_and_mutation( return row_key, mutation +@pytest.mark.usefixtures("table") @pytest_asyncio.fixture(scope="function") async def temp_rows(table): builder = TempRowBuilder(table) @@ -238,6 +145,8 @@ async def temp_rows(table): await builder.delete_rows() +@pytest.mark.usefixtures("table") +@pytest.mark.usefixtures("client") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=10) @pytest.mark.asyncio async def test_ping_and_warm_gapic(client, table): @@ -249,6 +158,8 @@ async def test_ping_and_warm_gapic(client, table): await client._gapic_client.ping_and_warm(request) +@pytest.mark.usefixtures("table") +@pytest.mark.usefixtures("client") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_ping_and_warm(client, table): @@ -265,6 +176,7 @@ async def test_ping_and_warm(client, table): assert results[0] is None +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutation_set_cell(table, temp_rows): @@ -282,9 +194,11 @@ async def test_mutation_set_cell(table, temp_rows): assert (await _retrieve_cell_value(table, row_key)) == new_value +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio -async def test_sample_row_keys(client, table, temp_rows): +async def test_sample_row_keys(client, table, temp_rows, column_split_config): """ Sample keys should return a single sample in small test tables """ @@ -292,12 +206,18 @@ async def test_sample_row_keys(client, table, temp_rows): await temp_rows.add_row(b"row_key_2") results = await table.sample_row_keys() - assert len(results) == 1 - sample = results[0] - assert isinstance(sample[0], bytes) - assert isinstance(sample[1], int) - - + assert len(results) == len(column_split_config) + 1 + # first keys should match the split config + for idx in range(len(column_split_config)): + assert results[idx][0] == column_split_config[idx] + assert isinstance(results[idx][1], int) + # last sample should be empty key + assert results[-1][0] == b"" + assert isinstance(results[-1][1], int) + + +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_bulk_mutations_set_cell(client, table, temp_rows): """ @@ -317,6 +237,8 @@ async def test_bulk_mutations_set_cell(client, table, temp_rows): assert (await _retrieve_cell_value(table, row_key)) == new_value +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_context_manager(client, table, temp_rows): @@ -343,6 +265,8 @@ async def test_mutations_batcher_context_manager(client, table, temp_rows): assert len(batcher._staged_entries) == 0 +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_timer_flush(client, table, temp_rows): @@ -367,6 +291,8 @@ async def test_mutations_batcher_timer_flush(client, table, temp_rows): assert (await _retrieve_cell_value(table, row_key)) == new_value +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_count_flush(client, table, temp_rows): @@ -401,6 +327,8 @@ async def test_mutations_batcher_count_flush(client, table, temp_rows): assert (await _retrieve_cell_value(table, row_key2)) == new_value2 +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_bytes_flush(client, table, temp_rows): @@ -436,6 +364,8 @@ async def test_mutations_batcher_bytes_flush(client, table, temp_rows): assert (await _retrieve_cell_value(table, row_key2)) == new_value2 +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_no_flush(client, table, temp_rows): @@ -472,6 +402,8 @@ async def test_mutations_batcher_no_flush(client, table, temp_rows): assert (await _retrieve_cell_value(table, row_key2)) == start_value +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @pytest.mark.parametrize( "start,increment,expected", [ @@ -512,6 +444,8 @@ async def test_read_modify_write_row_increment( assert (await _retrieve_cell_value(table, row_key)) == result[0].value +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @pytest.mark.parametrize( "start,append,expected", [ @@ -549,6 +483,8 @@ async def test_read_modify_write_row_append( assert (await _retrieve_cell_value(table, row_key)) == result[0].value +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_read_modify_write_row_chained(client, table, temp_rows): """ @@ -585,6 +521,8 @@ async def test_read_modify_write_row_chained(client, table, temp_rows): assert (await _retrieve_cell_value(table, row_key)) == result[0].value +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") @pytest.mark.parametrize( "start_val,predicate_range,expected_result", [ @@ -631,6 +569,7 @@ async def test_check_and_mutate( assert (await _retrieve_cell_value(table, row_key)) == expected_value +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_stream(table, temp_rows): @@ -650,6 +589,7 @@ async def test_read_rows_stream(table, temp_rows): await generator.__anext__() +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows(table, temp_rows): @@ -665,6 +605,7 @@ async def test_read_rows(table, temp_rows): assert row_list[1].row_key == b"row_key_2" +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_sharded_simple(table, temp_rows): @@ -687,6 +628,7 @@ async def test_read_rows_sharded_simple(table, temp_rows): assert row_list[3].row_key == b"d" +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_sharded_from_sample(table, temp_rows): @@ -711,6 +653,7 @@ async def test_read_rows_sharded_from_sample(table, temp_rows): assert row_list[2].row_key == b"d" +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_sharded_filters_limits(table, temp_rows): @@ -739,6 +682,7 @@ async def test_read_rows_sharded_filters_limits(table, temp_rows): assert row_list[2][0].labels == ["second"] +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_range_query(table, temp_rows): @@ -760,6 +704,7 @@ async def test_read_rows_range_query(table, temp_rows): assert row_list[1].row_key == b"c" +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_single_key_query(table, temp_rows): @@ -780,6 +725,7 @@ async def test_read_rows_single_key_query(table, temp_rows): assert row_list[1].row_key == b"c" +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_with_filter(table, temp_rows): @@ -803,46 +749,29 @@ async def test_read_rows_with_filter(table, temp_rows): assert row[0].labels == [expected_label] +@pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_read_rows_stream_close(table, temp_rows): """ Ensure that the read_rows_stream can be closed """ + from google.cloud.bigtable.data import ReadRowsQuery + await temp_rows.add_row(b"row_key_1") await temp_rows.add_row(b"row_key_2") - # full table scan - generator = await table.read_rows_stream({}) + query = ReadRowsQuery() + generator = await table.read_rows_stream(query) + # grab first row first_row = await generator.__anext__() assert first_row.row_key == b"row_key_1" + # close stream early await generator.aclose() - assert generator.active is False - with pytest.raises(StopAsyncIteration) as e: - await generator.__anext__() - assert "closed" in str(e) - - -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) -@pytest.mark.asyncio -async def test_read_rows_stream_inactive_timer(table, temp_rows): - """ - Ensure that the read_rows_stream method works - """ - from google.cloud.bigtable.data.exceptions import IdleTimeout - - await temp_rows.add_row(b"row_key_1") - await temp_rows.add_row(b"row_key_2") - - generator = await table.read_rows_stream({}) - await generator._start_idle_timer(0.05) - await asyncio.sleep(0.2) - assert generator.active is False - with pytest.raises(IdleTimeout) as e: + with pytest.raises(StopAsyncIteration): await generator.__anext__() - assert "inactivity" in str(e) - assert "idle_timeout=0.1" in str(e) +@pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_read_row(table, temp_rows): """ @@ -857,6 +786,7 @@ async def test_read_row(table, temp_rows): assert row.cells[0].value == b"value" +@pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_read_row_missing(table): """ @@ -872,6 +802,7 @@ async def test_read_row_missing(table): assert "Row key must be non-empty" in str(e) +@pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_read_row_w_filter(table, temp_rows): """ @@ -890,6 +821,7 @@ async def test_read_row_w_filter(table, temp_rows): assert row.cells[0].labels == [expected_label] +@pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_row_exists(table, temp_rows): from google.api_core import exceptions @@ -909,6 +841,7 @@ async def test_row_exists(table, temp_rows): assert "Row kest must be non-empty" in str(e) +@pytest.mark.usefixtures("table") @retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.parametrize( "cell_value,filter_input,expect_match", diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index 0388d62e9..08422abca 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -185,7 +185,7 @@ async def test_mutate_rows_attempt_exception(self, exc_type): except Exception as e: found_exc = e assert client.mutate_rows.call_count == 1 - assert type(found_exc) == exc_type + assert type(found_exc) is exc_type assert found_exc == expected_exception assert len(instance.errors) == 2 assert len(instance.remaining_indices) == 0 diff --git a/tests/unit/data/_async/test__read_rows.py b/tests/unit/data/_async/test__read_rows.py index 76e1148de..200defbbf 100644 --- a/tests/unit/data/_async/test__read_rows.py +++ b/tests/unit/data/_async/test__read_rows.py @@ -12,8 +12,6 @@ # limitations under the License. import pytest -import sys -import asyncio from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync @@ -47,37 +45,18 @@ def _get_target_class(): def _make_one(self, *args, **kwargs): return self._get_target_class()(*args, **kwargs) - def test_ctor_defaults(self): - request = {} - client = mock.Mock() - client.read_rows = mock.Mock() - client.read_rows.return_value = None - default_operation_timeout = 600 - time_gen_mock = mock.Mock() - with mock.patch( - "google.cloud.bigtable.data._async._read_rows._attempt_timeout_generator", - time_gen_mock, - ): - instance = self._make_one(request, client) - assert time_gen_mock.call_count == 1 - time_gen_mock.assert_called_once_with(None, default_operation_timeout) - assert instance.transient_errors == [] - assert instance._last_emitted_row_key is None - assert instance._emit_count == 0 - assert instance.operation_timeout == default_operation_timeout - retryable_fn = instance._partial_retryable - assert retryable_fn.func == instance._read_rows_retryable_attempt - assert retryable_fn.args[0] == client.read_rows - assert retryable_fn.args[1] == time_gen_mock.return_value - assert retryable_fn.args[2] == 0 - assert client.read_rows.call_count == 0 - def test_ctor(self): + from google.cloud.bigtable.data import ReadRowsQuery + row_limit = 91 - request = {"rows_limit": row_limit} + query = ReadRowsQuery(limit=row_limit) client = mock.Mock() client.read_rows = mock.Mock() client.read_rows.return_value = None + table = mock.Mock() + table._client = client + table.table_name = "test_table" + table.app_profile_id = "test_profile" expected_operation_timeout = 42 expected_request_timeout = 44 time_gen_mock = mock.Mock() @@ -86,8 +65,8 @@ def test_ctor(self): time_gen_mock, ): instance = self._make_one( - request, - client, + query, + table, operation_timeout=expected_operation_timeout, attempt_timeout=expected_request_timeout, ) @@ -95,39 +74,19 @@ def test_ctor(self): time_gen_mock.assert_called_once_with( expected_request_timeout, expected_operation_timeout ) - assert instance.transient_errors == [] - assert instance._last_emitted_row_key is None - assert instance._emit_count == 0 + assert instance._last_yielded_row_key is None + assert instance._remaining_count == row_limit assert instance.operation_timeout == expected_operation_timeout - retryable_fn = instance._partial_retryable - assert retryable_fn.func == instance._read_rows_retryable_attempt - assert retryable_fn.args[0] == client.read_rows - assert retryable_fn.args[1] == time_gen_mock.return_value - assert retryable_fn.args[2] == row_limit assert client.read_rows.call_count == 0 - - def test___aiter__(self): - request = {} - client = mock.Mock() - client.read_rows = mock.Mock() - instance = self._make_one(request, client) - assert instance.__aiter__() is instance - - @pytest.mark.asyncio - async def test_transient_error_capture(self): - from google.api_core import exceptions as core_exceptions - - client = mock.Mock() - client.read_rows = mock.Mock() - test_exc = core_exceptions.Aborted("test") - test_exc2 = core_exceptions.DeadlineExceeded("test") - client.read_rows.side_effect = [test_exc, test_exc2] - instance = self._make_one({}, client) - with pytest.raises(RuntimeError): - await instance.__anext__() - assert len(instance.transient_errors) == 2 - assert instance.transient_errors[0] == test_exc - assert instance.transient_errors[1] == test_exc2 + assert instance._metadata == [ + ( + "x-goog-request-params", + "table_name=test_table&app_profile_id=test_profile", + ) + ] + assert instance.request.table_name == table.table_name + assert instance.request.app_profile_id == table.app_profile_id + assert instance.request.rows_limit == row_limit @pytest.mark.parametrize( "in_keys,last_key,expected", @@ -140,11 +99,18 @@ async def test_transient_error_capture(self): ], ) def test_revise_request_rowset_keys(self, in_keys, last_key, expected): - sample_range = {"start_key_open": last_key} - row_set = {"row_keys": in_keys, "row_ranges": [sample_range]} + from google.cloud.bigtable_v2.types import RowSet as RowSetPB + from google.cloud.bigtable_v2.types import RowRange as RowRangePB + + in_keys = [key.encode("utf-8") for key in in_keys] + expected = [key.encode("utf-8") for key in expected] + last_key = last_key.encode("utf-8") + + sample_range = RowRangePB(start_key_open=last_key) + row_set = RowSetPB(row_keys=in_keys, row_ranges=[sample_range]) revised = self._get_target_class()._revise_request_rowset(row_set, last_key) - assert revised["row_keys"] == expected - assert revised["row_ranges"] == [sample_range] + assert revised.row_keys == expected + assert revised.row_ranges == [sample_range] @pytest.mark.parametrize( "in_ranges,last_key,expected", @@ -192,31 +158,52 @@ def test_revise_request_rowset_keys(self, in_keys, last_key, expected): ], ) def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected): - next_key = last_key + "a" - row_set = {"row_keys": [next_key], "row_ranges": in_ranges} + from google.cloud.bigtable_v2.types import RowSet as RowSetPB + from google.cloud.bigtable_v2.types import RowRange as RowRangePB + + # convert to protobuf + next_key = (last_key + "a").encode("utf-8") + last_key = last_key.encode("utf-8") + in_ranges = [ + RowRangePB(**{k: v.encode("utf-8") for k, v in r.items()}) + for r in in_ranges + ] + expected = [ + RowRangePB(**{k: v.encode("utf-8") for k, v in r.items()}) for r in expected + ] + + row_set = RowSetPB(row_ranges=in_ranges, row_keys=[next_key]) revised = self._get_target_class()._revise_request_rowset(row_set, last_key) - assert revised["row_keys"] == [next_key] - assert revised["row_ranges"] == expected + assert revised.row_keys == [next_key] + assert revised.row_ranges == expected @pytest.mark.parametrize("last_key", ["a", "b", "c"]) def test_revise_request_full_table(self, last_key): - row_set = {"row_keys": [], "row_ranges": []} + from google.cloud.bigtable_v2.types import RowSet as RowSetPB + from google.cloud.bigtable_v2.types import RowRange as RowRangePB + + # convert to protobuf + last_key = last_key.encode("utf-8") + row_set = RowSetPB() for selected_set in [row_set, None]: revised = self._get_target_class()._revise_request_rowset( selected_set, last_key ) - assert revised["row_keys"] == [] - assert len(revised["row_ranges"]) == 1 - assert revised["row_ranges"][0]["start_key_open"] == last_key + assert revised.row_keys == [] + assert len(revised.row_ranges) == 1 + assert revised.row_ranges[0] == RowRangePB(start_key_open=last_key) def test_revise_to_empty_rowset(self): """revising to an empty rowset should raise error""" from google.cloud.bigtable.data.exceptions import _RowSetComplete + from google.cloud.bigtable_v2.types import RowSet as RowSetPB + from google.cloud.bigtable_v2.types import RowRange as RowRangePB - row_keys = ["a", "b", "c"] - row_set = {"row_keys": row_keys, "row_ranges": [{"end_key_open": "c"}]} + row_keys = [b"a", b"b", b"c"] + row_range = RowRangePB(end_key_open=b"c") + row_set = RowSetPB(row_keys=row_keys, row_ranges=[row_range]) with pytest.raises(_RowSetComplete): - self._get_target_class()._revise_request_rowset(row_set, "d") + self._get_target_class()._revise_request_rowset(row_set, b"d") @pytest.mark.parametrize( "start_limit,emit_num,expected_limit", @@ -224,8 +211,8 @@ def test_revise_to_empty_rowset(self): (10, 0, 10), (10, 1, 9), (10, 10, 0), - (0, 10, 0), - (0, 0, 0), + (None, 10, None), + (None, 0, None), (4, 2, 2), ], ) @@ -238,27 +225,26 @@ async def test_revise_limit(self, start_limit, emit_num, expected_limit): - if the number emitted exceeds the new limit, an exception should should be raised (tested in test_revise_limit_over_limit) """ - import itertools - - request = {"rows_limit": start_limit} - instance = self._make_one(request, mock.Mock()) - instance._emit_count = emit_num - instance._last_emitted_row_key = "a" - gapic_mock = mock.Mock() - gapic_mock.side_effect = [GeneratorExit("stop_fn")] - mock_timeout_gen = itertools.repeat(5) - - attempt = instance._read_rows_retryable_attempt( - gapic_mock, mock_timeout_gen, start_limit - ) - if start_limit != 0 and expected_limit == 0: - # if we emitted the expected number of rows, we should receive a StopAsyncIteration - with pytest.raises(StopAsyncIteration): - await attempt.__anext__() - else: - with pytest.raises(GeneratorExit): - await attempt.__anext__() - assert request["rows_limit"] == expected_limit + from google.cloud.bigtable.data import ReadRowsQuery + + async def mock_stream(): + for i in range(emit_num): + yield i + + query = ReadRowsQuery(limit=start_limit) + table = mock.Mock() + table.table_name = "table_name" + table.app_profile_id = "app_profile_id" + with mock.patch.object( + _ReadRowsOperationAsync, "_read_rows_attempt" + ) as mock_attempt: + mock_attempt.return_value = mock_stream() + instance = self._make_one(query, table, 10, 10) + assert instance._remaining_count == start_limit + # read emit_num rows + async for val in instance.start_operation(): + pass + assert instance._remaining_count == expected_limit @pytest.mark.parametrize("start_limit,emit_num", [(5, 10), (3, 9), (1, 10)]) @pytest.mark.asyncio @@ -267,137 +253,91 @@ async def test_revise_limit_over_limit(self, start_limit, emit_num): Should raise runtime error if we get in state where emit_num > start_num (unless start_num == 0, which represents unlimited) """ - import itertools - - request = {"rows_limit": start_limit} - instance = self._make_one(request, mock.Mock()) - instance._emit_count = emit_num - instance._last_emitted_row_key = "a" - mock_timeout_gen = itertools.repeat(5) - attempt = instance._read_rows_retryable_attempt( - mock.Mock(), mock_timeout_gen, start_limit - ) - with pytest.raises(RuntimeError) as e: - await attempt.__anext__() - assert "emit count exceeds row limit" in str(e.value) - - @pytest.mark.asyncio - async def test_aclose(self): - import asyncio - - instance = self._make_one({}, mock.Mock()) - await instance.aclose() - assert instance._stream is None - assert instance._last_emitted_row_key is None - with pytest.raises(asyncio.InvalidStateError): - await instance.__anext__() - # try calling a second time - await instance.aclose() - - @pytest.mark.parametrize("limit", [1, 3, 10]) - @pytest.mark.asyncio - async def test_retryable_attempt_hit_limit(self, limit): - """ - Stream should end after hitting the limit - """ - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - import itertools - - instance = self._make_one({}, mock.Mock()) - - async def mock_gapic(*args, **kwargs): - # continuously return a single row - async def gen(): - for i in range(limit * 2): - chunk = ReadRowsResponse.CellChunk( - row_key=str(i).encode(), - family_name="family_name", - qualifier=b"qualifier", - commit_row=True, - ) - yield ReadRowsResponse(chunks=[chunk]) + from google.cloud.bigtable.data import ReadRowsQuery - return gen() + async def mock_stream(): + for i in range(emit_num): + yield i - mock_timeout_gen = itertools.repeat(5) - gen = instance._read_rows_retryable_attempt(mock_gapic, mock_timeout_gen, limit) - # should yield values up to the limit - for i in range(limit): - await gen.__anext__() - # next value should be StopAsyncIteration - with pytest.raises(StopAsyncIteration): - await gen.__anext__() + query = ReadRowsQuery(limit=start_limit) + table = mock.Mock() + table.table_name = "table_name" + table.app_profile_id = "app_profile_id" + with mock.patch.object( + _ReadRowsOperationAsync, "_read_rows_attempt" + ) as mock_attempt: + mock_attempt.return_value = mock_stream() + instance = self._make_one(query, table, 10, 10) + assert instance._remaining_count == start_limit + with pytest.raises(RuntimeError) as e: + # read emit_num rows + async for val in instance.start_operation(): + pass + assert "emit count exceeds row limit" in str(e.value) @pytest.mark.asyncio - async def test_retryable_ignore_repeated_rows(self): + async def test_aclose(self): """ - Duplicate rows should cause an invalid chunk error + should be able to close a stream safely with aclose. + Closed generators should raise StopAsyncIteration on next yield """ - from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync - from google.cloud.bigtable.data.row import Row - from google.cloud.bigtable.data.exceptions import InvalidChunk async def mock_stream(): while True: - yield Row(b"dup_key", cells=[]) - yield Row(b"dup_key", cells=[]) + yield 1 with mock.patch.object( - _ReadRowsOperationAsync, "merge_row_response_stream" - ) as mock_stream_fn: - mock_stream_fn.return_value = mock_stream() - instance = self._make_one({}, mock.AsyncMock()) - first_row = await instance.__anext__() - assert first_row.row_key == b"dup_key" - with pytest.raises(InvalidChunk) as exc: - await instance.__anext__() - assert "Last emitted row key out of order" in str(exc.value) + _ReadRowsOperationAsync, "_read_rows_attempt" + ) as mock_attempt: + instance = self._make_one(mock.Mock(), mock.Mock(), 1, 1) + wrapped_gen = mock_stream() + mock_attempt.return_value = wrapped_gen + gen = instance.start_operation() + # read one row + await gen.__anext__() + await gen.aclose() + with pytest.raises(StopAsyncIteration): + await gen.__anext__() + # try calling a second time + await gen.aclose() + # ensure close was propagated to wrapped generator + with pytest.raises(StopAsyncIteration): + await wrapped_gen.__anext__() @pytest.mark.asyncio - async def test_retryable_ignore_last_scanned_rows(self): + async def test_retryable_ignore_repeated_rows(self): """ - Last scanned rows should not be emitted + Duplicate rows should cause an invalid chunk error """ from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync - from google.cloud.bigtable.data.row import Row, _LastScannedRow - - async def mock_stream(): - while True: - yield Row(b"key1", cells=[]) - yield _LastScannedRow(b"key2_ignored") - yield Row(b"key3", cells=[]) + from google.cloud.bigtable.data.exceptions import InvalidChunk + from google.cloud.bigtable_v2.types import ReadRowsResponse - with mock.patch.object( - _ReadRowsOperationAsync, "merge_row_response_stream" - ) as mock_stream_fn: - mock_stream_fn.return_value = mock_stream() - instance = self._make_one({}, mock.AsyncMock()) - first_row = await instance.__anext__() - assert first_row.row_key == b"key1" - second_row = await instance.__anext__() - assert second_row.row_key == b"key3" + row_key = b"duplicate" - @pytest.mark.asyncio - async def test_retryable_cancel_on_close(self): - """Underlying gapic call should be cancelled when stream is closed""" - from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync - from google.cloud.bigtable.data.row import Row + async def mock_awaitable_stream(): + async def mock_stream(): + while True: + yield ReadRowsResponse( + chunks=[ + ReadRowsResponse.CellChunk(row_key=row_key, commit_row=True) + ] + ) + yield ReadRowsResponse( + chunks=[ + ReadRowsResponse.CellChunk(row_key=row_key, commit_row=True) + ] + ) - async def mock_stream(): - while True: - yield Row(b"key1", cells=[]) + return mock_stream() - with mock.patch.object( - _ReadRowsOperationAsync, "merge_row_response_stream" - ) as mock_stream_fn: - mock_stream_fn.return_value = mock_stream() - mock_gapic = mock.AsyncMock() - mock_call = await mock_gapic.read_rows() - instance = self._make_one({}, mock_gapic) - await instance.__anext__() - assert mock_call.cancel.call_count == 0 - await instance.aclose() - assert mock_call.cancel.call_count == 1 + instance = mock.Mock() + instance._last_yielded_row_key = None + stream = _ReadRowsOperationAsync.chunk_stream(instance, mock_awaitable_stream()) + await stream.__anext__() + with pytest.raises(InvalidChunk) as exc: + await stream.__anext__() + assert "row keys should be strictly increasing" in str(exc.value) class MockStream(_ReadRowsOperationAsync): @@ -427,199 +367,3 @@ async def __anext__(self): async def aclose(self): pass - - -class TestReadRowsAsyncIterator: - async def mock_stream(self, size=10): - for i in range(size): - yield i - - def _make_one(self, *args, **kwargs): - from google.cloud.bigtable.data._async._read_rows import ReadRowsAsyncIterator - - stream = MockStream(*args, **kwargs) - return ReadRowsAsyncIterator(stream) - - def test_ctor(self): - with mock.patch("time.monotonic", return_value=0): - iterator = self._make_one() - assert iterator._last_interaction_time == 0 - assert iterator._idle_timeout_task is None - assert iterator.active is True - - def test___aiter__(self): - iterator = self._make_one() - assert iterator.__aiter__() is iterator - - @pytest.mark.skipif( - sys.version_info < (3, 8), reason="mock coroutine requires python3.8 or higher" - ) - @pytest.mark.asyncio - async def test__start_idle_timer(self): - """Should start timer coroutine""" - iterator = self._make_one() - expected_timeout = 10 - with mock.patch("time.monotonic", return_value=1): - with mock.patch.object(iterator, "_idle_timeout_coroutine") as mock_coro: - await iterator._start_idle_timer(expected_timeout) - assert mock_coro.call_count == 1 - assert mock_coro.call_args[0] == (expected_timeout,) - assert iterator._last_interaction_time == 1 - assert iterator._idle_timeout_task is not None - - @pytest.mark.skipif( - sys.version_info < (3, 8), reason="mock coroutine requires python3.8 or higher" - ) - @pytest.mark.asyncio - async def test__start_idle_timer_duplicate(self): - """Multiple calls should replace task""" - iterator = self._make_one() - with mock.patch.object(iterator, "_idle_timeout_coroutine") as mock_coro: - await iterator._start_idle_timer(1) - first_task = iterator._idle_timeout_task - await iterator._start_idle_timer(2) - second_task = iterator._idle_timeout_task - assert mock_coro.call_count == 2 - - assert first_task is not None - assert first_task != second_task - # old tasks hould be cancelled - with pytest.raises(asyncio.CancelledError): - await first_task - # new task should not be cancelled - await second_task - - @pytest.mark.asyncio - async def test__idle_timeout_coroutine(self): - from google.cloud.bigtable.data.exceptions import IdleTimeout - - iterator = self._make_one() - await iterator._idle_timeout_coroutine(0.05) - await asyncio.sleep(0.1) - assert iterator.active is False - with pytest.raises(IdleTimeout): - await iterator.__anext__() - - @pytest.mark.asyncio - async def test__idle_timeout_coroutine_extensions(self): - """touching the generator should reset the idle timer""" - iterator = self._make_one(items=list(range(100))) - await iterator._start_idle_timer(0.05) - for i in range(10): - # will not expire as long as it is in use - assert iterator.active is True - await iterator.__anext__() - await asyncio.sleep(0.03) - # now let it expire - await asyncio.sleep(0.5) - assert iterator.active is False - - @pytest.mark.asyncio - async def test___anext__(self): - num_rows = 10 - iterator = self._make_one(items=list(range(num_rows))) - for i in range(num_rows): - assert await iterator.__anext__() == i - with pytest.raises(StopAsyncIteration): - await iterator.__anext__() - - @pytest.mark.asyncio - async def test___anext__with_deadline_error(self): - """ - RetryErrors mean a deadline has been hit. - Should be wrapped in a DeadlineExceeded exception - """ - from google.api_core import exceptions as core_exceptions - - items = [1, core_exceptions.RetryError("retry error", None)] - expected_timeout = 99 - iterator = self._make_one(items=items, operation_timeout=expected_timeout) - assert await iterator.__anext__() == 1 - with pytest.raises(core_exceptions.DeadlineExceeded) as exc: - await iterator.__anext__() - assert f"operation_timeout of {expected_timeout:0.1f}s exceeded" in str( - exc.value - ) - assert exc.value.__cause__ is None - - @pytest.mark.asyncio - async def test___anext__with_deadline_error_with_cause(self): - """ - Transient errors should be exposed as an error group - """ - from google.api_core import exceptions as core_exceptions - from google.cloud.bigtable.data.exceptions import RetryExceptionGroup - - items = [1, core_exceptions.RetryError("retry error", None)] - expected_timeout = 99 - errors = [RuntimeError("error1"), ValueError("error2")] - iterator = self._make_one( - items=items, operation_timeout=expected_timeout, errors=errors - ) - assert await iterator.__anext__() == 1 - with pytest.raises(core_exceptions.DeadlineExceeded) as exc: - await iterator.__anext__() - assert f"operation_timeout of {expected_timeout:0.1f}s exceeded" in str( - exc.value - ) - error_group = exc.value.__cause__ - assert isinstance(error_group, RetryExceptionGroup) - assert len(error_group.exceptions) == 2 - assert error_group.exceptions[0] is errors[0] - assert error_group.exceptions[1] is errors[1] - assert "2 failed attempts" in str(error_group) - - @pytest.mark.asyncio - async def test___anext__with_error(self): - """ - Other errors should be raised as-is - """ - from google.api_core import exceptions as core_exceptions - - items = [1, core_exceptions.InternalServerError("mock error")] - iterator = self._make_one(items=items) - assert await iterator.__anext__() == 1 - with pytest.raises(core_exceptions.InternalServerError) as exc: - await iterator.__anext__() - assert exc.value is items[1] - assert iterator.active is False - # next call should raise same error - with pytest.raises(core_exceptions.InternalServerError) as exc: - await iterator.__anext__() - - @pytest.mark.asyncio - async def test__finish_with_error(self): - iterator = self._make_one() - await iterator._start_idle_timer(10) - timeout_task = iterator._idle_timeout_task - assert await iterator.__anext__() == 0 - assert iterator.active is True - err = ZeroDivisionError("mock error") - await iterator._finish_with_error(err) - assert iterator.active is False - assert iterator._error is err - assert iterator._idle_timeout_task is None - with pytest.raises(ZeroDivisionError) as exc: - await iterator.__anext__() - assert exc.value is err - # timeout task should be cancelled - with pytest.raises(asyncio.CancelledError): - await timeout_task - - @pytest.mark.asyncio - async def test_aclose(self): - iterator = self._make_one() - await iterator._start_idle_timer(10) - timeout_task = iterator._idle_timeout_task - assert await iterator.__anext__() == 0 - assert iterator.active is True - await iterator.aclose() - assert iterator.active is False - assert isinstance(iterator._error, StopAsyncIteration) - assert iterator._idle_timeout_task is None - with pytest.raises(StopAsyncIteration) as e: - await iterator.__anext__() - assert "closed" in str(e.value) - # timeout task should be cancelled - with pytest.raises(asyncio.CancelledError): - await timeout_task diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 5a96c6c16..48d9085c6 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1233,6 +1233,7 @@ async def test_read_rows_stream(self): @pytest.mark.asyncio async def test_read_rows_query_matches_request(self, include_app_profile): from google.cloud.bigtable.data import RowRange + from google.cloud.bigtable.data.row_filters import PassAllFilter app_profile_id = "app_profile_id" if include_app_profile else None async with self._make_table(app_profile_id=app_profile_id) as table: @@ -1240,7 +1241,7 @@ async def test_read_rows_query_matches_request(self, include_app_profile): read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream([]) row_keys = [b"test_1", "test_2"] row_ranges = RowRange("1start", "2end") - filter_ = {"test": "filter"} + filter_ = PassAllFilter(True) limit = 99 query = ReadRowsQuery( row_keys=row_keys, @@ -1252,22 +1253,8 @@ async def test_read_rows_query_matches_request(self, include_app_profile): results = await table.read_rows(query, operation_timeout=3) assert len(results) == 0 call_request = read_rows.call_args_list[0][0][0] - query_dict = query._to_dict() - if include_app_profile: - assert set(call_request.keys()) == set(query_dict.keys()) | { - "table_name", - "app_profile_id", - } - else: - assert set(call_request.keys()) == set(query_dict.keys()) | { - "table_name" - } - assert call_request["rows"] == query_dict["rows"] - assert call_request["filter"] == filter_ - assert call_request["rows_limit"] == limit - assert call_request["table_name"] == table.table_name - if include_app_profile: - assert call_request["app_profile_id"] == app_profile_id + query_pb = query._to_pb(table) + assert call_request == query_pb @pytest.mark.parametrize("operation_timeout", [0.001, 0.023, 0.1]) @pytest.mark.asyncio @@ -1333,7 +1320,7 @@ async def test_read_rows_attempt_timeout( if expected_num == 0: assert retry_exc is None else: - assert type(retry_exc) == RetryExceptionGroup + assert type(retry_exc) is RetryExceptionGroup assert f"{expected_num} failed attempts" in str(retry_exc) assert len(retry_exc.exceptions) == expected_num for sub_exc in retry_exc.exceptions: @@ -1351,55 +1338,6 @@ async def test_read_rows_attempt_timeout( < 0.05 ) - @pytest.mark.asyncio - async def test_read_rows_idle_timeout(self): - from google.cloud.bigtable.data._async.client import ReadRowsAsyncIterator - from google.cloud.bigtable_v2.services.bigtable.async_client import ( - BigtableAsyncClient, - ) - from google.cloud.bigtable.data.exceptions import IdleTimeout - from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync - - chunks = [ - self._make_chunk(row_key=b"test_1"), - self._make_chunk(row_key=b"test_2"), - ] - with mock.patch.object(BigtableAsyncClient, "read_rows") as read_rows: - read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( - chunks - ) - with mock.patch.object( - ReadRowsAsyncIterator, "_start_idle_timer" - ) as start_idle_timer: - client = self._make_client() - table = client.get_table("instance", "table") - query = ReadRowsQuery() - gen = await table.read_rows_stream(query) - # should start idle timer on creation - start_idle_timer.assert_called_once() - with mock.patch.object( - _ReadRowsOperationAsync, "aclose", AsyncMock() - ) as aclose: - # start idle timer with our own value - await gen._start_idle_timer(0.1) - # should timeout after being abandoned - await gen.__anext__() - await asyncio.sleep(0.2) - # generator should be expired - assert not gen.active - assert type(gen._error) == IdleTimeout - assert gen._idle_timeout_task is None - await client.close() - with pytest.raises(IdleTimeout) as e: - await gen.__anext__() - - expected_msg = ( - "Timed out waiting for next Row to be consumed. (idle_timeout=0.1s)" - ) - assert e.value.message == expected_msg - aclose.assert_called_once() - aclose.assert_awaited() - @pytest.mark.parametrize( "exc_type", [ @@ -1422,7 +1360,7 @@ async def test_read_rows_retryable_error(self, exc_type): except core_exceptions.DeadlineExceeded as e: retry_exc = e.__cause__ root_cause = retry_exc.exceptions[0] - assert type(root_cause) == exc_type + assert type(root_cause) is exc_type assert root_cause == expected_error @pytest.mark.parametrize( @@ -1460,32 +1398,33 @@ async def test_read_rows_revise_request(self): """ from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync from google.cloud.bigtable.data.exceptions import InvalidChunk + from google.cloud.bigtable_v2.types import RowSet + return_val = RowSet() with mock.patch.object( _ReadRowsOperationAsync, "_revise_request_rowset" ) as revise_rowset: - with mock.patch.object(_ReadRowsOperationAsync, "aclose"): - revise_rowset.return_value = "modified" - async with self._make_table() as table: - read_rows = table.client._gapic_client.read_rows - read_rows.side_effect = ( - lambda *args, **kwargs: self._make_gapic_stream(chunks) - ) - row_keys = [b"test_1", b"test_2", b"test_3"] - query = ReadRowsQuery(row_keys=row_keys) - chunks = [ - self._make_chunk(row_key=b"test_1"), - core_exceptions.Aborted("mock retryable error"), - ] - try: - await table.read_rows(query) - except InvalidChunk: - revise_rowset.assert_called() - revise_call_kwargs = revise_rowset.call_args_list[0].kwargs - assert revise_call_kwargs["row_set"] == query._to_dict()["rows"] - assert revise_call_kwargs["last_seen_row_key"] == b"test_1" - read_rows_request = read_rows.call_args_list[1].args[0] - assert read_rows_request["rows"] == "modified" + revise_rowset.return_value = return_val + async with self._make_table() as table: + read_rows = table.client._gapic_client.read_rows + read_rows.side_effect = lambda *args, **kwargs: self._make_gapic_stream( + chunks + ) + row_keys = [b"test_1", b"test_2", b"test_3"] + query = ReadRowsQuery(row_keys=row_keys) + chunks = [ + self._make_chunk(row_key=b"test_1"), + core_exceptions.Aborted("mock retryable error"), + ] + try: + await table.read_rows(query) + except InvalidChunk: + revise_rowset.assert_called() + first_call_kwargs = revise_rowset.call_args_list[0].kwargs + assert first_call_kwargs["row_set"] == query._to_pb(table).rows + assert first_call_kwargs["last_seen_row_key"] == b"test_1" + revised_call = read_rows.call_args_list[1].args[0] + assert revised_call.rows == return_val @pytest.mark.asyncio async def test_read_rows_default_timeouts(self): @@ -1559,10 +1498,10 @@ async def test_read_row(self): assert kwargs["attempt_timeout"] == expected_req_timeout assert len(args) == 1 assert isinstance(args[0], ReadRowsQuery) - assert args[0]._to_dict() == { - "rows": {"row_keys": [row_key], "row_ranges": []}, - "rows_limit": 1, - } + query = args[0] + assert query.row_keys == [row_key] + assert query.row_ranges == [] + assert query.limit == 1 @pytest.mark.asyncio async def test_read_row_w_filter(self): @@ -1591,11 +1530,11 @@ async def test_read_row_w_filter(self): assert kwargs["attempt_timeout"] == expected_req_timeout assert len(args) == 1 assert isinstance(args[0], ReadRowsQuery) - assert args[0]._to_dict() == { - "rows": {"row_keys": [row_key], "row_ranges": []}, - "rows_limit": 1, - "filter": expected_filter, - } + query = args[0] + assert query.row_keys == [row_key] + assert query.row_ranges == [] + assert query.limit == 1 + assert query.filter == expected_filter @pytest.mark.asyncio async def test_read_row_no_response(self): @@ -1619,20 +1558,10 @@ async def test_read_row_no_response(self): assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout assert isinstance(args[0], ReadRowsQuery) - assert args[0]._to_dict() == { - "rows": {"row_keys": [row_key], "row_ranges": []}, - "rows_limit": 1, - } - - @pytest.mark.parametrize("input_row", [None, 5, object()]) - @pytest.mark.asyncio - async def test_read_row_w_invalid_input(self, input_row): - """Should raise error when passed None""" - async with self._make_client() as client: - table = client.get_table("instance", "table") - with pytest.raises(ValueError) as e: - await table.read_row(input_row) - assert "must be string or bytes" in e + query = args[0] + assert query.row_keys == [row_key] + assert query.row_ranges == [] + assert query.limit == 1 @pytest.mark.parametrize( "return_value,expected_result", @@ -1672,21 +1601,11 @@ async def test_row_exists(self, return_value, expected_result): ] } } - assert args[0]._to_dict() == { - "rows": {"row_keys": [row_key], "row_ranges": []}, - "rows_limit": 1, - "filter": expected_filter, - } - - @pytest.mark.parametrize("input_row", [None, 5, object()]) - @pytest.mark.asyncio - async def test_row_exists_w_invalid_input(self, input_row): - """Should raise error when passed None""" - async with self._make_client() as client: - table = client.get_table("instance", "table") - with pytest.raises(ValueError) as e: - await table.row_exists(input_row) - assert "must be string or bytes" in e + query = args[0] + assert query.row_keys == [row_key] + assert query.row_ranges == [] + assert query.limit == 1 + assert query.filter._to_dict() == expected_filter @pytest.mark.parametrize("include_app_profile", [True, False]) @pytest.mark.asyncio @@ -1739,7 +1658,7 @@ async def test_read_rows_sharded_multiple_queries(self): lambda *args, **kwargs: TestReadRows._make_gapic_stream( [ TestReadRows._make_chunk(row_key=k) - for k in args[0]["rows"]["row_keys"] + for k in args[0].rows.row_keys ] ) ) diff --git a/tests/unit/data/test__helpers.py b/tests/unit/data/test__helpers.py index e3bfd4750..08bc397c3 100644 --- a/tests/unit/data/test__helpers.py +++ b/tests/unit/data/test__helpers.py @@ -103,39 +103,56 @@ class TestConvertRetryDeadline: """ @pytest.mark.asyncio - async def test_no_error(self): - async def test_func(): + @pytest.mark.parametrize("is_async", [True, False]) + async def test_no_error(self, is_async): + def test_func(): return 1 - wrapped = _helpers._convert_retry_deadline(test_func, 0.1) - assert await wrapped() == 1 + async def test_async(): + return test_func() + + func = test_async if is_async else test_func + wrapped = _helpers._convert_retry_deadline(func, 0.1, is_async) + result = await wrapped() if is_async else wrapped() + assert result == 1 @pytest.mark.asyncio @pytest.mark.parametrize("timeout", [0.1, 2.0, 30.0]) - async def test_retry_error(self, timeout): + @pytest.mark.parametrize("is_async", [True, False]) + async def test_retry_error(self, timeout, is_async): from google.api_core.exceptions import RetryError, DeadlineExceeded - async def test_func(): + def test_func(): raise RetryError("retry error", None) - wrapped = _helpers._convert_retry_deadline(test_func, timeout) + async def test_async(): + return test_func() + + func = test_async if is_async else test_func + wrapped = _helpers._convert_retry_deadline(func, timeout, is_async=is_async) with pytest.raises(DeadlineExceeded) as e: - await wrapped() + await wrapped() if is_async else wrapped() assert e.value.__cause__ is None assert f"operation_timeout of {timeout}s exceeded" in str(e.value) @pytest.mark.asyncio - async def test_with_retry_errors(self): + @pytest.mark.parametrize("is_async", [True, False]) + async def test_with_retry_errors(self, is_async): from google.api_core.exceptions import RetryError, DeadlineExceeded timeout = 10.0 - async def test_func(): + def test_func(): raise RetryError("retry error", None) + async def test_async(): + return test_func() + + func = test_async if is_async else test_func + associated_errors = [RuntimeError("error1"), ZeroDivisionError("other")] wrapped = _helpers._convert_retry_deadline( - test_func, timeout, associated_errors + func, timeout, associated_errors, is_async ) with pytest.raises(DeadlineExceeded) as e: await wrapped() diff --git a/tests/unit/data/test__read_rows_state_machine.py b/tests/unit/data/test__read_rows_state_machine.py deleted file mode 100644 index 0d1ee6b06..000000000 --- a/tests/unit/data/test__read_rows_state_machine.py +++ /dev/null @@ -1,663 +0,0 @@ -# Licensed 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. - -import unittest -import pytest - -from google.cloud.bigtable.data.exceptions import InvalidChunk -from google.cloud.bigtable.data._read_rows_state_machine import AWAITING_NEW_ROW -from google.cloud.bigtable.data._read_rows_state_machine import AWAITING_NEW_CELL -from google.cloud.bigtable.data._read_rows_state_machine import AWAITING_CELL_VALUE - -# try/except added for compatibility with python < 3.8 -try: - from unittest import mock - from unittest.mock import AsyncMock # type: ignore -except ImportError: # pragma: NO COVER - import mock # type: ignore - from mock import AsyncMock # type: ignore # noqa F401 - -TEST_FAMILY = "family_name" -TEST_QUALIFIER = b"qualifier" -TEST_TIMESTAMP = 123456789 -TEST_LABELS = ["label1", "label2"] - - -class TestStateMachine(unittest.TestCase): - @staticmethod - def _get_target_class(): - from google.cloud.bigtable.data._read_rows_state_machine import _StateMachine - - return _StateMachine - - def _make_one(self, *args, **kwargs): - return self._get_target_class()(*args, **kwargs) - - def test_ctor(self): - from google.cloud.bigtable.data._read_rows_state_machine import _RowBuilder - - instance = self._make_one() - assert instance.last_seen_row_key is None - assert instance.current_state == AWAITING_NEW_ROW - assert instance.current_family is None - assert instance.current_qualifier is None - assert isinstance(instance.adapter, _RowBuilder) - assert instance.adapter.current_key is None - assert instance.adapter.working_cell is None - assert instance.adapter.working_value is None - assert instance.adapter.completed_cells == [] - - def test_is_terminal_state(self): - - instance = self._make_one() - assert instance.is_terminal_state() is True - instance.current_state = AWAITING_NEW_ROW - assert instance.is_terminal_state() is True - instance.current_state = AWAITING_NEW_CELL - assert instance.is_terminal_state() is False - instance.current_state = AWAITING_CELL_VALUE - assert instance.is_terminal_state() is False - - def test__reset_row(self): - instance = self._make_one() - instance.current_state = mock.Mock() - instance.current_family = "family" - instance.current_qualifier = "qualifier" - instance.adapter = mock.Mock() - instance._reset_row() - assert instance.current_state == AWAITING_NEW_ROW - assert instance.current_family is None - assert instance.current_qualifier is None - assert instance.adapter.reset.call_count == 1 - - def test_handle_last_scanned_row_wrong_state(self): - from google.cloud.bigtable.data.exceptions import InvalidChunk - - instance = self._make_one() - instance.current_state = AWAITING_NEW_CELL - with pytest.raises(InvalidChunk) as e: - instance.handle_last_scanned_row("row_key") - assert e.value.args[0] == "Last scanned row key received in invalid state" - instance.current_state = AWAITING_CELL_VALUE - with pytest.raises(InvalidChunk) as e: - instance.handle_last_scanned_row("row_key") - assert e.value.args[0] == "Last scanned row key received in invalid state" - - def test_handle_last_scanned_row_out_of_order(self): - from google.cloud.bigtable.data.exceptions import InvalidChunk - - instance = self._make_one() - instance.last_seen_row_key = b"b" - with pytest.raises(InvalidChunk) as e: - instance.handle_last_scanned_row(b"a") - assert e.value.args[0] == "Last scanned row key is out of order" - with pytest.raises(InvalidChunk) as e: - instance.handle_last_scanned_row(b"b") - assert e.value.args[0] == "Last scanned row key is out of order" - - def test_handle_last_scanned_row(self): - from google.cloud.bigtable.data.row import _LastScannedRow - - instance = self._make_one() - instance.adapter = mock.Mock() - instance.last_seen_row_key = b"a" - output_row = instance.handle_last_scanned_row(b"b") - assert instance.last_seen_row_key == b"b" - assert isinstance(output_row, _LastScannedRow) - assert output_row.row_key == b"b" - assert instance.current_state == AWAITING_NEW_ROW - assert instance.current_family is None - assert instance.current_qualifier is None - assert instance.adapter.reset.call_count == 1 - - def test__handle_complete_row(self): - from google.cloud.bigtable.data.row import Row - - instance = self._make_one() - instance.current_state = mock.Mock() - instance.current_family = "family" - instance.current_qualifier = "qualifier" - instance.adapter = mock.Mock() - instance._handle_complete_row(Row(b"row_key", {})) - assert instance.last_seen_row_key == b"row_key" - assert instance.current_state == AWAITING_NEW_ROW - assert instance.current_family is None - assert instance.current_qualifier is None - assert instance.adapter.reset.call_count == 1 - - def test__handle_reset_chunk_errors(self): - from google.cloud.bigtable.data.exceptions import InvalidChunk - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - - instance = self._make_one() - with pytest.raises(InvalidChunk) as e: - instance._handle_reset_chunk(mock.Mock()) - instance.current_state = mock.Mock() - assert e.value.args[0] == "Reset chunk received when not processing row" - with pytest.raises(InvalidChunk) as e: - instance._handle_reset_chunk( - ReadRowsResponse.CellChunk(row_key=b"row_key")._pb - ) - assert e.value.args[0] == "Reset chunk has a row key" - with pytest.raises(InvalidChunk) as e: - instance._handle_reset_chunk( - ReadRowsResponse.CellChunk(family_name="family")._pb - ) - assert e.value.args[0] == "Reset chunk has a family name" - with pytest.raises(InvalidChunk) as e: - instance._handle_reset_chunk( - ReadRowsResponse.CellChunk(qualifier=b"qualifier")._pb - ) - assert e.value.args[0] == "Reset chunk has a qualifier" - with pytest.raises(InvalidChunk) as e: - instance._handle_reset_chunk( - ReadRowsResponse.CellChunk(timestamp_micros=1)._pb - ) - assert e.value.args[0] == "Reset chunk has a timestamp" - with pytest.raises(InvalidChunk) as e: - instance._handle_reset_chunk(ReadRowsResponse.CellChunk(value=b"value")._pb) - assert e.value.args[0] == "Reset chunk has a value" - with pytest.raises(InvalidChunk) as e: - instance._handle_reset_chunk( - ReadRowsResponse.CellChunk(labels=["label"])._pb - ) - assert e.value.args[0] == "Reset chunk has labels" - - def test_handle_chunk_out_of_order(self): - from google.cloud.bigtable.data.exceptions import InvalidChunk - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - - instance = self._make_one() - instance.last_seen_row_key = b"b" - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(row_key=b"a")._pb - instance.handle_chunk(chunk) - assert "increasing" in e.value.args[0] - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(row_key=b"b")._pb - instance.handle_chunk(chunk) - assert "increasing" in e.value.args[0] - - def test_handle_chunk_reset(self): - """Should call _handle_reset_chunk when a chunk with reset_row is encountered""" - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - - instance = self._make_one() - with mock.patch.object(type(instance), "_handle_reset_chunk") as mock_reset: - chunk = ReadRowsResponse.CellChunk(reset_row=True)._pb - output = instance.handle_chunk(chunk) - assert output is None - assert mock_reset.call_count == 1 - - @pytest.mark.parametrize("state", [AWAITING_NEW_ROW, AWAITING_CELL_VALUE]) - def handle_chunk_with_commit_wrong_state(self, state): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - - instance = self._make_one() - with mock.patch.object( - type(instance.current_state), "handle_chunk" - ) as mock_state_handle: - mock_state_handle.return_value = state(mock.Mock()) - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(commit_row=True)._pb - instance.handle_chunk(mock.Mock(), chunk) - assert instance.current_state == state - assert e.value.args[0] == "Commit chunk received with in invalid state" - - def test_handle_chunk_with_commit(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data.row import Row - - instance = self._make_one() - with mock.patch.object(type(instance), "_reset_row") as mock_reset: - chunk = ReadRowsResponse.CellChunk( - row_key=b"row_key", family_name="f", qualifier=b"q", commit_row=True - )._pb - output = instance.handle_chunk(chunk) - assert isinstance(output, Row) - assert output.row_key == b"row_key" - assert output[0].family == "f" - assert output[0].qualifier == b"q" - assert instance.last_seen_row_key == b"row_key" - assert mock_reset.call_count == 1 - - def test_handle_chunk_with_commit_empty_strings(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data.row import Row - - instance = self._make_one() - with mock.patch.object(type(instance), "_reset_row") as mock_reset: - chunk = ReadRowsResponse.CellChunk( - row_key=b"row_key", family_name="", qualifier=b"", commit_row=True - )._pb - output = instance.handle_chunk(chunk) - assert isinstance(output, Row) - assert output.row_key == b"row_key" - assert output[0].family == "" - assert output[0].qualifier == b"" - assert instance.last_seen_row_key == b"row_key" - assert mock_reset.call_count == 1 - - def handle_chunk_incomplete(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - - instance = self._make_one() - chunk = ReadRowsResponse.CellChunk( - row_key=b"row_key", family_name="f", qualifier=b"q", commit_row=False - )._pb - output = instance.handle_chunk(chunk) - assert output is None - assert isinstance(instance.current_state, AWAITING_CELL_VALUE) - assert instance.current_family == "f" - assert instance.current_qualifier == b"q" - - -class TestState(unittest.TestCase): - def test_AWAITING_NEW_ROW_empty_key(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - - instance = AWAITING_NEW_ROW - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(row_key=b"")._pb - instance.handle_chunk(mock.Mock(), chunk) - assert "missing a row key" in e.value.args[0] - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk()._pb - instance.handle_chunk(mock.Mock(), chunk) - assert "missing a row key" in e.value.args[0] - - def test_AWAITING_NEW_ROW(self): - """ - AWAITING_NEW_ROW should start a RowBuilder row, then - delegate the call to AWAITING_NEW_CELL - """ - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - - instance = AWAITING_NEW_ROW - state_machine = mock.Mock() - with mock.patch.object(AWAITING_NEW_CELL, "handle_chunk") as mock_delegate: - chunk = ReadRowsResponse.CellChunk(row_key=b"row_key")._pb - instance.handle_chunk(state_machine, chunk) - assert state_machine.adapter.start_row.call_count == 1 - assert state_machine.adapter.start_row.call_args[0][0] == b"row_key" - mock_delegate.assert_called_once_with(state_machine, chunk) - - def test_AWAITING_NEW_CELL_family_without_qualifier(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - state_machine.current_qualifier = b"q" - instance = AWAITING_NEW_CELL - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(family_name="fam")._pb - instance.handle_chunk(state_machine, chunk) - assert "New family must specify qualifier" in e.value.args[0] - - def test_AWAITING_NEW_CELL_qualifier_without_family(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - instance = AWAITING_NEW_CELL - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(qualifier=b"q")._pb - instance.handle_chunk(state_machine, chunk) - assert "Family not found" in e.value.args[0] - - def test_AWAITING_NEW_CELL_no_row_state(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - instance = AWAITING_NEW_CELL - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk()._pb - instance.handle_chunk(state_machine, chunk) - assert "Missing family for new cell" in e.value.args[0] - state_machine.current_family = "fam" - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk()._pb - instance.handle_chunk(state_machine, chunk) - assert "Missing qualifier for new cell" in e.value.args[0] - - def test_AWAITING_NEW_CELL_invalid_row_key(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - instance = AWAITING_NEW_CELL - state_machine.adapter.current_key = b"abc" - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(row_key=b"123")._pb - instance.handle_chunk(state_machine, chunk) - assert "Row key changed mid row" in e.value.args[0] - - def test_AWAITING_NEW_CELL_success_no_split(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - state_machine.adapter = mock.Mock() - instance = AWAITING_NEW_CELL - row_key = b"row_key" - family = "fam" - qualifier = b"q" - labels = ["label"] - timestamp = 123 - value = b"value" - chunk = ReadRowsResponse.CellChunk( - row_key=row_key, - family_name=family, - qualifier=qualifier, - timestamp_micros=timestamp, - value=value, - labels=labels, - )._pb - state_machine.adapter.current_key = row_key - new_state = instance.handle_chunk(state_machine, chunk) - assert state_machine.adapter.start_cell.call_count == 1 - kwargs = state_machine.adapter.start_cell.call_args[1] - assert kwargs["family"] == family - assert kwargs["qualifier"] == qualifier - assert kwargs["timestamp_micros"] == timestamp - assert kwargs["labels"] == labels - assert state_machine.adapter.cell_value.call_count == 1 - assert state_machine.adapter.cell_value.call_args[0][0] == value - assert state_machine.adapter.finish_cell.call_count == 1 - assert new_state == AWAITING_NEW_CELL - - def test_AWAITING_NEW_CELL_success_with_split(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - state_machine.adapter = mock.Mock() - instance = AWAITING_NEW_CELL - row_key = b"row_key" - family = "fam" - qualifier = b"q" - labels = ["label"] - timestamp = 123 - value = b"value" - chunk = ReadRowsResponse.CellChunk( - value_size=1, - row_key=row_key, - family_name=family, - qualifier=qualifier, - timestamp_micros=timestamp, - value=value, - labels=labels, - )._pb - state_machine.adapter.current_key = row_key - new_state = instance.handle_chunk(state_machine, chunk) - assert state_machine.adapter.start_cell.call_count == 1 - kwargs = state_machine.adapter.start_cell.call_args[1] - assert kwargs["family"] == family - assert kwargs["qualifier"] == qualifier - assert kwargs["timestamp_micros"] == timestamp - assert kwargs["labels"] == labels - assert state_machine.adapter.cell_value.call_count == 1 - assert state_machine.adapter.cell_value.call_args[0][0] == value - assert state_machine.adapter.finish_cell.call_count == 0 - assert new_state == AWAITING_CELL_VALUE - - def test_AWAITING_CELL_VALUE_w_row_key(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - instance = AWAITING_CELL_VALUE - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(row_key=b"123")._pb - instance.handle_chunk(state_machine, chunk) - assert "In progress cell had a row key" in e.value.args[0] - - def test_AWAITING_CELL_VALUE_w_family(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - instance = AWAITING_CELL_VALUE - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(family_name="")._pb - instance.handle_chunk(state_machine, chunk) - assert "In progress cell had a family name" in e.value.args[0] - - def test_AWAITING_CELL_VALUE_w_qualifier(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - instance = AWAITING_CELL_VALUE - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(qualifier=b"")._pb - instance.handle_chunk(state_machine, chunk) - assert "In progress cell had a qualifier" in e.value.args[0] - - def test_AWAITING_CELL_VALUE_w_timestamp(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - instance = AWAITING_CELL_VALUE - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(timestamp_micros=123)._pb - instance.handle_chunk(state_machine, chunk) - assert "In progress cell had a timestamp" in e.value.args[0] - - def test_AWAITING_CELL_VALUE_w_labels(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - instance = AWAITING_CELL_VALUE - with pytest.raises(InvalidChunk) as e: - chunk = ReadRowsResponse.CellChunk(labels=[""])._pb - instance.handle_chunk(state_machine, chunk) - assert "In progress cell had labels" in e.value.args[0] - - def test_AWAITING_CELL_VALUE_continuation(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - state_machine.adapter = mock.Mock() - instance = AWAITING_CELL_VALUE - value = b"value" - chunk = ReadRowsResponse.CellChunk(value=value, value_size=1)._pb - new_state = instance.handle_chunk(state_machine, chunk) - assert state_machine.adapter.cell_value.call_count == 1 - assert state_machine.adapter.cell_value.call_args[0][0] == value - assert state_machine.adapter.finish_cell.call_count == 0 - assert new_state == AWAITING_CELL_VALUE - - def test_AWAITING_CELL_VALUE_final_chunk(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._async._read_rows import _StateMachine - - state_machine = _StateMachine() - state_machine.adapter = mock.Mock() - instance = AWAITING_CELL_VALUE - value = b"value" - chunk = ReadRowsResponse.CellChunk(value=value, value_size=0)._pb - new_state = instance.handle_chunk(state_machine, chunk) - assert state_machine.adapter.cell_value.call_count == 1 - assert state_machine.adapter.cell_value.call_args[0][0] == value - assert state_machine.adapter.finish_cell.call_count == 1 - assert new_state == AWAITING_NEW_CELL - - -class TestRowBuilder(unittest.TestCase): - @staticmethod - def _get_target_class(): - from google.cloud.bigtable.data._read_rows_state_machine import _RowBuilder - - return _RowBuilder - - def _make_one(self, *args, **kwargs): - return self._get_target_class()(*args, **kwargs) - - def test_ctor(self): - with mock.patch.object(self._get_target_class(), "reset") as reset: - self._make_one() - reset.assert_called_once() - row_builder = self._make_one() - self.assertIsNone(row_builder.current_key) - self.assertIsNone(row_builder.working_cell) - self.assertIsNone(row_builder.working_value) - self.assertEqual(row_builder.completed_cells, []) - - def test_start_row(self): - row_builder = self._make_one() - row_builder.start_row(b"row_key") - self.assertEqual(row_builder.current_key, b"row_key") - row_builder.start_row(b"row_key2") - self.assertEqual(row_builder.current_key, b"row_key2") - - def test_start_cell(self): - # test with no family - with self.assertRaises(InvalidChunk) as e: - self._make_one().start_cell("", TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) - self.assertEqual(str(e.exception), "Missing family for a new cell") - # test with no row - with self.assertRaises(InvalidChunk) as e: - row_builder = self._make_one() - row_builder.start_cell( - TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS - ) - self.assertEqual(str(e.exception), "start_cell called without a row") - # test with valid row - row_builder = self._make_one() - row_builder.start_row(b"row_key") - row_builder.start_cell(TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) - self.assertEqual(row_builder.working_cell.family, TEST_FAMILY) - self.assertEqual(row_builder.working_cell.qualifier, TEST_QUALIFIER) - self.assertEqual(row_builder.working_cell.timestamp_micros, TEST_TIMESTAMP) - self.assertEqual(row_builder.working_cell.labels, TEST_LABELS) - self.assertEqual(row_builder.working_value, b"") - - def test_cell_value(self): - row_builder = self._make_one() - row_builder.start_row(b"row_key") - with self.assertRaises(InvalidChunk): - # start_cell must be called before cell_value - row_builder.cell_value(b"cell_value") - row_builder.start_cell(TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) - row_builder.cell_value(b"cell_value") - self.assertEqual(row_builder.working_value, b"cell_value") - # should be able to continuously append to the working value - row_builder.cell_value(b"appended") - self.assertEqual(row_builder.working_value, b"cell_valueappended") - - def test_finish_cell(self): - row_builder = self._make_one() - row_builder.start_row(b"row_key") - row_builder.start_cell(TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) - row_builder.finish_cell() - self.assertEqual(len(row_builder.completed_cells), 1) - self.assertEqual(row_builder.completed_cells[0].family, TEST_FAMILY) - self.assertEqual(row_builder.completed_cells[0].qualifier, TEST_QUALIFIER) - self.assertEqual( - row_builder.completed_cells[0].timestamp_micros, TEST_TIMESTAMP - ) - self.assertEqual(row_builder.completed_cells[0].labels, TEST_LABELS) - self.assertEqual(row_builder.completed_cells[0].value, b"") - self.assertEqual(row_builder.working_cell, None) - self.assertEqual(row_builder.working_value, None) - # add additional cell with value - row_builder.start_cell(TEST_FAMILY, TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) - row_builder.cell_value(b"cell_value") - row_builder.cell_value(b"appended") - row_builder.finish_cell() - self.assertEqual(len(row_builder.completed_cells), 2) - self.assertEqual(row_builder.completed_cells[1].family, TEST_FAMILY) - self.assertEqual(row_builder.completed_cells[1].qualifier, TEST_QUALIFIER) - self.assertEqual( - row_builder.completed_cells[1].timestamp_micros, TEST_TIMESTAMP - ) - self.assertEqual(row_builder.completed_cells[1].labels, TEST_LABELS) - self.assertEqual(row_builder.completed_cells[1].value, b"cell_valueappended") - self.assertEqual(row_builder.working_cell, None) - self.assertEqual(row_builder.working_value, None) - - def test_finish_cell_no_cell(self): - with self.assertRaises(InvalidChunk) as e: - self._make_one().finish_cell() - self.assertEqual(str(e.exception), "finish_cell called before start_cell") - with self.assertRaises(InvalidChunk) as e: - row_builder = self._make_one() - row_builder.start_row(b"row_key") - row_builder.finish_cell() - self.assertEqual(str(e.exception), "finish_cell called before start_cell") - - def test_finish_row(self): - row_builder = self._make_one() - row_builder.start_row(b"row_key") - for i in range(3): - row_builder.start_cell(str(i), TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) - row_builder.cell_value(b"cell_value: ") - row_builder.cell_value(str(i).encode("utf-8")) - row_builder.finish_cell() - self.assertEqual(len(row_builder.completed_cells), i + 1) - output = row_builder.finish_row() - self.assertEqual(row_builder.current_key, None) - self.assertEqual(row_builder.working_cell, None) - self.assertEqual(row_builder.working_value, None) - self.assertEqual(len(row_builder.completed_cells), 0) - - self.assertEqual(output.row_key, b"row_key") - self.assertEqual(len(output), 3) - for i in range(3): - self.assertEqual(output[i].family, str(i)) - self.assertEqual(output[i].qualifier, TEST_QUALIFIER) - self.assertEqual(output[i].timestamp_micros, TEST_TIMESTAMP) - self.assertEqual(output[i].labels, TEST_LABELS) - self.assertEqual(output[i].value, b"cell_value: " + str(i).encode("utf-8")) - - def test_finish_row_no_row(self): - with self.assertRaises(InvalidChunk) as e: - self._make_one().finish_row() - self.assertEqual(str(e.exception), "No row in progress") - - def test_reset(self): - row_builder = self._make_one() - row_builder.start_row(b"row_key") - for i in range(3): - row_builder.start_cell(str(i), TEST_QUALIFIER, TEST_TIMESTAMP, TEST_LABELS) - row_builder.cell_value(b"cell_value: ") - row_builder.cell_value(str(i).encode("utf-8")) - row_builder.finish_cell() - self.assertEqual(len(row_builder.completed_cells), i + 1) - row_builder.reset() - self.assertEqual(row_builder.current_key, None) - self.assertEqual(row_builder.working_cell, None) - self.assertEqual(row_builder.working_value, None) - self.assertEqual(len(row_builder.completed_cells), 0) - - -class TestChunkHasField: - def test__chunk_has_field_empty(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._read_rows_state_machine import _chunk_has_field - - chunk = ReadRowsResponse.CellChunk()._pb - assert not _chunk_has_field(chunk, "family_name") - assert not _chunk_has_field(chunk, "qualifier") - - def test__chunk_has_field_populated_empty_strings(self): - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - from google.cloud.bigtable.data._read_rows_state_machine import _chunk_has_field - - chunk = ReadRowsResponse.CellChunk(qualifier=b"", family_name="")._pb - assert _chunk_has_field(chunk, "family_name") - assert _chunk_has_field(chunk, "qualifier") diff --git a/tests/unit/data/test_read_rows_acceptance.py b/tests/unit/data/test_read_rows_acceptance.py index 804e4e0fb..15680984b 100644 --- a/tests/unit/data/test_read_rows_acceptance.py +++ b/tests/unit/data/test_read_rows_acceptance.py @@ -24,7 +24,6 @@ from google.cloud.bigtable.data._async.client import BigtableDataClientAsync from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync -from google.cloud.bigtable.data._read_rows_state_machine import _StateMachine from google.cloud.bigtable.data.row import Row from ..v2_client.test_row_merger import ReadRowsTest, TestFile @@ -66,11 +65,15 @@ async def _scenerio_stream(): yield ReadRowsResponse(chunks=[chunk]) try: - state = _StateMachine() results = [] - async for row in _ReadRowsOperationAsync.merge_row_response_stream( - _scenerio_stream(), state - ): + instance = mock.Mock() + instance._last_yielded_row_key = None + instance._remaining_count = None + chunker = _ReadRowsOperationAsync.chunk_stream( + instance, _coro_wrapper(_scenerio_stream()) + ) + merger = _ReadRowsOperationAsync.merge_rows(chunker) + async for row in merger: for cell in row: cell_result = ReadRowsTest.Result( row_key=cell.row_key, @@ -81,8 +84,6 @@ async def _scenerio_stream(): label=cell.labels[0] if cell.labels else "", ) results.append(cell_result) - if not state.is_terminal_state(): - raise InvalidChunk("state machine has partial frame after reading") except InvalidChunk: results.append(ReadRowsTest.Result(error=True)) for expected, actual in zip_longest(test_case.results, results): @@ -148,12 +149,15 @@ async def test_out_of_order_rows(): async def _row_stream(): yield ReadRowsResponse(last_scanned_row_key=b"a") - state = _StateMachine() - state.last_seen_row_key = b"a" + instance = mock.Mock() + instance._remaining_count = None + instance._last_yielded_row_key = b"b" + chunker = _ReadRowsOperationAsync.chunk_stream( + instance, _coro_wrapper(_row_stream()) + ) + merger = _ReadRowsOperationAsync.merge_rows(chunker) with pytest.raises(InvalidChunk): - async for _ in _ReadRowsOperationAsync.merge_row_response_stream( - _row_stream(), state - ): + async for _ in merger: pass @@ -304,14 +308,22 @@ async def test_mid_cell_labels_change(): ) +async def _coro_wrapper(stream): + return stream + + async def _process_chunks(*chunks): async def _row_stream(): yield ReadRowsResponse(chunks=chunks) - state = _StateMachine() + instance = mock.Mock() + instance._remaining_count = None + instance._last_yielded_row_key = None + chunker = _ReadRowsOperationAsync.chunk_stream( + instance, _coro_wrapper(_row_stream()) + ) + merger = _ReadRowsOperationAsync.merge_rows(chunker) results = [] - async for row in _ReadRowsOperationAsync.merge_row_response_stream( - _row_stream(), state - ): + async for row in merger: results.append(row) return results diff --git a/tests/unit/data/test_read_rows_query.py b/tests/unit/data/test_read_rows_query.py index 1e4e27d36..ba3b0468b 100644 --- a/tests/unit/data/test_read_rows_query.py +++ b/tests/unit/data/test_read_rows_query.py @@ -32,10 +32,6 @@ def _make_one(self, *args, **kwargs): def test_ctor_start_end(self): row_range = self._make_one("test_row", "test_row2") - assert row_range._start.key == "test_row".encode() - assert row_range._end.key == "test_row2".encode() - assert row_range._start.is_inclusive is True - assert row_range._end.is_inclusive is False assert row_range.start_key == "test_row".encode() assert row_range.end_key == "test_row2".encode() assert row_range.start_is_inclusive is True @@ -60,8 +56,6 @@ def test_ctor_empty_strings(self): empty strings should be treated as None """ row_range = self._make_one("", "") - assert row_range._start is None - assert row_range._end is None assert row_range.start_key is None assert row_range.end_key is None assert row_range.start_is_inclusive is True @@ -79,20 +73,6 @@ def test_ctor_defaults(self): assert row_range.start_key is None assert row_range.end_key is None - def test_ctor_flags_only(self): - with pytest.raises(ValueError) as exc: - self._make_one(start_is_inclusive=True, end_is_inclusive=True) - assert str(exc.value) == "start_is_inclusive must be set with start_key" - with pytest.raises(ValueError) as exc: - self._make_one(start_is_inclusive=False, end_is_inclusive=False) - assert str(exc.value) == "start_is_inclusive must be set with start_key" - with pytest.raises(ValueError) as exc: - self._make_one(start_is_inclusive=False) - assert str(exc.value) == "start_is_inclusive must be set with start_key" - with pytest.raises(ValueError) as exc: - self._make_one(end_is_inclusive=True) - assert str(exc.value) == "end_is_inclusive must be set with end_key" - def test_ctor_invalid_keys(self): # test with invalid keys with pytest.raises(ValueError) as exc: @@ -105,138 +85,6 @@ def test_ctor_invalid_keys(self): self._make_one("2", "1") assert str(exc.value) == "start_key must be less than or equal to end_key" - def test__to_dict_defaults(self): - row_range = self._make_one("test_row", "test_row2") - expected = { - "start_key_closed": b"test_row", - "end_key_open": b"test_row2", - } - assert row_range._to_dict() == expected - - def test__to_dict_inclusive_flags(self): - row_range = self._make_one("test_row", "test_row2", False, True) - expected = { - "start_key_open": b"test_row", - "end_key_closed": b"test_row2", - } - assert row_range._to_dict() == expected - - @pytest.mark.parametrize( - "input_dict,expected_start,expected_end,start_is_inclusive,end_is_inclusive", - [ - ( - {"start_key_closed": "test_row", "end_key_open": "test_row2"}, - b"test_row", - b"test_row2", - True, - False, - ), - ( - {"start_key_closed": b"test_row", "end_key_open": b"test_row2"}, - b"test_row", - b"test_row2", - True, - False, - ), - ( - {"start_key_open": "test_row", "end_key_closed": "test_row2"}, - b"test_row", - b"test_row2", - False, - True, - ), - ({"start_key_open": b"a"}, b"a", None, False, None), - ({"end_key_closed": b"b"}, None, b"b", None, True), - ({"start_key_closed": "a"}, b"a", None, True, None), - ({"end_key_open": b"b"}, None, b"b", None, False), - ({}, None, None, None, None), - ], - ) - def test__from_dict( - self, - input_dict, - expected_start, - expected_end, - start_is_inclusive, - end_is_inclusive, - ): - from google.cloud.bigtable.data.read_rows_query import RowRange - - row_range = RowRange._from_dict(input_dict) - assert row_range._to_dict().keys() == input_dict.keys() - found_start = row_range._start - found_end = row_range._end - if expected_start is None: - assert found_start is None - assert start_is_inclusive is None - else: - assert found_start.key == expected_start - assert found_start.is_inclusive == start_is_inclusive - if expected_end is None: - assert found_end is None - assert end_is_inclusive is None - else: - assert found_end.key == expected_end - assert found_end.is_inclusive == end_is_inclusive - - @pytest.mark.parametrize( - "dict_repr", - [ - {"start_key_closed": "test_row", "end_key_open": "test_row2"}, - {"start_key_closed": b"test_row", "end_key_open": b"test_row2"}, - {"start_key_open": "test_row", "end_key_closed": "test_row2"}, - {"start_key_open": b"a"}, - {"end_key_closed": b"b"}, - {"start_key_closed": "a"}, - {"end_key_open": b"b"}, - {}, - ], - ) - def test__from_points(self, dict_repr): - from google.cloud.bigtable.data.read_rows_query import RowRange - - row_range_from_dict = RowRange._from_dict(dict_repr) - row_range_from_points = RowRange._from_points( - row_range_from_dict._start, row_range_from_dict._end - ) - assert row_range_from_points._to_dict() == row_range_from_dict._to_dict() - - @pytest.mark.parametrize( - "first_dict,second_dict,should_match", - [ - ( - {"start_key_closed": "a", "end_key_open": "b"}, - {"start_key_closed": "a", "end_key_open": "b"}, - True, - ), - ( - {"start_key_closed": "a", "end_key_open": "b"}, - {"start_key_closed": "a", "end_key_open": "c"}, - False, - ), - ( - {"start_key_closed": "a", "end_key_open": "b"}, - {"start_key_closed": "a", "end_key_closed": "b"}, - False, - ), - ( - {"start_key_closed": b"a", "end_key_open": b"b"}, - {"start_key_closed": "a", "end_key_open": "b"}, - True, - ), - ({}, {}, True), - ({"start_key_closed": "a"}, {}, False), - ({"start_key_closed": "a"}, {"start_key_closed": "a"}, True), - ({"start_key_closed": "a"}, {"start_key_open": "a"}, False), - ], - ) - def test___hash__(self, first_dict, second_dict, should_match): - from google.cloud.bigtable.data.read_rows_query import RowRange - - row_range1 = RowRange._from_dict(first_dict) - row_range2 = RowRange._from_dict(second_dict) - assert (hash(row_range1) == hash(row_range2)) == should_match - @pytest.mark.parametrize( "dict_repr,expected", [ @@ -352,8 +200,8 @@ def _make_one(self, *args, **kwargs): def test_ctor_defaults(self): query = self._make_one() - assert query.row_keys == set() - assert query.row_ranges == set() + assert query.row_keys == list() + assert query.row_ranges == list() assert query.filter is None assert query.limit is None @@ -396,9 +244,6 @@ def test_set_filter(self): assert query.filter is None query.filter = RowFilterChain() assert query.filter == RowFilterChain() - with pytest.raises(ValueError) as exc: - query.filter = 1 - assert str(exc.value) == "row_filter must be a RowFilter or dict" def test_set_limit(self): query = self._make_one() @@ -408,7 +253,7 @@ def test_set_limit(self): query.limit = 9 assert query.limit == 9 query.limit = 0 - assert query.limit == 0 + assert query.limit is None with pytest.raises(ValueError) as exc: query.limit = -1 assert str(exc.value) == "limit must be >= 0" @@ -418,7 +263,7 @@ def test_set_limit(self): def test_add_key_str(self): query = self._make_one() - assert query.row_keys == set() + assert query.row_keys == list() input_str = "test_row" query.add_key(input_str) assert len(query.row_keys) == 1 @@ -431,7 +276,7 @@ def test_add_key_str(self): def test_add_key_bytes(self): query = self._make_one() - assert query.row_keys == set() + assert query.row_keys == list() input_bytes = b"test_row" query.add_key(input_bytes) assert len(query.row_keys) == 1 @@ -444,7 +289,7 @@ def test_add_key_bytes(self): def test_add_rows_batch(self): query = self._make_one() - assert query.row_keys == set() + assert query.row_keys == list() input_batch = ["test_row", b"test_row2", "test_row3"] for k in input_batch: query.add_key(k) @@ -471,24 +316,11 @@ def test_add_key_invalid(self): query.add_key(["s"]) assert str(exc.value) == "row_key must be string or bytes" - def test_duplicate_rows(self): - # should only hold one of each input key - key_1 = b"test_row" - key_2 = b"test_row2" - query = self._make_one(row_keys=[key_1, key_1, key_2]) - assert len(query.row_keys) == 2 - assert key_1 in query.row_keys - assert key_2 in query.row_keys - key_3 = "test_row3" - for i in range(10): - query.add_key(key_3) - assert len(query.row_keys) == 3 - def test_add_range(self): from google.cloud.bigtable.data.read_rows_query import RowRange query = self._make_one() - assert query.row_ranges == set() + assert query.row_ranges == list() input_range = RowRange(start_key=b"test_row") query.add_range(input_range) assert len(query.row_ranges) == 1 @@ -498,83 +330,6 @@ def test_add_range(self): assert len(query.row_ranges) == 2 assert input_range in query.row_ranges assert input_range2 in query.row_ranges - query.add_range(input_range2) - assert len(query.row_ranges) == 2 - - def test_add_range_dict(self): - from google.cloud.bigtable.data.read_rows_query import RowRange - - query = self._make_one() - assert query.row_ranges == set() - input_range = {"start_key_closed": b"test_row"} - query.add_range(input_range) - assert len(query.row_ranges) == 1 - range_obj = RowRange._from_dict(input_range) - assert range_obj in query.row_ranges - - def test_to_dict_rows_default(self): - # dictionary should be in rowset proto format - from google.cloud.bigtable_v2.types.bigtable import ReadRowsRequest - - query = self._make_one() - output = query._to_dict() - assert isinstance(output, dict) - assert len(output.keys()) == 1 - expected = {"rows": {"row_keys": [], "row_ranges": []}} - assert output == expected - - request_proto = ReadRowsRequest(**output) - assert request_proto.rows.row_keys == [] - assert request_proto.rows.row_ranges == [] - assert not request_proto.filter - assert request_proto.rows_limit == 0 - - def test_to_dict_rows_populated(self): - # dictionary should be in rowset proto format - from google.cloud.bigtable_v2.types.bigtable import ReadRowsRequest - from google.cloud.bigtable.data.row_filters import PassAllFilter - from google.cloud.bigtable.data.read_rows_query import RowRange - - row_filter = PassAllFilter(False) - query = self._make_one(limit=100, row_filter=row_filter) - query.add_range(RowRange("test_row", "test_row2")) - query.add_range(RowRange("test_row3")) - query.add_range(RowRange(start_key=None, end_key="test_row5")) - query.add_range(RowRange(b"test_row6", b"test_row7", False, True)) - query.add_range({}) - query.add_key("test_row") - query.add_key(b"test_row2") - query.add_key("test_row3") - query.add_key(b"test_row3") - query.add_key(b"test_row4") - output = query._to_dict() - assert isinstance(output, dict) - request_proto = ReadRowsRequest(**output) - rowset_proto = request_proto.rows - # check rows - assert len(rowset_proto.row_keys) == 4 - assert rowset_proto.row_keys[0] == b"test_row" - assert rowset_proto.row_keys[1] == b"test_row2" - assert rowset_proto.row_keys[2] == b"test_row3" - assert rowset_proto.row_keys[3] == b"test_row4" - # check ranges - assert len(rowset_proto.row_ranges) == 5 - assert { - "start_key_closed": b"test_row", - "end_key_open": b"test_row2", - } in output["rows"]["row_ranges"] - assert {"start_key_closed": b"test_row3"} in output["rows"]["row_ranges"] - assert {"end_key_open": b"test_row5"} in output["rows"]["row_ranges"] - assert { - "start_key_open": b"test_row6", - "end_key_closed": b"test_row7", - } in output["rows"]["row_ranges"] - assert {} in output["rows"]["row_ranges"] - # check limit - assert request_proto.rows_limit == 100 - # check filter - filter_proto = request_proto.filter - assert filter_proto == row_filter._to_pb() def _parse_query_string(self, query_string): from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery, RowRange @@ -781,7 +536,6 @@ def test_shard_limit_exception(self): ((), ("a",), False), (("a",), (), False), (("a",), ("a",), True), - (("a",), (["a", b"a"],), True), # duplicate keys ((["a"],), (["a", "b"],), False), ((["a", "b"],), (["a", "b"],), True), ((["a", b"b"],), ([b"a", "b"],), True), @@ -792,7 +546,6 @@ def test_shard_limit_exception(self): (("a", ["b"]), ("a", ["b", "c"]), False), (("a", ["b", "c"]), ("a", [b"b", "c"]), True), (("a", ["b", "c"], 1), ("a", ["b", b"c"], 1), True), - (("a", ["b"], 1), ("a", ["b", b"b", "b"], 1), True), # duplicate ranges (("a", ["b"], 1), ("a", ["b"], 2), False), (("a", ["b"], 1, {"a": "b"}), ("a", ["b"], 1, {"a": "b"}), True), (("a", ["b"], 1, {"a": "b"}), ("a", ["b"], 1), False), @@ -833,4 +586,4 @@ def test___repr__(self): def test_empty_row_set(self): """Empty strings should be treated as keys inputs""" query = self._make_one(row_keys="") - assert query.row_keys == {b""} + assert query.row_keys == [b""] diff --git a/tests/unit/data/test_row.py b/tests/unit/data/test_row.py index df2fc72c0..10b5bdb23 100644 --- a/tests/unit/data/test_row.py +++ b/tests/unit/data/test_row.py @@ -481,20 +481,6 @@ def test_get_column_components(self): row_response._get_column_components(), [(TEST_FAMILY_ID, TEST_QUALIFIER)] ) - def test_index_of(self): - # given a cell, should find index in underlying list - cell_list = [self._make_cell(value=str(i).encode()) for i in range(10)] - sorted(cell_list) - row_response = self._make_one(TEST_ROW_KEY, cell_list) - - self.assertEqual(row_response.index(cell_list[0]), 0) - self.assertEqual(row_response.index(cell_list[5]), 5) - self.assertEqual(row_response.index(cell_list[9]), 9) - with self.assertRaises(ValueError): - row_response.index(self._make_cell()) - with self.assertRaises(ValueError): - row_response.index(None) - class TestCell(unittest.TestCase): @staticmethod From 0b3606fd14e706f53fb71a657b2d7d512e0dbc74 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 17 Aug 2023 12:00:21 -0600 Subject: [PATCH 20/56] chore: add user agent suffix (#842) --- google/cloud/bigtable/data/_async/client.py | 9 ++++++++- tests/unit/data/_async/test_client.py | 2 +- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index ff323b9bc..1ca9e0f7b 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -122,7 +122,7 @@ def __init__( BigtableClientMeta._transport_registry[transport_str] = transport # set up client info headers for veneer library client_info = DEFAULT_CLIENT_INFO - client_info.client_library_version = client_info.gapic_version + client_info.client_library_version = self._client_version() # parse client options if type(client_options) is dict: client_options = client_options_lib.from_dict(client_options) @@ -163,6 +163,13 @@ def __init__( stacklevel=2, ) + @staticmethod + def _client_version() -> str: + """ + Helper function to return the client version string for this client + """ + return f"{google.cloud.bigtable.__version__}-data-async" + def _start_background_channel_refresh(self) -> None: """ Starts a background task to ping and warm each channel in the pool diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 48d9085c6..c2c4b0615 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -39,7 +39,7 @@ from mock import AsyncMock # type: ignore VENEER_HEADER_REGEX = re.compile( - r"gapic\/[0-9]+\.[\w.-]+ gax\/[0-9]+\.[\w.-]+ gccl\/[0-9]+\.[\w.-]+ gl-python\/[0-9]+\.[\w.-]+ grpc\/[0-9]+\.[\w.-]+" + r"gapic\/[0-9]+\.[\w.-]+ gax\/[0-9]+\.[\w.-]+ gccl\/[0-9]+\.[\w.-]+-data-async gl-python\/[0-9]+\.[\w.-]+ grpc\/[0-9]+\.[\w.-]+" ) From b6d232ab532e2a6250483a419cdda701ed7f29de Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 17 Aug 2023 12:56:01 -0600 Subject: [PATCH 21/56] feat: optimize retries (#854) --- .../cloud/bigtable/data/_async/_read_rows.py | 71 +++++++++--------- google/cloud/bigtable/data/_helpers.py | 3 + python-api-core | 2 +- setup.py | 2 +- testing/constraints-3.7.txt | 2 +- tests/unit/data/_async/test__read_rows.py | 74 ++++++++++++------- 6 files changed, 90 insertions(+), 64 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 6edb72858..20b5618ea 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -31,7 +31,7 @@ from google.cloud.bigtable.data._helpers import _make_metadata from google.api_core import retry_async as retries -from google.api_core.retry_streaming_async import AsyncRetryableGenerator +from google.api_core.retry_streaming_async import retry_target_stream from google.api_core.retry import exponential_sleep_generator from google.api_core import exceptions as core_exceptions @@ -100,35 +100,17 @@ def __init__( self._last_yielded_row_key: bytes | None = None self._remaining_count: int | None = self.request.rows_limit or None - async def start_operation(self) -> AsyncGenerator[Row, None]: + def start_operation(self) -> AsyncGenerator[Row, None]: """ Start the read_rows operation, retrying on retryable errors. """ - transient_errors = [] - - def on_error_fn(exc): - if self._predicate(exc): - transient_errors.append(exc) - - retry_gen = AsyncRetryableGenerator( + return retry_target_stream( self._read_rows_attempt, self._predicate, exponential_sleep_generator(0.01, 60, multiplier=2), self.operation_timeout, - on_error_fn, + exception_factory=self._build_exception, ) - try: - async for row in retry_gen: - yield row - if self._remaining_count is not None: - self._remaining_count -= 1 - if self._remaining_count < 0: - raise RuntimeError("emit count exceeds row limit") - except core_exceptions.RetryError: - self._raise_retry_error(transient_errors) - except GeneratorExit: - # propagate close to wrapped generator - await retry_gen.aclose() def _read_rows_attempt(self) -> AsyncGenerator[Row, None]: """ @@ -202,6 +184,10 @@ async def chunk_stream( elif c.commit_row: # update row state after each commit self._last_yielded_row_key = current_key + if self._remaining_count is not None: + self._remaining_count -= 1 + if self._remaining_count < 0: + raise InvalidChunk("emit count exceeds row limit") current_key = None @staticmethod @@ -354,19 +340,34 @@ def _revise_request_rowset( raise _RowSetComplete() return RowSetPB(row_keys=adjusted_keys, row_ranges=adjusted_ranges) - def _raise_retry_error(self, transient_errors: list[Exception]) -> None: + @staticmethod + def _build_exception( + exc_list: list[Exception], is_timeout: bool, timeout_val: float + ) -> tuple[Exception, Exception | None]: """ - If the retryable deadline is hit, wrap the raised exception - in a RetryExceptionGroup + Build retry error based on exceptions encountered during operation + + Args: + - exc_list: list of exceptions encountered during operation + - is_timeout: whether the operation failed due to timeout + - timeout_val: the operation timeout value in seconds, for constructing + the error message + Returns: + - tuple of the exception to raise, and a cause exception if applicable """ - timeout_value = self.operation_timeout - timeout_str = f" of {timeout_value:.1f}s" if timeout_value is not None else "" - error_str = f"operation_timeout{timeout_str} exceeded" - new_exc = core_exceptions.DeadlineExceeded( - error_str, + if is_timeout: + # if failed due to timeout, raise deadline exceeded as primary exception + source_exc: Exception = core_exceptions.DeadlineExceeded( + f"operation_timeout of {timeout_val} exceeded" + ) + elif exc_list: + # otherwise, raise non-retryable error as primary exception + source_exc = exc_list.pop() + else: + source_exc = RuntimeError("failed with unspecified exception") + # use the retry exception group as the cause of the exception + cause_exc: Exception | None = ( + RetryExceptionGroup(exc_list) if exc_list else None ) - source_exc = None - if transient_errors: - source_exc = RetryExceptionGroup(transient_errors) - new_exc.__cause__ = source_exc - raise new_exc from source_exc + source_exc.__cause__ = cause_exc + return source_exc, cause_exc diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index b13b670d4..1f8a63d21 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -62,6 +62,9 @@ def _attempt_timeout_generator( yield max(0, min(per_request_timeout, deadline - time.monotonic())) +# TODO:replace this function with an exception_factory passed into the retry when +# feature is merged: +# https://github.com/googleapis/python-bigtable/blob/ea5b4f923e42516729c57113ddbe28096841b952/google/cloud/bigtable/data/_async/_read_rows.py#L130 def _convert_retry_deadline( func: Callable[..., Any], timeout_value: float | None = None, diff --git a/python-api-core b/python-api-core index a526d6593..a8cfa66b8 160000 --- a/python-api-core +++ b/python-api-core @@ -1 +1 @@ -Subproject commit a526d659320939cd7f47ee775b250e8a3e3ab16b +Subproject commit a8cfa66b8d6001da56823c6488b5da4957e5702b diff --git a/setup.py b/setup.py index e05b37c79..e5efc9937 100644 --- a/setup.py +++ b/setup.py @@ -37,7 +37,7 @@ # 'Development Status :: 5 - Production/Stable' release_status = "Development Status :: 5 - Production/Stable" dependencies = [ - "google-api-core[grpc] == 2.12.0.dev0", # TODO: change to >= after streaming retries is merged + "google-api-core[grpc] == 2.12.0.dev1", # TODO: change to >= after streaming retries is merged "google-cloud-core >= 1.4.1, <3.0.0dev", "grpc-google-iam-v1 >= 0.12.4, <1.0.0dev", "proto-plus >= 1.22.0, <2.0.0dev", diff --git a/testing/constraints-3.7.txt b/testing/constraints-3.7.txt index 92b616563..9f23121d1 100644 --- a/testing/constraints-3.7.txt +++ b/testing/constraints-3.7.txt @@ -5,7 +5,7 @@ # # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 -google-api-core==2.12.0.dev0 +google-api-core==2.12.0.dev1 google-cloud-core==2.3.2 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 diff --git a/tests/unit/data/_async/test__read_rows.py b/tests/unit/data/_async/test__read_rows.py index 200defbbf..4e7797c6d 100644 --- a/tests/unit/data/_async/test__read_rows.py +++ b/tests/unit/data/_async/test__read_rows.py @@ -226,24 +226,34 @@ async def test_revise_limit(self, start_limit, emit_num, expected_limit): should be raised (tested in test_revise_limit_over_limit) """ from google.cloud.bigtable.data import ReadRowsQuery + from google.cloud.bigtable_v2.types import ReadRowsResponse - async def mock_stream(): - for i in range(emit_num): - yield i + async def awaitable_stream(): + async def mock_stream(): + for i in range(emit_num): + yield ReadRowsResponse( + chunks=[ + ReadRowsResponse.CellChunk( + row_key=str(i).encode(), + family_name="b", + qualifier=b"c", + value=b"d", + commit_row=True, + ) + ] + ) + + return mock_stream() query = ReadRowsQuery(limit=start_limit) table = mock.Mock() table.table_name = "table_name" table.app_profile_id = "app_profile_id" - with mock.patch.object( - _ReadRowsOperationAsync, "_read_rows_attempt" - ) as mock_attempt: - mock_attempt.return_value = mock_stream() - instance = self._make_one(query, table, 10, 10) - assert instance._remaining_count == start_limit - # read emit_num rows - async for val in instance.start_operation(): - pass + instance = self._make_one(query, table, 10, 10) + assert instance._remaining_count == start_limit + # read emit_num rows + async for val in instance.chunk_stream(awaitable_stream()): + pass assert instance._remaining_count == expected_limit @pytest.mark.parametrize("start_limit,emit_num", [(5, 10), (3, 9), (1, 10)]) @@ -254,26 +264,37 @@ async def test_revise_limit_over_limit(self, start_limit, emit_num): (unless start_num == 0, which represents unlimited) """ from google.cloud.bigtable.data import ReadRowsQuery + from google.cloud.bigtable_v2.types import ReadRowsResponse + from google.cloud.bigtable.data.exceptions import InvalidChunk - async def mock_stream(): - for i in range(emit_num): - yield i + async def awaitable_stream(): + async def mock_stream(): + for i in range(emit_num): + yield ReadRowsResponse( + chunks=[ + ReadRowsResponse.CellChunk( + row_key=str(i).encode(), + family_name="b", + qualifier=b"c", + value=b"d", + commit_row=True, + ) + ] + ) + + return mock_stream() query = ReadRowsQuery(limit=start_limit) table = mock.Mock() table.table_name = "table_name" table.app_profile_id = "app_profile_id" - with mock.patch.object( - _ReadRowsOperationAsync, "_read_rows_attempt" - ) as mock_attempt: - mock_attempt.return_value = mock_stream() - instance = self._make_one(query, table, 10, 10) - assert instance._remaining_count == start_limit - with pytest.raises(RuntimeError) as e: - # read emit_num rows - async for val in instance.start_operation(): - pass - assert "emit count exceeds row limit" in str(e.value) + instance = self._make_one(query, table, 10, 10) + assert instance._remaining_count == start_limit + with pytest.raises(InvalidChunk) as e: + # read emit_num rows + async for val in instance.chunk_stream(awaitable_stream()): + pass + assert "emit count exceeds row limit" in str(e.value) @pytest.mark.asyncio async def test_aclose(self): @@ -333,6 +354,7 @@ async def mock_stream(): instance = mock.Mock() instance._last_yielded_row_key = None + instance._remaining_count = None stream = _ReadRowsOperationAsync.chunk_stream(instance, mock_awaitable_stream()) await stream.__anext__() with pytest.raises(InvalidChunk) as exc: From 8708a2568d71e730372a2f717143a6c0910823fa Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 18 Aug 2023 13:54:49 -0600 Subject: [PATCH 22/56] feat: add test proxy (#836) --- .kokoro/presubmit/conformance.cfg | 6 + .../bigtable/data/_async/_mutate_rows.py | 8 + google/cloud/bigtable/data/_async/client.py | 34 +- noxfile.py | 21 + test_proxy/README.md | 60 +++ test_proxy/handlers/client_handler_data.py | 214 +++++++++ test_proxy/handlers/client_handler_legacy.py | 132 ++++++ test_proxy/handlers/grpc_handler.py | 148 ++++++ test_proxy/noxfile.py | 80 ++++ test_proxy/protos/bigtable_pb2.py | 145 ++++++ test_proxy/protos/bigtable_pb2_grpc.py | 363 +++++++++++++++ test_proxy/protos/data_pb2.py | 68 +++ test_proxy/protos/data_pb2_grpc.py | 4 + test_proxy/protos/request_stats_pb2.py | 33 ++ test_proxy/protos/request_stats_pb2_grpc.py | 4 + test_proxy/protos/test_proxy_pb2.py | 71 +++ test_proxy/protos/test_proxy_pb2_grpc.py | 433 ++++++++++++++++++ test_proxy/run_tests.sh | 47 ++ test_proxy/test_proxy.py | 193 ++++++++ 19 files changed, 2057 insertions(+), 7 deletions(-) create mode 100644 .kokoro/presubmit/conformance.cfg create mode 100644 test_proxy/README.md create mode 100644 test_proxy/handlers/client_handler_data.py create mode 100644 test_proxy/handlers/client_handler_legacy.py create mode 100644 test_proxy/handlers/grpc_handler.py create mode 100644 test_proxy/noxfile.py create mode 100644 test_proxy/protos/bigtable_pb2.py create mode 100644 test_proxy/protos/bigtable_pb2_grpc.py create mode 100644 test_proxy/protos/data_pb2.py create mode 100644 test_proxy/protos/data_pb2_grpc.py create mode 100644 test_proxy/protos/request_stats_pb2.py create mode 100644 test_proxy/protos/request_stats_pb2_grpc.py create mode 100644 test_proxy/protos/test_proxy_pb2.py create mode 100644 test_proxy/protos/test_proxy_pb2_grpc.py create mode 100755 test_proxy/run_tests.sh create mode 100644 test_proxy/test_proxy.py diff --git a/.kokoro/presubmit/conformance.cfg b/.kokoro/presubmit/conformance.cfg new file mode 100644 index 000000000..4f44e8a78 --- /dev/null +++ b/.kokoro/presubmit/conformance.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "NOX_SESSION" + value: "conformance" +} diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index d4f5bc215..baae205d9 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -15,6 +15,7 @@ from __future__ import annotations from typing import TYPE_CHECKING +import asyncio import functools from google.api_core import exceptions as core_exceptions @@ -183,6 +184,13 @@ async def _run_attempt(self): self._handle_entry_error(orig_idx, entry_error) # remove processed entry from active list del active_request_indices[result.index] + except asyncio.CancelledError: + # when retry wrapper timeout expires, the operation is cancelled + # make sure incomplete indices are tracked, + # but don't record exception (it will be raised by wrapper) + # TODO: remove asyncio.wait_for in retry wrapper. Let grpc call handle expiration + self.remaining_indices.extend(active_request_indices.values()) + raise except Exception as exc: # add this exception to list for each mutation that wasn't # already handled, and update remaining_indices if mutation is retryable diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 1ca9e0f7b..8524cd9aa 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -30,6 +30,7 @@ import warnings import sys import random +import os from collections import namedtuple @@ -38,10 +39,12 @@ from google.cloud.bigtable_v2.services.bigtable.async_client import DEFAULT_CLIENT_INFO from google.cloud.bigtable_v2.services.bigtable.transports.pooled_grpc_asyncio import ( PooledBigtableGrpcAsyncIOTransport, + PooledChannel, ) from google.cloud.bigtable_v2.types.bigtable import PingAndWarmRequest from google.cloud.client import ClientWithProject from google.api_core.exceptions import GoogleAPICallError +from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore from google.api_core import retry_async as retries from google.api_core import exceptions as core_exceptions from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync @@ -150,18 +153,35 @@ def __init__( # keep track of table objects associated with each instance # only remove instance from _active_instances when all associated tables remove it self._instance_owners: dict[_WarmedInstanceKey, Set[int]] = {} - # attempt to start background tasks self._channel_init_time = time.monotonic() self._channel_refresh_tasks: list[asyncio.Task[None]] = [] - try: - self._start_background_channel_refresh() - except RuntimeError: + self._emulator_host = os.getenv(BIGTABLE_EMULATOR) + if self._emulator_host is not None: + # connect to an emulator host warnings.warn( - f"{self.__class__.__name__} should be started in an " - "asyncio event loop. Channel refresh will not be started", + "Connecting to Bigtable emulator at {}".format(self._emulator_host), RuntimeWarning, stacklevel=2, ) + self.transport._grpc_channel = PooledChannel( + pool_size=pool_size, + host=self._emulator_host, + insecure=True, + ) + # refresh cached stubs to use emulator pool + self.transport._stubs = {} + self.transport._prep_wrapped_messages(client_info) + else: + # attempt to start background channel refresh tasks + try: + self._start_background_channel_refresh() + except RuntimeError: + warnings.warn( + f"{self.__class__.__name__} should be started in an " + "asyncio event loop. Channel refresh will not be started", + RuntimeWarning, + stacklevel=2, + ) @staticmethod def _client_version() -> str: @@ -176,7 +196,7 @@ def _start_background_channel_refresh(self) -> None: Raises: - RuntimeError if not called in an asyncio event loop """ - if not self._channel_refresh_tasks: + if not self._channel_refresh_tasks and not self._emulator_host: # raise RuntimeError if there is no event loop asyncio.get_running_loop() for channel_idx in range(self.transport.pool_size): diff --git a/noxfile.py b/noxfile.py index 16447778e..4b57e617f 100644 --- a/noxfile.py +++ b/noxfile.py @@ -278,6 +278,27 @@ def system_emulated(session): os.killpg(os.getpgid(p.pid), signal.SIGKILL) +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) +def conformance(session): + """ + Run the set of shared bigtable conformance tests + """ + TEST_REPO_URL = "https://github.com/googleapis/cloud-bigtable-clients-test.git" + CLONE_REPO_DIR = "cloud-bigtable-clients-test" + # install dependencies + constraints_path = str( + CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt" + ) + install_unittest_dependencies(session, "-c", constraints_path) + with session.chdir("test_proxy"): + # download the conformance test suite + clone_dir = os.path.join(CURRENT_DIRECTORY, CLONE_REPO_DIR) + if not os.path.exists(clone_dir): + print("downloading copy of test repo") + session.run("git", "clone", TEST_REPO_URL, CLONE_REPO_DIR, external=True) + session.run("bash", "-e", "run_tests.sh", external=True) + + @nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) def system(session): """Run the system test suite.""" diff --git a/test_proxy/README.md b/test_proxy/README.md new file mode 100644 index 000000000..08741fd5d --- /dev/null +++ b/test_proxy/README.md @@ -0,0 +1,60 @@ +# CBT Python Test Proxy + +The CBT test proxy is intended for running conformance tests for Cloud Bigtable Python Client. + +## Option 1: Run Tests with Nox + +You can run the conformance tests in a single line by calling `nox -s conformance` from the repo root + + +``` +cd python-bigtable/test_proxy +nox -s conformance +``` + +## Option 2: Run processes manually + +### Start test proxy + +You can use `test_proxy.py` to launch a new test proxy process directly + +``` +cd python-bigtable/test_proxy +python test_proxy.py +``` + +The port can be set by passing in an extra positional argument + +``` +cd python-bigtable/test_proxy +python test_proxy.py --port 8080 +``` + +You can run the test proxy against the previous `v2` client by running it with the `--legacy-client` flag: + +``` +python test_proxy.py --legacy-client +``` + +### Run the test cases + +Prerequisites: +- If you have not already done so, [install golang](https://go.dev/doc/install). +- Before running tests, [launch an instance of the test proxy](#start-test-proxy) +in a separate shell session, and make note of the port + + +Clone and navigate to the go test library: + +``` +git clone https://github.com/googleapis/cloud-bigtable-clients-test.git +cd cloud-bigtable-clients-test/tests +``` + + +Launch the tests + +``` +go test -v -proxy_addr=:50055 +``` + diff --git a/test_proxy/handlers/client_handler_data.py b/test_proxy/handlers/client_handler_data.py new file mode 100644 index 000000000..43ff5d634 --- /dev/null +++ b/test_proxy/handlers/client_handler_data.py @@ -0,0 +1,214 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +""" +This module contains the client handler process for proxy_server.py. +""" +import os + +from google.cloud.environment_vars import BIGTABLE_EMULATOR +from google.cloud.bigtable.data import BigtableDataClientAsync + + +def error_safe(func): + """ + Catch and pass errors back to the grpc_server_process + Also check if client is closed before processing requests + """ + async def wrapper(self, *args, **kwargs): + try: + if self.closed: + raise RuntimeError("client is closed") + return await func(self, *args, **kwargs) + except (Exception, NotImplementedError) as e: + # exceptions should be raised in grpc_server_process + return encode_exception(e) + + return wrapper + + +def encode_exception(exc): + """ + Encode an exception or chain of exceptions to pass back to grpc_handler + """ + from google.api_core.exceptions import GoogleAPICallError + error_msg = f"{type(exc).__name__}: {exc}" + result = {"error": error_msg} + if exc.__cause__: + result["cause"] = encode_exception(exc.__cause__) + if hasattr(exc, "exceptions"): + result["subexceptions"] = [encode_exception(e) for e in exc.exceptions] + if hasattr(exc, "index"): + result["index"] = exc.index + if isinstance(exc, GoogleAPICallError): + if exc.grpc_status_code is not None: + result["code"] = exc.grpc_status_code.value[0] + elif exc.code is not None: + result["code"] = int(exc.code) + else: + result["code"] = -1 + elif result.get("cause", {}).get("code", None): + # look for code code in cause + result["code"] = result["cause"]["code"] + elif result.get("subexceptions", None): + # look for code in subexceptions + for subexc in result["subexceptions"]: + if subexc.get("code", None): + result["code"] = subexc["code"] + return result + + +class TestProxyClientHandler: + """ + Implements the same methods as the grpc server, but handles the client + library side of the request. + + Requests received in TestProxyGrpcServer are converted to a dictionary, + and supplied to the TestProxyClientHandler methods as kwargs. + The client response is then returned back to the TestProxyGrpcServer + """ + + def __init__( + self, + data_target=None, + project_id=None, + instance_id=None, + app_profile_id=None, + per_operation_timeout=None, + **kwargs, + ): + self.closed = False + # use emulator + os.environ[BIGTABLE_EMULATOR] = data_target + self.client = BigtableDataClientAsync(project=project_id) + self.instance_id = instance_id + self.app_profile_id = app_profile_id + self.per_operation_timeout = per_operation_timeout + + def close(self): + # TODO: call self.client.close() + self.closed = True + + @error_safe + async def ReadRows(self, request, **kwargs): + table_id = request.pop("table_name").split("/")[-1] + app_profile_id = self.app_profile_id or request.get("app_profile_id", None) + table = self.client.get_table(self.instance_id, table_id, app_profile_id) + kwargs["operation_timeout"] = kwargs.get("operation_timeout", self.per_operation_timeout) or 20 + result_list = await table.read_rows(request, **kwargs) + # pack results back into protobuf-parsable format + serialized_response = [row._to_dict() for row in result_list] + return serialized_response + + @error_safe + async def ReadRow(self, row_key, **kwargs): + table_id = kwargs.pop("table_name").split("/")[-1] + app_profile_id = self.app_profile_id or kwargs.get("app_profile_id", None) + table = self.client.get_table(self.instance_id, table_id, app_profile_id) + kwargs["operation_timeout"] = kwargs.get("operation_timeout", self.per_operation_timeout) or 20 + result_row = await table.read_row(row_key, **kwargs) + # pack results back into protobuf-parsable format + if result_row: + return result_row._to_dict() + else: + return "None" + + @error_safe + async def MutateRow(self, request, **kwargs): + from google.cloud.bigtable.data.mutations import Mutation + table_id = request["table_name"].split("/")[-1] + app_profile_id = self.app_profile_id or request.get("app_profile_id", None) + table = self.client.get_table(self.instance_id, table_id, app_profile_id) + kwargs["operation_timeout"] = kwargs.get("operation_timeout", self.per_operation_timeout) or 20 + row_key = request["row_key"] + mutations = [Mutation._from_dict(d) for d in request["mutations"]] + await table.mutate_row(row_key, mutations, **kwargs) + return "OK" + + @error_safe + async def BulkMutateRows(self, request, **kwargs): + from google.cloud.bigtable.data.mutations import RowMutationEntry + table_id = request["table_name"].split("/")[-1] + app_profile_id = self.app_profile_id or request.get("app_profile_id", None) + table = self.client.get_table(self.instance_id, table_id, app_profile_id) + kwargs["operation_timeout"] = kwargs.get("operation_timeout", self.per_operation_timeout) or 20 + entry_list = [RowMutationEntry._from_dict(entry) for entry in request["entries"]] + await table.bulk_mutate_rows(entry_list, **kwargs) + return "OK" + + @error_safe + async def CheckAndMutateRow(self, request, **kwargs): + from google.cloud.bigtable.data.mutations import Mutation, SetCell + table_id = request["table_name"].split("/")[-1] + app_profile_id = self.app_profile_id or request.get("app_profile_id", None) + table = self.client.get_table(self.instance_id, table_id, app_profile_id) + kwargs["operation_timeout"] = kwargs.get("operation_timeout", self.per_operation_timeout) or 20 + row_key = request["row_key"] + # add default values for incomplete dicts, so they can still be parsed to objects + true_mutations = [] + for mut_dict in request.get("true_mutations", []): + try: + true_mutations.append(Mutation._from_dict(mut_dict)) + except ValueError: + # invalid mutation type. Conformance test may be sending generic empty request + mutation = SetCell("", "", "", 0) + true_mutations.append(mutation) + false_mutations = [] + for mut_dict in request.get("false_mutations", []): + try: + false_mutations.append(Mutation._from_dict(mut_dict)) + except ValueError: + # invalid mutation type. Conformance test may be sending generic empty request + false_mutations.append(SetCell("", "", "", 0)) + predicate_filter = request.get("predicate_filter", None) + result = await table.check_and_mutate_row( + row_key, + predicate_filter, + true_case_mutations=true_mutations, + false_case_mutations=false_mutations, + **kwargs, + ) + return result + + @error_safe + async def ReadModifyWriteRow(self, request, **kwargs): + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule + table_id = request["table_name"].split("/")[-1] + app_profile_id = self.app_profile_id or request.get("app_profile_id", None) + table = self.client.get_table(self.instance_id, table_id, app_profile_id) + kwargs["operation_timeout"] = kwargs.get("operation_timeout", self.per_operation_timeout) or 20 + row_key = request["row_key"] + rules = [] + for rule_dict in request.get("rules", []): + qualifier = rule_dict["column_qualifier"] + if "append_value" in rule_dict: + new_rule = AppendValueRule(rule_dict["family_name"], qualifier, rule_dict["append_value"]) + else: + new_rule = IncrementRule(rule_dict["family_name"], qualifier, rule_dict["increment_amount"]) + rules.append(new_rule) + result = await table.read_modify_write_row(row_key, rules, **kwargs) + # pack results back into protobuf-parsable format + if result: + return result._to_dict() + else: + return "None" + + @error_safe + async def SampleRowKeys(self, request, **kwargs): + table_id = request["table_name"].split("/")[-1] + app_profile_id = self.app_profile_id or request.get("app_profile_id", None) + table = self.client.get_table(self.instance_id, table_id, app_profile_id) + kwargs["operation_timeout"] = kwargs.get("operation_timeout", self.per_operation_timeout) or 20 + result = await table.sample_row_keys(**kwargs) + return result diff --git a/test_proxy/handlers/client_handler_legacy.py b/test_proxy/handlers/client_handler_legacy.py new file mode 100644 index 000000000..b423165f1 --- /dev/null +++ b/test_proxy/handlers/client_handler_legacy.py @@ -0,0 +1,132 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +""" +This module contains the client handler process for proxy_server.py. +""" +import os + +from google.cloud.environment_vars import BIGTABLE_EMULATOR +from google.cloud.bigtable.client import Client + +import client_handler_data as client_handler + +import warnings +warnings.filterwarnings("ignore", category=DeprecationWarning) + + +class LegacyTestProxyClientHandler(client_handler.TestProxyClientHandler): + + def __init__( + self, + data_target=None, + project_id=None, + instance_id=None, + app_profile_id=None, + per_operation_timeout=None, + **kwargs, + ): + self.closed = False + # use emulator + os.environ[BIGTABLE_EMULATOR] = data_target + self.client = Client(project=project_id) + self.instance_id = instance_id + self.app_profile_id = app_profile_id + self.per_operation_timeout = per_operation_timeout + + async def close(self): + self.closed = True + + @client_handler.error_safe + async def ReadRows(self, request, **kwargs): + table_id = request["table_name"].split("/")[-1] + # app_profile_id = self.app_profile_id or request.get("app_profile_id", None) + instance = self.client.instance(self.instance_id) + table = instance.table(table_id) + + limit = request.get("rows_limit", None) + start_key = request.get("rows", {}).get("row_keys", [None])[0] + end_key = request.get("rows", {}).get("row_keys", [None])[-1] + end_inclusive = request.get("rows", {}).get("row_ranges", [{}])[-1].get("end_key_closed", True) + + row_list = [] + for row in table.read_rows(start_key=start_key, end_key=end_key, limit=limit, end_inclusive=end_inclusive): + # parse results into proto formatted dict + dict_val = {"row_key": row.row_key} + for family, family_cells in row.cells.items(): + family_dict = {"name": family} + for qualifier, qualifier_cells in family_cells.items(): + column_dict = {"qualifier": qualifier} + for cell in qualifier_cells: + cell_dict = { + "value": cell.value, + "timestamp_micros": cell.timestamp.timestamp() * 1000000, + "labels": cell.labels, + } + column_dict.setdefault("cells", []).append(cell_dict) + family_dict.setdefault("columns", []).append(column_dict) + dict_val.setdefault("families", []).append(family_dict) + row_list.append(dict_val) + return row_list + + @client_handler.error_safe + async def MutateRow(self, request, **kwargs): + from google.cloud.bigtable.row import DirectRow + table_id = request["table_name"].split("/")[-1] + instance = self.client.instance(self.instance_id) + table = instance.table(table_id) + row_key = request["row_key"] + new_row = DirectRow(row_key, table) + for m_dict in request.get("mutations", []): + if m_dict.get("set_cell"): + details = m_dict["set_cell"] + new_row.set_cell(details["family_name"], details["column_qualifier"], details["value"], timestamp=details["timestamp_micros"]) + elif m_dict.get("delete_from_column"): + details = m_dict["delete_from_column"] + new_row.delete_cell(details["family_name"], details["column_qualifier"], timestamp=details["timestamp_micros"]) + elif m_dict.get("delete_from_family"): + details = m_dict["delete_from_family"] + new_row.delete_cells(details["family_name"], timestamp=details["timestamp_micros"]) + elif m_dict.get("delete_from_row"): + new_row.delete() + async with self.measure_call(): + table.mutate_rows([new_row]) + return "OK" + + @client_handler.error_safe + async def BulkMutateRows(self, request, **kwargs): + from google.cloud.bigtable.row import DirectRow + table_id = request["table_name"].split("/")[-1] + instance = self.client.instance(self.instance_id) + table = instance.table(table_id) + rows = [] + for entry in request.get("entries", []): + row_key = entry["row_key"] + new_row = DirectRow(row_key, table) + for m_dict in entry.get("mutations", {}): + if m_dict.get("set_cell"): + details = m_dict["set_cell"] + new_row.set_cell(details["family_name"], details["column_qualifier"], details["value"], timestamp=details.get("timestamp_micros",None)) + elif m_dict.get("delete_from_column"): + details = m_dict["delete_from_column"] + new_row.delete_cell(details["family_name"], details["column_qualifier"], timestamp=details["timestamp_micros"]) + elif m_dict.get("delete_from_family"): + details = m_dict["delete_from_family"] + new_row.delete_cells(details["family_name"], timestamp=details["timestamp_micros"]) + elif m_dict.get("delete_from_row"): + new_row.delete() + rows.append(new_row) + async with self.measure_call(): + table.mutate_rows(rows) + return "OK" + diff --git a/test_proxy/handlers/grpc_handler.py b/test_proxy/handlers/grpc_handler.py new file mode 100644 index 000000000..2c70778dd --- /dev/null +++ b/test_proxy/handlers/grpc_handler.py @@ -0,0 +1,148 @@ + +import time + +import test_proxy_pb2 +import test_proxy_pb2_grpc +import data_pb2 +import bigtable_pb2 +from google.rpc.status_pb2 import Status +from google.protobuf import json_format + + +class TestProxyGrpcServer(test_proxy_pb2_grpc.CloudBigtableV2TestProxyServicer): + """ + Implements a grpc server that proxies conformance test requests to the client library + + Due to issues with using protoc-compiled protos and client-library + proto-plus objects in the same process, this server defers requests to + matching methods in a TestProxyClientHandler instance in a separate + process. + This happens invisbly in the decorator @delegate_to_client_handler, with the + results attached to each request as a client_response kwarg + """ + + def __init__(self, request_q, queue_pool): + self.open_queues = list(range(len(queue_pool))) + self.queue_pool = queue_pool + self.request_q = request_q + + def delegate_to_client_handler(func, timeout_seconds=300): + """ + Decorator that transparently passes a request to the client + handler process, and then attaches the resonse to the wrapped call + """ + + def wrapper(self, request, context, **kwargs): + deadline = time.time() + timeout_seconds + json_dict = json_format.MessageToDict(request) + out_idx = self.open_queues.pop() + json_dict["proxy_request"] = func.__name__ + json_dict["response_queue_idx"] = out_idx + out_q = self.queue_pool[out_idx] + self.request_q.put(json_dict) + # wait for response + while time.time() < deadline: + if not out_q.empty(): + response = out_q.get() + self.open_queues.append(out_idx) + if isinstance(response, Exception): + raise response + else: + return func( + self, + request, + context, + client_response=response, + **kwargs, + ) + time.sleep(1e-4) + + return wrapper + + + @delegate_to_client_handler + def CreateClient(self, request, context, client_response=None): + return test_proxy_pb2.CreateClientResponse() + + @delegate_to_client_handler + def CloseClient(self, request, context, client_response=None): + return test_proxy_pb2.CloseClientResponse() + + @delegate_to_client_handler + def RemoveClient(self, request, context, client_response=None): + return test_proxy_pb2.RemoveClientResponse() + + @delegate_to_client_handler + def ReadRows(self, request, context, client_response=None): + status = Status() + rows = [] + if isinstance(client_response, dict) and "error" in client_response: + status = Status(code=5, message=client_response["error"]) + else: + rows = [data_pb2.Row(**d) for d in client_response] + result = test_proxy_pb2.RowsResult(row=rows, status=status) + return result + + @delegate_to_client_handler + def ReadRow(self, request, context, client_response=None): + status = Status() + row = None + if isinstance(client_response, dict) and "error" in client_response: + status=Status(code=client_response.get("code", 5), message=client_response.get("error")) + elif client_response != "None": + row = data_pb2.Row(**client_response) + result = test_proxy_pb2.RowResult(row=row, status=status) + return result + + @delegate_to_client_handler + def MutateRow(self, request, context, client_response=None): + status = Status() + if isinstance(client_response, dict) and "error" in client_response: + status = Status(code=client_response.get("code", 5), message=client_response["error"]) + return test_proxy_pb2.MutateRowResult(status=status) + + @delegate_to_client_handler + def BulkMutateRows(self, request, context, client_response=None): + status = Status() + entries = [] + if isinstance(client_response, dict) and "error" in client_response: + entries = [bigtable_pb2.MutateRowsResponse.Entry(index=exc_dict.get("index",1), status=Status(code=exc_dict.get("code", 5))) for exc_dict in client_response.get("subexceptions", [])] + if not entries: + # only return failure on the overall request if there are failed entries + status = Status(code=client_response.get("code", 5), message=client_response["error"]) + # TODO: protos were updated. entry is now entries: https://github.com/googleapis/cndb-client-testing-protos/commit/e6205a2bba04acc10d12421a1402870b4a525fb3 + response = test_proxy_pb2.MutateRowsResult(status=status, entry=entries) + return response + + @delegate_to_client_handler + def CheckAndMutateRow(self, request, context, client_response=None): + if isinstance(client_response, dict) and "error" in client_response: + status = Status(code=client_response.get("code", 5), message=client_response["error"]) + response = test_proxy_pb2.CheckAndMutateRowResult(status=status) + else: + result = bigtable_pb2.CheckAndMutateRowResponse(predicate_matched=client_response) + response = test_proxy_pb2.CheckAndMutateRowResult(result=result, status=Status()) + return response + + @delegate_to_client_handler + def ReadModifyWriteRow(self, request, context, client_response=None): + status = Status() + row = None + if isinstance(client_response, dict) and "error" in client_response: + status = Status(code=client_response.get("code", 5), message=client_response.get("error")) + elif client_response != "None": + row = data_pb2.Row(**client_response) + result = test_proxy_pb2.RowResult(row=row, status=status) + return result + + @delegate_to_client_handler + def SampleRowKeys(self, request, context, client_response=None): + status = Status() + sample_list = [] + if isinstance(client_response, dict) and "error" in client_response: + status = Status(code=client_response.get("code", 5), message=client_response.get("error")) + else: + for sample in client_response: + sample_list.append(bigtable_pb2.SampleRowKeysResponse(offset_bytes=sample[1], row_key=sample[0])) + # TODO: protos were updated. sample is now samples: https://github.com/googleapis/cndb-client-testing-protos/commit/e6205a2bba04acc10d12421a1402870b4a525fb3 + return test_proxy_pb2.SampleRowKeysResult(status=status, sample=sample_list) diff --git a/test_proxy/noxfile.py b/test_proxy/noxfile.py new file mode 100644 index 000000000..bebf247b7 --- /dev/null +++ b/test_proxy/noxfile.py @@ -0,0 +1,80 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +from __future__ import absolute_import +import os +import pathlib +import re +from colorlog.escape_codes import parse_colors + +import nox + + +DEFAULT_PYTHON_VERSION = "3.10" + +PROXY_SERVER_PORT=os.environ.get("PROXY_SERVER_PORT", "50055") +PROXY_CLIENT_VERSION=os.environ.get("PROXY_CLIENT_VERSION", None) + +CURRENT_DIRECTORY = pathlib.Path(__file__).parent.absolute() +REPO_ROOT_DIRECTORY = CURRENT_DIRECTORY.parent + +nox.options.sessions = ["run_proxy", "conformance_tests"] + +TEST_REPO_URL = "https://github.com/googleapis/cloud-bigtable-clients-test.git" +CLONE_REPO_DIR = "cloud-bigtable-clients-test" + +# Error if a python version is missing +nox.options.error_on_missing_interpreters = True + + +def default(session): + """ + if nox is run directly, run the test_proxy session + """ + test_proxy(session) + + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def conformance_tests(session): + """ + download and run the conformance test suite against the test proxy + """ + import subprocess + import time + # download the conformance test suite + clone_dir = os.path.join(CURRENT_DIRECTORY, CLONE_REPO_DIR) + if not os.path.exists(clone_dir): + print("downloading copy of test repo") + session.run("git", "clone", TEST_REPO_URL, CLONE_REPO_DIR) + # start tests + with session.chdir(f"{clone_dir}/tests"): + session.run("go", "test", "-v", f"-proxy_addr=:{PROXY_SERVER_PORT}") + +@nox.session(python=DEFAULT_PYTHON_VERSION) +def test_proxy(session): + """Start up the test proxy""" + # Install all dependencies, then install this package into the + # virtualenv's dist-packages. + # session.install( + # "grpcio", + # ) + if PROXY_CLIENT_VERSION is not None: + # install released version of the library + session.install(f"python-bigtable=={PROXY_CLIENT_VERSION}") + else: + # install the library from the source + session.install("-e", str(REPO_ROOT_DIRECTORY)) + session.install("-e", str(REPO_ROOT_DIRECTORY / "python-api-core")) + + session.run("python", "test_proxy.py", "--port", PROXY_SERVER_PORT, *session.posargs,) diff --git a/test_proxy/protos/bigtable_pb2.py b/test_proxy/protos/bigtable_pb2.py new file mode 100644 index 000000000..936a4ed55 --- /dev/null +++ b/test_proxy/protos/bigtable_pb2.py @@ -0,0 +1,145 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: google/bigtable/v2/bigtable.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.api import annotations_pb2 as google_dot_api_dot_annotations__pb2 +from google.api import client_pb2 as google_dot_api_dot_client__pb2 +from google.api import field_behavior_pb2 as google_dot_api_dot_field__behavior__pb2 +from google.api import resource_pb2 as google_dot_api_dot_resource__pb2 +from google.api import routing_pb2 as google_dot_api_dot_routing__pb2 +import data_pb2 as google_dot_bigtable_dot_v2_dot_data__pb2 +import request_stats_pb2 as google_dot_bigtable_dot_v2_dot_request__stats__pb2 +from google.protobuf import duration_pb2 as google_dot_protobuf_dot_duration__pb2 +from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 +from google.protobuf import wrappers_pb2 as google_dot_protobuf_dot_wrappers__pb2 +from google.rpc import status_pb2 as google_dot_rpc_dot_status__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n!google/bigtable/v2/bigtable.proto\x12\x12google.bigtable.v2\x1a\x1cgoogle/api/annotations.proto\x1a\x17google/api/client.proto\x1a\x1fgoogle/api/field_behavior.proto\x1a\x19google/api/resource.proto\x1a\x18google/api/routing.proto\x1a\x1dgoogle/bigtable/v2/data.proto\x1a&google/bigtable/v2/request_stats.proto\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1egoogle/protobuf/wrappers.proto\x1a\x17google/rpc/status.proto\"\x90\x03\n\x0fReadRowsRequest\x12>\n\ntable_name\x18\x01 \x01(\tB*\xe0\x41\x02\xfa\x41$\n\"bigtableadmin.googleapis.com/Table\x12\x16\n\x0e\x61pp_profile_id\x18\x05 \x01(\t\x12(\n\x04rows\x18\x02 \x01(\x0b\x32\x1a.google.bigtable.v2.RowSet\x12-\n\x06\x66ilter\x18\x03 \x01(\x0b\x32\x1d.google.bigtable.v2.RowFilter\x12\x12\n\nrows_limit\x18\x04 \x01(\x03\x12P\n\x12request_stats_view\x18\x06 \x01(\x0e\x32\x34.google.bigtable.v2.ReadRowsRequest.RequestStatsView\"f\n\x10RequestStatsView\x12\"\n\x1eREQUEST_STATS_VIEW_UNSPECIFIED\x10\x00\x12\x16\n\x12REQUEST_STATS_NONE\x10\x01\x12\x16\n\x12REQUEST_STATS_FULL\x10\x02\"\xb1\x03\n\x10ReadRowsResponse\x12>\n\x06\x63hunks\x18\x01 \x03(\x0b\x32..google.bigtable.v2.ReadRowsResponse.CellChunk\x12\x1c\n\x14last_scanned_row_key\x18\x02 \x01(\x0c\x12\x37\n\rrequest_stats\x18\x03 \x01(\x0b\x32 .google.bigtable.v2.RequestStats\x1a\x85\x02\n\tCellChunk\x12\x0f\n\x07row_key\x18\x01 \x01(\x0c\x12\x31\n\x0b\x66\x61mily_name\x18\x02 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\tqualifier\x18\x03 \x01(\x0b\x32\x1b.google.protobuf.BytesValue\x12\x18\n\x10timestamp_micros\x18\x04 \x01(\x03\x12\x0e\n\x06labels\x18\x05 \x03(\t\x12\r\n\x05value\x18\x06 \x01(\x0c\x12\x12\n\nvalue_size\x18\x07 \x01(\x05\x12\x13\n\treset_row\x18\x08 \x01(\x08H\x00\x12\x14\n\ncommit_row\x18\t \x01(\x08H\x00\x42\x0c\n\nrow_status\"n\n\x14SampleRowKeysRequest\x12>\n\ntable_name\x18\x01 \x01(\tB*\xe0\x41\x02\xfa\x41$\n\"bigtableadmin.googleapis.com/Table\x12\x16\n\x0e\x61pp_profile_id\x18\x02 \x01(\t\">\n\x15SampleRowKeysResponse\x12\x0f\n\x07row_key\x18\x01 \x01(\x0c\x12\x14\n\x0coffset_bytes\x18\x02 \x01(\x03\"\xb6\x01\n\x10MutateRowRequest\x12>\n\ntable_name\x18\x01 \x01(\tB*\xe0\x41\x02\xfa\x41$\n\"bigtableadmin.googleapis.com/Table\x12\x16\n\x0e\x61pp_profile_id\x18\x04 \x01(\t\x12\x14\n\x07row_key\x18\x02 \x01(\x0c\x42\x03\xe0\x41\x02\x12\x34\n\tmutations\x18\x03 \x03(\x0b\x32\x1c.google.bigtable.v2.MutationB\x03\xe0\x41\x02\"\x13\n\x11MutateRowResponse\"\xfe\x01\n\x11MutateRowsRequest\x12>\n\ntable_name\x18\x01 \x01(\tB*\xe0\x41\x02\xfa\x41$\n\"bigtableadmin.googleapis.com/Table\x12\x16\n\x0e\x61pp_profile_id\x18\x03 \x01(\t\x12\x41\n\x07\x65ntries\x18\x02 \x03(\x0b\x32+.google.bigtable.v2.MutateRowsRequest.EntryB\x03\xe0\x41\x02\x1aN\n\x05\x45ntry\x12\x0f\n\x07row_key\x18\x01 \x01(\x0c\x12\x34\n\tmutations\x18\x02 \x03(\x0b\x32\x1c.google.bigtable.v2.MutationB\x03\xe0\x41\x02\"\x8f\x01\n\x12MutateRowsResponse\x12=\n\x07\x65ntries\x18\x01 \x03(\x0b\x32,.google.bigtable.v2.MutateRowsResponse.Entry\x1a:\n\x05\x45ntry\x12\r\n\x05index\x18\x01 \x01(\x03\x12\"\n\x06status\x18\x02 \x01(\x0b\x32\x12.google.rpc.Status\"\xae\x02\n\x18\x43heckAndMutateRowRequest\x12>\n\ntable_name\x18\x01 \x01(\tB*\xe0\x41\x02\xfa\x41$\n\"bigtableadmin.googleapis.com/Table\x12\x16\n\x0e\x61pp_profile_id\x18\x07 \x01(\t\x12\x14\n\x07row_key\x18\x02 \x01(\x0c\x42\x03\xe0\x41\x02\x12\x37\n\x10predicate_filter\x18\x06 \x01(\x0b\x32\x1d.google.bigtable.v2.RowFilter\x12\x34\n\x0etrue_mutations\x18\x04 \x03(\x0b\x32\x1c.google.bigtable.v2.Mutation\x12\x35\n\x0f\x66\x61lse_mutations\x18\x05 \x03(\x0b\x32\x1c.google.bigtable.v2.Mutation\"6\n\x19\x43heckAndMutateRowResponse\x12\x19\n\x11predicate_matched\x18\x01 \x01(\x08\"i\n\x12PingAndWarmRequest\x12;\n\x04name\x18\x01 \x01(\tB-\xe0\x41\x02\xfa\x41\'\n%bigtableadmin.googleapis.com/Instance\x12\x16\n\x0e\x61pp_profile_id\x18\x02 \x01(\t\"\x15\n\x13PingAndWarmResponse\"\xc6\x01\n\x19ReadModifyWriteRowRequest\x12>\n\ntable_name\x18\x01 \x01(\tB*\xe0\x41\x02\xfa\x41$\n\"bigtableadmin.googleapis.com/Table\x12\x16\n\x0e\x61pp_profile_id\x18\x04 \x01(\t\x12\x14\n\x07row_key\x18\x02 \x01(\x0c\x42\x03\xe0\x41\x02\x12;\n\x05rules\x18\x03 \x03(\x0b\x32\'.google.bigtable.v2.ReadModifyWriteRuleB\x03\xe0\x41\x02\"B\n\x1aReadModifyWriteRowResponse\x12$\n\x03row\x18\x01 \x01(\x0b\x32\x17.google.bigtable.v2.Row\"\x86\x01\n,GenerateInitialChangeStreamPartitionsRequest\x12>\n\ntable_name\x18\x01 \x01(\tB*\xe0\x41\x02\xfa\x41$\n\"bigtableadmin.googleapis.com/Table\x12\x16\n\x0e\x61pp_profile_id\x18\x02 \x01(\t\"g\n-GenerateInitialChangeStreamPartitionsResponse\x12\x36\n\tpartition\x18\x01 \x01(\x0b\x32#.google.bigtable.v2.StreamPartition\"\x9b\x03\n\x17ReadChangeStreamRequest\x12>\n\ntable_name\x18\x01 \x01(\tB*\xe0\x41\x02\xfa\x41$\n\"bigtableadmin.googleapis.com/Table\x12\x16\n\x0e\x61pp_profile_id\x18\x02 \x01(\t\x12\x36\n\tpartition\x18\x03 \x01(\x0b\x32#.google.bigtable.v2.StreamPartition\x12\x30\n\nstart_time\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.TimestampH\x00\x12K\n\x13\x63ontinuation_tokens\x18\x06 \x01(\x0b\x32,.google.bigtable.v2.StreamContinuationTokensH\x00\x12,\n\x08\x65nd_time\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x35\n\x12heartbeat_duration\x18\x07 \x01(\x0b\x32\x19.google.protobuf.DurationB\x0c\n\nstart_from\"\xeb\t\n\x18ReadChangeStreamResponse\x12N\n\x0b\x64\x61ta_change\x18\x01 \x01(\x0b\x32\x37.google.bigtable.v2.ReadChangeStreamResponse.DataChangeH\x00\x12K\n\theartbeat\x18\x02 \x01(\x0b\x32\x36.google.bigtable.v2.ReadChangeStreamResponse.HeartbeatH\x00\x12P\n\x0c\x63lose_stream\x18\x03 \x01(\x0b\x32\x38.google.bigtable.v2.ReadChangeStreamResponse.CloseStreamH\x00\x1a\xf4\x01\n\rMutationChunk\x12X\n\nchunk_info\x18\x01 \x01(\x0b\x32\x44.google.bigtable.v2.ReadChangeStreamResponse.MutationChunk.ChunkInfo\x12.\n\x08mutation\x18\x02 \x01(\x0b\x32\x1c.google.bigtable.v2.Mutation\x1aY\n\tChunkInfo\x12\x1a\n\x12\x63hunked_value_size\x18\x01 \x01(\x05\x12\x1c\n\x14\x63hunked_value_offset\x18\x02 \x01(\x05\x12\x12\n\nlast_chunk\x18\x03 \x01(\x08\x1a\xc6\x03\n\nDataChange\x12J\n\x04type\x18\x01 \x01(\x0e\x32<.google.bigtable.v2.ReadChangeStreamResponse.DataChange.Type\x12\x19\n\x11source_cluster_id\x18\x02 \x01(\t\x12\x0f\n\x07row_key\x18\x03 \x01(\x0c\x12\x34\n\x10\x63ommit_timestamp\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x12\n\ntiebreaker\x18\x05 \x01(\x05\x12J\n\x06\x63hunks\x18\x06 \x03(\x0b\x32:.google.bigtable.v2.ReadChangeStreamResponse.MutationChunk\x12\x0c\n\x04\x64one\x18\x08 \x01(\x08\x12\r\n\x05token\x18\t \x01(\t\x12;\n\x17\x65stimated_low_watermark\x18\n \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"P\n\x04Type\x12\x14\n\x10TYPE_UNSPECIFIED\x10\x00\x12\x08\n\x04USER\x10\x01\x12\x16\n\x12GARBAGE_COLLECTION\x10\x02\x12\x10\n\x0c\x43ONTINUATION\x10\x03\x1a\x91\x01\n\tHeartbeat\x12G\n\x12\x63ontinuation_token\x18\x01 \x01(\x0b\x32+.google.bigtable.v2.StreamContinuationToken\x12;\n\x17\x65stimated_low_watermark\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x1a{\n\x0b\x43loseStream\x12\"\n\x06status\x18\x01 \x01(\x0b\x32\x12.google.rpc.Status\x12H\n\x13\x63ontinuation_tokens\x18\x02 \x03(\x0b\x32+.google.bigtable.v2.StreamContinuationTokenB\x0f\n\rstream_record2\xd7\x18\n\x08\x42igtable\x12\x9b\x02\n\x08ReadRows\x12#.google.bigtable.v2.ReadRowsRequest\x1a$.google.bigtable.v2.ReadRowsResponse\"\xc1\x01\x82\xd3\xe4\x93\x02>\"9/v2/{table_name=projects/*/instances/*/tables/*}:readRows:\x01*\x8a\xd3\xe4\x93\x02N\x12:\n\ntable_name\x12,{table_name=projects/*/instances/*/tables/*}\x12\x10\n\x0e\x61pp_profile_id\xda\x41\ntable_name\xda\x41\x19table_name,app_profile_id0\x01\x12\xac\x02\n\rSampleRowKeys\x12(.google.bigtable.v2.SampleRowKeysRequest\x1a).google.bigtable.v2.SampleRowKeysResponse\"\xc3\x01\x82\xd3\xe4\x93\x02@\x12>/v2/{table_name=projects/*/instances/*/tables/*}:sampleRowKeys\x8a\xd3\xe4\x93\x02N\x12:\n\ntable_name\x12,{table_name=projects/*/instances/*/tables/*}\x12\x10\n\x0e\x61pp_profile_id\xda\x41\ntable_name\xda\x41\x19table_name,app_profile_id0\x01\x12\xc1\x02\n\tMutateRow\x12$.google.bigtable.v2.MutateRowRequest\x1a%.google.bigtable.v2.MutateRowResponse\"\xe6\x01\x82\xd3\xe4\x93\x02?\":/v2/{table_name=projects/*/instances/*/tables/*}:mutateRow:\x01*\x8a\xd3\xe4\x93\x02N\x12:\n\ntable_name\x12,{table_name=projects/*/instances/*/tables/*}\x12\x10\n\x0e\x61pp_profile_id\xda\x41\x1ctable_name,row_key,mutations\xda\x41+table_name,row_key,mutations,app_profile_id\x12\xb3\x02\n\nMutateRows\x12%.google.bigtable.v2.MutateRowsRequest\x1a&.google.bigtable.v2.MutateRowsResponse\"\xd3\x01\x82\xd3\xe4\x93\x02@\";/v2/{table_name=projects/*/instances/*/tables/*}:mutateRows:\x01*\x8a\xd3\xe4\x93\x02N\x12:\n\ntable_name\x12,{table_name=projects/*/instances/*/tables/*}\x12\x10\n\x0e\x61pp_profile_id\xda\x41\x12table_name,entries\xda\x41!table_name,entries,app_profile_id0\x01\x12\xad\x03\n\x11\x43heckAndMutateRow\x12,.google.bigtable.v2.CheckAndMutateRowRequest\x1a-.google.bigtable.v2.CheckAndMutateRowResponse\"\xba\x02\x82\xd3\xe4\x93\x02G\"B/v2/{table_name=projects/*/instances/*/tables/*}:checkAndMutateRow:\x01*\x8a\xd3\xe4\x93\x02N\x12:\n\ntable_name\x12,{table_name=projects/*/instances/*/tables/*}\x12\x10\n\x0e\x61pp_profile_id\xda\x41\x42table_name,row_key,predicate_filter,true_mutations,false_mutations\xda\x41Qtable_name,row_key,predicate_filter,true_mutations,false_mutations,app_profile_id\x12\xee\x01\n\x0bPingAndWarm\x12&.google.bigtable.v2.PingAndWarmRequest\x1a\'.google.bigtable.v2.PingAndWarmResponse\"\x8d\x01\x82\xd3\xe4\x93\x02+\"&/v2/{name=projects/*/instances/*}:ping:\x01*\x8a\xd3\xe4\x93\x02\x39\x12%\n\x04name\x12\x1d{name=projects/*/instances/*}\x12\x10\n\x0e\x61pp_profile_id\xda\x41\x04name\xda\x41\x13name,app_profile_id\x12\xdd\x02\n\x12ReadModifyWriteRow\x12-.google.bigtable.v2.ReadModifyWriteRowRequest\x1a..google.bigtable.v2.ReadModifyWriteRowResponse\"\xe7\x01\x82\xd3\xe4\x93\x02H\"C/v2/{table_name=projects/*/instances/*/tables/*}:readModifyWriteRow:\x01*\x8a\xd3\xe4\x93\x02N\x12:\n\ntable_name\x12,{table_name=projects/*/instances/*/tables/*}\x12\x10\n\x0e\x61pp_profile_id\xda\x41\x18table_name,row_key,rules\xda\x41\'table_name,row_key,rules,app_profile_id\x12\xbb\x02\n%GenerateInitialChangeStreamPartitions\x12@.google.bigtable.v2.GenerateInitialChangeStreamPartitionsRequest\x1a\x41.google.bigtable.v2.GenerateInitialChangeStreamPartitionsResponse\"\x8a\x01\x82\xd3\xe4\x93\x02[\"V/v2/{table_name=projects/*/instances/*/tables/*}:generateInitialChangeStreamPartitions:\x01*\xda\x41\ntable_name\xda\x41\x19table_name,app_profile_id0\x01\x12\xe6\x01\n\x10ReadChangeStream\x12+.google.bigtable.v2.ReadChangeStreamRequest\x1a,.google.bigtable.v2.ReadChangeStreamResponse\"u\x82\xd3\xe4\x93\x02\x46\"A/v2/{table_name=projects/*/instances/*/tables/*}:readChangeStream:\x01*\xda\x41\ntable_name\xda\x41\x19table_name,app_profile_id0\x01\x1a\xdb\x02\xca\x41\x17\x62igtable.googleapis.com\xd2\x41\xbd\x02https://www.googleapis.com/auth/bigtable.data,https://www.googleapis.com/auth/bigtable.data.readonly,https://www.googleapis.com/auth/cloud-bigtable.data,https://www.googleapis.com/auth/cloud-bigtable.data.readonly,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/cloud-platform.read-onlyB\xeb\x02\n\x16\x63om.google.bigtable.v2B\rBigtableProtoP\x01Z:google.golang.org/genproto/googleapis/bigtable/v2;bigtable\xaa\x02\x18Google.Cloud.Bigtable.V2\xca\x02\x18Google\\Cloud\\Bigtable\\V2\xea\x02\x1bGoogle::Cloud::Bigtable::V2\xea\x41P\n%bigtableadmin.googleapis.com/Instance\x12\'projects/{project}/instances/{instance}\xea\x41\\\n\"bigtableadmin.googleapis.com/Table\x12\x36projects/{project}/instances/{instance}/tables/{table}b\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.bigtable.v2.bigtable_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'\n\026com.google.bigtable.v2B\rBigtableProtoP\001Z:google.golang.org/genproto/googleapis/bigtable/v2;bigtable\252\002\030Google.Cloud.Bigtable.V2\312\002\030Google\\Cloud\\Bigtable\\V2\352\002\033Google::Cloud::Bigtable::V2\352AP\n%bigtableadmin.googleapis.com/Instance\022\'projects/{project}/instances/{instance}\352A\\\n\"bigtableadmin.googleapis.com/Table\0226projects/{project}/instances/{instance}/tables/{table}' + _READROWSREQUEST.fields_by_name['table_name']._options = None + _READROWSREQUEST.fields_by_name['table_name']._serialized_options = b'\340A\002\372A$\n\"bigtableadmin.googleapis.com/Table' + _SAMPLEROWKEYSREQUEST.fields_by_name['table_name']._options = None + _SAMPLEROWKEYSREQUEST.fields_by_name['table_name']._serialized_options = b'\340A\002\372A$\n\"bigtableadmin.googleapis.com/Table' + _MUTATEROWREQUEST.fields_by_name['table_name']._options = None + _MUTATEROWREQUEST.fields_by_name['table_name']._serialized_options = b'\340A\002\372A$\n\"bigtableadmin.googleapis.com/Table' + _MUTATEROWREQUEST.fields_by_name['row_key']._options = None + _MUTATEROWREQUEST.fields_by_name['row_key']._serialized_options = b'\340A\002' + _MUTATEROWREQUEST.fields_by_name['mutations']._options = None + _MUTATEROWREQUEST.fields_by_name['mutations']._serialized_options = b'\340A\002' + _MUTATEROWSREQUEST_ENTRY.fields_by_name['mutations']._options = None + _MUTATEROWSREQUEST_ENTRY.fields_by_name['mutations']._serialized_options = b'\340A\002' + _MUTATEROWSREQUEST.fields_by_name['table_name']._options = None + _MUTATEROWSREQUEST.fields_by_name['table_name']._serialized_options = b'\340A\002\372A$\n\"bigtableadmin.googleapis.com/Table' + _MUTATEROWSREQUEST.fields_by_name['entries']._options = None + _MUTATEROWSREQUEST.fields_by_name['entries']._serialized_options = b'\340A\002' + _CHECKANDMUTATEROWREQUEST.fields_by_name['table_name']._options = None + _CHECKANDMUTATEROWREQUEST.fields_by_name['table_name']._serialized_options = b'\340A\002\372A$\n\"bigtableadmin.googleapis.com/Table' + _CHECKANDMUTATEROWREQUEST.fields_by_name['row_key']._options = None + _CHECKANDMUTATEROWREQUEST.fields_by_name['row_key']._serialized_options = b'\340A\002' + _PINGANDWARMREQUEST.fields_by_name['name']._options = None + _PINGANDWARMREQUEST.fields_by_name['name']._serialized_options = b'\340A\002\372A\'\n%bigtableadmin.googleapis.com/Instance' + _READMODIFYWRITEROWREQUEST.fields_by_name['table_name']._options = None + _READMODIFYWRITEROWREQUEST.fields_by_name['table_name']._serialized_options = b'\340A\002\372A$\n\"bigtableadmin.googleapis.com/Table' + _READMODIFYWRITEROWREQUEST.fields_by_name['row_key']._options = None + _READMODIFYWRITEROWREQUEST.fields_by_name['row_key']._serialized_options = b'\340A\002' + _READMODIFYWRITEROWREQUEST.fields_by_name['rules']._options = None + _READMODIFYWRITEROWREQUEST.fields_by_name['rules']._serialized_options = b'\340A\002' + _GENERATEINITIALCHANGESTREAMPARTITIONSREQUEST.fields_by_name['table_name']._options = None + _GENERATEINITIALCHANGESTREAMPARTITIONSREQUEST.fields_by_name['table_name']._serialized_options = b'\340A\002\372A$\n\"bigtableadmin.googleapis.com/Table' + _READCHANGESTREAMREQUEST.fields_by_name['table_name']._options = None + _READCHANGESTREAMREQUEST.fields_by_name['table_name']._serialized_options = b'\340A\002\372A$\n\"bigtableadmin.googleapis.com/Table' + _BIGTABLE._options = None + _BIGTABLE._serialized_options = b'\312A\027bigtable.googleapis.com\322A\275\002https://www.googleapis.com/auth/bigtable.data,https://www.googleapis.com/auth/bigtable.data.readonly,https://www.googleapis.com/auth/cloud-bigtable.data,https://www.googleapis.com/auth/cloud-bigtable.data.readonly,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/cloud-platform.read-only' + _BIGTABLE.methods_by_name['ReadRows']._options = None + _BIGTABLE.methods_by_name['ReadRows']._serialized_options = b'\202\323\344\223\002>\"9/v2/{table_name=projects/*/instances/*/tables/*}:readRows:\001*\212\323\344\223\002N\022:\n\ntable_name\022,{table_name=projects/*/instances/*/tables/*}\022\020\n\016app_profile_id\332A\ntable_name\332A\031table_name,app_profile_id' + _BIGTABLE.methods_by_name['SampleRowKeys']._options = None + _BIGTABLE.methods_by_name['SampleRowKeys']._serialized_options = b'\202\323\344\223\002@\022>/v2/{table_name=projects/*/instances/*/tables/*}:sampleRowKeys\212\323\344\223\002N\022:\n\ntable_name\022,{table_name=projects/*/instances/*/tables/*}\022\020\n\016app_profile_id\332A\ntable_name\332A\031table_name,app_profile_id' + _BIGTABLE.methods_by_name['MutateRow']._options = None + _BIGTABLE.methods_by_name['MutateRow']._serialized_options = b'\202\323\344\223\002?\":/v2/{table_name=projects/*/instances/*/tables/*}:mutateRow:\001*\212\323\344\223\002N\022:\n\ntable_name\022,{table_name=projects/*/instances/*/tables/*}\022\020\n\016app_profile_id\332A\034table_name,row_key,mutations\332A+table_name,row_key,mutations,app_profile_id' + _BIGTABLE.methods_by_name['MutateRows']._options = None + _BIGTABLE.methods_by_name['MutateRows']._serialized_options = b'\202\323\344\223\002@\";/v2/{table_name=projects/*/instances/*/tables/*}:mutateRows:\001*\212\323\344\223\002N\022:\n\ntable_name\022,{table_name=projects/*/instances/*/tables/*}\022\020\n\016app_profile_id\332A\022table_name,entries\332A!table_name,entries,app_profile_id' + _BIGTABLE.methods_by_name['CheckAndMutateRow']._options = None + _BIGTABLE.methods_by_name['CheckAndMutateRow']._serialized_options = b'\202\323\344\223\002G\"B/v2/{table_name=projects/*/instances/*/tables/*}:checkAndMutateRow:\001*\212\323\344\223\002N\022:\n\ntable_name\022,{table_name=projects/*/instances/*/tables/*}\022\020\n\016app_profile_id\332ABtable_name,row_key,predicate_filter,true_mutations,false_mutations\332AQtable_name,row_key,predicate_filter,true_mutations,false_mutations,app_profile_id' + _BIGTABLE.methods_by_name['PingAndWarm']._options = None + _BIGTABLE.methods_by_name['PingAndWarm']._serialized_options = b'\202\323\344\223\002+\"&/v2/{name=projects/*/instances/*}:ping:\001*\212\323\344\223\0029\022%\n\004name\022\035{name=projects/*/instances/*}\022\020\n\016app_profile_id\332A\004name\332A\023name,app_profile_id' + _BIGTABLE.methods_by_name['ReadModifyWriteRow']._options = None + _BIGTABLE.methods_by_name['ReadModifyWriteRow']._serialized_options = b'\202\323\344\223\002H\"C/v2/{table_name=projects/*/instances/*/tables/*}:readModifyWriteRow:\001*\212\323\344\223\002N\022:\n\ntable_name\022,{table_name=projects/*/instances/*/tables/*}\022\020\n\016app_profile_id\332A\030table_name,row_key,rules\332A\'table_name,row_key,rules,app_profile_id' + _BIGTABLE.methods_by_name['GenerateInitialChangeStreamPartitions']._options = None + _BIGTABLE.methods_by_name['GenerateInitialChangeStreamPartitions']._serialized_options = b'\202\323\344\223\002[\"V/v2/{table_name=projects/*/instances/*/tables/*}:generateInitialChangeStreamPartitions:\001*\332A\ntable_name\332A\031table_name,app_profile_id' + _BIGTABLE.methods_by_name['ReadChangeStream']._options = None + _BIGTABLE.methods_by_name['ReadChangeStream']._serialized_options = b'\202\323\344\223\002F\"A/v2/{table_name=projects/*/instances/*/tables/*}:readChangeStream:\001*\332A\ntable_name\332A\031table_name,app_profile_id' + _READROWSREQUEST._serialized_start=392 + _READROWSREQUEST._serialized_end=792 + _READROWSREQUEST_REQUESTSTATSVIEW._serialized_start=690 + _READROWSREQUEST_REQUESTSTATSVIEW._serialized_end=792 + _READROWSRESPONSE._serialized_start=795 + _READROWSRESPONSE._serialized_end=1228 + _READROWSRESPONSE_CELLCHUNK._serialized_start=967 + _READROWSRESPONSE_CELLCHUNK._serialized_end=1228 + _SAMPLEROWKEYSREQUEST._serialized_start=1230 + _SAMPLEROWKEYSREQUEST._serialized_end=1340 + _SAMPLEROWKEYSRESPONSE._serialized_start=1342 + _SAMPLEROWKEYSRESPONSE._serialized_end=1404 + _MUTATEROWREQUEST._serialized_start=1407 + _MUTATEROWREQUEST._serialized_end=1589 + _MUTATEROWRESPONSE._serialized_start=1591 + _MUTATEROWRESPONSE._serialized_end=1610 + _MUTATEROWSREQUEST._serialized_start=1613 + _MUTATEROWSREQUEST._serialized_end=1867 + _MUTATEROWSREQUEST_ENTRY._serialized_start=1789 + _MUTATEROWSREQUEST_ENTRY._serialized_end=1867 + _MUTATEROWSRESPONSE._serialized_start=1870 + _MUTATEROWSRESPONSE._serialized_end=2013 + _MUTATEROWSRESPONSE_ENTRY._serialized_start=1955 + _MUTATEROWSRESPONSE_ENTRY._serialized_end=2013 + _CHECKANDMUTATEROWREQUEST._serialized_start=2016 + _CHECKANDMUTATEROWREQUEST._serialized_end=2318 + _CHECKANDMUTATEROWRESPONSE._serialized_start=2320 + _CHECKANDMUTATEROWRESPONSE._serialized_end=2374 + _PINGANDWARMREQUEST._serialized_start=2376 + _PINGANDWARMREQUEST._serialized_end=2481 + _PINGANDWARMRESPONSE._serialized_start=2483 + _PINGANDWARMRESPONSE._serialized_end=2504 + _READMODIFYWRITEROWREQUEST._serialized_start=2507 + _READMODIFYWRITEROWREQUEST._serialized_end=2705 + _READMODIFYWRITEROWRESPONSE._serialized_start=2707 + _READMODIFYWRITEROWRESPONSE._serialized_end=2773 + _GENERATEINITIALCHANGESTREAMPARTITIONSREQUEST._serialized_start=2776 + _GENERATEINITIALCHANGESTREAMPARTITIONSREQUEST._serialized_end=2910 + _GENERATEINITIALCHANGESTREAMPARTITIONSRESPONSE._serialized_start=2912 + _GENERATEINITIALCHANGESTREAMPARTITIONSRESPONSE._serialized_end=3015 + _READCHANGESTREAMREQUEST._serialized_start=3018 + _READCHANGESTREAMREQUEST._serialized_end=3429 + _READCHANGESTREAMRESPONSE._serialized_start=3432 + _READCHANGESTREAMRESPONSE._serialized_end=4691 + _READCHANGESTREAMRESPONSE_MUTATIONCHUNK._serialized_start=3700 + _READCHANGESTREAMRESPONSE_MUTATIONCHUNK._serialized_end=3944 + _READCHANGESTREAMRESPONSE_MUTATIONCHUNK_CHUNKINFO._serialized_start=3855 + _READCHANGESTREAMRESPONSE_MUTATIONCHUNK_CHUNKINFO._serialized_end=3944 + _READCHANGESTREAMRESPONSE_DATACHANGE._serialized_start=3947 + _READCHANGESTREAMRESPONSE_DATACHANGE._serialized_end=4401 + _READCHANGESTREAMRESPONSE_DATACHANGE_TYPE._serialized_start=4321 + _READCHANGESTREAMRESPONSE_DATACHANGE_TYPE._serialized_end=4401 + _READCHANGESTREAMRESPONSE_HEARTBEAT._serialized_start=4404 + _READCHANGESTREAMRESPONSE_HEARTBEAT._serialized_end=4549 + _READCHANGESTREAMRESPONSE_CLOSESTREAM._serialized_start=4551 + _READCHANGESTREAMRESPONSE_CLOSESTREAM._serialized_end=4674 + _BIGTABLE._serialized_start=4694 + _BIGTABLE._serialized_end=7853 +# @@protoc_insertion_point(module_scope) diff --git a/test_proxy/protos/bigtable_pb2_grpc.py b/test_proxy/protos/bigtable_pb2_grpc.py new file mode 100644 index 000000000..9ce87d869 --- /dev/null +++ b/test_proxy/protos/bigtable_pb2_grpc.py @@ -0,0 +1,363 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +import bigtable_pb2 as google_dot_bigtable_dot_v2_dot_bigtable__pb2 + + +class BigtableStub(object): + """Service for reading from and writing to existing Bigtable tables. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.ReadRows = channel.unary_stream( + '/google.bigtable.v2.Bigtable/ReadRows', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadRowsRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadRowsResponse.FromString, + ) + self.SampleRowKeys = channel.unary_stream( + '/google.bigtable.v2.Bigtable/SampleRowKeys', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.SampleRowKeysRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.SampleRowKeysResponse.FromString, + ) + self.MutateRow = channel.unary_unary( + '/google.bigtable.v2.Bigtable/MutateRow', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowResponse.FromString, + ) + self.MutateRows = channel.unary_stream( + '/google.bigtable.v2.Bigtable/MutateRows', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowsRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowsResponse.FromString, + ) + self.CheckAndMutateRow = channel.unary_unary( + '/google.bigtable.v2.Bigtable/CheckAndMutateRow', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.CheckAndMutateRowRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.CheckAndMutateRowResponse.FromString, + ) + self.PingAndWarm = channel.unary_unary( + '/google.bigtable.v2.Bigtable/PingAndWarm', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.PingAndWarmRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.PingAndWarmResponse.FromString, + ) + self.ReadModifyWriteRow = channel.unary_unary( + '/google.bigtable.v2.Bigtable/ReadModifyWriteRow', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadModifyWriteRowRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadModifyWriteRowResponse.FromString, + ) + self.GenerateInitialChangeStreamPartitions = channel.unary_stream( + '/google.bigtable.v2.Bigtable/GenerateInitialChangeStreamPartitions', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.GenerateInitialChangeStreamPartitionsRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.GenerateInitialChangeStreamPartitionsResponse.FromString, + ) + self.ReadChangeStream = channel.unary_stream( + '/google.bigtable.v2.Bigtable/ReadChangeStream', + request_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadChangeStreamRequest.SerializeToString, + response_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadChangeStreamResponse.FromString, + ) + + +class BigtableServicer(object): + """Service for reading from and writing to existing Bigtable tables. + """ + + def ReadRows(self, request, context): + """Streams back the contents of all requested rows in key order, optionally + applying the same Reader filter to each. Depending on their size, + rows and cells may be broken up across multiple responses, but + atomicity of each row will still be preserved. See the + ReadRowsResponse documentation for details. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def SampleRowKeys(self, request, context): + """Returns a sample of row keys in the table. The returned row keys will + delimit contiguous sections of the table of approximately equal size, + which can be used to break up the data for distributed tasks like + mapreduces. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def MutateRow(self, request, context): + """Mutates a row atomically. Cells already present in the row are left + unchanged unless explicitly changed by `mutation`. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def MutateRows(self, request, context): + """Mutates multiple rows in a batch. Each individual row is mutated + atomically as in MutateRow, but the entire batch is not executed + atomically. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def CheckAndMutateRow(self, request, context): + """Mutates a row atomically based on the output of a predicate Reader filter. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def PingAndWarm(self, request, context): + """Warm up associated instance metadata for this connection. + This call is not required but may be useful for connection keep-alive. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ReadModifyWriteRow(self, request, context): + """Modifies a row atomically on the server. The method reads the latest + existing timestamp and value from the specified columns and writes a new + entry based on pre-defined read/modify/write rules. The new value for the + timestamp is the greater of the existing timestamp or the current server + time. The method returns the new contents of all modified cells. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def GenerateInitialChangeStreamPartitions(self, request, context): + """NOTE: This API is intended to be used by Apache Beam BigtableIO. + Returns the current list of partitions that make up the table's + change stream. The union of partitions will cover the entire keyspace. + Partitions can be read with `ReadChangeStream`. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ReadChangeStream(self, request, context): + """NOTE: This API is intended to be used by Apache Beam BigtableIO. + Reads changes from a table's change stream. Changes will + reflect both user-initiated mutations and mutations that are caused by + garbage collection. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_BigtableServicer_to_server(servicer, server): + rpc_method_handlers = { + 'ReadRows': grpc.unary_stream_rpc_method_handler( + servicer.ReadRows, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadRowsRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadRowsResponse.SerializeToString, + ), + 'SampleRowKeys': grpc.unary_stream_rpc_method_handler( + servicer.SampleRowKeys, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.SampleRowKeysRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.SampleRowKeysResponse.SerializeToString, + ), + 'MutateRow': grpc.unary_unary_rpc_method_handler( + servicer.MutateRow, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowResponse.SerializeToString, + ), + 'MutateRows': grpc.unary_stream_rpc_method_handler( + servicer.MutateRows, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowsRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowsResponse.SerializeToString, + ), + 'CheckAndMutateRow': grpc.unary_unary_rpc_method_handler( + servicer.CheckAndMutateRow, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.CheckAndMutateRowRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.CheckAndMutateRowResponse.SerializeToString, + ), + 'PingAndWarm': grpc.unary_unary_rpc_method_handler( + servicer.PingAndWarm, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.PingAndWarmRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.PingAndWarmResponse.SerializeToString, + ), + 'ReadModifyWriteRow': grpc.unary_unary_rpc_method_handler( + servicer.ReadModifyWriteRow, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadModifyWriteRowRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadModifyWriteRowResponse.SerializeToString, + ), + 'GenerateInitialChangeStreamPartitions': grpc.unary_stream_rpc_method_handler( + servicer.GenerateInitialChangeStreamPartitions, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.GenerateInitialChangeStreamPartitionsRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.GenerateInitialChangeStreamPartitionsResponse.SerializeToString, + ), + 'ReadChangeStream': grpc.unary_stream_rpc_method_handler( + servicer.ReadChangeStream, + request_deserializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadChangeStreamRequest.FromString, + response_serializer=google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadChangeStreamResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'google.bigtable.v2.Bigtable', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class Bigtable(object): + """Service for reading from and writing to existing Bigtable tables. + """ + + @staticmethod + def ReadRows(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream(request, target, '/google.bigtable.v2.Bigtable/ReadRows', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadRowsRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadRowsResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def SampleRowKeys(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream(request, target, '/google.bigtable.v2.Bigtable/SampleRowKeys', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.SampleRowKeysRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.SampleRowKeysResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def MutateRow(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.v2.Bigtable/MutateRow', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def MutateRows(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream(request, target, '/google.bigtable.v2.Bigtable/MutateRows', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowsRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.MutateRowsResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def CheckAndMutateRow(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.v2.Bigtable/CheckAndMutateRow', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.CheckAndMutateRowRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.CheckAndMutateRowResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def PingAndWarm(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.v2.Bigtable/PingAndWarm', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.PingAndWarmRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.PingAndWarmResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def ReadModifyWriteRow(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.v2.Bigtable/ReadModifyWriteRow', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadModifyWriteRowRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadModifyWriteRowResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def GenerateInitialChangeStreamPartitions(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream(request, target, '/google.bigtable.v2.Bigtable/GenerateInitialChangeStreamPartitions', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.GenerateInitialChangeStreamPartitionsRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.GenerateInitialChangeStreamPartitionsResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def ReadChangeStream(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream(request, target, '/google.bigtable.v2.Bigtable/ReadChangeStream', + google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadChangeStreamRequest.SerializeToString, + google_dot_bigtable_dot_v2_dot_bigtable__pb2.ReadChangeStreamResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/test_proxy/protos/data_pb2.py b/test_proxy/protos/data_pb2.py new file mode 100644 index 000000000..fff212034 --- /dev/null +++ b/test_proxy/protos/data_pb2.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: google/bigtable/v2/data.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1dgoogle/bigtable/v2/data.proto\x12\x12google.bigtable.v2\"@\n\x03Row\x12\x0b\n\x03key\x18\x01 \x01(\x0c\x12,\n\x08\x66\x61milies\x18\x02 \x03(\x0b\x32\x1a.google.bigtable.v2.Family\"C\n\x06\x46\x61mily\x12\x0c\n\x04name\x18\x01 \x01(\t\x12+\n\x07\x63olumns\x18\x02 \x03(\x0b\x32\x1a.google.bigtable.v2.Column\"D\n\x06\x43olumn\x12\x11\n\tqualifier\x18\x01 \x01(\x0c\x12\'\n\x05\x63\x65lls\x18\x02 \x03(\x0b\x32\x18.google.bigtable.v2.Cell\"?\n\x04\x43\x65ll\x12\x18\n\x10timestamp_micros\x18\x01 \x01(\x03\x12\r\n\x05value\x18\x02 \x01(\x0c\x12\x0e\n\x06labels\x18\x03 \x03(\t\"\x8a\x01\n\x08RowRange\x12\x1a\n\x10start_key_closed\x18\x01 \x01(\x0cH\x00\x12\x18\n\x0estart_key_open\x18\x02 \x01(\x0cH\x00\x12\x16\n\x0c\x65nd_key_open\x18\x03 \x01(\x0cH\x01\x12\x18\n\x0e\x65nd_key_closed\x18\x04 \x01(\x0cH\x01\x42\x0b\n\tstart_keyB\t\n\x07\x65nd_key\"L\n\x06RowSet\x12\x10\n\x08row_keys\x18\x01 \x03(\x0c\x12\x30\n\nrow_ranges\x18\x02 \x03(\x0b\x32\x1c.google.bigtable.v2.RowRange\"\xc6\x01\n\x0b\x43olumnRange\x12\x13\n\x0b\x66\x61mily_name\x18\x01 \x01(\t\x12 \n\x16start_qualifier_closed\x18\x02 \x01(\x0cH\x00\x12\x1e\n\x14start_qualifier_open\x18\x03 \x01(\x0cH\x00\x12\x1e\n\x14\x65nd_qualifier_closed\x18\x04 \x01(\x0cH\x01\x12\x1c\n\x12\x65nd_qualifier_open\x18\x05 \x01(\x0cH\x01\x42\x11\n\x0fstart_qualifierB\x0f\n\rend_qualifier\"N\n\x0eTimestampRange\x12\x1e\n\x16start_timestamp_micros\x18\x01 \x01(\x03\x12\x1c\n\x14\x65nd_timestamp_micros\x18\x02 \x01(\x03\"\x98\x01\n\nValueRange\x12\x1c\n\x12start_value_closed\x18\x01 \x01(\x0cH\x00\x12\x1a\n\x10start_value_open\x18\x02 \x01(\x0cH\x00\x12\x1a\n\x10\x65nd_value_closed\x18\x03 \x01(\x0cH\x01\x12\x18\n\x0e\x65nd_value_open\x18\x04 \x01(\x0cH\x01\x42\r\n\x0bstart_valueB\x0b\n\tend_value\"\xdf\x08\n\tRowFilter\x12\x34\n\x05\x63hain\x18\x01 \x01(\x0b\x32#.google.bigtable.v2.RowFilter.ChainH\x00\x12>\n\ninterleave\x18\x02 \x01(\x0b\x32(.google.bigtable.v2.RowFilter.InterleaveH\x00\x12<\n\tcondition\x18\x03 \x01(\x0b\x32\'.google.bigtable.v2.RowFilter.ConditionH\x00\x12\x0e\n\x04sink\x18\x10 \x01(\x08H\x00\x12\x19\n\x0fpass_all_filter\x18\x11 \x01(\x08H\x00\x12\x1a\n\x10\x62lock_all_filter\x18\x12 \x01(\x08H\x00\x12\x1e\n\x14row_key_regex_filter\x18\x04 \x01(\x0cH\x00\x12\x1b\n\x11row_sample_filter\x18\x0e \x01(\x01H\x00\x12\"\n\x18\x66\x61mily_name_regex_filter\x18\x05 \x01(\tH\x00\x12\'\n\x1d\x63olumn_qualifier_regex_filter\x18\x06 \x01(\x0cH\x00\x12>\n\x13\x63olumn_range_filter\x18\x07 \x01(\x0b\x32\x1f.google.bigtable.v2.ColumnRangeH\x00\x12\x44\n\x16timestamp_range_filter\x18\x08 \x01(\x0b\x32\".google.bigtable.v2.TimestampRangeH\x00\x12\x1c\n\x12value_regex_filter\x18\t \x01(\x0cH\x00\x12<\n\x12value_range_filter\x18\x0f \x01(\x0b\x32\x1e.google.bigtable.v2.ValueRangeH\x00\x12%\n\x1b\x63\x65lls_per_row_offset_filter\x18\n \x01(\x05H\x00\x12$\n\x1a\x63\x65lls_per_row_limit_filter\x18\x0b \x01(\x05H\x00\x12\'\n\x1d\x63\x65lls_per_column_limit_filter\x18\x0c \x01(\x05H\x00\x12!\n\x17strip_value_transformer\x18\r \x01(\x08H\x00\x12!\n\x17\x61pply_label_transformer\x18\x13 \x01(\tH\x00\x1a\x37\n\x05\x43hain\x12.\n\x07\x66ilters\x18\x01 \x03(\x0b\x32\x1d.google.bigtable.v2.RowFilter\x1a<\n\nInterleave\x12.\n\x07\x66ilters\x18\x01 \x03(\x0b\x32\x1d.google.bigtable.v2.RowFilter\x1a\xad\x01\n\tCondition\x12\x37\n\x10predicate_filter\x18\x01 \x01(\x0b\x32\x1d.google.bigtable.v2.RowFilter\x12\x32\n\x0btrue_filter\x18\x02 \x01(\x0b\x32\x1d.google.bigtable.v2.RowFilter\x12\x33\n\x0c\x66\x61lse_filter\x18\x03 \x01(\x0b\x32\x1d.google.bigtable.v2.RowFilterB\x08\n\x06\x66ilter\"\xc9\x04\n\x08Mutation\x12\x38\n\x08set_cell\x18\x01 \x01(\x0b\x32$.google.bigtable.v2.Mutation.SetCellH\x00\x12K\n\x12\x64\x65lete_from_column\x18\x02 \x01(\x0b\x32-.google.bigtable.v2.Mutation.DeleteFromColumnH\x00\x12K\n\x12\x64\x65lete_from_family\x18\x03 \x01(\x0b\x32-.google.bigtable.v2.Mutation.DeleteFromFamilyH\x00\x12\x45\n\x0f\x64\x65lete_from_row\x18\x04 \x01(\x0b\x32*.google.bigtable.v2.Mutation.DeleteFromRowH\x00\x1a\x61\n\x07SetCell\x12\x13\n\x0b\x66\x61mily_name\x18\x01 \x01(\t\x12\x18\n\x10\x63olumn_qualifier\x18\x02 \x01(\x0c\x12\x18\n\x10timestamp_micros\x18\x03 \x01(\x03\x12\r\n\x05value\x18\x04 \x01(\x0c\x1ay\n\x10\x44\x65leteFromColumn\x12\x13\n\x0b\x66\x61mily_name\x18\x01 \x01(\t\x12\x18\n\x10\x63olumn_qualifier\x18\x02 \x01(\x0c\x12\x36\n\ntime_range\x18\x03 \x01(\x0b\x32\".google.bigtable.v2.TimestampRange\x1a\'\n\x10\x44\x65leteFromFamily\x12\x13\n\x0b\x66\x61mily_name\x18\x01 \x01(\t\x1a\x0f\n\rDeleteFromRowB\n\n\x08mutation\"\x80\x01\n\x13ReadModifyWriteRule\x12\x13\n\x0b\x66\x61mily_name\x18\x01 \x01(\t\x12\x18\n\x10\x63olumn_qualifier\x18\x02 \x01(\x0c\x12\x16\n\x0c\x61ppend_value\x18\x03 \x01(\x0cH\x00\x12\x1a\n\x10increment_amount\x18\x04 \x01(\x03H\x00\x42\x06\n\x04rule\"B\n\x0fStreamPartition\x12/\n\trow_range\x18\x01 \x01(\x0b\x32\x1c.google.bigtable.v2.RowRange\"W\n\x18StreamContinuationTokens\x12;\n\x06tokens\x18\x01 \x03(\x0b\x32+.google.bigtable.v2.StreamContinuationToken\"`\n\x17StreamContinuationToken\x12\x36\n\tpartition\x18\x01 \x01(\x0b\x32#.google.bigtable.v2.StreamPartition\x12\r\n\x05token\x18\x02 \x01(\tB\xb5\x01\n\x16\x63om.google.bigtable.v2B\tDataProtoP\x01Z:google.golang.org/genproto/googleapis/bigtable/v2;bigtable\xaa\x02\x18Google.Cloud.Bigtable.V2\xca\x02\x18Google\\Cloud\\Bigtable\\V2\xea\x02\x1bGoogle::Cloud::Bigtable::V2b\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.bigtable.v2.data_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'\n\026com.google.bigtable.v2B\tDataProtoP\001Z:google.golang.org/genproto/googleapis/bigtable/v2;bigtable\252\002\030Google.Cloud.Bigtable.V2\312\002\030Google\\Cloud\\Bigtable\\V2\352\002\033Google::Cloud::Bigtable::V2' + _ROW._serialized_start=53 + _ROW._serialized_end=117 + _FAMILY._serialized_start=119 + _FAMILY._serialized_end=186 + _COLUMN._serialized_start=188 + _COLUMN._serialized_end=256 + _CELL._serialized_start=258 + _CELL._serialized_end=321 + _ROWRANGE._serialized_start=324 + _ROWRANGE._serialized_end=462 + _ROWSET._serialized_start=464 + _ROWSET._serialized_end=540 + _COLUMNRANGE._serialized_start=543 + _COLUMNRANGE._serialized_end=741 + _TIMESTAMPRANGE._serialized_start=743 + _TIMESTAMPRANGE._serialized_end=821 + _VALUERANGE._serialized_start=824 + _VALUERANGE._serialized_end=976 + _ROWFILTER._serialized_start=979 + _ROWFILTER._serialized_end=2098 + _ROWFILTER_CHAIN._serialized_start=1795 + _ROWFILTER_CHAIN._serialized_end=1850 + _ROWFILTER_INTERLEAVE._serialized_start=1852 + _ROWFILTER_INTERLEAVE._serialized_end=1912 + _ROWFILTER_CONDITION._serialized_start=1915 + _ROWFILTER_CONDITION._serialized_end=2088 + _MUTATION._serialized_start=2101 + _MUTATION._serialized_end=2686 + _MUTATION_SETCELL._serialized_start=2396 + _MUTATION_SETCELL._serialized_end=2493 + _MUTATION_DELETEFROMCOLUMN._serialized_start=2495 + _MUTATION_DELETEFROMCOLUMN._serialized_end=2616 + _MUTATION_DELETEFROMFAMILY._serialized_start=2618 + _MUTATION_DELETEFROMFAMILY._serialized_end=2657 + _MUTATION_DELETEFROMROW._serialized_start=2659 + _MUTATION_DELETEFROMROW._serialized_end=2674 + _READMODIFYWRITERULE._serialized_start=2689 + _READMODIFYWRITERULE._serialized_end=2817 + _STREAMPARTITION._serialized_start=2819 + _STREAMPARTITION._serialized_end=2885 + _STREAMCONTINUATIONTOKENS._serialized_start=2887 + _STREAMCONTINUATIONTOKENS._serialized_end=2974 + _STREAMCONTINUATIONTOKEN._serialized_start=2976 + _STREAMCONTINUATIONTOKEN._serialized_end=3072 +# @@protoc_insertion_point(module_scope) diff --git a/test_proxy/protos/data_pb2_grpc.py b/test_proxy/protos/data_pb2_grpc.py new file mode 100644 index 000000000..2daafffeb --- /dev/null +++ b/test_proxy/protos/data_pb2_grpc.py @@ -0,0 +1,4 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + diff --git a/test_proxy/protos/request_stats_pb2.py b/test_proxy/protos/request_stats_pb2.py new file mode 100644 index 000000000..95fcc6e0f --- /dev/null +++ b/test_proxy/protos/request_stats_pb2.py @@ -0,0 +1,33 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: google/bigtable/v2/request_stats.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import duration_pb2 as google_dot_protobuf_dot_duration__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n&google/bigtable/v2/request_stats.proto\x12\x12google.bigtable.v2\x1a\x1egoogle/protobuf/duration.proto\"\x82\x01\n\x12ReadIterationStats\x12\x17\n\x0frows_seen_count\x18\x01 \x01(\x03\x12\x1b\n\x13rows_returned_count\x18\x02 \x01(\x03\x12\x18\n\x10\x63\x65lls_seen_count\x18\x03 \x01(\x03\x12\x1c\n\x14\x63\x65lls_returned_count\x18\x04 \x01(\x03\"Q\n\x13RequestLatencyStats\x12:\n\x17\x66rontend_server_latency\x18\x01 \x01(\x0b\x32\x19.google.protobuf.Duration\"\xa1\x01\n\x11\x46ullReadStatsView\x12\x44\n\x14read_iteration_stats\x18\x01 \x01(\x0b\x32&.google.bigtable.v2.ReadIterationStats\x12\x46\n\x15request_latency_stats\x18\x02 \x01(\x0b\x32\'.google.bigtable.v2.RequestLatencyStats\"c\n\x0cRequestStats\x12\x45\n\x14\x66ull_read_stats_view\x18\x01 \x01(\x0b\x32%.google.bigtable.v2.FullReadStatsViewH\x00\x42\x0c\n\nstats_viewB\xbd\x01\n\x16\x63om.google.bigtable.v2B\x11RequestStatsProtoP\x01Z:google.golang.org/genproto/googleapis/bigtable/v2;bigtable\xaa\x02\x18Google.Cloud.Bigtable.V2\xca\x02\x18Google\\Cloud\\Bigtable\\V2\xea\x02\x1bGoogle::Cloud::Bigtable::V2b\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'google.bigtable.v2.request_stats_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'\n\026com.google.bigtable.v2B\021RequestStatsProtoP\001Z:google.golang.org/genproto/googleapis/bigtable/v2;bigtable\252\002\030Google.Cloud.Bigtable.V2\312\002\030Google\\Cloud\\Bigtable\\V2\352\002\033Google::Cloud::Bigtable::V2' + _READITERATIONSTATS._serialized_start=95 + _READITERATIONSTATS._serialized_end=225 + _REQUESTLATENCYSTATS._serialized_start=227 + _REQUESTLATENCYSTATS._serialized_end=308 + _FULLREADSTATSVIEW._serialized_start=311 + _FULLREADSTATSVIEW._serialized_end=472 + _REQUESTSTATS._serialized_start=474 + _REQUESTSTATS._serialized_end=573 +# @@protoc_insertion_point(module_scope) diff --git a/test_proxy/protos/request_stats_pb2_grpc.py b/test_proxy/protos/request_stats_pb2_grpc.py new file mode 100644 index 000000000..2daafffeb --- /dev/null +++ b/test_proxy/protos/request_stats_pb2_grpc.py @@ -0,0 +1,4 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + diff --git a/test_proxy/protos/test_proxy_pb2.py b/test_proxy/protos/test_proxy_pb2.py new file mode 100644 index 000000000..8c7817b14 --- /dev/null +++ b/test_proxy/protos/test_proxy_pb2.py @@ -0,0 +1,71 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: test_proxy.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.api import client_pb2 as google_dot_api_dot_client__pb2 +import bigtable_pb2 as google_dot_bigtable_dot_v2_dot_bigtable__pb2 +import data_pb2 as google_dot_bigtable_dot_v2_dot_data__pb2 +from google.protobuf import duration_pb2 as google_dot_protobuf_dot_duration__pb2 +from google.rpc import status_pb2 as google_dot_rpc_dot_status__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x10test_proxy.proto\x12\x19google.bigtable.testproxy\x1a\x17google/api/client.proto\x1a!google/bigtable/v2/bigtable.proto\x1a\x1dgoogle/bigtable/v2/data.proto\x1a\x1egoogle/protobuf/duration.proto\x1a\x17google/rpc/status.proto\"\xb8\x01\n\x13\x43reateClientRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x13\n\x0b\x64\x61ta_target\x18\x02 \x01(\t\x12\x12\n\nproject_id\x18\x03 \x01(\t\x12\x13\n\x0binstance_id\x18\x04 \x01(\t\x12\x16\n\x0e\x61pp_profile_id\x18\x05 \x01(\t\x12\x38\n\x15per_operation_timeout\x18\x06 \x01(\x0b\x32\x19.google.protobuf.Duration\"\x16\n\x14\x43reateClientResponse\"\'\n\x12\x43loseClientRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\"\x15\n\x13\x43loseClientResponse\"(\n\x13RemoveClientRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\"\x16\n\x14RemoveClientResponse\"w\n\x0eReadRowRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x12\n\ntable_name\x18\x04 \x01(\t\x12\x0f\n\x07row_key\x18\x02 \x01(\t\x12-\n\x06\x66ilter\x18\x03 \x01(\x0b\x32\x1d.google.bigtable.v2.RowFilter\"U\n\tRowResult\x12\"\n\x06status\x18\x01 \x01(\x0b\x32\x12.google.rpc.Status\x12$\n\x03row\x18\x02 \x01(\x0b\x32\x17.google.bigtable.v2.Row\"u\n\x0fReadRowsRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x34\n\x07request\x18\x02 \x01(\x0b\x32#.google.bigtable.v2.ReadRowsRequest\x12\x19\n\x11\x63\x61ncel_after_rows\x18\x03 \x01(\x05\"V\n\nRowsResult\x12\"\n\x06status\x18\x01 \x01(\x0b\x32\x12.google.rpc.Status\x12$\n\x03row\x18\x02 \x03(\x0b\x32\x17.google.bigtable.v2.Row\"\\\n\x10MutateRowRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x35\n\x07request\x18\x02 \x01(\x0b\x32$.google.bigtable.v2.MutateRowRequest\"5\n\x0fMutateRowResult\x12\"\n\x06status\x18\x01 \x01(\x0b\x32\x12.google.rpc.Status\"^\n\x11MutateRowsRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x36\n\x07request\x18\x02 \x01(\x0b\x32%.google.bigtable.v2.MutateRowsRequest\"s\n\x10MutateRowsResult\x12\"\n\x06status\x18\x01 \x01(\x0b\x32\x12.google.rpc.Status\x12;\n\x05\x65ntry\x18\x02 \x03(\x0b\x32,.google.bigtable.v2.MutateRowsResponse.Entry\"l\n\x18\x43heckAndMutateRowRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12=\n\x07request\x18\x02 \x01(\x0b\x32,.google.bigtable.v2.CheckAndMutateRowRequest\"|\n\x17\x43heckAndMutateRowResult\x12\"\n\x06status\x18\x01 \x01(\x0b\x32\x12.google.rpc.Status\x12=\n\x06result\x18\x02 \x01(\x0b\x32-.google.bigtable.v2.CheckAndMutateRowResponse\"d\n\x14SampleRowKeysRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x39\n\x07request\x18\x02 \x01(\x0b\x32(.google.bigtable.v2.SampleRowKeysRequest\"t\n\x13SampleRowKeysResult\x12\"\n\x06status\x18\x01 \x01(\x0b\x32\x12.google.rpc.Status\x12\x39\n\x06sample\x18\x02 \x03(\x0b\x32).google.bigtable.v2.SampleRowKeysResponse\"n\n\x19ReadModifyWriteRowRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12>\n\x07request\x18\x02 \x01(\x0b\x32-.google.bigtable.v2.ReadModifyWriteRowRequest2\xa4\t\n\x18\x43loudBigtableV2TestProxy\x12q\n\x0c\x43reateClient\x12..google.bigtable.testproxy.CreateClientRequest\x1a/.google.bigtable.testproxy.CreateClientResponse\"\x00\x12n\n\x0b\x43loseClient\x12-.google.bigtable.testproxy.CloseClientRequest\x1a..google.bigtable.testproxy.CloseClientResponse\"\x00\x12q\n\x0cRemoveClient\x12..google.bigtable.testproxy.RemoveClientRequest\x1a/.google.bigtable.testproxy.RemoveClientResponse\"\x00\x12\\\n\x07ReadRow\x12).google.bigtable.testproxy.ReadRowRequest\x1a$.google.bigtable.testproxy.RowResult\"\x00\x12_\n\x08ReadRows\x12*.google.bigtable.testproxy.ReadRowsRequest\x1a%.google.bigtable.testproxy.RowsResult\"\x00\x12\x66\n\tMutateRow\x12+.google.bigtable.testproxy.MutateRowRequest\x1a*.google.bigtable.testproxy.MutateRowResult\"\x00\x12m\n\x0e\x42ulkMutateRows\x12,.google.bigtable.testproxy.MutateRowsRequest\x1a+.google.bigtable.testproxy.MutateRowsResult\"\x00\x12~\n\x11\x43heckAndMutateRow\x12\x33.google.bigtable.testproxy.CheckAndMutateRowRequest\x1a\x32.google.bigtable.testproxy.CheckAndMutateRowResult\"\x00\x12r\n\rSampleRowKeys\x12/.google.bigtable.testproxy.SampleRowKeysRequest\x1a..google.bigtable.testproxy.SampleRowKeysResult\"\x00\x12r\n\x12ReadModifyWriteRow\x12\x34.google.bigtable.testproxy.ReadModifyWriteRowRequest\x1a$.google.bigtable.testproxy.RowResult\"\x00\x1a\x34\xca\x41\x31\x62igtable-test-proxy-not-accessible.googleapis.comB6\n#com.google.cloud.bigtable.testproxyP\x01Z\r./testproxypbb\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'test_proxy_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'\n#com.google.cloud.bigtable.testproxyP\001Z\r./testproxypb' + _CLOUDBIGTABLEV2TESTPROXY._options = None + _CLOUDBIGTABLEV2TESTPROXY._serialized_options = b'\312A1bigtable-test-proxy-not-accessible.googleapis.com' + _CREATECLIENTREQUEST._serialized_start=196 + _CREATECLIENTREQUEST._serialized_end=380 + _CREATECLIENTRESPONSE._serialized_start=382 + _CREATECLIENTRESPONSE._serialized_end=404 + _CLOSECLIENTREQUEST._serialized_start=406 + _CLOSECLIENTREQUEST._serialized_end=445 + _CLOSECLIENTRESPONSE._serialized_start=447 + _CLOSECLIENTRESPONSE._serialized_end=468 + _REMOVECLIENTREQUEST._serialized_start=470 + _REMOVECLIENTREQUEST._serialized_end=510 + _REMOVECLIENTRESPONSE._serialized_start=512 + _REMOVECLIENTRESPONSE._serialized_end=534 + _READROWREQUEST._serialized_start=536 + _READROWREQUEST._serialized_end=655 + _ROWRESULT._serialized_start=657 + _ROWRESULT._serialized_end=742 + _READROWSREQUEST._serialized_start=744 + _READROWSREQUEST._serialized_end=861 + _ROWSRESULT._serialized_start=863 + _ROWSRESULT._serialized_end=949 + _MUTATEROWREQUEST._serialized_start=951 + _MUTATEROWREQUEST._serialized_end=1043 + _MUTATEROWRESULT._serialized_start=1045 + _MUTATEROWRESULT._serialized_end=1098 + _MUTATEROWSREQUEST._serialized_start=1100 + _MUTATEROWSREQUEST._serialized_end=1194 + _MUTATEROWSRESULT._serialized_start=1196 + _MUTATEROWSRESULT._serialized_end=1311 + _CHECKANDMUTATEROWREQUEST._serialized_start=1313 + _CHECKANDMUTATEROWREQUEST._serialized_end=1421 + _CHECKANDMUTATEROWRESULT._serialized_start=1423 + _CHECKANDMUTATEROWRESULT._serialized_end=1547 + _SAMPLEROWKEYSREQUEST._serialized_start=1549 + _SAMPLEROWKEYSREQUEST._serialized_end=1649 + _SAMPLEROWKEYSRESULT._serialized_start=1651 + _SAMPLEROWKEYSRESULT._serialized_end=1767 + _READMODIFYWRITEROWREQUEST._serialized_start=1769 + _READMODIFYWRITEROWREQUEST._serialized_end=1879 + _CLOUDBIGTABLEV2TESTPROXY._serialized_start=1882 + _CLOUDBIGTABLEV2TESTPROXY._serialized_end=3070 +# @@protoc_insertion_point(module_scope) diff --git a/test_proxy/protos/test_proxy_pb2_grpc.py b/test_proxy/protos/test_proxy_pb2_grpc.py new file mode 100644 index 000000000..60214a584 --- /dev/null +++ b/test_proxy/protos/test_proxy_pb2_grpc.py @@ -0,0 +1,433 @@ +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc + +import test_proxy_pb2 as test__proxy__pb2 + + +class CloudBigtableV2TestProxyStub(object): + """Note that all RPCs are unary, even when the equivalent client binding call + may be streaming. This is an intentional simplification. + + Most methods have sync (default) and async variants. For async variants, + the proxy is expected to perform the async operation, then wait for results + before delivering them back to the driver client. + + Operations that may have interesting concurrency characteristics are + represented explicitly in the API (see ReadRowsRequest.cancel_after_rows). + We include such operations only when they can be meaningfully performed + through client bindings. + + Users should generally avoid setting deadlines for requests to the Proxy + because operations are not cancelable. If the deadline is set anyway, please + understand that the underlying operation will continue to be executed even + after the deadline expires. + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.CreateClient = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/CreateClient', + request_serializer=test__proxy__pb2.CreateClientRequest.SerializeToString, + response_deserializer=test__proxy__pb2.CreateClientResponse.FromString, + ) + self.CloseClient = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/CloseClient', + request_serializer=test__proxy__pb2.CloseClientRequest.SerializeToString, + response_deserializer=test__proxy__pb2.CloseClientResponse.FromString, + ) + self.RemoveClient = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/RemoveClient', + request_serializer=test__proxy__pb2.RemoveClientRequest.SerializeToString, + response_deserializer=test__proxy__pb2.RemoveClientResponse.FromString, + ) + self.ReadRow = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/ReadRow', + request_serializer=test__proxy__pb2.ReadRowRequest.SerializeToString, + response_deserializer=test__proxy__pb2.RowResult.FromString, + ) + self.ReadRows = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/ReadRows', + request_serializer=test__proxy__pb2.ReadRowsRequest.SerializeToString, + response_deserializer=test__proxy__pb2.RowsResult.FromString, + ) + self.MutateRow = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/MutateRow', + request_serializer=test__proxy__pb2.MutateRowRequest.SerializeToString, + response_deserializer=test__proxy__pb2.MutateRowResult.FromString, + ) + self.BulkMutateRows = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/BulkMutateRows', + request_serializer=test__proxy__pb2.MutateRowsRequest.SerializeToString, + response_deserializer=test__proxy__pb2.MutateRowsResult.FromString, + ) + self.CheckAndMutateRow = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/CheckAndMutateRow', + request_serializer=test__proxy__pb2.CheckAndMutateRowRequest.SerializeToString, + response_deserializer=test__proxy__pb2.CheckAndMutateRowResult.FromString, + ) + self.SampleRowKeys = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/SampleRowKeys', + request_serializer=test__proxy__pb2.SampleRowKeysRequest.SerializeToString, + response_deserializer=test__proxy__pb2.SampleRowKeysResult.FromString, + ) + self.ReadModifyWriteRow = channel.unary_unary( + '/google.bigtable.testproxy.CloudBigtableV2TestProxy/ReadModifyWriteRow', + request_serializer=test__proxy__pb2.ReadModifyWriteRowRequest.SerializeToString, + response_deserializer=test__proxy__pb2.RowResult.FromString, + ) + + +class CloudBigtableV2TestProxyServicer(object): + """Note that all RPCs are unary, even when the equivalent client binding call + may be streaming. This is an intentional simplification. + + Most methods have sync (default) and async variants. For async variants, + the proxy is expected to perform the async operation, then wait for results + before delivering them back to the driver client. + + Operations that may have interesting concurrency characteristics are + represented explicitly in the API (see ReadRowsRequest.cancel_after_rows). + We include such operations only when they can be meaningfully performed + through client bindings. + + Users should generally avoid setting deadlines for requests to the Proxy + because operations are not cancelable. If the deadline is set anyway, please + understand that the underlying operation will continue to be executed even + after the deadline expires. + """ + + def CreateClient(self, request, context): + """Client management: + + Creates a client in the proxy. + Each client has its own dedicated channel(s), and can be used concurrently + and independently with other clients. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def CloseClient(self, request, context): + """Closes a client in the proxy, making it not accept new requests. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def RemoveClient(self, request, context): + """Removes a client in the proxy, making it inaccessible. Client closing + should be done by CloseClient() separately. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ReadRow(self, request, context): + """Bigtable operations: for each operation, you should use the synchronous or + asynchronous variant of the client method based on the `use_async_method` + setting of the client instance. For starters, you can choose to implement + one variant, and return UNIMPLEMENTED status for the other. + + Reads a row with the client instance. + The result row may not be present in the response. + Callers should check for it (e.g. calling has_row() in C++). + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ReadRows(self, request, context): + """Reads rows with the client instance. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def MutateRow(self, request, context): + """Writes a row with the client instance. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def BulkMutateRows(self, request, context): + """Writes multiple rows with the client instance. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def CheckAndMutateRow(self, request, context): + """Performs a check-and-mutate-row operation with the client instance. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def SampleRowKeys(self, request, context): + """Obtains a row key sampling with the client instance. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ReadModifyWriteRow(self, request, context): + """Performs a read-modify-write operation with the client. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_CloudBigtableV2TestProxyServicer_to_server(servicer, server): + rpc_method_handlers = { + 'CreateClient': grpc.unary_unary_rpc_method_handler( + servicer.CreateClient, + request_deserializer=test__proxy__pb2.CreateClientRequest.FromString, + response_serializer=test__proxy__pb2.CreateClientResponse.SerializeToString, + ), + 'CloseClient': grpc.unary_unary_rpc_method_handler( + servicer.CloseClient, + request_deserializer=test__proxy__pb2.CloseClientRequest.FromString, + response_serializer=test__proxy__pb2.CloseClientResponse.SerializeToString, + ), + 'RemoveClient': grpc.unary_unary_rpc_method_handler( + servicer.RemoveClient, + request_deserializer=test__proxy__pb2.RemoveClientRequest.FromString, + response_serializer=test__proxy__pb2.RemoveClientResponse.SerializeToString, + ), + 'ReadRow': grpc.unary_unary_rpc_method_handler( + servicer.ReadRow, + request_deserializer=test__proxy__pb2.ReadRowRequest.FromString, + response_serializer=test__proxy__pb2.RowResult.SerializeToString, + ), + 'ReadRows': grpc.unary_unary_rpc_method_handler( + servicer.ReadRows, + request_deserializer=test__proxy__pb2.ReadRowsRequest.FromString, + response_serializer=test__proxy__pb2.RowsResult.SerializeToString, + ), + 'MutateRow': grpc.unary_unary_rpc_method_handler( + servicer.MutateRow, + request_deserializer=test__proxy__pb2.MutateRowRequest.FromString, + response_serializer=test__proxy__pb2.MutateRowResult.SerializeToString, + ), + 'BulkMutateRows': grpc.unary_unary_rpc_method_handler( + servicer.BulkMutateRows, + request_deserializer=test__proxy__pb2.MutateRowsRequest.FromString, + response_serializer=test__proxy__pb2.MutateRowsResult.SerializeToString, + ), + 'CheckAndMutateRow': grpc.unary_unary_rpc_method_handler( + servicer.CheckAndMutateRow, + request_deserializer=test__proxy__pb2.CheckAndMutateRowRequest.FromString, + response_serializer=test__proxy__pb2.CheckAndMutateRowResult.SerializeToString, + ), + 'SampleRowKeys': grpc.unary_unary_rpc_method_handler( + servicer.SampleRowKeys, + request_deserializer=test__proxy__pb2.SampleRowKeysRequest.FromString, + response_serializer=test__proxy__pb2.SampleRowKeysResult.SerializeToString, + ), + 'ReadModifyWriteRow': grpc.unary_unary_rpc_method_handler( + servicer.ReadModifyWriteRow, + request_deserializer=test__proxy__pb2.ReadModifyWriteRowRequest.FromString, + response_serializer=test__proxy__pb2.RowResult.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'google.bigtable.testproxy.CloudBigtableV2TestProxy', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + + + # This class is part of an EXPERIMENTAL API. +class CloudBigtableV2TestProxy(object): + """Note that all RPCs are unary, even when the equivalent client binding call + may be streaming. This is an intentional simplification. + + Most methods have sync (default) and async variants. For async variants, + the proxy is expected to perform the async operation, then wait for results + before delivering them back to the driver client. + + Operations that may have interesting concurrency characteristics are + represented explicitly in the API (see ReadRowsRequest.cancel_after_rows). + We include such operations only when they can be meaningfully performed + through client bindings. + + Users should generally avoid setting deadlines for requests to the Proxy + because operations are not cancelable. If the deadline is set anyway, please + understand that the underlying operation will continue to be executed even + after the deadline expires. + """ + + @staticmethod + def CreateClient(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/CreateClient', + test__proxy__pb2.CreateClientRequest.SerializeToString, + test__proxy__pb2.CreateClientResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def CloseClient(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/CloseClient', + test__proxy__pb2.CloseClientRequest.SerializeToString, + test__proxy__pb2.CloseClientResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def RemoveClient(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/RemoveClient', + test__proxy__pb2.RemoveClientRequest.SerializeToString, + test__proxy__pb2.RemoveClientResponse.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def ReadRow(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/ReadRow', + test__proxy__pb2.ReadRowRequest.SerializeToString, + test__proxy__pb2.RowResult.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def ReadRows(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/ReadRows', + test__proxy__pb2.ReadRowsRequest.SerializeToString, + test__proxy__pb2.RowsResult.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def MutateRow(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/MutateRow', + test__proxy__pb2.MutateRowRequest.SerializeToString, + test__proxy__pb2.MutateRowResult.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def BulkMutateRows(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/BulkMutateRows', + test__proxy__pb2.MutateRowsRequest.SerializeToString, + test__proxy__pb2.MutateRowsResult.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def CheckAndMutateRow(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/CheckAndMutateRow', + test__proxy__pb2.CheckAndMutateRowRequest.SerializeToString, + test__proxy__pb2.CheckAndMutateRowResult.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def SampleRowKeys(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/SampleRowKeys', + test__proxy__pb2.SampleRowKeysRequest.SerializeToString, + test__proxy__pb2.SampleRowKeysResult.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) + + @staticmethod + def ReadModifyWriteRow(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary(request, target, '/google.bigtable.testproxy.CloudBigtableV2TestProxy/ReadModifyWriteRow', + test__proxy__pb2.ReadModifyWriteRowRequest.SerializeToString, + test__proxy__pb2.RowResult.FromString, + options, channel_credentials, + insecure, call_credentials, compression, wait_for_ready, timeout, metadata) diff --git a/test_proxy/run_tests.sh b/test_proxy/run_tests.sh new file mode 100755 index 000000000..15b146b03 --- /dev/null +++ b/test_proxy/run_tests.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# Copyright 2023 Google LLC +# +# Licensed 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 +# +# https://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 + +# attempt download golang if not found +if [[ ! -x "$(command -v go)" ]]; then + echo "Downloading golang..." + wget https://go.dev/dl/go1.20.2.linux-amd64.tar.gz + tar -xzf go1.20.2.linux-amd64.tar.gz + export GOROOT=$(pwd)/go + export PATH=$GOROOT/bin:$PATH + export GOPATH=$HOME/go + go version +fi + +# ensure the working dir is the script's folder +SCRIPT_DIR=$(realpath $(dirname "$0")) +cd $SCRIPT_DIR + +export PROXY_SERVER_PORT=50055 + +# download test suite +if [ ! -d "cloud-bigtable-clients-test" ]; then + git clone https://github.com/googleapis/cloud-bigtable-clients-test.git +fi + +# start proxy +python test_proxy.py --port $PROXY_SERVER_PORT & +PROXY_PID=$! +function finish { + kill $PROXY_PID +} +trap finish EXIT + +# run tests +pushd cloud-bigtable-clients-test/tests +go test -v -proxy_addr=:$PROXY_SERVER_PORT diff --git a/test_proxy/test_proxy.py b/test_proxy/test_proxy.py new file mode 100644 index 000000000..a0cf2f1f0 --- /dev/null +++ b/test_proxy/test_proxy.py @@ -0,0 +1,193 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +""" +The Python implementation of the `cloud-bigtable-clients-test` proxy server. + +https://github.com/googleapis/cloud-bigtable-clients-test + +This server is intended to be used to test the correctness of Bigtable +clients across languages. + +Contributor Note: the proxy implementation is split across TestProxyClientHandler +and TestProxyGrpcServer. This is due to the fact that generated protos and proto-plus +objects cannot be used in the same process, so we had to make use of the +multiprocessing module to allow them to work together. +""" + +import multiprocessing +import argparse +import sys +import os +sys.path.append("handlers") + + +def grpc_server_process(request_q, queue_pool, port=50055): + """ + Defines a process that hosts a grpc server + proxies requests to a client_handler_process + """ + sys.path.append("protos") + from concurrent import futures + + import grpc + import test_proxy_pb2_grpc + import grpc_handler + + # Start gRPC server + server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) + test_proxy_pb2_grpc.add_CloudBigtableV2TestProxyServicer_to_server( + grpc_handler.TestProxyGrpcServer(request_q, queue_pool), server + ) + server.add_insecure_port("[::]:" + port) + server.start() + print("grpc_server_process started, listening on " + port) + server.wait_for_termination() + + +async def client_handler_process_async(request_q, queue_pool, use_legacy_client=False): + """ + Defines a process that recives Bigtable requests from a grpc_server_process, + and runs the request using a client library instance + """ + import base64 + import re + import asyncio + import warnings + import client_handler_data + import client_handler_legacy + warnings.filterwarnings("ignore", category=RuntimeWarning, message=".*Bigtable emulator.*") + + def camel_to_snake(str): + return re.sub(r"(? Date: Mon, 16 Oct 2023 09:08:07 -0700 Subject: [PATCH 23/56] chore(tests): add conformance tests to CI for v3 (#870) --- .github/sync-repo-settings.yaml | 18 +++ .github/workflows/conformance.yaml | 54 +++++++ .kokoro/conformance.sh | 52 +++++++ google/cloud/bigtable/data/_async/client.py | 5 +- test_proxy/handlers/client_handler_legacy.py | 139 ++++++++++++++++--- testing/constraints-3.8.txt | 13 ++ 6 files changed, 262 insertions(+), 19 deletions(-) create mode 100644 .github/workflows/conformance.yaml create mode 100644 .kokoro/conformance.sh diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml index a0d3362c9..a8cc5b33b 100644 --- a/.github/sync-repo-settings.yaml +++ b/.github/sync-repo-settings.yaml @@ -31,6 +31,24 @@ branchProtectionRules: - 'Kokoro' - 'Kokoro system-3.8' - 'cla/google' +- pattern: experimental_v3 + # Can admins overwrite branch protection. + # Defaults to `true` + isAdminEnforced: false + # Number of approving reviews required to update matching branches. + # Defaults to `1` + requiredApprovingReviewCount: 1 + # Are reviews from code owners required to update matching branches. + # Defaults to `false` + requiresCodeOwnerReviews: false + # Require up to date branches + requiresStrictStatusChecks: false + # List of required status check contexts that must pass for commits to be accepted to matching branches. + requiredStatusCheckContexts: + - 'Kokoro' + - 'Kokoro system-3.8' + - 'cla/google' + - 'Conformance / Async v3 Client / Python 3.8' # List of explicit permissions to add (additive only) permissionRules: # Team slug to add to repository permissions diff --git a/.github/workflows/conformance.yaml b/.github/workflows/conformance.yaml new file mode 100644 index 000000000..bffbf68cc --- /dev/null +++ b/.github/workflows/conformance.yaml @@ -0,0 +1,54 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. +# Github action job to test core java library features on +# downstream client libraries before they are released. +on: + push: + branches: + - main + pull_request: +name: Conformance +jobs: + conformance: + runs-on: ubuntu-latest + strategy: + matrix: + py-version: [ 3.8 ] + client-type: [ "Async v3", "Legacy" ] + name: "${{ matrix.client-type }} Client / Python ${{ matrix.py-version }}" + steps: + - uses: actions/checkout@v3 + name: "Checkout python-bigtable" + - uses: actions/checkout@v3 + name: "Checkout conformance tests" + with: + repository: googleapis/cloud-bigtable-clients-test + ref: main + path: cloud-bigtable-clients-test + - uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.py-version }} + - uses: actions/setup-go@v4 + with: + go-version: '>=1.20.2' + - run: chmod +x .kokoro/conformance.sh + - run: pip install -e . + name: "Install python-bigtable from HEAD" + - run: go version + - run: .kokoro/conformance.sh + name: "Run tests" + env: + CLIENT_TYPE: ${{ matrix.client-type }} + PYTHONUNBUFFERED: 1 + diff --git a/.kokoro/conformance.sh b/.kokoro/conformance.sh new file mode 100644 index 000000000..1c0b3ee0d --- /dev/null +++ b/.kokoro/conformance.sh @@ -0,0 +1,52 @@ +#!/bin/bash + +# Copyright 2023 Google LLC +# +# Licensed 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. + +set -eo pipefail + +## cd to the parent directory, i.e. the root of the git repo +cd $(dirname $0)/.. + +PROXY_ARGS="" +TEST_ARGS="" +if [[ "${CLIENT_TYPE^^}" == "LEGACY" ]]; then + echo "Using legacy client" + PROXY_ARGS="--legacy-client" + # legacy client does not expose mutate_row. Disable those tests + TEST_ARGS="-skip TestMutateRow_" +fi + +# Build and start the proxy in a separate process +PROXY_PORT=9999 +pushd test_proxy +nohup python test_proxy.py --port $PROXY_PORT $PROXY_ARGS & +proxyPID=$! +popd + +# Kill proxy on exit +function cleanup() { + echo "Cleanup testbench"; + kill $proxyPID +} +trap cleanup EXIT + +# Run the conformance test +pushd cloud-bigtable-clients-test/tests +eval "go test -v -proxy_addr=:$PROXY_PORT $TEST_ARGS" +RETURN_CODE=$? +popd + +echo "exiting with ${RETURN_CODE}" +exit ${RETURN_CODE} diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 8524cd9aa..e5be1b2d3 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -132,6 +132,10 @@ def __init__( client_options = cast( Optional[client_options_lib.ClientOptions], client_options ) + self._emulator_host = os.getenv(BIGTABLE_EMULATOR) + if self._emulator_host is not None and credentials is None: + # use insecure channel if emulator is set + credentials = google.auth.credentials.AnonymousCredentials() # initialize client ClientWithProject.__init__( self, @@ -155,7 +159,6 @@ def __init__( self._instance_owners: dict[_WarmedInstanceKey, Set[int]] = {} self._channel_init_time = time.monotonic() self._channel_refresh_tasks: list[asyncio.Task[None]] = [] - self._emulator_host = os.getenv(BIGTABLE_EMULATOR) if self._emulator_host is not None: # connect to an emulator host warnings.warn( diff --git a/test_proxy/handlers/client_handler_legacy.py b/test_proxy/handlers/client_handler_legacy.py index b423165f1..400f618b5 100644 --- a/test_proxy/handlers/client_handler_legacy.py +++ b/test_proxy/handlers/client_handler_legacy.py @@ -44,7 +44,7 @@ def __init__( self.app_profile_id = app_profile_id self.per_operation_timeout = per_operation_timeout - async def close(self): + def close(self): self.closed = True @client_handler.error_safe @@ -79,8 +79,33 @@ async def ReadRows(self, request, **kwargs): row_list.append(dict_val) return row_list + @client_handler.error_safe + async def ReadRow(self, row_key, **kwargs): + table_id = kwargs["table_name"].split("/")[-1] + instance = self.client.instance(self.instance_id) + table = instance.table(table_id) + + row = table.read_row(row_key) + # parse results into proto formatted dict + dict_val = {"row_key": row.row_key} + for family, family_cells in row.cells.items(): + family_dict = {"name": family} + for qualifier, qualifier_cells in family_cells.items(): + column_dict = {"qualifier": qualifier} + for cell in qualifier_cells: + cell_dict = { + "value": cell.value, + "timestamp_micros": cell.timestamp.timestamp() * 1000000, + "labels": cell.labels, + } + column_dict.setdefault("cells", []).append(cell_dict) + family_dict.setdefault("columns", []).append(column_dict) + dict_val.setdefault("families", []).append(family_dict) + return dict_val + @client_handler.error_safe async def MutateRow(self, request, **kwargs): + from datetime import datetime from google.cloud.bigtable.row import DirectRow table_id = request["table_name"].split("/")[-1] instance = self.client.instance(self.instance_id) @@ -88,24 +113,23 @@ async def MutateRow(self, request, **kwargs): row_key = request["row_key"] new_row = DirectRow(row_key, table) for m_dict in request.get("mutations", []): + details = m_dict.get("set_cell") or m_dict.get("delete_from_column") or m_dict.get("delete_from_family") or m_dict.get("delete_from_row") + timestamp = datetime.fromtimestamp(details.get("timestamp_micros")) if details.get("timestamp_micros") else None if m_dict.get("set_cell"): - details = m_dict["set_cell"] - new_row.set_cell(details["family_name"], details["column_qualifier"], details["value"], timestamp=details["timestamp_micros"]) + new_row.set_cell(details["family_name"], details["column_qualifier"], details["value"], timestamp=timestamp) elif m_dict.get("delete_from_column"): - details = m_dict["delete_from_column"] - new_row.delete_cell(details["family_name"], details["column_qualifier"], timestamp=details["timestamp_micros"]) + new_row.delete_cell(details["family_name"], details["column_qualifier"], timestamp=timestamp) elif m_dict.get("delete_from_family"): - details = m_dict["delete_from_family"] - new_row.delete_cells(details["family_name"], timestamp=details["timestamp_micros"]) + new_row.delete_cells(details["family_name"], timestamp=timestamp) elif m_dict.get("delete_from_row"): new_row.delete() - async with self.measure_call(): - table.mutate_rows([new_row]) + table.mutate_rows([new_row]) return "OK" @client_handler.error_safe async def BulkMutateRows(self, request, **kwargs): from google.cloud.bigtable.row import DirectRow + from datetime import datetime table_id = request["table_name"].split("/")[-1] instance = self.client.instance(self.instance_id) table = instance.table(table_id) @@ -113,20 +137,99 @@ async def BulkMutateRows(self, request, **kwargs): for entry in request.get("entries", []): row_key = entry["row_key"] new_row = DirectRow(row_key, table) - for m_dict in entry.get("mutations", {}): + for m_dict in entry.get("mutations"): + details = m_dict.get("set_cell") or m_dict.get("delete_from_column") or m_dict.get("delete_from_family") or m_dict.get("delete_from_row") + timestamp = datetime.fromtimestamp(details.get("timestamp_micros")) if details.get("timestamp_micros") else None if m_dict.get("set_cell"): - details = m_dict["set_cell"] - new_row.set_cell(details["family_name"], details["column_qualifier"], details["value"], timestamp=details.get("timestamp_micros",None)) + new_row.set_cell(details["family_name"], details["column_qualifier"], details["value"], timestamp=timestamp) elif m_dict.get("delete_from_column"): - details = m_dict["delete_from_column"] - new_row.delete_cell(details["family_name"], details["column_qualifier"], timestamp=details["timestamp_micros"]) + new_row.delete_cell(details["family_name"], details["column_qualifier"], timestamp=timestamp) elif m_dict.get("delete_from_family"): - details = m_dict["delete_from_family"] - new_row.delete_cells(details["family_name"], timestamp=details["timestamp_micros"]) + new_row.delete_cells(details["family_name"], timestamp=timestamp) elif m_dict.get("delete_from_row"): new_row.delete() rows.append(new_row) - async with self.measure_call(): - table.mutate_rows(rows) + table.mutate_rows(rows) return "OK" + @client_handler.error_safe + async def CheckAndMutateRow(self, request, **kwargs): + from google.cloud.bigtable.row import ConditionalRow + from google.cloud.bigtable.row_filters import PassAllFilter + table_id = request["table_name"].split("/")[-1] + instance = self.client.instance(self.instance_id) + table = instance.table(table_id) + + predicate_filter = request.get("predicate_filter", PassAllFilter(True)) + new_row = ConditionalRow(request["row_key"], table, predicate_filter) + + combined_mutations = [{"state": True, **m} for m in request.get("true_mutations", [])] + combined_mutations.extend([{"state": False, **m} for m in request.get("false_mutations", [])]) + for mut_dict in combined_mutations: + if "set_cell" in mut_dict: + details = mut_dict["set_cell"] + new_row.set_cell( + details.get("family_name", ""), + details.get("column_qualifier", ""), + details.get("value", ""), + timestamp=details.get("timestamp_micros", None), + state=mut_dict["state"], + ) + elif "delete_from_column" in mut_dict: + details = mut_dict["delete_from_column"] + new_row.delete_cell( + details.get("family_name", ""), + details.get("column_qualifier", ""), + timestamp=details.get("timestamp_micros", None), + state=mut_dict["state"], + ) + elif "delete_from_family" in mut_dict: + details = mut_dict["delete_from_family"] + new_row.delete_cells( + details.get("family_name", ""), + timestamp=details.get("timestamp_micros", None), + state=mut_dict["state"], + ) + elif "delete_from_row" in mut_dict: + new_row.delete(state=mut_dict["state"]) + else: + raise RuntimeError(f"Unknown mutation type: {mut_dict}") + return new_row.commit() + + @client_handler.error_safe + async def ReadModifyWriteRow(self, request, **kwargs): + from google.cloud.bigtable.row import AppendRow + from google.cloud._helpers import _microseconds_from_datetime + table_id = request["table_name"].split("/")[-1] + instance = self.client.instance(self.instance_id) + table = instance.table(table_id) + row_key = request["row_key"] + new_row = AppendRow(row_key, table) + for rule_dict in request.get("rules", []): + qualifier = rule_dict["column_qualifier"] + family = rule_dict["family_name"] + if "append_value" in rule_dict: + new_row.append_cell_value(family, qualifier, rule_dict["append_value"]) + else: + new_row.increment_cell_value(family, qualifier, rule_dict["increment_amount"]) + raw_result = new_row.commit() + result_families = [] + for family, column_dict in raw_result.items(): + result_columns = [] + for column, cell_list in column_dict.items(): + result_cells = [] + for cell_tuple in cell_list: + cell_dict = {"value": cell_tuple[0], "timestamp_micros": _microseconds_from_datetime(cell_tuple[1])} + result_cells.append(cell_dict) + result_columns.append({"qualifier": column, "cells": result_cells}) + result_families.append({"name": family, "columns": result_columns}) + return {"key": row_key, "families": result_families} + + @client_handler.error_safe + async def SampleRowKeys(self, request, **kwargs): + table_id = request["table_name"].split("/")[-1] + instance = self.client.instance(self.instance_id) + table = instance.table(table_id) + response = list(table.sample_row_keys()) + tuple_response = [(s.row_key, s.offset_bytes) for s in response] + return tuple_response diff --git a/testing/constraints-3.8.txt b/testing/constraints-3.8.txt index e69de29bb..7045a2894 100644 --- a/testing/constraints-3.8.txt +++ b/testing/constraints-3.8.txt @@ -0,0 +1,13 @@ +# This constraints file is used to check that lower bounds +# are correct in setup.py +# List *all* library dependencies and extras in this file. +# Pin the version to the lower bound. +# +# e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", +# Then this file should have foo==1.14.0 +google-api-core==2.12.0.dev1 +google-cloud-core==2.3.2 +grpc-google-iam-v1==0.12.4 +proto-plus==1.22.0 +libcst==0.2.5 +protobuf==3.19.5 From 50531e5b8f100a817cff521323974998f77b76fd Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 26 Oct 2023 14:23:49 -0700 Subject: [PATCH 24/56] chore(tests): turn off fast fail for conformance tets (#882) --- .github/workflows/conformance.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/conformance.yaml b/.github/workflows/conformance.yaml index bffbf68cc..69350b18d 100644 --- a/.github/workflows/conformance.yaml +++ b/.github/workflows/conformance.yaml @@ -26,6 +26,7 @@ jobs: matrix: py-version: [ 3.8 ] client-type: [ "Async v3", "Legacy" ] + fail-fast: false name: "${{ matrix.client-type }} Client / Python ${{ matrix.py-version }}" steps: - uses: actions/checkout@v3 From 8ff12166a7018170cd7a3472e4543e6f3a497a85 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 26 Oct 2023 16:19:45 -0700 Subject: [PATCH 25/56] feat: add TABLE_DEFAULTS enum for table method arguments (#880) --- google/cloud/bigtable/data/__init__.py | 12 +- google/cloud/bigtable/data/_async/client.py | 175 +++++++----------- .../bigtable/data/_async/mutations_batcher.py | 27 ++- google/cloud/bigtable/data/_helpers.py | 76 +++++++- noxfile.py | 2 +- tests/unit/data/test__helpers.py | 65 +++++++ 6 files changed, 229 insertions(+), 128 deletions(-) diff --git a/google/cloud/bigtable/data/__init__.py b/google/cloud/bigtable/data/__init__.py index 4b01d0e6b..a68be5417 100644 --- a/google/cloud/bigtable/data/__init__.py +++ b/google/cloud/bigtable/data/__init__.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # - -from typing import List, Tuple - from google.cloud.bigtable import gapic_version as package_version from google.cloud.bigtable.data._async.client import BigtableDataClientAsync @@ -44,10 +41,10 @@ from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup -# Type alias for the output of sample_keys -RowKeySamples = List[Tuple[bytes, int]] -# type alias for the output of query.shard() -ShardedQuery = List[ReadRowsQuery] +from google.cloud.bigtable.data._helpers import TABLE_DEFAULT +from google.cloud.bigtable.data._helpers import RowKeySamples +from google.cloud.bigtable.data._helpers import ShardedQuery + __version__: str = package_version.__version__ @@ -74,4 +71,5 @@ "MutationsExceptionGroup", "ShardedReadRowsExceptionGroup", "ShardedQuery", + "TABLE_DEFAULT", ) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index e5be1b2d3..c6637581c 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -32,7 +32,6 @@ import random import os -from collections import namedtuple from google.cloud.bigtable_v2.services.bigtable.client import BigtableClientMeta from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient @@ -59,30 +58,26 @@ from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync +from google.cloud.bigtable.data._helpers import TABLE_DEFAULT +from google.cloud.bigtable.data._helpers import _WarmedInstanceKey +from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT from google.cloud.bigtable.data._helpers import _make_metadata from google.cloud.bigtable.data._helpers import _convert_retry_deadline from google.cloud.bigtable.data._helpers import _validate_timeouts +from google.cloud.bigtable.data._helpers import _get_timeouts +from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync from google.cloud.bigtable.data._async.mutations_batcher import _MB_SIZE -from google.cloud.bigtable.data._helpers import _attempt_timeout_generator - from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule from google.cloud.bigtable.data.row_filters import RowFilter from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter from google.cloud.bigtable.data.row_filters import RowFilterChain -if TYPE_CHECKING: - from google.cloud.bigtable.data import RowKeySamples - from google.cloud.bigtable.data import ShardedQuery - -# used by read_rows_sharded to limit how many requests are attempted in parallel -_CONCURRENCY_LIMIT = 10 -# used to register instance data with the client for channel warming -_WarmedInstanceKey = namedtuple( - "_WarmedInstanceKey", ["instance_name", "table_name", "app_profile_id"] -) +if TYPE_CHECKING: + from google.cloud.bigtable.data._helpers import RowKeySamples + from google.cloud.bigtable.data._helpers import ShardedQuery class BigtableDataClientAsync(ClientWithProject): @@ -525,8 +520,8 @@ async def read_rows_stream( self, query: ReadRowsQuery, *, - operation_timeout: float | None = None, - attempt_timeout: float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> AsyncIterable[Row]: """ Read a set of rows from the table, based on the specified query. @@ -538,12 +533,12 @@ async def read_rows_stream( - query: contains details about which rows to return - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - If None, defaults to the Table's default_read_rows_operation_timeout + Defaults to the Table's default_read_rows_operation_timeout - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_read_rows_attempt_timeout, - or the operation_timeout if that is also None. + Defaults to the Table's default_read_rows_attempt_timeout. + If None, defaults to operation_timeout. Returns: - an asynchronous iterator that yields rows returned by the query Raises: @@ -553,15 +548,9 @@ async def read_rows_stream( - GoogleAPIError: raised if the request encounters an unrecoverable error - IdleTimeout: if iterator was abandoned """ - operation_timeout = ( - operation_timeout or self.default_read_rows_operation_timeout - ) - attempt_timeout = ( - attempt_timeout - or self.default_read_rows_attempt_timeout - or operation_timeout + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self ) - _validate_timeouts(operation_timeout, attempt_timeout) row_merger = _ReadRowsOperationAsync( query, @@ -575,8 +564,8 @@ async def read_rows( self, query: ReadRowsQuery, *, - operation_timeout: float | None = None, - attempt_timeout: float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> list[Row]: """ Read a set of rows from the table, based on the specified query. @@ -589,12 +578,12 @@ async def read_rows( - query: contains details about which rows to return - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - If None, defaults to the Table's default_read_rows_operation_timeout + Defaults to the Table's default_read_rows_operation_timeout - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_read_rows_attempt_timeout, - or the operation_timeout if that is also None. + Defaults to the Table's default_read_rows_attempt_timeout. + If None, defaults to operation_timeout. Returns: - a list of Rows returned by the query Raises: @@ -615,8 +604,8 @@ async def read_row( row_key: str | bytes, *, row_filter: RowFilter | None = None, - operation_timeout: int | float | None = None, - attempt_timeout: int | float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> Row | None: """ Read a single row from the table, based on the specified key. @@ -627,12 +616,12 @@ async def read_row( - query: contains details about which rows to return - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - If None, defaults to the Table's default_read_rows_operation_timeout + Defaults to the Table's default_read_rows_operation_timeout - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_read_rows_attempt_timeout, or the operation_timeout - if that is also None. + Defaults to the Table's default_read_rows_attempt_timeout. + If None, defaults to operation_timeout. Returns: - a Row object if the row exists, otherwise None Raises: @@ -657,8 +646,8 @@ async def read_rows_sharded( self, sharded_query: ShardedQuery, *, - operation_timeout: int | float | None = None, - attempt_timeout: int | float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> list[Row]: """ Runs a sharded query in parallel, then return the results in a single list. @@ -677,12 +666,12 @@ async def read_rows_sharded( - sharded_query: a sharded query to execute - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - If None, defaults to the Table's default_read_rows_operation_timeout + Defaults to the Table's default_read_rows_operation_timeout - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_read_rows_attempt_timeout, or the operation_timeout - if that is also None. + Defaults to the Table's default_read_rows_attempt_timeout. + If None, defaults to operation_timeout. Raises: - ShardedReadRowsExceptionGroup: if any of the queries failed - ValueError: if the query_list is empty @@ -690,15 +679,9 @@ async def read_rows_sharded( if not sharded_query: raise ValueError("empty sharded_query") # reduce operation_timeout between batches - operation_timeout = ( - operation_timeout or self.default_read_rows_operation_timeout + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self ) - attempt_timeout = ( - attempt_timeout - or self.default_read_rows_attempt_timeout - or operation_timeout - ) - _validate_timeouts(operation_timeout, attempt_timeout) timeout_generator = _attempt_timeout_generator( operation_timeout, operation_timeout ) @@ -744,8 +727,8 @@ async def row_exists( self, row_key: str | bytes, *, - operation_timeout: int | float | None = None, - attempt_timeout: int | float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> bool: """ Return a boolean indicating whether the specified row exists in the table. @@ -754,12 +737,12 @@ async def row_exists( - row_key: the key of the row to check - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - If None, defaults to the Table's default_read_rows_operation_timeout + Defaults to the Table's default_read_rows_operation_timeout - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_read_rows_attempt_timeout, or the operation_timeout - if that is also None. + Defaults to the Table's default_read_rows_attempt_timeout. + If None, defaults to operation_timeout. Returns: - a bool indicating whether the row exists Raises: @@ -785,8 +768,8 @@ async def row_exists( async def sample_row_keys( self, *, - operation_timeout: float | None = None, - attempt_timeout: float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, + attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, ) -> RowKeySamples: """ Return a set of RowKeySamples that delimit contiguous sections of the table of @@ -801,13 +784,13 @@ async def sample_row_keys( Args: - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will be retried within the budget. - If None, defaults to the Table's default_operation_timeout + Failed requests will be retried within the budget.i + Defaults to the Table's default_operation_timeout - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_attempt_timeout, or the operation_timeout - if that is also None. + Defaults to the Table's default_attempt_timeout. + If None, defaults to operation_timeout. Returns: - a set of RowKeySamples the delimit contiguous sections of the table Raises: @@ -817,12 +800,9 @@ async def sample_row_keys( - GoogleAPIError: raised if the request encounters an unrecoverable error """ # prepare timeouts - operation_timeout = operation_timeout or self.default_operation_timeout - attempt_timeout = ( - attempt_timeout or self.default_attempt_timeout or operation_timeout + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self ) - _validate_timeouts(operation_timeout, attempt_timeout) - attempt_timeout_gen = _attempt_timeout_generator( attempt_timeout, operation_timeout ) @@ -873,8 +853,8 @@ def mutations_batcher( flush_limit_bytes: int = 20 * _MB_SIZE, flow_control_max_mutation_count: int = 100_000, flow_control_max_bytes: int = 100 * _MB_SIZE, - batch_operation_timeout: float | None = None, - batch_attempt_timeout: float | None = None, + batch_operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, + batch_attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, ) -> MutationsBatcherAsync: """ Returns a new mutations batcher instance. @@ -890,11 +870,11 @@ def mutations_batcher( - flush_limit_bytes: Flush immediately after flush_limit_bytes bytes are added. - flow_control_max_mutation_count: Maximum number of inflight mutations. - flow_control_max_bytes: Maximum number of inflight bytes. - - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None, - table default_mutate_rows_operation_timeout will be used - - batch_attempt_timeout: timeout for each individual request, in seconds. If None, - table default_mutate_rows_attempt_timeout will be used, or batch_operation_timeout - if that is also None. + - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. + Defaults to the Table's default_mutate_rows_operation_timeout + - batch_attempt_timeout: timeout for each individual request, in seconds. + Defaults to the Table's default_mutate_rows_attempt_timeout. + If None, defaults to batch_operation_timeout. Returns: - a MutationsBatcherAsync context manager that can batch requests """ @@ -914,8 +894,8 @@ async def mutate_row( row_key: str | bytes, mutations: list[Mutation] | Mutation, *, - operation_timeout: float | None = None, - attempt_timeout: float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, + attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, ): """ Mutates a row atomically. @@ -931,12 +911,12 @@ async def mutate_row( - mutations: the set of mutations to apply to the row - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - If None, defaults to the Table's default_operation_timeout + Defaults to the Table's default_operation_timeout - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_attempt_timeout, or the operation_timeout - if that is also None. + Defaults to the Table's default_attempt_timeout. + If None, defaults to operation_timeout. Raises: - DeadlineExceeded: raised after operation timeout will be chained with a RetryExceptionGroup containing all @@ -944,11 +924,9 @@ async def mutate_row( - GoogleAPIError: raised on non-idempotent operations that cannot be safely retried. """ - operation_timeout = operation_timeout or self.default_operation_timeout - attempt_timeout = ( - attempt_timeout or self.default_attempt_timeout or operation_timeout + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self ) - _validate_timeouts(operation_timeout, attempt_timeout) if isinstance(row_key, str): row_key = row_key.encode("utf-8") @@ -1000,8 +978,8 @@ async def bulk_mutate_rows( self, mutation_entries: list[RowMutationEntry], *, - operation_timeout: float | None = None, - attempt_timeout: float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, + attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, ): """ Applies mutations for multiple rows in a single batched request. @@ -1021,25 +999,19 @@ async def bulk_mutate_rows( in arbitrary order - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will be retried within the budget. - If None, defaults to the Table's default_mutate_rows_operation_timeout + Defaults to the Table's default_mutate_rows_operation_timeout - attempt_timeout: the time budget for an individual network request, in seconds. If it takes longer than this time to complete, the request will be cancelled with a DeadlineExceeded exception, and a retry will be attempted. - If None, defaults to the Table's default_mutate_rows_attempt_timeout, - or the operation_timeout if that is also None. + Defaults to the Table's default_mutate_rows_attempt_timeout. + If None, defaults to operation_timeout. Raises: - MutationsExceptionGroup if one or more mutations fails Contains details about any failed entries in .exceptions """ - operation_timeout = ( - operation_timeout or self.default_mutate_rows_operation_timeout - ) - attempt_timeout = ( - attempt_timeout - or self.default_mutate_rows_attempt_timeout - or operation_timeout + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self ) - _validate_timeouts(operation_timeout, attempt_timeout) operation = _MutateRowsOperationAsync( self.client._gapic_client, @@ -1057,7 +1029,7 @@ async def check_and_mutate_row( *, true_case_mutations: Mutation | list[Mutation] | None = None, false_case_mutations: Mutation | list[Mutation] | None = None, - operation_timeout: int | float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, ) -> bool: """ Mutates a row atomically based on the output of a predicate filter @@ -1086,15 +1058,12 @@ async def check_and_mutate_row( `true_case_mutations is empty, and at most 100000. - operation_timeout: the time budget for the entire operation, in seconds. Failed requests will not be retried. Defaults to the Table's default_operation_timeout - if None. Returns: - bool indicating whether the predicate was true or false Raises: - GoogleAPIError exceptions from grpc call """ - operation_timeout = operation_timeout or self.default_operation_timeout - if operation_timeout <= 0: - raise ValueError("operation_timeout must be greater than 0") + operation_timeout, _ = _get_timeouts(operation_timeout, None, self) row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key if true_case_mutations is not None and not isinstance( true_case_mutations, list @@ -1128,7 +1097,7 @@ async def read_modify_write_row( row_key: str | bytes, rules: ReadModifyWriteRule | list[ReadModifyWriteRule], *, - operation_timeout: int | float | None = None, + operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, ) -> Row: """ Reads and modifies a row atomically according to input ReadModifyWriteRules, @@ -1145,15 +1114,15 @@ async def read_modify_write_row( Rules are applied in order, meaning that earlier rules will affect the results of later ones. - operation_timeout: the time budget for the entire operation, in seconds. - Failed requests will not be retried. Defaults to the Table's default_operation_timeout - if None. + Failed requests will not be retried. + Defaults to the Table's default_operation_timeout. Returns: - Row: containing cell data that was modified as part of the operation Raises: - GoogleAPIError exceptions from grpc call """ - operation_timeout = operation_timeout or self.default_operation_timeout + operation_timeout, _ = _get_timeouts(operation_timeout, None, self) row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key if operation_timeout <= 0: raise ValueError("operation_timeout must be greater than 0") diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index 34e1bfb5d..7ff5f9a0b 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -23,7 +23,8 @@ from google.cloud.bigtable.data.mutations import RowMutationEntry from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup from google.cloud.bigtable.data.exceptions import FailedMutationEntryError -from google.cloud.bigtable.data._helpers import _validate_timeouts +from google.cloud.bigtable.data._helpers import _get_timeouts +from google.cloud.bigtable.data._helpers import TABLE_DEFAULT from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync from google.cloud.bigtable.data._async._mutate_rows import ( @@ -189,8 +190,8 @@ def __init__( flush_limit_bytes: int = 20 * _MB_SIZE, flow_control_max_mutation_count: int = 100_000, flow_control_max_bytes: int = 100 * _MB_SIZE, - batch_operation_timeout: float | None = None, - batch_attempt_timeout: float | None = None, + batch_operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, + batch_attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, ): """ Args: @@ -202,21 +203,15 @@ def __init__( - flush_limit_bytes: Flush immediately after flush_limit_bytes bytes are added. - flow_control_max_mutation_count: Maximum number of inflight mutations. - flow_control_max_bytes: Maximum number of inflight bytes. - - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None, - table default_mutate_rows_operation_timeout will be used - - batch_attempt_timeout: timeout for each individual request, in seconds. If None, - table default_mutate_rows_attempt_timeout will be used, or batch_operation_timeout - if that is also None. + - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. + If TABLE_DEFAULT, defaults to the Table's default_mutate_rows_operation_timeout. + - batch_attempt_timeout: timeout for each individual request, in seconds. + If TABLE_DEFAULT, defaults to the Table's default_mutate_rows_attempt_timeout. + If None, defaults to batch_operation_timeout. """ - self._operation_timeout: float = ( - batch_operation_timeout or table.default_mutate_rows_operation_timeout + self._operation_timeout, self._attempt_timeout = _get_timeouts( + batch_operation_timeout, batch_attempt_timeout, table ) - self._attempt_timeout: float = ( - batch_attempt_timeout - or table.default_mutate_rows_attempt_timeout - or self._operation_timeout - ) - _validate_timeouts(self._operation_timeout, self._attempt_timeout) self.closed: bool = False self._table = table self._staged_entries: list[RowMutationEntry] = [] diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 1f8a63d21..1d56926ff 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -13,8 +13,11 @@ # from __future__ import annotations -from typing import Callable, Any +from typing import Callable, List, Tuple, Any import time +import enum +from collections import namedtuple +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.api_core import exceptions as core_exceptions from google.cloud.bigtable.data.exceptions import RetryExceptionGroup @@ -23,6 +26,30 @@ Helper functions used in various places in the library. """ +# Type alias for the output of sample_keys +RowKeySamples = List[Tuple[bytes, int]] + +# type alias for the output of query.shard() +ShardedQuery = List[ReadRowsQuery] + +# used by read_rows_sharded to limit how many requests are attempted in parallel +_CONCURRENCY_LIMIT = 10 + +# used to register instance data with the client for channel warming +_WarmedInstanceKey = namedtuple( + "_WarmedInstanceKey", ["instance_name", "table_name", "app_profile_id"] +) + + +# enum used on method calls when table defaults should be used +class TABLE_DEFAULT(enum.Enum): + # default for mutate_row, sample_row_keys, check_and_mutate_row, and read_modify_write_row + DEFAULT = "DEFAULT" + # default for read_rows, read_rows_stream, read_rows_sharded, row_exists, and read_row + READ_ROWS = "READ_ROWS_DEFAULT" + # default for bulk_mutate_rows and mutations_batcher + MUTATE_ROWS = "MUTATE_ROWS_DEFAULT" + def _make_metadata( table_name: str, app_profile_id: str | None @@ -114,6 +141,51 @@ def wrapper(*args, **kwargs): return wrapper_async if is_async else wrapper +def _get_timeouts( + operation: float | TABLE_DEFAULT, attempt: float | None | TABLE_DEFAULT, table +) -> tuple[float, float]: + """ + Convert passed in timeout values to floats, using table defaults if necessary. + + attempt will use operation value if None, or if larger than operation. + + Will call _validate_timeouts on the outputs, and raise ValueError if the + resulting timeouts are invalid. + + Args: + - operation: The timeout value to use for the entire operation, in seconds. + - attempt: The timeout value to use for each attempt, in seconds. + - table: The table to use for default values. + Returns: + - A tuple of (operation_timeout, attempt_timeout) + """ + # load table defaults if necessary + if operation == TABLE_DEFAULT.DEFAULT: + final_operation = table.default_operation_timeout + elif operation == TABLE_DEFAULT.READ_ROWS: + final_operation = table.default_read_rows_operation_timeout + elif operation == TABLE_DEFAULT.MUTATE_ROWS: + final_operation = table.default_mutate_rows_operation_timeout + else: + final_operation = operation + if attempt == TABLE_DEFAULT.DEFAULT: + attempt = table.default_attempt_timeout + elif attempt == TABLE_DEFAULT.READ_ROWS: + attempt = table.default_read_rows_attempt_timeout + elif attempt == TABLE_DEFAULT.MUTATE_ROWS: + attempt = table.default_mutate_rows_attempt_timeout + + if attempt is None: + # no timeout specified, use operation timeout for both + final_attempt = final_operation + else: + # cap attempt timeout at operation timeout + final_attempt = min(attempt, final_operation) if final_operation else attempt + + _validate_timeouts(final_operation, final_attempt, allow_none=False) + return final_operation, final_attempt + + def _validate_timeouts( operation_timeout: float, attempt_timeout: float | None, allow_none: bool = False ): @@ -128,6 +200,8 @@ def _validate_timeouts( Raises: - ValueError if operation_timeout or attempt_timeout are invalid. """ + if operation_timeout is None: + raise ValueError("operation_timeout cannot be None") if operation_timeout <= 0: raise ValueError("operation_timeout must be greater than 0") if not allow_none and attempt_timeout is None: diff --git a/noxfile.py b/noxfile.py index 4b57e617f..e1d2f4acc 100644 --- a/noxfile.py +++ b/noxfile.py @@ -460,7 +460,7 @@ def prerelease_deps(session): # Exclude version 1.52.0rc1 which has a known issue. See https://github.com/grpc/grpc/issues/32163 "grpcio!=1.52.0rc1", "grpcio-status", - "google-api-core", + "google-api-core==2.12.0.dev1", # TODO: remove this once streaming retries is merged "proto-plus", "google-cloud-testutils", # dependencies of google-cloud-testutils" diff --git a/tests/unit/data/test__helpers.py b/tests/unit/data/test__helpers.py index 08bc397c3..6c11fa86a 100644 --- a/tests/unit/data/test__helpers.py +++ b/tests/unit/data/test__helpers.py @@ -14,6 +14,7 @@ import pytest import google.cloud.bigtable.data._helpers as _helpers +from google.cloud.bigtable.data._helpers import TABLE_DEFAULT import google.cloud.bigtable.data.exceptions as bigtable_exceptions import mock @@ -199,3 +200,67 @@ def test_validate_with_inputs(self, args, expected): except ValueError: pass assert success == expected + + +class TestGetTimeouts: + @pytest.mark.parametrize( + "input_times,input_table,expected", + [ + ((2, 1), {}, (2, 1)), + ((2, 4), {}, (2, 2)), + ((2, None), {}, (2, 2)), + ( + (TABLE_DEFAULT.DEFAULT, TABLE_DEFAULT.DEFAULT), + {"operation": 3, "attempt": 2}, + (3, 2), + ), + ( + (TABLE_DEFAULT.READ_ROWS, TABLE_DEFAULT.READ_ROWS), + {"read_rows_operation": 3, "read_rows_attempt": 2}, + (3, 2), + ), + ( + (TABLE_DEFAULT.MUTATE_ROWS, TABLE_DEFAULT.MUTATE_ROWS), + {"mutate_rows_operation": 3, "mutate_rows_attempt": 2}, + (3, 2), + ), + ((10, TABLE_DEFAULT.DEFAULT), {"attempt": None}, (10, 10)), + ((10, TABLE_DEFAULT.DEFAULT), {"attempt": 5}, (10, 5)), + ((10, TABLE_DEFAULT.DEFAULT), {"attempt": 100}, (10, 10)), + ((TABLE_DEFAULT.DEFAULT, 10), {"operation": 12}, (12, 10)), + ((TABLE_DEFAULT.DEFAULT, 10), {"operation": 3}, (3, 3)), + ], + ) + def test_get_timeouts(self, input_times, input_table, expected): + """ + test input/output mappings for a variety of valid inputs + """ + fake_table = mock.Mock() + for key in input_table.keys(): + # set the default fields in our fake table mock + setattr(fake_table, f"default_{key}_timeout", input_table[key]) + t1, t2 = _helpers._get_timeouts(input_times[0], input_times[1], fake_table) + assert t1 == expected[0] + assert t2 == expected[1] + + @pytest.mark.parametrize( + "input_times,input_table", + [ + ([0, 1], {}), + ([1, 0], {}), + ([None, 1], {}), + ([TABLE_DEFAULT.DEFAULT, 1], {"operation": None}), + ([TABLE_DEFAULT.DEFAULT, 1], {"operation": 0}), + ([1, TABLE_DEFAULT.DEFAULT], {"attempt": 0}), + ], + ) + def test_get_timeouts_invalid(self, input_times, input_table): + """ + test with inputs that should raise error during validation step + """ + fake_table = mock.Mock() + for key in input_table.keys(): + # set the default fields in our fake table mock + setattr(fake_table, f"default_{key}_timeout", input_table[key]) + with pytest.raises(ValueError): + _helpers._get_timeouts(input_times[0], input_times[1], fake_table) From 94bfe664c5e9daa3bc76d6f423dd12b7949b2f76 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 26 Oct 2023 17:04:47 -0700 Subject: [PATCH 26/56] fix: pass None for retry in gapic calls (#881) --- google/cloud/bigtable/data/_async/_mutate_rows.py | 1 + google/cloud/bigtable/data/_async/_read_rows.py | 1 + google/cloud/bigtable/data/_async/client.py | 3 +++ .../cloud/bigtable_v2/services/bigtable/async_client.py | 4 ++-- google/cloud/bigtable_v2/services/bigtable/client.py | 4 ++-- tests/unit/data/_async/test__mutate_rows.py | 3 ++- tests/unit/data/_async/test_client.py | 9 ++++++++- 7 files changed, 19 insertions(+), 6 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index baae205d9..be84fac17 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -79,6 +79,7 @@ def __init__( table_name=table.table_name, app_profile_id=table.app_profile_id, metadata=metadata, + retry=None, ) # create predicate for determining which errors are retryable self.is_retryable = retries.if_exception_type( diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 20b5618ea..90cc7e87c 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -140,6 +140,7 @@ def _read_rows_attempt(self) -> AsyncGenerator[Row, None]: self.request, timeout=next(self.attempt_timeout_gen), metadata=self._metadata, + retry=None, ) chunked_stream = self.chunk_stream(gapic_stream) return self.merge_rows(chunked_stream) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index c6637581c..90939927e 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -837,6 +837,7 @@ async def execute_rpc(): app_profile_id=self.app_profile_id, timeout=next(attempt_timeout_gen), metadata=metadata, + retry=None, ) return [(s.row_key, s.offset_bytes) async for s in results] @@ -1089,6 +1090,7 @@ async def check_and_mutate_row( }, metadata=metadata, timeout=operation_timeout, + retry=None, ) return result.predicate_matched @@ -1142,6 +1144,7 @@ async def read_modify_write_row( }, metadata=metadata, timeout=operation_timeout, + retry=None, ) # construct Row from result return Row._from_pb(result.row) diff --git a/google/cloud/bigtable_v2/services/bigtable/async_client.py b/google/cloud/bigtable_v2/services/bigtable/async_client.py index abd82d4d8..a80be70af 100644 --- a/google/cloud/bigtable_v2/services/bigtable/async_client.py +++ b/google/cloud/bigtable_v2/services/bigtable/async_client.py @@ -40,9 +40,9 @@ from google.oauth2 import service_account # type: ignore try: - OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault] + OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None] except AttributeError: # pragma: NO COVER - OptionalRetry = Union[retries.Retry, object] # type: ignore + OptionalRetry = Union[retries.Retry, object, None] # type: ignore from google.cloud.bigtable_v2.types import bigtable from google.cloud.bigtable_v2.types import data diff --git a/google/cloud/bigtable_v2/services/bigtable/client.py b/google/cloud/bigtable_v2/services/bigtable/client.py index b0efc8a0b..1c2e7b822 100644 --- a/google/cloud/bigtable_v2/services/bigtable/client.py +++ b/google/cloud/bigtable_v2/services/bigtable/client.py @@ -43,9 +43,9 @@ from google.oauth2 import service_account # type: ignore try: - OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault] + OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None] except AttributeError: # pragma: NO COVER - OptionalRetry = Union[retries.Retry, object] # type: ignore + OptionalRetry = Union[retries.Retry, object, None] # type: ignore from google.cloud.bigtable_v2.types import bigtable from google.cloud.bigtable_v2.types import data diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index 08422abca..eae3483ed 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -93,9 +93,10 @@ def test_ctor(self): assert client.mutate_rows.call_count == 1 # gapic_fn should call with table details inner_kwargs = client.mutate_rows.call_args[1] - assert len(inner_kwargs) == 3 + assert len(inner_kwargs) == 4 assert inner_kwargs["table_name"] == table.table_name assert inner_kwargs["app_profile_id"] == table.app_profile_id + assert inner_kwargs["retry"] is None metadata = inner_kwargs["metadata"] assert len(metadata) == 1 assert metadata[0][0] == "x-goog-request-params" diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index c2c4b0615..4ae46da6e 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1329,6 +1329,7 @@ async def test_read_rows_attempt_timeout( # check timeouts for _, call_kwargs in read_rows.call_args_list[:-1]: assert call_kwargs["timeout"] == per_request_t + assert call_kwargs["retry"] is None # last timeout should be adjusted to account for the time spent assert ( abs( @@ -1884,6 +1885,7 @@ async def test_sample_row_keys_default_timeout(self): _, kwargs = sample_row_keys.call_args assert abs(kwargs["timeout"] - expected_timeout) < 0.1 assert result == [] + assert kwargs["retry"] is None @pytest.mark.asyncio async def test_sample_row_keys_gapic_params(self): @@ -1905,11 +1907,12 @@ async def test_sample_row_keys_gapic_params(self): await table.sample_row_keys(attempt_timeout=expected_timeout) args, kwargs = sample_row_keys.call_args assert len(args) == 0 - assert len(kwargs) == 4 + assert len(kwargs) == 5 assert kwargs["timeout"] == expected_timeout assert kwargs["app_profile_id"] == expected_profile assert kwargs["table_name"] == table.table_name assert kwargs["metadata"] is not None + assert kwargs["retry"] is None @pytest.mark.parametrize("include_app_profile", [True, False]) @pytest.mark.asyncio @@ -2231,6 +2234,7 @@ async def test_bulk_mutate_rows(self, mutation_arg): ) assert kwargs["entries"] == [bulk_mutation._to_dict()] assert kwargs["timeout"] == expected_attempt_timeout + assert kwargs["retry"] is None @pytest.mark.asyncio async def test_bulk_mutate_rows_multiple_entries(self): @@ -2595,6 +2599,7 @@ async def test_check_and_mutate(self, gapic_result): ] assert request["app_profile_id"] == app_profile assert kwargs["timeout"] == operation_timeout + assert kwargs["retry"] is None @pytest.mark.asyncio async def test_check_and_mutate_bad_timeout(self): @@ -2678,6 +2683,7 @@ async def test_check_and_mutate_predicate_object(self): kwargs = mock_gapic.call_args[1] assert kwargs["request"]["predicate_filter"] == predicate_dict assert mock_predicate._to_dict.call_count == 1 + assert kwargs["retry"] is None @pytest.mark.asyncio async def test_check_and_mutate_mutations_parsing(self): @@ -2781,6 +2787,7 @@ async def test_read_modify_write_call_rule_args(self, call_rules, expected_rules assert mock_gapic.call_count == 1 found_kwargs = mock_gapic.call_args_list[0][1] assert found_kwargs["request"]["rules"] == expected_rules + assert found_kwargs["retry"] is None @pytest.mark.parametrize("rules", [[], None]) @pytest.mark.asyncio From 3ac80a958638d89281a54eaeddbef28e9d2aee87 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 22 Nov 2023 14:25:46 -0800 Subject: [PATCH 27/56] feat: replace internal dictionaries with protos in gapic calls (#875) --- .../bigtable/data/_async/_mutate_rows.py | 22 +++- google/cloud/bigtable/data/_async/client.py | 59 ++++----- .../bigtable/data/_async/mutations_batcher.py | 3 - google/cloud/bigtable/data/mutations.py | 15 +++ .../bigtable/data/read_modify_write_rules.py | 11 +- tests/unit/data/_async/test__mutate_rows.py | 6 +- tests/unit/data/_async/test_client.py | 112 +++++++++--------- tests/unit/data/test_mutations.py | 84 +++++++++++++ .../unit/data/test_read_modify_write_rules.py | 42 +++++++ 9 files changed, 253 insertions(+), 101 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index be84fac17..5bf759151 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -16,10 +16,12 @@ from typing import TYPE_CHECKING import asyncio +from dataclasses import dataclass import functools from google.api_core import exceptions as core_exceptions from google.api_core import retry_async as retries +import google.cloud.bigtable_v2.types.bigtable as types_pb import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _make_metadata from google.cloud.bigtable.data._helpers import _convert_retry_deadline @@ -36,6 +38,16 @@ from google.cloud.bigtable.data._async.client import TableAsync +@dataclass +class _EntryWithProto: + """ + A dataclass to hold a RowMutationEntry and its corresponding proto representation. + """ + + entry: RowMutationEntry + proto: types_pb.MutateRowsRequest.Entry + + class _MutateRowsOperationAsync: """ MutateRowsOperation manages the logic of sending a set of row mutations, @@ -105,7 +117,7 @@ def __init__( self.timeout_generator = _attempt_timeout_generator( attempt_timeout, operation_timeout ) - self.mutations = mutation_entries + self.mutations = [_EntryWithProto(m, m._to_pb()) for m in mutation_entries] self.remaining_indices = list(range(len(self.mutations))) self.errors: dict[int, list[Exception]] = {} @@ -136,7 +148,7 @@ async def start(self): cause_exc = exc_list[0] else: cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) - entry = self.mutations[idx] + entry = self.mutations[idx].entry all_errors.append( bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) ) @@ -154,9 +166,7 @@ async def _run_attempt(self): retry after the attempt is complete - GoogleAPICallError: if the gapic rpc fails """ - request_entries = [ - self.mutations[idx]._to_dict() for idx in self.remaining_indices - ] + request_entries = [self.mutations[idx].proto for idx in self.remaining_indices] # track mutations in this request that have not been finalized yet active_request_indices = { req_idx: orig_idx for req_idx, orig_idx in enumerate(self.remaining_indices) @@ -214,7 +224,7 @@ def _handle_entry_error(self, idx: int, exc: Exception): - idx: the index of the mutation that failed - exc: the exception to add to the list """ - entry = self.mutations[idx] + entry = self.mutations[idx].entry self.errors.setdefault(idx, []).append(exc) if ( entry.is_idempotent() diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 90939927e..ab8cc48f8 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -924,22 +924,17 @@ async def mutate_row( GoogleAPIError exceptions from any retries that failed - GoogleAPIError: raised on non-idempotent operations that cannot be safely retried. + - ValueError if invalid arguments are provided """ operation_timeout, attempt_timeout = _get_timeouts( operation_timeout, attempt_timeout, self ) - if isinstance(row_key, str): - row_key = row_key.encode("utf-8") - request = {"table_name": self.table_name, "row_key": row_key} - if self.app_profile_id: - request["app_profile_id"] = self.app_profile_id + if not mutations: + raise ValueError("No mutations provided") + mutations_list = mutations if isinstance(mutations, list) else [mutations] - if isinstance(mutations, Mutation): - mutations = [mutations] - request["mutations"] = [mutation._to_dict() for mutation in mutations] - - if all(mutation.is_idempotent() for mutation in mutations): + if all(mutation.is_idempotent() for mutation in mutations_list): # mutations are all idempotent and safe to retry predicate = retries.if_exception_type( core_exceptions.DeadlineExceeded, @@ -972,7 +967,13 @@ def on_error_fn(exc): metadata = _make_metadata(self.table_name, self.app_profile_id) # trigger rpc await deadline_wrapped( - request, timeout=attempt_timeout, metadata=metadata, retry=None + row_key=row_key.encode("utf-8") if isinstance(row_key, str) else row_key, + mutations=[mutation._to_pb() for mutation in mutations_list], + table_name=self.table_name, + app_profile_id=self.app_profile_id, + timeout=attempt_timeout, + metadata=metadata, + retry=None, ) async def bulk_mutate_rows( @@ -1009,6 +1010,7 @@ async def bulk_mutate_rows( Raises: - MutationsExceptionGroup if one or more mutations fails Contains details about any failed entries in .exceptions + - ValueError if invalid arguments are provided """ operation_timeout, attempt_timeout = _get_timeouts( operation_timeout, attempt_timeout, self @@ -1065,29 +1067,24 @@ async def check_and_mutate_row( - GoogleAPIError exceptions from grpc call """ operation_timeout, _ = _get_timeouts(operation_timeout, None, self) - row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key if true_case_mutations is not None and not isinstance( true_case_mutations, list ): true_case_mutations = [true_case_mutations] - true_case_dict = [m._to_dict() for m in true_case_mutations or []] + true_case_list = [m._to_pb() for m in true_case_mutations or []] if false_case_mutations is not None and not isinstance( false_case_mutations, list ): false_case_mutations = [false_case_mutations] - false_case_dict = [m._to_dict() for m in false_case_mutations or []] + false_case_list = [m._to_pb() for m in false_case_mutations or []] metadata = _make_metadata(self.table_name, self.app_profile_id) result = await self.client._gapic_client.check_and_mutate_row( - request={ - "predicate_filter": predicate._to_dict() - if predicate is not None - else None, - "true_mutations": true_case_dict, - "false_mutations": false_case_dict, - "table_name": self.table_name, - "row_key": row_key, - "app_profile_id": self.app_profile_id, - }, + true_mutations=true_case_list, + false_mutations=false_case_list, + predicate_filter=predicate._to_pb() if predicate is not None else None, + row_key=row_key.encode("utf-8") if isinstance(row_key, str) else row_key, + table_name=self.table_name, + app_profile_id=self.app_profile_id, metadata=metadata, timeout=operation_timeout, retry=None, @@ -1123,25 +1120,21 @@ async def read_modify_write_row( operation Raises: - GoogleAPIError exceptions from grpc call + - ValueError if invalid arguments are provided """ operation_timeout, _ = _get_timeouts(operation_timeout, None, self) - row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key if operation_timeout <= 0: raise ValueError("operation_timeout must be greater than 0") if rules is not None and not isinstance(rules, list): rules = [rules] if not rules: raise ValueError("rules must contain at least one item") - # concert to dict representation - rules_dict = [rule._to_dict() for rule in rules] metadata = _make_metadata(self.table_name, self.app_profile_id) result = await self.client._gapic_client.read_modify_write_row( - request={ - "rules": rules_dict, - "table_name": self.table_name, - "row_key": row_key, - "app_profile_id": self.app_profile_id, - }, + rules=[rule._to_pb() for rule in rules], + row_key=row_key.encode("utf-8") if isinstance(row_key, str) else row_key, + table_name=self.table_name, + app_profile_id=self.app_profile_id, metadata=metadata, timeout=operation_timeout, retry=None, diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index 7ff5f9a0b..91d2b11e1 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -342,9 +342,6 @@ async def _execute_mutate_rows( - list of FailedMutationEntryError objects for mutations that failed. FailedMutationEntryError objects will not contain index information """ - request = {"table_name": self._table.table_name} - if self._table.app_profile_id: - request["app_profile_id"] = self._table.app_profile_id try: operation = _MutateRowsOperationAsync( self._table.client._gapic_client, diff --git a/google/cloud/bigtable/data/mutations.py b/google/cloud/bigtable/data/mutations.py index 06db21879..b5729d25e 100644 --- a/google/cloud/bigtable/data/mutations.py +++ b/google/cloud/bigtable/data/mutations.py @@ -19,9 +19,12 @@ from abc import ABC, abstractmethod from sys import getsizeof +import google.cloud.bigtable_v2.types.bigtable as types_pb +import google.cloud.bigtable_v2.types.data as data_pb from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE + # special value for SetCell mutation timestamps. If set, server will assign a timestamp _SERVER_SIDE_TIMESTAMP = -1 @@ -36,6 +39,12 @@ class Mutation(ABC): def _to_dict(self) -> dict[str, Any]: raise NotImplementedError + def _to_pb(self) -> data_pb.Mutation: + """ + Convert the mutation to protobuf + """ + return data_pb.Mutation(**self._to_dict()) + def is_idempotent(self) -> bool: """ Check if the mutation is idempotent @@ -221,6 +230,12 @@ def _to_dict(self) -> dict[str, Any]: "mutations": [mutation._to_dict() for mutation in self.mutations], } + def _to_pb(self) -> types_pb.MutateRowsRequest.Entry: + return types_pb.MutateRowsRequest.Entry( + row_key=self.row_key, + mutations=[mutation._to_pb() for mutation in self.mutations], + ) + def is_idempotent(self) -> bool: """Check if the mutation is idempotent""" return all(mutation.is_idempotent() for mutation in self.mutations) diff --git a/google/cloud/bigtable/data/read_modify_write_rules.py b/google/cloud/bigtable/data/read_modify_write_rules.py index 3a3eb3752..f43dbe79f 100644 --- a/google/cloud/bigtable/data/read_modify_write_rules.py +++ b/google/cloud/bigtable/data/read_modify_write_rules.py @@ -16,6 +16,8 @@ import abc +import google.cloud.bigtable_v2.types.data as data_pb + # value must fit in 64-bit signed integer _MAX_INCREMENT_VALUE = (1 << 63) - 1 @@ -29,9 +31,12 @@ def __init__(self, family: str, qualifier: bytes | str): self.qualifier = qualifier @abc.abstractmethod - def _to_dict(self): + def _to_dict(self) -> dict[str, str | bytes | int]: raise NotImplementedError + def _to_pb(self) -> data_pb.ReadModifyWriteRule: + return data_pb.ReadModifyWriteRule(**self._to_dict()) + class IncrementRule(ReadModifyWriteRule): def __init__(self, family: str, qualifier: bytes | str, increment_amount: int = 1): @@ -44,7 +49,7 @@ def __init__(self, family: str, qualifier: bytes | str, increment_amount: int = super().__init__(family, qualifier) self.increment_amount = increment_amount - def _to_dict(self): + def _to_dict(self) -> dict[str, str | bytes | int]: return { "family_name": self.family, "column_qualifier": self.qualifier, @@ -64,7 +69,7 @@ def __init__(self, family: str, qualifier: bytes | str, append_value: bytes | st super().__init__(family, qualifier) self.append_value = append_value - def _to_dict(self): + def _to_dict(self) -> dict[str, str | bytes | int]: return { "family_name": self.family, "column_qualifier": self.qualifier, diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index eae3483ed..89a153af2 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -75,6 +75,7 @@ def test_ctor(self): """ test that constructor sets all the attributes correctly """ + from google.cloud.bigtable.data._async._mutate_rows import _EntryWithProto from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import ServiceUnavailable @@ -103,7 +104,8 @@ def test_ctor(self): assert str(table.table_name) in metadata[0][1] assert str(table.app_profile_id) in metadata[0][1] # entries should be passed down - assert instance.mutations == entries + entries_w_pb = [_EntryWithProto(e, e._to_pb()) for e in entries] + assert instance.mutations == entries_w_pb # timeout_gen should generate per-attempt timeout assert next(instance.timeout_generator) == attempt_timeout # ensure predicate is set @@ -306,7 +308,7 @@ async def test_run_attempt_single_entry_success(self): assert mock_gapic_fn.call_count == 1 _, kwargs = mock_gapic_fn.call_args assert kwargs["timeout"] == expected_timeout - assert kwargs["entries"] == [mutation._to_dict()] + assert kwargs["entries"] == [mutation._to_pb()] @pytest.mark.asyncio async def test_run_attempt_empty_request(self): diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 4ae46da6e..7afecc5b0 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -2032,18 +2032,17 @@ async def test_mutate_row(self, mutation_arg): ) assert mock_gapic.call_count == 1 kwargs = mock_gapic.call_args_list[0].kwargs - request = mock_gapic.call_args[0][0] assert ( - request["table_name"] + kwargs["table_name"] == "projects/project/instances/instance/tables/table" ) - assert request["row_key"] == b"row_key" + assert kwargs["row_key"] == b"row_key" formatted_mutations = ( - [mutation._to_dict() for mutation in mutation_arg] + [mutation._to_pb() for mutation in mutation_arg] if isinstance(mutation_arg, list) - else [mutation_arg._to_dict()] + else [mutation_arg._to_pb()] ) - assert request["mutations"] == formatted_mutations + assert kwargs["mutations"] == formatted_mutations assert kwargs["timeout"] == expected_attempt_timeout # make sure gapic layer is not retrying assert kwargs["retry"] is None @@ -2146,7 +2145,7 @@ async def test_mutate_row_metadata(self, include_app_profile): with mock.patch.object( client._gapic_client, "mutate_row", AsyncMock() ) as read_rows: - await table.mutate_row("rk", {}) + await table.mutate_row("rk", mock.Mock()) kwargs = read_rows.call_args_list[0].kwargs metadata = kwargs["metadata"] goog_metadata = None @@ -2160,6 +2159,15 @@ async def test_mutate_row_metadata(self, include_app_profile): else: assert "app_profile_id=" not in goog_metadata + @pytest.mark.parametrize("mutations", [[], None]) + @pytest.mark.asyncio + async def test_mutate_row_no_mutations(self, mutations): + async with self._make_client() as client: + async with client.get_table("instance", "table") as table: + with pytest.raises(ValueError) as e: + await table.mutate_row("key", mutations=mutations) + assert e.value.args[0] == "No mutations provided" + class TestBulkMutateRows: def _make_client(self, *args, **kwargs): @@ -2232,7 +2240,7 @@ async def test_bulk_mutate_rows(self, mutation_arg): kwargs["table_name"] == "projects/project/instances/instance/tables/table" ) - assert kwargs["entries"] == [bulk_mutation._to_dict()] + assert kwargs["entries"] == [bulk_mutation._to_pb()] assert kwargs["timeout"] == expected_attempt_timeout assert kwargs["retry"] is None @@ -2257,8 +2265,8 @@ async def test_bulk_mutate_rows_multiple_entries(self): kwargs["table_name"] == "projects/project/instances/instance/tables/table" ) - assert kwargs["entries"][0] == entry_1._to_dict() - assert kwargs["entries"][1] == entry_2._to_dict() + assert kwargs["entries"][0] == entry_1._to_pb() + assert kwargs["entries"][1] == entry_2._to_pb() @pytest.mark.asyncio @pytest.mark.parametrize( @@ -2587,17 +2595,16 @@ async def test_check_and_mutate(self, gapic_result): ) assert found == gapic_result kwargs = mock_gapic.call_args[1] - request = kwargs["request"] - assert request["table_name"] == table.table_name - assert request["row_key"] == row_key - assert request["predicate_filter"] == predicate - assert request["true_mutations"] == [ - m._to_dict() for m in true_mutations + assert kwargs["table_name"] == table.table_name + assert kwargs["row_key"] == row_key + assert kwargs["predicate_filter"] == predicate + assert kwargs["true_mutations"] == [ + m._to_pb() for m in true_mutations ] - assert request["false_mutations"] == [ - m._to_dict() for m in false_mutations + assert kwargs["false_mutations"] == [ + m._to_pb() for m in false_mutations ] - assert request["app_profile_id"] == app_profile + assert kwargs["app_profile_id"] == app_profile assert kwargs["timeout"] == operation_timeout assert kwargs["retry"] is None @@ -2655,9 +2662,8 @@ async def test_check_and_mutate_single_mutations(self): false_case_mutations=false_mutation, ) kwargs = mock_gapic.call_args[1] - request = kwargs["request"] - assert request["true_mutations"] == [true_mutation._to_dict()] - assert request["false_mutations"] == [false_mutation._to_dict()] + assert kwargs["true_mutations"] == [true_mutation._to_pb()] + assert kwargs["false_mutations"] == [false_mutation._to_pb()] @pytest.mark.asyncio async def test_check_and_mutate_predicate_object(self): @@ -2665,8 +2671,8 @@ async def test_check_and_mutate_predicate_object(self): from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse mock_predicate = mock.Mock() - predicate_dict = {"predicate": "dict"} - mock_predicate._to_dict.return_value = predicate_dict + predicate_pb = {"predicate": "dict"} + mock_predicate._to_pb.return_value = predicate_pb async with self._make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( @@ -2681,19 +2687,19 @@ async def test_check_and_mutate_predicate_object(self): false_case_mutations=[mock.Mock()], ) kwargs = mock_gapic.call_args[1] - assert kwargs["request"]["predicate_filter"] == predicate_dict - assert mock_predicate._to_dict.call_count == 1 + assert kwargs["predicate_filter"] == predicate_pb + assert mock_predicate._to_pb.call_count == 1 assert kwargs["retry"] is None @pytest.mark.asyncio async def test_check_and_mutate_mutations_parsing(self): - """mutations objects should be converted to dicts""" + """mutations objects should be converted to protos""" from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse from google.cloud.bigtable.data.mutations import DeleteAllFromRow mutations = [mock.Mock() for _ in range(5)] for idx, mutation in enumerate(mutations): - mutation._to_dict.return_value = {"fake": idx} + mutation._to_pb.return_value = f"fake {idx}" mutations.append(DeleteAllFromRow()) async with self._make_client() as client: async with client.get_table("instance", "table") as table: @@ -2709,16 +2715,16 @@ async def test_check_and_mutate_mutations_parsing(self): true_case_mutations=mutations[0:2], false_case_mutations=mutations[2:], ) - kwargs = mock_gapic.call_args[1]["request"] - assert kwargs["true_mutations"] == [{"fake": 0}, {"fake": 1}] + kwargs = mock_gapic.call_args[1] + assert kwargs["true_mutations"] == ["fake 0", "fake 1"] assert kwargs["false_mutations"] == [ - {"fake": 2}, - {"fake": 3}, - {"fake": 4}, - {"delete_from_row": {}}, + "fake 2", + "fake 3", + "fake 4", + DeleteAllFromRow()._to_pb(), ] assert all( - mutation._to_dict.call_count == 1 for mutation in mutations[:5] + mutation._to_pb.call_count == 1 for mutation in mutations[:5] ) @pytest.mark.parametrize("include_app_profile", [True, False]) @@ -2757,18 +2763,18 @@ def _make_client(self, *args, **kwargs): [ ( AppendValueRule("f", "c", b"1"), - [AppendValueRule("f", "c", b"1")._to_dict()], + [AppendValueRule("f", "c", b"1")._to_pb()], ), ( [AppendValueRule("f", "c", b"1")], - [AppendValueRule("f", "c", b"1")._to_dict()], + [AppendValueRule("f", "c", b"1")._to_pb()], ), - (IncrementRule("f", "c", 1), [IncrementRule("f", "c", 1)._to_dict()]), + (IncrementRule("f", "c", 1), [IncrementRule("f", "c", 1)._to_pb()]), ( [AppendValueRule("f", "c", b"1"), IncrementRule("f", "c", 1)], [ - AppendValueRule("f", "c", b"1")._to_dict(), - IncrementRule("f", "c", 1)._to_dict(), + AppendValueRule("f", "c", b"1")._to_pb(), + IncrementRule("f", "c", 1)._to_pb(), ], ), ], @@ -2786,7 +2792,7 @@ async def test_read_modify_write_call_rule_args(self, call_rules, expected_rules await table.read_modify_write_row("key", call_rules) assert mock_gapic.call_count == 1 found_kwargs = mock_gapic.call_args_list[0][1] - assert found_kwargs["request"]["rules"] == expected_rules + assert found_kwargs["rules"] == expected_rules assert found_kwargs["retry"] is None @pytest.mark.parametrize("rules", [[], None]) @@ -2811,15 +2817,14 @@ async def test_read_modify_write_call_defaults(self): ) as mock_gapic: await table.read_modify_write_row(row_key, mock.Mock()) assert mock_gapic.call_count == 1 - found_kwargs = mock_gapic.call_args_list[0][1] - request = found_kwargs["request"] + kwargs = mock_gapic.call_args_list[0][1] assert ( - request["table_name"] + kwargs["table_name"] == f"projects/{project}/instances/{instance}/tables/{table_id}" ) - assert request["app_profile_id"] is None - assert request["row_key"] == row_key.encode() - assert found_kwargs["timeout"] > 1 + assert kwargs["app_profile_id"] is None + assert kwargs["row_key"] == row_key.encode() + assert kwargs["timeout"] > 1 @pytest.mark.asyncio async def test_read_modify_write_call_overrides(self): @@ -2839,11 +2844,10 @@ async def test_read_modify_write_call_overrides(self): operation_timeout=expected_timeout, ) assert mock_gapic.call_count == 1 - found_kwargs = mock_gapic.call_args_list[0][1] - request = found_kwargs["request"] - assert request["app_profile_id"] is profile_id - assert request["row_key"] == row_key - assert found_kwargs["timeout"] == expected_timeout + kwargs = mock_gapic.call_args_list[0][1] + assert kwargs["app_profile_id"] is profile_id + assert kwargs["row_key"] == row_key + assert kwargs["timeout"] == expected_timeout @pytest.mark.asyncio async def test_read_modify_write_string_key(self): @@ -2855,8 +2859,8 @@ async def test_read_modify_write_string_key(self): ) as mock_gapic: await table.read_modify_write_row(row_key, mock.Mock()) assert mock_gapic.call_count == 1 - found_kwargs = mock_gapic.call_args_list[0][1] - assert found_kwargs["request"]["row_key"] == row_key.encode() + kwargs = mock_gapic.call_args_list[0][1] + assert kwargs["row_key"] == row_key.encode() @pytest.mark.asyncio async def test_read_modify_write_row_building(self): diff --git a/tests/unit/data/test_mutations.py b/tests/unit/data/test_mutations.py index 8680a8da9..485c86e42 100644 --- a/tests/unit/data/test_mutations.py +++ b/tests/unit/data/test_mutations.py @@ -307,6 +307,42 @@ def test__to_dict_server_timestamp(self): assert got_inner_dict["value"] == expected_value assert len(got_inner_dict.keys()) == 4 + def test__to_pb(self): + """ensure proto representation is as expected""" + import google.cloud.bigtable_v2.types.data as data_pb + + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + expected_value = b"test-value" + expected_timestamp = 123456789 + instance = self._make_one( + expected_family, expected_qualifier, expected_value, expected_timestamp + ) + got_pb = instance._to_pb() + assert isinstance(got_pb, data_pb.Mutation) + assert got_pb.set_cell.family_name == expected_family + assert got_pb.set_cell.column_qualifier == expected_qualifier + assert got_pb.set_cell.timestamp_micros == expected_timestamp + assert got_pb.set_cell.value == expected_value + + def test__to_pb_server_timestamp(self): + """test with server side timestamp -1 value""" + import google.cloud.bigtable_v2.types.data as data_pb + + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + expected_value = b"test-value" + expected_timestamp = -1 + instance = self._make_one( + expected_family, expected_qualifier, expected_value, expected_timestamp + ) + got_pb = instance._to_pb() + assert isinstance(got_pb, data_pb.Mutation) + assert got_pb.set_cell.family_name == expected_family + assert got_pb.set_cell.column_qualifier == expected_qualifier + assert got_pb.set_cell.timestamp_micros == expected_timestamp + assert got_pb.set_cell.value == expected_value + @pytest.mark.parametrize( "timestamp,expected_value", [ @@ -406,6 +442,18 @@ def test__to_dict(self, start, end): if end is not None: assert time_range_dict["end_timestamp_micros"] == end + def test__to_pb(self): + """ensure proto representation is as expected""" + import google.cloud.bigtable_v2.types.data as data_pb + + expected_family = "test-family" + expected_qualifier = b"test-qualifier" + instance = self._make_one(expected_family, expected_qualifier) + got_pb = instance._to_pb() + assert isinstance(got_pb, data_pb.Mutation) + assert got_pb.delete_from_column.family_name == expected_family + assert got_pb.delete_from_column.column_qualifier == expected_qualifier + def test_is_idempotent(self): """is_idempotent is always true""" instance = self._make_one( @@ -445,6 +493,16 @@ def test__to_dict(self): assert len(got_inner_dict.keys()) == 1 assert got_inner_dict["family_name"] == expected_family + def test__to_pb(self): + """ensure proto representation is as expected""" + import google.cloud.bigtable_v2.types.data as data_pb + + expected_family = "test-family" + instance = self._make_one(expected_family) + got_pb = instance._to_pb() + assert isinstance(got_pb, data_pb.Mutation) + assert got_pb.delete_from_family.family_name == expected_family + def test_is_idempotent(self): """is_idempotent is always true""" instance = self._make_one("test-family") @@ -477,6 +535,15 @@ def test__to_dict(self): assert list(got_dict.keys()) == ["delete_from_row"] assert len(got_dict["delete_from_row"].keys()) == 0 + def test__to_pb(self): + """ensure proto representation is as expected""" + import google.cloud.bigtable_v2.types.data as data_pb + + instance = self._make_one() + got_pb = instance._to_pb() + assert isinstance(got_pb, data_pb.Mutation) + assert "delete_from_row" in str(got_pb) + def test_is_idempotent(self): """is_idempotent is always true""" instance = self._make_one() @@ -550,6 +617,23 @@ def test__to_dict(self): assert instance._to_dict() == expected_result assert mutation_mock._to_dict.call_count == n_mutations + def test__to_pb(self): + from google.cloud.bigtable_v2.types.bigtable import MutateRowsRequest + from google.cloud.bigtable_v2.types.data import Mutation + + expected_key = "row_key" + mutation_mock = mock.Mock() + n_mutations = 3 + expected_mutations = [mutation_mock for i in range(n_mutations)] + for mock_mutations in expected_mutations: + mock_mutations._to_pb.return_value = Mutation() + instance = self._make_one(expected_key, expected_mutations) + pb_result = instance._to_pb() + assert isinstance(pb_result, MutateRowsRequest.Entry) + assert pb_result.row_key == b"row_key" + assert pb_result.mutations == [Mutation()] * n_mutations + assert mutation_mock._to_pb.call_count == n_mutations + @pytest.mark.parametrize( "mutations,result", [ diff --git a/tests/unit/data/test_read_modify_write_rules.py b/tests/unit/data/test_read_modify_write_rules.py index aeb41f19c..1f67da13b 100644 --- a/tests/unit/data/test_read_modify_write_rules.py +++ b/tests/unit/data/test_read_modify_write_rules.py @@ -36,6 +36,9 @@ def test_abstract(self): self._target_class()(family="foo", qualifier=b"bar") def test__to_dict(self): + """ + to_dict not implemented in base class + """ with pytest.raises(NotImplementedError): self._target_class()._to_dict(mock.Mock()) @@ -97,6 +100,27 @@ def test__to_dict(self, args, expected): } assert instance._to_dict() == expected + @pytest.mark.parametrize( + "args,expected", + [ + (("fam", b"qual", 1), ("fam", b"qual", 1)), + (("fam", b"qual", -12), ("fam", b"qual", -12)), + (("fam", "qual", 1), ("fam", b"qual", 1)), + (("fam", "qual", 0), ("fam", b"qual", 0)), + (("", "", 0), ("", b"", 0)), + (("f", b"q"), ("f", b"q", 1)), + ], + ) + def test__to_pb(self, args, expected): + import google.cloud.bigtable_v2.types.data as data_pb + + instance = self._target_class()(*args) + pb_result = instance._to_pb() + assert isinstance(pb_result, data_pb.ReadModifyWriteRule) + assert pb_result.family_name == expected[0] + assert pb_result.column_qualifier == expected[1] + assert pb_result.increment_amount == expected[2] + class TestAppendValueRule: def _target_class(self): @@ -142,3 +166,21 @@ def test__to_dict(self, args, expected): "append_value": expected[2], } assert instance._to_dict() == expected + + @pytest.mark.parametrize( + "args,expected", + [ + (("fam", b"qual", b"val"), ("fam", b"qual", b"val")), + (("fam", "qual", b"val"), ("fam", b"qual", b"val")), + (("", "", b""), ("", b"", b"")), + ], + ) + def test__to_pb(self, args, expected): + import google.cloud.bigtable_v2.types.data as data_pb + + instance = self._target_class()(*args) + pb_result = instance._to_pb() + assert isinstance(pb_result, data_pb.ReadModifyWriteRule) + assert pb_result.family_name == expected[0] + assert pb_result.column_qualifier == expected[1] + assert pb_result.append_value == expected[2] From b1914517bdc2b42d4c4d8071359d9ed70ac79785 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 1 Dec 2023 13:09:59 -0800 Subject: [PATCH 28/56] chore: optimize gapic calls (#863) --- .../services/bigtable/async_client.py | 98 ++++++++----------- .../bigtable/transports/grpc_asyncio.py | 62 ++++++++++++ tests/unit/data/_async/test_client.py | 2 +- 3 files changed, 105 insertions(+), 57 deletions(-) diff --git a/google/cloud/bigtable_v2/services/bigtable/async_client.py b/google/cloud/bigtable_v2/services/bigtable/async_client.py index a80be70af..d325564c0 100644 --- a/google/cloud/bigtable_v2/services/bigtable/async_client.py +++ b/google/cloud/bigtable_v2/services/bigtable/async_client.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import functools from collections import OrderedDict import functools import re @@ -272,7 +273,8 @@ def read_rows( "the individual field arguments should be set." ) - request = bigtable.ReadRowsRequest(request) + if not isinstance(request, bigtable.ReadRowsRequest): + request = bigtable.ReadRowsRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. @@ -283,12 +285,9 @@ def read_rows( # Wrap the RPC method; this adds retry and timeout information, # and friendly error handling. - rpc = gapic_v1.method_async.wrap_method( - self._client._transport.read_rows, - default_timeout=43200.0, - client_info=DEFAULT_CLIENT_INFO, - ) - + rpc = self._client._transport._wrapped_methods[ + self._client._transport.read_rows + ] # Certain fields should be provided within the metadata header; # add these here. metadata = tuple(metadata) + ( @@ -367,7 +366,8 @@ def sample_row_keys( "the individual field arguments should be set." ) - request = bigtable.SampleRowKeysRequest(request) + if not isinstance(request, bigtable.SampleRowKeysRequest): + request = bigtable.SampleRowKeysRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. @@ -378,12 +378,9 @@ def sample_row_keys( # Wrap the RPC method; this adds retry and timeout information, # and friendly error handling. - rpc = gapic_v1.method_async.wrap_method( - self._client._transport.sample_row_keys, - default_timeout=60.0, - client_info=DEFAULT_CLIENT_INFO, - ) - + rpc = self._client._transport._wrapped_methods[ + self._client._transport.sample_row_keys + ] # Certain fields should be provided within the metadata header; # add these here. metadata = tuple(metadata) + ( @@ -479,7 +476,8 @@ async def mutate_row( "the individual field arguments should be set." ) - request = bigtable.MutateRowRequest(request) + if not isinstance(request, bigtable.MutateRowRequest): + request = bigtable.MutateRowRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. @@ -494,21 +492,9 @@ async def mutate_row( # Wrap the RPC method; this adds retry and timeout information, # and friendly error handling. - rpc = gapic_v1.method_async.wrap_method( - self._client._transport.mutate_row, - default_retry=retries.Retry( - initial=0.01, - maximum=60.0, - multiplier=2, - predicate=retries.if_exception_type( - core_exceptions.DeadlineExceeded, - core_exceptions.ServiceUnavailable, - ), - deadline=60.0, - ), - default_timeout=60.0, - client_info=DEFAULT_CLIENT_INFO, - ) + rpc = self._client._transport._wrapped_methods[ + self._client._transport.mutate_row + ] # Certain fields should be provided within the metadata header; # add these here. @@ -601,7 +587,8 @@ def mutate_rows( "the individual field arguments should be set." ) - request = bigtable.MutateRowsRequest(request) + if not isinstance(request, bigtable.MutateRowsRequest): + request = bigtable.MutateRowsRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. @@ -614,11 +601,9 @@ def mutate_rows( # Wrap the RPC method; this adds retry and timeout information, # and friendly error handling. - rpc = gapic_v1.method_async.wrap_method( - self._client._transport.mutate_rows, - default_timeout=600.0, - client_info=DEFAULT_CLIENT_INFO, - ) + rpc = self._client._transport._wrapped_methods[ + self._client._transport.mutate_rows + ] # Certain fields should be provided within the metadata header; # add these here. @@ -749,7 +734,8 @@ async def check_and_mutate_row( "the individual field arguments should be set." ) - request = bigtable.CheckAndMutateRowRequest(request) + if not isinstance(request, bigtable.CheckAndMutateRowRequest): + request = bigtable.CheckAndMutateRowRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. @@ -768,11 +754,9 @@ async def check_and_mutate_row( # Wrap the RPC method; this adds retry and timeout information, # and friendly error handling. - rpc = gapic_v1.method_async.wrap_method( - self._client._transport.check_and_mutate_row, - default_timeout=20.0, - client_info=DEFAULT_CLIENT_INFO, - ) + rpc = self._client._transport._wrapped_methods[ + self._client._transport.check_and_mutate_row + ] # Certain fields should be provided within the metadata header; # add these here. @@ -851,7 +835,8 @@ async def ping_and_warm( "the individual field arguments should be set." ) - request = bigtable.PingAndWarmRequest(request) + if not isinstance(request, bigtable.PingAndWarmRequest): + request = bigtable.PingAndWarmRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. @@ -862,11 +847,9 @@ async def ping_and_warm( # Wrap the RPC method; this adds retry and timeout information, # and friendly error handling. - rpc = gapic_v1.method_async.wrap_method( - self._client._transport.ping_and_warm, - default_timeout=None, - client_info=DEFAULT_CLIENT_INFO, - ) + rpc = self._client._transport._wrapped_methods[ + self._client._transport.ping_and_warm + ] # Certain fields should be provided within the metadata header; # add these here. @@ -968,7 +951,8 @@ async def read_modify_write_row( "the individual field arguments should be set." ) - request = bigtable.ReadModifyWriteRowRequest(request) + if not isinstance(request, bigtable.ReadModifyWriteRowRequest): + request = bigtable.ReadModifyWriteRowRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. @@ -983,11 +967,9 @@ async def read_modify_write_row( # Wrap the RPC method; this adds retry and timeout information, # and friendly error handling. - rpc = gapic_v1.method_async.wrap_method( - self._client._transport.read_modify_write_row, - default_timeout=20.0, - client_info=DEFAULT_CLIENT_INFO, - ) + rpc = self._client._transport._wrapped_methods[ + self._client._transport.read_modify_write_row + ] # Certain fields should be provided within the metadata header; # add these here. @@ -1076,7 +1058,10 @@ def generate_initial_change_stream_partitions( "the individual field arguments should be set." ) - request = bigtable.GenerateInitialChangeStreamPartitionsRequest(request) + if not isinstance( + request, bigtable.GenerateInitialChangeStreamPartitionsRequest + ): + request = bigtable.GenerateInitialChangeStreamPartitionsRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. @@ -1174,7 +1159,8 @@ def read_change_stream( "the individual field arguments should be set." ) - request = bigtable.ReadChangeStreamRequest(request) + if not isinstance(request, bigtable.ReadChangeStreamRequest): + request = bigtable.ReadChangeStreamRequest(request) # If we have keyword arguments corresponding to fields on the # request, apply these. diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py b/google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py index 8bf02ce77..3450d4969 100644 --- a/google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py +++ b/google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py @@ -18,6 +18,8 @@ from google.api_core import gapic_v1 from google.api_core import grpc_helpers_async +from google.api_core import exceptions as core_exceptions +from google.api_core import retry as retries from google.auth import credentials as ga_credentials # type: ignore from google.auth.transport.grpc import SslCredentials # type: ignore @@ -512,6 +514,66 @@ def read_change_stream( ) return self._stubs["read_change_stream"] + def _prep_wrapped_messages(self, client_info): + # Precompute the wrapped methods. + self._wrapped_methods = { + self.read_rows: gapic_v1.method_async.wrap_method( + self.read_rows, + default_timeout=43200.0, + client_info=client_info, + ), + self.sample_row_keys: gapic_v1.method_async.wrap_method( + self.sample_row_keys, + default_timeout=60.0, + client_info=client_info, + ), + self.mutate_row: gapic_v1.method_async.wrap_method( + self.mutate_row, + default_retry=retries.Retry( + initial=0.01, + maximum=60.0, + multiplier=2, + predicate=retries.if_exception_type( + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + ), + deadline=60.0, + ), + default_timeout=60.0, + client_info=client_info, + ), + self.mutate_rows: gapic_v1.method_async.wrap_method( + self.mutate_rows, + default_timeout=600.0, + client_info=client_info, + ), + self.check_and_mutate_row: gapic_v1.method_async.wrap_method( + self.check_and_mutate_row, + default_timeout=20.0, + client_info=client_info, + ), + self.ping_and_warm: gapic_v1.method_async.wrap_method( + self.ping_and_warm, + default_timeout=None, + client_info=client_info, + ), + self.read_modify_write_row: gapic_v1.method_async.wrap_method( + self.read_modify_write_row, + default_timeout=20.0, + client_info=client_info, + ), + self.generate_initial_change_stream_partitions: gapic_v1.method_async.wrap_method( + self.generate_initial_change_stream_partitions, + default_timeout=60.0, + client_info=client_info, + ), + self.read_change_stream: gapic_v1.method_async.wrap_method( + self.read_change_stream, + default_timeout=43200.0, + client_info=client_info, + ), + } + def close(self): return self.grpc_channel.close() diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 7afecc5b0..7718246fc 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -141,7 +141,7 @@ async def test_ctor_dict_options(self): async def test_veneer_grpc_headers(self): # client_info should be populated with headers to # detect as a veneer client - patch = mock.patch("google.api_core.gapic_v1.method.wrap_method") + patch = mock.patch("google.api_core.gapic_v1.method_async.wrap_method") with patch as gapic_mock: client = self._make_one(project="project-id") wrapped_call_list = gapic_mock.call_args_list From 285cdd3bf74a871e5fa8134029e3ac060d2622f5 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 1 Dec 2023 14:04:31 -0800 Subject: [PATCH 29/56] feat: expose retryable error codes to users (#879) --- .../bigtable/data/_async/_mutate_rows.py | 6 +- .../cloud/bigtable/data/_async/_read_rows.py | 15 +- google/cloud/bigtable/data/_async/client.py | 140 +++++-- .../bigtable/data/_async/mutations_batcher.py | 12 +- google/cloud/bigtable/data/_helpers.py | 31 +- tests/unit/data/_async/test__mutate_rows.py | 26 +- tests/unit/data/_async/test_client.py | 358 ++++++++++-------- .../data/_async/test_mutations_batcher.py | 82 +++- tests/unit/data/test__helpers.py | 48 +++ 9 files changed, 514 insertions(+), 204 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 5bf759151..5971a9894 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -14,7 +14,7 @@ # from __future__ import annotations -from typing import TYPE_CHECKING +from typing import Sequence, TYPE_CHECKING import asyncio from dataclasses import dataclass import functools @@ -66,6 +66,7 @@ def __init__( mutation_entries: list["RowMutationEntry"], operation_timeout: float, attempt_timeout: float | None, + retryable_exceptions: Sequence[type[Exception]] = (), ): """ Args: @@ -96,8 +97,7 @@ def __init__( # create predicate for determining which errors are retryable self.is_retryable = retries.if_exception_type( # RPC level errors - core_exceptions.DeadlineExceeded, - core_exceptions.ServiceUnavailable, + *retryable_exceptions, # Entry level errors bt_exceptions._MutateRowsIncomplete, ) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 90cc7e87c..ad1f7b84d 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -15,7 +15,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING, AsyncGenerator, AsyncIterable, Awaitable +from typing import ( + TYPE_CHECKING, + AsyncGenerator, + AsyncIterable, + Awaitable, + Sequence, +) from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB @@ -74,6 +80,7 @@ def __init__( table: "TableAsync", operation_timeout: float, attempt_timeout: float, + retryable_exceptions: Sequence[type[Exception]] = (), ): self.attempt_timeout_gen = _attempt_timeout_generator( attempt_timeout, operation_timeout @@ -88,11 +95,7 @@ def __init__( else: self.request = query._to_pb(table) self.table = table - self._predicate = retries.if_exception_type( - core_exceptions.DeadlineExceeded, - core_exceptions.ServiceUnavailable, - core_exceptions.Aborted, - ) + self._predicate = retries.if_exception_type(*retryable_exceptions) self._metadata = _make_metadata( table.table_name, table.app_profile_id, diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index ab8cc48f8..a79ead7f8 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -21,6 +21,7 @@ AsyncIterable, Optional, Set, + Sequence, TYPE_CHECKING, ) @@ -45,7 +46,9 @@ from google.api_core.exceptions import GoogleAPICallError from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore from google.api_core import retry_async as retries -from google.api_core import exceptions as core_exceptions +from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import ServiceUnavailable +from google.api_core.exceptions import Aborted from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync import google.auth.credentials @@ -64,6 +67,7 @@ from google.cloud.bigtable.data._helpers import _make_metadata from google.cloud.bigtable.data._helpers import _convert_retry_deadline from google.cloud.bigtable.data._helpers import _validate_timeouts +from google.cloud.bigtable.data._helpers import _get_retryable_errors from google.cloud.bigtable.data._helpers import _get_timeouts from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync @@ -366,21 +370,10 @@ async def _remove_instance_registration( except KeyError: return False - def get_table( - self, - instance_id: str, - table_id: str, - app_profile_id: str | None = None, - *, - default_read_rows_operation_timeout: float = 600, - default_read_rows_attempt_timeout: float | None = None, - default_mutate_rows_operation_timeout: float = 600, - default_mutate_rows_attempt_timeout: float | None = None, - default_operation_timeout: float = 60, - default_attempt_timeout: float | None = None, - ) -> TableAsync: + def get_table(self, instance_id: str, table_id: str, *args, **kwargs) -> TableAsync: """ - Returns a table instance for making data API requests + Returns a table instance for making data API requests. All arguments are passed + directly to the TableAsync constructor. Args: instance_id: The Bigtable instance ID to associate with this client. @@ -402,15 +395,17 @@ def get_table( seconds. If not set, defaults to 60 seconds default_attempt_timeout: The default timeout for all other individual rpc requests, in seconds. If not set, defaults to 20 seconds + default_read_rows_retryable_errors: a list of errors that will be retried + if encountered during read_rows and related operations. + Defaults to 4 (DeadlineExceeded), 14 (ServiceUnavailable), and 10 (Aborted) + default_mutate_rows_retryable_errors: a list of errors that will be retried + if encountered during mutate_rows and related operations. + Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable) + default_retryable_errors: a list of errors that will be retried if + encountered during all other operations. + Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable) """ - return TableAsync( - self, - instance_id, - table_id, - app_profile_id, - default_operation_timeout=default_operation_timeout, - default_attempt_timeout=default_attempt_timeout, - ) + return TableAsync(self, instance_id, table_id, *args, **kwargs) async def __aenter__(self): self._start_background_channel_refresh() @@ -442,6 +437,19 @@ def __init__( default_mutate_rows_attempt_timeout: float | None = 60, default_operation_timeout: float = 60, default_attempt_timeout: float | None = 20, + default_read_rows_retryable_errors: Sequence[type[Exception]] = ( + DeadlineExceeded, + ServiceUnavailable, + Aborted, + ), + default_mutate_rows_retryable_errors: Sequence[type[Exception]] = ( + DeadlineExceeded, + ServiceUnavailable, + ), + default_retryable_errors: Sequence[type[Exception]] = ( + DeadlineExceeded, + ServiceUnavailable, + ), ): """ Initialize a Table instance @@ -468,9 +476,20 @@ def __init__( seconds. If not set, defaults to 60 seconds default_attempt_timeout: The default timeout for all other individual rpc requests, in seconds. If not set, defaults to 20 seconds + default_read_rows_retryable_errors: a list of errors that will be retried + if encountered during read_rows and related operations. + Defaults to 4 (DeadlineExceeded), 14 (ServiceUnavailable), and 10 (Aborted) + default_mutate_rows_retryable_errors: a list of errors that will be retried + if encountered during mutate_rows and related operations. + Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable) + default_retryable_errors: a list of errors that will be retried if + encountered during all other operations. + Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable) Raises: - RuntimeError if called outside of an async context (no running event loop) """ + # NOTE: any changes to the signature of this method should also be reflected + # in client.get_table() # validate timeouts _validate_timeouts( default_operation_timeout, default_attempt_timeout, allow_none=True @@ -506,6 +525,14 @@ def __init__( ) self.default_mutate_rows_attempt_timeout = default_mutate_rows_attempt_timeout + self.default_read_rows_retryable_errors = ( + default_read_rows_retryable_errors or () + ) + self.default_mutate_rows_retryable_errors = ( + default_mutate_rows_retryable_errors or () + ) + self.default_retryable_errors = default_retryable_errors or () + # raises RuntimeError if called outside of an async context (no running event loop) try: self._register_instance_task = asyncio.create_task( @@ -522,12 +549,15 @@ async def read_rows_stream( *, operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> AsyncIterable[Row]: """ Read a set of rows from the table, based on the specified query. Returns an iterator to asynchronously stream back row data. - Failed requests within operation_timeout will be retried. + Failed requests within operation_timeout will be retried based on the + retryable_errors list until operation_timeout is reached. Args: - query: contains details about which rows to return @@ -539,6 +569,8 @@ async def read_rows_stream( a DeadlineExceeded exception, and a retry will be attempted. Defaults to the Table's default_read_rows_attempt_timeout. If None, defaults to operation_timeout. + - retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_read_rows_retryable_errors Returns: - an asynchronous iterator that yields rows returned by the query Raises: @@ -551,12 +583,14 @@ async def read_rows_stream( operation_timeout, attempt_timeout = _get_timeouts( operation_timeout, attempt_timeout, self ) + retryable_excs = _get_retryable_errors(retryable_errors, self) row_merger = _ReadRowsOperationAsync( query, self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, + retryable_exceptions=retryable_excs, ) return row_merger.start_operation() @@ -566,13 +600,16 @@ async def read_rows( *, operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> list[Row]: """ Read a set of rows from the table, based on the specified query. Retruns results as a list of Row objects when the request is complete. For streamed results, use read_rows_stream. - Failed requests within operation_timeout will be retried. + Failed requests within operation_timeout will be retried based on the + retryable_errors list until operation_timeout is reached. Args: - query: contains details about which rows to return @@ -584,6 +621,10 @@ async def read_rows( a DeadlineExceeded exception, and a retry will be attempted. Defaults to the Table's default_read_rows_attempt_timeout. If None, defaults to operation_timeout. + If None, defaults to the Table's default_read_rows_attempt_timeout, + or the operation_timeout if that is also None. + - retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_read_rows_retryable_errors. Returns: - a list of Rows returned by the query Raises: @@ -596,6 +637,7 @@ async def read_rows( query, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, + retryable_errors=retryable_errors, ) return [row async for row in row_generator] @@ -606,11 +648,14 @@ async def read_row( row_filter: RowFilter | None = None, operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> Row | None: """ Read a single row from the table, based on the specified key. - Failed requests within operation_timeout will be retried. + Failed requests within operation_timeout will be retried based on the + retryable_errors list until operation_timeout is reached. Args: - query: contains details about which rows to return @@ -622,6 +667,8 @@ async def read_row( a DeadlineExceeded exception, and a retry will be attempted. Defaults to the Table's default_read_rows_attempt_timeout. If None, defaults to operation_timeout. + - retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_read_rows_retryable_errors. Returns: - a Row object if the row exists, otherwise None Raises: @@ -637,6 +684,7 @@ async def read_row( query, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, + retryable_errors=retryable_errors, ) if len(results) == 0: return None @@ -648,6 +696,8 @@ async def read_rows_sharded( *, operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> list[Row]: """ Runs a sharded query in parallel, then return the results in a single list. @@ -672,6 +722,8 @@ async def read_rows_sharded( a DeadlineExceeded exception, and a retry will be attempted. Defaults to the Table's default_read_rows_attempt_timeout. If None, defaults to operation_timeout. + - retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_read_rows_retryable_errors. Raises: - ShardedReadRowsExceptionGroup: if any of the queries failed - ValueError: if the query_list is empty @@ -701,6 +753,7 @@ async def read_rows_sharded( query, operation_timeout=batch_operation_timeout, attempt_timeout=min(attempt_timeout, batch_operation_timeout), + retryable_errors=retryable_errors, ) for query in batch ] @@ -729,10 +782,13 @@ async def row_exists( *, operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, + retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.READ_ROWS, ) -> bool: """ Return a boolean indicating whether the specified row exists in the table. uses the filters: chain(limit cells per row = 1, strip value) + Args: - row_key: the key of the row to check - operation_timeout: the time budget for the entire operation, in seconds. @@ -743,6 +799,8 @@ async def row_exists( a DeadlineExceeded exception, and a retry will be attempted. Defaults to the Table's default_read_rows_attempt_timeout. If None, defaults to operation_timeout. + - retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_read_rows_retryable_errors. Returns: - a bool indicating whether the row exists Raises: @@ -762,6 +820,7 @@ async def row_exists( query, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, + retryable_errors=retryable_errors, ) return len(results) > 0 @@ -770,6 +829,8 @@ async def sample_row_keys( *, operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, + retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, ) -> RowKeySamples: """ Return a set of RowKeySamples that delimit contiguous sections of the table of @@ -791,6 +852,8 @@ async def sample_row_keys( a DeadlineExceeded exception, and a retry will be attempted. Defaults to the Table's default_attempt_timeout. If None, defaults to operation_timeout. + - retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_retryable_errors. Returns: - a set of RowKeySamples the delimit contiguous sections of the table Raises: @@ -807,10 +870,8 @@ async def sample_row_keys( attempt_timeout, operation_timeout ) # prepare retryable - predicate = retries.if_exception_type( - core_exceptions.DeadlineExceeded, - core_exceptions.ServiceUnavailable, - ) + retryable_excs = _get_retryable_errors(retryable_errors, self) + predicate = retries.if_exception_type(*retryable_excs) transient_errors = [] def on_error_fn(exc): @@ -856,6 +917,8 @@ def mutations_batcher( flow_control_max_bytes: int = 100 * _MB_SIZE, batch_operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, batch_attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, + batch_retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, ) -> MutationsBatcherAsync: """ Returns a new mutations batcher instance. @@ -876,6 +939,8 @@ def mutations_batcher( - batch_attempt_timeout: timeout for each individual request, in seconds. Defaults to the Table's default_mutate_rows_attempt_timeout. If None, defaults to batch_operation_timeout. + - batch_retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_mutate_rows_retryable_errors. Returns: - a MutationsBatcherAsync context manager that can batch requests """ @@ -888,6 +953,7 @@ def mutations_batcher( flow_control_max_bytes=flow_control_max_bytes, batch_operation_timeout=batch_operation_timeout, batch_attempt_timeout=batch_attempt_timeout, + batch_retryable_errors=batch_retryable_errors, ) async def mutate_row( @@ -897,6 +963,8 @@ async def mutate_row( *, operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, + retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.DEFAULT, ): """ Mutates a row atomically. @@ -918,6 +986,9 @@ async def mutate_row( a DeadlineExceeded exception, and a retry will be attempted. Defaults to the Table's default_attempt_timeout. If None, defaults to operation_timeout. + - retryable_errors: a list of errors that will be retried if encountered. + Only idempotent mutations will be retried. Defaults to the Table's + default_retryable_errors. Raises: - DeadlineExceeded: raised after operation timeout will be chained with a RetryExceptionGroup containing all @@ -937,8 +1008,7 @@ async def mutate_row( if all(mutation.is_idempotent() for mutation in mutations_list): # mutations are all idempotent and safe to retry predicate = retries.if_exception_type( - core_exceptions.DeadlineExceeded, - core_exceptions.ServiceUnavailable, + *_get_retryable_errors(retryable_errors, self) ) else: # mutations should not be retried @@ -982,6 +1052,8 @@ async def bulk_mutate_rows( *, operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, + retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, ): """ Applies mutations for multiple rows in a single batched request. @@ -1007,6 +1079,8 @@ async def bulk_mutate_rows( a DeadlineExceeded exception, and a retry will be attempted. Defaults to the Table's default_mutate_rows_attempt_timeout. If None, defaults to operation_timeout. + - retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_mutate_rows_retryable_errors Raises: - MutationsExceptionGroup if one or more mutations fails Contains details about any failed entries in .exceptions @@ -1015,6 +1089,7 @@ async def bulk_mutate_rows( operation_timeout, attempt_timeout = _get_timeouts( operation_timeout, attempt_timeout, self ) + retryable_excs = _get_retryable_errors(retryable_errors, self) operation = _MutateRowsOperationAsync( self.client._gapic_client, @@ -1022,6 +1097,7 @@ async def bulk_mutate_rows( mutation_entries, operation_timeout, attempt_timeout, + retryable_exceptions=retryable_excs, ) await operation.start() diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index 91d2b11e1..b2da30040 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -14,7 +14,7 @@ # from __future__ import annotations -from typing import Any, TYPE_CHECKING +from typing import Any, Sequence, TYPE_CHECKING import asyncio import atexit import warnings @@ -23,6 +23,7 @@ from google.cloud.bigtable.data.mutations import RowMutationEntry from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup from google.cloud.bigtable.data.exceptions import FailedMutationEntryError +from google.cloud.bigtable.data._helpers import _get_retryable_errors from google.cloud.bigtable.data._helpers import _get_timeouts from google.cloud.bigtable.data._helpers import TABLE_DEFAULT @@ -192,6 +193,8 @@ def __init__( flow_control_max_bytes: int = 100 * _MB_SIZE, batch_operation_timeout: float | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, batch_attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, + batch_retryable_errors: Sequence[type[Exception]] + | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, ): """ Args: @@ -208,10 +211,16 @@ def __init__( - batch_attempt_timeout: timeout for each individual request, in seconds. If TABLE_DEFAULT, defaults to the Table's default_mutate_rows_attempt_timeout. If None, defaults to batch_operation_timeout. + - batch_retryable_errors: a list of errors that will be retried if encountered. + Defaults to the Table's default_mutate_rows_retryable_errors. """ self._operation_timeout, self._attempt_timeout = _get_timeouts( batch_operation_timeout, batch_attempt_timeout, table ) + self._retryable_errors: list[type[Exception]] = _get_retryable_errors( + batch_retryable_errors, table + ) + self.closed: bool = False self._table = table self._staged_entries: list[RowMutationEntry] = [] @@ -349,6 +358,7 @@ async def _execute_mutate_rows( batch, operation_timeout=self._operation_timeout, attempt_timeout=self._attempt_timeout, + retryable_exceptions=self._retryable_errors, ) await operation.start() except MutationsExceptionGroup as e: diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 1d56926ff..96ea1d1ce 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -11,9 +11,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # +""" +Helper functions used in various places in the library. +""" from __future__ import annotations -from typing import Callable, List, Tuple, Any +from typing import Callable, Sequence, List, Tuple, Any, TYPE_CHECKING import time import enum from collections import namedtuple @@ -22,6 +25,10 @@ from google.api_core import exceptions as core_exceptions from google.cloud.bigtable.data.exceptions import RetryExceptionGroup +if TYPE_CHECKING: + import grpc + from google.cloud.bigtable.data import TableAsync + """ Helper functions used in various places in the library. """ @@ -142,7 +149,9 @@ def wrapper(*args, **kwargs): def _get_timeouts( - operation: float | TABLE_DEFAULT, attempt: float | None | TABLE_DEFAULT, table + operation: float | TABLE_DEFAULT, + attempt: float | None | TABLE_DEFAULT, + table: "TableAsync", ) -> tuple[float, float]: """ Convert passed in timeout values to floats, using table defaults if necessary. @@ -209,3 +218,21 @@ def _validate_timeouts( elif attempt_timeout is not None: if attempt_timeout <= 0: raise ValueError("attempt_timeout must be greater than 0") + + +def _get_retryable_errors( + call_codes: Sequence["grpc.StatusCode" | int | type[Exception]] | TABLE_DEFAULT, + table: "TableAsync", +) -> list[type[Exception]]: + # load table defaults if necessary + if call_codes == TABLE_DEFAULT.DEFAULT: + call_codes = table.default_retryable_errors + elif call_codes == TABLE_DEFAULT.READ_ROWS: + call_codes = table.default_read_rows_retryable_errors + elif call_codes == TABLE_DEFAULT.MUTATE_ROWS: + call_codes = table.default_mutate_rows_retryable_errors + + return [ + e if isinstance(e, type) else type(core_exceptions.from_grpc_status(e, "")) + for e in call_codes + ] diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index 89a153af2..d41929518 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -46,9 +46,10 @@ def _make_one(self, *args, **kwargs): if not args: kwargs["gapic_client"] = kwargs.pop("gapic_client", mock.Mock()) kwargs["table"] = kwargs.pop("table", AsyncMock()) - kwargs["mutation_entries"] = kwargs.pop("mutation_entries", []) kwargs["operation_timeout"] = kwargs.pop("operation_timeout", 5) kwargs["attempt_timeout"] = kwargs.pop("attempt_timeout", 0.1) + kwargs["retryable_exceptions"] = kwargs.pop("retryable_exceptions", ()) + kwargs["mutation_entries"] = kwargs.pop("mutation_entries", []) return self._target_class()(*args, **kwargs) async def _mock_stream(self, mutation_list, error_dict): @@ -78,15 +79,21 @@ def test_ctor(self): from google.cloud.bigtable.data._async._mutate_rows import _EntryWithProto from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete from google.api_core.exceptions import DeadlineExceeded - from google.api_core.exceptions import ServiceUnavailable + from google.api_core.exceptions import Aborted client = mock.Mock() table = mock.Mock() entries = [_make_mutation(), _make_mutation()] operation_timeout = 0.05 attempt_timeout = 0.01 + retryable_exceptions = () instance = self._make_one( - client, table, entries, operation_timeout, attempt_timeout + client, + table, + entries, + operation_timeout, + attempt_timeout, + retryable_exceptions, ) # running gapic_fn should trigger a client call assert client.mutate_rows.call_count == 0 @@ -110,8 +117,8 @@ def test_ctor(self): assert next(instance.timeout_generator) == attempt_timeout # ensure predicate is set assert instance.is_retryable is not None - assert instance.is_retryable(DeadlineExceeded("")) is True - assert instance.is_retryable(ServiceUnavailable("")) is True + assert instance.is_retryable(DeadlineExceeded("")) is False + assert instance.is_retryable(Aborted("")) is False assert instance.is_retryable(_MutateRowsIncomplete("")) is True assert instance.is_retryable(RuntimeError("")) is False assert instance.remaining_indices == list(range(len(entries))) @@ -232,7 +239,7 @@ async def test_mutate_rows_exception(self, exc_type): @pytest.mark.parametrize( "exc_type", - [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable], + [core_exceptions.DeadlineExceeded, RuntimeError], ) @pytest.mark.asyncio async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): @@ -256,7 +263,12 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): ) as attempt_mock: attempt_mock.side_effect = [expected_cause] * num_retries + [None] instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, + table, + entries, + operation_timeout, + operation_timeout, + retryable_exceptions=(exc_type,), ) await instance.start() assert attempt_mock.call_count == num_retries + 1 diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 7718246fc..54bbb6158 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -26,6 +26,8 @@ from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.api_core import exceptions as core_exceptions from google.cloud.bigtable.data.exceptions import InvalidChunk +from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete +from google.cloud.bigtable.data import TABLE_DEFAULT from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule @@ -841,6 +843,39 @@ async def test_get_table(self): assert client._instance_owners[instance_key] == {id(table)} await client.close() + @pytest.mark.asyncio + async def test_get_table_arg_passthrough(self): + """ + All arguments passed in get_table should be sent to constructor + """ + async with self._make_one(project="project-id") as client: + with mock.patch( + "google.cloud.bigtable.data._async.client.TableAsync.__init__", + ) as mock_constructor: + mock_constructor.return_value = None + assert not client._active_instances + expected_table_id = "table-id" + expected_instance_id = "instance-id" + expected_app_profile_id = "app-profile-id" + expected_args = (1, "test", {"test": 2}) + expected_kwargs = {"hello": "world", "test": 2} + + client.get_table( + expected_instance_id, + expected_table_id, + expected_app_profile_id, + *expected_args, + **expected_kwargs, + ) + mock_constructor.assert_called_once_with( + client, + expected_instance_id, + expected_table_id, + expected_app_profile_id, + *expected_args, + **expected_kwargs, + ) + @pytest.mark.asyncio async def test_get_table_context_manager(self): from google.cloud.bigtable.data._async.client import TableAsync @@ -1099,6 +1134,173 @@ def test_table_ctor_sync(self): TableAsync(client, "instance-id", "table-id") assert e.match("TableAsync must be created within an async event loop context.") + @pytest.mark.asyncio + # iterate over all retryable rpcs + @pytest.mark.parametrize( + "fn_name,fn_args,retry_fn_path,extra_retryables", + [ + ( + "read_rows_stream", + (ReadRowsQuery(),), + "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + (), + ), + ( + "read_rows", + (ReadRowsQuery(),), + "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + (), + ), + ( + "read_row", + (b"row_key",), + "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + (), + ), + ( + "read_rows_sharded", + ([ReadRowsQuery()],), + "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + (), + ), + ( + "row_exists", + (b"row_key",), + "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + (), + ), + ("sample_row_keys", (), "google.api_core.retry_async.retry_target", ()), + ( + "mutate_row", + (b"row_key", []), + "google.api_core.retry_async.retry_target", + (), + ), + ( + "bulk_mutate_rows", + ([mutations.RowMutationEntry(b"key", [mock.Mock()])],), + "google.api_core.retry_async.retry_target", + (_MutateRowsIncomplete,), + ), + ], + ) + # test different inputs for retryable exceptions + @pytest.mark.parametrize( + "input_retryables,expected_retryables", + [ + ( + TABLE_DEFAULT.READ_ROWS, + [ + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + core_exceptions.Aborted, + ], + ), + ( + TABLE_DEFAULT.DEFAULT, + [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable], + ), + ( + TABLE_DEFAULT.MUTATE_ROWS, + [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable], + ), + ([], []), + ([4], [core_exceptions.DeadlineExceeded]), + ], + ) + async def test_customizable_retryable_errors( + self, + input_retryables, + expected_retryables, + fn_name, + fn_args, + retry_fn_path, + extra_retryables, + ): + """ + Test that retryable functions support user-configurable arguments, and that the configured retryables are passed + down to the gapic layer. + """ + from google.cloud.bigtable.data import BigtableDataClientAsync + + with mock.patch( + "google.api_core.retry_async.if_exception_type" + ) as predicate_builder_mock: + with mock.patch(retry_fn_path) as retry_fn_mock: + async with BigtableDataClientAsync() as client: + table = client.get_table("instance-id", "table-id") + expected_predicate = lambda a: a in expected_retryables # noqa + predicate_builder_mock.return_value = expected_predicate + retry_fn_mock.side_effect = RuntimeError("stop early") + with pytest.raises(Exception): + # we expect an exception from attempting to call the mock + test_fn = table.__getattribute__(fn_name) + await test_fn(*fn_args, retryable_errors=input_retryables) + # passed in errors should be used to build the predicate + predicate_builder_mock.assert_called_once_with( + *expected_retryables, *extra_retryables + ) + retry_call_args = retry_fn_mock.call_args_list[0].args + # output of if_exception_type should be sent in to retry constructor + assert retry_call_args[1] is expected_predicate + + @pytest.mark.parametrize( + "fn_name,fn_args,gapic_fn", + [ + ("read_rows_stream", (ReadRowsQuery(),), "read_rows"), + ("read_rows", (ReadRowsQuery(),), "read_rows"), + ("read_row", (b"row_key",), "read_rows"), + ("read_rows_sharded", ([ReadRowsQuery()],), "read_rows"), + ("row_exists", (b"row_key",), "read_rows"), + ("sample_row_keys", (), "sample_row_keys"), + ("mutate_row", (b"row_key", []), "mutate_row"), + ( + "bulk_mutate_rows", + ([mutations.RowMutationEntry(b"key", [mock.Mock()])],), + "mutate_rows", + ), + ("check_and_mutate_row", (b"row_key", None), "check_and_mutate_row"), + ( + "read_modify_write_row", + (b"row_key", mock.Mock()), + "read_modify_write_row", + ), + ], + ) + @pytest.mark.parametrize("include_app_profile", [True, False]) + @pytest.mark.asyncio + async def test_call_metadata(self, include_app_profile, fn_name, fn_args, gapic_fn): + """check that all requests attach proper metadata headers""" + from google.cloud.bigtable.data import TableAsync + from google.cloud.bigtable.data import BigtableDataClientAsync + + profile = "profile" if include_app_profile else None + with mock.patch( + f"google.cloud.bigtable_v2.BigtableAsyncClient.{gapic_fn}", mock.AsyncMock() + ) as gapic_mock: + gapic_mock.side_effect = RuntimeError("stop early") + async with BigtableDataClientAsync() as client: + table = TableAsync(client, "instance-id", "table-id", profile) + try: + test_fn = table.__getattribute__(fn_name) + maybe_stream = await test_fn(*fn_args) + [i async for i in maybe_stream] + except Exception: + # we expect an exception from attempting to call the mock + pass + kwargs = gapic_mock.call_args_list[0].kwargs + metadata = kwargs["metadata"] + goog_metadata = None + for key, value in metadata: + if key == "x-goog-request-params": + goog_metadata = value + assert goog_metadata is not None, "x-goog-request-params not found" + assert "table_name=" + table.table_name in goog_metadata + if include_app_profile: + assert "app_profile_id=profile" in goog_metadata + else: + assert "app_profile_id=" not in goog_metadata + class TestReadRows: """ @@ -1608,28 +1810,6 @@ async def test_row_exists(self, return_value, expected_result): assert query.limit == 1 assert query.filter._to_dict() == expected_filter - @pytest.mark.parametrize("include_app_profile", [True, False]) - @pytest.mark.asyncio - async def test_read_rows_metadata(self, include_app_profile): - """request should attach metadata headers""" - profile = "profile" if include_app_profile else None - async with self._make_table(app_profile_id=profile) as table: - read_rows = table.client._gapic_client.read_rows - read_rows.return_value = self._make_gapic_stream([]) - await table.read_rows(ReadRowsQuery()) - kwargs = read_rows.call_args_list[0].kwargs - metadata = kwargs["metadata"] - goog_metadata = None - for key, value in metadata: - if key == "x-goog-request-params": - goog_metadata = value - assert goog_metadata is not None, "x-goog-request-params not found" - assert "table_name=" + table.table_name in goog_metadata - if include_app_profile: - assert "app_profile_id=profile" in goog_metadata - else: - assert "app_profile_id=" not in goog_metadata - class TestReadRowsSharded: def _make_client(self, *args, **kwargs): @@ -1735,30 +1915,6 @@ async def mock_call(*args, **kwargs): # if run in sequence, we would expect this to take 1 second assert call_time < 0.2 - @pytest.mark.parametrize("include_app_profile", [True, False]) - @pytest.mark.asyncio - async def test_read_rows_sharded_metadata(self, include_app_profile): - """request should attach metadata headers""" - profile = "profile" if include_app_profile else None - async with self._make_client() as client: - async with client.get_table("i", "t", app_profile_id=profile) as table: - with mock.patch.object( - client._gapic_client, "read_rows", AsyncMock() - ) as read_rows: - await table.read_rows_sharded([ReadRowsQuery()]) - kwargs = read_rows.call_args_list[0].kwargs - metadata = kwargs["metadata"] - goog_metadata = None - for key, value in metadata: - if key == "x-goog-request-params": - goog_metadata = value - assert goog_metadata is not None, "x-goog-request-params not found" - assert "table_name=" + table.table_name in goog_metadata - if include_app_profile: - assert "app_profile_id=profile" in goog_metadata - else: - assert "app_profile_id=" not in goog_metadata - @pytest.mark.asyncio async def test_read_rows_sharded_batching(self): """ @@ -1875,7 +2031,10 @@ async def test_sample_row_keys_default_timeout(self): expected_timeout = 99 async with self._make_client() as client: async with client.get_table( - "i", "t", default_operation_timeout=expected_timeout + "i", + "t", + default_operation_timeout=expected_timeout, + default_attempt_timeout=expected_timeout, ) as table: with mock.patch.object( table.client._gapic_client, "sample_row_keys", AsyncMock() @@ -1914,30 +2073,6 @@ async def test_sample_row_keys_gapic_params(self): assert kwargs["metadata"] is not None assert kwargs["retry"] is None - @pytest.mark.parametrize("include_app_profile", [True, False]) - @pytest.mark.asyncio - async def test_sample_row_keys_metadata(self, include_app_profile): - """request should attach metadata headers""" - profile = "profile" if include_app_profile else None - async with self._make_client() as client: - async with client.get_table("i", "t", app_profile_id=profile) as table: - with mock.patch.object( - client._gapic_client, "sample_row_keys", AsyncMock() - ) as read_rows: - await table.sample_row_keys() - kwargs = read_rows.call_args_list[0].kwargs - metadata = kwargs["metadata"] - goog_metadata = None - for key, value in metadata: - if key == "x-goog-request-params": - goog_metadata = value - assert goog_metadata is not None, "x-goog-request-params not found" - assert "table_name=" + table.table_name in goog_metadata - if include_app_profile: - assert "app_profile_id=profile" in goog_metadata - else: - assert "app_profile_id=" not in goog_metadata - @pytest.mark.parametrize( "retryable_exception", [ @@ -2525,39 +2660,6 @@ async def test_bulk_mutate_error_index(self): assert isinstance(cause.exceptions[1], DeadlineExceeded) assert isinstance(cause.exceptions[2], FailedPrecondition) - @pytest.mark.parametrize("include_app_profile", [True, False]) - @pytest.mark.asyncio - async def test_bulk_mutate_row_metadata(self, include_app_profile): - """request should attach metadata headers""" - profile = "profile" if include_app_profile else None - async with self._make_client() as client: - async with client.get_table("i", "t", app_profile_id=profile) as table: - with mock.patch.object( - client._gapic_client, "mutate_rows", AsyncMock() - ) as mutate_rows: - mutate_rows.side_effect = core_exceptions.Aborted("mock") - mutation = mock.Mock() - mutation.size.return_value = 1 - entry = mock.Mock() - entry.mutations = [mutation] - try: - await table.bulk_mutate_rows([entry]) - except Exception: - # exception used to end early - pass - kwargs = mutate_rows.call_args_list[0].kwargs - metadata = kwargs["metadata"] - goog_metadata = None - for key, value in metadata: - if key == "x-goog-request-params": - goog_metadata = value - assert goog_metadata is not None, "x-goog-request-params not found" - assert "table_name=" + table.table_name in goog_metadata - if include_app_profile: - assert "app_profile_id=profile" in goog_metadata - else: - assert "app_profile_id=" not in goog_metadata - class TestCheckAndMutateRow: def _make_client(self, *args, **kwargs): @@ -2727,30 +2829,6 @@ async def test_check_and_mutate_mutations_parsing(self): mutation._to_pb.call_count == 1 for mutation in mutations[:5] ) - @pytest.mark.parametrize("include_app_profile", [True, False]) - @pytest.mark.asyncio - async def test_check_and_mutate_metadata(self, include_app_profile): - """request should attach metadata headers""" - profile = "profile" if include_app_profile else None - async with self._make_client() as client: - async with client.get_table("i", "t", app_profile_id=profile) as table: - with mock.patch.object( - client._gapic_client, "check_and_mutate_row", AsyncMock() - ) as mock_gapic: - await table.check_and_mutate_row(b"key", mock.Mock()) - kwargs = mock_gapic.call_args_list[0].kwargs - metadata = kwargs["metadata"] - goog_metadata = None - for key, value in metadata: - if key == "x-goog-request-params": - goog_metadata = value - assert goog_metadata is not None, "x-goog-request-params not found" - assert "table_name=" + table.table_name in goog_metadata - if include_app_profile: - assert "app_profile_id=profile" in goog_metadata - else: - assert "app_profile_id=" not in goog_metadata - class TestReadModifyWriteRow: def _make_client(self, *args, **kwargs): @@ -2882,27 +2960,3 @@ async def test_read_modify_write_row_building(self): await table.read_modify_write_row("key", mock.Mock()) assert constructor_mock.call_count == 1 constructor_mock.assert_called_once_with(mock_response.row) - - @pytest.mark.parametrize("include_app_profile", [True, False]) - @pytest.mark.asyncio - async def test_read_modify_write_metadata(self, include_app_profile): - """request should attach metadata headers""" - profile = "profile" if include_app_profile else None - async with self._make_client() as client: - async with client.get_table("i", "t", app_profile_id=profile) as table: - with mock.patch.object( - client._gapic_client, "read_modify_write_row", AsyncMock() - ) as mock_gapic: - await table.read_modify_write_row("key", mock.Mock()) - kwargs = mock_gapic.call_args_list[0].kwargs - metadata = kwargs["metadata"] - goog_metadata = None - for key, value in metadata: - if key == "x-goog-request-params": - goog_metadata = value - assert goog_metadata is not None, "x-goog-request-params not found" - assert "table_name=" + table.table_name in goog_metadata - if include_app_profile: - assert "app_profile_id=profile" in goog_metadata - else: - assert "app_profile_id=" not in goog_metadata diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py index f95b53271..17bd8d420 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -14,6 +14,9 @@ import pytest import asyncio +import google.api_core.exceptions as core_exceptions +from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete +from google.cloud.bigtable.data import TABLE_DEFAULT # try/except added for compatibility with python < 3.8 try: @@ -286,10 +289,17 @@ def _get_target_class(self): return MutationsBatcherAsync def _make_one(self, table=None, **kwargs): + from google.api_core.exceptions import DeadlineExceeded + from google.api_core.exceptions import ServiceUnavailable + if table is None: table = mock.Mock() table.default_mutate_rows_operation_timeout = 10 table.default_mutate_rows_attempt_timeout = 10 + table.default_mutate_rows_retryable_errors = ( + DeadlineExceeded, + ServiceUnavailable, + ) return self._get_target_class()(table, **kwargs) @@ -302,6 +312,7 @@ async def test_ctor_defaults(self, flush_timer_mock): table = mock.Mock() table.default_mutate_rows_operation_timeout = 10 table.default_mutate_rows_attempt_timeout = 8 + table.default_mutate_rows_retryable_errors = [Exception] async with self._make_one(table) as instance: assert instance._table == table assert instance.closed is False @@ -323,6 +334,9 @@ async def test_ctor_defaults(self, flush_timer_mock): assert ( instance._attempt_timeout == table.default_mutate_rows_attempt_timeout ) + assert ( + instance._retryable_errors == table.default_mutate_rows_retryable_errors + ) await asyncio.sleep(0) assert flush_timer_mock.call_count == 1 assert flush_timer_mock.call_args[0][0] == 5 @@ -343,6 +357,7 @@ async def test_ctor_explicit(self, flush_timer_mock): flow_control_max_bytes = 12 operation_timeout = 11 attempt_timeout = 2 + retryable_errors = [Exception] async with self._make_one( table, flush_interval=flush_interval, @@ -352,6 +367,7 @@ async def test_ctor_explicit(self, flush_timer_mock): flow_control_max_bytes=flow_control_max_bytes, batch_operation_timeout=operation_timeout, batch_attempt_timeout=attempt_timeout, + batch_retryable_errors=retryable_errors, ) as instance: assert instance._table == table assert instance.closed is False @@ -371,6 +387,7 @@ async def test_ctor_explicit(self, flush_timer_mock): assert instance._entries_processed_since_last_raise == 0 assert instance._operation_timeout == operation_timeout assert instance._attempt_timeout == attempt_timeout + assert instance._retryable_errors == retryable_errors await asyncio.sleep(0) assert flush_timer_mock.call_count == 1 assert flush_timer_mock.call_args[0][0] == flush_interval @@ -386,6 +403,7 @@ async def test_ctor_no_flush_limits(self, flush_timer_mock): table = mock.Mock() table.default_mutate_rows_operation_timeout = 10 table.default_mutate_rows_attempt_timeout = 8 + table.default_mutate_rows_retryable_errors = () flush_interval = None flush_limit_count = None flush_limit_bytes = None @@ -442,7 +460,7 @@ def test_default_argument_consistency(self): batcher_init_signature.pop("table") # both should have same number of arguments assert len(get_batcher_signature.keys()) == len(batcher_init_signature.keys()) - assert len(get_batcher_signature) == 7 # update if expected params change + assert len(get_batcher_signature) == 8 # update if expected params change # both should have same argument names assert set(get_batcher_signature.keys()) == set(batcher_init_signature.keys()) # both should have same default values @@ -882,6 +900,7 @@ async def test__execute_mutate_rows(self, mutate_rows): table.app_profile_id = "test-app-profile" table.default_mutate_rows_operation_timeout = 17 table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () async with self._make_one(table) as instance: batch = [_make_mutation()] result = await instance._execute_mutate_rows(batch) @@ -911,6 +930,7 @@ async def test__execute_mutate_rows_returns_errors(self, mutate_rows): table = mock.Mock() table.default_mutate_rows_operation_timeout = 17 table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () async with self._make_one(table) as instance: batch = [_make_mutation()] result = await instance._execute_mutate_rows(batch) @@ -1102,3 +1122,63 @@ def test__add_exceptions(self, limit, in_e, start_e, end_e): # then, the newest slots should be filled with the last items of the input list for i in range(1, newest_list_diff + 1): assert mock_batcher._newest_exceptions[-i] == input_list[-i] + + @pytest.mark.asyncio + # test different inputs for retryable exceptions + @pytest.mark.parametrize( + "input_retryables,expected_retryables", + [ + ( + TABLE_DEFAULT.READ_ROWS, + [ + core_exceptions.DeadlineExceeded, + core_exceptions.ServiceUnavailable, + core_exceptions.Aborted, + ], + ), + ( + TABLE_DEFAULT.DEFAULT, + [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable], + ), + ( + TABLE_DEFAULT.MUTATE_ROWS, + [core_exceptions.DeadlineExceeded, core_exceptions.ServiceUnavailable], + ), + ([], []), + ([4], [core_exceptions.DeadlineExceeded]), + ], + ) + async def test_customizable_retryable_errors( + self, input_retryables, expected_retryables + ): + """ + Test that retryable functions support user-configurable arguments, and that the configured retryables are passed + down to the gapic layer. + """ + from google.cloud.bigtable.data._async.client import TableAsync + + with mock.patch( + "google.api_core.retry_async.if_exception_type" + ) as predicate_builder_mock: + with mock.patch( + "google.api_core.retry_async.retry_target" + ) as retry_fn_mock: + table = None + with mock.patch("asyncio.create_task"): + table = TableAsync(mock.Mock(), "instance", "table") + async with self._make_one( + table, batch_retryable_errors=input_retryables + ) as instance: + assert instance._retryable_errors == expected_retryables + expected_predicate = lambda a: a in expected_retryables # noqa + predicate_builder_mock.return_value = expected_predicate + retry_fn_mock.side_effect = RuntimeError("stop early") + mutation = _make_mutation(count=1, size=1) + await instance._execute_mutate_rows([mutation]) + # passed in errors should be used to build the predicate + predicate_builder_mock.assert_called_once_with( + *expected_retryables, _MutateRowsIncomplete + ) + retry_call_args = retry_fn_mock.call_args_list[0].args + # output of if_exception_type should be sent in to retry constructor + assert retry_call_args[1] is expected_predicate diff --git a/tests/unit/data/test__helpers.py b/tests/unit/data/test__helpers.py index 6c11fa86a..b9c1dc2bb 100644 --- a/tests/unit/data/test__helpers.py +++ b/tests/unit/data/test__helpers.py @@ -13,6 +13,8 @@ # import pytest +import grpc +from google.api_core import exceptions as core_exceptions import google.cloud.bigtable.data._helpers as _helpers from google.cloud.bigtable.data._helpers import TABLE_DEFAULT import google.cloud.bigtable.data.exceptions as bigtable_exceptions @@ -264,3 +266,49 @@ def test_get_timeouts_invalid(self, input_times, input_table): setattr(fake_table, f"default_{key}_timeout", input_table[key]) with pytest.raises(ValueError): _helpers._get_timeouts(input_times[0], input_times[1], fake_table) + + +class TestGetRetryableErrors: + @pytest.mark.parametrize( + "input_codes,input_table,expected", + [ + ((), {}, []), + ((Exception,), {}, [Exception]), + (TABLE_DEFAULT.DEFAULT, {"default": [Exception]}, [Exception]), + ( + TABLE_DEFAULT.READ_ROWS, + {"default_read_rows": (RuntimeError, ValueError)}, + [RuntimeError, ValueError], + ), + ( + TABLE_DEFAULT.MUTATE_ROWS, + {"default_mutate_rows": (ValueError,)}, + [ValueError], + ), + ((4,), {}, [core_exceptions.DeadlineExceeded]), + ( + [grpc.StatusCode.DEADLINE_EXCEEDED], + {}, + [core_exceptions.DeadlineExceeded], + ), + ( + (14, grpc.StatusCode.ABORTED, RuntimeError), + {}, + [ + core_exceptions.ServiceUnavailable, + core_exceptions.Aborted, + RuntimeError, + ], + ), + ], + ) + def test_get_retryable_errors(self, input_codes, input_table, expected): + """ + test input/output mappings for a variety of valid inputs + """ + fake_table = mock.Mock() + for key in input_table.keys(): + # set the default fields in our fake table mock + setattr(fake_table, f"{key}_retryable_errors", input_table[key]) + result = _helpers._get_retryable_errors(input_codes, fake_table) + assert result == expected From 9342e2703cdabc174943dd67102b00c1119506b1 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 15 Dec 2023 12:30:54 -0800 Subject: [PATCH 30/56] chore: update api_core submodule (#897) --- google/__init__.py | 6 -- google/cloud/__init__.py | 6 -- google/cloud/bigtable/data/__init__.py | 2 - .../bigtable/data/_async/_mutate_rows.py | 33 +++------ .../cloud/bigtable/data/_async/_read_rows.py | 42 +---------- google/cloud/bigtable/data/_async/client.py | 71 +++++++----------- .../bigtable/data/_async/mutations_batcher.py | 3 + google/cloud/bigtable/data/_helpers.py | 74 +++++++------------ google/cloud/bigtable/data/exceptions.py | 9 --- mypy.ini | 2 +- noxfile.py | 13 ++-- owlbot.py | 3 +- python-api-core | 2 +- setup.py | 12 +-- testing/constraints-3.7.txt | 2 +- testing/constraints-3.8.txt | 3 +- tests/system/data/test_system.py | 54 +++++++------- tests/unit/data/_async/test__mutate_rows.py | 12 +-- tests/unit/data/_async/test_client.py | 40 +++++----- .../data/_async/test_mutations_batcher.py | 4 +- tests/unit/data/test__helpers.py | 66 ----------------- 21 files changed, 142 insertions(+), 317 deletions(-) delete mode 100644 google/__init__.py delete mode 100644 google/cloud/__init__.py diff --git a/google/__init__.py b/google/__init__.py deleted file mode 100644 index a5ba80656..000000000 --- a/google/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -try: - import pkg_resources - - pkg_resources.declare_namespace(__name__) -except ImportError: - pass diff --git a/google/cloud/__init__.py b/google/cloud/__init__.py deleted file mode 100644 index a5ba80656..000000000 --- a/google/cloud/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -try: - import pkg_resources - - pkg_resources.declare_namespace(__name__) -except ImportError: - pass diff --git a/google/cloud/bigtable/data/__init__.py b/google/cloud/bigtable/data/__init__.py index a68be5417..5229f8021 100644 --- a/google/cloud/bigtable/data/__init__.py +++ b/google/cloud/bigtable/data/__init__.py @@ -32,7 +32,6 @@ from google.cloud.bigtable.data.mutations import DeleteAllFromFamily from google.cloud.bigtable.data.mutations import DeleteAllFromRow -from google.cloud.bigtable.data.exceptions import IdleTimeout from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable.data.exceptions import FailedMutationEntryError from google.cloud.bigtable.data.exceptions import FailedQueryShardError @@ -63,7 +62,6 @@ "DeleteAllFromRow", "Row", "Cell", - "IdleTimeout", "InvalidChunk", "FailedMutationEntryError", "FailedQueryShardError", diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 5971a9894..d4ffdee22 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -15,17 +15,16 @@ from __future__ import annotations from typing import Sequence, TYPE_CHECKING -import asyncio from dataclasses import dataclass import functools from google.api_core import exceptions as core_exceptions -from google.api_core import retry_async as retries +from google.api_core import retry as retries import google.cloud.bigtable_v2.types.bigtable as types_pb import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _make_metadata -from google.cloud.bigtable.data._helpers import _convert_retry_deadline from google.cloud.bigtable.data._helpers import _attempt_timeout_generator +from google.cloud.bigtable.data._helpers import _retry_exception_factory # mutate_rows requests are limited to this number of mutations from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT @@ -101,17 +100,13 @@ def __init__( # Entry level errors bt_exceptions._MutateRowsIncomplete, ) - # build retryable operation - retry = retries.AsyncRetry( - predicate=self.is_retryable, - timeout=operation_timeout, - initial=0.01, - multiplier=2, - maximum=60, - ) - retry_wrapped = retry(self._run_attempt) - self._operation = _convert_retry_deadline( - retry_wrapped, operation_timeout, is_async=True + sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) + self._operation = retries.retry_target_async( + self._run_attempt, + self.is_retryable, + sleep_generator, + operation_timeout, + exception_factory=_retry_exception_factory, ) # initialize state self.timeout_generator = _attempt_timeout_generator( @@ -130,7 +125,7 @@ async def start(self): """ try: # trigger mutate_rows - await self._operation() + await self._operation except Exception as exc: # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations incomplete_indices = self.remaining_indices.copy() @@ -180,6 +175,7 @@ async def _run_attempt(self): result_generator = await self._gapic_fn( timeout=next(self.timeout_generator), entries=request_entries, + retry=None, ) async for result_list in result_generator: for result in result_list.entries: @@ -195,13 +191,6 @@ async def _run_attempt(self): self._handle_entry_error(orig_idx, entry_error) # remove processed entry from active list del active_request_indices[result.index] - except asyncio.CancelledError: - # when retry wrapper timeout expires, the operation is cancelled - # make sure incomplete indices are tracked, - # but don't record exception (it will be raised by wrapper) - # TODO: remove asyncio.wait_for in retry wrapper. Let grpc call handle expiration - self.remaining_indices.extend(active_request_indices.values()) - raise except Exception as exc: # add this exception to list for each mutation that wasn't # already handled, and update remaining_indices if mutation is retryable diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index ad1f7b84d..9e0fd78e1 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -31,15 +31,13 @@ from google.cloud.bigtable.data.row import Row, Cell from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data.exceptions import InvalidChunk -from google.cloud.bigtable.data.exceptions import RetryExceptionGroup from google.cloud.bigtable.data.exceptions import _RowSetComplete from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _make_metadata +from google.cloud.bigtable.data._helpers import _retry_exception_factory -from google.api_core import retry_async as retries -from google.api_core.retry_streaming_async import retry_target_stream +from google.api_core import retry as retries from google.api_core.retry import exponential_sleep_generator -from google.api_core import exceptions as core_exceptions if TYPE_CHECKING: from google.cloud.bigtable.data._async.client import TableAsync @@ -107,12 +105,12 @@ def start_operation(self) -> AsyncGenerator[Row, None]: """ Start the read_rows operation, retrying on retryable errors. """ - return retry_target_stream( + return retries.retry_target_stream_async( self._read_rows_attempt, self._predicate, exponential_sleep_generator(0.01, 60, multiplier=2), self.operation_timeout, - exception_factory=self._build_exception, + exception_factory=_retry_exception_factory, ) def _read_rows_attempt(self) -> AsyncGenerator[Row, None]: @@ -343,35 +341,3 @@ def _revise_request_rowset( # this will avoid an unwanted full table scan raise _RowSetComplete() return RowSetPB(row_keys=adjusted_keys, row_ranges=adjusted_ranges) - - @staticmethod - def _build_exception( - exc_list: list[Exception], is_timeout: bool, timeout_val: float - ) -> tuple[Exception, Exception | None]: - """ - Build retry error based on exceptions encountered during operation - - Args: - - exc_list: list of exceptions encountered during operation - - is_timeout: whether the operation failed due to timeout - - timeout_val: the operation timeout value in seconds, for constructing - the error message - Returns: - - tuple of the exception to raise, and a cause exception if applicable - """ - if is_timeout: - # if failed due to timeout, raise deadline exceeded as primary exception - source_exc: Exception = core_exceptions.DeadlineExceeded( - f"operation_timeout of {timeout_val} exceeded" - ) - elif exc_list: - # otherwise, raise non-retryable error as primary exception - source_exc = exc_list.pop() - else: - source_exc = RuntimeError("failed with unspecified exception") - # use the retry exception group as the cause of the exception - cause_exc: Exception | None = ( - RetryExceptionGroup(exc_list) if exc_list else None - ) - source_exc.__cause__ = cause_exc - return source_exc, cause_exc diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index a79ead7f8..d0578ff1a 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -33,6 +33,7 @@ import random import os +from functools import partial from google.cloud.bigtable_v2.services.bigtable.client import BigtableClientMeta from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient @@ -43,9 +44,8 @@ ) from google.cloud.bigtable_v2.types.bigtable import PingAndWarmRequest from google.cloud.client import ClientWithProject -from google.api_core.exceptions import GoogleAPICallError from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore -from google.api_core import retry_async as retries +from google.api_core import retry as retries from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import ServiceUnavailable from google.api_core.exceptions import Aborted @@ -65,7 +65,7 @@ from google.cloud.bigtable.data._helpers import _WarmedInstanceKey from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT from google.cloud.bigtable.data._helpers import _make_metadata -from google.cloud.bigtable.data._helpers import _convert_retry_deadline +from google.cloud.bigtable.data._helpers import _retry_exception_factory from google.cloud.bigtable.data._helpers import _validate_timeouts from google.cloud.bigtable.data._helpers import _get_retryable_errors from google.cloud.bigtable.data._helpers import _get_timeouts @@ -223,7 +223,7 @@ async def close(self, timeout: float = 2.0): async def _ping_and_warm_instances( self, channel: grpc.aio.Channel, instance_key: _WarmedInstanceKey | None = None - ) -> list[GoogleAPICallError | None]: + ) -> list[BaseException | None]: """ Prepares the backend for requests on a channel @@ -578,7 +578,6 @@ async def read_rows_stream( will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions from any retries that failed - GoogleAPIError: raised if the request encounters an unrecoverable error - - IdleTimeout: if iterator was abandoned """ operation_timeout, attempt_timeout = _get_timeouts( operation_timeout, attempt_timeout, self @@ -761,6 +760,9 @@ async def read_rows_sharded( for result in batch_result: if isinstance(result, Exception): error_dict[shard_idx] = result + elif isinstance(result, BaseException): + # BaseException not expected; raise immediately + raise result else: results_list.extend(result) shard_idx += 1 @@ -872,22 +874,8 @@ async def sample_row_keys( # prepare retryable retryable_excs = _get_retryable_errors(retryable_errors, self) predicate = retries.if_exception_type(*retryable_excs) - transient_errors = [] - def on_error_fn(exc): - # add errors to list if retryable - if predicate(exc): - transient_errors.append(exc) - - retry = retries.AsyncRetry( - predicate=predicate, - timeout=operation_timeout, - initial=0.01, - multiplier=2, - maximum=60, - on_error=on_error_fn, - is_stream=False, - ) + sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) # prepare request metadata = _make_metadata(self.table_name, self.app_profile_id) @@ -902,10 +890,13 @@ async def execute_rpc(): ) return [(s.row_key, s.offset_bytes) async for s in results] - wrapped_fn = _convert_retry_deadline( - retry(execute_rpc), operation_timeout, transient_errors, is_async=True + return await retries.retry_target_async( + execute_rpc, + predicate, + sleep_generator, + operation_timeout, + exception_factory=_retry_exception_factory, ) - return await wrapped_fn() def mutations_batcher( self, @@ -1014,37 +1005,25 @@ async def mutate_row( # mutations should not be retried predicate = retries.if_exception_type() - transient_errors = [] - - def on_error_fn(exc): - if predicate(exc): - transient_errors.append(exc) + sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - retry = retries.AsyncRetry( - predicate=predicate, - on_error=on_error_fn, - timeout=operation_timeout, - initial=0.01, - multiplier=2, - maximum=60, - ) - # wrap rpc in retry logic - retry_wrapped = retry(self.client._gapic_client.mutate_row) - # convert RetryErrors from retry wrapper into DeadlineExceeded errors - deadline_wrapped = _convert_retry_deadline( - retry_wrapped, operation_timeout, transient_errors, is_async=True - ) - metadata = _make_metadata(self.table_name, self.app_profile_id) - # trigger rpc - await deadline_wrapped( + target = partial( + self.client._gapic_client.mutate_row, row_key=row_key.encode("utf-8") if isinstance(row_key, str) else row_key, mutations=[mutation._to_pb() for mutation in mutations_list], table_name=self.table_name, app_profile_id=self.app_profile_id, timeout=attempt_timeout, - metadata=metadata, + metadata=_make_metadata(self.table_name, self.app_profile_id), retry=None, ) + return await retries.retry_target_async( + target, + predicate, + sleep_generator, + operation_timeout, + exception_factory=_retry_exception_factory, + ) async def bulk_mutate_rows( self, diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index b2da30040..5d5dd535e 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -489,6 +489,9 @@ async def _wait_for_batch_results( if isinstance(result, Exception): # will receive direct Exception objects if request task fails found_errors.append(result) + elif isinstance(result, BaseException): + # BaseException not expected from grpc calls. Raise immediately + raise result elif result: # completed requests will return a list of FailedMutationEntryError for e in result: diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 96ea1d1ce..a0b13cbaf 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -16,13 +16,14 @@ """ from __future__ import annotations -from typing import Callable, Sequence, List, Tuple, Any, TYPE_CHECKING +from typing import Sequence, List, Tuple, TYPE_CHECKING import time import enum from collections import namedtuple from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.api_core import exceptions as core_exceptions +from google.api_core.retry import RetryFailureReason from google.cloud.bigtable.data.exceptions import RetryExceptionGroup if TYPE_CHECKING: @@ -96,56 +97,37 @@ def _attempt_timeout_generator( yield max(0, min(per_request_timeout, deadline - time.monotonic())) -# TODO:replace this function with an exception_factory passed into the retry when -# feature is merged: -# https://github.com/googleapis/python-bigtable/blob/ea5b4f923e42516729c57113ddbe28096841b952/google/cloud/bigtable/data/_async/_read_rows.py#L130 -def _convert_retry_deadline( - func: Callable[..., Any], - timeout_value: float | None = None, - retry_errors: list[Exception] | None = None, - is_async: bool = False, -): +def _retry_exception_factory( + exc_list: list[Exception], + reason: RetryFailureReason, + timeout_val: float | None, +) -> tuple[Exception, Exception | None]: """ - Decorator to convert RetryErrors raised by api_core.retry into - DeadlineExceeded exceptions, indicating that the underlying retries have - exhaused the timeout value. - Optionally attaches a RetryExceptionGroup to the DeadlineExceeded.__cause__, - detailing the failed exceptions associated with each retry. - - Supports both sync and async function wrapping. + Build retry error based on exceptions encountered during operation Args: - - func: The function to decorate - - timeout_value: The timeout value to display in the DeadlineExceeded error message - - retry_errors: An optional list of exceptions to attach as a RetryExceptionGroup to the DeadlineExceeded.__cause__ + - exc_list: list of exceptions encountered during operation + - is_timeout: whether the operation failed due to timeout + - timeout_val: the operation timeout value in seconds, for constructing + the error message + Returns: + - tuple of the exception to raise, and a cause exception if applicable """ - timeout_str = f" of {timeout_value:.1f}s" if timeout_value is not None else "" - error_str = f"operation_timeout{timeout_str} exceeded" - - def handle_error(): - new_exc = core_exceptions.DeadlineExceeded( - error_str, + if reason == RetryFailureReason.TIMEOUT: + timeout_val_str = f"of {timeout_val:0.1f}s " if timeout_val is not None else "" + # if failed due to timeout, raise deadline exceeded as primary exception + source_exc: Exception = core_exceptions.DeadlineExceeded( + f"operation_timeout{timeout_val_str} exceeded" ) - source_exc = None - if retry_errors: - source_exc = RetryExceptionGroup(retry_errors) - new_exc.__cause__ = source_exc - raise new_exc from source_exc - - # separate wrappers for async and sync functions - async def wrapper_async(*args, **kwargs): - try: - return await func(*args, **kwargs) - except core_exceptions.RetryError: - handle_error() - - def wrapper(*args, **kwargs): - try: - return func(*args, **kwargs) - except core_exceptions.RetryError: - handle_error() - - return wrapper_async if is_async else wrapper + elif exc_list: + # otherwise, raise non-retryable error as primary exception + source_exc = exc_list.pop() + else: + source_exc = RuntimeError("failed with unspecified exception") + # use the retry exception group as the cause of the exception + cause_exc: Exception | None = RetryExceptionGroup(exc_list) if exc_list else None + source_exc.__cause__ = cause_exc + return source_exc, cause_exc def _get_timeouts( diff --git a/google/cloud/bigtable/data/exceptions.py b/google/cloud/bigtable/data/exceptions.py index 7344874df..3c73ec4e9 100644 --- a/google/cloud/bigtable/data/exceptions.py +++ b/google/cloud/bigtable/data/exceptions.py @@ -28,15 +28,6 @@ from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery -class IdleTimeout(core_exceptions.DeadlineExceeded): - """ - Exception raised by ReadRowsIterator when the generator - has been idle for longer than the internal idle_timeout. - """ - - pass - - class InvalidChunk(core_exceptions.GoogleAPICallError): """Exception raised to invalid chunk data from back-end.""" diff --git a/mypy.ini b/mypy.ini index f12ed46fc..3a17a37c6 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1,5 +1,5 @@ [mypy] -python_version = 3.6 +python_version = 3.7 namespace_packages = True exclude = tests/unit/gapic/ diff --git a/noxfile.py b/noxfile.py index e1d2f4acc..2e053ffcf 100644 --- a/noxfile.py +++ b/noxfile.py @@ -39,9 +39,7 @@ "pytest-cov", "pytest-asyncio", ] -UNIT_TEST_EXTERNAL_DEPENDENCIES = [ - # "git+https://github.com/googleapis/python-api-core.git@retry_generators" -] +UNIT_TEST_EXTERNAL_DEPENDENCIES = [] UNIT_TEST_LOCAL_DEPENDENCIES = [] UNIT_TEST_DEPENDENCIES = [] UNIT_TEST_EXTRAS = [] @@ -54,9 +52,7 @@ "pytest-asyncio", "google-cloud-testutils", ] -SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [ - # "git+https://github.com/googleapis/python-api-core.git@retry_generators" -] +SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [] SYSTEM_TEST_LOCAL_DEPENDENCIES = [] UNIT_TEST_DEPENDENCIES = [] SYSTEM_TEST_DEPENDENCIES = [] @@ -138,7 +134,8 @@ def mypy(session): session.install("google-cloud-testutils") session.run( "mypy", - "google/cloud/bigtable/data", + "-p", + "google.cloud.bigtable.data", "--check-untyped-defs", "--warn-unreachable", "--disallow-any-generics", @@ -460,7 +457,7 @@ def prerelease_deps(session): # Exclude version 1.52.0rc1 which has a known issue. See https://github.com/grpc/grpc/issues/32163 "grpcio!=1.52.0rc1", "grpcio-status", - "google-api-core==2.12.0.dev1", # TODO: remove this once streaming retries is merged + "google-api-core==2.16.0rc0", # TODO: remove pin once streaming retries is merged "proto-plus", "google-cloud-testutils", # dependencies of google-cloud-testutils" diff --git a/owlbot.py b/owlbot.py index b542b3246..9ca859fb9 100644 --- a/owlbot.py +++ b/owlbot.py @@ -170,7 +170,8 @@ def mypy(session): session.install("google-cloud-testutils") session.run( "mypy", - "google/cloud/bigtable", + "-p", + "google.cloud.bigtable", "--check-untyped-defs", "--warn-unreachable", "--disallow-any-generics", diff --git a/python-api-core b/python-api-core index a8cfa66b8..17ff5f1d8 160000 --- a/python-api-core +++ b/python-api-core @@ -1 +1 @@ -Subproject commit a8cfa66b8d6001da56823c6488b5da4957e5702b +Subproject commit 17ff5f1d83a9a6f50a0226fb0e794634bd584f17 diff --git a/setup.py b/setup.py index e5efc9937..0bce3a5d6 100644 --- a/setup.py +++ b/setup.py @@ -37,8 +37,8 @@ # 'Development Status :: 5 - Production/Stable' release_status = "Development Status :: 5 - Production/Stable" dependencies = [ - "google-api-core[grpc] == 2.12.0.dev1", # TODO: change to >= after streaming retries is merged - "google-cloud-core >= 1.4.1, <3.0.0dev", + "google-api-core[grpc] >= 2.16.0rc0", + "google-cloud-core >= 1.4.4, <3.0.0dev", "grpc-google-iam-v1 >= 0.12.4, <1.0.0dev", "proto-plus >= 1.22.0, <2.0.0dev", "proto-plus >= 1.22.2, <2.0.0dev; python_version>='3.11'", @@ -59,15 +59,10 @@ # benchmarks, etc. packages = [ package - for package in setuptools.PEP420PackageFinder.find() + for package in setuptools.find_namespace_packages() if package.startswith("google") ] -# Determine which namespaces are needed. -namespaces = ["google"] -if "google.cloud" in packages: - namespaces.append("google.cloud") - setuptools.setup( name=name, @@ -93,7 +88,6 @@ ], platforms="Posix; MacOS X; Windows", packages=packages, - namespace_packages=namespaces, install_requires=dependencies, extras_require=extras, scripts=[ diff --git a/testing/constraints-3.7.txt b/testing/constraints-3.7.txt index 9f23121d1..83bfe4577 100644 --- a/testing/constraints-3.7.txt +++ b/testing/constraints-3.7.txt @@ -5,7 +5,7 @@ # # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 -google-api-core==2.12.0.dev1 +google-api-core==2.16.0rc0 google-cloud-core==2.3.2 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 diff --git a/testing/constraints-3.8.txt b/testing/constraints-3.8.txt index 7045a2894..505ba9934 100644 --- a/testing/constraints-3.8.txt +++ b/testing/constraints-3.8.txt @@ -5,9 +5,10 @@ # # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 -google-api-core==2.12.0.dev1 +google-api-core==2.16.0rc0 google-cloud-core==2.3.2 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 libcst==0.2.5 protobuf==3.19.5 +pytest-asyncio==0.21.1 diff --git a/tests/system/data/test_system.py b/tests/system/data/test_system.py index 6bd21f386..fb0d9eb82 100644 --- a/tests/system/data/test_system.py +++ b/tests/system/data/test_system.py @@ -92,14 +92,15 @@ async def add_row( self.rows.append(row_key) async def delete_rows(self): - request = { - "table_name": self.table.table_name, - "entries": [ - {"row_key": row, "mutations": [{"delete_from_row": {}}]} - for row in self.rows - ], - } - await self.table.client._gapic_client.mutate_rows(request) + if self.rows: + request = { + "table_name": self.table.table_name, + "entries": [ + {"row_key": row, "mutations": [{"delete_from_row": {}}]} + for row in self.rows + ], + } + await self.table.client._gapic_client.mutate_rows(request) @pytest.mark.usefixtures("table") @@ -147,7 +148,7 @@ async def temp_rows(table): @pytest.mark.usefixtures("table") @pytest.mark.usefixtures("client") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=10) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=10) @pytest.mark.asyncio async def test_ping_and_warm_gapic(client, table): """ @@ -160,7 +161,7 @@ async def test_ping_and_warm_gapic(client, table): @pytest.mark.usefixtures("table") @pytest.mark.usefixtures("client") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_ping_and_warm(client, table): """ @@ -176,9 +177,9 @@ async def test_ping_and_warm(client, table): assert results[0] is None -@pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio +@pytest.mark.usefixtures("table") +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) async def test_mutation_set_cell(table, temp_rows): """ Ensure cells can be set properly @@ -196,7 +197,7 @@ async def test_mutation_set_cell(table, temp_rows): @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_sample_row_keys(client, table, temp_rows, column_split_config): """ @@ -239,7 +240,7 @@ async def test_bulk_mutations_set_cell(client, table, temp_rows): @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_context_manager(client, table, temp_rows): """ @@ -267,7 +268,7 @@ async def test_mutations_batcher_context_manager(client, table, temp_rows): @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_timer_flush(client, table, temp_rows): """ @@ -293,7 +294,7 @@ async def test_mutations_batcher_timer_flush(client, table, temp_rows): @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_count_flush(client, table, temp_rows): """ @@ -329,7 +330,7 @@ async def test_mutations_batcher_count_flush(client, table, temp_rows): @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_bytes_flush(client, table, temp_rows): """ @@ -366,7 +367,6 @@ async def test_mutations_batcher_bytes_flush(client, table, temp_rows): @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_mutations_batcher_no_flush(client, table, temp_rows): """ @@ -570,7 +570,7 @@ async def test_check_and_mutate( @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_stream(table, temp_rows): """ @@ -590,7 +590,7 @@ async def test_read_rows_stream(table, temp_rows): @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows(table, temp_rows): """ @@ -606,7 +606,7 @@ async def test_read_rows(table, temp_rows): @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_sharded_simple(table, temp_rows): """ @@ -629,7 +629,7 @@ async def test_read_rows_sharded_simple(table, temp_rows): @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_sharded_from_sample(table, temp_rows): """ @@ -654,7 +654,7 @@ async def test_read_rows_sharded_from_sample(table, temp_rows): @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_sharded_filters_limits(table, temp_rows): """ @@ -683,7 +683,7 @@ async def test_read_rows_sharded_filters_limits(table, temp_rows): @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_range_query(table, temp_rows): """ @@ -705,7 +705,7 @@ async def test_read_rows_range_query(table, temp_rows): @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_single_key_query(table, temp_rows): """ @@ -726,7 +726,7 @@ async def test_read_rows_single_key_query(table, temp_rows): @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio async def test_read_rows_with_filter(table, temp_rows): """ @@ -842,7 +842,7 @@ async def test_row_exists(table, temp_rows): @pytest.mark.usefixtures("table") -@retry.Retry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) +@retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.parametrize( "cell_value,filter_input,expect_match", [ diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index d41929518..e03028c45 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -164,11 +164,13 @@ async def test_mutate_rows_operation(self): table = mock.Mock() entries = [_make_mutation(), _make_mutation()] operation_timeout = 0.05 - instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout - ) - with mock.patch.object(instance, "_operation", AsyncMock()) as attempt_mock: - attempt_mock.return_value = None + cls = self._target_class() + with mock.patch( + f"{cls.__module__}.{cls.__name__}._run_attempt", AsyncMock() + ) as attempt_mock: + instance = self._make_one( + client, table, entries, operation_timeout, operation_timeout + ) await instance.start() assert attempt_mock.call_count == 1 diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 54bbb6158..60a305bcb 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1142,44 +1142,44 @@ def test_table_ctor_sync(self): ( "read_rows_stream", (ReadRowsQuery(),), - "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + "google.api_core.retry.retry_target_stream_async", (), ), ( "read_rows", (ReadRowsQuery(),), - "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + "google.api_core.retry.retry_target_stream_async", (), ), ( "read_row", (b"row_key",), - "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + "google.api_core.retry.retry_target_stream_async", (), ), ( "read_rows_sharded", ([ReadRowsQuery()],), - "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + "google.api_core.retry.retry_target_stream_async", (), ), ( "row_exists", (b"row_key",), - "google.cloud.bigtable.data._async._read_rows.retry_target_stream", + "google.api_core.retry.retry_target_stream_async", (), ), - ("sample_row_keys", (), "google.api_core.retry_async.retry_target", ()), + ("sample_row_keys", (), "google.api_core.retry.retry_target_async", ()), ( "mutate_row", - (b"row_key", []), - "google.api_core.retry_async.retry_target", + (b"row_key", [mock.Mock()]), + "google.api_core.retry.retry_target_async", (), ), ( "bulk_mutate_rows", ([mutations.RowMutationEntry(b"key", [mock.Mock()])],), - "google.api_core.retry_async.retry_target", + "google.api_core.retry.retry_target_async", (_MutateRowsIncomplete,), ), ], @@ -1223,15 +1223,15 @@ async def test_customizable_retryable_errors( """ from google.cloud.bigtable.data import BigtableDataClientAsync - with mock.patch( - "google.api_core.retry_async.if_exception_type" - ) as predicate_builder_mock: - with mock.patch(retry_fn_path) as retry_fn_mock: - async with BigtableDataClientAsync() as client: - table = client.get_table("instance-id", "table-id") - expected_predicate = lambda a: a in expected_retryables # noqa + with mock.patch(retry_fn_path) as retry_fn_mock: + async with BigtableDataClientAsync() as client: + table = client.get_table("instance-id", "table-id") + expected_predicate = lambda a: a in expected_retryables # noqa + retry_fn_mock.side_effect = RuntimeError("stop early") + with mock.patch( + "google.api_core.retry.if_exception_type" + ) as predicate_builder_mock: predicate_builder_mock.return_value = expected_predicate - retry_fn_mock.side_effect = RuntimeError("stop early") with pytest.raises(Exception): # we expect an exception from attempting to call the mock test_fn = table.__getattribute__(fn_name) @@ -1253,10 +1253,10 @@ async def test_customizable_retryable_errors( ("read_rows_sharded", ([ReadRowsQuery()],), "read_rows"), ("row_exists", (b"row_key",), "read_rows"), ("sample_row_keys", (), "sample_row_keys"), - ("mutate_row", (b"row_key", []), "mutate_row"), + ("mutate_row", (b"row_key", [mock.Mock()]), "mutate_row"), ( "bulk_mutate_rows", - ([mutations.RowMutationEntry(b"key", [mock.Mock()])],), + ([mutations.RowMutationEntry(b"key", [mutations.DeleteAllFromRow()])],), "mutate_rows", ), ("check_and_mutate_row", (b"row_key", None), "check_and_mutate_row"), @@ -2204,7 +2204,7 @@ async def test_mutate_row_retryable_errors(self, retryable_exception): mutation = mutations.DeleteAllFromRow() assert mutation.is_idempotent() is True await table.mutate_row( - "row_key", mutation, operation_timeout=0.05 + "row_key", mutation, operation_timeout=0.01 ) cause = e.value.__cause__ assert isinstance(cause, RetryExceptionGroup) diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py index 17bd8d420..446cd822e 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -1158,10 +1158,10 @@ async def test_customizable_retryable_errors( from google.cloud.bigtable.data._async.client import TableAsync with mock.patch( - "google.api_core.retry_async.if_exception_type" + "google.api_core.retry.if_exception_type" ) as predicate_builder_mock: with mock.patch( - "google.api_core.retry_async.retry_target" + "google.api_core.retry.retry_target_async" ) as retry_fn_mock: table = None with mock.patch("asyncio.create_task"): diff --git a/tests/unit/data/test__helpers.py b/tests/unit/data/test__helpers.py index b9c1dc2bb..5a9c500ed 100644 --- a/tests/unit/data/test__helpers.py +++ b/tests/unit/data/test__helpers.py @@ -17,7 +17,6 @@ from google.api_core import exceptions as core_exceptions import google.cloud.bigtable.data._helpers as _helpers from google.cloud.bigtable.data._helpers import TABLE_DEFAULT -import google.cloud.bigtable.data.exceptions as bigtable_exceptions import mock @@ -100,71 +99,6 @@ def test_attempt_timeout_w_sleeps(self): expected_value -= sleep_time -class TestConvertRetryDeadline: - """ - Test _convert_retry_deadline wrapper - """ - - @pytest.mark.asyncio - @pytest.mark.parametrize("is_async", [True, False]) - async def test_no_error(self, is_async): - def test_func(): - return 1 - - async def test_async(): - return test_func() - - func = test_async if is_async else test_func - wrapped = _helpers._convert_retry_deadline(func, 0.1, is_async) - result = await wrapped() if is_async else wrapped() - assert result == 1 - - @pytest.mark.asyncio - @pytest.mark.parametrize("timeout", [0.1, 2.0, 30.0]) - @pytest.mark.parametrize("is_async", [True, False]) - async def test_retry_error(self, timeout, is_async): - from google.api_core.exceptions import RetryError, DeadlineExceeded - - def test_func(): - raise RetryError("retry error", None) - - async def test_async(): - return test_func() - - func = test_async if is_async else test_func - wrapped = _helpers._convert_retry_deadline(func, timeout, is_async=is_async) - with pytest.raises(DeadlineExceeded) as e: - await wrapped() if is_async else wrapped() - assert e.value.__cause__ is None - assert f"operation_timeout of {timeout}s exceeded" in str(e.value) - - @pytest.mark.asyncio - @pytest.mark.parametrize("is_async", [True, False]) - async def test_with_retry_errors(self, is_async): - from google.api_core.exceptions import RetryError, DeadlineExceeded - - timeout = 10.0 - - def test_func(): - raise RetryError("retry error", None) - - async def test_async(): - return test_func() - - func = test_async if is_async else test_func - - associated_errors = [RuntimeError("error1"), ZeroDivisionError("other")] - wrapped = _helpers._convert_retry_deadline( - func, timeout, associated_errors, is_async - ) - with pytest.raises(DeadlineExceeded) as e: - await wrapped() - cause = e.value.__cause__ - assert isinstance(cause, bigtable_exceptions.RetryExceptionGroup) - assert cause.exceptions == tuple(associated_errors) - assert f"operation_timeout of {timeout}s exceeded" in str(e.value) - - class TestValidateTimeouts: def test_validate_timeouts_error_messages(self): with pytest.raises(ValueError) as e: From 858b93aee7512c1dd45f57eea8b500864eeedae6 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 15 Dec 2023 15:30:15 -0800 Subject: [PATCH 31/56] chore: merge main into experimental_v3 (#900) --- .coveragerc | 8 - .flake8 | 2 +- .github/.OwlBot.lock.yaml | 4 +- .github/auto-label.yaml | 2 +- .github/workflows/docs.yml | 10 +- .github/workflows/lint.yml | 4 +- .github/workflows/mypy.yml | 4 +- .github/workflows/system_emulated.yml | 8 +- .github/workflows/unittest.yml | 10 +- .gitignore | 1 + .kokoro/build.sh | 2 +- .kokoro/docker/docs/Dockerfile | 2 +- .kokoro/populate-secrets.sh | 2 +- .kokoro/publish-docs.sh | 2 +- .kokoro/release.sh | 2 +- .kokoro/release/common.cfg | 9 + .kokoro/requirements.txt | 535 +++++---- .kokoro/samples/python3.12/common.cfg | 40 + .kokoro/samples/python3.12/continuous.cfg | 6 + .kokoro/samples/python3.12/periodic-head.cfg | 11 + .kokoro/samples/python3.12/periodic.cfg | 6 + .kokoro/samples/python3.12/presubmit.cfg | 6 + .kokoro/test-samples-against-head.sh | 2 +- .kokoro/test-samples-impl.sh | 2 +- .kokoro/test-samples.sh | 2 +- .kokoro/trampoline.sh | 2 +- .kokoro/trampoline_v2.sh | 2 +- .pre-commit-config.yaml | 6 +- .release-please-manifest.json | 2 +- .trampolinerc | 4 +- CHANGELOG.md | 54 + CONTRIBUTING.rst | 6 +- MANIFEST.in | 2 +- docs/conf.py | 2 +- docs/snippets.py | 1 - docs/snippets_table.py | 1 - google/cloud/bigtable/batcher.py | 129 +- google/cloud/bigtable/gapic_version.py | 2 +- google/cloud/bigtable_admin/__init__.py | 6 +- google/cloud/bigtable_admin/gapic_version.py | 2 +- google/cloud/bigtable_admin_v2/__init__.py | 6 +- .../bigtable_admin_v2/gapic_metadata.json | 15 + .../cloud/bigtable_admin_v2/gapic_version.py | 2 +- .../bigtable_admin_v2/services/__init__.py | 2 +- .../bigtable_instance_admin/__init__.py | 2 +- .../bigtable_instance_admin/async_client.py | 148 +-- .../bigtable_instance_admin/client.py | 72 +- .../bigtable_instance_admin/pagers.py | 2 +- .../transports/__init__.py | 2 +- .../transports/base.py | 2 +- .../transports/grpc.py | 2 +- .../transports/grpc_asyncio.py | 2 +- .../transports/rest.py | 176 +-- .../services/bigtable_table_admin/__init__.py | 2 +- .../bigtable_table_admin/async_client.py | 294 +++-- .../services/bigtable_table_admin/client.py | 218 ++-- .../services/bigtable_table_admin/pagers.py | 2 +- .../transports/__init__.py | 2 +- .../bigtable_table_admin/transports/base.py | 16 +- .../bigtable_table_admin/transports/grpc.py | 38 +- .../transports/grpc_asyncio.py | 40 +- .../bigtable_table_admin/transports/rest.py | 313 +++-- .../cloud/bigtable_admin_v2/types/__init__.py | 6 +- .../types/bigtable_instance_admin.py | 2 +- .../types/bigtable_table_admin.py | 107 +- .../cloud/bigtable_admin_v2/types/common.py | 2 +- .../cloud/bigtable_admin_v2/types/instance.py | 67 +- google/cloud/bigtable_admin_v2/types/table.py | 61 +- google/cloud/bigtable_v2/__init__.py | 2 +- google/cloud/bigtable_v2/gapic_version.py | 2 +- google/cloud/bigtable_v2/services/__init__.py | 2 +- .../bigtable_v2/services/bigtable/__init__.py | 2 +- .../services/bigtable/async_client.py | 26 +- .../bigtable_v2/services/bigtable/client.py | 2 +- .../services/bigtable/transports/__init__.py | 2 +- .../services/bigtable/transports/base.py | 2 +- .../services/bigtable/transports/grpc.py | 2 +- .../bigtable/transports/grpc_asyncio.py | 2 +- .../services/bigtable/transports/rest.py | 2 +- google/cloud/bigtable_v2/types/__init__.py | 2 +- google/cloud/bigtable_v2/types/bigtable.py | 21 +- google/cloud/bigtable_v2/types/data.py | 3 +- .../cloud/bigtable_v2/types/feature_flags.py | 48 +- .../cloud/bigtable_v2/types/request_stats.py | 3 +- .../bigtable_v2/types/response_params.py | 2 +- mypy.ini | 2 +- noxfile.py | 47 +- owlbot.py | 2 - samples/beam/noxfile.py | 2 +- samples/beam/requirements-test.txt | 2 +- samples/beam/requirements.txt | 6 +- samples/hello/noxfile.py | 2 +- samples/hello/requirements-test.txt | 2 +- samples/hello/requirements.txt | 4 +- samples/hello_happybase/noxfile.py | 2 +- samples/hello_happybase/requirements-test.txt | 2 +- samples/hello_happybase/requirements.txt | 1 + samples/instanceadmin/noxfile.py | 2 +- samples/instanceadmin/requirements-test.txt | 2 +- samples/instanceadmin/requirements.txt | 2 +- samples/metricscaler/noxfile.py | 2 +- samples/metricscaler/requirements-test.txt | 4 +- samples/metricscaler/requirements.txt | 4 +- samples/quickstart/noxfile.py | 2 +- samples/quickstart/requirements-test.txt | 2 +- samples/quickstart/requirements.txt | 2 +- samples/quickstart_happybase/noxfile.py | 2 +- .../requirements-test.txt | 2 +- samples/quickstart_happybase/requirements.txt | 1 + samples/snippets/deletes/noxfile.py | 2 +- .../snippets/deletes/requirements-test.txt | 2 +- samples/snippets/deletes/requirements.txt | 2 +- samples/snippets/filters/noxfile.py | 2 +- .../snippets/filters/requirements-test.txt | 2 +- samples/snippets/filters/requirements.txt | 2 +- samples/snippets/reads/noxfile.py | 2 +- samples/snippets/reads/requirements-test.txt | 2 +- samples/snippets/reads/requirements.txt | 2 +- samples/snippets/writes/noxfile.py | 2 +- samples/snippets/writes/requirements-test.txt | 2 +- samples/snippets/writes/requirements.txt | 2 +- samples/tableadmin/noxfile.py | 2 +- samples/tableadmin/requirements-test.txt | 4 +- samples/tableadmin/requirements.txt | 2 +- scripts/decrypt-secrets.sh | 2 +- scripts/fixup_bigtable_admin_v2_keywords.py | 3 +- scripts/fixup_bigtable_v2_keywords.py | 4 +- scripts/readme-gen/readme_gen.py | 18 +- setup.cfg | 2 +- setup.py | 6 +- testing/constraints-3.12.txt | 0 testing/constraints-3.7.txt | 2 +- testing/constraints-3.8.txt | 2 +- tests/__init__.py | 2 +- tests/system/v2_client/test_data_api.py | 36 + tests/system/v2_client/test_instance_admin.py | 1 - tests/unit/data/_async/test_client.py | 1 - tests/unit/data/test_exceptions.py | 2 - tests/unit/gapic/__init__.py | 2 +- .../unit/gapic/bigtable_admin_v2/__init__.py | 2 +- .../test_bigtable_instance_admin.py | 629 +++++++--- .../test_bigtable_table_admin.py | 1051 ++++++++++++++--- tests/unit/gapic/bigtable_v2/__init__.py | 2 +- tests/unit/gapic/bigtable_v2/test_bigtable.py | 143 ++- tests/unit/test_packaging.py | 37 + tests/unit/v2_client/test_batcher.py | 31 +- tests/unit/v2_client/test_cluster.py | 4 - tests/unit/v2_client/test_column_family.py | 1 - tests/unit/v2_client/test_instance.py | 1 - tests/unit/v2_client/test_row_data.py | 26 +- tests/unit/v2_client/test_table.py | 9 +- 151 files changed, 3275 insertions(+), 1508 deletions(-) create mode 100644 .kokoro/samples/python3.12/common.cfg create mode 100644 .kokoro/samples/python3.12/continuous.cfg create mode 100644 .kokoro/samples/python3.12/periodic-head.cfg create mode 100644 .kokoro/samples/python3.12/periodic.cfg create mode 100644 .kokoro/samples/python3.12/presubmit.cfg create mode 100644 testing/constraints-3.12.txt create mode 100644 tests/unit/test_packaging.py diff --git a/.coveragerc b/.coveragerc index 702b85681..f12d4dc21 100644 --- a/.coveragerc +++ b/.coveragerc @@ -18,8 +18,6 @@ [run] branch = True omit = - google/cloud/__init__.py - google/__init__.py google/cloud/bigtable_admin/__init__.py google/cloud/bigtable_admin/gapic_version.py @@ -33,11 +31,5 @@ exclude_lines = def __repr__ # Ignore abstract methods raise NotImplementedError - # Ignore setuptools-less fallback - except pkg_resources.DistributionNotFound: omit = - */gapic/*.py - */proto/*.py - */core/*.py */site-packages/*.py - google/cloud/__init__.py diff --git a/.flake8 b/.flake8 index 2e4387498..87f6e408c 100644 --- a/.flake8 +++ b/.flake8 @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.github/.OwlBot.lock.yaml b/.github/.OwlBot.lock.yaml index 1b3cb6c52..40bf99731 100644 --- a/.github/.OwlBot.lock.yaml +++ b/.github/.OwlBot.lock.yaml @@ -13,5 +13,5 @@ # limitations under the License. docker: image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest - digest: sha256:ddf4551385d566771dc713090feb7b4c1164fb8a698fe52bbe7670b24236565b -# created: 2023-06-27T13:04:21.96690344Z + digest: sha256:230f7fe8a0d2ed81a519cfc15c6bb11c5b46b9fb449b8b1219b3771bcb520ad2 +# created: 2023-12-09T15:16:25.430769578Z diff --git a/.github/auto-label.yaml b/.github/auto-label.yaml index 41bff0b53..b2016d119 100644 --- a/.github/auto-label.yaml +++ b/.github/auto-label.yaml @@ -1,4 +1,4 @@ -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index e97d89e48..698fbc5c9 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -8,9 +8,9 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: "3.9" - name: Install nox @@ -24,11 +24,11 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: - python-version: "3.9" + python-version: "3.10" - name: Install nox run: | python -m pip install --upgrade setuptools pip wheel diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 16d5a9e90..4866193af 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -8,9 +8,9 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: "3.8" - name: Install nox diff --git a/.github/workflows/mypy.yml b/.github/workflows/mypy.yml index c63242630..3915cddd3 100644 --- a/.github/workflows/mypy.yml +++ b/.github/workflows/mypy.yml @@ -8,9 +8,9 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: "3.8" - name: Install nox diff --git a/.github/workflows/system_emulated.yml b/.github/workflows/system_emulated.yml index e1f43fd40..7669901c9 100644 --- a/.github/workflows/system_emulated.yml +++ b/.github/workflows/system_emulated.yml @@ -7,20 +7,20 @@ on: jobs: run-systests: - runs-on: ubuntu-20.04 + runs-on: ubuntu-22.04 steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: '3.8' - name: Setup GCloud SDK - uses: google-github-actions/setup-gcloud@v1.1.0 + uses: google-github-actions/setup-gcloud@v2.0.0 - name: Install / run Nox run: | diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index 8057a7691..d6ca65627 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -8,12 +8,12 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python: ['3.7', '3.8', '3.9', '3.10', '3.11'] + python: ['3.7', '3.8', '3.9', '3.10', '3.11', '3.12'] steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python }} - name: Install nox @@ -37,9 +37,9 @@ jobs: - unit steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Setup Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: "3.8" - name: Install coverage diff --git a/.gitignore b/.gitignore index b4243ced7..d083ea1dd 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,7 @@ docs.metadata # Virtual environment env/ +venv/ # Test logs coverage.xml diff --git a/.kokoro/build.sh b/.kokoro/build.sh index 2ab1155b2..dec6b66a7 100755 --- a/.kokoro/build.sh +++ b/.kokoro/build.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2018 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/docker/docs/Dockerfile b/.kokoro/docker/docs/Dockerfile index f8137d0ae..8e39a2cc4 100644 --- a/.kokoro/docker/docs/Dockerfile +++ b/.kokoro/docker/docs/Dockerfile @@ -1,4 +1,4 @@ -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/populate-secrets.sh b/.kokoro/populate-secrets.sh index f52514257..6f3972140 100755 --- a/.kokoro/populate-secrets.sh +++ b/.kokoro/populate-secrets.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2020 Google LLC. +# Copyright 2023 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/publish-docs.sh b/.kokoro/publish-docs.sh index 1c4d62370..9eafe0be3 100755 --- a/.kokoro/publish-docs.sh +++ b/.kokoro/publish-docs.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/release.sh b/.kokoro/release.sh index 6b594c813..2e1cbfa81 100755 --- a/.kokoro/release.sh +++ b/.kokoro/release.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/release/common.cfg b/.kokoro/release/common.cfg index 8477e4ca6..2a8fd970c 100644 --- a/.kokoro/release/common.cfg +++ b/.kokoro/release/common.cfg @@ -38,3 +38,12 @@ env_vars: { key: "SECRET_MANAGER_KEYS" value: "releasetool-publish-reporter-app,releasetool-publish-reporter-googleapis-installation,releasetool-publish-reporter-pem" } + +# Store the packages we uploaded to PyPI. That way, we have a record of exactly +# what we published, which we can use to generate SBOMs and attestations. +action { + define_artifacts { + regex: "github/python-bigtable/**/*.tar.gz" + strip_prefix: "github/python-bigtable" + } +} diff --git a/.kokoro/requirements.txt b/.kokoro/requirements.txt index c7929db6d..e5c1ffca9 100644 --- a/.kokoro/requirements.txt +++ b/.kokoro/requirements.txt @@ -4,91 +4,75 @@ # # pip-compile --allow-unsafe --generate-hashes requirements.in # -argcomplete==2.0.0 \ - --hash=sha256:6372ad78c89d662035101418ae253668445b391755cfe94ea52f1b9d22425b20 \ - --hash=sha256:cffa11ea77999bb0dd27bb25ff6dc142a6796142f68d45b1a26b11f58724561e +argcomplete==3.1.4 \ + --hash=sha256:72558ba729e4c468572609817226fb0a6e7e9a0a7d477b882be168c0b4a62b94 \ + --hash=sha256:fbe56f8cda08aa9a04b307d8482ea703e96a6a801611acb4be9bf3942017989f # via nox -attrs==22.1.0 \ - --hash=sha256:29adc2665447e5191d0e7c568fde78b21f9672d344281d0c6e1ab085429b22b6 \ - --hash=sha256:86efa402f67bf2df34f51a335487cf46b1ec130d02b8d39fd248abfd30da551c +attrs==23.1.0 \ + --hash=sha256:1f28b4522cdc2fb4256ac1a020c78acf9cba2c6b461ccd2c126f3aa8e8335d04 \ + --hash=sha256:6279836d581513a26f1bf235f9acd333bc9115683f14f7e8fae46c98fc50e015 # via gcp-releasetool -bleach==5.0.1 \ - --hash=sha256:085f7f33c15bd408dd9b17a4ad77c577db66d76203e5984b1bd59baeee948b2a \ - --hash=sha256:0d03255c47eb9bd2f26aa9bb7f2107732e7e8fe195ca2f64709fcf3b0a4a085c - # via readme-renderer -cachetools==5.2.0 \ - --hash=sha256:6a94c6402995a99c3970cc7e4884bb60b4a8639938157eeed436098bf9831757 \ - --hash=sha256:f9f17d2aec496a9aa6b76f53e3b614c965223c061982d434d160f930c698a9db +cachetools==5.3.2 \ + --hash=sha256:086ee420196f7b2ab9ca2db2520aca326318b68fe5ba8bc4d49cca91add450f2 \ + --hash=sha256:861f35a13a451f94e301ce2bec7cac63e881232ccce7ed67fab9b5df4d3beaa1 # via google-auth -certifi==2022.12.7 \ - --hash=sha256:35824b4c3a97115964b408844d64aa14db1cc518f6562e8d7261699d1350a9e3 \ - --hash=sha256:4ad3232f5e926d6718ec31cfc1fcadfde020920e278684144551c91769c7bc18 +certifi==2023.7.22 \ + --hash=sha256:539cc1d13202e33ca466e88b2807e29f4c13049d6d87031a3c110744495cb082 \ + --hash=sha256:92d6037539857d8206b8f6ae472e8b77db8058fec5937a1ef3f54304089edbb9 # via requests -cffi==1.15.1 \ - --hash=sha256:00a9ed42e88df81ffae7a8ab6d9356b371399b91dbdf0c3cb1e84c03a13aceb5 \ - --hash=sha256:03425bdae262c76aad70202debd780501fabeaca237cdfddc008987c0e0f59ef \ - --hash=sha256:04ed324bda3cda42b9b695d51bb7d54b680b9719cfab04227cdd1e04e5de3104 \ - --hash=sha256:0e2642fe3142e4cc4af0799748233ad6da94c62a8bec3a6648bf8ee68b1c7426 \ - --hash=sha256:173379135477dc8cac4bc58f45db08ab45d228b3363adb7af79436135d028405 \ - --hash=sha256:198caafb44239b60e252492445da556afafc7d1e3ab7a1fb3f0584ef6d742375 \ - --hash=sha256:1e74c6b51a9ed6589199c787bf5f9875612ca4a8a0785fb2d4a84429badaf22a \ - --hash=sha256:2012c72d854c2d03e45d06ae57f40d78e5770d252f195b93f581acf3ba44496e \ - --hash=sha256:21157295583fe8943475029ed5abdcf71eb3911894724e360acff1d61c1d54bc \ - --hash=sha256:2470043b93ff09bf8fb1d46d1cb756ce6132c54826661a32d4e4d132e1977adf \ - --hash=sha256:285d29981935eb726a4399badae8f0ffdff4f5050eaa6d0cfc3f64b857b77185 \ - --hash=sha256:30d78fbc8ebf9c92c9b7823ee18eb92f2e6ef79b45ac84db507f52fbe3ec4497 \ - --hash=sha256:320dab6e7cb2eacdf0e658569d2575c4dad258c0fcc794f46215e1e39f90f2c3 \ - --hash=sha256:33ab79603146aace82c2427da5ca6e58f2b3f2fb5da893ceac0c42218a40be35 \ - --hash=sha256:3548db281cd7d2561c9ad9984681c95f7b0e38881201e157833a2342c30d5e8c \ - --hash=sha256:3799aecf2e17cf585d977b780ce79ff0dc9b78d799fc694221ce814c2c19db83 \ - --hash=sha256:39d39875251ca8f612b6f33e6b1195af86d1b3e60086068be9cc053aa4376e21 \ - --hash=sha256:3b926aa83d1edb5aa5b427b4053dc420ec295a08e40911296b9eb1b6170f6cca \ - --hash=sha256:3bcde07039e586f91b45c88f8583ea7cf7a0770df3a1649627bf598332cb6984 \ - --hash=sha256:3d08afd128ddaa624a48cf2b859afef385b720bb4b43df214f85616922e6a5ac \ - --hash=sha256:3eb6971dcff08619f8d91607cfc726518b6fa2a9eba42856be181c6d0d9515fd \ - --hash=sha256:40f4774f5a9d4f5e344f31a32b5096977b5d48560c5592e2f3d2c4374bd543ee \ - --hash=sha256:4289fc34b2f5316fbb762d75362931e351941fa95fa18789191b33fc4cf9504a \ - --hash=sha256:470c103ae716238bbe698d67ad020e1db9d9dba34fa5a899b5e21577e6d52ed2 \ - --hash=sha256:4f2c9f67e9821cad2e5f480bc8d83b8742896f1242dba247911072d4fa94c192 \ - --hash=sha256:50a74364d85fd319352182ef59c5c790484a336f6db772c1a9231f1c3ed0cbd7 \ - --hash=sha256:54a2db7b78338edd780e7ef7f9f6c442500fb0d41a5a4ea24fff1c929d5af585 \ - --hash=sha256:5635bd9cb9731e6d4a1132a498dd34f764034a8ce60cef4f5319c0541159392f \ - --hash=sha256:59c0b02d0a6c384d453fece7566d1c7e6b7bae4fc5874ef2ef46d56776d61c9e \ - --hash=sha256:5d598b938678ebf3c67377cdd45e09d431369c3b1a5b331058c338e201f12b27 \ - --hash=sha256:5df2768244d19ab7f60546d0c7c63ce1581f7af8b5de3eb3004b9b6fc8a9f84b \ - --hash=sha256:5ef34d190326c3b1f822a5b7a45f6c4535e2f47ed06fec77d3d799c450b2651e \ - --hash=sha256:6975a3fac6bc83c4a65c9f9fcab9e47019a11d3d2cf7f3c0d03431bf145a941e \ - --hash=sha256:6c9a799e985904922a4d207a94eae35c78ebae90e128f0c4e521ce339396be9d \ - --hash=sha256:70df4e3b545a17496c9b3f41f5115e69a4f2e77e94e1d2a8e1070bc0c38c8a3c \ - --hash=sha256:7473e861101c9e72452f9bf8acb984947aa1661a7704553a9f6e4baa5ba64415 \ - --hash=sha256:8102eaf27e1e448db915d08afa8b41d6c7ca7a04b7d73af6514df10a3e74bd82 \ - --hash=sha256:87c450779d0914f2861b8526e035c5e6da0a3199d8f1add1a665e1cbc6fc6d02 \ - --hash=sha256:8b7ee99e510d7b66cdb6c593f21c043c248537a32e0bedf02e01e9553a172314 \ - --hash=sha256:91fc98adde3d7881af9b59ed0294046f3806221863722ba7d8d120c575314325 \ - --hash=sha256:94411f22c3985acaec6f83c6df553f2dbe17b698cc7f8ae751ff2237d96b9e3c \ - --hash=sha256:98d85c6a2bef81588d9227dde12db8a7f47f639f4a17c9ae08e773aa9c697bf3 \ - --hash=sha256:9ad5db27f9cabae298d151c85cf2bad1d359a1b9c686a275df03385758e2f914 \ - --hash=sha256:a0b71b1b8fbf2b96e41c4d990244165e2c9be83d54962a9a1d118fd8657d2045 \ - --hash=sha256:a0f100c8912c114ff53e1202d0078b425bee3649ae34d7b070e9697f93c5d52d \ - --hash=sha256:a591fe9e525846e4d154205572a029f653ada1a78b93697f3b5a8f1f2bc055b9 \ - --hash=sha256:a5c84c68147988265e60416b57fc83425a78058853509c1b0629c180094904a5 \ - --hash=sha256:a66d3508133af6e8548451b25058d5812812ec3798c886bf38ed24a98216fab2 \ - --hash=sha256:a8c4917bd7ad33e8eb21e9a5bbba979b49d9a97acb3a803092cbc1133e20343c \ - --hash=sha256:b3bbeb01c2b273cca1e1e0c5df57f12dce9a4dd331b4fa1635b8bec26350bde3 \ - --hash=sha256:cba9d6b9a7d64d4bd46167096fc9d2f835e25d7e4c121fb2ddfc6528fb0413b2 \ - --hash=sha256:cc4d65aeeaa04136a12677d3dd0b1c0c94dc43abac5860ab33cceb42b801c1e8 \ - --hash=sha256:ce4bcc037df4fc5e3d184794f27bdaab018943698f4ca31630bc7f84a7b69c6d \ - --hash=sha256:cec7d9412a9102bdc577382c3929b337320c4c4c4849f2c5cdd14d7368c5562d \ - --hash=sha256:d400bfb9a37b1351253cb402671cea7e89bdecc294e8016a707f6d1d8ac934f9 \ - --hash=sha256:d61f4695e6c866a23a21acab0509af1cdfd2c013cf256bbf5b6b5e2695827162 \ - --hash=sha256:db0fbb9c62743ce59a9ff687eb5f4afbe77e5e8403d6697f7446e5f609976f76 \ - --hash=sha256:dd86c085fae2efd48ac91dd7ccffcfc0571387fe1193d33b6394db7ef31fe2a4 \ - --hash=sha256:e00b098126fd45523dd056d2efba6c5a63b71ffe9f2bbe1a4fe1716e1d0c331e \ - --hash=sha256:e229a521186c75c8ad9490854fd8bbdd9a0c9aa3a524326b55be83b54d4e0ad9 \ - --hash=sha256:e263d77ee3dd201c3a142934a086a4450861778baaeeb45db4591ef65550b0a6 \ - --hash=sha256:ed9cb427ba5504c1dc15ede7d516b84757c3e3d7868ccc85121d9310d27eed0b \ - --hash=sha256:fa6693661a4c91757f4412306191b6dc88c1703f780c8234035eac011922bc01 \ - --hash=sha256:fcd131dd944808b5bdb38e6f5b53013c5aa4f334c5cad0c72742f6eba4b73db0 +cffi==1.16.0 \ + --hash=sha256:0c9ef6ff37e974b73c25eecc13952c55bceed9112be2d9d938ded8e856138bcc \ + --hash=sha256:131fd094d1065b19540c3d72594260f118b231090295d8c34e19a7bbcf2e860a \ + --hash=sha256:1b8ebc27c014c59692bb2664c7d13ce7a6e9a629be20e54e7271fa696ff2b417 \ + --hash=sha256:2c56b361916f390cd758a57f2e16233eb4f64bcbeee88a4881ea90fca14dc6ab \ + --hash=sha256:2d92b25dbf6cae33f65005baf472d2c245c050b1ce709cc4588cdcdd5495b520 \ + --hash=sha256:31d13b0f99e0836b7ff893d37af07366ebc90b678b6664c955b54561fc36ef36 \ + --hash=sha256:32c68ef735dbe5857c810328cb2481e24722a59a2003018885514d4c09af9743 \ + --hash=sha256:3686dffb02459559c74dd3d81748269ffb0eb027c39a6fc99502de37d501faa8 \ + --hash=sha256:582215a0e9adbe0e379761260553ba11c58943e4bbe9c36430c4ca6ac74b15ed \ + --hash=sha256:5b50bf3f55561dac5438f8e70bfcdfd74543fd60df5fa5f62d94e5867deca684 \ + --hash=sha256:5bf44d66cdf9e893637896c7faa22298baebcd18d1ddb6d2626a6e39793a1d56 \ + --hash=sha256:6602bc8dc6f3a9e02b6c22c4fc1e47aa50f8f8e6d3f78a5e16ac33ef5fefa324 \ + --hash=sha256:673739cb539f8cdaa07d92d02efa93c9ccf87e345b9a0b556e3ecc666718468d \ + --hash=sha256:68678abf380b42ce21a5f2abde8efee05c114c2fdb2e9eef2efdb0257fba1235 \ + --hash=sha256:68e7c44931cc171c54ccb702482e9fc723192e88d25a0e133edd7aff8fcd1f6e \ + --hash=sha256:6b3d6606d369fc1da4fd8c357d026317fbb9c9b75d36dc16e90e84c26854b088 \ + --hash=sha256:748dcd1e3d3d7cd5443ef03ce8685043294ad6bd7c02a38d1bd367cfd968e000 \ + --hash=sha256:7651c50c8c5ef7bdb41108b7b8c5a83013bfaa8a935590c5d74627c047a583c7 \ + --hash=sha256:7b78010e7b97fef4bee1e896df8a4bbb6712b7f05b7ef630f9d1da00f6444d2e \ + --hash=sha256:7e61e3e4fa664a8588aa25c883eab612a188c725755afff6289454d6362b9673 \ + --hash=sha256:80876338e19c951fdfed6198e70bc88f1c9758b94578d5a7c4c91a87af3cf31c \ + --hash=sha256:8895613bcc094d4a1b2dbe179d88d7fb4a15cee43c052e8885783fac397d91fe \ + --hash=sha256:88e2b3c14bdb32e440be531ade29d3c50a1a59cd4e51b1dd8b0865c54ea5d2e2 \ + --hash=sha256:8f8e709127c6c77446a8c0a8c8bf3c8ee706a06cd44b1e827c3e6a2ee6b8c098 \ + --hash=sha256:9cb4a35b3642fc5c005a6755a5d17c6c8b6bcb6981baf81cea8bfbc8903e8ba8 \ + --hash=sha256:9f90389693731ff1f659e55c7d1640e2ec43ff725cc61b04b2f9c6d8d017df6a \ + --hash=sha256:a09582f178759ee8128d9270cd1344154fd473bb77d94ce0aeb2a93ebf0feaf0 \ + --hash=sha256:a6a14b17d7e17fa0d207ac08642c8820f84f25ce17a442fd15e27ea18d67c59b \ + --hash=sha256:a72e8961a86d19bdb45851d8f1f08b041ea37d2bd8d4fd19903bc3083d80c896 \ + --hash=sha256:abd808f9c129ba2beda4cfc53bde801e5bcf9d6e0f22f095e45327c038bfe68e \ + --hash=sha256:ac0f5edd2360eea2f1daa9e26a41db02dd4b0451b48f7c318e217ee092a213e9 \ + --hash=sha256:b29ebffcf550f9da55bec9e02ad430c992a87e5f512cd63388abb76f1036d8d2 \ + --hash=sha256:b2ca4e77f9f47c55c194982e10f058db063937845bb2b7a86c84a6cfe0aefa8b \ + --hash=sha256:b7be2d771cdba2942e13215c4e340bfd76398e9227ad10402a8767ab1865d2e6 \ + --hash=sha256:b84834d0cf97e7d27dd5b7f3aca7b6e9263c56308ab9dc8aae9784abb774d404 \ + --hash=sha256:b86851a328eedc692acf81fb05444bdf1891747c25af7529e39ddafaf68a4f3f \ + --hash=sha256:bcb3ef43e58665bbda2fb198698fcae6776483e0c4a631aa5647806c25e02cc0 \ + --hash=sha256:c0f31130ebc2d37cdd8e44605fb5fa7ad59049298b3f745c74fa74c62fbfcfc4 \ + --hash=sha256:c6a164aa47843fb1b01e941d385aab7215563bb8816d80ff3a363a9f8448a8dc \ + --hash=sha256:d8a9d3ebe49f084ad71f9269834ceccbf398253c9fac910c4fd7053ff1386936 \ + --hash=sha256:db8e577c19c0fda0beb7e0d4e09e0ba74b1e4c092e0e40bfa12fe05b6f6d75ba \ + --hash=sha256:dc9b18bf40cc75f66f40a7379f6a9513244fe33c0e8aa72e2d56b0196a7ef872 \ + --hash=sha256:e09f3ff613345df5e8c3667da1d918f9149bd623cd9070c983c013792a9a62eb \ + --hash=sha256:e4108df7fe9b707191e55f33efbcb2d81928e10cea45527879a4749cbe472614 \ + --hash=sha256:e6024675e67af929088fda399b2094574609396b1decb609c55fa58b028a32a1 \ + --hash=sha256:e70f54f1796669ef691ca07d046cd81a29cb4deb1e5f942003f401c0c4a2695d \ + --hash=sha256:e715596e683d2ce000574bae5d07bd522c781a822866c20495e52520564f0969 \ + --hash=sha256:e760191dd42581e023a68b758769e2da259b5d52e3103c6060ddc02c9edb8d7b \ + --hash=sha256:ed86a35631f7bfbb28e108dd96773b9d5a6ce4811cf6ea468bb6a359b256b1e4 \ + --hash=sha256:ee07e47c12890ef248766a6e55bd38ebfb2bb8edd4142d56db91b21ea68b7627 \ + --hash=sha256:fa3a0128b152627161ce47201262d3140edb5a5c3da88d73a1b790a959126956 \ + --hash=sha256:fcc8eb6d5902bb1cf6dc4f187ee3ea80a1eba0a89aba40a5cb20a5087d961357 # via cryptography charset-normalizer==2.1.1 \ --hash=sha256:5a3d016c7c547f69d6f81fb0db9449ce888b418b5b9952cc5e6e66843e9dd845 \ @@ -109,74 +93,74 @@ colorlog==6.7.0 \ # via # gcp-docuploader # nox -commonmark==0.9.1 \ - --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \ - --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9 - # via rich -cryptography==41.0.0 \ - --hash=sha256:0ddaee209d1cf1f180f1efa338a68c4621154de0afaef92b89486f5f96047c55 \ - --hash=sha256:14754bcdae909d66ff24b7b5f166d69340ccc6cb15731670435efd5719294895 \ - --hash=sha256:344c6de9f8bda3c425b3a41b319522ba3208551b70c2ae00099c205f0d9fd3be \ - --hash=sha256:34d405ea69a8b34566ba3dfb0521379b210ea5d560fafedf9f800a9a94a41928 \ - --hash=sha256:3680248309d340fda9611498a5319b0193a8dbdb73586a1acf8109d06f25b92d \ - --hash=sha256:3c5ef25d060c80d6d9f7f9892e1d41bb1c79b78ce74805b8cb4aa373cb7d5ec8 \ - --hash=sha256:4ab14d567f7bbe7f1cdff1c53d5324ed4d3fc8bd17c481b395db224fb405c237 \ - --hash=sha256:5c1f7293c31ebc72163a9a0df246f890d65f66b4a40d9ec80081969ba8c78cc9 \ - --hash=sha256:6b71f64beeea341c9b4f963b48ee3b62d62d57ba93eb120e1196b31dc1025e78 \ - --hash=sha256:7d92f0248d38faa411d17f4107fc0bce0c42cae0b0ba5415505df72d751bf62d \ - --hash=sha256:8362565b3835ceacf4dc8f3b56471a2289cf51ac80946f9087e66dc283a810e0 \ - --hash=sha256:84a165379cb9d411d58ed739e4af3396e544eac190805a54ba2e0322feb55c46 \ - --hash=sha256:88ff107f211ea696455ea8d911389f6d2b276aabf3231bf72c8853d22db755c5 \ - --hash=sha256:9f65e842cb02550fac96536edb1d17f24c0a338fd84eaf582be25926e993dde4 \ - --hash=sha256:a4fc68d1c5b951cfb72dfd54702afdbbf0fb7acdc9b7dc4301bbf2225a27714d \ - --hash=sha256:b7f2f5c525a642cecad24ee8670443ba27ac1fab81bba4cc24c7b6b41f2d0c75 \ - --hash=sha256:b846d59a8d5a9ba87e2c3d757ca019fa576793e8758174d3868aecb88d6fc8eb \ - --hash=sha256:bf8fc66012ca857d62f6a347007e166ed59c0bc150cefa49f28376ebe7d992a2 \ - --hash=sha256:f5d0bf9b252f30a31664b6f64432b4730bb7038339bd18b1fafe129cfc2be9be +cryptography==41.0.6 \ + --hash=sha256:068bc551698c234742c40049e46840843f3d98ad7ce265fd2bd4ec0d11306596 \ + --hash=sha256:0f27acb55a4e77b9be8d550d762b0513ef3fc658cd3eb15110ebbcbd626db12c \ + --hash=sha256:2132d5865eea673fe6712c2ed5fb4fa49dba10768bb4cc798345748380ee3660 \ + --hash=sha256:3288acccef021e3c3c10d58933f44e8602cf04dba96d9796d70d537bb2f4bbc4 \ + --hash=sha256:35f3f288e83c3f6f10752467c48919a7a94b7d88cc00b0668372a0d2ad4f8ead \ + --hash=sha256:398ae1fc711b5eb78e977daa3cbf47cec20f2c08c5da129b7a296055fbb22aed \ + --hash=sha256:422e3e31d63743855e43e5a6fcc8b4acab860f560f9321b0ee6269cc7ed70cc3 \ + --hash=sha256:48783b7e2bef51224020efb61b42704207dde583d7e371ef8fc2a5fb6c0aabc7 \ + --hash=sha256:4d03186af98b1c01a4eda396b137f29e4e3fb0173e30f885e27acec8823c1b09 \ + --hash=sha256:5daeb18e7886a358064a68dbcaf441c036cbdb7da52ae744e7b9207b04d3908c \ + --hash=sha256:60e746b11b937911dc70d164060d28d273e31853bb359e2b2033c9e93e6f3c43 \ + --hash=sha256:742ae5e9a2310e9dade7932f9576606836ed174da3c7d26bc3d3ab4bd49b9f65 \ + --hash=sha256:7e00fb556bda398b99b0da289ce7053639d33b572847181d6483ad89835115f6 \ + --hash=sha256:85abd057699b98fce40b41737afb234fef05c67e116f6f3650782c10862c43da \ + --hash=sha256:8efb2af8d4ba9dbc9c9dd8f04d19a7abb5b49eab1f3694e7b5a16a5fc2856f5c \ + --hash=sha256:ae236bb8760c1e55b7a39b6d4d32d2279bc6c7c8500b7d5a13b6fb9fc97be35b \ + --hash=sha256:afda76d84b053923c27ede5edc1ed7d53e3c9f475ebaf63c68e69f1403c405a8 \ + --hash=sha256:b27a7fd4229abef715e064269d98a7e2909ebf92eb6912a9603c7e14c181928c \ + --hash=sha256:b648fe2a45e426aaee684ddca2632f62ec4613ef362f4d681a9a6283d10e079d \ + --hash=sha256:c5a550dc7a3b50b116323e3d376241829fd326ac47bc195e04eb33a8170902a9 \ + --hash=sha256:da46e2b5df770070412c46f87bac0849b8d685c5f2679771de277a422c7d0b86 \ + --hash=sha256:f39812f70fc5c71a15aa3c97b2bbe213c3f2a460b79bd21c40d033bb34a9bf36 \ + --hash=sha256:ff369dd19e8fe0528b02e8df9f2aeb2479f89b1270d90f96a63500afe9af5cae # via # gcp-releasetool # secretstorage -distlib==0.3.6 \ - --hash=sha256:14bad2d9b04d3a36127ac97f30b12a19268f211063d8f8ee4f47108896e11b46 \ - --hash=sha256:f35c4b692542ca110de7ef0bea44d73981caeb34ca0b9b6b2e6d7790dda8f80e +distlib==0.3.7 \ + --hash=sha256:2e24928bc811348f0feb63014e97aaae3037f2cf48712d51ae61df7fd6075057 \ + --hash=sha256:9dafe54b34a028eafd95039d5e5d4851a13734540f1331060d31c9916e7147a8 # via virtualenv -docutils==0.19 \ - --hash=sha256:33995a6753c30b7f577febfc2c50411fec6aac7f7ffeb7c4cfe5991072dcf9e6 \ - --hash=sha256:5e1de4d849fee02c63b040a4a3fd567f4ab104defd8a5511fbbc24a8a017efbc +docutils==0.20.1 \ + --hash=sha256:96f387a2c5562db4476f09f13bbab2192e764cac08ebbf3a34a95d9b1e4a59d6 \ + --hash=sha256:f08a4e276c3a1583a86dce3e34aba3fe04d02bba2dd51ed16106244e8a923e3b # via readme-renderer -filelock==3.8.0 \ - --hash=sha256:55447caa666f2198c5b6b13a26d2084d26fa5b115c00d065664b2124680c4edc \ - --hash=sha256:617eb4e5eedc82fc5f47b6d61e4d11cb837c56cb4544e39081099fa17ad109d4 +filelock==3.13.1 \ + --hash=sha256:521f5f56c50f8426f5e03ad3b281b490a87ef15bc6c526f168290f0c7148d44e \ + --hash=sha256:57dbda9b35157b05fb3e58ee91448612eb674172fab98ee235ccb0b5bee19a1c # via virtualenv -gcp-docuploader==0.6.4 \ - --hash=sha256:01486419e24633af78fd0167db74a2763974765ee8078ca6eb6964d0ebd388af \ - --hash=sha256:70861190c123d907b3b067da896265ead2eeb9263969d6955c9e0bb091b5ccbf +gcp-docuploader==0.6.5 \ + --hash=sha256:30221d4ac3e5a2b9c69aa52fdbef68cc3f27d0e6d0d90e220fc024584b8d2318 \ + --hash=sha256:b7458ef93f605b9d46a4bf3a8dc1755dad1f31d030c8679edf304e343b347eea # via -r requirements.in -gcp-releasetool==1.10.5 \ - --hash=sha256:174b7b102d704b254f2a26a3eda2c684fd3543320ec239baf771542a2e58e109 \ - --hash=sha256:e29d29927fe2ca493105a82958c6873bb2b90d503acac56be2c229e74de0eec9 +gcp-releasetool==1.16.0 \ + --hash=sha256:27bf19d2e87aaa884096ff941aa3c592c482be3d6a2bfe6f06afafa6af2353e3 \ + --hash=sha256:a316b197a543fd036209d0caba7a8eb4d236d8e65381c80cbc6d7efaa7606d63 # via -r requirements.in -google-api-core==2.10.2 \ - --hash=sha256:10c06f7739fe57781f87523375e8e1a3a4674bf6392cd6131a3222182b971320 \ - --hash=sha256:34f24bd1d5f72a8c4519773d99ca6bf080a6c4e041b4e9f024fe230191dda62e +google-api-core==2.12.0 \ + --hash=sha256:c22e01b1e3c4dcd90998494879612c38d0a3411d1f7b679eb89e2abe3ce1f553 \ + --hash=sha256:ec6054f7d64ad13b41e43d96f735acbd763b0f3b695dabaa2d579673f6a6e160 # via # google-cloud-core # google-cloud-storage -google-auth==2.14.1 \ - --hash=sha256:ccaa901f31ad5cbb562615eb8b664b3dd0bf5404a67618e642307f00613eda4d \ - --hash=sha256:f5d8701633bebc12e0deea4df8abd8aff31c28b355360597f7f2ee60f2e4d016 +google-auth==2.23.4 \ + --hash=sha256:79905d6b1652187def79d491d6e23d0cbb3a21d3c7ba0dbaa9c8a01906b13ff3 \ + --hash=sha256:d4bbc92fe4b8bfd2f3e8d88e5ba7085935da208ee38a134fc280e7ce682a05f2 # via # gcp-releasetool # google-api-core # google-cloud-core # google-cloud-storage -google-cloud-core==2.3.2 \ - --hash=sha256:8417acf6466be2fa85123441696c4badda48db314c607cf1e5d543fa8bdc22fe \ - --hash=sha256:b9529ee7047fd8d4bf4a2182de619154240df17fbe60ead399078c1ae152af9a +google-cloud-core==2.3.3 \ + --hash=sha256:37b80273c8d7eee1ae816b3a20ae43585ea50506cb0e60f3cf5be5f87f1373cb \ + --hash=sha256:fbd11cad3e98a7e5b0343dc07cb1039a5ffd7a5bb96e1f1e27cee4bda4a90863 # via google-cloud-storage -google-cloud-storage==2.6.0 \ - --hash=sha256:104ca28ae61243b637f2f01455cc8a05e8f15a2a18ced96cb587241cdd3820f5 \ - --hash=sha256:4ad0415ff61abdd8bb2ae81c1f8f7ec7d91a1011613f2db87c614c550f97bfe9 +google-cloud-storage==2.13.0 \ + --hash=sha256:ab0bf2e1780a1b74cf17fccb13788070b729f50c252f0c94ada2aae0ca95437d \ + --hash=sha256:f62dc4c7b6cd4360d072e3deb28035fbdad491ac3d9b0b1815a12daea10f37c7 # via gcp-docuploader google-crc32c==1.5.0 \ --hash=sha256:024894d9d3cfbc5943f8f230e23950cd4906b2fe004c72e29b209420a1e6b05a \ @@ -247,29 +231,31 @@ google-crc32c==1.5.0 \ --hash=sha256:f583edb943cf2e09c60441b910d6a20b4d9d626c75a36c8fcac01a6c96c01183 \ --hash=sha256:fd8536e902db7e365f49e7d9029283403974ccf29b13fc7028b97e2295b33556 \ --hash=sha256:fe70e325aa68fa4b5edf7d1a4b6f691eb04bbccac0ace68e34820d283b5f80d4 - # via google-resumable-media -google-resumable-media==2.4.0 \ - --hash=sha256:2aa004c16d295c8f6c33b2b4788ba59d366677c0a25ae7382436cb30f776deaa \ - --hash=sha256:8d5518502f92b9ecc84ac46779bd4f09694ecb3ba38a3e7ca737a86d15cbca1f + # via + # google-cloud-storage + # google-resumable-media +google-resumable-media==2.6.0 \ + --hash=sha256:972852f6c65f933e15a4a210c2b96930763b47197cdf4aa5f5bea435efb626e7 \ + --hash=sha256:fc03d344381970f79eebb632a3c18bb1828593a2dc5572b5f90115ef7d11e81b # via google-cloud-storage -googleapis-common-protos==1.57.0 \ - --hash=sha256:27a849d6205838fb6cc3c1c21cb9800707a661bb21c6ce7fb13e99eb1f8a0c46 \ - --hash=sha256:a9f4a1d7f6d9809657b7f1316a1aa527f6664891531bcfcc13b6696e685f443c +googleapis-common-protos==1.61.0 \ + --hash=sha256:22f1915393bb3245343f6efe87f6fe868532efc12aa26b391b15132e1279f1c0 \ + --hash=sha256:8a64866a97f6304a7179873a465d6eee97b7a24ec6cfd78e0f575e96b821240b # via google-api-core idna==3.4 \ --hash=sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4 \ --hash=sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2 # via requests -importlib-metadata==5.0.0 \ - --hash=sha256:da31db32b304314d044d3c12c79bd59e307889b287ad12ff387b3500835fc2ab \ - --hash=sha256:ddb0e35065e8938f867ed4928d0ae5bf2a53b7773871bfe6bcc7e4fcdc7dea43 +importlib-metadata==6.8.0 \ + --hash=sha256:3ebb78df84a805d7698245025b975d9d67053cd94c79245ba4b3eb694abe68bb \ + --hash=sha256:dbace7892d8c0c4ac1ad096662232f831d4e64f4c4545bd53016a3e9d4654743 # via # -r requirements.in # keyring # twine -jaraco-classes==3.2.3 \ - --hash=sha256:2353de3288bc6b82120752201c6b1c1a14b058267fa424ed5ce5984e3b922158 \ - --hash=sha256:89559fa5c1d3c34eff6f631ad80bb21f378dbcbb35dd161fd2c6b93f5be2f98a +jaraco-classes==3.3.0 \ + --hash=sha256:10afa92b6743f25c0cf5f37c6bb6e18e2c5bb84a16527ccfc0040ea377e7aaeb \ + --hash=sha256:c063dd08e89217cee02c8d5e5ec560f2c8ce6cdc2fcdc2e68f7b2e5547ed3621 # via keyring jeepney==0.8.0 \ --hash=sha256:5efe48d255973902f6badc3ce55e2aa6c5c3b3bc642059ef3a91247bcfcc5806 \ @@ -281,75 +267,121 @@ jinja2==3.1.2 \ --hash=sha256:31351a702a408a9e7595a8fc6150fc3f43bb6bf7e319770cbc0db9df9437e852 \ --hash=sha256:6088930bfe239f0e6710546ab9c19c9ef35e29792895fed6e6e31a023a182a61 # via gcp-releasetool -keyring==23.11.0 \ - --hash=sha256:3dd30011d555f1345dec2c262f0153f2f0ca6bca041fb1dc4588349bb4c0ac1e \ - --hash=sha256:ad192263e2cdd5f12875dedc2da13534359a7e760e77f8d04b50968a821c2361 +keyring==24.2.0 \ + --hash=sha256:4901caaf597bfd3bbd78c9a0c7c4c29fcd8310dab2cffefe749e916b6527acd6 \ + --hash=sha256:ca0746a19ec421219f4d713f848fa297a661a8a8c1504867e55bfb5e09091509 # via # gcp-releasetool # twine -markupsafe==2.1.1 \ - --hash=sha256:0212a68688482dc52b2d45013df70d169f542b7394fc744c02a57374a4207003 \ - --hash=sha256:089cf3dbf0cd6c100f02945abeb18484bd1ee57a079aefd52cffd17fba910b88 \ - --hash=sha256:10c1bfff05d95783da83491be968e8fe789263689c02724e0c691933c52994f5 \ - --hash=sha256:33b74d289bd2f5e527beadcaa3f401e0df0a89927c1559c8566c066fa4248ab7 \ - --hash=sha256:3799351e2336dc91ea70b034983ee71cf2f9533cdff7c14c90ea126bfd95d65a \ - --hash=sha256:3ce11ee3f23f79dbd06fb3d63e2f6af7b12db1d46932fe7bd8afa259a5996603 \ - --hash=sha256:421be9fbf0ffe9ffd7a378aafebbf6f4602d564d34be190fc19a193232fd12b1 \ - --hash=sha256:43093fb83d8343aac0b1baa75516da6092f58f41200907ef92448ecab8825135 \ - --hash=sha256:46d00d6cfecdde84d40e572d63735ef81423ad31184100411e6e3388d405e247 \ - --hash=sha256:4a33dea2b688b3190ee12bd7cfa29d39c9ed176bda40bfa11099a3ce5d3a7ac6 \ - --hash=sha256:4b9fe39a2ccc108a4accc2676e77da025ce383c108593d65cc909add5c3bd601 \ - --hash=sha256:56442863ed2b06d19c37f94d999035e15ee982988920e12a5b4ba29b62ad1f77 \ - --hash=sha256:671cd1187ed5e62818414afe79ed29da836dde67166a9fac6d435873c44fdd02 \ - --hash=sha256:694deca8d702d5db21ec83983ce0bb4b26a578e71fbdbd4fdcd387daa90e4d5e \ - --hash=sha256:6a074d34ee7a5ce3effbc526b7083ec9731bb3cbf921bbe1d3005d4d2bdb3a63 \ - --hash=sha256:6d0072fea50feec76a4c418096652f2c3238eaa014b2f94aeb1d56a66b41403f \ - --hash=sha256:6fbf47b5d3728c6aea2abb0589b5d30459e369baa772e0f37a0320185e87c980 \ - --hash=sha256:7f91197cc9e48f989d12e4e6fbc46495c446636dfc81b9ccf50bb0ec74b91d4b \ - --hash=sha256:86b1f75c4e7c2ac2ccdaec2b9022845dbb81880ca318bb7a0a01fbf7813e3812 \ - --hash=sha256:8dc1c72a69aa7e082593c4a203dcf94ddb74bb5c8a731e4e1eb68d031e8498ff \ - --hash=sha256:8e3dcf21f367459434c18e71b2a9532d96547aef8a871872a5bd69a715c15f96 \ - --hash=sha256:8e576a51ad59e4bfaac456023a78f6b5e6e7651dcd383bcc3e18d06f9b55d6d1 \ - --hash=sha256:96e37a3dc86e80bf81758c152fe66dbf60ed5eca3d26305edf01892257049925 \ - --hash=sha256:97a68e6ada378df82bc9f16b800ab77cbf4b2fada0081794318520138c088e4a \ - --hash=sha256:99a2a507ed3ac881b975a2976d59f38c19386d128e7a9a18b7df6fff1fd4c1d6 \ - --hash=sha256:a49907dd8420c5685cfa064a1335b6754b74541bbb3706c259c02ed65b644b3e \ - --hash=sha256:b09bf97215625a311f669476f44b8b318b075847b49316d3e28c08e41a7a573f \ - --hash=sha256:b7bd98b796e2b6553da7225aeb61f447f80a1ca64f41d83612e6139ca5213aa4 \ - --hash=sha256:b87db4360013327109564f0e591bd2a3b318547bcef31b468a92ee504d07ae4f \ - --hash=sha256:bcb3ed405ed3222f9904899563d6fc492ff75cce56cba05e32eff40e6acbeaa3 \ - --hash=sha256:d4306c36ca495956b6d568d276ac11fdd9c30a36f1b6eb928070dc5360b22e1c \ - --hash=sha256:d5ee4f386140395a2c818d149221149c54849dfcfcb9f1debfe07a8b8bd63f9a \ - --hash=sha256:dda30ba7e87fbbb7eab1ec9f58678558fd9a6b8b853530e176eabd064da81417 \ - --hash=sha256:e04e26803c9c3851c931eac40c695602c6295b8d432cbe78609649ad9bd2da8a \ - --hash=sha256:e1c0b87e09fa55a220f058d1d49d3fb8df88fbfab58558f1198e08c1e1de842a \ - --hash=sha256:e72591e9ecd94d7feb70c1cbd7be7b3ebea3f548870aa91e2732960fa4d57a37 \ - --hash=sha256:e8c843bbcda3a2f1e3c2ab25913c80a3c5376cd00c6e8c4a86a89a28c8dc5452 \ - --hash=sha256:efc1913fd2ca4f334418481c7e595c00aad186563bbc1ec76067848c7ca0a933 \ - --hash=sha256:f121a1420d4e173a5d96e47e9a0c0dcff965afdf1626d28de1460815f7c4ee7a \ - --hash=sha256:fc7b548b17d238737688817ab67deebb30e8073c95749d55538ed473130ec0c7 +markdown-it-py==3.0.0 \ + --hash=sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1 \ + --hash=sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb + # via rich +markupsafe==2.1.3 \ + --hash=sha256:05fb21170423db021895e1ea1e1f3ab3adb85d1c2333cbc2310f2a26bc77272e \ + --hash=sha256:0a4e4a1aff6c7ac4cd55792abf96c915634c2b97e3cc1c7129578aa68ebd754e \ + --hash=sha256:10bbfe99883db80bdbaff2dcf681dfc6533a614f700da1287707e8a5d78a8431 \ + --hash=sha256:134da1eca9ec0ae528110ccc9e48041e0828d79f24121a1a146161103c76e686 \ + --hash=sha256:14ff806850827afd6b07a5f32bd917fb7f45b046ba40c57abdb636674a8b559c \ + --hash=sha256:1577735524cdad32f9f694208aa75e422adba74f1baee7551620e43a3141f559 \ + --hash=sha256:1b40069d487e7edb2676d3fbdb2b0829ffa2cd63a2ec26c4938b2d34391b4ecc \ + --hash=sha256:1b8dd8c3fd14349433c79fa8abeb573a55fc0fdd769133baac1f5e07abf54aeb \ + --hash=sha256:1f67c7038d560d92149c060157d623c542173016c4babc0c1913cca0564b9939 \ + --hash=sha256:282c2cb35b5b673bbcadb33a585408104df04f14b2d9b01d4c345a3b92861c2c \ + --hash=sha256:2c1b19b3aaacc6e57b7e25710ff571c24d6c3613a45e905b1fde04d691b98ee0 \ + --hash=sha256:2ef12179d3a291be237280175b542c07a36e7f60718296278d8593d21ca937d4 \ + --hash=sha256:338ae27d6b8745585f87218a3f23f1512dbf52c26c28e322dbe54bcede54ccb9 \ + --hash=sha256:3c0fae6c3be832a0a0473ac912810b2877c8cb9d76ca48de1ed31e1c68386575 \ + --hash=sha256:3fd4abcb888d15a94f32b75d8fd18ee162ca0c064f35b11134be77050296d6ba \ + --hash=sha256:42de32b22b6b804f42c5d98be4f7e5e977ecdd9ee9b660fda1a3edf03b11792d \ + --hash=sha256:47d4f1c5f80fc62fdd7777d0d40a2e9dda0a05883ab11374334f6c4de38adffd \ + --hash=sha256:504b320cd4b7eff6f968eddf81127112db685e81f7e36e75f9f84f0df46041c3 \ + --hash=sha256:525808b8019e36eb524b8c68acdd63a37e75714eac50e988180b169d64480a00 \ + --hash=sha256:56d9f2ecac662ca1611d183feb03a3fa4406469dafe241673d521dd5ae92a155 \ + --hash=sha256:5bbe06f8eeafd38e5d0a4894ffec89378b6c6a625ff57e3028921f8ff59318ac \ + --hash=sha256:65c1a9bcdadc6c28eecee2c119465aebff8f7a584dd719facdd9e825ec61ab52 \ + --hash=sha256:68e78619a61ecf91e76aa3e6e8e33fc4894a2bebe93410754bd28fce0a8a4f9f \ + --hash=sha256:69c0f17e9f5a7afdf2cc9fb2d1ce6aabdb3bafb7f38017c0b77862bcec2bbad8 \ + --hash=sha256:6b2b56950d93e41f33b4223ead100ea0fe11f8e6ee5f641eb753ce4b77a7042b \ + --hash=sha256:715d3562f79d540f251b99ebd6d8baa547118974341db04f5ad06d5ea3eb8007 \ + --hash=sha256:787003c0ddb00500e49a10f2844fac87aa6ce977b90b0feaaf9de23c22508b24 \ + --hash=sha256:7ef3cb2ebbf91e330e3bb937efada0edd9003683db6b57bb108c4001f37a02ea \ + --hash=sha256:8023faf4e01efadfa183e863fefde0046de576c6f14659e8782065bcece22198 \ + --hash=sha256:8758846a7e80910096950b67071243da3e5a20ed2546e6392603c096778d48e0 \ + --hash=sha256:8afafd99945ead6e075b973fefa56379c5b5c53fd8937dad92c662da5d8fd5ee \ + --hash=sha256:8c41976a29d078bb235fea9b2ecd3da465df42a562910f9022f1a03107bd02be \ + --hash=sha256:8e254ae696c88d98da6555f5ace2279cf7cd5b3f52be2b5cf97feafe883b58d2 \ + --hash=sha256:8f9293864fe09b8149f0cc42ce56e3f0e54de883a9de90cd427f191c346eb2e1 \ + --hash=sha256:9402b03f1a1b4dc4c19845e5c749e3ab82d5078d16a2a4c2cd2df62d57bb0707 \ + --hash=sha256:962f82a3086483f5e5f64dbad880d31038b698494799b097bc59c2edf392fce6 \ + --hash=sha256:9aad3c1755095ce347e26488214ef77e0485a3c34a50c5a5e2471dff60b9dd9c \ + --hash=sha256:9dcdfd0eaf283af041973bff14a2e143b8bd64e069f4c383416ecd79a81aab58 \ + --hash=sha256:aa57bd9cf8ae831a362185ee444e15a93ecb2e344c8e52e4d721ea3ab6ef1823 \ + --hash=sha256:aa7bd130efab1c280bed0f45501b7c8795f9fdbeb02e965371bbef3523627779 \ + --hash=sha256:ab4a0df41e7c16a1392727727e7998a467472d0ad65f3ad5e6e765015df08636 \ + --hash=sha256:ad9e82fb8f09ade1c3e1b996a6337afac2b8b9e365f926f5a61aacc71adc5b3c \ + --hash=sha256:af598ed32d6ae86f1b747b82783958b1a4ab8f617b06fe68795c7f026abbdcad \ + --hash=sha256:b076b6226fb84157e3f7c971a47ff3a679d837cf338547532ab866c57930dbee \ + --hash=sha256:b7ff0f54cb4ff66dd38bebd335a38e2c22c41a8ee45aa608efc890ac3e3931bc \ + --hash=sha256:bfce63a9e7834b12b87c64d6b155fdd9b3b96191b6bd334bf37db7ff1fe457f2 \ + --hash=sha256:c011a4149cfbcf9f03994ec2edffcb8b1dc2d2aede7ca243746df97a5d41ce48 \ + --hash=sha256:c9c804664ebe8f83a211cace637506669e7890fec1b4195b505c214e50dd4eb7 \ + --hash=sha256:ca379055a47383d02a5400cb0d110cef0a776fc644cda797db0c5696cfd7e18e \ + --hash=sha256:cb0932dc158471523c9637e807d9bfb93e06a95cbf010f1a38b98623b929ef2b \ + --hash=sha256:cd0f502fe016460680cd20aaa5a76d241d6f35a1c3350c474bac1273803893fa \ + --hash=sha256:ceb01949af7121f9fc39f7d27f91be8546f3fb112c608bc4029aef0bab86a2a5 \ + --hash=sha256:d080e0a5eb2529460b30190fcfcc4199bd7f827663f858a226a81bc27beaa97e \ + --hash=sha256:dd15ff04ffd7e05ffcb7fe79f1b98041b8ea30ae9234aed2a9168b5797c3effb \ + --hash=sha256:df0be2b576a7abbf737b1575f048c23fb1d769f267ec4358296f31c2479db8f9 \ + --hash=sha256:e09031c87a1e51556fdcb46e5bd4f59dfb743061cf93c4d6831bf894f125eb57 \ + --hash=sha256:e4dd52d80b8c83fdce44e12478ad2e85c64ea965e75d66dbeafb0a3e77308fcc \ + --hash=sha256:f698de3fd0c4e6972b92290a45bd9b1536bffe8c6759c62471efaa8acb4c37bc \ + --hash=sha256:fec21693218efe39aa7f8599346e90c705afa52c5b31ae019b2e57e8f6542bb2 \ + --hash=sha256:ffcc3f7c66b5f5b7931a5aa68fc9cecc51e685ef90282f4a82f0f5e9b704ad11 # via jinja2 -more-itertools==9.0.0 \ - --hash=sha256:250e83d7e81d0c87ca6bd942e6aeab8cc9daa6096d12c5308f3f92fa5e5c1f41 \ - --hash=sha256:5a6257e40878ef0520b1803990e3e22303a41b5714006c32a3fd8304b26ea1ab +mdurl==0.1.2 \ + --hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \ + --hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba + # via markdown-it-py +more-itertools==10.1.0 \ + --hash=sha256:626c369fa0eb37bac0291bce8259b332fd59ac792fa5497b59837309cd5b114a \ + --hash=sha256:64e0735fcfdc6f3464ea133afe8ea4483b1c5fe3a3d69852e6503b43a0b222e6 # via jaraco-classes -nox==2022.11.21 \ - --hash=sha256:0e41a990e290e274cb205a976c4c97ee3c5234441a8132c8c3fd9ea3c22149eb \ - --hash=sha256:e21c31de0711d1274ca585a2c5fde36b1aa962005ba8e9322bf5eeed16dcd684 +nh3==0.2.14 \ + --hash=sha256:116c9515937f94f0057ef50ebcbcc10600860065953ba56f14473ff706371873 \ + --hash=sha256:18415df36db9b001f71a42a3a5395db79cf23d556996090d293764436e98e8ad \ + --hash=sha256:203cac86e313cf6486704d0ec620a992c8bc164c86d3a4fd3d761dd552d839b5 \ + --hash=sha256:2b0be5c792bd43d0abef8ca39dd8acb3c0611052ce466d0401d51ea0d9aa7525 \ + --hash=sha256:377aaf6a9e7c63962f367158d808c6a1344e2b4f83d071c43fbd631b75c4f0b2 \ + --hash=sha256:525846c56c2bcd376f5eaee76063ebf33cf1e620c1498b2a40107f60cfc6054e \ + --hash=sha256:5529a3bf99402c34056576d80ae5547123f1078da76aa99e8ed79e44fa67282d \ + --hash=sha256:7771d43222b639a4cd9e341f870cee336b9d886de1ad9bec8dddab22fe1de450 \ + --hash=sha256:88c753efbcdfc2644a5012938c6b9753f1c64a5723a67f0301ca43e7b85dcf0e \ + --hash=sha256:93a943cfd3e33bd03f77b97baa11990148687877b74193bf777956b67054dcc6 \ + --hash=sha256:9be2f68fb9a40d8440cbf34cbf40758aa7f6093160bfc7fb018cce8e424f0c3a \ + --hash=sha256:a0c509894fd4dccdff557068e5074999ae3b75f4c5a2d6fb5415e782e25679c4 \ + --hash=sha256:ac8056e937f264995a82bf0053ca898a1cb1c9efc7cd68fa07fe0060734df7e4 \ + --hash=sha256:aed56a86daa43966dd790ba86d4b810b219f75b4bb737461b6886ce2bde38fd6 \ + --hash=sha256:e8986f1dd3221d1e741fda0a12eaa4a273f1d80a35e31a1ffe579e7c621d069e \ + --hash=sha256:f99212a81c62b5f22f9e7c3e347aa00491114a5647e1f13bbebd79c3e5f08d75 + # via readme-renderer +nox==2023.4.22 \ + --hash=sha256:0b1adc619c58ab4fa57d6ab2e7823fe47a32e70202f287d78474adcc7bda1891 \ + --hash=sha256:46c0560b0dc609d7d967dc99e22cb463d3c4caf54a5fda735d6c11b5177e3a9f # via -r requirements.in -packaging==21.3 \ - --hash=sha256:dd47c42927d89ab911e606518907cc2d3a1f38bbd026385970643f9c5b8ecfeb \ - --hash=sha256:ef103e05f519cdc783ae24ea4e2e0f508a9c99b2d4969652eed6a2e1ea5bd522 +packaging==23.2 \ + --hash=sha256:048fb0e9405036518eaaf48a55953c750c11e1a1b68e0dd1a9d62ed0c092cfc5 \ + --hash=sha256:8c491190033a9af7e1d931d0b5dacc2ef47509b34dd0de67ed209b5203fc88c7 # via # gcp-releasetool # nox -pkginfo==1.8.3 \ - --hash=sha256:848865108ec99d4901b2f7e84058b6e7660aae8ae10164e015a6dcf5b242a594 \ - --hash=sha256:a84da4318dd86f870a9447a8c98340aa06216bfc6f2b7bdc4b8766984ae1867c +pkginfo==1.9.6 \ + --hash=sha256:4b7a555a6d5a22169fcc9cf7bfd78d296b0361adad412a346c1226849af5e546 \ + --hash=sha256:8fd5896e8718a4372f0ea9cc9d96f6417c9b986e23a4d116dda26b62cc29d046 # via twine -platformdirs==2.5.4 \ - --hash=sha256:1006647646d80f16130f052404c6b901e80ee4ed6bef6792e1f238a8969106f7 \ - --hash=sha256:af0276409f9a02373d540bf8480021a048711d572745aef4b7842dad245eba10 +platformdirs==3.11.0 \ + --hash=sha256:cf8ee52a3afdb965072dcc652433e0c7e3e40cf5ea1477cd4b3b1d2eb75495b3 \ + --hash=sha256:e9d171d00af68be50e9202731309c4e658fd8bc76f55c11c7dd760d023bda68e # via virtualenv protobuf==3.20.3 \ --hash=sha256:03038ac1cfbc41aa21f6afcbcd357281d7521b4157926f30ebecc8d4ea59dcb7 \ @@ -378,34 +410,31 @@ protobuf==3.20.3 \ # gcp-docuploader # gcp-releasetool # google-api-core -pyasn1==0.4.8 \ - --hash=sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d \ - --hash=sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba + # googleapis-common-protos +pyasn1==0.5.0 \ + --hash=sha256:87a2121042a1ac9358cabcaf1d07680ff97ee6404333bacca15f76aa8ad01a57 \ + --hash=sha256:97b7290ca68e62a832558ec3976f15cbf911bf5d7c7039d8b861c2a0ece69fde # via # pyasn1-modules # rsa -pyasn1-modules==0.2.8 \ - --hash=sha256:905f84c712230b2c592c19470d3ca8d552de726050d1d1716282a1f6146be65e \ - --hash=sha256:a50b808ffeb97cb3601dd25981f6b016cbb3d31fbf57a8b8a87428e6158d0c74 +pyasn1-modules==0.3.0 \ + --hash=sha256:5bd01446b736eb9d31512a30d46c1ac3395d676c6f3cafa4c03eb54b9925631c \ + --hash=sha256:d3ccd6ed470d9ffbc716be08bd90efbd44d0734bc9303818f7336070984a162d # via google-auth pycparser==2.21 \ --hash=sha256:8ee45429555515e1f6b185e78100aea234072576aa43ab53aefcae078162fca9 \ --hash=sha256:e644fdec12f7872f86c58ff790da456218b10f863970249516d60a5eaca77206 # via cffi -pygments==2.13.0 \ - --hash=sha256:56a8508ae95f98e2b9bdf93a6be5ae3f7d8af858b43e02c5a2ff083726be40c1 \ - --hash=sha256:f643f331ab57ba3c9d89212ee4a2dabc6e94f117cf4eefde99a0574720d14c42 +pygments==2.16.1 \ + --hash=sha256:13fc09fa63bc8d8671a6d247e1eb303c4b343eaee81d861f3404db2935653692 \ + --hash=sha256:1daff0494820c69bc8941e407aa20f577374ee88364ee10a98fdbe0aece96e29 # via # readme-renderer # rich -pyjwt==2.6.0 \ - --hash=sha256:69285c7e31fc44f68a1feb309e948e0df53259d579295e6cfe2b1792329f05fd \ - --hash=sha256:d83c3d892a77bbb74d3e1a2cfa90afaadb60945205d1095d9221f04466f64c14 +pyjwt==2.8.0 \ + --hash=sha256:57e28d156e3d5c10088e0c68abb90bfac3df82b40a71bd0daa20c65ccd5c23de \ + --hash=sha256:59127c392cc44c2da5bb3192169a91f429924e17aff6534d70fdc02ab3e04320 # via gcp-releasetool -pyparsing==3.0.9 \ - --hash=sha256:2b020ecf7d21b687f219b71ecad3631f644a47f01403fa1d1036b0c6416d70fb \ - --hash=sha256:5026bae9a10eeaefb61dab2f09052b9f4307d44aee4eda64b309723d8d206bbc - # via packaging pyperclip==1.8.2 \ --hash=sha256:105254a8b04934f0bc84e9c24eb360a591aaf6535c9def5f29d92af107a9bf57 # via gcp-releasetool @@ -413,9 +442,9 @@ python-dateutil==2.8.2 \ --hash=sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86 \ --hash=sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9 # via gcp-releasetool -readme-renderer==37.3 \ - --hash=sha256:cd653186dfc73055656f090f227f5cb22a046d7f71a841dfa305f55c9a513273 \ - --hash=sha256:f67a16caedfa71eef48a31b39708637a6f4664c4394801a7b0d6432d13907343 +readme-renderer==42.0 \ + --hash=sha256:13d039515c1f24de668e2c93f2e877b9dbe6c6c32328b90a40a49d8b2b85f36d \ + --hash=sha256:2d55489f83be4992fe4454939d1a051c33edbab778e82761d060c9fc6b308cd1 # via twine requests==2.31.0 \ --hash=sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f \ @@ -426,17 +455,17 @@ requests==2.31.0 \ # google-cloud-storage # requests-toolbelt # twine -requests-toolbelt==0.10.1 \ - --hash=sha256:18565aa58116d9951ac39baa288d3adb5b3ff975c4f25eee78555d89e8f247f7 \ - --hash=sha256:62e09f7ff5ccbda92772a29f394a49c3ad6cb181d568b1337626b2abb628a63d +requests-toolbelt==1.0.0 \ + --hash=sha256:7681a0a3d047012b5bdc0ee37d7f8f07ebe76ab08caeccfc3921ce23c88d5bc6 \ + --hash=sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06 # via twine rfc3986==2.0.0 \ --hash=sha256:50b1502b60e289cb37883f3dfd34532b8873c7de9f49bb546641ce9cbd256ebd \ --hash=sha256:97aacf9dbd4bfd829baad6e6309fa6573aaf1be3f6fa735c8ab05e46cecb261c # via twine -rich==12.6.0 \ - --hash=sha256:a4eb26484f2c82589bd9a17c73d32a010b1e29d89f1604cd9bf3a2097b81bb5e \ - --hash=sha256:ba3a3775974105c221d31141f2c116f4fd65c5ceb0698657a11e9f295ec93fd0 +rich==13.6.0 \ + --hash=sha256:2b38e2fe9ca72c9a00170a1a2d20c63c790d0e10ef1fe35eba76e1e7b1d7d245 \ + --hash=sha256:5c14d22737e6d5084ef4771b62d5d4363165b403455a30a1c8ca39dc7b644bef # via twine rsa==4.9 \ --hash=sha256:90260d9058e514786967344d0ef75fa8727eed8a7d2e43ce9f4bcf1b536174f7 \ @@ -450,43 +479,37 @@ six==1.16.0 \ --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \ --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254 # via - # bleach # gcp-docuploader - # google-auth # python-dateutil -twine==4.0.1 \ - --hash=sha256:42026c18e394eac3e06693ee52010baa5313e4811d5a11050e7d48436cf41b9e \ - --hash=sha256:96b1cf12f7ae611a4a40b6ae8e9570215daff0611828f5fe1f37a16255ab24a0 +twine==4.0.2 \ + --hash=sha256:929bc3c280033347a00f847236564d1c52a3e61b1ac2516c97c48f3ceab756d8 \ + --hash=sha256:9e102ef5fdd5a20661eb88fad46338806c3bd32cf1db729603fe3697b1bc83c8 # via -r requirements.in -typing-extensions==4.4.0 \ - --hash=sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa \ - --hash=sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e +typing-extensions==4.8.0 \ + --hash=sha256:8f92fc8806f9a6b641eaa5318da32b44d401efaac0f6678c9bc448ba3605faa0 \ + --hash=sha256:df8e4339e9cb77357558cbdbceca33c303714cf861d1eef15e1070055ae8b7ef # via -r requirements.in -urllib3==1.26.12 \ - --hash=sha256:3fa96cf423e6987997fc326ae8df396db2a8b7c667747d47ddd8ecba91f4a74e \ - --hash=sha256:b930dd878d5a8afb066a637fbb35144fe7901e3b209d1cd4f524bd0e9deee997 +urllib3==2.0.7 \ + --hash=sha256:c97dfde1f7bd43a71c8d2a58e369e9b2bf692d1334ea9f9cae55add7d0dd0f84 \ + --hash=sha256:fdb6d215c776278489906c2f8916e6e7d4f5a9b602ccbcfdf7f016fc8da0596e # via # requests # twine -virtualenv==20.16.7 \ - --hash=sha256:8691e3ff9387f743e00f6bb20f70121f5e4f596cae754531f2b3b3a1b1ac696e \ - --hash=sha256:efd66b00386fdb7dbe4822d172303f40cd05e50e01740b19ea42425cbe653e29 +virtualenv==20.24.6 \ + --hash=sha256:02ece4f56fbf939dbbc33c0715159951d6bf14aaf5457b092e4548e1382455af \ + --hash=sha256:520d056652454c5098a00c0f073611ccbea4c79089331f60bf9d7ba247bb7381 # via nox -webencodings==0.5.1 \ - --hash=sha256:a0af1213f3c2226497a97e2b3aa01a7e4bee4f403f95be16fc9acd2947514a78 \ - --hash=sha256:b36a1c245f2d304965eb4e0a82848379241dc04b865afcc4aab16748587e1923 - # via bleach -wheel==0.38.4 \ - --hash=sha256:965f5259b566725405b05e7cf774052044b1ed30119b5d586b2703aafe8719ac \ - --hash=sha256:b60533f3f5d530e971d6737ca6d58681ee434818fab630c83a734bb10c083ce8 +wheel==0.41.3 \ + --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \ + --hash=sha256:4d4987ce51a49370ea65c0bfd2234e8ce80a12780820d9dc462597a6e60d0841 # via -r requirements.in -zipp==3.10.0 \ - --hash=sha256:4fcb6f278987a6605757302a6e40e896257570d11c51628968ccb2a47e80c6c1 \ - --hash=sha256:7a7262fd930bd3e36c50b9a64897aec3fafff3dfdeec9623ae22b40e93f99bb8 +zipp==3.17.0 \ + --hash=sha256:0e923e726174922dce09c53c59ad483ff7bbb8e572e00c7f7c46b88556409f31 \ + --hash=sha256:84e64a1c28cf7e91ed2078bb8cc8c259cb19b76942096c8d7b84947690cabaf0 # via importlib-metadata # The following packages are considered to be unsafe in a requirements file: -setuptools==65.5.1 \ - --hash=sha256:d0b9a8433464d5800cbe05094acf5c6d52a91bfac9b52bcfc4d41382be5d5d31 \ - --hash=sha256:e197a19aa8ec9722928f2206f8de752def0e4c9fc6953527360d1c36d94ddb2f +setuptools==68.2.2 \ + --hash=sha256:4ac1475276d2f1c48684874089fefcd83bd7162ddaafb81fac866ba0db282a87 \ + --hash=sha256:b454a35605876da60632df1a60f736524eb73cc47bbc9f3f1ef1b644de74fd2a # via -r requirements.in diff --git a/.kokoro/samples/python3.12/common.cfg b/.kokoro/samples/python3.12/common.cfg new file mode 100644 index 000000000..34e0a95f3 --- /dev/null +++ b/.kokoro/samples/python3.12/common.cfg @@ -0,0 +1,40 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Build logs will be here +action { + define_artifacts { + regex: "**/*sponge_log.xml" + } +} + +# Specify which tests to run +env_vars: { + key: "RUN_TESTS_SESSION" + value: "py-3.12" +} + +# Declare build specific Cloud project. +env_vars: { + key: "BUILD_SPECIFIC_GCLOUD_PROJECT" + value: "python-docs-samples-tests-312" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigtable/.kokoro/test-samples.sh" +} + +# Configure the docker image for kokoro-trampoline. +env_vars: { + key: "TRAMPOLINE_IMAGE" + value: "gcr.io/cloud-devrel-kokoro-resources/python-samples-testing-docker" +} + +# Download secrets for samples +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/python-docs-samples" + +# Download trampoline resources. +gfile_resources: "/bigstore/cloud-devrel-kokoro-resources/trampoline" + +# Use the trampoline script to run in docker. +build_file: "python-bigtable/.kokoro/trampoline_v2.sh" \ No newline at end of file diff --git a/.kokoro/samples/python3.12/continuous.cfg b/.kokoro/samples/python3.12/continuous.cfg new file mode 100644 index 000000000..a1c8d9759 --- /dev/null +++ b/.kokoro/samples/python3.12/continuous.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/samples/python3.12/periodic-head.cfg b/.kokoro/samples/python3.12/periodic-head.cfg new file mode 100644 index 000000000..be25a34f9 --- /dev/null +++ b/.kokoro/samples/python3.12/periodic-head.cfg @@ -0,0 +1,11 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} + +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/python-bigtable/.kokoro/test-samples-against-head.sh" +} diff --git a/.kokoro/samples/python3.12/periodic.cfg b/.kokoro/samples/python3.12/periodic.cfg new file mode 100644 index 000000000..71cd1e597 --- /dev/null +++ b/.kokoro/samples/python3.12/periodic.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "False" +} diff --git a/.kokoro/samples/python3.12/presubmit.cfg b/.kokoro/samples/python3.12/presubmit.cfg new file mode 100644 index 000000000..a1c8d9759 --- /dev/null +++ b/.kokoro/samples/python3.12/presubmit.cfg @@ -0,0 +1,6 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +env_vars: { + key: "INSTALL_LIBRARY_FROM_SOURCE" + value: "True" +} \ No newline at end of file diff --git a/.kokoro/test-samples-against-head.sh b/.kokoro/test-samples-against-head.sh index ba3a707b0..63ac41dfa 100755 --- a/.kokoro/test-samples-against-head.sh +++ b/.kokoro/test-samples-against-head.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/test-samples-impl.sh b/.kokoro/test-samples-impl.sh index 2c6500cae..5a0f5fab6 100755 --- a/.kokoro/test-samples-impl.sh +++ b/.kokoro/test-samples-impl.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2021 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/test-samples.sh b/.kokoro/test-samples.sh index 11c042d34..50b35a48c 100755 --- a/.kokoro/test-samples.sh +++ b/.kokoro/test-samples.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/trampoline.sh b/.kokoro/trampoline.sh index a4241db23..5c7c8633a 100755 --- a/.kokoro/trampoline.sh +++ b/.kokoro/trampoline.sh @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2017 Google Inc. +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.kokoro/trampoline_v2.sh b/.kokoro/trampoline_v2.sh index 4af6cdc26..59a7cf3a9 100755 --- a/.kokoro/trampoline_v2.sh +++ b/.kokoro/trampoline_v2.sh @@ -1,5 +1,5 @@ #!/usr/bin/env bash -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5405cc8ff..6a8e16950 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,4 +1,4 @@ -# Copyright 2021 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,10 +22,10 @@ repos: - id: end-of-file-fixer - id: check-yaml - repo: https://github.com/psf/black - rev: 22.3.0 + rev: 23.7.0 hooks: - id: black - repo: https://github.com/pycqa/flake8 - rev: 3.9.2 + rev: 6.1.0 hooks: - id: flake8 diff --git a/.release-please-manifest.json b/.release-please-manifest.json index b7f666a68..a5ab48803 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,3 +1,3 @@ { - ".": "2.19.0" + ".": "2.22.0" } \ No newline at end of file diff --git a/.trampolinerc b/.trampolinerc index 0eee72ab6..a7dfeb42c 100644 --- a/.trampolinerc +++ b/.trampolinerc @@ -1,4 +1,4 @@ -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Template for .trampolinerc - # Add required env vars here. required_envvars+=( ) diff --git a/CHANGELOG.md b/CHANGELOG.md index dc80386a4..5f86fdd88 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,60 @@ [1]: https://pypi.org/project/google-cloud-bigtable/#history +## [2.22.0](https://github.com/googleapis/python-bigtable/compare/v2.21.0...v2.22.0) (2023-12-12) + + +### Features + +* Add support for Cloud Bigtable Request Priorities in App Profiles ([#871](https://github.com/googleapis/python-bigtable/issues/871)) ([a4d551e](https://github.com/googleapis/python-bigtable/commit/a4d551e34006202ee96a395a2107d7acdc5881de)) +* Add support for Python 3.12 ([#888](https://github.com/googleapis/python-bigtable/issues/888)) ([4f050aa](https://github.com/googleapis/python-bigtable/commit/4f050aa5aed9a9dcf209779d5c10e5de8e2ff19e)) +* Introduce compatibility with native namespace packages ([#893](https://github.com/googleapis/python-bigtable/issues/893)) ([d218f4e](https://github.com/googleapis/python-bigtable/commit/d218f4ebd4ed6705721dca9318df955b40b0d0ac)) +* Publish CopyBackup protos to external customers ([#855](https://github.com/googleapis/python-bigtable/issues/855)) ([4105df7](https://github.com/googleapis/python-bigtable/commit/4105df762f1318c49bba030063897f0c50e4daee)) + + +### Bug Fixes + +* Add feature flag for improved mutate rows throttling ([e5af359](https://github.com/googleapis/python-bigtable/commit/e5af3597f45fc4c094c59abca876374f5a866c1b)) +* Add lock to flow control ([#899](https://github.com/googleapis/python-bigtable/issues/899)) ([e4e63c7](https://github.com/googleapis/python-bigtable/commit/e4e63c7b5b91273b3aae04fda59cc5a21c848de2)) +* Mutations batcher race condition ([#896](https://github.com/googleapis/python-bigtable/issues/896)) ([fe58f61](https://github.com/googleapis/python-bigtable/commit/fe58f617c7364d7e99e2ec50abd5f080852bf033)) +* Require google-cloud-core 1.4.4 ([#866](https://github.com/googleapis/python-bigtable/issues/866)) ([09f8a46](https://github.com/googleapis/python-bigtable/commit/09f8a4667d8b68a9f2048ba1aa57db4f775a2c03)) +* Use `retry_async` instead of `retry` in async client ([597efd1](https://github.com/googleapis/python-bigtable/commit/597efd11d15f20549010b4301be4d9768326e6a2)) + + +### Documentation + +* Minor formatting ([e5af359](https://github.com/googleapis/python-bigtable/commit/e5af3597f45fc4c094c59abca876374f5a866c1b)) + +## [2.21.0](https://github.com/googleapis/python-bigtable/compare/v2.20.0...v2.21.0) (2023-08-02) + + +### Features + +* Add last_scanned_row_responses to FeatureFlags ([#845](https://github.com/googleapis/python-bigtable/issues/845)) ([14a6739](https://github.com/googleapis/python-bigtable/commit/14a673901f82fa247c8027730a0bba41e0ec4757)) + + +### Documentation + +* Minor formatting ([#851](https://github.com/googleapis/python-bigtable/issues/851)) ([5ebe231](https://github.com/googleapis/python-bigtable/commit/5ebe2312dab70210811fca68c6625d2546442afd)) + +## [2.20.0](https://github.com/googleapis/python-bigtable/compare/v2.19.0...v2.20.0) (2023-07-17) + + +### Features + +* Add experimental reverse scan for public preview ([d5720f8](https://github.com/googleapis/python-bigtable/commit/d5720f8f5b5a81572f31d40051b3ec0f1d104304)) +* Increase the maximum retention period for a Cloud Bigtable backup from 30 days to 90 days ([d5720f8](https://github.com/googleapis/python-bigtable/commit/d5720f8f5b5a81572f31d40051b3ec0f1d104304)) + + +### Bug Fixes + +* Add async context manager return types ([#828](https://github.com/googleapis/python-bigtable/issues/828)) ([475a160](https://github.com/googleapis/python-bigtable/commit/475a16072f3ad41357bdb765fff608a39141ec00)) + + +### Documentation + +* Fix formatting for reversed order field example ([#831](https://github.com/googleapis/python-bigtable/issues/831)) ([fddd0ba](https://github.com/googleapis/python-bigtable/commit/fddd0ba97155e112af92a98fd8f20e59b139d177)) + ## [2.19.0](https://github.com/googleapis/python-bigtable/compare/v2.18.1...v2.19.0) (2023-06-08) diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 504fb3742..947c129b7 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -22,7 +22,7 @@ In order to add a feature: documentation. - The feature must work fully on the following CPython versions: - 3.7, 3.8, 3.9, 3.10 and 3.11 on both UNIX and Windows. + 3.7, 3.8, 3.9, 3.10, 3.11 and 3.12 on both UNIX and Windows. - The feature must not add unnecessary dependencies (where "unnecessary" is of course subjective, but new dependencies should @@ -72,7 +72,7 @@ We use `nox `__ to instrument our tests. - To run a single unit test:: - $ nox -s unit-3.11 -- -k + $ nox -s unit-3.12 -- -k .. note:: @@ -226,12 +226,14 @@ We support: - `Python 3.9`_ - `Python 3.10`_ - `Python 3.11`_ +- `Python 3.12`_ .. _Python 3.7: https://docs.python.org/3.7/ .. _Python 3.8: https://docs.python.org/3.8/ .. _Python 3.9: https://docs.python.org/3.9/ .. _Python 3.10: https://docs.python.org/3.10/ .. _Python 3.11: https://docs.python.org/3.11/ +.. _Python 3.12: https://docs.python.org/3.12/ Supported versions can be found in our ``noxfile.py`` `config`_. diff --git a/MANIFEST.in b/MANIFEST.in index e783f4c62..e0a667053 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/docs/conf.py b/docs/conf.py index 34f3a4d08..b5a870f58 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2021 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/docs/snippets.py b/docs/snippets.py index 1d93fdf12..fa3aa3627 100644 --- a/docs/snippets.py +++ b/docs/snippets.py @@ -448,7 +448,6 @@ def test_bigtable_create_table(): def test_bigtable_list_tables(): - # [START bigtable_api_list_tables] from google.cloud.bigtable import Client diff --git a/docs/snippets_table.py b/docs/snippets_table.py index f27260425..893135275 100644 --- a/docs/snippets_table.py +++ b/docs/snippets_table.py @@ -964,7 +964,6 @@ def test_bigtable_create_family_gc_nested(): def test_bigtable_row_data_cells_cell_value_cell_values(): - value = b"value_in_col1" row = Config.TABLE.row(b"row_key_1") row.set_cell( diff --git a/google/cloud/bigtable/batcher.py b/google/cloud/bigtable/batcher.py index a6eb806e9..f9b85386d 100644 --- a/google/cloud/bigtable/batcher.py +++ b/google/cloud/bigtable/batcher.py @@ -53,12 +53,19 @@ def __init__(self, max_mutation_bytes=MAX_MUTATION_SIZE, flush_count=FLUSH_COUNT self.flush_count = flush_count def get(self): - """Retrieve an item from the queue. Recalculate queue size.""" - row = self._queue.get() - mutation_size = row.get_mutations_size() - self.total_mutation_count -= len(row._get_mutations()) - self.total_size -= mutation_size - return row + """ + Retrieve an item from the queue. Recalculate queue size. + + If the queue is empty, return None. + """ + try: + row = self._queue.get_nowait() + mutation_size = row.get_mutations_size() + self.total_mutation_count -= len(row._get_mutations()) + self.total_size -= mutation_size + return row + except queue.Empty: + return None def put(self, item): """Insert an item to the queue. Recalculate queue size.""" @@ -79,9 +86,6 @@ def full(self): return True return False - def empty(self): - return self._queue.empty() - @dataclass class _BatchInfo: @@ -110,6 +114,7 @@ def __init__( self.inflight_size = 0 self.event = threading.Event() self.event.set() + self._lock = threading.Lock() def is_blocked(self): """Returns True if: @@ -128,8 +133,9 @@ def control_flow(self, batch_info): Calculate the resources used by this batch """ - self.inflight_mutations += batch_info.mutations_count - self.inflight_size += batch_info.mutations_size + with self._lock: + self.inflight_mutations += batch_info.mutations_count + self.inflight_size += batch_info.mutations_size self.set_flow_control_status() def wait(self): @@ -154,8 +160,9 @@ def release(self, batch_info): Release the resources. Decrement the row size to allow enqueued mutations to be run. """ - self.inflight_mutations -= batch_info.mutations_count - self.inflight_size -= batch_info.mutations_size + with self._lock: + self.inflight_mutations -= batch_info.mutations_count + self.inflight_size -= batch_info.mutations_size self.set_flow_control_status() @@ -292,8 +299,10 @@ def flush(self): * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. """ rows_to_flush = [] - while not self._rows.empty(): - rows_to_flush.append(self._rows.get()) + row = self._rows.get() + while row is not None: + rows_to_flush.append(row) + row = self._rows.get() response = self._flush_rows(rows_to_flush) return response @@ -303,58 +312,68 @@ def _flush_async(self): :raises: * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. """ - - rows_to_flush = [] - mutations_count = 0 - mutations_size = 0 - rows_count = 0 - batch_info = _BatchInfo() - - while not self._rows.empty(): - row = self._rows.get() - mutations_count += len(row._get_mutations()) - mutations_size += row.get_mutations_size() - rows_count += 1 - rows_to_flush.append(row) - batch_info.mutations_count = mutations_count - batch_info.rows_count = rows_count - batch_info.mutations_size = mutations_size - - if ( - rows_count >= self.flush_count - or mutations_size >= self.max_row_bytes - or mutations_count >= self.flow_control.max_mutations - or mutations_size >= self.flow_control.max_mutation_bytes - or self._rows.empty() # submit when it reached the end of the queue + next_row = self._rows.get() + while next_row is not None: + # start a new batch + rows_to_flush = [next_row] + batch_info = _BatchInfo( + mutations_count=len(next_row._get_mutations()), + rows_count=1, + mutations_size=next_row.get_mutations_size(), + ) + # fill up batch with rows + next_row = self._rows.get() + while next_row is not None and self._row_fits_in_batch( + next_row, batch_info ): - # wait for resources to become available, before submitting any new batch - self.flow_control.wait() - # once unblocked, submit a batch - # event flag will be set by control_flow to block subsequent thread, but not blocking this one - self.flow_control.control_flow(batch_info) - future = self._executor.submit(self._flush_rows, rows_to_flush) - self.futures_mapping[future] = batch_info - future.add_done_callback(self._batch_completed_callback) - - # reset and start a new batch - rows_to_flush = [] - mutations_size = 0 - rows_count = 0 - mutations_count = 0 - batch_info = _BatchInfo() + rows_to_flush.append(next_row) + batch_info.mutations_count += len(next_row._get_mutations()) + batch_info.rows_count += 1 + batch_info.mutations_size += next_row.get_mutations_size() + next_row = self._rows.get() + # send batch over network + # wait for resources to become available + self.flow_control.wait() + # once unblocked, submit the batch + # event flag will be set by control_flow to block subsequent thread, but not blocking this one + self.flow_control.control_flow(batch_info) + future = self._executor.submit(self._flush_rows, rows_to_flush) + # schedule release of resources from flow control + self.futures_mapping[future] = batch_info + future.add_done_callback(self._batch_completed_callback) def _batch_completed_callback(self, future): """Callback for when the mutation has finished to clean up the current batch and release items from the flow controller. - Raise exceptions if there's any. Release the resources locked by the flow control and allow enqueued tasks to be run. """ - processed_rows = self.futures_mapping[future] self.flow_control.release(processed_rows) del self.futures_mapping[future] + def _row_fits_in_batch(self, row, batch_info): + """Checks if a row can fit in the current batch. + + :type row: class + :param row: :class:`~google.cloud.bigtable.row.DirectRow`. + + :type batch_info: :class:`_BatchInfo` + :param batch_info: Information about the current batch. + + :rtype: bool + :returns: True if the row can fit in the current batch. + """ + new_rows_count = batch_info.rows_count + 1 + new_mutations_count = batch_info.mutations_count + len(row._get_mutations()) + new_mutations_size = batch_info.mutations_size + row.get_mutations_size() + return ( + new_rows_count <= self.flush_count + and new_mutations_size <= self.max_row_bytes + and new_mutations_count <= self.flow_control.max_mutations + and new_mutations_size <= self.flow_control.max_mutation_bytes + ) + def _flush_rows(self, rows_to_flush): """Mutate the specified rows. diff --git a/google/cloud/bigtable/gapic_version.py b/google/cloud/bigtable/gapic_version.py index 0f1a446f3..03d6d0200 100644 --- a/google/cloud/bigtable/gapic_version.py +++ b/google/cloud/bigtable/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.19.0" # {x-release-please-version} +__version__ = "2.22.0" # {x-release-please-version} diff --git a/google/cloud/bigtable_admin/__init__.py b/google/cloud/bigtable_admin/__init__.py index 0ba93ec63..d26d79b3c 100644 --- a/google/cloud/bigtable_admin/__init__.py +++ b/google/cloud/bigtable_admin/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -115,6 +115,8 @@ from google.cloud.bigtable_admin_v2.types.bigtable_table_admin import ( CheckConsistencyResponse, ) +from google.cloud.bigtable_admin_v2.types.bigtable_table_admin import CopyBackupMetadata +from google.cloud.bigtable_admin_v2.types.bigtable_table_admin import CopyBackupRequest from google.cloud.bigtable_admin_v2.types.bigtable_table_admin import ( CreateBackupMetadata, ) @@ -242,6 +244,8 @@ "UpdateInstanceMetadata", "CheckConsistencyRequest", "CheckConsistencyResponse", + "CopyBackupMetadata", + "CopyBackupRequest", "CreateBackupMetadata", "CreateBackupRequest", "CreateTableFromSnapshotMetadata", diff --git a/google/cloud/bigtable_admin/gapic_version.py b/google/cloud/bigtable_admin/gapic_version.py index 0f1a446f3..03d6d0200 100644 --- a/google/cloud/bigtable_admin/gapic_version.py +++ b/google/cloud/bigtable_admin/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.19.0" # {x-release-please-version} +__version__ = "2.22.0" # {x-release-please-version} diff --git a/google/cloud/bigtable_admin_v2/__init__.py b/google/cloud/bigtable_admin_v2/__init__.py index c030ec1bd..811b956e0 100644 --- a/google/cloud/bigtable_admin_v2/__init__.py +++ b/google/cloud/bigtable_admin_v2/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -51,6 +51,8 @@ from .types.bigtable_instance_admin import UpdateInstanceMetadata from .types.bigtable_table_admin import CheckConsistencyRequest from .types.bigtable_table_admin import CheckConsistencyResponse +from .types.bigtable_table_admin import CopyBackupMetadata +from .types.bigtable_table_admin import CopyBackupRequest from .types.bigtable_table_admin import CreateBackupMetadata from .types.bigtable_table_admin import CreateBackupRequest from .types.bigtable_table_admin import CreateTableFromSnapshotMetadata @@ -116,6 +118,8 @@ "CheckConsistencyResponse", "Cluster", "ColumnFamily", + "CopyBackupMetadata", + "CopyBackupRequest", "CreateAppProfileRequest", "CreateBackupMetadata", "CreateBackupRequest", diff --git a/google/cloud/bigtable_admin_v2/gapic_metadata.json b/google/cloud/bigtable_admin_v2/gapic_metadata.json index d797338cc..9b3426470 100644 --- a/google/cloud/bigtable_admin_v2/gapic_metadata.json +++ b/google/cloud/bigtable_admin_v2/gapic_metadata.json @@ -349,6 +349,11 @@ "check_consistency" ] }, + "CopyBackup": { + "methods": [ + "copy_backup" + ] + }, "CreateBackup": { "methods": [ "create_backup" @@ -474,6 +479,11 @@ "check_consistency" ] }, + "CopyBackup": { + "methods": [ + "copy_backup" + ] + }, "CreateBackup": { "methods": [ "create_backup" @@ -599,6 +609,11 @@ "check_consistency" ] }, + "CopyBackup": { + "methods": [ + "copy_backup" + ] + }, "CreateBackup": { "methods": [ "create_backup" diff --git a/google/cloud/bigtable_admin_v2/gapic_version.py b/google/cloud/bigtable_admin_v2/gapic_version.py index 0f1a446f3..03d6d0200 100644 --- a/google/cloud/bigtable_admin_v2/gapic_version.py +++ b/google/cloud/bigtable_admin_v2/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.19.0" # {x-release-please-version} +__version__ = "2.22.0" # {x-release-please-version} diff --git a/google/cloud/bigtable_admin_v2/services/__init__.py b/google/cloud/bigtable_admin_v2/services/__init__.py index e8e1c3845..89a37dc92 100644 --- a/google/cloud/bigtable_admin_v2/services/__init__.py +++ b/google/cloud/bigtable_admin_v2/services/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/__init__.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/__init__.py index 1fb10736e..40631d1b4 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/__init__.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/async_client.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/async_client.py index 12811bcea..e4c4639af 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/async_client.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/async_client.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,14 +33,14 @@ from google.api_core.client_options import ClientOptions from google.api_core import exceptions as core_exceptions from google.api_core import gapic_v1 -from google.api_core import retry as retries +from google.api_core import retry_async as retries from google.auth import credentials as ga_credentials # type: ignore from google.oauth2 import service_account # type: ignore try: - OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault] + OptionalRetry = Union[retries.AsyncRetry, gapic_v1.method._MethodDefault] except AttributeError: # pragma: NO COVER - OptionalRetry = Union[retries.Retry, object] # type: ignore + OptionalRetry = Union[retries.AsyncRetry, object] # type: ignore from google.api_core import operation # type: ignore from google.api_core import operation_async # type: ignore @@ -305,7 +305,7 @@ async def create_instance( This corresponds to the ``clusters`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -404,7 +404,7 @@ async def get_instance( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -440,7 +440,7 @@ async def get_instance( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.get_instance, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -496,7 +496,7 @@ async def list_instances( This corresponds to the ``parent`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -529,7 +529,7 @@ async def list_instances( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.list_instances, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -581,7 +581,7 @@ async def update_instance( served from all [Clusters][google.bigtable.admin.v2.Cluster] in the instance. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -603,7 +603,7 @@ async def update_instance( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.update_instance, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -669,7 +669,7 @@ async def partial_update_instance( This corresponds to the ``update_mask`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -709,7 +709,7 @@ async def partial_update_instance( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.partial_update_instance, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -775,7 +775,7 @@ async def delete_instance( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -869,7 +869,7 @@ async def create_cluster( This corresponds to the ``cluster`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -964,7 +964,7 @@ async def get_cluster( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -999,7 +999,7 @@ async def get_cluster( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.get_cluster, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1057,7 +1057,7 @@ async def list_clusters( This corresponds to the ``parent`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1090,7 +1090,7 @@ async def list_clusters( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.list_clusters, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1141,7 +1141,7 @@ async def update_cluster( location, capable of serving all [Tables][google.bigtable.admin.v2.Table] in the parent [Instance][google.bigtable.admin.v2.Instance]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1164,7 +1164,7 @@ async def update_cluster( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.update_cluster, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1248,7 +1248,7 @@ async def partial_update_cluster( This corresponds to the ``update_mask`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1343,7 +1343,7 @@ async def delete_cluster( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1431,7 +1431,7 @@ async def create_app_profile( This corresponds to the ``app_profile`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1515,7 +1515,7 @@ async def get_app_profile( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1549,7 +1549,7 @@ async def get_app_profile( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.get_app_profile, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1608,7 +1608,7 @@ async def list_app_profiles( This corresponds to the ``parent`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1644,7 +1644,7 @@ async def list_app_profiles( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.list_app_profiles, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1717,7 +1717,7 @@ async def update_app_profile( This corresponds to the ``update_mask`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1754,7 +1754,7 @@ async def update_app_profile( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.update_app_profile, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1820,7 +1820,7 @@ async def delete_app_profile( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1890,7 +1890,7 @@ async def get_iam_policy( This corresponds to the ``resource`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1919,42 +1919,11 @@ async def get_iam_policy( **JSON example:** - { - "bindings": [ - { - "role": - "roles/resourcemanager.organizationAdmin", - "members": [ "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - - }, { "role": - "roles/resourcemanager.organizationViewer", - "members": [ "user:eve@example.com" ], - "condition": { "title": "expirable access", - "description": "Does not grant access after - Sep 2020", "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", } } - - ], "etag": "BwWWja0YfJA=", "version": 3 - - } + :literal:`\` { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 }`\ \` **YAML example:** - bindings: - members: - user:\ mike@example.com - - group:\ admins@example.com - domain:google.com - - serviceAccount:\ my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - user:\ eve@example.com role: - roles/resourcemanager.organizationViewer - condition: title: expirable access description: - Does not grant access after Sep 2020 expression: - request.time < - timestamp('2020-10-01T00:00:00.000Z') etag: - BwWWja0YfJA= version: 3 + :literal:`\` bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') etag: BwWWja0YfJA= version: 3`\ \` For a description of IAM and its features, see the [IAM @@ -1984,7 +1953,7 @@ async def get_iam_policy( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.get_iam_policy, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -2039,7 +2008,7 @@ async def set_iam_policy( This corresponds to the ``resource`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2068,42 +2037,11 @@ async def set_iam_policy( **JSON example:** - { - "bindings": [ - { - "role": - "roles/resourcemanager.organizationAdmin", - "members": [ "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - - }, { "role": - "roles/resourcemanager.organizationViewer", - "members": [ "user:eve@example.com" ], - "condition": { "title": "expirable access", - "description": "Does not grant access after - Sep 2020", "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", } } - - ], "etag": "BwWWja0YfJA=", "version": 3 - - } + :literal:`\` { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 }`\ \` **YAML example:** - bindings: - members: - user:\ mike@example.com - - group:\ admins@example.com - domain:google.com - - serviceAccount:\ my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - user:\ eve@example.com role: - roles/resourcemanager.organizationViewer - condition: title: expirable access description: - Does not grant access after Sep 2020 expression: - request.time < - timestamp('2020-10-01T00:00:00.000Z') etag: - BwWWja0YfJA= version: 3 + :literal:`\` bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') etag: BwWWja0YfJA= version: 3`\ \` For a description of IAM and its features, see the [IAM @@ -2188,7 +2126,7 @@ async def test_iam_permissions( This corresponds to the ``permissions`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2222,7 +2160,7 @@ async def test_iam_permissions( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.test_iam_permissions, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -2279,7 +2217,7 @@ async def list_hot_tablets( This corresponds to the ``parent`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2315,7 +2253,7 @@ async def list_hot_tablets( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.list_hot_tablets, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -2355,7 +2293,7 @@ async def list_hot_tablets( # Done; return the response. return response - async def __aenter__(self): + async def __aenter__(self) -> "BigtableInstanceAdminAsyncClient": return self async def __aexit__(self, exc_type, exc, tb): diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/client.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/client.py index ecc9bf1e2..52c61ea4f 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/client.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/client.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -2143,42 +2143,11 @@ def get_iam_policy( **JSON example:** - { - "bindings": [ - { - "role": - "roles/resourcemanager.organizationAdmin", - "members": [ "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - - }, { "role": - "roles/resourcemanager.organizationViewer", - "members": [ "user:eve@example.com" ], - "condition": { "title": "expirable access", - "description": "Does not grant access after - Sep 2020", "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", } } - - ], "etag": "BwWWja0YfJA=", "version": 3 - - } + :literal:`\` { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 }`\ \` **YAML example:** - bindings: - members: - user:\ mike@example.com - - group:\ admins@example.com - domain:google.com - - serviceAccount:\ my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - user:\ eve@example.com role: - roles/resourcemanager.organizationViewer - condition: title: expirable access description: - Does not grant access after Sep 2020 expression: - request.time < - timestamp('2020-10-01T00:00:00.000Z') etag: - BwWWja0YfJA= version: 3 + :literal:`\` bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') etag: BwWWja0YfJA= version: 3`\ \` For a description of IAM and its features, see the [IAM @@ -2279,42 +2248,11 @@ def set_iam_policy( **JSON example:** - { - "bindings": [ - { - "role": - "roles/resourcemanager.organizationAdmin", - "members": [ "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - - }, { "role": - "roles/resourcemanager.organizationViewer", - "members": [ "user:eve@example.com" ], - "condition": { "title": "expirable access", - "description": "Does not grant access after - Sep 2020", "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", } } - - ], "etag": "BwWWja0YfJA=", "version": 3 - - } + :literal:`\` { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 }`\ \` **YAML example:** - bindings: - members: - user:\ mike@example.com - - group:\ admins@example.com - domain:google.com - - serviceAccount:\ my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - user:\ eve@example.com role: - roles/resourcemanager.organizationViewer - condition: title: expirable access description: - Does not grant access after Sep 2020 expression: - request.time < - timestamp('2020-10-01T00:00:00.000Z') etag: - BwWWja0YfJA= version: 3 + :literal:`\` bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') etag: BwWWja0YfJA= version: 3`\ \` For a description of IAM and its features, see the [IAM diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/pagers.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/pagers.py index bfcbbf23d..0d646a96e 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/pagers.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/pagers.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/__init__.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/__init__.py index e5637c0da..62da28c88 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/__init__.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/base.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/base.py index bd45f319f..d92d25453 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/base.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/base.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/grpc.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/grpc.py index f037f5a44..eca37957d 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/grpc.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/grpc.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/grpc_asyncio.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/grpc_asyncio.py index 82b03b0bb..145aa427d 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/grpc_asyncio.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/grpc_asyncio.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/rest.py b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/rest.py index e9b94cf78..9d5502b7e 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/rest.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_instance_admin/transports/rest.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -44,8 +44,8 @@ from google.cloud.bigtable_admin_v2.types import instance from google.iam.v1 import iam_policy_pb2 # type: ignore from google.iam.v1 import policy_pb2 # type: ignore -from google.longrunning import operations_pb2 # type: ignore from google.protobuf import empty_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore from .base import ( BigtableInstanceAdminTransport, @@ -1612,54 +1612,54 @@ def __call__( :: - { - "bindings": [ - { - "role": "roles/resourcemanager.organizationAdmin", - "members": [ - "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - }, - { - "role": "roles/resourcemanager.organizationViewer", - "members": [ - "user:eve@example.com" - ], - "condition": { - "title": "expirable access", - "description": "Does not grant access after Sep 2020", - "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", - } - } - ], - "etag": "BwWWja0YfJA=", - "version": 3 - } + { + "bindings": [ + { + "role": "roles/resourcemanager.organizationAdmin", + "members": [ + "user:mike@example.com", + "group:admins@example.com", + "domain:google.com", + "serviceAccount:my-project-id@appspot.gserviceaccount.com" + ] + }, + { + "role": "roles/resourcemanager.organizationViewer", + "members": [ + "user:eve@example.com" + ], + "condition": { + "title": "expirable access", + "description": "Does not grant access after Sep 2020", + "expression": "request.time < + timestamp('2020-10-01T00:00:00.000Z')", + } + } + ], + "etag": "BwWWja0YfJA=", + "version": 3 + } **YAML example:** :: - bindings: - - members: - - user:mike@example.com - - group:admins@example.com - - domain:google.com - - serviceAccount:my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - - user:eve@example.com - role: roles/resourcemanager.organizationViewer - condition: - title: expirable access - description: Does not grant access after Sep 2020 - expression: request.time < timestamp('2020-10-01T00:00:00.000Z') - etag: BwWWja0YfJA= - version: 3 + bindings: + - members: + - user:mike@example.com + - group:admins@example.com + - domain:google.com + - serviceAccount:my-project-id@appspot.gserviceaccount.com + role: roles/resourcemanager.organizationAdmin + - members: + - user:eve@example.com + role: roles/resourcemanager.organizationViewer + condition: + title: expirable access + description: Does not grant access after Sep 2020 + expression: request.time < timestamp('2020-10-01T00:00:00.000Z') + etag: BwWWja0YfJA= + version: 3 For a description of IAM and its features, see the `IAM documentation `__. @@ -2439,54 +2439,54 @@ def __call__( :: - { - "bindings": [ - { - "role": "roles/resourcemanager.organizationAdmin", - "members": [ - "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - }, - { - "role": "roles/resourcemanager.organizationViewer", - "members": [ - "user:eve@example.com" - ], - "condition": { - "title": "expirable access", - "description": "Does not grant access after Sep 2020", - "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", - } - } - ], - "etag": "BwWWja0YfJA=", - "version": 3 - } + { + "bindings": [ + { + "role": "roles/resourcemanager.organizationAdmin", + "members": [ + "user:mike@example.com", + "group:admins@example.com", + "domain:google.com", + "serviceAccount:my-project-id@appspot.gserviceaccount.com" + ] + }, + { + "role": "roles/resourcemanager.organizationViewer", + "members": [ + "user:eve@example.com" + ], + "condition": { + "title": "expirable access", + "description": "Does not grant access after Sep 2020", + "expression": "request.time < + timestamp('2020-10-01T00:00:00.000Z')", + } + } + ], + "etag": "BwWWja0YfJA=", + "version": 3 + } **YAML example:** :: - bindings: - - members: - - user:mike@example.com - - group:admins@example.com - - domain:google.com - - serviceAccount:my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - - user:eve@example.com - role: roles/resourcemanager.organizationViewer - condition: - title: expirable access - description: Does not grant access after Sep 2020 - expression: request.time < timestamp('2020-10-01T00:00:00.000Z') - etag: BwWWja0YfJA= - version: 3 + bindings: + - members: + - user:mike@example.com + - group:admins@example.com + - domain:google.com + - serviceAccount:my-project-id@appspot.gserviceaccount.com + role: roles/resourcemanager.organizationAdmin + - members: + - user:eve@example.com + role: roles/resourcemanager.organizationViewer + condition: + title: expirable access + description: Does not grant access after Sep 2020 + expression: request.time < timestamp('2020-10-01T00:00:00.000Z') + etag: BwWWja0YfJA= + version: 3 For a description of IAM and its features, see the `IAM documentation `__. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/__init__.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/__init__.py index 515696537..544649e90 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/__init__.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/async_client.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/async_client.py index 1663c16eb..5a4435bde 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/async_client.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/async_client.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,14 +33,14 @@ from google.api_core.client_options import ClientOptions from google.api_core import exceptions as core_exceptions from google.api_core import gapic_v1 -from google.api_core import retry as retries +from google.api_core import retry_async as retries from google.auth import credentials as ga_credentials # type: ignore from google.oauth2 import service_account # type: ignore try: - OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault] + OptionalRetry = Union[retries.AsyncRetry, gapic_v1.method._MethodDefault] except AttributeError: # pragma: NO COVER - OptionalRetry = Union[retries.Retry, object] # type: ignore + OptionalRetry = Union[retries.AsyncRetry, object] # type: ignore from google.api_core import operation # type: ignore from google.api_core import operation_async # type: ignore @@ -282,7 +282,7 @@ async def create_table( This corresponds to the ``table`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -358,6 +358,7 @@ async def create_table_from_snapshot( r"""Creates a new table from the specified snapshot. The target table must not exist. The snapshot and the table must be in the same instance. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -401,7 +402,7 @@ async def create_table_from_snapshot( This corresponds to the ``source_snapshot`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -493,7 +494,7 @@ async def list_tables( This corresponds to the ``parent`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -529,7 +530,7 @@ async def list_tables( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.list_tables, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -592,7 +593,7 @@ async def get_table( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -627,7 +628,7 @@ async def get_table( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.get_table, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -700,7 +701,7 @@ async def update_table( This corresponds to the ``update_mask`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -793,7 +794,7 @@ async def delete_table( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -864,7 +865,7 @@ async def undelete_table( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -975,7 +976,7 @@ async def modify_column_families( This corresponds to the ``modifications`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1050,7 +1051,7 @@ async def drop_row_range( request (Optional[Union[google.cloud.bigtable_admin_v2.types.DropRowRangeRequest, dict]]): The request object. Request message for [google.bigtable.admin.v2.BigtableTableAdmin.DropRowRange][google.bigtable.admin.v2.BigtableTableAdmin.DropRowRange] - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1110,7 +1111,7 @@ async def generate_consistency_token( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1143,7 +1144,7 @@ async def generate_consistency_token( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.generate_consistency_token, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1210,7 +1211,7 @@ async def check_consistency( This corresponds to the ``consistency_token`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1245,7 +1246,7 @@ async def check_consistency( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.check_consistency, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1293,6 +1294,7 @@ async def snapshot_table( r"""Creates a new snapshot in the specified cluster from the specified source table. The cluster and the table must be in the same instance. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1342,7 +1344,7 @@ async def snapshot_table( This corresponds to the ``description`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1457,7 +1459,7 @@ async def get_snapshot( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1469,6 +1471,7 @@ async def get_snapshot( time. A snapshot can be used as a checkpoint for data restoration or a data source for a new table. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud @@ -1500,7 +1503,7 @@ async def get_snapshot( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.get_snapshot, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1573,7 +1576,7 @@ async def list_snapshots( This corresponds to the ``parent`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1616,7 +1619,7 @@ async def list_snapshots( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.list_snapshots, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1668,6 +1671,7 @@ async def delete_snapshot( metadata: Sequence[Tuple[str, str]] = (), ) -> None: r"""Permanently deletes the specified snapshot. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1694,7 +1698,7 @@ async def delete_snapshot( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1790,7 +1794,7 @@ async def create_backup( This corresponds to the ``backup`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1882,7 +1886,7 @@ async def get_backup( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1913,7 +1917,7 @@ async def get_backup( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.get_backup, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -1983,7 +1987,7 @@ async def update_backup( This corresponds to the ``update_mask`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2062,7 +2066,7 @@ async def delete_backup( This corresponds to the ``name`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2134,7 +2138,7 @@ async def list_backups( This corresponds to the ``parent`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2143,7 +2147,7 @@ async def list_backups( Returns: google.cloud.bigtable_admin_v2.services.bigtable_table_admin.pagers.ListBackupsAsyncPager: The response for - [ListBackups][google.bigtable.admin.v2.BigtableTableAdmin.ListBackups]. + [ListBackups][google.bigtable.admin.v2.BigtableTableAdmin.ListBackups]. Iterating over this object will yield results and resolve additional pages automatically. @@ -2170,7 +2174,7 @@ async def list_backups( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.list_backups, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -2218,9 +2222,8 @@ async def restore_table( timeout: Union[float, object] = gapic_v1.method.DEFAULT, metadata: Sequence[Tuple[str, str]] = (), ) -> operation_async.AsyncOperation: - r"""Create a new table by restoring from a completed backup. The new - table must be in the same project as the instance containing the - backup. The returned table [long-running + r"""Create a new table by restoring from a completed backup. The + returned table [long-running operation][google.longrunning.Operation] can be used to track the progress of the operation, and to cancel it. The [metadata][google.longrunning.Operation.metadata] field type is @@ -2232,7 +2235,7 @@ async def restore_table( request (Optional[Union[google.cloud.bigtable_admin_v2.types.RestoreTableRequest, dict]]): The request object. The request for [RestoreTable][google.bigtable.admin.v2.BigtableTableAdmin.RestoreTable]. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2283,6 +2286,141 @@ async def restore_table( # Done; return the response. return response + async def copy_backup( + self, + request: Optional[Union[bigtable_table_admin.CopyBackupRequest, dict]] = None, + *, + parent: Optional[str] = None, + backup_id: Optional[str] = None, + source_backup: Optional[str] = None, + expire_time: Optional[timestamp_pb2.Timestamp] = None, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operation_async.AsyncOperation: + r"""Copy a Cloud Bigtable backup to a new backup in the + destination cluster located in the destination instance + and project. + + Args: + request (Optional[Union[google.cloud.bigtable_admin_v2.types.CopyBackupRequest, dict]]): + The request object. The request for + [CopyBackup][google.bigtable.admin.v2.BigtableTableAdmin.CopyBackup]. + parent (:class:`str`): + Required. The name of the destination cluster that will + contain the backup copy. The cluster must already + exists. Values are of the form: + ``projects/{project}/instances/{instance}/clusters/{cluster}``. + + This corresponds to the ``parent`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + backup_id (:class:`str`): + Required. The id of the new backup. The ``backup_id`` + along with ``parent`` are combined as + {parent}/backups/{backup_id} to create the full backup + name, of the form: + ``projects/{project}/instances/{instance}/clusters/{cluster}/backups/{backup_id}``. + This string must be between 1 and 50 characters in + length and match the regex [*a-zA-Z0-9][-*.a-zA-Z0-9]*. + + This corresponds to the ``backup_id`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + source_backup (:class:`str`): + Required. The source backup to be copied from. The + source backup needs to be in READY state for it to be + copied. Copying a copied backup is not allowed. Once + CopyBackup is in progress, the source backup cannot be + deleted or cleaned up on expiration until CopyBackup is + finished. Values are of the form: + ``projects//instances//clusters//backups/``. + + This corresponds to the ``source_backup`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + expire_time (:class:`google.protobuf.timestamp_pb2.Timestamp`): + Required. Required. The expiration time of the copied + backup with microsecond granularity that must be at + least 6 hours and at most 30 days from the time the + request is received. Once the ``expire_time`` has + passed, Cloud Bigtable will delete the backup and free + the resources used by the backup. + + This corresponds to the ``expire_time`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.api_core.operation_async.AsyncOperation: + An object representing a long-running operation. + + The result type for the operation will be + :class:`google.cloud.bigtable_admin_v2.types.Backup` A + backup of a Cloud Bigtable table. + + """ + # Create or coerce a protobuf request object. + # Quick check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([parent, backup_id, source_backup, expire_time]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = bigtable_table_admin.CopyBackupRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + if parent is not None: + request.parent = parent + if backup_id is not None: + request.backup_id = backup_id + if source_backup is not None: + request.source_backup = source_backup + if expire_time is not None: + request.expire_time = expire_time + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.copy_backup, + default_timeout=None, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("parent", request.parent),)), + ) + + # Send the request. + response = await rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Wrap the response in an operation future. + response = operation_async.from_gapic( + response, + self._client._transport.operations_client, + table.Backup, + metadata_type=bigtable_table_admin.CopyBackupMetadata, + ) + + # Done; return the response. + return response + async def get_iam_policy( self, request: Optional[Union[iam_policy_pb2.GetIamPolicyRequest, dict]] = None, @@ -2308,7 +2446,7 @@ async def get_iam_policy( This corresponds to the ``resource`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2337,42 +2475,11 @@ async def get_iam_policy( **JSON example:** - { - "bindings": [ - { - "role": - "roles/resourcemanager.organizationAdmin", - "members": [ "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - - }, { "role": - "roles/resourcemanager.organizationViewer", - "members": [ "user:eve@example.com" ], - "condition": { "title": "expirable access", - "description": "Does not grant access after - Sep 2020", "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", } } - - ], "etag": "BwWWja0YfJA=", "version": 3 - - } + :literal:`\` { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 }`\ \` **YAML example:** - bindings: - members: - user:\ mike@example.com - - group:\ admins@example.com - domain:google.com - - serviceAccount:\ my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - user:\ eve@example.com role: - roles/resourcemanager.organizationViewer - condition: title: expirable access description: - Does not grant access after Sep 2020 expression: - request.time < - timestamp('2020-10-01T00:00:00.000Z') etag: - BwWWja0YfJA= version: 3 + :literal:`\` bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') etag: BwWWja0YfJA= version: 3`\ \` For a description of IAM and its features, see the [IAM @@ -2402,7 +2509,7 @@ async def get_iam_policy( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.get_iam_policy, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -2457,7 +2564,7 @@ async def set_iam_policy( This corresponds to the ``resource`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2486,42 +2593,11 @@ async def set_iam_policy( **JSON example:** - { - "bindings": [ - { - "role": - "roles/resourcemanager.organizationAdmin", - "members": [ "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - - }, { "role": - "roles/resourcemanager.organizationViewer", - "members": [ "user:eve@example.com" ], - "condition": { "title": "expirable access", - "description": "Does not grant access after - Sep 2020", "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", } } - - ], "etag": "BwWWja0YfJA=", "version": 3 - - } + :literal:`\` { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 }`\ \` **YAML example:** - bindings: - members: - user:\ mike@example.com - - group:\ admins@example.com - domain:google.com - - serviceAccount:\ my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - user:\ eve@example.com role: - roles/resourcemanager.organizationViewer - condition: title: expirable access description: - Does not grant access after Sep 2020 expression: - request.time < - timestamp('2020-10-01T00:00:00.000Z') etag: - BwWWja0YfJA= version: 3 + :literal:`\` bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') etag: BwWWja0YfJA= version: 3`\ \` For a description of IAM and its features, see the [IAM @@ -2606,7 +2682,7 @@ async def test_iam_permissions( This corresponds to the ``permissions`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -2640,7 +2716,7 @@ async def test_iam_permissions( # and friendly error handling. rpc = gapic_v1.method_async.wrap_method( self._client._transport.test_iam_permissions, - default_retry=retries.Retry( + default_retry=retries.AsyncRetry( initial=1.0, maximum=60.0, multiplier=2, @@ -2671,7 +2747,7 @@ async def test_iam_permissions( # Done; return the response. return response - async def __aenter__(self): + async def __aenter__(self) -> "BigtableTableAdminAsyncClient": return self async def __aexit__(self, exc_type, exc, tb): diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/client.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/client.py index e043aa224..d0c04ed11 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/client.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/client.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -685,6 +685,7 @@ def create_table_from_snapshot( r"""Creates a new table from the specified snapshot. The target table must not exist. The snapshot and the table must be in the same instance. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1587,6 +1588,7 @@ def snapshot_table( r"""Creates a new snapshot in the specified cluster from the specified source table. The cluster and the table must be in the same instance. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -1763,6 +1765,7 @@ def get_snapshot( time. A snapshot can be used as a checkpoint for data restoration or a data source for a new table. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud @@ -1942,6 +1945,7 @@ def delete_snapshot( metadata: Sequence[Tuple[str, str]] = (), ) -> None: r"""Permanently deletes the specified snapshot. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -2407,7 +2411,7 @@ def list_backups( Returns: google.cloud.bigtable_admin_v2.services.bigtable_table_admin.pagers.ListBackupsPager: The response for - [ListBackups][google.bigtable.admin.v2.BigtableTableAdmin.ListBackups]. + [ListBackups][google.bigtable.admin.v2.BigtableTableAdmin.ListBackups]. Iterating over this object will yield results and resolve additional pages automatically. @@ -2472,9 +2476,8 @@ def restore_table( timeout: Union[float, object] = gapic_v1.method.DEFAULT, metadata: Sequence[Tuple[str, str]] = (), ) -> operation.Operation: - r"""Create a new table by restoring from a completed backup. The new - table must be in the same project as the instance containing the - backup. The returned table [long-running + r"""Create a new table by restoring from a completed backup. The + returned table [long-running operation][google.longrunning.Operation] can be used to track the progress of the operation, and to cancel it. The [metadata][google.longrunning.Operation.metadata] field type is @@ -2538,6 +2541,141 @@ def restore_table( # Done; return the response. return response + def copy_backup( + self, + request: Optional[Union[bigtable_table_admin.CopyBackupRequest, dict]] = None, + *, + parent: Optional[str] = None, + backup_id: Optional[str] = None, + source_backup: Optional[str] = None, + expire_time: Optional[timestamp_pb2.Timestamp] = None, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operation.Operation: + r"""Copy a Cloud Bigtable backup to a new backup in the + destination cluster located in the destination instance + and project. + + Args: + request (Union[google.cloud.bigtable_admin_v2.types.CopyBackupRequest, dict]): + The request object. The request for + [CopyBackup][google.bigtable.admin.v2.BigtableTableAdmin.CopyBackup]. + parent (str): + Required. The name of the destination cluster that will + contain the backup copy. The cluster must already + exists. Values are of the form: + ``projects/{project}/instances/{instance}/clusters/{cluster}``. + + This corresponds to the ``parent`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + backup_id (str): + Required. The id of the new backup. The ``backup_id`` + along with ``parent`` are combined as + {parent}/backups/{backup_id} to create the full backup + name, of the form: + ``projects/{project}/instances/{instance}/clusters/{cluster}/backups/{backup_id}``. + This string must be between 1 and 50 characters in + length and match the regex [*a-zA-Z0-9][-*.a-zA-Z0-9]*. + + This corresponds to the ``backup_id`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + source_backup (str): + Required. The source backup to be copied from. The + source backup needs to be in READY state for it to be + copied. Copying a copied backup is not allowed. Once + CopyBackup is in progress, the source backup cannot be + deleted or cleaned up on expiration until CopyBackup is + finished. Values are of the form: + ``projects//instances//clusters//backups/``. + + This corresponds to the ``source_backup`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + expire_time (google.protobuf.timestamp_pb2.Timestamp): + Required. Required. The expiration time of the copied + backup with microsecond granularity that must be at + least 6 hours and at most 30 days from the time the + request is received. Once the ``expire_time`` has + passed, Cloud Bigtable will delete the backup and free + the resources used by the backup. + + This corresponds to the ``expire_time`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.api_core.operation.Operation: + An object representing a long-running operation. + + The result type for the operation will be + :class:`google.cloud.bigtable_admin_v2.types.Backup` A + backup of a Cloud Bigtable table. + + """ + # Create or coerce a protobuf request object. + # Quick check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([parent, backup_id, source_backup, expire_time]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a bigtable_table_admin.CopyBackupRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, bigtable_table_admin.CopyBackupRequest): + request = bigtable_table_admin.CopyBackupRequest(request) + # If we have keyword arguments corresponding to fields on the + # request, apply these. + if parent is not None: + request.parent = parent + if backup_id is not None: + request.backup_id = backup_id + if source_backup is not None: + request.source_backup = source_backup + if expire_time is not None: + request.expire_time = expire_time + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.copy_backup] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("parent", request.parent),)), + ) + + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Wrap the response in an operation future. + response = operation.from_gapic( + response, + self._transport.operations_client, + table.Backup, + metadata_type=bigtable_table_admin.CopyBackupMetadata, + ) + + # Done; return the response. + return response + def get_iam_policy( self, request: Optional[Union[iam_policy_pb2.GetIamPolicyRequest, dict]] = None, @@ -2592,42 +2730,11 @@ def get_iam_policy( **JSON example:** - { - "bindings": [ - { - "role": - "roles/resourcemanager.organizationAdmin", - "members": [ "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - - }, { "role": - "roles/resourcemanager.organizationViewer", - "members": [ "user:eve@example.com" ], - "condition": { "title": "expirable access", - "description": "Does not grant access after - Sep 2020", "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", } } - - ], "etag": "BwWWja0YfJA=", "version": 3 - - } + :literal:`\` { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 }`\ \` **YAML example:** - bindings: - members: - user:\ mike@example.com - - group:\ admins@example.com - domain:google.com - - serviceAccount:\ my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - user:\ eve@example.com role: - roles/resourcemanager.organizationViewer - condition: title: expirable access description: - Does not grant access after Sep 2020 expression: - request.time < - timestamp('2020-10-01T00:00:00.000Z') etag: - BwWWja0YfJA= version: 3 + :literal:`\` bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') etag: BwWWja0YfJA= version: 3`\ \` For a description of IAM and its features, see the [IAM @@ -2728,42 +2835,11 @@ def set_iam_policy( **JSON example:** - { - "bindings": [ - { - "role": - "roles/resourcemanager.organizationAdmin", - "members": [ "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - - }, { "role": - "roles/resourcemanager.organizationViewer", - "members": [ "user:eve@example.com" ], - "condition": { "title": "expirable access", - "description": "Does not grant access after - Sep 2020", "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", } } - - ], "etag": "BwWWja0YfJA=", "version": 3 - - } + :literal:`\` { "bindings": [ { "role": "roles/resourcemanager.organizationAdmin", "members": [ "user:mike@example.com", "group:admins@example.com", "domain:google.com", "serviceAccount:my-project-id@appspot.gserviceaccount.com" ] }, { "role": "roles/resourcemanager.organizationViewer", "members": [ "user:eve@example.com" ], "condition": { "title": "expirable access", "description": "Does not grant access after Sep 2020", "expression": "request.time < timestamp('2020-10-01T00:00:00.000Z')", } } ], "etag": "BwWWja0YfJA=", "version": 3 }`\ \` **YAML example:** - bindings: - members: - user:\ mike@example.com - - group:\ admins@example.com - domain:google.com - - serviceAccount:\ my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - user:\ eve@example.com role: - roles/resourcemanager.organizationViewer - condition: title: expirable access description: - Does not grant access after Sep 2020 expression: - request.time < - timestamp('2020-10-01T00:00:00.000Z') etag: - BwWWja0YfJA= version: 3 + :literal:`\` bindings: - members: - user:mike@example.com - group:admins@example.com - domain:google.com - serviceAccount:my-project-id@appspot.gserviceaccount.com role: roles/resourcemanager.organizationAdmin - members: - user:eve@example.com role: roles/resourcemanager.organizationViewer condition: title: expirable access description: Does not grant access after Sep 2020 expression: request.time < timestamp('2020-10-01T00:00:00.000Z') etag: BwWWja0YfJA= version: 3`\ \` For a description of IAM and its features, see the [IAM diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/pagers.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/pagers.py index e639227df..331647b4c 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/pagers.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/pagers.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/__init__.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/__init__.py index 585b4e437..be4aa8d2a 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/__init__.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/base.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/base.py index cade1335b..c3cf01a96 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/base.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/base.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -322,6 +322,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=60.0, client_info=client_info, ), + self.copy_backup: gapic_v1.method.wrap_method( + self.copy_backup, + default_timeout=None, + client_info=client_info, + ), self.get_iam_policy: gapic_v1.method.wrap_method( self.get_iam_policy, default_retry=retries.Retry( @@ -577,6 +582,15 @@ def restore_table( ]: raise NotImplementedError() + @property + def copy_backup( + self, + ) -> Callable[ + [bigtable_table_admin.CopyBackupRequest], + Union[operations_pb2.Operation, Awaitable[operations_pb2.Operation]], + ]: + raise NotImplementedError() + @property def get_iam_policy( self, diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/grpc.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/grpc.py index f8cf9f834..d765869cd 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/grpc.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/grpc.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -294,6 +294,7 @@ def create_table_from_snapshot( Creates a new table from the specified snapshot. The target table must not exist. The snapshot and the table must be in the same instance. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -591,6 +592,7 @@ def snapshot_table( Creates a new snapshot in the specified cluster from the specified source table. The cluster and the table must be in the same instance. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -692,6 +694,7 @@ def delete_snapshot( r"""Return a callable for the delete snapshot method over gRPC. Permanently deletes the specified snapshot. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -866,9 +869,8 @@ def restore_table( ) -> Callable[[bigtable_table_admin.RestoreTableRequest], operations_pb2.Operation]: r"""Return a callable for the restore table method over gRPC. - Create a new table by restoring from a completed backup. The new - table must be in the same project as the instance containing the - backup. The returned table [long-running + Create a new table by restoring from a completed backup. The + returned table [long-running operation][google.longrunning.Operation] can be used to track the progress of the operation, and to cancel it. The [metadata][google.longrunning.Operation.metadata] field type is @@ -894,6 +896,34 @@ def restore_table( ) return self._stubs["restore_table"] + @property + def copy_backup( + self, + ) -> Callable[[bigtable_table_admin.CopyBackupRequest], operations_pb2.Operation]: + r"""Return a callable for the copy backup method over gRPC. + + Copy a Cloud Bigtable backup to a new backup in the + destination cluster located in the destination instance + and project. + + Returns: + Callable[[~.CopyBackupRequest], + ~.Operation]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "copy_backup" not in self._stubs: + self._stubs["copy_backup"] = self.grpc_channel.unary_unary( + "/google.bigtable.admin.v2.BigtableTableAdmin/CopyBackup", + request_serializer=bigtable_table_admin.CopyBackupRequest.serialize, + response_deserializer=operations_pb2.Operation.FromString, + ) + return self._stubs["copy_backup"] + @property def get_iam_policy( self, diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/grpc_asyncio.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/grpc_asyncio.py index 54eb7e524..b60a7351c 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/grpc_asyncio.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/grpc_asyncio.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -302,6 +302,7 @@ def create_table_from_snapshot( Creates a new table from the specified snapshot. The target table must not exist. The snapshot and the table must be in the same instance. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -607,6 +608,7 @@ def snapshot_table( Creates a new snapshot in the specified cluster from the specified source table. The cluster and the table must be in the same instance. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -710,6 +712,7 @@ def delete_snapshot( r"""Return a callable for the delete snapshot method over gRPC. Permanently deletes the specified snapshot. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be @@ -890,9 +893,8 @@ def restore_table( ]: r"""Return a callable for the restore table method over gRPC. - Create a new table by restoring from a completed backup. The new - table must be in the same project as the instance containing the - backup. The returned table [long-running + Create a new table by restoring from a completed backup. The + returned table [long-running operation][google.longrunning.Operation] can be used to track the progress of the operation, and to cancel it. The [metadata][google.longrunning.Operation.metadata] field type is @@ -918,6 +920,36 @@ def restore_table( ) return self._stubs["restore_table"] + @property + def copy_backup( + self, + ) -> Callable[ + [bigtable_table_admin.CopyBackupRequest], Awaitable[operations_pb2.Operation] + ]: + r"""Return a callable for the copy backup method over gRPC. + + Copy a Cloud Bigtable backup to a new backup in the + destination cluster located in the destination instance + and project. + + Returns: + Callable[[~.CopyBackupRequest], + Awaitable[~.Operation]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "copy_backup" not in self._stubs: + self._stubs["copy_backup"] = self.grpc_channel.unary_unary( + "/google.bigtable.admin.v2.BigtableTableAdmin/CopyBackup", + request_serializer=bigtable_table_admin.CopyBackupRequest.serialize, + response_deserializer=operations_pb2.Operation.FromString, + ) + return self._stubs["copy_backup"] + @property def get_iam_policy( self, diff --git a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/rest.py b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/rest.py index 4d5b2ed1c..41b893eb7 100644 --- a/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/rest.py +++ b/google/cloud/bigtable_admin_v2/services/bigtable_table_admin/transports/rest.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -45,8 +45,8 @@ from google.cloud.bigtable_admin_v2.types import table as gba_table from google.iam.v1 import iam_policy_pb2 # type: ignore from google.iam.v1 import policy_pb2 # type: ignore -from google.longrunning import operations_pb2 # type: ignore from google.protobuf import empty_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore from .base import ( BigtableTableAdminTransport, @@ -84,6 +84,14 @@ def post_check_consistency(self, response): logging.log(f"Received response: {response}") return response + def pre_copy_backup(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + def post_copy_backup(self, response): + logging.log(f"Received response: {response}") + return response + def pre_create_backup(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -281,6 +289,29 @@ def post_check_consistency( """ return response + def pre_copy_backup( + self, + request: bigtable_table_admin.CopyBackupRequest, + metadata: Sequence[Tuple[str, str]], + ) -> Tuple[bigtable_table_admin.CopyBackupRequest, Sequence[Tuple[str, str]]]: + """Pre-rpc interceptor for copy_backup + + Override in a subclass to manipulate the request or metadata + before they are sent to the BigtableTableAdmin server. + """ + return request, metadata + + def post_copy_backup( + self, response: operations_pb2.Operation + ) -> operations_pb2.Operation: + """Post-rpc interceptor for copy_backup + + Override in a subclass to manipulate the response + after it is returned by the BigtableTableAdmin server but before + it is returned to user code. + """ + return response + def pre_create_backup( self, request: bigtable_table_admin.CreateBackupRequest, @@ -1010,6 +1041,103 @@ def __call__( resp = self._interceptor.post_check_consistency(resp) return resp + class _CopyBackup(BigtableTableAdminRestStub): + def __hash__(self): + return hash("CopyBackup") + + __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = {} + + @classmethod + def _get_unset_required_fields(cls, message_dict): + return { + k: v + for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items() + if k not in message_dict + } + + def __call__( + self, + request: bigtable_table_admin.CopyBackupRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operations_pb2.Operation: + r"""Call the copy backup method over HTTP. + + Args: + request (~.bigtable_table_admin.CopyBackupRequest): + The request object. The request for + [CopyBackup][google.bigtable.admin.v2.BigtableTableAdmin.CopyBackup]. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + ~.operations_pb2.Operation: + This resource represents a + long-running operation that is the + result of a network API call. + + """ + + http_options: List[Dict[str, str]] = [ + { + "method": "post", + "uri": "/v2/{parent=projects/*/instances/*/clusters/*}/backups:copy", + "body": "*", + }, + ] + request, metadata = self._interceptor.pre_copy_backup(request, metadata) + pb_request = bigtable_table_admin.CopyBackupRequest.pb(request) + transcoded_request = path_template.transcode(http_options, pb_request) + + # Jsonify the request body + + body = json_format.MessageToJson( + transcoded_request["body"], + including_default_value_fields=False, + use_integers_for_enums=True, + ) + uri = transcoded_request["uri"] + method = transcoded_request["method"] + + # Jsonify the query params + query_params = json.loads( + json_format.MessageToJson( + transcoded_request["query_params"], + including_default_value_fields=False, + use_integers_for_enums=True, + ) + ) + query_params.update(self._get_unset_required_fields(query_params)) + + query_params["$alt"] = "json;enum-encoding=int" + + # Send the request + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = getattr(self._session, method)( + "{host}{uri}".format(host=self._host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + data=body, + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + # Return the response + resp = operations_pb2.Operation() + json_format.Parse(response.content, resp, ignore_unknown_fields=True) + resp = self._interceptor.post_copy_backup(resp) + return resp + class _CreateBackup(BigtableTableAdminRestStub): def __hash__(self): return hash("CreateBackup") @@ -1881,54 +2009,54 @@ def __call__( :: - { - "bindings": [ - { - "role": "roles/resourcemanager.organizationAdmin", - "members": [ - "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - }, - { - "role": "roles/resourcemanager.organizationViewer", - "members": [ - "user:eve@example.com" - ], - "condition": { - "title": "expirable access", - "description": "Does not grant access after Sep 2020", - "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", - } - } - ], - "etag": "BwWWja0YfJA=", - "version": 3 - } + { + "bindings": [ + { + "role": "roles/resourcemanager.organizationAdmin", + "members": [ + "user:mike@example.com", + "group:admins@example.com", + "domain:google.com", + "serviceAccount:my-project-id@appspot.gserviceaccount.com" + ] + }, + { + "role": "roles/resourcemanager.organizationViewer", + "members": [ + "user:eve@example.com" + ], + "condition": { + "title": "expirable access", + "description": "Does not grant access after Sep 2020", + "expression": "request.time < + timestamp('2020-10-01T00:00:00.000Z')", + } + } + ], + "etag": "BwWWja0YfJA=", + "version": 3 + } **YAML example:** :: - bindings: - - members: - - user:mike@example.com - - group:admins@example.com - - domain:google.com - - serviceAccount:my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - - user:eve@example.com - role: roles/resourcemanager.organizationViewer - condition: - title: expirable access - description: Does not grant access after Sep 2020 - expression: request.time < timestamp('2020-10-01T00:00:00.000Z') - etag: BwWWja0YfJA= - version: 3 + bindings: + - members: + - user:mike@example.com + - group:admins@example.com + - domain:google.com + - serviceAccount:my-project-id@appspot.gserviceaccount.com + role: roles/resourcemanager.organizationAdmin + - members: + - user:eve@example.com + role: roles/resourcemanager.organizationViewer + condition: + title: expirable access + description: Does not grant access after Sep 2020 + expression: request.time < timestamp('2020-10-01T00:00:00.000Z') + etag: BwWWja0YfJA= + version: 3 For a description of IAM and its features, see the `IAM documentation `__. @@ -2044,6 +2172,7 @@ def __call__( time. A snapshot can be used as a checkpoint for data restoration or a data source for a new table. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud @@ -2733,54 +2862,54 @@ def __call__( :: - { - "bindings": [ - { - "role": "roles/resourcemanager.organizationAdmin", - "members": [ - "user:mike@example.com", - "group:admins@example.com", - "domain:google.com", - "serviceAccount:my-project-id@appspot.gserviceaccount.com" - ] - }, - { - "role": "roles/resourcemanager.organizationViewer", - "members": [ - "user:eve@example.com" - ], - "condition": { - "title": "expirable access", - "description": "Does not grant access after Sep 2020", - "expression": "request.time < - timestamp('2020-10-01T00:00:00.000Z')", - } - } - ], - "etag": "BwWWja0YfJA=", - "version": 3 - } + { + "bindings": [ + { + "role": "roles/resourcemanager.organizationAdmin", + "members": [ + "user:mike@example.com", + "group:admins@example.com", + "domain:google.com", + "serviceAccount:my-project-id@appspot.gserviceaccount.com" + ] + }, + { + "role": "roles/resourcemanager.organizationViewer", + "members": [ + "user:eve@example.com" + ], + "condition": { + "title": "expirable access", + "description": "Does not grant access after Sep 2020", + "expression": "request.time < + timestamp('2020-10-01T00:00:00.000Z')", + } + } + ], + "etag": "BwWWja0YfJA=", + "version": 3 + } **YAML example:** :: - bindings: - - members: - - user:mike@example.com - - group:admins@example.com - - domain:google.com - - serviceAccount:my-project-id@appspot.gserviceaccount.com - role: roles/resourcemanager.organizationAdmin - - members: - - user:eve@example.com - role: roles/resourcemanager.organizationViewer - condition: - title: expirable access - description: Does not grant access after Sep 2020 - expression: request.time < timestamp('2020-10-01T00:00:00.000Z') - etag: BwWWja0YfJA= - version: 3 + bindings: + - members: + - user:mike@example.com + - group:admins@example.com + - domain:google.com + - serviceAccount:my-project-id@appspot.gserviceaccount.com + role: roles/resourcemanager.organizationAdmin + - members: + - user:eve@example.com + role: roles/resourcemanager.organizationViewer + condition: + title: expirable access + description: Does not grant access after Sep 2020 + expression: request.time < timestamp('2020-10-01T00:00:00.000Z') + etag: BwWWja0YfJA= + version: 3 For a description of IAM and its features, see the `IAM documentation `__. @@ -3360,6 +3489,14 @@ def check_consistency( # In C++ this would require a dynamic_cast return self._CheckConsistency(self._session, self._host, self._interceptor) # type: ignore + @property + def copy_backup( + self, + ) -> Callable[[bigtable_table_admin.CopyBackupRequest], operations_pb2.Operation]: + # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here. + # In C++ this would require a dynamic_cast + return self._CopyBackup(self._session, self._host, self._interceptor) # type: ignore + @property def create_backup( self, diff --git a/google/cloud/bigtable_admin_v2/types/__init__.py b/google/cloud/bigtable_admin_v2/types/__init__.py index 69153c9fc..a2fefffc8 100644 --- a/google/cloud/bigtable_admin_v2/types/__init__.py +++ b/google/cloud/bigtable_admin_v2/types/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -44,6 +44,8 @@ from .bigtable_table_admin import ( CheckConsistencyRequest, CheckConsistencyResponse, + CopyBackupMetadata, + CopyBackupRequest, CreateBackupMetadata, CreateBackupRequest, CreateTableFromSnapshotMetadata, @@ -130,6 +132,8 @@ "UpdateInstanceMetadata", "CheckConsistencyRequest", "CheckConsistencyResponse", + "CopyBackupMetadata", + "CopyBackupRequest", "CreateBackupMetadata", "CreateBackupRequest", "CreateTableFromSnapshotMetadata", diff --git a/google/cloud/bigtable_admin_v2/types/bigtable_instance_admin.py b/google/cloud/bigtable_admin_v2/types/bigtable_instance_admin.py index a22543354..87332a351 100644 --- a/google/cloud/bigtable_admin_v2/types/bigtable_instance_admin.py +++ b/google/cloud/bigtable_admin_v2/types/bigtable_instance_admin.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/types/bigtable_table_admin.py b/google/cloud/bigtable_admin_v2/types/bigtable_table_admin.py index 4c4b9e9e2..6a3b31a1e 100644 --- a/google/cloud/bigtable_admin_v2/types/bigtable_table_admin.py +++ b/google/cloud/bigtable_admin_v2/types/bigtable_table_admin.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -62,6 +62,8 @@ "DeleteBackupRequest", "ListBackupsRequest", "ListBackupsResponse", + "CopyBackupRequest", + "CopyBackupMetadata", }, ) @@ -76,8 +78,7 @@ class RestoreTableRequest(proto.Message): Attributes: parent (str): Required. The name of the instance in which to create the - restored table. This instance must be in the same project as - the source backup. Values are of the form + restored table. Values are of the form ``projects//instances/``. table_id (str): Required. The id of the table to create and restore to. This @@ -359,7 +360,7 @@ class ListTablesRequest(proto.Message): should be listed. Values are of the form ``projects/{project}/instances/{instance}``. view (google.cloud.bigtable_admin_v2.types.Table.View): - The view to be applied to the returned tables' fields. Only + The view to be applied to the returned tables' fields. NAME_ONLY view (default) and REPLICATION_VIEW are supported. page_size (int): Maximum number of results per page. @@ -917,6 +918,7 @@ class DeleteSnapshotRequest(proto.Message): class SnapshotTableMetadata(proto.Message): r"""The metadata for the Operation returned by SnapshotTable. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be changed in @@ -1192,8 +1194,15 @@ class ListBackupsRequest(proto.Message): fields in [Backup][google.bigtable.admin.v2.Backup]. The full syntax is described at https://aip.dev/132#ordering. - Fields supported are: \* name \* source_table \* expire_time - \* start_time \* end_time \* size_bytes \* state + Fields supported are: + + - name + - source_table + - expire_time + - start_time + - end_time + - size_bytes + - state For example, "start_time". The default sorting order is ascending. To specify descending order for the field, a @@ -1266,4 +1275,90 @@ def raw_page(self): ) +class CopyBackupRequest(proto.Message): + r"""The request for + [CopyBackup][google.bigtable.admin.v2.BigtableTableAdmin.CopyBackup]. + + Attributes: + parent (str): + Required. The name of the destination cluster that will + contain the backup copy. The cluster must already exists. + Values are of the form: + ``projects/{project}/instances/{instance}/clusters/{cluster}``. + backup_id (str): + Required. The id of the new backup. The ``backup_id`` along + with ``parent`` are combined as {parent}/backups/{backup_id} + to create the full backup name, of the form: + ``projects/{project}/instances/{instance}/clusters/{cluster}/backups/{backup_id}``. + This string must be between 1 and 50 characters in length + and match the regex [*a-zA-Z0-9][-*.a-zA-Z0-9]*. + source_backup (str): + Required. The source backup to be copied from. The source + backup needs to be in READY state for it to be copied. + Copying a copied backup is not allowed. Once CopyBackup is + in progress, the source backup cannot be deleted or cleaned + up on expiration until CopyBackup is finished. Values are of + the form: + ``projects//instances//clusters//backups/``. + expire_time (google.protobuf.timestamp_pb2.Timestamp): + Required. Required. The expiration time of the copied backup + with microsecond granularity that must be at least 6 hours + and at most 30 days from the time the request is received. + Once the ``expire_time`` has passed, Cloud Bigtable will + delete the backup and free the resources used by the backup. + """ + + parent: str = proto.Field( + proto.STRING, + number=1, + ) + backup_id: str = proto.Field( + proto.STRING, + number=2, + ) + source_backup: str = proto.Field( + proto.STRING, + number=3, + ) + expire_time: timestamp_pb2.Timestamp = proto.Field( + proto.MESSAGE, + number=4, + message=timestamp_pb2.Timestamp, + ) + + +class CopyBackupMetadata(proto.Message): + r"""Metadata type for the google.longrunning.Operation returned by + [CopyBackup][google.bigtable.admin.v2.BigtableTableAdmin.CopyBackup]. + + Attributes: + name (str): + The name of the backup being created through the copy + operation. Values are of the form + ``projects//instances//clusters//backups/``. + source_backup_info (google.cloud.bigtable_admin_v2.types.BackupInfo): + Information about the source backup that is + being copied from. + progress (google.cloud.bigtable_admin_v2.types.OperationProgress): + The progress of the + [CopyBackup][google.bigtable.admin.v2.BigtableTableAdmin.CopyBackup] + operation. + """ + + name: str = proto.Field( + proto.STRING, + number=1, + ) + source_backup_info: gba_table.BackupInfo = proto.Field( + proto.MESSAGE, + number=2, + message=gba_table.BackupInfo, + ) + progress: common.OperationProgress = proto.Field( + proto.MESSAGE, + number=3, + message=common.OperationProgress, + ) + + __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/bigtable_admin_v2/types/common.py b/google/cloud/bigtable_admin_v2/types/common.py index 2cc71fc43..959b9deb1 100644 --- a/google/cloud/bigtable_admin_v2/types/common.py +++ b/google/cloud/bigtable_admin_v2/types/common.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_admin_v2/types/instance.py b/google/cloud/bigtable_admin_v2/types/instance.py index 2b5d81636..78efd711b 100644 --- a/google/cloud/bigtable_admin_v2/types/instance.py +++ b/google/cloud/bigtable_admin_v2/types/instance.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -173,7 +173,7 @@ class AutoscalingTargets(proto.Message): The storage utilization that the Autoscaler should be trying to achieve. This number is limited between 2560 (2.5TiB) and 5120 (5TiB) for a SSD cluster and between 8192 (8TiB) and - 16384 (16TiB) for an HDD cluster; otherwise it will return + 16384 (16TiB) for an HDD cluster, otherwise it will return INVALID_ARGUMENT error. If this value is set to 0, it will be treated as if it were set to the default value: 2560 for SSD, 8192 for HDD. @@ -419,8 +419,43 @@ class AppProfile(proto.Message): Use a single-cluster routing policy. This field is a member of `oneof`_ ``routing_policy``. + priority (google.cloud.bigtable_admin_v2.types.AppProfile.Priority): + This field has been deprecated in favor of + ``standard_isolation.priority``. If you set this field, + ``standard_isolation.priority`` will be set instead. + + The priority of requests sent using this app profile. + + This field is a member of `oneof`_ ``isolation``. + standard_isolation (google.cloud.bigtable_admin_v2.types.AppProfile.StandardIsolation): + The standard options used for isolating this + app profile's traffic from other use cases. + + This field is a member of `oneof`_ ``isolation``. """ + class Priority(proto.Enum): + r"""Possible priorities for an app profile. Note that higher + priority writes can sometimes queue behind lower priority writes + to the same tablet, as writes must be strictly sequenced in the + durability log. + + Values: + PRIORITY_UNSPECIFIED (0): + Default value. Mapped to PRIORITY_HIGH (the legacy behavior) + on creation. + PRIORITY_LOW (1): + No description available. + PRIORITY_MEDIUM (2): + No description available. + PRIORITY_HIGH (3): + No description available. + """ + PRIORITY_UNSPECIFIED = 0 + PRIORITY_LOW = 1 + PRIORITY_MEDIUM = 2 + PRIORITY_HIGH = 3 + class MultiClusterRoutingUseAny(proto.Message): r"""Read/write requests are routed to the nearest cluster in the instance, and will fail over to the nearest cluster that is @@ -466,6 +501,22 @@ class SingleClusterRouting(proto.Message): number=2, ) + class StandardIsolation(proto.Message): + r"""Standard options for isolating this app profile's traffic + from other use cases. + + Attributes: + priority (google.cloud.bigtable_admin_v2.types.AppProfile.Priority): + The priority of requests sent using this app + profile. + """ + + priority: "AppProfile.Priority" = proto.Field( + proto.ENUM, + number=1, + enum="AppProfile.Priority", + ) + name: str = proto.Field( proto.STRING, number=1, @@ -490,6 +541,18 @@ class SingleClusterRouting(proto.Message): oneof="routing_policy", message=SingleClusterRouting, ) + priority: Priority = proto.Field( + proto.ENUM, + number=7, + oneof="isolation", + enum=Priority, + ) + standard_isolation: StandardIsolation = proto.Field( + proto.MESSAGE, + number=11, + oneof="isolation", + message=StandardIsolation, + ) class HotTablet(proto.Message): diff --git a/google/cloud/bigtable_admin_v2/types/table.py b/google/cloud/bigtable_admin_v2/types/table.py index 16d136e16..57bd1b00f 100644 --- a/google/cloud/bigtable_admin_v2/types/table.py +++ b/google/cloud/bigtable_admin_v2/types/table.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -124,7 +124,8 @@ class Table(proto.Message): ``REPLICATION_VIEW``, ``ENCRYPTION_VIEW``, ``FULL`` column_families (MutableMapping[str, google.cloud.bigtable_admin_v2.types.ColumnFamily]): The column families configured for this table, mapped by - column family ID. Views: ``SCHEMA_VIEW``, ``FULL`` + column family ID. Views: ``SCHEMA_VIEW``, ``STATS_VIEW``, + ``FULL`` granularity (google.cloud.bigtable_admin_v2.types.Table.TimestampGranularity): Immutable. The granularity (i.e. ``MILLIS``) at which timestamps are stored in this table. Timestamps not matching @@ -141,14 +142,16 @@ class Table(proto.Message): this table. Otherwise, the change stream is disabled and the change stream is not retained. deletion_protection (bool): - Set to true to make the table protected - against data loss. i.e. deleting the following - resources through Admin APIs are prohibited: - - The table. - - The column families in the table. - - The instance containing the table. - Note one can still delete the data stored in the - table through Data APIs. + Set to true to make the table protected against data loss. + i.e. deleting the following resources through Admin APIs are + prohibited: + + - The table. + - The column families in the table. + - The instance containing the table. + + Note one can still delete the data stored in the table + through Data APIs. """ class TimestampGranularity(proto.Enum): @@ -308,6 +311,7 @@ class ColumnFamily(proto.Message): gc_rule (google.cloud.bigtable_admin_v2.types.GcRule): Garbage collection rule specified as a protobuf. Must serialize to at most 500 bytes. + NOTE: Garbage collection executes opportunistically in the background, and so it's possible for reads to return a cell even if it @@ -478,6 +482,7 @@ class Snapshot(proto.Message): r"""A snapshot of a table at a particular time. A snapshot can be used as a checkpoint for data restoration or a data source for a new table. + Note: This is a private alpha release of Cloud Bigtable snapshots. This feature is not currently available to most Cloud Bigtable customers. This feature might be changed in @@ -486,8 +491,7 @@ class Snapshot(proto.Message): Attributes: name (str): - Output only. The unique name of the snapshot. Values are of - the form + The unique name of the snapshot. Values are of the form ``projects/{project}/instances/{instance}/clusters/{cluster}/snapshots/{snapshot}``. source_table (google.cloud.bigtable_admin_v2.types.Table): Output only. The source table at the time the @@ -502,16 +506,15 @@ class Snapshot(proto.Message): Output only. The time when the snapshot is created. delete_time (google.protobuf.timestamp_pb2.Timestamp): - Output only. The time when the snapshot will - be deleted. The maximum amount of time a - snapshot can stay active is 365 days. If 'ttl' - is not specified, the default maximum of 365 - days will be used. + The time when the snapshot will be deleted. + The maximum amount of time a snapshot can stay + active is 365 days. If 'ttl' is not specified, + the default maximum of 365 days will be used. state (google.cloud.bigtable_admin_v2.types.Snapshot.State): Output only. The current state of the snapshot. description (str): - Output only. Description of the snapshot. + Description of the snapshot. """ class State(proto.Enum): @@ -587,10 +590,16 @@ class Backup(proto.Message): backup was created. This needs to be in the same instance as the backup. Values are of the form ``projects/{project}/instances/{instance}/tables/{source_table}``. + source_backup (str): + Output only. Name of the backup from which + this backup was copied. If a backup is not + created by copying a backup, this field will be + empty. Values are of the form: + projects//instances//backups/. expire_time (google.protobuf.timestamp_pb2.Timestamp): Required. The expiration time of the backup, with microseconds granularity that must be at least 6 hours and - at most 30 days from the time the request is received. Once + at most 90 days from the time the request is received. Once the ``expire_time`` has passed, Cloud Bigtable will delete the backup and free the resources used by the backup. start_time (google.protobuf.timestamp_pb2.Timestamp): @@ -636,6 +645,10 @@ class State(proto.Enum): proto.STRING, number=2, ) + source_backup: str = proto.Field( + proto.STRING, + number=10, + ) expire_time: timestamp_pb2.Timestamp = proto.Field( proto.MESSAGE, number=3, @@ -684,6 +697,12 @@ class BackupInfo(proto.Message): source_table (str): Output only. Name of the table the backup was created from. + source_backup (str): + Output only. Name of the backup from which + this backup was copied. If a backup is not + created by copying a backup, this field will be + empty. Values are of the form: + projects//instances//backups/. """ backup: str = proto.Field( @@ -704,6 +723,10 @@ class BackupInfo(proto.Message): proto.STRING, number=4, ) + source_backup: str = proto.Field( + proto.STRING, + number=10, + ) __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/bigtable_v2/__init__.py b/google/cloud/bigtable_v2/__init__.py index ee3bd8c0c..80bd4ec09 100644 --- a/google/cloud/bigtable_v2/__init__.py +++ b/google/cloud/bigtable_v2/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/gapic_version.py b/google/cloud/bigtable_v2/gapic_version.py index 0f1a446f3..03d6d0200 100644 --- a/google/cloud/bigtable_v2/gapic_version.py +++ b/google/cloud/bigtable_v2/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "2.19.0" # {x-release-please-version} +__version__ = "2.22.0" # {x-release-please-version} diff --git a/google/cloud/bigtable_v2/services/__init__.py b/google/cloud/bigtable_v2/services/__init__.py index e8e1c3845..89a37dc92 100644 --- a/google/cloud/bigtable_v2/services/__init__.py +++ b/google/cloud/bigtable_v2/services/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/services/bigtable/__init__.py b/google/cloud/bigtable_v2/services/bigtable/__init__.py index cfce7b6b8..f10a68e5b 100644 --- a/google/cloud/bigtable_v2/services/bigtable/__init__.py +++ b/google/cloud/bigtable_v2/services/bigtable/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/services/bigtable/async_client.py b/google/cloud/bigtable_v2/services/bigtable/async_client.py index d325564c0..e497ff25b 100644 --- a/google/cloud/bigtable_v2/services/bigtable/async_client.py +++ b/google/cloud/bigtable_v2/services/bigtable/async_client.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -41,9 +41,9 @@ from google.oauth2 import service_account # type: ignore try: - OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None] + OptionalRetry = Union[retries.AsyncRetry, gapic_v1.method._MethodDefault, None] except AttributeError: # pragma: NO COVER - OptionalRetry = Union[retries.Retry, object, None] # type: ignore + OptionalRetry = Union[retries.AsyncRetry, object, None] # type: ignore from google.cloud.bigtable_v2.types import bigtable from google.cloud.bigtable_v2.types import data @@ -251,7 +251,7 @@ def read_rows( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -344,7 +344,7 @@ def sample_row_keys( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -454,7 +454,7 @@ async def mutate_row( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -565,7 +565,7 @@ def mutate_rows( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -703,7 +703,7 @@ async def check_and_mutate_row( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -812,7 +812,7 @@ async def ping_and_warm( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -929,7 +929,7 @@ async def read_modify_write_row( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1034,7 +1034,7 @@ def generate_initial_change_stream_partitions( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1136,7 +1136,7 @@ def read_change_stream( This corresponds to the ``app_profile_id`` field on the ``request`` instance; if ``request`` is provided, this should not be set. - retry (google.api_core.retry.Retry): Designation of what errors, if any, + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, should be retried. timeout (float): The timeout for this request. metadata (Sequence[Tuple[str, str]]): Strings which should be @@ -1196,7 +1196,7 @@ def read_change_stream( # Done; return the response. return response - async def __aenter__(self): + async def __aenter__(self) -> "BigtableAsyncClient": return self async def __aexit__(self, exc_type, exc, tb): diff --git a/google/cloud/bigtable_v2/services/bigtable/client.py b/google/cloud/bigtable_v2/services/bigtable/client.py index 1c2e7b822..54ba6af43 100644 --- a/google/cloud/bigtable_v2/services/bigtable/client.py +++ b/google/cloud/bigtable_v2/services/bigtable/client.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/__init__.py b/google/cloud/bigtable_v2/services/bigtable/transports/__init__.py index e8796bb8c..6a9eb0e58 100644 --- a/google/cloud/bigtable_v2/services/bigtable/transports/__init__.py +++ b/google/cloud/bigtable_v2/services/bigtable/transports/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/base.py b/google/cloud/bigtable_v2/services/bigtable/transports/base.py index 5b4580c18..b580bbca7 100644 --- a/google/cloud/bigtable_v2/services/bigtable/transports/base.py +++ b/google/cloud/bigtable_v2/services/bigtable/transports/base.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/grpc.py b/google/cloud/bigtable_v2/services/bigtable/transports/grpc.py index b9e073e8a..8ba04e761 100644 --- a/google/cloud/bigtable_v2/services/bigtable/transports/grpc.py +++ b/google/cloud/bigtable_v2/services/bigtable/transports/grpc.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py b/google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py index 3450d4969..1d0a2bc4c 100644 --- a/google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py +++ b/google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/services/bigtable/transports/rest.py b/google/cloud/bigtable_v2/services/bigtable/transports/rest.py index 4343fbb90..31d230f94 100644 --- a/google/cloud/bigtable_v2/services/bigtable/transports/rest.py +++ b/google/cloud/bigtable_v2/services/bigtable/transports/rest.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/types/__init__.py b/google/cloud/bigtable_v2/types/__init__.py index 9f15efaf5..f266becb9 100644 --- a/google/cloud/bigtable_v2/types/__init__.py +++ b/google/cloud/bigtable_v2/types/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/google/cloud/bigtable_v2/types/bigtable.py b/google/cloud/bigtable_v2/types/bigtable.py index 13f6ac0db..57f806408 100644 --- a/google/cloud/bigtable_v2/types/bigtable.py +++ b/google/cloud/bigtable_v2/types/bigtable.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -80,6 +80,21 @@ class ReadRowsRequest(proto.Message): request_stats_view (google.cloud.bigtable_v2.types.ReadRowsRequest.RequestStatsView): The view into RequestStats, as described above. + reversed (bool): + Experimental API - Please note that this API is currently + experimental and can change in the future. + + Return rows in lexiographical descending order of the row + keys. The row contents will not be affected by this flag. + + Example result set: + + :: + + [ + {key: "k2", "f:col1": "v1", "f:col2": "v1"}, + {key: "k1", "f:col1": "v2", "f:col2": "v2"} + ] """ class RequestStatsView(proto.Enum): @@ -131,6 +146,10 @@ class RequestStatsView(proto.Enum): number=6, enum=RequestStatsView, ) + reversed: bool = proto.Field( + proto.BOOL, + number=7, + ) class ReadRowsResponse(proto.Message): diff --git a/google/cloud/bigtable_v2/types/data.py b/google/cloud/bigtable_v2/types/data.py index 515e167df..e37644a76 100644 --- a/google/cloud/bigtable_v2/types/data.py +++ b/google/cloud/bigtable_v2/types/data.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -697,6 +697,7 @@ class Condition(proto.Message): r"""A RowFilter which evaluates one of two possible RowFilters, depending on whether or not a predicate RowFilter outputs any cells from the input row. + IMPORTANT NOTE: The predicate filter does not execute atomically with the true and false filters, which may lead to inconsistent or unexpected results. Additionally, Condition filters have poor diff --git a/google/cloud/bigtable_v2/types/feature_flags.py b/google/cloud/bigtable_v2/types/feature_flags.py index 1b5f76e24..92ac5023d 100644 --- a/google/cloud/bigtable_v2/types/feature_flags.py +++ b/google/cloud/bigtable_v2/types/feature_flags.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -29,26 +29,54 @@ class FeatureFlags(proto.Message): - r"""Feature flags supported by a client. This is intended to be sent as - part of request metadata to assure the server that certain behaviors - are safe to enable. This proto is meant to be serialized and - websafe-base64 encoded under the ``bigtable-features`` metadata key. - The value will remain constant for the lifetime of a client and due - to HTTP2's HPACK compression, the request overhead will be tiny. - This is an internal implementation detail and should not be used by - endusers directly. + r"""Feature flags supported or enabled by a client. This is intended to + be sent as part of request metadata to assure the server that + certain behaviors are safe to enable. This proto is meant to be + serialized and websafe-base64 encoded under the + ``bigtable-features`` metadata key. The value will remain constant + for the lifetime of a client and due to HTTP2's HPACK compression, + the request overhead will be tiny. This is an internal + implementation detail and should not be used by end users directly. Attributes: + reverse_scans (bool): + Notify the server that the client supports + reverse scans. The server will reject + ReadRowsRequests with the reverse bit set when + this is absent. mutate_rows_rate_limit (bool): Notify the server that the client enables batch write flow control by requesting - RateLimitInfo from MutateRowsResponse. + RateLimitInfo from MutateRowsResponse. Due to + technical reasons, this disables partial + retries. + mutate_rows_rate_limit2 (bool): + Notify the server that the client enables + batch write flow control by requesting + RateLimitInfo from MutateRowsResponse. With + partial retries enabled. + last_scanned_row_responses (bool): + Notify the server that the client supports the + last_scanned_row field in ReadRowsResponse for long-running + scans. """ + reverse_scans: bool = proto.Field( + proto.BOOL, + number=1, + ) mutate_rows_rate_limit: bool = proto.Field( proto.BOOL, number=3, ) + mutate_rows_rate_limit2: bool = proto.Field( + proto.BOOL, + number=5, + ) + last_scanned_row_responses: bool = proto.Field( + proto.BOOL, + number=4, + ) __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/bigtable_v2/types/request_stats.py b/google/cloud/bigtable_v2/types/request_stats.py index d72ba8694..61cce9491 100644 --- a/google/cloud/bigtable_v2/types/request_stats.py +++ b/google/cloud/bigtable_v2/types/request_stats.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -86,6 +86,7 @@ class RequestLatencyStats(proto.Message): response. For more context on the component that is measuring this latency, see: https://cloud.google.com/bigtable/docs/overview + Note: This value may be slightly shorter than the value reported into aggregate latency metrics in Monitoring for this request diff --git a/google/cloud/bigtable_v2/types/response_params.py b/google/cloud/bigtable_v2/types/response_params.py index 2532e64e2..98e3a67db 100644 --- a/google/cloud/bigtable_v2/types/response_params.py +++ b/google/cloud/bigtable_v2/types/response_params.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/mypy.ini b/mypy.ini index 3a17a37c6..31cc24223 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1,5 +1,5 @@ [mypy] -python_version = 3.7 +python_version = 3.8 namespace_packages = True exclude = tests/unit/gapic/ diff --git a/noxfile.py b/noxfile.py index 2e053ffcf..49e40ddf6 100644 --- a/noxfile.py +++ b/noxfile.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # -# Copyright 2018 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,21 +17,24 @@ # Generated by synthtool. DO NOT EDIT! from __future__ import absolute_import + import os import pathlib import re import shutil +from typing import Dict, List import warnings import nox -BLACK_VERSION = "black==22.3.0" -ISORT_VERSION = "isort==5.10.1" +FLAKE8_VERSION = "flake8==6.1.0" +BLACK_VERSION = "black[jupyter]==23.7.0" +ISORT_VERSION = "isort==5.11.0" LINT_PATHS = ["docs", "google", "tests", "noxfile.py", "setup.py"] DEFAULT_PYTHON_VERSION = "3.8" -UNIT_TEST_PYTHON_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +UNIT_TEST_PYTHON_VERSIONS: List[str] = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] UNIT_TEST_STANDARD_DEPENDENCIES = [ "mock", "asyncmock", @@ -39,25 +42,24 @@ "pytest-cov", "pytest-asyncio", ] -UNIT_TEST_EXTERNAL_DEPENDENCIES = [] -UNIT_TEST_LOCAL_DEPENDENCIES = [] -UNIT_TEST_DEPENDENCIES = [] -UNIT_TEST_EXTRAS = [] -UNIT_TEST_EXTRAS_BY_PYTHON = {} - -SYSTEM_TEST_PYTHON_VERSIONS = ["3.8"] -SYSTEM_TEST_STANDARD_DEPENDENCIES = [ +UNIT_TEST_EXTERNAL_DEPENDENCIES: List[str] = [] +UNIT_TEST_LOCAL_DEPENDENCIES: List[str] = [] +UNIT_TEST_DEPENDENCIES: List[str] = [] +UNIT_TEST_EXTRAS: List[str] = [] +UNIT_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = {} + +SYSTEM_TEST_PYTHON_VERSIONS: List[str] = ["3.8"] +SYSTEM_TEST_STANDARD_DEPENDENCIES: List[str] = [ "mock", "pytest", "pytest-asyncio", "google-cloud-testutils", ] -SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [] -SYSTEM_TEST_LOCAL_DEPENDENCIES = [] -UNIT_TEST_DEPENDENCIES = [] -SYSTEM_TEST_DEPENDENCIES = [] -SYSTEM_TEST_EXTRAS = [] -SYSTEM_TEST_EXTRAS_BY_PYTHON = {} +SYSTEM_TEST_EXTERNAL_DEPENDENCIES: List[str] = [] +SYSTEM_TEST_LOCAL_DEPENDENCIES: List[str] = [] +SYSTEM_TEST_DEPENDENCIES: List[str] = [] +SYSTEM_TEST_EXTRAS: List[str] = [] +SYSTEM_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = {} CURRENT_DIRECTORY = pathlib.Path(__file__).parent.absolute() @@ -72,6 +74,7 @@ "lint_setup_py", "blacken", "docs", + "format", ] # Error if a python version is missing @@ -85,7 +88,7 @@ def lint(session): Returns a failure if the linters find linting errors or sufficiently serious code quality issues. """ - session.install("flake8", BLACK_VERSION) + session.install(FLAKE8_VERSION, BLACK_VERSION) session.run( "black", "--check", @@ -218,7 +221,6 @@ def unit(session): def install_systemtest_dependencies(session, *constraints): - # Use pre-release gRPC for system tests. # Exclude version 1.52.0rc1 which has a known issue. # See https://github.com/grpc/grpc/issues/32163 @@ -378,7 +380,7 @@ def docs(session): ) -@nox.session(python="3.9") +@nox.session(python="3.10") def docfx(session): """Build the docfx yaml files for this library.""" @@ -458,6 +460,7 @@ def prerelease_deps(session): "grpcio!=1.52.0rc1", "grpcio-status", "google-api-core==2.16.0rc0", # TODO: remove pin once streaming retries is merged + "google-auth", "proto-plus", "google-cloud-testutils", # dependencies of google-cloud-testutils" @@ -470,7 +473,6 @@ def prerelease_deps(session): # Remaining dependencies other_deps = [ "requests", - "google-auth", ] session.install(*other_deps) @@ -479,6 +481,7 @@ def prerelease_deps(session): "python", "-c", "import google.protobuf; print(google.protobuf.__version__)" ) session.run("python", "-c", "import grpc; print(grpc.__version__)") + session.run("python", "-c", "import google.auth; print(google.auth.__version__)") session.run("py.test", "tests/unit") diff --git a/owlbot.py b/owlbot.py index 9ca859fb9..be28fa2a1 100644 --- a/owlbot.py +++ b/owlbot.py @@ -176,8 +176,6 @@ def mypy(session): "--warn-unreachable", "--disallow-any-generics", "--exclude", - "google/cloud/bigtable/deprecated", - "--exclude", "tests/system/v2_client", "--exclude", "tests/unit/v2_client", diff --git a/samples/beam/noxfile.py b/samples/beam/noxfile.py index 3d4395024..80ffdb178 100644 --- a/samples/beam/noxfile.py +++ b/samples/beam/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/beam/requirements-test.txt b/samples/beam/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/beam/requirements-test.txt +++ b/samples/beam/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/beam/requirements.txt b/samples/beam/requirements.txt index 8be9b98e0..813fc8d2b 100644 --- a/samples/beam/requirements.txt +++ b/samples/beam/requirements.txt @@ -1,3 +1,3 @@ -apache-beam==2.46.0 -google-cloud-bigtable==2.17.0 -google-cloud-core==2.3.2 +apache-beam==2.52.0 +google-cloud-bigtable==2.22.0 +google-cloud-core==2.4.1 diff --git a/samples/hello/noxfile.py b/samples/hello/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/hello/noxfile.py +++ b/samples/hello/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/hello/requirements-test.txt b/samples/hello/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/hello/requirements-test.txt +++ b/samples/hello/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/hello/requirements.txt b/samples/hello/requirements.txt index 199541ffe..68419fbcb 100644 --- a/samples/hello/requirements.txt +++ b/samples/hello/requirements.txt @@ -1,2 +1,2 @@ -google-cloud-bigtable==2.17.0 -google-cloud-core==2.3.2 +google-cloud-bigtable==2.22.0 +google-cloud-core==2.4.1 diff --git a/samples/hello_happybase/noxfile.py b/samples/hello_happybase/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/hello_happybase/noxfile.py +++ b/samples/hello_happybase/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/hello_happybase/requirements-test.txt b/samples/hello_happybase/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/hello_happybase/requirements-test.txt +++ b/samples/hello_happybase/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/hello_happybase/requirements.txt b/samples/hello_happybase/requirements.txt index a144f03e1..d3368cd0f 100644 --- a/samples/hello_happybase/requirements.txt +++ b/samples/hello_happybase/requirements.txt @@ -1 +1,2 @@ google-cloud-happybase==0.33.0 +six==1.16.0 # See https://github.com/googleapis/google-cloud-python-happybase/issues/128 diff --git a/samples/instanceadmin/noxfile.py b/samples/instanceadmin/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/instanceadmin/noxfile.py +++ b/samples/instanceadmin/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/instanceadmin/requirements-test.txt b/samples/instanceadmin/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/instanceadmin/requirements-test.txt +++ b/samples/instanceadmin/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/instanceadmin/requirements.txt b/samples/instanceadmin/requirements.txt index 04e476254..a01a0943c 100644 --- a/samples/instanceadmin/requirements.txt +++ b/samples/instanceadmin/requirements.txt @@ -1,2 +1,2 @@ -google-cloud-bigtable==2.17.0 +google-cloud-bigtable==2.22.0 backoff==2.2.1 diff --git a/samples/metricscaler/noxfile.py b/samples/metricscaler/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/metricscaler/noxfile.py +++ b/samples/metricscaler/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/metricscaler/requirements-test.txt b/samples/metricscaler/requirements-test.txt index 761227068..80ef7d3d0 100644 --- a/samples/metricscaler/requirements-test.txt +++ b/samples/metricscaler/requirements-test.txt @@ -1,3 +1,3 @@ -pytest==7.3.1 -mock==5.0.2 +pytest==7.4.3 +mock==5.1.0 google-cloud-testutils diff --git a/samples/metricscaler/requirements.txt b/samples/metricscaler/requirements.txt index 02e08b4c8..38c355ce3 100644 --- a/samples/metricscaler/requirements.txt +++ b/samples/metricscaler/requirements.txt @@ -1,2 +1,2 @@ -google-cloud-bigtable==2.17.0 -google-cloud-monitoring==2.14.2 +google-cloud-bigtable==2.22.0 +google-cloud-monitoring==2.18.0 diff --git a/samples/quickstart/noxfile.py b/samples/quickstart/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/quickstart/noxfile.py +++ b/samples/quickstart/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/quickstart/requirements-test.txt b/samples/quickstart/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/quickstart/requirements-test.txt +++ b/samples/quickstart/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/quickstart/requirements.txt b/samples/quickstart/requirements.txt index 909f8c365..6dc985893 100644 --- a/samples/quickstart/requirements.txt +++ b/samples/quickstart/requirements.txt @@ -1 +1 @@ -google-cloud-bigtable==2.17.0 +google-cloud-bigtable==2.22.0 diff --git a/samples/quickstart_happybase/noxfile.py b/samples/quickstart_happybase/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/quickstart_happybase/noxfile.py +++ b/samples/quickstart_happybase/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/quickstart_happybase/requirements-test.txt b/samples/quickstart_happybase/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/quickstart_happybase/requirements-test.txt +++ b/samples/quickstart_happybase/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/quickstart_happybase/requirements.txt b/samples/quickstart_happybase/requirements.txt index a144f03e1..d3368cd0f 100644 --- a/samples/quickstart_happybase/requirements.txt +++ b/samples/quickstart_happybase/requirements.txt @@ -1 +1,2 @@ google-cloud-happybase==0.33.0 +six==1.16.0 # See https://github.com/googleapis/google-cloud-python-happybase/issues/128 diff --git a/samples/snippets/deletes/noxfile.py b/samples/snippets/deletes/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/snippets/deletes/noxfile.py +++ b/samples/snippets/deletes/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/snippets/deletes/requirements-test.txt b/samples/snippets/deletes/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/snippets/deletes/requirements-test.txt +++ b/samples/snippets/deletes/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/snippets/deletes/requirements.txt b/samples/snippets/deletes/requirements.txt index 200665631..ae10593d2 100644 --- a/samples/snippets/deletes/requirements.txt +++ b/samples/snippets/deletes/requirements.txt @@ -1,2 +1,2 @@ -google-cloud-bigtable==2.17.0 +google-cloud-bigtable==2.22.0 snapshottest==0.6.0 \ No newline at end of file diff --git a/samples/snippets/filters/noxfile.py b/samples/snippets/filters/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/snippets/filters/noxfile.py +++ b/samples/snippets/filters/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/snippets/filters/requirements-test.txt b/samples/snippets/filters/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/snippets/filters/requirements-test.txt +++ b/samples/snippets/filters/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/snippets/filters/requirements.txt b/samples/snippets/filters/requirements.txt index 200665631..ae10593d2 100644 --- a/samples/snippets/filters/requirements.txt +++ b/samples/snippets/filters/requirements.txt @@ -1,2 +1,2 @@ -google-cloud-bigtable==2.17.0 +google-cloud-bigtable==2.22.0 snapshottest==0.6.0 \ No newline at end of file diff --git a/samples/snippets/reads/noxfile.py b/samples/snippets/reads/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/snippets/reads/noxfile.py +++ b/samples/snippets/reads/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/snippets/reads/requirements-test.txt b/samples/snippets/reads/requirements-test.txt index c4d04a08d..f9708e4b7 100644 --- a/samples/snippets/reads/requirements-test.txt +++ b/samples/snippets/reads/requirements-test.txt @@ -1 +1 @@ -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/snippets/reads/requirements.txt b/samples/snippets/reads/requirements.txt index 200665631..ae10593d2 100644 --- a/samples/snippets/reads/requirements.txt +++ b/samples/snippets/reads/requirements.txt @@ -1,2 +1,2 @@ -google-cloud-bigtable==2.17.0 +google-cloud-bigtable==2.22.0 snapshottest==0.6.0 \ No newline at end of file diff --git a/samples/snippets/writes/noxfile.py b/samples/snippets/writes/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/snippets/writes/noxfile.py +++ b/samples/snippets/writes/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/snippets/writes/requirements-test.txt b/samples/snippets/writes/requirements-test.txt index 96aa71dab..908e344b5 100644 --- a/samples/snippets/writes/requirements-test.txt +++ b/samples/snippets/writes/requirements-test.txt @@ -1,2 +1,2 @@ backoff==2.2.1 -pytest==7.3.1 +pytest==7.4.3 diff --git a/samples/snippets/writes/requirements.txt b/samples/snippets/writes/requirements.txt index 32cead029..07b0a191d 100644 --- a/samples/snippets/writes/requirements.txt +++ b/samples/snippets/writes/requirements.txt @@ -1 +1 @@ -google-cloud-bigtable==2.17.0 \ No newline at end of file +google-cloud-bigtable==2.22.0 \ No newline at end of file diff --git a/samples/tableadmin/noxfile.py b/samples/tableadmin/noxfile.py index 7c8a63994..483b55901 100644 --- a/samples/tableadmin/noxfile.py +++ b/samples/tableadmin/noxfile.py @@ -89,7 +89,7 @@ def get_pytest_env_vars() -> Dict[str, str]: # DO NOT EDIT - automatically generated. # All versions used to test samples. -ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11"] +ALL_VERSIONS = ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12"] # Any default versions that should be ignored. IGNORED_VERSIONS = TEST_CONFIG["ignored_versions"] diff --git a/samples/tableadmin/requirements-test.txt b/samples/tableadmin/requirements-test.txt index ca1f33bd3..39d590005 100644 --- a/samples/tableadmin/requirements-test.txt +++ b/samples/tableadmin/requirements-test.txt @@ -1,2 +1,2 @@ -pytest==7.3.1 -google-cloud-testutils==1.3.3 +pytest==7.4.3 +google-cloud-testutils==1.4.0 diff --git a/samples/tableadmin/requirements.txt b/samples/tableadmin/requirements.txt index 909f8c365..6dc985893 100644 --- a/samples/tableadmin/requirements.txt +++ b/samples/tableadmin/requirements.txt @@ -1 +1 @@ -google-cloud-bigtable==2.17.0 +google-cloud-bigtable==2.22.0 diff --git a/scripts/decrypt-secrets.sh b/scripts/decrypt-secrets.sh index 21f6d2a26..0018b421d 100755 --- a/scripts/decrypt-secrets.sh +++ b/scripts/decrypt-secrets.sh @@ -1,6 +1,6 @@ #!/bin/bash -# Copyright 2015 Google Inc. All rights reserved. +# Copyright 2023 Google LLC All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/scripts/fixup_bigtable_admin_v2_keywords.py b/scripts/fixup_bigtable_admin_v2_keywords.py index 17be56f2f..6882feaf6 100644 --- a/scripts/fixup_bigtable_admin_v2_keywords.py +++ b/scripts/fixup_bigtable_admin_v2_keywords.py @@ -1,6 +1,6 @@ #! /usr/bin/env python3 # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -40,6 +40,7 @@ class bigtable_adminCallTransformer(cst.CSTTransformer): CTRL_PARAMS: Tuple[str] = ('retry', 'timeout', 'metadata') METHOD_TO_PARAMS: Dict[str, Tuple[str]] = { 'check_consistency': ('name', 'consistency_token', ), + 'copy_backup': ('parent', 'backup_id', 'source_backup', 'expire_time', ), 'create_app_profile': ('parent', 'app_profile_id', 'app_profile', 'ignore_warnings', ), 'create_backup': ('parent', 'backup_id', 'backup', ), 'create_cluster': ('parent', 'cluster_id', 'cluster', ), diff --git a/scripts/fixup_bigtable_v2_keywords.py b/scripts/fixup_bigtable_v2_keywords.py index 11ffed53f..8d32e5b70 100644 --- a/scripts/fixup_bigtable_v2_keywords.py +++ b/scripts/fixup_bigtable_v2_keywords.py @@ -1,6 +1,6 @@ #! /usr/bin/env python3 # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,7 +46,7 @@ class bigtableCallTransformer(cst.CSTTransformer): 'ping_and_warm': ('name', 'app_profile_id', ), 'read_change_stream': ('table_name', 'app_profile_id', 'partition', 'start_time', 'continuation_tokens', 'end_time', 'heartbeat_duration', ), 'read_modify_write_row': ('table_name', 'row_key', 'rules', 'app_profile_id', ), - 'read_rows': ('table_name', 'app_profile_id', 'rows', 'filter', 'rows_limit', 'request_stats_view', ), + 'read_rows': ('table_name', 'app_profile_id', 'rows', 'filter', 'rows_limit', 'request_stats_view', 'reversed', ), 'sample_row_keys': ('table_name', 'app_profile_id', ), } diff --git a/scripts/readme-gen/readme_gen.py b/scripts/readme-gen/readme_gen.py index 91b59676b..1acc11983 100644 --- a/scripts/readme-gen/readme_gen.py +++ b/scripts/readme-gen/readme_gen.py @@ -1,6 +1,6 @@ #!/usr/bin/env python -# Copyright 2016 Google Inc +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,17 +33,17 @@ autoescape=True, ) -README_TMPL = jinja_env.get_template('README.tmpl.rst') +README_TMPL = jinja_env.get_template("README.tmpl.rst") def get_help(file): - return subprocess.check_output(['python', file, '--help']).decode() + return subprocess.check_output(["python", file, "--help"]).decode() def main(): parser = argparse.ArgumentParser() - parser.add_argument('source') - parser.add_argument('--destination', default='README.rst') + parser.add_argument("source") + parser.add_argument("--destination", default="README.rst") args = parser.parse_args() @@ -51,9 +51,9 @@ def main(): root = os.path.dirname(source) destination = os.path.join(root, args.destination) - jinja_env.globals['get_help'] = get_help + jinja_env.globals["get_help"] = get_help - with io.open(source, 'r') as f: + with io.open(source, "r") as f: config = yaml.load(f) # This allows get_help to execute in the right directory. @@ -61,9 +61,9 @@ def main(): output = README_TMPL.render(config) - with io.open(destination, 'w') as f: + with io.open(destination, "w") as f: f.write(output) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/setup.cfg b/setup.cfg index c3a2b39f6..052350089 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- # -# Copyright 2020 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/setup.py b/setup.py index 0bce3a5d6..3aeb275ae 100644 --- a/setup.py +++ b/setup.py @@ -37,8 +37,8 @@ # 'Development Status :: 5 - Production/Stable' release_status = "Development Status :: 5 - Production/Stable" dependencies = [ - "google-api-core[grpc] >= 2.16.0rc0", - "google-cloud-core >= 1.4.4, <3.0.0dev", + "google-api-core[grpc] >= 2.16.0rc0, <3.0.0dev,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,!=2.8.*,!=2.9.*,!=2.10.*", + "google-cloud-core >= 2.0.0, <3.0.0dev", "grpc-google-iam-v1 >= 0.12.4, <1.0.0dev", "proto-plus >= 1.22.0, <2.0.0dev", "proto-plus >= 1.22.2, <2.0.0dev; python_version>='3.11'", @@ -63,7 +63,6 @@ if package.startswith("google") ] - setuptools.setup( name=name, version=version, @@ -83,6 +82,7 @@ "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Operating System :: OS Independent", "Topic :: Internet", ], diff --git a/testing/constraints-3.12.txt b/testing/constraints-3.12.txt new file mode 100644 index 000000000..e69de29bb diff --git a/testing/constraints-3.7.txt b/testing/constraints-3.7.txt index 83bfe4577..b87fca3e6 100644 --- a/testing/constraints-3.7.txt +++ b/testing/constraints-3.7.txt @@ -6,7 +6,7 @@ # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 google-api-core==2.16.0rc0 -google-cloud-core==2.3.2 +google-cloud-core==2.0.0 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 libcst==0.2.5 diff --git a/testing/constraints-3.8.txt b/testing/constraints-3.8.txt index 505ba9934..42a424b6a 100644 --- a/testing/constraints-3.8.txt +++ b/testing/constraints-3.8.txt @@ -6,7 +6,7 @@ # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 google-api-core==2.16.0rc0 -google-cloud-core==2.3.2 +google-cloud-core==2.0.0 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 libcst==0.2.5 diff --git a/tests/__init__.py b/tests/__init__.py index e8e1c3845..89a37dc92 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/system/v2_client/test_data_api.py b/tests/system/v2_client/test_data_api.py index 2ca7e1504..579837e34 100644 --- a/tests/system/v2_client/test_data_api.py +++ b/tests/system/v2_client/test_data_api.py @@ -381,3 +381,39 @@ def test_access_with_non_admin_client(data_client, data_instance_id, data_table_ instance = data_client.instance(data_instance_id) table = instance.table(data_table_id) assert table.read_row("nonesuch") is None # no raise + + +def test_mutations_batcher_threading(data_table, rows_to_delete): + """ + Test the mutations batcher by sending a bunch of mutations using different + flush methods + """ + import mock + import time + from google.cloud.bigtable.batcher import MutationsBatcher + + num_sent = 20 + all_results = [] + + def callback(results): + all_results.extend(results) + + # override flow control max elements + with mock.patch("google.cloud.bigtable.batcher.MAX_OUTSTANDING_ELEMENTS", 2): + with MutationsBatcher( + data_table, + flush_count=5, + flush_interval=0.07, + batch_completed_callback=callback, + ) as batcher: + # send mutations in a way that timed flushes and count flushes interleave + for i in range(num_sent): + row = data_table.direct_row("row{}".format(i)) + row.set_cell( + COLUMN_FAMILY_ID1, COL_NAME1, "val{}".format(i).encode("utf-8") + ) + rows_to_delete.append(row) + batcher.mutate(row) + time.sleep(0.01) + # ensure all mutations were sent + assert len(all_results) == num_sent diff --git a/tests/system/v2_client/test_instance_admin.py b/tests/system/v2_client/test_instance_admin.py index e5e311213..bd5c7e912 100644 --- a/tests/system/v2_client/test_instance_admin.py +++ b/tests/system/v2_client/test_instance_admin.py @@ -28,7 +28,6 @@ def _create_app_profile_helper( allow_transactional_writes=None, ignore_warnings=None, ): - app_profile = instance.app_profile( app_profile_id=app_profile_id, routing_policy_type=routing_policy_type, diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 60a305bcb..46080e497 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1461,7 +1461,6 @@ async def test_read_rows_query_matches_request(self, include_app_profile): @pytest.mark.parametrize("operation_timeout", [0.001, 0.023, 0.1]) @pytest.mark.asyncio async def test_read_rows_timeout(self, operation_timeout): - async with self._make_table() as table: read_rows = table.client._gapic_client.read_rows query = ReadRowsQuery() diff --git a/tests/unit/data/test_exceptions.py b/tests/unit/data/test_exceptions.py index 2defffc86..bc921717e 100644 --- a/tests/unit/data/test_exceptions.py +++ b/tests/unit/data/test_exceptions.py @@ -457,7 +457,6 @@ def _get_class(self): return FailedMutationEntryError def _make_one(self, idx=9, entry=mock.Mock(), cause=RuntimeError("mock")): - return self._get_class()(idx, entry, cause) def test_raise(self): @@ -516,7 +515,6 @@ def _get_class(self): return FailedQueryShardError def _make_one(self, idx=9, query=mock.Mock(), cause=RuntimeError("mock")): - return self._get_class()(idx, query, cause) def test_raise(self): diff --git a/tests/unit/gapic/__init__.py b/tests/unit/gapic/__init__.py index e8e1c3845..89a37dc92 100644 --- a/tests/unit/gapic/__init__.py +++ b/tests/unit/gapic/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/unit/gapic/bigtable_admin_v2/__init__.py b/tests/unit/gapic/bigtable_admin_v2/__init__.py index e8e1c3845..89a37dc92 100644 --- a/tests/unit/gapic/bigtable_admin_v2/__init__.py +++ b/tests/unit/gapic/bigtable_admin_v2/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/unit/gapic/bigtable_admin_v2/test_bigtable_instance_admin.py b/tests/unit/gapic/bigtable_admin_v2/test_bigtable_instance_admin.py index 76715f1ed..ddbf0032f 100644 --- a/tests/unit/gapic/bigtable_admin_v2/test_bigtable_instance_admin.py +++ b/tests/unit/gapic/bigtable_admin_v2/test_bigtable_instance_admin.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -63,7 +63,7 @@ from google.iam.v1 import iam_policy_pb2 # type: ignore from google.iam.v1 import options_pb2 # type: ignore from google.iam.v1 import policy_pb2 # type: ignore -from google.longrunning import operations_pb2 +from google.longrunning import operations_pb2 # type: ignore from google.oauth2 import service_account from google.protobuf import field_mask_pb2 # type: ignore from google.protobuf import timestamp_pb2 # type: ignore @@ -2425,11 +2425,6 @@ def test_get_cluster(request_type, transport: str = "grpc"): state=instance.Cluster.State.READY, serve_nodes=1181, default_storage_type=common.StorageType.SSD, - cluster_config=instance.Cluster.ClusterConfig( - cluster_autoscaling_config=instance.Cluster.ClusterAutoscalingConfig( - autoscaling_limits=instance.AutoscalingLimits(min_serve_nodes=1600) - ) - ), ) response = client.get_cluster(request) @@ -3529,9 +3524,7 @@ def test_create_app_profile(request_type, transport: str = "grpc"): name="name_value", etag="etag_value", description="description_value", - multi_cluster_routing_use_any=instance.AppProfile.MultiClusterRoutingUseAny( - cluster_ids=["cluster_ids_value"] - ), + priority=instance.AppProfile.Priority.PRIORITY_LOW, ) response = client.create_app_profile(request) @@ -3801,9 +3794,7 @@ def test_get_app_profile(request_type, transport: str = "grpc"): name="name_value", etag="etag_value", description="description_value", - multi_cluster_routing_use_any=instance.AppProfile.MultiClusterRoutingUseAny( - cluster_ids=["cluster_ids_value"] - ), + priority=instance.AppProfile.Priority.PRIORITY_LOW, ) response = client.get_app_profile(request) @@ -4456,9 +4447,11 @@ async def test_list_app_profiles_async_pages(): RuntimeError, ) pages = [] - async for page_ in ( + # Workaround issue in python 3.9 related to code coverage by adding `# pragma: no branch` + # See https://github.com/googleapis/gapic-generator-python/pull/1174#issuecomment-1025132372 + async for page_ in ( # pragma: no branch await client.list_app_profiles(request={}) - ).pages: # pragma: no branch + ).pages: pages.append(page_) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token @@ -6138,9 +6131,11 @@ async def test_list_hot_tablets_async_pages(): RuntimeError, ) pages = [] - async for page_ in ( + # Workaround issue in python 3.9 related to code coverage by adding `# pragma: no branch` + # See https://github.com/googleapis/gapic-generator-python/pull/1174#issuecomment-1025132372 + async for page_ in ( # pragma: no branch await client.list_hot_tablets(request={}) - ).pages: # pragma: no branch + ).pages: pages.append(page_) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token @@ -6459,8 +6454,9 @@ def test_get_instance_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.Instance.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.Instance.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -6539,8 +6535,9 @@ def test_get_instance_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = instance.Instance.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.Instance.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -6663,8 +6660,9 @@ def test_get_instance_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.Instance.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.Instance.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -6728,8 +6726,9 @@ def test_list_instances_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListInstancesResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListInstancesResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -6809,10 +6808,11 @@ def test_list_instances_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListInstancesResponse.pb( + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListInstancesResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -6939,8 +6939,9 @@ def test_list_instances_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListInstancesResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListInstancesResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7007,8 +7008,9 @@ def test_update_instance_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.Instance.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.Instance.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7086,8 +7088,9 @@ def test_update_instance_rest_required_fields(request_type=instance.Instance): response_value = Response() response_value.status_code = 200 - pb_return_value = instance.Instance.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.Instance.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7215,6 +7218,75 @@ def test_partial_update_instance_rest(request_type): "create_time": {"seconds": 751, "nanos": 543}, "satisfies_pzs": True, } + # The version of a generated dependency at test runtime may differ from the version used during generation. + # Delete any fields which are not present in the current runtime dependency + # See https://github.com/googleapis/gapic-generator-python/issues/1748 + + # Determine if the message type is proto-plus or protobuf + test_field = bigtable_instance_admin.PartialUpdateInstanceRequest.meta.fields[ + "instance" + ] + + def get_message_fields(field): + # Given a field which is a message (composite type), return a list with + # all the fields of the message. + # If the field is not a composite type, return an empty list. + message_fields = [] + + if hasattr(field, "message") and field.message: + is_field_type_proto_plus_type = not hasattr(field.message, "DESCRIPTOR") + + if is_field_type_proto_plus_type: + message_fields = field.message.meta.fields.values() + # Add `# pragma: NO COVER` because there may not be any `*_pb2` field types + else: # pragma: NO COVER + message_fields = field.message.DESCRIPTOR.fields + return message_fields + + runtime_nested_fields = [ + (field.name, nested_field.name) + for field in get_message_fields(test_field) + for nested_field in get_message_fields(field) + ] + + subfields_not_in_runtime = [] + + # For each item in the sample request, create a list of sub fields which are not present at runtime + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for field, value in request_init["instance"].items(): # pragma: NO COVER + result = None + is_repeated = False + # For repeated fields + if isinstance(value, list) and len(value): + is_repeated = True + result = value[0] + # For fields where the type is another message + if isinstance(value, dict): + result = value + + if result and hasattr(result, "keys"): + for subfield in result.keys(): + if (field, subfield) not in runtime_nested_fields: + subfields_not_in_runtime.append( + { + "field": field, + "subfield": subfield, + "is_repeated": is_repeated, + } + ) + + # Remove fields from the sample request which are not present in the runtime version of the dependency + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for subfield_to_delete in subfields_not_in_runtime: # pragma: NO COVER + field = subfield_to_delete.get("field") + field_repeated = subfield_to_delete.get("is_repeated") + subfield = subfield_to_delete.get("subfield") + if subfield: + if field_repeated: + for i in range(0, len(request_init["instance"][field])): + del request_init["instance"][field][i][subfield] + else: + del request_init["instance"][field][subfield] request = request_type(**request_init) # Mock the http request call within the method and fake a response. @@ -7395,15 +7467,6 @@ def test_partial_update_instance_rest_bad_request( # send a request that will satisfy transcoding request_init = {"instance": {"name": "projects/sample1/instances/sample2"}} - request_init["instance"] = { - "name": "projects/sample1/instances/sample2", - "display_name": "display_name_value", - "state": 1, - "type_": 1, - "labels": {}, - "create_time": {"seconds": 751, "nanos": 543}, - "satisfies_pzs": True, - } request = request_type(**request_init) # Mock the http request call within the method and fake a BadRequest error. @@ -7766,6 +7829,73 @@ def test_create_cluster_rest(request_type): "default_storage_type": 1, "encryption_config": {"kms_key_name": "kms_key_name_value"}, } + # The version of a generated dependency at test runtime may differ from the version used during generation. + # Delete any fields which are not present in the current runtime dependency + # See https://github.com/googleapis/gapic-generator-python/issues/1748 + + # Determine if the message type is proto-plus or protobuf + test_field = bigtable_instance_admin.CreateClusterRequest.meta.fields["cluster"] + + def get_message_fields(field): + # Given a field which is a message (composite type), return a list with + # all the fields of the message. + # If the field is not a composite type, return an empty list. + message_fields = [] + + if hasattr(field, "message") and field.message: + is_field_type_proto_plus_type = not hasattr(field.message, "DESCRIPTOR") + + if is_field_type_proto_plus_type: + message_fields = field.message.meta.fields.values() + # Add `# pragma: NO COVER` because there may not be any `*_pb2` field types + else: # pragma: NO COVER + message_fields = field.message.DESCRIPTOR.fields + return message_fields + + runtime_nested_fields = [ + (field.name, nested_field.name) + for field in get_message_fields(test_field) + for nested_field in get_message_fields(field) + ] + + subfields_not_in_runtime = [] + + # For each item in the sample request, create a list of sub fields which are not present at runtime + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for field, value in request_init["cluster"].items(): # pragma: NO COVER + result = None + is_repeated = False + # For repeated fields + if isinstance(value, list) and len(value): + is_repeated = True + result = value[0] + # For fields where the type is another message + if isinstance(value, dict): + result = value + + if result and hasattr(result, "keys"): + for subfield in result.keys(): + if (field, subfield) not in runtime_nested_fields: + subfields_not_in_runtime.append( + { + "field": field, + "subfield": subfield, + "is_repeated": is_repeated, + } + ) + + # Remove fields from the sample request which are not present in the runtime version of the dependency + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for subfield_to_delete in subfields_not_in_runtime: # pragma: NO COVER + field = subfield_to_delete.get("field") + field_repeated = subfield_to_delete.get("is_repeated") + subfield = subfield_to_delete.get("subfield") + if subfield: + if field_repeated: + for i in range(0, len(request_init["cluster"][field])): + del request_init["cluster"][field][i][subfield] + else: + del request_init["cluster"][field][subfield] request = request_type(**request_init) # Mock the http request call within the method and fake a response. @@ -7964,26 +8094,6 @@ def test_create_cluster_rest_bad_request( # send a request that will satisfy transcoding request_init = {"parent": "projects/sample1/instances/sample2"} - request_init["cluster"] = { - "name": "name_value", - "location": "location_value", - "state": 1, - "serve_nodes": 1181, - "cluster_config": { - "cluster_autoscaling_config": { - "autoscaling_limits": { - "min_serve_nodes": 1600, - "max_serve_nodes": 1602, - }, - "autoscaling_targets": { - "cpu_utilization_percent": 2483, - "storage_utilization_gib_per_node": 3404, - }, - } - }, - "default_storage_type": 1, - "encryption_config": {"kms_key_name": "kms_key_name_value"}, - } request = request_type(**request_init) # Mock the http request call within the method and fake a BadRequest error. @@ -8088,18 +8198,14 @@ def test_get_cluster_rest(request_type): state=instance.Cluster.State.READY, serve_nodes=1181, default_storage_type=common.StorageType.SSD, - cluster_config=instance.Cluster.ClusterConfig( - cluster_autoscaling_config=instance.Cluster.ClusterAutoscalingConfig( - autoscaling_limits=instance.AutoscalingLimits(min_serve_nodes=1600) - ) - ), ) # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.Cluster.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.Cluster.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8178,8 +8284,9 @@ def test_get_cluster_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = instance.Cluster.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.Cluster.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8302,8 +8409,9 @@ def test_get_cluster_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.Cluster.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.Cluster.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8368,8 +8476,9 @@ def test_list_clusters_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListClustersResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListClustersResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8449,10 +8558,9 @@ def test_list_clusters_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListClustersResponse.pb( - return_value - ) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListClustersResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8579,8 +8687,9 @@ def test_list_clusters_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListClustersResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListClustersResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8776,6 +8885,75 @@ def test_partial_update_cluster_rest(request_type): "default_storage_type": 1, "encryption_config": {"kms_key_name": "kms_key_name_value"}, } + # The version of a generated dependency at test runtime may differ from the version used during generation. + # Delete any fields which are not present in the current runtime dependency + # See https://github.com/googleapis/gapic-generator-python/issues/1748 + + # Determine if the message type is proto-plus or protobuf + test_field = bigtable_instance_admin.PartialUpdateClusterRequest.meta.fields[ + "cluster" + ] + + def get_message_fields(field): + # Given a field which is a message (composite type), return a list with + # all the fields of the message. + # If the field is not a composite type, return an empty list. + message_fields = [] + + if hasattr(field, "message") and field.message: + is_field_type_proto_plus_type = not hasattr(field.message, "DESCRIPTOR") + + if is_field_type_proto_plus_type: + message_fields = field.message.meta.fields.values() + # Add `# pragma: NO COVER` because there may not be any `*_pb2` field types + else: # pragma: NO COVER + message_fields = field.message.DESCRIPTOR.fields + return message_fields + + runtime_nested_fields = [ + (field.name, nested_field.name) + for field in get_message_fields(test_field) + for nested_field in get_message_fields(field) + ] + + subfields_not_in_runtime = [] + + # For each item in the sample request, create a list of sub fields which are not present at runtime + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for field, value in request_init["cluster"].items(): # pragma: NO COVER + result = None + is_repeated = False + # For repeated fields + if isinstance(value, list) and len(value): + is_repeated = True + result = value[0] + # For fields where the type is another message + if isinstance(value, dict): + result = value + + if result and hasattr(result, "keys"): + for subfield in result.keys(): + if (field, subfield) not in runtime_nested_fields: + subfields_not_in_runtime.append( + { + "field": field, + "subfield": subfield, + "is_repeated": is_repeated, + } + ) + + # Remove fields from the sample request which are not present in the runtime version of the dependency + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for subfield_to_delete in subfields_not_in_runtime: # pragma: NO COVER + field = subfield_to_delete.get("field") + field_repeated = subfield_to_delete.get("is_repeated") + subfield = subfield_to_delete.get("subfield") + if subfield: + if field_repeated: + for i in range(0, len(request_init["cluster"][field])): + del request_init["cluster"][field][i][subfield] + else: + del request_init["cluster"][field][subfield] request = request_type(**request_init) # Mock the http request call within the method and fake a response. @@ -8958,26 +9136,6 @@ def test_partial_update_cluster_rest_bad_request( request_init = { "cluster": {"name": "projects/sample1/instances/sample2/clusters/sample3"} } - request_init["cluster"] = { - "name": "projects/sample1/instances/sample2/clusters/sample3", - "location": "location_value", - "state": 1, - "serve_nodes": 1181, - "cluster_config": { - "cluster_autoscaling_config": { - "autoscaling_limits": { - "min_serve_nodes": 1600, - "max_serve_nodes": 1602, - }, - "autoscaling_targets": { - "cpu_utilization_percent": 2483, - "storage_utilization_gib_per_node": 3404, - }, - } - }, - "default_storage_type": 1, - "encryption_config": {"kms_key_name": "kms_key_name_value"}, - } request = request_type(**request_init) # Mock the http request call within the method and fake a BadRequest error. @@ -9335,7 +9493,78 @@ def test_create_app_profile_rest(request_type): "cluster_id": "cluster_id_value", "allow_transactional_writes": True, }, + "priority": 1, + "standard_isolation": {"priority": 1}, } + # The version of a generated dependency at test runtime may differ from the version used during generation. + # Delete any fields which are not present in the current runtime dependency + # See https://github.com/googleapis/gapic-generator-python/issues/1748 + + # Determine if the message type is proto-plus or protobuf + test_field = bigtable_instance_admin.CreateAppProfileRequest.meta.fields[ + "app_profile" + ] + + def get_message_fields(field): + # Given a field which is a message (composite type), return a list with + # all the fields of the message. + # If the field is not a composite type, return an empty list. + message_fields = [] + + if hasattr(field, "message") and field.message: + is_field_type_proto_plus_type = not hasattr(field.message, "DESCRIPTOR") + + if is_field_type_proto_plus_type: + message_fields = field.message.meta.fields.values() + # Add `# pragma: NO COVER` because there may not be any `*_pb2` field types + else: # pragma: NO COVER + message_fields = field.message.DESCRIPTOR.fields + return message_fields + + runtime_nested_fields = [ + (field.name, nested_field.name) + for field in get_message_fields(test_field) + for nested_field in get_message_fields(field) + ] + + subfields_not_in_runtime = [] + + # For each item in the sample request, create a list of sub fields which are not present at runtime + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for field, value in request_init["app_profile"].items(): # pragma: NO COVER + result = None + is_repeated = False + # For repeated fields + if isinstance(value, list) and len(value): + is_repeated = True + result = value[0] + # For fields where the type is another message + if isinstance(value, dict): + result = value + + if result and hasattr(result, "keys"): + for subfield in result.keys(): + if (field, subfield) not in runtime_nested_fields: + subfields_not_in_runtime.append( + { + "field": field, + "subfield": subfield, + "is_repeated": is_repeated, + } + ) + + # Remove fields from the sample request which are not present in the runtime version of the dependency + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for subfield_to_delete in subfields_not_in_runtime: # pragma: NO COVER + field = subfield_to_delete.get("field") + field_repeated = subfield_to_delete.get("is_repeated") + subfield = subfield_to_delete.get("subfield") + if subfield: + if field_repeated: + for i in range(0, len(request_init["app_profile"][field])): + del request_init["app_profile"][field][i][subfield] + else: + del request_init["app_profile"][field][subfield] request = request_type(**request_init) # Mock the http request call within the method and fake a response. @@ -9345,16 +9574,15 @@ def test_create_app_profile_rest(request_type): name="name_value", etag="etag_value", description="description_value", - multi_cluster_routing_use_any=instance.AppProfile.MultiClusterRoutingUseAny( - cluster_ids=["cluster_ids_value"] - ), + priority=instance.AppProfile.Priority.PRIORITY_LOW, ) # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.AppProfile.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.AppProfile.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9446,8 +9674,9 @@ def test_create_app_profile_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = instance.AppProfile.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.AppProfile.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9554,18 +9783,6 @@ def test_create_app_profile_rest_bad_request( # send a request that will satisfy transcoding request_init = {"parent": "projects/sample1/instances/sample2"} - request_init["app_profile"] = { - "name": "name_value", - "etag": "etag_value", - "description": "description_value", - "multi_cluster_routing_use_any": { - "cluster_ids": ["cluster_ids_value1", "cluster_ids_value2"] - }, - "single_cluster_routing": { - "cluster_id": "cluster_id_value", - "allow_transactional_writes": True, - }, - } request = request_type(**request_init) # Mock the http request call within the method and fake a BadRequest error. @@ -9605,8 +9822,9 @@ def test_create_app_profile_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.AppProfile.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.AppProfile.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9670,16 +9888,15 @@ def test_get_app_profile_rest(request_type): name="name_value", etag="etag_value", description="description_value", - multi_cluster_routing_use_any=instance.AppProfile.MultiClusterRoutingUseAny( - cluster_ids=["cluster_ids_value"] - ), + priority=instance.AppProfile.Priority.PRIORITY_LOW, ) # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.AppProfile.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.AppProfile.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9756,8 +9973,9 @@ def test_get_app_profile_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = instance.AppProfile.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.AppProfile.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9882,8 +10100,9 @@ def test_get_app_profile_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = instance.AppProfile.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = instance.AppProfile.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9949,10 +10168,9 @@ def test_list_app_profiles_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListAppProfilesResponse.pb( - return_value - ) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListAppProfilesResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10035,10 +10253,11 @@ def test_list_app_profiles_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListAppProfilesResponse.pb( + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListAppProfilesResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10173,10 +10392,9 @@ def test_list_app_profiles_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListAppProfilesResponse.pb( - return_value - ) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListAppProfilesResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10301,7 +10519,78 @@ def test_update_app_profile_rest(request_type): "cluster_id": "cluster_id_value", "allow_transactional_writes": True, }, + "priority": 1, + "standard_isolation": {"priority": 1}, } + # The version of a generated dependency at test runtime may differ from the version used during generation. + # Delete any fields which are not present in the current runtime dependency + # See https://github.com/googleapis/gapic-generator-python/issues/1748 + + # Determine if the message type is proto-plus or protobuf + test_field = bigtable_instance_admin.UpdateAppProfileRequest.meta.fields[ + "app_profile" + ] + + def get_message_fields(field): + # Given a field which is a message (composite type), return a list with + # all the fields of the message. + # If the field is not a composite type, return an empty list. + message_fields = [] + + if hasattr(field, "message") and field.message: + is_field_type_proto_plus_type = not hasattr(field.message, "DESCRIPTOR") + + if is_field_type_proto_plus_type: + message_fields = field.message.meta.fields.values() + # Add `# pragma: NO COVER` because there may not be any `*_pb2` field types + else: # pragma: NO COVER + message_fields = field.message.DESCRIPTOR.fields + return message_fields + + runtime_nested_fields = [ + (field.name, nested_field.name) + for field in get_message_fields(test_field) + for nested_field in get_message_fields(field) + ] + + subfields_not_in_runtime = [] + + # For each item in the sample request, create a list of sub fields which are not present at runtime + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for field, value in request_init["app_profile"].items(): # pragma: NO COVER + result = None + is_repeated = False + # For repeated fields + if isinstance(value, list) and len(value): + is_repeated = True + result = value[0] + # For fields where the type is another message + if isinstance(value, dict): + result = value + + if result and hasattr(result, "keys"): + for subfield in result.keys(): + if (field, subfield) not in runtime_nested_fields: + subfields_not_in_runtime.append( + { + "field": field, + "subfield": subfield, + "is_repeated": is_repeated, + } + ) + + # Remove fields from the sample request which are not present in the runtime version of the dependency + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for subfield_to_delete in subfields_not_in_runtime: # pragma: NO COVER + field = subfield_to_delete.get("field") + field_repeated = subfield_to_delete.get("is_repeated") + subfield = subfield_to_delete.get("subfield") + if subfield: + if field_repeated: + for i in range(0, len(request_init["app_profile"][field])): + del request_init["app_profile"][field][i][subfield] + else: + del request_init["app_profile"][field][subfield] request = request_type(**request_init) # Mock the http request call within the method and fake a response. @@ -10496,18 +10785,6 @@ def test_update_app_profile_rest_bad_request( "name": "projects/sample1/instances/sample2/appProfiles/sample3" } } - request_init["app_profile"] = { - "name": "projects/sample1/instances/sample2/appProfiles/sample3", - "etag": "etag_value", - "description": "description_value", - "multi_cluster_routing_use_any": { - "cluster_ids": ["cluster_ids_value1", "cluster_ids_value2"] - }, - "single_cluster_routing": { - "cluster_id": "cluster_id_value", - "allow_transactional_writes": True, - }, - } request = request_type(**request_init) # Mock the http request call within the method and fake a BadRequest error. @@ -10896,8 +11173,7 @@ def test_get_iam_policy_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10974,8 +11250,7 @@ def test_get_iam_policy_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11096,8 +11371,7 @@ def test_get_iam_policy_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11163,8 +11437,7 @@ def test_set_iam_policy_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11241,8 +11514,7 @@ def test_set_iam_policy_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11371,8 +11643,7 @@ def test_set_iam_policy_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11437,8 +11708,7 @@ def test_test_iam_permissions_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11518,8 +11788,7 @@ def test_test_iam_permissions_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11651,8 +11920,7 @@ def test_test_iam_permissions_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11718,10 +11986,9 @@ def test_list_hot_tablets_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListHotTabletsResponse.pb( - return_value - ) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListHotTabletsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11805,10 +12072,11 @@ def test_list_hot_tablets_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListHotTabletsResponse.pb( + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListHotTabletsResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11947,10 +12215,9 @@ def test_list_hot_tablets_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_instance_admin.ListHotTabletsResponse.pb( - return_value - ) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_instance_admin.ListHotTabletsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value diff --git a/tests/unit/gapic/bigtable_admin_v2/test_bigtable_table_admin.py b/tests/unit/gapic/bigtable_admin_v2/test_bigtable_table_admin.py index 8498e4fa5..b29dc5106 100644 --- a/tests/unit/gapic/bigtable_admin_v2/test_bigtable_table_admin.py +++ b/tests/unit/gapic/bigtable_admin_v2/test_bigtable_table_admin.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -62,7 +62,7 @@ from google.iam.v1 import iam_policy_pb2 # type: ignore from google.iam.v1 import options_pb2 # type: ignore from google.iam.v1 import policy_pb2 # type: ignore -from google.longrunning import operations_pb2 +from google.longrunning import operations_pb2 # type: ignore from google.oauth2 import service_account from google.protobuf import any_pb2 # type: ignore from google.protobuf import duration_pb2 # type: ignore @@ -1691,9 +1691,11 @@ async def test_list_tables_async_pages(): RuntimeError, ) pages = [] - async for page_ in ( + # Workaround issue in python 3.9 related to code coverage by adding `# pragma: no branch` + # See https://github.com/googleapis/gapic-generator-python/pull/1174#issuecomment-1025132372 + async for page_ in ( # pragma: no branch await client.list_tables(request={}) - ).pages: # pragma: no branch + ).pages: pages.append(page_) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token @@ -4457,9 +4459,11 @@ async def test_list_snapshots_async_pages(): RuntimeError, ) pages = [] - async for page_ in ( + # Workaround issue in python 3.9 related to code coverage by adding `# pragma: no branch` + # See https://github.com/googleapis/gapic-generator-python/pull/1174#issuecomment-1025132372 + async for page_ in ( # pragma: no branch await client.list_snapshots(request={}) - ).pages: # pragma: no branch + ).pages: pages.append(page_) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token @@ -4956,6 +4960,7 @@ def test_get_backup(request_type, transport: str = "grpc"): call.return_value = table.Backup( name="name_value", source_table="source_table_value", + source_backup="source_backup_value", size_bytes=1089, state=table.Backup.State.CREATING, ) @@ -4970,6 +4975,7 @@ def test_get_backup(request_type, transport: str = "grpc"): assert isinstance(response, table.Backup) assert response.name == "name_value" assert response.source_table == "source_table_value" + assert response.source_backup == "source_backup_value" assert response.size_bytes == 1089 assert response.state == table.Backup.State.CREATING @@ -5010,6 +5016,7 @@ async def test_get_backup_async( table.Backup( name="name_value", source_table="source_table_value", + source_backup="source_backup_value", size_bytes=1089, state=table.Backup.State.CREATING, ) @@ -5025,6 +5032,7 @@ async def test_get_backup_async( assert isinstance(response, table.Backup) assert response.name == "name_value" assert response.source_table == "source_table_value" + assert response.source_backup == "source_backup_value" assert response.size_bytes == 1089 assert response.state == table.Backup.State.CREATING @@ -5196,6 +5204,7 @@ def test_update_backup(request_type, transport: str = "grpc"): call.return_value = table.Backup( name="name_value", source_table="source_table_value", + source_backup="source_backup_value", size_bytes=1089, state=table.Backup.State.CREATING, ) @@ -5210,6 +5219,7 @@ def test_update_backup(request_type, transport: str = "grpc"): assert isinstance(response, table.Backup) assert response.name == "name_value" assert response.source_table == "source_table_value" + assert response.source_backup == "source_backup_value" assert response.size_bytes == 1089 assert response.state == table.Backup.State.CREATING @@ -5251,6 +5261,7 @@ async def test_update_backup_async( table.Backup( name="name_value", source_table="source_table_value", + source_backup="source_backup_value", size_bytes=1089, state=table.Backup.State.CREATING, ) @@ -5266,6 +5277,7 @@ async def test_update_backup_async( assert isinstance(response, table.Backup) assert response.name == "name_value" assert response.source_table == "source_table_value" + assert response.source_backup == "source_backup_value" assert response.size_bytes == 1089 assert response.state == table.Backup.State.CREATING @@ -6058,9 +6070,11 @@ async def test_list_backups_async_pages(): RuntimeError, ) pages = [] - async for page_ in ( + # Workaround issue in python 3.9 related to code coverage by adding `# pragma: no branch` + # See https://github.com/googleapis/gapic-generator-python/pull/1174#issuecomment-1025132372 + async for page_ in ( # pragma: no branch await client.list_backups(request={}) - ).pages: # pragma: no branch + ).pages: pages.append(page_) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token @@ -6211,6 +6225,262 @@ async def test_restore_table_field_headers_async(): ) in kw["metadata"] +@pytest.mark.parametrize( + "request_type", + [ + bigtable_table_admin.CopyBackupRequest, + dict, + ], +) +def test_copy_backup(request_type, transport: str = "grpc"): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.copy_backup), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation(name="operations/spam") + response = client.copy_backup(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == bigtable_table_admin.CopyBackupRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, future.Future) + + +def test_copy_backup_empty_call(): + # This test is a coverage failsafe to make sure that totally empty calls, + # i.e. request == None and no flattened fields passed, work. + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.copy_backup), "__call__") as call: + client.copy_backup() + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == bigtable_table_admin.CopyBackupRequest() + + +@pytest.mark.asyncio +async def test_copy_backup_async( + transport: str = "grpc_asyncio", request_type=bigtable_table_admin.CopyBackupRequest +): + client = BigtableTableAdminAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.copy_backup), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation(name="operations/spam") + ) + response = await client.copy_backup(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == bigtable_table_admin.CopyBackupRequest() + + # Establish that the response is the type that we expect. + assert isinstance(response, future.Future) + + +@pytest.mark.asyncio +async def test_copy_backup_async_from_dict(): + await test_copy_backup_async(request_type=dict) + + +def test_copy_backup_field_headers(): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = bigtable_table_admin.CopyBackupRequest() + + request.parent = "parent_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.copy_backup), "__call__") as call: + call.return_value = operations_pb2.Operation(name="operations/op") + client.copy_backup(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "parent=parent_value", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_copy_backup_field_headers_async(): + client = BigtableTableAdminAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = bigtable_table_admin.CopyBackupRequest() + + request.parent = "parent_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.copy_backup), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation(name="operations/op") + ) + await client.copy_backup(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "parent=parent_value", + ) in kw["metadata"] + + +def test_copy_backup_flattened(): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.copy_backup), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation(name="operations/op") + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.copy_backup( + parent="parent_value", + backup_id="backup_id_value", + source_backup="source_backup_value", + expire_time=timestamp_pb2.Timestamp(seconds=751), + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + arg = args[0].parent + mock_val = "parent_value" + assert arg == mock_val + arg = args[0].backup_id + mock_val = "backup_id_value" + assert arg == mock_val + arg = args[0].source_backup + mock_val = "source_backup_value" + assert arg == mock_val + assert TimestampRule().to_proto(args[0].expire_time) == timestamp_pb2.Timestamp( + seconds=751 + ) + + +def test_copy_backup_flattened_error(): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.copy_backup( + bigtable_table_admin.CopyBackupRequest(), + parent="parent_value", + backup_id="backup_id_value", + source_backup="source_backup_value", + expire_time=timestamp_pb2.Timestamp(seconds=751), + ) + + +@pytest.mark.asyncio +async def test_copy_backup_flattened_async(): + client = BigtableTableAdminAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.copy_backup), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation(name="operations/op") + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation(name="operations/spam") + ) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.copy_backup( + parent="parent_value", + backup_id="backup_id_value", + source_backup="source_backup_value", + expire_time=timestamp_pb2.Timestamp(seconds=751), + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + arg = args[0].parent + mock_val = "parent_value" + assert arg == mock_val + arg = args[0].backup_id + mock_val = "backup_id_value" + assert arg == mock_val + arg = args[0].source_backup + mock_val = "source_backup_value" + assert arg == mock_val + assert TimestampRule().to_proto(args[0].expire_time) == timestamp_pb2.Timestamp( + seconds=751 + ) + + +@pytest.mark.asyncio +async def test_copy_backup_flattened_error_async(): + client = BigtableTableAdminAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.copy_backup( + bigtable_table_admin.CopyBackupRequest(), + parent="parent_value", + backup_id="backup_id_value", + source_backup="source_backup_value", + expire_time=timestamp_pb2.Timestamp(seconds=751), + ) + + @pytest.mark.parametrize( "request_type", [ @@ -7015,8 +7285,9 @@ def test_create_table_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = gba_table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = gba_table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7098,8 +7369,9 @@ def test_create_table_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = gba_table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = gba_table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7233,8 +7505,9 @@ def test_create_table_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = gba_table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = gba_table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7586,8 +7859,9 @@ def test_list_tables_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListTablesResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListTablesResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7670,8 +7944,9 @@ def test_list_tables_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListTablesResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListTablesResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7805,8 +8080,9 @@ def test_list_tables_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListTablesResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListTablesResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -7929,8 +8205,9 @@ def test_get_table_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8009,8 +8286,9 @@ def test_get_table_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8133,8 +8411,9 @@ def test_get_table_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -8200,11 +8479,79 @@ def test_update_table_rest(request_type): "start_time": {"seconds": 751, "nanos": 543}, "end_time": {}, "source_table": "source_table_value", + "source_backup": "source_backup_value", }, }, "change_stream_config": {"retention_period": {"seconds": 751, "nanos": 543}}, "deletion_protection": True, } + # The version of a generated dependency at test runtime may differ from the version used during generation. + # Delete any fields which are not present in the current runtime dependency + # See https://github.com/googleapis/gapic-generator-python/issues/1748 + + # Determine if the message type is proto-plus or protobuf + test_field = bigtable_table_admin.UpdateTableRequest.meta.fields["table"] + + def get_message_fields(field): + # Given a field which is a message (composite type), return a list with + # all the fields of the message. + # If the field is not a composite type, return an empty list. + message_fields = [] + + if hasattr(field, "message") and field.message: + is_field_type_proto_plus_type = not hasattr(field.message, "DESCRIPTOR") + + if is_field_type_proto_plus_type: + message_fields = field.message.meta.fields.values() + # Add `# pragma: NO COVER` because there may not be any `*_pb2` field types + else: # pragma: NO COVER + message_fields = field.message.DESCRIPTOR.fields + return message_fields + + runtime_nested_fields = [ + (field.name, nested_field.name) + for field in get_message_fields(test_field) + for nested_field in get_message_fields(field) + ] + + subfields_not_in_runtime = [] + + # For each item in the sample request, create a list of sub fields which are not present at runtime + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for field, value in request_init["table"].items(): # pragma: NO COVER + result = None + is_repeated = False + # For repeated fields + if isinstance(value, list) and len(value): + is_repeated = True + result = value[0] + # For fields where the type is another message + if isinstance(value, dict): + result = value + + if result and hasattr(result, "keys"): + for subfield in result.keys(): + if (field, subfield) not in runtime_nested_fields: + subfields_not_in_runtime.append( + { + "field": field, + "subfield": subfield, + "is_repeated": is_repeated, + } + ) + + # Remove fields from the sample request which are not present in the runtime version of the dependency + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for subfield_to_delete in subfields_not_in_runtime: # pragma: NO COVER + field = subfield_to_delete.get("field") + field_repeated = subfield_to_delete.get("is_repeated") + subfield = subfield_to_delete.get("subfield") + if subfield: + if field_repeated: + for i in range(0, len(request_init["table"][field])): + del request_init["table"][field][i][subfield] + else: + del request_init["table"][field][subfield] request = request_type(**request_init) # Mock the http request call within the method and fake a response. @@ -8386,23 +8733,6 @@ def test_update_table_rest_bad_request( request_init = { "table": {"name": "projects/sample1/instances/sample2/tables/sample3"} } - request_init["table"] = { - "name": "projects/sample1/instances/sample2/tables/sample3", - "cluster_states": {}, - "column_families": {}, - "granularity": 1, - "restore_info": { - "source_type": 1, - "backup_info": { - "backup": "backup_value", - "start_time": {"seconds": 751, "nanos": 543}, - "end_time": {}, - "source_table": "source_table_value", - }, - }, - "change_stream_config": {"retention_period": {"seconds": 751, "nanos": 543}}, - "deletion_protection": True, - } request = request_type(**request_init) # Mock the http request call within the method and fake a BadRequest error. @@ -9027,8 +9357,9 @@ def test_modify_column_families_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9106,8 +9437,9 @@ def test_modify_column_families_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9244,8 +9576,9 @@ def test_modify_column_families_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Table.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Table.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9514,10 +9847,11 @@ def test_generate_consistency_token_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.GenerateConsistencyTokenResponse.pb( + # Convert return value to protobuf type + return_value = bigtable_table_admin.GenerateConsistencyTokenResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9593,10 +9927,11 @@ def test_generate_consistency_token_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.GenerateConsistencyTokenResponse.pb( + # Convert return value to protobuf type + return_value = bigtable_table_admin.GenerateConsistencyTokenResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9724,10 +10059,11 @@ def test_generate_consistency_token_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.GenerateConsistencyTokenResponse.pb( + # Convert return value to protobuf type + return_value = bigtable_table_admin.GenerateConsistencyTokenResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9792,8 +10128,9 @@ def test_check_consistency_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.CheckConsistencyResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.CheckConsistencyResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -9873,10 +10210,11 @@ def test_check_consistency_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.CheckConsistencyResponse.pb( + # Convert return value to protobuf type + return_value = bigtable_table_admin.CheckConsistencyResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10012,8 +10350,9 @@ def test_check_consistency_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.CheckConsistencyResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.CheckConsistencyResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10371,8 +10710,9 @@ def test_get_snapshot_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Snapshot.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Snapshot.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10450,8 +10790,9 @@ def test_get_snapshot_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = table.Snapshot.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Snapshot.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10578,8 +10919,9 @@ def test_get_snapshot_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Snapshot.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Snapshot.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10644,8 +10986,9 @@ def test_list_snapshots_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListSnapshotsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListSnapshotsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10727,10 +11070,9 @@ def test_list_snapshots_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListSnapshotsResponse.pb( - return_value - ) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListSnapshotsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -10865,8 +11207,9 @@ def test_list_snapshots_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListSnapshotsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListSnapshotsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11240,6 +11583,7 @@ def test_create_backup_rest(request_type): request_init["backup"] = { "name": "name_value", "source_table": "source_table_value", + "source_backup": "source_backup_value", "expire_time": {"seconds": 751, "nanos": 543}, "start_time": {}, "end_time": {}, @@ -11260,10 +11604,77 @@ def test_create_backup_rest(request_type): "kms_key_version": "kms_key_version_value", }, } - request = request_type(**request_init) + # The version of a generated dependency at test runtime may differ from the version used during generation. + # Delete any fields which are not present in the current runtime dependency + # See https://github.com/googleapis/gapic-generator-python/issues/1748 + + # Determine if the message type is proto-plus or protobuf + test_field = bigtable_table_admin.CreateBackupRequest.meta.fields["backup"] + + def get_message_fields(field): + # Given a field which is a message (composite type), return a list with + # all the fields of the message. + # If the field is not a composite type, return an empty list. + message_fields = [] + + if hasattr(field, "message") and field.message: + is_field_type_proto_plus_type = not hasattr(field.message, "DESCRIPTOR") + + if is_field_type_proto_plus_type: + message_fields = field.message.meta.fields.values() + # Add `# pragma: NO COVER` because there may not be any `*_pb2` field types + else: # pragma: NO COVER + message_fields = field.message.DESCRIPTOR.fields + return message_fields + + runtime_nested_fields = [ + (field.name, nested_field.name) + for field in get_message_fields(test_field) + for nested_field in get_message_fields(field) + ] - # Mock the http request call within the method and fake a response. - with mock.patch.object(type(client.transport._session), "request") as req: + subfields_not_in_runtime = [] + + # For each item in the sample request, create a list of sub fields which are not present at runtime + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for field, value in request_init["backup"].items(): # pragma: NO COVER + result = None + is_repeated = False + # For repeated fields + if isinstance(value, list) and len(value): + is_repeated = True + result = value[0] + # For fields where the type is another message + if isinstance(value, dict): + result = value + + if result and hasattr(result, "keys"): + for subfield in result.keys(): + if (field, subfield) not in runtime_nested_fields: + subfields_not_in_runtime.append( + { + "field": field, + "subfield": subfield, + "is_repeated": is_repeated, + } + ) + + # Remove fields from the sample request which are not present in the runtime version of the dependency + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for subfield_to_delete in subfields_not_in_runtime: # pragma: NO COVER + field = subfield_to_delete.get("field") + field_repeated = subfield_to_delete.get("is_repeated") + subfield = subfield_to_delete.get("subfield") + if subfield: + if field_repeated: + for i in range(0, len(request_init["backup"][field])): + del request_init["backup"][field][i][subfield] + else: + del request_init["backup"][field][subfield] + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. return_value = operations_pb2.Operation(name="operations/spam") @@ -11458,29 +11869,6 @@ def test_create_backup_rest_bad_request( # send a request that will satisfy transcoding request_init = {"parent": "projects/sample1/instances/sample2/clusters/sample3"} - request_init["backup"] = { - "name": "name_value", - "source_table": "source_table_value", - "expire_time": {"seconds": 751, "nanos": 543}, - "start_time": {}, - "end_time": {}, - "size_bytes": 1089, - "state": 1, - "encryption_info": { - "encryption_type": 1, - "encryption_status": { - "code": 411, - "message": "message_value", - "details": [ - { - "type_url": "type.googleapis.com/google.protobuf.Duration", - "value": b"\x08\x0c\x10\xdb\x07", - } - ], - }, - "kms_key_version": "kms_key_version_value", - }, - } request = request_type(**request_init) # Mock the http request call within the method and fake a BadRequest error. @@ -11587,6 +11975,7 @@ def test_get_backup_rest(request_type): return_value = table.Backup( name="name_value", source_table="source_table_value", + source_backup="source_backup_value", size_bytes=1089, state=table.Backup.State.CREATING, ) @@ -11594,8 +11983,9 @@ def test_get_backup_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Backup.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Backup.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11605,6 +11995,7 @@ def test_get_backup_rest(request_type): assert isinstance(response, table.Backup) assert response.name == "name_value" assert response.source_table == "source_table_value" + assert response.source_backup == "source_backup_value" assert response.size_bytes == 1089 assert response.state == table.Backup.State.CREATING @@ -11673,8 +12064,9 @@ def test_get_backup_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = table.Backup.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Backup.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11801,8 +12193,9 @@ def test_get_backup_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Backup.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Backup.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11862,6 +12255,7 @@ def test_update_backup_rest(request_type): request_init["backup"] = { "name": "projects/sample1/instances/sample2/clusters/sample3/backups/sample4", "source_table": "source_table_value", + "source_backup": "source_backup_value", "expire_time": {"seconds": 751, "nanos": 543}, "start_time": {}, "end_time": {}, @@ -11882,6 +12276,73 @@ def test_update_backup_rest(request_type): "kms_key_version": "kms_key_version_value", }, } + # The version of a generated dependency at test runtime may differ from the version used during generation. + # Delete any fields which are not present in the current runtime dependency + # See https://github.com/googleapis/gapic-generator-python/issues/1748 + + # Determine if the message type is proto-plus or protobuf + test_field = bigtable_table_admin.UpdateBackupRequest.meta.fields["backup"] + + def get_message_fields(field): + # Given a field which is a message (composite type), return a list with + # all the fields of the message. + # If the field is not a composite type, return an empty list. + message_fields = [] + + if hasattr(field, "message") and field.message: + is_field_type_proto_plus_type = not hasattr(field.message, "DESCRIPTOR") + + if is_field_type_proto_plus_type: + message_fields = field.message.meta.fields.values() + # Add `# pragma: NO COVER` because there may not be any `*_pb2` field types + else: # pragma: NO COVER + message_fields = field.message.DESCRIPTOR.fields + return message_fields + + runtime_nested_fields = [ + (field.name, nested_field.name) + for field in get_message_fields(test_field) + for nested_field in get_message_fields(field) + ] + + subfields_not_in_runtime = [] + + # For each item in the sample request, create a list of sub fields which are not present at runtime + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for field, value in request_init["backup"].items(): # pragma: NO COVER + result = None + is_repeated = False + # For repeated fields + if isinstance(value, list) and len(value): + is_repeated = True + result = value[0] + # For fields where the type is another message + if isinstance(value, dict): + result = value + + if result and hasattr(result, "keys"): + for subfield in result.keys(): + if (field, subfield) not in runtime_nested_fields: + subfields_not_in_runtime.append( + { + "field": field, + "subfield": subfield, + "is_repeated": is_repeated, + } + ) + + # Remove fields from the sample request which are not present in the runtime version of the dependency + # Add `# pragma: NO COVER` because this test code will not run if all subfields are present at runtime + for subfield_to_delete in subfields_not_in_runtime: # pragma: NO COVER + field = subfield_to_delete.get("field") + field_repeated = subfield_to_delete.get("is_repeated") + subfield = subfield_to_delete.get("subfield") + if subfield: + if field_repeated: + for i in range(0, len(request_init["backup"][field])): + del request_init["backup"][field][i][subfield] + else: + del request_init["backup"][field][subfield] request = request_type(**request_init) # Mock the http request call within the method and fake a response. @@ -11890,6 +12351,7 @@ def test_update_backup_rest(request_type): return_value = table.Backup( name="name_value", source_table="source_table_value", + source_backup="source_backup_value", size_bytes=1089, state=table.Backup.State.CREATING, ) @@ -11897,8 +12359,9 @@ def test_update_backup_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Backup.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Backup.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -11908,6 +12371,7 @@ def test_update_backup_rest(request_type): assert isinstance(response, table.Backup) assert response.name == "name_value" assert response.source_table == "source_table_value" + assert response.source_backup == "source_backup_value" assert response.size_bytes == 1089 assert response.state == table.Backup.State.CREATING @@ -11974,8 +12438,9 @@ def test_update_backup_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = table.Backup.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Backup.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -12073,29 +12538,6 @@ def test_update_backup_rest_bad_request( "name": "projects/sample1/instances/sample2/clusters/sample3/backups/sample4" } } - request_init["backup"] = { - "name": "projects/sample1/instances/sample2/clusters/sample3/backups/sample4", - "source_table": "source_table_value", - "expire_time": {"seconds": 751, "nanos": 543}, - "start_time": {}, - "end_time": {}, - "size_bytes": 1089, - "state": 1, - "encryption_info": { - "encryption_type": 1, - "encryption_status": { - "code": 411, - "message": "message_value", - "details": [ - { - "type_url": "type.googleapis.com/google.protobuf.Duration", - "value": b"\x08\x0c\x10\xdb\x07", - } - ], - }, - "kms_key_version": "kms_key_version_value", - }, - } request = request_type(**request_init) # Mock the http request call within the method and fake a BadRequest error. @@ -12138,8 +12580,9 @@ def test_update_backup_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = table.Backup.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = table.Backup.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -12464,8 +12907,9 @@ def test_list_backups_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListBackupsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListBackupsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -12549,8 +12993,9 @@ def test_list_backups_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListBackupsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListBackupsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -12687,8 +13132,9 @@ def test_list_backups_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable_table_admin.ListBackupsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable_table_admin.ListBackupsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13006,6 +13452,296 @@ def test_restore_table_rest_error(): ) +@pytest.mark.parametrize( + "request_type", + [ + bigtable_table_admin.CopyBackupRequest, + dict, + ], +) +def test_copy_backup_rest(request_type): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # send a request that will satisfy transcoding + request_init = {"parent": "projects/sample1/instances/sample2/clusters/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + response = client.copy_backup(request) + + # Establish that the response is the type that we expect. + assert response.operation.name == "operations/spam" + + +def test_copy_backup_rest_required_fields( + request_type=bigtable_table_admin.CopyBackupRequest, +): + transport_class = transports.BigtableTableAdminRestTransport + + request_init = {} + request_init["parent"] = "" + request_init["backup_id"] = "" + request_init["source_backup"] = "" + request = request_type(**request_init) + pb_request = request_type.pb(request) + jsonified_request = json.loads( + json_format.MessageToJson( + pb_request, + including_default_value_fields=False, + use_integers_for_enums=False, + ) + ) + + # verify fields with default values are dropped + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).copy_backup._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with default values are now present + + jsonified_request["parent"] = "parent_value" + jsonified_request["backupId"] = "backup_id_value" + jsonified_request["sourceBackup"] = "source_backup_value" + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).copy_backup._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with non-default values are left alone + assert "parent" in jsonified_request + assert jsonified_request["parent"] == "parent_value" + assert "backupId" in jsonified_request + assert jsonified_request["backupId"] == "backup_id_value" + assert "sourceBackup" in jsonified_request + assert jsonified_request["sourceBackup"] == "source_backup_value" + + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + request = request_type(**request_init) + + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + # Mock the http request call within the method and fake a response. + with mock.patch.object(Session, "request") as req: + # We need to mock transcode() because providing default values + # for required fields will fail the real version if the http_options + # expect actual values for those fields. + with mock.patch.object(path_template, "transcode") as transcode: + # A uri without fields and an empty body will force all the + # request fields to show up in the query_params. + pb_request = request_type.pb(request) + transcode_result = { + "uri": "v1/sample_method", + "method": "post", + "query_params": pb_request, + } + transcode_result["body"] = pb_request + transcode.return_value = transcode_result + + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + + response = client.copy_backup(request) + + expected_params = [("$alt", "json;enum-encoding=int")] + actual_params = req.call_args.kwargs["params"] + assert expected_params == actual_params + + +def test_copy_backup_rest_unset_required_fields(): + transport = transports.BigtableTableAdminRestTransport( + credentials=ga_credentials.AnonymousCredentials + ) + + unset_fields = transport.copy_backup._get_unset_required_fields({}) + assert set(unset_fields) == ( + set(()) + & set( + ( + "parent", + "backupId", + "sourceBackup", + "expireTime", + ) + ) + ) + + +@pytest.mark.parametrize("null_interceptor", [True, False]) +def test_copy_backup_rest_interceptors(null_interceptor): + transport = transports.BigtableTableAdminRestTransport( + credentials=ga_credentials.AnonymousCredentials(), + interceptor=None + if null_interceptor + else transports.BigtableTableAdminRestInterceptor(), + ) + client = BigtableTableAdminClient(transport=transport) + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + operation.Operation, "_set_result_from_operation" + ), mock.patch.object( + transports.BigtableTableAdminRestInterceptor, "post_copy_backup" + ) as post, mock.patch.object( + transports.BigtableTableAdminRestInterceptor, "pre_copy_backup" + ) as pre: + pre.assert_not_called() + post.assert_not_called() + pb_message = bigtable_table_admin.CopyBackupRequest.pb( + bigtable_table_admin.CopyBackupRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = Response() + req.return_value.status_code = 200 + req.return_value.request = PreparedRequest() + req.return_value._content = json_format.MessageToJson( + operations_pb2.Operation() + ) + + request = bigtable_table_admin.CopyBackupRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = operations_pb2.Operation() + + client.copy_backup( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + + +def test_copy_backup_rest_bad_request( + transport: str = "rest", request_type=bigtable_table_admin.CopyBackupRequest +): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # send a request that will satisfy transcoding + request_init = {"parent": "projects/sample1/instances/sample2/clusters/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(Session, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 400 + response_value.request = Request() + req.return_value = response_value + client.copy_backup(request) + + +def test_copy_backup_rest_flattened(): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + + # get arguments that satisfy an http rule for this method + sample_request = { + "parent": "projects/sample1/instances/sample2/clusters/sample3" + } + + # get truthy value for each flattened field + mock_args = dict( + parent="parent_value", + backup_id="backup_id_value", + source_backup="source_backup_value", + expire_time=timestamp_pb2.Timestamp(seconds=751), + ) + mock_args.update(sample_request) + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + + client.copy_backup(**mock_args) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(req.mock_calls) == 1 + _, args, _ = req.mock_calls[0] + assert path_template.validate( + "%s/v2/{parent=projects/*/instances/*/clusters/*}/backups:copy" + % client.transport._host, + args[1], + ) + + +def test_copy_backup_rest_flattened_error(transport: str = "rest"): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.copy_backup( + bigtable_table_admin.CopyBackupRequest(), + parent="parent_value", + backup_id="backup_id_value", + source_backup="source_backup_value", + expire_time=timestamp_pb2.Timestamp(seconds=751), + ) + + +def test_copy_backup_rest_error(): + client = BigtableTableAdminClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + + @pytest.mark.parametrize( "request_type", [ @@ -13034,8 +13770,7 @@ def test_get_iam_policy_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13112,8 +13847,7 @@ def test_get_iam_policy_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13236,8 +13970,7 @@ def test_get_iam_policy_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13303,8 +14036,7 @@ def test_set_iam_policy_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13381,8 +14113,7 @@ def test_set_iam_policy_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13513,8 +14244,7 @@ def test_set_iam_policy_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13579,8 +14309,7 @@ def test_test_iam_permissions_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13660,8 +14389,7 @@ def test_test_iam_permissions_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13795,8 +14523,7 @@ def test_test_iam_permissions_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = return_value - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -13995,6 +14722,7 @@ def test_bigtable_table_admin_base_transport(): "delete_backup", "list_backups", "restore_table", + "copy_backup", "get_iam_policy", "set_iam_policy", "test_iam_permissions", @@ -14371,6 +15099,9 @@ def test_bigtable_table_admin_client_transport_session_collision(transport_name) session1 = client1.transport.restore_table._session session2 = client2.transport.restore_table._session assert session1 != session2 + session1 = client1.transport.copy_backup._session + session2 = client2.transport.copy_backup._session + assert session1 != session2 session1 = client1.transport.get_iam_policy._session session2 = client2.transport.get_iam_policy._session assert session1 != session2 diff --git a/tests/unit/gapic/bigtable_v2/__init__.py b/tests/unit/gapic/bigtable_v2/__init__.py index e8e1c3845..89a37dc92 100644 --- a/tests/unit/gapic/bigtable_v2/__init__.py +++ b/tests/unit/gapic/bigtable_v2/__init__.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/unit/gapic/bigtable_v2/test_bigtable.py b/tests/unit/gapic/bigtable_v2/test_bigtable.py index 03ba3044f..2319306d7 100644 --- a/tests/unit/gapic/bigtable_v2/test_bigtable.py +++ b/tests/unit/gapic/bigtable_v2/test_bigtable.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -# Copyright 2022 Google LLC +# Copyright 2023 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -3004,8 +3004,9 @@ def test_read_rows_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadRowsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadRowsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) @@ -3086,8 +3087,9 @@ def test_read_rows_rest_required_fields(request_type=bigtable.ReadRowsRequest): response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadRowsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadRowsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") @@ -3215,8 +3217,9 @@ def test_read_rows_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadRowsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadRowsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -3286,8 +3289,9 @@ def test_sample_row_keys_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.SampleRowKeysResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.SampleRowKeysResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) @@ -3372,8 +3376,9 @@ def test_sample_row_keys_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.SampleRowKeysResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.SampleRowKeysResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") @@ -3501,8 +3506,9 @@ def test_sample_row_keys_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.SampleRowKeysResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.SampleRowKeysResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -3569,8 +3575,9 @@ def test_mutate_row_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.MutateRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.MutateRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -3647,8 +3654,9 @@ def test_mutate_row_rest_required_fields(request_type=bigtable.MutateRowRequest) response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.MutateRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.MutateRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -3787,8 +3795,9 @@ def test_mutate_row_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.MutateRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.MutateRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -3858,8 +3867,9 @@ def test_mutate_rows_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.MutateRowsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.MutateRowsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) @@ -3939,8 +3949,9 @@ def test_mutate_rows_rest_required_fields(request_type=bigtable.MutateRowsReques response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.MutateRowsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.MutateRowsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") @@ -4077,8 +4088,9 @@ def test_mutate_rows_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.MutateRowsResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.MutateRowsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4148,8 +4160,9 @@ def test_check_and_mutate_row_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.CheckAndMutateRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.CheckAndMutateRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4229,8 +4242,9 @@ def test_check_and_mutate_row_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.CheckAndMutateRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.CheckAndMutateRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4386,8 +4400,9 @@ def test_check_and_mutate_row_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.CheckAndMutateRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.CheckAndMutateRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4473,8 +4488,9 @@ def test_ping_and_warm_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.PingAndWarmResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.PingAndWarmResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4547,8 +4563,9 @@ def test_ping_and_warm_rest_required_fields(request_type=bigtable.PingAndWarmReq response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.PingAndWarmResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.PingAndWarmResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4670,8 +4687,9 @@ def test_ping_and_warm_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.PingAndWarmResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.PingAndWarmResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4733,8 +4751,9 @@ def test_read_modify_write_row_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadModifyWriteRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadModifyWriteRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4813,8 +4832,9 @@ def test_read_modify_write_row_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadModifyWriteRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadModifyWriteRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -4951,8 +4971,9 @@ def test_read_modify_write_row_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadModifyWriteRowResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadModifyWriteRowResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -5018,10 +5039,11 @@ def test_generate_initial_change_stream_partitions_rest(request_type): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.GenerateInitialChangeStreamPartitionsResponse.pb( + # Convert return value to protobuf type + return_value = bigtable.GenerateInitialChangeStreamPartitionsResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) @@ -5107,10 +5129,11 @@ def test_generate_initial_change_stream_partitions_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.GenerateInitialChangeStreamPartitionsResponse.pb( + # Convert return value to protobuf type + return_value = bigtable.GenerateInitialChangeStreamPartitionsResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") @@ -5249,10 +5272,11 @@ def test_generate_initial_change_stream_partitions_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.GenerateInitialChangeStreamPartitionsResponse.pb( + # Convert return value to protobuf type + return_value = bigtable.GenerateInitialChangeStreamPartitionsResponse.pb( return_value ) - json_return_value = json_format.MessageToJson(pb_return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value @@ -5316,17 +5340,14 @@ def test_read_change_stream_rest(request_type): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = bigtable.ReadChangeStreamResponse( - data_change=bigtable.ReadChangeStreamResponse.DataChange( - type_=bigtable.ReadChangeStreamResponse.DataChange.Type.USER - ), - ) + return_value = bigtable.ReadChangeStreamResponse() # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadChangeStreamResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadChangeStreamResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) @@ -5408,8 +5429,9 @@ def test_read_change_stream_rest_required_fields( response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadChangeStreamResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadChangeStreamResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") @@ -5539,8 +5561,9 @@ def test_read_change_stream_rest_flattened(): # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 - pb_return_value = bigtable.ReadChangeStreamResponse.pb(return_value) - json_return_value = json_format.MessageToJson(pb_return_value) + # Convert return value to protobuf type + return_value = bigtable.ReadChangeStreamResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) json_return_value = "[{}]".format(json_return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value diff --git a/tests/unit/test_packaging.py b/tests/unit/test_packaging.py new file mode 100644 index 000000000..93fa4d1c3 --- /dev/null +++ b/tests/unit/test_packaging.py @@ -0,0 +1,37 @@ +# Copyright 2023 Google LLC +# +# Licensed 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. + +import os +import subprocess +import sys + + +def test_namespace_package_compat(tmp_path): + # The ``google`` namespace package should not be masked + # by the presence of ``google-cloud-bigtable``. + google = tmp_path / "google" + google.mkdir() + google.joinpath("othermod.py").write_text("") + env = dict(os.environ, PYTHONPATH=str(tmp_path)) + cmd = [sys.executable, "-m", "google.othermod"] + subprocess.check_call(cmd, env=env) + + # The ``google.cloud`` namespace package should not be masked + # by the presence of ``google-cloud-bigtable``. + google_cloud = tmp_path / "google" / "cloud" + google_cloud.mkdir() + google_cloud.joinpath("othermod.py").write_text("") + env = dict(os.environ, PYTHONPATH=str(tmp_path)) + cmd = [sys.executable, "-m", "google.cloud.othermod"] + subprocess.check_call(cmd, env=env) diff --git a/tests/unit/v2_client/test_batcher.py b/tests/unit/v2_client/test_batcher.py index ab511e030..fcf606972 100644 --- a/tests/unit/v2_client/test_batcher.py +++ b/tests/unit/v2_client/test_batcher.py @@ -59,7 +59,6 @@ def callback_fn(response): def test_mutation_batcher_mutate_row(): table = _Table(TABLE_NAME) with MutationsBatcher(table=table) as mutation_batcher: - rows = [ DirectRow(row_key=b"row_key"), DirectRow(row_key=b"row_key_2"), @@ -75,7 +74,6 @@ def test_mutation_batcher_mutate_row(): def test_mutation_batcher_mutate(): table = _Table(TABLE_NAME) with MutationsBatcher(table=table) as mutation_batcher: - row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", 1) row.set_cell("cf1", b"c2", 2) @@ -98,7 +96,6 @@ def test_mutation_batcher_flush_w_no_rows(): def test_mutation_batcher_mutate_w_max_flush_count(): table = _Table(TABLE_NAME) with MutationsBatcher(table=table, flush_count=3) as mutation_batcher: - row_1 = DirectRow(row_key=b"row_key_1") row_2 = DirectRow(row_key=b"row_key_2") row_3 = DirectRow(row_key=b"row_key_3") @@ -114,7 +111,6 @@ def test_mutation_batcher_mutate_w_max_flush_count(): def test_mutation_batcher_mutate_w_max_mutations(): table = _Table(TABLE_NAME) with MutationsBatcher(table=table) as mutation_batcher: - row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", 1) row.set_cell("cf1", b"c2", 2) @@ -130,7 +126,6 @@ def test_mutation_batcher_mutate_w_max_row_bytes(): with MutationsBatcher( table=table, max_row_bytes=3 * 1024 * 1024 ) as mutation_batcher: - number_of_bytes = 1 * 1024 * 1024 max_value = b"1" * number_of_bytes @@ -168,7 +163,6 @@ def test_mutations_batcher_context_manager_flushed_when_closed(): with MutationsBatcher( table=table, max_row_bytes=3 * 1024 * 1024 ) as mutation_batcher: - number_of_bytes = 1 * 1024 * 1024 max_value = b"1" * number_of_bytes @@ -204,21 +198,22 @@ def test_mutations_batcher_response_with_error_codes(): mocked_response = [Status(code=1), Status(code=5)] - table = mock.Mock() - mutation_batcher = MutationsBatcher(table=table) + with mock.patch("tests.unit.v2_client.test_batcher._Table") as mocked_table: + table = mocked_table.return_value + mutation_batcher = MutationsBatcher(table=table) - row1 = DirectRow(row_key=b"row_key") - row2 = DirectRow(row_key=b"row_key") - table.mutate_rows.return_value = mocked_response + row1 = DirectRow(row_key=b"row_key") + row2 = DirectRow(row_key=b"row_key") + table.mutate_rows.return_value = mocked_response - mutation_batcher.mutate_rows([row1, row2]) - with pytest.raises(MutationsBatchError) as exc: - mutation_batcher.close() - assert exc.value.message == "Errors in batch mutations." - assert len(exc.value.exc) == 2 + mutation_batcher.mutate_rows([row1, row2]) + with pytest.raises(MutationsBatchError) as exc: + mutation_batcher.close() + assert exc.value.message == "Errors in batch mutations." + assert len(exc.value.exc) == 2 - assert exc.value.exc[0].message == mocked_response[0].message - assert exc.value.exc[1].message == mocked_response[1].message + assert exc.value.exc[0].message == mocked_response[0].message + assert exc.value.exc[1].message == mocked_response[1].message def test_flow_control_event_is_set_when_not_blocked(): diff --git a/tests/unit/v2_client/test_cluster.py b/tests/unit/v2_client/test_cluster.py index cb0312b0c..65ed47437 100644 --- a/tests/unit/v2_client/test_cluster.py +++ b/tests/unit/v2_client/test_cluster.py @@ -752,7 +752,6 @@ def test_cluster_update_w_partial_autoscaling_config(): }, ] for config in cluster_config: - cluster = _make_cluster( CLUSTER_ID, instance, @@ -927,7 +926,6 @@ def test_cluster_disable_autoscaling(): def test_create_cluster_with_both_manual_and_autoscaling(): - from google.cloud.bigtable.instance import Instance from google.cloud.bigtable.enums import StorageType @@ -955,7 +953,6 @@ def test_create_cluster_with_both_manual_and_autoscaling(): def test_create_cluster_with_partial_autoscaling_config(): - from google.cloud.bigtable.instance import Instance from google.cloud.bigtable.enums import StorageType @@ -996,7 +993,6 @@ def test_create_cluster_with_partial_autoscaling_config(): def test_create_cluster_with_no_scaling_config(): - from google.cloud.bigtable.instance import Instance from google.cloud.bigtable.enums import StorageType diff --git a/tests/unit/v2_client/test_column_family.py b/tests/unit/v2_client/test_column_family.py index b164b2fc1..e4f74e264 100644 --- a/tests/unit/v2_client/test_column_family.py +++ b/tests/unit/v2_client/test_column_family.py @@ -595,7 +595,6 @@ def test__gc_rule_from_pb_unknown_field_name(): from google.cloud.bigtable.column_family import _gc_rule_from_pb class MockProto(object): - names = [] _pb = {} diff --git a/tests/unit/v2_client/test_instance.py b/tests/unit/v2_client/test_instance.py index c577adca5..797e4bd9c 100644 --- a/tests/unit/v2_client/test_instance.py +++ b/tests/unit/v2_client/test_instance.py @@ -67,7 +67,6 @@ def _make_instance(*args, **kwargs): def test_instance_constructor_defaults(): - client = object() instance = _make_instance(INSTANCE_ID, client) assert instance.instance_id == INSTANCE_ID diff --git a/tests/unit/v2_client/test_row_data.py b/tests/unit/v2_client/test_row_data.py index fba69ceba..7c2987b56 100644 --- a/tests/unit/v2_client/test_row_data.py +++ b/tests/unit/v2_client/test_row_data.py @@ -362,6 +362,30 @@ def test__retry_read_rows_exception_deadline_exceeded_wrapped_in_grpc(): assert _retry_read_rows_exception(exception) +def test_partial_cell_data(): + from google.cloud.bigtable.row_data import PartialCellData + + expected_key = b"row-key" + expected_family_name = b"family-name" + expected_qualifier = b"qualifier" + expected_timestamp = 1234 + instance = PartialCellData( + expected_key, expected_family_name, expected_qualifier, expected_timestamp + ) + assert instance.row_key == expected_key + assert instance.family_name == expected_family_name + assert instance.qualifier == expected_qualifier + assert instance.timestamp_micros == expected_timestamp + assert instance.value == b"" + assert instance.labels == () + # test updating value + added_value = b"added-value" + instance.append_value(added_value) + assert instance.value == added_value + instance.append_value(added_value) + assert instance.value == added_value + added_value + + def _make_partial_rows_data(*args, **kwargs): from google.cloud.bigtable.row_data import PartialRowsData @@ -1118,7 +1142,6 @@ def test_RRRM_build_updated_request_row_ranges_valid(): class _MockCancellableIterator(object): - cancel_calls = 0 def __init__(self, *values): @@ -1199,5 +1222,4 @@ def _read_rows_retry_exception(exc): class _Client(object): - data_stub = None diff --git a/tests/unit/v2_client/test_table.py b/tests/unit/v2_client/test_table.py index 3d7d2e8ee..032363bd7 100644 --- a/tests/unit/v2_client/test_table.py +++ b/tests/unit/v2_client/test_table.py @@ -1067,8 +1067,8 @@ def test_table_yield_retry_rows(): for row in table.yield_rows(start_key=ROW_KEY_1, end_key=ROW_KEY_2): rows.append(row) - assert len(warned) == 1 - assert warned[0].category is DeprecationWarning + assert len(warned) >= 1 + assert DeprecationWarning in [w.category for w in warned] result = rows[1] assert result.row_key == ROW_KEY_2 @@ -1140,8 +1140,8 @@ def test_table_yield_rows_with_row_set(): for row in table.yield_rows(row_set=row_set): rows.append(row) - assert len(warned) == 1 - assert warned[0].category is DeprecationWarning + assert len(warned) >= 1 + assert DeprecationWarning in [w.category for w in warned] assert rows[0].row_key == ROW_KEY_1 assert rows[1].row_key == ROW_KEY_2 @@ -1689,7 +1689,6 @@ def _do_mutate_retryable_rows_helper( expected_entries = [] for row, prior_status in zip(rows, worker.responses_statuses): - if prior_status is None or prior_status.code in RETRYABLES: mutations = row._get_mutations().copy() # row clears on success entry = data_messages_v2_pb2.MutateRowsRequest.Entry( From cc79d8ccf8feae2fc9afada386c7aef2530bb021 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 15 Dec 2023 15:38:11 -0800 Subject: [PATCH 32/56] chore: pin conformance tests to v0.0.2 (#903) --- .github/workflows/conformance.yaml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/conformance.yaml b/.github/workflows/conformance.yaml index 69350b18d..63023d162 100644 --- a/.github/workflows/conformance.yaml +++ b/.github/workflows/conformance.yaml @@ -24,10 +24,11 @@ jobs: runs-on: ubuntu-latest strategy: matrix: + test-version: [ "v0.0.2" ] py-version: [ 3.8 ] client-type: [ "Async v3", "Legacy" ] fail-fast: false - name: "${{ matrix.client-type }} Client / Python ${{ matrix.py-version }}" + name: "${{ matrix.client-type }} Client / Python ${{ matrix.py-version }} / Test Tag ${{ matrix.test-version }}" steps: - uses: actions/checkout@v3 name: "Checkout python-bigtable" @@ -35,7 +36,7 @@ jobs: name: "Checkout conformance tests" with: repository: googleapis/cloud-bigtable-clients-test - ref: main + ref: ${{ matrix.test-version }} path: cloud-bigtable-clients-test - uses: actions/setup-python@v4 with: From f0d75de0a9f4d5a7ff6cda96b911626bd41538ce Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 19 Dec 2023 11:33:18 -0800 Subject: [PATCH 33/56] fix: bulk mutation eventual success (#909) --- .../bigtable/data/_async/_mutate_rows.py | 3 +++ tests/system/data/test_system.py | 22 +++++++++++++++++ tests/unit/data/_async/test_client.py | 24 +++++++++++++++++++ 3 files changed, 49 insertions(+) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index d4ffdee22..7d1144553 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -189,6 +189,9 @@ async def _run_attempt(self): if result.status.code != 0: # mutation failed; update error list (and remaining_indices if retryable) self._handle_entry_error(orig_idx, entry_error) + elif orig_idx in self.errors: + # mutation succeeded; remove from error list + del self.errors[orig_idx] # remove processed entry from active list del active_request_indices[result.index] except Exception as exc: diff --git a/tests/system/data/test_system.py b/tests/system/data/test_system.py index fb0d9eb82..b8423279a 100644 --- a/tests/system/data/test_system.py +++ b/tests/system/data/test_system.py @@ -238,6 +238,28 @@ async def test_bulk_mutations_set_cell(client, table, temp_rows): assert (await _retrieve_cell_value(table, row_key)) == new_value +@pytest.mark.asyncio +async def test_bulk_mutations_raise_exception(client, table): + """ + If an invalid mutation is passed, an exception should be raised + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.exceptions import FailedMutationEntryError + + row_key = uuid.uuid4().hex.encode() + mutation = SetCell(family="nonexistent", qualifier=b"test-qualifier", new_value=b"") + bulk_mutation = RowMutationEntry(row_key, [mutation]) + + with pytest.raises(MutationsExceptionGroup) as exc: + await table.bulk_mutate_rows([bulk_mutation]) + assert len(exc.value.exceptions) == 1 + entry_error = exc.value.exceptions[0] + assert isinstance(entry_error, FailedMutationEntryError) + assert entry_error.index == 0 + assert entry_error.entry == bulk_mutation + + @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") @retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 46080e497..9a12abe9b 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -2659,6 +2659,30 @@ async def test_bulk_mutate_error_index(self): assert isinstance(cause.exceptions[1], DeadlineExceeded) assert isinstance(cause.exceptions[2], FailedPrecondition) + @pytest.mark.asyncio + async def test_bulk_mutate_error_recovery(self): + """ + If an error occurs, then resolves, no exception should be raised + """ + from google.api_core.exceptions import DeadlineExceeded + + async with self._make_client(project="project") as client: + table = client.get_table("instance", "table") + with mock.patch.object(client._gapic_client, "mutate_rows") as mock_gapic: + # fail with a retryable error, then a non-retryable one + mock_gapic.side_effect = [ + self._mock_response([DeadlineExceeded("mock")]), + self._mock_response([None]), + ] + mutation = mutations.SetCell( + "family", b"qualifier", b"value", timestamp_micros=123 + ) + entries = [ + mutations.RowMutationEntry((f"row_key_{i}").encode(), [mutation]) + for i in range(3) + ] + await table.bulk_mutate_rows(entries, operation_timeout=1000) + class TestCheckAndMutateRow: def _make_client(self, *args, **kwargs): From 5e111804b3c6f87361bd63fe407b252cc14de958 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 17 Jan 2024 16:36:46 -0800 Subject: [PATCH 34/56] chore: move batcher file --- tests/unit/test_batcher.py | 269 --------------------------- tests/unit/v2_client/test_batcher.py | 2 +- 2 files changed, 1 insertion(+), 270 deletions(-) delete mode 100644 tests/unit/test_batcher.py diff --git a/tests/unit/test_batcher.py b/tests/unit/test_batcher.py deleted file mode 100644 index 741d9f282..000000000 --- a/tests/unit/test_batcher.py +++ /dev/null @@ -1,269 +0,0 @@ -# Copyright 2018 Google LLC -# -# Licensed 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. - - -import mock -import time - -import pytest - -from google.cloud.bigtable.row import DirectRow -from google.cloud.bigtable.batcher import ( - _FlowControl, - MutationsBatcher, - MutationsBatchError, -) - -TABLE_ID = "table-id" -TABLE_NAME = "/tables/" + TABLE_ID - - -def test_mutation_batcher_constructor(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table) as mutation_batcher: - assert table is mutation_batcher.table - - -def test_mutation_batcher_w_user_callback(): - table = _Table(TABLE_NAME) - - def callback_fn(response): - callback_fn.count = len(response) - - with MutationsBatcher( - table, flush_count=1, batch_completed_callback=callback_fn - ) as mutation_batcher: - rows = [ - DirectRow(row_key=b"row_key"), - DirectRow(row_key=b"row_key_2"), - DirectRow(row_key=b"row_key_3"), - DirectRow(row_key=b"row_key_4"), - ] - - mutation_batcher.mutate_rows(rows) - - assert callback_fn.count == 4 - - -def test_mutation_batcher_mutate_row(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table) as mutation_batcher: - rows = [ - DirectRow(row_key=b"row_key"), - DirectRow(row_key=b"row_key_2"), - DirectRow(row_key=b"row_key_3"), - DirectRow(row_key=b"row_key_4"), - ] - - mutation_batcher.mutate_rows(rows) - - assert table.mutation_calls == 1 - - -def test_mutation_batcher_mutate(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table) as mutation_batcher: - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", 1) - row.set_cell("cf1", b"c2", 2) - row.set_cell("cf1", b"c3", 3) - row.set_cell("cf1", b"c4", 4) - - mutation_batcher.mutate(row) - - assert table.mutation_calls == 1 - - -def test_mutation_batcher_flush_w_no_rows(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table) as mutation_batcher: - mutation_batcher.flush() - - assert table.mutation_calls == 0 - - -def test_mutation_batcher_mutate_w_max_flush_count(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table, flush_count=3) as mutation_batcher: - row_1 = DirectRow(row_key=b"row_key_1") - row_2 = DirectRow(row_key=b"row_key_2") - row_3 = DirectRow(row_key=b"row_key_3") - - mutation_batcher.mutate(row_1) - mutation_batcher.mutate(row_2) - mutation_batcher.mutate(row_3) - - assert table.mutation_calls == 1 - - -@mock.patch("google.cloud.bigtable.batcher.MAX_OUTSTANDING_ELEMENTS", new=3) -def test_mutation_batcher_mutate_w_max_mutations(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table) as mutation_batcher: - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", 1) - row.set_cell("cf1", b"c2", 2) - row.set_cell("cf1", b"c3", 3) - - mutation_batcher.mutate(row) - - assert table.mutation_calls == 1 - - -def test_mutation_batcher_mutate_w_max_row_bytes(): - table = _Table(TABLE_NAME) - with MutationsBatcher( - table=table, max_row_bytes=3 * 1024 * 1024 - ) as mutation_batcher: - number_of_bytes = 1 * 1024 * 1024 - max_value = b"1" * number_of_bytes - - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", max_value) - row.set_cell("cf1", b"c2", max_value) - row.set_cell("cf1", b"c3", max_value) - - mutation_batcher.mutate(row) - - assert table.mutation_calls == 1 - - -def test_mutations_batcher_flushed_when_closed(): - table = _Table(TABLE_NAME) - mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024) - - number_of_bytes = 1 * 1024 * 1024 - max_value = b"1" * number_of_bytes - - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", max_value) - row.set_cell("cf1", b"c2", max_value) - - mutation_batcher.mutate(row) - assert table.mutation_calls == 0 - - mutation_batcher.close() - - assert table.mutation_calls == 1 - - -def test_mutations_batcher_context_manager_flushed_when_closed(): - table = _Table(TABLE_NAME) - with MutationsBatcher( - table=table, max_row_bytes=3 * 1024 * 1024 - ) as mutation_batcher: - number_of_bytes = 1 * 1024 * 1024 - max_value = b"1" * number_of_bytes - - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", max_value) - row.set_cell("cf1", b"c2", max_value) - - mutation_batcher.mutate(row) - - assert table.mutation_calls == 1 - - -@mock.patch("google.cloud.bigtable.batcher.MutationsBatcher.flush") -def test_mutations_batcher_flush_interval(mocked_flush): - table = _Table(TABLE_NAME) - flush_interval = 0.5 - mutation_batcher = MutationsBatcher(table=table, flush_interval=flush_interval) - - assert mutation_batcher._timer.interval == flush_interval - mocked_flush.assert_not_called() - - time.sleep(0.4) - mocked_flush.assert_not_called() - - time.sleep(0.1) - mocked_flush.assert_called_once_with() - - mutation_batcher.close() - - -def test_mutations_batcher_response_with_error_codes(): - from google.rpc.status_pb2 import Status - - mocked_response = [Status(code=1), Status(code=5)] - - with mock.patch("tests.unit.test_batcher._Table") as mocked_table: - table = mocked_table.return_value - mutation_batcher = MutationsBatcher(table=table) - - row1 = DirectRow(row_key=b"row_key") - row2 = DirectRow(row_key=b"row_key") - table.mutate_rows.return_value = mocked_response - - mutation_batcher.mutate_rows([row1, row2]) - with pytest.raises(MutationsBatchError) as exc: - mutation_batcher.close() - assert exc.value.message == "Errors in batch mutations." - assert len(exc.value.exc) == 2 - - assert exc.value.exc[0].message == mocked_response[0].message - assert exc.value.exc[1].message == mocked_response[1].message - - -def test_flow_control_event_is_set_when_not_blocked(): - flow_control = _FlowControl() - - flow_control.set_flow_control_status() - assert flow_control.event.is_set() - - -def test_flow_control_event_is_not_set_when_blocked(): - flow_control = _FlowControl() - - flow_control.inflight_mutations = flow_control.max_mutations - flow_control.inflight_size = flow_control.max_mutation_bytes - - flow_control.set_flow_control_status() - assert not flow_control.event.is_set() - - -@mock.patch("concurrent.futures.ThreadPoolExecutor.submit") -def test_flush_async_batch_count(mocked_executor_submit): - table = _Table(TABLE_NAME) - mutation_batcher = MutationsBatcher(table=table, flush_count=2) - - number_of_bytes = 1 * 1024 * 1024 - max_value = b"1" * number_of_bytes - for index in range(5): - row = DirectRow(row_key=f"row_key_{index}") - row.set_cell("cf1", b"c1", max_value) - mutation_batcher.mutate(row) - mutation_batcher._flush_async() - - # 3 batches submitted. 2 batches of 2 items, and the last one a single item batch. - assert mocked_executor_submit.call_count == 3 - - -class _Instance(object): - def __init__(self, client=None): - self._client = client - - -class _Table(object): - def __init__(self, name, client=None): - self.name = name - self._instance = _Instance(client) - self.mutation_calls = 0 - - def mutate_rows(self, rows): - from google.rpc.status_pb2 import Status - - self.mutation_calls += 1 - - return [Status(code=0) for _ in rows] diff --git a/tests/unit/v2_client/test_batcher.py b/tests/unit/v2_client/test_batcher.py index fcf606972..741d9f282 100644 --- a/tests/unit/v2_client/test_batcher.py +++ b/tests/unit/v2_client/test_batcher.py @@ -198,7 +198,7 @@ def test_mutations_batcher_response_with_error_codes(): mocked_response = [Status(code=1), Status(code=5)] - with mock.patch("tests.unit.v2_client.test_batcher._Table") as mocked_table: + with mock.patch("tests.unit.test_batcher._Table") as mocked_table: table = mocked_table.return_value mutation_batcher = MutationsBatcher(table=table) From 54dbc0814ecde6f43020cc0a335ca0b7a83dd61e Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 17 Jan 2024 16:52:38 -0800 Subject: [PATCH 35/56] chore(tests): fixed failing test --- tests/unit/v2_client/test_batcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/v2_client/test_batcher.py b/tests/unit/v2_client/test_batcher.py index 741d9f282..fcf606972 100644 --- a/tests/unit/v2_client/test_batcher.py +++ b/tests/unit/v2_client/test_batcher.py @@ -198,7 +198,7 @@ def test_mutations_batcher_response_with_error_codes(): mocked_response = [Status(code=1), Status(code=5)] - with mock.patch("tests.unit.test_batcher._Table") as mocked_table: + with mock.patch("tests.unit.v2_client.test_batcher._Table") as mocked_table: table = mocked_table.return_value mutation_batcher = MutationsBatcher(table=table) From fd4fba3d58aaf97db727d4c7fb357d7045b68adb Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 14:14:20 -0800 Subject: [PATCH 36/56] chore: downgraded system_emulated --- .github/workflows/system_emulated.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/system_emulated.yml b/.github/workflows/system_emulated.yml index ceb4e0c4d..7669901c9 100644 --- a/.github/workflows/system_emulated.yml +++ b/.github/workflows/system_emulated.yml @@ -20,7 +20,7 @@ jobs: python-version: '3.8' - name: Setup GCloud SDK - uses: google-github-actions/setup-gcloud@v2.0.1 + uses: google-github-actions/setup-gcloud@v2.0.0 - name: Install / run Nox run: | From f4ac54fc233c651a76c352058bf39ff8e6a47766 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 14:26:06 -0800 Subject: [PATCH 37/56] fix: use default project in emulator mode --- google/cloud/bigtable/data/_async/client.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index d0578ff1a..da54b37cb 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -54,6 +54,7 @@ import google.auth.credentials import google.auth._default from google.api_core import client_options as client_options_lib +from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT from google.cloud.bigtable.data.row import Row from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data.exceptions import FailedQueryShardError @@ -132,9 +133,12 @@ def __init__( Optional[client_options_lib.ClientOptions], client_options ) self._emulator_host = os.getenv(BIGTABLE_EMULATOR) - if self._emulator_host is not None and credentials is None: + if self._emulator_host is not None: # use insecure channel if emulator is set - credentials = google.auth.credentials.AnonymousCredentials() + if credentials is None: + credentials = google.auth.credentials.AnonymousCredentials() + if project is None: + project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT # initialize client ClientWithProject.__init__( self, From 79157c97320a95f19f7bb1eae7d2b598bb236fe3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 14:38:28 -0800 Subject: [PATCH 38/56] chore(tests): system tests support emulator --- tests/system/data/setup_fixtures.py | 36 +++++++++++------------------ 1 file changed, 14 insertions(+), 22 deletions(-) diff --git a/tests/system/data/setup_fixtures.py b/tests/system/data/setup_fixtures.py index 7a7faa5f5..42eb68df7 100644 --- a/tests/system/data/setup_fixtures.py +++ b/tests/system/data/setup_fixtures.py @@ -32,25 +32,17 @@ def event_loop(): @pytest.fixture(scope="session") -def instance_admin_client(): - """Client for interacting with the Instance Admin API.""" - from google.cloud.bigtable_admin_v2 import BigtableInstanceAdminClient - - with BigtableInstanceAdminClient() as client: - yield client - - -@pytest.fixture(scope="session") -def table_admin_client(): - """Client for interacting with the Table Admin API.""" - from google.cloud.bigtable_admin_v2 import BigtableTableAdminClient - - with BigtableTableAdminClient() as client: - yield client +def admin_client(): + """ + Client for interacting with Table and Instance admin APIs + """ + from google.cloud.bigtable.client import Client + client = Client(admin=True) + yield client @pytest.fixture(scope="session") -def instance_id(instance_admin_client, project_id, cluster_config): +def instance_id(admin_client, project_id, cluster_config): """ Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session """ @@ -67,7 +59,7 @@ def instance_id(instance_admin_client, project_id, cluster_config): # create a new temporary test instance instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" try: - operation = instance_admin_client.create_instance( + operation = admin_client.instance_admin_client.create_instance( parent=f"projects/{project_id}", instance_id=instance_id, instance=types.Instance( @@ -80,7 +72,7 @@ def instance_id(instance_admin_client, project_id, cluster_config): except exceptions.AlreadyExists: pass yield instance_id - instance_admin_client.delete_instance( + admin_client.instance_admin_client.delete_instance( name=f"projects/{project_id}/instances/{instance_id}" ) @@ -95,7 +87,7 @@ def column_split_config(): @pytest.fixture(scope="session") def table_id( - table_admin_client, + admin_client, project_id, instance_id, column_family_config, @@ -106,7 +98,7 @@ def table_id( Returns BIGTABLE_TEST_TABLE if set, otherwise creates a new temporary table for the test session Args: - - table_admin_client: Client for interacting with the Table Admin API. Supplied by the table_admin_client fixture. + - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture. - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture. - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture. - init_column_families: A list of column families to initialize the table with, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. @@ -131,7 +123,7 @@ def table_id( try: parent_path = f"projects/{project_id}/instances/{instance_id}" print(f"Creating table: {parent_path}/tables/{init_table_id}") - table_admin_client.create_table( + admin_client.table_admin_client.create_table( request={ "parent": parent_path, "table_id": init_table_id, @@ -145,7 +137,7 @@ def table_id( yield init_table_id print(f"Deleting table: {parent_path}/tables/{init_table_id}") try: - table_admin_client.delete_table(name=f"{parent_path}/tables/{init_table_id}") + admin_client.table_admin_client.delete_table(name=f"{parent_path}/tables/{init_table_id}") except exceptions.NotFound: print(f"Table {init_table_id} not found, skipping deletion") From b5c533b935a9b0776732d7f09500a0729209845c Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 15:30:20 -0800 Subject: [PATCH 39/56] chore(tests): use legacy client to spin up instance in system tests --- tests/system/data/setup_fixtures.py | 42 ++++++++++++----------------- 1 file changed, 17 insertions(+), 25 deletions(-) diff --git a/tests/system/data/setup_fixtures.py b/tests/system/data/setup_fixtures.py index 42eb68df7..075465397 100644 --- a/tests/system/data/setup_fixtures.py +++ b/tests/system/data/setup_fixtures.py @@ -41,12 +41,12 @@ def admin_client(): client = Client(admin=True) yield client + @pytest.fixture(scope="session") def instance_id(admin_client, project_id, cluster_config): """ Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session """ - from google.cloud.bigtable_admin_v2 import types from google.api_core import exceptions # use user-specified instance if available @@ -57,24 +57,18 @@ def instance_id(admin_client, project_id, cluster_config): return # create a new temporary test instance - instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" try: - operation = admin_client.instance_admin_client.create_instance( - parent=f"projects/{project_id}", - instance_id=instance_id, - instance=types.Instance( - display_name="Test Instance", - labels={"python-system-test": "true"}, - ), - clusters=cluster_config, - ) - operation.result(timeout=240) + instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" + instance = admin_client.instance(instance_id, labels={"python-system-test":"true"}) + cluster_id = list(cluster_config.keys())[0] + cluster_obj = cluster_config[cluster_id] + location_id = cluster_obj.location.split("/")[-1] + cluster = instance.cluster(cluster_id, location_id=location_id, serve_nodes=cluster_obj.serve_nodes) + instance.create(clusters=[cluster]) except exceptions.AlreadyExists: pass yield instance_id - admin_client.instance_admin_client.delete_instance( - name=f"projects/{project_id}/instances/{instance_id}" - ) + instance.delete() @pytest.fixture(scope="session") @@ -109,6 +103,8 @@ def table_id( """ from google.api_core import exceptions from google.api_core import retry + from google.cloud.bigtable.column_family import MaxVersionsGCRule + import time # use user-specified instance if available user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") @@ -123,21 +119,17 @@ def table_id( try: parent_path = f"projects/{project_id}/instances/{instance_id}" print(f"Creating table: {parent_path}/tables/{init_table_id}") - admin_client.table_admin_client.create_table( - request={ - "parent": parent_path, - "table_id": init_table_id, - "table": {"column_families": column_family_config}, - "initial_splits": [{"key": key} for key in column_split_config], - }, - retry=retry, - ) + instance = admin_client.instance(instance_id) + table = instance.table(init_table_id) + gc_rule = MaxVersionsGCRule(10) + # run rpc with retry on FailedPrecondition (when cluster is not ready) + retry(table.create)(column_families={k:gc_rule for k,_ in column_family_config.items()}, initial_split_keys=column_split_config) except exceptions.AlreadyExists: pass yield init_table_id print(f"Deleting table: {parent_path}/tables/{init_table_id}") try: - admin_client.table_admin_client.delete_table(name=f"{parent_path}/tables/{init_table_id}") + table.delete() except exceptions.NotFound: print(f"Table {init_table_id} not found, skipping deletion") From 54462e29650297c313a82823f56fb45e5a98cac5 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 15:36:47 -0800 Subject: [PATCH 40/56] chore(tests): revert back to gapic clients This reverts commit b5c533b935a9b0776732d7f09500a0729209845c. --- tests/system/data/setup_fixtures.py | 42 +++++++++++++++++------------ 1 file changed, 25 insertions(+), 17 deletions(-) diff --git a/tests/system/data/setup_fixtures.py b/tests/system/data/setup_fixtures.py index 075465397..42eb68df7 100644 --- a/tests/system/data/setup_fixtures.py +++ b/tests/system/data/setup_fixtures.py @@ -41,12 +41,12 @@ def admin_client(): client = Client(admin=True) yield client - @pytest.fixture(scope="session") def instance_id(admin_client, project_id, cluster_config): """ Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session """ + from google.cloud.bigtable_admin_v2 import types from google.api_core import exceptions # use user-specified instance if available @@ -57,18 +57,24 @@ def instance_id(admin_client, project_id, cluster_config): return # create a new temporary test instance + instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" try: - instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" - instance = admin_client.instance(instance_id, labels={"python-system-test":"true"}) - cluster_id = list(cluster_config.keys())[0] - cluster_obj = cluster_config[cluster_id] - location_id = cluster_obj.location.split("/")[-1] - cluster = instance.cluster(cluster_id, location_id=location_id, serve_nodes=cluster_obj.serve_nodes) - instance.create(clusters=[cluster]) + operation = admin_client.instance_admin_client.create_instance( + parent=f"projects/{project_id}", + instance_id=instance_id, + instance=types.Instance( + display_name="Test Instance", + labels={"python-system-test": "true"}, + ), + clusters=cluster_config, + ) + operation.result(timeout=240) except exceptions.AlreadyExists: pass yield instance_id - instance.delete() + admin_client.instance_admin_client.delete_instance( + name=f"projects/{project_id}/instances/{instance_id}" + ) @pytest.fixture(scope="session") @@ -103,8 +109,6 @@ def table_id( """ from google.api_core import exceptions from google.api_core import retry - from google.cloud.bigtable.column_family import MaxVersionsGCRule - import time # use user-specified instance if available user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") @@ -119,17 +123,21 @@ def table_id( try: parent_path = f"projects/{project_id}/instances/{instance_id}" print(f"Creating table: {parent_path}/tables/{init_table_id}") - instance = admin_client.instance(instance_id) - table = instance.table(init_table_id) - gc_rule = MaxVersionsGCRule(10) - # run rpc with retry on FailedPrecondition (when cluster is not ready) - retry(table.create)(column_families={k:gc_rule for k,_ in column_family_config.items()}, initial_split_keys=column_split_config) + admin_client.table_admin_client.create_table( + request={ + "parent": parent_path, + "table_id": init_table_id, + "table": {"column_families": column_family_config}, + "initial_splits": [{"key": key} for key in column_split_config], + }, + retry=retry, + ) except exceptions.AlreadyExists: pass yield init_table_id print(f"Deleting table: {parent_path}/tables/{init_table_id}") try: - table.delete() + admin_client.table_admin_client.delete_table(name=f"{parent_path}/tables/{init_table_id}") except exceptions.NotFound: print(f"Table {init_table_id} not found, skipping deletion") From aeea5c6384ed4df369f37889404759fd18c4f09b Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 15:59:29 -0800 Subject: [PATCH 41/56] fixed support for system_emulated --- tests/system/data/setup_fixtures.py | 37 ++++++++++++++++------------- tests/system/data/test_system.py | 9 +++++-- 2 files changed, 28 insertions(+), 18 deletions(-) diff --git a/tests/system/data/setup_fixtures.py b/tests/system/data/setup_fixtures.py index 42eb68df7..a2858a29d 100644 --- a/tests/system/data/setup_fixtures.py +++ b/tests/system/data/setup_fixtures.py @@ -48,6 +48,7 @@ def instance_id(admin_client, project_id, cluster_config): """ from google.cloud.bigtable_admin_v2 import types from google.api_core import exceptions + from google.cloud.environment_vars import BIGTABLE_EMULATOR # use user-specified instance if available user_specified_instance = os.getenv("BIGTABLE_TEST_INSTANCE") @@ -58,23 +59,27 @@ def instance_id(admin_client, project_id, cluster_config): # create a new temporary test instance instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" - try: - operation = admin_client.instance_admin_client.create_instance( - parent=f"projects/{project_id}", - instance_id=instance_id, - instance=types.Instance( - display_name="Test Instance", - labels={"python-system-test": "true"}, - ), - clusters=cluster_config, + if os.getenv(BIGTABLE_EMULATOR): + # don't create instance if in emulator mode + yield instance_id + else: + try: + operation = admin_client.instance_admin_client.create_instance( + parent=f"projects/{project_id}", + instance_id=instance_id, + instance=types.Instance( + display_name="Test Instance", + # labels={"python-system-test": "true"}, + ), + clusters=cluster_config, + ) + operation.result(timeout=240) + except exceptions.AlreadyExists: + pass + yield instance_id + admin_client.instance_admin_client.delete_instance( + name=f"projects/{project_id}/instances/{instance_id}" ) - operation.result(timeout=240) - except exceptions.AlreadyExists: - pass - yield instance_id - admin_client.instance_admin_client.delete_instance( - name=f"projects/{project_id}/instances/{instance_id}" - ) @pytest.fixture(scope="session") diff --git a/tests/system/data/test_system.py b/tests/system/data/test_system.py index b8423279a..2a7adee35 100644 --- a/tests/system/data/test_system.py +++ b/tests/system/data/test_system.py @@ -16,10 +16,12 @@ import pytest_asyncio import asyncio import uuid +import os from google.api_core import retry from google.api_core.exceptions import ClientError from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE +from google.cloud.environment_vars import BIGTABLE_EMULATOR TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" @@ -195,6 +197,7 @@ async def test_mutation_set_cell(table, temp_rows): assert (await _retrieve_cell_value(table, row_key)) == new_value +@pytest.mark.skipif(bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't use splits") @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") @retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @@ -808,6 +811,7 @@ async def test_read_row(table, temp_rows): assert row.cells[0].value == b"value" +@pytest.mark.skipif(bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't raise InvalidArgument") @pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_read_row_missing(table): @@ -821,7 +825,7 @@ async def test_read_row_missing(table): assert result is None with pytest.raises(exceptions.InvalidArgument) as e: await table.read_row("") - assert "Row key must be non-empty" in str(e) + assert "Row keys must be non-empty" in str(e) @pytest.mark.usefixtures("table") @@ -843,6 +847,7 @@ async def test_read_row_w_filter(table, temp_rows): assert row.cells[0].labels == [expected_label] +@pytest.mark.skipif(bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't raise InvalidArgument") @pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_row_exists(table, temp_rows): @@ -860,7 +865,7 @@ async def test_row_exists(table, temp_rows): assert await table.row_exists(b"3") is True with pytest.raises(exceptions.InvalidArgument) as e: await table.row_exists("") - assert "Row kest must be non-empty" in str(e) + assert "Row keys must be non-empty" in str(e) @pytest.mark.usefixtures("table") From ef29512d3c3b0a2192293d9e1f8748b6f79033c5 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 16:00:20 -0800 Subject: [PATCH 42/56] chore: ran blacken --- tests/system/data/setup_fixtures.py | 5 ++++- tests/system/data/test_system.py | 14 +++++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/system/data/setup_fixtures.py b/tests/system/data/setup_fixtures.py index a2858a29d..77086b7f3 100644 --- a/tests/system/data/setup_fixtures.py +++ b/tests/system/data/setup_fixtures.py @@ -41,6 +41,7 @@ def admin_client(): client = Client(admin=True) yield client + @pytest.fixture(scope="session") def instance_id(admin_client, project_id, cluster_config): """ @@ -142,7 +143,9 @@ def table_id( yield init_table_id print(f"Deleting table: {parent_path}/tables/{init_table_id}") try: - admin_client.table_admin_client.delete_table(name=f"{parent_path}/tables/{init_table_id}") + admin_client.table_admin_client.delete_table( + name=f"{parent_path}/tables/{init_table_id}" + ) except exceptions.NotFound: print(f"Table {init_table_id} not found, skipping deletion") diff --git a/tests/system/data/test_system.py b/tests/system/data/test_system.py index 2a7adee35..4c977f00a 100644 --- a/tests/system/data/test_system.py +++ b/tests/system/data/test_system.py @@ -197,7 +197,9 @@ async def test_mutation_set_cell(table, temp_rows): assert (await _retrieve_cell_value(table, row_key)) == new_value -@pytest.mark.skipif(bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't use splits") +@pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't use splits" +) @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("table") @retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @@ -811,7 +813,10 @@ async def test_read_row(table, temp_rows): assert row.cells[0].value == b"value" -@pytest.mark.skipif(bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't raise InvalidArgument") +@pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't raise InvalidArgument", +) @pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_read_row_missing(table): @@ -847,7 +852,10 @@ async def test_read_row_w_filter(table, temp_rows): assert row.cells[0].labels == [expected_label] -@pytest.mark.skipif(bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't raise InvalidArgument") +@pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't raise InvalidArgument", +) @pytest.mark.usefixtures("table") @pytest.mark.asyncio async def test_row_exists(table, temp_rows): From 3b5b033c28468ce0aaa3b224e89e7958b5b1ea2e Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 16:08:19 -0800 Subject: [PATCH 43/56] chore(tests): use emulator mode in github actions unit tests --- .github/workflows/unittest.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index f4a337c49..739017beb 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -23,6 +23,8 @@ jobs: - name: Run unit tests env: COVERAGE_FILE: .coverage-${{ matrix.python }} + # use emulator mode to disable automatic credential detection + BIGTABLE_EMULATOR_HOST: "none" run: | nox -s unit-${{ matrix.python }} - name: Upload coverage results From 73171ebfa8078f27fb2124e611633ef08793fa99 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 16:37:51 -0800 Subject: [PATCH 44/56] run unit tests in emulator mode by default --- .github/workflows/unittest.yml | 2 -- tests/unit/data/_async/test_client.py | 34 +++++++++++++------- tests/unit/data/test_read_rows_acceptance.py | 4 ++- 3 files changed, 25 insertions(+), 15 deletions(-) diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index 739017beb..f4a337c49 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -23,8 +23,6 @@ jobs: - name: Run unit tests env: COVERAGE_FILE: .coverage-${{ matrix.python }} - # use emulator mode to disable automatic credential detection - BIGTABLE_EMULATOR_HOST: "none" run: | nox -s unit-${{ matrix.python }} - name: Upload coverage results diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 9a12abe9b..2bc77af50 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -51,8 +51,15 @@ def _get_target_class(self): return BigtableDataClientAsync - def _make_one(self, *args, **kwargs): - return self._get_target_class()(*args, **kwargs) + def _make_one(self, *args, use_emulator=True, **kwargs): + import os + env_mask = {} + # by default, use emulator mode to avoid auth issues in CI + # emulator mode must be disabled by tests that check channel pooling/refresh background tasks + if use_emulator: + env_mask["BIGTABLE_EMULATOR_HOST"] = "localhost" + with mock.patch.dict(os.environ, env_mask): + return self._get_target_class()(*args, **kwargs) @pytest.mark.asyncio async def test_ctor(self): @@ -63,8 +70,9 @@ async def test_ctor(self): project="project-id", pool_size=expected_pool_size, credentials=expected_credentials, + use_emulator=False, ) - await asyncio.sleep(0.1) + await asyncio.sleep(0) assert client.project == expected_project assert len(client.transport._grpc_channel._pool) == expected_pool_size assert not client._active_instances @@ -98,6 +106,7 @@ async def test_ctor_super_inits(self): pool_size=pool_size, credentials=credentials, client_options=options_parsed, + use_emulator=False, ) except AttributeError: pass @@ -135,7 +144,7 @@ async def test_ctor_dict_options(self): with mock.patch.object( self._get_target_class(), "_start_background_channel_refresh" ) as start_background_refresh: - client = self._make_one(client_options=client_options) + client = self._make_one(client_options=client_options, use_emulator=False) start_background_refresh.assert_called_once() await client.close() @@ -235,14 +244,15 @@ async def test_channel_pool_replace(self): @pytest.mark.filterwarnings("ignore::RuntimeWarning") def test__start_background_channel_refresh_sync(self): # should raise RuntimeError if called in a sync context - client = self._make_one(project="project-id") + client = self._make_one(project="project-id", use_emulator=False) with pytest.raises(RuntimeError): client._start_background_channel_refresh() @pytest.mark.asyncio async def test__start_background_channel_refresh_tasks_exist(self): # if tasks exist, should do nothing - client = self._make_one(project="project-id") + client = self._make_one(project="project-id", use_emulator=False) + assert len(client._channel_refresh_tasks) > 0 with mock.patch.object(asyncio, "create_task") as create_task: client._start_background_channel_refresh() create_task.assert_not_called() @@ -252,7 +262,7 @@ async def test__start_background_channel_refresh_tasks_exist(self): @pytest.mark.parametrize("pool_size", [1, 3, 7]) async def test__start_background_channel_refresh(self, pool_size): # should create background tasks for each channel - client = self._make_one(project="project-id", pool_size=pool_size) + client = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) ping_and_warm = AsyncMock() client._ping_and_warm_instances = ping_and_warm client._start_background_channel_refresh() @@ -272,7 +282,7 @@ async def test__start_background_channel_refresh(self, pool_size): async def test__start_background_channel_refresh_tasks_names(self): # if tasks exist, should do nothing pool_size = 3 - client = self._make_one(project="project-id", pool_size=pool_size) + client = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) for i in range(pool_size): name = client._channel_refresh_tasks[i].get_name() assert str(i) in name @@ -537,7 +547,7 @@ async def test__manage_channel_refresh(self, num_cycles): grpc_helpers_async, "create_channel" ) as create_channel: create_channel.return_value = new_channel - client = self._make_one(project="project-id") + client = self._make_one(project="project-id", use_emulator=False) create_channel.reset_mock() try: await client._manage_channel( @@ -919,9 +929,9 @@ async def test_multiple_pool_sizes(self): # should be able to create multiple clients with different pool sizes without issue pool_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256] for pool_size in pool_sizes: - client = self._make_one(project="project-id", pool_size=pool_size) + client = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) assert len(client._channel_refresh_tasks) == pool_size - client_duplicate = self._make_one(project="project-id", pool_size=pool_size) + client_duplicate = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) assert len(client_duplicate._channel_refresh_tasks) == pool_size assert str(pool_size) in str(client.transport) await client.close() @@ -934,7 +944,7 @@ async def test_close(self): ) pool_size = 7 - client = self._make_one(project="project-id", pool_size=pool_size) + client = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) assert len(client._channel_refresh_tasks) == pool_size tasks_list = list(client._channel_refresh_tasks) for task in client._channel_refresh_tasks: diff --git a/tests/unit/data/test_read_rows_acceptance.py b/tests/unit/data/test_read_rows_acceptance.py index 15680984b..7cb3c08dc 100644 --- a/tests/unit/data/test_read_rows_acceptance.py +++ b/tests/unit/data/test_read_rows_acceptance.py @@ -119,7 +119,9 @@ def cancel(self): return mock_stream(chunk_list) try: - client = BigtableDataClientAsync() + with mock.patch.dict(os.environ, {"BIGTABLE_EMULATOR_HOST": "localhost"}): + # use emulator mode to avoid auth issues in CI + client = BigtableDataClientAsync() table = client.get_table("instance", "table") results = [] with mock.patch.object(table.client._gapic_client, "read_rows") as read_rows: From 6fa49ab7a1a5d8c15827f5ee6b31630ed9150699 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 17:15:41 -0800 Subject: [PATCH 45/56] iterating on tests --- tests/system/data/test_system.py | 21 +++++++++++++++++++++ tests/unit/data/_async/test_client.py | 20 ++++---------------- 2 files changed, 25 insertions(+), 16 deletions(-) diff --git a/tests/system/data/test_system.py b/tests/system/data/test_system.py index 4c977f00a..87467f27c 100644 --- a/tests/system/data/test_system.py +++ b/tests/system/data/test_system.py @@ -595,6 +595,27 @@ async def test_check_and_mutate( expected_value = true_mutation_value if expected_result else false_mutation_value assert (await _retrieve_cell_value(table, row_key)) == expected_value +@pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't raise InvalidArgument", +) +@pytest.mark.usefixtures("client") +@pytest.mark.usefixtures("table") +@pytest.mark.asyncio +async def test_check_and_mutate_empty_request(client, table): + """ + check_and_mutate with no true or fale mutations should raise an error + """ + from google.api_core import exceptions + + with pytest.raises(exceptions.InvalidArgument) as e: + await table.check_and_mutate_row( + b'row_key', + None, + true_case_mutations=None, + false_case_mutations=None + ) + assert "No mutations provided" in str(e.value) @pytest.mark.usefixtures("table") @retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 2bc77af50..907def6af 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -58,6 +58,10 @@ def _make_one(self, *args, use_emulator=True, **kwargs): # emulator mode must be disabled by tests that check channel pooling/refresh background tasks if use_emulator: env_mask["BIGTABLE_EMULATOR_HOST"] = "localhost" + else: + # set some default values + kwargs["credentials"] = kwargs.get("credentials", AnonymousCredentials()) + kwargs["project"] = kwargs.get("project", "project-id") with mock.patch.dict(os.environ, env_mask): return self._get_target_class()(*args, **kwargs) @@ -2758,22 +2762,6 @@ async def test_check_and_mutate_bad_timeout(self): ) assert str(e.value) == "operation_timeout must be greater than 0" - @pytest.mark.asyncio - async def test_check_and_mutate_no_mutations(self): - """Requests require either true_case_mutations or false_case_mutations""" - from google.api_core.exceptions import InvalidArgument - - async with self._make_client() as client: - async with client.get_table("instance", "table") as table: - with pytest.raises(InvalidArgument) as e: - await table.check_and_mutate_row( - b"row_key", - None, - true_case_mutations=None, - false_case_mutations=None, - ) - assert "No mutations provided" in str(e.value) - @pytest.mark.asyncio async def test_check_and_mutate_single_mutations(self): """if single mutations are passed, they should be internally wrapped in a list""" From 52ec52adb539c5c9b246c0db4373f697fcd02313 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 17:22:45 -0800 Subject: [PATCH 46/56] use the same client init function in unit tests --- tests/unit/data/_async/test_client.py | 142 +++++++++++--------------- 1 file changed, 59 insertions(+), 83 deletions(-) diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 907def6af..26660ad18 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -45,25 +45,30 @@ ) +def _make_client(*args, use_emulator=True, **kwargs): + import os + from google.cloud.bigtable.data._async.client import BigtableDataClientAsync + env_mask = {} + # by default, use emulator mode to avoid auth issues in CI + # emulator mode must be disabled by tests that check channel pooling/refresh background tasks + if use_emulator: + env_mask["BIGTABLE_EMULATOR_HOST"] = "localhost" + else: + # set some default values + kwargs["credentials"] = kwargs.get("credentials", AnonymousCredentials()) + kwargs["project"] = kwargs.get("project", "project-id") + with mock.patch.dict(os.environ, env_mask): + return BigtableDataClientAsync(*args, **kwargs) + + class TestBigtableDataClientAsync: def _get_target_class(self): from google.cloud.bigtable.data._async.client import BigtableDataClientAsync return BigtableDataClientAsync - def _make_one(self, *args, use_emulator=True, **kwargs): - import os - env_mask = {} - # by default, use emulator mode to avoid auth issues in CI - # emulator mode must be disabled by tests that check channel pooling/refresh background tasks - if use_emulator: - env_mask["BIGTABLE_EMULATOR_HOST"] = "localhost" - else: - # set some default values - kwargs["credentials"] = kwargs.get("credentials", AnonymousCredentials()) - kwargs["project"] = kwargs.get("project", "project-id") - with mock.patch.dict(os.environ, env_mask): - return self._get_target_class()(*args, **kwargs) + def _make_one(self, *args, **kwargs): + return _make_client(*args, **kwargs) @pytest.mark.asyncio async def test_ctor(self): @@ -1321,11 +1326,6 @@ class TestReadRows: Tests for table.read_rows and related methods. """ - def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - - return BigtableDataClientAsync(*args, **kwargs) - def _make_table(self, *args, **kwargs): from google.cloud.bigtable.data._async.client import TableAsync @@ -1694,7 +1694,7 @@ async def test_read_rows_default_timeout_override(self): @pytest.mark.asyncio async def test_read_row(self): """Test reading a single row""" - async with self._make_client() as client: + async with _make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" with mock.patch.object(table, "read_rows") as read_rows: @@ -1722,7 +1722,7 @@ async def test_read_row(self): @pytest.mark.asyncio async def test_read_row_w_filter(self): """Test reading a single row with an added filter""" - async with self._make_client() as client: + async with _make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" with mock.patch.object(table, "read_rows") as read_rows: @@ -1755,7 +1755,7 @@ async def test_read_row_w_filter(self): @pytest.mark.asyncio async def test_read_row_no_response(self): """should return None if row does not exist""" - async with self._make_client() as client: + async with _make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" with mock.patch.object(table, "read_rows") as read_rows: @@ -1790,7 +1790,7 @@ async def test_read_row_no_response(self): @pytest.mark.asyncio async def test_row_exists(self, return_value, expected_result): """Test checking for row existence""" - async with self._make_client() as client: + async with _make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" with mock.patch.object(table, "read_rows") as read_rows: @@ -1825,14 +1825,10 @@ async def test_row_exists(self, return_value, expected_result): class TestReadRowsSharded: - def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - - return BigtableDataClientAsync(*args, **kwargs) @pytest.mark.asyncio async def test_read_rows_sharded_empty_query(self): - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with pytest.raises(ValueError) as exc: await table.read_rows_sharded([]) @@ -1843,7 +1839,7 @@ async def test_read_rows_sharded_multiple_queries(self): """ Test with multiple queries. Should return results from both """ - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( table.client._gapic_client, "read_rows" @@ -1869,7 +1865,7 @@ async def test_read_rows_sharded_multiple_queries_calls(self, n_queries): """ Each query should trigger a separate read_rows call """ - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object(table, "read_rows") as read_rows: query_list = [ReadRowsQuery() for _ in range(n_queries)] @@ -1884,7 +1880,7 @@ async def test_read_rows_sharded_errors(self): from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup from google.cloud.bigtable.data.exceptions import FailedQueryShardError - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object(table, "read_rows") as read_rows: read_rows.side_effect = RuntimeError("mock error") @@ -1915,7 +1911,7 @@ async def mock_call(*args, **kwargs): await asyncio.sleep(0.1) return [mock.Mock()] - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object(table, "read_rows") as read_rows: read_rows.side_effect = mock_call @@ -1988,10 +1984,6 @@ async def test_read_rows_sharded_batching(self): class TestSampleRowKeys: - def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - - return BigtableDataClientAsync(*args, **kwargs) async def _make_gapic_stream(self, sample_list: list[tuple[bytes, int]]): from google.cloud.bigtable_v2.types import SampleRowKeysResponse @@ -2009,7 +2001,7 @@ async def test_sample_row_keys(self): (b"test_2", 100), (b"test_3", 200), ] - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( table.client._gapic_client, "sample_row_keys", AsyncMock() @@ -2029,7 +2021,7 @@ async def test_sample_row_keys_bad_timeout(self): """ should raise error if timeout is negative """ - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with pytest.raises(ValueError) as e: await table.sample_row_keys(operation_timeout=-1) @@ -2042,7 +2034,7 @@ async def test_sample_row_keys_bad_timeout(self): async def test_sample_row_keys_default_timeout(self): """Should fallback to using table default operation_timeout""" expected_timeout = 99 - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table( "i", "t", @@ -2068,7 +2060,7 @@ async def test_sample_row_keys_gapic_params(self): expected_profile = "test1" instance = "instance_name" table_id = "my_table" - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table( instance, table_id, app_profile_id=expected_profile ) as table: @@ -2101,7 +2093,7 @@ async def test_sample_row_keys_retryable_errors(self, retryable_exception): from google.api_core.exceptions import DeadlineExceeded from google.cloud.bigtable.data.exceptions import RetryExceptionGroup - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( table.client._gapic_client, "sample_row_keys", AsyncMock() @@ -2130,7 +2122,7 @@ async def test_sample_row_keys_non_retryable_errors(self, non_retryable_exceptio """ non-retryable errors should cause a raise """ - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( table.client._gapic_client, "sample_row_keys", AsyncMock() @@ -2141,10 +2133,6 @@ async def test_sample_row_keys_non_retryable_errors(self, non_retryable_exceptio class TestMutateRow: - def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - - return BigtableDataClientAsync(*args, **kwargs) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -2167,7 +2155,7 @@ def _make_client(self, *args, **kwargs): async def test_mutate_row(self, mutation_arg): """Test mutations with no errors""" expected_attempt_timeout = 19 - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_row" @@ -2207,7 +2195,7 @@ async def test_mutate_row_retryable_errors(self, retryable_exception): from google.api_core.exceptions import DeadlineExceeded from google.cloud.bigtable.data.exceptions import RetryExceptionGroup - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_row" @@ -2237,7 +2225,7 @@ async def test_mutate_row_non_idempotent_retryable_errors( """ Non-idempotent mutations should not be retried """ - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_row" @@ -2265,7 +2253,7 @@ async def test_mutate_row_non_idempotent_retryable_errors( ) @pytest.mark.asyncio async def test_mutate_row_non_retryable_errors(self, non_retryable_exception): - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_row" @@ -2288,7 +2276,7 @@ async def test_mutate_row_non_retryable_errors(self, non_retryable_exception): async def test_mutate_row_metadata(self, include_app_profile): """request should attach metadata headers""" profile = "profile" if include_app_profile else None - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("i", "t", app_profile_id=profile) as table: with mock.patch.object( client._gapic_client, "mutate_row", AsyncMock() @@ -2310,7 +2298,7 @@ async def test_mutate_row_metadata(self, include_app_profile): @pytest.mark.parametrize("mutations", [[], None]) @pytest.mark.asyncio async def test_mutate_row_no_mutations(self, mutations): - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with pytest.raises(ValueError) as e: await table.mutate_row("key", mutations=mutations) @@ -2318,10 +2306,6 @@ async def test_mutate_row_no_mutations(self, mutations): class TestBulkMutateRows: - def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - - return BigtableDataClientAsync(*args, **kwargs) async def _mock_response(self, response_list): from google.cloud.bigtable_v2.types import MutateRowsResponse @@ -2371,7 +2355,7 @@ async def generator(): async def test_bulk_mutate_rows(self, mutation_arg): """Test mutations with no errors""" expected_attempt_timeout = 19 - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_rows" @@ -2395,7 +2379,7 @@ async def test_bulk_mutate_rows(self, mutation_arg): @pytest.mark.asyncio async def test_bulk_mutate_rows_multiple_entries(self): """Test mutations with no errors""" - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_rows" @@ -2436,7 +2420,7 @@ async def test_bulk_mutate_rows_idempotent_mutation_error_retryable( MutationsExceptionGroup, ) - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_rows" @@ -2482,7 +2466,7 @@ async def test_bulk_mutate_rows_idempotent_mutation_error_non_retryable( MutationsExceptionGroup, ) - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_rows" @@ -2522,7 +2506,7 @@ async def test_bulk_mutate_idempotent_retryable_request_errors( MutationsExceptionGroup, ) - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_rows" @@ -2560,7 +2544,7 @@ async def test_bulk_mutate_rows_non_idempotent_retryable_errors( MutationsExceptionGroup, ) - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_rows" @@ -2602,7 +2586,7 @@ async def test_bulk_mutate_rows_non_retryable_errors(self, non_retryable_excepti MutationsExceptionGroup, ) - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_rows" @@ -2638,7 +2622,7 @@ async def test_bulk_mutate_error_index(self): MutationsExceptionGroup, ) - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "mutate_rows" @@ -2680,7 +2664,7 @@ async def test_bulk_mutate_error_recovery(self): """ from google.api_core.exceptions import DeadlineExceeded - async with self._make_client(project="project") as client: + async with _make_client(project="project") as client: table = client.get_table("instance", "table") with mock.patch.object(client._gapic_client, "mutate_rows") as mock_gapic: # fail with a retryable error, then a non-retryable one @@ -2699,10 +2683,6 @@ async def test_bulk_mutate_error_recovery(self): class TestCheckAndMutateRow: - def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - - return BigtableDataClientAsync(*args, **kwargs) @pytest.mark.parametrize("gapic_result", [True, False]) @pytest.mark.asyncio @@ -2710,7 +2690,7 @@ async def test_check_and_mutate(self, gapic_result): from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse app_profile = "app_profile_id" - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table( "instance", "table", app_profile_id=app_profile ) as table: @@ -2750,7 +2730,7 @@ async def test_check_and_mutate(self, gapic_result): @pytest.mark.asyncio async def test_check_and_mutate_bad_timeout(self): """Should raise error if operation_timeout < 0""" - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with pytest.raises(ValueError) as e: await table.check_and_mutate_row( @@ -2768,7 +2748,7 @@ async def test_check_and_mutate_single_mutations(self): from google.cloud.bigtable.data.mutations import SetCell from google.cloud.bigtable_v2.types import CheckAndMutateRowResponse - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "check_and_mutate_row" @@ -2796,7 +2776,7 @@ async def test_check_and_mutate_predicate_object(self): mock_predicate = mock.Mock() predicate_pb = {"predicate": "dict"} mock_predicate._to_pb.return_value = predicate_pb - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "check_and_mutate_row" @@ -2824,7 +2804,7 @@ async def test_check_and_mutate_mutations_parsing(self): for idx, mutation in enumerate(mutations): mutation._to_pb.return_value = f"fake {idx}" mutations.append(DeleteAllFromRow()) - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "check_and_mutate_row" @@ -2852,10 +2832,6 @@ async def test_check_and_mutate_mutations_parsing(self): class TestReadModifyWriteRow: - def _make_client(self, *args, **kwargs): - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync - - return BigtableDataClientAsync(*args, **kwargs) @pytest.mark.parametrize( "call_rules,expected_rules", @@ -2883,7 +2859,7 @@ async def test_read_modify_write_call_rule_args(self, call_rules, expected_rules """ Test that the gapic call is called with given rules """ - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with mock.patch.object( client._gapic_client, "read_modify_write_row" @@ -2897,7 +2873,7 @@ async def test_read_modify_write_call_rule_args(self, call_rules, expected_rules @pytest.mark.parametrize("rules", [[], None]) @pytest.mark.asyncio async def test_read_modify_write_no_rules(self, rules): - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table") as table: with pytest.raises(ValueError) as e: await table.read_modify_write_row("key", rules=rules) @@ -2909,7 +2885,7 @@ async def test_read_modify_write_call_defaults(self): table_id = "table1" project = "project1" row_key = "row_key1" - async with self._make_client(project=project) as client: + async with _make_client(project=project) as client: async with client.get_table(instance, table_id) as table: with mock.patch.object( client._gapic_client, "read_modify_write_row" @@ -2930,7 +2906,7 @@ async def test_read_modify_write_call_overrides(self): row_key = b"row_key1" expected_timeout = 12345 profile_id = "profile1" - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table( "instance", "table_id", app_profile_id=profile_id ) as table: @@ -2951,7 +2927,7 @@ async def test_read_modify_write_call_overrides(self): @pytest.mark.asyncio async def test_read_modify_write_string_key(self): row_key = "string_row_key1" - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table_id") as table: with mock.patch.object( client._gapic_client, "read_modify_write_row" @@ -2971,7 +2947,7 @@ async def test_read_modify_write_row_building(self): from google.cloud.bigtable_v2.types import Row as RowPB mock_response = ReadModifyWriteRowResponse(row=RowPB()) - async with self._make_client() as client: + async with _make_client() as client: async with client.get_table("instance", "table_id") as table: with mock.patch.object( client._gapic_client, "read_modify_write_row" From b2bf56f58a0ef3c569d3071849645b3408ba7fce Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 17:34:09 -0800 Subject: [PATCH 47/56] use _make_client in more places --- tests/system/data/test_system.py | 7 ++-- tests/unit/data/_async/test_client.py | 46 +++++++++++++-------------- 2 files changed, 25 insertions(+), 28 deletions(-) diff --git a/tests/system/data/test_system.py b/tests/system/data/test_system.py index 87467f27c..aeb08fc1a 100644 --- a/tests/system/data/test_system.py +++ b/tests/system/data/test_system.py @@ -595,6 +595,7 @@ async def test_check_and_mutate( expected_value = true_mutation_value if expected_result else false_mutation_value assert (await _retrieve_cell_value(table, row_key)) == expected_value + @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't raise InvalidArgument", @@ -610,13 +611,11 @@ async def test_check_and_mutate_empty_request(client, table): with pytest.raises(exceptions.InvalidArgument) as e: await table.check_and_mutate_row( - b'row_key', - None, - true_case_mutations=None, - false_case_mutations=None + b"row_key", None, true_case_mutations=None, false_case_mutations=None ) assert "No mutations provided" in str(e.value) + @pytest.mark.usefixtures("table") @retry.AsyncRetry(predicate=retry.if_exception_type(ClientError), initial=1, maximum=5) @pytest.mark.asyncio diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 26660ad18..52d53d09e 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -48,6 +48,7 @@ def _make_client(*args, use_emulator=True, **kwargs): import os from google.cloud.bigtable.data._async.client import BigtableDataClientAsync + env_mask = {} # by default, use emulator mode to avoid auth issues in CI # emulator mode must be disabled by tests that check channel pooling/refresh background tasks @@ -271,7 +272,9 @@ async def test__start_background_channel_refresh_tasks_exist(self): @pytest.mark.parametrize("pool_size", [1, 3, 7]) async def test__start_background_channel_refresh(self, pool_size): # should create background tasks for each channel - client = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) + client = self._make_one( + project="project-id", pool_size=pool_size, use_emulator=False + ) ping_and_warm = AsyncMock() client._ping_and_warm_instances = ping_and_warm client._start_background_channel_refresh() @@ -291,7 +294,9 @@ async def test__start_background_channel_refresh(self, pool_size): async def test__start_background_channel_refresh_tasks_names(self): # if tasks exist, should do nothing pool_size = 3 - client = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) + client = self._make_one( + project="project-id", pool_size=pool_size, use_emulator=False + ) for i in range(pool_size): name = client._channel_refresh_tasks[i].get_name() assert str(i) in name @@ -938,9 +943,13 @@ async def test_multiple_pool_sizes(self): # should be able to create multiple clients with different pool sizes without issue pool_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256] for pool_size in pool_sizes: - client = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) + client = self._make_one( + project="project-id", pool_size=pool_size, use_emulator=False + ) assert len(client._channel_refresh_tasks) == pool_size - client_duplicate = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) + client_duplicate = self._make_one( + project="project-id", pool_size=pool_size, use_emulator=False + ) assert len(client_duplicate._channel_refresh_tasks) == pool_size assert str(pool_size) in str(client.transport) await client.close() @@ -953,7 +962,9 @@ async def test_close(self): ) pool_size = 7 - client = self._make_one(project="project-id", pool_size=pool_size, use_emulator=False) + client = self._make_one( + project="project-id", pool_size=pool_size, use_emulator=False + ) assert len(client._channel_refresh_tasks) == pool_size tasks_list = list(client._channel_refresh_tasks) for task in client._channel_refresh_tasks: @@ -1003,10 +1014,9 @@ async def test_context_manager(self): def test_client_ctor_sync(self): # initializing client in a sync context should raise RuntimeError - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync with pytest.warns(RuntimeWarning) as warnings: - client = BigtableDataClientAsync(project="project-id") + client = _make_client(project="project-id") expected_warning = [w for w in warnings if "client.py" in w.filename] assert len(expected_warning) == 1 assert ( @@ -1020,7 +1030,6 @@ def test_client_ctor_sync(self): class TestTableAsync: @pytest.mark.asyncio async def test_table_ctor(self): - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync from google.cloud.bigtable.data._async.client import TableAsync from google.cloud.bigtable.data._async.client import _WarmedInstanceKey @@ -1033,7 +1042,7 @@ async def test_table_ctor(self): expected_read_rows_attempt_timeout = 0.5 expected_mutate_rows_operation_timeout = 2.5 expected_mutate_rows_attempt_timeout = 0.75 - client = BigtableDataClientAsync() + client = _make_client() assert not client._active_instances table = TableAsync( @@ -1088,12 +1097,11 @@ async def test_table_ctor_defaults(self): """ should provide default timeout values and app_profile_id """ - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync from google.cloud.bigtable.data._async.client import TableAsync expected_table_id = "table-id" expected_instance_id = "instance-id" - client = BigtableDataClientAsync() + client = _make_client() assert not client._active_instances table = TableAsync( @@ -1119,10 +1127,9 @@ async def test_table_ctor_invalid_timeout_values(self): """ bad timeout values should raise ValueError """ - from google.cloud.bigtable.data._async.client import BigtableDataClientAsync from google.cloud.bigtable.data._async.client import TableAsync - client = BigtableDataClientAsync() + client = _make_client() timeout_pairs = [ ("default_operation_timeout", "default_attempt_timeout"), @@ -1240,10 +1247,8 @@ async def test_customizable_retryable_errors( Test that retryable functions support user-configurable arguments, and that the configured retryables are passed down to the gapic layer. """ - from google.cloud.bigtable.data import BigtableDataClientAsync - with mock.patch(retry_fn_path) as retry_fn_mock: - async with BigtableDataClientAsync() as client: + async with _make_client() as client: table = client.get_table("instance-id", "table-id") expected_predicate = lambda a: a in expected_retryables # noqa retry_fn_mock.side_effect = RuntimeError("stop early") @@ -1291,14 +1296,13 @@ async def test_customizable_retryable_errors( async def test_call_metadata(self, include_app_profile, fn_name, fn_args, gapic_fn): """check that all requests attach proper metadata headers""" from google.cloud.bigtable.data import TableAsync - from google.cloud.bigtable.data import BigtableDataClientAsync profile = "profile" if include_app_profile else None with mock.patch( f"google.cloud.bigtable_v2.BigtableAsyncClient.{gapic_fn}", mock.AsyncMock() ) as gapic_mock: gapic_mock.side_effect = RuntimeError("stop early") - async with BigtableDataClientAsync() as client: + async with _make_client() as client: table = TableAsync(client, "instance-id", "table-id", profile) try: test_fn = table.__getattribute__(fn_name) @@ -1825,7 +1829,6 @@ async def test_row_exists(self, return_value, expected_result): class TestReadRowsSharded: - @pytest.mark.asyncio async def test_read_rows_sharded_empty_query(self): async with _make_client() as client: @@ -1984,7 +1987,6 @@ async def test_read_rows_sharded_batching(self): class TestSampleRowKeys: - async def _make_gapic_stream(self, sample_list: list[tuple[bytes, int]]): from google.cloud.bigtable_v2.types import SampleRowKeysResponse @@ -2133,7 +2135,6 @@ async def test_sample_row_keys_non_retryable_errors(self, non_retryable_exceptio class TestMutateRow: - @pytest.mark.asyncio @pytest.mark.parametrize( "mutation_arg", @@ -2306,7 +2307,6 @@ async def test_mutate_row_no_mutations(self, mutations): class TestBulkMutateRows: - async def _mock_response(self, response_list): from google.cloud.bigtable_v2.types import MutateRowsResponse from google.rpc import status_pb2 @@ -2683,7 +2683,6 @@ async def test_bulk_mutate_error_recovery(self): class TestCheckAndMutateRow: - @pytest.mark.parametrize("gapic_result", [True, False]) @pytest.mark.asyncio async def test_check_and_mutate(self, gapic_result): @@ -2832,7 +2831,6 @@ async def test_check_and_mutate_mutations_parsing(self): class TestReadModifyWriteRow: - @pytest.mark.parametrize( "call_rules,expected_rules", [ From 79310695d939ad6f7e77039b925a85667b600043 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 17:38:44 -0800 Subject: [PATCH 48/56] iterating on tests --- tests/unit/data/_async/test_client.py | 2 +- tests/unit/data/_async/test_mutations_batcher.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 52d53d09e..a0019947d 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1016,7 +1016,7 @@ def test_client_ctor_sync(self): # initializing client in a sync context should raise RuntimeError with pytest.warns(RuntimeWarning) as warnings: - client = _make_client(project="project-id") + client = _make_client(project="project-id", use_emulator=False) expected_warning = [w for w in warnings if "client.py" in w.filename] assert len(expected_warning) == 1 assert ( diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py index 446cd822e..cca7c9824 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -726,7 +726,7 @@ async def mock_call(*args, **kwargs): assert len(instance._oldest_exceptions) == 0 assert len(instance._newest_exceptions) == 0 # if flushes were sequential, total duration would be 1s - assert duration < 0.25 + assert duration < 0.5 assert op_mock.call_count == num_calls @pytest.mark.asyncio From 8b886b347e063b102c5eb4992375e1f6c90a38c1 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Jan 2024 17:41:01 -0800 Subject: [PATCH 49/56] changed cover requirement --- .github/workflows/unittest.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index f4a337c49..87d08602f 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -54,4 +54,4 @@ jobs: run: | find .coverage-results -type f -name '*.zip' -exec unzip {} \; coverage combine .coverage-results/**/.coverage* - coverage report --show-missing --fail-under=100 + coverage report --show-missing --fail-under=99 From c3ed5aa29c06abc3b50278028c23ab96a242a1c3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 19 Jan 2024 13:00:25 -0800 Subject: [PATCH 50/56] updated README --- README.rst | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/README.rst b/README.rst index 5f7d5809d..1c9897eba 100644 --- a/README.rst +++ b/README.rst @@ -105,3 +105,13 @@ with the same interface as the popular `HappyBase `__ library. Unlike HappyBase, ``google-cloud-happybase`` uses ``google-cloud-bigtable`` under the covers, rather than Apache HBase. + +``experimental async data client`` +---------------------------------- + +`v2.23.0` includes a preview release of a new data client, accessible at `google.cloud.bigtable.data.BigtableDataClientAsync`. + +This new client supports asyncio, with a corresponding synchronous surface coming soon. + +The new client is currently in preview, and is not recommended for production use. + From 158ceedc9a6db7cb27f43bb8b85d607c197ab830 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 19 Jan 2024 13:24:06 -0800 Subject: [PATCH 51/56] updated owlbot file --- owlbot.py | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/owlbot.py b/owlbot.py index be28fa2a1..6a5fafe5f 100644 --- a/owlbot.py +++ b/owlbot.py @@ -90,6 +90,9 @@ def get_staging_dirs( split_system_tests=True, microgenerator=True, cov_level=99, + system_test_external_dependencies=[ + "pytest-asyncio", + ], ) s.move(templated_files, excludes=[".coveragerc", "README.rst", ".github/release-please.yml"]) @@ -142,7 +145,35 @@ def system_emulated(session): escape="()" ) -# add system_emulated nox session +conformance_session = """ +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) +def conformance(session): + TEST_REPO_URL = "https://github.com/googleapis/cloud-bigtable-clients-test.git" + CLONE_REPO_DIR = "cloud-bigtable-clients-test" + # install dependencies + constraints_path = str( + CURRENT_DIRECTORY / "testing" / f"constraints-{session.python}.txt" + ) + install_unittest_dependencies(session, "-c", constraints_path) + with session.chdir("test_proxy"): + # download the conformance test suite + clone_dir = os.path.join(CURRENT_DIRECTORY, CLONE_REPO_DIR) + if not os.path.exists(clone_dir): + print("downloading copy of test repo") + session.run("git", "clone", TEST_REPO_URL, CLONE_REPO_DIR, external=True) + session.run("bash", "-e", "run_tests.sh", external=True) + +""" + +place_before( + "noxfile.py", + "@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS)\n" + "def system(session):", + conformance_session, + escape="()" +) + +# add system_emulated and mypy and conformance to nox session s.replace("noxfile.py", """nox.options.sessions = \[ "unit", @@ -151,6 +182,7 @@ def system_emulated(session): "unit", "system_emulated", "system", + "conformance", "mypy",""", ) @@ -171,7 +203,7 @@ def mypy(session): session.run( "mypy", "-p", - "google.cloud.bigtable", + "google.cloud.bigtable.data", "--check-untyped-defs", "--warn-unreachable", "--disallow-any-generics", From 22e19473484f50e5e26138d2477dbd4e301ac7f4 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 19 Jan 2024 13:25:13 -0800 Subject: [PATCH 52/56] updated api_core version in setup.py --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 3aeb275ae..887b0199b 100644 --- a/setup.py +++ b/setup.py @@ -37,7 +37,7 @@ # 'Development Status :: 5 - Production/Stable' release_status = "Development Status :: 5 - Production/Stable" dependencies = [ - "google-api-core[grpc] >= 2.16.0rc0, <3.0.0dev,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,!=2.8.*,!=2.9.*,!=2.10.*", + "google-api-core[grpc] >= 2.16.0, <3.0.0dev,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,!=2.8.*,!=2.9.*,!=2.10.*", "google-cloud-core >= 2.0.0, <3.0.0dev", "grpc-google-iam-v1 >= 0.12.4, <1.0.0dev", "proto-plus >= 1.22.0, <2.0.0dev", From 9e60999de96afe07f028ecfd058bdef57af2a80f Mon Sep 17 00:00:00 2001 From: Owl Bot Date: Fri, 19 Jan 2024 21:27:55 +0000 Subject: [PATCH 53/56] =?UTF-8?q?=F0=9F=A6=89=20Updates=20from=20OwlBot=20?= =?UTF-8?q?post-processor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --- .kokoro/trampoline.sh | 2 +- noxfile.py | 35 ++++++++++++++--------------------- 2 files changed, 15 insertions(+), 22 deletions(-) diff --git a/.kokoro/trampoline.sh b/.kokoro/trampoline.sh index 5c7c8633a..d85b1f267 100755 --- a/.kokoro/trampoline.sh +++ b/.kokoro/trampoline.sh @@ -25,4 +25,4 @@ function cleanup() { trap cleanup EXIT $(dirname $0)/populate-secrets.sh # Secret Manager secrets. -python3 "${KOKORO_GFILE_DIR}/trampoline_v1.py" +python3 "${KOKORO_GFILE_DIR}/trampoline_v1.py" \ No newline at end of file diff --git a/noxfile.py b/noxfile.py index 010e9b143..72378ccec 100644 --- a/noxfile.py +++ b/noxfile.py @@ -52,10 +52,11 @@ SYSTEM_TEST_STANDARD_DEPENDENCIES: List[str] = [ "mock", "pytest", - "pytest-asyncio", "google-cloud-testutils", ] -SYSTEM_TEST_EXTERNAL_DEPENDENCIES: List[str] = [] +SYSTEM_TEST_EXTERNAL_DEPENDENCIES: List[str] = [ + "pytest-asyncio", +] SYSTEM_TEST_LOCAL_DEPENDENCIES: List[str] = [] SYSTEM_TEST_DEPENDENCIES: List[str] = [] SYSTEM_TEST_EXTRAS: List[str] = [] @@ -68,6 +69,7 @@ "unit", "system_emulated", "system", + "conformance", "mypy", "cover", "lint", @@ -160,8 +162,16 @@ def install_unittest_dependencies(session, *constraints): standard_deps = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_DEPENDENCIES session.install(*standard_deps, *constraints) + if UNIT_TEST_EXTERNAL_DEPENDENCIES: + warnings.warn( + "'unit_test_external_dependencies' is deprecated. Instead, please " + "use 'unit_test_dependencies' or 'unit_test_local_dependencies'.", + DeprecationWarning, + ) + session.install(*UNIT_TEST_EXTERNAL_DEPENDENCIES, *constraints) + if UNIT_TEST_LOCAL_DEPENDENCIES: - session.install("-e", *UNIT_TEST_LOCAL_DEPENDENCIES, *constraints) + session.install(*UNIT_TEST_LOCAL_DEPENDENCIES, *constraints) if UNIT_TEST_EXTRAS_BY_PYTHON: extras = UNIT_TEST_EXTRAS_BY_PYTHON.get(session.python, []) @@ -175,20 +185,6 @@ def install_unittest_dependencies(session, *constraints): else: session.install("-e", ".", *constraints) - if UNIT_TEST_EXTERNAL_DEPENDENCIES: - warnings.warn( - "'unit_test_external_dependencies' is deprecated. Instead, please " - "use 'unit_test_dependencies' or 'unit_test_local_dependencies'.", - DeprecationWarning, - ) - session.install( - "--upgrade", - "--no-deps", - "--force-reinstall", - *UNIT_TEST_EXTERNAL_DEPENDENCIES, - *constraints, - ) - def default(session): # Install all test dependencies, then install this package in-place. @@ -279,9 +275,6 @@ def system_emulated(session): @nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) def conformance(session): - """ - Run the set of shared bigtable conformance tests - """ TEST_REPO_URL = "https://github.com/googleapis/cloud-bigtable-clients-test.git" CLONE_REPO_DIR = "cloud-bigtable-clients-test" # install dependencies @@ -477,7 +470,7 @@ def prerelease_deps(session): # Exclude version 1.52.0rc1 which has a known issue. See https://github.com/grpc/grpc/issues/32163 "grpcio!=1.52.0rc1", "grpcio-status", - "google-api-core==2.16.0rc0", # TODO: remove pin once streaming retries is merged + "google-api-core", "google-auth", "proto-plus", "google-cloud-testutils", From 1a32de2d7c8e73027dd1981d6fff705788b2c35d Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 26 Jan 2024 16:33:34 -0800 Subject: [PATCH 54/56] Revert "updated README" This reverts commit c3ed5aa29c06abc3b50278028c23ab96a242a1c3. --- README.rst | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/README.rst b/README.rst index 1c9897eba..5f7d5809d 100644 --- a/README.rst +++ b/README.rst @@ -105,13 +105,3 @@ with the same interface as the popular `HappyBase `__ library. Unlike HappyBase, ``google-cloud-happybase`` uses ``google-cloud-bigtable`` under the covers, rather than Apache HBase. - -``experimental async data client`` ----------------------------------- - -`v2.23.0` includes a preview release of a new data client, accessible at `google.cloud.bigtable.data.BigtableDataClientAsync`. - -This new client supports asyncio, with a corresponding synchronous surface coming soon. - -The new client is currently in preview, and is not recommended for production use. - From cab3694ba9822657f2b85a0db4fb7fc008f65c68 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 29 Jan 2024 14:33:36 -0800 Subject: [PATCH 55/56] fixed versions --- setup.py | 4 ++-- testing/constraints-3.7.txt | 2 +- testing/constraints-3.8.txt | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index 887b0199b..8b698a35b 100644 --- a/setup.py +++ b/setup.py @@ -37,8 +37,8 @@ # 'Development Status :: 5 - Production/Stable' release_status = "Development Status :: 5 - Production/Stable" dependencies = [ - "google-api-core[grpc] >= 2.16.0, <3.0.0dev,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,!=2.8.*,!=2.9.*,!=2.10.*", - "google-cloud-core >= 2.0.0, <3.0.0dev", + "google-api-core[grpc] >= 2.16.0, <3.0.0dev", + "google-cloud-core >= 1.4.4, <3.0.0dev", "grpc-google-iam-v1 >= 0.12.4, <1.0.0dev", "proto-plus >= 1.22.0, <2.0.0dev", "proto-plus >= 1.22.2, <2.0.0dev; python_version>='3.11'", diff --git a/testing/constraints-3.7.txt b/testing/constraints-3.7.txt index b87fca3e6..c684ca534 100644 --- a/testing/constraints-3.7.txt +++ b/testing/constraints-3.7.txt @@ -5,7 +5,7 @@ # # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 -google-api-core==2.16.0rc0 +google-api-core==2.16.0 google-cloud-core==2.0.0 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 diff --git a/testing/constraints-3.8.txt b/testing/constraints-3.8.txt index 42a424b6a..ee858c3ec 100644 --- a/testing/constraints-3.8.txt +++ b/testing/constraints-3.8.txt @@ -5,7 +5,7 @@ # # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", # Then this file should have foo==1.14.0 -google-api-core==2.16.0rc0 +google-api-core==2.16.0 google-cloud-core==2.0.0 grpc-google-iam-v1==0.12.4 proto-plus==1.22.0 From 952d91c28e202dfc6e163861548806ba82189607 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 29 Jan 2024 14:56:44 -0800 Subject: [PATCH 56/56] remove conformance from kokoro main run --- noxfile.py | 1 - owlbot.py | 1 - 2 files changed, 2 deletions(-) diff --git a/noxfile.py b/noxfile.py index 72378ccec..daf730a9a 100644 --- a/noxfile.py +++ b/noxfile.py @@ -69,7 +69,6 @@ "unit", "system_emulated", "system", - "conformance", "mypy", "cover", "lint", diff --git a/owlbot.py b/owlbot.py index 6a5fafe5f..3fb079396 100644 --- a/owlbot.py +++ b/owlbot.py @@ -182,7 +182,6 @@ def conformance(session): "unit", "system_emulated", "system", - "conformance", "mypy",""", )