From cb472826211c4bd3e9293b517f20a5d026256789 Mon Sep 17 00:00:00 2001
From: Michael DM Dryden <mk.dryden@utoronto.ca>
Date: Thu, 4 Jun 2020 02:24:24 -0400
Subject: [PATCH] PostgresStore: Create database if missing

---
 poetry.lock                       | 32 ++++++++++++++++++++++++++++++-
 pyproject.toml                    |  1 +
 telegram_stats_bot/log_storage.py | 22 ++++++++++++++++++++-
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 37ed1ac..119e09c 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -238,6 +238,33 @@ postgresql_psycopg2binary = ["psycopg2-binary"]
 postgresql_psycopg2cffi = ["psycopg2cffi"]
 pymysql = ["pymysql"]
 
+[[package]]
+category = "main"
+description = "Various utility functions for SQLAlchemy."
+name = "sqlalchemy-utils"
+optional = false
+python-versions = "*"
+version = "0.36.6"
+
+[package.dependencies]
+SQLAlchemy = ">=1.0"
+six = "*"
+
+[package.extras]
+anyjson = ["anyjson (>=0.3.3)"]
+arrow = ["arrow (>=0.3.4)"]
+babel = ["Babel (>=1.3)"]
+color = ["colour (>=0.0.4)"]
+encrypted = ["cryptography (>=0.6)"]
+intervals = ["intervals (>=0.7.1)"]
+password = ["passlib (>=1.6,<2.0)"]
+pendulum = ["pendulum (>=2.0.5)"]
+phone = ["phonenumbers (>=5.9.2)"]
+test = ["pytest (>=2.7.1)", "Pygments (>=1.2)", "Jinja2 (>=2.3)", "docutils (>=0.10)", "flexmock (>=0.9.7)", "mock (2.0.0)", "psycopg2 (>=2.5.1)", "psycopg2cffi (>=2.8.1)", "pg8000 (>=1.12.4)", "pytz (>=2014.2)", "python-dateutil (>=2.6)", "pymysql", "flake8 (>=2.4.0)", "isort (>=4.2.2)", "pyodbc"]
+test_all = ["anyjson (>=0.3.3)", "arrow (>=0.3.4)", "Babel (>=1.3)", "colour (>=0.0.4)", "cryptography (>=0.6)", "intervals (>=0.7.1)", "passlib (>=1.6,<2.0)", "pendulum (>=2.0.5)", "phonenumbers (>=5.9.2)", "pytest (>=2.7.1)", "Pygments (>=1.2)", "Jinja2 (>=2.3)", "docutils (>=0.10)", "flexmock (>=0.9.7)", "mock (2.0.0)", "psycopg2 (>=2.5.1)", "psycopg2cffi (>=2.8.1)", "pg8000 (>=1.12.4)", "pytz (>=2014.2)", "python-dateutil (>=2.6)", "pymysql", "flake8 (>=2.4.0)", "isort (>=4.2.2)", "pyodbc", "python-dateutil", "furl (>=0.4.1)"]
+timezone = ["python-dateutil"]
+url = ["furl (>=0.4.1)"]
+
 [[package]]
 category = "main"
 description = "Tornado is a Python web framework and asynchronous networking library, originally developed at FriendFeed."
@@ -256,7 +283,7 @@ python-versions = "*"
 version = "3.7.4.2"
 
 [metadata]
-content-hash = "5293594d2a128d025336950f15a771b97b505200dc53f61bb8f4b38b0322e2f0"
+content-hash = "1948c56aa84c247f9cf0051c330763a7f5a17ea0cd8a7a398d7d2fc71792da52"
 python-versions = "^3.7"
 
 [metadata.files]
@@ -518,6 +545,9 @@ sqlalchemy = [
     {file = "SQLAlchemy-1.3.17-cp38-cp38-win_amd64.whl", hash = "sha256:ed375a79f06cad285166e5be74745df1ed6845c5624aafadec4b7a29c25866ef"},
     {file = "SQLAlchemy-1.3.17.tar.gz", hash = "sha256:156a27548ba4e1fed944ff9fcdc150633e61d350d673ae7baaf6c25c04ac1f71"},
 ]
+sqlalchemy-utils = [
+    {file = "SQLAlchemy-Utils-0.36.6.tar.gz", hash = "sha256:7a7fab14bed80df065412bbf71a0a9b0bfeb4b7c111c2d9bffe57283082f3a6b"},
+]
 tornado = [
     {file = "tornado-6.0.4-cp35-cp35m-win32.whl", hash = "sha256:5217e601700f24e966ddab689f90b7ea4bd91ff3357c3600fa1045e26d68e55d"},
     {file = "tornado-6.0.4-cp35-cp35m-win_amd64.whl", hash = "sha256:c98232a3ac391f5faea6821b53db8db461157baa788f5d6222a193e9456e1740"},
diff --git a/pyproject.toml b/pyproject.toml
index e2c1c37..bd2700a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,6 +10,7 @@ readme = "README.rst"
 python = "^3.7"
 python-telegram-bot = "^12.7"
 sqlalchemy = {version = "^1.3.17", extras = ["postgresql_psycopg2binary"]}
+sqlalchemy-utils = "^0.36.6"
 pandas = "^1.0.4"
 seaborn = "^0.10.1"
 numpy = "^1.18.4"
diff --git a/telegram_stats_bot/log_storage.py b/telegram_stats_bot/log_storage.py
index 2a1dac1..26eda72 100644
--- a/telegram_stats_bot/log_storage.py
+++ b/telegram_stats_bot/log_storage.py
@@ -18,6 +18,7 @@
 #
 # You should have received a copy of the GNU Public License
 # along with this program. If not, see [http://www.gnu.org/licenses/].
+from copy import copy
 
 import datetime
 import logging
@@ -25,6 +26,8 @@ import json
 import os
 
 from sqlalchemy import MetaData, Table, Column, create_engine, BigInteger, TIMESTAMP, Text
+from sqlalchemy_utils import database_exists
+from sqlalchemy_utils.functions.orm import quote
 
 from .parse import MessageDict
 from .db import init_dbs
@@ -69,7 +72,24 @@ class JSONStore(object):
 
 class PostgresStore(object):
     def __init__(self, connection_url: str):
-        self.engine = create_engine(connection_url, echo=False)
+        self.engine = create_engine(connection_url, echo=False, isolation_level="AUTOCOMMIT")
+        if not database_exists(self.engine.url):
+            text = f"""
+                    CREATE DATABASE {quote(self.engine, self.engine.url.database)}
+                    ENCODING 'utf-8'
+                    TEMPLATE {quote(self.engine, 'template1')}
+                    """
+
+            url = copy(self.engine.url)
+            url.database = 'postgres'
+
+            engine = create_engine(url, echo=False, isolation_level="AUTOCOMMIT")
+            result_proxy = engine.execute(text)
+
+            if result_proxy is not None:
+                result_proxy.close()
+            engine.dispose()
+
         init_dbs(self.engine)
 
     def append_data(self, name: str, data: MessageDict):