def test_get_managed_views_for_dataset_map_empty_list(self) -> None: # Arrange walker = BigQueryViewDagWalker(self.empty_view_list) # Act result_dict = get_managed_view_and_materialized_table_addresses_by_dataset( walker ) # Assert expected_result: Dict[str, Set[BigQueryAddress]] = {} self.assertEqual(expected_result, result_dict)
def main() -> None: """Executes the main flow of the script.""" parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( "--project-id", choices=[GCP_PROJECT_STAGING, GCP_PROJECT_PRODUCTION], help= "Used to select which GCP project against which to run this script.", required=True, ) parser.add_argument( "--dry-run", default=True, type=str_to_bool, help= "Runs delete in dry-run mode, only prints the views/tables it would delete.", ) args = parser.parse_args() logging.getLogger().setLevel(logging.INFO) with local_project_id_override(args.project_id): views = [] for view_builder in DEPLOYED_VIEW_BUILDERS: if view_builder.dataset_id in VIEW_SOURCE_TABLE_DATASETS: raise ValueError( f"Found view [{view_builder.view_id}] in source-table-only dataset [{view_builder.dataset_id}]" ) try: view = view_builder.build() except BigQueryViewBuilderShouldNotBuildError: logging.warning( "Condition failed for view builder %s in dataset %s. Continuing without it.", view_builder.view_id, view_builder.dataset_id, ) continue views.append(view) dag_walker = BigQueryViewDagWalker(views) managed_views_map = ( get_managed_view_and_materialized_table_addresses_by_dataset( dag_walker)) cleanup_datasets_and_delete_unmanaged_views( bq_client=BigQueryClientImpl(), managed_views_map=managed_views_map, dry_run=args.dry_run, )
def test_get_managed_views_for_dataset_map_all_views_same_dataset(self) -> None: # Arrange walker = BigQueryViewDagWalker(self.all_views_same_dataset) # Act result_dict = get_managed_view_and_materialized_table_addresses_by_dataset( walker ) # Assert expected_result: Dict[str, Set[BigQueryAddress]] = { "dataset_1": { BigQueryAddress(dataset_id="dataset_1", table_id="table_1"), BigQueryAddress(dataset_id="dataset_1", table_id="table_2"), BigQueryAddress(dataset_id="dataset_1", table_id="table_3"), }, } self.assertEqual(expected_result, result_dict)
def test_get_managed_views_for_dataset_map_x_shaped_dag(self) -> None: # Arrange walker = BigQueryViewDagWalker(self.x_shaped_dag_views_list) # Act result_dict = get_managed_view_and_materialized_table_addresses_by_dataset( walker ) # Assert expected_result: Dict[str, Set[BigQueryAddress]] = { "dataset_1": {BigQueryAddress(dataset_id="dataset_1", table_id="table_1")}, "dataset_2": {BigQueryAddress(dataset_id="dataset_2", table_id="table_2")}, "dataset_3": {BigQueryAddress(dataset_id="dataset_3", table_id="table_3")}, "dataset_4": {BigQueryAddress(dataset_id="dataset_4", table_id="table_4")}, "dataset_5": {BigQueryAddress(dataset_id="dataset_5", table_id="table_5")}, } self.assertEqual(expected_result, result_dict)
def _create_managed_dataset_and_deploy_views( views_to_update: List[BigQueryView], bq_region_override: Optional[str], force_materialize: bool, set_temp_dataset_table_expiration: bool = False, ) -> None: """Create and update the given views and their parent datasets. Cleans up unmanaged views and datasets For each dataset key in the given dictionary, creates the dataset if it does not exist, and creates or updates the underlying views mapped to that dataset. If a view has a set materialized_address field, materializes the view into a table. Then, cleans up BigQuery by deleting unmanaged datasets and unmanaged views within managed datasets. This is not performed if a temporary dataset table expiration is already set. Args: views_to_update: A list of view objects to be created or updated. set_temp_dataset_table_expiration: If True, new datasets will be created with an expiration of TEMP_DATASET_DEFAULT_TABLE_EXPIRATION_MS. """ bq_client = BigQueryClientImpl(region_override=bq_region_override) dag_walker = BigQueryViewDagWalker(views_to_update) managed_views_map = get_managed_view_and_materialized_table_addresses_by_dataset( dag_walker) managed_dataset_ids = list(managed_views_map.keys()) _create_all_datasets_if_necessary(bq_client, managed_dataset_ids, set_temp_dataset_table_expiration) if not set_temp_dataset_table_expiration: # We don't want to be deleting unmanaged views/tables if we're creating sandbox datasets cleanup_datasets_and_delete_unmanaged_views(bq_client, managed_views_map, dry_run=False) def process_fn(v: BigQueryView, parent_results: Dict[BigQueryView, bool]) -> bool: """Returns True if this view or any of its parents were updated.""" return _create_or_update_view_and_materialize_if_necessary( bq_client, v, parent_results, force_materialize) dag_walker.process_dag(process_fn)
def test_cleanup_datasets_and_delete_unmanaged_views_dry_run( self, mock_get_datasets_that_have_ever_been_managed: mock.MagicMock ) -> None: mock_get_datasets_that_have_ever_been_managed.return_value = { "dataset_1", "dataset_2", "dataset_unmanaged", } dataset = bigquery.dataset.DatasetReference(self.project_id, "dataset_1") dataset_2 = bigquery.dataset.DatasetReference(self.project_id, "dataset_2") dataset_3 = bigquery.dataset.DatasetReference( self.project_id, "dataset_unmanaged" ) mock_view_builders = [ SimpleBigQueryViewBuilder( dataset_id="dataset_1", view_id="my_fake_view", description="my_fake_view description", view_query_template="SELECT NULL LIMIT 0", should_materialize=False, ), SimpleBigQueryViewBuilder( dataset_id="dataset_2", view_id="my_fake_view_2", description="my_fake_view_2 description", view_query_template="SELECT NULL LIMIT 0", should_materialize=False, ), ] mock_table_resource_ds_1_table = { "tableReference": { "projectId": self.project_id, "datasetId": "dataset_1", "tableId": "my_fake_view", }, } mock_table_resource_ds_1_table_bogus = { "tableReference": { "projectId": self.project_id, "datasetId": "dataset_1", "tableId": "bogus_view_1", }, } mock_table_resource_ds_2_table = { "tableReference": { "projectId": self.project_id, "datasetId": "dataset_2", "tableId": "my_fake_view_2", }, } mock_table_resource_ds_2_table_bogus = { "tableReference": { "projectId": self.project_id, "datasetId": "dataset_2", "tableId": "bogus_view_2", }, } def mock_list_tables(dataset_id: str) -> bigquery.table.TableListItem: if dataset_id == dataset.dataset_id: return [ bigquery.table.TableListItem(mock_table_resource_ds_1_table), bigquery.table.TableListItem(mock_table_resource_ds_1_table_bogus), ] if dataset_id == dataset_2.dataset_id: return [ bigquery.table.TableListItem(mock_table_resource_ds_2_table), bigquery.table.TableListItem(mock_table_resource_ds_2_table_bogus), ] raise ValueError(f"No tables for id: {dataset_id}") self.mock_client.list_tables.side_effect = mock_list_tables def get_dataset_ref(dataset_id: str) -> bigquery.dataset.DatasetReference: if dataset_id == dataset.dataset_id: return dataset if dataset_id == dataset_2.dataset_id: return dataset_2 if dataset_id == dataset_3.dataset_id: return dataset_3 raise ValueError(f"No dataset for id: {dataset_id}") self.mock_client.dataset_ref_for_id.side_effect = get_dataset_ref self.mock_client.dataset_exists.return_value = True views_to_update = [view_builder.build() for view_builder in mock_view_builders] dag_walker = BigQueryViewDagWalker(views_to_update) managed_views_map = ( get_managed_view_and_materialized_table_addresses_by_dataset(dag_walker) ) cleanup_datasets_and_delete_unmanaged_views( self.mock_client, managed_views_map, dry_run=True ) self.mock_client.delete_dataset.assert_not_called() self.mock_client.list_tables.assert_called() self.mock_client.delete_table.assert_not_called()
def test_cleanup_datasets_and_delete_unmanaged_views_unmanaged_dataset_and_dataset_not_in_BigQuery( self, mock_get_datasets_that_have_ever_been_managed: mock.MagicMock ) -> None: mock_get_datasets_that_have_ever_been_managed.return_value = { "dataset_1", "dataset_unmanaged", } dataset = bigquery.dataset.DatasetReference(self.project_id, "dataset_1") dataset_2 = bigquery.dataset.DatasetReference( self.project_id, "dataset_unmanaged" ) mock_view_builders = [ SimpleBigQueryViewBuilder( dataset_id="dataset_1", view_id="my_fake_view", description="my_fake_view description", view_query_template="SELECT NULL LIMIT 0", should_materialize=False, ) ] mock_table_resource_ds_1_table_1 = { "tableReference": { "projectId": self.project_id, "datasetId": "dataset_1", "tableId": "my_fake_view", }, } self.mock_client.list_tables.return_value = [ bigquery.table.TableListItem(mock_table_resource_ds_1_table_1), ] def get_dataset_ref(dataset_id: str) -> bigquery.dataset.DatasetReference: if dataset_id == dataset.dataset_id: return dataset if dataset_id == dataset_2.dataset_id: return dataset_2 raise ValueError(f"No dataset for id: {dataset_id}") self.mock_client.dataset_ref_for_id.side_effect = get_dataset_ref def mock_dataset_exists(dataset_ref: bigquery.dataset.DatasetReference) -> bool: if dataset_ref == dataset: return True return False self.mock_client.dataset_exists.side_effect = mock_dataset_exists views_to_update = [view_builder.build() for view_builder in mock_view_builders] dag_walker = BigQueryViewDagWalker(views_to_update) managed_views_map = ( get_managed_view_and_materialized_table_addresses_by_dataset(dag_walker) ) with self.assertLogs() as captured_log: cleanup_datasets_and_delete_unmanaged_views( self.mock_client, managed_views_map, dry_run=False ) self.assertEqual(len(captured_log.records), 1) # check that there is only one log message self.assertEqual( captured_log.records[0].getMessage(), "Dataset dataset_unmanaged isn't being managed and no longer exists in BigQuery. It can be safely removed from the list DATASETS_THAT_HAVE_EVER_BEEN_MANAGED.", ) self.mock_client.delete_dataset.assert_not_called() self.mock_client.delete_table.assert_not_called()
def test_cleanup_datasets_and_delete_unmanaged_views_unmanaged_dataset( self, mock_get_datasets_that_have_ever_been_managed: mock.MagicMock ) -> None: mock_get_datasets_that_have_ever_been_managed.return_value = { "dataset_1", "bogus_dataset", } dataset = bigquery.dataset.DatasetReference(self.project_id, "dataset_1") dataset_2 = bigquery.dataset.DatasetReference(self.project_id, "bogus_dataset") sample_views = [ { "view_id": "my_fake_view", "view_query": "SELECT NULL LIMIT 0", } ] mock_view_builders = [ SimpleBigQueryViewBuilder( dataset_id="dataset_1", description=f"{view['view_id']} description", view_query_template="a", should_materialize=False, materialized_address_override=None, should_build_predicate=None, **view, ) for view in sample_views ] mock_table_resource_ds_1_table_1 = { "tableReference": { "projectId": self.project_id, "datasetId": "dataset_1", "tableId": "my_fake_view", }, } self.mock_client.list_tables.return_value = [ bigquery.table.TableListItem(mock_table_resource_ds_1_table_1), ] def get_dataset_ref(dataset_id: str) -> bigquery.dataset.DatasetReference: if dataset_id == dataset.dataset_id: return dataset if dataset_id == dataset_2.dataset_id: return dataset_2 raise ValueError(f"No dataset for id: {dataset_id}") self.mock_client.dataset_ref_for_id.side_effect = get_dataset_ref self.mock_client.dataset_exists.return_value = True views_to_update = [view_builder.build() for view_builder in mock_view_builders] dag_walker = BigQueryViewDagWalker(views_to_update) managed_views_map = ( get_managed_view_and_materialized_table_addresses_by_dataset(dag_walker) ) cleanup_datasets_and_delete_unmanaged_views( self.mock_client, managed_views_map, dry_run=False ) self.mock_client.delete_dataset.assert_called_with( self.mock_client.dataset_ref_for_id("bogus_dataset"), delete_contents=True ) self.mock_client.list_tables.assert_called() self.mock_client.delete_table.assert_not_called()
def rematerialize_views( views_to_update: List[BigQueryView], all_view_builders: Sequence[BigQueryViewBuilder], view_source_table_datasets: Set[str], dataset_overrides: Optional[Dict[str, str]] = None, skip_missing_views: bool = False, bq_region_override: Optional[str] = None, ) -> None: """For all views in the provided |views_to_update| list, re-materializes any materialized views. This should be called only when we want to refresh the data in the materialized view(s), not when we want to update the underlying query of the view(s). Args: views_to_update: List of views to re-materialize all_view_builders: Superset of the views_to_update that contains all views that either depend on or are dependents of the list of input views. view_source_table_datasets: Set of datasets containing tables that can be treated as root nodes in the view dependency graph. dataset_overrides: A dictionary mapping dataset_ids to the dataset name they should be replaced with for the given list of views_to_update. skip_missing_views: If True, ignores any input views that do not exist. If False, crashes if tries to materialize a view that does not exist. bq_region_override: If set, overrides the region (e.g. us-east1) associated with all BigQuery operations. """ set_default_table_expiration_for_new_datasets = bool(dataset_overrides) if set_default_table_expiration_for_new_datasets: logging.info( "Found non-empty dataset overrides. New datasets created in this process will have a " "default table expiration of 24 hours.") try: bq_client = BigQueryClientImpl(region_override=bq_region_override) all_views_dag_walker = BigQueryViewDagWalker( _build_views_to_update( view_source_table_datasets=view_source_table_datasets, candidate_view_builders=all_view_builders, dataset_overrides=dataset_overrides, )) dataset_map = get_managed_view_and_materialized_table_addresses_by_dataset( all_views_dag_walker) _create_all_datasets_if_necessary( bq_client, list(dataset_map.keys()), set_default_table_expiration_for_new_datasets, ) # Limit DAG to only ancestor views and the set of views to update ancestors_dag_walker = all_views_dag_walker.get_ancestors_sub_dag( views_to_update) def _materialize_view( v: BigQueryView, _parent_results: Dict[BigQueryView, None]) -> None: if not v.materialized_address: logging.info("Skipping non-materialized view [%s.%s].", v.dataset_id, v.view_id) return if skip_missing_views and not bq_client.table_exists( bq_client.dataset_ref_for_id(dataset_id=v.dataset_id), v.view_id): logging.info( "Skipping materialization of view [%s.%s] which does not exist", v.dataset_id, v.view_id, ) return bq_client.materialize_view_to_table(v) ancestors_dag_walker.process_dag(_materialize_view) except Exception as e: with monitoring.measurements() as measurements: measurements.measure_int_put(m_failed_view_update, 1) raise e from e