# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/integrations/Automation.ipynb.

# %% auto 0
__all__ = ['remove_partition_by_x_days', 'get_company_domains']

# %% ../../nbs/integrations/Automation.ipynb 2
import datetime as dt
import pandas as pd

import domolibrary.client.DomoAuth as dmda
import domolibrary.classes.DomoDataset as dmds

# %% ../../nbs/integrations/Automation.ipynb 3
async def remove_partition_by_x_days(
    auth: dmda.DomoAuth,
    dataset_id: str,
    x_last_days: int = 0,
    separator: str = None,
    date_index: int = 0,
    date_format: str = "%Y-%m-%d",
):
    domo_ds = dmds.DomoDataset(auth=auth, id=dataset_id)

    list_partition = await domo_ds.list_partitions(auth=auth, dataset_id=dataset_id)

    today = dt.date.today()
    days_ago = today - dt.timedelta(days=x_last_days)
    for i in list_partition:
        compare_date = ""
        if separator is not None and separator != "":
            compare_date = i["partitionId"].split(separator)[date_index]
        else:
            compare_date = i["partitionId"]

        try:
            d = dt.datetime.strptime(compare_date, date_format).date()
        except ValueError:
            d = None
        if d is not None and d < days_ago:
            print(
                auth.domo_instance,
                ": 🚀  Removing partition key : ",
                (i["partitionId"]),
                " in ",
                dataset_id,
            )
            await domo_ds.delete_partition(
                dataset_partition_id=i["partitionId"], dataset_id=dataset_id, auth=auth
            )

# %% ../../nbs/integrations/Automation.ipynb 4
async def get_company_domains(
    auth: dmda.DomoAuth,
    dataset_id: str,
    handle_err_fn: callable,
    sql: str = "select domain from table",
    global_admin_username: str = None,
    global_admin_password: str = None,
    execution_env: str = None,
    debug_api: bool = False,
) -> pd.DataFrame:
    ds = await dmds.DomoDataset.get_from_id(
        auth=auth, id=dataset_id, debug_api=debug_api
    )

    print(f"⚙️ START - Retrieving company list \n{ds.display_url()}")
    print(f"⚙️ SQL = {sql}")

    df = await ds.query_dataset_private(
        auth=auth,
        dataset_id=dataset_id,
        sql=sql,
        loop_until_end=True,
        debug_api=debug_api,
    )

    df["domo_instance"] = df["domain"].apply(lambda x: x.replace(".domo.com", ""))

    if global_admin_username:
        df["domo_username"] = global_admin_username
    if global_admin_password:
        df["domo_password"] = global_admin_password

    if execution_env:
        df["env"] = execution_env or "manual"

    if df.empty:
        raise Exception("no companies retrieved")
        return False

    print(
        f"\n⚙️ SUCCESS 🎉 Retrieved company list \nThere are {len(df.index)} companies to update"
    )

    return df
