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

@@ -882,9 +882,9 @@ 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']] = r
                 ids.append(r['uuid'])
         elif pid:
             pid = self.primary_id(table)
@@ -944,13 +944,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,31 +960,33 @@ 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'], '*')
@@ -999,8 +995,7 @@ def do_test(r):
                     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)]
 

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['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__,

superduper/base/apply.py

@@ -251,8 +251,8 @@ def _apply(
     deprecated_context: str | None = None,
 ):
 
-    if object.status == STATUS_UNINITIALIZED:
-        object.status, object.details = pending_status()
+    # if object.status == STATUS_UNINITIALIZED:
+    #     object.status, object.details = pending_status()
 
     processed_components = processed_components or set()
     if context is None:
@@ -311,7 +311,6 @@ def wrapper(child):
         if current.hash == object.hash:
             apply_status = 'same'
             object.version = current.version
-            object.status, object.details = running_status()
         elif current.uuid == object.uuid:
             apply_status = 'update'
             object.version = current.version
@@ -361,16 +360,23 @@ def wrapper(child):
             jobs=list(job_events.values()),
             context=context,
         )
-        # TODO - add multiple services (potentially empty)
-        for service in object.services:
-            put_events[f'{object.huuid}/{service}'] = PutComponent(
-                component=object.component,
-                identifier=object.identifier,
-                uuid=object.uuid,
-                context=context,
-                version=object.version,
-                service=service,
-            )
+        filter = getattr(object, 'filter_deployment', None)
+        if filter:
+            filter = getattr(object, filter)
+
+        branch, parent = object.get_branch_and_parent()
+        put_events[object.huuid] = PutComponent(
+            component=object.component,
+            identifier=object.identifier,
+            uuid=object.uuid,
+            context=context,
+            version=object.version,
+            services=object.services,
+            branch=branch,
+            parent=parent,
+            tags={tag: getattr(object, tag) for tag in object.tags},
+            filter=filter,
+        )
 
     elif apply_status == 'breaking':
         metadata_event = Create(
@@ -387,15 +393,23 @@ def wrapper(child):
             jobs=list(job_events.values()),
             context=context,
         )
-        for service in object.services:
-            put_events[f'{object.huuid}/{service}'] = PutComponent(
-                component=object.component,
-                identifier=object.identifier,
-                uuid=object.uuid,
-                context=context,
-                version=object.version,
-                service=service,
-            )
+        filter = getattr(object, 'filter_deployment', None)
+        if filter:
+            filter = getattr(object, filter)
+
+        branch, parent = object.get_branch_and_parent()
+        put_events[object.huuid] = PutComponent(
+            component=object.component,
+            identifier=object.identifier,
+            uuid=object.uuid,
+            context=context,
+            version=object.version,
+            services=object.services,
+            branch=branch,
+            tags={tag: getattr(object, tag) for tag in object.tags},
+            parent=parent,
+            filter=filter,
+        )
 
         d = db['Deployment']
         assert deprecated_context is not None

superduper/base/base.py

@@ -20,8 +20,26 @@
 class _UniqueRegistry:
     def __init__(self, d):
         self.d = d
+
+    def __contains__(self, item):
+        if '.' in item:
+            return item in self.d
+        else:
+            return item in [k.split('.')[-1] for k in self.d]
+
+    def drop(self):
+        for k in list(self.d.keys()):
+            del self.d[k]
+
     def __getitem__(self, item):
-        return self.d[item]
+        if '.' in item:
+            return self.d[item]
+        else:
+            try:
+                key = next(k for k in self.d if k.split('.')[-1] == item)
+            except StopIteration:
+                raise KeyError(item)
+            return self.d[key]
 
     def __setitem__(self, key, value):
         matching = [x for x in self.d.keys() if x.split('.')[-1] == key.split('.')[-1]]
@@ -108,7 +126,7 @@ class Base(metaclass=BaseMeta):
     """Base class for all superduper classes."""
 
     verbosity: t.ClassVar[int] = 0
-    set_post_init: t.ClassVar[t.Sequence[str]] = ()
+    primary_id: t.ClassVar[str] = 'uuid'
 
     def __init_subclass__(cls):
         full_path = f"{cls.__module__}.{cls.__name__}"
@@ -151,16 +169,14 @@ def class_schema(cls):
 
     @lazy_classproperty
     def table(cls):
-        from superduper import Component
         from superduper.components.table import Table
-        from superduper.misc.importing import isreallyinstance
 
+        if cls.__name__ == 'Base':
+            return None
         return Table(
             identifier=cls.__name__,
             fields=cls.class_schema.fields,
-            path=cls.__module__ + '.' + cls.__name__,
-            primary_id='uuid',
-            is_component=isreallyinstance(cls, Component),
+            primary_id=cls.primary_id,
         )
 
     @staticmethod
@@ -298,10 +314,8 @@ def decode(cls, r, db: t.Optional['Datalayer'] = None):
         """
         from superduper.base.document import Document
 
-        if '_path' in r:
-            from superduper.misc.importing import import_object
-
-            cls = import_object(r['_path'])
+        if 'component' in r:
+            cls = REGISTRY[r['component']]
 
         r = Document.decode(r, schema=cls.class_schema, db=db)
         return cls.from_dict(r, db=db)
@@ -334,6 +348,7 @@ def encode(
                             del r[k]
         if r is None:
             r = self.dict(metadata=context.metadata)
+            r['component'] = self.__class__.__name__
 
         if not context.defaults:
             for k, v in list(r.items()):
@@ -443,7 +458,6 @@ def dict(self, metadata: bool = True) -> t.Dict[str, t.Any]:
         from superduper import Document
 
         r = asdict(self)
-        r['_path'] = self.__class__.__module__ + '.' + self.__class__.__name__
         if metadata:
             metadata = getattr(self, 'metadata', {})
             for k, v in metadata.items():

superduper/base/datalayer.py

@@ -11,7 +11,7 @@
 from superduper.backends.base.data_backend import BaseDataBackend
 from superduper.base import apply, exceptions
 from superduper.base.artifacts import ArtifactStore
-from superduper.base.base import Base
+from superduper.base.base import REGISTRY, Base
 from superduper.base.config import Config
 from superduper.base.datatype import BaseType, ComponentType
 from superduper.base.document import Document
@@ -95,8 +95,6 @@ def insert(self, items: t.List[Base]):
         """
         table = self.pre_insert(items)
         data = [x.dict() for x in items]
-        for r in data:
-            del r['_path']
         return self[table.identifier].insert(data)
 
     def replace(self, condition: t.Dict, item: Base):
@@ -309,7 +307,9 @@ def pre_insert(
             return self.metadata.create(type(items[0]))
 
     def _post_query(self, table: str, ids: t.Sequence[str], type_: str):
-        if table in metaclasses or self.metadata.is_component(table):
+        if table in metaclasses or (
+            table in REGISTRY and issubclass(REGISTRY[table], Component)
+        ):
             return
         if (
             not table.startswith(CFG.output_prefix)