better mongodb handling

9 files changed, 72 insertions, 13 deletions

diff --git a/data/collector_container/Dockerfile b/data/collector_container/Dockerfile
index a9bb4e5..05bae7a 100644
--- a/data/collector_container/Dockerfile
+++ b/data/collector_container/Dockerfile
@@ -30,8 +30,8 @@ WORKDIR /app/
 USER collector
 
 # Add healthcheck
-HEALTHCHECK --interval=30s --timeout=15s --retries=1 --start-period=30s \
-    CMD sh healthcheck.sh || bash -c 'kill -s 15 1 && (sleep 7; kill -s 9 1)'
+HEALTHCHECK --interval=2m --timeout=15s --retries=1 --start-period=30s \
+    CMD sh ./src/collector/healthcheck.sh COLLECTOR || bash -c 'kill -s 15 1 && (sleep 7; kill -s 9 1)'
 
 ENTRYPOINT ["uvicorn", "src.collector.main:app", "--host", "0.0.0.0", "--workers", "1", "--header", "server:collector"]
 
diff --git a/data/healthcheck-mongodb.js b/data/healthcheck-mongodb.js
new file mode 100644
index 0000000..64071e6
--- /dev/null
+++ b/data/healthcheck-mongodb.js
@@ -0,0 +1,7 @@
+// Get data from DB, check if data is equal to test
+var healthcheck_status = db.healthcheck.findOne({"healthcheck": "ok"});
+print(healthcheck_status);
+if (healthcheck_status["healthcheck"] == "ok"){
+    exit(0);
+}
+exit(1);
diff --git a/data/init-mongodb.js b/data/init-mongodb.js
index 6057d84..054dfd1 100644
--- a/data/init-mongodb.js
+++ b/data/init-mongodb.js
@@ -1,8 +1,3 @@
-
-// To help improve our products, anonymous usage data is collected and sent to MongoDB periodically (https://www.mongodb.com/legal/privacy-policy).
-// You can opt-out by running the disableTelemetry() command.
-disableTelemetry()
-
 // Create the DB by inserting some data
 db.REPLACE_COLLECTION.insertOne({init_key: "init_data"})
 
@@ -23,6 +18,9 @@ db.createUser(
 // Delete the init data
 db.REPLACE_COLLECTION.deleteOne({init_key: "init_data"})
 
+// Add healthcheck
+db.healthcheck.insertOne({healthcheck: "ok"})
+
 // Disable the ad about monitoring
 db.disableFreeMonitoring()
 
diff --git a/data/mongodb_container/Dockerfile b/data/mongodb_container/Dockerfile
index 32ee43b..f37d2a3 100644
--- a/data/mongodb_container/Dockerfile
+++ b/data/mongodb_container/Dockerfile
@@ -19,9 +19,15 @@ RUN find / -xdev -perm /6000 -type f -exec chmod a-s {} \; || true
 
 COPY ./data/mongodb_entrypoint.sh /mongodb_entrypoint.sh
 COPY ./data/init-mongodb.js /init-mongodb.js
+COPY ./data/healthcheck-mongodb.js /healthcheck-mongodb.js
+COPY ./src/collector/healthcheck.sh /healthcheck.sh
 
 USER mongodb
 
 WORKDIR /data/db
 
+# Add healthcheck
+HEALTHCHECK --interval=30s --timeout=15s --retries=1 --start-period=30s \
+    CMD sh /healthcheck.sh MONGODB || bash -c 'kill -s 15 1 && (sleep 7; kill -s 9 1)'
+
 ENTRYPOINT ["bash", "/mongodb_entrypoint.sh"]
diff --git a/data/mongodb_entrypoint.sh b/data/mongodb_entrypoint.sh
index 7a81abc..81ed619 100755
--- a/data/mongodb_entrypoint.sh
+++ b/data/mongodb_entrypoint.sh
@@ -11,7 +11,8 @@ then
   sed -i "s/REPLACE_COLLECTION/$MONGODB_COLLECTION/g" /data/db/init-mongodb.js
 
   # Update and shutdown our DB with changes
-  /usr/bin/mongosh localhost:27015/production /data/db/init-mongodb.js
+  /usr/bin/mongosh --eval 'disableTelemetry()' localhost:27015/production /data/db/init-mongodb.js
+  # /usr/bin/mongosh localhost:27015/production /data/db/init-mongodb.js
   sleep 1 # Allow DB to shutdown
   /usr/bin/touch /data/db/user_exist
   rm /data/db/init-mongodb.js
@@ -19,3 +20,4 @@ fi
 
 # Startup normally now with our user
 exec /usr/bin/mongod --nounixsocket --bind_ip_all --auth
+
diff --git a/src/collector/db.py b/src/collector/db.py
index 2f16e12..54cdc21 100644
--- a/src/collector/db.py
+++ b/src/collector/db.py
@@ -49,7 +49,7 @@ class DBClient:
                 await self.collection.find_one({"_id": ObjectId("507f1f77bcf86cd799439011")})
                 print("Connection to DB - OK")
                 break
-            except Exception:  # pylint: disable=braod-except
+            except Exception:  # pylint: disable=broad-except
                 print(f"WARNING failed to connect to DB - {i} / 4", flush=True)
             sleep(1)
         else:
diff --git a/src/collector/healthcheck.py b/src/collector/healthcheck.py
new file mode 100644
index 0000000..7d336fd
--- /dev/null
+++ b/src/collector/healthcheck.py
@@ -0,0 +1,38 @@
+"""
+Send a healthcheck request
+"""
+import sys
+import time
+import json
+
+import requests
+
+
+def check_collector() -> bool:
+    """Check our collector using /info
+
+    :return: bool
+    """
+    time.sleep(2)  # Prevent race condition with redis container healthcheck
+
+    req = requests.get(
+        "http://localhost:8000/info",
+        timeout=3,
+        # verify="./rootCA.crt",
+    )
+
+    if req.status_code != 200:
+        return False
+
+    data = json.loads(req.text)
+    if isinstance(data["estimated document count"], int) and data["estimated document count"] >= 0:
+        return req.status_code == 200
+
+    return False
+
+
+if __name__ == "__main__":
+    if sys.argv[1] == "COLLECTOR":
+        if check_collector():
+            sys.exit(0)
+    sys.exit(1)
diff --git a/src/collector/healthcheck.sh b/src/collector/healthcheck.sh
new file mode 100755
index 0000000..3c5e042
--- /dev/null
+++ b/src/collector/healthcheck.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# If mongodb container
+if [ $1 = "MONGODB" ]
+then
+  /usr/bin/mongosh --eval 'disableTelemetry()' -u $MONGODB_USERNAME -p $MONGODB_PASSWORD localhost:27017/production /healthcheck-mongodb.js
+  exit $?
+fi
+
+# If collector
+/usr/bin/python3 ./src/collector/healthcheck.py $1 || exit 1
diff --git a/src/collector/main.py b/src/collector/main.py
index 36a8ebd..beaf2ea 100755
--- a/src/collector/main.py
+++ b/src/collector/main.py
@@ -10,8 +10,6 @@ from bson import (
     ObjectId,
     json_util,
 )
-from dotenv import load_dotenv
-
 from .db import (
     DBClient,
     SearchInput,
@@ -19,7 +17,6 @@ from .db import (
 from .schema import valid_schema
 
 
-load_dotenv()
 # Get credentials
 if "MONGODB_USERNAME" not in environ or "MONGODB_PASSWORD" not in environ or "MONGODB_COLLECTION" not in environ:
     print("Missing MONGODB_USERNAME or MONGODB_PASSWORD or MONGODB_COLLECTION in env")
@@ -167,4 +164,4 @@ async def info() -> JSONResponse:
     """
 
     count = await db.estimated_document_count()
-    return JSONResponse(content={"status": "success", "Estimated document count": count})
+    return JSONResponse(content={"status": "success", "estimated document count": count})