diff --git a/hikka/main.py b/hikka/main.py index 5169e18..5d1eedf 100755 --- a/hikka/main.py +++ b/hikka/main.py @@ -385,7 +385,12 @@ class Hikka: importlib.invalidate_caches() self._get_api_token() - async def save_client_session(self, client: TelegramClient): + async def save_client_session( + self, + client: TelegramClient, + heroku_config: "ConfigVars" = None, # type: ignore + heroku_app: "App" = None, # type: ignore + ): if hasattr(client, "_tg_id"): telegram_id = client._tg_id else: @@ -412,8 +417,9 @@ class Hikka: if "DYNO" not in os.environ: session.save() else: - config = heroku.get_app(self.api_token)[1] - config["hikka_session"] = session.save() + heroku_config["hikka_session"] = session.save() + heroku_app.update_config(heroku_config) + # Heroku will restart the app after updating config client.session = session # Set db attribute to this client in order to save diff --git a/hikka/web/root.py b/hikka/web/root.py index ea21db0..cf1ac9a 100644 --- a/hikka/web/root.py +++ b/hikka/web/root.py @@ -96,6 +96,7 @@ class Web: "tg_done": bool(self.client_data), "okteto": "OKTETO" in os.environ, "lavhost": "LAVHOST" in os.environ, + "heroku": "DYNO" in os.environ, } async def check_session(self, request): @@ -192,15 +193,12 @@ class Web: return web.Response(status=400) if "DYNO" not in os.environ: + # On Heroku it'll be saved later 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, @@ -293,7 +291,9 @@ class Web: # so it doesn't create bot immediately. That's why we only save its session # in case user closes web early. It will be handled on restart # If user finishes login further, client will be passed to main - await main.hikka.save_client_session(self._pending_client) + # To prevent Heroku from restarting too soon, we'll do it after setting bot + if "DYNO" not in os.environ: + await main.hikka.save_client_session(self._pending_client) return web.Response() @@ -304,6 +304,18 @@ class Web: if not self._pending_client: return web.Response(status=400) + if "DYNO" in os.environ: + app, config = heroku.get_app() + config["api_id"] = self.api_token.ID + config["api_hash"] = self.api_token.HASH + await main.hikka.save_client_session( + self._pending_client, + heroku_config=config, + heroku_app=app, + ) + # We don't care what happens next, bc Heroku will restart anyway + return + first_session = not bool(main.hikka.clients) # Client is ready to pass in to dispatcher diff --git a/web-resources/root.jinja2 b/web-resources/root.jinja2 index 96da5c6..23cc682 100755 --- a/web-resources/root.jinja2 +++ b/web-resources/root.jinja2 @@ -67,6 +67,10 @@ color: #fff; padding: 15px 0; } + + .finish_block { + display: none; + } {% endblock %} @@ -80,7 +84,17 @@