Get antfly-sdk runnable

src/antfly/client.py

@@ -1,41 +1,44 @@
 """Main client interface for Antfly SDK."""
 
-from typing import Any, Optional, cast
+from typing import Any, Optional, Union, cast
 
 from httpx import Timeout
 
 from antfly.client_generated import Client
-from antfly.client_generated.api.api_table import (
-    batch,
+from antfly.client_generated.api.data_operations import batch_write, lookup_key
+from antfly.client_generated.api.table_management import (
     create_table,
     drop_table,
     get_table,
     list_tables,
-    lookup_key,
-    query_table,
 )
 from antfly.client_generated.client import AuthenticatedClient
 from antfly.client_generated.models import (
     BatchRequest,
     BatchRequestInserts,
+    BatchRequestInsertsAdditionalProperty,
     CreateTableRequest,
     CreateTableRequestIndexes,
     Error,
-    QueryRequest,
-    QueryRequestFullTextSearch,
-    QueryResponses,
     Table,
     TableSchema,
     TableStatus,
 )
-from antfly.client_generated.types import UNSET
+from antfly.client_generated.types import UNSET, Unset
 
 from .exceptions import AntflyException
 
+# Type alias: generated API functions type-hint AuthenticatedClient but work with
+# Client too since both have identical interfaces (get_httpx_client, etc.)
+# We use Client with basic auth via httpx_args instead of token-based auth.
+ApiClient = Union[Client, AuthenticatedClient]
+
 
 class AntflyClient:
     """High-level client for interacting with Antfly database."""
 
+    _client: ApiClient
+
     def __init__(
         self,
         base_url: str,
@@ -58,7 +61,7 @@ def __init__(
         if username and password:
             httpx_args["auth"] = (username, password)
 
-        self._client = Client(
+        self._client: ApiClient = Client(
             base_url=self.base_url,
             timeout=Timeout(timeout),
             httpx_args=httpx_args,
@@ -96,7 +99,7 @@ def create_table(
 
         response = create_table.sync(
             table_name=name,
-            client=cast(AuthenticatedClient, self._client),
+            client=self._client,  # type: ignore[arg-type]
             body=request,
         )
 
@@ -117,7 +120,7 @@ def list_tables(self) -> list[TableStatus]:
         Raises:
             AntflyException: If listing tables fails
         """
-        response = list_tables.sync(client=cast(AuthenticatedClient, self._client))
+        response = list_tables.sync(client=self._client)  # type: ignore[arg-type]
 
         if isinstance(response, Error):
             raise AntflyException(f"Failed to list tables: {response.error}")
@@ -141,7 +144,7 @@ def get_table(self, name: str) -> TableStatus:
         """
         response = get_table.sync(
             table_name=name,
-            client=cast(AuthenticatedClient, self._client),
+            client=self._client,  # type: ignore[arg-type]
         )
 
         if isinstance(response, Error):
@@ -163,77 +166,15 @@ def drop_table(self, name: str) -> None:
         """
         response = drop_table.sync(
             table_name=name,
-            client=cast(AuthenticatedClient, self._client),
+            client=self._client,  # type: ignore[arg-type]
         )
 
         if isinstance(response, Error):
             raise AntflyException(f"Failed to drop table '{name}': {response.error}")
         if response is None:
             raise AntflyException(f"Failed to drop table '{name}'")
 
-    # Query operations
-
-    def query(
-        self,
-        table: Optional[str] = None,
-        full_text_search: Optional[dict[str, Any]] = None,
-        semantic_search: Optional[str] = None,
-        filter_prefix: Optional[str] = None,
-        limit: int = 10,
-        offset: int = 0,
-        **kwargs: Any,
-    ) -> QueryResponses:
-        """
-        Query a table or perform global query.
-
-        Args:
-            table: Table name (optional for global query)
-            full_text_search: Full-text search query
-            semantic_search: Semantic search query
-            filter_prefix: Key prefix filter
-            limit: Maximum number of results
-            offset: Number of results to skip
-            **kwargs: Additional query parameters
-
-        Returns:
-            Query result object
-
-        Raises:
-            AntflyException: If query fails
-        """
-        request = QueryRequest(
-            table=table if table is not None else UNSET,
-            full_text_search=(
-                cast(QueryRequestFullTextSearch, full_text_search) if full_text_search is not None else UNSET
-            ),
-            semantic_search=semantic_search if semantic_search is not None else UNSET,
-            filter_prefix=filter_prefix if filter_prefix is not None else UNSET,
-            limit=limit,
-            offset=offset,
-            **kwargs,
-        )
-
-        if table:
-            response = query_table.sync(
-                table_name=table,
-                client=cast(AuthenticatedClient, self._client),
-                body=request,
-            )
-        else:
-            # Use global query endpoint
-            from antfly.client_generated.api.api_table import global_query
-
-            response = global_query.sync(
-                client=cast(AuthenticatedClient, self._client),
-                body=request,
-            )
-
-        if isinstance(response, Error):
-            raise AntflyException(f"Query failed: {response.error}")
-        if response is None:
-            raise AntflyException("Query failed")
-
-        return response
+    # Data operations
 
     def get(self, table: str, key: str) -> dict[str, Any]:
         """
@@ -252,7 +193,7 @@ def get(self, table: str, key: str) -> dict[str, Any]:
         response = lookup_key.sync(
             table_name=table,
             key=key,
-            client=cast(AuthenticatedClient, self._client),
+            client=self._client,  # type: ignore[arg-type]
         )
 
         if isinstance(response, Error):
@@ -279,14 +220,23 @@ def batch(
         Raises:
             AntflyException: If batch operation fails
         """
+        # Convert plain dict to proper BatchRequestInserts model
+        inserts_model: BatchRequestInserts | Unset = UNSET
+        if inserts is not None:
+            inserts_model = BatchRequestInserts()
+            for key, value in inserts.items():
+                prop = BatchRequestInsertsAdditionalProperty()
+                prop.additional_properties = value
+                inserts_model[key] = prop
+
         request = BatchRequest(
-            inserts=cast(BatchRequestInserts, inserts) if inserts is not None else UNSET,
-            deletes=deletes or [],
+            inserts=inserts_model,
+            deletes=deletes if deletes is not None else UNSET,
         )
 
-        response = batch.sync(
+        response = batch_write.sync(
             table_name=table,
-            client=cast(AuthenticatedClient, self._client),
+            client=self._client,  # type: ignore[arg-type]
             body=request,
         )
 

src/antfly/client_generated/api/data_operations/batch_write.py

@@ -34,10 +34,9 @@ def _get_kwargs(
 def _parse_response(
     *, client: Union[AuthenticatedClient, Client], response: httpx.Response
 ) -> Optional[Union[BatchResponse, Error]]:
-    if response.status_code == 201:
-        response_201 = BatchResponse.from_dict(response.json())
-
-        return response_201
+    # Accept both 200 and 201 as success - server may return either
+    if response.status_code in (200, 201):
+        return BatchResponse.from_dict(response.json())
 
     if response.status_code == 400:
         response_400 = Error.from_dict(response.json())