From 8766e36aa6f4c3cadf5b751219790380e32642c7 Mon Sep 17 00:00:00 2001
From: mitsuha_s <stmiroslavsckaya@yandex.ru>
Date: Sun, 21 Aug 2022 16:03:00 +0300
Subject: [PATCH] add database migrations

---
 database.py                       | 30 +++++++++---------------------
 migrations/0000.initial-schema.py | 18 ++++++++++++++++++
 requirements.txt                  |  1 +
 3 files changed, 28 insertions(+), 21 deletions(-)
 create mode 100644 migrations/0000.initial-schema.py

diff --git a/database.py b/database.py
index 4ea6a67..9034de3 100644
--- a/database.py
+++ b/database.py
@@ -4,6 +4,8 @@ from psycopg2.extensions import connection
 from psycopg2.extras import DictCursor, DictRow
 from exceptions import DisplayableException
 from rss import FeedItem
+from yoyo import read_migrations
+from yoyo import get_backend
 
 
 class Database:
@@ -15,7 +17,7 @@ class Database:
         self.log.debug('Database.__init__(DSN=\'%s\')', dsn)
         self.conn: connection = psycopg2.connect(dsn)
         self.cur: DictCursor = self.conn.cursor(cursor_factory=DictCursor)
-        self.__init_schema()
+        self.__migration(dsn)
 
     def add_user(self, telegram_id: int) -> int:
         """Add a user's telegram id to the database and return its database id."""
@@ -156,26 +158,12 @@ class Database:
             'INSERT INTO feeds_last_items (feed_id, url, guid) VALUES (%s, %s, %s)', new_items)
         self.conn.commit()
 
-    def __init_schema(self) -> None:
-        self.log.debug('__init_schema()')
-        self.cur.execute(
-            'CREATE TABLE IF NOT EXISTS users (id SERIAL PRIMARY KEY, telegram_id INTEGER NOT NULL UNIQUE)'
-        )
-        self.cur.execute('CREATE TABLE IF NOT EXISTS feeds (id SERIAL PRIMARY KEY, url TEXT NOT NULL UNIQUE)')
-        self.cur.execute(
-            'CREATE TABLE IF NOT EXISTS subscriptions ('
-            '   user_id INTEGER REFERENCES users,'
-            '   feed_id INTEGER REFERENCES feeds,'
-            '   UNIQUE (user_id, feed_id)'
-            ')'
-        )
-        self.cur.execute(
-            'CREATE TABLE IF NOT EXISTS feeds_last_items ('
-            '   feed_id INTEGER REFERENCES feeds ON DELETE CASCADE,'
-            '   url TEXT NOT NULL,'
-            '   guid TEXT'
-            ')'
-        )
+    def __migration(self, dsn: str) -> None:
+        backend = get_backend(dsn)
+        migrations = read_migrations('./migrations')
+
+        with backend.lock():
+            backend.apply_migrations(backend.to_apply(migrations))
 
     @staticmethod
     def __dictrow_to_dict_list(rows: list[DictRow]) -> list[dict]:
diff --git a/migrations/0000.initial-schema.py b/migrations/0000.initial-schema.py
new file mode 100644
index 0000000..1553418
--- /dev/null
+++ b/migrations/0000.initial-schema.py
@@ -0,0 +1,18 @@
+from yoyo import step
+
+steps = [
+    step('CREATE TABLE IF NOT EXISTS users (id SERIAL PRIMARY KEY, telegram_id INTEGER NOT NULL UNIQUE)'),
+    step('CREATE TABLE IF NOT EXISTS feeds (id SERIAL PRIMARY KEY, url TEXT NOT NULL UNIQUE)'),
+    step('CREATE TABLE IF NOT EXISTS subscriptions ('
+            '   user_id INTEGER REFERENCES users,'
+            '   feed_id INTEGER REFERENCES feeds,'
+            '   UNIQUE (user_id, feed_id)'
+            ')'
+    ),
+    step('CREATE TABLE IF NOT EXISTS feeds_last_items ('
+            '   feed_id INTEGER REFERENCES feeds ON DELETE CASCADE,'
+            '   url TEXT NOT NULL,'
+            '   guid TEXT'
+            ')'
+    )
+]
diff --git a/requirements.txt b/requirements.txt
index cc5e556..fd8e157 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,3 +13,4 @@ six==1.16.0
 urllib3==1.26.9
 validators==0.19.0
 webencodings==0.5.1
+yoyo-migrations==7.3.2