diff options
-rw-r--r-- | custodia/store/interface.py | 2 | ||||
-rw-r--r-- | custodia/store/sqlite.py | 17 |
2 files changed, 16 insertions, 3 deletions
diff --git a/custodia/store/interface.py b/custodia/store/interface.py index c804670..6ca3ea4 100644 --- a/custodia/store/interface.py +++ b/custodia/store/interface.py @@ -10,7 +10,7 @@ class CSStore(object): def get(self, key): raise NotImplementedError - def set(self, key, value): + def set(self, key, value, replace=False): raise NotImplementedError def list(self, keyfilter=None): diff --git a/custodia/store/sqlite.py b/custodia/store/sqlite.py index 5f33a2a..0ed8223 100644 --- a/custodia/store/sqlite.py +++ b/custodia/store/sqlite.py @@ -42,8 +42,12 @@ class SqliteStore(CSStore): "(key PRIMARY KEY UNIQUE, value)" % self.table cur.execute(create) - def set(self, key, value): - setdata = "INSERT OR REPLACE into %s VALUES (?, ?)" % self.table + def set(self, key, value, replace=False): + if replace: + query = "INSERT OR REPLACE into %s VALUES (?, ?)" + else: + query = "INSERT into %s VALUES (?, ?)" + setdata = query % (self.table,) try: conn = sqlite3.connect(self.dburi) with conn: @@ -160,3 +164,12 @@ class SqliteStoreTests(unittest.TestCase): value = self.store.get('/sub%') self.assertEqual(value, None) + + def test_5_set_replace(self): + with self.assertRaises(CSStoreError): + self.store.set('key', 'replaced') + + self.store.set('key', 'replaced', replace=True) + + value = self.store.get('key') + self.assertEqual(value, 'replaced') |