From 34adb2db2f85529117693ce02da340f6d8b88d7f Mon Sep 17 00:00:00 2001 From: Laurent Bachelier Date: Fri, 15 Mar 2019 16:38:11 +0100 Subject: [PATCH] config: Add SQLiteConfig --- weboob/tools/config/sqliteconfig.py | 242 ++++++++++++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 weboob/tools/config/sqliteconfig.py diff --git a/weboob/tools/config/sqliteconfig.py b/weboob/tools/config/sqliteconfig.py new file mode 100644 index 0000000000..714e445b0f --- /dev/null +++ b/weboob/tools/config/sqliteconfig.py @@ -0,0 +1,242 @@ +# -*- coding: utf-8 -*- + +# Copyright(C) 2019 Laurent Bachelier +# +# This file is part of weboob. +# +# weboob is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# weboob is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with weboob. If not, see . + + +import os +import sqlite3 +import tempfile +from collections import Mapping, MutableMapping + +import yaml +from six import PY2 + +from .extra import time_buffer +from .iconfig import ConfigError, IConfig +from .yamlconfig import WeboobDumper + +try: + from yaml import CLoader as Loader +except ImportError: + from yaml import Loader + + +__all__ = ['SQLiteConfig'] + + +class VirtualRootDict(Mapping): + def __init__(self, config): + self.config = config + + def __getitem__(self, base): + if base in self.config._tables: + return VirtualDict(self.config, base) + raise KeyError('%s table not found' % base) + + def __iter__(self): + for base in self.config._tables: + yield base + + def __len__(self): + return len(self.config._tables) + + +class VirtualDict(MutableMapping): + def __init__(self, config, base): + self.config = config + self.base = base + + def __getitem__(self, key): + try: + return self.config.get(self.base, key) + except ConfigError: + raise KeyError('%s key in %s table not found' % (key, self.base)) + + def __contains__(self, key): + return self.config.has(self.base, key) + + def __iter__(self): + for key in self.config.keys(self.base): + yield key + + def iteritems(self): + return self.config.items(self.base) + + def items(self): + if PY2: + return list(self.iteritems()) + else: + return self.iteritems() + + def __len__(self): + return self.config.count(self.base) + + def __delitem__(self, key): + try: + self.config.delete(self, self.base, key) + except ConfigError: + raise KeyError('%s key in %s table not found' % (key, self.base)) + + def __setitem__(self, key, value): + self.config.set(self.base, key, value) + + +class SQLiteConfig(IConfig): + commit_since_seconds = 3600 + dump_since_seconds = 600 + + def __init__(self, path, commit_since_seconds=None, dump_since_seconds=None, last_run=True, logger=None): + self.path = path + if commit_since_seconds: + self.commit_since_seconds = commit_since_seconds + if dump_since_seconds: + self.dump_since_seconds = dump_since_seconds + if self.commit_since_seconds: + self.commit = time_buffer(since_seconds=self.commit_since_seconds, last_run=last_run, logger=logger)(self.commit) + if self.dump_since_seconds: + self.dump = time_buffer(since_seconds=self.dump_since_seconds, last_run=last_run, logger=logger)(self.dump) + + def load(self, default={}, optimize=True): + self.storage = sqlite3.connect(self.path) + self.storage.execute('PRAGMA page_size = 4096') + if optimize: + self.storage.execute('VACUUM') + self.storage.execute('REINDEX') + self._tables = set(self.tables()) + self.values = VirtualRootDict(self) + + def save(self, commit_since_seconds=None, dump_since_seconds=None): + self.commit(since_seconds=commit_since_seconds) + # No one would want immediate dumps, assume it means no dumps + if self.dump_since_seconds: + self.dump(since_seconds=dump_since_seconds) + + def force_save(self): + self.save(commit_since_seconds=False, dump_since_seconds=False) + + def __exit__(self, t, v, tb): + self.force_save() + super(SQLiteConfig, self).__exit__(t, v, tb) + + def commit(self, **kwargs): + kwargs.pop('since_seconds', None) + self.storage.commit() + + def dump(self, **kwargs): + kwargs.pop('since_seconds', None) + target = os.path.splitext(self.path)[0] + '.sql' + with tempfile.NamedTemporaryFile(dir=os.path.dirname(self.path), delete=False) as f: + for line in self.storage.iterdump(): + f.write(line) + f.write(b'\n') + os.rename(f.name, target) + + def ensure_table(self, name): + if name not in self._tables: + self.storage.execute('''CREATE TABLE IF NOT EXISTS %s ( + key text PRIMARY KEY, + value text + );''' % name) + self._tables.add(name) + + def tables(self): + cur = self.storage.cursor() + cur.execute('''SELECT name FROM sqlite_master + WHERE type="table" AND name NOT LIKE "sqlite_%";''') + return [k[0] for k in cur.fetchall()] + + def items(self, table, size=100): + """ + Low memory way of listing all items. + The size parameters alters how many items are fetched at a time. + """ + cur = self.storage.cursor() + cur.execute('SELECT key, value FROM %s;' % table) + items = cur.fetchmany(size) + while items: + for key, strvalue in items: + yield key, yaml.load(strvalue, Loader=Loader) + items = cur.fetchmany(size) + + def keys(self, table, size=200): + """ + Low memory way of listing all keys. + The size parameters alters how many items are fetched at a time. + """ + cur = self.storage.cursor() + cur.execute('SELECT key FROM %s;' % table) + items = cur.fetchmany(size) + while items: + for item in items: + yield item[0] + items = cur.fetchmany(size) + + def count(self, table): + cur = self.storage.cursor() + cur.execute('SELECT count(*) FROM %s;' % table) + return cur.fetchone()[0] + + def get(self, *args, **kwargs): + table = args[0] + key = '.'.join(args[1:]) + self.ensure_table(table) + try: + cur = self.storage.cursor() + cur.execute('SELECT value FROM %s WHERE key=?;' % table, (key, )) + row = cur.fetchone() + if row is None: + if 'default' in kwargs: + value = kwargs.get('default') + else: + raise ConfigError() + else: + strvalue = row[0] + value = yaml.load(strvalue, Loader=Loader) + except TypeError: + raise ConfigError() + return value + + def set(self, *args): + table = args[0] + key = '.'.join(args[1:-1]) + value = args[-1] + self.ensure_table(table) + try: + strvalue = yaml.dump(value, None, Dumper=WeboobDumper, default_flow_style=False) + cur = self.storage.cursor() + cur.execute('''INSERT OR REPLACE INTO %s VALUES (?, ?)''' % table, (key, strvalue)) + except KeyError: + raise ConfigError() + except TypeError: + raise ConfigError() + + def delete(self, *args): + table = args[0] + self.ensure_table(table) + key = '.'.join(args[1:]) + cur = self.storage.cursor() + cur.execute('DELETE FROM %s WHERE key=?;' % table, (key, )) + if not cur.rowcount: + raise ConfigError() + + def has(self, *args): + table = args[0] + key = '.'.join(args[1:]) + cur = self.storage.cursor() + cur.execute('SELECT count(*) FROM %s WHERE key=?;' % table, (key, )) + return cur.fetchone()[0] > 0 -- GitLab