@@ -153,7 +153,9 @@ abstract class Avram::Database
153153
154154 # :nodoc:
155155 def run
156- yield current_connection || db
156+ with_connection do |conn |
157+ yield conn
158+ end
157159 end
158160
159161 # :nodoc:
@@ -173,16 +175,31 @@ abstract class Avram::Database
173175 end
174176 end
175177
176- private def current_connection : DB ::Connection
177- connections[object_id] ||= db.checkout
178+ # singular place to retrieve a DB::Connection
179+ # must be passed a block and we
180+ # try to release the connection back to the pool
181+ # once the block is finished
182+ private def with_connection
183+ key = object_id
184+ connections[key] ||= db.checkout
185+ connection = connections[key]
186+
187+ begin
188+ yield connection
189+ ensure
190+ if ! connection._avram_in_transaction?
191+ connection.release
192+ connections.delete(key)
193+ end
194+ end
178195 end
179196
180197 private def object_id : UInt64
181198 self .class.lock_id || Fiber .current.object_id
182199 end
183200
184- private def current_transaction : DB ::Transaction ?
185- current_connection ._avram_stack.last?
201+ private def current_transaction ( connection : DB :: Connection ) : DB ::Transaction ?
202+ connection ._avram_stack.last?
186203 end
187204
188205 protected def truncate
@@ -199,12 +216,14 @@ abstract class Avram::Database
199216
200217 # :nodoc:
201218 def transaction : Bool
202- if current_transaction.try(& ._avram_joinable?)
203- yield
204- true
205- else
206- wrap_in_transaction do
219+ with_connection do |conn |
220+ if current_transaction(conn).try(& ._avram_joinable?)
207221 yield
222+ true
223+ else
224+ wrap_in_transaction(conn) do
225+ yield
226+ end
208227 end
209228 end
210229 end
@@ -213,18 +232,13 @@ abstract class Avram::Database
213232 self .class.connections
214233 end
215234
216- private def wrap_in_transaction
217- (current_transaction || current_connection ).transaction do
235+ private def wrap_in_transaction ( conn )
236+ (current_transaction(conn) || conn ).transaction do
218237 yield
219238 end
220239 true
221240 rescue e : Avram ::Rollback
222241 false
223- ensure
224- if ! current_connection._avram_in_transaction?
225- current_connection.release
226- connections.delete(object_id)
227- end
228242 end
229243
230244 class DatabaseCleaner
0 commit comments