Skip to content

Instantly share code, notes, and snippets.

@iKunalChhabra
Last active June 27, 2023 00:52
Show Gist options
  • Save iKunalChhabra/665955a222e70a92753fb768d78a884b to your computer and use it in GitHub Desktop.
Save iKunalChhabra/665955a222e70a92753fb768d78a884b to your computer and use it in GitHub Desktop.
A Utility to store key-value, configs, secrets, properties with python dict like interface that uses sqlite3 as a backend.
import sqlite3
import os
import json
from cryptography import fernet
class KeyValStore:
"""
KeyValStore is a simple key-value store that uses sqlite3 as a backend.
It is not intended to be used as a production database, but rather as a
simple key-value store for storing configuration data.
The KeyValStore class is a wrapper around sqlite3 that provides a dict-like
interface for storing and retrieving data.
KeyValStore requires keys and values to be of type string.
Example usage:
>>> from data_store import KeyValStore
>>> store = KeyValStore('test_store')
>>> store['foo'] = 'bar'
>>> store['foo']
'bar'
>>> store['foo'] = 'baz'
>>> store['foo']
'baz'
>>> store.close()
>>> with KeyValStore('test_store') as store:
... store['foo'] = 'bar'
... data = store['foo']
>>> data
'bar'
"""
def __init__(self, store_name, store_path='./.data'):
"""
:param store_name: filename of the store
:param store_path: [optional] path to the store. Defaults to './.data'
"""
self.__conn = self.__create_store(store_name, store_path)
self.__cursor = self.__conn.cursor()
self.__table_name = '_sql_data_store'
self.__initialize_store()
def __initialize_store(self):
self.__cursor.execute(f"CREATE TABLE IF NOT EXISTS {self.__table_name} (key text PRIMARY KEY, value text)")
self.__conn.commit()
def __create_store(self, store_name, store_path):
if not os.path.exists(store_path):
os.makedirs(store_path)
self.__store_path = os.path.join(store_path, store_name) + '.db'
return sqlite3.connect(f'{self.__store_path}')
def __validate_key(self, key):
if type(key) is not str:
raise TypeError('Key must be of type string')
def __validate_value(self, value):
if type(value) is not str:
raise TypeError('Value must be of type string')
def __setitem__(self, key,value):
self.__validate_key(key)
self.__validate_value(value)
existing_value = self.get(key)
if existing_value is None:
self.__cursor.execute(f"INSERT INTO {self.__table_name} VALUES (?, ?)", (key, value))
self.__conn.commit()
else:
self.__cursor.execute(f"UPDATE {self.__table_name} SET value = ? WHERE key = ?", (value, key))
self.__conn.commit()
# note: upsert runs slower than select + insert or update
# self.__cursor.execute(f"INSERT OR REPLACE INTO {self.__table_name} VALUES (?, ?)", (key, value))
# self.__conn.commit()
def __getitem__(self, key):
self.__validate_key(key)
self.__cursor.execute(f"SELECT value FROM {self.__table_name} WHERE key = ?", (key,))
result = self.__cursor.fetchone()
if result is None:
raise KeyError(f'Key {key} not found')
return result[0]
def __delitem__(self, key):
self.__cursor.execute(f"DELETE FROM {self.__table_name} WHERE key = ?", (key,))
self.__conn.commit()
def __len__(self):
self.__cursor.execute(f"SELECT COUNT(*) FROM {self.__table_name}")
return self.__cursor.fetchone()[0]
def __iter__(self):
self.__cursor.execute(f"SELECT key FROM {self.__table_name}")
for row in self.__cursor.fetchall():
yield row[0]
def __contains__(self, key):
self.__cursor.execute(f"SELECT COUNT(*) FROM {self.__table_name} WHERE key = ?", (key,))
return self.__cursor.fetchone()[0] > 0
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
self.close()
def __repr__(self):
return f'<{self.__class__.__name__}: {self.__store_path}>'
def __str__(self):
return self.__repr__()
def keys(self):
yield from self.__iter__()
def values(self):
self.__cursor.execute(f"SELECT value FROM {self.__table_name}")
for row in self.__cursor.fetchall():
yield row[0]
def items(self):
self.__cursor.execute(f"SELECT key, value FROM {self.__table_name}")
for row in self.__cursor.fetchall():
yield row
def clear(self):
self.__cursor.execute(f"DELETE FROM {self.__table_name}")
self.__conn.commit()
def pop(self, key):
value = self.__getitem__(key)
self.__delitem__(key)
return value
def popitem(self):
self.__cursor.execute(f"SELECT key, value FROM {self.__table_name} LIMIT 1")
result = self.__cursor.fetchone()
if result is None:
raise KeyError('popitem(): store is empty')
self.__delitem__(result[0])
return result
def get(self, key, default=None):
try:
return self.__getitem__(key)
except KeyError:
return default
def close(self):
self.__conn.close()
class ConfigStore(KeyValStore):
"""
ConfigStore is a wrapper around KeyValStore that provides a dict-like
interface for storing and retrieving configuration data.
ConfigStore requires keys to be of type string and values to be of type
dict.
Example usage:
>>> from data_store import ConfigStore
>>> store = ConfigStore('test_store')
>>> store['foo'] = {'bar': 'baz'}
>>> store['foo']
{'bar': 'baz'}
>>> store['foo'] = {'bar': 'qux'}
>>> store['foo']
{'bar': 'qux'}
>>> store.close()
>>> with ConfigStore('test_store') as store:
... store['foo'] = {'bar': 'baz'}
>>> with ConfigStore('test_store') as store:
... store['foo']
{'bar': 'baz'}
"""
def __init__(self, store_name, store_path='./.data'):
"""
:param store_name: name of the store
:param store_path: [optional] path to the store. Defaults to './.data'
"""
super().__init__(store_name, store_path)
def __validate_value(self, value):
if type(value) is not dict:
raise ValueError('Value must be of type dict')
def __setitem__(self, key, value):
self.__validate_value(value)
super().__setitem__(key, json.dumps(value))
def __getitem__(self, key):
return json.loads(super().__getitem__(key))
class SecretStore(KeyValStore):
"""
SecretStore is a wrapper around KeyValStore that provides a dict-like
interface for storing and retrieving secret data.
SecretStore requires keys and values to be of type string.
Example usage:
1. Generate unique key and keep it safe. This key is master password for the store.
>>> from data_store import SecretStore
>>> key = SecretStore.generate_fernet_key()
2. Create store and pass the key.
>>> from data_store import SecretStore
>>> store = SecretStore('test_store', key)
>>> store['foo'] = 'bar'
>>> store['foo']
'bar'
>>> store['foo'] = 'baz'
>>> store['foo']
'baz'
>>> store.close()
>>> with SecretStore('test_store', key) as store:
... store['foo'] = 'bar'
>>> with SecretStore('test_store', key) as store:
... store['foo']
'bar'
"""
def __init__(self, store_name, fernet_key, store_path='./.data'):
"""
:param store_name: name of the store
:param fernet_key: key for encryption/decryption
:param store_path: [optional] path to the store. Defaults to './.data'
"""
super().__init__(store_name, store_path)
self.__fernet = fernet.Fernet(fernet_key.encode())
self['__fernet_key__'] = fernet_key
@staticmethod
def generate_fernet_key():
return fernet.Fernet.generate_key().decode()
def __validate_value(self, value):
if type(value) is not str:
raise ValueError('Value must be of type string')
def __setitem__(self, key, value):
self.__validate_value(value)
super().__setitem__(key, self.__fernet.encrypt(value.encode()).decode())
def __getitem__(self, key):
try:
value = self.__fernet.decrypt(super().__getitem__(key).encode()).decode()
except fernet.InvalidToken:
raise ValueError('Invalid fernet key for this store') from None
return value
def clear(self):
raise NotImplementedError(f'clear() is not supported for {self.__class__.__name__}')
def pop(self, key):
raise NotImplementedError(f'pop() is not supported for {self.__class__.__name__}')
def popitem(self):
raise NotImplementedError(f'popitem() is not supported for {self.__class__.__name__}')
def __delitem__(self, key):
raise NotImplementedError(f'delitem() is not supported for {self.__class__.__name__}')
def values(self):
raise NotImplementedError(f'values() is not supported for {self.__class__.__name__}')
def items(self):
raise NotImplementedError(f'items() is not supported for {self.__class__.__name__}')
class PropertyStore(ConfigStore):
"""
PropertyStore is a wrapper around ConfigStore that provides a dict-like
interface for storing and retrieving property data.
PropertyStore requires keys to be of type string and values to be of type
int, float, bool, str, list or dict.
Example usage:
>>> from data_store import PropertyStore
>>> store = PropertyStore('test_store')
>>> store['foo'] = 1
>>> store['foo']
1
>>> store['foo'] = 2.0
>>> store['foo']
2.0
>>> store['foo'] = True
>>> store['foo']
True
>>> store['foo'] = 'bar'
>>> store['foo']
'bar'
>>> store['foo'] = ['bar', 'baz']
>>> store['foo']
['bar', 'baz']
>>> store['foo'] = {'bar': 'baz'}
>>> store['foo']
{'bar': 'baz'}
>>> store.close()
>>> with PropertyStore('test_store') as store:
... store['foo'] = 1
>>> with PropertyStore('test_store') as store:
... store['foo']
1
"""
def __init__(self, store_name, store_path='./.data'):
"""
:param store_name: name of the store
:param store_path: [optional] path to the store. Defaults to './.data'
"""
self.__property_key = '_v'
super().__init__(store_name, store_path)
def __validate_value(self, value):
if type(value) not in [int, float, bool, str, list, dict]:
raise ValueError(f'Value must be of type int, float, bool, str, list or dict.\
\nProvided type: {type(value).__name__}')
def __setitem__(self, key, value):
self.__validate_value(value)
super().__setitem__(key, {self.__property_key: value})
def __getitem__(self, key):
return super().__getitem__(key)[self.__property_key]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment