From f2a0975724f586badef24333ad5bc62a59fdfc5e Mon Sep 17 00:00:00 2001
From: n atta kusi adusei <attakusiadusei@gmail.com>
Date: Wed, 27 Dec 2023 22:11:03 +0000
Subject: [PATCH] Retry lost connections (#990)

Fixes `DB::ConnectionLost` errors.
---
 src/avram/database.cr | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/avram/database.cr b/src/avram/database.cr
index d8e85b7e8..c843b26ea 100644
--- a/src/avram/database.cr
+++ b/src/avram/database.cr
@@ -187,15 +187,18 @@ abstract class Avram::Database
   # once the block is finished
   private def with_connection(&)
     key = object_id
-    connections[key] ||= db.checkout
-    connection = connections[key]
-
-    begin
-      yield connection
-    ensure
-      if !connection._avram_in_transaction?
-        connection.release
-        connections.delete(key)
+
+    db.retry do
+      connections[key] ||= db.checkout
+      connection = connections[key]
+
+      begin
+        yield connection
+      ensure
+        if !connection._avram_in_transaction?
+          connection.release
+          connections.delete(key)
+        end
       end
     end
   end