feat: make AuthEntryReference typed

python/cocoindex/auth_registry.py

@@ -3,16 +3,19 @@
 """
 
 from dataclasses import dataclass
+from typing import Generic, TypeVar
 
 from . import _engine
 from .convert import dump_engine_object
 
+T = TypeVar("T")
+
 @dataclass
-class AuthEntryReference:
+class AuthEntryReference(Generic[T]):
     """Reference an auth entry by its key."""
     key: str
 
-def add_auth_entry(key: str, value) -> AuthEntryReference:
+def add_auth_entry(key: str, value: T) -> AuthEntryReference[T]:
     """Add an auth entry to the registry. Returns its reference."""
     _engine.add_auth_entry(key, dump_engine_object(value))
     return AuthEntryReference(key)

python/cocoindex/storages.py

@@ -5,11 +5,11 @@
 from . import op
 from . import index
 from .auth_registry import AuthEntryReference
+from .setting import DatabaseConnectionSpec
 
 class Postgres(op.StorageSpec):
     """Storage powered by Postgres and pgvector."""
-
-    database: AuthEntryReference | None = None
+    database: AuthEntryReference[DatabaseConnectionSpec] | None = None
     table_name: str | None = None
 
 @dataclass
@@ -72,15 +72,14 @@ class Relationships:
 
 class Neo4j(op.StorageSpec):
     """Graph storage powered by Neo4j."""
-
-    connection: AuthEntryReference
+    connection: AuthEntryReference[Neo4jConnection]
     mapping: Nodes | Relationships
 
 class Neo4jDeclaration(op.DeclarationSpec):
     """Declarations for Neo4j."""
 
     kind = "Neo4j"
-    connection: AuthEntryReference
+    connection: AuthEntryReference[Neo4jConnection]
     nodes_label: str
     primary_key_fields: Sequence[str]
     vector_indexes: Sequence[index.VectorIndexDef] = ()