mi
31 lines
722 B
Python
31 lines
722 B
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import asyncpg
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
|
|
|
|
|
def asyncpg_dsn(url: str) -> str:
|
|
if url.startswith("postgresql+asyncpg://"):
|
|
return url.replace("postgresql+asyncpg://", "postgresql://", 1)
|
|
return url
|
|
|
|
|
|
def create_postgres_engine(
|
|
url: str,
|
|
*,
|
|
server_settings: dict[str, str] | None = None,
|
|
**engine_options: Any,
|
|
) -> AsyncEngine:
|
|
dsn = asyncpg_dsn(url)
|
|
|
|
async def connect():
|
|
return await asyncpg.connect(dsn=dsn, server_settings=server_settings)
|
|
|
|
return create_async_engine(
|
|
"postgresql+asyncpg://",
|
|
async_creator=connect,
|
|
**engine_options,
|
|
)
|