Class: Dataflow::Adapters::MongoDbAdapter

Inherits:
Object
  • Object
show all
Defined in:
lib/dataflow/adapters/mongo_db_adapter.rb

Overview

Interface between a data node and mongodb. We use mongodb to perform all the store/retrieve operations.

Constant Summary collapse

SYSTEM_ID = 
'_id'

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(args) ⇒ MongoDbAdapter

Returns a new instance of MongoDbAdapter.



35
36
37
38
39
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 35

def initialize(args)
  update_settings(args)
  @client = MongoDbAdapter.client(settings)
  @admin_client = MongoDbAdapter.admin_client(settings)
end

Instance Attribute Details

#client ⇒ Object (readonly)

Returns the value of attribute client.



33
34
35
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 33

def client
  @client
end

#settings ⇒ Object (readonly)

Returns the value of attribute settings.



32
33
34
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 32

def settings
  @settings
end

Class Method Details

.admin_client(settings) ⇒ Object 



19
20
21
22
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 19

def admin_client(settings)
  return @admin_client if @admin_client
  @admin_client = client(settings, db_name: 'admin')
end

.client(settings, db_name: nil) ⇒ Object 



10
11
12
13
14
15
16
17
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 10

def client(settings, db_name: nil)
  @clients ||= {}

  settings.adapter_type = 'mongodb'
  connection_uri = settings.connection_uri_or_default
  db_name ||= settings.db_name
  @clients["#{connection_uri}.#{db_name}"] ||= Mongo::Client.new([connection_uri], database: db_name)
end

.disconnect_clients ⇒ Object

Force the clients to disconnect their connections. Use before forking.



26
27
28
29
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 26

def disconnect_clients
  @clients ||= {}
  @clients.values.each(&:close)
end

Instance Method Details

#all(where: {}, fields: [], sort: {}, offset: 0, limit: 0) ⇒ Object

retrieve all elements from a data node



51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 51

def all(where: {}, fields: [], sort: {}, offset: 0, limit: 0)
  projection = fields.map { |field| [field, 1] }

  unless fields.map(&:to_s).include?(SYSTEM_ID)
    # by default, do not select the _id field
    projection << [SYSTEM_ID, 0].freeze
  end

  opts = transform_to_query(where)
  res = client[read_dataset_name].find(opts)
  res = res.projection(projection.to_h)

  res = res.sort(sort)   if sort
  res = res.skip(offset) if offset > 0
  res = res.limit(limit) if limit > 0

  if block_given?
    yield res
  else
    res.to_a
  end
end

#all_paginated(where: {}, fields: [], cursor: nil) ⇒ Hash

Helper that supports paginating through the whole dataset at fixed performance. Unlike using offset/skip which requires to read through the skipped content (high usage of CPU), we use the internal mongo cursor to get batch of results.

Returns:

  • (Hash) —

    with 2 fields: data and next_cursor for the next call



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
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 79

def all_paginated(where: {}, fields: [], cursor: nil)
  cursor = cursor.to_i
  data = []

  # If there is no cursor, we make the initial query
  # get the first batch of data and get the cursor id.
  if cursor.zero?
    all(where: where, fields: fields) do |res|
      results = res.initial_query
      data = results.documents
      cursor = res.cursor.id
    end
  end

  # The first query's result batch is a small 101 set of results
  # so we want to get one more batch of data.
  # However, there might be queries whose results are very small
  # and the resulting cursor is 0. In such case there is no more
  # data to be fetched.
  unless cursor.zero?
    # send a getMore command on the cursor id
    command = { getMore: cursor, collection: read_dataset_name }
    result = client.database.command(command).documents[0]
    cursor = result['cursor']['id']
    data += result['cursor']['nextBatch']
  end

  # We want to return the cursor as a string.
  # If there is no cursor (zero) then make it empty
  cursor = '' if cursor.zero?

  { 'data' => data, 'next_cursor' => cursor.to_s }
rescue Mongo::Error::OperationFailure
  { 'data' => data, 'next_cursor' => '' }
end

#count(where: {}) ⇒ Object

count the number of records



133
134
135
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 133

def count(where: {})
  client[read_dataset_name].count(transform_to_query(where))
end

#create_indexes(dataset: nil, type: :all, drop_retry_on_error: true) ⇒ Object

Create the indexes on this dataset.

Parameters:

  • dataset (String) (defaults to: nil) —

    Specify on which dataset the operation will be performed. Default: the adatpter’s settings’ dataset.

  • type (Symbol) (defaults to: :all) —

    select which indexes type to create. Can be :all (default), :unique_only, :non_unique_only



190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 190

def create_indexes(dataset: nil, type: :all, drop_retry_on_error: true)
  dataset ||= write_dataset_name
  return unless settings.indexes.present?

  indexes = (settings.indexes || [])

  case type
  when :unique_only
    indexes = indexes.select { |idx| idx['unique'] }
  when :non_unique_only
    indexes = indexes.reject { |idx| idx['unique'] }
  end

  indexes = indexes.map { |x| format_index(x) }
  client[dataset].indexes.create_many(indexes)
rescue Mongo::Error::OperationFailure => e
  raise e unless drop_retry_on_error
  client[dataset].indexes.drop_all
  create_indexes(drop_retry_on_error: false)
end

#delete(where: {}) ⇒ Object

Delete records that match the options.

Parameters:

  • where (defaults to: {}) —

    query to apply on the delete operation.



168
169
170
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 168

def delete(where: {})
  client[read_dataset_name].delete_many(transform_to_query(where))
end

#drop_dataset(dataset) ⇒ Object 



180
181
182
183
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 180

def drop_dataset(dataset)
  collection = client[dataset]
  collection.drop
end

#dump(base_folder:) ⇒ Object 



228
229
230
231
232
233
234
235
236
237
238
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 228

def dump(base_folder:)
  archive_path = "#{base_folder}/#{@settings.db_name}.#{@settings.dataset_name}.gz"
  options = "--archive=#{archive_path} --db=#{@settings.db_name} --collection=#{read_dataset_name}"
  options += "--host=#{@settings.db_host}" if @settings.db_host.present?
  options += "--port=#{@settings.db_port}" if @settings.db_port.present?
  options += "--username=#{@settings.db_user}" if @settings.db_user.present?
  options += "--password=#{@settings.db_password}" if @settings.db_password.present?
  `mkdir -p #{base_folder}`
  `mongodump #{options} --gzip`
  archive_path
end

#find(where: {}, fields: [], sort: {}, offset: 0) ⇒ Object

retrieve a single element from a data node



46
47
48
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 46

def find(where: {}, fields: [], sort: {}, offset: 0)
  all(where: where, fields: fields, sort: sort, offset: offset, limit: 1).first
end

#ordered_system_id_queries(batch_size:, where: {}) ⇒ Object

Create queries that permit processing the whole dataset in parallel without using offsets.



116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 116

def ordered_system_id_queries(batch_size:, where: {})
  ids = all(fields: [SYSTEM_ID], where: where, sort: { SYSTEM_ID => 1 }).map { |x| x[SYSTEM_ID].to_s }
  queries_count = (ids.size / batch_size.to_f).ceil
  Array.new(queries_count) do |i|
    from = ids[i * batch_size]
    to = ids[(i + 1) * batch_size] || ids[-1]
    is_last = i == queries_count - 1

    where_query = { SYSTEM_ID => { '>=' => from } }
    operator = is_last ? '<=' : '<'
    where_query[SYSTEM_ID][operator] = to

    where_query
  end
end

#recreate_dataset(dataset: nil) ⇒ Object

recreate the table/collection



173
174
175
176
177
178
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 173

def recreate_dataset(dataset: nil)
  dataset ||= write_dataset_name
  drop_dataset(dataset)
  collection = client[dataset]
  collection.create
end

#restore(filepath:) ⇒ Object 



240
241
242
243
244
245
246
247
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 240

def restore(filepath:)
  options = "--archive=#{filepath}  --db=#{@settings.db_name} --collection=#{read_dataset_name}"
  options += "--host=#{@settings.db_host}" if @settings.db_host.present?
  options += "--port=#{@settings.db_port}" if @settings.db_port.present?
  options += "--username=#{@settings.db_user}" if @settings.db_user.present?
  options += "--password=#{@settings.db_password}" if @settings.db_password.present?
  `mongorestore #{options} --gzip`
end

#save(records:, replace_by: nil) ⇒ Object

Save the given records.

Parameters:

  • replace_by (Array) (defaults to: nil) —

    if the replace_by key is provided, it will try to replace records with the matching key, or insert if none is found.



141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 141

def save(records:, replace_by: nil)
  if replace_by.present?
    replace_keys = Array(replace_by)
    bulk_ops = records.map do |record|
      filter = replace_keys.map { |x| [x, record[x]] }.to_h
      {
        replace_one: {
          filter: filter,
          replacement: record,
          upsert: true
        }
      }
    end
    client[write_dataset_name].bulk_write(bulk_ops, ordered: false)
  else
    client[write_dataset_name].insert_many(records, ordered: false)
  end
rescue Mongo::Error::BulkWriteError => e
  dup_key_error = e.result['writeErrors'].all? { |x| x['code'] == 11_000 }
  # don't raise if it is errors about duplicated keys
  unless dup_key_error
    raise e
  end
end

#transform_to_query(opts) ⇒ Object 



249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 249

def transform_to_query(opts)
  sanitized_opts = {}
  opts.each do |k, v|
    if v.is_a? Array
      # e.g. { 'id' => [1,2] } transform to mongodb IN clauses
      sanitized_opts[k] = { '$in' => v.map { |value| try_cast_value(k, value) } }
    elsif v.is_a? Hash
      sanitized_opts[k] = {}
      v.each do |operator, value|
        case operator.to_s
        when '!='
          # we still need to check and transform into
          if value.is_a? Array
            # { '$nin' => [value] }
            sanitized_opts[k]['$nin'] = value.map { |x| try_cast_value(k, x) }
          else
            # or {'$ne' => value }
            sanitized_opts[k]['$ne'] = try_cast_value(k, value)
          end
        when '<'
          sanitized_opts[k]['$lt'] = try_cast_value(k, value)
        when '<='
          sanitized_opts[k]['$lte'] = try_cast_value(k, value)
        when '>'
          sanitized_opts[k]['$gt'] = try_cast_value(k, value)
        when '>='
          sanitized_opts[k]['$gte'] = try_cast_value(k, value)
        when '~*' # match regex /regex/i (case insensitive)
          sanitized_opts[k]['$regex'] = /#{value}/i
        when '~'  # match regex /regex/  (case sensitive)
          sanitized_opts[k]['$regex'] = /#{value}/
        end
      end
    else
      sanitized_opts[k] = try_cast_value(k, v)
    end
  end
  sanitized_opts
end

#update_settings(args) ⇒ Object 



41
42
43
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 41

def update_settings(args)
  @settings = Dataflow::Adapters::Settings.new(args)
end

#usage(dataset:) ⇒ Object 



211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
 
# File 'lib/dataflow/adapters/mongo_db_adapter.rb', line 211

def usage(dataset:)
  indexes = retrieve_collection_indexes(dataset)
  command = { collstats: dataset }
  result = client.database.command(command).documents[0]
  {
    memory: result['size'],
    storage: result['storageSize'],
    effective_indexes: indexes
  }
rescue Mongo::Error::OperationFailure, Mongo::Error::InvalidCollectionName
  {
    memory: 0,
    storage: 0,
    effective_indexes: indexes
  }
end