bug fixed

metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_lineage_core.py

@@ -7,7 +7,7 @@
 from datahub.utilities.urns.dataset_urn import DatasetUrn
 
 from datahub_airflow_plugin.client.airflow_generator import AirflowGenerator
-from datahub_airflow_plugin.entities import _Entity
+from datahub.api.entities.dataset import _Entity
 
 if TYPE_CHECKING:
     from airflow import DAG

metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py

@@ -10,13 +10,13 @@
 import airflow.configuration
 import airflow.version
 import datahub.emitter.mce_builder as builder
+from datahub.api.entities.dataset import Dataset, Urn
 import packaging.version
 import pytest
 from airflow.lineage import apply_lineage, prepare_lineage
 from airflow.models import DAG, Connection, DagBag, DagRun, TaskInstance
 from datahub_provider import get_provider_info
 from datahub_provider._airflow_shims import AIRFLOW_PATCHED, EmptyOperator
-from datahub_provider.entities import Dataset, Urn
 from datahub_provider.hooks.datahub import DatahubKafkaHook, DatahubRestHook
 from datahub_provider.operators.datahub import DatahubEmitterOperator
 
@@ -79,7 +79,8 @@ def test_dags_load_with_no_errors(pytestconfig: pytest.Config) -> None:
 
     # Note: the .airflowignore file skips the snowflake DAG.
     dag_bag = DagBag(dag_folder=str(airflow_examples_folder), include_examples=False)
-
+    import pdb
+    pdb.set_trace()
     import_errors = dag_bag.import_errors
 
     assert len(import_errors) == 0