ci: Add bigtable cleanup script

Signed-off-by: Danny C <[email protected]>

.github/workflows/nightly-ci.yml

@@ -27,7 +27,7 @@ jobs:
   cleanup_dynamo_tables:
     if: github.repository == 'feast-dev/feast'
     runs-on: ubuntu-latest
-    name: Cleanup dynamo tables which can fail to cleanup
+    name: Cleanup Bigtable / Dynamo tables which can fail to cleanup
     steps:
       - uses: actions/checkout@v2
         with:
@@ -47,9 +47,20 @@ jobs:
       - name: Install Python dependencies
         run: |
           pip install boto3
+          pip install google-cloud-bigtable
           pip install tqdm
-      - name: Run DynamoDB cleanup script
-        run: python infra/scripts/cleanup_dynamo_ci.py
+      - name: Authenticate to Google Cloud
+        uses: 'google-github-actions/auth@v1'
+        with:
+          credentials_json: '${{ secrets.GCP_SA_KEY }}'
+      - name: Set up gcloud SDK
+        uses: google-github-actions/setup-gcloud@v1
+        with:
+          project_id: ${{ secrets.GCP_PROJECT_ID }}
+      - name: Use gcloud CLI
+        run: gcloud info
+      - name: Run DynamoDB / Bigtable cleanup script
+        run: python infra/scripts/cleanup_ci.py
   build-docker-image:
     if: github.repository == 'feast-dev/feast'
     needs: [check_date]

infra/scripts/cleanup_ci.py

@@ -0,0 +1,50 @@
+from time import sleep
+import boto3
+from tqdm import tqdm
+from google.cloud import bigtable
+from google.cloud.bigtable import enums
+
+
+def cleanup_dynamo_ci():
+    db = boto3.resource("dynamodb")
+
+    num_to_delete = 0
+    all_tables = db.tables.all()
+    for table in all_tables:
+        if "integration_test" in table.name:
+            num_to_delete += 1
+    with tqdm(total=num_to_delete) as progress:
+        for table in all_tables:
+            if "integration_test" in table.name:
+                table.delete()
+                progress.update()
+    print(f"Deleted {num_to_delete} CI DynamoDB tables")
+
+
+def cleanup_bigtable_ci():
+    client = bigtable.Client(project="kf-feast", admin=True)
+    instance = client.instance("feast-integration-tests")
+    if instance.exists():
+        print(f"Deleted Bigtable CI instance")
+        instance.delete()
+
+    location_id = "us-central1-f"
+    serve_nodes = 1
+    storage_type = enums.StorageType.SSD
+    cluster = instance.cluster(
+        "feast-integration-tests-c1",
+        location_id=location_id,
+        serve_nodes=serve_nodes,
+        default_storage_type=storage_type,
+    )
+    instance.create(clusters=[cluster])
+    print(f"Created new Bigtable CI tables")
+
+
+def main() -> None:
+    cleanup_dynamo_ci()
+    cleanup_bigtable_ci()
+
+
+if __name__ == "__main__":
+    main()