diff --git a/backend/app/models/db/static_data.py b/backend/app/models/db/static_data.py index 0cfa727..3c866e8 100644 --- a/backend/app/models/db/static_data.py +++ b/backend/app/models/db/static_data.py @@ -25,7 +25,7 @@ class StaticData(Base): # Constraints and indexes __table_args__ = ( CheckConstraint( - "category IN ('constellation', 'galaxy', 'star', 'nebula', 'cluster')", + "category IN ('constellation', 'galaxy', 'star', 'nebula', 'cluster', 'asteroid_belt', 'kuiper_belt')", name="chk_category", ), UniqueConstraint("category", "name", name="uq_category_name"), diff --git a/backend/scripts/run_sql.py b/backend/scripts/run_sql.py new file mode 100644 index 0000000..98ee8d0 --- /dev/null +++ b/backend/scripts/run_sql.py @@ -0,0 +1,40 @@ +import asyncio +import sys +from sqlalchemy import text +from app.database import get_db, init_db +import logging + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +async def run_sql_file(sql_file_path): + await init_db() + + try: + with open(sql_file_path, 'r') as f: + sql_content = f.read() + + # Split by semicolon to handle multiple statements if needed + # But sqlalchemy text() might handle it. Let's try executing as one block if possible, + # or split manually if it's simple. + statements = [s.strip() for s in sql_content.split(';') if s.strip()] + + async for session in get_db(): + for stmt in statements: + logger.info(f"Executing: {stmt[:50]}...") + await session.execute(text(stmt)) + await session.commit() + logger.info("SQL execution completed successfully.") + + except FileNotFoundError: + logger.error(f"File not found: {sql_file_path}") + except Exception as e: + logger.error(f"Error executing SQL: {e}") + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: python -m scripts.run_sql ") + sys.exit(1) + + sql_file = sys.argv[1] + asyncio.run(run_sql_file(sql_file)) diff --git a/backend/scripts/seed_asteroid_belts.py b/backend/scripts/seed_asteroid_belts.py new file mode 100644 index 0000000..e94039b --- /dev/null +++ b/backend/scripts/seed_asteroid_belts.py @@ -0,0 +1,83 @@ +import asyncio +from sqlalchemy.ext.asyncio import AsyncSession +from app.database import get_db, init_db +from app.models.db import StaticData +from datetime import datetime +import logging + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +async def seed_asteroid_belts(): + await init_db() # Ensure database is initialized + + async for session in get_db(): # Use async for to get the session + logger.info("Seeding asteroid and Kuiper belt static data...") + + belts_data = [ + { + "category": "asteroid_belt", + "name": "Main Asteroid Belt", + "name_zh": "主小行星带", + "data": { + "innerRadiusAU": 2.2, + "outerRadiusAU": 3.2, + "count": 1500, + "color": "#665544", + "size": 0.1, + "opacity": 0.4, + "heightScale": 0.05, + "rotationSpeed": 0.02 + } + }, + { + "category": "kuiper_belt", + "name": "Kuiper Belt", + "name_zh": "柯伊伯带", + "data": { + "innerRadiusAU": 30, + "outerRadiusAU": 50, + "count": 2500, + "color": "#AABBDD", + "size": 0.2, + "opacity": 0.3, + "heightScale": 0.1, + "rotationSpeed": 0.005 + } + } + ] + + for belt_item in belts_data: + # Check if an item with the same category and name already exists + existing_item = await session.execute( + StaticData.__table__.select().where( + StaticData.category == belt_item["category"], + StaticData.name == belt_item["name"] + ) + ) + if existing_item.scalar_one_or_none(): + logger.info(f"Static data for {belt_item['name']} already exists. Updating...") + stmt = StaticData.__table__.update().where( + StaticData.category == belt_item["category"], + StaticData.name == belt_item["name"] + ).values( + name_zh=belt_item["name_zh"], + data=belt_item["data"], + updated_at=datetime.utcnow() + ) + await session.execute(stmt) + else: + logger.info(f"Adding static data for {belt_item['name']}...") + static_data_entry = StaticData( + category=belt_item["category"], + name=belt_item["name"], + name_zh=belt_item["name_zh"], + data=belt_item["data"] + ) + session.add(static_data_entry) + + await session.commit() + logger.info("Asteroid and Kuiper belt static data seeding complete.") + +if __name__ == "__main__": + asyncio.run(seed_asteroid_belts()) diff --git a/backend/scripts/update_category_constraint.sql b/backend/scripts/update_category_constraint.sql new file mode 100644 index 0000000..c571f5f --- /dev/null +++ b/backend/scripts/update_category_constraint.sql @@ -0,0 +1,6 @@ +-- Remove the old constraint +ALTER TABLE static_data DROP CONSTRAINT IF EXISTS chk_category; + +-- Add the updated constraint +ALTER TABLE static_data ADD CONSTRAINT chk_category +CHECK (category IN ('constellation', 'galaxy', 'star', 'nebula', 'cluster', 'asteroid_belt', 'kuiper_belt'));