69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
import platform
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from sqlalchemy.engine.url import make_url
|
||
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
||
|
|
|
||
|
|
_PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
if str(_PROJECT_ROOT) not in sys.path:
|
||
|
|
sys.path.insert(0, str(_PROJECT_ROOT))
|
||
|
|
|
||
|
|
from backend.config.settings import load_settings
|
||
|
|
from backend.database.schema import init_schema, init_timescaledb
|
||
|
|
|
||
|
|
|
||
|
|
def _redact_url(url: str) -> str:
|
||
|
|
try:
|
||
|
|
parsed = make_url(url)
|
||
|
|
if parsed.password:
|
||
|
|
parsed = parsed.set(password="***")
|
||
|
|
return str(parsed)
|
||
|
|
except Exception:
|
||
|
|
return url
|
||
|
|
|
||
|
|
|
||
|
|
async def _run(url: str, enable_timescaledb: bool) -> None:
|
||
|
|
engine = create_async_engine(url, echo=False, pool_pre_ping=True)
|
||
|
|
try:
|
||
|
|
await init_schema(engine)
|
||
|
|
if enable_timescaledb:
|
||
|
|
await init_timescaledb(engine)
|
||
|
|
finally:
|
||
|
|
await engine.dispose()
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--url", default=None)
|
||
|
|
parser.add_argument("--timescaledb", action="store_true")
|
||
|
|
parser.add_argument("--no-timescaledb", action="store_true")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
settings = load_settings()
|
||
|
|
url = (args.url or os.getenv("SMARTEDT_DATABASE_URL") or settings.database.url).strip()
|
||
|
|
enable_timescaledb = settings.database.timescaledb
|
||
|
|
if args.timescaledb:
|
||
|
|
enable_timescaledb = True
|
||
|
|
if args.no_timescaledb:
|
||
|
|
enable_timescaledb = False
|
||
|
|
|
||
|
|
print(f"Connecting to DB: {_redact_url(url)}")
|
||
|
|
print(f"Init schema: yes; init timescaledb: {'yes' if enable_timescaledb else 'no'}")
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
if platform.system() == "Windows" and hasattr(asyncio, "WindowsSelectorEventLoopPolicy"):
|
||
|
|
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
|
||
|
|
|
||
|
|
asyncio.run(_run(url, enable_timescaledb))
|
||
|
|
print("✅ 初始化完成")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|