def dataset_id(bigquery_client: bigquery.Client, project_id: str):
dataset_id = prefixer.create_prefix()
full_dataset_id = f"{project_id}.{dataset_id}"
dataset = bigquery.Dataset(full_dataset_id)
bigquery_client.create_dataset(dataset)
yield dataset_id
bigquery_client.delete_dataset(dataset, delete_contents=True, not_found_ok=True)
def dataset_id(client: bigquery.Client):
project_id = client.project
dataset_id = prefixer.create_prefix()
dataset = bigquery.Dataset(f"{project_id}.{dataset_id}")
dataset = client.create_dataset(dataset)
yield dataset_id
client.delete_dataset(dataset_id, delete_contents=True)
def cleanup_datasets(bigquery_client: bigquery.Client):
yesterday = datetime.datetime.utcnow() - datetime.timedelta(days=1)
for dataset in bigquery_client.list_datasets():
if (dataset.dataset_id.startswith(RESOURCE_PREFIX)
and resource_name_to_date(dataset.dataset_id) < yesterday):
bigquery_client.delete_dataset(dataset,
delete_contents=True,
not_found_ok=True)
def dataset(bq: bigquery.Client, dataset_id: str):
"""Context manager for creating and deleting the BigQuery dataset for a test."""
try:
result = bq.get_dataset(dataset_id)
except NotFound:
result = bq.create_dataset(dataset_id)
try:
yield result.reference
finally:
bq.delete_dataset(dataset_id, delete_contents=True)
def main(prefixes):
client = Client()
pattern = re.compile('|'.join('^{}.*$'.format(prefix)
for prefix in prefixes))
ds_items = list(client.list_datasets())
for dataset in ds_items:
ds_id = dataset.dataset_id
if pattern.match(ds_id):
print("Deleting dataset: {}".format(ds_id))
try:
client.delete_dataset(dataset.reference, delete_contents=True)
except NotFound:
print(" NOT FOUND")
def delete_dataset(client: bigquery.Client, dataset_id: str, delete_contents: bool = True, ):
"""
Deletes dataset from the current project
Args:
client: BQ API client
dataset_id: dataset to be deleted
delete_contents: (default is True)
Returns:
deletes dataset
"""
# TODO: ADD CHECK OF EXISTENCE
dataset_ref = client.dataset(dataset_id)
client.delete_dataset(dataset_ref, delete_contents)
def bigquery_dataset(bigquery_client: bigquery.Client,
bigquery_schema: List[bigquery.SchemaField]):
project_id = bigquery_client.project
dataset_id = prefixer.create_prefix()
dataset = bigquery.Dataset(f"{project_id}.{dataset_id}")
dataset = bigquery_client.create_dataset(dataset)
sample_table_id = f"{project_id}.{dataset_id}.sample"
job1 = load_sample_data(sample_table_id, bigquery_client, bigquery_schema)
job1.result()
one_row_table_id = f"{project_id}.{dataset_id}.sample_one_row"
job2 = load_sample_data(
one_row_table_id,
bigquery_client,
bigquery_schema,
filename="sample_one_row.json",
)
job2.result()
view = bigquery.Table(f"{project_id}.{dataset_id}.sample_view", )
view.view_query = f"SELECT string FROM `{dataset_id}.sample`"
bigquery_client.create_table(view)
yield dataset_id
bigquery_client.delete_dataset(dataset_id, delete_contents=True)
def cleanup_datasets(bq_client: bigquery.Client):
for dataset in bq_client.list_datasets():
if prefixer.should_cleanup(dataset.dataset_id):
bq_client.delete_dataset(dataset,
delete_contents=True,
not_found_ok=True)
def delete_dataset_if_exists(client: bq.Client, dataset_id: str):
dataset_full_id = get_full_dataset_name(client, dataset_id)
client.delete_dataset(
dataset=dataset_full_id, delete_contents=True, not_found_ok=True
)
