From 5ac2414d416f34f15d73f6665ba1c5846e8ab738 Mon Sep 17 00:00:00 2001
From: Patrick Walukagga <pwalukagga@gmail.com>
Date: Fri, 3 Oct 2025 12:29:59 +0300
Subject: [PATCH] Integrate Postgres and Postgis

---
 scripts/load_to_postgis.py   |  4 +++-
 utils/GeospatialAnalyzer2.py | 19 +++++++++++++++++++
 utils/factory.py             |  4 ++--
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/scripts/load_to_postgis.py b/scripts/load_to_postgis.py
index 99b3195..01641b2 100644
--- a/scripts/load_to_postgis.py
+++ b/scripts/load_to_postgis.py
@@ -30,6 +30,7 @@
 import pandas as pd
 import geopandas as gpd
 from sqlalchemy import create_engine, text
+from utils.GeospatialAnalyzer2 import get_database_url
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("load_to_postgis")
@@ -166,9 +167,10 @@ def main(data_dir: str, db_uri: str, ogr2ogr_path: str = "ogr2ogr"):
 
 
 if __name__ == "__main__":
+    DB_URI = get_database_url()
     p = argparse.ArgumentParser()
     p.add_argument("--data-dir", default="data", help="Path to data directory to scan and load")
-    p.add_argument("--db-uri", default=DEFAULT_DB_URI, help="SQLAlchemy DB URI for PostGIS")
+    p.add_argument("--db-uri", default=DB_URI, help="SQLAlchemy DB URI for PostGIS")
     p.add_argument("--ogr2ogr", default="ogr2ogr", help="Path to ogr2ogr binary")
     args = p.parse_args()
 
diff --git a/utils/GeospatialAnalyzer2.py b/utils/GeospatialAnalyzer2.py
index 577e7a0..2f2032a 100644
--- a/utils/GeospatialAnalyzer2.py
+++ b/utils/GeospatialAnalyzer2.py
@@ -31,15 +31,34 @@
 from shapely import wkt, wkb
 from shapely.geometry import base, Point, Polygon
 from sqlalchemy import create_engine, text
+from dotenv import load_dotenv, find_dotenv
+
 
 from configs.stats import DEFAULT_TILE_STAT_COLUMNS
 
 # Configure logging for this module
 logger = logging.getLogger("GeospatialAnalyzer2")
 logger.setLevel(logging.INFO)
+load_dotenv(find_dotenv())
 
 # Default DB URI when running the recommended local PostGIS docker
 DEFAULT_DB_URI = "postgresql+psycopg://pguser:pgpass@localhost:5432/suntrace"
+DATABASE_URL = os.getenv("SUNTRACE_DATABASE_URI", DEFAULT_DB_URI)
+
+
+def get_database_url() -> str:
+    """
+    Get and potentially modify the database URL.
+
+    This function checks if the DATABASE_URL starts with 'postgres://' and
+    replaces it with 'postgresql+psycopg://' if so.
+
+    Returns:
+        str: The potentially modified database URL.
+    """
+    if DATABASE_URL and DATABASE_URL.startswith("postgres://"):
+        return DATABASE_URL.replace("postgres://", "postgresql+psycopg://", 1)
+    return DATABASE_URL
 
 
 class GeospatialAnalyzer2:
diff --git a/utils/factory.py b/utils/factory.py
index 8962f82..2f68afc 100644
--- a/utils/factory.py
+++ b/utils/factory.py
@@ -9,7 +9,7 @@
 
 from configs import paths
 from utils.GeospatialAnalyzer import GeospatialAnalyzer
-from utils.GeospatialAnalyzer2 import GeospatialAnalyzer2, DEFAULT_DB_URI
+from utils.GeospatialAnalyzer2 import GeospatialAnalyzer2, DEFAULT_DB_URI, get_database_url
 
 logger = logging.getLogger(__name__)
 
@@ -56,7 +56,7 @@ def create_geospatial_analyzer(
             use_postgis = bool(os.getenv("SUNTRACE_DATABASE_URI"))
 
     if database_uri is None:
-        database_uri = os.getenv("SUNTRACE_DATABASE_URI")
+        database_uri = get_database_url()
 
     if use_postgis:
         uri = database_uri or DEFAULT_DB_URI