Fix dagster-openlineage imports for latest Dagster

Author: cmpadden

libraries/dagster-openlineage/dagster_openlineage/compat.py

@@ -14,11 +14,17 @@
 STEP_EVENTS = getattr(dagster, "STEP_EVENTS", None)
 
 if PIPELINE_EVENTS is None or STEP_EVENTS is None:
-    _core_events_spec = importlib.util.find_spec("dagster.core.events")
-    if _core_events_spec is not None:
-        _core_events_module = importlib.import_module("dagster.core.events")
-        PIPELINE_EVENTS = getattr(_core_events_module, "PIPELINE_EVENTS", set())
-        STEP_EVENTS = getattr(_core_events_module, "STEP_EVENTS", set())
+    for _core_events_module_name in ("dagster._core.events", "dagster.core.events"):
+        try:
+            _core_events_spec = importlib.util.find_spec(_core_events_module_name)
+        except ModuleNotFoundError:
+            _core_events_spec = None
+
+        if _core_events_spec is not None:
+            _core_events_module = importlib.import_module(_core_events_module_name)
+            PIPELINE_EVENTS = getattr(_core_events_module, "PIPELINE_EVENTS", set())
+            STEP_EVENTS = getattr(_core_events_module, "STEP_EVENTS", set())
+            break
     else:
         PIPELINE_EVENTS = set()
         STEP_EVENTS = set()

libraries/dagster-openlineage/dagster_openlineage_tests/test_e2e_asset_emission.py

@@ -51,7 +51,7 @@
 from dagster._core.execution.plan.objects import StepFailureData
 from dagster._core.instance import DagsterInstance
 from dagster._core.storage.event_log import InMemoryEventLogStorage
-from dagster.core.test_utils import instance_for_test
+from dagster._core.test_utils import instance_for_test
 from openlineage.client import OpenLineageClient
 from openlineage.client.event_v2 import DatasetEvent, RunEvent, RunState
 from openlineage.client.generated.schema_dataset import SchemaDatasetFacet
@@ -385,7 +385,7 @@ def test_sensor_emits_asset_ol_events_from_real_execution():
             temp_dir=temp_dir,
             overrides={
                 "event_log_storage": {
-                    "module": "dagster.core.storage.event_log",
+                    "module": "dagster._core.storage.event_log",
                     "class": "ConsolidatedSqliteEventLogStorage",
                     "config": {"base_dir": temp_dir},
                 },

libraries/dagster-openlineage/dagster_openlineage_tests/test_sensor_cursor.py

@@ -14,7 +14,7 @@
 
 from dagster_openlineage.compat import DagsterEventType, SensorDefinition
 from dagster import build_sensor_context
-from dagster.core.test_utils import instance_for_test
+from dagster._core.test_utils import instance_for_test
 
 from .conftest import make_test_event_log_record