diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
index 749616d716..34ca585c98 100644
--- a/.github/workflows/docker.yaml
+++ b/.github/workflows/docker.yaml
@@ -12,5 +12,5 @@ jobs:
     steps:
       - uses: actions/checkout@v2
       - name: Build Docker Image
-        run: docker build . --file docker/Dockerfile --tag inventree:$(date +%s)
+        run: cd docker && docker build . --tag inventree:$(date +%s)
     
\ No newline at end of file
diff --git a/InvenTree/InvenTree/management/commands/wait_for_db.py b/InvenTree/InvenTree/management/commands/wait_for_db.py
index 9019a168ef..198c5dcf73 100644
--- a/InvenTree/InvenTree/management/commands/wait_for_db.py
+++ b/InvenTree/InvenTree/management/commands/wait_for_db.py
@@ -7,6 +7,8 @@ from django.core.management.base import BaseCommand
 from django.db import connections
 from django.db.utils import OperationalError
 
+import psycopg2
+
 import time
 
 
@@ -26,6 +28,6 @@ class Command(BaseCommand):
                 db_conn = connections['default']
                 # prints success messge in green
                 self.stdout.write(self.style.SUCCESS('InvenTree database connected'))
-            except OperationalError:
+            except (OperationalError, psycopg2.OperationalError):
                 self.stdout.write(self.style.ERROR("Database unavailable, waiting 5 seconds ..."))
                 time.sleep(5)
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 073698adcd..a86fd09e45 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,14 +1,18 @@
 FROM python:alpine as production
 
 # GitHub source
-ARG INVENTREE_REPO="https://github.com/inventree/InvenTree.git"
-ARG INVENTREE_VERSION="master"
+ARG repository="https://github.com/inventree/InvenTree.git"
+ARG branch="master"
 
 ENV PYTHONUNBUFFERED 1
 
 # InvenTree key settings
 ENV INVENTREE_HOME="/home/inventree"
 
+# GitHub settings
+ENV INVENTREE_REPO="${repository}"
+ENV INVENTREE_BRANCH="${branch}"
+
 ENV INVENTREE_LOG_LEVEL="INFO"
 
 # InvenTree paths
@@ -57,7 +61,7 @@ RUN apk add --no-cache postgresql-contrib postgresql-dev libpq
 RUN apk add --no-cache mariadb-connector-c mariadb-dev
 
 # Clone source code
-RUN git clone --branch ${INVENTREE_VERSION} --depth 1 ${INVENTREE_REPO} ${INVENTREE_SRC_DIR}
+RUN git clone --branch ${INVENTREE_BRANCH} --depth 1 ${INVENTREE_REPO} ${INVENTREE_SRC_DIR}
 
 # Setup Python virtual environment
 RUN python3 -m venv ${INVENTREE_VENV}