~junaidali/charms/trusty/plumgrid-edge/docker-oil

« back to all changes in this revision

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

  • Committer: bbaqar at plumgrid
  • Date: 2016-04-25 09:18:40 UTC
  • mfrom: (26.1.3 plumgrid-edge)
  • Revision ID: bbaqar@plumgrid.com-20160425091840-sirw6bbzalts677s
Merge: Liberty/Mitaka support

Show diffs side-by-side

added added

removed removed

Lines of Context:
152
152
import collections
153
153
import contextlib
154
154
import datetime
 
155
import itertools
155
156
import json
156
157
import os
157
158
import pprint
164
165
class Storage(object):
165
166
    """Simple key value database for local unit state within charms.
166
167
 
167
 
    Modifications are automatically committed at hook exit. That's
168
 
    currently regardless of exit code.
 
168
    Modifications are not persisted unless :meth:`flush` is called.
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
 
            self.db_path = os.path.join(
177
 
                os.environ.get('CHARM_DIR', ''), '.unit-state.db')
 
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')
178
181
        self.conn = sqlite3.connect('%s' % self.db_path)
179
182
        self.cursor = self.conn.cursor()
180
183
        self.revision = None
189
192
        self.conn.close()
190
193
        self._closed = True
191
194
 
192
 
    def _scoped_query(self, stmt, params=None):
193
 
        if params is None:
194
 
            params = []
195
 
        return stmt, params
196
 
 
197
195
    def get(self, key, default=None, record=False):
198
 
        self.cursor.execute(
199
 
            *self._scoped_query(
200
 
                'select data from kv where key=?', [key]))
 
196
        self.cursor.execute('select data from kv where key=?', [key])
201
197
        result = self.cursor.fetchone()
202
198
        if not result:
203
199
            return default
206
202
        return json.loads(result[0])
207
203
 
208
204
    def getrange(self, key_prefix, strip=False):
209
 
        stmt = "select key, data from kv where key like '%s%%'" % key_prefix
210
 
        self.cursor.execute(*self._scoped_query(stmt))
 
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])
211
216
        result = self.cursor.fetchall()
212
217
 
213
218
        if not result:
214
 
            return None
 
219
            return {}
215
220
        if not strip:
216
221
            key_prefix = ''
217
222
        return dict([
218
223
            (k[len(key_prefix):], json.loads(v)) for k, v in result])
219
224
 
220
225
    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
        """
221
233
        for k, v in mapping.items():
222
234
            self.set("%s%s" % (prefix, k), v)
223
235
 
224
236
    def unset(self, key):
 
237
        """
 
238
        Remove a key from the database entirely.
 
239
        """
225
240
        self.cursor.execute('delete from kv where key=?', [key])
226
241
        if self.revision and self.cursor.rowcount:
227
242
            self.cursor.execute(
228
243
                'insert into kv_revisions values (?, ?, ?)',
229
244
                [key, self.revision, json.dumps('DELETED')])
230
245
 
 
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
 
231
270
    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
        """
232
277
        serialized = json.dumps(value)
233
278
 
234
 
        self.cursor.execute(
235
 
            'select data from kv where key=?', [key])
 
279
        self.cursor.execute('select data from kv where key=?', [key])
236
280
        exists = self.cursor.fetchone()
237
281
 
238
282
        # Skip mutations to the same value