operations/postgres: enable modifying existing postgres databases (#1319)

Author: wowi42

pyinfra/facts/postgres.py

@@ -144,7 +144,7 @@ class PostgresDatabases(PostgresFactBase):
     """
 
     default = dict
-    psql_command = "SELECT pg_catalog.pg_encoding_to_char(encoding), * FROM pg_catalog.pg_database"
+    psql_command = "SELECT pg_catalog.pg_encoding_to_char(encoding), *, pg_catalog.pg_get_userbyid(datdba) AS owner FROM pg_catalog.pg_database"  # noqa: E501
 
     @override
     def process(self, output):
@@ -161,7 +161,7 @@ def process(self, output):
 
         for details in rows:
             details["encoding"] = details.pop("pg_encoding_to_char")
-
+            details["owner"] = details.pop("owner")
             for key, value in list(details.items()):
                 if key.endswith("id") or key in (
                     "dba",

pyinfra/operations/postgres.py

@@ -199,9 +199,8 @@ def database(
     + psql_*: global module arguments, see above
 
     Updates:
-        pyinfra will not attempt to change existing databases - it will either
-        create or drop databases, but not alter them (if the db exists this
-        operation will make no changes).
+        pyinfra will change existing databases - but some parameters are not
+        changeable (template, encoding, lc_collate and lc_ctype).
 
     **Example:**
 
@@ -212,7 +211,7 @@ def database(
             database="pyinfra_stuff",
             owner="pyinfra",
             encoding="UTF8",
-            sudo_user="postgres",
+            _sudo_user="postgres",
         )
 
     """
@@ -267,7 +266,51 @@ def database(
             database=psql_database,
         )
     else:
-        host.noop("postgresql database {0} exists".format(database))
+        current_db = current_databases[database]
+
+        for key, value, current_value in (
+            ("TEMPLATE", template, current_db.get("istemplate")),
+            ("ENCODING", encoding, current_db.get("encoding")),
+            ("LC_COLLATE", lc_collate, None),
+            ("LC_CTYPE", lc_ctype, None),
+        ):
+            if value and (current_value is None or current_value != value):
+                host.noop(
+                    "postgresql database {0} already exists, skipping {1}".format(database, key)
+                )
+
+        sql_bits = []
+
+        if owner and "owner" in current_db and current_db["owner"] != owner:
+            sql_bits.append('ALTER DATABASE "{0}" OWNER TO "{1}";'.format(database, owner))
+
+        if tablespace and "tablespace" in current_db and current_db["tablespace"] != tablespace:
+            sql_bits.append(
+                'ALTER DATABASE "{0}" SET TABLESPACE "{1}";'.format(database, tablespace)
+            )
+
+        if (
+            connection_limit
+            and "connlimit" in current_db
+            and current_db["connlimit"] != connection_limit
+        ):
+            sql_bits.append(
+                'ALTER DATABASE "{0}" CONNECTION LIMIT {1};'.format(database, connection_limit)
+            )
+
+        if len(sql_bits) > 0:
+            yield make_execute_psql_command(
+                StringCommand(*sql_bits),
+                user=psql_user,
+                password=psql_password,
+                host=psql_host,
+                port=psql_port,
+                database=psql_database,
+            )
+        else:
+            host.noop(
+                "postgresql database {0} already exists with the same parameters".format(database)
+            )
 
 
 @operation(is_idempotent=False)

tests/facts/postgresql.PostgresqlDatabases/multiple.json

@@ -1,10 +1,10 @@
 {
-    "command": "psql -Ac 'SELECT pg_catalog.pg_encoding_to_char(encoding), * FROM pg_catalog.pg_database'",
+    "command": "psql -Ac 'SELECT pg_catalog.pg_encoding_to_char(encoding), *, pg_catalog.pg_get_userbyid(datdba) AS owner FROM pg_catalog.pg_database'",
     "requires_command": "psql",
     "output": [
-        "pg_encoding_to_char|datname|datdba|encoding|datcollate|datctype|datistemplate|datallowconn|datconnlimit|datlastsysoid|datfrozenxid|datminmxid|dattablespace|datacl",
-"UTF8|postgres|10|6|en_US.UTF-8|en_US.UTF-8|f|t|-1|13052|548|1|1663|",
-"UTF8|pyinfra_stuff|16384|6|en_US.UTF-8|en_US.UTF-8|f|t|-1|13052|548|1|1663|",
+        "pg_encoding_to_char|datname|datdba|encoding|datcollate|datctype|datistemplate|datallowconn|datconnlimit|datlastsysoid|datfrozenxid|datminmxid|dattablespace|datacl|owner",
+"UTF8|postgres|10|6|en_US.UTF-8|en_US.UTF-8|f|t|-1|13052|548|1|1663||postgres",
+"UTF8|pyinfra_stuff|16384|6|en_US.UTF-8|en_US.UTF-8|f|t|-1|13052|548|1|1663||pyinfra",
         "(4 rows)"
     ],
     "fact": {
@@ -20,7 +20,8 @@
             "frozenxid": 548,
             "minmxid": 1,
             "tablespace": 1663,
-            "acl": null
+            "acl": null,
+            "owner": "postgres"
         },
         "pyinfra_stuff": {
             "dba": 16384,
@@ -34,7 +35,8 @@
             "frozenxid": 548,
             "minmxid": 1,
             "tablespace": 1663,
-            "acl": null
+            "acl": null,
+            "owner": "pyinfra"
         }
     }
 }

tests/operations/postgresql.database/update_database.json

@@ -0,0 +1,17 @@
+{
+    "args": ["somedb"],
+    "kwargs": {
+        "owner": "someowner",
+        "tablespace": "someothertablespace"
+    },
+    "facts" :{
+        "postgres.PostgresDatabases": {
+            "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": {
+                "somedb": { "owner": "someotherowner", "tablespace": "sometablespace" }
+            }
+        }
+    },
+    "commands": [
+        "psql -Ac 'ALTER DATABASE \"somedb\" OWNER TO \"someowner\"; ALTER DATABASE \"somedb\" SET TABLESPACE \"someothertablespace\";'"
+    ]
+}

tests/operations/postgresql.database/add_database_noop.json renamed to tests/operations/postgresql.database/update_database_noop.json

@@ -1,12 +1,15 @@
 {
     "args": ["somedb"],
+    "kwargs": {
+        "owner": "someowner"
+    },
     "facts" :{
         "postgres.PostgresDatabases": {
             "psql_database=None, psql_host=None, psql_password=None, psql_port=None, psql_user=None": {
-                "somedb": {}
+                "somedb": { "owner": "someowner" }
             }
         }
     },
     "commands": [],
-    "noop_description": "postgresql database somedb exists"
+    "noop_description": "postgresql database somedb already exists with the same parameters"
 }