pair unmatched requests with same result

Author: aIbrahiim

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

@@ -147,7 +147,6 @@ def __init__(
     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))
@@ -199,41 +198,34 @@ 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)
-      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
+      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)
 
       responses_dict = self._execute_query(query)
       unmatched_requests = requests_map.copy()
       if responses_dict:
-        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()
+        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 unmatched_requests and 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}")
@@ -249,13 +241,10 @@ 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:
-        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
+        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)
       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,7 +62,7 @@ pipelines:
               project: apache-beam-testing
               table_name: "{BQ_TABLE}"
               fields: ['label']
-              row_restriction_template: "label = '{{}}'"
+              row_restriction_template: "label = '37a'"
             timeout: 30
 
         - type: MapToFields