diff options
author | 2022-07-09 21:48:28 +0100 | |
---|---|---|
committer | 2022-07-09 21:48:28 +0100 | |
commit | b0589b9aa254b623f9bd8dcb40a5c67bea0b7d00 (patch) | |
tree | fbfc23f9fdd6d29f8d3475dab00a356bd6e01338 | |
parent | Install botcore and bump deps (diff) |
Move TZDateTime to avoid circular import
-rw-r--r-- | metricity/database.py | 36 | ||||
-rw-r--r-- | metricity/models.py | 3 | ||||
-rw-r--r-- | metricity/utils.py | 35 |
3 files changed, 36 insertions, 38 deletions
diff --git a/metricity/database.py b/metricity/database.py index 1534e2c..a4d953e 100644 --- a/metricity/database.py +++ b/metricity/database.py @@ -1,7 +1,11 @@ -"""Methods for connecting and interacting with the database.""" +"""General utility functions and classes for Metricity.""" + import logging +from datetime import datetime, timezone import gino +from sqlalchemy.engine import Dialect +from sqlalchemy.types import DateTime, TypeDecorator from metricity.config import DatabaseConfig @@ -26,3 +30,33 @@ async def connect() -> None: log.info("Initiating connection to the database") await db.set_bind(build_db_uri()) log.info("Database connection established") + + +class TZDateTime(TypeDecorator): + """ + A db type that supports the use of aware datetimes in user-land. + + Source from SQLAlchemy docs: + https://docs.sqlalchemy.org/en/14/core/custom_types.html#store-timezone-aware-timestamps-as-timezone-naive-utc + + Edited to include docstrings and type hints. + """ + + impl = DateTime + cache_ok = True + + def process_bind_param(self, value: datetime, dialect: Dialect) -> datetime: + """Convert the value to aware before saving to db.""" + if value is not None: + if not value.tzinfo: + raise TypeError("tzinfo is required") + value = value.astimezone(timezone.utc).replace( + tzinfo=None + ) + return value + + def process_result_value(self, value: datetime, dialect: Dialect) -> datetime: + """Convert the value to aware before passing back to user-land.""" + if value is not None: + value = value.replace(tzinfo=timezone.utc) + return value diff --git a/metricity/models.py b/metricity/models.py index 87be19b..4f136de 100644 --- a/metricity/models.py +++ b/metricity/models.py @@ -5,8 +5,7 @@ from typing import Any, Dict, List from sqlalchemy.dialects.postgresql import insert -from metricity.database import db -from metricity.utils import TZDateTime +from metricity.database import TZDateTime, db class Category(db.Model): diff --git a/metricity/utils.py b/metricity/utils.py deleted file mode 100644 index 3f3547c..0000000 --- a/metricity/utils.py +++ /dev/null @@ -1,35 +0,0 @@ -"""General utility functions and classes for Metricity.""" -from datetime import datetime, timezone - -from sqlalchemy.engine import Dialect -from sqlalchemy.types import DateTime, TypeDecorator - - -class TZDateTime(TypeDecorator): - """ - A db type that supports the use of aware datetimes in user-land. - - Source from SQLAlchemy docs: - https://docs.sqlalchemy.org/en/14/core/custom_types.html#store-timezone-aware-timestamps-as-timezone-naive-utc - - Editted to include docstrings and type hints. - """ - - impl = DateTime - cache_ok = True - - def process_bind_param(self, value: datetime, dialect: Dialect) -> datetime: - """Convert the value to aware before saving to db.""" - if value is not None: - if not value.tzinfo: - raise TypeError("tzinfo is required") - value = value.astimezone(timezone.utc).replace( - tzinfo=None - ) - return value - - def process_result_value(self, value: datetime, dialect: Dialect) -> datetime: - """Convert the value to aware before passing back to user-land.""" - if value is not None: - value = value.replace(tzinfo=timezone.utc) - return value |