From 1e7c0201d99a7d5d2e2c3ebac91c5ea97ed4da48 Mon Sep 17 00:00:00 2001 From: hikariatama Date: Mon, 30 May 2022 20:58:20 +0000 Subject: [PATCH] Heroku multiple accounts database, fixes --- hikka/database.py | 9 ++++++--- hikka/heroku.py | 2 +- hikka/main.py | 7 ++++--- hikka/modules/test.py | 4 ++-- hikka/modules/updater.py | 9 +++++++-- hikka/web/root.py | 19 ++++++++++++------- requirements.txt | 3 ++- 7 files changed, 34 insertions(+), 19 deletions(-) diff --git a/hikka/database.py b/hikka/database.py index 9144836..6fa912c 100755 --- a/hikka/database.py +++ b/hikka/database.py @@ -68,8 +68,8 @@ class Database(dict): await asyncio.sleep(5) self._postgre.execute( - "DELETE FROM hikka WHERE id = 1; INSERT INTO hikka (id, data) VALUES (1, %s);", - (json.dumps(self),), + "DELETE FROM hikka WHERE id = %s; INSERT INTO hikka (id, data) VALUES (%s, %s);", + (self._client._tg_id, self._client._tg_id, json.dumps(self)), ) self._postgre.connection.commit() @@ -152,7 +152,10 @@ class Database(dict): """Read database and stores it in self""" if self._postgre: try: - self._postgre.execute("SELECT data FROM hikka") + self._postgre.execute( + "SELECT data FROM hikka WHERE id=%s;", + (self._client._tg_id,), + ) self.update(**json.loads(self._postgre.fetchall()[0][0])) except Exception: logger.exception("Error reading postgresql database") diff --git a/hikka/heroku.py b/hikka/heroku.py index 4acb40a..a807f78 100644 --- a/hikka/heroku.py +++ b/hikka/heroku.py @@ -24,7 +24,7 @@ def publish( app, config = get_app(key, api_token, create_new) # Will be configured later in app - config["HIKKA_SESSION"] = "None" + config["hikka_session"] = None config["heroku_api_token"] = key if api_token is not None: diff --git a/hikka/main.py b/hikka/main.py index 02553d3..4183714 100755 --- a/hikka/main.py +++ b/hikka/main.py @@ -328,8 +328,8 @@ class Hikka: ) ] - if os.environ.get("HIKKA_SESSION") not in {None, "None"}: - self.sessions += [StringSession(os.environ.get("HIKKA_SESSION"))] + if os.environ.get("hikka_session"): + self.sessions += [StringSession(os.environ.get("hikka_session"))] def _get_api_token(self): """Get API Token from disk or environment""" @@ -413,7 +413,8 @@ class Hikka: if "DYNO" not in os.environ: session.save() else: - os.environ["HIKKA_SESSION"] = session.save() + config = heroku.get_app(os.environ["heroku_api_token"])[1] + config["hikka_session"] = session.save() client.session = session # Set db attribute to this client in order to save diff --git a/hikka/modules/test.py b/hikka/modules/test.py index 143ba9a..6d930d6 100755 --- a/hikka/modules/test.py +++ b/hikka/modules/test.py @@ -413,8 +413,8 @@ class TestMod(loader.Module): self._logchat = int(f"-100{chat.id}") - if not is_new or all( - participant.id != self.inline.bot_id + if not is_new and any( + participant.id == self.inline.bot_id for participant in (await self._client.get_participants(chat, limit=3)) ): logging.getLogger().handlers[0].install_tg_log(self) diff --git a/hikka/modules/updater.py b/hikka/modules/updater.py index bcbca35..6b93187 100755 --- a/hikka/modules/updater.py +++ b/hikka/modules/updater.py @@ -41,7 +41,7 @@ from telethon.tl.functions.messages import ( ) from telethon.tl.types import DialogFilter, Message -from .. import loader, utils +from .. import loader, utils, heroku from ..inline.types import InlineCall logger = logging.getLogger(__name__) @@ -177,6 +177,11 @@ class UpdaterMod(loader.Module): os.system("lavhost restart") return + if "DYNO" in os.environ: + app = heroku.get_app(os.environ["heroku_api_token"])[0] + app.dynos()[0].restart() + return + atexit.register(restart, *sys.argv[1:]) handler = logging.getLogger().handlers[0] handler.setLevel(logging.CRITICAL) @@ -415,7 +420,7 @@ class UpdaterMod(loader.Module): self.set("do_not_create", True) - async def update_complete(self, client: "TelegramClient"): # noqa: F821 + async def update_complete(self, client: "TelegramClient"): # type: ignore logger.debug("Self update successful! Edit message") start = self.get("restart_ts") try: diff --git a/hikka/web/root.py b/hikka/web/root.py index 4b91a21..413dc11 100644 --- a/hikka/web/root.py +++ b/hikka/web/root.py @@ -42,7 +42,7 @@ import re import requests import time -from .. import utils, main, database +from .. import utils, main, database, heroku from aiogram.types import InlineKeyboardMarkup, InlineKeyboardButton from telethon.errors.rpcerrorlist import YouBlockedUserError @@ -116,7 +116,7 @@ class Web: async def _check_bot( self, - client: "TelegramClient", # noqa: F821 + client: "TelegramClient", # type: ignore username: str, ) -> bool: async with client.conversation("@BotFather", exclusive=False) as conv: @@ -191,11 +191,16 @@ class Web: ): return web.Response(status=400) - with open( - os.path.join(self.data_root or DATA_DIR, "api_token.txt"), - "w", - ) as f: - f.write(api_id + "\n" + api_hash) + if "DYNO" not in os.environ: + with open( + os.path.join(self.data_root or DATA_DIR, "api_token.txt"), + "w", + ) as f: + f.write(api_id + "\n" + api_hash) + else: + config = heroku.get_app(os.environ["heroku_api_token"])[1] + config["api_id"] = api_id + config["api_hash"] = api_hash self.api_token = collections.namedtuple("api_token", ("ID", "HASH"))( api_id, diff --git a/requirements.txt b/requirements.txt index 7a5d4d8..9a1a8a6 100755 --- a/requirements.txt +++ b/requirements.txt @@ -9,6 +9,7 @@ requests==2.27.1 aiogram==2.19 websockets==10.2 grapheme==0.6.0 -heroku3 +heroku3==5.1.4 +psycopg2-binary==2.9.3 # Python 3.8+