272 lines
9.0 KiB
Python
272 lines
9.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Script to seed the Firestore database with initial data.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import asyncio
|
|
import logging
|
|
import argparse
|
|
from datetime import datetime, timedelta
|
|
import secrets
|
|
import hashlib
|
|
from bson import ObjectId
|
|
from pydantic import HttpUrl
|
|
|
|
# Add the parent directory to the path so we can import from src
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
|
|
|
from src.db.models.team import TeamModel
|
|
from src.db.models.user import UserModel
|
|
from src.db.models.api_key import ApiKeyModel
|
|
from src.db.models.image import ImageModel
|
|
from src.db.providers.firestore_provider import firestore_db
|
|
from src.db.repositories.firestore_team_repository import firestore_team_repository
|
|
from src.db.repositories.firestore_user_repository import firestore_user_repository
|
|
from src.db.repositories.firestore_api_key_repository import firestore_api_key_repository
|
|
from src.db.repositories.firestore_image_repository import firestore_image_repository
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def generate_api_key(length=32):
|
|
"""Generate a random API key"""
|
|
return secrets.token_hex(length)
|
|
|
|
def hash_api_key(api_key):
|
|
"""Hash an API key for storage"""
|
|
return hashlib.sha256(api_key.encode()).hexdigest()
|
|
|
|
async def seed_teams():
|
|
"""Seed the database with team data"""
|
|
logger.info("Seeding teams...")
|
|
|
|
teams_data = [
|
|
{
|
|
"name": "Sereact Development",
|
|
"description": "Internal development team"
|
|
},
|
|
{
|
|
"name": "Marketing Team",
|
|
"description": "Marketing and design team"
|
|
},
|
|
{
|
|
"name": "Customer Support",
|
|
"description": "Customer support and success team"
|
|
}
|
|
]
|
|
|
|
team_ids = []
|
|
for team_data in teams_data:
|
|
team = TeamModel(**team_data)
|
|
created_team = await firestore_team_repository.create(team)
|
|
team_ids.append(created_team.id)
|
|
logger.info(f"Created team: {created_team.name} (ID: {created_team.id})")
|
|
|
|
return team_ids
|
|
|
|
async def seed_users(team_ids):
|
|
"""Seed the database with user data"""
|
|
logger.info("Seeding users...")
|
|
|
|
users_data = [
|
|
{
|
|
"email": "admin@sereact.com",
|
|
"name": "Admin User",
|
|
"team_id": team_ids[0],
|
|
"is_admin": True
|
|
},
|
|
{
|
|
"email": "developer@sereact.com",
|
|
"name": "Developer User",
|
|
"team_id": team_ids[0]
|
|
},
|
|
{
|
|
"email": "marketing@sereact.com",
|
|
"name": "Marketing User",
|
|
"team_id": team_ids[1]
|
|
},
|
|
{
|
|
"email": "support@sereact.com",
|
|
"name": "Support User",
|
|
"team_id": team_ids[2]
|
|
}
|
|
]
|
|
|
|
user_ids = []
|
|
for user_data in users_data:
|
|
user = UserModel(**user_data)
|
|
created_user = await firestore_user_repository.create(user)
|
|
user_ids.append(created_user.id)
|
|
logger.info(f"Created user: {created_user.name} (ID: {created_user.id})")
|
|
|
|
return user_ids
|
|
|
|
async def seed_api_keys(user_ids, team_ids):
|
|
"""Seed the database with API key data"""
|
|
logger.info("Seeding API keys...")
|
|
|
|
api_keys_data = [
|
|
{
|
|
"user_id": user_ids[0],
|
|
"team_id": team_ids[0],
|
|
"name": "Admin Key",
|
|
"description": "API key for admin user"
|
|
},
|
|
{
|
|
"user_id": user_ids[1],
|
|
"team_id": team_ids[0],
|
|
"name": "Development Key",
|
|
"description": "API key for development user"
|
|
},
|
|
{
|
|
"user_id": user_ids[2],
|
|
"team_id": team_ids[1],
|
|
"name": "Marketing Key",
|
|
"description": "API key for marketing user"
|
|
},
|
|
{
|
|
"user_id": user_ids[3],
|
|
"team_id": team_ids[2],
|
|
"name": "Support Key",
|
|
"description": "API key for support user"
|
|
}
|
|
]
|
|
|
|
generated_keys = []
|
|
for api_key_data in api_keys_data:
|
|
# Generate a unique API key
|
|
api_key = generate_api_key()
|
|
key_hash = hash_api_key(api_key)
|
|
|
|
# Create API key object
|
|
api_key_data["key_hash"] = key_hash
|
|
api_key_data["expiry_date"] = datetime.utcnow() + timedelta(days=365)
|
|
|
|
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,
|
|
"key": api_key,
|
|
"name": created_api_key.name
|
|
})
|
|
|
|
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):")
|
|
for key in generated_keys:
|
|
logger.info(f"Name: {key['name']}, Key: {key['key']}")
|
|
|
|
return generated_keys
|
|
|
|
async def seed_images(team_ids, user_ids):
|
|
"""Seed the database with image metadata"""
|
|
logger.info("Seeding images...")
|
|
|
|
images_data = [
|
|
{
|
|
"filename": "image1.jpg",
|
|
"original_filename": "product_photo.jpg",
|
|
"file_size": 1024 * 1024, # 1MB
|
|
"content_type": "image/jpeg",
|
|
"storage_path": "teams/{}/images/image1.jpg".format(team_ids[0]),
|
|
"public_url": "https://storage.googleapis.com/example-bucket/teams/{}/images/image1.jpg".format(team_ids[0]),
|
|
"team_id": team_ids[0],
|
|
"uploader_id": user_ids[0],
|
|
"description": "Product photo for marketing",
|
|
"tags": ["product", "marketing", "high-resolution"],
|
|
"metadata": {
|
|
"width": 1920,
|
|
"height": 1080,
|
|
"color_space": "sRGB"
|
|
}
|
|
},
|
|
{
|
|
"filename": "image2.png",
|
|
"original_filename": "logo.png",
|
|
"file_size": 512 * 1024, # 512KB
|
|
"content_type": "image/png",
|
|
"storage_path": "teams/{}/images/image2.png".format(team_ids[1]),
|
|
"public_url": "https://storage.googleapis.com/example-bucket/teams/{}/images/image2.png".format(team_ids[1]),
|
|
"team_id": team_ids[1],
|
|
"uploader_id": user_ids[2],
|
|
"description": "Company logo",
|
|
"tags": ["logo", "branding"],
|
|
"metadata": {
|
|
"width": 800,
|
|
"height": 600,
|
|
"color_space": "sRGB"
|
|
}
|
|
},
|
|
{
|
|
"filename": "image3.jpg",
|
|
"original_filename": "support_screenshot.jpg",
|
|
"file_size": 256 * 1024, # 256KB
|
|
"content_type": "image/jpeg",
|
|
"storage_path": "teams/{}/images/image3.jpg".format(team_ids[2]),
|
|
"public_url": "https://storage.googleapis.com/example-bucket/teams/{}/images/image3.jpg".format(team_ids[2]),
|
|
"team_id": team_ids[2],
|
|
"uploader_id": user_ids[3],
|
|
"description": "Screenshot for support ticket",
|
|
"tags": ["support", "screenshot", "bug"],
|
|
"metadata": {
|
|
"width": 1280,
|
|
"height": 720,
|
|
"color_space": "sRGB"
|
|
}
|
|
}
|
|
]
|
|
|
|
image_ids = []
|
|
for image_data in images_data:
|
|
image = ImageModel(**image_data)
|
|
created_image = await firestore_image_repository.create(image)
|
|
image_ids.append(created_image.id)
|
|
logger.info(f"Created image: {created_image.filename} (ID: {created_image.id})")
|
|
|
|
return image_ids
|
|
|
|
async def seed_database():
|
|
"""Seed the database with initial data"""
|
|
try:
|
|
# Connect to Firestore
|
|
firestore_db.connect()
|
|
|
|
# Seed teams first
|
|
team_ids = await seed_teams()
|
|
|
|
# Seed users with team IDs
|
|
user_ids = await seed_users(team_ids)
|
|
|
|
# Seed API keys with user and team IDs
|
|
api_keys = await seed_api_keys(user_ids, team_ids)
|
|
|
|
# Seed images with team and user IDs
|
|
image_ids = await seed_images(team_ids, user_ids)
|
|
|
|
logger.info("Database seeding completed successfully!")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error seeding database: {e}")
|
|
raise
|
|
finally:
|
|
# Disconnect from Firestore
|
|
firestore_db.disconnect()
|
|
|
|
def main():
|
|
"""Main entry point"""
|
|
parser = argparse.ArgumentParser(description="Seed the Firestore database with initial data")
|
|
parser.add_argument("--force", action="store_true", help="Force seeding even if data exists")
|
|
args = parser.parse_args()
|
|
|
|
asyncio.run(seed_database())
|
|
|
|
if __name__ == "__main__":
|
|
main() |