disable feast if not installed

apache · Oct 3, 2024 · c647e47 · c647e47
1 parent ee6258e
commit c647e47
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/sdks/python/apache_beam/yaml/yaml_enrichment.py b/sdks/python/apache_beam/yaml/yaml_enrichment.py
@@ -20,12 +20,17 @@
 from typing import Optional
 
 import apache_beam as beam
+from apache_beam.yaml import options
 from apache_beam.transforms.enrichment import Enrichment
 from apache_beam.transforms.enrichment_handlers.bigquery import BigQueryEnrichmentHandler
 from apache_beam.transforms.enrichment_handlers.bigtable import BigTableEnrichmentHandler
-from apache_beam.transforms.enrichment_handlers.feast_feature_store import FeastFeatureStoreEnrichmentHandler
 from apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store import VertexAIFeatureStoreEnrichmentHandler
 
+try:
+  from apache_beam.transforms.enrichment_handlers.feast_feature_store import FeastFeatureStoreEnrichmentHandler
+except ImportError:
+  FeastFeatureStoreEnrichmentHandler = None
+
 
 @beam.ptransform.ptransform_fn
 def enrichment_transform(
@@ -87,6 +92,12 @@ def enrichment_transform(
             timeout: 30
 
     """
+  options.YamlOptions.check_enabled(pcoll.pipeline, 'Enrichment')
+  if (enrichment_handler == 'FeastFeatureStore' and
+      not FeastFeatureStoreEnrichmentHandler):
+    raise ValueError(
+        "FeastFeatureStore handler requires 'feast' package to be installed. " +
+        "Please install using 'pip install feast[gcp]' and try again.")
   handler_map = {
       'BigQuery': BigQueryEnrichmentHandler,
       'BigTable': BigTableEnrichmentHandler,