Merge pull request #210 from DalgoT4D/get-scheduled-flow-runs

get scheduled flow-runs for a deployment

Author: fatchat

proxy/main.py

@@ -28,6 +28,7 @@
     get_flow_runs_by_name,
     set_deployment_schedule,
     get_deployment,
+    get_deployment_scheduled_flow_runs,
     get_flow_run,
     retry_flow_run,
     create_secret_block,
@@ -489,6 +490,21 @@ def put_dataflow_v1(deployment_id, payload: DeploymentUpdate2):
     return {"success": 1}
 
 
+@app.get("/proxy/v1/deployments/get_scheduled_flow_runs")
+def get_dataflow_scheduled_flow_runs(deployment_id: str):
+    """fetch scheduled flow-runs for a deployment"""
+    if not isinstance(deployment_id, str):
+        raise TypeError("deployment_id must be a string")
+    try:
+        res = get_deployment_scheduled_flow_runs(deployment_id)
+    except Exception as error:
+        logger.exception(error)
+        raise HTTPException(
+            status_code=400, detail="failed to fetch scheduled flow-runs for deployment"
+        ) from error
+    return {"flow_runs": res}
+
+
 @app.post("/proxy/flow_run/")
 async def get_flowrun(payload: FlowRunRequest):
     """look up a flow run by name and return id if found"""

proxy/service.py

@@ -725,6 +725,19 @@ def get_deployment(deployment_id: str) -> dict:
     return res
 
 
+def get_deployment_scheduled_flow_runs(deployment_id: str) -> dict:
+    """fetch scheduled flow-runs for a deployment"""
+    if not isinstance(deployment_id, str):
+        raise TypeError("deployment_id must be a string")
+    res = prefect_post(
+        "deployments/get_scheduled_flow_runs",
+        {
+            "deployment_ids": [deployment_id],
+        },
+    )
+    return res
+
+
 def update_flow_run_final_state(flow_run: dict) -> dict:
     """
     fetch tasks of the flow_run & checks for the custom flow_run state

tests/test_main.py

@@ -35,6 +35,7 @@
     post_run_dbtcore_flow_v1,
     post_dataflow_v1,
     put_dataflow_v1,
+    get_dataflow_scheduled_flow_runs,
     get_long_running_flows,
     patch_dbt_cloud_creds,
     get_dbt_cloud_creds,
@@ -588,6 +589,15 @@ def test_put_dataflow_v1_success():
         assert result == {"success": 1}
 
 
+def test_get_dataflow_scheduled_flow_runs():
+    with patch(
+        "proxy.main.get_deployment_scheduled_flow_runs"
+    ) as mock_get_deployment_scheduled_flow_runs:
+        mock_get_deployment_scheduled_flow_runs.return_value = [{"id": "flow_run_id"}]
+        result = get_dataflow_scheduled_flow_runs("deployment-id")
+        assert result == {"flow_runs": [{"id": "flow_run_id"}]}
+
+
 def test_get_flow_run_by_id_badparams():
     with pytest.raises(TypeError) as excinfo:
         get_flow_run_by_id(123)

tests/test_service.py

@@ -52,6 +52,7 @@
     post_deployment_v1,
     put_deployment_v1,
     get_deployment,
+    get_deployment_scheduled_flow_runs,
     CronSchedule,
     post_deployment_flow_run,
     create_secret_block,
@@ -1091,6 +1092,19 @@ def test_get_deployment(mock_get: Mock):
     assert response == "retval"
 
 
+@patch("proxy.service.prefect_post")
+def test_get_deployment_scheduled_flow_runs(mock_post: Mock):
+    mock_post.return_value = [{"id": "flow_run_id"}]
+    response = get_deployment_scheduled_flow_runs("deployment-id")
+    mock_post.assert_called_once_with(
+        "deployments/get_scheduled_flow_runs",
+        {
+            "deployment_ids": ["deployment-id"],
+        },
+    )
+    assert response == [{"id": "flow_run_id"}]
+
+
 def test_get_flow_runs_by_deployment_id_type_error():
     with pytest.raises(TypeError):
         get_flow_runs_by_deployment_id(123, 10, "")