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

hdx/database/__init__.py

@@ -1,32 +1,34 @@
 from ._version import version as __version__  # noqa: F401
+from .dburi import get_connection_uri
 
 """Database utilities"""
 import logging
-from typing import Any, Dict, Optional, Union
-from urllib.parse import urlsplit
+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
@@ -34,7 +36,9 @@ class Database:
         port (Union[int, str, None]): Database port
         username (Optional[str]): Username to log into database
         password (Optional[str]): Password to log into database
-        driver (str): Database driver. Defaults to "postgresql".
+        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.
@@ -52,7 +56,9 @@ class Database:
         port: Union[int, str, None] = None,
         username: Optional[str] = None,
         password: Optional[str] = None,
-        driver: str = "postgresql",
+        dialect: str = "postgresql",
+        driver: Optional[str] = None,
+        db_has_tz: bool = False,
         **kwargs: Any,
     ) -> None:
         if port is not None:
@@ -76,12 +82,22 @@ class Database:
             port = self.server.local_bind_port
         else:
             self.server = None
-        if driver == "postgresql":
-            wait_for_postgresql(database, host, port, username, password)
-        db_url = self.get_sqlalchemy_url(
-            database, host, port, username, password, driver=driver
+        db_uri = get_connection_uri(
+            database,
+            host,
+            port,
+            username,
+            password,
+            dialect=dialect,
+            driver=driver,
         )
-        self.session = self.get_session(db_url)
+        if dialect == "postgresql":
+            wait_for_postgresql(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
@@ -92,74 +108,20 @@ class Database:
             self.server.stop()
 
     @staticmethod
-    def get_session(db_url: 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_url (str): SQLAlchemy url
+            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_url, poolclass=NullPool, echo=False)
+        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()
-
-    @staticmethod
-    def get_params_from_sqlalchemy_url(db_url: str) -> Dict[str, Any]:
-        """Gets PostgreSQL database connection parameters from SQLAlchemy url
-
-        Args:
-            db_url (str): SQLAlchemy url
-
-        Returns:
-            Dict[str,Any]: Dictionary of database connection parameters
-        """
-        result = urlsplit(db_url)
-        return {
-            "database": result.path[1:],
-            "host": result.hostname,
-            "port": result.port,
-            "username": result.username,
-            "password": result.password,
-            "driver": result.scheme,
-        }
-
-    @staticmethod
-    def get_sqlalchemy_url(
-        database: Optional[str] = None,
-        host: Optional[str] = None,
-        port: Union[int, str, None] = None,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
-        driver: str = "postgresql",
-    ) -> str:
-        """Gets SQLAlchemy url from database connection parameters
-
-        Args:
-            database (Optional[str]): Database name
-            host (Optional[str]): Host where database is located
-            port (Union[int, str, None]): Database port
-            username (Optional[str]): Username to log into database
-            password (Optional[str]): Password to log into database
-            driver (str): Database driver. Defaults to "postgresql".
-
-        Returns:
-            db_url (str): SQLAlchemy url
-        """
-        strings = [f"{driver}://"]
-        if username:
-            strings.append(username)
-            if password:
-                strings.append(f":{password}@")
-            else:
-                strings.append("@")
-        if host:
-            strings.append(host)
-        if port is not None:
-            strings.append(f":{port}")
-        if database:
-            strings.append(f"/{database}")
-        return "".join(strings)

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.0'
-__version_tuple__ = version_tuple = (1, 2, 0)
+__version__ = version = '1.2.2'
+__version_tuple__ = version_tuple = (1, 2, 2)

hdx/database/dburi.py

@@ -0,0 +1,99 @@
+"""Connection URI utilities"""
+
+from typing import Any, Dict, Optional, Union
+from urllib.parse import urlsplit
+
+
+def get_params_from_connection_uri(
+    db_uri: str,
+    include_driver: bool = True,
+) -> Dict[str, Any]:
+    """Gets PostgreSQL database connection parameters from connection URI
+
+    Args:
+        db_uri (str): Connection URI
+        include_driver (bool): Whether to include driver in params. Defaults to True.
+
+    Returns:
+        Dict[str,Any]: Dictionary of database connection parameters
+    """
+    result = urlsplit(db_uri)
+    dialectdriver = result.scheme.split("+")
+    if len(dialectdriver) == 1:
+        dialect = dialectdriver[0]
+        driver = None
+    else:
+        dialect, driver = dialectdriver
+    result = {
+        "database": result.path[1:],
+        "host": result.hostname,
+        "port": result.port,
+        "username": result.username,
+        "password": result.password,
+        "dialect": dialect,
+    }
+    if include_driver:
+        result["driver"] = driver
+    return result
+
+
+def get_connection_uri(
+    database: Optional[str] = None,
+    host: Optional[str] = None,
+    port: Union[int, str, None] = None,
+    username: Optional[str] = None,
+    password: Optional[str] = None,
+    dialect: str = "postgresql",
+    driver: Optional[str] = None,
+    include_driver: bool = True,
+) -> str:
+    """Gets connection uri from database connection parameters
+
+    Args:
+        database (Optional[str]): Database name
+        host (Optional[str]): Host where database is located
+        port (Union[int, str, None]): Database port
+        username (Optional[str]): Username to log into 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)
+        include_driver (bool): Whether to include driver in uri. Defaults to True.
+
+    Returns:
+        str: Connection URI
+    """
+    strings = [dialect]
+    if include_driver:
+        if dialect == "postgresql" and driver is None:
+            driver = "psycopg"
+        if driver:
+            strings.append(f"+{driver}")
+    strings.append("://")
+    if username:
+        strings.append(username)
+        if password:
+            strings.append(f":{password}@")
+        else:
+            strings.append("@")
+    if host:
+        strings.append(host)
+    if port is not None:
+        strings.append(f":{port}")
+    if database:
+        strings.append(f"/{database}")
+    return "".join(strings)
+
+
+def remove_driver_from_uri(db_uri: str) -> str:
+    """Remove driver from connection URI
+
+    Args:
+        db_uri (str): Connection URI
+
+    Returns:
+        str: Connection URI with driver removed
+    """
+    result = urlsplit(db_uri)
+    dialectdriver = result.scheme
+    dialect = dialectdriver.split("+")[0]
+    return db_uri.replace(dialectdriver, dialect)

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

@@ -1,56 +1,40 @@
 """PostgreSQL specific utilities"""
 import logging
 import time
-from typing import Optional, Union
+
+from .dburi import remove_driver_from_uri
 
 try:
-    import psycopg2
+    import psycopg
 except ImportError:
+    psycopg = None
     # dependency missing, issue a warning
     import warnings
 
     warnings.warn(
-        "psycopg2 not found! Please install hdx-python-database[postgresql] to enable."
+        "psycopg not found! Please install hdx-python-database[postgresql] to enable."
     )
 
 logger = logging.getLogger(__name__)
 
 
-def wait_for_postgresql(
-    database: Optional[str],
-    host: Optional[str],
-    port: Union[int, str, None],
-    username: Optional[str],
-    password: Optional[str],
-) -> None:
+def wait_for_postgresql(db_uri: str) -> None:
     """Waits for PostgreSQL database to be up
 
     Args:
-        database (Optional[str]): Database name
-        host (Optional[str]): Host where database is located
-        port (Union[int, str, None]): Database port
-        username (Optional[str]): Username to log into database
-        password (Optional[str]): Password to log into database
+        db_uri (str): Connection URI
 
     Returns:
         None
     """
     connecting_string = "Checking for PostgreSQL..."
-    if port is not None:
-        port = int(port)
+    db_uri_nd = remove_driver_from_uri(db_uri)
     while True:
         try:
             logger.info(connecting_string)
-            connection = psycopg2.connect(
-                database=database,
-                host=host,
-                port=port,
-                user=username,
-                password=password,
-                connect_timeout=3,
-            )
+            connection = psycopg.connect(db_uri_nd, connect_timeout=3)
             connection.close()
             logger.info("PostgreSQL is running!")
             break
-        except psycopg2.OperationalError:
+        except psycopg.OperationalError:
             time.sleep(1)

{hdx_python_database-1.2.0.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.0
+Version: 1.2.2
 Summary: HDX Python database utilities
 Home-page: https://github.com/OCHA-DAP/hdx-python-database
 Author: Michael Rans
@@ -13,7 +13,6 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.7
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
@@ -26,15 +25,15 @@ Classifier: Operating System :: POSIX :: Linux
 Classifier: Operating System :: Unix
 Classifier: Operating System :: MacOS
 Classifier: Operating System :: Microsoft :: Windows
-Requires-Python: >=3.7
+Requires-Python: >=3.8
 Description-Content-Type: text/markdown; charset=UTF-8
 License-File: LICENSE
 Requires-Dist: sqlalchemy (>=2)
 Requires-Dist: sshtunnel
 Provides-Extra: postgresql
-Requires-Dist: psycopg2-binary ; extra == 'postgresql'
+Requires-Dist: psycopg[binary] ; extra == 'postgresql'
 
-[![Build Status](https://github.com/OCHA-DAP/hdx-python-database/workflows/build/badge.svg)](https://github.com/OCHA-DAP/hdx-python-database/actions?query=workflow%3Abuild)
+[![Build Status](https://github.com/OCHA-DAP/hdx-python-database/actions/workflows/run-python-tests.yml/badge.svg)](https://github.com/OCHA-DAP/hdx-python-database/actions/workflows/run-python-tests.yml)
 [![Coverage Status](https://codecov.io/gh/OCHA-DAP/hdx-python-database/branch/main/graph/badge.svg?token=JpWZc5js4y)](https://codecov.io/gh/OCHA-DAP/hdx-python-database)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 [![Imports: isort](https://img.shields.io/badge/%20imports-isort-%231674b1?style=flat&labelColor=ef8336)](https://pycqa.github.io/isort/)

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.0.dist-info/RECORD

@@ -1,9 +0,0 @@
-hdx/database/__init__.py,sha256=-X6o1GRxotnDS36Av4Ek9-9qwsH38auCgMYYPw0N2DM,5443
-hdx/database/_version.py,sha256=CilGHycSRzmSEnCMQvTRJ5J_LWxMs---UqkVbOmQ2rA,160
-hdx/database/postgresql.py,sha256=mvekzE3OE_LgoUyoJSGXwYh9K4DwzgSNi-6MsaHQ9RE,1511
-hdx_python_database-1.2.0.dist-info/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
-hdx_python_database-1.2.0.dist-info/METADATA,sha256=TS90b_YLxudhzOc3wTPyBj0voEpFA-uZRls-nZcqxHY,2483
-hdx_python_database-1.2.0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-hdx_python_database-1.2.0.dist-info/top_level.txt,sha256=cfo3TyDjnu-B1jb_itPPFQCH00c5VT5GuZ30eNrGmlI,4
-hdx_python_database-1.2.0.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-hdx_python_database-1.2.0.dist-info/RECORD,,