tasks.py

tasks.py

import os
import  time
import json

import slack
from xplenty import XplentyClient

ACCOUNT_ID = os.getenv('ACCOUNT_ID')  # Your account ID
API_KEY = os.getenv('API_KEY')        # Your account's API Key
SLACK_TOKEN = os.getenv('SLACK_TOKEN') # Slack Token for sending alerts


def create_client():
    return XplentyClient(ACCOUNT_ID, API_KEY)


# Creates a cluster, creates a job with package configuration and created cluster's configuration
def run_package(**kwargs):
    cluster_definition = kwargs['cluster']
    package_definition = kwargs['package']
    client = create_client()
    cluster = client.create_cluster(cluster_definition['cluster_type'], cluster_definition['nodes'],
                                    cluster_definition['name'], cluster_definition['description'],
                                    cluster_definition['terminate_on_idle'], cluster_definition['time_to_idle'])

    job = client.add_job(cluster.id, package_definition['package_id'], kwargs['variables'])
    job_id = job.id
    print('Job Added successfully')
    # Wait until the job completes
    while True:
        job = client.get_job(job_id)
        if job.status not in ['failed', 'completed']:
            time.sleep(60)
        else:
            break
    print('Job completed with status {}'.format(job.status))
    kwargs['ti'].xcom_push(key='xplenty_job_id', value=job.id)
    if job.status == 'failed':
        raise Exception('Job Failed')


def send_slack_alert(**kwargs):
    job_id = kwargs['ti'].xcom_pull(key=None, task_ids='run_package')
    job = create_client().get_job(job_id)
    client = slack.WebClient(token=SLACK_TOKEN)
    response = client.chat_postMessage(
        channel='#xplenty-alerts',
        text="Job {} Completed with Status {}".format(job.name, job.status))
    if not response.get('ok'):
        raise Exception('Failed to send Slack alert')