[dagster-dataform] Import module objects into __init__.py (#228)

* [dagster-dataform] Initial README

* [dagster-dataform] Root files

* [dagster-dataform] Integration modules

* [dagster-dataform] Integration unit tests

* [dagster-dataform] Add prereq for freshness checks

* [dagster-dataform] import module objects into __init__.py

* [dagster-dataform] fix import names

* [dagster-dataform] Resolve linting warnings.

* [dagster-dataform] Add # noqa for stubborn type warnings

---------

Co-authored-by: 2705815 <Justin.Cohen2@aetna.com>

Author: Jcohen010

libraries/dagster-dataform/dagster_dataform/__init__.py

@@ -1,7 +1,18 @@
 from dagster._core.libraries import DagsterLibraryRegistry
+from dagster_dataform.resources import (
+    DataformRepositoryResource as DataformRepositoryResource,
+    load_dataform_assets as load_dataform_assets,
+    load_dataform_asset_check_specs as load_dataform_asset_check_specs,
+)
+from dagster_dataform.dataform_polling_sensor import (
+    create_dataform_workflow_invocation_sensor as create_dataform_workflow_invocation_sensor,
+)
+from dagster_dataform.dataform_orchestration_schedule import (
+    create_dataform_orchestration_schedule as create_dataform_orchestration_schedule,
+)
 
 __version__ = "0.0.1"
 
 DagsterLibraryRegistry.register(
-    "example-integration", __version__, is_dagster_package=False
+    "dagster-dataform", __version__, is_dagster_package=False
 )

libraries/dagster-dataform/dagster_dataform/resources.py

@@ -148,7 +148,7 @@ def query_workflow_invocation(
 
         if not name:
             self.logger.error("No workflow invocation name available")
-            return []
+            return []  # noqa
 
         self.logger.info(f"Querying workflow invocation: {name}")
 
@@ -163,7 +163,7 @@ def query_workflow_invocation(
         # self.logger.info(f"Found {len(response)} workflow invocation actions")
 
         # Handle the response
-        return response
+        return response  # noqa
 
     def create_workflow_invocation(
         self, compilation_result_name: str
@@ -213,7 +213,7 @@ def load_dataform_assets(
         try:
             spec = AssetSpec(
                 key=asset.target.name,
-                kinds=["bigquery"],
+                kinds=["bigquery"],  # noqa
                 metadata={
                     "Project ID": asset.target.database,
                     "Dataset": asset.target.schema,

libraries/dagster-dataform/dagster_dataform_tests/conftest.py

@@ -81,6 +81,7 @@ def create_compilation_result(
     ),
 )
 
+
 @pytest.fixture
 def mock_dataform_client(request):
     # This allows us to parametrize the test with different values for the git commitish, default database, default schema, default location, and assertion schema if necesessary (see test_dataform_repository_resource_get_latest_compilation_result_name_wrong_environment)

libraries/dagster-dataform/dagster_dataform_tests/test_dataform_orchestration_schedule.py

@@ -1,19 +1,20 @@
-from dagster_dataform.dataform_orchestration_schedule import create_dataform_orchestration_schedule
+from dagster_dataform.dataform_orchestration_schedule import (
+    create_dataform_orchestration_schedule,
+)
 from dagster_dataform.resources import DataformRepositoryResource
 from dagster_dataform_tests.conftest import mock_dataform_client
 import dagster as dg
 from dagster import build_schedule_context
 import pytest
-import json
-from google.cloud import dataform_v1
+
 
 def test_dataform_orchestration_schedule_creates_schedule_and_job():
     resource = DataformRepositoryResource(
         project_id="test-project",
         repository_id="test-repo",
         location="us-central1",
         environment="dev",
-        client=mock_dataform_client,
+        client=mock_dataform_client,  # noqa
     )
 
     schedule = create_dataform_orchestration_schedule(
@@ -29,6 +30,7 @@ def test_dataform_orchestration_schedule_creates_schedule_and_job():
     assert schedule.job is not None
     assert schedule.job.name == "dataform_workflow_invocation_job"
 
+
 @pytest.mark.parametrize(
     "mock_dataform_client",
     [
@@ -63,6 +65,3 @@ def test_dataform_orchestration_schedule_tick_creates_run_request(mock_dataform_
     assert result is not None
     assert result.run_requests is not None
     assert len(result.run_requests) == 1
-
-
-

libraries/dagster-dataform/dagster_dataform_tests/test_dataform_polling_sensor.py

@@ -1,18 +1,21 @@
-from dagster_dataform.dataform_polling_sensor import create_dataform_workflow_invocation_sensor
+from dagster_dataform.dataform_polling_sensor import (
+    create_dataform_workflow_invocation_sensor,
+)
 from dagster_dataform.resources import DataformRepositoryResource
 from dagster_dataform_tests.conftest import mock_dataform_client
 import dagster as dg
 from dagster import build_sensor_context
 import pytest
 import json
 
+
 def test_dataform_polling_sensor_creates_sensor_and_job():
     resource = DataformRepositoryResource(
         project_id="test-project",
         repository_id="test-repo",
         location="us-central1",
         environment="dev",
-        client=mock_dataform_client,
+        client=mock_dataform_client,  # noqa
     )
 
     sensor = create_dataform_workflow_invocation_sensor(
@@ -33,7 +36,7 @@ def test_dataform_polling_sensor_creates_sensor_and_job_when_passed_job():
         repository_id="test-repo",
         location="us-central1",
         environment="dev",
-        client=mock_dataform_client,
+        client=mock_dataform_client,  # noqa
     )
 
     @dg.job
@@ -66,7 +69,9 @@ def test_job():
     ],
     indirect=True,
 )
-def test_dataform_polling_sensor_returns_sensor_result_cursor_updated(mock_dataform_client):
+def test_dataform_polling_sensor_returns_sensor_result_cursor_updated(
+    mock_dataform_client,
+):
     resource = DataformRepositoryResource(
         project_id="test-project",
         repository_id="test-repo",
@@ -80,13 +85,9 @@ def test_dataform_polling_sensor_returns_sensor_result_cursor_updated(mock_dataf
         minutes_ago=10,
     )
 
-    cursor = json.dumps({
-        "test-asset" : 1723957400
-    })
+    cursor = json.dumps({"test-asset": 1723957400})
 
-    expected_new_cursor = json.dumps({
-        "test-asset" : 1723958400
-    })
+    expected_new_cursor = json.dumps({"test-asset": 1723958400})
 
     context = build_sensor_context(cursor=cursor)
     result = sensor.evaluate_tick(context)
@@ -100,6 +101,7 @@ def test_dataform_polling_sensor_returns_sensor_result_cursor_updated(mock_dataf
     assert result.asset_events[0].metadata["Invocation SQL Query"] is not None
     assert result.asset_events[0].metadata["BigQuery JobID"] is not None
 
+
 @pytest.mark.parametrize(
     "mock_dataform_client",
     [
@@ -113,7 +115,9 @@ def test_dataform_polling_sensor_returns_sensor_result_cursor_updated(mock_dataf
     ],
     indirect=True,
 )
-def test_dataform_polling_sensor_returns_sensor_result_cursor_not_updated(mock_dataform_client):
+def test_dataform_polling_sensor_returns_sensor_result_cursor_not_updated(
+    mock_dataform_client,
+):
     resource = DataformRepositoryResource(
         project_id="test-project",
         repository_id="test-repo",
@@ -127,13 +131,9 @@ def test_dataform_polling_sensor_returns_sensor_result_cursor_not_updated(mock_d
         minutes_ago=10,
     )
 
-    cursor = json.dumps({
-        "test-asset" : 1723958400
-    })
+    cursor = json.dumps({"test-asset": 1723958400})
 
-    expected_new_cursor = json.dumps({
-        "test-asset" : 1723958400
-    })
+    expected_new_cursor = json.dumps({"test-asset": 1723958400})
 
     context = build_sensor_context(cursor=cursor)
     result = sensor.evaluate_tick(context)
@@ -143,6 +143,7 @@ def test_dataform_polling_sensor_returns_sensor_result_cursor_not_updated(mock_d
     assert result.asset_events is not None
     assert len(result.asset_events) == 0
 
+
 @pytest.mark.parametrize(
     "mock_dataform_client",
     [
@@ -152,12 +153,14 @@ def test_dataform_polling_sensor_returns_sensor_result_cursor_not_updated(mock_d
             "default_schema": "test-schema",
             "default_location": "us-central1",
             "assertion_schema": "test-assertion-schema",
-            "workflow_invocation_passed": False
+            "workflow_invocation_passed": False,
         }
     ],
     indirect=True,
 )
-def test_dataform_polling_sensor_returns_sensor_result_invocation_failed(mock_dataform_client):
+def test_dataform_polling_sensor_returns_sensor_result_invocation_failed(
+    mock_dataform_client,
+):
     resource = DataformRepositoryResource(
         project_id="test-project",
         repository_id="test-repo",
@@ -171,13 +174,9 @@ def test_dataform_polling_sensor_returns_sensor_result_invocation_failed(mock_da
         minutes_ago=10,
     )
 
-    cursor = json.dumps({
-        "test-asset" : 1723957400
-    })
+    cursor = json.dumps({"test-asset": 1723957400})
 
-    expected_new_cursor = json.dumps({
-        "test-asset" : 1723958400
-    })
+    expected_new_cursor = json.dumps({"test-asset": 1723958400})
 
     context = build_sensor_context(cursor=cursor)
     result = sensor.evaluate_tick(context)
@@ -192,13 +191,3 @@ def test_dataform_polling_sensor_returns_sensor_result_invocation_failed(mock_da
     assert result.asset_events[0].metadata["BigQuery JobID"] is not None
     assert result.run_requests is not None
     assert len(result.run_requests) == 1
-
-
-
-
-
-
-
-
-
-

libraries/dagster-dataform/dagster_dataform_tests/test_resources.py

@@ -2,16 +2,9 @@
 from dagster_dataform.utils import get_epoch_time_ago
 import pytest
 from google.cloud import dataform_v1
-from dagster_dataform_tests.conftest import (
-    create_compilation_result,
-    MOCK_COMPILATION_RESULT_ACTION,
-    MOCK_WORKFLOW_INVOCATION,
-    MOCK_WORKFLOW_INVOCATION_ACTION_FAILED,
-    MOCK_WORKFLOW_INVOCATION_ACTION_PASSED,
-    mock_dataform_client,
-)
 from dagster import AssetSpec
 
+
 @pytest.mark.parametrize(
     "mock_dataform_client",
     [

libraries/dagster-dataform/dagster_dataform_tests/test_utils.py

@@ -1,46 +1,47 @@
-import pytest
 from unittest.mock import patch
 from dagster_dataform.utils import get_epoch_time_ago
 
 
-@patch('dagster_dataform.utils.datetime')
+@patch("dagster_dataform.utils.datetime")
 def test_get_epoch_time_ago_basic(mock_datetime):
     """Test basic functionality with 5 minutes."""
     from datetime import datetime, timezone, timedelta
-    
+
     # Mock current time
     mock_now = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
     mock_datetime.now.return_value = mock_now
-    
+
     result = get_epoch_time_ago(5)
-    
+
     # Should return timestamp 5 minutes ago
     expected = int((mock_now - timedelta(minutes=5)).timestamp())
     assert result == expected
 
 
-@patch('dagster_dataform.utils.datetime')
+@patch("dagster_dataform.utils.datetime")
 def test_get_epoch_time_ago_zero(mock_datetime):
     """Test with 0 minutes."""
     from datetime import datetime, timezone
-    
+
     mock_now = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
     mock_datetime.now.return_value = mock_now
-    
+
     result = get_epoch_time_ago(0)
-    
+
     # Should return current timestamp
     expected = int(mock_now.timestamp())
     assert result == expected
 
 
-@patch('dagster_dataform.utils.datetime')
+@patch("dagster_dataform.utils.datetime")
 def test_get_epoch_time_ago_returns_integer(mock_datetime):
     """Test that function returns integer."""
     from datetime import datetime, timezone
-    
-    mock_datetime.now.return_value = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
-    
+
+    mock_datetime.now.return_value = datetime(
+        2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc
+    )
+
     result = get_epoch_time_ago(10)
-    
+
     assert isinstance(result, int)