initial version

.dockerignore

@@ -0,0 +1,2 @@
+.git
+.idea

.gitignore

@@ -0,0 +1 @@
+.idea

Dockerfile

@@ -0,0 +1,28 @@
+# ------------------------------------------------------------
+# create requirements.txt
+# ------------------------------------------------------------
+FROM python:3 AS pipenv
+RUN pip install pipenv
+COPY Pipfile* /tmp/
+RUN cd /tmp && pipenv lock --requirements > requirements.txt
+
+# ------------------------------------------------------------
+# build actual container
+# ------------------------------------------------------------
+FROM python:3-slim
+
+ENV TRIES=0 \
+    RABBITMQ_URI="" \
+    MONGO_URI="" \
+    URL=""
+
+# Install requirements
+WORKDIR /usr/src
+COPY --from=pipenv /tmp/requirements.txt /usr/src/
+RUN pip install -r /usr/src/requirements.txt
+
+# Copy application
+COPY *.py checks.json /usr/src/
+
+# Start application
+CMD python /usr/src/main.py

Pipfile

@@ -0,0 +1,14 @@
+[[source]]
+name = "pypi"
+url = "https://pypi.org/simple"
+verify_ssl = true
+
+[dev-packages]
+
+[packages]
+pika = "*"
+pymongo = "*"
+psycopg2-binary = "*"
+
+[requires]
+python_version = "3.7"

README.md

@@ -0,0 +1,103 @@
+# NCSA CHECKS
+
+The goal of this project is to have a simple set of functions that
+can be used as init containers in your kubernetes deployment. Each
+of the checks will see if the service is ready to accept
+connections.
+
+For example the following snippet will wait for rabbitmq and 
+postgresql to be ready before starting the actual container:
+
+```yaml
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: mywebapp
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: mywebapp
+  template:
+    metadata:
+      labels:
+        app: mywebapp
+    spec:
+      initContainers:
+        - name: check-rabbitmq
+          image: "ncsa/check:1.0.0"
+          env:
+            - name: RABBITMQ_URI
+              value: "amqp://user:pass@rabbitmq/%2F"
+        - name: check-postgresql
+          image: "ncsa/check:1.0.0"
+          env:
+            - name: PGHOST
+              value: "postgres"
+            - name: PGUSER
+              value: "postgres"
+            - name: PGTABLE
+              value: "users"
+      containers:
+        - name: mywebapp
+          image: mywebapp:1.0
+          ports:
+            - containerPort: 80
+          env:
+            - name: RABBITMQ_URI
+              value: "amqp://user:pass@rabbitmq/%2F"
+            - name: PGHOST
+              value: "postgres"
+            - name: PGUSER
+              value: "postgres"
+```
+
+## Tests Supported
+
+RabbitMQ
+- description : checks to see if can connect to RabbitMQ Server
+- exit code : 1
+- parameters
+  - RABBITMQ_URI [required] : URI for RabbitMQ server 
+
+Mongo
+- description : checks to see if can connect to MongoDB Server
+- exit code : 2
+- parameters
+  - MONGO_URI [required] : URI for MongoDB server 
+
+URL
+- description : checks to see if URL is reachable and returns status 
+  code of 200. If optional text is given the text needs to be in the
+  returned body.
+- exit code : 3
+- parameters
+  - URL [required] : URL to check
+  - URL_TEXT [optional] : text to be found in body returned 
+  
+PostgreSQL
+- description : checks to see if database is up. Will check to see
+  if the optional table exists
+- exit code : 4
+- parameters
+  - PG_URI [required*] : URI to connect to postgresql.
+  - PGTABLE [optional] : table that should exist 
+
+## Version History
+
+### 1.0 - 2020-08-07
+
+#### Added
+- support for RabbitMQ
+- support for MongoDB
+- support for URL
+- support for PostgreSQL
+
+
+## TODO:
+
+check_postgres: support for standard postgresql parameters
+check_service: check if host:port is reachable
+
+create simple tests for each check
+github actions

checks.json

@@ -0,0 +1,28 @@
+[
+    {
+        "function": "checks.check_rabbitmq",
+        "parameters": {
+            "rabbitmq_uri": "RABBITMQ_URI"
+        }
+    },
+    {
+        "function": "checks.check_mongo",
+        "parameters": {
+            "mongo_uri": "MONGO_URI"
+        }
+    },
+    {
+        "function": "checks.check_url",
+        "parameters": {
+            "url": "URL",
+            "text": "URL_TEXT"
+        }
+    },
+    {
+        "function": "checks.check_postgresql",
+        "parameters": {
+            "pg_uri": "PG_URI",
+            "pg_table": "PG_TABLE"
+        }
+    }
+]

checks.py

@@ -0,0 +1,84 @@
+import urllib.request
+
+import pika
+import psycopg2
+import pymongo
+
+
+# ----------------------------------------------------------------------
+# RABBITMQ
+# ----------------------------------------------------------------------
+def check_rabbitmq(rabbitmq_uri):
+    params = pika.URLParameters(rabbitmq_uri)
+    connection = pika.BlockingConnection(params)
+    if connection.is_open:
+        print('Connected to RabbitMQ.')
+        connection.close()
+        return True
+    return False
+
+
+# ----------------------------------------------------------------------
+# MONGO
+# ----------------------------------------------------------------------
+def check_mongo(mongo_uri):
+    client = pymongo.MongoClient(mongo_uri,
+                                 connectTimeoutMS=1000,
+                                 socketTimeoutMS=1000,
+                                 serverSelectionTimeoutMS=1000)
+    result = client.admin.command('ismaster')
+    print("Connected to Mongo, master = ", result)
+    return True
+
+
+# ----------------------------------------------------------------------
+# URL
+# ----------------------------------------------------------------------
+def check_url(url, text=None):
+    response = urllib.request.urlopen(url)
+    if response.code >= 200 or response <= 299:
+        print("Connected to URL ", url)
+        if text:
+            for line in response.readlines():
+                if text in line.decode("utf-8"):
+                    print(line)
+                    response.close()
+                    print(f"Text '{text}' found in response.")
+                    return True
+            response.close()
+            print(f"Text '{text}' not found in response.")
+            return False
+        else:
+            response.close()
+            return True
+    response.close()
+    return False
+
+
+# ----------------------------------------------------------------------
+# RABBITMQ
+# ----------------------------------------------------------------------
+def check_postgresql(pg_uri=None, pg_table=None):
+    connection = None
+    if pg_uri:
+        connection = psycopg2.connect(pg_uri)
+    else:
+        exit(4)
+    if not connection:
+        return False
+    cursor = connection.cursor()
+    print("Connected to database")
+    if pg_table:
+        cursor.execute('SELECT count(*) FROM "%s";' % pg_table)
+        row = cursor.fetchone()
+        cursor.close()
+        connection.close()
+        if row:
+            print("Found %d rows in table %s." % (row[0], pg_table))
+            return True
+        else:
+            print("Did not find table %s." % pg_table)
+    else:
+        cursor.close()
+        connection.close()
+        return True