Queer European MD passionate about IT
Browse Source

Provide bot with an associated database (SQLite by default)

Davte 4 years ago
parent
commit
589e27bf0e
2 changed files with 58 additions and 3 deletions
  1. 7 3
      davtelepot/bot.py
  2. 51 0
      davtelepot/database.py

+ 7 - 3
davtelepot/bot.py

@@ -15,13 +15,14 @@ from aiohttp import web
 
 # Project modules
 from api import TelegramBot, TelegramError
+from database import ObjectWithDatabase
 from utilities import escape_html_chars, get_secure_key, make_lines_of_buttons
 
 # Do not log aiohttp `INFO` and `DEBUG` levels
 logging.getLogger('aiohttp').setLevel(logging.WARNING)
 
 
-class Bot(TelegramBot):
+class Bot(TelegramBot, ObjectWithDatabase):
     """Simple Bot object, providing methods corresponding to Telegram bot API.
 
     Multiple Bot() instances may be run together, along with a aiohttp web app.
@@ -41,7 +42,7 @@ class Bot(TelegramBot):
 
     def __init__(
         self, token, hostname='', certificate=None, max_connections=40,
-        allowed_updates=[]
+        allowed_updates=[], database_url='bot.db'
     ):
         """Init a bot instance.
 
@@ -56,8 +57,11 @@ class Bot(TelegramBot):
         allowed_updates : List(str)
             Allowed update types (empty list to allow all).
         """
+        # Append `self` to class list of instances
         self.__class__.bots.append(self)
-        super().__init__(token)
+        # Call superclasses constructors with proper arguments
+        TelegramBot.__init__(self, token)
+        ObjectWithDatabase.__init__(self, database_url=database_url)
         self._path = None
         self._offset = 0
         self._hostname = hostname

+ 51 - 0
davtelepot/database.py

@@ -0,0 +1,51 @@
+"""Provide any inheriting object with a dataset-powered database management."""
+
+# Standard library modules
+import logging
+
+# Third party modules
+import dataset
+
+
+class ObjectWithDatabase(object):
+    """Objects inheriting from this class will have a `.db` method.
+
+    Using `subclass_instance.db` will open a SQL transaction.
+    To perform multiple SQL queries in a single transaction use a with
+        statement as in this simple example:
+        ```
+        with my_object.db as db:
+            if db['my_table'].find_one(id=14):
+                db['fourteen_exists'].insert(
+                    {'exists': True}
+                )
+        ```
+    """
+
+    def __init__(self, database_url=None):
+        """Instantiate object and open connection with database."""
+        if database_url is None:
+            database_url = 'database.db'
+        if ':///' not in database_url:
+            # Default database engine is sqlite, which operates on a
+            # single-file database having `.db` extension
+            if not database_url.endswith('.db'):
+                database_url += '.db'
+            database_url = f'sqlite:///{database_url}'
+        self._database_url = database_url
+        try:
+            self._database = dataset.connect(self.db_url)
+        except Exception as e:
+            self._database_url = None
+            self._database = None
+            logging.error(f"{e}")
+
+    @property
+    def db_url(self):
+        """Return complete path to database."""
+        return self._database_url
+
+    @property
+    def db(self):
+        """Return the dataset.Database instance related to `self`."""
+        return self._database