cosmo_backend/scripts/populate_resources.py

144 lines
6.4 KiB
Python

"""
Populate resources table with texture and model files
"""
import asyncio
import os
from pathlib import Path
from sqlalchemy.dialects.postgresql import insert as pg_insert
from app.database import get_db
from app.models.db.resource import Resource
# Mapping of texture files to celestial body IDs (use numeric Horizons IDs)
TEXTURE_MAPPING = {
"2k_sun.jpg": {"body_id": "10", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_mercury.jpg": {"body_id": "199", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_venus_surface.jpg": {"body_id": "299", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_venus_atmosphere.jpg": {"body_id": "299", "resource_type": "texture", "mime_type": "image/jpeg", "extra_data": {"layer": "atmosphere"}},
"2k_earth_daymap.jpg": {"body_id": "399", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_earth_nightmap.jpg": {"body_id": "399", "resource_type": "texture", "mime_type": "image/jpeg", "extra_data": {"layer": "night"}},
"2k_moon.jpg": {"body_id": "301", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_mars.jpg": {"body_id": "499", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_jupiter.jpg": {"body_id": "599", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_saturn.jpg": {"body_id": "699", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_saturn_ring_alpha.png": {"body_id": "699", "resource_type": "texture", "mime_type": "image/png", "extra_data": {"layer": "ring"}},
"2k_uranus.jpg": {"body_id": "799", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_neptune.jpg": {"body_id": "899", "resource_type": "texture", "mime_type": "image/jpeg"},
"2k_stars_milky_way.jpg": {"body_id": None, "resource_type": "texture", "mime_type": "image/jpeg", "extra_data": {"usage": "skybox"}},
}
# Mapping of model files to celestial body IDs (use numeric probe IDs)
MODEL_MAPPING = {
"voyager_1.glb": {"body_id": "-31", "resource_type": "model", "mime_type": "model/gltf-binary"},
"voyager_2.glb": {"body_id": "-32", "resource_type": "model", "mime_type": "model/gltf-binary"},
"juno.glb": {"body_id": "-61", "resource_type": "model", "mime_type": "model/gltf-binary"},
"parker_solar_probe.glb": {"body_id": "-96", "resource_type": "model", "mime_type": "model/gltf-binary"},
"cassini.glb": {"body_id": "-82", "resource_type": "model", "mime_type": "model/gltf-binary"},
}
async def populate_resources():
"""Populate resources table with texture and model files"""
# Get upload directory path
upload_dir = Path(__file__).parent.parent / "upload"
texture_dir = upload_dir / "texture"
model_dir = upload_dir / "model"
print(f"📂 Scanning upload directory: {upload_dir}")
print(f"📂 Texture directory: {texture_dir}")
print(f"📂 Model directory: {model_dir}")
async for session in get_db():
try:
# Process textures
print("\n🖼️ Processing textures...")
texture_count = 0
for filename, mapping in TEXTURE_MAPPING.items():
file_path = texture_dir / filename
if not file_path.exists():
print(f"⚠️ Warning: Texture file not found: {filename}")
continue
file_size = file_path.stat().st_size
# Prepare resource data
resource_data = {
"body_id": mapping["body_id"],
"resource_type": mapping["resource_type"],
"file_path": f"texture/{filename}",
"file_size": file_size,
"mime_type": mapping["mime_type"],
"extra_data": mapping.get("extra_data"),
}
# Use upsert to avoid duplicates
stmt = pg_insert(Resource).values(**resource_data)
stmt = stmt.on_conflict_do_update(
index_elements=['body_id', 'resource_type', 'file_path'],
set_={
'file_size': file_size,
'mime_type': mapping["mime_type"],
'extra_data': mapping.get("extra_data"),
}
)
await session.execute(stmt)
texture_count += 1
print(f"{filename} -> {mapping['body_id'] or 'global'} ({file_size} bytes)")
# Process models
print("\n🚀 Processing models...")
model_count = 0
for filename, mapping in MODEL_MAPPING.items():
file_path = model_dir / filename
if not file_path.exists():
print(f"⚠️ Warning: Model file not found: {filename}")
continue
file_size = file_path.stat().st_size
# Prepare resource data
resource_data = {
"body_id": mapping["body_id"],
"resource_type": mapping["resource_type"],
"file_path": f"model/{filename}",
"file_size": file_size,
"mime_type": mapping["mime_type"],
"extra_data": mapping.get("extra_data"),
}
# Use upsert to avoid duplicates
stmt = pg_insert(Resource).values(**resource_data)
stmt = stmt.on_conflict_do_update(
index_elements=['body_id', 'resource_type', 'file_path'],
set_={
'file_size': file_size,
'mime_type': mapping["mime_type"],
'extra_data': mapping.get("extra_data"),
}
)
await session.execute(stmt)
model_count += 1
print(f"{filename} -> {mapping['body_id']} ({file_size} bytes)")
# Commit all changes
await session.commit()
print(f"\n✨ Successfully populated resources table:")
print(f" 📊 Textures: {texture_count}")
print(f" 📊 Models: {model_count}")
print(f" 📊 Total: {texture_count + model_count}")
except Exception as e:
print(f"❌ Error populating resources: {e}")
await session.rollback()
raise
finally:
break
if __name__ == "__main__":
asyncio.run(populate_resources())