Encode additional parameters in Component.data and store in 1 table

Author: blythed

.github/workflows/ci_code.yml

@@ -103,6 +103,7 @@ jobs:
 
     - name: Unit Testing
       run: |
+        sqlite3 test.db "create table t(f int); drop table t;"
         make unit_testing pytest_arguments="--cov=superduper --cov-report=xml" SUPERDUPER_CONFIG=test/configs/${{ matrix.config }}
 
     - name: Usecase Testing

plugins/sql/superduper_sql/data_backend.py

@@ -359,7 +359,10 @@ class SQLDatabackend(IbisDataBackend):
 
     def __init__(self, uri, plugin, flavour=None):
         super().__init__(uri, plugin, flavour)
-        self._create_sqlalchemy_engine()
+        if 'sqlite://./' in uri:
+            self._create_sqlalchemy_engine(uri.replace('./', '//'))
+        else:
+            self._create_sqlalchemy_engine(uri)
         self.sm = sessionmaker(bind=self.alchemy_engine)
 
     @property
@@ -374,6 +377,8 @@ def update(self, table, condition, key, value):
         with self.sm() as session:
             metadata = MetaData()
 
+            assert table in self.list_tables()
+
             metadata.reflect(bind=session.bind)
             table = Table(table, metadata, autoload_with=session.bind)
 
@@ -422,16 +427,16 @@ def delete(self, table, condition):
         except NoSuchTableError:
             raise exceptions.NotFound("Table", table)
 
-    def _create_sqlalchemy_engine(self):
+    def _create_sqlalchemy_engine(self, uri):
         with self.connection_manager.get_connection() as conn:
-            self.alchemy_engine = create_engine(self.uri, creator=lambda: conn.con)
+            self.alchemy_engine = create_engine(uri, creator=lambda: conn.con)
             if not self._test_engine():
                 logging.warn(
                     "Unable to reuse the ibis connection "
                     "to create the SQLAlchemy engine. "
                     "Creating a new connection with the URI."
                 )
-                self.alchemy_engine = create_engine(self.uri)
+                self.alchemy_engine = create_engine(uri)
 
     def _test_engine(self):
         """Test the engine."""

superduper/backends/base/data_backend.py

@@ -1,3 +1,4 @@
+import copy
 import functools
 import hashlib
 import json
@@ -643,11 +644,7 @@ def delete(self, table, condition):
         :param table: The table to delete from.
         :param condition: The condition to delete.
         """
-        r_table = self._get_with_component_identifier('Table', table)
-        if r_table is None:
-            raise exceptions.NotFound("Table", table)
-
-        if not r_table['is_component']:
+        if table not in {'Deployment', 'Component'}:
             pid = self.primary_id(table)
             if pid in condition:
                 docs = self.get_many(table, condition[pid])
@@ -660,13 +657,16 @@ def delete(self, table, condition):
             if 'uuid' in condition:
                 docs = self.get_many(table, '*', condition['uuid'])
             elif 'identifier' in condition:
-                docs = self.get_many(table, condition['identifier'], '*')
+                assert 'component' in condition
+                docs = self.get_many(
+                    table, condition['component'], condition['identifier'], '*'
+                )
                 docs = self._do_filter(docs, condition)
             else:
-                docs = self.get_many(table, '*', '*')
+                docs = self.get_many(table, '*', '*', '*')
                 docs = self._do_filter(docs, condition)
             for r in docs:
-                del self[table, r['identifier'], r['uuid']]
+                del self[table, r['component'], r['identifier'], r['uuid']]
 
     def drop_table(self, table):
         """Drop data from table.
@@ -735,9 +735,7 @@ def replace(self, table, condition, r):
         :param condition: The condition to update.
         :param r: The document to replace.
         """
-        r_table = self._get_with_component_identifier('Table', table)
-
-        if not r_table['is_component']:
+        if table not in {'Component', 'Deployment'}:
             pid = self.primary_id(table)
             if pid in condition:
                 docs = self.get_many(table, condition[pid])
@@ -749,18 +747,21 @@ def replace(self, table, condition, r):
                 self[table, s[pid]] = r
         else:
             if 'uuid' in condition:
-                s = self.get_many(table, '*', condition['uuid'])[0]
-                self[table, s['identifier'], condition['uuid']] = r
+                s = self.get_many(table, '*', '*', condition['uuid'])[0]
+                self[table, s['component'], s['identifier'], condition['uuid']] = r
             elif 'identifier' in condition:
-                docs = self.get_many(table, condition['identifier'], '*')
+                assert 'component' in condition
+                docs = self.get_many(
+                    table, condition['component'], condition['identifier'], '*'
+                )
                 docs = self._do_filter(docs, condition)
                 for s in docs:
-                    self[table, s['identifier'], s['uuid']] = r
+                    self[table, s['component'], s['identifier'], s['uuid']] = r
             else:
-                docs = self.get_many(table, '*', '*')
+                docs = self.get_many(table, '*', '*', '*')
                 docs = self._do_filter(docs, condition)
                 for s in docs:
-                    self[table, s['identifier'], s['uuid']] = r
+                    self[table, s['component'], s['identifier'], s['uuid']] = r
 
     def update(self, table, condition, key, value):
         """Update data in the database.
@@ -770,9 +771,7 @@ def update(self, table, condition, key, value):
         :param key: The key to update.
         :param value: The value to update.
         """
-        r_table = self._get_with_component_identifier('Table', table)
-
-        if not r_table['is_component']:
+        if table not in {'Component', 'Deployment'}:
             pid = self.primary_id(table)
             if pid in condition:
                 docs = self.get_many(table, '*', condition[pid]) + self.get_many(
@@ -788,19 +787,21 @@ def update(self, table, condition, key, value):
             if 'uuid' in condition:
                 s = self.get_many(table, '*', condition['uuid'])[0]
                 s[key] = value
-                self[table, s['identifier'], condition['uuid']] = s
+                self[table, s['component'], s['identifier'], condition['uuid']] = s
             elif 'identifier' in condition:
-                docs = self.get_many(table, condition['identifier'], '*')
+                docs = self.get_many(
+                    table, condition['component'], condition['identifier'], '*'
+                )
                 docs = self._do_filter(docs, condition)
                 for s in docs:
                     s[key] = value
-                    self[table, s['identifier'], s['uuid']] = s
+                    self[table, s['component'], s['identifier'], s['uuid']] = s
             else:
-                docs = self.get_many(table, '*', '*')
+                docs = self.get_many(table, '*', '*', '*')
                 docs = self._do_filter(docs, condition)
                 for s in docs:
                     s[key] = value
-                    self[table, s['identifier'], s['uuid']] = s
+                    self[table, s['component'], s['identifier'], s['uuid']] = s
 
     @abstractmethod
     def keys(self, *pattern) -> t.List[t.Tuple[str, str, str]]:
@@ -882,9 +883,11 @@ def insert(self, table, documents):
         except exceptions.NotFound:
             pid = None
 
-        if ('uuid' == pid or not pid) and "uuid" in documents[0]:
+        if table in {'Component', 'Deployment'}:
             for r in documents:
-                self[table, r['identifier'], r['uuid']] = r
+                self[table, r['component'], r['identifier'], r['uuid']] = copy.deepcopy(
+                    r
+                )
                 ids.append(r['uuid'])
         elif pid:
             pid = self.primary_id(table)
@@ -944,13 +947,7 @@ def do_test(r):
                         return False
                 return True
 
-        tables = self.get_many('Table', query.table, '*')
-        if not tables:
-            raise exceptions.NotFound("Table", query.table)
-
-        is_component = max(tables, key=lambda x: x['version'])['is_component']
-
-        if not is_component:
+        if query.table not in {'Component', 'Deployment'}:
             pid = self.primary_id(query.table)
 
             if pid in filter_kwargs:
@@ -966,41 +963,47 @@ def do_test(r):
         else:
 
             if not filter_kwargs:
-                keys = self.keys(query.table, '*', '*')
+                keys = self.keys(query.table, '*', '*', '*')
                 docs = [self[k] for k in keys]
             elif set(filter_kwargs.keys()) == {'uuid'}:
-                keys = self.keys(query.table, '*', filter_kwargs['uuid']['value'])
+                keys = self.keys(query.table, '*', '*', filter_kwargs['uuid']['value'])
                 docs = [self[k] for k in keys]
             elif set(filter_kwargs.keys()) == {'identifier'}:
                 assert filter_kwargs['identifier']['op'] == '=='
-
                 keys = self.keys(query.table, filter_kwargs['identifier']['value'], '*')
                 docs = [self[k] for k in keys]
-            elif set(filter_kwargs.keys()) == {'identifier', 'uuid'}:
+            elif set(filter_kwargs.keys()) == {'component', 'identifier', 'uuid'}:
                 assert filter_kwargs['identifier']['op'] == '=='
                 assert filter_kwargs['uuid']['op'] == '=='
+                assert filter_kwargs['component']['op'] == '=='
 
                 r = self[
                     query.table,
+                    filter_kwargs['component']['value'],
                     filter_kwargs['identifier']['value'],
                     filter_kwargs['uuid']['value'],
                 ]
                 if r is None:
                     docs = []
                 else:
                     docs = [r]
-            elif set(filter_kwargs.keys()) == {'identifier', 'version'}:
+            elif set(filter_kwargs.keys()) == {'component', 'identifier', 'version'}:
                 assert filter_kwargs['identifier']['op'] == '=='
+                assert filter_kwargs['component']['op'] == '=='
                 assert filter_kwargs['version']['op'] == '=='
 
-                keys = self.keys(query.table, filter_kwargs['identifier']['value'], '*')
+                keys = self.keys(
+                    query.table,
+                    filter_kwargs['component']['value'],
+                    filter_kwargs['identifier']['value'],
+                    '*',
+                )
                 docs = [self[k] for k in keys]
                 docs = [
                     r for r in docs if r['version'] == filter_kwargs['version']['value']
                 ]
             else:
-
-                keys = self.keys(query.table, '*', '*')
+                keys = self.keys(query.table, '*', '*', '*')
                 docs = [self[k] for k in keys]
                 docs = [r for r in docs if do_test(r)]
 
@@ -1011,4 +1014,4 @@ def do_test(r):
             cols = query.decomposition.select.args
             for i, r in enumerate(docs):
                 docs[i] = {k: v for k, v in r.items() if k in cols}
-        return docs
+        return [copy.deepcopy(r) for r in docs]

superduper/backends/local/vector_search.py

@@ -42,38 +42,31 @@ def build_tool(self, component):
 
     def initialize(self):
         """Initialize the vector search."""
-        components = []
-        from superduper import VectorIndex
-
-        for cls in self.db.show('Table'):
-            t = self.db.load('Table', identifier=cls)
-            if t.is_component and t.cls is not None:
-                if issubclass(t.cls, VectorIndex):
-                    components.append(t.identifier)
+        t = self.db.metadata.db['Deployment']
+        components = (
+            t.filter(t['parent'] == 'VectorIndex')
+            .select('component', 'uuid', 'identifier')
+            .execute()
+        )
+
         for component in components:
             try:
-                for identifier in self.db.show(component):
-                    try:
-                        vector_index = self.db.load(component, identifier=identifier)
-                        self.put_component(component, vector_index.uuid)
-                        vectors = vector_index.get_vectors()
-                        vectors = [VectorItem(**vector) for vector in vectors]
-                        self.get_tool(vector_index.uuid).add(vectors)
-
-                    except FileNotFoundError:
-                        logging.error(
-                            f'Could not load vector index: {identifier} '
-                            'Is the artifact store correctly configured?'
-                        )
-                        continue
-                    except TypeError as e:
-                        import traceback
-
-                        logging.error(
-                            f'Could not load vector index: {identifier} ' f'{e}'
-                        )
-                        logging.error(traceback.format_exc())
-                        continue
+                self.put_component(component['component'], uuid=component['uuid'])
+
+            except FileNotFoundError:
+                logging.error(
+                    f'Could not load vector index: {component["identifier"]} '
+                    'Is the artifact store correctly configured?'
+                )
+                continue
+            except TypeError as e:
+                import traceback
+
+                logging.error(
+                    f'Could not load vector index: {component["identifier"]} ' f'{e}'
+                )
+                logging.error(traceback.format_exc())
+                continue
             except exceptions.NotFound:
                 pass
 

superduper/base/annotations.py

@@ -49,7 +49,7 @@ def decorated(self, *, context: str = '', job: bool = False, **kwargs):
 
             if job:
                 return Job(
-                    component=self.__class__.__name__,
+                    parent_component=self.__class__.__name__,
                     identifier=self.identifier,
                     uuid=self.uuid,
                     method=f.__name__,