support redis db

README.md

@@ -12,9 +12,10 @@ Bot uses [Telegram Bot API](https://core.telegram.org/bots/api) to post messages
 - Run `pip3 install -r requirements.txt` to install dependencies
 - Create your bot via [BotFather](https://t.me/BotFather)
 - Rename `conf/config-sample.yaml` to `conf/config.yaml` and
+  - set `db_type` to `sqlite` or `redis`
   - replace `tg_chat_id` with your channel/chat/group id
   - replace `tg_token` with your bot token
-  - you can also change the `db_name`, `log_name` and `log_file`
+  - you can also change the `db_name`, `redis_url`, `log_name` and `log_file`
 - Run `python3 app.py` in the project folder
 - Use misc/stackernews.service to run the bot as a service
 

app.py

@@ -2,13 +2,14 @@
 # -*- coding: utf-8 -*-
 
 import datetime
-from apscheduler.schedulers.background import BackgroundScheduler
+from apscheduler.schedulers.blocking import BlockingScheduler
 from stackernews import StackerNews
 
 stacker = StackerNews()
 
-scheduler = BackgroundScheduler(timezone='Asia/Hong_Kong')
-scheduler.add_job(stacker.run, 'interval', minutes=10, id='stacker', next_run_time=datetime.datetime.now())
+scheduler = BlockingScheduler(timezone='Asia/Hong_Kong')
+scheduler.add_job(stacker.run, 'interval', minutes=1, id='stacker',
+                  max_instances=1, next_run_time=datetime.datetime.now())
 
 
 if __name__ == '__main__':

conf/config-sample.yaml

@@ -1,5 +1,7 @@
 
+db_type: sqlite  # or redis
 db_name: sqlite:///db/stacker.db
+redis_url: redis://localhost:6379/0
 
 log_name: StackerNews
 log_file: log/stacker-news-top.log

lib/__init__.py

@@ -1,5 +1 @@
-# -*- coding: utf-8 -*-
-
-from .utils import session, logger, config, replace
-from .tg_bot import telegram_bot_send_text
-from .db import Database
+# -*- coding: utf-8 -*-

lib/config.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+
+import logging
+import os
+import yaml
+
+
+def create_logger() -> logging.Logger:
+    config = load_config()
+    _logger = logging.getLogger(config['log_name'])
+    _logger.setLevel(logging.INFO)
+    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    file_handler = logging.FileHandler(config['log_file'])
+    file_handler.setFormatter(formatter)
+    _logger.addHandler(file_handler)
+    return _logger
+
+
+def load_config():
+    """
+    Load config from conf/config.yaml
+    {'db_type': 'sqlite',
+     'db_name': 'sqlite:///db/stacker.db',
+     'redis_url': 'redis://localhost:6379/0',
+     'log_name': 'StackerNews',
+     'log_file': 'log/stacker-news-top.log',
+     'tg_chat_id': -1000000000001,
+     'tg_token': '123456789:TELEGRAM_BOT_TOKEN_SAMPLE'
+    }
+    """
+    try:
+        with open('conf/config.yaml', 'r') as f:
+            _config = yaml.safe_load(f)
+    except FileNotFoundError:
+        return load_config_from_env({})
+    except yaml.YAMLError:
+        return load_config_from_env({})
+    except Exception:
+        return load_config_from_env({})
+
+    return load_config_from_env(_config)
+
+
+def load_config_from_env(_config):
+    """
+    Load config from environment variables.
+    """
+    _config_env = {
+        'db_type': os.getenv('DB_TYPE', _config.get('db_type')),
+        'db_name': os.environ.get('DB_NAME', _config.get('db_name', 'sqlite:///db/stacker.db')),
+        'redis_url': os.environ.get('REDIS_URL', _config.get('redis_url')),
+        'log_name': os.environ.get('LOG_NAME', _config.get('log_name', 'StackerNews')),
+        'log_file': os.environ.get('LOG_FILE', _config.get('log_file', 'log/stacker-news-top.log')),
+        'tg_chat_id': os.environ.get('TG_CHAT_ID', _config.get('tg_chat_id')),
+        'tg_token': os.environ.get('TG_TOKEN', _config.get('tg_token'))
+    }
+
+    if _config_env['db_type'] not in ['sqlite', 'redis']:
+        raise Exception('Invalid database type.')
+    if not _config_env['tg_chat_id'] or not _config_env['tg_token']:
+        raise ValueError('Missing Telegram chat_id or token.')
+
+    return _config_env
+

lib/db.py

@@ -1,78 +1,18 @@
 # -*- coding: utf-8 -*-
 
-import datetime
-from sqlalchemy import create_engine
-from sqlalchemy import Column, Integer, String, DateTime
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker
-from lib import config
-
-Base = declarative_base()
-
-
-class Stacker(Base):
-    """
-    The Database model for the Stacker table.
-    """
-    __tablename__ = 'threads'
-
-    id = Column(Integer, primary_key=True)
-    thread_id = Column(Integer)
-    title = Column(String)
-    sats = Column(Integer)
-    comments = Column(Integer)
-    create_at = Column(DateTime)
-
-    def __init__(self, thread_id, title, sats, comments, create_at):
-        self.thread_id = thread_id
-        self.title = title
-        self.sats = sats
-        self.comments = comments
-        self.create_at = create_at
+from .db_redis import RedisDatabase
+from .db_sqlite import SqliteDatabase
+from lib.config import load_config
 
 
 class Database(object):
-    def __init__(self):
-        self.engine = create_engine(config['db_name'], echo=False)
-        # create tables
-        Base.metadata.create_all(self.engine)
-
-        # create a Session
-        session_made = sessionmaker(bind=self.engine)
-        self.session = session_made()
-
-    def add_thread(self, thread_id: int, title: str, sats: int, comments: int):
-        """
-        Add a thread to the database.
-        :param thread_id: The thread id to add.
-        :param title: The thread title to add.
-        :param sats: The thread sats to add.
-        :param comments: The thread comments to add.
-        """
-        thread_db = self.get_thread(thread_id)
-        if thread_db:
-            return
-        thread = Stacker(thread_id, title, sats, comments, datetime.datetime.now())
-        self.session.add(thread)
-        self.session.commit()
-
-    def get_thread(self, thread_id: int) -> Stacker:
-        """
-        Get a thread from the database.
-        :param thread_id: The thread id to get.
-        """
-        # for thread in self.session.query(Stacker).filter(Stacker.title == 'Hello World'):
-        #     print(thread.thread_id, thread.title, thread.sats, thread.comments, thread.create_at)
-        thread = self.session.query(Stacker).filter_by(thread_id=thread_id).first()
-        return thread
-
-    def del_thread(self, thread_id: int):
-        """
-        Delete a thread from the database.
-        :param thread_id: The thread id to delete.
-        """
-        thread = self.get_thread(thread_id)
-        if not thread:
-            return
-        self.session.delete(thread)
-        self.session.commit()
+    def init_db(self):
+        config = load_config()
+        if config['db_type'] == 'redis':
+            self.db = RedisDatabase(config['redis_url'])
+        elif config['db_type'] == 'sqlite':
+            self.db = SqliteDatabase(config['db_name'])
+        else:
+            raise Exception('Database type not supported.')
+
+        return self.db

lib/db_redis.py

@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+
+import redis
+
+
+class RedisDatabase(object):
+    def __init__(self, redis_url):
+        self.r = redis.Redis.from_url(redis_url)
+
+    def add_thread(self, thread_id: int, title: str, sats: int, comments: int):
+        """
+        Add a thread to the database.
+        :param thread_id: The thread id to add.
+        :param title: The thread title to add.
+        :param sats: The thread sats to add.
+        :param comments: The thread comments to add.
+        """
+        self.r.hmset(thread_id, {'title': title, 'sats': sats, 'comments': comments})
+
+    def get_thread(self, thread_id: int) -> dict:
+        """
+        Get a thread from the database.
+        :param thread_id: The thread id to get.
+        :return: The thread object.
+        """
+        return self.r.hgetall(thread_id)
+
+    def del_thread(self, thread_id: int):
+        """
+        Delete a thread from the database.
+        :param thread_id: The thread id to delete.
+        """
+        self.r.delete(thread_id)

lib/db_sqlite.py

@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+
+import datetime
+from sqlalchemy import create_engine
+from sqlalchemy import Column, Integer, String, DateTime
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+Base = declarative_base()
+
+
+class Stacker(Base):
+    """
+    The Database model for the Stacker table.
+    """
+    __tablename__ = 'threads'
+
+    id = Column(Integer, primary_key=True)
+    thread_id = Column(Integer)
+    title = Column(String)
+    sats = Column(Integer)
+    comments = Column(Integer)
+    create_at = Column(DateTime)
+
+    def __init__(self, thread_id, title, sats, comments, create_at):
+        self.thread_id = thread_id
+        self.title = title
+        self.sats = sats
+        self.comments = comments
+        self.create_at = create_at
+
+
+class SqliteDatabase(object):
+    def __init__(self, db_name):
+        self.engine = create_engine(db_name, echo=False)
+        # create tables
+        Base.metadata.create_all(self.engine)
+
+        # create a Session
+        session_made = sessionmaker(bind=self.engine)
+        self.session = session_made()
+
+    def add_thread(self, thread_id: int, title: str, sats: int, comments: int):
+        """
+        Add a thread to the database.
+        :param thread_id: The thread id to add.
+        :param title: The thread title to add.
+        :param sats: The thread sats to add.
+        :param comments: The thread comments to add.
+        """
+        thread_db = self.get_thread(thread_id)
+        if thread_db:
+            return
+        thread = Stacker(thread_id, title, sats, comments, datetime.datetime.now())
+        self.session.add(thread)
+        self.session.commit()
+
+    def get_thread(self, thread_id: int) -> Stacker:
+        """
+        Get a thread from the database.
+        :param thread_id: The thread id to get.
+        """
+        # for thread in self.session.query(Stacker).filter(Stacker.title == 'Hello World'):
+        #     print(thread.thread_id, thread.title, thread.sats, thread.comments, thread.create_at)
+        thread = self.session.query(Stacker).filter_by(thread_id=thread_id).first()
+        return thread
+
+    def del_thread(self, thread_id: int):
+        """
+        Delete a thread from the database.
+        :param thread_id: The thread id to delete.
+        """
+        thread = self.get_thread(thread_id)
+        if not thread:
+            return
+        self.session.delete(thread)
+        self.session.commit()

lib/tg_bot.py

@@ -1,16 +1,17 @@
 # -*- coding: utf-8 -*-
 
-from lib import session, config
+from lib.utils import session
+from lib.config import load_config
 
-url = 'https://api.telegram.org/bot{}/sendMessage'.format(config['tg_token'])
 
+class TelegramBot(object):
+    def __init__(self):
+        self.config = load_config()
+        self.url = 'https://api.telegram.org/bot{}/sendMessage'.format(self.config['tg_token'])
 
-def telegram_bot_send_text(message):
-    """
-    Send message to telegram channel/chat_id/group via telegram bot.
-    """
-    data = {'chat_id': config['tg_chat_id'],
-            'text': message,
-            'parse_mode': 'MarkdownV2'}
-    response = session.post(url, json=data, headers={'Accept': 'application/json'})
-    return response.json()
+    def send(self, message):
+        data = {'chat_id': self.config['tg_chat_id'],
+                'text': message,
+                'parse_mode': 'MarkdownV2'}
+        response = session.post(self.url, json=data, headers={'Accept': 'application/json'})
+        return response.json()