From a8eb80aad2e497e8bb9752b48b8f2b95bb786eba Mon Sep 17 00:00:00 2001 From: Kumi Date: Wed, 15 May 2024 22:41:26 +0200 Subject: [PATCH] feat: Initial implementation of a PostgreSQL connection exporter --- .gitignore | 5 + LICENSE | 19 +++ README.md | 32 ++++ config.dist.yaml | 8 + pyproject.toml | 32 ++++ src/postgres_connection_exporter/__init__.py | 0 src/postgres_connection_exporter/__main__.py | 166 +++++++++++++++++++ 7 files changed, 262 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README.md create mode 100644 config.dist.yaml create mode 100644 pyproject.toml create mode 100644 src/postgres_connection_exporter/__init__.py create mode 100644 src/postgres_connection_exporter/__main__.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..91c9bd0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +.venv +config.yaml +__pycache__ +*.pyc +venv \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..32bb202 --- /dev/null +++ b/LICENSE @@ -0,0 +1,19 @@ +Copyright (c) 2024 Kumi Mitterer + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..78b0e8e --- /dev/null +++ b/README.md @@ -0,0 +1,32 @@ +# PostgreSQL Connection Exporter for Prometheus + +This is a simple server that exports PostgreSQL connection metrics in a format that can be scraped by Prometheus. + +It outputs the following metrics: + +- The number of connections per database +- The number of connections per user +- The number of connections per client address +- The number of connections per state + +## Installation + +You can install the exporter from PyPI. Within a virtual environment, run: + +```bash +pip install postgres-connection-exporter +``` + +## Usage + +The exporter is configured using a `config.yaml`. You can find an example in [config.dist.yaml](config.dist.yaml). + +After you have created your `config.yaml`, you can start the exporter with: + +```bash +postgres-connection-exporter +``` + +## License + +This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details. \ No newline at end of file diff --git a/config.dist.yaml b/config.dist.yaml new file mode 100644 index 0000000..e166c40 --- /dev/null +++ b/config.dist.yaml @@ -0,0 +1,8 @@ +hosts: # List of database hosts + # Each host must have a user, password, host, and optionally port + - user: "user1" + password: "password1" + host: "host1" + port: "5432" +exporter: + port: 8989 # Port on which the exporter will listen diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..27a97a3 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,32 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "postgres-connection-exporter" +version = "0.0.1" +authors = [ + { name="Kumi Mitterer", email="postgres-connection-exporter@kumi.email" }, +] +description = "A Prometheus exporter for PostgreSQL connection metrics" +readme = "README.md" +license = { file="LICENSE" } +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent" +] +dependencies = [ + "prometheus-client", + "psycopg2-binary", + "pyyaml", +] + +[project.scripts] +postgres-connection-exporter = "postgres_connection_exporter.__main__:main" + +[project.urls] +"Homepage" = "https://git.private.coffee/kumi/postgres-connection-exporter" +"Bug Tracker" = "https://git.private.coffee/kumi/postgres-connection-exporter/issues" +"Source Code" = "https://git.private.coffee/kumi/postgres-connection-exporter" \ No newline at end of file diff --git a/src/postgres_connection_exporter/__init__.py b/src/postgres_connection_exporter/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/postgres_connection_exporter/__main__.py b/src/postgres_connection_exporter/__main__.py new file mode 100644 index 0000000..f83b53f --- /dev/null +++ b/src/postgres_connection_exporter/__main__.py @@ -0,0 +1,166 @@ +import yaml +from prometheus_client import start_http_server, Gauge +import psycopg2 +import time + +CONNECTIONS_PER_DB = Gauge( + "db_connections_per_database", + "Number of current connections per database", + ["database", "host"], +) +CONNECTIONS_PER_USER = Gauge( + "db_connections_per_user", + "Number of current connections per user", + ["user", "host"], +) +CONNECTIONS_BY_STATE = Gauge( + "db_connections_by_state", + "Number of current connections by state", + ["state", "host"], +) +CONNECTIONS_PER_SOURCE = Gauge( + "db_connections_per_source", + "Number of current connections per source", + ["source", "host"], +) + + +def load_config(config_file="config.yaml"): + with open(config_file, "r") as f: + return yaml.safe_load(f) + + +def get_all_databases(db_config): + try: + conn = psycopg2.connect( + dbname="postgres", + user=db_config["user"], + password=db_config["password"], + host=db_config["host"], + port=db_config.get("port", 5432), + ) + cursor = conn.cursor() + + cursor.execute("SELECT datname FROM pg_database WHERE datistemplate = false;") + databases = cursor.fetchall() + + cursor.close() + conn.close() + + return [db[0] for db in databases] + + except Exception as e: + print(f"Error retrieving list of databases: {e}") + return [] + + +def get_db_connections(db_config): + try: + # Connect to the PostgreSQL database + conn = psycopg2.connect( + dbname="postgres", + user=db_config["user"], + password=db_config["password"], + host=db_config["host"], + port=db_config["port"], + ) + cursor = conn.cursor() + + # Query to get the number of connections per database + cursor.execute( + """ + SELECT datname, COUNT(*) + FROM pg_stat_activity + GROUP BY datname; + """ + ) + db_connections = cursor.fetchall() + for db, count in db_connections: + CONNECTIONS_PER_DB.labels(database=db, host=db_config["host"]).set(count) + + # Query to get the number of connections per user + cursor.execute( + """ + SELECT usename, COUNT(*) + FROM pg_stat_activity + GROUP BY usename; + """ + ) + user_connections = cursor.fetchall() + for user, count in user_connections: + CONNECTIONS_PER_USER.labels(user=user, host=db_config["host"]).set(count) + + # Query to get the number of connections by state + cursor.execute( + """ + SELECT state, COUNT(*) + FROM pg_stat_activity + GROUP BY state; + """ + ) + state_connections = cursor.fetchall() + for state, count in state_connections: + CONNECTIONS_BY_STATE.labels(state=state, host=db_config["host"]).set(count) + + # Query to get the number of connections per source + cursor.execute( + """ + SELECT client_addr, COUNT(*) + FROM pg_stat_activity + GROUP BY client_addr; + """ + ) + source_connections = cursor.fetchall() + for source, count in source_connections: + CONNECTIONS_PER_SOURCE.labels(source=source, host=db_config["host"]).set(count) + + cursor.close() + conn.close() + except Exception as e: + print(f"Error retrieving data from {db_config['host']}: {e}") + + +def main(): + config = load_config() + + if not ("hosts" in config and config["hosts"]): + print("No database hosts specified in the configuration file.") + exit(1) + + databases_to_query = [] + + for host in config["hosts"]: + if not all(key in host for key in ["user", "password", "host", "port"]): + print("Database configuration is missing required fields.") + exit(1) + + db_config = { + "user": host["user"], + "password": host["password"], + "host": host["host"], + "port": host["port"], + } + + databases_to_query.append(db_config) + + if not databases_to_query: + print("No databases to query.") + exit(1) + + exporter_port = ( + config["exporter"]["port"] + if "exporter" in config and "port" in config["exporter"] + else 8989 + ) + + start_http_server(exporter_port) + print(f"Prometheus exporter started on port {exporter_port}") + + while True: + for db in databases_to_query: + get_db_connections(db) + time.sleep(15) + + +if __name__ == "__main__": + main()