# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/CLI_DataBlob.ipynb.

# %% auto 0
__all__ = ['logger']

# %% ../../notebooks/CLI_DataBlob.ipynb 3
from typing import *

# %% ../../notebooks/CLI_DataBlob.ipynb 4
import os
import json

import typer
from typer import echo
from tabulate import tabulate
import datetime as dt
import pandas as pd

from airt.client import Client
from airt.cli import helper
from airt.logger import get_logger, set_level
from airt.constant import CLIENT_DB_USERNAME, CLIENT_DB_PASSWORD

# %% ../../notebooks/CLI_DataBlob.ipynb 6
app = typer.Typer(
    help="""A set of commands for importing and processing data from sources such as CSV/parquet files, databases, AWS S3 buckets, and Azure Blob Storage."""
)

# %% ../../notebooks/CLI_DataBlob.ipynb 8
logger = get_logger(__name__)

# %% ../../notebooks/CLI_DataBlob.ipynb 18
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def details(
    uuid: str = typer.Argument(
        ...,
        help="Datablob uuid.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Return details of a datablob."""

    from airt.client import DataBlob

    db = DataBlob(uuid=uuid)
    df = db.details()

    df["pulled_on"] = helper.humanize_date(df["pulled_on"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    return {"df": df}

# %% ../../notebooks/CLI_DataBlob.ipynb 22
@app.command("from-s3")
@helper.requires_auth_token
def from_s3(
    uri: str = typer.Argument(..., help="The AWS S3 bucket uri."),
    access_key: Optional[str] = typer.Option(
        None,
        help="Access key for the S3 bucket. If **None** (default value), then the value from **AWS_ACCESS_KEY_ID** environment variable is used.",
    ),
    secret_key: Optional[str] = typer.Option(
        None,
        help="Secret key for the S3 bucket. If **None** (default value), then the value from **AWS_SECRET_ACCESS_KEY** environment variable is used.",
    ),
    cloud_provider: Optional[str] = typer.Option(
        None,
        "--cloud-provider",
        "-cp",
        help="The destination cloud storage provider's name to store the datablob. Currently, the API only supports **aws** and **azure** as cloud storage providers. If **None** (default value), then **aws**  will be used as the cloud storage provider.",
    ),
    region: Optional[str] = typer.Option(
        None,
        "--region",
        "-r",
        help="The destination cloud provider's region to save your datablob. If **None** (default value) then the default region will be assigned based on the cloud provider. "
        "In the case of **aws**, the datablob's source bucket region will be used and in the case of **azure**, **westeurope** will be used. "
        "The supported AWS regions are: ap-northeast-1, ap-northeast-2, ap-south-1, ap-southeast-1, ap-southeast-2, ca-central-1, eu-central-1, "
        "eu-north-1, eu-west-1, eu-west-2, eu-west-3, sa-east-1, us-east-1, us-east-2, us-west-1, us-west-2. The supported Azure Blob Storage "
        "regions are: australiacentral, australiacentral2, australiaeast, australiasoutheast, brazilsouth, canadacentral, canadaeast, centralindia, "
        "centralus, eastasia, eastus, eastus2, francecentral, francesouth, germanynorth, germanywestcentral, japaneast, japanwest, koreacentral, koreasouth, "
        "northcentralus, northeurope, norwayeast, norwaywest, southafricanorth, southafricawest, southcentralus, southeastasia, southindia, switzerlandnorth, "
        "switzerlandwest, uaecentral, uaenorth, uksouth, ukwest, westcentralus, westeurope, westindia, westus, westus2.",
    ),
    tag: Optional[str] = typer.Option(
        None,
        "--tag",
        "-t",
        help="A string to tag the datablob. If not passed, then the tag **latest** will be assigned to the datablob.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output datablob uuid only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Create and return a datablob that encapsulates the data from an AWS S3 bucket."""

    from airt.client import DataBlob

    db = DataBlob.from_s3(
        uri=uri,
        access_key=access_key,
        secret_key=secret_key,
        cloud_provider=cloud_provider,
        region=region,
        tag=tag,
    )

    if quiet:
        db.wait()

        typer.echo(f"{db.uuid}")
    else:
        typer.echo(f"Pulling datablob uuid: {db.uuid}")

        db.progress_bar()

# %% ../../notebooks/CLI_DataBlob.ipynb 26
@app.command("from-azure-blob-storage")
@helper.requires_auth_token
def from_azure_blob_storage(
    uri: str = typer.Argument(..., help="Azure Blob Storage URI of the source file."),
    credential: str = typer.Option(
        ...,
        "--credential",
        "-c",
        help="Credential to access the Azure Blob Storage.",
    ),
    cloud_provider: Optional[str] = typer.Option(
        None,
        "--cloud-provider",
        "-cp",
        help="The destination cloud storage provider's name to store the datablob. Currently, the API only supports **aws** and **azure** as cloud storage providers. If **None** (default value), then **azure**  will be used as the cloud storage provider.",
    ),
    region: Optional[str] = typer.Option(
        None,
        "--region",
        "-r",
        help="The destination cloud provider's region to save your datablob. If **None** (default value) then the default region will be assigned based on the cloud provider. "
        "In the case of **aws**, **eu-west-1** will be used and in the case of **azure**, **westeurope** will be used. "
        "The supported AWS regions are: ap-northeast-1, ap-northeast-2, ap-south-1, ap-southeast-1, ap-southeast-2, ca-central-1, eu-central-1, "
        "eu-north-1, eu-west-1, eu-west-2, eu-west-3, sa-east-1, us-east-1, us-east-2, us-west-1, us-west-2. The supported Azure Blob Storage "
        "regions are: australiacentral, australiacentral2, australiaeast, australiasoutheast, brazilsouth, canadacentral, canadaeast, centralindia, "
        "centralus, eastasia, eastus, eastus2, francecentral, francesouth, germanynorth, germanywestcentral, japaneast, japanwest, koreacentral, koreasouth, "
        "northcentralus, northeurope, norwayeast, norwaywest, southafricanorth, southafricawest, southcentralus, southeastasia, southindia, switzerlandnorth, "
        "switzerlandwest, uaecentral, uaenorth, uksouth, ukwest, westcentralus, westeurope, westindia, westus, westus2.",
    ),
    tag: Optional[str] = typer.Option(
        None,
        "--tag",
        "-t",
        help="A string to tag the datablob. If not passed, then the tag **latest** will be assigned to the datablob.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output datablob uuid only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Create and return a datablob that encapsulates the data from an Azure Blob Storage."""

    from airt.client import DataBlob

    db = DataBlob.from_azure_blob_storage(
        uri=uri,
        credential=credential,
        cloud_provider=cloud_provider,
        region=region,
        tag=tag,
    )

    if quiet:
        db.wait()

        typer.echo(f"{db.uuid}")
    else:
        typer.echo(f"Pulling datablob uuid: {db.uuid}")

        db.progress_bar()

# %% ../../notebooks/CLI_DataBlob.ipynb 30
@app.command("from-mysql")
@helper.requires_auth_token
def from_mysql(
    host: str = typer.Option(..., help="Remote database host name."),
    database: str = typer.Option(..., help="Database name."),
    table: str = typer.Option(..., help="Table name."),
    port: int = typer.Option(
        3306,
        help="Host port number. If not passed, then the default value **3306** will be used.",
    ),
    cloud_provider: Optional[str] = typer.Option(
        None,
        "--cloud-provider",
        "-cp",
        help="The destination cloud storage provider's name to store the datablob. Currently, the API only supports **aws** and **azure** as cloud storage providers. If **None** (default value), then **aws**  will be used as the cloud storage provider.",
    ),
    region: Optional[str] = typer.Option(
        None,
        "--region",
        "-r",
        help="The destination cloud provider's region to save your datablob. If **None** (default value) then the default region will be assigned based on the cloud provider. "
        "In the case of **aws**, **eu-west-1** will be used and in the case of **azure**, **westeurope** will be used. "
        "The supported AWS regions are: ap-northeast-1, ap-northeast-2, ap-south-1, ap-southeast-1, ap-southeast-2, ca-central-1, eu-central-1, "
        "eu-north-1, eu-west-1, eu-west-2, eu-west-3, sa-east-1, us-east-1, us-east-2, us-west-1, us-west-2. The supported Azure Blob Storage "
        "regions are: australiacentral, australiacentral2, australiaeast, australiasoutheast, brazilsouth, canadacentral, canadaeast, centralindia, "
        "centralus, eastasia, eastus, eastus2, francecentral, francesouth, germanynorth, germanywestcentral, japaneast, japanwest, koreacentral, koreasouth, "
        "northcentralus, northeurope, norwayeast, norwaywest, southafricanorth, southafricawest, southcentralus, southeastasia, southindia, switzerlandnorth, "
        "switzerlandwest, uaecentral, uaenorth, uksouth, ukwest, westcentralus, westeurope, westindia, westus, westus2.",
    ),
    username: Optional[str] = typer.Option(
        None,
        "--username",
        "-u",
        help=f'Database username. If not passed, the default value "root" will be used unless the value is explicitly set in the environment variable **{CLIENT_DB_USERNAME}**.',
    ),
    password: Optional[str] = typer.Option(
        None,
        "--password",
        "-p",
        help=f'Database password. If not passed, the default value "" will be used unless the value is explicitly set in the environment variable **{CLIENT_DB_PASSWORD}**.',
    ),
    tag: Optional[str] = typer.Option(
        None,
        "--tag",
        "-t",
        help="A string to tag the datablob. If not passed, then the tag **latest** will be assigned to the datablob.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output datablob uuid only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Create and return a datablob that encapsulates the data from a mysql database.

    If the database requires authentication, pass the username/password as commandline arguments or store it in
    the **AIRT_CLIENT_DB_USERNAME** and **AIRT_CLIENT_DB_PASSWORD** environment variables.
    """

    from airt.client import DataBlob

    db = DataBlob.from_mysql(
        host=host,
        database=database,
        port=port,
        table=table,
        username=username,
        password=password,
        cloud_provider=cloud_provider,
        region=region,
        tag=tag,
    )

    if quiet:
        db.wait()
        typer.echo(f"{db.uuid}")
    else:
        typer.echo(f"Pulling datablob uuid: {db.uuid}")
        db.progress_bar()

# %% ../../notebooks/CLI_DataBlob.ipynb 34
@app.command("to-datasource")
@helper.requires_auth_token
def to_datasource(
    uuid: str = typer.Option(
        ...,
        help="Datablob uuid.",
    ),
    file_type: str = typer.Option(
        ...,
        help='The file type of the datablob. Currently, the API only supports "csv" and "parquet" as file types.',
    ),
    index_column: str = typer.Option(
        ...,
        help="The column to use as index (row labels).",
    ),
    sort_by: str = typer.Option(
        ...,
        help="The column(s) to sort the data. Can either be a string or a JSON encoded list of strings.",
    ),
    deduplicate_data: bool = typer.Option(
        False,
        help="If set to **True** (default value **False**), the datasource will be created with duplicate rows removed.",
    ),
    blocksize: str = typer.Option(
        "256MB",
        help="The number of bytes used to split larger files. If None, then the default value **256MB** will be used.",
    ),
    kwargs_json: Optional[str] = typer.Option(
        None,
        help="Additional JSON encoded dict arguments to use while processing the data.e.g: To skip 100 lines from the bottom of file, pass '{"
        '"skipfooter"'
        ": 100}'",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output datasource uuid only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Process the datablob and return a datasource object."""

    from airt.client import DataBlob

    kwargs = json.loads(kwargs_json) if kwargs_json else {}

    try:
        sort_by = json.loads(sort_by)

    except json.JSONDecodeError as e:
        pass

    db = DataBlob(uuid=uuid)
    ds = db.to_datasource(
        file_type=file_type,
        index_column=index_column,
        sort_by=sort_by,
        deduplicate_data=deduplicate_data,
        blocksize=blocksize,
        **kwargs,
    )

    if quiet:
        ds.wait()
        typer.echo(f"{ds.uuid}")
    else:
        typer.echo(f"Processing and pulling the datasource uuid: {ds.uuid}")

        ds.progress_bar()

# %% ../../notebooks/CLI_DataBlob.ipynb 38
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def ls(
    offset: int = typer.Option(
        0,
        "--offset",
        "-o",
        help="The number of datablobs to offset at the beginning. If **None**, then the default value **0** will be used.",
    ),
    limit: int = typer.Option(
        100,
        "--limit",
        "-l",
        help="The maximum number of datablobs to return from the server. If **None**, then the default value **100** will be used.",
    ),
    disabled: bool = typer.Option(
        False,
        "--disabled",
        help="If set to **True**, then only the deleted datablobs will be returned."
        "Else, the default value **False** will be used to return only the list"
        "of active datablobs.",
    ),
    completed: bool = typer.Option(
        False,
        "--completed",
        help="If set to **True**, then only the datablobs that are successfully downloaded"
        "to the server will be returned. Else, the default value **False** will be used to"
        "return all the datablobs.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output only datablob uuids separated by space",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Return the list of datablobs."""

    from airt.client import DataBlob

    dbx = DataBlob.ls(
        offset=offset, limit=limit, disabled=disabled, completed=completed
    )

    df = DataBlob.as_df(dbx)

    df["pulled_on"] = helper.humanize_date(df["pulled_on"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    return {"df": df, "quite_column_name": "datablob_uuid"}

# %% ../../notebooks/CLI_DataBlob.ipynb 43
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def rm(
    uuid: str = typer.Argument(
        ...,
        help="Datablob uuid.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output the deleted datablob uuid only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Delete a datablob from the server."""

    from airt.client import DataBlob

    db = DataBlob(uuid=uuid)
    df = db.delete()

    df["pulled_on"] = helper.humanize_date(df["pulled_on"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    return {"df": df, "quite_column_name": "datablob_uuid"}

# %% ../../notebooks/CLI_DataBlob.ipynb 47
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def tag(
    uuid: str = typer.Option(
        ...,
        "--datablob_uuid",
        "-uuid",
        help="Datablob uuid in the server.",
    ),
    name: str = typer.Option(
        ...,
        "--name",
        "-n",
        help="A string to tag the datablob.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Tag an existing datablob in the server."""

    from airt.client import DataBlob

    db = DataBlob(uuid=uuid)
    df = db.tag(name=name)

    df["pulled_on"] = helper.humanize_date(df["pulled_on"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    return {"df": df}

# %% ../../notebooks/CLI_DataBlob.ipynb 50
@app.command("from-local")
@helper.requires_auth_token
def from_local(
    path: str = typer.Option(
        ...,
        "--path",
        "-p",
        help="The relative or absolute path to a local CSV/parquet file or to a directory containing the CSV/parquet files.",
    ),
    cloud_provider: Optional[str] = typer.Option(
        None,
        "--cloud-provider",
        "-cp",
        help="The destination cloud storage provider's name to store the datablob. Currently, the API only supports **aws** and **azure** as cloud storage providers. If **None** (default value), then **aws**  will be used as the cloud storage provider.",
    ),
    region: Optional[str] = typer.Option(
        None,
        "--region",
        "-r",
        help="The destination cloud provider's region to save your datablob. If **None** (default value) then the default region will be assigned based on the cloud provider. "
        "In the case of **aws**, **eu-west-1** will be used and in the case of **azure**, **westeurope** will be used. "
        "The supported AWS regions are: ap-northeast-1, ap-northeast-2, ap-south-1, ap-southeast-1, ap-southeast-2, ca-central-1, eu-central-1, "
        "eu-north-1, eu-west-1, eu-west-2, eu-west-3, sa-east-1, us-east-1, us-east-2, us-west-1, us-west-2. The supported Azure Blob Storage "
        "regions are: australiacentral, australiacentral2, australiaeast, australiasoutheast, brazilsouth, canadacentral, canadaeast, centralindia, "
        "centralus, eastasia, eastus, eastus2, francecentral, francesouth, germanynorth, germanywestcentral, japaneast, japanwest, koreacentral, koreasouth, "
        "northcentralus, northeurope, norwayeast, norwaywest, southafricanorth, southafricawest, southcentralus, southeastasia, southindia, switzerlandnorth, "
        "switzerlandwest, uaecentral, uaenorth, uksouth, ukwest, westcentralus, westeurope, westindia, westus, westus2.",
    ),
    tag: Optional[str] = typer.Option(
        None,
        "--tag",
        "-t",
        help="A string to tag the datablob. If not passed, then the tag **latest** will be assigned to the datablob.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output data id only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Create and return a datablob from local csv file.

    The API currently allows users to create datablobs from CSV or Parquet files. We intend to support additional file formats in future releases.
    """

    from airt.client import DataBlob

    if quiet:
        db = DataBlob.from_local(
            path=path,
            cloud_provider=cloud_provider,
            region=region,
            tag=tag,
            show_progress=False,
        )
        typer.echo(f"{db.uuid}")
    else:
        db = DataBlob.from_local(
            path=path, cloud_provider=cloud_provider, region=region, tag=tag
        )
        typer.echo(f"Successfully pulled the datablob uuid: {db.uuid}.")

# %% ../../notebooks/CLI_DataBlob.ipynb 54
@app.command("from-clickhouse")
@helper.requires_auth_token
def from_clickhouse(
    host: str = typer.Option(..., help="Remote database host name."),
    database: str = typer.Option(..., help="Database name."),
    table: str = typer.Option(..., help="Table name."),
    protocol: str = typer.Option(
        ..., help='Protocol to use. The valid values are "native" and "http".'
    ),
    index_column: str = typer.Option(
        ..., help="The column to use as index (row labels)."
    ),
    timestamp_column: str = typer.Option(
        ..., help="Timestamp column name in the tabel."
    ),
    port: int = typer.Option(
        0,
        help="Host port number. If not passed, then the default value **0** will be used.",
    ),
    cloud_provider: Optional[str] = typer.Option(
        None,
        "--cloud-provider",
        "-cp",
        help="The destination cloud storage provider's name to store the datablob. Currently, the API only supports **aws** and **azure** as cloud storage providers. If **None** (default value), then **aws**  will be used as the cloud storage provider.",
    ),
    region: Optional[str] = typer.Option(
        None,
        "--region",
        "-r",
        help="The destination cloud provider's region to save your datablob. If **None** (default value) then the default region will be assigned based on the cloud provider. "
        "In the case of **aws**, **eu-west-1** will be used and in the case of **azure**, **westeurope** will be used. "
        "The supported AWS regions are: ap-northeast-1, ap-northeast-2, ap-south-1, ap-southeast-1, ap-southeast-2, ca-central-1, eu-central-1, "
        "eu-north-1, eu-west-1, eu-west-2, eu-west-3, sa-east-1, us-east-1, us-east-2, us-west-1, us-west-2. The supported Azure Blob Storage "
        "regions are: australiacentral, australiacentral2, australiaeast, australiasoutheast, brazilsouth, canadacentral, canadaeast, centralindia, "
        "centralus, eastasia, eastus, eastus2, francecentral, francesouth, germanynorth, germanywestcentral, japaneast, japanwest, koreacentral, koreasouth, "
        "northcentralus, northeurope, norwayeast, norwaywest, southafricanorth, southafricawest, southcentralus, southeastasia, southindia, switzerlandnorth, "
        "switzerlandwest, uaecentral, uaenorth, uksouth, ukwest, westcentralus, westeurope, westindia, westus, westus2.",
    ),
    username: Optional[str] = typer.Option(
        None,
        "--username",
        "-u",
        help="Database username. If not passed, the default value 'root' will be used unless the value is explicitly set in the environment variable **CLICKHOUSE_USERNAME**.",
    ),
    password: Optional[str] = typer.Option(
        None,
        "--password",
        "-p",
        help="Database password. If not passed, the default value '' will be used unless the value is explicitly set in the environment variable **CLICKHOUSE_PASSWORD**.",
    ),
    filters_json: Optional[str] = typer.Option(
        None,
        "--filters-json",
        "-f",
        help="Additional parameters to be used when importing data. For example, if you want to filter and extract data only for a specific user_id, pass '{"
        '"user_id"'
        ": 1}'.",
    ),
    tag: Optional[str] = typer.Option(
        None,
        "--tag",
        "-t",
        help="A string to tag the datablob. If not passed, then the tag **latest** will be assigned to the datablob.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output datablob uuid only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Create and return a datablob that encapsulates the data from a ClickHouse database.

    If the database requires authentication, pass the username/password as commandline arguments or store it in
    the **CLICKHOUSE_USERNAME** and **CLICKHOUSE_PASSWORD** environment variables.
    """

    filters = json.loads(filters_json) if filters_json else None

    from airt.client import DataBlob

    db = DataBlob.from_clickhouse(
        host=host,
        database=database,
        table=table,
        protocol=protocol,
        index_column=index_column,
        timestamp_column=timestamp_column,
        port=port,
        username=username,
        password=password,
        filters=filters,
        cloud_provider=cloud_provider,
        region=region,
        tag=tag,
    )

    if quiet:
        db.wait()
        typer.echo(f"{db.uuid}")
    else:
        typer.echo(f"Pulling datablob uuid: {db.uuid}")
        db.progress_bar()
