athena-intelligence 0.1.101__tar.gz → 0.1.102__tar.gz

{athena_intelligence-0.1.101 → athena_intelligence-0.1.102}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: athena-intelligence
-Version: 0.1.101
+Version: 0.1.102
 Summary: Athena Intelligence Python Library
 Requires-Python: >=3.8,<4.0
 Classifier: Intended Audience :: Developers

{athena_intelligence-0.1.101 → athena_intelligence-0.1.102}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "athena-intelligence"
-version = "0.1.101"
+version = "0.1.102"
 description = "Athena Intelligence Python Library"
 readme = "README.md"
 authors = []

{athena_intelligence-0.1.101 → athena_intelligence-0.1.102}/src/athena/client.py

@@ -1,14 +1,17 @@
+import enum
 import io
 import os
 import typing
 import warnings
 
 import httpx
+from typing import cast, List, Union
 from typing_extensions import TypeVar, ParamSpec
 
 from .base_client import BaseAthena, AsyncBaseAthena
 from .environment import AthenaEnvironment
 from .tools.client import AsyncToolsClient, ToolsClient
+from .query.client import AsyncQueryClient, QueryClient
 from .types.data_frame_request_out import DataFrameRequestOut
 
 if typing.TYPE_CHECKING:
@@ -32,6 +35,10 @@ def _inherit_signature_and_doc(
     return decorator
 
 
+class SpecialEnvironments(enum.Enum):
+    AUTODETECT_ENVIRONMENT = 'AUTO'
+
+
 class WrappedToolsClient(ToolsClient):
 
     def get_file(self, asset_id: str) -> io.BytesIO:
@@ -89,6 +96,25 @@ class WrappedToolsClient(ToolsClient):
         return _to_pandas_df(file_bytes, *args, **kwargs)
 
 
+class WrappedQueryClient(QueryClient):
+
+    @_inherit_signature_and_doc(QueryClient.execute, {'DataFrameRequestOut': 'pd.DataFrame'})
+    def execute(self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs) -> 'pd.DataFrame':
+        _check_pandas_installed()
+        model = super().execute(
+            sql_command=sql_command,
+            database_asset_ids=database_asset_ids,
+            **kwargs
+        )
+        return _read_json_frame(model)
+
+    @_inherit_signature_and_doc(QueryClient.execute_snippet, {'DataFrameRequestOut': 'pd.DataFrame'})
+    def execute_snippet(self, *, snippet_asset_id: str, **kwargs) -> 'pd.DataFrame':
+        _check_pandas_installed()
+        model = super().execute_snippet(snippet_asset_id=snippet_asset_id, **kwargs)
+        return _read_json_frame(model)
+
+
 def _add_docs_for_async_variant(obj):
     def decorator(decorated):
         doc = obj.__doc__
@@ -123,6 +149,25 @@ class WrappedAsyncToolsClient(AsyncToolsClient):
         return _to_pandas_df(file_bytes, *args, **kwargs)
 
 
+class WrappedAsyncQueryClient(AsyncQueryClient):
+
+    @_inherit_signature_and_doc(AsyncQueryClient.execute, {'DataFrameRequestOut': 'pd.DataFrame'})
+    async def execute(self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs) -> 'pd.DataFrame':
+        _check_pandas_installed()
+        model = await super().execute(
+            sql_command=sql_command,
+            database_asset_ids=database_asset_ids,
+            **kwargs
+        )
+        return _read_json_frame(model)
+
+    @_inherit_signature_and_doc(AsyncQueryClient.execute_snippet, {'DataFrameRequestOut': 'pd.DataFrame'})
+    async def execute_snippet(self, *, snippet_asset_id: str, **kwargs) -> 'pd.DataFrame':
+        _check_pandas_installed()
+        model = await super().execute_snippet(snippet_asset_id=snippet_asset_id, **kwargs)
+        return _read_json_frame(model)
+
+
 class Athena(BaseAthena):
     """
     Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propogate to these functions.
@@ -133,13 +178,9 @@ class Athena(BaseAthena):
         The base url to use for requests from the client.
 
     environment : AthenaEnvironment
-        The environment to use for requests from the client. from .environment import AthenaEnvironment
-
-
-
-        Defaults to AthenaEnvironment.PRODUCTION
-
+        The environment to use for requests from the client.
 
+        Defaults to `AthenaEnvironment.PRODUCTION` when outside of athena notebook environment.
 
     api_key: typing.Optional[str].  The API key. Required when used outside of the athena notebook environment.
     timeout : typing.Optional[float]
@@ -163,7 +204,7 @@ class Athena(BaseAthena):
         self,
         *,
         base_url: typing.Optional[str] = None,
-        environment: AthenaEnvironment = AthenaEnvironment.PRODUCTION,
+        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,    # type: ignore[arg-type]
         api_key: typing.Optional[str] = None,
         timeout: typing.Optional[float] = 60,
         httpx_client: typing.Optional[httpx.Client] = None
@@ -176,14 +217,24 @@ class Athena(BaseAthena):
                     "Athena() missing 1 required keyword-only argument: 'api_key'"
                     ' (ATHENA_API_KEY environment variable not found)'
                 )
+        if environment == SpecialEnvironments.AUTODETECT_ENVIRONMENT:
+            if 'ATHENA_API_URL' in os.environ:
+
+                class AutodetectedEnvironments:
+                    CURRENT = os.environ['ATHENA_API_URL']
+
+                environment = cast(AthenaEnvironment, AutodetectedEnvironments.CURRENT)
+            else:
+                environment = AthenaEnvironment.PRODUCTION
         super().__init__(
             base_url=base_url,
-            environment=environment,
+            environment=environment,    # type: ignore[arg-type]
             api_key=api_key,
             timeout=timeout,
             httpx_client=httpx_client,
         )
         self.tools = WrappedToolsClient(client_wrapper=self._client_wrapper)
+        self.query = WrappedQueryClient(client_wrapper=self._client_wrapper)
 
 
 class AsyncAthena(AsyncBaseAthena):
@@ -197,13 +248,9 @@ class AsyncAthena(AsyncBaseAthena):
         The base url to use for requests from the client.
 
     environment : AthenaEnvironment
-        The environment to use for requests from the client. from .environment import AthenaEnvironment
-
-
-
-        Defaults to AthenaEnvironment.PRODUCTION
-
+        The environment to use for requests from the client.
 
+        Defaults to `AthenaEnvironment.PRODUCTION` when outside of athena notebook environment.
 
     api_key: typing.Optional[str].  The API key. Required when used outside of the athena notebook environment.
     timeout : typing.Optional[float]
@@ -227,7 +274,7 @@ class AsyncAthena(AsyncBaseAthena):
         self,
         *,
         base_url: typing.Optional[str] = None,
-        environment: AthenaEnvironment = AthenaEnvironment.PRODUCTION,
+        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,    # type: ignore[arg-type]
         api_key: typing.Optional[str] = None,
         timeout: typing.Optional[float] = 60,
         httpx_client: typing.Optional[httpx.AsyncClient] = None
@@ -240,14 +287,24 @@ class AsyncAthena(AsyncBaseAthena):
                     "AsyncAthena() missing 1 required keyword-only argument: 'api_key'"
                     ' (ATHENA_API_KEY environment variable not found)'
                 )
+        if environment == SpecialEnvironments.AUTODETECT_ENVIRONMENT:
+            if 'ATHENA_API_URL' in os.environ:
+
+                class AutodetectedEnvironments:
+                    CURRENT = os.environ['ATHENA_API_URL']
+
+                environment = cast(AthenaEnvironment, AutodetectedEnvironments.CURRENT)
+            else:
+                environment = AthenaEnvironment.PRODUCTION
         super().__init__(
             base_url=base_url,
-            environment=environment,
+            environment=environment,    # type: ignore[arg-type]
             api_key=api_key,
             timeout=timeout,
             httpx_client=httpx_client,
         )
         self.tools = WrappedAsyncToolsClient(client_wrapper=self._client_wrapper)
+        self.query = WrappedAsyncQueryClient(client_wrapper=self._client_wrapper)
 
 
 def _read_json_frame(model: DataFrameRequestOut) -> 'pd.DataFrame':

{athena_intelligence-0.1.101 → athena_intelligence-0.1.102}/src/athena/core/client_wrapper.py

@@ -17,7 +17,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "athena-intelligence",
-            "X-Fern-SDK-Version": "0.1.101",
+            "X-Fern-SDK-Version": "0.1.102",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

{athena_intelligence-0.1.101 → athena_intelligence-0.1.102}/src/athena/types/sql_unauthorized_error.py

@@ -8,6 +8,8 @@ from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
 
 
 class SqlUnauthorizedError(pydantic_v1.BaseModel):
+    message: str
+
     def json(self, **kwargs: typing.Any) -> str:
         kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
         return super().json(**kwargs_with_defaults)