1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
|
# frozen-string-literal: true
#
# The query_blocker extension adds Database#block_queries.
# Inside the block passed to #block_queries, any attempts to
# execute a query/statement on the database will raise a
# Sequel::QueryBlocker::BlockedQuery exception.
#
# DB.extension :query_blocker
# DB.block_queries do
# ds = DB[:table] # No exception
# ds = ds.where(column: 1) # No exception
# ds.all # Exception raised
# end
#
# To handle concurrency, you can pass a :scope option:
#
# # Current Thread
# DB.block_queries(scope: :thread){}
#
# # Current Fiber
# DB.block_queries(scope: :fiber){}
#
# # Specific Thread
# DB.block_queries(scope: Thread.current){}
#
# # Specific Fiber
# DB.block_queries(scope: Fiber.current){}
#
# Database#block_queries is useful for blocking queries inside
# the block. However, there may be cases where you want to
# allow queries in specific places inside a block_queries block.
# You can use Database#allow_queries for that:
#
# DB.block_queries do
# DB.allow_queries do
# DB[:table].all # Query allowed
# end
#
# DB[:table].all # Exception raised
# end
#
# When mixing block_queries and allow_queries with scopes, the
# narrowest scope has priority. So if you are blocking with
# :thread scope, and allowing with :fiber scope, queries in the
# current fiber will be allowed, but queries in different fibers of
# the current thread will be blocked.
#
# Note that this should catch all queries executed through the
# Database instance. Whether it catches queries executed directly
# on a connection object depends on the adapter in use.
#
# Related module: Sequel::QueryBlocker
require "fiber"
#
module Sequel
module QueryBlocker
# Exception class raised if there is an attempt to execute a
# query/statement on the database inside a block passed to
# block_queries.
class BlockedQuery < Sequel::Error
end
def self.extended(db)
db.instance_exec do
@blocked_query_scopes ||= {}
end
end
# If checking a connection for validity, and a BlockedQuery exception is
# raised, treat it as a valid connection. You cannot check whether the
# connection is valid without issuing a query, and if queries are blocked,
# you need to assume it is valid or assume it is not. Since it most cases
# it will be valid, this assumes validity.
def valid_connection?(conn)
super
rescue BlockedQuery
true
end
# Check whether queries are blocked before executing them.
def log_connection_yield(sql, conn, args=nil)
# All database adapters should be calling this method around
# query execution (otherwise the queries would not get logged),
# ensuring the blocking is checked. Any database adapter issuing
# a query without calling this method is considered buggy.
check_blocked_queries!
super
end
# Whether queries are currently blocked.
def block_queries?
b = @blocked_query_scopes
b.fetch(Fiber.current) do
b.fetch(Thread.current) do
b.fetch(:global, false)
end
end
end
# Allow queries inside the block. Only useful if they are already blocked
# for the same scope. Useful for blocking queries generally, and only allowing
# them in specific places. Takes the same :scope option as #block_queries.
def allow_queries(opts=OPTS, &block)
_allow_or_block_queries(false, opts, &block)
end
# Reject (raise an BlockedQuery exception) if there is an attempt to execute
# a query/statement inside the block.
#
# The :scope option indicates which queries are rejected inside the block:
#
# :global :: This is the default, and rejects all queries.
# :thread :: Reject all queries in the current thread.
# :fiber :: Reject all queries in the current fiber.
# Thread :: Reject all queries in the given thread.
# Fiber :: Reject all queries in the given fiber.
def block_queries(opts=OPTS, &block)
_allow_or_block_queries(true, opts, &block)
end
private
# Internals of block_queries and allow_queries.
def _allow_or_block_queries(value, opts)
scope = query_blocker_scope(opts)
prev_value = nil
scopes = @blocked_query_scopes
begin
Sequel.synchronize do
prev_value = scopes[scope]
scopes[scope] = value
end
yield
ensure
Sequel.synchronize do
if prev_value.nil?
scopes.delete(scope)
else
scopes[scope] = prev_value
end
end
end
end
# The scope for the query block, either :global, or a Thread or Fiber instance.
def query_blocker_scope(opts)
case scope = opts[:scope]
when nil
:global
when :global, Thread, Fiber
scope
when :thread
Thread.current
when :fiber
Fiber.current
else
raise Sequel::Error, "invalid scope given to block_queries: #{scope.inspect}"
end
end
# Raise a BlockQuery exception if queries are currently blocked.
def check_blocked_queries!
raise BlockedQuery, "cannot execute query inside a block_queries block" if block_queries?
end
end
Database.register_extension(:query_blocker, QueryBlocker)
end
|
