src/tools/kickstarter.py - cord-workflow-airflow - Gitiles

 #!/usr/bin/env python3

 # Copyright 2019-present Open Networking Foundation
 #
 # Licensed 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.

 """
 Workflow Kickstarter

 This module kickstarts Airflow workflows for requests from Workflow Controller
 """

 import json
 import os
 import os.path
 import argparse
 import pyfiglet
 import traceback
 import socket
 import time
 import subprocess

 from multistructlog import create_logger
 from cord_workflow_controller_client.manager import Manager

 # We can't use experimental APIs for managing workflows/workflow runs of Airflow
 # - REST API does not provide sufficient features at this version
 # - API_Client does not work if a caller is not in main thread

 # from importlib import import_module
 # from airflow import configuration as AirflowConf
 # from airflow import api
 # from airflow.models import DagRun

 try:
     from urllib.parse import urlparse
 except ImportError:
     from urlparse import urlparse

 log = create_logger()
 manager = None
 # airflow_client = None

 airflow_bin = os.getenv('AIRFLOW_BIN', '/usr/local/bin')

 progargs = {
     'controller_url': 'http://localhost:3030',
     'airflow_bin': airflow_bin,
     'logging': None
 }

 DEFAULT_CONFIG_FILE_PATH = '/etc/cord_workflow_airflow_extensions/config.json'
 SOCKET_CONNECTION_TEST_TIMEOUT = 5
 DEFAULT_CONNECTION_TEST_DELAY = 5
 DEFAULT_CONNECTION_TEST_RETRY = 999999


 def print_graffiti():
     result = pyfiglet.figlet_format("CORD\nWorkflow\nKickstarter", font="graffiti")
     print(result)


 def get_arg_parser():
     parser = argparse.ArgumentParser(description='CORD Workflow Kickstarter Daemon.', prog='kickstarter')
     parser.add_argument('--config', help='locate a configuration file')
     parser.add_argument('--controller', help='CORD Workflow Controller URL')
     return parser


 def read_config(path):
     if os.path.exists(path):
         with open(path) as json_config_file:
             data = json.load(json_config_file)
             return data
     return {}


 def pretty_format_json(j):
     dumps = json.dumps(j, sort_keys=True, indent=4, separators=(',', ': '))
     return dumps


 def is_port_open(url, timeout):
     o = urlparse(url)
     hostname = o.hostname
     port = o.port

     if (not port) or port <= 0:
         if o.scheme.lower() == 'http':
             port = 80
         elif o.scheme.lower() == 'https':
             port = 443

     succeed = False
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     sock.settimeout(timeout)
     try:
         sock.connect((hostname, int(port)))
         sock.shutdown(socket.SHUT_RDWR)
         succeed = True
     except BaseException:
         pass
     finally:
         sock.close()

     return succeed


 def check_web_live(url,
                    retry=DEFAULT_CONNECTION_TEST_RETRY,
                    delay=DEFAULT_CONNECTION_TEST_DELAY,
                    timeout=SOCKET_CONNECTION_TEST_TIMEOUT):
     ipup = False
     for _ in range(retry):
         if is_port_open(url, timeout):
             ipup = True
             break
         else:
             time.sleep(delay)
     return ipup


 def get_airflow_cli():
     return os.path.join(progargs['airflow_bin'], 'airflow')


 def check_airflow_live():
     try:
         subprocess.check_call([get_airflow_cli(), 'list_dags'])
         return True
     except subprocess.CalledProcessError as e:
         log.error(e)
         return False


 def on_kickstart(workflow_id, workflow_run_id):
     # if manager and airflow_client:
     if manager:
         try:
             created = False
             log.info('> Kickstarting a workflow (%s) => workflow run (%s)' % (workflow_id, workflow_run_id))
             # message = airflow_client.trigger_dag(
             #     dag_id=workflow_id,
             #     run_id=workflow_run_id
             # )
             # log.info('> Airflow Response: %s' % message)

             output = subprocess.Popen(
                 [get_airflow_cli(), 'trigger_dag', '-r', workflow_run_id, workflow_id],
                 stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                 encoding='utf8'
             )

             for line in output.stdout.readlines():
                 if 'Created <DagRun ' in line:
                     created = True
                     break

             if created:
                 # let controller know that the new workflow run is created
                 log.info('> Notifying a new workflow run (%s) for a workflow (%s)' % (workflow_run_id, workflow_id))
                 manager.report_new_workflow_run(workflow_id, workflow_run_id)
         except subprocess.CalledProcessError as e:
             # when shell exited with non-zero code
             log.error('> Error : %s' % e)
         except Exception as e:
             log.error('> Error : %s' % e)
             log.debug(traceback.format_exc())


 def on_check_status(workflow_id, workflow_run_id):
     # if manager and airflow_client:
     if manager:
         try:
             status = 'unknown'
             log.info('> Checking status of workflow run (%s)' % (workflow_run_id))

             # run = DagRun.find(dag_id=workflow_id, run_id=workflow_run_id)
             # status = 'unknown'
             # if run:
             #     # run is an array
             #     # this should be one of ['success', 'running', 'failed']
             #     status = run[0].state
             # else:
             #     log.error(
             #         'Cannot retrieve status of a workflow run (%s, %s)' %
             #         (workflow_id, workflow_run_id)
             #     )
             #     status = 'unknown'

             output = subprocess.Popen(
                 [get_airflow_cli(), 'list_dag_runs', workflow_id],
                 stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                 encoding='utf8'
             )

             title = False
             body = False
             for line in output.stdout.readlines():
                 if 'DAG RUNS' in line:
                     title = True
                 elif title and ('--------' in line):
                     body = True
                 elif body:
                     # id  | run_id | state | execution_date | state_date |
                     if workflow_run_id in line:
                         # found the line
                         # 1 | ssss | running | 2019-07-25T21:35:06+00:00 |
                         # 2019-07-25T21:35:06.242130+00:00 |
                         fields = line.split('|')
                         status = fields[2].strip()
                         break

             log.info('> status : %s' % status)

             # let controller know the status of the workflow run
             log.info(
                 '> Updating status of a workflow run (%s) - status : %s' %
                 (workflow_run_id, status)
             )
             manager.report_workflow_run_status(workflow_id, workflow_run_id, status)
         except subprocess.CalledProcessError as e:
             # when shell exited with non-zero code
             log.error('> Error : %s' % e)
         except Exception as e:
             log.error('> Error : %s' % e)
             log.debug(traceback.format_exc())


 def on_check_status_bulk(requests):
     # if manager and airflow_client:
     if requests:
         req = {}
         for req in requests:
             workflow_id = req['workflow_id']
             workflow_run_id = req['workflow_run_id']

             if workflow_id not in req:
                 req[workflow_id] = []

             req[workflow_id].append(workflow_run_id)

         if manager:
             try:
                 log.info('> Checking status of workflow runs')

                 statuses = []
                 for workflow_id in req:
                     workflow_run_ids = req[workflow_id]

                     output = subprocess.Popen(
                         [get_airflow_cli(), 'list_dag_runs', workflow_id],
                         stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                         encoding='utf8'
                     )

                     title = False
                     body = False
                     for line in output.stdout.readlines():
                         if 'DAG RUNS' in line:
                             title = True
                         elif title and ('--------' in line):
                             body = True
                         elif body:
                             # id  | run_id | state | execution_date | state_date |
                             for workflow_run_id in workflow_run_ids:
                                 if workflow_run_id in line:
                                     # found the line
                                     # 1 | ssss | running | 2019-07-25T21:35:06+00:00 |
                                     # 2019-07-25T21:35:06.242130+00:00 |
                                     fields = line.split('|')
                                     status = fields[2].strip()

                                     log.info('> status of a workflow run (%s) : %s' % (workflow_run_id, status))
                                     statuses.append({
                                         'workflow_id': workflow_id,
                                         'workflow_run_id': workflow_run_id,
                                         'status': status
                                     })

                 # let controller know statuses of workflow runs
                 log.info('> Updating status of workflow runs')
                 manager.report_workflow_run_status_bulk(statuses)
             except subprocess.CalledProcessError as e:
                 # when shell exited with non-zero code
                 log.error('> Error : %s' % e)
             except Exception as e:
                 log.error('> Error : %s' % e)
                 log.debug(traceback.format_exc())


 # for command-line execution
 def main(args):
     print_graffiti()

     # check if config path is set
     config_file_path = DEFAULT_CONFIG_FILE_PATH
     if args.config:
         config_file_path = args.config

     if os.path.exists(config_file_path):
         # read config
         config = read_config(config_file_path)
         if config:
             global progargs
             for k in progargs:
                 # overwrite
                 if k in config:
                     progargs[k] = config[k]

     global log
     log = create_logger(progargs["logging"])

     if args.controller:
         progargs['controller_url'] = args.controller

     print('=CONFIG=')
     config_json_string = pretty_format_json(progargs)
     print(config_json_string)
     print('\n')

     # checking controller and airflow web interface
     log.info('Checking if Workflow Controller (%s) is live...' % progargs['controller_url'])
     controller_live = check_web_live(progargs['controller_url'])
     if not controller_live:
         log.error('Controller (%s) appears to be down' % progargs['controller_url'])
         raise IOError('Controller (%s) appears to be down' % progargs['controller_url'])

     airflow_live = check_airflow_live()
     if not airflow_live:
         log.error('Airflow appears to be down')
         raise IOError('Airflow appears to be down')

     # connect to workflow controller
     log.info('Connecting to Workflow Controller (%s)...' % progargs['controller_url'])
     global manager
     manager = Manager(logger=log)
     manager.connect(progargs['controller_url'])
     manager.set_handlers({
         'kickstart': on_kickstart,
         'check_status': on_check_status,
         'check_status_bulk': on_check_status_bulk
     })

     # connect to airflow
     # global airflow_client
     # log.info('Connecting to Airflow...')

     # api.load_auth()
     # api_module = import_module(AirflowConf.get('cli', 'api_client'))
     # airflow_client = api_module.Client(
     #     api_base_url=AirflowConf.get('cli', 'endpoint_url'),
     #     auth=api.api_auth.client_auth
     # )

     log.info('Waiting for kickstart events from Workflow Controller...')
     try:
         manager.wait()
     finally:
         log.info('Terminating the program...')
         manager.disconnect()


 if __name__ == "__main__":
     parser = get_arg_parser()
     args = parser.parse_args()
     main(args)
	#!/usr/bin/env python3

	# Copyright 2019-present Open Networking Foundation
	#
	# Licensed 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.

	"""
	Workflow Kickstarter

	This module kickstarts Airflow workflows for requests from Workflow Controller
	"""

	import json
	import os
	import os.path
	import argparse
	import pyfiglet
	import traceback
	import socket
	import time
	import subprocess

	from multistructlog import create_logger
	from cord_workflow_controller_client.manager import Manager

	# We can't use experimental APIs for managing workflows/workflow runs of Airflow
	# - REST API does not provide sufficient features at this version
	# - API_Client does not work if a caller is not in main thread

	# from importlib import import_module
	# from airflow import configuration as AirflowConf
	# from airflow import api
	# from airflow.models import DagRun

	try:
	from urllib.parse import urlparse
	except ImportError:
	from urlparse import urlparse

	log = create_logger()
	manager = None
	# airflow_client = None

	airflow_bin = os.getenv('AIRFLOW_BIN', '/usr/local/bin')

	progargs = {
	'controller_url': 'http://localhost:3030',
	'airflow_bin': airflow_bin,
	'logging': None
	}

	DEFAULT_CONFIG_FILE_PATH = '/etc/cord_workflow_airflow_extensions/config.json'
	SOCKET_CONNECTION_TEST_TIMEOUT = 5
	DEFAULT_CONNECTION_TEST_DELAY = 5
	DEFAULT_CONNECTION_TEST_RETRY = 999999


	def print_graffiti():
	result = pyfiglet.figlet_format("CORD\nWorkflow\nKickstarter", font="graffiti")
	print(result)


	def get_arg_parser():
	parser = argparse.ArgumentParser(description='CORD Workflow Kickstarter Daemon.', prog='kickstarter')
	parser.add_argument('--config', help='locate a configuration file')
	parser.add_argument('--controller', help='CORD Workflow Controller URL')
	return parser


	def read_config(path):
	if os.path.exists(path):
	with open(path) as json_config_file:
	data = json.load(json_config_file)
	return data
	return {}


	def pretty_format_json(j):
	dumps = json.dumps(j, sort_keys=True, indent=4, separators=(',', ': '))
	return dumps


	def is_port_open(url, timeout):
	o = urlparse(url)
	hostname = o.hostname
	port = o.port

	if (not port) or port <= 0:
	if o.scheme.lower() == 'http':
	port = 80
	elif o.scheme.lower() == 'https':
	port = 443

	succeed = False
	sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	sock.settimeout(timeout)
	try:
	sock.connect((hostname, int(port)))
	sock.shutdown(socket.SHUT_RDWR)
	succeed = True
	except BaseException:
	pass
	finally:
	sock.close()

	return succeed


	def check_web_live(url,
	retry=DEFAULT_CONNECTION_TEST_RETRY,
	delay=DEFAULT_CONNECTION_TEST_DELAY,
	timeout=SOCKET_CONNECTION_TEST_TIMEOUT):
	ipup = False
	for _ in range(retry):
	if is_port_open(url, timeout):
	ipup = True
	break
	else:
	time.sleep(delay)
	return ipup


	def get_airflow_cli():
	return os.path.join(progargs['airflow_bin'], 'airflow')


	def check_airflow_live():
	try:
	subprocess.check_call([get_airflow_cli(), 'list_dags'])
	return True
	except subprocess.CalledProcessError as e:
	log.error(e)
	return False


	def on_kickstart(workflow_id, workflow_run_id):
	# if manager and airflow_client:
	if manager:
	try:
	created = False
	log.info('> Kickstarting a workflow (%s) => workflow run (%s)' % (workflow_id, workflow_run_id))
	# message = airflow_client.trigger_dag(
	# dag_id=workflow_id,
	# run_id=workflow_run_id
	# )
	# log.info('> Airflow Response: %s' % message)

	output = subprocess.Popen(
	[get_airflow_cli(), 'trigger_dag', '-r', workflow_run_id, workflow_id],
	stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
	encoding='utf8'
	)

	for line in output.stdout.readlines():
	if 'Created <DagRun ' in line:
	created = True
	break

	if created:
	# let controller know that the new workflow run is created
	log.info('> Notifying a new workflow run (%s) for a workflow (%s)' % (workflow_run_id, workflow_id))
	manager.report_new_workflow_run(workflow_id, workflow_run_id)
	except subprocess.CalledProcessError as e:
	# when shell exited with non-zero code
	log.error('> Error : %s' % e)
	except Exception as e:
	log.error('> Error : %s' % e)
	log.debug(traceback.format_exc())


	def on_check_status(workflow_id, workflow_run_id):
	# if manager and airflow_client:
	if manager:
	try:
	status = 'unknown'
	log.info('> Checking status of workflow run (%s)' % (workflow_run_id))

	# run = DagRun.find(dag_id=workflow_id, run_id=workflow_run_id)
	# status = 'unknown'
	# if run:
	# # run is an array
	# # this should be one of ['success', 'running', 'failed']
	# status = run[0].state
	# else:
	# log.error(
	# 'Cannot retrieve status of a workflow run (%s, %s)' %
	# (workflow_id, workflow_run_id)
	# )
	# status = 'unknown'

	output = subprocess.Popen(
	[get_airflow_cli(), 'list_dag_runs', workflow_id],
	stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
	encoding='utf8'
	)

	title = False
	body = False
	for line in output.stdout.readlines():
	if 'DAG RUNS' in line:
	title = True
	elif title and ('--------' in line):
	body = True
	elif body:
	# id \| run_id \| state \| execution_date \| state_date \|
	if workflow_run_id in line:
	# found the line
	# 1 \| ssss \| running \| 2019-07-25T21:35:06+00:00 \|
	# 2019-07-25T21:35:06.242130+00:00 \|
	fields = line.split('\|')
	status = fields[2].strip()
	break

	log.info('> status : %s' % status)

	# let controller know the status of the workflow run
	log.info(
	'> Updating status of a workflow run (%s) - status : %s' %
	(workflow_run_id, status)
	)
	manager.report_workflow_run_status(workflow_id, workflow_run_id, status)
	except subprocess.CalledProcessError as e:
	# when shell exited with non-zero code
	log.error('> Error : %s' % e)
	except Exception as e:
	log.error('> Error : %s' % e)
	log.debug(traceback.format_exc())


	def on_check_status_bulk(requests):
	# if manager and airflow_client:
	if requests:
	req = {}
	for req in requests:
	workflow_id = req['workflow_id']
	workflow_run_id = req['workflow_run_id']

	if workflow_id not in req:
	req[workflow_id] = []

	req[workflow_id].append(workflow_run_id)

	if manager:
	try:
	log.info('> Checking status of workflow runs')

	statuses = []
	for workflow_id in req:
	workflow_run_ids = req[workflow_id]

	output = subprocess.Popen(
	[get_airflow_cli(), 'list_dag_runs', workflow_id],
	stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
	encoding='utf8'
	)

	title = False
	body = False
	for line in output.stdout.readlines():
	if 'DAG RUNS' in line:
	title = True
	elif title and ('--------' in line):
	body = True
	elif body:
	# id \| run_id \| state \| execution_date \| state_date \|
	for workflow_run_id in workflow_run_ids:
	if workflow_run_id in line:
	# found the line
	# 1 \| ssss \| running \| 2019-07-25T21:35:06+00:00 \|
	# 2019-07-25T21:35:06.242130+00:00 \|
	fields = line.split('\|')
	status = fields[2].strip()

	log.info('> status of a workflow run (%s) : %s' % (workflow_run_id, status))
	statuses.append({
	'workflow_id': workflow_id,
	'workflow_run_id': workflow_run_id,
	'status': status
	})

	# let controller know statuses of workflow runs
	log.info('> Updating status of workflow runs')
	manager.report_workflow_run_status_bulk(statuses)
	except subprocess.CalledProcessError as e:
	# when shell exited with non-zero code
	log.error('> Error : %s' % e)
	except Exception as e:
	log.error('> Error : %s' % e)
	log.debug(traceback.format_exc())


	# for command-line execution
	def main(args):
	print_graffiti()

	# check if config path is set
	config_file_path = DEFAULT_CONFIG_FILE_PATH
	if args.config:
	config_file_path = args.config

	if os.path.exists(config_file_path):
	# read config
	config = read_config(config_file_path)
	if config:
	global progargs
	for k in progargs:
	# overwrite
	if k in config:
	progargs[k] = config[k]

	global log
	log = create_logger(progargs["logging"])

	if args.controller:
	progargs['controller_url'] = args.controller

	print('=CONFIG=')
	config_json_string = pretty_format_json(progargs)
	print(config_json_string)
	print('\n')

	# checking controller and airflow web interface
	log.info('Checking if Workflow Controller (%s) is live...' % progargs['controller_url'])
	controller_live = check_web_live(progargs['controller_url'])
	if not controller_live:
	log.error('Controller (%s) appears to be down' % progargs['controller_url'])
	raise IOError('Controller (%s) appears to be down' % progargs['controller_url'])

	airflow_live = check_airflow_live()
	if not airflow_live:
	log.error('Airflow appears to be down')
	raise IOError('Airflow appears to be down')

	# connect to workflow controller
	log.info('Connecting to Workflow Controller (%s)...' % progargs['controller_url'])
	global manager
	manager = Manager(logger=log)
	manager.connect(progargs['controller_url'])
	manager.set_handlers({
	'kickstart': on_kickstart,
	'check_status': on_check_status,
	'check_status_bulk': on_check_status_bulk
	})

	# connect to airflow
	# global airflow_client
	# log.info('Connecting to Airflow...')

	# api.load_auth()
	# api_module = import_module(AirflowConf.get('cli', 'api_client'))
	# airflow_client = api_module.Client(
	# api_base_url=AirflowConf.get('cli', 'endpoint_url'),
	# auth=api.api_auth.client_auth
	# )

	log.info('Waiting for kickstart events from Workflow Controller...')
	try:
	manager.wait()
	finally:
	log.info('Terminating the program...')
	manager.disconnect()


	if __name__ == "__main__":
	parser = get_arg_parser()
	args = parser.parse_args()
	main(args)