Integrate Postgres and Postgis

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()
 

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:

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