[yaml] Add use cases for Enrichment transform in YAML

sdks/python/apache_beam/io/gcp/bigquery_tools.py

@@ -351,7 +351,7 @@ class BigQueryWrapper(object):
 
   HISTOGRAM_METRIC_LOGGER = MetricLogger()
 
-  def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
+  def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None, project_id = None):
     self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions())
     self.gcp_bq_client = client or gcp_bigquery.Client(
         client_info=ClientInfo(

sdks/python/apache_beam/yaml/examples/bigtable_enrichment.yaml

@@ -0,0 +1,55 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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.
+#
+
+pipeline:
+  type: chain
+  transforms:
+
+  # Step 1: Creating a collection of elements that needs
+  # to be enriched. Here we are simulating sales data
+    - type: Create
+      config:
+        elements:
+        - sale_id: 1
+          customer_id: 1
+          product_id: 1
+          quantity: 1
+
+  # Step 2: Enriching the data with Bigtable
+  # This specific bigtable stores product data in the below format
+  # product:product_id, product:product_name, product:product_stock
+    - type: Enrichment
+      config:
+        enrichment_handler: 'BigTable'
+        handler_config:
+          project_id: 'apache-beam-testing'
+          instance_id: 'beam-test'
+          table_id: 'bigtable-enrichment-test'
+          row_key: 'product_id'
+        timeout: 30
+
+  # Step 3: Logging for testing
+  # This is a simple way to view the enriched data
+  # We can also store it somewhere like a json file
+    - type: LogForTesting
+
+options:
+  yaml_experimental_features: Enrichment
+
+# Expected:
+#  Row(sale_id=1, customer_id=1, product_id=1, quantity=1, product={'product_id': '1', 'product_name': 'pixel 5', 'product_stock': '2'})

sdks/python/apache_beam/yaml/examples/enrich_spanner_with_bigquery.yaml

@@ -0,0 +1,99 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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.
+#
+
+pipeline:
+  transforms:
+    # Step 1: Read orders details from Spanner
+    - type: ReadFromSpanner
+      name: ReadOrders
+      config:
+        project_id: 'apache-beam-testing'
+        instance_id: 'orders-test'
+        database_id: 'order-database'
+        query: 'SELECT customer_id, product_id, order_date, order_amount FROM orders'
+
+    # Step 2: Enrich order details with customers details from BigQuery
+    - type: Enrichment
+      name: Enriched
+      input: ReadOrders
+      config:
+        enrichment_handler: 'BigQuery'
+        handler_config:
+          project: "apache-beam-testing"
+          table_name: "apache-beam-testing.ALL_TEST.customers"
+          row_restriction_template: "customer_id = 1001 or customer_id = 1003"
+          fields: ["customer_id"]
+
+    # Step 3: Map enriched values to Beam schema
+    # TODO: This should be removed when schema'd enrichment is available
+    - type: MapToFields
+      name: MapEnrichedValues
+      input: Enriched
+      config:
+        language: python
+        fields:
+          customer_id:
+            callable: 'lambda x: x.customer_id'
+            output_type: integer
+          customer_name:
+            callable: 'lambda x: x.customer_name'
+            output_type: string
+          customer_email:
+            callable: 'lambda x: x.customer_email'
+            output_type: string 
+          product_id:
+            callable: 'lambda x: x.product_id'
+            output_type: integer
+          order_date:
+            callable: 'lambda x: x.order_date'
+            output_type: string
+          order_amount:
+            callable: 'lambda x: x.order_amount'
+            output_type: integer
+
+    # Step 4: Filter orders with amount greater than 100
+    - type: Filter
+      name: FilterHighValueOrders
+      input: MapEnrichedValues
+      config:
+        keep: "order_amount > 100"
+        language: "python"
+
+
+    # Step 6: Write processed order to another spanner table
+    - type: WriteToSpanner
+      name: WriteProcessedOrders
+      input: FilterHighValueOrders
+      config:
+        project_id: 'apache-beam-testing'
+        instance_id: 'orders-test'
+        database_id: 'order-database'
+        table_id: 'orders_with_customers'
+        error_handling:
+          output: my_error_output
+
+    # Step 7: Handle write errors by writing to JSON
+    - type: WriteToJson
+      name: WriteErrorsToJson
+      input: WriteProcessedOrders.my_error_output
+      config:
+        path: 'errors.json'
+
+# Expected:
+#  Row(customer_id=1001, customer_name='Alice', customer_email='[email protected]', product_id=2001, order_date='24-03-24', order_amount=150)
+#  Row(customer_id=1003, customer_name='Claire', customer_email='[email protected]', product_id=2003, order_date='7-05-24', order_amount=110)

sdks/python/apache_beam/yaml/examples/testing/examples_test.py

@@ -155,8 +155,14 @@ def _wordcount_test_preprocessor(
       env.input_file('kinglear.txt', '\n'.join(lines)))
 
 
-@YamlExamplesTestSuite.register_test_preprocessor(
-    ['test_simple_filter_yaml', 'test_simple_filter_and_combine_yaml'])
+@YamlExamplesTestSuite.register_test_preprocessor([
+    'test_simple_filter_yaml',
+    'test_simple_filter_and_combine_yaml',
+    'test_spanner_read_yaml',
+    'test_spanner_write_yaml',
+    'test_bigtable_enrichment_yaml',
+    'test_enrich_spanner_with_bigquery_yaml'
+])
 def _file_io_write_test_preprocessor(
     test_spec: dict, expected: List[str], env: TestEnvironment):