diff --git a/scripts/seed_firestore.py b/scripts/seed_firestore.py index db260f1..92c585d 100644 --- a/scripts/seed_firestore.py +++ b/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() \ No newline at end of file