feat: Refactor Metadata, add system tests, remove preview warning (#1099)

Author: jackdingilian

google/cloud/bigtable/data/_async/client.py

@@ -610,10 +610,6 @@ async def execute_query(
             google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
                 a parameter is passed without an explicit type, and the type cannot be infered
         """
-        warnings.warn(
-            "ExecuteQuery is in preview and may change in the future.",
-            category=RuntimeWarning,
-        )
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
         converted_param_types = _to_param_types(parameters, parameter_types)
         prepare_request = {

google/cloud/bigtable/data/_sync_autogen/client.py

@@ -471,10 +471,6 @@ def execute_query(
             google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
                 a parameter is passed without an explicit type, and the type cannot be infered
         """
-        warnings.warn(
-            "ExecuteQuery is in preview and may change in the future.",
-            category=RuntimeWarning,
-        )
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
         converted_param_types = _to_param_types(parameters, parameter_types)
         prepare_request = {

google/cloud/bigtable/data/execute_query/__init__.py

@@ -20,7 +20,6 @@
 )
 from google.cloud.bigtable.data.execute_query.metadata import (
     Metadata,
-    ProtoMetadata,
     SqlType,
 )
 from google.cloud.bigtable.data.execute_query.values import (
@@ -39,7 +38,6 @@
     "QueryResultRow",
     "Struct",
     "Metadata",
-    "ProtoMetadata",
     "ExecuteQueryIteratorAsync",
     "ExecuteQueryIterator",
 ]

google/cloud/bigtable/data/execute_query/_reader.py

@@ -19,7 +19,6 @@
     Iterable,
     Optional,
     Sequence,
-    cast,
 )
 from abc import ABC, abstractmethod
 
@@ -32,7 +31,7 @@
 from google.cloud.bigtable.helpers import batched
 
 from google.cloud.bigtable.data.execute_query.values import QueryResultRow
-from google.cloud.bigtable.data.execute_query.metadata import Metadata, ProtoMetadata
+from google.cloud.bigtable.data.execute_query.metadata import Metadata
 
 
 T = TypeVar("T")
@@ -90,7 +89,7 @@ def _parse_proto_rows(self, bytes_to_parse: bytes) -> Iterable[PBValue]:
         return proto_rows.values
 
     def _construct_query_result_row(
-        self, values: Sequence[PBValue], metadata: ProtoMetadata
+        self, values: Sequence[PBValue], metadata: Metadata
     ) -> QueryResultRow:
         result = QueryResultRow()
         columns = metadata.columns
@@ -107,16 +106,13 @@ def _construct_query_result_row(
     def consume(
         self, batches_to_consume: List[bytes], metadata: Metadata
     ) -> Optional[Iterable[QueryResultRow]]:
-        proto_metadata = cast(ProtoMetadata, metadata)
-        num_columns = len(proto_metadata.columns)
+        num_columns = len(metadata.columns)
         rows = []
         for batch_bytes in batches_to_consume:
             values = self._parse_proto_rows(batch_bytes)
             for row_data in batched(values, n=num_columns):
                 if len(row_data) == num_columns:
-                    rows.append(
-                        self._construct_query_result_row(row_data, proto_metadata)
-                    )
+                    rows.append(self._construct_query_result_row(row_data, metadata))
                 else:
                     raise ValueError(
                         "Unexpected error, recieved bad number of values. "

google/cloud/bigtable/data/execute_query/metadata.py

@@ -298,14 +298,6 @@ def _to_value_pb_dict(self, value: Any) -> Dict[str, Any]:
 
 
 class Metadata:
-    """
-    Base class for metadata returned by the ExecuteQuery operation.
-    """
-
-    pass
-
-
-class ProtoMetadata(Metadata):
     """
     Metadata class for the ExecuteQuery operation.
 
@@ -335,7 +327,7 @@ def columns(self) -> List[Column]:
     def __init__(
         self, columns: Optional[List[Tuple[Optional[str], SqlType.Type]]] = None
     ):
-        self._columns: List[ProtoMetadata.Column] = []
+        self._columns: List[Metadata.Column] = []
         self._column_indexes: Dict[str, List[int]] = defaultdict(list)
         self._duplicate_names: Set[str] = set()
 
@@ -345,7 +337,7 @@ def __init__(
                     if column_name in self._column_indexes:
                         self._duplicate_names.add(column_name)
                     self._column_indexes[column_name].append(len(self._columns))
-                self._columns.append(ProtoMetadata.Column(column_name, column_type))
+                self._columns.append(Metadata.Column(column_name, column_type))
 
     def __getitem__(self, index_or_name: Union[str, int]) -> Column:
         if isinstance(index_or_name, str):
@@ -381,7 +373,7 @@ def _pb_metadata_to_metadata_types(
             fields.append(
                 (column_metadata.name, _pb_type_to_metadata_type(column_metadata.type))
             )
-        return ProtoMetadata(fields)
+        return Metadata(fields)
     raise ValueError("Invalid ResultSetMetadata object received.")
 
 

tests/system/data/test_system_async.py

@@ -1063,6 +1063,36 @@ async def test_execute_query_simple(self, client, table_id, instance_id):
         assert row["a"] == 1
         assert row["b"] == "foo"
 
+    @CrossSync.pytest
+    @pytest.mark.usefixtures("table")
+    @CrossSync.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    async def test_execute_against_table(
+        self, client, instance_id, table_id, temp_rows
+    ):
+        await temp_rows.add_row(b"row_key_1")
+        result = await client.execute_query(
+            "SELECT * FROM `" + table_id + "`", instance_id
+        )
+        rows = [r async for r in result]
+
+        assert len(rows) == 1
+        assert rows[0]["_key"] == b"row_key_1"
+        family_map = rows[0][TEST_FAMILY]
+        assert len(family_map) == 1
+        assert family_map[b"q"] == b"test-value"
+        assert len(rows[0][TEST_FAMILY_2]) == 0
+        md = result.metadata
+        assert len(md) == 3
+        assert md["_key"].column_type == SqlType.Bytes()
+        assert md[TEST_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Bytes()
+        )
+        assert md[TEST_FAMILY_2].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Bytes()
+        )
+
     @CrossSync.pytest
     @pytest.mark.usefixtures("client")
     @CrossSync.Retry(
@@ -1149,3 +1179,28 @@ async def test_execute_query_params(self, client, table_id, instance_id):
             date_pb2.Date(year=2025, month=1, day=17),
             None,
         ]
+
+    @CrossSync.pytest
+    @pytest.mark.usefixtures("table")
+    @CrossSync.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    async def test_execute_metadata_on_empty_response(
+        self, client, instance_id, table_id, temp_rows
+    ):
+        await temp_rows.add_row(b"row_key_1")
+        result = await client.execute_query(
+            "SELECT * FROM `" + table_id + "` WHERE _key='non-existent'", instance_id
+        )
+        rows = [r async for r in result]
+
+        assert len(rows) == 0
+        md = result.metadata
+        assert len(md) == 3
+        assert md["_key"].column_type == SqlType.Bytes()
+        assert md[TEST_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Bytes()
+        )
+        assert md[TEST_FAMILY_2].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Bytes()
+        )

tests/system/data/test_system_autogen.py

@@ -869,6 +869,30 @@ def test_execute_query_simple(self, client, table_id, instance_id):
         assert row["a"] == 1
         assert row["b"] == "foo"
 
+    @pytest.mark.usefixtures("table")
+    @CrossSync._Sync_Impl.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    def test_execute_against_table(self, client, instance_id, table_id, temp_rows):
+        temp_rows.add_row(b"row_key_1")
+        result = client.execute_query("SELECT * FROM `" + table_id + "`", instance_id)
+        rows = [r for r in result]
+        assert len(rows) == 1
+        assert rows[0]["_key"] == b"row_key_1"
+        family_map = rows[0][TEST_FAMILY]
+        assert len(family_map) == 1
+        assert family_map[b"q"] == b"test-value"
+        assert len(rows[0][TEST_FAMILY_2]) == 0
+        md = result.metadata
+        assert len(md) == 3
+        assert md["_key"].column_type == SqlType.Bytes()
+        assert md[TEST_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Bytes()
+        )
+        assert md[TEST_FAMILY_2].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Bytes()
+        )
+
     @pytest.mark.usefixtures("client")
     @CrossSync._Sync_Impl.Retry(
         predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
@@ -945,3 +969,26 @@ def test_execute_query_params(self, client, table_id, instance_id):
             date_pb2.Date(year=2025, month=1, day=17),
             None,
         ]
+
+    @pytest.mark.usefixtures("table")
+    @CrossSync._Sync_Impl.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    def test_execute_metadata_on_empty_response(
+        self, client, instance_id, table_id, temp_rows
+    ):
+        temp_rows.add_row(b"row_key_1")
+        result = client.execute_query(
+            "SELECT * FROM `" + table_id + "` WHERE _key='non-existent'", instance_id
+        )
+        rows = [r for r in result]
+        assert len(rows) == 0
+        md = result.metadata
+        assert len(md) == 3
+        assert md["_key"].column_type == SqlType.Bytes()
+        assert md[TEST_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Bytes()
+        )
+        assert md[TEST_FAMILY_2].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Bytes()
+        )

tests/unit/data/execute_query/test_query_result_row_reader.py

@@ -18,7 +18,7 @@
 from google.cloud.bigtable.data.execute_query._reader import _QueryResultRowReader
 
 from google.cloud.bigtable.data.execute_query.metadata import (
-    ProtoMetadata,
+    Metadata,
     SqlType,
     _pb_metadata_to_metadata_types,
 )
@@ -37,9 +37,7 @@
 
 class TestQueryResultRowReader:
     def test__single_values_received(self):
-        metadata = ProtoMetadata(
-            [("test1", SqlType.Int64()), ("test2", SqlType.Int64())]
-        )
+        metadata = Metadata([("test1", SqlType.Int64()), ("test2", SqlType.Int64())])
         values = [
             proto_rows_bytes(int_val(1), int_val(2)),
             proto_rows_bytes(int_val(3), int_val(4)),
@@ -61,9 +59,7 @@ def test__multiple_rows_received(self):
             proto_rows_bytes(int_val(7), int_val(8)),
         ]
 
-        metadata = ProtoMetadata(
-            [("test1", SqlType.Int64()), ("test2", SqlType.Int64())]
-        )
+        metadata = Metadata([("test1", SqlType.Int64()), ("test2", SqlType.Int64())])
         reader = _QueryResultRowReader()
 
         result = reader.consume(values[0:1], metadata)
@@ -89,9 +85,7 @@ def test__multiple_rows_received(self):
         assert result[0][1] == result[0]["test2"] == 8
 
     def test__received_values_are_passed_to_parser_in_batches(self):
-        metadata = ProtoMetadata(
-            [("test1", SqlType.Int64()), ("test2", SqlType.Int64())]
-        )
+        metadata = Metadata([("test1", SqlType.Int64()), ("test2", SqlType.Int64())])
 
         # TODO move to a SqlType test
         assert SqlType.Struct([("a", SqlType.Int64())]) == SqlType.Struct(
@@ -128,7 +122,7 @@ def test__received_values_are_passed_to_parser_in_batches(self):
             )
 
     def test__parser_errors_are_forwarded(self):
-        metadata = ProtoMetadata([("test1", SqlType.Int64())])
+        metadata = Metadata([("test1", SqlType.Int64())])
 
         values = [str_val("test")]
 
@@ -236,7 +230,7 @@ def test_multiple_batches(self):
         ]
         results = reader.consume(
             batches,
-            ProtoMetadata([("test1", SqlType.Int64()), ("test2", SqlType.Int64())]),
+            Metadata([("test1", SqlType.Int64()), ("test2", SqlType.Int64())]),
         )
         assert len(results) == 4
         [row1, row2, row3, row4] = results
@@ -250,9 +244,9 @@ def test_multiple_batches(self):
         assert row4["test2"] == 8
 
 
-class TestProtoMetadata:
+class TestMetadata:
     def test__duplicate_column_names(self):
-        metadata = ProtoMetadata(
+        metadata = Metadata(
             [
                 ("test1", SqlType.Int64()),
                 ("test2", SqlType.Bytes()),