Change env var setup

Author: davidsklar99

services/usage-stats/src/insert.py

@@ -7,7 +7,15 @@
 
 load_dotenv()
 
-DATABASE_URL = os.getenv("DATABASE_URL")
+raw_url = os.getenv("DATABASE_URL", "postgresql://user:pass@localhost:5432/dbname")
+
+# 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,15 @@
 
 load_dotenv()
 
-DATABASE_URL = os.getenv("DATABASE_URL")
+raw_url = os.getenv("DATABASE_URL", "postgresql://user:pass@localhost:5432/dbname")
+
+# 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,11 +9,18 @@
 
 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", "mysql://user:password@localhost:3306/database"
 )
 
+# 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,11 +10,18 @@
 
 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", "mysql://user:password@localhost:3306/database"
 )
 
+# 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)