474 lines
12 KiB
Python
474 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
import os
|
|
|
|
try:
|
|
import ujson as json
|
|
except ImportError:
|
|
import json
|
|
|
|
import irc3
|
|
import shelve
|
|
|
|
__doc__ = '''
|
|
==========================================
|
|
:mod:`irc3.plugins.storage` Storage plugin
|
|
==========================================
|
|
|
|
Add a ``db`` attribute to the bot
|
|
|
|
..
|
|
>>> from irc3.testing import IrcBot
|
|
>>> from irc3.testing import ini2config
|
|
>>> import tempfile
|
|
>>> fd = tempfile.NamedTemporaryFile(prefix='irc3', suffix='.db')
|
|
>>> db_file = fd.name
|
|
>>> fd.close()
|
|
>>> fd = tempfile.NamedTemporaryFile(prefix='irc3', suffix='.json')
|
|
>>> json_file = fd.name
|
|
>>> fd.close()
|
|
>>> fd = tempfile.NamedTemporaryFile(prefix='irc3', suffix='.sqlite')
|
|
>>> sqlite_file = fd.name
|
|
>>> fd.close()
|
|
|
|
Usage::
|
|
|
|
>>> config = ini2config("""
|
|
... [bot]
|
|
... includes =
|
|
... irc3.plugins.storage
|
|
... storage = json://%s
|
|
... """ % json_file)
|
|
>>> bot = IrcBot(**config)
|
|
|
|
Then use it::
|
|
|
|
>>> bot.db['mykey'] = dict(key='value')
|
|
>>> 'mykey' in bot.db
|
|
True
|
|
>>> bot.db['mykey']
|
|
{'key': 'value'}
|
|
>>> bot.db.setdefault('mykey', key='default')
|
|
{'key': 'value'}
|
|
>>> bot.db.setdefault('mykey', item='default')
|
|
{'item': 'default'}
|
|
>>> bot.db.set('mykey', item='value')
|
|
>>> bot.db.setdefault('mykey', item='default')
|
|
{'item': 'value'}
|
|
>>> del bot.db['mykey']
|
|
>>> bot.db.get('mykey')
|
|
>>> bot.db.get('mykey', 'default')
|
|
'default'
|
|
>>> bot.db['mykey']
|
|
Traceback (most recent call last):
|
|
...
|
|
KeyError: 'mykey'
|
|
>>> 'mykey' in bot.db
|
|
False
|
|
>>> bot.db.setlist('mylist', ['foo', 'bar'])
|
|
>>> bot.db.getlist('mylist')
|
|
['foo', 'bar']
|
|
>>> del bot.db['mylist']
|
|
|
|
You can use an instance as key::
|
|
|
|
>>> class MyPlugin:
|
|
... pass
|
|
>>> plugin = MyPlugin()
|
|
>>> bot.db[plugin] = dict(key='value')
|
|
>>> bot.db[plugin]
|
|
{'key': 'value'}
|
|
>>> del bot.db[plugin]
|
|
>>> bot.db.get(plugin)
|
|
|
|
..
|
|
>>> bot.db.SIGINT()
|
|
|
|
You can also use shelve::
|
|
|
|
>>> config = ini2config("""
|
|
... [bot]
|
|
... includes =
|
|
... irc3.plugins.storage
|
|
... storage = shelve://%s
|
|
... """ % db_file)
|
|
>>> bot = IrcBot(**config)
|
|
>>> bot.db['mykey'] = dict(key='value')
|
|
>>> bot.db['mykey']
|
|
{'key': 'value'}
|
|
>>> del bot.db['mykey']
|
|
>>> bot.db.get('mykey')
|
|
>>> bot.db.setlist('mylist', ['foo', 'bar'])
|
|
>>> bot.db.getlist('mylist')
|
|
['foo', 'bar']
|
|
>>> del bot.db['mylist']
|
|
|
|
..
|
|
>>> bot.db.getlist('mylist', ['foo', 'bar'])
|
|
['foo', 'bar']
|
|
>>> bot.db.setlist('mylist', ['foo', 'bar'])
|
|
>>> bot.db.setlist('mylist', ['foo', 'bar'])
|
|
>>> del bot.db['mylist']
|
|
>>> bot.db.SIGINT()
|
|
|
|
|
|
Or redis::
|
|
|
|
>>> config = ini2config("""
|
|
... [bot]
|
|
... includes =
|
|
... irc3.plugins.storage
|
|
... storage = redis://localhost:6379/10
|
|
... """)
|
|
>>> bot = IrcBot(**config)
|
|
|
|
..
|
|
>>> bot.db.backend.flushdb() # require redis
|
|
>>> bot.db.SIGINT()
|
|
|
|
Then use it::
|
|
|
|
>>> bot.db['mykey'] = dict(key='value')
|
|
>>> bot.db['mykey']
|
|
{'key': 'value'}
|
|
>>> del bot.db['mykey']
|
|
>>> bot.db.get('mykey')
|
|
>>> bot.db['mykey']
|
|
Traceback (most recent call last):
|
|
...
|
|
KeyError: 'mykey'
|
|
>>> bot.db.setlist('mylist', ['foo', 'bar'])
|
|
>>> bot.db.getlist('mylist')
|
|
['foo', 'bar']
|
|
>>> del bot.db['mylist']
|
|
|
|
Or sqlite::
|
|
|
|
>>> config = ini2config("""
|
|
... [bot]
|
|
... includes =
|
|
... irc3.plugins.storage
|
|
... storage = sqlite://%s
|
|
... """ % sqlite_file)
|
|
>>> bot = IrcBot(**config)
|
|
|
|
..
|
|
>>> bot.db.backend.flushdb() # require redis
|
|
>>> bot.db.SIGINT()
|
|
|
|
Then use it::
|
|
|
|
>>> bot.db['mykey'] = dict(key='value')
|
|
>>> bot.db['mykey']
|
|
{'key': 'value'}
|
|
>>> del bot.db['mykey']
|
|
>>> bot.db.get('mykey')
|
|
>>> bot.db['mykey']
|
|
Traceback (most recent call last):
|
|
...
|
|
KeyError: 'mykey'
|
|
>>> bot.db.setlist('mylist', ['foo', 'bar'])
|
|
>>> bot.db.getlist('mylist')
|
|
['foo', 'bar']
|
|
>>> del bot.db['mylist']
|
|
|
|
Api
|
|
===
|
|
|
|
.. autoclass:: Storage
|
|
:members: __getitem__,__setitem__,__delitem__,__contains__,get,set,setdefault
|
|
|
|
'''
|
|
|
|
|
|
class Shelve:
|
|
|
|
def __init__(self, uri=None, **kwargs):
|
|
self.filename = uri[9:]
|
|
self.db = shelve.open(self.filename)
|
|
|
|
def set(self, key, value):
|
|
self.db[key] = value
|
|
self.db.sync()
|
|
|
|
def get(self, key):
|
|
return self.db[key]
|
|
|
|
def delete(self, key):
|
|
del self.db[key]
|
|
self.sync()
|
|
|
|
def contains(self, key):
|
|
return key in self.db
|
|
|
|
def sync(self):
|
|
self.db.sync()
|
|
|
|
def close(self):
|
|
self.db.close()
|
|
|
|
|
|
class JSON:
|
|
|
|
def __init__(self, uri=None, **kwargs):
|
|
self.filename = uri[7:]
|
|
if os.path.isfile(self.filename): # pragma: no cover
|
|
with open(self.filename) as fd:
|
|
self.db = json.load(fd)
|
|
else:
|
|
self.db = {}
|
|
|
|
def set(self, key, value):
|
|
self.db[key] = value
|
|
self.sync()
|
|
|
|
def get(self, key):
|
|
return self.db[key]
|
|
|
|
def delete(self, key):
|
|
del self.db[key]
|
|
self.sync()
|
|
|
|
def contains(self, key):
|
|
return key in self.db
|
|
|
|
def sync(self):
|
|
with open(self.filename, 'w') as fd:
|
|
json.dump(self.db, fd, indent=2, sort_keys=True)
|
|
|
|
def close(self):
|
|
self.sync()
|
|
|
|
|
|
class Redis:
|
|
|
|
def __init__(self, uri=None, **kwargs):
|
|
ConnectionPool = irc3.utils.maybedotted(
|
|
'redis.connection.ConnectionPool')
|
|
pool = ConnectionPool.from_url(uri)
|
|
StrictRedis = irc3.utils.maybedotted('redis.client.StrictRedis')
|
|
self.db = StrictRedis(connection_pool=pool)
|
|
|
|
def set(self, key, value):
|
|
self.db.hmset(key, value)
|
|
|
|
def get(self, key):
|
|
keys = self.db.hkeys(key)
|
|
if not keys:
|
|
raise KeyError()
|
|
values = self.db.hmget(key, keys)
|
|
keys = [k.decode('utf8') for k in keys]
|
|
values = [v.decode('utf8') for v in values]
|
|
values = dict(zip(keys, values))
|
|
return values
|
|
|
|
def delete(self, key):
|
|
self.db.delete(key)
|
|
|
|
def contains(self, key):
|
|
return self.db.exists(key)
|
|
|
|
def flushdb(self):
|
|
self.db.flushdb()
|
|
|
|
def sync(self):
|
|
self.db.save()
|
|
|
|
def close(self):
|
|
self.sync()
|
|
|
|
|
|
class SQLite:
|
|
|
|
CREATE_TABLE = """
|
|
CREATE TABLE IF NOT EXISTS
|
|
irc3_storage (
|
|
key text not null,
|
|
value text default '',
|
|
PRIMARY KEY (key)
|
|
);
|
|
"""
|
|
UPSERT = """
|
|
INSERT OR REPLACE INTO irc3_storage(key,value) VALUES(?, ?);
|
|
"""
|
|
|
|
def __init__(self, uri=None, **kwargs):
|
|
self.sqlite = irc3.utils.maybedotted('sqlite3')
|
|
self.uri = uri.split('://')[-1]
|
|
conn = self.sqlite.connect(self.uri)
|
|
cursor = conn.cursor()
|
|
cursor.execute(self.CREATE_TABLE)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def set(self, key, value):
|
|
conn = self.sqlite.connect(self.uri)
|
|
cursor = conn.cursor()
|
|
cursor.execute(self.UPSERT, (key, json.dumps(value)))
|
|
cursor.fetchall()
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def get(self, key):
|
|
value = None
|
|
conn = self.sqlite.connect(self.uri)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT value FROM irc3_storage where key=?;", (key,))
|
|
for row in cursor.fetchall():
|
|
value = json.loads(row[0])
|
|
break
|
|
cursor.close()
|
|
conn.close()
|
|
if value is None:
|
|
raise KeyError(key)
|
|
return value
|
|
|
|
def delete(self, key):
|
|
conn = self.sqlite.connect(self.uri)
|
|
cursor = conn.cursor()
|
|
cursor.execute("DELETE FROM irc3_storage where key=?;", (key,))
|
|
cursor.close()
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def contains(self, key):
|
|
conn = self.sqlite.connect(self.uri)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT value FROM irc3_storage where key=?;", (key,))
|
|
res = False
|
|
if len(list(cursor.fetchall())) == 1:
|
|
res = True
|
|
cursor.close()
|
|
conn.close()
|
|
return res
|
|
|
|
def flushdb(self):
|
|
conn = self.sqlite.connect(self.uri)
|
|
cursor = conn.cursor()
|
|
cursor.execute("DROP TABLE IF EXISTS irc3_storage;")
|
|
cursor.execute(self.CREATE_TABLE)
|
|
cursor.close()
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def sync(self):
|
|
pass
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
|
|
@irc3.plugin
|
|
class Storage:
|
|
|
|
backends = {
|
|
'shelve': Shelve,
|
|
'json': JSON,
|
|
'unix': Redis,
|
|
'redis': Redis,
|
|
'rediss': Redis,
|
|
'sqlite': SQLite,
|
|
}
|
|
|
|
def __init__(self, context):
|
|
uri = context.config.storage
|
|
name = uri.split('://', 1)[0]
|
|
try:
|
|
factory = self.backends[name]
|
|
except KeyError: # pragma: no cover
|
|
raise LookupError('No such backend %s' % name)
|
|
self.backend = factory(uri)
|
|
self.context = context
|
|
self.context.db = self
|
|
|
|
def setdefault(self, key_, **kwargs):
|
|
"""Update storage value for key with kwargs iif the keys doesn't
|
|
exist. Return stored values"""
|
|
stored = self[key_]
|
|
changed = False
|
|
for k, v in kwargs.items():
|
|
if k not in stored:
|
|
stored[k] = v
|
|
changed = True
|
|
else:
|
|
kwargs[k] = stored[k]
|
|
if changed:
|
|
self[key_] = stored
|
|
return kwargs
|
|
|
|
def get(self, key_, default=None):
|
|
"""Get storage value for key or return default"""
|
|
if key_ not in self:
|
|
return default
|
|
else:
|
|
return self[key_]
|
|
|
|
def getlist(self, key_, default=None):
|
|
"""Get storage value (as list) for key or return default"""
|
|
if key_ not in self:
|
|
return default
|
|
else:
|
|
value = self[key_]
|
|
value = [(int(i), v) for i, v in value.items()]
|
|
return [v for k, v in sorted(value)]
|
|
|
|
def set(self, key_, **kwargs):
|
|
"""Update storage value for key with kwargs"""
|
|
stored = self.get(key_, dict())
|
|
changed = False
|
|
for k, v in kwargs.items():
|
|
if k not in stored or stored[k] != v:
|
|
stored[k] = v
|
|
changed = True
|
|
if changed:
|
|
self[key_] = stored
|
|
|
|
def setlist(self, key_, value):
|
|
"""Update storage value (as list)"""
|
|
value = dict([(str(i), v) for i, v in enumerate(value)])
|
|
if key_ in self:
|
|
del self[key_]
|
|
self.set(key_, **value)
|
|
|
|
def __setitem__(self, key, value):
|
|
"""Set storage value for key"""
|
|
key = getattr(key, '__module__', key)
|
|
if not isinstance(value, dict): # pragma: no cover
|
|
raise TypeError('value must be a dict')
|
|
try:
|
|
return self.backend.set(key, value)
|
|
except Exception as e: # pragma: no cover
|
|
self.context.log.exception(e)
|
|
raise
|
|
|
|
def __getitem__(self, key):
|
|
"""Get storage value for key"""
|
|
key = getattr(key, '__module__', key)
|
|
try:
|
|
return self.backend.get(key)
|
|
except KeyError:
|
|
raise KeyError(key)
|
|
except Exception as e: # pragma: no cover
|
|
self.context.log.exception(e)
|
|
raise
|
|
|
|
def __delitem__(self, key):
|
|
"""Delete key in storage"""
|
|
key = getattr(key, '__module__', key)
|
|
try:
|
|
self.backend.delete(key)
|
|
except Exception as e: # pragma: no cover
|
|
self.context.log.exception(e)
|
|
raise
|
|
|
|
def __contains__(self, key):
|
|
"""Return True if storage contains key"""
|
|
key = getattr(key, '__module__', key)
|
|
try:
|
|
return self.backend.contains(key)
|
|
except Exception as e: # pragma: no cover
|
|
self.context.log.exception(e)
|
|
raise
|
|
|
|
def SIGINT(self):
|
|
self.backend.close()
|