vgi-python 0.31.2__py3-none-any.whl → 0.31.3__py3-none-any.whl

vgi/client/catalog_mixin.py

@@ -56,6 +56,7 @@ from vgi.catalog import (
     ColumnStatistics,
     CopyFromFormatInfo,
     FunctionInfo,
+    IndexInfo,
     MacroInfo,
     MacroType,
     OnConflict,
@@ -113,15 +114,19 @@ class CatalogClientMixin:
 
     # Type hints for attributes expected from Client
     server_path: str | Sequence[str]
-    _transport: Literal["subprocess", "http", "tcp", "launch"]
+    _transport: Literal["subprocess", "http", "tcp", "iroh", "httpi", "launch"]
+    _iroh_endpoint: str | None
+    _iroh_options: dict[str, Any]
     _base_url: str | None
     _tcp_host: str | None
     _tcp_port: int | None
+    _tcp_proxy: str | None
     _launch_argv: tuple[str, ...] | None
     _launch_idle_timeout: float
     _launch_state_dir: str | None
     _launch_socket_path: str | None
     _external_location: Any | None
+    _accepted_max_response_bytes: int | None
 
     def _worker_argv(self) -> list[str]:
         """The worker command as an argv list, however ``server_path`` was given.
@@ -162,26 +167,53 @@ class CatalogClientMixin:
         """
         try:
             transport = getattr(self, "_transport", "subprocess")
-            if transport == "http":
+            if transport in ("http", "httpi"):
                 from vgi_rpc.http import http_connect
 
                 httpx_client = self._get_or_create_httpx_client()
-                with http_connect(
+                if transport == "httpi":
+                    connection = http_connect(
+                        VgiProtocol,  # type: ignore[type-abstract]
+                        base_url=self._base_url,
+                        client=httpx_client,
+                        prefix=getattr(self, "_http_prefix", ""),
+                        external_location=getattr(self, "_external_location", None),
+                        accepted_max_response_bytes=self._accepted_max_response_bytes,
+                    )
+                else:
+                    connection = http_connect(
+                        VgiProtocol,  # type: ignore[type-abstract]
+                        base_url=self._base_url,
+                        client=httpx_client,
+                        external_location=getattr(self, "_external_location", None),
+                        accepted_max_response_bytes=self._accepted_max_response_bytes,
+                    )
+                with connection as proxy:
+                    yield proxy
+            elif transport == "iroh":
+                from vgi_rpc.iroh import iroh_connect
+
+                assert self._iroh_endpoint is not None
+                with iroh_connect(
                     VgiProtocol,  # type: ignore[type-abstract]
-                    base_url=self._base_url,
-                    client=httpx_client,
+                    self._iroh_endpoint,
                     external_location=getattr(self, "_external_location", None),
+                    **self._iroh_options,
                 ) as proxy:
                     yield proxy
             elif transport == "tcp":
                 from vgi_rpc.rpc import tcp_connect
 
                 assert self._tcp_host is not None and self._tcp_port is not None
+                proxy_options: dict[str, Any] = {}
+                if self._tcp_proxy is not None:
+                    proxy_options["proxy"] = self._tcp_proxy
                 with tcp_connect(
                     VgiProtocol,  # type: ignore[type-abstract]
                     self._tcp_host,
                     self._tcp_port,
                     external_location=getattr(self, "_external_location", None),
+                    **proxy_options,
                 ) as proxy:
                     yield proxy
             elif transport == "launch":
@@ -523,6 +555,16 @@ class CatalogClientMixin:
         type: Literal[SchemaObjectType.SCALAR_MACRO, SchemaObjectType.TABLE_MACRO],
     ) -> Sequence[MacroInfo]: ...
 
+    @overload
+    def schema_contents(
+        self,
+        *,
+        attach_opaque_data: AttachOpaqueData,
+        transaction_opaque_data: TransactionOpaqueData | None = None,
+        name: str,
+        type: Literal[SchemaObjectType.INDEX],
+    ) -> Sequence[IndexInfo]: ...
+
     @overload
     def schema_contents(
         self,
@@ -531,7 +573,7 @@ class CatalogClientMixin:
         transaction_opaque_data: TransactionOpaqueData | None = None,
         name: str,
         type: SchemaObjectType,
-    ) -> Sequence[TableInfo | ViewInfo | FunctionInfo | MacroInfo]: ...
+    ) -> Sequence[TableInfo | ViewInfo | FunctionInfo | MacroInfo | IndexInfo]: ...
 
     def schema_contents(
         self,
@@ -540,8 +582,8 @@ class CatalogClientMixin:
         transaction_opaque_data: TransactionOpaqueData | None = None,
         name: str,
         type: SchemaObjectType,
-    ) -> Sequence[TableInfo | ViewInfo | FunctionInfo | MacroInfo]:
-        """List contents of a schema (tables, views, functions, macros).
+    ) -> Sequence[TableInfo | ViewInfo | FunctionInfo | MacroInfo | IndexInfo]:
+        """List contents of a schema (tables, views, functions, macros, indexes).
 
         Args:
             attach_opaque_data: The attachment ID from catalog_attach.
@@ -554,9 +596,11 @@ class CatalogClientMixin:
                 - `SchemaObjectType.TABLE_FUNCTION`: Return only table functions
                 - `SchemaObjectType.SCALAR_MACRO`: Return only scalar macros
                 - `SchemaObjectType.TABLE_MACRO`: Return only table macros
+                - `SchemaObjectType.INDEX`: Return only indexes
 
         Returns:
-            List of [`TableInfo`][], [`ViewInfo`][], [`FunctionInfo`][], or [`MacroInfo`][] depending on the type.
+            List of [`TableInfo`][], [`ViewInfo`][], [`FunctionInfo`][], [`MacroInfo`][],
+            or [`IndexInfo`][] depending on the type.
 
         """
         with self._catalog_connect() as proxy:
@@ -579,6 +623,17 @@ class CatalogClientMixin:
                     type=type,
                     transaction_opaque_data=transaction_opaque_data,
                 ).to_infos()
+            elif type == SchemaObjectType.INDEX:
+                # Its own RPC, not a variant of catalog_schema_contents_functions
+                # (the `else` branch below) -- an index has no FunctionInfo-shaped
+                # representation at all. Before this branch existed, `type=INDEX`
+                # silently fell through to the functions RPC and returned
+                # nonsense (or an empty list) instead of raising or working.
+                return proxy.catalog_schema_contents_indexes(
+                    attach_opaque_data=attach_opaque_data,
+                    name=name,
+                    transaction_opaque_data=transaction_opaque_data,
+                ).to_infos()
             else:
                 return proxy.catalog_schema_contents_functions(
                     attach_opaque_data=attach_opaque_data,