@@ -153,7 +153,9 @@ abstract class Avram::Database
153
153
154
154
# :nodoc:
155
155
def run
156
- yield current_connection || db
156
+ with_connection do |conn |
157
+ yield conn
158
+ end
157
159
end
158
160
159
161
# :nodoc:
@@ -173,16 +175,31 @@ abstract class Avram::Database
173
175
end
174
176
end
175
177
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
178
195
end
179
196
180
197
private def object_id : UInt64
181
198
self .class.lock_id || Fiber .current.object_id
182
199
end
183
200
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?
186
203
end
187
204
188
205
protected def truncate
@@ -199,12 +216,14 @@ abstract class Avram::Database
199
216
200
217
# :nodoc:
201
218
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?)
207
221
yield
222
+ true
223
+ else
224
+ wrap_in_transaction(conn) do
225
+ yield
226
+ end
208
227
end
209
228
end
210
229
end
@@ -213,18 +232,13 @@ abstract class Avram::Database
213
232
self .class.connections
214
233
end
215
234
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
218
237
yield
219
238
end
220
239
true
221
240
rescue e : Avram ::Rollback
222
241
false
223
- ensure
224
- if ! current_connection._avram_in_transaction?
225
- current_connection.release
226
- connections.delete(object_id)
227
- end
228
242
end
229
243
230
244
class DatabaseCleaner
0 commit comments