Skip to content

sqlmodel

apex_utils.sqlmodel #

URLConfig #

Bases: BaseModel

A Pydantic model for the configuration of a SQLAlchemy URL.

Attributes:

Name Type Description
drivername str

The name of the database driver, e.g. "postgresql".
username str

The username to connect to the database.
password str

The password to connect to the database.
host str

The host of the database.
port int

The port of the database.
database str

The name of the database to connect to.

to_url #

to_url() -> URL

Convert the URLConfig model to a SQLAlchemy URL.

Returns:

Name Type Description
URL URL

The SQLAlchemy URL object.
Source code in apex_utils/sqlmodel.py 
27
28
29
30
31
32
33
34
35
def to_url(self) -> URL:
    """Convert the URLConfig model to a SQLAlchemy URL.

    Returns:
        URL: The SQLAlchemy URL object.
    """
    from sqlalchemy import URL  # noqa: PLC0415

    return URL.create(**self.model_dump())

initialize_schemas #

initialize_schemas(engine: Engine | None = None) -> None

Initialize the schemas for the Wright Way database.

This function is used to create the necessary schemas in the database if they do not already exist. It checks each table's schema and creates it if it is not present. It is inspired by sqlalchemy/sqlalchemy #3914.

Parameters:

Name Type Description Default
engine Engine | None

SQLAlchemy engine, if not provided, the default engine will be used. Defaults to None.

 None
Source code in apex_utils/sqlmodel.py 
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def initialize_schemas(engine: Engine | None = None) -> None:
    """Initialize the schemas for the Wright Way database.

    This function is used to create the necessary schemas in the database if they do not
    already exist. It checks each table's schema and creates it if it is not present. It
    is inspired by [sqlalchemy/sqlalchemy #3914](https://github.com/sqlalchemy/sqlalchemy/issues/3914#issuecomment-441936279).

    Args:
        engine (Engine | None, optional): SQLAlchemy engine, if not provided, the
            default engine will be used. Defaults to None.
    """
    from sqlalchemy import inspect  # noqa: PLC0415
    from sqlalchemy.schema import CreateSchema  # noqa: PLC0415
    from sqlmodel import SQLModel  # noqa: PLC0415

    engine = resolve_engine(engine)
    for table in SQLModel.metadata.tables.values():
        if table.schema is None:
            msg = f"Table `{table.name}` has no schema, skipping schema creation"
            warn(message=msg, stacklevel=1)
            continue
        if not inspect(engine).has_schema(table.schema):
            with engine.connect() as conn:
                conn.execute(CreateSchema(table.schema, if_not_exists=True))
                conn.commit()

resolve_engine #

resolve_engine(engine: Engine | None = None, config: URLConfig | dict | None = None) -> Engine

Build a duckdb engine for the Wright Way database.

If the engine is provided, it is returned. If a config is provided, it is used as to build the sqlalchemy.engine.URL object. If neither are provided, the config is attempted to be pulled from the environment variable APEX_DB_CONFIG. This could also be from APEX_DB_CONFIG in Doppler.

Parameters:

Name Type Description Default
engine Engine | None

An existing engine to use. Defaults to None.

 None
config URLConfig | dict | None

The configuration for the engine. Defaults to None.

 None
Source code in apex_utils/sqlmodel.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def resolve_engine(
    engine: Engine | None = None,
    config: URLConfig | dict | None = None,
) -> Engine:
    """Build a duckdb engine for the Wright Way database.

    If the engine is provided, it is returned. If a config is provided, it is used as to
    build the `sqlalchemy.engine.URL` object. If neither are provided, the config is
    attempted to be pulled from the environment variable `APEX_DB_CONFIG`. This could
    also be from `APEX_DB_CONFIG` in Doppler.

    Args:
        engine (Engine | None, optional): An existing engine to use. Defaults to None.
        config (URLConfig | dict | None, optional): The configuration for the engine.
            Defaults to None.
    """
    from sqlmodel import create_engine  # noqa: PLC0415

    from apex_utils.doppler import pull_secret  # noqa: PLC0415

    if engine is not None:
        return engine

    if config is None:
        config = json.loads(pull_secret(name="APEX_DB_CONFIG"))

    if not isinstance(config, (URLConfig | dict)):
        msg = "The config must be a dict or `APEX_DB_CONFIG` must point to a dict"
        raise TypeError(msg)

    if isinstance(config, URLConfig):
        return create_engine(config.to_url())

    return create_engine(URLConfig(**config).to_url())