def test_of_non_existing_table(self):
with pytest.raises(FileNotFoundError):
Metadata.of_table(
"test",
"no_metadata",
"v1",
TEST_DIR / "data" / "test_sql" / "moz-fx-data-test-project",
)
def test_invalid_review_bugs(self):
assert Metadata("Test", "Description", ["*****@*****.**"],
{"review_bugs": [123456]})
with pytest.raises(ValueError):
Metadata(
"Test",
"Description",
["*****@*****.**"],
{"review_bugs": 123456},
)
def test_is_valid_label(self):
assert Metadata.is_valid_label("valid_label")
assert Metadata.is_valid_label("valid-label1")
assert Metadata.is_valid_label("1231")
assert Metadata.is_valid_label("1231-21")
assert Metadata.is_valid_label("a" * 63)
assert Metadata.is_valid_label("låbel") is False
assert Metadata.is_valid_label("a" * 64) is False
assert Metadata.is_valid_label("INVALID") is False
assert Metadata.is_valid_label("invalid.label") is False
assert Metadata.is_valid_label("") is False
def test_invalid_label(self):
with pytest.raises(ValueError):
Metadata(
"Test metadata",
"test description",
["*****@*****.**"],
{"INVALID-KEY": "foo"},
)
with pytest.raises(ValueError):
Metadata(
"Test metadata",
"test description",
["*****@*****.**"],
{"foo": "INVALID-VALUE"},
)
def test_of_table(self):
metadata = Metadata.of_table("test", "non_incremental_query", "v1",
TEST_DIR / "data" / "test_sql")
assert metadata.friendly_name == "Test table for a non-incremental query"
assert metadata.description == "Test table for a non-incremental query"
assert metadata.review_bug() == "1999999"
def test_add_tasks(self):
dag = Dag("bqetl_test_dag", "daily", self.default_args)
query_file = (TEST_DIR / "data" / "test_sql" /
"moz-fx-data-test-project" / "test" /
"incremental_query_v1" / "query.sql")
scheduling = {
"dag_name": "bqetl_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
"task_name": "custom_task_name",
}
metadata = Metadata("test", "test", ["*****@*****.**"], {},
scheduling)
task1 = Task.of_query(query_file)
task2 = Task.of_query(query_file, metadata)
assert dag.tasks == []
dag.add_tasks([task1, task2])
assert len(dag.tasks) == 2
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_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_task_depends_on(self):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"incremental_query_v1" / "query.sql")
scheduling = {
"dag_name":
"bqetl_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
"depends_on": [
{
"dag_name": "external_dag",
"task_id": "external_task"
},
{
"dag_name": "external_dag2",
"task_id": "external_task2",
"execution_delta": "15m",
},
],
}
metadata = Metadata("test", "test", ["*****@*****.**"], {},
scheduling)
task = Task.of_query(query_file, metadata)
assert task.dag_name == "bqetl_test_dag"
assert len(task.depends_on) == 2
assert task.depends_on[0].dag_name == "external_dag"
assert task.depends_on[0].task_id == "external_task"
assert task.depends_on[1].dag_name == "external_dag2"
assert task.depends_on[1].task_id == "external_task2"
assert task.depends_on[1].execution_delta == "15m"
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_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_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 main():
"""Publish query data as JSON to GCS."""
args, query_arguments = parser.parse_known_args()
try:
metadata = Metadata.of_sql_file(args.query_file)
except FileNotFoundError:
print("No metadata file for: {}".format(args.query_file))
return
# check if the data should be published as JSON
if not metadata.is_public_json():
return
if not validate_public_data(metadata, args.query_file):
sys.exit(1)
storage_client = storage.Client()
client = bigquery.Client(args.public_project_id)
publisher = JsonPublisher(
client,
storage_client,
args.public_project_id,
args.query_file,
args.api_version,
args.target_bucket,
args.parameter,
)
publisher.publish_json()
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_validate_custom_task_name(self):
query_file = (TEST_DIR / "data" / "test_sql" /
"moz-fx-data-test-project" / "test" /
"incremental_query_v1" / "query.sql")
scheduling = {
"dag_name": "bqetl_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
"task_name": "a" * 63,
}
metadata = Metadata("test", "test", ["*****@*****.**"], {},
scheduling)
with pytest.raises(ValueError):
Task.of_query(query_file, metadata)
scheduling = {
"dag_name": "bqetl_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
"task_name": "",
}
metadata = Metadata("test", "test", ["*****@*****.**"], {},
scheduling)
with pytest.raises(ValueError):
Task.of_query(query_file, metadata)
scheduling = {
"dag_name": "bqetl_test_dag",
"default_args": {
"owner": "*****@*****.**"
},
"task_name": "a" * 62,
}
metadata = Metadata("test", "test", ["*****@*****.**"], {},
scheduling)
task = Task.of_query(query_file, metadata)
assert task.task_name == "a" * 62
def test_no_dag_name(self):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"incremental_query_v1" / "query.sql")
metadata = Metadata("test", "test", {}, {"foo": "bar"})
with pytest.raises(TaskParseException):
Task(query_file, metadata)
def of_query(cls, query_file):
"""
Create task that schedules the corresponding query in Airflow.
Raises FileNotFoundError if not metadata file exists for query.
"""
metadata = Metadata.of_sql_file(query_file)
return cls(query_file, metadata)
def test_of_sql_file(self):
metadata_file = (TEST_DIR / "data" / "test_sql" / "test" /
"non_incremental_query_v1" / "query.sql")
metadata = Metadata.of_sql_file(metadata_file)
assert metadata.friendly_name == "Test table for a non-incremental query"
assert metadata.description == "Test table for a non-incremental query"
assert metadata.review_bug() == "1999999"
def test_no_dag_name(self):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"incremental_query_v1" / "query.sql")
metadata = Metadata("test", "test", [], {}, self.default_scheduling)
with pytest.raises(TaskParseException):
Task.of_query(query_file, metadata)
def test_unscheduled_task(self, tmp_path):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"incremental_query_v1" / "query.sql")
metadata = Metadata("test", "test", [], {}, {})
with pytest.raises(UnscheduledTask):
Task.of_query(query_file, metadata)
def test_metadata_instantiation(self):
metadata = Metadata("Test metadata", "test description",
["*****@*****.**"], {})
assert metadata.friendly_name == "Test metadata"
assert metadata.description == "test description"
assert metadata.owners == ["*****@*****.**"]
assert metadata.labels == {}
assert metadata.scheduling == {}
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_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_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_of_table(self):
metadata = Metadata.of_table(
"test",
"non_incremental_query",
"v1",
TEST_DIR / "data" / "test_sql" / "moz-fx-data-test-project",
)
assert metadata.friendly_name == "Test table for a non-incremental query"
assert metadata.description == "Test table for a non-incremental query"
assert metadata.review_bugs() == ["1999999", "12121212"]
def test_is_valid_public_data(self):
metadata_not_public = Metadata("No public data", "No public data", {},
{})
assert validate_public_data(metadata_not_public,
"test/path/metadata.yaml")
metadata_valid_public = Metadata(
"Public json data",
"Public json data",
[],
{
"public_json": True,
"review_bug": 123456
},
{},
)
assert validate_public_data(metadata_valid_public,
"test/path/metadata.yaml")
metadata_valid_public = Metadata(
"Public BigQuery data",
"Public BigQuery data",
[],
{
"public_bigquery": True,
"review_bug": 123456
},
{},
)
assert validate_public_data(metadata_valid_public,
"test/path/metadata.yaml")
metadata_invalid_public = Metadata(
"Public BigQuery data",
"Public BigQuery data",
[],
{"public_bigquery": True},
{},
)
assert (validate_public_data(metadata_invalid_public,
"test/path/metadata.yaml") is False)
def test_task_instantiation(self):
query_file = (TEST_DIR / "data" / "test_sql" / "test" /
"incremental_query_v1" / "query.sql")
metadata = Metadata("test", "test", ["*****@*****.**"], {},
self.default_scheduling)
task = Task.of_query(query_file, metadata)
assert task.dag_name == "bqetl_test_dag"
assert task.depends_on_past is False
assert task.task_name == "test__incremental_query__v1"
assert task.public_json is False
def run(
query_file,
dataset_id,
destination_table,
query_arguments,
public_project_id=PUBLIC_PROJECT_ID,
):
"""Execute bq to run a query."""
if dataset_id is not None:
# dataset ID was parsed by argparse but needs to be passed as parameter
# when running the query
query_arguments.append("--dataset_id={}".format(dataset_id))
use_public_table = False
try:
metadata = Metadata.of_query_file(query_file)
if metadata.is_public_bigquery():
if not validate_public_data(metadata, query_file):
sys.exit(1)
# change the destination table to write results to the public dataset;
# a view to the public table in the internal dataset is created
# when CI runs
if (dataset_id is not None and destination_table is not None
and re.match(DESTINATION_TABLE_RE, destination_table)):
destination_table = "{}:{}.{}".format(public_project_id,
dataset_id,
destination_table)
query_arguments.append(
"--destination_table={}".format(destination_table))
use_public_table = True
else:
print("ERROR: Cannot run public dataset query. Parameters"
" --destination_table=<table without dataset ID> and"
" --dataset_id=<dataset> required")
sys.exit(1)
except yaml.YAMLError as e:
print(e)
sys.exit(1)
except FileNotFoundError:
print("INFO: No metadata.yaml found for {}", query_file)
if not use_public_table and destination_table is not None:
# destination table was parsed by argparse, however if it wasn't modified to
# point to a public table it needs to be passed as parameter for the query
query_arguments.append(
"--destination_table={}".format(destination_table))
with open(query_file) as query_stream:
# run the query as shell command so that passed parameters can be used as is
subprocess.check_call(["bq"] + query_arguments, stdin=query_stream)
def test_of_query_file(self):
metadata_file = (TEST_DIR / "data" / "test_sql" /
"moz-fx-data-test-project" / "test" /
"non_incremental_query_v1" / "query.sql")
metadata = Metadata.of_query_file(metadata_file)
assert metadata.friendly_name == "Test table for a non-incremental query"
assert metadata.description == "Test table for a non-incremental query"
assert metadata.review_bugs() == ["1999999", "12121212"]
assert metadata.workgroup_access is not None
assert metadata.workgroup_access[0].role == "roles/bigquery.dataViewer"
assert (metadata.workgroup_access[0].members[0] ==
"workgroup:dataops-managed/taar")
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 == []
