update cells iterator to handle repeated records

Change-Id: I368b9ec7737e6cad018728b329d636de0a678816

Author: psnelG

core/bigquery.py

@@ -28,7 +28,8 @@
 from google.cloud.bigquery_storage import ReadSession
 from google.cloud.exceptions import NotFound
 
-from .auth import Credentials
+from core.auth import Credentials
+from core.helpers import parse_column_path
 
 BQ_SCOPES = ['https://www.googleapis.com/auth/bigquery']
 
@@ -189,7 +190,6 @@ def get_readrows_iterator(
         read_session=requested_session,
         max_stream_count=1,
     )
-
     # Use 0th stream because because max_stream_count=1
     stream_name = session.streams[0].name
 
@@ -201,38 +201,69 @@ def get_readrows_iterator(
     return cast(Iterable[Mapping], rows)
 
 
-def get_cells_iterator(bq_read_client: BigQueryReadClient,
-                       table_metadata: TableMetadata,
-                       column: str) -> Generator[Any, None, None]:
-    """
-    Get an Iterator of cell values with the requested columns of the table,
-    using an authenticated BigQuery Storage API client.
-
-    Note: Does support nested columns.
-
-    Args:
-        * bq_read_client: BigQuery Storage API Read client
-        * table_metadata: TableMetadata object
-        * column : Column name to select
-
-    Returns:
-        * Iterator of cell values
-    """
-    nested_columns = column.split('.')
-    parent_column = nested_columns[0]
-
-    rows = get_readrows_iterator(bq_read_client,
-                                 table_metadata, [parent_column],
-                                 data_format=DataFormat.AVRO)
-
-    for row in rows:
-        value = row
-        for nested_column in nested_columns:
-            try:
-                value = value[nested_column]
-            except KeyError:
-                raise KeyError(f'{nested_column} was not found.')
-        yield value
+def get_cells_iterator(
+    bq_read_client: BigQueryReadClient,
+    table_metadata: TableMetadata,
+    column: str,
+) -> Generator[Any, None, None]:
+    """Retrieves an iterator of cell values for a specified column, optimized
+    for both simple and nested column
+
+  access, including handling special value structures with dynamic value types
+  for nested columns.
+
+  Args:
+      bq_read_client (BigQueryReadClient): The BigQuery Storage API Read client.
+      table_metadata (TableMetadata): The table's metadata.
+      column (str): The column name, supporting nested fields and array indices
+        for complex cases.
+
+  Returns:
+      Generator[Any, None, None]: An iterator over cell values.
+  """
+    # Check if the column path indicates a simple column access
+    if "." not in column and "[" not in column:
+        # Optimize data retrieval for simple column access
+        rows = get_readrows_iterator(bq_read_client,
+                                     table_metadata, [column],
+                                     data_format=DataFormat.AVRO)
+        for row in rows:
+            yield row.get(column)
+    else:
+        # Handle nested and special value structures as before
+        nested_columns = parse_column_path(column)
+        parent_column = nested_columns[0][0]
+
+        rows = get_readrows_iterator(
+            bq_read_client,
+            table_metadata,
+            [parent_column],
+            data_format=DataFormat.AVRO,
+        )
+
+        for current_value in rows:
+            for column_name, key in nested_columns:
+                if isinstance(current_value, dict):
+                    current_value = current_value.get(column_name)
+                elif isinstance(current_value, list) and key:
+                    current_value = next(
+                        (item for item in current_value
+                         if item.get(column_name) == key),
+                        None,
+                    )
+
+                if isinstance(current_value, dict) and "value" in current_value:
+                    extracted_value = next(
+                        (value for key, value in current_value["value"].items()
+                         if value is not None),
+                        None,
+                    )
+                    current_value = (extracted_value if extracted_value
+                                     is not None else current_value)
+
+                if current_value is None:
+                    break
+            yield current_value
 
 
 def get_table(bq_legacy_client: BigQueryLegacyClient,
@@ -293,8 +324,8 @@ def upload_rows(bq_legacy_client: BigQueryLegacyClient,
         * List of errors, if any
     """
     if len(rows) > 0:
-        result = bq_legacy_client.insert_rows_json(
-            table_metadata.full_table_id, rows)
+        result = bq_legacy_client.insert_rows_json(table_metadata.full_table_id,
+                                                   rows)
         # result is empty if no errors occurred
         for row in result:
             # Note:

core/helpers.py

@@ -13,6 +13,7 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+import re
 from typing import Any, Callable, Generic, List, NoReturn, TypeVar, Union
 
 
@@ -30,6 +31,33 @@ def get_function_name(function: Any) -> str:
     return str(function.__qualname__.split('.')[0])
 
 
+PATTERN = re.compile(r"(^.*?)\[([^\[\]]*?)\](.*?)$")
+
+
+def parse_column_path(column: str) -> list:
+    """Parse the column string to extract nested fields and array indices.
+
+  Args:
+      column (str): The column string with potential nested fields and array
+        indices.
+
+  Returns:
+      list: A list of tuples, where each tuple contains the column name and the
+      key/index.
+  """
+    nested_columns = []
+    for col in column.split("."):
+        parts = PATTERN.match(col)
+        if parts:
+            column_name, key = parts.groups()[0], parts.groups()[1]
+        else:
+            column_name, key = col, None
+        if not column_name:
+            raise ValueError(f"Invalid column path: {column}")
+        nested_columns.append((column_name, key))
+    return nested_columns
+
+
 T = TypeVar('T')
 FlushFunction = Callable[[List[T]], Any]
 

docs/config.md

@@ -87,7 +87,7 @@ An example is provided in `deployment/config_template.json`, and below:
     * `table_name`: BigQuery table name
 
 * BigQuery table `columns`: (mapping)
-  * key: BigQuery table column name
+  * key: BigQuery table column name. Supports nested (Repeated and Nullable records BQ data type) columns. More details [here]()
   * value:
     * `parser`: [Parser name](#parsers)
     * `rules`: (list)
@@ -112,6 +112,40 @@ Advanced:
 * Grant Service Account Token Creator (`roles/iam.serviceAccountTokenCreator`) to the DQM Service Account.
 * Specify `service_account_email` in your config.
 
+### Nested column logic for column name
+
+The DQM configuration file allows you to target specific columns within complex data structures in your BigQuery tables. Here's how to handle nested columns and extract values from arrays of structs:
+
+**1. Nested Columns**
+
+* Use dot notation to reference columns within nested structures (STRUCTs in BigQuery).
+
+**Example:**
+
+```
+column2.nested_column1 
+```
+
+*  This targets `nested_column1`  within the `column2` struct.
+
+**2. Key-Value Logic within Arrays**
+
+* Employ the `key[some_value]` syntax to extract values from arrays of structs where the elements have a `key` field.
+
+**Example:**
+
+```
+event_params.key[ga_session_id]
+```
+
+* This targets the `value` where the corresponding `key` is "ga_session_number" within the `event_params` array.
+
+**Key Points:**
+
+* **Matching Logic:** The `key[some_value]` syntax will extract the **first** matching value within the array.
+* **Parsers:** Choose parsers that are compatible with the data type of the values you are extracting (e.g., `parse_int` for integers, `parse_str` for strings).
+* **Custom Rules:** For more complex logic or to process multiple values with the same key, you can write custom Python rules.
+
 ## Parsers
 
 DQM always treats values from BigQuery as non-typed, i.e. the required Type

test/core/test_bigquery.py

@@ -17,6 +17,8 @@
 from collections.abc import Iterable
 import os
 import unittest
+from unittest.mock import MagicMock
+from unittest.mock import patch
 
 from google.cloud.bigquery import Client as BigQueryLegacyClient
 from google.cloud.bigquery_storage import BigQueryReadClient
@@ -27,6 +29,7 @@
 from core.auth import OAuthCredentials
 from core.bigquery import get_bq_legacy_client
 from core.bigquery import get_bq_read_client
+from core.bigquery import get_cells_iterator
 from core.bigquery import get_readrows_iterator
 from core.bigquery import TableMetadata
 
@@ -92,3 +95,155 @@ def test_get_readrows_iterator(self):
                                        table_name=os.environ['TEST_TABLE_NAME'])
         rows = get_readrows_iterator(self.bqs_client, table_metadata)
         self.assertTrue(isinstance(rows, Iterable))
+
+
+class TestGetCellsIterator(unittest.TestCase):
+
+    @patch('core.bigquery.get_readrows_iterator')
+    def test_get_cells_iterator_single_column(self, mock_get_readrows_iterator):
+        mock_get_readrows_iterator.return_value = iter([
+            {
+                "column_name": 10
+            },
+            {
+                "column_name": 20
+            },
+        ])
+        mock_table_metadata = TableMetadata(project_id="test-project",
+                                            dataset_id="test-dataset",
+                                            table_name="test-table")
+        result = list(
+            get_cells_iterator(MagicMock(spec=BigQueryReadClient),
+                               mock_table_metadata, "column_name"))
+        self.assertEqual(result, [10, 20])
+
+    @patch('core.bigquery.get_readrows_iterator')
+    def test_get_cells_iterator_nested_column(self, mock_get_readrows_iterator):
+        mock_get_readrows_iterator.return_value = iter([
+            {
+                "c": {
+                    "nested": {
+                        "column_name": "value1"
+                    }
+                }
+            },
+            {
+                "c": {
+                    "nested": {
+                        "column_name": "value2"
+                    }
+                }
+            },
+        ])
+        mock_table_metadata = TableMetadata(project_id="test-project",
+                                            dataset_id="test-dataset",
+                                            table_name="test-table")
+        result = list(
+            get_cells_iterator(MagicMock(spec=BigQueryReadClient),
+                               mock_table_metadata, "c.nested.column_name"))
+        self.assertEqual(result, ["value1", "value2"])
+
+    @patch('core.bigquery.get_readrows_iterator')
+    def test_get_cells_iterator_extract_special_value(
+            self, mock_get_readrows_iterator):
+        mock_data = [{
+            "event_params": [{
+                "key": "ga_session_number",
+                "value": {
+                    'string_value': None,
+                    'int_value': 1,
+                    'float_value': None,
+                    'double_value': None
+                }
+            },]
+        }]
+        mock_get_readrows_iterator.return_value = iter(mock_data)
+        mock_table_metadata = TableMetadata(project_id="test-project",
+                                            dataset_id="test-dataset",
+                                            table_name="test-table")
+        result = list(
+            get_cells_iterator(MagicMock(spec=BigQueryReadClient),
+                               mock_table_metadata,
+                               "event_params.key[ga_session_number]"))
+        self.assertEqual(result, [1])
+        mock_get_readrows_iterator.assert_called_once()
+
+    @patch('core.bigquery.get_readrows_iterator')
+    def test_get_cells_iterator_extract_double_nested_value(
+            self, mock_get_readrows_iterator):
+        mock_data = [{
+            "event_params": [{
+                "key": "ga_session_number",
+                "value": {
+                    'string_value': {
+                        "deep_column": 1
+                    }
+                }
+            },]
+        }]
+        mock_get_readrows_iterator.return_value = iter(mock_data)
+        mock_table_metadata = TableMetadata(project_id="test-project",
+                                            dataset_id="test-dataset",
+                                            table_name="test-table")
+        result = list(
+            get_cells_iterator(
+                MagicMock(spec=BigQueryReadClient), mock_table_metadata,
+                "event_params.key[ga_session_number].deep_column"))
+        self.assertEqual(result, [1])
+        mock_get_readrows_iterator.assert_called_once()
+
+    @patch('core.bigquery.get_readrows_iterator')
+    def test_get_cells_iterator_key_not_found(self, mock_get_readrows_iterator):
+        mock_data = [
+            {
+                "event_params": [{
+                    "key": "some_other_key",
+                    "value": {
+                        'int_value': 2
+                    }
+                }]
+            },
+        ]
+        mock_get_readrows_iterator.return_value = iter(mock_data)
+        mock_table_metadata = TableMetadata(project_id="test-project",
+                                            dataset_id="test-dataset",
+                                            table_name="test-table")
+        result = list(
+            get_cells_iterator(MagicMock(spec=BigQueryReadClient),
+                               mock_table_metadata,
+                               "event_params.key[non_existent_key]"))
+        self.assertEqual(result, [None])
+        mock_get_readrows_iterator.assert_called_once()
+
+    @patch('core.bigquery.get_readrows_iterator')
+    def test_get_cells_iterator_return_full_record(self,
+                                                   mock_get_readrows_iterator):
+        # Simulate data where the specified key is not present
+        mock_data = [
+            {
+                "event_params": [{
+                    "key": "some_other_key",
+                    "value": {
+                        'int_value': 2
+                    }
+                }]
+            },
+        ]
+        mock_get_readrows_iterator.return_value = iter(mock_data)
+        mock_table_metadata = TableMetadata(project_id="test-project",
+                                            dataset_id="test-dataset",
+                                            table_name="test-table")
+        result = list(
+            get_cells_iterator(MagicMock(spec=BigQueryReadClient),
+                               mock_table_metadata, "event_params"))
+        self.assertEqual(result, [[{
+            "key": "some_other_key",
+            "value": {
+                'int_value': 2
+            }
+        }]])
+        mock_get_readrows_iterator.assert_called_once()
+
+
+# TODO(psnel) add support for highly nested usecases alternating between
+# repeated and nullable structures.