From 47cf63e0e68e9d797239f68f028d6124456b171c Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Wed, 7 Aug 2019 00:00:59 +0300
Subject: [PATCH] Add psycopg2 as an optional dependency and throw error at
 startup if db is not initialized

---
 mautrix_telegram/__main__.py    | 5 +++++
 mautrix_telegram/db/__init__.py | 4 +++-
 optional-requirements.txt       | 1 +
 setup.py                        | 3 ++-
 4 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/mautrix_telegram/__main__.py b/mautrix_telegram/__main__.py
index 3ab61092..ab3b0b26 100644
--- a/mautrix_telegram/__main__.py
+++ b/mautrix_telegram/__main__.py
@@ -14,6 +14,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from itertools import chain
+import sys
 
 from alchemysession import AlchemySessionContainer
 
@@ -57,6 +58,10 @@ class TelegramBridge(Bridge):
 
     def prepare_db(self) -> None:
         super().prepare_db()
+        if not self.db.has_table("alembic_version"):
+            self.log.critical("alembic_version table not found. "
+                              "Did you forget to `alembic upgrade head`?")
+            sys.exit(10)
         init_db(self.db)
         self.session_container = AlchemySessionContainer(
             engine=self.db, table_base=Base, session=False,
diff --git a/mautrix_telegram/db/__init__.py b/mautrix_telegram/db/__init__.py
index e66a6545..28106767 100644
--- a/mautrix_telegram/db/__init__.py
+++ b/mautrix_telegram/db/__init__.py
@@ -13,6 +13,8 @@
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from sqlalchemy.engine.base import Engine
+
 from mautrix.bridge.db import UserProfile, RoomState
 
 from .bot_chat import BotChat
@@ -23,7 +25,7 @@ from .telegram_file import TelegramFile
 from .user import User, UserPortal, Contact
 
 
-def init(db_engine) -> None:
+def init(db_engine: Engine) -> None:
     for table in (Portal, Message, User, Contact, UserPortal, Puppet, TelegramFile, UserProfile,
                   RoomState, BotChat):
         table.db = db_engine
diff --git a/optional-requirements.txt b/optional-requirements.txt
index b8751b97..9e0a0ff0 100644
--- a/optional-requirements.txt
+++ b/optional-requirements.txt
@@ -2,3 +2,4 @@ cryptg
 Pillow
 moviepy
 prometheus_client
+psycopg2-binary
diff --git a/setup.py b/setup.py
index 82e26bcb..f31b21b5 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,8 @@ extras = {
     "fast_crypto": ["cryptg>=0.1,<0.3"],
     "webp_convert": ["Pillow>=4.3.0,<7"],
     "hq_thumbnails": ["moviepy>=1.0,<2.0"],
-    "metrics": ["prometheus-client>=0.6.0,<0.8.0"],
+    "metrics": ["prometheus_client>=0.6.0,<0.8.0"],
+    "postgres": ["psycopg2-binary>=2,<3"],
 }
 extras["all"] = list({dep for deps in extras.values() for dep in deps})