# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
import json
import re

from airflow.hooks.http_hook import HttpHook
from airflow.exceptions import AirflowException

class DiscordWebhookHook(HttpHook):
    """
    This hook allows you to post messages to Discord using incoming webhooks.
    Takes a Discord connection ID with a default relative webhook endpoint. The
    default endpoint can be overridden using the webhook_endpoint parameter
    (https://discordapp.com/developers/docs/resources/webhook).

    Each Discord webhook can be pre-configured to use a specific username and
    avatar_url. You can override these defaults in this hook.

    :param http_conn_id: Http connection ID with host as "https://discord.com/api/" and
                         default webhook endpoint in the extra field in the form of
                         {"webhook_endpoint": "webhooks/{webhook.id}/{webhook.token}"}
    :type http_conn_id: str
    :param webhook_endpoint: Discord webhook endpoint in the form of
                             "webhooks/{webhook.id}/{webhook.token}"
    :type webhook_endpoint: str
    :param message: The message you want to send to your Discord channel
                    (max 2000 characters)
    :type message: str
    :param username: Override the default username of the webhook
    :type username: str
    :param avatar_url: Override the default avatar of the webhook
    :type avatar_url: str
    :param tts: Is a text-to-speech message
    :type tts: bool
    :param proxy: Proxy to use to make the Discord webhook call
    :type proxy: str
    """

    def __init__(self,
                 http_conn_id=None,
                 webhook_endpoint=None,
                 message="",
                 username=None,
                 avatar_url=None,
                 tts=False,
                 proxy=None,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.http_conn_id = http_conn_id
        self.webhook_endpoint = self._get_webhook_endpoint(http_conn_id, webhook_endpoint)
        self.message = message
        self.username = username
        self.avatar_url = avatar_url
        self.tts = tts
        self.proxy = proxy

    def _get_webhook_endpoint(self, http_conn_id, webhook_endpoint):
        """
        Given a Discord http_conn_id, return the default webhook endpoint or override if a
        webhook_endpoint is manually supplied.

        :param http_conn_id: The provided connection ID
        :param webhook_endpoint: The manually provided webhook endpoint
        :return: Webhook endpoint (str) to use
        """
        if webhook_endpoint:
            endpoint = webhook_endpoint
        elif http_conn_id:
            conn = self.get_connection(http_conn_id)
            extra = conn.extra_dejson
            endpoint = extra.get('webhook_endpoint', '')
        else:
            raise AirflowException('Cannot get webhook endpoint: No valid Discord '
                                   'webhook endpoint or http_conn_id supplied.')

        # make sure endpoint matches the expected Discord webhook format
        if not re.match('^webhooks/[0-9]+/[a-zA-Z0-9_-]+$', endpoint):
            raise AirflowException('Expected Discord webhook endpoint in the form '
                                   'of "webhooks/{webhook.id}/{webhook.token}".')

        return endpoint

    def _build_discord_payload(self):
        """
        Construct the Discord JSON payload. All relevant parameters are combined here
        to a valid Discord JSON payload.

        :return: Discord payload (str) to send
        """
        payload = {}

        if self.username:
            payload['username'] = self.username
        if self.avatar_url:
            payload['avatar_url'] = self.avatar_url

        payload['tts'] = self.tts

        if len(self.message) <= 2000:
            payload['content'] = self.message
        else:
            raise AirflowException('Discord message length must be 2000 or fewer '
                                   'characters.')

        return json.dumps(payload)

    def execute(self):
        """
        Execute the Discord webhook call
        """
        proxies = {}
        if self.proxy:
            # we only need https proxy for Discord
            proxies = {'https': self.proxy}

        discord_payload = self._build_discord_payload()

        self.run(endpoint=self.webhook_endpoint,
                 data=discord_payload,
                 headers={'Content-type': 'application/json'},
                 extra_options={'proxies': proxies})


# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

import base64

from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook

from googleapiclient.discovery import build

def _b64encode(s):
    """ Base 64 encodes a bytes object to a string """
    return base64.b64encode(s).decode('ascii')

def _b64decode(s):
    """ Base 64 decodes a string to bytes. """
    return base64.b64decode(s.encode('utf-8'))

class GoogleCloudKMSHook(GoogleCloudBaseHook):
    """
    Interact with Google Cloud KMS. This hook uses the Google Cloud Platform
    connection.
    """

    def __init__(self, gcp_conn_id='google_cloud_default', delegate_to=None):
        super().__init__(gcp_conn_id, delegate_to=delegate_to)
        self.num_retries = self._get_field('num_retries', 5)

    def get_conn(self):
        """
        Returns a KMS service object.

        :rtype: googleapiclient.discovery.Resource
        """
        http_authorized = self._authorize()
        return build(
            'cloudkms', 'v1', http=http_authorized, cache_discovery=False)

    def encrypt(self, key_name, plaintext, authenticated_data=None):
        """
        Encrypts a plaintext message using Google Cloud KMS.

        :param key_name: The Resource Name for the key (or key version)
                         to be used for encyption. Of the form
                         ``projects/*/locations/*/keyRings/*/cryptoKeys/**``
        :type key_name: str
        :param plaintext: The message to be encrypted.
        :type plaintext: bytes
        :param authenticated_data: Optional additional authenticated data that
                                   must also be provided to decrypt the message.
        :type authenticated_data: bytes
        :return: The base 64 encoded ciphertext of the original message.
        :rtype: str
        """
        keys = self.get_conn().projects().locations().keyRings().cryptoKeys()
        body = {'plaintext': _b64encode(plaintext)}
        if authenticated_data:
            body['additionalAuthenticatedData'] = _b64encode(authenticated_data)

        request = keys.encrypt(name=key_name, body=body)
        response = request.execute(num_retries=self.num_retries)

        ciphertext = response['ciphertext']
        return ciphertext

    def decrypt(self, key_name, ciphertext, authenticated_data=None):
        """
        Decrypts a ciphertext message using Google Cloud KMS.

        :param key_name: The Resource Name for the key to be used for decyption.
                         Of the form ``projects/*/locations/*/keyRings/*/cryptoKeys/**``
        :type key_name: str
        :param ciphertext: The message to be decrypted.
        :type ciphertext: str
        :param authenticated_data: Any additional authenticated data that was
                                   provided when encrypting the message.
        :type authenticated_data: bytes
        :return: The original message.
        :rtype: bytes
        """
        keys = self.get_conn().projects().locations().keyRings().cryptoKeys()
        body = {'ciphertext': ciphertext}
        if authenticated_data:
            body['additionalAuthenticatedData'] = _b64encode(authenticated_data)

        request = keys.decrypt(name=key_name, body=body)
        response = request.execute(num_retries=self.num_retries)

        plaintext = _b64decode(response['plaintext'])
        return plaintext


# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

"""
This module contains a sqoop 1.x hook
"""
import subprocess

from airflow.exceptions import AirflowException
from airflow.hooks.base_hook import BaseHook
from copy import deepcopy

class SqoopHook(BaseHook):
    """
    This hook is a wrapper around the sqoop 1 binary. To be able to use the hook
    it is required that "sqoop" is in the PATH.

    Additional arguments that can be passed via the 'extra' JSON field of the
    sqoop connection:

        * ``job_tracker``: Job tracker local|jobtracker:port.
        * ``namenode``: Namenode.
        * ``lib_jars``: Comma separated jar files to include in the classpath.
        * ``files``: Comma separated files to be copied to the map reduce cluster.
        * ``archives``: Comma separated archives to be unarchived on the compute
            machines.
        * ``password_file``: Path to file containing the password.

    :param conn_id: Reference to the sqoop connection.
    :type conn_id: str
    :param verbose: Set sqoop to verbose.
    :type verbose: bool
    :param num_mappers: Number of map tasks to import in parallel.
    :type num_mappers: int
    :param properties: Properties to set via the -D argument
    :type properties: dict
    """

    def __init__(self, conn_id='sqoop_default', verbose=False,
                 num_mappers=None, hcatalog_database=None,
                 hcatalog_table=None, properties=None):
        # No mutable types in the default parameters
        self.conn = self.get_connection(conn_id)
        connection_parameters = self.conn.extra_dejson
        self.job_tracker = connection_parameters.get('job_tracker', None)
        self.namenode = connection_parameters.get('namenode', None)
        self.libjars = connection_parameters.get('libjars', None)
        self.files = connection_parameters.get('files', None)
        self.archives = connection_parameters.get('archives', None)
        self.password_file = connection_parameters.get('password_file', None)
        self.hcatalog_database = hcatalog_database
        self.hcatalog_table = hcatalog_table
        self.verbose = verbose
        self.num_mappers = num_mappers
        self.properties = properties or {}
        self.log.info(
            "Using connection to: {}:{}/{}".format(
                self.conn.host, self.conn.port, self.conn.schema
            )
        )

    def get_conn(self):
        return self.conn

    def cmd_mask_password(self, cmd_orig):
        cmd = deepcopy(cmd_orig)
        try:
            password_index = cmd.index('--password')
            cmd[password_index + 1] = 'MASKED'
        except ValueError:
            self.log.debug("No password in sqoop cmd")
        return cmd

    def Popen(self, cmd, **kwargs):
        """
        Remote Popen

        :param cmd: command to remotely execute
        :param kwargs: extra arguments to Popen (see subprocess.Popen)
        :return: handle to subprocess
        """
        masked_cmd = ' '.join(self.cmd_mask_password(cmd))
        self.log.info("Executing command: {}".format(masked_cmd))
        self.sp = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            **kwargs)

        for line in iter(self.sp.stdout):
            self.log.info(line.strip())

        self.sp.wait()

        self.log.info("Command exited with return code %s", self.sp.returncode)

        if self.sp.returncode:
            raise AirflowException("Sqoop command failed: {}".format(masked_cmd))

    def _prepare_command(self, export=False):
        sqoop_cmd_type = "export" if export else "import"
        connection_cmd = ["sqoop", sqoop_cmd_type]

        for key, value in self.properties.items():
            connection_cmd += ["-D", "{}={}".format(key, value)]

        if self.namenode:
            connection_cmd += ["-fs", self.namenode]
        if self.job_tracker:
            connection_cmd += ["-jt", self.job_tracker]
        if self.libjars:
            connection_cmd += ["-libjars", self.libjars]
        if self.files:
            connection_cmd += ["-files", self.files]
        if self.archives:
            connection_cmd += ["-archives", self.archives]
        if self.conn.login:
            connection_cmd += ["--username", self.conn.login]
        if self.conn.password:
            connection_cmd += ["--password", self.conn.password]
        if self.password_file:
            connection_cmd += ["--password-file", self.password_file]
        if self.verbose:
            connection_cmd += ["--verbose"]
        if self.num_mappers:
            connection_cmd += ["--num-mappers", str(self.num_mappers)]
        if self.hcatalog_database:
            connection_cmd += ["--hcatalog-database", self.hcatalog_database]
        if self.hcatalog_table:
            connection_cmd += ["--hcatalog-table", self.hcatalog_table]

        connection_cmd += ["--connect", "{}:{}/{}".format(
            self.conn.host,
            self.conn.port,
            self.conn.schema
        )]

        return connection_cmd

    @staticmethod
    def _get_export_format_argument(file_type='text'):
        if file_type == "avro":
            return ["--as-avrodatafile"]
        elif file_type == "sequence":
            return ["--as-sequencefile"]
        elif file_type == "parquet":
            return ["--as-parquetfile"]
        elif file_type == "text":
            return ["--as-textfile"]
        else:
            raise AirflowException("Argument file_type should be 'avro', "
                                   "'sequence', 'parquet' or 'text'.")

    def _import_cmd(self, target_dir, append, file_type, split_by, direct,
                    driver, extra_import_options):

        cmd = self._prepare_command(export=False)

        if target_dir:
            cmd += ["--target-dir", target_dir]

        if append:
            cmd += ["--append"]

        cmd += self._get_export_format_argument(file_type)

        if split_by:
            cmd += ["--split-by", split_by]

        if direct:
            cmd += ["--direct"]

        if driver:
            cmd += ["--driver", driver]

        if extra_import_options:
            for key, value in extra_import_options.items():
                cmd += ['--{}'.format(key)]
                if value:
                    cmd += [str(value)]

        return cmd

    def import_table(self, table, target_dir=None, append=False, file_type="text",
                     columns=None, split_by=None, where=None, direct=False,
                     driver=None, extra_import_options=None):
        """
        Imports table from remote location to target dir. Arguments are
        copies of direct sqoop command line arguments

        :param table: Table to read
        :param target_dir: HDFS destination dir
        :param append: Append data to an existing dataset in HDFS
        :param file_type: "avro", "sequence", "text" or "parquet".
            Imports data to into the specified format. Defaults to text.
        :param columns: <col,col,col…> Columns to import from table
        :param split_by: Column of the table used to split work units
        :param where: WHERE clause to use during import
        :param direct: Use direct connector if exists for the database
        :param driver: Manually specify JDBC driver class to use
        :param extra_import_options: Extra import options to pass as dict.
            If a key doesn't have a value, just pass an empty string to it.
            Don't include prefix of -- for sqoop options.
        """
        cmd = self._import_cmd(target_dir, append, file_type, split_by, direct,
                               driver, extra_import_options)

        cmd += ["--table", table]

        if columns:
            cmd += ["--columns", columns]
        if where:
            cmd += ["--where", where]

        self.Popen(cmd)

    def import_query(self, query, target_dir, append=False, file_type="text",
                     split_by=None, direct=None, driver=None, extra_import_options=None):
        """
        Imports a specific query from the rdbms to hdfs

        :param query: Free format query to run
        :param target_dir: HDFS destination dir
        :param append: Append data to an existing dataset in HDFS
        :param file_type: "avro", "sequence", "text" or "parquet"
            Imports data to hdfs into the specified format. Defaults to text.
        :param split_by: Column of the table used to split work units
        :param direct: Use direct import fast path
        :param driver: Manually specify JDBC driver class to use
        :param extra_import_options: Extra import options to pass as dict.
            If a key doesn't have a value, just pass an empty string to it.
            Don't include prefix of -- for sqoop options.
        """
        cmd = self._import_cmd(target_dir, append, file_type, split_by, direct,
                               driver, extra_import_options)
        cmd += ["--query", query]

        self.Popen(cmd)

    def _export_cmd(self, table, export_dir, input_null_string,
                    input_null_non_string, staging_table, clear_staging_table,
                    enclosed_by, escaped_by, input_fields_terminated_by,
                    input_lines_terminated_by, input_optionally_enclosed_by,
                    batch, relaxed_isolation, extra_export_options):

        cmd = self._prepare_command(export=True)

        if input_null_string:
            cmd += ["--input-null-string", input_null_string]

        if input_null_non_string:
            cmd += ["--input-null-non-string", input_null_non_string]

        if staging_table:
            cmd += ["--staging-table", staging_table]

        if clear_staging_table:
            cmd += ["--clear-staging-table"]

        if enclosed_by:
            cmd += ["--enclosed-by", enclosed_by]

        if escaped_by:
            cmd += ["--escaped-by", escaped_by]

        if input_fields_terminated_by:
            cmd += ["--input-fields-terminated-by", input_fields_terminated_by]

        if input_lines_terminated_by:
            cmd += ["--input-lines-terminated-by", input_lines_terminated_by]

        if input_optionally_enclosed_by:
            cmd += ["--input-optionally-enclosed-by",
                    input_optionally_enclosed_by]

        if batch:
            cmd += ["--batch"]

        if relaxed_isolation:
            cmd += ["--relaxed-isolation"]

        if export_dir:
            cmd += ["--export-dir", export_dir]

        if extra_export_options:
            for key, value in extra_export_options.items():
                cmd += ['--{}'.format(key)]
                if value:
                    cmd += [str(value)]

        # The required option
        cmd += ["--table", table]

        return cmd

    def export_table(self, table, export_dir, input_null_string,
                     input_null_non_string, staging_table,
                     clear_staging_table, enclosed_by,
                     escaped_by, input_fields_terminated_by,
                     input_lines_terminated_by,
                     input_optionally_enclosed_by, batch,
                     relaxed_isolation, extra_export_options=None):
        """
        Exports Hive table to remote location. Arguments are copies of direct
        sqoop command line Arguments

        :param table: Table remote destination
        :param export_dir: Hive table to export
        :param input_null_string: The string to be interpreted as null for
            string columns
        :param input_null_non_string: The string to be interpreted as null
            for non-string columns
        :param staging_table: The table in which data will be staged before
            being inserted into the destination table
        :param clear_staging_table: Indicate that any data present in the
            staging table can be deleted
        :param enclosed_by: Sets a required field enclosing character
        :param escaped_by: Sets the escape character
        :param input_fields_terminated_by: Sets the field separator character
        :param input_lines_terminated_by: Sets the end-of-line character
        :param input_optionally_enclosed_by: Sets a field enclosing character
        :param batch: Use batch mode for underlying statement execution
        :param relaxed_isolation: Transaction isolation to read uncommitted
            for the mappers
        :param extra_export_options: Extra export options to pass as dict.
            If a key doesn't have a value, just pass an empty string to it.
            Don't include prefix of -- for sqoop options.
        """
        cmd = self._export_cmd(table, export_dir, input_null_string,
                               input_null_non_string, staging_table,
                               clear_staging_table, enclosed_by, escaped_by,
                               input_fields_terminated_by,
                               input_lines_terminated_by,
                               input_optionally_enclosed_by, batch,
                               relaxed_isolation, extra_export_options)

        self.Popen(cmd)


# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from google.cloud.texttospeech_v1 import TextToSpeechClient

from airflow.contrib.hooks.gcp_api_base_hook import GoogleCloudBaseHook

class GCPTextToSpeechHook(GoogleCloudBaseHook):
    """
    Hook for Google Cloud Text to Speech API.

    :param gcp_conn_id: The connection ID to use when fetching connection info.
    :type gcp_conn_id: str
    :param delegate_to: The account to impersonate, if any.
        For this to work, the service account making the request must have
        domain-wide delegation enabled.
    :type delegate_to: str
    """

    _client = None

    def __init__(self, gcp_conn_id="google_cloud_default", delegate_to=None):
        super().__init__(gcp_conn_id, delegate_to)

    def get_conn(self):
        """
        Retrieves connection to Cloud Text to Speech.

        :return: Google Cloud Text to Speech client object.
        :rtype: google.cloud.texttospeech_v1.TextToSpeechClient
        """
        if not self._client:
            self._client = TextToSpeechClient(credentials=self._get_credentials())
        return self._client

    def synthesize_speech(self, input_data, voice, audio_config, retry=None, timeout=None):
        """
        Synthesizes text input

        :param input_data: text input to be synthesized. See more:
            https://googleapis.github.io/google-cloud-python/latest/texttospeech/gapic/v1/types.html#google.cloud.texttospeech_v1.types.SynthesisInput
        :type input_data: dict or google.cloud.texttospeech_v1.types.SynthesisInput
        :param voice: configuration of voice to be used in synthesis. See more:
            https://googleapis.github.io/google-cloud-python/latest/texttospeech/gapic/v1/types.html#google.cloud.texttospeech_v1.types.VoiceSelectionParams
        :type voice: dict or google.cloud.texttospeech_v1.types.VoiceSelectionParams
        :param audio_config: configuration of the synthesized audio. See more:
            https://googleapis.github.io/google-cloud-python/latest/texttospeech/gapic/v1/types.html#google.cloud.texttospeech_v1.types.AudioConfig
        :type audio_config: dict or google.cloud.texttospeech_v1.types.AudioConfig
        :return: SynthesizeSpeechResponse See more:
            https://googleapis.github.io/google-cloud-python/latest/texttospeech/gapic/v1/types.html#google.cloud.texttospeech_v1.types.SynthesizeSpeechResponse
        :rtype: object
        :param retry: (Optional) A retry object used to retry requests. If None is specified,
                requests will not be retried.
        :type retry: google.api_core.retry.Retry
        :param timeout: (Optional) The amount of time, in seconds, to wait for the request to complete.
            Note that if retry is specified, the timeout applies to each individual attempt.
        :type timeout: float
        """
        client = self.get_conn()
        self.log.info("Synthesizing input: %s" % input_data)
        return client.synthesize_speech(
            input_=input_data, voice=voice, audio_config=audio_config, retry=retry, timeout=timeout
        )


# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
import os

from cached_property import cached_property

from airflow import configuration
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.log.file_task_handler import FileTaskHandler

class S3TaskHandler(FileTaskHandler, LoggingMixin):
    """
    S3TaskHandler is a python log handler that handles and reads
    task instance logs. It extends airflow FileTaskHandler and
    uploads to and reads from S3 remote storage.
    """
    def __init__(self, base_log_folder, s3_log_folder, filename_template):
        super().__init__(base_log_folder, filename_template)
        self.remote_base = s3_log_folder
        self.log_relative_path = ''
        self._hook = None
        self.closed = False
        self.upload_on_close = True

    @cached_property
    def hook(self):
        remote_conn_id = configuration.conf.get('core', 'REMOTE_LOG_CONN_ID')
        try:
            from airflow.hooks.S3_hook import S3Hook
            return S3Hook(remote_conn_id)
        except Exception:
            self.log.error(
                'Could not create an S3Hook with connection id "%s". '
                'Please make sure that airflow[aws] is installed and '
                'the S3 connection exists.', remote_conn_id
            )

    def set_context(self, ti):
        super().set_context(ti)
        # Local location and remote location is needed to open and
        # upload local log file to S3 remote storage.
        self.log_relative_path = self._render_filename(ti, ti.try_number)
        self.upload_on_close = not ti.raw

    def close(self):
        """
        Close and upload local log file to remote storage S3.
        """
        # When application exit, system shuts down all handlers by
        # calling close method. Here we check if logger is already
        # closed to prevent uploading the log to remote storage multiple
        # times when `logging.shutdown` is called.
        if self.closed:
            return

        super().close()

        if not self.upload_on_close:
            return

        local_loc = os.path.join(self.local_base, self.log_relative_path)
        remote_loc = os.path.join(self.remote_base, self.log_relative_path)
        if os.path.exists(local_loc):
            # read log and remove old logs to get just the latest additions
            with open(local_loc, 'r') as logfile:
                log = logfile.read()
            self.s3_write(log, remote_loc)

        # Mark closed so we don't double write if close is called twice
        self.closed = True

    def _read(self, ti, try_number, metadata=None):
        """
        Read logs of given task instance and try_number from S3 remote storage.
        If failed, read the log from task instance host machine.
        :param ti: task instance object
        :param try_number: task instance try_number to read logs from
        :param metadata: log metadata,
                         can be used for steaming log reading and auto-tailing.
        """
        # Explicitly getting log relative path is necessary as the given
        # task instance might be different than task instance passed in
        # in set_context method.
        log_relative_path = self._render_filename(ti, try_number)
        remote_loc = os.path.join(self.remote_base, log_relative_path)

        if self.s3_log_exists(remote_loc):
            # If S3 remote file exists, we do not fetch logs from task instance
            # local machine even if there are errors reading remote logs, as
            # returned remote_log will contain error messages.
            remote_log = self.s3_read(remote_loc, return_error=True)
            log = '*** Reading remote log from {}.\n{}\n'.format(
                remote_loc, remote_log)
            return log, {'end_of_log': True}
        else:
            return super()._read(ti, try_number)

    def s3_log_exists(self, remote_log_location):
        """
        Check if remote_log_location exists in remote storage
        :param remote_log_location: log's location in remote storage
        :return: True if location exists else False
        """
        try:
            return self.hook.get_key(remote_log_location) is not None
        except Exception:
            pass
        return False

    def s3_read(self, remote_log_location, return_error=False):
        """
        Returns the log found at the remote_log_location. Returns '' if no
        logs are found or there is an error.
        :param remote_log_location: the log's location in remote storage
        :type remote_log_location: str (path)
        :param return_error: if True, returns a string error message if an
            error occurs. Otherwise returns '' when an error occurs.
        :type return_error: bool
        """
        try:
            return self.hook.read_key(remote_log_location)
        except Exception:
            msg = 'Could not read logs from {}'.format(remote_log_location)
            self.log.exception(msg)
            # return error if needed
            if return_error:
                return msg

    def s3_write(self, log, remote_log_location, append=True):
        """
        Writes the log to the remote_log_location. Fails silently if no hook
        was created.
        :param log: the log to write to the remote_log_location
        :type log: str
        :param remote_log_location: the log's location in remote storage
        :type remote_log_location: str (path)
        :param append: if False, any existing log file is overwritten. If True,
            the new log is appended to any existing logs.
        :type append: bool
        """
        if append and self.s3_log_exists(remote_log_location):
            old_log = self.s3_read(remote_log_location)
            log = '\n'.join([old_log, log]) if old_log else log

        try:
            self.hook.load_string(
                log,
                key=remote_log_location,
                replace=True,
                encrypt=configuration.conf.getboolean('core', 'ENCRYPT_S3_LOGS'),
            )
        except Exception:
            self.log.exception('Could not write logs to %s', remote_log_location)


# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import os
import six

from airflow.configuration import conf
from airflow.contrib.kubernetes.pod import Pod, Resources
from airflow.contrib.kubernetes.secret import Secret
from airflow.utils.log.logging_mixin import LoggingMixin

class WorkerConfiguration(LoggingMixin):
    """Contains Kubernetes Airflow Worker configuration logic"""

    dags_volume_name = 'airflow-dags'
    logs_volume_name = 'airflow-logs'
    git_sync_ssh_secret_volume_name = 'git-sync-ssh-key'
    git_ssh_key_secret_key = 'gitSshKey'
    git_sync_ssh_known_hosts_volume_name = 'git-sync-known-hosts'
    git_ssh_known_hosts_configmap_key = 'known_hosts'

    def __init__(self, kube_config):
        self.kube_config = kube_config
        self.worker_airflow_home = self.kube_config.airflow_home
        self.worker_airflow_dags = self.kube_config.dags_folder
        self.worker_airflow_logs = self.kube_config.base_log_folder

        super().__init__()

    def _get_init_containers(self):
        """When using git to retrieve the DAGs, use the GitSync Init Container"""
        # If we're using volume claims to mount the dags, no init container is needed
        if self.kube_config.dags_volume_claim or \
           self.kube_config.dags_volume_host or self.kube_config.dags_in_image:
            return []

        # Otherwise, define a git-sync init container
        init_environment = [{
            'name': 'GIT_SYNC_REPO',
            'value': self.kube_config.git_repo
        }, {
            'name': 'GIT_SYNC_BRANCH',
            'value': self.kube_config.git_branch
        }, {
            'name': 'GIT_SYNC_ROOT',
            'value': self.kube_config.git_sync_root
        }, {
            'name': 'GIT_SYNC_DEST',
            'value': self.kube_config.git_sync_dest
        }, {
            'name': 'GIT_SYNC_DEPTH',
            'value': '1'
        }, {
            'name': 'GIT_SYNC_ONE_TIME',
            'value': 'true'
        }]
        if self.kube_config.git_user:
            init_environment.append({
                'name': 'GIT_SYNC_USERNAME',
                'value': self.kube_config.git_user
            })
        if self.kube_config.git_password:
            init_environment.append({
                'name': 'GIT_SYNC_PASSWORD',
                'value': self.kube_config.git_password
            })

        volume_mounts = [{
            'mountPath': self.kube_config.git_sync_root,
            'name': self.dags_volume_name,
            'readOnly': False
        }]
        if self.kube_config.git_ssh_key_secret_name:
            volume_mounts.append({
                'name': self.git_sync_ssh_secret_volume_name,
                'mountPath': '/etc/git-secret/ssh',
                'subPath': 'ssh'
            })
            init_environment.extend([
                {
                    'name': 'GIT_SSH_KEY_FILE',
                    'value': '/etc/git-secret/ssh'
                },
                {
                    'name': 'GIT_SYNC_SSH',
                    'value': 'true'
                }])
        if self.kube_config.git_ssh_known_hosts_configmap_name:
            volume_mounts.append({
                'name': self.git_sync_ssh_known_hosts_volume_name,
                'mountPath': '/etc/git-secret/known_hosts',
                'subPath': 'known_hosts'
            })
            init_environment.extend([
                {
                    'name': 'GIT_KNOWN_HOSTS',
                    'value': 'true'
                },
                {
                    'name': 'GIT_SSH_KNOWN_HOSTS_FILE',
                    'value': '/etc/git-secret/known_hosts'
                }
            ])
        else:
            init_environment.append({
                'name': 'GIT_KNOWN_HOSTS',
                'value': 'false'
            })

        return [{
            'name': self.kube_config.git_sync_init_container_name,
            'image': self.kube_config.git_sync_container,
            'securityContext': {'runAsUser': 65533},  # git-sync user
            'env': init_environment,
            'volumeMounts': volume_mounts
        }]

    def _get_environment(self):
        """Defines any necessary environment variables for the pod executor"""
        env = {}

        for env_var_name, env_var_val in six.iteritems(self.kube_config.kube_env_vars):
            env[env_var_name] = env_var_val

        env["AIRFLOW__CORE__EXECUTOR"] = "LocalExecutor"

        if self.kube_config.airflow_configmap:
            env['AIRFLOW_HOME'] = self.worker_airflow_home
            env['AIRFLOW__CORE__DAGS_FOLDER'] = self.worker_airflow_dags
        if (not self.kube_config.airflow_configmap and
                'AIRFLOW__CORE__SQL_ALCHEMY_CONN' not in self.kube_config.kube_secrets):
            env['AIRFLOW__CORE__SQL_ALCHEMY_CONN'] = conf.get("core", "SQL_ALCHEMY_CONN")
        if self.kube_config.git_dags_folder_mount_point:
            # /root/airflow/dags/repo/dags
            dag_volume_mount_path = os.path.join(
                self.kube_config.git_dags_folder_mount_point,
                self.kube_config.git_sync_dest,  # repo
                self.kube_config.git_subpath     # dags
            )
            env['AIRFLOW__CORE__DAGS_FOLDER'] = dag_volume_mount_path
        return env

    def _get_configmaps(self):
        """Extracts any configmapRefs to envFrom"""
        if not self.kube_config.env_from_configmap_ref:
            return []
        return self.kube_config.env_from_configmap_ref.split(',')

    def _get_secrets(self):
        """Defines any necessary secrets for the pod executor"""
        worker_secrets = []

        for env_var_name, obj_key_pair in six.iteritems(self.kube_config.kube_secrets):
            k8s_secret_obj, k8s_secret_key = obj_key_pair.split('=')
            worker_secrets.append(
                Secret('env', env_var_name, k8s_secret_obj, k8s_secret_key)
            )

        if self.kube_config.env_from_secret_ref:
            for secret_ref in self.kube_config.env_from_secret_ref.split(','):
                worker_secrets.append(
                    Secret('env', None, secret_ref)
                )

        return worker_secrets

    def _get_image_pull_secrets(self):
        """Extracts any image pull secrets for fetching container(s)"""
        if not self.kube_config.image_pull_secrets:
            return []
        return self.kube_config.image_pull_secrets.split(',')

    def _get_security_context(self):
        """Defines the security context"""
        security_context = {}

        if self.kube_config.worker_run_as_user:
            security_context['runAsUser'] = self.kube_config.worker_run_as_user

        if self.kube_config.worker_fs_group:
            security_context['fsGroup'] = self.kube_config.worker_fs_group

        # set fs_group to 65533 if not explicitly specified and using git ssh keypair auth
        if self.kube_config.git_ssh_key_secret_name and security_context.get('fsGroup') is None:
            security_context['fsGroup'] = 65533

        return security_context

    def _get_labels(self, labels):
        copy = self.kube_config.kube_labels.copy()
        copy.update(labels)
        return copy

    def _get_volumes_and_mounts(self):
        def _construct_volume(name, claim, host):
            volume = {
                'name': name
            }
            if claim:
                volume['persistentVolumeClaim'] = {
                    'claimName': claim
                }
            elif host:
                volume['hostPath'] = {
                    'path': host,
                    'type': ''
                }
            else:
                volume['emptyDir'] = {}
            return volume

        volumes = {
            self.dags_volume_name: _construct_volume(
                self.dags_volume_name,
                self.kube_config.dags_volume_claim,
                self.kube_config.dags_volume_host
            ),
            self.logs_volume_name: _construct_volume(
                self.logs_volume_name,
                self.kube_config.logs_volume_claim,
                self.kube_config.logs_volume_host
            )
        }

        volume_mounts = {
            self.dags_volume_name: {
                'name': self.dags_volume_name,
                'mountPath': self.generate_dag_volume_mount_path(),
                'readOnly': True,
            },
            self.logs_volume_name: {
                'name': self.logs_volume_name,
                'mountPath': self.worker_airflow_logs,
            }
        }

        if self.kube_config.dags_volume_subpath:
            volume_mounts[self.dags_volume_name]['subPath'] = self.kube_config.dags_volume_subpath

        if self.kube_config.logs_volume_subpath:
            volume_mounts[self.logs_volume_name]['subPath'] = self.kube_config.logs_volume_subpath

        if self.kube_config.dags_in_image:
            del volumes[self.dags_volume_name]
            del volume_mounts[self.dags_volume_name]

        # Get the SSH key from secrets as a volume
        if self.kube_config.git_ssh_key_secret_name:
            volumes[self.git_sync_ssh_secret_volume_name] = {
                'name': self.git_sync_ssh_secret_volume_name,
                'secret': {
                    'secretName': self.kube_config.git_ssh_key_secret_name,
                    'items': [{
                        'key': self.git_ssh_key_secret_key,
                        'path': 'ssh',
                        'mode': 0o440
                    }]
                }
            }

        if self.kube_config.git_ssh_known_hosts_configmap_name:
            volumes[self.git_sync_ssh_known_hosts_volume_name] = {
                'name': self.git_sync_ssh_known_hosts_volume_name,
                'configMap': {
                    'name': self.kube_config.git_ssh_known_hosts_configmap_name
                },
                'mode': 0o440
            }

        # Mount the airflow.cfg file via a configmap the user has specified
        if self.kube_config.airflow_configmap:
            config_volume_name = 'airflow-config'
            config_path = '{}/airflow.cfg'.format(self.worker_airflow_home)
            volumes[config_volume_name] = {
                'name': config_volume_name,
                'configMap': {
                    'name': self.kube_config.airflow_configmap
                }
            }
            volume_mounts[config_volume_name] = {
                'name': config_volume_name,
                'mountPath': config_path,
                'subPath': 'airflow.cfg',
                'readOnly': True
            }

        return volumes, volume_mounts

    def generate_dag_volume_mount_path(self):
        if self.kube_config.dags_volume_claim or self.kube_config.dags_volume_host:
            dag_volume_mount_path = self.worker_airflow_dags
        else:
            dag_volume_mount_path = self.kube_config.git_dags_folder_mount_point

        return dag_volume_mount_path

    def make_pod(self, namespace, worker_uuid, pod_id, dag_id, task_id, execution_date,
                 try_number, airflow_command, kube_executor_config):
        volumes_dict, volume_mounts_dict = self._get_volumes_and_mounts()
        worker_init_container_spec = self._get_init_containers()
        resources = Resources(
            request_memory=kube_executor_config.request_memory,
            request_cpu=kube_executor_config.request_cpu,
            limit_memory=kube_executor_config.limit_memory,
            limit_cpu=kube_executor_config.limit_cpu
        )
        gcp_sa_key = kube_executor_config.gcp_service_account_key
        annotations = dict(kube_executor_config.annotations) or self.kube_config.kube_annotations
        if gcp_sa_key:
            annotations['iam.cloud.google.com/service-account'] = gcp_sa_key

        volumes = [value for value in volumes_dict.values()] + kube_executor_config.volumes
        volume_mounts = [value for value in volume_mounts_dict.values()] + kube_executor_config.volume_mounts

        affinity = kube_executor_config.affinity or self.kube_config.kube_affinity
        tolerations = kube_executor_config.tolerations or self.kube_config.kube_tolerations

        return Pod(
            namespace=namespace,
            name=pod_id,
            image=kube_executor_config.image or self.kube_config.kube_image,
            image_pull_policy=(kube_executor_config.image_pull_policy or
                               self.kube_config.kube_image_pull_policy),
            cmds=airflow_command,
            labels=self._get_labels({
                'airflow-worker': worker_uuid,
                'dag_id': dag_id,
                'task_id': task_id,
                'execution_date': execution_date,
                'try_number': str(try_number),
            }),
            envs=self._get_environment(),
            secrets=self._get_secrets(),
            service_account_name=self.kube_config.worker_service_account_name,
            image_pull_secrets=self.kube_config.image_pull_secrets,
            init_containers=worker_init_container_spec,
            volumes=volumes,
            volume_mounts=volume_mounts,
            resources=resources,
            annotations=annotations,
            node_selectors=(kube_executor_config.node_selectors or
                            self.kube_config.kube_node_selectors),
            affinity=affinity,
            tolerations=tolerations,
            security_context=self._get_security_context(),
            configmaps=self._get_configmaps()
        )


# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#

import os
import time
import datetime
import six
import re

from airflow.exceptions import AirflowException
from airflow.hooks.base_hook import BaseHook
from airflow import configuration
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.state import State
from airflow.models import TaskInstance

from qds_sdk.qubole import Qubole
from qds_sdk.commands import Command, HiveCommand, PrestoCommand, HadoopCommand, \
    PigCommand, ShellCommand, SparkCommand, DbTapQueryCommand, DbExportCommand, \
    DbImportCommand, SqlCommand

COMMAND_CLASSES = {
    "hivecmd": HiveCommand,
    "prestocmd": PrestoCommand,
    "hadoopcmd": HadoopCommand,
    "shellcmd": ShellCommand,
    "pigcmd": PigCommand,
    "sparkcmd": SparkCommand,
    "dbtapquerycmd": DbTapQueryCommand,
    "dbexportcmd": DbExportCommand,
    "dbimportcmd": DbImportCommand,
    "sqlcmd": SqlCommand
}

POSITIONAL_ARGS = {
    'hadoopcmd': ['sub_command'],
    'shellcmd': ['parameters'],
    'pigcmd': ['parameters']
}

def flatten_list(list_of_lists):
    return [element for array in list_of_lists for element in array]

def filter_options(options):
    options_to_remove = ["help", "print-logs-live", "print-logs"]
    return [option for option in options if option not in options_to_remove]

def get_options_list(command_class):
    options_list = [option.get_opt_string().strip("--") for option in command_class.optparser.option_list]
    return filter_options(options_list)

def build_command_args():
    command_args, hyphen_args = {}, set()
    for cmd in COMMAND_CLASSES:

        # get all available options from the class
        opts_list = get_options_list(COMMAND_CLASSES[cmd])

        # append positional args if any for the command
        if cmd in POSITIONAL_ARGS:
            opts_list += POSITIONAL_ARGS[cmd]

        # get args with a hyphen and replace them with underscore
        for index, opt in enumerate(opts_list):
            if "-" in opt:
                opts_list[index] = opt.replace("-", "_")
                hyphen_args.add(opts_list[index])

        command_args[cmd] = opts_list
    return command_args, list(hyphen_args)

COMMAND_ARGS, HYPHEN_ARGS = build_command_args()

class QuboleHook(BaseHook):
    def __init__(self, *args, **kwargs):
        conn = self.get_connection(kwargs['qubole_conn_id'])
        Qubole.configure(api_token=conn.password, api_url=conn.host)
        self.task_id = kwargs['task_id']
        self.dag_id = kwargs['dag'].dag_id
        self.kwargs = kwargs
        self.cls = COMMAND_CLASSES[self.kwargs['command_type']]
        self.cmd = None
        self.task_instance = None

    @staticmethod
    def handle_failure_retry(context):
        ti = context['ti']
        cmd_id = ti.xcom_pull(key='qbol_cmd_id', task_ids=ti.task_id)

        if cmd_id is not None:
            cmd = Command.find(cmd_id)
            if cmd is not None:
                log = LoggingMixin().log
                if cmd.status == 'done':
                    log.info('Command ID: %s has been succeeded, hence marking this '
                             'TI as Success.', cmd_id)
                    ti.state = State.SUCCESS
                elif cmd.status == 'running':
                    log.info('Cancelling the Qubole Command Id: %s', cmd_id)
                    cmd.cancel()

    def execute(self, context):
        args = self.cls.parse(self.create_cmd_args(context))
        self.cmd = self.cls.create(**args)
        self.task_instance = context['task_instance']
        context['task_instance'].xcom_push(key='qbol_cmd_id', value=self.cmd.id)
        self.log.info(
            "Qubole command created with Id: %s and Status: %s",
            self.cmd.id, self.cmd.status
        )

        while not Command.is_done(self.cmd.status):
            time.sleep(Qubole.poll_interval)
            self.cmd = self.cls.find(self.cmd.id)
            self.log.info("Command Id: %s and Status: %s", self.cmd.id, self.cmd.status)

        if 'fetch_logs' in self.kwargs and self.kwargs['fetch_logs'] is True:
            self.log.info("Logs for Command Id: %s \n%s", self.cmd.id, self.cmd.get_log())

        if self.cmd.status != 'done':
            raise AirflowException('Command Id: {0} failed with Status: {1}'.format(
                                   self.cmd.id, self.cmd.status))

    def kill(self, ti):
        """
        Kill (cancel) a Qubole command
        :param ti: Task Instance of the dag, used to determine the Quboles command id
        :return: response from Qubole
        """
        if self.cmd is None:
            if not ti and not self.task_instance:
                raise Exception("Unable to cancel Qubole Command, context is unavailable!")
            elif not ti:
                ti = self.task_instance
            cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=ti.task_id)
            self.cmd = self.cls.find(cmd_id)
        if self.cls and self.cmd:
            self.log.info('Sending KILL signal to Qubole Command Id: %s', self.cmd.id)
            self.cmd.cancel()

    def get_results(self, ti=None, fp=None, inline=True, delim=None, fetch=True):
        """
        Get results (or just s3 locations) of a command from Qubole and save into a file
        :param ti: Task Instance of the dag, used to determine the Quboles command id
        :param fp: Optional file pointer, will create one and return if None passed
        :param inline: True to download actual results, False to get s3 locations only
        :param delim: Replaces the CTL-A chars with the given delim, defaults to ','
        :param fetch: when inline is True, get results directly from s3 (if large)
        :return: file location containing actual results or s3 locations of results
        """
        if fp is None:
            iso = datetime.datetime.utcnow().isoformat()
            logpath = os.path.expanduser(
                configuration.conf.get('core', 'BASE_LOG_FOLDER')
            )
            resultpath = logpath + '/' + self.dag_id + '/' + self.task_id + '/results'
            configuration.mkdir_p(resultpath)
            fp = open(resultpath + '/' + iso, 'wb')

        if self.cmd is None:
            cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=self.task_id)
            self.cmd = self.cls.find(cmd_id)

        self.cmd.get_results(fp, inline, delim, fetch)
        fp.flush()
        fp.close()
        return fp.name

    def get_log(self, ti):
        """
        Get Logs of a command from Qubole
        :param ti: Task Instance of the dag, used to determine the Quboles command id
        :return: command log as text
        """
        if self.cmd is None:
            cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=self.task_id)
        Command.get_log_id(self.cls, cmd_id)

    def get_jobs_id(self, ti):
        """
        Get jobs associated with a Qubole commands
        :param ti: Task Instance of the dag, used to determine the Quboles command id
        :return: Job information associated with command
        """
        if self.cmd is None:
            cmd_id = ti.xcom_pull(key="qbol_cmd_id", task_ids=self.task_id)
        Command.get_jobs_id(self.cls, cmd_id)

    def get_extra_links(self, operator, dttm):
        """
        Get link to qubole command result page.

        :param operator: operator
        :param dttm: datetime
        :return: url link
        """
        conn = BaseHook.get_connection(operator.kwargs['qubole_conn_id'])
        if conn and conn.host:
            host = re.sub(r'api$', 'v2/analyze?command_id=', conn.host)
        else:
            host = 'https://api.qubole.com/v2/analyze?command_id='

        ti = TaskInstance(task=operator, execution_date=dttm)
        qds_command_id = ti.xcom_pull(task_ids=operator.task_id, key='qbol_cmd_id')
        url = host + str(qds_command_id) if qds_command_id else ''
        return url

    def create_cmd_args(self, context):
        args = []
        cmd_type = self.kwargs['command_type']
        inplace_args = None
        tags = {self.dag_id, self.task_id, context['run_id']}
        positional_args_list = flatten_list(POSITIONAL_ARGS.values())

        for k, v in self.kwargs.items():
            if k in COMMAND_ARGS[cmd_type]:
                if k in HYPHEN_ARGS:
                    args.append("--{0}={1}".format(k.replace('_', '-'), v))
                elif k in positional_args_list:
                    inplace_args = v
                elif k == 'tags':
                    if isinstance(v, six.string_types):
                        tags.add(v)
                    elif isinstance(v, (list, tuple)):
                        for val in v:
                            tags.add(val)
                else:
                    args.append("--{0}={1}".format(k, v))

            if k == 'notify' and v is True:
                args.append("--notify")

        args.append("--tags={0}".format(','.join(filter(None, tags))))

        if inplace_args is not None:
            args += inplace_args.split(' ')

        return args