@@ -3,7 +3,8 @@ abstract class Avram::Database
33
44 @@db : DB ::Database ? = nil
55 @@lock = Mutex .new
6- class_getter transactions = {} of FiberId => DB ::Transaction
6+ class_getter connections = {} of FiberId => DB ::Connection
7+ class_property lock_id : UInt64 ?
78
89 macro inherited
910 Habitat .create do
@@ -24,6 +25,16 @@ abstract class Avram::Database
2425 % }
2526 end
2627
28+ def self.setup_connection (& block : DB ::Connection - > Nil )
29+ new.db.setup_connection do |conn |
30+ block.call conn
31+ end
32+ end
33+
34+ def self.verify_connection
35+ new.connection.open.close
36+ end
37+
2738 # Rollback the current transaction
2839 def self.rollback
2940 new.rollback
@@ -142,28 +153,36 @@ abstract class Avram::Database
142153
143154 # :nodoc:
144155 def run
145- yield current_transaction.try( & .connection) || db
156+ yield current_connection || db
146157 end
147158
148159 # :nodoc:
149160 def listen (* channels : String , & block : PQ ::Notification - > ) : Nil
150161 connection.connect_listen(* channels, & block)
151162 end
152163
153- private def connection : Avram ::Connection
164+ protected def connection : Avram ::Connection
154165 Avram ::Connection .new(url, database_class: self .class)
155166 end
156167
157- private def db : DB ::Database
168+ protected def db : DB ::Database
158169 @@db ||= @@lock .synchronize do
159170 # check @@db again because a previous request could have set it after
160171 # the first time it was checked
161172 @@db || connection.open
162173 end
163174 end
164175
176+ private def current_connection : DB ::Connection
177+ connections[object_id] ||= db.checkout
178+ end
179+
180+ private def object_id : UInt64
181+ self .class.lock_id || Fiber .current.object_id
182+ end
183+
165184 private def current_transaction : DB ::Transaction ?
166- transactions[ Fiber .current.object_id] ?
185+ current_connection._avram_stack.last ?
167186 end
168187
169188 protected def truncate
@@ -180,7 +199,7 @@ abstract class Avram::Database
180199
181200 # :nodoc:
182201 def transaction : Bool
183- if current_transaction
202+ if current_transaction.try( & ._avram_joinable?)
184203 yield
185204 true
186205 else
@@ -190,20 +209,23 @@ abstract class Avram::Database
190209 end
191210 end
192211
193- private def transactions
194- self .class.transactions
212+ private def connections
213+ self .class.connections
195214 end
196215
197216 private def wrap_in_transaction
198- db.transaction do |tx |
199- transactions[Fiber .current.object_id] ||= tx
217+ (current_transaction || current_connection).transaction do
200218 yield
201219 end
202220 true
203221 rescue e : Avram ::Rollback
204222 false
205223 ensure
206- transactions.delete(Fiber .current.object_id)
224+ # TODO: not sure of this
225+ if current_connection._avram_in_transaction?
226+ current_connection.release
227+ connections.delete(object_id)
228+ end
207229 end
208230
209231 class DatabaseCleaner
0 commit comments