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 == []