100 lines
5.2 KiB
Python
100 lines
5.2 KiB
Python
from dataclasses import dataclass, field
|
|
from collections.abc import Sequence
|
|
from typing import Any, Literal, ClassVar, overload
|
|
|
|
from sqlalchemy import create_engine, text, Row, RowMapping
|
|
from sqlalchemy.engine import Engine
|
|
|
|
from database.settings import Settings
|
|
from database.query import SQLQuery
|
|
|
|
@dataclass
|
|
class DatabaseConnection:
|
|
connection: str
|
|
engine: Engine = field(init=False)
|
|
|
|
_settings: ClassVar[Settings] = Settings()
|
|
|
|
def __post_init__(self) -> None:
|
|
self.engine = create_engine(
|
|
url=DatabaseConnection._settings.get_url_database(self.connection),
|
|
pool_recycle=3600,
|
|
pool_pre_ping=True,
|
|
echo=False
|
|
)
|
|
|
|
def execute(self, sql: str, params: dict[str, Any] | None = None) -> None:
|
|
"""Abre uma transação, executa e commita um INSERT/UPDATE/DELETE sem retorno."""
|
|
|
|
with self.engine.begin() as conn:
|
|
conn.execute(text(sql), params)
|
|
|
|
def execute_query(self, query: SQLQuery) -> None:
|
|
"""Mesmo que execute(), recebendo um SQLQuery em vez de sql/params separados."""
|
|
self.execute(query.sql, query.params)
|
|
|
|
@overload
|
|
def select_rows(self, sql: str, params: dict[str, Any] | None = None, *, one: Literal[True]) -> Row[Any] | None: ...
|
|
@overload
|
|
def select_rows(self, sql: str, params: dict[str, Any] | None = None, *, one: Literal[False] = False) -> Sequence[Row[Any]]: ...
|
|
def select_rows(self, sql: str, params: dict[str, Any] | None = None, *, one: bool = False) -> Row[Any] | Sequence[Row[Any]] | None:
|
|
"""Abre uma conexão e executa um SELECT que retorna tuplas (Row). \n
|
|
Parâmetro one indica se deve retornar um Row ou uma Sequence de Rows."""
|
|
with self.engine.connect() as conn:
|
|
result = conn.execute(text(sql), params)
|
|
return result.fetchone() if one else result.fetchall()
|
|
|
|
@overload
|
|
def select_rows_query(self, query: SQLQuery, *, one: Literal[True]) -> Row[Any] | None: ...
|
|
@overload
|
|
def select_rows_query(self, query: SQLQuery, *, one: Literal[False] = False) -> Sequence[Row[Any]]: ...
|
|
def select_rows_query(self, query: SQLQuery, *, one: bool = False) -> Row[Any] | Sequence[Row[Any]] | None:
|
|
"""Mesmo que select_rows(), recebendo um SQLQuery em vez de sql/params separados."""
|
|
return self.select_rows(query.sql, query.params, one=one)
|
|
|
|
@overload
|
|
def select_mappings(self, sql: str, params: dict[str, Any] | None = None, *, one: Literal[True]) -> RowMapping | None: ...
|
|
@overload
|
|
def select_mappings(self, sql: str, params: dict[str, Any] | None = None, *, one: Literal[False] = False) -> Sequence[RowMapping]: ...
|
|
def select_mappings(self, sql: str, params: dict[str, Any] | None = None, *, one: bool = False) -> RowMapping | Sequence[RowMapping] | None:
|
|
"""Abre uma conexão e executa um SELECT que retorna dicionários (RowMapping). \n
|
|
Parâmetro one indica se deve retornar um Row ou uma Sequence de Rows."""
|
|
with self.engine.connect() as conn:
|
|
result = conn.execute(text(sql), params).mappings()
|
|
return result.fetchone() if one else result.fetchall()
|
|
|
|
@overload
|
|
def select_mappings_query(self, query: SQLQuery, *, one: Literal[True]) -> RowMapping | None: ...
|
|
@overload
|
|
def select_mappings_query(self, query: SQLQuery, *, one: Literal[False] = False) -> Sequence[RowMapping]: ...
|
|
def select_mappings_query(self, query: SQLQuery, *, one: bool = False) -> RowMapping | Sequence[RowMapping] | None:
|
|
"""Mesmo que select_mappings(), recebendo um SQLQuery em vez de sql/params separados."""
|
|
return self.select_mappings(query.sql, query.params, one=one)
|
|
|
|
@overload
|
|
def execute_returning(self, sql: str, params: dict[str, Any] | None = None, *, one: Literal[True]) -> Row[Any] | None: ...
|
|
@overload
|
|
def execute_returning(self, sql: str, params: dict[str, Any] | None = None, *, one: Literal[False] = False) -> Sequence[Row[Any]]: ...
|
|
def execute_returning(self, sql: str, params: dict[str, Any] | None = None, *, one: bool = False) -> Row[Any] | Sequence[Row[Any]] | None:
|
|
"""Abre uma transação, executa e commita um INSERT/UPDATE/DELETE com retorno."""
|
|
with self.engine.begin() as conn:
|
|
result = conn.execute(text(sql), params)
|
|
return result.fetchone() if one else result.fetchall()
|
|
|
|
@overload
|
|
def execute_returning_query(self, query: SQLQuery, *, one: Literal[True]) -> Row[Any] | None: ...
|
|
@overload
|
|
def execute_returning_query(self, query: SQLQuery, *, one: Literal[False] = False) -> Sequence[Row[Any]]: ...
|
|
def execute_returning_query(self, query: SQLQuery, *, one: bool = False) -> Row[Any] | Sequence[Row[Any]] | None:
|
|
"""Mesmo que execute_returning(), recebendo um SQLQuery em vez de sql/params separados."""
|
|
return self.execute_returning(query.sql, query.params, one=one)
|
|
|
|
def close(self) -> None:
|
|
"""Fecha o pool de conexões do engine."""
|
|
self.engine.dispose()
|
|
|
|
def __enter__(self) -> "DatabaseConnection":
|
|
return self
|
|
|
|
def __exit__(self, *exc_info) -> None:
|
|
self.close() |