|
| 1 | +from time import sleep |
| 2 | +import boto3 |
| 3 | +from tqdm import tqdm |
| 4 | +from google.cloud import bigtable |
| 5 | +from google.cloud.bigtable import enums |
| 6 | + |
| 7 | + |
| 8 | +def cleanup_dynamo_ci(): |
| 9 | + db = boto3.resource("dynamodb") |
| 10 | + |
| 11 | + num_to_delete = 0 |
| 12 | + all_tables = db.tables.all() |
| 13 | + for table in all_tables: |
| 14 | + if "integration_test" in table.name: |
| 15 | + num_to_delete += 1 |
| 16 | + with tqdm(total=num_to_delete) as progress: |
| 17 | + for table in all_tables: |
| 18 | + if "integration_test" in table.name: |
| 19 | + table.delete() |
| 20 | + progress.update() |
| 21 | + print(f"Deleted {num_to_delete} CI DynamoDB tables") |
| 22 | + |
| 23 | + |
| 24 | +def cleanup_bigtable_ci(): |
| 25 | + client = bigtable.Client(project="kf-feast", admin=True) |
| 26 | + instance = client.instance("feast-integration-tests") |
| 27 | + if instance.exists(): |
| 28 | + print(f"Deleted Bigtable CI instance") |
| 29 | + instance.delete() |
| 30 | + |
| 31 | + location_id = "us-central1-f" |
| 32 | + serve_nodes = 1 |
| 33 | + storage_type = enums.StorageType.SSD |
| 34 | + cluster = instance.cluster( |
| 35 | + "feast-integration-tests-c1", |
| 36 | + location_id=location_id, |
| 37 | + serve_nodes=serve_nodes, |
| 38 | + default_storage_type=storage_type, |
| 39 | + ) |
| 40 | + instance.create(clusters=[cluster]) |
| 41 | + print(f"Created new Bigtable CI tables") |
| 42 | + |
| 43 | + |
| 44 | +def main() -> None: |
| 45 | + cleanup_dynamo_ci() |
| 46 | + cleanup_bigtable_ci() |
| 47 | + |
| 48 | + |
| 49 | +if __name__ == "__main__": |
| 50 | + main() |
0 commit comments