"""
    Airflow API

    Airflow API. All endpoints located under ``/api/v2`` can be used safely, are stable and backward compatible. Endpoints located under ``/ui`` are dedicated to the UI and are subject to breaking change depending on the need of the frontend. Users should not rely on those but use the public ones instead.

    The version of the OpenAPI document: 2
    Generated by OpenAPI Generator (https://openapi-generator.tech)

    Do not edit the class manually.
"""  # noqa: E501


import warnings
from pydantic import validate_call, Field, StrictFloat, StrictStr, StrictInt
from typing import Any, Dict, List, Optional, Tuple, Union
from typing_extensions import Annotated

from pydantic import Field, StrictBool, StrictStr
from typing import Any, Optional
from typing_extensions import Annotated
from airflow_client.client.models.task_state_store_body import TaskStateStoreBody
from airflow_client.client.models.task_state_store_collection_response import TaskStateStoreCollectionResponse
from airflow_client.client.models.task_state_store_patch_body import TaskStateStorePatchBody
from airflow_client.client.models.task_state_store_response import TaskStateStoreResponse

from airflow_client.client.api_client import ApiClient, RequestSerialized
from airflow_client.client.api_response import ApiResponse
from airflow_client.client.rest import RESTResponseType


class TaskStateStoreApi:
    """NOTE: This class is auto generated by OpenAPI Generator
    Ref: https://openapi-generator.tech

    Do not edit the class manually.
    """

    def __init__(self, api_client=None) -> None:
        if api_client is None:
            api_client = ApiClient.get_default()
        self.api_client = api_client


    @validate_call
    def clear_task_state_store(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        all_map_indices: Optional[StrictBool] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> None:
        """Clear Task State Store

        Delete all task state store keys for a task instance.  When ``all_map_indices=true``, state store is cleared for every map index of the task and the ``map_index`` parameter is ignored.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param map_index:
        :type map_index: int
        :param all_map_indices:
        :type all_map_indices: bool
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._clear_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            map_index=map_index,
            all_map_indices=all_map_indices,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        ).data


    @validate_call
    def clear_task_state_store_with_http_info(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        all_map_indices: Optional[StrictBool] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> ApiResponse[None]:
        """Clear Task State Store

        Delete all task state store keys for a task instance.  When ``all_map_indices=true``, state store is cleared for every map index of the task and the ``map_index`` parameter is ignored.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param map_index:
        :type map_index: int
        :param all_map_indices:
        :type all_map_indices: bool
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._clear_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            map_index=map_index,
            all_map_indices=all_map_indices,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        )


    @validate_call
    def clear_task_state_store_without_preload_content(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        all_map_indices: Optional[StrictBool] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> RESTResponseType:
        """Clear Task State Store

        Delete all task state store keys for a task instance.  When ``all_map_indices=true``, state store is cleared for every map index of the task and the ``map_index`` parameter is ignored.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param map_index:
        :type map_index: int
        :param all_map_indices:
        :type all_map_indices: bool
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._clear_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            map_index=map_index,
            all_map_indices=all_map_indices,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        return response_data.response


    def _clear_task_state_store_serialize(
        self,
        dag_id,
        dag_run_id,
        task_id,
        map_index,
        all_map_indices,
        _request_auth,
        _content_type,
        _headers,
        _host_index,
    ) -> RequestSerialized:

        _host = None

        _collection_formats: Dict[str, str] = {
        }

        _path_params: Dict[str, str] = {}
        _query_params: List[Tuple[str, str]] = []
        _header_params: Dict[str, Optional[str]] = _headers or {}
        _form_params: List[Tuple[str, str]] = []
        _files: Dict[
            str, Union[str, bytes, List[str], List[bytes], List[Tuple[str, bytes]]]
        ] = {}
        _body_params: Optional[bytes] = None

        # process the path parameters
        if dag_id is not None:
            _path_params['dag_id'] = dag_id
        if dag_run_id is not None:
            _path_params['dag_run_id'] = dag_run_id
        if task_id is not None:
            _path_params['task_id'] = task_id
        # process the query parameters
        if map_index is not None:
            
            _query_params.append(('map_index', map_index))
            
        if all_map_indices is not None:
            
            _query_params.append(('all_map_indices', all_map_indices))
            
        # process the header parameters
        # process the form parameters
        # process the body parameter


        # set the HTTP header `Accept`
        if 'Accept' not in _header_params:
            _header_params['Accept'] = self.api_client.select_header_accept(
                [
                    'application/json'
                ]
            )


        # authentication setting
        _auth_settings: List[str] = [
            'OAuth2PasswordBearer', 
            'HTTPBearer'
        ]

        return self.api_client.param_serialize(
            method='DELETE',
            resource_path='/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store',
            path_params=_path_params,
            query_params=_query_params,
            header_params=_header_params,
            body=_body_params,
            post_params=_form_params,
            files=_files,
            auth_settings=_auth_settings,
            collection_formats=_collection_formats,
            _host=_host,
            _request_auth=_request_auth
        )




    @validate_call
    def delete_task_state_store(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> None:
        """Delete Task State Store

        Delete a single task state store key. No-op if the key does not exist.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._delete_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        ).data


    @validate_call
    def delete_task_state_store_with_http_info(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> ApiResponse[None]:
        """Delete Task State Store

        Delete a single task state store key. No-op if the key does not exist.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._delete_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        )


    @validate_call
    def delete_task_state_store_without_preload_content(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> RESTResponseType:
        """Delete Task State Store

        Delete a single task state store key. No-op if the key does not exist.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._delete_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        return response_data.response


    def _delete_task_state_store_serialize(
        self,
        dag_id,
        dag_run_id,
        task_id,
        key,
        map_index,
        _request_auth,
        _content_type,
        _headers,
        _host_index,
    ) -> RequestSerialized:

        _host = None

        _collection_formats: Dict[str, str] = {
        }

        _path_params: Dict[str, str] = {}
        _query_params: List[Tuple[str, str]] = []
        _header_params: Dict[str, Optional[str]] = _headers or {}
        _form_params: List[Tuple[str, str]] = []
        _files: Dict[
            str, Union[str, bytes, List[str], List[bytes], List[Tuple[str, bytes]]]
        ] = {}
        _body_params: Optional[bytes] = None

        # process the path parameters
        if dag_id is not None:
            _path_params['dag_id'] = dag_id
        if dag_run_id is not None:
            _path_params['dag_run_id'] = dag_run_id
        if task_id is not None:
            _path_params['task_id'] = task_id
        if key is not None:
            _path_params['key'] = key
        # process the query parameters
        if map_index is not None:
            
            _query_params.append(('map_index', map_index))
            
        # process the header parameters
        # process the form parameters
        # process the body parameter


        # set the HTTP header `Accept`
        if 'Accept' not in _header_params:
            _header_params['Accept'] = self.api_client.select_header_accept(
                [
                    'application/json'
                ]
            )


        # authentication setting
        _auth_settings: List[str] = [
            'OAuth2PasswordBearer', 
            'HTTPBearer'
        ]

        return self.api_client.param_serialize(
            method='DELETE',
            resource_path='/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}',
            path_params=_path_params,
            query_params=_query_params,
            header_params=_header_params,
            body=_body_params,
            post_params=_form_params,
            files=_files,
            auth_settings=_auth_settings,
            collection_formats=_collection_formats,
            _host=_host,
            _request_auth=_request_auth
        )




    @validate_call
    def get_task_state_store(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> TaskStateStoreResponse:
        """Get Task State Store

        Get a single task state store entry.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._get_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "TaskStateStoreResponse",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        ).data


    @validate_call
    def get_task_state_store_with_http_info(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> ApiResponse[TaskStateStoreResponse]:
        """Get Task State Store

        Get a single task state store entry.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._get_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "TaskStateStoreResponse",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        )


    @validate_call
    def get_task_state_store_without_preload_content(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> RESTResponseType:
        """Get Task State Store

        Get a single task state store entry.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._get_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "TaskStateStoreResponse",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        return response_data.response


    def _get_task_state_store_serialize(
        self,
        dag_id,
        dag_run_id,
        task_id,
        key,
        map_index,
        _request_auth,
        _content_type,
        _headers,
        _host_index,
    ) -> RequestSerialized:

        _host = None

        _collection_formats: Dict[str, str] = {
        }

        _path_params: Dict[str, str] = {}
        _query_params: List[Tuple[str, str]] = []
        _header_params: Dict[str, Optional[str]] = _headers or {}
        _form_params: List[Tuple[str, str]] = []
        _files: Dict[
            str, Union[str, bytes, List[str], List[bytes], List[Tuple[str, bytes]]]
        ] = {}
        _body_params: Optional[bytes] = None

        # process the path parameters
        if dag_id is not None:
            _path_params['dag_id'] = dag_id
        if dag_run_id is not None:
            _path_params['dag_run_id'] = dag_run_id
        if task_id is not None:
            _path_params['task_id'] = task_id
        if key is not None:
            _path_params['key'] = key
        # process the query parameters
        if map_index is not None:
            
            _query_params.append(('map_index', map_index))
            
        # process the header parameters
        # process the form parameters
        # process the body parameter


        # set the HTTP header `Accept`
        if 'Accept' not in _header_params:
            _header_params['Accept'] = self.api_client.select_header_accept(
                [
                    'application/json'
                ]
            )


        # authentication setting
        _auth_settings: List[str] = [
            'OAuth2PasswordBearer', 
            'HTTPBearer'
        ]

        return self.api_client.param_serialize(
            method='GET',
            resource_path='/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}',
            path_params=_path_params,
            query_params=_query_params,
            header_params=_header_params,
            body=_body_params,
            post_params=_form_params,
            files=_files,
            auth_settings=_auth_settings,
            collection_formats=_collection_formats,
            _host=_host,
            _request_auth=_request_auth
        )




    @validate_call
    def list_task_state_store(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        limit: Optional[Annotated[int, Field(strict=True, ge=0)]] = None,
        offset: Optional[Annotated[int, Field(strict=True, ge=0)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> TaskStateStoreCollectionResponse:
        """List Task State Store

        List all task state store entries for a task instance.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param map_index:
        :type map_index: int
        :param limit:
        :type limit: int
        :param offset:
        :type offset: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._list_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            map_index=map_index,
            limit=limit,
            offset=offset,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "TaskStateStoreCollectionResponse",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        ).data


    @validate_call
    def list_task_state_store_with_http_info(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        limit: Optional[Annotated[int, Field(strict=True, ge=0)]] = None,
        offset: Optional[Annotated[int, Field(strict=True, ge=0)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> ApiResponse[TaskStateStoreCollectionResponse]:
        """List Task State Store

        List all task state store entries for a task instance.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param map_index:
        :type map_index: int
        :param limit:
        :type limit: int
        :param offset:
        :type offset: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._list_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            map_index=map_index,
            limit=limit,
            offset=offset,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "TaskStateStoreCollectionResponse",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        )


    @validate_call
    def list_task_state_store_without_preload_content(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        limit: Optional[Annotated[int, Field(strict=True, ge=0)]] = None,
        offset: Optional[Annotated[int, Field(strict=True, ge=0)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> RESTResponseType:
        """List Task State Store

        List all task state store entries for a task instance.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param map_index:
        :type map_index: int
        :param limit:
        :type limit: int
        :param offset:
        :type offset: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._list_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            map_index=map_index,
            limit=limit,
            offset=offset,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "TaskStateStoreCollectionResponse",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        return response_data.response


    def _list_task_state_store_serialize(
        self,
        dag_id,
        dag_run_id,
        task_id,
        map_index,
        limit,
        offset,
        _request_auth,
        _content_type,
        _headers,
        _host_index,
    ) -> RequestSerialized:

        _host = None

        _collection_formats: Dict[str, str] = {
        }

        _path_params: Dict[str, str] = {}
        _query_params: List[Tuple[str, str]] = []
        _header_params: Dict[str, Optional[str]] = _headers or {}
        _form_params: List[Tuple[str, str]] = []
        _files: Dict[
            str, Union[str, bytes, List[str], List[bytes], List[Tuple[str, bytes]]]
        ] = {}
        _body_params: Optional[bytes] = None

        # process the path parameters
        if dag_id is not None:
            _path_params['dag_id'] = dag_id
        if dag_run_id is not None:
            _path_params['dag_run_id'] = dag_run_id
        if task_id is not None:
            _path_params['task_id'] = task_id
        # process the query parameters
        if map_index is not None:
            
            _query_params.append(('map_index', map_index))
            
        if limit is not None:
            
            _query_params.append(('limit', limit))
            
        if offset is not None:
            
            _query_params.append(('offset', offset))
            
        # process the header parameters
        # process the form parameters
        # process the body parameter


        # set the HTTP header `Accept`
        if 'Accept' not in _header_params:
            _header_params['Accept'] = self.api_client.select_header_accept(
                [
                    'application/json'
                ]
            )


        # authentication setting
        _auth_settings: List[str] = [
            'OAuth2PasswordBearer', 
            'HTTPBearer'
        ]

        return self.api_client.param_serialize(
            method='GET',
            resource_path='/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store',
            path_params=_path_params,
            query_params=_query_params,
            header_params=_header_params,
            body=_body_params,
            post_params=_form_params,
            files=_files,
            auth_settings=_auth_settings,
            collection_formats=_collection_formats,
            _host=_host,
            _request_auth=_request_auth
        )




    @validate_call
    def patch_task_state_store(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        task_state_store_patch_body: TaskStateStorePatchBody,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> object:
        """Patch Task State Store

        Update the value of an existing task state store key.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param task_state_store_patch_body: (required)
        :type task_state_store_patch_body: TaskStateStorePatchBody
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._patch_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            task_state_store_patch_body=task_state_store_patch_body,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "object",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        ).data


    @validate_call
    def patch_task_state_store_with_http_info(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        task_state_store_patch_body: TaskStateStorePatchBody,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> ApiResponse[object]:
        """Patch Task State Store

        Update the value of an existing task state store key.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param task_state_store_patch_body: (required)
        :type task_state_store_patch_body: TaskStateStorePatchBody
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._patch_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            task_state_store_patch_body=task_state_store_patch_body,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "object",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        )


    @validate_call
    def patch_task_state_store_without_preload_content(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        task_state_store_patch_body: TaskStateStorePatchBody,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> RESTResponseType:
        """Patch Task State Store

        Update the value of an existing task state store key.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param task_state_store_patch_body: (required)
        :type task_state_store_patch_body: TaskStateStorePatchBody
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._patch_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            task_state_store_patch_body=task_state_store_patch_body,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '200': "object",
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        return response_data.response


    def _patch_task_state_store_serialize(
        self,
        dag_id,
        dag_run_id,
        task_id,
        key,
        task_state_store_patch_body,
        map_index,
        _request_auth,
        _content_type,
        _headers,
        _host_index,
    ) -> RequestSerialized:

        _host = None

        _collection_formats: Dict[str, str] = {
        }

        _path_params: Dict[str, str] = {}
        _query_params: List[Tuple[str, str]] = []
        _header_params: Dict[str, Optional[str]] = _headers or {}
        _form_params: List[Tuple[str, str]] = []
        _files: Dict[
            str, Union[str, bytes, List[str], List[bytes], List[Tuple[str, bytes]]]
        ] = {}
        _body_params: Optional[bytes] = None

        # process the path parameters
        if dag_id is not None:
            _path_params['dag_id'] = dag_id
        if dag_run_id is not None:
            _path_params['dag_run_id'] = dag_run_id
        if task_id is not None:
            _path_params['task_id'] = task_id
        if key is not None:
            _path_params['key'] = key
        # process the query parameters
        if map_index is not None:
            
            _query_params.append(('map_index', map_index))
            
        # process the header parameters
        # process the form parameters
        # process the body parameter
        if task_state_store_patch_body is not None:
            _body_params = task_state_store_patch_body


        # set the HTTP header `Accept`
        if 'Accept' not in _header_params:
            _header_params['Accept'] = self.api_client.select_header_accept(
                [
                    'application/json'
                ]
            )

        # set the HTTP header `Content-Type`
        if _content_type:
            _header_params['Content-Type'] = _content_type
        else:
            _default_content_type = (
                self.api_client.select_header_content_type(
                    [
                        'application/json'
                    ]
                )
            )
            if _default_content_type is not None:
                _header_params['Content-Type'] = _default_content_type

        # authentication setting
        _auth_settings: List[str] = [
            'OAuth2PasswordBearer', 
            'HTTPBearer'
        ]

        return self.api_client.param_serialize(
            method='PATCH',
            resource_path='/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}',
            path_params=_path_params,
            query_params=_query_params,
            header_params=_header_params,
            body=_body_params,
            post_params=_form_params,
            files=_files,
            auth_settings=_auth_settings,
            collection_formats=_collection_formats,
            _host=_host,
            _request_auth=_request_auth
        )




    @validate_call
    def set_task_state_store(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        task_state_store_body: TaskStateStoreBody,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> None:
        """Set Task State Store

        Set a task state store value. Creates or overwrites the key.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param task_state_store_body: (required)
        :type task_state_store_body: TaskStateStoreBody
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._set_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            task_state_store_body=task_state_store_body,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        ).data


    @validate_call
    def set_task_state_store_with_http_info(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        task_state_store_body: TaskStateStoreBody,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> ApiResponse[None]:
        """Set Task State Store

        Set a task state store value. Creates or overwrites the key.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param task_state_store_body: (required)
        :type task_state_store_body: TaskStateStoreBody
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._set_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            task_state_store_body=task_state_store_body,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        response_data.read()
        return self.api_client.response_deserialize(
            response_data=response_data,
            response_types_map=_response_types_map,
        )


    @validate_call
    def set_task_state_store_without_preload_content(
        self,
        dag_id: StrictStr,
        dag_run_id: StrictStr,
        task_id: StrictStr,
        key: StrictStr,
        task_state_store_body: TaskStateStoreBody,
        map_index: Optional[Annotated[int, Field(strict=True, ge=-1)]] = None,
        _request_timeout: Union[
            None,
            Annotated[StrictFloat, Field(gt=0)],
            Tuple[
                Annotated[StrictFloat, Field(gt=0)],
                Annotated[StrictFloat, Field(gt=0)]
            ]
        ] = None,
        _request_auth: Optional[Dict[StrictStr, Any]] = None,
        _content_type: Optional[StrictStr] = None,
        _headers: Optional[Dict[StrictStr, Any]] = None,
        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
    ) -> RESTResponseType:
        """Set Task State Store

        Set a task state store value. Creates or overwrites the key.

        :param dag_id: (required)
        :type dag_id: str
        :param dag_run_id: (required)
        :type dag_run_id: str
        :param task_id: (required)
        :type task_id: str
        :param key: (required)
        :type key: str
        :param task_state_store_body: (required)
        :type task_state_store_body: TaskStateStoreBody
        :param map_index:
        :type map_index: int
        :param _request_timeout: timeout setting for this request. If one
                                 number provided, it will be total request
                                 timeout. It can also be a pair (tuple) of
                                 (connection, read) timeouts.
        :type _request_timeout: int, tuple(int, int), optional
        :param _request_auth: set to override the auth_settings for an a single
                              request; this effectively ignores the
                              authentication in the spec for a single request.
        :type _request_auth: dict, optional
        :param _content_type: force content-type for the request.
        :type _content_type: str, Optional
        :param _headers: set to override the headers for a single
                         request; this effectively ignores the headers
                         in the spec for a single request.
        :type _headers: dict, optional
        :param _host_index: set to override the host_index for a single
                            request; this effectively ignores the host_index
                            in the spec for a single request.
        :type _host_index: int, optional
        :return: Returns the result object.
        """ # noqa: E501

        _param = self._set_task_state_store_serialize(
            dag_id=dag_id,
            dag_run_id=dag_run_id,
            task_id=task_id,
            key=key,
            task_state_store_body=task_state_store_body,
            map_index=map_index,
            _request_auth=_request_auth,
            _content_type=_content_type,
            _headers=_headers,
            _host_index=_host_index
        )

        _response_types_map: Dict[str, Optional[str]] = {
            '204': None,
            '401': "HTTPExceptionResponse",
            '403': "HTTPExceptionResponse",
            '404': "HTTPExceptionResponse",
            '422': "HTTPValidationError",
        }
        response_data = self.api_client.call_api(
            *_param,
            _request_timeout=_request_timeout
        )
        return response_data.response


    def _set_task_state_store_serialize(
        self,
        dag_id,
        dag_run_id,
        task_id,
        key,
        task_state_store_body,
        map_index,
        _request_auth,
        _content_type,
        _headers,
        _host_index,
    ) -> RequestSerialized:

        _host = None

        _collection_formats: Dict[str, str] = {
        }

        _path_params: Dict[str, str] = {}
        _query_params: List[Tuple[str, str]] = []
        _header_params: Dict[str, Optional[str]] = _headers or {}
        _form_params: List[Tuple[str, str]] = []
        _files: Dict[
            str, Union[str, bytes, List[str], List[bytes], List[Tuple[str, bytes]]]
        ] = {}
        _body_params: Optional[bytes] = None

        # process the path parameters
        if dag_id is not None:
            _path_params['dag_id'] = dag_id
        if dag_run_id is not None:
            _path_params['dag_run_id'] = dag_run_id
        if task_id is not None:
            _path_params['task_id'] = task_id
        if key is not None:
            _path_params['key'] = key
        # process the query parameters
        if map_index is not None:
            
            _query_params.append(('map_index', map_index))
            
        # process the header parameters
        # process the form parameters
        # process the body parameter
        if task_state_store_body is not None:
            _body_params = task_state_store_body


        # set the HTTP header `Accept`
        if 'Accept' not in _header_params:
            _header_params['Accept'] = self.api_client.select_header_accept(
                [
                    'application/json'
                ]
            )

        # set the HTTP header `Content-Type`
        if _content_type:
            _header_params['Content-Type'] = _content_type
        else:
            _default_content_type = (
                self.api_client.select_header_content_type(
                    [
                        'application/json'
                    ]
                )
            )
            if _default_content_type is not None:
                _header_params['Content-Type'] = _default_content_type

        # authentication setting
        _auth_settings: List[str] = [
            'OAuth2PasswordBearer', 
            'HTTPBearer'
        ]

        return self.api_client.param_serialize(
            method='PUT',
            resource_path='/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}',
            path_params=_path_params,
            query_params=_query_params,
            header_params=_header_params,
            body=_body_params,
            post_params=_form_params,
            files=_files,
            auth_settings=_auth_settings,
            collection_formats=_collection_formats,
            _host=_host,
            _request_auth=_request_auth
        )


