Fix BigQueryEnrichmentHandler

Author: aIbrahiim

sdks/python/apache_beam/transforms/enrichment_handlers/bigquery.py

@@ -141,11 +141,12 @@ def __init__(
     self.project = project
     self.column_names = column_names
     self.select_fields = ",".join(column_names) if column_names else '*'
-    self.row_restriction_template = row_restriction_template
+    self.row_restriction_template = row_restriction_template.replace('{{}}', '{}')
     self.table_name = table_name
     self.fields = fields if fields else []
     self.condition_value_fn = condition_value_fn
     self.query_fn = query_fn
+    self._has_placeholders = '{}' in self.row_restriction_template
     self.query_template = (
         "SELECT %s FROM %s WHERE %s" %
         (self.select_fields, self.table_name, self.row_restriction_template))
@@ -197,29 +198,41 @@ def __call__(self, request: Union[beam.Row, list[beam.Row]], *args, **kwargs):
             [fr'({self.row_restriction_template})'] * batch_size)
         raw_query = self.query_template.replace(
             self.row_restriction_template, batched_condition_template)
-      for req in request:
-        request_dict = req._asdict()
-        try:
-          current_values = (
-              self.condition_value_fn(req) if self.condition_value_fn else
-              [request_dict[field] for field in self.fields])
-        except KeyError as e:
-          raise KeyError(
-              "Make sure the values passed in `fields` are the "
-              "keys in the input `beam.Row`." + str(e))
-        values.extend(current_values)
-        requests_map[self.create_row_key(req)] = req
-      query = raw_query.format(*values)
+      if self._has_placeholders:
+        for req in request:
+          request_dict = req._asdict()
+          try:
+            current_values = (
+                self.condition_value_fn(req) if self.condition_value_fn else
+                [request_dict[field] for field in self.fields])
+          except KeyError as e:
+            raise KeyError(
+                "Make sure the values passed in `fields` are the "
+                "keys in the input `beam.Row`." + str(e))
+          values.extend(current_values)
+          requests_map[self.create_row_key(req)] = req
+        query = raw_query.format(*values)
+      else:
+        for req in request:
+          requests_map[id(req)] = req  # Use object id as key
+        query = raw_query
 
       responses_dict = self._execute_query(query)
       unmatched_requests = requests_map.copy()
       if responses_dict:
-        for response in responses_dict:
-          response_row = beam.Row(**response)
-          response_key = self.create_row_key(response_row)
-          if response_key in unmatched_requests:
-            req = unmatched_requests.pop(response_key)
-            responses.append((req, response_row))
+        if self._has_placeholders:
+          for response in responses_dict:
+            response_row = beam.Row(**response)
+            response_key = self.create_row_key(response_row)
+            if response_key in unmatched_requests:
+              req = unmatched_requests.pop(response_key)
+              responses.append((req, response_row))
+        else:
+          if responses_dict:
+            response_row = beam.Row(**responses_dict[0])
+            for req in unmatched_requests.values():
+              responses.append((req, response_row))
+            unmatched_requests.clear()
       if unmatched_requests:
         if self.throw_exception_on_empty_results:
           raise ValueError(f"no matching row found for query: {query}")
@@ -235,11 +248,13 @@ def __call__(self, request: Union[beam.Row, list[beam.Row]], *args, **kwargs):
         # that should be populated into the query template string.
         query = self.query_fn(request)
       else:
-        values = (
-            self.condition_value_fn(request) if self.condition_value_fn else
-            list(map(request_dict.get, self.fields)))
-        # construct the query.
-        query = self.query_template.format(*values)
+        if self._has_placeholders:
+          values = (
+              self.condition_value_fn(request) if self.condition_value_fn else
+              list(map(request_dict.get, self.fields)))
+          query = self.query_template.format(*values)
+        else:
+          query = self.query_template
       response_dict = self._execute_query(query)
       if response_dict is None:
         if self.throw_exception_on_empty_results:

sdks/python/apache_beam/yaml/extended_tests/data/enrichment.yaml

@@ -62,9 +62,9 @@ pipelines:
               project: apache-beam-testing
               table_name: "{BQ_TABLE}"
               fields: ['label']
-              row_restriction_template: "label = '37a'"
+              row_restriction_template: "label = '{{}}'"
             timeout: 30
-          
+
         - type: MapToFields
           config:
             language: python
@@ -82,6 +82,8 @@ pipelines:
         - type: AssertEqual
           config:
             elements:
+              - {label: '11a', rank: 0, name: 'S1'}
               - {label: '37a', rank: 1, name: 'S2'}
+              - {label: '389a', rank: 2, name: 'S3'}
     options:
-      yaml_experimental_features: [ 'Enrichment' ]
+      yaml_experimental_features: [ 'Enrichment' ]