def clean()

in bigquery_etl/cli/stage.py [0:0]

• 24 lines of code
• 14 McCabe index (conditional complexity)

def clean(project_id, dataset_suffix, delete_expired):
    """Reset the stage environment."""
    client = bigquery.Client(project_id)

    dataset_filter = (
        None if delete_expired is True else f"labels.suffix:{dataset_suffix}"
    )
    datasets = client.list_datasets(filter=dataset_filter)

    current_timestamp = int(
        (datetime.utcnow() - datetime(1970, 1, 1)).total_seconds() * 1000
    )

    for dataset in datasets:
        if dataset.labels:
            # remove datasets that either match the suffix or are expired
            for label, value in dataset.labels.items():
                if (
                    dataset_suffix and label == "suffix" and value == dataset_suffix
                ) or (
                    delete_expired
                    and label == "expires_on"
                    and int(value) < current_timestamp
                ):
                    click.echo(f"Deleting dataset {dataset.full_dataset_id}")
                    try:
                        client.delete_dataset(dataset, delete_contents=True)
                    except NotFound as e:
                        # account for other concurrent CI runs
                        click.echo(f"Failed to delete: {e.message}")