feat(engine): Prod time

Use postgres if prod
TracecatHQ · Mar 21, 2024 · 197296f · 197296f
1 parent de33678
commit 197296f
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 10 deletions.
diff --git a/aws/stack.py b/aws/stack.py
@@ -130,6 +130,9 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
             "TRACECAT__DB_ENCRYPTION_KEY": ecs.Secret.from_secrets_manager(
                 tracecat_secret, field="db-encryption-key"
             ),
+            "TRACECAT__DB_URI": ecs.Secret.from_secrets_manager(
+                tracecat_secret, field="db-uri"
+            ),
         }
         api_secrets = {
             **shared_secrets,
@@ -182,8 +185,9 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
                 build_args={"API_MODULE": "tracecat.api.app:app"},
             ),
             cpu=256,
-            memory_limit_mib=512,
+            memory_limit_mib=1024,
             environment={
+                "TRACECAT__APP_ENV": TRACECAT__APP_ENV,
                 "API_MODULE": "tracecat.api.app:app",
                 "SUPABASE_JWT_ALGORITHM": "HS256",
             },
@@ -251,8 +255,12 @@ def __init__(self, scope: Construct, id: str, **kwargs) -> None:
                 build_args={"API_MODULE": "tracecat.runner.app:app", "PORT": "8001"},
             ),
             cpu=256,
-            memory_limit_mib=512,
-            environment={"API_MODULE": "tracecat.runner.app:app", "PORT": "8001"},
+            memory_limit_mib=1024,
+            environment={
+                "TRACECAT__APP_ENV": TRACECAT__APP_ENV,
+                "API_MODULE": "tracecat.runner.app:app",
+                "PORT": "8001",
+            },
             secrets=runner_secrets,
             port_mappings=[ecs.PortMapping(container_port=8001)],
             logging=ecs.LogDrivers.aws_logs(

diff --git a/pyproject.toml b/pyproject.toml
@@ -22,23 +22,24 @@ classifiers = [
     "Topic :: System :: Systems Administration",
 ]
 dependencies = [
+    "adbc-driver-sqlite",
     "colorlog",
+    "cryptography",
     "fastapi",
     "lancedb==0.6.3",
     "openai",
     "orjson",
     "polars",
     "psycopg[binary]",
+    "psycopg2-binary",
+    "pyarrow",
     "pydantic==2.6.1",
     "python-jose[cryptography]",
     "python-slugify",
     "sqlmodel",
     "tantivy",
     "tenacity",
     "uvicorn",
-    "cryptography",
-    "adbc-driver-sqlite",
-    "pyarrow",
 ]
 dynamic = ["version"]
 

diff --git a/tracecat/config.py b/tracecat/config.py
@@ -4,6 +4,7 @@
 HTTP_MAX_RETRIES = 6
 LLM_MAX_RETRIES = 3
 
+TRACECAT__APP_ENV = os.environ.get("TRACECAT__APP_ENV", "dev")
 TRACECAT__OAUTH2_GMAIL_PATH = (
     Path("~/tracecat-runner-client-secret.json").expanduser().resolve()
 )

diff --git a/tracecat/db.py b/tracecat/db.py
@@ -13,6 +13,7 @@
 
 from tracecat import auth
 from tracecat.auth import decrypt_key, encrypt_key
+from tracecat.config import TRACECAT__APP_ENV
 from tracecat.labels.mitre import get_mitre_tactics_techniques
 
 STORAGE_PATH = Path(
@@ -27,7 +28,10 @@
     "Sinkholed",
 ]
 STORAGE_PATH.mkdir(parents=True, exist_ok=True)
-TRACECAT__SQLITE_URI = f"sqlite:////{STORAGE_PATH}/database.db"
+if TRACECAT__APP_ENV == "prod":
+    TRACECAT__DB_URI = os.environ["TRACECAT__DB_URI"]
+else:
+    TRACECAT__DB_URI = f"sqlite:////{STORAGE_PATH}/database.db"
 
 
 class User(SQLModel, table=True):
@@ -196,9 +200,15 @@ def secret(self) -> str:
 
 
 def create_db_engine() -> Engine:
-    engine = create_engine(
-        TRACECAT__SQLITE_URI, echo=True, connect_args={"check_same_thread": False}
-    )
+    if TRACECAT__APP_ENV == "prod":
+        engine_kwargs = {
+            "pool_timeout": 30,
+            "pool_recycle": 3600,
+            "connect_args": {"sslmode": "require"},
+        }
+    else:
+        engine_kwargs = {"connect_args": {"check_same_thread": False}}
+    engine = create_engine(TRACECAT__DB_URI, echo=True, **engine_kwargs)
     return engine