summaryrefslogtreecommitdiff
path: root/puppet/modules/couchdb/files/couch-doc-update
blob: a137e7ff9a4dbf57ed2a8d30db21bf7a3a9aca5e (plain)
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
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
#!/usr/bin/ruby
require 'syslog'

#
# This script will delete or update the values of a particular couchdb document. The benefit of this little script over
# using a simple curl command for updating a document is this:
#
#   * exit non-zero status if document was not updated.
#   * updates existing documents easily, taking care of the _rev id for you.
#   * if document doesn't exist, it is created
#
# REQUIREMENTS
#
#   gem 'couchrest'
#
# USAGE
#
#   see the ouput of
#
#     couch-doc-update
#
#   the content of <file> will be merged with the data provided.
#   If you only want the file content use --data '{}'
#
# EXAMPLE
#
#   create a new user:
#     couch-doc-update --db _users --id org.couchdb.user:ca_daemon --data '{"type": "user", "name": "ca_daemon", "roles": ["certs"], "password": "sshhhh"}'
#
#   update a user:
#     couch-doc-update --db _users --id org.couchdb.user:ca_daemon --data '{"password":"sssshhh"}'
#
#   To update the _users DB on bigcouch, you must connect to port 5986 instead of the default couchdb port 5984
#
#   delete a doc:
#     couch-doc-update --delete --db invite_codes --id dfaf0ee65670c16d5a9161dc86f3bff8
#

begin; require 'rubygems'; rescue LoadError; end # optionally load rubygems
require 'couchrest'

def main
  db, id, data, delete = process_options

  result = if delete
    delete_document(db, id)
  else
    set_document(db, id, data)
  end

  exit 0 if result['ok']
  raise StandardError.new(result.inspect)
rescue StandardError => exc
  db_without_password = db.to_s.sub(/:[^\/]*@/, ':PASSWORD_HIDDEN@')
  indent = "       "
  log "ERROR: " + exc.to_s
  log indent + $@[0..4].join("\n#{indent}")
  log indent + "Failed writing to #{db_without_password}/#{id}"
  exit 1
end

def log(message)
  $stderr.puts message
  Syslog.open('couch-doc-update') do |logger|
    logger.log(Syslog::LOG_CRIT, message)
  end
end

def process_options
  #
  # parse options
  #
  host       = nil
  db_name    = nil
  doc_id     = nil
  new_data   = nil
  filename   = nil
  netrc_file = nil
  delete     = false
  loop do
    case ARGV[0]
      when '--host' then ARGV.shift; host     = ARGV.shift
      when '--db'   then ARGV.shift; db_name  = ARGV.shift
      when '--id'   then ARGV.shift; doc_id   = ARGV.shift
      when '--data' then ARGV.shift; new_data = ARGV.shift
      when '--file' then ARGV.shift; filename = ARGV.shift
      when '--netrc-file' then ARGV.shift; netrc_file = ARGV.shift
      when '--delete' then ARGV.shift; delete = true
      when /^-/     then usage("Unknown option: #{ARGV[0].inspect}")
      else break
    end
  end
  usage("Missing required option") unless db_name && doc_id && (new_data || delete)

  unless delete
    new_data = MultiJson.load(new_data)
    new_data.merge!(read_file(filename)) if filename
  end
  db  = CouchRest.database(connection_string(db_name, host, netrc_file))
  return db, doc_id, new_data, delete
end

def read_file(filename)
  data = MultiJson.load( IO.read(filename) )
  # strip off _id and _rev to avoid conflicts
  data.delete_if {|k,v| k.start_with? '_'}
end

  #
  # update document
  #
def set_document(db, id, data)
  attempt ||= 1
  doc = get_document(db, id)
  if doc
    doc.id ||= id
    update_document(db, doc, data)
  else
    create_document(db, id, data)
  end
rescue RestClient::Conflict
  # retry once, reraise if that does not work
  raise if attempt > 1
  attempt += 1
  retry
end

COUCH_RESPONSE_OK = { 'ok' => true }

# Deletes document, if exists, with retry
def delete_document(db, id)
  attempts ||= 1
  doc = get_document(db, id)
  if doc
    db.delete_doc(doc)
  else
    COUCH_RESPONSE_OK
  end
rescue RestClient::ExceptionWithResponse => e
  if attempts < 6 && !e.response.nil? && RETRY_CODES.include?(e.response.code)
    attempts += 1
    sleep 10
    retry
  else
    raise e
  end
end

def get_document(db, doc_id)
  begin
    db.get(doc_id)
  rescue RestClient::ResourceNotFound
    nil
  end
end

# if the response status code is one of these
# then retry instead of failing.
RETRY_CODES = [500, 422].freeze

def update_document(db, doc, data)
  attempts ||= 1
  doc.reject! {|k,v| !["_id", "_rev"].include? k}
  doc.merge! data
  db.save_doc(doc)
rescue RestClient::ExceptionWithResponse => e
  if attempts < 6 && !e.response.nil? && RETRY_CODES.include?(e.response.code)
    attempts += 1
    sleep 10
    retry
  else
    raise e
  end
end

def create_document(db, doc_id, data)
  attempts ||= 1
  data["_id"] = doc_id
  db.save_doc(data)
rescue RestClient::ExceptionWithResponse => e
  if attempts < 6 && !e.response.nil? && RETRY_CODES.include?(e.response.code)
    attempts += 1
    sleep 10
    retry
  else
    raise e
  end
end

def connection_string(database, host, netrc_file = nil)
  protocol  = "http"
  #hostname  = "127.0.0.1"
  port      = "5984"
  username  = "admin"
  password  = ""

  netrc = File.read(netrc_file || '/etc/couchdb/couchdb.netrc')
  netrc.scan(/\w+ [\w\.]+/).each do |key_value|
    key, value = key_value.split ' '
    case key
      when "machine"  then host ||= value + ':' + port
      when "login"    then username = value
      when "password" then password = value
    end
  end

  host ||= '127.0.0.1:5984'

  "%s://%s:%s@%s/%s" % [protocol, username, password, host, database]
end

def usage(s)
  $stderr.puts(s)
  $stderr.puts("Usage: #{File.basename($0)} --host <host> --db <db> --id <doc_id> --data <json> [--file <file>] [--netrc-file <netrc-file>]")
  $stderr.puts("       #{File.basename($0)} --host <host> --db <db> --id <doc_id> --delete [--netrc-file <netrc-file>]")
  exit(2)
end

main()