shaped 2.0.2__py3-none-any.whl → 2.0.4__py3-none-any.whl

shaped/client.py

@@ -3,45 +3,43 @@ import socket
 from typing import Any, Dict, List, Optional, Union
 
 from shaped.autogen.api.engine_api import EngineApi
+from shaped.autogen.api.query_api import QueryApi
 from shaped.autogen.api.table_api import TableApi
 from shaped.autogen.api.view_api import ViewApi
-from shaped.autogen.api.query_api import QueryApi
-from shaped.autogen.models.query_request import QueryRequest
-from shaped.autogen.models.rank_query_config import RankQueryConfig
-from shaped.autogen.models.saved_query_request import SavedQueryRequest
+from shaped.autogen.api_client import ApiClient
+from shaped.autogen.configuration import Configuration
+from shaped.autogen.models.ai_enrichment_view_config import AIEnrichmentViewConfig
+from shaped.autogen.models.create_table_response import CreateTableResponse
+from shaped.autogen.models.create_view_response import CreateViewResponse
+from shaped.autogen.models.delete_engine_response import DeleteEngineResponse
+from shaped.autogen.models.delete_table_response import DeleteTableResponse
+from shaped.autogen.models.delete_view_response import DeleteViewResponse
 from shaped.autogen.models.engine_config_v2 import EngineConfigV2
-from shaped.autogen.models.setup_engine_response import SetupEngineResponse
-from shaped.autogen.models.list_engines_response import ListEnginesResponse
 from shaped.autogen.models.engine_details_response import EngineDetailsResponse
-from shaped.autogen.models.delete_engine_response import DeleteEngineResponse
-from shaped.autogen.models.request import Request as TableRequest
-from shaped.autogen.models.create_table_response import CreateTableResponse
-from shaped.autogen.models.update_table_response import UpdateTableResponse
+from shaped.autogen.models.list_engines_response import ListEnginesResponse
 from shaped.autogen.models.list_tables_response import ListTablesResponse
-from shaped.autogen.models.table_insert_arguments import TableInsertArguments
-from shaped.autogen.models.table_insert_response import TableInsertResponse
-from shaped.autogen.models.delete_table_response import DeleteTableResponse
-from shaped.autogen.models.request1 import Request1 as ViewRequest
-from shaped.autogen.models.create_view_response import CreateViewResponse
-from shaped.autogen.models.update_view_response import UpdateViewResponse
 from shaped.autogen.models.list_views_response import ListViewsResponse
+from shaped.autogen.models.parameters_value import ParametersValue
+from shaped.autogen.models.query1 import Query1
+from shaped.autogen.models.query_request import QueryRequest
+from shaped.autogen.models.query_result import QueryResult
+from shaped.autogen.models.rank_query_config import RankQueryConfig
+from shaped.autogen.models.request import Request as TableRequest
+from shaped.autogen.models.request1 import Request1 as ViewRequest
 from shaped.autogen.models.response_get_view_details_views_view_name_get import (
     ResponseGetViewDetailsViewsViewNameGet,
 )
-from shaped.autogen.models.delete_view_response import DeleteViewResponse
-from shaped.autogen.models.query_result import QueryResult
 from shaped.autogen.models.saved_query_info_response import SavedQueryInfoResponse
 from shaped.autogen.models.saved_query_list_response import SavedQueryListResponse
+from shaped.autogen.models.saved_query_request import SavedQueryRequest
+from shaped.autogen.models.setup_engine_response import SetupEngineResponse
 from shaped.autogen.models.sql_view_config import SQLViewConfig
-from shaped.autogen.models.ai_enrichment_view_config import AIEnrichmentViewConfig
-from shaped.autogen.models.parameters_value import ParametersValue
-from shaped.autogen.models.query1 import Query1
-
-from shaped.autogen.api_client import ApiClient
-from shaped.autogen.configuration import Configuration
-from shaped.query_builder import RankQueryBuilder
+from shaped.autogen.models.table_insert_arguments import TableInsertArguments
+from shaped.autogen.models.table_insert_response import TableInsertResponse
+from shaped.autogen.models.update_table_response import UpdateTableResponse
+from shaped.autogen.models.update_view_response import UpdateViewResponse
 from shaped.config_builders import Engine
-
+from shaped.query_builder import RankQueryBuilder
 
 # Type aliases for better developer experience.
 ViewConfig = Union[SQLViewConfig, AIEnrichmentViewConfig]
@@ -88,7 +86,7 @@ class Client:
 
         # Enable response compression for high-performance query endpoints.
         # Reduces bandwidth and improves performance for large query results.
-        self._api_client.set_default_header('Accept-Encoding', 'gzip, deflate')
+        self._api_client.set_default_header("Accept-Encoding", "gzip, deflate")
 
         # Initialize V2 API instances.
         self._engine_api = EngineApi(self._api_client)
@@ -100,16 +98,12 @@ class Client:
         if not self._logger.hasHandlers():
             console_handler = logging.StreamHandler()
             console_handler.setLevel(logging.INFO)
-            console_format = logging.Formatter(
-                "%(name)s - %(levelname)s - %(message)s"
-            )
+            console_format = logging.Formatter("%(name)s - %(levelname)s - %(message)s")
             console_handler.setFormatter(console_format)
             self._logger.addHandler(console_handler)
             self._logger.setLevel(logging.INFO)
 
-    def _query_api_call(
-        self, method: Any, *args: Any, **kwargs: Any
-    ) -> Any:
+    def _query_api_call(self, method: Any, *args: Any, **kwargs: Any) -> Any:
         """
         Helper for Query API calls that require headers.
 
@@ -219,7 +213,9 @@ class Client:
 
     # Table API Methods
 
-    def create_table(self, table_config: Union[TableRequest, Dict[str, Any]]) -> CreateTableResponse:
+    def create_table(
+        self, table_config: Union[TableRequest, Dict[str, Any]]
+    ) -> CreateTableResponse:
         """
         Create a new table.
 
@@ -247,7 +243,9 @@ class Client:
             x_api_key=self._api_key, request=request
         )
 
-    def update_table(self, table_config: Union[TableRequest, Dict[str, Any]]) -> UpdateTableResponse:
+    def update_table(
+        self, table_config: Union[TableRequest, Dict[str, Any]]
+    ) -> UpdateTableResponse:
         """
         Update an existing table.
 
@@ -298,7 +296,9 @@ class Client:
         self._logger.debug("Inserting rows into table %s", table_name)
         insert_args = TableInsertArguments(data=rows)
         return self._table_api.post_table_insert_tables_table_name_insert_post(
-            table_name=table_name, x_api_key=self._api_key, table_insert_arguments=insert_args
+            table_name=table_name,
+            x_api_key=self._api_key,
+            table_insert_arguments=insert_args,
         )
 
     def delete_table(self, table_name: str) -> DeleteTableResponse:
@@ -318,7 +318,9 @@ class Client:
 
     # View API Methods
 
-    def create_view(self, view_config: Union[ViewConfig, Dict[str, Any]]) -> CreateViewResponse:
+    def create_view(
+        self, view_config: Union[ViewConfig, Dict[str, Any]]
+    ) -> CreateViewResponse:
         """
         Create a new view.
 
@@ -344,7 +346,9 @@ class Client:
             x_api_key=self._api_key, request1=request
         )
 
-    def update_view(self, view_config: Union[ViewConfig, Dict[str, Any]]) -> UpdateViewResponse:
+    def update_view(
+        self, view_config: Union[ViewConfig, Dict[str, Any]]
+    ) -> UpdateViewResponse:
         """
         Update an existing view.
 
@@ -379,9 +383,7 @@ class Client:
         self._logger.debug("Listing views")
         return self._view_api.get_views_views_get(x_api_key=self._api_key)
 
-    def get_view(
-        self, view_name: str
-    ) -> ResponseGetViewDetailsViewsViewNameGet:
+    def get_view(self, view_name: str) -> ResponseGetViewDetailsViewsViewNameGet:
         """
         Get details for a specific view.
 
@@ -498,6 +500,7 @@ class Client:
                 query_dict = query.model_dump(by_alias=True)
             else:
                 query_dict = dict(query)
+
             # Convert sets to lists for JSON serialization.
             def convert_sets_to_lists(obj):
                 """Recursively convert sets to lists for JSON serialization."""
@@ -512,7 +515,9 @@ class Client:
             query_dict_cleaned = convert_sets_to_lists(query_dict)
             query_obj = Query1.from_dict(query_dict_cleaned)
         else:
-            raise TypeError(f"Invalid query type: {type(query)}. Expected RankQueryConfig, dict, or str.")
+            raise TypeError(
+                f"Invalid query type: {type(query)}. Expected RankQueryConfig, dict, or str."
+            )
 
         query_request = QueryRequest(
             query=query_obj,
@@ -600,7 +605,9 @@ class Client:
         Returns:
             SavedQueryInfoResponse with query information.
         """
-        self._logger.debug("Getting saved query info for %s in engine %s", query_name, engine_name)
+        self._logger.debug(
+            "Getting saved query info for %s in engine %s", query_name, engine_name
+        )
         return self._query_api_call(
             self._query_api.get_saved_query_info_queries_query_name_get,
             engine_name=engine_name,