Source code for graphbrain.memory.sqlite

import json

from sqlite3 import connect

from graphbrain.hyperedge import hedge
from graphbrain.memory.keyvalue import KeyValue
from graphbrain.memory.permutations import str_plus_1


def _encode_attributes(attributes):
    return json.dumps(attributes, ensure_ascii=False, check_circular=False, separators=(',', ':'))


def _decode_attributes(value):
    return json.loads(value)


[docs] class SQLite(KeyValue): """Implements SQLite hypergraph storage.""" def __init__(self, locator_string): super().__init__(locator_string) self.conn = connect(self.locator_string, isolation_level=None) self.cur = None # self.conn.execute('PRAGMA synchronous = OFF') # self.conn.execute('PRAGMA journal_mode = MEMORY') self.conn.execute('CREATE TABLE IF NOT EXISTS v (key TEXT PRIMARY KEY, value TEXT)') self.conn.execute('CREATE TABLE IF NOT EXISTS p (key TEXT PRIMARY KEY)') # =================================== # Implementation of interface methods # ===================================
[docs] def close(self): if self.conn: self.conn.close() self.conn = None
[docs] def destroy(self): cur = self.conn.cursor() cur.execute('DELETE FROM v') cur.execute('DELETE FROM p')
[docs] def all(self): cur = self.conn.cursor() for key, _ in cur.execute('SELECT * FROM v'): edge = hedge(key) if edge is not None: yield edge
[docs] def all_attributes(self): cur = self.conn.cursor() for key, value in cur.execute('SELECT * FROM v'): edge = hedge(key) if edge is not None: attributes = _decode_attributes(value) yield edge, attributes
def begin_transaction(self): if self.batch_mode: return self.cur = self.conn.cursor() self.cur.execute('BEGIN TRANSACTION') def end_transaction(self): if self.batch_mode: return self.conn.commit() self.cur = None # ========================================== # Implementation of private abstract methods # ========================================== def _edge2key(self, edge): return edge.to_str() def _exists_key(self, key): """Checks if the given key exists.""" cur = self.conn.cursor() for key, _ in cur.execute('SELECT * FROM v WHERE key = ?', (key, )): return True return False def _add_key(self, key, attributes): """Adds the given edge, given its key.""" value = _encode_attributes(attributes) self.cur.execute('INSERT OR REPLACE INTO v (key, value) VALUES(?, ?)', (key, value)) def _attribute_key(self, key): cur = self.conn.cursor() for key, value in cur.execute('SELECT * FROM v WHERE key = ?', (key,)): return _decode_attributes(value) return None def _write_edge_permutation(self, perm): """Writes a given permutation.""" self.cur.execute('INSERT OR IGNORE INTO p (key) VALUES(?)', (perm,)) def _remove_edge_permutation(self, perm): """Removes a given permutation.""" self.cur.execute('DELETE FROM p WHERE key = ?', (perm,)) def _remove_key(self, key): """Removes an edge, given its key.""" self.cur.execute('DELETE FROM v WHERE key = ?', (key,)) def _permutations_with_prefix(self, prefix): end_str = str_plus_1(prefix) cur = self.conn.cursor() for row in cur.execute('SELECT * FROM p WHERE key >= ? AND key < ?', (prefix, end_str)): yield row[0] def _edges_with_prefix(self, prefix): end_str = str_plus_1(prefix) cur = self.conn.cursor() for key, _ in cur.execute('SELECT * FROM v WHERE key >= ? AND key < ?', (prefix, end_str)): yield hedge(key)