"""Query helpers."""
import boto3

from dataexport.conf import config
from dataexport.conf import constants
from dataexport.dtos import Job


SF_TABLE = "{DB}.{SCHEMA}.DATA_EXPORT_STATUS".format(
    DB=config.SNOWFLAKE_DATABASE, SCHEMA=config.SNOWFLAKE_SCHEMA)


def format_query(query: str, job_id: str) -> str:
    credentials = boto3.Session().get_credentials()
    return """
    CALL DEV_ENGINEERING.TEST_SCHEMA.export_to_s3(
        '{job_id}',
        '{query}',
        'LargeDataExport',
        '{aws_key_id}',
        '{aws_secret}',
        '{aws_token}'
    );
    """.format(
        job_id=job_id,
        query=query,
        aws_key_id=credentials.access_key,
        aws_secret=credentials.secret_key,
        aws_token=credentials.token
    )


def format_status_query(job: Job) -> str:
    if job.status == constants.JOB_STATUS_STARTED:
        return "INSERT INTO {table} (job_id, query_string, job_status) VALUES ('{id}','{query}','{status}')".format(
            table=SF_TABLE, status=job.status, id=job.id, query=job.query_string)
    else:
        return "UPDATE {table} SET job_status = '{status}' WHERE job_id = '{id}'".format(
            table=SF_TABLE, status=job.status, id=job.id)