2024-12-26 14:42:24 +02:00

43 lines
1.1 KiB
Python

import os.path
import time
from typing import Sequence, Type
import orjson
from tortoise import Model, Tortoise, run_async
from tortoise.transactions import atomic
from service.config import TORTOISE_ORM
from service.database.models import User
models: Sequence[Type[Model]] = (User, )
@atomic()
async def import_data():
for filename, model in ((model._meta.db_table, model) for model in models):
ts = time.time()
_fn = f"fixtures/{filename}.json"
if not os.path.exists(_fn):
continue
with open(_fn, "r") as f:
data = orjson.loads(f.read())
await model.bulk_create((model(**row) for row in data), ignore_conflicts=True)
te = time.time()
print(f"{model.__name__} imported ({len(data)} rows in {te - ts:.5f}s)")
async def do_import():
await Tortoise.init(TORTOISE_ORM)
await import_data()
def main():
start_ts = time.time()
run_async(do_import())
end_ts = time.time()
print(f"Import completed in {end_ts - start_ts:.3f}s")
if __name__ == "__main__":
main()