← Back to branch summary

~tribaal/charms/trusty/nova-compute/enable-api-rate-limiting

« back to all changes in this revision

Viewing changes to hooks/charmhelpers/core/unitdata.py

  • Committer: Edward Hope-Morley
  • Date: 2015-03-16 15:38:04 UTC
  • mto: This revision was merged to the branch mainline in revision 108.
  • Revision ID: edward.hope-morley@canonical.com-20150316153804-du1szvv0gwfve4g6
cleanup

Show diffs side-by-side

added added

removed removed

Lines of Context:
hooks/charmhelpers/core/unitdata.py
152
152
import collections
153
153
import contextlib
154
154
import datetime
155
 
import itertools
156
155
import json
157
156
import os
158
157
import pprint
165
164
class Storage(object):
166
165
    """Simple key value database for local unit state within charms.
167
166
 
168
 
    Modifications are not persisted unless :meth:`flush` is called.
 
167
    Modifications are automatically committed at hook exit. That's
 
168
    currently regardless of exit code.
169
169
 
170
170
    To support dicts, lists, integer, floats, and booleans values
171
171
    are automatically json encoded/decoded.
173
173
    def __init__(self, path=None):
174
174
        self.db_path = path
175
175
        if path is None:
176
 
            if 'UNIT_STATE_DB' in os.environ:
177
 
                self.db_path = os.environ['UNIT_STATE_DB']
178
 
            else:
179
 
                self.db_path = os.path.join(
180
 
                    os.environ.get('CHARM_DIR', ''), '.unit-state.db')
 
176
            self.db_path = os.path.join(
 
177
                os.environ.get('CHARM_DIR', ''), '.unit-state.db')
181
178
        self.conn = sqlite3.connect('%s' % self.db_path)
182
179
        self.cursor = self.conn.cursor()
183
180
        self.revision = None
192
189
        self.conn.close()
193
190
        self._closed = True
194
191
 
 
192
    def _scoped_query(self, stmt, params=None):
 
193
        if params is None:
 
194
            params = []
 
195
        return stmt, params
 
196
 
195
197
    def get(self, key, default=None, record=False):
196
 
        self.cursor.execute('select data from kv where key=?', [key])
 
198
        self.cursor.execute(
 
199
            *self._scoped_query(
 
200
                'select data from kv where key=?', [key]))
197
201
        result = self.cursor.fetchone()
198
202
        if not result:
199
203
            return default
202
206
        return json.loads(result[0])
203
207
 
204
208
    def getrange(self, key_prefix, strip=False):
205
 
        """
206
 
        Get a range of keys starting with a common prefix as a mapping of
207
 
        keys to values.
208
 
 
209
 
        :param str key_prefix: Common prefix among all keys
210
 
        :param bool strip: Optionally strip the common prefix from the key
211
 
            names in the returned dict
212
 
        :return dict: A (possibly empty) dict of key-value mappings
213
 
        """
214
 
        self.cursor.execute("select key, data from kv where key like ?",
215
 
                            ['%s%%' % key_prefix])
 
209
        stmt = "select key, data from kv where key like '%s%%'" % key_prefix
 
210
        self.cursor.execute(*self._scoped_query(stmt))
216
211
        result = self.cursor.fetchall()
217
212
 
218
213
        if not result:
219
 
            return {}
 
214
            return None
220
215
        if not strip:
221
216
            key_prefix = ''
222
217
        return dict([
223
218
            (k[len(key_prefix):], json.loads(v)) for k, v in result])
224
219
 
225
220
    def update(self, mapping, prefix=""):
226
 
        """
227
 
        Set the values of multiple keys at once.
228
 
 
229
 
        :param dict mapping: Mapping of keys to values
230
 
        :param str prefix: Optional prefix to apply to all keys in `mapping`
231
 
            before setting
232
 
        """
233
221
        for k, v in mapping.items():
234
222
            self.set("%s%s" % (prefix, k), v)
235
223
 
236
224
    def unset(self, key):
237
 
        """
238
 
        Remove a key from the database entirely.
239
 
        """
240
225
        self.cursor.execute('delete from kv where key=?', [key])
241
226
        if self.revision and self.cursor.rowcount:
242
227
            self.cursor.execute(
243
228
                'insert into kv_revisions values (?, ?, ?)',
244
229
                [key, self.revision, json.dumps('DELETED')])
245
230
 
246
 
    def unsetrange(self, keys=None, prefix=""):
247
 
        """
248
 
        Remove a range of keys starting with a common prefix, from the database
249
 
        entirely.
250
 
 
251
 
        :param list keys: List of keys to remove.
252
 
        :param str prefix: Optional prefix to apply to all keys in ``keys``
253
 
            before removing.
254
 
        """
255
 
        if keys is not None:
256
 
            keys = ['%s%s' % (prefix, key) for key in keys]
257
 
            self.cursor.execute('delete from kv where key in (%s)' % ','.join(['?'] * len(keys)), keys)
258
 
            if self.revision and self.cursor.rowcount:
259
 
                self.cursor.execute(
260
 
                    'insert into kv_revisions values %s' % ','.join(['(?, ?, ?)'] * len(keys)),
261
 
                    list(itertools.chain.from_iterable((key, self.revision, json.dumps('DELETED')) for key in keys)))
262
 
        else:
263
 
            self.cursor.execute('delete from kv where key like ?',
264
 
                                ['%s%%' % prefix])
265
 
            if self.revision and self.cursor.rowcount:
266
 
                self.cursor.execute(
267
 
                    'insert into kv_revisions values (?, ?, ?)',
268
 
                    ['%s%%' % prefix, self.revision, json.dumps('DELETED')])
269
 
 
270
231
    def set(self, key, value):
271
 
        """
272
 
        Set a value in the database.
273
 
 
274
 
        :param str key: Key to set the value for
275
 
        :param value: Any JSON-serializable value to be set
276
 
        """
277
232
        serialized = json.dumps(value)
278
233
 
279
 
        self.cursor.execute('select data from kv where key=?', [key])
 
234
        self.cursor.execute(
 
235
            'select data from kv where key=?', [key])
280
236
        exists = self.cursor.fetchone()
281
237
 
282
238
        # Skip mutations to the same value
487
443
        data = hookenv.execution_environment()
488
444
        self.conf = conf_delta = self.kv.delta(data['conf'], 'config')
489
445
        self.rels = rels_delta = self.kv.delta(data['rels'], 'rels')
490
 
        self.kv.set('env', dict(data['env']))
 
446
        self.kv.set('env', data['env'])
491
447
        self.kv.set('unit', data['unit'])
492
448
        self.kv.set('relid', data.get('relid'))
493
449
        return conf_delta, rels_delta