hdx-python-database 1.3.1__py3-none-any.whl → 1.3.3__py3-none-any.whl

hdx/database/__init__.py

@@ -6,15 +6,14 @@ from typing import Any, Dict, List, Optional, Tuple, Type, Union
 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 sqlalchemy.util import Properties
 
 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
 
@@ -41,6 +40,9 @@ class Database:
     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).
@@ -63,6 +65,7 @@ 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.
@@ -95,9 +98,11 @@ class Database:
             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)
@@ -144,6 +149,8 @@ class Database:
                 )
         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
@@ -155,20 +162,41 @@ class Database:
         if recreate_schema:
             self.recreate_schema(engine, schema_name)
         self.prepare_results = prepare_fn()
-        self.session, self.reflected_classes = self.create_session(
+        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) -> "Database":
-        return self
+    def cleanup(self) -> None:
+        """Cleanup SQLAlchemy.
 
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        Returns:
+            sqlalchemy.Engine: SQLAlchemy engine
+        """
         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) -> None:
+        """Drop all tables.
+
+        Returns:
+            None
+        """
+        self.base.metadata.drop_all(self.engine)
+
     def get_engine(self) -> Engine:
         """Returns SQLAlchemy engine.
 
@@ -207,7 +235,7 @@ class Database:
         db_uri: Optional[str] = None,
         table_base: Type[DeclarativeBase] = NoTZBase,
         reflect: bool = False,
-    ) -> Tuple[Session, Optional[Properties]]:
+    ) -> 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,
@@ -222,22 +250,19 @@ class Database:
             reflect (bool): Whether to reflect existing tables. Defaults to False.
 
         Returns:
-            Tuple[Session, Optional[Properties]]: (SQLAlchemy session, reflected classes if available)
+            Tuple[Session, Any]: (SQLAlchemy session, base)
         """
         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)
-        Session = sessionmaker(bind=engine)
         if reflect:
             Base = automap_base(declarative_base=table_base)
             Base.prepare(autoload_with=engine)
-            reflected_classes = Base.classes
+            table_base = Base
         else:
             table_base.metadata.create_all(engine)
-            reflected_classes = None
-        session = Session()
-        return session, reflected_classes
+        return Session(engine), table_base
 
     @staticmethod
     def recreate_schema(engine: Engine, schema_name: str = "public") -> bool:
@@ -285,6 +310,7 @@ class Database:
 
         Args:
             view_params_list (List[Dict]): List of dictionaries with view parameters
+
         Returns:
             List[TableClause]: SQLAlchemy Views
         """

hdx/database/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.3.1'
-__version_tuple__ = version_tuple = (1, 3, 1)
+__version__ = version = '1.3.3'
+__version_tuple__ = version_tuple = (1, 3, 3)

hdx/database/postgresql.py

@@ -1,9 +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
@@ -19,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
 
@@ -39,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_python_database-1.3.1.dist-info → hdx_python_database-1.3.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: hdx-python-database
-Version: 1.3.1
+Version: 1.3.3
 Summary: HDX Python database utilities
 Project-URL: Homepage, https://github.com/OCHA-DAP/hdx-python-database
 Author-email: Michael Rans <rans@email.com>

hdx_python_database-1.3.3.dist-info/RECORD

@@ -0,0 +1,12 @@
+hdx/database/__init__.py,sha256=EYVc16LZCPFULDd9cHHt0S1MKREzE94DtF_MkBhFVng,11589
+hdx/database/_version.py,sha256=VriGPi1kVXIBM0YGAuhpE803XR-FNq1JvTW1Kz2us08,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.3.dist-info/METADATA,sha256=zpM5pfKKaDgOwzDxEfcjctN1me8xY0cKYi88klSrxAM,2691
+hdx_python_database-1.3.3.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+hdx_python_database-1.3.3.dist-info/licenses/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
+hdx_python_database-1.3.3.dist-info/RECORD,,

{hdx_python_database-1.3.1.dist-info → hdx_python_database-1.3.3.dist-info}/WHEEL

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

hdx_python_database-1.3.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-hdx/database/__init__.py,sha256=ApMtv6Ul4dwZCeS1Wc_Igm9ZnYKkWo_HBi9DBKFHx3Q,10957
-hdx/database/_version.py,sha256=cOVPCvD2h2G_2KB6G3ddreYkIQfAnS6WqgAkF_qgGOQ,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=idtRdJ3zbnpInsjf-_JvWSxtDvU8NQPYcvgHq25ju-k,970
-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.1.dist-info/METADATA,sha256=qxQBQ56bNMZP_DoLVHidyBvpjA0Dm2Qbm0dWtsRfIwk,2691
-hdx_python_database-1.3.1.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
-hdx_python_database-1.3.1.dist-info/licenses/LICENSE,sha256=wc-4GpMn-ODs-U_bTe1YCiPVgvcjzrpYOx2wPuyAeII,1079
-hdx_python_database-1.3.1.dist-info/RECORD,,