registry storage configuration

basis-company · Nov 15, 2023 · bacc319 · bacc319
1 parent b3a4300
commit bacc319
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 17 deletions.
diff --git a/registry/registry.py b/registry/registry.py
@@ -19,15 +19,15 @@ class QueryContext:
 
 
 class Registry:
-    def __init__(self) -> None:
+    def __init__(self, storages: Optional[list[Storage]] = None) -> None:
         self.ready: bool = False
         self.repositories: dict[type[Repository], Repository] = {
             repository: repository()
             for repository in Repository.__subclasses__()
         }
-        self.storages: list[Storage] = [
+        self.storages: list[Storage] = storages or [
             Storage(
-                id=1,
+                id=StorageRepository.storage_id,
                 storage_class=StorageClass.MEMORY,
                 driver=StorageDriver.MEMORY,
                 dsn=''
@@ -40,21 +40,27 @@ def get_repository(self, cls: type[Repository]) -> Repository:
 
         return self.repositories[cls]
 
-    async def bootstrap(self, source: Optional[Storage] = None):
+    async def bootstrap(self):
         if self.ready:
             return
 
+        primary: Optional[Storage] = None
         self.ready = True
-        if not source:
-            source = self.storages[0]
 
-        driver = await get_driver(source.driver, source.dsn)
+        for candidate in self.storages:
+            if candidate.id == StorageRepository.storage_id:
+                primary = self.storages[0]
+
+        if not primary:
+            raise LookupError('primary storage not found')
+
+        driver = await get_driver(primary.driver, primary.dsn)
 
         for repository in self.repositories.values():
             if isinstance(repository, BucketRepository):
                 await repository.bootstrap(driver)
             if isinstance(repository, StorageRepository):
-                await repository.bootstrap(driver, source)
+                await repository.bootstrap(driver, primary)
 
     async def find_or_create(
         self,

diff --git a/registry/repository.py b/registry/repository.py
@@ -84,27 +84,31 @@ class BucketRepository(Repository):
     async def bootstrap(self, driver: Driver) -> None:
         bucket_row = await driver.find_or_create(
             entity=Bucket,
-            query={'id': BucketRepository.bucket_id},
+            query={
+                'id': BucketRepository.bucket_id
+            },
             data={
                 'bucket_id': BucketRepository.bucket_id,
                 'id': BucketRepository.bucket_id,
                 'key': '',
                 'repository': BucketRepository,
                 'status': BucketStatus.READY,
-                'storage_id': 1,
+                'storage_id': StorageRepository.storage_id,
             }
         )
 
         storage_row = await driver.find_or_create(
             entity=Bucket,
-            query={'id': StorageRepository.bucket_id},
+            query={
+                'id': StorageRepository.bucket_id
+            },
             data={
                 'bucket_id': BucketRepository.bucket_id,
                 'id': StorageRepository.bucket_id,
                 'key': '',
                 'repository': StorageRepository,
                 'status': BucketStatus.READY,
-                'storage_id': 1,
+                'storage_id': StorageRepository.storage_id,
             }
         )
 
@@ -114,13 +118,19 @@ async def bootstrap(self, driver: Driver) -> None:
 
 class StorageRepository(Repository):
     bucket_id: int = 2
+    storage_id: int = 1
     entities = [Storage]
 
     async def bootstrap(self, driver: Driver, storage: Storage) -> None:
+        if storage.id != self.storage_id:
+            raise ValueError(f'Invalid storage_id: {storage.id}')
         await driver.find_or_create(
             entity=Storage,
-            query={'id': storage.id},
+            query=dict(
+                id=storage.id,
+            ),
             data=dict(
-                bucket_id=StorageRepository.bucket_id, **storage.__dict__
+                bucket_id=StorageRepository.bucket_id,
+                **storage.__dict__
             ),
         )
diff --git a/tests/test_repository.py b/tests/test_repository.py
@@ -3,9 +3,10 @@
 from pytest import mark
 
 from registry.drivers import get_driver
-from registry.entity import Entity
+from registry.entity import Entity, Storage
 from registry.registry import Registry
 from registry.repository import Index, Repository
+from registry.schema import StorageClass, StorageDriver
 
 
 @dataclass
@@ -30,8 +31,12 @@ class ActionRepository(Repository):
 
 
 @mark.asyncio
-async def test_hello():
-    registry = Registry()
+@mark.parametrize("storage", [
+    Storage(1, StorageClass.HOT, StorageDriver.MEMORY, '1'),
+    Storage(1, StorageClass.HOT, StorageDriver.MEMORY, '2'),
+])
+async def test_hello(storage: Storage):
+    registry = Registry([storage])
     assert len(await registry.find(Action)) == 0
 
     # create two actions