hdx-python-database 1.2.1__py2.py3-none-any.whl → 1.2.2__py2.py3-none-any.whl

hdx/database/__init__.py

@@ -3,30 +3,32 @@ from .dburi import get_connection_uri
 
 """Database utilities"""
 import logging
-from typing import Any, Optional, Union
+from typing import Any, Optional, Type, Union
 
 from sqlalchemy import create_engine
-from sqlalchemy.ext.declarative import declarative_base, declared_attr
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.ext.declarative import declared_attr
+from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
 from sqlalchemy.pool import NullPool
 from sshtunnel import SSHTunnelForwarder
 
-from hdx.database.postgresql import wait_for_postgresql
+from .no_timezone import Base as NoTZBase
+from .postgresql import wait_for_postgresql
 
 logger = logging.getLogger(__name__)
 
 
-class HDXBase:
+class Base(DeclarativeBase):
     @declared_attr.directive
     def __tablename__(cls):
         return f"{cls.__name__.lower()}s"
 
 
-Base = declarative_base(cls=HDXBase)
-
-
 class Database:
-    """Database helper class to handle ssh tunnels, waiting for PostgreSQL to be up etc.
+    """Database helper class to handle ssh tunnels, waiting for PostgreSQL to
+    be up etc. Can be used in a with statement returning a Session object.
+    db_has_tz which defaults to False indicates whether database datetime
+    columns have timezones. If not, conversion occurs between Python datetimes
+    with timezones to timezoneless database columns.
 
     Args:
         database (Optional[str]): Database name
@@ -36,6 +38,7 @@ class Database:
         password (Optional[str]): Password to log into database
         dialect (str): Database dialect. Defaults to "postgresql".
         driver (Optional[str]): Database driver. Defaults to None (psycopg if postgresql or None)
+        db_has_tz (bool): True if db datetime columns have timezone. Defaults to False.
         **kwargs: See below
         ssh_host (str): SSH host (the server to connect to)
         ssh_port (int): SSH port. Defaults to 22.
@@ -55,6 +58,7 @@ class Database:
         password: Optional[str] = None,
         dialect: str = "postgresql",
         driver: Optional[str] = None,
+        db_has_tz: bool = False,
         **kwargs: Any,
     ) -> None:
         if port is not None:
@@ -89,7 +93,11 @@ class Database:
         )
         if dialect == "postgresql":
             wait_for_postgresql(db_uri)
-        self.session = self.get_session(db_uri)
+        if db_has_tz:
+            table_base = Base
+        else:
+            table_base = NoTZBase
+        self.session = self.get_session(db_uri, table_base=table_base)
 
     def __enter__(self) -> Session:
         return self.session
@@ -100,17 +108,20 @@ class Database:
             self.server.stop()
 
     @staticmethod
-    def get_session(db_uri: str) -> Session:
-        """Gets SQLAlchemy session given url. Your tables must inherit
-        from Base in hdx.utilities.database.
+    def get_session(
+        db_uri: str, table_base: Type[DeclarativeBase] = NoTZBase
+    ) -> Session:
+        """Gets SQLAlchemy session given url. Tables must inherit from Base in
+        hdx.utilities.database unless base is defined.
 
         Args:
             db_uri (str): Connection URI
+            table_base (Type[DeclarativeBase]): Base database table class. Defaults to NoTZBase.
 
         Returns:
             sqlalchemy.orm.Session: SQLAlchemy session
         """
         engine = create_engine(db_uri, poolclass=NullPool, echo=False)
         Session = sessionmaker(bind=engine)
-        Base.metadata.create_all(engine)
+        table_base.metadata.create_all(engine)
         return Session()

hdx/database/_version.py

@@ -1,4 +1,4 @@
 # file generated by setuptools_scm
 # don't change, don't track in version control
-__version__ = version = '1.2.1'
-__version_tuple__ = version_tuple = (1, 2, 1)
+__version__ = version = '1.2.2'
+__version_tuple__ = version_tuple = (1, 2, 2)

hdx/database/no_timezone.py

@@ -0,0 +1,40 @@
+"""Utilities for database datetime columns without timezone"""
+from datetime import datetime, timezone
+from typing import Optional
+
+from sqlalchemy import DateTime, TypeDecorator
+from sqlalchemy.orm import DeclarativeBase, declared_attr
+
+
+class ConversionNoTZ(TypeDecorator):
+    """Convert from/to datetime with timezone from database columns that don't
+    use timezone"""
+
+    impl = DateTime
+    cache_ok = True
+
+    def process_bind_param(self, value: Optional[datetime], _):
+        if value is None:
+            return value
+        if value.tzinfo:
+            value = value.astimezone(timezone.utc)
+
+        return value.replace(tzinfo=None)
+
+    def process_result_value(self, value, _):
+        if value is None:
+            return value
+        if value.tzinfo is None:
+            return value.replace(tzinfo=timezone.utc)
+
+        return value.astimezone(timezone.utc)
+
+
+class Base(DeclarativeBase):
+    type_annotation_map = {
+        datetime: ConversionNoTZ,
+    }
+
+    @declared_attr.directive
+    def __tablename__(cls):
+        return f"{cls.__name__.lower()}s"

hdx/database/postgresql.py

@@ -2,7 +2,7 @@
 import logging
 import time
 
-from hdx.database.dburi import remove_driver_from_uri
+from .dburi import remove_driver_from_uri
 
 try:
     import psycopg

{hdx_python_database-1.2.1.dist-info → hdx_python_database-1.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hdx-python-database
-Version: 1.2.1
+Version: 1.2.2
 Summary: HDX Python database utilities
 Home-page: https://github.com/OCHA-DAP/hdx-python-database
 Author: Michael Rans

hdx_python_database-1.2.2.dist-info/RECORD

@@ -0,0 +1,11 @@
+hdx/database/__init__.py,sha256=HPKyvzWk-Gqd63sE3kcoyEO9n5h_9l3IeuzhzFnMWfQ,4458
+hdx/database/_version.py,sha256=O2qpyKb1ilx0S5H2I5jnx1k8SSt3R0rHyv19gH5I8Ws,160
+hdx/database/dburi.py,sha256=hNksrjaCC7AtD_M5ef0DCjkS0rJCofQqdXVD2BHBIGc,2938
+hdx/database/no_timezone.py,sha256=7Gnidmu317AWAWx4aPY4oIWBio9dqqSsiUTndEi5hMY,1082
+hdx/database/postgresql.py,sha256=9VlcdHhBURZPDaazj1kEq8gA4hUFJY-5hIN9uzibOlM,969
+hdx_python_database-1.2.2.dist-info/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
+hdx_python_database-1.2.2.dist-info/METADATA,sha256=ykW44dWeZ6DVkzUerqYbitQNoJE5WiBOkhBZXo3Rtzg,2464
+hdx_python_database-1.2.2.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+hdx_python_database-1.2.2.dist-info/top_level.txt,sha256=cfo3TyDjnu-B1jb_itPPFQCH00c5VT5GuZ30eNrGmlI,4
+hdx_python_database-1.2.2.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+hdx_python_database-1.2.2.dist-info/RECORD,,

hdx_python_database-1.2.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-hdx/database/__init__.py,sha256=OcQxMp3lTyDyw_nwI_du5rz_MfJNizVe7mY815ZI-24,3753
-hdx/database/_version.py,sha256=osFGs6qmlEth0Xu2MHIHoChfXcj8NUW4hLWDEu-1_K4,160
-hdx/database/dburi.py,sha256=hNksrjaCC7AtD_M5ef0DCjkS0rJCofQqdXVD2BHBIGc,2938
-hdx/database/postgresql.py,sha256=mV31vmWCa8Dnqz_82zVkB8x021_HBNbz1rkFzFoKQh0,981
-hdx_python_database-1.2.1.dist-info/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
-hdx_python_database-1.2.1.dist-info/METADATA,sha256=oRowXwXZgOMa5MDNdFtqRVUoQ_1Ob-tQBuH_tWJ7w3E,2464
-hdx_python_database-1.2.1.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-hdx_python_database-1.2.1.dist-info/top_level.txt,sha256=cfo3TyDjnu-B1jb_itPPFQCH00c5VT5GuZ30eNrGmlI,4
-hdx_python_database-1.2.1.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-hdx_python_database-1.2.1.dist-info/RECORD,,