42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
from __future__ import annotations
|
|
|
|
from sqlalchemy import Engine, create_engine
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
from cloud.sql_repository import SQLAlchemyCloudRepository
|
|
|
|
|
|
def normalize_database_url(database_url: str) -> str:
|
|
if database_url.startswith("postgresql://"):
|
|
return "postgresql+psycopg://" + database_url.removeprefix("postgresql://")
|
|
return database_url
|
|
|
|
|
|
def create_database_engine(database_url: str) -> Engine:
|
|
normalized_url = normalize_database_url(database_url)
|
|
options: dict[str, object] = {"pool_pre_ping": True}
|
|
if normalized_url.startswith("sqlite:///"):
|
|
options["connect_args"] = {"check_same_thread": False}
|
|
if normalized_url == "sqlite:///:memory:":
|
|
options["poolclass"] = StaticPool
|
|
return create_engine(normalized_url, **options)
|
|
|
|
|
|
class CloudDatabase:
|
|
"""Own the cloud database engine and repository lifecycle."""
|
|
|
|
def __init__(
|
|
self,
|
|
database_url: str,
|
|
*,
|
|
create_schema: bool = True,
|
|
) -> None:
|
|
self.engine = create_database_engine(database_url)
|
|
self.repository = SQLAlchemyCloudRepository(
|
|
self.engine,
|
|
create_schema=create_schema,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
self.repository.close()
|