Merge branch '8189-handle-large-collections-better' refs #8189
[arvados.git] / apps / workbench / app / models / arvados_resource_list.rb
1 class ArvadosResourceList
2   include ArvadosApiClientHelper
3   include Enumerable
4
5   attr_reader :resource_class
6
7   def initialize resource_class=nil
8     @resource_class = resource_class
9     @fetch_multiple_pages = true
10     @arvados_api_token = Thread.current[:arvados_api_token]
11     @reader_tokens = Thread.current[:reader_tokens]
12   end
13
14   def eager(bool=true)
15     @eager = bool
16     self
17   end
18
19   def distinct(bool=true)
20     @distinct = bool
21     self
22   end
23
24   def limit(max_results)
25     if not max_results.nil? and not max_results.is_a? Integer
26       raise ArgumentError("argument to limit() must be an Integer or nil")
27     end
28     @limit = max_results
29     self
30   end
31
32   def offset(skip)
33     @offset = skip
34     self
35   end
36
37   def order(orderby_spec)
38     @orderby_spec = orderby_spec
39     self
40   end
41
42   def select(columns=nil)
43     # If no column arguments were given, invoke Enumerable#select.
44     if columns.nil?
45       super()
46     else
47       @select ||= []
48       @select += columns
49       self
50     end
51   end
52
53   def filter _filters
54     @filters ||= []
55     @filters += _filters
56     self
57   end
58
59   def where(cond)
60     @cond = cond.dup
61     @cond.keys.each do |uuid_key|
62       if @cond[uuid_key] and (@cond[uuid_key].is_a? Array or
63                              @cond[uuid_key].is_a? ArvadosBase)
64         # Coerce cond[uuid_key] to an array of uuid strings.  This
65         # allows caller the convenience of passing an array of real
66         # objects and uuids in cond[uuid_key].
67         if !@cond[uuid_key].is_a? Array
68           @cond[uuid_key] = [@cond[uuid_key]]
69         end
70         @cond[uuid_key] = @cond[uuid_key].collect do |item|
71           if item.is_a? ArvadosBase
72             item.uuid
73           else
74             item
75           end
76         end
77       end
78     end
79     @cond.keys.select { |x| x.match /_kind$/ }.each do |kind_key|
80       if @cond[kind_key].is_a? Class
81         @cond = @cond.merge({ kind_key => 'arvados#' + arvados_api_client.class_kind(@cond[kind_key]) })
82       end
83     end
84     self
85   end
86
87   def fetch_multiple_pages(f)
88     @fetch_multiple_pages = f
89     self
90   end
91
92   def results
93     if !@results
94       @results = []
95       self.each_page do |r|
96         @results.concat r
97       end
98     end
99     @results
100   end
101
102   def results=(r)
103     @results = r
104     @items_available = r.items_available if r.respond_to? :items_available
105     @result_limit = r.limit if r.respond_to? :limit
106     @result_offset = r.offset if r.respond_to? :offset
107     @results
108   end
109
110   def to_ary
111     results
112   end
113
114   def each(&block)
115     if not @results.nil?
116       @results.each &block
117     else
118       self.each_page do |items|
119         items.each do |i|
120           block.call i
121         end
122       end
123     end
124     self
125   end
126
127   def first
128     results.first
129   end
130
131   def last
132     results.last
133   end
134
135   def [](*x)
136     results.send('[]', *x)
137   end
138
139   def |(x)
140     if x.is_a? Hash
141       self.to_hash | x
142     else
143       results | x.to_ary
144     end
145   end
146
147   def to_hash
148     Hash[self.collect { |x| [x.uuid, x] }]
149   end
150
151   def empty?
152     self.first.nil?
153   end
154
155   def items_available
156     results
157     @items_available
158   end
159
160   def result_limit
161     results
162     @result_limit
163   end
164
165   def result_offset
166     results
167     @result_offset
168   end
169
170   # Obsolete method retained during api transition.
171   def links_for item_or_uuid, link_class=false
172     []
173   end
174
175   protected
176
177   def each_page
178     api_params = {
179       _method: 'GET'
180     }
181     api_params[:where] = @cond if @cond
182     api_params[:eager] = '1' if @eager
183     api_params[:select] = @select if @select
184     api_params[:order] = @orderby_spec if @orderby_spec
185     api_params[:filters] = @filters if @filters
186     api_params[:distinct] = @distinct if @distinct
187     if @fetch_multiple_pages
188       # Default limit to (effectively) api server's MAX_LIMIT
189       api_params[:limit] = 2**(0.size*8 - 1) - 1
190     end
191
192     item_count = 0
193     offset = @offset || 0
194     @result_limit = nil
195     @result_offset = nil
196
197     begin
198       api_params[:offset] = offset
199       api_params[:limit] = (@limit - item_count) if @limit
200
201       res = arvados_api_client.api(@resource_class, '', api_params,
202                                    arvados_api_token: @arvados_api_token,
203                                    reader_tokens: @reader_tokens)
204       items = arvados_api_client.unpack_api_response res
205
206       @items_available = items.items_available if items.respond_to?(:items_available)
207       @result_limit = items.limit if (@fetch_multiple_pages == false) and items.respond_to?(:limit)
208       @result_offset = items.offset if (@fetch_multiple_pages == false) and items.respond_to?(:offset)
209
210       break if items.nil? or not items.any?
211
212       item_count += items.size
213       if items.respond_to?(:offset)
214         offset = items.offset + items.size
215       else
216         offset = item_count
217       end
218
219       yield items
220
221       break if @limit and item_count >= @limit
222       break if items.respond_to? :items_available and offset >= items.items_available
223     end while @fetch_multiple_pages
224     self
225   end
226
227 end