Merge pull request #202 from UW-Macrostrat/stats-post

Change env var setup

Author: davidsklar99

services/usage-stats/sample.env

@@ -1,5 +1,7 @@
+# For both URLs, the correct driver is added by the container
+
 # Matomo connection
-MARIADB_URL=<url>
+MARIADB_URL=<"mysql://user:password@localhost:3306/database">
 
 # Macrostrat connection
-DATABASE_URL=<url>
+DATABASE_URL=<"postgresql://user:pass@localhost:5432/dbname">

services/usage-stats/src/insert.py

@@ -7,7 +7,17 @@
 
 load_dotenv()
 
-DATABASE_URL = os.getenv("DATABASE_URL")
+raw_url = os.getenv("DATABASE_URL")
+
+# Ensure it uses asyncpg
+if raw_url.startswith("postgresql://"):
+    raw_url = raw_url.replace("postgresql://", "postgresql+asyncpg://", 1)
+elif not raw_url.startswith("postgresql+asyncpg://"):
+    raise ValueError(
+        "Invalid DATABASE_URL: must start with postgresql:// or postgresql+asyncpg://"
+    )
+
+DATABASE_URL = raw_url
 
 engine = create_async_engine(DATABASE_URL, echo=True)
 

services/usage-stats/src/last_id.py

@@ -7,7 +7,17 @@
 
 load_dotenv()
 
-DATABASE_URL = os.getenv("DATABASE_URL")
+raw_url = os.getenv("DATABASE_URL")
+
+# Ensure it uses asyncpg
+if raw_url.startswith("postgresql://"):
+    raw_url = raw_url.replace("postgresql://", "postgresql+asyncpg://", 1)
+elif not raw_url.startswith("postgresql+asyncpg://"):
+    raise ValueError(
+        "Invalid DATABASE_URL: must start with postgresql:// or postgresql+asyncpg://"
+    )
+
+DATABASE_URL = raw_url
 
 engine = create_async_engine(DATABASE_URL, echo=True)
 

services/usage-stats/src/macrostrat.py

@@ -9,10 +9,17 @@
 
 BATCH_SIZE = 1000  # Adjust as needed
 
-# Database URL format: mysql+asyncmy://user:password@host:port/database
-DATABASE_URL = os.getenv(
-    "MARIADB_URL", "mysql+asyncmy://user:password@localhost:3306/database"
-)
+raw_url = os.getenv("MARIADB_URL")
+
+# Ensure the URL uses asyncmy driver
+if raw_url.startswith("mysql://"):
+    raw_url = raw_url.replace("mysql://", "mysql+asyncmy://", 1)
+elif not raw_url.startswith("mysql+asyncmy://"):
+    raise ValueError(
+        "Invalid DATABASE_URL: must start with mysql:// or mysql+asyncmy://"
+    )
+
+DATABASE_URL = raw_url
 
 # Create async engine
 engine = create_async_engine(DATABASE_URL, echo=True)

services/usage-stats/src/rockd.py

@@ -10,10 +10,17 @@
 
 BATCH_SIZE = 1000  # Adjust batch size as needed
 
-# Database URL format: mysql+asyncmy://user:password@host:port/database
-DATABASE_URL = os.getenv(
-    "MARIADB_URL", "mysql+asyncmy://user:password@localhost:3306/database"
-)
+raw_url = os.getenv("MARIADB_URL")
+
+# Ensure the URL uses asyncmy driver
+if raw_url.startswith("mysql://"):
+    raw_url = raw_url.replace("mysql://", "mysql+asyncmy://", 1)
+elif not raw_url.startswith("mysql+asyncmy://"):
+    raise ValueError(
+        "Invalid DATABASE_URL: must start with mysql:// or mysql+asyncmy://"
+    )
+
+DATABASE_URL = raw_url
 
 # Create async SQLAlchemy engine
 engine = create_async_engine(DATABASE_URL, echo=True)