84 lines
3.0 KiB
Python
84 lines
3.0 KiB
Python
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())
|