athena-intelligence 0.1.100__py3-none-any.whl → 0.1.102__py3-none-any.whl

athena/__init__.py

@@ -18,6 +18,7 @@ from .errors import (
 )
 from . import query, tools
 from .environment import AthenaEnvironment
+from .query import QueryExecuteRequestDatabaseAssetIds
 from .tools import ToolsDataFrameRequestColumnsItem
 from .version import __version__
 
@@ -31,6 +32,7 @@ __all__ = [
     "FileFetchError",
     "InternalServerError",
     "NotFoundError",
+    "QueryExecuteRequestDatabaseAssetIds",
     "SqlUnauthorizedError",
     "ToolsDataFrameRequestColumnsItem",
     "UnauthorizedError",

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/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.100",
+            "X-Fern-SDK-Version": "0.1.102",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

athena/query/__init__.py

@@ -1,2 +1,5 @@
 # This file was auto-generated by Fern from our API Definition.
 
+from .types import QueryExecuteRequestDatabaseAssetIds
+
+__all__ = ["QueryExecuteRequestDatabaseAssetIds"]

athena/query/client.py

@@ -5,6 +5,7 @@ from json.decoder import JSONDecodeError
 
 from ..core.api_error import ApiError
 from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
+from ..core.jsonable_encoder import jsonable_encoder
 from ..core.pydantic_utilities import pydantic_v1
 from ..core.request_options import RequestOptions
 from ..errors.internal_server_error import InternalServerError
@@ -12,6 +13,7 @@ from ..errors.unauthorized_error import UnauthorizedError
 from ..errors.unprocessable_entity_error import UnprocessableEntityError
 from ..types.data_frame_request_out import DataFrameRequestOut
 from ..types.sql_unauthorized_error import SqlUnauthorizedError
+from .types.query_execute_request_database_asset_ids import QueryExecuteRequestDatabaseAssetIds
 
 
 class QueryClient:
@@ -19,6 +21,64 @@ class QueryClient:
         self._client_wrapper = client_wrapper
 
     def execute(
+        self,
+        *,
+        sql_command: str,
+        database_asset_ids: QueryExecuteRequestDatabaseAssetIds,
+        request_options: typing.Optional[RequestOptions] = None
+    ) -> DataFrameRequestOut:
+        """
+        Get the result of an SQL query over given assets.
+
+        Parameters
+        ----------
+        sql_command : str
+            SQL query string
+
+        database_asset_ids : QueryExecuteRequestDatabaseAssetIds
+            Single ID or list of asset IDs
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        DataFrameRequestOut
+            Successful Response
+
+        Examples
+        --------
+        from athena.client import Athena
+
+        client = Athena(
+            api_key="YOUR_API_KEY",
+        )
+        client.query.execute(
+            sql_command="sql_command",
+            database_asset_ids="database_asset_ids",
+        )
+        """
+        _response = self._client_wrapper.httpx_client.request(
+            "api/v0/query/sql/code/execute",
+            method="GET",
+            params={"sql_command": sql_command, "database_asset_ids": jsonable_encoder(database_asset_ids)},
+            request_options=request_options,
+        )
+        if 200 <= _response.status_code < 300:
+            return pydantic_v1.parse_obj_as(DataFrameRequestOut, _response.json())  # type: ignore
+        if _response.status_code == 401:
+            raise UnauthorizedError(pydantic_v1.parse_obj_as(SqlUnauthorizedError, _response.json()))  # type: ignore
+        if _response.status_code == 422:
+            raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
+        if _response.status_code == 500:
+            raise InternalServerError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
+        try:
+            _response_json = _response.json()
+        except JSONDecodeError:
+            raise ApiError(status_code=_response.status_code, body=_response.text)
+        raise ApiError(status_code=_response.status_code, body=_response_json)
+
+    def execute_snippet(
         self, *, snippet_asset_id: str, request_options: typing.Optional[RequestOptions] = None
     ) -> DataFrameRequestOut:
         """
@@ -43,7 +103,7 @@ class QueryClient:
         client = Athena(
             api_key="YOUR_API_KEY",
         )
-        client.query.execute(
+        client.query.execute_snippet(
             snippet_asset_id="snippet_asset_id",
         )
         """
@@ -73,6 +133,64 @@ class AsyncQueryClient:
         self._client_wrapper = client_wrapper
 
     async def execute(
+        self,
+        *,
+        sql_command: str,
+        database_asset_ids: QueryExecuteRequestDatabaseAssetIds,
+        request_options: typing.Optional[RequestOptions] = None
+    ) -> DataFrameRequestOut:
+        """
+        Get the result of an SQL query over given assets.
+
+        Parameters
+        ----------
+        sql_command : str
+            SQL query string
+
+        database_asset_ids : QueryExecuteRequestDatabaseAssetIds
+            Single ID or list of asset IDs
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        DataFrameRequestOut
+            Successful Response
+
+        Examples
+        --------
+        from athena.client import AsyncAthena
+
+        client = AsyncAthena(
+            api_key="YOUR_API_KEY",
+        )
+        await client.query.execute(
+            sql_command="sql_command",
+            database_asset_ids="database_asset_ids",
+        )
+        """
+        _response = await self._client_wrapper.httpx_client.request(
+            "api/v0/query/sql/code/execute",
+            method="GET",
+            params={"sql_command": sql_command, "database_asset_ids": jsonable_encoder(database_asset_ids)},
+            request_options=request_options,
+        )
+        if 200 <= _response.status_code < 300:
+            return pydantic_v1.parse_obj_as(DataFrameRequestOut, _response.json())  # type: ignore
+        if _response.status_code == 401:
+            raise UnauthorizedError(pydantic_v1.parse_obj_as(SqlUnauthorizedError, _response.json()))  # type: ignore
+        if _response.status_code == 422:
+            raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
+        if _response.status_code == 500:
+            raise InternalServerError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
+        try:
+            _response_json = _response.json()
+        except JSONDecodeError:
+            raise ApiError(status_code=_response.status_code, body=_response.text)
+        raise ApiError(status_code=_response.status_code, body=_response_json)
+
+    async def execute_snippet(
         self, *, snippet_asset_id: str, request_options: typing.Optional[RequestOptions] = None
     ) -> DataFrameRequestOut:
         """
@@ -97,7 +215,7 @@ class AsyncQueryClient:
         client = AsyncAthena(
             api_key="YOUR_API_KEY",
         )
-        await client.query.execute(
+        await client.query.execute_snippet(
             snippet_asset_id="snippet_asset_id",
         )
         """

athena/query/types/__init__.py

@@ -0,0 +1,5 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from .query_execute_request_database_asset_ids import QueryExecuteRequestDatabaseAssetIds
+
+__all__ = ["QueryExecuteRequestDatabaseAssetIds"]

athena/query/types/query_execute_request_database_asset_ids.py

@@ -0,0 +1,5 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+QueryExecuteRequestDatabaseAssetIds = typing.Union[str, typing.List[str]]

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)

{athena_intelligence-0.1.100.dist-info → athena_intelligence-0.1.102.dist-info}/METADATA

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

{athena_intelligence-0.1.100.dist-info → athena_intelligence-0.1.102.dist-info}/RECORD

@@ -1,9 +1,9 @@
-athena/__init__.py,sha256=_1vtiajlFbHjHlAngtGg_C_XyqinPbX9lbn2D-zjgWc,1094
+athena/__init__.py,sha256=lGo6ASUYoueZCu_oWrD7ruCsa0nLRNo9qZ_-2RV4WxY,1192
 athena/base_client.py,sha256=TlP599mOBvj7Tk8IpFe5dgrDM98GJu3lEQh_zwl4vtA,5439
-athena/client.py,sha256=2jxJ8mZWLa9LOH9XCxbxwLV6qVPXw53gi7GiYUTcpZg,9346
+athena/client.py,sha256=z58lccIOrRq-Qp_Vs3rxqSmBtESUwhHrLlhAZd-O0ao,12352
 athena/core/__init__.py,sha256=UFXpYzcGxWQUucU1TkjOQ9mGWN3A5JohluOIWVYKU4I,973
 athena/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-athena/core/client_wrapper.py,sha256=De-MIVXRtOZtBkHGrEavMifAIjhtpKVddLmCkL5sCzM,1806
+athena/core/client_wrapper.py,sha256=Q4TSQemyLKOkvsjr4up2880ddYGycZ-8PHxNfijtNow,1806
 athena/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 athena/core/file.py,sha256=sy1RUGZ3aJYuw998bZytxxo6QdgKmlnlgBaMvwEKCGg,1480
 athena/core/http_client.py,sha256=Z4NuAsJD-51yqmoME17O5sxwx5orSp1wsnd6bPyKcgA,17768
@@ -20,8 +20,10 @@ athena/errors/unauthorized_error.py,sha256=6qGByG08vA3iFHmK5ReHUK2AiGeCsK3Vf50C5
 athena/errors/unprocessable_entity_error.py,sha256=OztAzRntOJVxbZtXwoR4epahhG6chKOB5chVLUC8D9Q,257
 athena/errors/unsupported_media_type_error.py,sha256=fQ7TYQ3QYcT_YzFzO4f8_bLger7UQfZmuF22I9jYxFA,340
 athena/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-athena/query/__init__.py,sha256=FTtvy8EDg9nNNg9WCatVgKTRYV8-_v1roeGPAKoa_pw,65
-athena/query/client.py,sha256=q2ucGRJX5TY0sEJvYVzKKZcuDjYzhgUS-SoPQ4MIRfQ,4634
+athena/query/__init__.py,sha256=EIeVs3aDVNAqLyWetSrQldzc2yUdJkf_cJXyrLdMDj0,171
+athena/query/client.py,sha256=3i_yiGD7JM204TsTwY7R1CvGRs0odCh6QkYeNggHbc8,9147
+athena/query/types/__init__.py,sha256=WX-Or2h5NY2sv93ojrZsHcmiFHGYdqd0yxNo-5iGHR4,206
+athena/query/types/query_execute_request_database_asset_ids.py,sha256=aoVl5Xb34Q27hYGuVTnByGIxtHkL67wAwzXh7eJctew,154
 athena/tools/__init__.py,sha256=3n7oOoMebo06MAQqYRE2CX9Q0fTNnKBYE0cTlh1MPkM,165
 athena/tools/client.py,sha256=sMV_--PzUReXXILfKJeS8g9GQzYxxPDZI59wlos_rSo,10565
 athena/tools/types/__init__.py,sha256=cA-ZQm6veQAP3_vKu9KkZpISsQqgTBN_Z--FGY1c2iA,197
@@ -33,8 +35,8 @@ athena/types/data_frame_request_out_data_item_item.py,sha256=KMTJRr-1bdKDNMbUeri
 athena/types/data_frame_request_out_index_item.py,sha256=bW7oe912trpkYKodj-I_AiTXXy61yWzliekcsUZkZE0,141
 athena/types/data_frame_unknown_format_error.py,sha256=lbgAUArEgIYZt3P5Ji4Clp-xyXnKJR7-v5VzXwKu0J8,1168
 athena/types/file_fetch_error.py,sha256=iR7IXjj2C4lJFcWV3aS2gfhmT3CE-_hhBDSWFaCktHU,1162
-athena/types/sql_unauthorized_error.py,sha256=c3LR7lDgg66km9J0BYDCRu0hdDYaPqpnfcE3YXXKkVM,1105
+athena/types/sql_unauthorized_error.py,sha256=QSQbBlOBVYYjvcmSLZbsr4twcanPeMxIl74sKUU3Hac,1123
 athena/version.py,sha256=8aYAOJtVLaJLpRp6mTiEIhnl8gXA7yE0aDtZ-3mKQ4k,87
-athena_intelligence-0.1.100.dist-info/METADATA,sha256=JInGWwIplFb6naHD9zUXn3e5KXrpXlsGvdFya138JCg,5274
-athena_intelligence-0.1.100.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
-athena_intelligence-0.1.100.dist-info/RECORD,,
+athena_intelligence-0.1.102.dist-info/METADATA,sha256=p7kjPFWHqAU4uvLgawWmeBMiKo1NPX6eqb5RHZptpuw,5274
+athena_intelligence-0.1.102.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
+athena_intelligence-0.1.102.dist-info/RECORD,,