from montecarlodata.config import Config
from montecarlodata.errors import manage_errors
from montecarlodata.integrations.onboarding.base import BaseOnboardingService
from montecarlodata.integrations.onboarding.fields import EXPECTED_GENERIC_DB_GQL_RESPONSE_FIELD, \
    EXPECTED_HIVE_S3_GQL_RESPONSE_FIELD, EXPECTED_HIVE_SQL_GQL_RESPONSE_FIELD, HIVE_MYSQL_CONNECTION_TYPE, \
    HIVE_S3_CONNECTION_TYPE, HIVE_SQL_CONNECTION_TYPE, QL_JOB_TYPE, AWS_RDS_CA_CERT, PRESTO_CATALOG_KEY, \
    HIVE_GET_PARTS_KEY, HIVE_MAX_PARTS_KEY, HIVE_MAX_PARTS_DEFAULT_VALUE
from montecarlodata.queries.onboarding import TEST_DATABASE_CRED_MUTATION, TEST_S3_CRED_MUTATION, \
    TEST_HIVE_SQL_CRED_MUTATION


class HiveOnboardingService(BaseOnboardingService):

    def __init__(self, config: Config, **kwargs):
        super().__init__(config, **kwargs)

    @manage_errors
    def onboard_hive_mysql(self, **kwargs) -> None:
        """
        Onboard a hive-mysql connection by validating and adding a connection
        """
        if kwargs['use_ssl']:
            kwargs['ssl_options'] = {'ca': AWS_RDS_CA_CERT}
            kwargs.pop('use_ssl')

        catalog = kwargs.pop('catalog', None)
        job_limits = {HIVE_GET_PARTS_KEY: True, HIVE_MAX_PARTS_KEY: HIVE_MAX_PARTS_DEFAULT_VALUE}

        if catalog:
            job_limits[PRESTO_CATALOG_KEY] = catalog

        kwargs['connectionType'] = HIVE_MYSQL_CONNECTION_TYPE
        self.onboard(validation_query=TEST_DATABASE_CRED_MUTATION,
                     validation_response=EXPECTED_GENERIC_DB_GQL_RESPONSE_FIELD,
                     connection_type=HIVE_MYSQL_CONNECTION_TYPE, job_limits=job_limits, **kwargs)

    @manage_errors
    def onboard_hive_s3(self, **kwargs) -> None:
        """
        Onboard a presto-s3 connection by validating and adding a connection
        """
        kwargs['connectionType'] = HIVE_S3_CONNECTION_TYPE
        self.onboard(validation_query=TEST_S3_CRED_MUTATION,
                     validation_response=EXPECTED_HIVE_S3_GQL_RESPONSE_FIELD,
                     connection_type=HIVE_S3_CONNECTION_TYPE, job_types=QL_JOB_TYPE, **kwargs)

    @manage_errors
    def onboard_hive_sql(self, **kwargs) -> None:
        """
        Onboard a hive-sql connection by validating and adding a connection
        """
        self.onboard(validation_query=TEST_HIVE_SQL_CRED_MUTATION,
                     validation_response=EXPECTED_HIVE_SQL_GQL_RESPONSE_FIELD,
                     connection_type=HIVE_SQL_CONNECTION_TYPE, **kwargs)
