app_quickstart3.py

""" PyMotyc application, the "complex" one.

Here we try to reproduce more or less real usage scenario:
    - MotycModel as a base to have id model field, which represents _id document field of ObjectId type,
    - many models with related ones,
    - discriminated unions of models to store in one collection,
    - collections of related models,
    - different strategies of identity management,
    - complex search and update queries.

For basic usage see app_quickstart.py, app_quickstart2.py

"""

import asyncio
import uuid
from datetime import datetime
from typing import Union, List

from bson import ObjectId
from motor.motor_asyncio import AsyncIOMotorClient
from pydantic import BaseModel
from typing_extensions import Literal

from pymotyc import M, MotycModel, Engine, Collection, WithId

engine = Engine()


# ====================================================
# Models

# ----------------------------------------------------
# Employee

class Employee(MotycModel, WithId):
    name: str
    age: int


# ----------------------------------------------------
# Product (discriminated union)

ProductId = str


class ProductBase(MotycModel):
    kind: str  # key for discriminated union
    product_id: ProductId = None  # identity, UUID generated by PyMotyc


class Book(ProductBase):
    kind: Literal['book'] = 'book'
    title: str
    pages: int


class Computer(ProductBase):
    kind: Literal['computer'] = 'computer'
    vendor: str


Product = Union[Book, Computer]


# ----------------------------------------------------
# Supplier (with nested model Address)

class Address(BaseModel):
    index: int
    address: str


class Supplier(MotycModel):
    name: str  # mandatory, should be unique, will be used as identity
    location: Address


# ----------------------------------------------------
# Order (with collections of related models)

class OrderHistoryItem(BaseModel):
    timestamp: datetime
    status: Literal['created', 'processed', 'delivered']


class Order(MotycModel, WithId):
    history: List[OrderHistoryItem] = []
    employee_id: ObjectId
    supplier: str
    products: List[ProductId]


# ====================================================
# Database

@engine.database
class Warehouse:
    employees: Collection[Employee]
    products: Collection[Product] = Collection(identity='product_id', indexes=['kind'])
    suppliers: Collection[Supplier] = Collection(identity='name', indexes=['location.address'])  # todo: MotycFields?
    orders: Collection[Order]


# ====================================================

async def main():
    # ----------------------------------------------------
    # Initialization

    motor = AsyncIOMotorClient("mongodb://127.0.0.1:27017")
    await engine.bind(motor=motor, inject_motyc_fields=True)

    # todo: database drop
    await Warehouse.employees.collection.drop()
    await Warehouse.employees.create_indexes()

    await Warehouse.products.collection.drop()
    await Warehouse.products.create_indexes()

    await Warehouse.suppliers.collection.drop()
    await Warehouse.suppliers.create_indexes()

    await Warehouse.orders.collection.drop()
    await Warehouse.orders.create_indexes()

    # ----------------------------------------------------
    # Filling

    await Warehouse.employees.save(Employee(name='Vasya Pupkin', age=42))
    await Warehouse.employees.save(Employee(name='Frosya Taburetkina', age=22))

    await Warehouse.products.save(Book(title='Hamlet', pages=42))
    await Warehouse.products.save(Book(title='Harry Potter', pages=442))
    await Warehouse.products.save(Computer(vendor='apple'))
    await Warehouse.products.save(Computer(vendor='hp'))

    await Warehouse.suppliers.save(Supplier(name='GUM', location=Address(index=109012, address="Red sq. 3, Moscow, Russia")))
    await Warehouse.suppliers.save(Supplier(name='CUM', location=Address(index=125009, address="Petrovka st. 2, Moscow, Russia")))

    # ----------------------------------------------------
    # Exploration

    # W/o query builder
    employee = await Warehouse.employees.find_one({Employee.name: 'Vasya Pupkin'})
    assert isinstance(employee.id, ObjectId)  # Mongo's ObjectId

    # Discriminated union collection, with query builder
    product = await Warehouse.products.find_one((M(ProductBase.kind) == 'book') & (M(Book.title) == 'Hamlet'))
    assert uuid.UUID(product.product_id)  # UUID generated by PyMotyc
    assert isinstance(product, Book)
    assert product.pages == 42

    # Query builder for nested model
    supplier = await Warehouse.suppliers.find_one(M(Supplier.location.address).regex('Red'))
    assert supplier.name == 'GUM'

    # ----------------------------------------------------
    # Order management

    order = await Warehouse.orders.save(Order(
        employee_id=employee.id,
        supplier=supplier.name,
        products=[product.product_id],
        history=[OrderHistoryItem(timestamp=datetime.now(), status='created')]
    ))

    # Modify and save the model instance
    order.history.append(OrderHistoryItem(timestamp=datetime.now(), status='processed'))
    await order.save()  # can be used like this because of MotycModel

    # Or modify inplace
    await Warehouse.orders.update_one(
        {Order.id: order.id},
        update={'$push': {Order.history: OrderHistoryItem(timestamp=datetime.now(), status='delivered')}}
    )


if __name__ == "__main__":
    asyncio.run(main())