import time
import os
from typing import Optional

from sqlalchemy.exc import OperationalError
from sqlalchemy.engine import URL
from sqlmodel import Field, Session, SQLModel, create_engine, select

from loguru import logger


class Hero(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    secret_name: str
    age: Optional[int] = None


def main():
    hero_1 = Hero(name="Deadpond", secret_name="<NAME>")
    hero_2 = Hero(name="Spider-Boy", secret_name="<NAME>")
    hero_3 = Hero(name="Rusty-Man", secret_name="<NAME>", age=48)

    host_name = "postgres" if os.environ.get("IS_INSIDE_DOCKER") else "localhost"
    url = URL.create(drivername="postgresql", username="postgres", password="<PASSWORD>", host=host_name, port=5432)
    engine = create_engine(url)
    for _ in range(5):
        try:
            SQLModel.metadata.create_all(engine)
            break
        except OperationalError:
            logger.error("Is postgres database running?")
            time.sleep(2)

    with Session(engine) as session:
        session.add_all([hero_1, hero_2])
        session.add(hero_3)
        session.commit()

    with Session(engine) as session:
        statement = select(Hero).where(Hero.name == "Spider-Boy")
        hero = session.exec(statement).first()
        logger.info(hero)


if __name__ == '__main__':
    main()

sqlmodel.create_engine loguru.logger.info sqlalchemy.engine.URL.create os.environ.get time.sleep loguru.logger.error sqlmodel.select sqlmodel.Session sqlmodel.SQLModel.metadata.create_all sqlmodel.Field
[((285, 322), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (290, 322), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((677, 788), 'sqlalchemy.engine.URL.create', 'URL.create', ([], {'drivername': '"""postgresql"""', 'username': '"""postgres"""', 'password': '"""<PASSWORD>"""', 'host': 'host_name', 'port': '(5432)'}), "(drivername='postgresql', username='postgres', password=\n    '<PASSWORD>', host=host_name, port=5432)\n", (687, 788), False, 'from sqlalchemy.engine import URL\n'), ((797, 815), 'sqlmodel.create_engine', 'create_engine', (['url'], {}), '(url)\n', (810, 815), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((615, 649), 'os.environ.get', 'os.environ.get', (['"""IS_INSIDE_DOCKER"""'], {}), "('IS_INSIDE_DOCKER')\n", (629, 649), False, 'import os\n'), ((1046, 1061), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (1053, 1061), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1179, 1194), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (1186, 1194), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((1328, 1345), 'loguru.logger.info', 'logger.info', (['hero'], {}), '(hero)\n', (1339, 1345), False, 'from loguru import logger\n'), ((864, 900), 'sqlmodel.SQLModel.metadata.create_all', 'SQLModel.metadata.create_all', (['engine'], {}), '(engine)\n', (892, 900), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n'), ((964, 1009), 'loguru.logger.error', 'logger.error', (['"""Is postgres database running?"""'], {}), "('Is postgres database running?')\n", (976, 1009), False, 'from loguru import logger\n'), ((1022, 1035), 'time.sleep', 'time.sleep', (['(2)'], {}), '(2)\n', (1032, 1035), False, 'import time\n'), ((1227, 1239), 'sqlmodel.select', 'select', (['Hero'], {}), '(Hero)\n', (1233, 1239), False, 'from sqlmodel import Field, Session, SQLModel, create_engine, select\n')]