fix(ingest): replace ImportError with ModuleNotFoundError (#2498)

Using the more specific exception will prevent us from accidentally ignoring errors that should be handled.
datahub-project · May 5, 2021 · 7f0656f · 7f0656f
1 parent b4457af
commit 7f0656f
Show file tree

Hide file tree

Showing 8 changed files with 9 additions and 8 deletions.
diff --git a/metadata-ingestion/examples/airflow/generic_recipe_sample_dag.py b/metadata-ingestion/examples/airflow/generic_recipe_sample_dag.py
@@ -11,8 +11,9 @@
 
 try:
     from airflow.operators.python import PythonOperator
-except ImportError:
+except ModuleNotFoundError:
     from airflow.operators.python_operator import PythonOperator
+
 from airflow.utils.dates import days_ago
 
 from datahub.ingestion.run.pipeline import Pipeline

diff --git a/metadata-ingestion/examples/airflow/lineage_backend_demo.py b/metadata-ingestion/examples/airflow/lineage_backend_demo.py
@@ -10,7 +10,7 @@
 
 try:
     from airflow.operators.bash import BashOperator
-except ImportError:
+except ModuleNotFoundError:
     from airflow.operators.bash_operator import BashOperator
 
 from datahub.integrations.airflow.entities import Dataset

diff --git a/metadata-ingestion/examples/airflow/mysql_sample_dag.py b/metadata-ingestion/examples/airflow/mysql_sample_dag.py
@@ -12,7 +12,7 @@
 
 try:
     from airflow.operators.python import PythonOperator
-except ImportError:
+except ModuleNotFoundError:
     from airflow.operators.python_operator import PythonOperator
 
 from datahub.ingestion.run.pipeline import Pipeline

diff --git a/metadata-ingestion/src/datahub/ingestion/api/registry.py b/metadata-ingestion/src/datahub/ingestion/api/registry.py
@@ -54,7 +54,7 @@ def load(self, entry_point_key: str) -> None:
 
             try:
                 plugin_class = entry_point.load()
-            except ImportError as e:
+            except ModuleNotFoundError as e:
                 self.register_disabled(name, e)
                 continue
 

diff --git a/metadata-ingestion/src/datahub/integrations/airflow/__init__.py b/metadata-ingestion/src/datahub/integrations/airflow/__init__.py
@@ -2,6 +2,6 @@
     from datahub.integrations.airflow.lineage_backend import (
         DatahubAirflowLineageBackend,
     )
-except ImportError:
+except ModuleNotFoundError:
     # Compat for Airflow 2.x.
     pass
diff --git a/metadata-ingestion/src/datahub/integrations/airflow/hooks.py b/metadata-ingestion/src/datahub/integrations/airflow/hooks.py
@@ -6,7 +6,7 @@
     from airflow.hooks.base import BaseHook
 
     AIRFLOW_1 = False
-except ImportError:
+except ModuleNotFoundError:
     from airflow.hooks.base_hook import BaseHook
 
     AIRFLOW_1 = True

diff --git a/metadata-ingestion/tests/unit/test_airflow.py b/metadata-ingestion/tests/unit/test_airflow.py
@@ -14,7 +14,7 @@
 
 try:
     from airflow.operators.dummy import DummyOperator
-except ImportError:
+except ModuleNotFoundError:
     from airflow.operators.dummy_operator import DummyOperator
 
 import datahub.emitter.mce_builder as builder

diff --git a/metadata-ingestion/tests/unit/test_plugin_system.py b/metadata-ingestion/tests/unit/test_plugin_system.py
@@ -43,7 +43,7 @@ def test_registry():
     # Make a mini sink registry.
     fake_registry = Registry[Sink]()
     fake_registry.register("console", ConsoleSink)
-    fake_registry.register_disabled("disabled", ImportError("disabled sink"))
+    fake_registry.register_disabled("disabled", ModuleNotFoundError("disabled sink"))
 
     class DummyClass:
         pass