Heroku multiple accounts database, fixes

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")

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:

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

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)

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:

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,

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+