From 458d145546f5f2aaf6e1d005298b5d71b0e4ad42 Mon Sep 17 00:00:00 2001
From: Kian-Tat Lim <ktl@slac.stanford.edu>
Date: Fri, 26 Apr 2024 21:51:11 -0700
Subject: [PATCH] Refactor Postgres setup.

---
 python/lsst/consdb/hinfo.py    | 16 ++--------------
 python/lsst/consdb/pqserver.py |  3 ++-
 python/lsst/consdb/utils.py    | 20 ++++++++++++++++++++
 3 files changed, 24 insertions(+), 15 deletions(-)
 create mode 100644 python/lsst/consdb/utils.py

diff --git a/python/lsst/consdb/hinfo.py b/python/lsst/consdb/hinfo.py
index ae63f2f6..fea5fe55 100644
--- a/python/lsst/consdb/hinfo.py
+++ b/python/lsst/consdb/hinfo.py
@@ -18,6 +18,7 @@
 from lsst.resources import ResourcePath
 from sqlalchemy import MetaData, Table, create_engine
 from sqlalchemy.dialects.postgresql import insert
+from utils import setup_postgres
 
 ###############################
 # Header Processing Functions #
@@ -287,20 +288,7 @@ def get_kafka_config() -> KafkaConfig:
         instrument_mapping = LSSTCAM_MAPPING
 logging.info(f"Instrument = {instrument}")
 
-host = os.environ.get("DB_HOST")
-passwd = os.environ.get("DB_PASS")
-user = os.environ.get("DB_USER")
-dbname = os.environ.get("DB_NAME")
-pg_url = ""
-if host and passwd and user and dbname:
-    logging.info(f"Connecting to {host} as {user} to {dbname}")
-    pg_url = f"postgresql://{user}:{passwd}@{host}/{dbname}"
-else:
-    pg_url = os.environ.get(
-        "POSTGRES_URL", "postgresql://usdf-butler.slac.stanford.edu:5432/lsstdb1"
-    )
-    logging.info(f"Using POSTGRES_URL {user} {host} {dbname}")
-engine = create_engine(pg_url)
+engine = setup_postgres()
 metadata_obj = MetaData(schema=f"cdb_{instrument.lower()}")
 exposure_table = Table("exposure", metadata_obj, autoload_with=engine)
 
diff --git a/python/lsst/consdb/pqserver.py b/python/lsst/consdb/pqserver.py
index 5344db69..d10bcdb3 100644
--- a/python/lsst/consdb/pqserver.py
+++ b/python/lsst/consdb/pqserver.py
@@ -2,9 +2,10 @@
 from sqlalchemy import create_engine, MetaData
 import sqlalchemy.exc
 
+from utils import setup_postgres
 
 app = Flask(__name__)
-engine = create_engine("postgresql://usdf-butler.slac.stanford.edu:5432/lsstdb1")
+engine = setup_postgres()
 metadata_obj = MetaData(schema="cdb_latiss")
 metadata_obj.reflect(engine)
 
diff --git a/python/lsst/consdb/utils.py b/python/lsst/consdb/utils.py
new file mode 100644
index 00000000..da966361
--- /dev/null
+++ b/python/lsst/consdb/utils.py
@@ -0,0 +1,20 @@
+import logging
+import os
+
+import sqlalchemy
+
+def setup_postgres() -> sqlalchemy.Engine:
+    host = os.environ.get("DB_HOST")
+    passwd = os.environ.get("DB_PASS")
+    user = os.environ.get("DB_USER")
+    dbname = os.environ.get("DB_NAME")
+    pg_url = ""
+    if host and passwd and user and dbname:
+        logging.info(f"Connecting to {host} as {user} to {dbname}")
+        pg_url = f"postgresql://{user}:{passwd}@{host}/{dbname}"
+    else:
+        pg_url = os.environ.get(
+            "POSTGRES_URL", "postgresql://usdf-butler.slac.stanford.edu:5432/lsstdb1"
+        )
+        logging.info(f"Using POSTGRES_URL {user} {host} {dbname}")
+    return create_engine(pg_url)