From f399a285e436f0a55266d52c9112d20eadd73bac Mon Sep 17 00:00:00 2001
From: Amadej Kastelic <amadejkastelic7@gmail.com>
Date: Mon, 15 Jul 2024 10:54:41 +0200
Subject: [PATCH] Recover from db error

---
 bot/common/utils.py                    | 10 +++++-----
 bot/management/commands/discord_bot.py |  2 +-
 bot/management/commands/purge_posts.py |  2 +-
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/bot/common/utils.py b/bot/common/utils.py
index 036823c..cd51a62 100644
--- a/bot/common/utils.py
+++ b/bot/common/utils.py
@@ -7,7 +7,7 @@
 import typing
 
 import magic
-from django.db import connections
+from django import db as django_db
 
 emoji = ['😼', '😺', '😸', '😹', '😻', '🙀', '😿', '😾', '😩', '🙈', '🙉', '🙊', '😳']
 
@@ -56,7 +56,7 @@ def random_emoji() -> str:
     return random.choice(emoji)
 
 
-def recover_from_db_error():
-    for conn in connections.all():
-        conn.close()
-        conn.connect()
+def recover_from_db_error(exc: Exception):
+    if isinstance(exc, django_db.OperationalError) and exc.args[0] in (2006, 2013):
+        django_db.reset_queries()
+        django_db.close_old_connections()
diff --git a/bot/management/commands/discord_bot.py b/bot/management/commands/discord_bot.py
index bbdce54..b596a5f 100644
--- a/bot/management/commands/discord_bot.py
+++ b/bot/management/commands/discord_bot.py
@@ -20,4 +20,4 @@ def handle(self, *args: typing.Any, **options: typing.Any) -> typing.NoReturn:
                 asyncio.run(bot.DiscordBot().run())
             except db.OperationalError as e:
                 logging.warning(f'DB Connection expired, reconnecting... {str(e)}')
-                utils.recover_from_db_error()
+                utils.recover_from_db_error(e)
diff --git a/bot/management/commands/purge_posts.py b/bot/management/commands/purge_posts.py
index e003bae..36f93ee 100644
--- a/bot/management/commands/purge_posts.py
+++ b/bot/management/commands/purge_posts.py
@@ -52,7 +52,7 @@ def handle(self, *args: typing.Any, **options: typing.Any) -> typing.NoReturn:
                 ).delete()
             except db.OperationalError as e:
                 logging.warning(f'DB Connection expired, reconnecting... {str(e)}')
-                utils.recover_from_db_error()
+                utils.recover_from_db_error(e)
                 time.sleep(1)
                 continue