98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
#
|
|
# 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.
|
|
|
|
import logging
|
|
# Using nosec to prevent Bandit blacklist reporting. Subprocess is used
|
|
# in a controlled way as part of this operator.
|
|
import subprocess # nosec
|
|
|
|
from airflow.exceptions import AirflowException
|
|
from airflow.models import BaseOperator
|
|
from airflow.plugins_manager import AirflowPlugin
|
|
from airflow.utils.decorators import apply_defaults
|
|
|
|
|
|
class K8sHealthCheckOperator(BaseOperator):
|
|
"""
|
|
Performs basic Kubernetes Health Check
|
|
"""
|
|
@apply_defaults
|
|
def __init__(self,
|
|
*args, **kwargs):
|
|
|
|
super(K8sHealthCheckOperator, self).__init__(*args, **kwargs)
|
|
|
|
def execute(self, context):
|
|
logging.info("Running Basic Kubernetes Cluster Health Check:")
|
|
|
|
# Location of Token
|
|
token_path = '/var/run/secrets/kubernetes.io/serviceaccount/token'
|
|
|
|
# Read content of token file as string
|
|
with open(token_path, "r") as token_file:
|
|
token = token_file.read()
|
|
|
|
# Location of CA Certificate
|
|
ca = '/var/run/secrets/kubernetes.io/serviceaccount/ca.crt'
|
|
|
|
# The current health checks only check to ensure that all
|
|
# pods are in 'running' state. The k8s health checks can
|
|
# be expanded in future if need be.
|
|
k8s_command = [
|
|
'kubectl', '--token', token,
|
|
'--certificate-authority', ca,
|
|
'get', 'pods', '--all-namespaces']
|
|
|
|
# Execute the Kubernetes Health Check Command
|
|
k8s_output = subprocess.Popen( # nosec
|
|
k8s_command,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT)
|
|
|
|
# Logs Output
|
|
logging.info("Output:")
|
|
|
|
line = ''
|
|
for line in iter(k8s_output.stdout.readline, b''):
|
|
line = line.strip()
|
|
logging.info(line)
|
|
# Skip the first line of output
|
|
if 'STATUS' in str(line, 'utf-8'):
|
|
continue
|
|
# Do nothing if pod is in 'Running' state
|
|
elif str(line, 'utf-8').split()[3] == 'Running':
|
|
continue
|
|
# Raise Exceptions if Pod is in state other than
|
|
# 'Running'
|
|
else:
|
|
pod = str(line, 'utf-8').split()[0]
|
|
pod_state = str(line, 'utf-8').split()[3]
|
|
logging.error('Pod %s is in %s state', pod, pod_state)
|
|
raise AirflowException("Kubernetes Health Checks Failed!")
|
|
|
|
# Wait for child process to terminate
|
|
# Set and return returncode attribute.
|
|
k8s_output.wait()
|
|
logging.info("Command exited with "
|
|
"return code {0}".format(k8s_output.returncode))
|
|
|
|
# Raise Execptions if Kubernetes Command Fails
|
|
if k8s_output.returncode:
|
|
raise AirflowException("Kubernetes Command Failed")
|
|
|
|
|
|
class K8sHealthCheckPlugin(AirflowPlugin):
|
|
name = "k8s_healthcheck_plugin"
|
|
operators = [K8sHealthCheckOperator]
|