284 lines
11 KiB
Python
284 lines
11 KiB
Python
"""
|
|
Populate Primary Stars for Star Systems
|
|
Phase 4.1: Data Migration
|
|
|
|
This script creates primary star records in celestial_bodies table
|
|
for all star systems in the star_systems table.
|
|
|
|
It does NOT fetch new data from NASA - all data already exists in star_systems.
|
|
"""
|
|
|
|
import asyncio
|
|
import sys
|
|
import os
|
|
import json
|
|
from datetime import datetime
|
|
|
|
# Add backend to path
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
|
|
|
from sqlalchemy import text
|
|
from app.database import get_db
|
|
|
|
async def populate_primary_stars():
|
|
"""Create primary star records for all star systems"""
|
|
|
|
print("=" * 70)
|
|
print("🌟 Phase 4.1: Populate Primary Stars")
|
|
print("=" * 70)
|
|
print()
|
|
|
|
async for session in get_db():
|
|
try:
|
|
# Step 1: Check current status
|
|
print("📊 Step 1: Checking current status...")
|
|
|
|
result = await session.execute(text(
|
|
"SELECT COUNT(*) FROM star_systems"
|
|
))
|
|
total_systems = result.scalar()
|
|
print(f" Total star systems: {total_systems}")
|
|
|
|
result = await session.execute(text(
|
|
"SELECT COUNT(*) FROM celestial_bodies WHERE type = 'star'"
|
|
))
|
|
existing_stars = result.scalar()
|
|
print(f" Existing stars in celestial_bodies: {existing_stars}")
|
|
print()
|
|
|
|
# Step 2: Fetch all star systems
|
|
print("📥 Step 2: Fetching star systems data...")
|
|
|
|
result = await session.execute(text("""
|
|
SELECT
|
|
id, name, name_zh, host_star_name,
|
|
spectral_type, radius_solar, mass_solar,
|
|
temperature_k, luminosity_solar, color,
|
|
description
|
|
FROM star_systems
|
|
ORDER BY id
|
|
"""))
|
|
systems = result.fetchall()
|
|
print(f" Fetched {len(systems)} star systems")
|
|
print()
|
|
|
|
# Step 3: Create primary star records
|
|
print("✨ Step 3: Creating primary star records...")
|
|
|
|
created_count = 0
|
|
updated_count = 0
|
|
skipped_count = 0
|
|
|
|
for system in systems:
|
|
star_id = f"star-{system.id}-primary"
|
|
|
|
# Derive star name from system info
|
|
# Remove "System" suffix from name_zh if present
|
|
star_name_zh = system.name_zh
|
|
if star_name_zh:
|
|
star_name_zh = star_name_zh.replace('系统', '').replace('System', '').strip()
|
|
|
|
# Create metadata JSON
|
|
metadata = {
|
|
"star_role": "primary",
|
|
"spectral_type": system.spectral_type,
|
|
"radius_solar": system.radius_solar,
|
|
"mass_solar": system.mass_solar,
|
|
"temperature_k": system.temperature_k,
|
|
"luminosity_solar": system.luminosity_solar,
|
|
"color": system.color
|
|
}
|
|
|
|
# Description
|
|
description = f"光谱类型: {system.spectral_type or 'Unknown'}"
|
|
if system.temperature_k:
|
|
description += f", 表面温度: {int(system.temperature_k)}K"
|
|
|
|
# Convert metadata to JSON string
|
|
metadata_json = json.dumps(metadata)
|
|
|
|
# Check if star already exists
|
|
check_result = await session.execute(
|
|
text("SELECT id FROM celestial_bodies WHERE id = :star_id").bindparams(star_id=star_id)
|
|
)
|
|
existing = check_result.fetchone()
|
|
|
|
if existing:
|
|
# Update existing record
|
|
await session.execute(
|
|
text("""
|
|
UPDATE celestial_bodies
|
|
SET name = :name,
|
|
name_zh = :name_zh,
|
|
type = 'star',
|
|
description = :description,
|
|
extra_data = CAST(:extra_data AS jsonb),
|
|
updated_at = NOW()
|
|
WHERE id = :star_id
|
|
""").bindparams(
|
|
name=system.host_star_name,
|
|
name_zh=star_name_zh,
|
|
description=description,
|
|
extra_data=metadata_json,
|
|
star_id=star_id
|
|
)
|
|
)
|
|
updated_count += 1
|
|
else:
|
|
# Insert new record
|
|
await session.execute(
|
|
text("""
|
|
INSERT INTO celestial_bodies (
|
|
id, system_id, name, name_zh, type,
|
|
description, extra_data, is_active,
|
|
created_at, updated_at
|
|
) VALUES (
|
|
:star_id, :system_id, :name, :name_zh, 'star',
|
|
:description, CAST(:extra_data AS jsonb), TRUE,
|
|
NOW(), NOW()
|
|
)
|
|
""").bindparams(
|
|
star_id=star_id,
|
|
system_id=system.id,
|
|
name=system.host_star_name,
|
|
name_zh=star_name_zh,
|
|
description=description,
|
|
extra_data=metadata_json
|
|
)
|
|
)
|
|
created_count += 1
|
|
|
|
# Progress indicator
|
|
if (created_count + updated_count) % 50 == 0:
|
|
print(f" Progress: {created_count + updated_count}/{len(systems)}")
|
|
|
|
await session.commit()
|
|
|
|
print(f" ✅ Created: {created_count}")
|
|
print(f" 🔄 Updated: {updated_count}")
|
|
print(f" ⏭️ Skipped: {skipped_count}")
|
|
print()
|
|
|
|
# Step 4: Create default positions (0, 0, 0) for all primary stars
|
|
print("📍 Step 4: Creating default positions...")
|
|
|
|
# First, check which stars don't have positions
|
|
result = await session.execute(text("""
|
|
SELECT cb.id
|
|
FROM celestial_bodies cb
|
|
WHERE cb.type = 'star'
|
|
AND cb.id LIKE 'star-%-primary'
|
|
AND NOT EXISTS (
|
|
SELECT 1 FROM positions p WHERE p.body_id = cb.id
|
|
)
|
|
"""))
|
|
stars_without_positions = result.fetchall()
|
|
|
|
print(f" Stars without positions: {len(stars_without_positions)}")
|
|
|
|
position_count = 0
|
|
for star_row in stars_without_positions:
|
|
star_id = star_row.id
|
|
|
|
# Create position at (0, 0, 0) - center of the system
|
|
await session.execute(
|
|
text("""
|
|
INSERT INTO positions (
|
|
body_id, time, x, y, z,
|
|
vx, vy, vz, source, created_at
|
|
) VALUES (
|
|
:body_id, NOW(), 0, 0, 0,
|
|
0, 0, 0, 'calculated', NOW()
|
|
)
|
|
""").bindparams(body_id=star_id)
|
|
)
|
|
position_count += 1
|
|
|
|
await session.commit()
|
|
|
|
print(f" ✅ Created {position_count} position records")
|
|
print()
|
|
|
|
# Step 5: Verification
|
|
print("🔍 Step 5: Verification...")
|
|
|
|
# Count stars by system
|
|
result = await session.execute(text("""
|
|
SELECT
|
|
COUNT(DISTINCT cb.system_id) as systems_with_stars,
|
|
COUNT(*) as total_stars
|
|
FROM celestial_bodies cb
|
|
WHERE cb.type = 'star' AND cb.id LIKE 'star-%-primary'
|
|
"""))
|
|
verification = result.fetchone()
|
|
|
|
print(f" Systems with primary stars: {verification.systems_with_stars}/{total_systems}")
|
|
print(f" Total primary star records: {verification.total_stars}")
|
|
|
|
# Check for systems without stars
|
|
result = await session.execute(text("""
|
|
SELECT ss.id, ss.name
|
|
FROM star_systems ss
|
|
WHERE NOT EXISTS (
|
|
SELECT 1 FROM celestial_bodies cb
|
|
WHERE cb.system_id = ss.id AND cb.type = 'star'
|
|
)
|
|
LIMIT 5
|
|
"""))
|
|
missing_stars = result.fetchall()
|
|
|
|
if missing_stars:
|
|
print(f" ⚠️ Systems without stars: {len(missing_stars)}")
|
|
for sys in missing_stars[:5]:
|
|
print(f" - {sys.name} (ID: {sys.id})")
|
|
else:
|
|
print(f" ✅ All systems have primary stars!")
|
|
|
|
print()
|
|
|
|
# Step 6: Sample data check
|
|
print("📋 Step 6: Sample data check...")
|
|
|
|
result = await session.execute(text("""
|
|
SELECT
|
|
cb.id, cb.name, cb.name_zh, cb.extra_data,
|
|
ss.name as system_name
|
|
FROM celestial_bodies cb
|
|
JOIN star_systems ss ON cb.system_id = ss.id
|
|
WHERE cb.type = 'star' AND cb.id LIKE 'star-%-primary'
|
|
ORDER BY ss.distance_pc
|
|
LIMIT 5
|
|
"""))
|
|
samples = result.fetchall()
|
|
|
|
print(" Nearest star systems:")
|
|
for sample in samples:
|
|
print(f" • {sample.name} ({sample.name_zh})")
|
|
print(f" System: {sample.system_name}")
|
|
print(f" Extra Data: {sample.extra_data}")
|
|
print()
|
|
|
|
print("=" * 70)
|
|
print("✅ Phase 4.1 Completed Successfully!")
|
|
print("=" * 70)
|
|
print()
|
|
print(f"Summary:")
|
|
print(f" • Total star systems: {total_systems}")
|
|
print(f" • Primary stars created: {created_count}")
|
|
print(f" • Primary stars updated: {updated_count}")
|
|
print(f" • Positions created: {position_count}")
|
|
print(f" • Coverage: {verification.systems_with_stars}/{total_systems} systems")
|
|
print()
|
|
|
|
except Exception as e:
|
|
await session.rollback()
|
|
print(f"\n❌ Error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(populate_primary_stars())
|