[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,52 @@
+# 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
+
+  # Expected:
+  #  Row(sale_id = 1, customer_id = 1, product_id = 1, quantity = 1, product = Row(product_id: "1", product_name = "pixel 5", product_stock = "2"))

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

@@ -0,0 +1,96 @@
+# 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'
+

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

@@ -156,7 +156,7 @@ def _wordcount_test_preprocessor(
 
 
 @YamlExamplesTestSuite.register_test_preprocessor(
-    ['test_simple_filter_yaml', 'test_simple_filter_and_combine_yaml'])
+    ['test_simple_filter_yaml', 'test_simple_filter_and_combine_yaml', 'test_spanner_read_yaml', 'test_spanner_write_yaml'])
 def _file_io_write_test_preprocessor(
     test_spec: dict, expected: List[str], env: TestEnvironment):