def get_airflow_dagruns_to_sync(
self,
last_seen_dag_run_id: Optional[int],
last_seen_log_id: Optional[int],
extra_dag_run_ids: Optional[List[int]],
dag_ids: Optional[str],
) -> AirflowDagRunsResponse:
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
dag_ids_list = dag_ids.split(",") if dag_ids else None
with self._get_session() as session:
data = get_new_dag_runs(
last_seen_dag_run_id=last_seen_dag_run_id,
last_seen_log_id=last_seen_log_id,
extra_dag_runs_ids=extra_dag_run_ids,
dag_ids=dag_ids_list,
include_subdags=True,
session=session,
)
self._raise_on_plugin_error_message(data,
"get_airflow_dagruns_to_sync")
json_data = json_conv(data)
self._on_data_received(json_data, "get_airflow_dagruns_to_sync")
return AirflowDagRunsResponse.from_dict(json_data)
def test_06_both_0(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(with_log=True)
result = get_new_dag_runs(0, 0, [])
self.validate_result(result, 1, 1, 1, expected_max_log_ids=[1])
def test_05_log_id_none(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, with_log=True)
result = get_new_dag_runs(1, None, [])
self.validate_result(result, 2, 3, 3)
def test_03_running(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, state="running", with_log=False)
result = get_new_dag_runs(None, None, [])
self.validate_result(result, 3, 3, None)
def test_07_big_run_id(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, with_log=False)
result = get_new_dag_runs(3, 0, [])
self.validate_result(result, 0, 3, None)
assert result
def test_13_fetch_in_chunks(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from dbnd_airflow.export_plugin.queries import (
MAX_PARAMETERS_INSIDE_IN_CLAUSE,
_find_dag_runs_by_list_in_chunks,
)
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
with mock.patch(
"dbnd_airflow.export_plugin.queries._find_dag_runs_by_list_in_chunks",
wraps=lambda a, b: _find_dag_runs_by_list_in_chunks(a, b),
) as m:
insert_dag_runs(dag_runs_count=MAX_PARAMETERS_INSIDE_IN_CLAUSE - 1,
with_log=True)
get_new_dag_runs(0, 0, [], [])
assert m.call_count == 0
insert_dag_runs(dag_runs_count=1, with_log=True)
get_new_dag_runs(0, 0, [], [])
assert m.call_count == 1
def test_10_running_paused(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import (
insert_dag_runs,
set_dag_is_paused,
)
insert_dag_runs(dag_runs_count=1, state="running", with_log=True)
set_dag_is_paused(is_paused=True)
result = get_new_dag_runs(1, 1, [])
self.validate_result(result, 0, 1, 1, True)
def test_09_paused(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import (
insert_dag_runs,
set_dag_is_paused,
)
insert_dag_runs(dag_runs_count=1, with_log=True)
set_dag_is_paused(is_paused=True)
result = get_new_dag_runs(0, 0, [])
self.validate_result(result, 1, 1, 1, True, expected_max_log_ids=[1])
def test_12_dag_ids(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, with_log=True)
insert_dag_runs(dag_id="plugin_other_dag",
dag_runs_count=3,
with_log=True)
result = get_new_dag_runs(0, 0, [], ["plugin_other_dag"])
self.validate_result(result, 3, 6, 6, expected_max_log_ids=[4, 5, 6])
def test_11_extra_dag_runs(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
from test_dbnd_airflow.export_plugin.db_data_generator import insert_dag_runs
insert_dag_runs(dag_runs_count=3, with_log=True)
result = get_new_dag_runs(2, 2, [1, 2])
self.validate_result(result,
3,
3,
3,
expected_max_log_ids=[None, None, 3])
def process_new_runs_request():
last_seen_dag_run_id = flask.request.values.get("last_seen_dag_run_id", type=int)
last_seen_log_id = flask.request.values.get("last_seen_log_id", type=int)
extra_dag_runs_ids = convert_url_param_value_to_list("extra_dag_runs_ids", int, [])
dag_ids = convert_url_param_value_to_list("dag_ids", str, None)
# default to true
include_subdags = flask.request.values.get("include_subdags", "").lower() != "false"
return json_response(
get_new_dag_runs(
last_seen_dag_run_id,
last_seen_log_id,
extra_dag_runs_ids,
dag_ids,
include_subdags,
).as_dict()
)
def test_01_empty_db(self):
from dbnd_airflow.export_plugin.api_functions import get_new_dag_runs
result = get_new_dag_runs(1, 1, [])
self.validate_result(result, 0, None, None)