hdx-python-database 1.3.0__py3-none-any.whl → 1.3.2__py3-none-any.whl

hdx/database/__init__.py

@@ -1,31 +1,59 @@
 """Database utilities"""
+
 import logging
-from typing import Any, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
-from sqlalchemy import create_engine
+from sqlalchemy import Engine, TableClause, create_engine
+from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.automap import automap_base
-from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
+from sqlalchemy.orm import DeclarativeBase, Session
 from sqlalchemy.pool import NullPool
+from sqlalchemy.sql.ddl import CreateSchema, DropSchema
 
 from ._version import version as __version__  # noqa: F401
 from .dburi import get_connection_uri
 from .no_timezone import Base as NoTZBase
-from .postgresql import wait_for_postgresql
+from .postgresql import restore_from_pgfile, wait_for_postgresql
+from .views import view
 from .with_timezone import Base as TZBase
 
 logger = logging.getLogger(__name__)
 
 
+class DatabaseError(Exception):
+    pass
+
+
+def do_nothing_fn():
+    return None
+
+
 class Database:
     """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 that
+    be up etc. Can be used in a with statement returning a Database object that
     if reflect is True will have a variable reflected_classes containing the
-    reflected classes. 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 (but not
-    when using reflection). If table_base is supplied, db_has_tz is ignored.
+    reflected classes. Either an SQLAlchemy engine, database uri or various
+    database parameters must be supplied.
+
+    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 (but not when using
+    reflection). If table_base is supplied, db_has_tz is ignored.
+
+    The database can be restored from a given pg_restore file supplied in the
+    parameter pg_restore_file.
+
+    There is an option to wipe and create an empty schema in the database by
+    setting recreate_schema to True and setting a schema_name ("public" is the
+    default).
+
+    If a prepare function is supplied in prepare_fn, it will be executed before
+    Base.metadata.create_all and the results of it returned in instance variable
+    prepare_results.
 
     Args:
+        engine (Optional[Engine]): SQLAlchemy engine to use.
+        db_uri (Optional[str]): Connection URI.
         database (Optional[str]): Database name
         host (Optional[str]): Host where database is located
         port (Union[int, str, None]): Database port
@@ -37,6 +65,10 @@ class Database:
         table_base (Optional[Type[DeclarativeBase]]): Override table base. Defaults to None.
         reflect (bool): Whether to reflect existing tables. Defaults to False.
         **kwargs: See below
+        pg_restore_file (str): Restore database from pg_restore file
+        recreate_schema (bool): Whether to recreate schema
+        schema_name (str): Database schema name. Defaults to "public".
+        prepare_fn (Callable[[], None]]): Function to call before Base.metadata.create_all.
         ssh_host (str): SSH host (the server to connect to)
         ssh_port (int): SSH port. Defaults to 22.
         ssh_username (str): SSH username
@@ -48,6 +80,8 @@ class Database:
 
     def __init__(
         self,
+        engine: Optional[Engine] = None,
+        db_uri: Optional[str] = None,
         database: Optional[str] = None,
         host: Optional[str] = None,
         port: Union[int, str, None] = None,
@@ -62,6 +96,16 @@ class Database:
     ) -> None:
         if port is not None:
             port = int(port)
+        schema_name = None
+        if len(kwargs) == 0:
+            pg_restore_file = None
+            recreate_schema = False
+            prepare_fn = do_nothing_fn
+        else:
+            pg_restore_file = kwargs.pop("pg_restore_file", None)
+            recreate_schema = kwargs.pop("recreate_schema", False)
+            schema_name = kwargs.pop("schema", "public")
+            prepare_fn = kwargs.pop("prepare_fn", do_nothing_fn)
         if len(kwargs) != 0:
             try:
                 import sshtunnel
@@ -89,64 +133,179 @@ class Database:
             port = self.server.local_bind_port
         else:
             self.server = None
-        db_uri = get_connection_uri(
-            database,
-            host,
-            port,
-            username,
-            password,
-            dialect=dialect,
-            driver=driver,
-        )
-        if dialect == "postgresql":
+        if not engine and not db_uri:
+            db_uri = get_connection_uri(
+                database,
+                host,
+                port,
+                username,
+                password,
+                dialect=dialect,
+                driver=driver,
+            )
+            if not db_uri:
+                raise DatabaseError(
+                    "No engine, database uri or connection parameters provided!"
+                )
+        if not engine and dialect == "postgresql":
             wait_for_postgresql(db_uri)
+        if pg_restore_file and db_uri:
+            restore_from_pgfile(db_uri, pg_restore_file)
         if not table_base:
             if db_has_tz:
                 table_base = TZBase
             else:
                 table_base = NoTZBase
-        self.session = self.get_session(
-            db_uri, table_base=table_base, reflect=reflect
+        if not engine:
+            engine = create_engine(db_uri, poolclass=NullPool, echo=False)
+        self.engine: Engine = engine
+        if recreate_schema:
+            self.recreate_schema(engine, schema_name)
+        self.prepare_results = prepare_fn()
+        self.session, self.base = self.create_session(
+            engine,
+            table_base=table_base,
+            reflect=reflect,
         )
+        if reflect:
+            self.reflected_classes = self.base.classes
+        else:
+            self.reflected_classes = None
 
-    def __enter__(self) -> Session:
-        return self.session
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+    def cleanup(self) -> None:
+        self.base.metadata.clear()
         self.session.close()
+        self.engine.dispose()
         if self.server is not None:
             self.server.stop()
 
+    def __enter__(self) -> "Database":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.cleanup()
+
+    def drop_all(self):
+        self.base.metadata.drop_all(self.engine)
+
+    def get_engine(self) -> Engine:
+        """Returns SQLAlchemy engine.
+
+        Returns:
+            sqlalchemy.Engine: SQLAlchemy engine
+        """
+        return self.engine
+
+    def get_session(self) -> Session:
+        """Returns SQLAlchemy session.
+
+        Returns:
+            sqlalchemy.orm.Session: SQLAlchemy session
+        """
+        return self.session
+
+    def get_reflected_classes(self) -> Any:
+        """Gets reflected classes.
+
+        Returns:
+            Any: Reflected classes
+        """
+        return self.reflected_classes
+
+    def get_prepare_results(self) -> Any:
+        """Returns results from prepare function.
+
+        Returns:
+            Any: Results from prepare function
+        """
+        return self.prepare_results
+
     @staticmethod
-    def get_session(
-        db_uri: str,
+    def create_session(
+        engine: Optional[Engine] = None,
+        db_uri: Optional[str] = None,
         table_base: Type[DeclarativeBase] = NoTZBase,
         reflect: bool = False,
-    ) -> Session:
-        """Gets SQLAlchemy session given url. Tables must inherit from Base in
+    ) -> Tuple[Session, Any]:
+        """Creates SQLAlchemy session given SQLAlchemy engine or database uri
+        (one of which must be supplied). Tables must inherit from Base in
         hdx.utilities.database unless base is defined. If reflect is True,
         classes will be reflected from an existing database and the reflected
-        classes are returned in a variable reflected_classes in the returned
-        Session object. Note that type annotation maps don't work with
-        reflection.
+        classes will be returned. Note that type annotation maps don't work
+        with reflection.
 
         Args:
-            db_uri (str): Connection URI
+            engine (Optional[Engine]): SQLAlchemy engine to use. Defaults to None (create from db_uri).
+            db_uri (Optional[str]): Connection URI. Defaults to None (use engine).
             table_base (Type[DeclarativeBase]): Base database table class. Defaults to NoTZBase.
             reflect (bool): Whether to reflect existing tables. Defaults to False.
 
         Returns:
-            sqlalchemy.orm.Session: SQLAlchemy session
+            Tuple[Session, Any]: (SQLAlchemy session, base)
         """
-        engine = create_engine(db_uri, poolclass=NullPool, echo=False)
-        Session = sessionmaker(bind=engine)
+        if not engine:
+            if db_uri is None:
+                raise DatabaseError("No engine or database uri provided!")
+            engine = create_engine(db_uri, poolclass=NullPool, echo=False)
         if reflect:
             Base = automap_base(declarative_base=table_base)
             Base.prepare(autoload_with=engine)
-            classes = Base.classes
+            table_base = Base
         else:
             table_base.metadata.create_all(engine)
-            classes = None
-        session = Session()
-        session.reflected_classes = classes
-        return session
+        return Session(engine), table_base
+
+    @staticmethod
+    def recreate_schema(engine: Engine, schema_name: str = "public") -> bool:
+        """Wipe and create empty schema in database using SQLAlchemy.
+
+        Args:
+            engine (Engine): SQLAlchemy engine to use.
+            schema_name (str): Schema name. Defaults to "public".
+
+        Returns:
+            bool: True if all successful, False if not
+        """
+        # Wipe and create an empty schema
+        try:
+            with engine.connect() as connection:
+                connection.execute(
+                    DropSchema(schema_name, cascade=True, if_exists=True)
+                )
+                connection.commit()
+                connection.execute(
+                    CreateSchema(schema_name, if_not_exists=True)
+                )
+                connection.commit()
+                return True
+        except SQLAlchemyError:
+            return False
+
+    @staticmethod
+    def prepare_view(view_params: Dict) -> TableClause:
+        """Prepare SQLAlchemy view from dictionary with keys: name, metadata and
+        selectable. Must be run before Base.metadata.create_all.
+
+        Args:
+            view_params (Dict): Dictionary with keys name, metadata, selectable
+
+        Returns:
+            TableClause: SQLAlchemy View
+        """
+        return view(**view_params)
+
+    @classmethod
+    def prepare_views(cls, view_params_list: List[Dict]) -> List[TableClause]:
+        """Prepare SQLAlchemy views from a list of dictionaries with keys: name,
+        metadata and selectable. Must be run before Base.metadata.create_all.
+
+        Args:
+            view_params_list (List[Dict]): List of dictionaries with view parameters
+
+        Returns:
+            List[TableClause]: SQLAlchemy Views
+        """
+        results = []
+        for view_params in view_params_list:
+            results.append(cls.prepare_view(view_params))
+        return results

hdx/database/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.3.0'
-__version_tuple__ = version_tuple = (1, 3, 0)
+__version__ = version = '1.3.2'
+__version_tuple__ = version_tuple = (1, 3, 2)

hdx/database/dburi.py

@@ -62,6 +62,8 @@ def get_connection_uri(
     Returns:
         str: Connection URI
     """
+    if not database:
+        return ""
     strings = [dialect]
     if include_driver:
         if dialect == "postgresql" and driver is None:

hdx/database/no_timezone.py

@@ -1,4 +1,5 @@
 """Utilities for database datetime columns without timezone"""
+
 from datetime import datetime, timezone
 from typing import Optional
 

hdx/database/postgresql.py

@@ -1,8 +1,11 @@
 """PostgreSQL specific utilities"""
+
 import logging
+import subprocess
 import time
+from os import environ
 
-from .dburi import remove_driver_from_uri
+from .dburi import get_params_from_connection_uri, remove_driver_from_uri
 
 try:
     import psycopg
@@ -18,6 +21,10 @@ except ImportError:
 logger = logging.getLogger(__name__)
 
 
+class PostgresError(Exception):
+    pass
+
+
 def wait_for_postgresql(db_uri: str) -> None:
     """Waits for PostgreSQL database to be up
 
@@ -38,3 +45,51 @@ def wait_for_postgresql(db_uri: str) -> None:
             break
         except psycopg.OperationalError:
             time.sleep(1)
+
+
+def restore_from_pgfile(db_uri: str, pg_restore_file: str) -> str:
+    """Restore database from a pg_restore file created by pg_backup using the
+    pg_restore command.
+
+    Args:
+        db_uri (str): Connection URI
+        pg_restore_file (str): Path to the pg_restore database file
+
+    Returns:
+        str: Output from the pg_restore command
+    """
+    db_params = get_params_from_connection_uri(db_uri)
+    subprocess_params = ["pg_restore", "-c"]
+    for key, value in db_params.items():
+        match key:
+            case "database":
+                subprocess_params.append("-d")
+            case "host":
+                subprocess_params.append("-h")
+            case "port":
+                subprocess_params.append("-p")
+            case "username":
+                subprocess_params.append("-U")
+            case _:
+                continue
+        subprocess_params.append(f"{value}")
+
+    subprocess_params.append(f"{pg_restore_file}")
+    env = environ.copy()
+    password = db_params.get("password")
+    if password:
+        env["PGPASSWORD"] = password
+    process = subprocess.run(
+        subprocess_params, env=env, capture_output=True, encoding="utf-8"
+    )
+    try:
+        process.check_returncode()
+    except subprocess.CalledProcessError as ex:
+        command = " ".join(subprocess_params)
+        raise PostgresError(
+            f"{command} failed. Return code: {process.returncode}"
+        ) from ex
+
+    for line in process.stdout.splitlines():
+        logger.info(line)
+    return process.stdout

hdx/database/views.py

@@ -1,10 +1,10 @@
 """
-Create a DB View.
+Prepare database views ins SQLAlchemy. Must be run before calling
+Base.metadata.create_all.
 
 Copied from:
 https://github.com/sqlalchemy/sqlalchemy/wiki/Views#sqlalchemy-14-20-version
 """
-from typing import Dict, List
 
 import sqlalchemy as sa
 from sqlalchemy import MetaData, Selectable, TableClause
@@ -46,7 +46,7 @@ def view_doesnt_exist(ddl, target, connection, **kw):
 
 
 def view(name: str, metadata: MetaData, selectable: Selectable) -> TableClause:
-    """Create SQLAlchemy view
+    """Prepare SQLAlchemy view. Must be run before Base.metadata.create_all.
 
     Args:
         name (str): View name
@@ -73,31 +73,3 @@ def view(name: str, metadata: MetaData, selectable: Selectable) -> TableClause:
         DropView(name).execute_if(callable_=view_exists),
     )
     return t
-
-
-def build_view(view_params: Dict) -> TableClause:
-    """Create SQLAlchemy view from dictionary with keys: name, metadata and
-    selectable
-
-    Args:
-        view_params (Dict): Dictionary with keys name, metadata, selectable
-
-    Returns:
-        TableClause: SQLAlchemy View
-    """
-    return view(**view_params)
-
-
-def build_views(view_params_list: List[Dict]) -> List[TableClause]:
-    """Create SQLAlchemy views from a list of dictionaries with keys: name,
-    metadata and selectable
-
-    Args:
-        view_params_list (List[Dict]): List of dictionaries with view parameters
-    Returns:
-        List[TableClause]: SQLAlchemy Views
-    """
-    results = []
-    for view_params in view_params_list:
-        results.append(build_view(view_params))
-    return results

{hdx_python_database-1.3.0.dist-info → hdx_python_database-1.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: hdx-python-database
-Version: 1.3.0
+Version: 1.3.2
 Summary: HDX Python database utilities
 Project-URL: Homepage, https://github.com/OCHA-DAP/hdx-python-database
 Author-email: Michael Rans <rans@email.com>
@@ -39,8 +39,7 @@ Description-Content-Type: text/markdown
 
 [![Build Status](https://github.com/OCHA-DAP/hdx-python-database/actions/workflows/run-python-tests.yaml/badge.svg)](https://github.com/OCHA-DAP/hdx-python-database/actions/workflows/run-python-tests.yaml)
 [![Coverage Status](https://coveralls.io/repos/github/OCHA-DAP/hdx-python-database/badge.svg?branch=main&ts=1)](https://coveralls.io/github/OCHA-DAP/hdx-python-database?branch=main)
-[![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/)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 [![Downloads](https://img.shields.io/pypi/dm/hdx-python-database.svg)](https://pypistats.org/packages/hdx-python-database)
 
 The HDX Python Database Library provides utilities for connecting to databases in a standard way including

hdx_python_database-1.3.2.dist-info/RECORD

@@ -0,0 +1,12 @@
+hdx/database/__init__.py,sha256=cOAnY4cOLcjxM556jAIfw9XqTIKk7mSwZCugIf3XpPA,11431
+hdx/database/_version.py,sha256=ik9WUFYTvRcmzFEUSqkI2H_62uw0XFdUT3mTLq8-_RY,411
+hdx/database/dburi.py,sha256=rw8RuLpU85vxYl5-Zk1SdZRAz-tqre_ot5aEnkjhxAQ,2977
+hdx/database/no_timezone.py,sha256=Lw087G3qOf8H70o8Hq_VvOMZ3ANei-lIA_sUlpFU_Uo,1130
+hdx/database/postgresql.py,sha256=BrEuRnuR2dt3HwJ-o_IzGFs8eT3iXKzDqMSNKPyQ18Y,2621
+hdx/database/utils.py,sha256=6lLymWriMFCq8BZgrDt4V5z14GR3PbvG_WXZeBc99Ws,415
+hdx/database/views.py,sha256=R_1Ni0U2CL3IlydxcB1vyZp7QoDZim6g-807jQIaWQU,1950
+hdx/database/with_timezone.py,sha256=GUUylmTwQHp22yedmCVKpVJzvjWgdS9lg9KeqciNJkY,235
+hdx_python_database-1.3.2.dist-info/METADATA,sha256=ZCPvpjSSs0PxElgVksIa2v7wBo2soTX52wfl4d9eau4,2691
+hdx_python_database-1.3.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+hdx_python_database-1.3.2.dist-info/licenses/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
+hdx_python_database-1.3.2.dist-info/RECORD,,

{hdx_python_database-1.3.0.dist-info → hdx_python_database-1.3.2.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.21.1
+Generator: hatchling 1.25.0
 Root-Is-Purelib: true
 Tag: py3-none-any

hdx_python_database-1.3.0.dist-info/RECORD

@@ -1,12 +0,0 @@
-hdx/database/__init__.py,sha256=0CpgEC9GAqCtmZtBsTvwoDsW3T0VHY0aCNAh53piP6g,5835
-hdx/database/_version.py,sha256=HGwtpza1HCPtlyqElUvIyH97K44TO13CYiYVZNezQ1M,411
-hdx/database/dburi.py,sha256=hNksrjaCC7AtD_M5ef0DCjkS0rJCofQqdXVD2BHBIGc,2938
-hdx/database/no_timezone.py,sha256=jawupasbCwNo-lDbX24CtxaSilb6LS-jLx-wZAa58nM,1129
-hdx/database/postgresql.py,sha256=9VlcdHhBURZPDaazj1kEq8gA4hUFJY-5hIN9uzibOlM,969
-hdx/database/utils.py,sha256=6lLymWriMFCq8BZgrDt4V5z14GR3PbvG_WXZeBc99Ws,415
-hdx/database/views.py,sha256=qJ_f6DpBUSzAGuBbgOXwQDZGni7b-DOUlAdAjSXKwhs,2635
-hdx/database/with_timezone.py,sha256=GUUylmTwQHp22yedmCVKpVJzvjWgdS9lg9KeqciNJkY,235
-hdx_python_database-1.3.0.dist-info/METADATA,sha256=8lJk615nDYiwB9vT9JUzfd-mpGVa6rIFA6kP9WVMgxI,2785
-hdx_python_database-1.3.0.dist-info/WHEEL,sha256=TJPnKdtrSue7xZ_AVGkp9YXcvDrobsjBds1du3Nx6dc,87
-hdx_python_database-1.3.0.dist-info/licenses/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
-hdx_python_database-1.3.0.dist-info/RECORD,,