def get_dags(project_id, dags_config):
"""Return all configured DAGs including associated tasks."""
tasks = []
dag_collection = DagCollection.from_file(dags_config)
for project_dir in project_dirs(project_id):
# parse metadata.yaml to retrieve scheduling information
if os.path.isdir(project_dir):
for root, dirs, files in os.walk(project_dir):
try:
if QUERY_FILE in files:
query_file = os.path.join(root, QUERY_FILE)
task = Task.of_query(query_file,
dag_collection=dag_collection)
elif QUERY_PART_FILE in files:
# multipart query
query_file = os.path.join(root, QUERY_PART_FILE)
task = Task.of_multipart_query(
query_file, dag_collection=dag_collection)
elif SCRIPT_FILE in files:
query_file = os.path.join(root, SCRIPT_FILE)
task = Task.of_script(query_file,
dag_collection=dag_collection)
elif PYTHON_SCRIPT_FILE in files:
query_file = os.path.join(root, PYTHON_SCRIPT_FILE)
task = Task.of_python_script(
query_file, dag_collection=dag_collection)
else:
continue
except FileNotFoundError:
# query has no metadata.yaml file; skip
pass
except UnscheduledTask:
# logging.debug(
# f"No scheduling information for {query_file}."
# )
#
# most tasks lack scheduling information for now
pass
except Exception as e:
# in the case that there was some other error, report the query
# that failed before exiting
logging.error(
f"Error processing task for query {query_file}")
raise e
else:
tasks.append(task)
else:
logging.error("""
Invalid project_dir: {}, project_dir must be a directory with
structure <sql>/<project>/<dataset>/<table>/metadata.yaml.
""".format(project_dir))
return dag_collection.with_tasks(tasks)
def test_of_python_script(self):
query_file = (TEST_DIR / "data" / "test_sql" /
"moz-fx-data-test-project" / "test" /
"incremental_query_v1" / "query.sql")
task = Task.of_python_script(query_file)
assert task.query_file == str(query_file)
assert task.dataset == "test"
assert task.project == "moz-fx-data-test-project"
assert task.table == "incremental_query"
assert task.version == "v1"
assert task.task_name == "test__incremental_query__v1"
assert task.dag_name == "bqetl_events"
assert task.depends_on_past is False
assert task.is_python_script
def test_python_script_to_airflow(self, tmp_path):
query_file = (TEST_DIR / "data" / "test_sql" /
"moz-fx-data-test-project" / "test" /
"python_script_query_v1" / "query.py")
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_test_dag",
"depends_on_past": True,
"arguments": ["--date", "{{ds}}"],
},
)
tasks = [Task.of_python_script(query_file, metadata)]
default_args = {
"depends_on_past": False,
"owner": "*****@*****.**",
"email": ["*****@*****.**"],
"start_date": "2020-01-01",
"retry_delay": "1h",
}
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": default_args,
}
}).with_tasks(tasks)
dags.to_airflow_dags(tmp_path)
result = (tmp_path / "bqetl_test_dag.py").read_text().strip()
expected = ((TEST_DIR / "data" / "dags" /
"python_script_test_dag").read_text().strip())
assert result == expected