# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/routes/dataset.ipynb.

# %% auto 0
__all__ = ['DatasetNotFoundError', 'QueryRequestError', 'query_dataset_public', 'query_dataset_private', 'get_dataset_by_id',
           'get_schema', 'set_dataset_tags', 'UploadDataError', 'upload_dataset_stage_1', 'upload_dataset_stage_2_file',
           'upload_dataset_stage_2_df', 'upload_dataset_stage_3', 'index_dataset', 'index_status',
           'generate_list_partitions_body', 'list_partitions', 'generate_create_dataset_body', 'create', 'delete']

# %% ../../nbs/routes/dataset.ipynb 3
from typing import Optional

import io
import pandas as pd

import httpx

import domolibrary.client.get_data as gd
import domolibrary.client.ResponseGetData as rgd
import domolibrary.client.DomoAuth as dmda


# %% ../../nbs/routes/dataset.ipynb 5
class DatasetNotFoundError(Exception):
    def __init__(self, dataset_id, domo_instance):
        message = f"dataset - {dataset_id} not found in {domo_instance}"

        super().__init__(message)

# %% ../../nbs/routes/dataset.ipynb 6
class QueryRequestError(Exception):
    def __init__(self, dataset_id, domo_instance, sql):
        message = f"dataset - {dataset_id} in {domo_instance} received a bad request.  Check your SQL \n {sql}"

        super().__init__(message)


# typically do not use
async def query_dataset_public(
    dev_auth: dmda.DomoDeveloperAuth,
    dataset_id: str,
    sql: str,
    session: httpx.AsyncClient,
    debug_api: bool = False,
):

    """query for hitting public apis, requires client_id and secret authentication"""

    url = f"https://api.domo.com/v1/datasets/query/execute/{dataset_id}?IncludeHeaders=true"

    body = {"sql": sql}

    return await gd.get_data(
        auth=dev_auth,
        url=url,
        method="POST",
        body=body,
        session=session,
        debug_api=debug_api,
    )


async def query_dataset_private(
    auth: dmda.DomoAuth,  # DomoFullAuth or DomoTokenAuth
    dataset_id: str,
    sql: str,
    session: Optional[httpx.AsyncClient] = None,
    loop_until_end: bool = False,  # retrieve all available rows
    limit=100,  # maximum rows to return per request.  refers to PAGINATION
    skip=0,
    maximum=100,  # equivalent to the LIMIT or TOP clause in SQL, the number of rows to return total
    debug_api: bool = False,
    debug_loop: bool = False,
    timeout :int = 10
):
    """execute SQL queries against private APIs, requires DomoFullAuth or DomoTokenAuth"""

    url = f"https://{auth.domo_instance}.domo.com/api/query/v1/execute/{dataset_id}"

    offset_params = {
        "offset": "offset",
        "limit": "limit",
    }

    def body_fn(skip, limit):
        return {"sql": f"{sql} limit {limit} offset {skip}"}

    def arr_fn(res) -> pd.DataFrame:
        rows_ls = res.response.get("rows")
        columns_ls = res.response.get("columns")
        output = []
        for row in rows_ls:
            new_row = {}
            for index, column in enumerate(columns_ls):
                new_row[column] = row[index]
            output.append(new_row)
            # pd.DataFrame(data=res.response.get('rows'), columns=res.response.get('columns'))
        return output

    res = await gd.looper(
        auth=auth,
        method="POST",
        url=url,
        arr_fn=arr_fn,
        offset_params=offset_params,
        limit=limit,
        skip=skip,
        maximum=maximum,
        session=session,
        body_fn=body_fn,
        debug_api=debug_api,
        debug_loop=debug_loop,
        loop_until_end=loop_until_end,
        timeout = timeout
    )

    if res.status == 404 and res.response == "Not Found":
        raise DatasetNotFoundError(
            dataset_id=dataset_id, domo_instance=auth.domo_instance
        )

    if res.status == 400 and res.response == "Bad Request":
        raise QueryRequestError(
            dataset_id=dataset_id, domo_instance=auth.domo_instance, sql=sql
        )

    return res

# %% ../../nbs/routes/dataset.ipynb 9
async def get_dataset_by_id(
    dataset_id: str,  # dataset id from URL
    auth: Optional[dmda.DomoAuth] = None,  # requires full authentication
    debug_api: bool = False,  # for troubleshooting API request
    session: Optional[httpx.AsyncClient] = None,
) -> rgd.ResponseGetData:  # returns metadata about a dataset
    """retrieve dataset metadata"""

    url = f"https://{auth.domo_instance}.domo.com/api/data/v3/datasources/{dataset_id}"

    res = await gd.get_data(
        auth=auth, url=url, method="GET", debug_api=debug_api, session=session
    )

    if res.status == 404 and res.response == "Not Found":
        raise DatasetNotFoundError(
            dataset_id=dataset_id, domo_instance=auth.domo_instance
        )

    return res

# %% ../../nbs/routes/dataset.ipynb 12
async def get_schema(
    auth: dmda.DomoAuth, dataset_id: str, debug_api: bool = False
) -> rgd.ResponseGetData:
    """retrieve the schema for a dataset"""

    url = f"https://{auth.domo_instance}.domo.com/api/query/v1/datasources/{dataset_id}/schema/indexed?includeHidden=false"

    return await gd.get_data(auth=auth, url=url, method="GET", debug_api=debug_api)

# %% ../../nbs/routes/dataset.ipynb 16
async def set_dataset_tags(
    auth: dmda.DomoFullAuth,
    tag_ls: [str],  # complete list of tags for dataset
    dataset_id: str,
    debug_api: bool = False,
    session: Optional[httpx.AsyncClient] = None,
    return_raw: bool = False,
):

    """REPLACE tags on this dataset with a new list"""

    url = f"https://{auth.domo_instance}.domo.com/api/data/ui/v3/datasources/{dataset_id}/tags"

    res = await gd.get_data(
        auth=auth,
        url=url,
        method="POST",
        debug_api=debug_api,
        body=tag_ls,
        session=session,
        return_raw=return_raw,
    )

    if return_raw:
        return res

    if res.status == 200:
        res.set_response(
            response=f'Dataset {dataset_id} tags updated to [{ ", ".join(tag_ls) }]'
        )

    return res

# %% ../../nbs/routes/dataset.ipynb 19
class UploadDataError(Exception):
    """raise if unable to upload data to Domo"""

    def __init__(self, stage_num: int, dataset_id: str, domo_instance: str):
        message = f"error uploading data to {dataset_id} during Stage { stage_num} in {domo_instance}"
        super().__init__(message)

# %% ../../nbs/routes/dataset.ipynb 20
async def upload_dataset_stage_1(
    auth: dmda.DomoAuth,
    dataset_id: str,
    #  restate_data_tag: str = None, # deprecated
    partition_tag: str = None,  # synonymous with data_tag
    session: Optional[httpx.AsyncClient] = None,
    debug_api: bool = False,
) -> rgd.ResponseGetData:

    """preps dataset for upload by creating an upload_id (upload session key) pass to stage 2 as a parameter"""

    url = f"https://{auth.domo_instance}.domo.com/api/data/v3/datasources/{dataset_id}/uploads"

    # base body assumes no paritioning
    body = {"action": None, "appendId": None}

    params = None

    if partition_tag:
        # params = {'dataTag': restate_data_tag or data_tag} # deprecated
        params = {"dataTag": partition_tag}
        body.update({"appendId": "latest"})

    res = await gd.get_data(
        auth=auth,
        url=url,
        method="POST",
        body=body,
        session=session,
        debug_api=debug_api,
        params=params,
    )

    if not res.is_success:
        raise UploadDataError(
            stage_num=1, dataset_id=dataset_id, domo_instance=auth.domo_instance
        )

    return res

# %% ../../nbs/routes/dataset.ipynb 21
async def upload_dataset_stage_2_file(
    auth: dmda.DomoAuth,
    dataset_id: str,
    upload_id: str,  # must originate from  a stage_1 upload response
    data_file: Optional[io.TextIOWrapper] = None,
    session: Optional[httpx.AsyncClient] = None,
    # only necessary if streaming multiple files into the same partition (multi-part upload)
    part_id: str = 2,
    debug_api: bool = False,
) -> rgd.ResponseGetData:

    url = f"https://{auth.domo_instance}.domo.com/api/data/v3/datasources/{dataset_id}/uploads/{upload_id}/parts/{part_id}"

    body = data_file

    res = await gd.get_data(
        url=url,
        method="PUT",
        auth=auth,
        content_type="text/csv",
        body=body,
        session=session,
        debug_api=debug_api,
    )
    if not res.is_success:
        raise UploadDataError(
            stage_num=2, dataset_id=dataset_id, domo_instance=auth.domo_instance
        )

    res.upload_id = upload_id
    res.dataset_id = dataset_id
    res.part_id = part_id

    return res

# %% ../../nbs/routes/dataset.ipynb 22
async def upload_dataset_stage_2_df(
    auth: dmda.DomoAuth,
    dataset_id: str,
    upload_id: str,  # must originate from  a stage_1 upload response
    upload_df: pd.DataFrame,
    session: Optional[httpx.AsyncClient] = None,
    part_id: str = 2,  # only necessary if streaming multiple files into the same partition (multi-part upload)
    debug_api: bool = False,
) -> rgd.ResponseGetData:

    url = f"https://{auth.domo_instance}.domo.com/api/data/v3/datasources/{dataset_id}/uploads/{upload_id}/parts/{part_id}"

    body = upload_df.to_csv(header=False, index=False)

    # if debug:
    #     print(body)

    res = await gd.get_data(
        url=url,
        method="PUT",
        auth=auth,
        content_type="text/csv",
        body=body,
        session=session,
        debug_api=debug_api,
    )

    if not res.is_success:
        raise UploadDataError(
            stage_num=2, dataset_id=dataset_id, domo_instance=auth.domo_instance
        )

    res.upload_id = upload_id
    res.dataset_id = dataset_id
    res.part_id = part_id

    return res

# %% ../../nbs/routes/dataset.ipynb 23
async def upload_dataset_stage_3(
    auth: dmda.DomoAuth,
    dataset_id: str,
    upload_id: str,  # must originate from  a stage_1 upload response
    session: Optional[httpx.AsyncClient] = None,
    update_method: str = "REPLACE",  # accepts REPLACE or APPEND
    #  restate_data_tag: str = None, # deprecated
    partition_tag: str = None,  # synonymous with data_tag
    is_index: bool = False,  # index after uploading
    debug_api: bool = False,
) -> rgd.ResponseGetData:

    """commit will close the upload session, upload_id.  this request defines how the data will be loaded into Adrenaline, update_method
    has optional flag for indexing dataset.
    """

    url = f"https://{auth.domo_instance}.domo.com/api/data/v3/datasources/{dataset_id}/uploads/{upload_id}/commit"

    body = {"index": is_index, "action": update_method}

    if partition_tag:

        body.update(
            {
                "action": "APPEND",
                #  'dataTag': restate_data_tag or data_tag,
                #  'appendId': 'latest' if (restate_data_tag or data_tag) else None,
                "dataTag": partition_tag,
                "appendId": "latest" if partition_tag else None,
                "index": is_index,
            }
        )

    res = await gd.get_data(
        auth=auth,
        method="PUT",
        url=url,
        body=body,
        session=session,
        debug_api=debug_api,
    )

    if not res.is_success:
        raise UploadDataError(
            stage_num=3, dataset_id=dataset_id, domo_instance=auth.domo_instance
        )

    res.upload_id = upload_id
    res.dataset_id = dataset_id

    return res

# %% ../../nbs/routes/dataset.ipynb 25
async def index_dataset(
    auth: dmda.DomoAuth,
    dataset_id: str,
    session: Optional[httpx.AsyncClient] = None,
    debug_api: bool = False,
) -> rgd.ResponseGetData:
    """manually index a dataset"""

    url = f"https://{auth.domo_instance}.domo.com/api/data/v3/datasources/{dataset_id}/indexes"

    body = {"dataIds": []}

    return await gd.get_data(
        auth=auth,
        method="POST",
        body=body,
        url=url,
        session=session,
        debug_api=debug_api,
    )

# %% ../../nbs/routes/dataset.ipynb 26
async def index_status(
    auth: dmda.DomoAuth,
    dataset_id: str,
    index_id: str,
    session: Optional[httpx.AsyncClient] = None,
    debug_api: bool = False,
) -> rgd.ResponseGetData:
    """get the completion status of an index"""

    url = f"https://{auth.domo_instance}.domo.com/api/data/v3/datasources/{dataset_id}/indexes/{index_id}/statuses"

    return await gd.get_data(
        auth=auth, method="GET", url=url, session=session, debug_api=debug_api
    )

# %% ../../nbs/routes/dataset.ipynb 28
def generate_list_partitions_body(limit=100, offset=0):
    return {
        "paginationFields": [
            {
                "fieldName": "datecompleted",
                "sortOrder": "DESC",
                "filterValues": {"MIN": None, "MAX": None},
            }
        ],
        "limit": limit,
        "offset": offset,
    }


async def list_partitions(
    auth: dmda.DomoAuth,
    dataset_id: str,
    body: dict = None,
    session: httpx.AsyncClient = None,
    debug_api: bool = False,
    debug_loop: bool = False,
):

    body = body or generate_list_partitions_body()

    url = f"https://{auth.domo_instance}.domo.com/api/query/v1/datasources/{dataset_id}/partition/list"

    offset_params = {
        "offset": "offset",
        "limit": "limit",
    }

    def arr_fn(res) -> list[dict]:
        return res.response

    res = await gd.looper(
        auth=auth,
        method="POST",
        url=url,
        arr_fn=arr_fn,
        body=body,
        offset_params_in_body=True,
        offset_params=offset_params,
        loop_until_end=True,
        session=session,
        debug_loop=debug_loop,
        debug_api=debug_api,
    )

    if res.status == 404 and res.response == "Not Found":
        raise DatasetNotFoundError(
            dataset_id=dataset_id, domo_instance=auth.domo_instance
        )
    return res

# %% ../../nbs/routes/dataset.ipynb 30
def generate_create_dataset_body(
    dataset_name: str, dataset_type: str = "API", schema: dict = None
):
    schema = schema or {
        "columns": [
            {"type": "STRING", "name": "Friend"},
            {"type": "STRING", "name": "Attending"},
        ]
    }

    return {
        "userDefinedType": dataset_type,
        "dataSourceName": dataset_name,
        "schema": schema,
    }


async def create(
    auth: dmda.DomoAuth,
    dataset_name: str,
    dataset_type: str = "api",
    session: httpx.AsyncClient = None,
    schema: dict = None,
    debug_api: bool = False,
):

    body = generate_create_dataset_body(
        dataset_name=dataset_name, dataset_type=dataset_type, schema=schema
    )

    url = f"https://{auth.domo_instance}.domo.com/api/data/v2/datasources"

    return await gd.get_data(
        auth=auth,
        method="POST",
        url=url,
        body=body,
        session=session,
        debug_api=debug_api,
    )

# %% ../../nbs/routes/dataset.ipynb 35
async def delete(
    auth: dmda.DomoAuth,
    dataset_id: str,
    session: httpx.AsyncClient = None,
    debug_api: bool = False,
):
    url = f"https://{auth.domo_instance}.domo.com/api/data/v3/datasources/{dataset_id}?deleteMethod=hard"

    return await gd.get_data(
        auth=auth, method="DELETE", url=url, session=session, debug_api=debug_api
    )
