def test_dags_with_invalid_tasks(self):
with pytest.raises(InvalidDag):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"incremental_query_v1" / "query.sql")
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_non_exisiting_dag",
"depends_on_past": True,
"param": "test_param",
},
)
tasks = [Task.of_query(query_file, metadata)]
DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": self.default_args,
}
}).with_tasks(tasks)
def test_to_airflow_duplicate_dependencies(self, tmp_path):
query_file = (
TEST_DIR
/ "data"
/ "test_sql"
/ "moz-fx-data-test-project"
/ "test"
/ "non_incremental_query_v1"
/ "query.sql"
)
query_file2 = (
TEST_DIR
/ "data"
/ "test_sql"
/ "moz-fx-data-test-project"
/ "test"
/ "no_metadata_query_v1"
/ "query.sql"
)
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_test_dag",
"depends_on_past": True,
"depends_on": [{"dag_name": "external", "task_id": "task1"}],
},
)
tasks = [
Task.of_query(query_file, metadata),
Task.of_query(query_file2, metadata),
]
default_args = {
"owner": "*****@*****.**",
"start_date": "2020-01-01",
}
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" / "test_dag_duplicate_dependencies")
.read_text()
.strip()
)
assert result == expected
def test_public_data_json_dag_add_task(self):
public_json_dag = PublicDataJsonDag("bqetl_public_data_json_dag",
"daily", self.default_args)
assert len(public_json_dag.tasks) == 0
query_file = (TEST_DIR / "data" / "test_sql" /
"moz-fx-data-test-project" / "test" /
"incremental_query_v1" / "query.sql")
task = Task.of_query(query_file)
dag = Dag(
"bqetl_events",
"0 1 * * *",
DagDefaultArgs("*****@*****.**", "2020-01-01"),
[task],
)
public_json_dag.add_export_tasks([task], DagCollection([dag]))
assert len(public_json_dag.tasks) == 1
assert (public_json_dag.tasks[0].task_name ==
"export_public_data_json_test__incremental_query__v1")
assert public_json_dag.tasks[
0].dag_name == "bqetl_public_data_json_dag"
assert len(public_json_dag.tasks[0].dependencies) == 1
def test_dags_with_tasks(self):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"incremental_query_v1" / "query.sql")
metadata = Metadata(
"test",
"test",
{},
{
"dag_name": "test_dag",
"depends_on_past": True,
"param": "test_param"
},
)
tasks = [Task(query_file, metadata)]
dags = DagCollection.from_dict({
"test_dag": {
"schedule_interval": "daily",
"default_args": {}
}
}).with_tasks(tasks)
assert len(dags.dags) == 1
dag = dags.dag_by_name("test_dag")
assert len(dag.tasks) == 1
assert dag.tasks[0].dag_name == "test_dag"
def test_public_json_dag_to_airflow(self, tmp_path):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"non_incremental_query_v1" / "query.sql")
tasks = [Task.of_query(query_file)]
default_args = {
"depends_on_past": False,
"owner": "*****@*****.**",
"email": ["*****@*****.**"],
"start_date": "2020-01-01",
"retry_delay": "1h",
}
dags = DagCollection.from_dict({
"bqetl_public_data_json": {
"schedule_interval": "daily",
"default_args": default_args,
},
"bqetl_core": {
"schedule_interval": "daily",
"default_args": default_args,
},
}).with_tasks(tasks)
dags.to_airflow_dags(tmp_path)
result = (tmp_path / "bqetl_public_data_json.py").read_text().strip()
expected_dag = ((TEST_DIR / "data" / "dags" /
"test_public_data_json_dag").read_text().strip())
assert result == expected_dag
def test_dags_with_tasks(self):
query_file = (TEST_DIR / "data" / "test_sql" /
"moz-fx-data-test-project" / "test" /
"incremental_query_v1" / "query.sql")
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_test_dag",
"depends_on_past": True
},
)
tasks = [Task.of_query(query_file, metadata)]
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": self.default_args,
}
}).with_tasks(tasks)
assert len(dags.dags) == 1
dag = dags.dag_by_name("bqetl_test_dag")
assert len(dag.tasks) == 1
assert dag.tasks[0].dag_name == "bqetl_test_dag"
def test_dags_from_dict(self):
dags = DagCollection.from_dict({
"test_dag1": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**"
},
},
"test_dag2": {
"schedule_interval": "daily",
"default_args": {}
},
})
assert len(dags.dags) == 2
assert dags.dag_by_name("test_dag1") is not None
assert dags.dag_by_name("test_dag2") is not None
dag1 = dags.dag_by_name("test_dag1")
assert len(dag1.tasks) == 0
assert dag1.schedule_interval == "daily"
assert dag1.default_args == {"owner": "*****@*****.**"}
dag2 = dags.dag_by_name("test_dag2")
assert len(dag2.tasks) == 0
assert dag2.schedule_interval == "daily"
assert dag2.default_args == {}
def get_dags(sql_dir, dags_config):
"""Return all configured DAGs including associated tasks."""
tasks = []
# parse metadata.yaml to retrieve scheduling information
if os.path.isdir(sql_dir):
for root, dirs, files in os.walk(sql_dir):
if QUERY_FILE in files:
query_file = os.path.join(root, QUERY_FILE)
try:
task = Task.of_query(query_file)
tasks.append(task)
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
else:
logging.error("""
Invalid sql_dir: {}, sql_dir must be a directory with
structure /<dataset>/<table>/metadata.yaml.
""".format(sql_dir))
return DagCollection.from_file(dags_config).with_tasks(tasks)
def test_task_for_table(self):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"incremental_query_v1" / "query.sql")
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_test_dag",
"depends_on_past": True
},
)
tasks = [Task.of_query(query_file, metadata)]
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": self.default_args,
}
}).with_tasks(tasks)
task = dags.task_for_table("test", "incremental_query_v1")
assert task
assert task.dag_name == "bqetl_test_dag"
def test_task_get_nested_view_dependencies(
self, tmp_path, bigquery_client, project_id, temporary_dataset
):
query_file_path = tmp_path / "sql" / temporary_dataset / "query_v1"
os.makedirs(query_file_path)
query_file = query_file_path / "query.sql"
query_file.write_text(
f"SELECT * FROM {project_id}.{temporary_dataset}.table1_v1 "
+ f"UNION ALL SELECT * FROM {project_id}.{temporary_dataset}.test_view"
)
schema = [bigquery.SchemaField("a", "STRING", mode="NULLABLE")]
table = bigquery.Table(
f"{project_id}.{temporary_dataset}.table1_v1", schema=schema
)
bigquery_client.create_table(table)
table = bigquery.Table(
f"{project_id}.{temporary_dataset}.table2_v1", schema=schema
)
bigquery_client.create_table(table)
view = bigquery.Table(f"{project_id}.{temporary_dataset}.test_view2")
view.view_query = f"SELECT * FROM {project_id}.{temporary_dataset}.table2_v1"
bigquery_client.create_table(view)
view = bigquery.Table(f"{project_id}.{temporary_dataset}.test_view")
view.view_query = f"SELECT * FROM {project_id}.{temporary_dataset}.test_view2"
bigquery_client.create_table(view)
metadata = Metadata(
"test", "test", ["*****@*****.**"], {}, self.default_scheduling
)
task = Task.of_query(query_file, metadata)
table_task1 = Task.of_query(
tmp_path / "sql" / temporary_dataset / "table1_v1" / "query.sql", metadata
)
table_task2 = Task.of_query(
tmp_path / "sql" / temporary_dataset / "table2_v1" / "query.sql", metadata
)
dags = DagCollection.from_dict(
{
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**",
"start_date": "2020-01-01",
},
}
}
).with_tasks([task, table_task1, table_task2])
task.with_dependencies(bigquery_client, dags)
result = task.dependencies
tables = [f"{t.dataset}__{t.table}__{t.version}" for t in result]
assert f"{temporary_dataset}__table1__v1" in tables
assert f"{temporary_dataset}__table2__v1" in tables
def test_task_for_non_existing_table(self):
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": self.default_args,
}
}).with_tasks([])
assert dags.task_for_table("test", "non_existing_table") is None
def test_dag_by_name(self):
dags = DagCollection.from_dict(
{"test_dag1": {
"schedule_interval": "daily",
"default_args": {}
}})
assert dags.dag_by_name("test_dag1") is not None
assert dags.dag_by_name("test_dag1").name == "test_dag1"
assert dags.dag_by_name("non_existing") is None
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_task_get_view_dependencies(self, tmp_path, bigquery_client,
project_id, temporary_dataset):
query_file_path = tmp_path / "sql" / temporary_dataset / "query_v1"
os.makedirs(query_file_path)
query_file = query_file_path / "query.sql"
query_file.write_text(
f"SELECT * FROM {project_id}.{temporary_dataset}.table1_v1 " +
f"UNION ALL SELECT * FROM {project_id}.{temporary_dataset}.test_view"
)
schema = [bigquery.SchemaField("a", "STRING", mode="NULLABLE")]
table = bigquery.Table(f"{project_id}.{temporary_dataset}.table1_v1",
schema=schema)
bigquery_client.create_table(table)
table = bigquery.Table(f"{project_id}.{temporary_dataset}.table2_v1",
schema=schema)
bigquery_client.create_table(table)
view = bigquery.Table(f"{project_id}.{temporary_dataset}.test_view")
view.view_query = f"SELECT * FROM {project_id}.{temporary_dataset}.table2_v1"
bigquery_client.create_table(view)
metadata = Metadata(
"test",
"test",
{},
{
"dag_name": "test_dag",
"depends_on_past": True,
"param": "test_param"
},
)
task = Task(query_file, metadata)
table_task1 = Task(
tmp_path / "sql" / temporary_dataset / "table1_v1" / "query.sql",
metadata)
table_task2 = Task(
tmp_path / "sql" / temporary_dataset / "table2_v1" / "query.sql",
metadata)
dags = DagCollection.from_dict({
"test_dag": {
"schedule_interval": "daily",
"default_args": {}
}
}).with_tasks([task, table_task1, table_task2])
result = task.get_dependencies(bigquery_client, dags)
tables = [f"{t.dataset}__{t.table}__{t.version}" for t in result]
assert f"{temporary_dataset}__table1__v1" in tables
assert f"{temporary_dataset}__table2__v1" in tables
def test_task_get_nested_view_dependencies(self, tmp_path):
query_file_path = tmp_path / "test-project" / "test" / "query_v1"
os.makedirs(query_file_path)
query_file = query_file_path / "query.sql"
query_file.write_text(
"SELECT * FROM `test-project`.test.table1_v1 "
"UNION ALL SELECT * FROM `test-project`.test.test_view")
view_file_path = tmp_path / "test-project" / "test" / "test_view"
os.makedirs(view_file_path)
view_file = view_file_path / "view.sql"
view_file.write_text(
"CREATE OR REPLACE VIEW `test-project`.test.test_view "
"AS SELECT * FROM `test-project`.test.test_view2")
view2_file_path = tmp_path / "test-project" / "test" / "test_view2"
os.makedirs(view2_file_path)
view2_file = view2_file_path / "view.sql"
view2_file.write_text(
"CREATE OR REPLACE VIEW `test-project`.test.test_view2 "
"AS SELECT * FROM `test-project`.test.table2_v1")
metadata = Metadata("test", "test", ["*****@*****.**"], {},
self.default_scheduling)
task = Task.of_query(query_file, metadata)
table_task1 = Task.of_query(
tmp_path / "test-project" / "test" / "table1_v1" / "query.sql",
metadata,
)
table_task2 = Task.of_query(
tmp_path / "test-project" / "test" / "table2_v1" / "query.sql",
metadata,
)
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**",
"start_date": "2020-01-01",
},
}
}).with_tasks([task, table_task1, table_task2])
task.with_dependencies(dags)
result = task.dependencies
tables = [t.task_id for t in result]
assert "test__table1__v1" in tables
assert "test__table2__v1" in tables
def test_multipart_task_get_dependencies(self, tmp_path, bigquery_client,
project_id, temporary_dataset):
query_file_path = tmp_path / project_id / temporary_dataset / "query_v1"
os.makedirs(query_file_path)
query_file_part1 = query_file_path / "part1.sql"
query_file_part1.write_text(
f"SELECT * FROM {project_id}.{temporary_dataset}.table1_v1")
query_file_part2 = query_file_path / "part2.sql"
query_file_part2.write_text(
f"SELECT * FROM {project_id}.{temporary_dataset}.table2_v1")
schema = [bigquery.SchemaField("a", "STRING", mode="NULLABLE")]
table = bigquery.Table(f"{project_id}.{temporary_dataset}.table1_v1",
schema=schema)
bigquery_client.create_table(table)
table = bigquery.Table(f"{project_id}.{temporary_dataset}.table2_v1",
schema=schema)
bigquery_client.create_table(table)
metadata = Metadata("test", "test", ["*****@*****.**"], {},
self.default_scheduling)
task = Task.of_multipart_query(query_file_part1, metadata)
table_task1 = Task.of_query(
tmp_path / project_id / temporary_dataset / "table1_v1" /
"query.sql",
metadata,
)
table_task2 = Task.of_query(
tmp_path / project_id / temporary_dataset / "table2_v1" /
"query.sql",
metadata,
)
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**",
"start_date": "2020-01-01",
},
}
}).with_tasks([task, table_task1, table_task2])
task.with_dependencies(dags)
result = task.dependencies
tables = [t.task_id for t in result]
assert f"{temporary_dataset}__table1__v1" in tables
assert f"{temporary_dataset}__table2__v1" in tables
def test_dags_from_file(self):
dags_file = TEST_DIR / "data" / "dags.yaml"
dags = DagCollection.from_file(dags_file)
assert len(dags.dags) == 2
assert dags.dag_by_name("not existing") is None
assert dags.dag_by_name("bqetl_events") is not None
assert dags.dag_by_name("bqetl_core") is not None
events_dag = dags.dag_by_name("bqetl_events")
assert len(events_dag.tasks) == 0
core_dag = dags.dag_by_name("bqetl_core")
assert len(core_dag.tasks) == 0
def test_task_get_dependencies_none(self, tmp_path, bigquery_client):
query_file_path = tmp_path / "sql" / "test" / "query_v1"
os.makedirs(query_file_path)
query_file = query_file_path / "query.sql"
query_file.write_text("SELECT 123423")
metadata = Metadata("test", "test", ["*****@*****.**"], {},
self.default_scheduling)
task = Task.of_query(query_file, metadata)
dags = DagCollection.from_dict({})
task.with_dependencies(bigquery_client, dags)
assert task.dependencies == []
def test_to_airflow(self, tmp_path):
query_file = (
TEST_DIR
/ "data"
/ "test_sql"
/ "moz-fx-data-test-project"
/ "test"
/ "non_incremental_query_v1"
/ "query.sql"
)
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_test_dag",
"depends_on_past": True,
"param": "test_param",
"arguments": ["--append_table"],
},
)
tasks = [Task.of_query(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" / "simple_test_dag").read_text().strip()
assert result == expected
def test_multipart_task_get_dependencies(self, tmp_path):
query_file_path = tmp_path / "test-project" / "test" / "query_v1"
os.makedirs(query_file_path)
query_file_part1 = query_file_path / "part1.sql"
query_file_part1.write_text(
"SELECT * FROM `test-project`.test.table1_v1")
query_file_part2 = query_file_path / "part2.sql"
query_file_part2.write_text(
"SELECT * FROM `test-project`.test.table2_v1")
metadata = Metadata("test", "test", ["*****@*****.**"], {},
self.default_scheduling)
task = Task.of_multipart_query(query_file_part1, metadata)
table_task1 = Task.of_query(
tmp_path / "test-project" / "test" / "table1_v1" / "query.sql",
metadata,
)
table_task2 = Task.of_query(
tmp_path / "test-project" / "test" / "table2_v1" / "query.sql",
metadata,
)
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**",
"start_date": "2020-01-01",
},
}
}).with_tasks([task, table_task1, table_task2])
task.with_dependencies(dags)
result = task.dependencies
tables = [t.task_id for t in result]
assert "test__table1__v1" in tables
assert "test__table2__v1" in tables
def test_task_get_dependencies_none(self, tmp_path, bigquery_client):
query_file_path = tmp_path / "sql" / "test" / "query_v1"
os.makedirs(query_file_path)
query_file = query_file_path / "query.sql"
query_file.write_text("SELECT 123423")
metadata = Metadata(
"test",
"test",
{},
{
"dag_name": "test_dag",
"depends_on_past": True,
"param": "test_param"
},
)
task = Task(query_file, metadata)
dags = DagCollection.from_dict({})
assert task.get_dependencies(bigquery_client, dags) == []
def test_to_airflow_with_dependencies(self, tmp_path):
query_file_path = tmp_path / "test-project" / "test" / "query_v1"
os.makedirs(query_file_path)
query_file = query_file_path / "query.sql"
query_file.write_text(
"SELECT * FROM `test-project`.test.table1_v1 "
"UNION ALL SELECT * FROM `test-project`.test.table2_v1 "
"UNION ALL SELECT * FROM `test-project`.test.external_table_v1")
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
},
)
task = Task.of_query(query_file, metadata)
table_task1 = Task.of_query(
tmp_path / "test-project" / "test" / "table1_v1" / "query.sql",
metadata,
)
os.makedirs(tmp_path / "test-project" / "test" / "table1_v1")
query_file = tmp_path / "test-project" / "test" / "table1_v1" / "query.sql"
query_file.write_text("SELECT 1")
table_task2 = Task.of_query(
tmp_path / "test-project" / "test" / "table2_v1" / "query.sql",
metadata,
)
os.makedirs(tmp_path / "test-project" / "test" / "table2_v1")
query_file = tmp_path / "test-project" / "test" / "table2_v1" / "query.sql"
query_file.write_text("SELECT 2")
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_external_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
},
)
external_table_task = Task.of_query(
tmp_path / "test-project" / "test" / "external_table_v1" /
"query.sql",
metadata,
)
os.makedirs(tmp_path / "test-project" / "test" / "external_table_v1")
query_file = (tmp_path / "test-project" / "test" /
"external_table_v1" / "query.sql")
query_file.write_text("SELECT 3")
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**",
"start_date": "2020-05-25",
},
},
"bqetl_external_test_dag": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**",
"start_date": "2020-05-25",
},
},
}).with_tasks([task, table_task1, table_task2, external_table_task])
dags.to_airflow_dags(tmp_path)
expected_dag_with_dependencies = (
(TEST_DIR / "data" / "dags" /
"test_dag_with_dependencies").read_text().strip())
expected_dag_external_dependency = (
(TEST_DIR / "data" / "dags" /
"test_dag_external_dependency").read_text().strip())
dag_with_dependencies = (tmp_path /
"bqetl_test_dag.py").read_text().strip()
dag_external_dependency = ((
tmp_path / "bqetl_external_test_dag.py").read_text().strip())
assert dag_with_dependencies == expected_dag_with_dependencies
assert dag_external_dependency == expected_dag_external_dependency
def test_dags_from_invalid_dict(self):
with pytest.raises(DagParseException):
DagCollection.from_dict({"foo": "bar"})
def test_dags_from_empty_dict(self):
dags = DagCollection.from_dict({})
assert len(dags.dags) == 0
def test_to_airflow_with_dependencies(self, tmp_path, project_id,
temporary_dataset, bigquery_client):
query_file_path = tmp_path / "sql" / temporary_dataset / "query_v1"
os.makedirs(query_file_path)
query_file = query_file_path / "query.sql"
query_file.write_text(
f"SELECT * FROM {project_id}.{temporary_dataset}.table1_v1 " +
f"UNION ALL SELECT * FROM {project_id}.{temporary_dataset}.table2_v1 "
+ "UNION ALL SELECT * FROM " +
f"{project_id}.{temporary_dataset}.external_table_v1")
schema = [bigquery.SchemaField("a", "STRING", mode="NULLABLE")]
table = bigquery.Table(f"{project_id}.{temporary_dataset}.table1_v1",
schema=schema)
bigquery_client.create_table(table)
table = bigquery.Table(f"{project_id}.{temporary_dataset}.table2_v1",
schema=schema)
bigquery_client.create_table(table)
table = bigquery.Table(
f"{project_id}.{temporary_dataset}.external_table_v1",
schema=schema)
bigquery_client.create_table(table)
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
},
)
task = Task.of_query(query_file, metadata)
table_task1 = Task.of_query(
tmp_path / "sql" / temporary_dataset / "table1_v1" / "query.sql",
metadata)
os.makedirs(tmp_path / "sql" / temporary_dataset / "table1_v1")
query_file = tmp_path / "sql" / temporary_dataset / "table1_v1" / "query.sql"
query_file.write_text("SELECT 1")
table_task2 = Task.of_query(
tmp_path / "sql" / temporary_dataset / "table2_v1" / "query.sql",
metadata)
os.makedirs(tmp_path / "sql" / temporary_dataset / "table2_v1")
query_file = tmp_path / "sql" / temporary_dataset / "table2_v1" / "query.sql"
query_file.write_text("SELECT 2")
metadata = Metadata(
"test",
"test",
["*****@*****.**"],
{},
{
"dag_name": "bqetl_external_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
},
)
external_table_task = Task.of_query(
tmp_path / "sql" / temporary_dataset / "external_table_v1" /
"query.sql",
metadata,
)
os.makedirs(tmp_path / "sql" / temporary_dataset / "external_table_v1")
query_file = (tmp_path / "sql" / temporary_dataset /
"external_table_v1" / "query.sql")
query_file.write_text("SELECT 3")
dags = DagCollection.from_dict({
"bqetl_test_dag": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**",
"start_date": "2020-05-25",
},
},
"bqetl_external_test_dag": {
"schedule_interval": "daily",
"default_args": {
"owner": "*****@*****.**",
"start_date": "2020-05-25",
},
},
}).with_tasks([task, table_task1, table_task2, external_table_task])
dags.to_airflow_dags(tmp_path)
# we need to use templates since the temporary dataset name changes between runs
env = Environment(loader=PackageLoader("tests", "data/dags"))
dag_template_with_dependencies = env.get_template(
"test_dag_with_dependencies")
dag_template_external_dependency = env.get_template(
"test_dag_external_dependency")
args = {"temporary_dataset": temporary_dataset}
expected_dag_with_dependencies = dag_template_with_dependencies.render(
args)
expected_dag_external_dependency = dag_template_external_dependency.render(
args)
dag_with_dependencies = (tmp_path /
"bqetl_test_dag.py").read_text().strip()
dag_external_dependency = ((
tmp_path / "bqetl_external_test_dag.py").read_text().strip())
assert dag_with_dependencies == expected_dag_with_dependencies
assert dag_external_dependency == expected_dag_external_dependency
def test_dags_from_empty_file(self, tmp_path):
dags_file = tmp_path / "dags.yaml"
dags_file.write_text("")
dags = DagCollection.from_file(dags_file)
assert len(dags.dags) == 0