diff --git a/src/pyinfra/facts/postgres.py b/src/pyinfra/facts/postgres.py index 552e8230b..92a4ec303 100644 --- a/src/pyinfra/facts/postgres.py +++ b/src/pyinfra/facts/postgres.py @@ -1,5 +1,7 @@ from __future__ import annotations +import re + from typing_extensions import override from pyinfra.api import FactBase, HiddenValue, QuoteString, StringCommand @@ -7,6 +9,20 @@ from .util.databases import parse_columns_and_rows +# PostgreSQL configuration parameter names are identifiers: a letter or underscore +# followed by letters, digits or underscores, optionally with a single dotted prefix +# for extension ("custom") parameters such as ``auto_explain.log_min_duration``. +SETTING_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)?$") + + +def quote_sql_literal(value: object) -> str: + """ + Render a value as a single-quoted PostgreSQL string literal, doubling any + embedded single quotes to prevent SQL injection. + """ + + return "'" + str(value).replace("'", "''") + "'" + def make_psql_command( database: str | None = None, @@ -182,3 +198,95 @@ def process(self, output): databases[details.pop("name")] = details return databases + + +def _process_settings_row(row: dict[str, str | None]) -> dict[str, str | None]: + return { + "value": row.get("setting"), + "unit": row.get("unit"), + "source": row.get("source"), + "context": row.get("context"), + } + + +class PostgresConfiguration(PostgresFactBase): + """ + Returns a dict of all PostgreSQL run-time configuration settings, as reported + by ``pg_settings``: + + .. code:: python + + { + "shared_buffers": { + "value": "16384", + "unit": "8kB", + "source": "configuration file", + "context": "postmaster", + }, + ... + } + + The ``value`` is the *running* value (``pg_settings.setting``), ``source`` is + where it came from and ``context`` says how a change takes effect (eg + ``postmaster`` needs a restart, ``sighup`` needs a reload). + """ + + default = dict + psql_command = "SELECT name, setting, unit, source, context FROM pg_settings ORDER BY name" + + @override + def process(self, output: list[str]) -> dict[str, dict[str, str | None]]: + # Remove the last line of the output (row count) + output = output[:-1] + rows = parse_columns_and_rows(output, "|") + return {row.pop("name"): _process_settings_row(row) for row in rows} + + +class PostgresSetting(PostgresFactBase): + """ + Returns a single PostgreSQL configuration setting as reported by + ``pg_settings``, or ``None`` when there is no such setting: + + .. code:: python + + { + "value": "4096", + "unit": "kB", + "source": "default", + "context": "user", + } + """ + + @override + def command( # type: ignore[override] + self, + setting: str, + psql_user: str | None = None, + psql_password: str | None = None, + psql_host: str | None = None, + psql_port: int | None = None, + psql_database: str | None = None, + ) -> StringCommand: + sql = ( + "SELECT name, setting, unit, source, context FROM pg_settings WHERE name = " + + quote_sql_literal(setting) + ) + return make_execute_psql_command( + sql, + user=psql_user, + password=psql_password, + host=psql_host, + port=psql_port, + database=psql_database, + ) + + @override + def process(self, output: list[str]) -> dict[str, str | None] | None: + # Remove the last line of the output (row count) + output = output[:-1] + rows = parse_columns_and_rows(output, "|") + if not rows: + return None + row = rows[0] + row.pop("name") + return _process_settings_row(row) diff --git a/src/pyinfra/operations/postgres.py b/src/pyinfra/operations/postgres.py index e37f12881..f66ce6773 100644 --- a/src/pyinfra/operations/postgres.py +++ b/src/pyinfra/operations/postgres.py @@ -16,13 +16,18 @@ from __future__ import annotations +import re + from pyinfra import host -from pyinfra.api import HiddenValue, QuoteString, StringCommand, operation +from pyinfra.api import HiddenValue, OperationError, QuoteString, StringCommand, operation from pyinfra.facts.postgres import ( + SETTING_NAME_RE, + PostgresConfiguration, PostgresDatabases, PostgresRoles, make_execute_psql_command, make_psql_command, + quote_sql_literal, ) @@ -440,3 +445,196 @@ def load( "<", QuoteString(src), ) + + +# Multipliers used to normalise PostgreSQL memory/time settings so an operation +# value such as "128MB" can be compared against the raw value + unit pg_settings +# reports (eg setting="16384", unit="8kB"). Memory is reduced to bytes and time +# to microseconds. +_MEMORY_UNITS = {"b": 1, "kb": 1024, "mb": 1024**2, "gb": 1024**3, "tb": 1024**4} +_TIME_UNITS = { + "us": 1, + "ms": 1000, + "s": 1_000_000, + "min": 60_000_000, + "h": 3_600_000_000, + "d": 86_400_000_000, +} +_UNIT_SCALES = {**_MEMORY_UNITS, **_TIME_UNITS} + +_TRUE_VALUES = {"on", "true", "yes", "1"} +_FALSE_VALUES = {"off", "false", "no", "0"} + +# A bare quantity ("128", "2.5") with an optional unit suffix ("128MB", "30s"). +_QUANTITY_RE = re.compile(r"^(-?\d+(?:\.\d+)?)\s*([a-zA-Z]*)$") +# A pg_settings unit, optionally prefixed with a block multiplier ("8kB", "kB"). +_UNIT_RE = re.compile(r"^(\d*)\s*([a-zA-Z]+)$") + + +def _to_bool(value: str) -> bool | None: + lowered = value.strip().lower() + if lowered in _TRUE_VALUES: + return True + if lowered in _FALSE_VALUES: + return False + return None + + +def _to_number(value: str) -> int | float | None: + try: + return int(value) + except ValueError: + try: + return float(value) + except ValueError: + return None + + +def _quantity_match(desired: str, current: str, unit: str) -> bool: + unit_match = _UNIT_RE.match(unit) + if unit_match is None: + return False + unit_multiplier = int(unit_match.group(1)) if unit_match.group(1) else 1 + unit_scale = _UNIT_SCALES.get(unit_match.group(2).lower()) + if unit_scale is None: + return False + + current_number = _to_number(current) + if current_number is None: + return False + current_base = current_number * unit_multiplier * unit_scale + + desired_match = _QUANTITY_RE.match(desired) + if desired_match is None: + return False + desired_number = _to_number(desired_match.group(1)) + if desired_number is None: + return False + suffix = desired_match.group(2) + if suffix: + suffix_scale = _UNIT_SCALES.get(suffix.lower()) + if suffix_scale is None: + return False + desired_base = desired_number * suffix_scale + else: + # A bare number is interpreted as a count in the setting's own unit. + desired_base = desired_number * unit_multiplier * unit_scale + + return current_base == desired_base + + +def _values_match(desired: object, current_value: str | None, unit: str | None) -> bool: + if current_value is None: + return False + + desired_str = str(desired).strip() + current_str = current_value.strip() + if desired_str == current_str: + return True + + if unit: + return _quantity_match(desired_str, current_str, unit) + + desired_bool = _to_bool(desired_str) + current_bool = _to_bool(current_str) + if desired_bool is not None and current_bool is not None: + return desired_bool == current_bool + + desired_number = _to_number(desired_str) + current_number = _to_number(current_str) + if desired_number is not None and current_number is not None: + return desired_number == current_number + + return desired_str.lower() == current_str.lower() + + +@operation( + idempotent_notice=( + "ALTER SYSTEM writes to postgresql.auto.conf; until the server reloads or " + "restarts, pg_settings still reports the old value, so this operation will " + "re-issue the change on each run until it is applied." + ), +) +def configuration( + setting: str, + value: str | int | bool | None = None, + present: bool = True, + # Details for speaking to PostgreSQL via `psql` CLI + psql_user: str | None = None, + psql_password: str | None = None, + psql_host: str | None = None, + psql_port: int | None = None, + psql_database: str | None = None, +): + """ + Set or reset a PostgreSQL server configuration parameter with ``ALTER SYSTEM``. + + + setting: name of the configuration parameter (eg ``work_mem``) + + value: desired value, required when ``present`` is ``True`` + + present: ``True`` to set ``value``, ``False`` to reset the setting to its + default with ``ALTER SYSTEM RESET`` + + psql_*: global module arguments, see above + + Reload/restart: + ``ALTER SYSTEM`` only writes ``postgresql.auto.conf``. Changes take effect + after the server reloads (``postgres.sql("SELECT pg_reload_conf()")``) or, + for settings whose ``context`` is ``postmaster``, after a full restart. + This operation never reloads or restarts for you, so it compares against + the *running* value and becomes a no-op once the change is live. + + **Example:** + + .. code:: python + + postgres.configuration( + name="Increase work_mem", + setting="work_mem", + value="8MB", + _sudo_user="postgres", + ) + """ + + if not SETTING_NAME_RE.match(setting): + raise OperationError(f"invalid PostgreSQL setting name: {setting}") + + if present and value is None: + raise OperationError("`value` is required when `present` is True") + + current = host.get_fact( + PostgresConfiguration, + psql_user=psql_user, + psql_password=psql_password, + psql_host=psql_host, + psql_port=psql_port, + psql_database=psql_database, + ).get(setting) + + if not present: + if current is None: + host.noop(f"postgresql setting {setting} is not set") + return + if current.get("source") == "default": + host.noop(f"postgresql setting {setting} is already at its default") + return + yield make_execute_psql_command( + f"ALTER SYSTEM RESET {setting}", + user=psql_user, + password=psql_password, + host=psql_host, + port=psql_port, + database=psql_database, + ) + return + + if current is not None and _values_match(value, current.get("value"), current.get("unit")): + host.noop(f"postgresql setting {setting} is already set to {value}") + return + + yield make_execute_psql_command( + f"ALTER SYSTEM SET {setting} = {quote_sql_literal(value)}", + user=psql_user, + password=psql_password, + host=psql_host, + port=psql_port, + database=psql_database, + ) diff --git a/tests/facts/postgres.PostgresConfiguration/custom_connection.yaml b/tests/facts/postgres.PostgresConfiguration/custom_connection.yaml new file mode 100644 index 000000000..2157f353c --- /dev/null +++ b/tests/facts/postgres.PostgresConfiguration/custom_connection.yaml @@ -0,0 +1,19 @@ +arg: + psql_user: myuser + psql_password: mypassword + psql_host: myhost + psql_port: 5432 +command: + raw: "PGPASSWORD=mypassword psql -U myuser -h myhost -p 5432 -Ac 'SELECT name, setting, unit, source, context FROM pg_settings ORDER BY name'" + masked: "PGPASSWORD='*MASKED*' psql -U myuser -h myhost -p 5432 -Ac 'SELECT name, setting, unit, source, context FROM pg_settings ORDER BY name'" +requires_command: psql +output: | + name|setting|unit|source|context + work_mem|4096|kB|default|user + (1 row) +fact: + work_mem: + value: "4096" + unit: kB + source: default + context: user diff --git a/tests/facts/postgres.PostgresConfiguration/settings.yaml b/tests/facts/postgres.PostgresConfiguration/settings.yaml new file mode 100644 index 000000000..ca89029bb --- /dev/null +++ b/tests/facts/postgres.PostgresConfiguration/settings.yaml @@ -0,0 +1,30 @@ +command: "psql -Ac 'SELECT name, setting, unit, source, context FROM pg_settings ORDER BY name'" +requires_command: psql +output: | + name|setting|unit|source|context + max_connections|100||default|postmaster + shared_buffers|16384|8kB|configuration file|postmaster + ssl|off||default|postmaster + work_mem|4096|kB|default|user + (4 rows) +fact: + max_connections: + value: "100" + unit: null + source: default + context: postmaster + shared_buffers: + value: "16384" + unit: "8kB" + source: configuration file + context: postmaster + ssl: + value: "off" + unit: null + source: default + context: postmaster + work_mem: + value: "4096" + unit: kB + source: default + context: user diff --git a/tests/facts/postgres.PostgresSetting/not_found.yaml b/tests/facts/postgres.PostgresSetting/not_found.yaml new file mode 100644 index 000000000..3d97ba429 --- /dev/null +++ b/tests/facts/postgres.PostgresSetting/not_found.yaml @@ -0,0 +1,8 @@ +arg: + setting: does_not_exist +command: "psql -Ac 'SELECT name, setting, unit, source, context FROM pg_settings WHERE name = '\"'\"'does_not_exist'\"'\"''" +requires_command: psql +output: | + name|setting|unit|source|context + (0 rows) +fact: null diff --git a/tests/facts/postgres.PostgresSetting/single.yaml b/tests/facts/postgres.PostgresSetting/single.yaml new file mode 100644 index 000000000..0bb9738d5 --- /dev/null +++ b/tests/facts/postgres.PostgresSetting/single.yaml @@ -0,0 +1,13 @@ +arg: + setting: work_mem +command: "psql -Ac 'SELECT name, setting, unit, source, context FROM pg_settings WHERE name = '\"'\"'work_mem'\"'\"''" +requires_command: psql +output: | + name|setting|unit|source|context + work_mem|4096|kB|default|user + (1 row) +fact: + value: "4096" + unit: kB + source: default + context: user diff --git a/tests/operations/postgres.configuration/boolean_change.yaml b/tests/operations/postgres.configuration/boolean_change.yaml new file mode 100644 index 000000000..523cc64e6 --- /dev/null +++ b/tests/operations/postgres.configuration/boolean_change.yaml @@ -0,0 +1,14 @@ +args: + - ssl +kwargs: + value: "on" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + ssl: + value: "off" + unit: null + source: default + context: postmaster +commands: + - "psql -Ac 'ALTER SYSTEM SET ssl = '\"'\"'on'\"'\"''" diff --git a/tests/operations/postgres.configuration/boolean_match_noop.yaml b/tests/operations/postgres.configuration/boolean_match_noop.yaml new file mode 100644 index 000000000..5c04fa142 --- /dev/null +++ b/tests/operations/postgres.configuration/boolean_match_noop.yaml @@ -0,0 +1,14 @@ +args: + - ssl +kwargs: + value: true +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + ssl: + value: "on" + unit: null + source: configuration file + context: postmaster +commands: [] +noop_description: "postgresql setting ssl is already set to True" diff --git a/tests/operations/postgres.configuration/invalid_name.yaml b/tests/operations/postgres.configuration/invalid_name.yaml new file mode 100644 index 000000000..2038cac7b --- /dev/null +++ b/tests/operations/postgres.configuration/invalid_name.yaml @@ -0,0 +1,7 @@ +args: + - "work_mem = 0; DROP DATABASE prod" +kwargs: + value: "8MB" +exception: + name: OperationError + message: "invalid PostgreSQL setting name: work_mem = 0; DROP DATABASE prod" diff --git a/tests/operations/postgres.configuration/memory_block_match_noop.yaml b/tests/operations/postgres.configuration/memory_block_match_noop.yaml new file mode 100644 index 000000000..1248e2b63 --- /dev/null +++ b/tests/operations/postgres.configuration/memory_block_match_noop.yaml @@ -0,0 +1,14 @@ +args: + - shared_buffers +kwargs: + value: "128MB" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + shared_buffers: + value: "16384" + unit: "8kB" + source: configuration file + context: postmaster +commands: [] +noop_description: "postgresql setting shared_buffers is already set to 128MB" diff --git a/tests/operations/postgres.configuration/memory_change.yaml b/tests/operations/postgres.configuration/memory_change.yaml new file mode 100644 index 000000000..c85a64535 --- /dev/null +++ b/tests/operations/postgres.configuration/memory_change.yaml @@ -0,0 +1,14 @@ +args: + - shared_buffers +kwargs: + value: "256MB" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + shared_buffers: + value: "16384" + unit: "8kB" + source: configuration file + context: postmaster +commands: + - "psql -Ac 'ALTER SYSTEM SET shared_buffers = '\"'\"'256MB'\"'\"''" diff --git a/tests/operations/postgres.configuration/missing_value.yaml b/tests/operations/postgres.configuration/missing_value.yaml new file mode 100644 index 000000000..1f0e7ac6f --- /dev/null +++ b/tests/operations/postgres.configuration/missing_value.yaml @@ -0,0 +1,5 @@ +args: + - work_mem +exception: + name: OperationError + message: "`value` is required when `present` is True" diff --git a/tests/operations/postgres.configuration/number_match_noop.yaml b/tests/operations/postgres.configuration/number_match_noop.yaml new file mode 100644 index 000000000..00b8cf272 --- /dev/null +++ b/tests/operations/postgres.configuration/number_match_noop.yaml @@ -0,0 +1,14 @@ +args: + - seq_page_cost +kwargs: + value: "1.0" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + seq_page_cost: + value: "1" + unit: null + source: default + context: user +commands: [] +noop_description: "postgresql setting seq_page_cost is already set to 1.0" diff --git a/tests/operations/postgres.configuration/reset.yaml b/tests/operations/postgres.configuration/reset.yaml new file mode 100644 index 000000000..cf1d7f1f5 --- /dev/null +++ b/tests/operations/postgres.configuration/reset.yaml @@ -0,0 +1,14 @@ +args: + - work_mem +kwargs: + present: false +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + work_mem: + value: "8192" + unit: kB + source: configuration file + context: user +commands: + - "psql -Ac 'ALTER SYSTEM RESET work_mem'" diff --git a/tests/operations/postgres.configuration/reset_absent_noop.yaml b/tests/operations/postgres.configuration/reset_absent_noop.yaml new file mode 100644 index 000000000..c25e18728 --- /dev/null +++ b/tests/operations/postgres.configuration/reset_absent_noop.yaml @@ -0,0 +1,9 @@ +args: + - some_custom.guc +kwargs: + present: false +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": {} +commands: [] +noop_description: "postgresql setting some_custom.guc is not set" diff --git a/tests/operations/postgres.configuration/reset_default_noop.yaml b/tests/operations/postgres.configuration/reset_default_noop.yaml new file mode 100644 index 000000000..9aff918fa --- /dev/null +++ b/tests/operations/postgres.configuration/reset_default_noop.yaml @@ -0,0 +1,14 @@ +args: + - work_mem +kwargs: + present: false +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + work_mem: + value: "4096" + unit: kB + source: default + context: user +commands: [] +noop_description: "postgresql setting work_mem is already at its default" diff --git a/tests/operations/postgres.configuration/set_custom_connection.yaml b/tests/operations/postgres.configuration/set_custom_connection.yaml new file mode 100644 index 000000000..2799c277b --- /dev/null +++ b/tests/operations/postgres.configuration/set_custom_connection.yaml @@ -0,0 +1,19 @@ +args: + - work_mem +kwargs: + value: "8MB" + psql_user: myuser + psql_password: mypassword + psql_host: myhost + psql_port: 5432 +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=myhost, psql_password=mypassword, psql_port=5432, psql_user=myuser": + work_mem: + value: "4096" + unit: kB + source: default + context: user +commands: + - raw: "PGPASSWORD=mypassword psql -U myuser -h myhost -p 5432 -Ac 'ALTER SYSTEM SET work_mem = '\"'\"'8MB'\"'\"''" + masked: "PGPASSWORD='*MASKED*' psql -U myuser -h myhost -p 5432 -Ac 'ALTER SYSTEM SET work_mem = '\"'\"'8MB'\"'\"''" diff --git a/tests/operations/postgres.configuration/set_new.yaml b/tests/operations/postgres.configuration/set_new.yaml new file mode 100644 index 000000000..69c91102e --- /dev/null +++ b/tests/operations/postgres.configuration/set_new.yaml @@ -0,0 +1,14 @@ +args: + - work_mem +kwargs: + value: "8MB" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + work_mem: + value: "4096" + unit: kB + source: default + context: user +commands: + - "psql -Ac 'ALTER SYSTEM SET work_mem = '\"'\"'8MB'\"'\"''" diff --git a/tests/operations/postgres.configuration/set_unchanged_noop.yaml b/tests/operations/postgres.configuration/set_unchanged_noop.yaml new file mode 100644 index 000000000..bc7824cca --- /dev/null +++ b/tests/operations/postgres.configuration/set_unchanged_noop.yaml @@ -0,0 +1,14 @@ +args: + - work_mem +kwargs: + value: "4MB" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + work_mem: + value: "4096" + unit: kB + source: configuration file + context: user +commands: [] +noop_description: "postgresql setting work_mem is already set to 4MB" diff --git a/tests/operations/postgres.configuration/set_value_with_quote.yaml b/tests/operations/postgres.configuration/set_value_with_quote.yaml new file mode 100644 index 000000000..da3a93bb2 --- /dev/null +++ b/tests/operations/postgres.configuration/set_value_with_quote.yaml @@ -0,0 +1,14 @@ +args: + - search_path +kwargs: + value: "a'b" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + search_path: + value: "\"$user\", public" + unit: null + source: default + context: user +commands: + - "psql -Ac 'ALTER SYSTEM SET search_path = '\"'\"'a'\"'\"''\"'\"'b'\"'\"''" diff --git a/tests/operations/postgres.configuration/time_match_noop.yaml b/tests/operations/postgres.configuration/time_match_noop.yaml new file mode 100644 index 000000000..8453dc7cc --- /dev/null +++ b/tests/operations/postgres.configuration/time_match_noop.yaml @@ -0,0 +1,14 @@ +args: + - idle_in_transaction_session_timeout +kwargs: + value: "30s" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": + idle_in_transaction_session_timeout: + value: "30000" + unit: ms + source: configuration file + context: user +commands: [] +noop_description: "postgresql setting idle_in_transaction_session_timeout is already set to 30s" diff --git a/tests/operations/postgres.configuration/unknown_setting_sets.yaml b/tests/operations/postgres.configuration/unknown_setting_sets.yaml new file mode 100644 index 000000000..79e4221f0 --- /dev/null +++ b/tests/operations/postgres.configuration/unknown_setting_sets.yaml @@ -0,0 +1,9 @@ +args: + - auto_explain.log_min_duration +kwargs: + value: "100" +facts: + postgres.PostgresConfiguration: + "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": {} +commands: + - "psql -Ac 'ALTER SYSTEM SET auto_explain.log_min_duration = '\"'\"'100'\"'\"''"