#409 merge SQLDict and SQLDictBinKey classes to one

proxy/indexer/indexer_base.py

@@ -6,10 +6,8 @@
 from typing import Dict, Union
 
 try:
-    from sql_dict_bin_key import SQLDictBinKey
     from sql_dict import SQLDict
 except ImportError:
-    from .sql_dict_bin_key import SQLDictBinKey
     from .sql_dict import SQLDict
 
 
@@ -42,6 +40,7 @@ def __init__(self,
 
         self.evm_loader_id = evm_loader_id
         self.client = Client(solana_url)
+        self.transaction_receipts = SQLDict(tablename="known_transactions", bin_key=True)
         self.last_slot = start_slot
         self.current_slot = 0
         self.counter_ = 0

proxy/indexer/sql_dict.py

@@ -28,13 +28,20 @@ def decode(obj):
     return loads(bytes(obj))
 
 
+def dummy(obj):
+    """Does nothing"""
+    return obj
+
+
 class SQLDict(MutableMapping):
     """Serialize an object using pickle to a binary format accepted by SQLite."""
 
-    def __init__(self, tablename='table'):
+    def __init__(self, tablename='table', bin_key=False):
         self.encode = encode
         self.decode = decode
-        self.tablename = tablename
+        self.key_encode = encode if bin_key else dummy
+        self.key_decode = decode if bin_key else dummy
+        self.tablename = tablename + ("_bin_key" if bin_key else "")
         self.conn = psycopg2.connect(
             dbname=POSTGRES_DB,
             user=POSTGRES_USER,
@@ -43,44 +50,45 @@ def __init__(self, tablename='table'):
         )
         self.conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
         cur = self.conn.cursor()
-        cur.execute('''
+        key_type = 'BYTEA' if bin_key else 'TEXT'
+        cur.execute(f'''
                 CREATE TABLE IF NOT EXISTS
-                {} (
-                    key TEXT UNIQUE,
+                {self.tablename} (
+                    key {key_type} UNIQUE,
                     value BYTEA
                 )
-            '''.format(self.tablename)
+            '''
         )
 
     def close(self):
         self.conn.close()
 
     def __len__(self):
         cur = self.conn.cursor()
-        cur.execute('SELECT COUNT(*) FROM {}'.format(self.tablename))
+        cur.execute(f'SELECT COUNT(*) FROM {self.tablename}')
         rows = cur.fetchone()[0]
         return rows if rows is not None else 0
 
     def iterkeys(self):
         cur = self.conn.cursor()
-        cur.execute('SELECT key FROM {}'.format(self.tablename))
+        cur.execute(f'SELECT key FROM {self.tablename}')
         rows = cur.fetchall()
         for row in rows:
-            yield row[0]
+            yield self.key_decode(row[0])
 
     def itervalues(self):
         cur = self.conn.cursor()
-        cur.execute('SELECT value FROM {}'.format(self.tablename))
+        cur.execute(f'SELECT value FROM {self.tablename}')
         rows = cur.fetchall()
-        for value in rows:
-            yield self.decode(value[0])
+        for row in rows:
+            yield self.decode(row[0])
 
     def iteritems(self):
         cur = self.conn.cursor()
-        cur.execute('SELECT key, value FROM {}'.format(self.tablename))
+        cur.execute(f'SELECT key, value FROM {self.tablename}')
         rows = cur.fetchall()
         for row in rows:
-            yield row[0], self.decode(row[1])
+            yield self.key_decode(row[0]), self.decode(row[1])
 
     def keys(self):
         return list(self.iterkeys())
@@ -92,35 +100,40 @@ def items(self):
         return list(self.iteritems())
 
     def __contains__(self, key):
+        bin_key = self.key_encode(key)
         cur = self.conn.cursor()
-        cur.execute('SELECT 1 FROM {} WHERE key = %s'.format(self.tablename), (key,))
+        cur.execute(f'SELECT 1 FROM {self.tablename} WHERE key = %s', (bin_key,))
         return cur.fetchone() is not None
 
     def __getitem__(self, key):
+        bin_key = self.key_encode(key)
         cur = self.conn.cursor()
-        cur.execute('SELECT value FROM {} WHERE key = %s'.format(self.tablename), (key,))
+        cur.execute(f'SELECT value FROM {self.tablename} WHERE key = %s', (bin_key,))
         item = cur.fetchone()
         if item is None:
             raise KeyError(key)
         return self.decode(item[0])
 
     def __setitem__(self, key, value):
+        bin_key = self.key_encode(key)
+        bin_value = self.encode(value)
         cur = self.conn.cursor()
-        cur.execute('''
-                INSERT INTO {} (key, value)
+        cur.execute(f'''
+                INSERT INTO {self.tablename} (key, value)
                 VALUES (%s,%s)
                 ON CONFLICT (key)
                 DO UPDATE SET
                 value = EXCLUDED.value
-            '''.format(self.tablename),
-            (key, self.encode(value))
+            ''',
+            (bin_key, bin_value)
         )
 
     def __delitem__(self, key):
+        bin_key = self.key_encode(key)
         cur = self.conn.cursor()
-        if key not in self:
+        if bin_key not in self:
             raise KeyError(key)
-        cur.execute('DELETE FROM {} WHERE key = %s'.format(self.tablename), (key,))
+        cur.execute(f'DELETE FROM {self.tablename} WHERE key = %s', (bin_key,))
 
     def __iter__(self):
         return self.iterkeys()