From a98daffa7f0cd10c63459751ed71f94eb9085280 Mon Sep 17 00:00:00 2001
From: Daven Quinn <code@davenquinn.com>
Date: Thu, 17 Oct 2024 15:50:04 -0500
Subject: [PATCH] Update database module to accept an engine as well as a URL

---
 database/CHANGELOG.md                    |  5 +++++
 database/macrostrat/database/__init__.py | 16 +++++++++-------
 database/pyproject.toml                  |  2 +-
 dinosaur/poetry.lock                     |  2 +-
 poetry.lock                              |  2 +-
 5 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/database/CHANGELOG.md b/database/CHANGELOG.md
index e74d744..70722f5 100644
--- a/database/CHANGELOG.md
+++ b/database/CHANGELOG.md
@@ -1,5 +1,10 @@
 # Changelog
 
+## [3.4.0] - 2024-10-17
+
+- Add capability to instantiate the `Database` wrapper class from an engine as
+  well as a URL.
+
 ## [3.3.0] - 2024-03-30
 
 This release focuses on nicer semantics for applying database fixtures.
diff --git a/database/macrostrat/database/__init__.py b/database/macrostrat/database/__init__.py
index 9bb7cf2..877fb3d 100644
--- a/database/macrostrat/database/__init__.py
+++ b/database/macrostrat/database/__init__.py
@@ -1,19 +1,17 @@
 import warnings
 from contextlib import contextmanager
-from enum import Enum
 from pathlib import Path
 from typing import Optional, Union
 
 from psycopg2.errors import InvalidSavepointSpecification
 from psycopg2.sql import Identifier
-from sqlalchemy import URL, MetaData, create_engine, inspect, text
+from sqlalchemy import URL, MetaData, create_engine, inspect, Engine
 from sqlalchemy.exc import IntegrityError, InternalError
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.orm import Session, scoped_session, sessionmaker
 from sqlalchemy.sql.expression import Insert
 
 from macrostrat.utils import get_logger
-
 from .mapper import DatabaseMapper
 from .postgresql import on_conflict, prefix_inserts  # noqa
 from .utils import (  # noqa
@@ -41,14 +39,14 @@ class Database(object):
 
     __inspector__ = None
 
-    def __init__(self, db_conn: Union[str, URL], *, echo_sql=False, **kwargs):
+    def __init__(self, db_conn: Union[str, URL, Engine], *, echo_sql=False, **kwargs):
         """
         Wrapper for interacting with a database using SQLAlchemy.
         Optimized for use with PostgreSQL, but usable with SQLite
         as well.
 
         Args:
-            db_conn (str): Connection string for the database.
+            db_conn (str | URL | Engine): Connection string or engine for the database.
 
         Keyword Args:
             echo_sql (bool): If True, will echo SQL commands to the
@@ -61,8 +59,12 @@ def __init__(self, db_conn: Union[str, URL], *, echo_sql=False, **kwargs):
 
         self.instance_params = kwargs.pop("instance_params", {})
 
-        log.info(f"Setting up database connection '{db_conn}'")
-        self.engine = create_engine(db_conn, echo=echo_sql, **kwargs)
+        if isinstance(db_conn, Engine):
+            log.info(f"Set up database connection with engine {db_conn.url}")
+            self.engine = db_conn
+        else:
+            log.info(f"Setting up database connection with URL '{db_conn}'")
+            self.engine = create_engine(db_conn, echo=echo_sql, **kwargs)
         self.metadata = kwargs.get("metadata", metadata)
 
         # Scoped session for database
diff --git a/database/pyproject.toml b/database/pyproject.toml
index 0123fa2..5eb6d53 100644
--- a/database/pyproject.toml
+++ b/database/pyproject.toml
@@ -3,7 +3,7 @@ authors = ["Daven Quinn <dev@davenquinn.com>"]
 description = "A SQLAlchemy-based database toolkit."
 name = "macrostrat.database"
 packages = [{ include = "macrostrat" }]
-version = "3.3.3"
+version = "3.4.0"
 
 [tool.poetry.dependencies]
 GeoAlchemy2 = "^0.14.0"
diff --git a/dinosaur/poetry.lock b/dinosaur/poetry.lock
index 2368f42..fb46f82 100644
--- a/dinosaur/poetry.lock
+++ b/dinosaur/poetry.lock
@@ -278,7 +278,7 @@ files = [
 
 [[package]]
 name = "macrostrat-database"
-version = "3.3.3"
+version = "3.4.0"
 description = "A SQLAlchemy-based database toolkit."
 optional = false
 python-versions = "^3.8"
diff --git a/poetry.lock b/poetry.lock
index dfd5c88..102901d 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -795,7 +795,7 @@ url = "auth-system"
 
 [[package]]
 name = "macrostrat-database"
-version = "3.3.3"
+version = "3.4.0"
 description = "A SQLAlchemy-based database toolkit."
 optional = false
 python-versions = "^3.8"