diff --git a/tubearchivist-metrics/environment.py b/tubearchivist-metrics/environment.py new file mode 100644 index 0000000..983effa --- /dev/null +++ b/tubearchivist-metrics/environment.py @@ -0,0 +1,32 @@ +""" +Functionality for setting up the environment for the metrics package. +Reads in environment variables for the application to use. +""" +import os + + +class AppConfig: + def __init__(self) -> None: + self.config = self.get_config() + + @staticmethod + def get_config(): + """ + Reads in environment variables for the application to use. + """ + + es_pass = os.environ.get("ES_PASSWORD") + es_user = os.environ.get("ES_USER", default="elastic") + es_url = os.environ.get("ES_URL", default="http://archivist-es:9200") + listen_port = os.environ.get("LISTEN_PORT", default="9934") + poll_interval = os.environ.get("POLL_INTERVAL", default="60") + + application = { + "es_url": es_url, + "es_user": es_user, + "es_pass": es_pass, + "listen_port": listen_port, + "poll_interval": poll_interval, + } + + return application diff --git a/tubearchivist-metrics/esconnect.py b/tubearchivist-metrics/esconnect.py new file mode 100644 index 0000000..094dc35 --- /dev/null +++ b/tubearchivist-metrics/esconnect.py @@ -0,0 +1,30 @@ +from elasticsearch import Elasticsearch +from environment import AppConfig + +config = AppConfig().config + + +class ElasticWrapper: + """ + makes calls to elastic search + returns response count + """ + + es_url = config["es_url"] + es_user = config["es_user"] + es_pass = config["es_pass"] + + es = Elasticsearch( + [es_url], + basic_auth=(es_user, es_pass), + timeout=30, + max_retries=10, + retry_on_timeout=True, + ) + + def get_count(index_name): + """ + Returns the number of documents in the index + """ + response = ElasticWrapper.es.count(index=index_name) + return response["count"] diff --git a/tubearchivist-metrics/getmetrics.py b/tubearchivist-metrics/getmetrics.py new file mode 100644 index 0000000..21497b0 --- /dev/null +++ b/tubearchivist-metrics/getmetrics.py @@ -0,0 +1,11 @@ +from esconnect import ElasticWrapper + + +class GetMetrics: + @staticmethod + def count(index_name): + + """Get count of documents from ES""" + result = ElasticWrapper.get_count(index_name) + print("Metric for " + index_name + ": " + str(result)) + return int(result) diff --git a/tubearchivist-metrics/main.py b/tubearchivist-metrics/main.py new file mode 100644 index 0000000..0d390ba --- /dev/null +++ b/tubearchivist-metrics/main.py @@ -0,0 +1,63 @@ +import time +from prometheus_client import start_http_server, Gauge, Enum + +from environment import AppConfig +from getmetrics import GetMetrics + +config = AppConfig().config + +# Print configuration on console when starting the application + +print("Configuration is currently set to:") +print("Elasticsearch URL: " + config["es_url"]) +print("Listen Port: " + config["listen_port"]) +print("Poll Interval (seconds): " + config["poll_interval"]) + + +class AppMetrics: + # fmt: off + def __init__(self, poll_interval=int(config["poll_interval"])): + + self.poll_interval = poll_interval + + # Metrics to expose + self.channel_count = Gauge("channel_count", "Number of channels") + self.playlist_count = Gauge("playlist_count", "Number of playlists") + self.download_count = Gauge("download_count", "Number of downloads") + self.download_queue = Gauge("download_queue", "Number of pending downloads") + self.subtitle_count = Gauge("subtitle_count", "Number of subtitles downloaded for videos") + # fmt: on + + def run_metrics_loop(self): + """ + Runs a loop that will update the metrics every second. + """ + while True: + self.retrieve_metrics() + time.sleep(self.poll_interval) + + def retrieve_metrics(self): + """ + Retrieves the metrics from the database and updates the metrics. + """ + self.channel_count.set(GetMetrics.count(index_name="ta_channel")) + self.playlist_count.set(GetMetrics.count(index_name="ta_playlist")) + self.download_queue.set(GetMetrics.count(index_name="ta_download")) + self.download_count.set(GetMetrics.count(index_name="ta_video")) + self.subtitle_count.set(GetMetrics.count(index_name="ta_subtitle")) + + +def main(): + """Main Entry Point""" + listen_port = int(config["listen_port"]) + poll_interval = int(config["poll_interval"]) + + app_metrics = AppMetrics( + poll_interval=poll_interval, + ) + start_http_server(listen_port) + app_metrics.run_metrics_loop() + + +if __name__ == "__main__": + main()