cp

scripts/seed_firestore.py

@@ -11,6 +11,7 @@ import argparse
 from datetime import datetime, timedelta
 import secrets
 import hashlib
+import json
 
 from dotenv import load_dotenv
 
@@ -37,6 +38,13 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Custom JSON encoder for handling ObjectId and datetime types
+class CustomJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if hasattr(obj, '__str__'):
+            return str(obj)
+        return super().default(obj)
+
 def generate_api_key(length=32):
     """Generate a random API key"""
     return secrets.token_hex(length)
@@ -45,6 +53,25 @@ def hash_api_key(api_key):
     """Hash an API key for storage"""
     return hashlib.sha256(api_key.encode()).hexdigest()
 
+async def clear_database():
+    """Clear all collections from the database"""
+    logger.info("Clearing database collections...")
+    
+    collections = ["teams", "users", "api_keys", "images"]
+    
+    for collection_name in collections:
+        try:
+            # Get all documents in collection
+            docs = firestore_db.get_collection(collection_name).stream()
+            for doc in docs:
+                doc.reference.delete()
+            logger.info(f"Cleared collection: {collection_name}")
+        except Exception as e:
+            logger.error(f"Error clearing collection {collection_name}: {e}")
+            raise
+    
+    logger.info("Database cleared successfully!")
+
 async def seed_teams():
     """Seed the database with team data"""
     logger.info("Seeding teams...")
@@ -154,18 +181,47 @@ async def seed_api_keys(user_ids, team_ids):
         api_key_obj = ApiKeyModel(**api_key_data)
         created_api_key = await firestore_api_key_repository.create(api_key_obj)
         
-        generated_keys.append({
-            "id": created_api_key.id,
+        # Convert ObjectId to string for JSON serialization
+        key_data = {
+            "id": str(created_api_key.id),
             "key": api_key,
-            "name": created_api_key.name
-        })
+            "name": created_api_key.name,
+            "user_id": str(created_api_key.user_id),
+            "team_id": str(created_api_key.team_id)
+        }
         
+        generated_keys.append(key_data)
         logger.info(f"Created API key: {created_api_key.name} (ID: {created_api_key.id})")
     
-    # Print the generated keys for reference
-    logger.info("\nGenerated API Keys (save these somewhere secure):")
+    # Save API keys to a file
+    api_keys_file = "api_keys.json"
+    with open(api_keys_file, "w") as f:
+        json.dump(generated_keys, f, indent=2, cls=CustomJSONEncoder)
+    
+    # Save as plain text too for easy access
+    with open("api_keys.txt", "w") as f:
+        f.write("API KEYS\n")
+        f.write("="*80 + "\n\n")
+        for key in generated_keys:
+            f.write(f"Name:        {key['name']}\n")
+            f.write(f"Key:         {key['key']}\n")
+            f.write(f"ID:          {key['id']}\n")
+            f.write("-"*80 + "\n\n")
+    
+    # Print the generated keys prominently
+    print("\n")
+    print("="*80)
+    print("                            GENERATED API KEYS")
+    print("="*80)
+    print("")
     for key in generated_keys:
-        logger.info(f"Name: {key['name']}, Key: {key['key']}")
+        print(f"Name:        {key['name']}")
+        print(f"Key:         {key['key']}")
+        print(f"ID:          {key['id']}")
+        print("-"*80)
+    print("\n")
+    
+    logger.info(f"API keys saved to {api_keys_file} and api_keys.txt")
     
     return generated_keys
 
@@ -236,12 +292,16 @@ async def seed_images(team_ids, user_ids):
     
     return image_ids
 
-async def seed_database():
+async def seed_database(clear=False):
     """Seed the database with initial data"""
     try:
         # Connect to Firestore
         firestore_db.connect()
         
+        # Clear database if requested
+        if clear:
+            await clear_database()
+        
         # Seed teams first
         team_ids = await seed_teams()
         
@@ -266,10 +326,11 @@ async def seed_database():
 def main():
     """Main entry point"""
     parser = argparse.ArgumentParser(description="Seed the Firestore database with initial data")
+    parser.add_argument("--clear", action="store_true", help="Clear all collections before seeding")
     parser.add_argument("--force", action="store_true", help="Force seeding even if data exists")
     args = parser.parse_args()
     
-    asyncio.run(seed_database())
+    asyncio.run(seed_database(clear=args.clear))
 
 if __name__ == "__main__":
     main()