diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml index b9e889f..f2156c6 100644 --- a/.github/workflows/unit-tests.yaml +++ b/.github/workflows/unit-tests.yaml @@ -22,4 +22,4 @@ jobs: - name: Run unit tests run: | - python -m unittest openshift_metrics/tests/test_utils.py + python -m unittest openshift_metrics/tests/test_* diff --git a/openshift_metrics/openshift_prometheus_metrics.py b/openshift_metrics/openshift_prometheus_metrics.py index 8daaf70..6aec568 100755 --- a/openshift_metrics/openshift_prometheus_metrics.py +++ b/openshift_metrics/openshift_prometheus_metrics.py @@ -15,6 +15,7 @@ import argparse from datetime import datetime, timedelta +from prometheus_client import PrometheusClient import os import sys import json @@ -77,24 +78,25 @@ def main(): token = os.environ.get("OPENSHIFT_TOKEN") + prom_client = PrometheusClient(openshift_url, token) metrics_dict = {} metrics_dict["start_date"] = report_start_date metrics_dict["end_date"] = report_end_date - cpu_request_metrics = utils.query_metric( - openshift_url, token, CPU_REQUEST, report_start_date, report_end_date + cpu_request_metrics = prom_client.query_metric( + CPU_REQUEST, report_start_date, report_end_date ) - memory_request_metrics = utils.query_metric( - openshift_url, token, MEMORY_REQUEST, report_start_date, report_end_date + memory_request_metrics = prom_client.query_metric( + MEMORY_REQUEST, report_start_date, report_end_date ) metrics_dict["cpu_metrics"] = cpu_request_metrics metrics_dict["memory_metrics"] = memory_request_metrics # because if nobody requests a GPU then we will get an empty set try: - gpu_request_metrics = utils.query_metric( - openshift_url, token, GPU_REQUEST, report_start_date, report_end_date + gpu_request_metrics = prom_client.query_metric( + GPU_REQUEST, report_start_date, report_end_date ) metrics_dict["gpu_metrics"] = gpu_request_metrics except utils.EmptyResultError: diff --git a/openshift_metrics/prometheus_client.py b/openshift_metrics/prometheus_client.py new file mode 100644 index 0000000..c64708a --- /dev/null +++ b/openshift_metrics/prometheus_client.py @@ -0,0 +1,41 @@ +import requests +import time + +from urllib3.util.retry import Retry +from requests.adapters import HTTPAdapter +from openshift_metrics.utils import EmptyResultError + +class PrometheusClient: + def __init__(self, prometheus_url: str, token: str, step_min: int=15): + self.prometheus_url = prometheus_url + self.token = token + self.step_min = step_min + + def query_metric(self, metric, start_date, end_date): + """Queries metric from the provided prometheus_url""" + data = None + headers = {"Authorization": f"Bearer {self.token}"} + day_url_vars = f"start={start_date}T00:00:00Z&end={end_date}T23:59:59Z" + url = f"{self.prometheus_url}/api/v1/query_range?query={metric}&{day_url_vars}&step={self.step_min}m" + + retries = Retry(total=3, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504]) + session = requests.Session() + session.mount("https://", HTTPAdapter(max_retries=retries)) + + print(f"Retrieving metric: {metric}") + + for _ in range(3): + response = session.get(url, headers=headers, verify=True) + + if response.status_code != 200: + print(f"{response.status_code} Response: {response.reason}") + else: + data = response.json()["data"]["result"] + if data: + break + print("Empty result set") + time.sleep(3) + + if not data: + raise EmptyResultError(f"Error retrieving metric: {metric}") + return data diff --git a/openshift_metrics/tests/test_prometheus_client.py b/openshift_metrics/tests/test_prometheus_client.py new file mode 100644 index 0000000..c5663f1 --- /dev/null +++ b/openshift_metrics/tests/test_prometheus_client.py @@ -0,0 +1,37 @@ +from requests.exceptions import ConnectionError +from unittest import TestCase, mock + +from openshift_metrics.prometheus_client import PrometheusClient + +class TestQueryMetric(TestCase): + + @mock.patch('requests.Session.get') + @mock.patch('time.sleep') + def test_query_metric(self, mock_sleep, mock_get): + mock_response = mock.Mock(status_code=200) + mock_response.json.return_value = {"data": { + "result": "this is data" + }} + mock_get.return_value = mock_response + prom_client = PrometheusClient('https://fake-url', 'fake-token') + metrics = prom_client.query_metric('fake-metric', '2022-03-14', '2022-03-14') + self.assertEqual(metrics, "this is data") + self.assertEqual(mock_get.call_count, 1) + + @mock.patch('requests.Session.get') + @mock.patch('time.sleep') + def test_query_metric_exception(self, mock_sleep, mock_get): + mock_get.return_value = mock.Mock(status_code=404) + prom_client = PrometheusClient('https://fake-url', 'fake-token') + self.assertRaises(Exception, prom_client.query_metric, + 'fake-metric', '2022-03-14', '2022-03-14') + self.assertEqual(mock_get.call_count, 3) + + @mock.patch('requests.Session.get') + @mock.patch('time.sleep') + def test_query_metric_connection_error(self, mock_sleep, mock_get): + mock_get.side_effect = [ConnectionError] + prom_client = PrometheusClient('https://fake-url', 'fake-token') + self.assertRaises(ConnectionError, prom_client.query_metric, + 'fake-metric', '2022-03-14', '2022-03-14') + self.assertEqual(mock_get.call_count, 1) diff --git a/openshift_metrics/tests/test_utils.py b/openshift_metrics/tests/test_utils.py index f46c65c..e9da675 100644 --- a/openshift_metrics/tests/test_utils.py +++ b/openshift_metrics/tests/test_utils.py @@ -17,39 +17,6 @@ from openshift_metrics import utils import os -class TestQueryMetric(TestCase): - - @mock.patch('requests.Session.get') - @mock.patch('time.sleep') - def test_query_metric(self, mock_sleep, mock_get): - mock_response = mock.Mock(status_code=200) - mock_response.json.return_value = {"data": { - "result": "this is data" - }} - mock_get.return_value = mock_response - - metrics = utils.query_metric('https://fake-url', 'fake-token', 'fake-metric', '2022-03-14', '2022-03-14') - self.assertEqual(metrics, "this is data") - self.assertEqual(mock_get.call_count, 1) - - @mock.patch('requests.Session.get') - @mock.patch('time.sleep') - def test_query_metric_exception(self, mock_sleep, mock_get): - mock_get.return_value = mock.Mock(status_code=404) - - self.assertRaises(Exception, utils.query_metric, 'https://fake-url', 'fake-token', - 'fake-metric', '2022-03-14', '2022-03-14') - self.assertEqual(mock_get.call_count, 3) - - @mock.patch('requests.Session.get') - @mock.patch('time.sleep') - def test_query_metric_connection_error(self, mock_sleep, mock_get): - mock_get.side_effect = [ConnectionError] - self.assertRaises(ConnectionError, utils.query_metric, 'https://fake-url', 'fake-token', - 'fake-metric', '2022-03-14', '2022-03-14') - self.assertEqual(mock_get.call_count, 1) - - class TestGetNamespaceAnnotations(TestCase): @mock.patch('openshift_metrics.utils.requests.post') diff --git a/openshift_metrics/utils.py b/openshift_metrics/utils.py index f54428e..a5b816e 100755 --- a/openshift_metrics/utils.py +++ b/openshift_metrics/utils.py @@ -111,36 +111,6 @@ def upload_to_s3(file, bucket, location): response = s3.upload_file(file, Bucket=bucket, Key=location) -def query_metric(openshift_url, token, metric, report_start_date, report_end_date): - """Queries metric from prometheus/thanos for the provided openshift_url""" - data = None - headers = {"Authorization": f"Bearer {token}"} - day_url_vars = f"start={report_start_date}T00:00:00Z&end={report_end_date}T23:59:59Z" - url = f"{openshift_url}/api/v1/query_range?query={metric}&{day_url_vars}&step={STEP_MIN}m" - - retries = Retry(total=3, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504]) - session = requests.Session() - session.mount("https://", HTTPAdapter(max_retries=retries)) - - print(f"Retrieving metric: {metric}") - - for _ in range(3): - response = session.get(url, headers=headers, verify=True) - - if response.status_code != 200: - print(f"{response.status_code} Response: {response.reason}") - else: - data = response.json()["data"]["result"] - if data: - break - print("Empty result set") - time.sleep(3) - - if not data: - raise EmptyResultError(f"Error retrieving metric: {metric}") - return data - - def get_namespace_attributes(): """ Returns allocation attributes from coldfront associated