toy.py

import asyncio
import time

from auth.model import TokenData
from pymongo import ReadPreference

from driver.schema import DriverCache
from driver_schedule.endpoints import get_driver_schedule_ep
from driver_schedule.schema import DriverScheduleRequest
from order.modelv3 import OrderV3
from store.model import StoreTankCache
from trailer.schema import TrailerCache


# async def main():
#     from app import init_db
#     await init_db()
#     StoreTankCache.populate()
#     DriverCache.populate()
#     TrailerCache.populate_if_needed()
#     payload = {
#         "shift": "am",
#         "date": "2024-04-17"
#     }
#     req = DriverScheduleRequest.model_validate(payload)
#     start = time.time()
#     await get_driver_schedule_ep(request=req, token_data=TokenData(username="username"))
#     print("Took", time.time() - start)


async def main():
    for x in range(20):
        from app import init_db

        await init_db()

        order_ids = [o.id async for o in OrderV3.find().sort("-_id").limit(500)]
        print(order_ids[:5])

        col = OrderV3.get_motor_collection()
        secondary_collection = col.with_options(
            read_preference=ReadPreference.SECONDARY
        )
        times = []
        for i in range(500):
            start = time.time()
            req = await secondary_collection.find({"_id": {"$in": order_ids}}).to_list(
                500
            )
            total_time = time.time() - start
            print("Took", total_time)

            times.append(total_time)
            if len(times) > 20:
                average = sum(times) / len(times)
                if average < 1:
                    print("SKIPPING FAST")
                    break


if __name__ == "__main__":
    import db

    db.connect()
    asyncio.run(main())