postgresql_db: added tablespace support (#56390)

lib/ansible/module_utils/database.py

@@ -38,7 +38,7 @@ class UnclosedQuoteError(SQLParseError):
 # maps a type of identifier to the maximum number of dot levels that are
 # allowed to specify that identifier.  For example, a database column can be
 # specified by up to 4 levels: database.schema.table.column
-_PG_IDENTIFIER_TO_DOT_LEVEL = dict(database=1, schema=2, table=3, column=4, role=1)
+_PG_IDENTIFIER_TO_DOT_LEVEL = dict(database=1, schema=2, table=3, column=4, role=1, tablespace=1)
 _MYSQL_IDENTIFIER_TO_DOT_LEVEL = dict(database=1, table=2, column=3, role=1, vars=1)
 
 

lib/ansible/modules/database/postgresql/postgresql_db.py

@@ -95,6 +95,14 @@ options:
       - Specifies the database connection limit.
     type: str
     version_added: '2.8'
+  tablespace:
+    description:
+      - The tablespace to set for the database
+        U(https://www.postgresql.org/docs/current/sql-alterdatabase.html).
+      - If you want to move the database back to the default tablespace,
+        explicitly set this to pg_default.
+    type: path
+    version_added: '2.9'
 notes:
 - State C(dump) and C(restore) don't require I(psycopg2) since version 2.8.
 author: "Ansible Core Team"
@@ -140,6 +148,14 @@ EXAMPLES = r'''
     state: dump
     target: /tmp/acme.sql
     target_opts: "-n public"
+
+# Note: In the example below, if database foo exists and has another tablespace
+# the tablespace will be changed to foo. Access to the database will be locked
+# until the copying of database files is finished.
+- name: Create a new database called foo in tablespace bar
+  postgresql_db:
+    name: foo
+    tablespace: bar
 '''
 
 import os
@@ -198,8 +214,11 @@ def get_db_info(cursor, db):
     query = """
     SELECT rolname AS owner,
     pg_encoding_to_char(encoding) AS encoding, encoding AS encoding_id,
-    datcollate AS lc_collate, datctype AS lc_ctype, pg_database.datconnlimit AS conn_limit
-    FROM pg_database JOIN pg_roles ON pg_roles.oid = pg_database.datdba
+    datcollate AS lc_collate, datctype AS lc_ctype, pg_database.datconnlimit AS conn_limit,
+    spcname AS tablespace
+    FROM pg_database
+    JOIN pg_roles ON pg_roles.oid = pg_database.datdba
+    JOIN pg_tablespace ON pg_tablespace.oid = pg_database.dattablespace
     WHERE datname = %(db)s
     """
     cursor.execute(query, {'db': db})
@@ -221,8 +240,8 @@ def db_delete(cursor, db):
         return False
 
 
-def db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit):
-    params = dict(enc=encoding, collate=lc_collate, ctype=lc_ctype, conn_limit=conn_limit)
+def db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit, tablespace):
+    params = dict(enc=encoding, collate=lc_collate, ctype=lc_ctype, conn_limit=conn_limit, tablespace=tablespace)
     if not db_exists(cursor, db):
         query_fragments = ['CREATE DATABASE %s' % pg_quote_identifier(db, 'database')]
         if owner:
@@ -235,6 +254,8 @@ def db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_
             query_fragments.append('LC_COLLATE %(collate)s')
         if lc_ctype:
             query_fragments.append('LC_CTYPE %(ctype)s')
+        if tablespace:
+            query_fragments.append('TABLESPACE %s' % pg_quote_identifier(tablespace, 'tablespace'))
         if conn_limit:
             query_fragments.append("CONNECTION LIMIT %(conn_limit)s" % {"conn_limit": conn_limit})
         query = ' '.join(query_fragments)
@@ -267,10 +288,13 @@ def db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_
             if conn_limit and conn_limit != str(db_info['conn_limit']):
                 changed = set_conn_limit(cursor, db, conn_limit)
 
+            if tablespace and tablespace != db_info['tablespace']:
+                changed = set_tablespace(cursor, db, tablespace)
+
             return changed
 
 
-def db_matches(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit):
+def db_matches(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit, tablespace):
     if not db_exists(cursor, db):
         return False
     else:
@@ -286,6 +310,8 @@ def db_matches(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn
             return False
         elif conn_limit and conn_limit != str(db_info['conn_limit']):
             return False
+        elif tablespace and tablespace != db_info['tablespace']:
+            return False
         else:
             return True
 
@@ -414,6 +440,14 @@ def do_with_password(module, cmd, password):
     rc, stderr, stdout = module.run_command(cmd, use_unsafe_shell=True, environ_update=env)
     return rc, stderr, stdout, cmd
 
+
+def set_tablespace(cursor, db, tablespace):
+    query = "ALTER DATABASE %s SET TABLESPACE %s" % (
+            pg_quote_identifier(db, 'database'),
+            pg_quote_identifier(tablespace, 'tablespace'))
+    cursor.execute(query)
+    return True
+
 # ===========================================
 # Module execution.
 #
@@ -433,7 +467,8 @@ def main():
         target_opts=dict(type='str', default=''),
         maintenance_db=dict(type='str', default="postgres"),
         session_role=dict(type='str'),
-        conn_limit=dict(type='str', default='')
+        conn_limit=dict(type='str', default=''),
+        tablespace=dict(type='path', default=''),
     )
 
     module = AnsibleModule(
@@ -454,6 +489,7 @@ def main():
     maintenance_db = module.params['maintenance_db']
     session_role = module.params["session_role"]
     conn_limit = module.params['conn_limit']
+    tablespace = module.params['tablespace']
 
     raw_connection = state in ("dump", "restore")
 
@@ -519,7 +555,7 @@ def main():
             if state == "absent":
                 changed = db_exists(cursor, db)
             elif state == "present":
-                changed = not db_matches(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit)
+                changed = not db_matches(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit, tablespace)
             module.exit_json(changed=changed, db=db)
 
         if state == "absent":
@@ -530,7 +566,7 @@ def main():
 
         elif state == "present":
             try:
-                changed = db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit)
+                changed = db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype, conn_limit, tablespace)
             except SQLParseError as e:
                 module.fail_json(msg=to_native(e), exception=traceback.format_exc())
 

test/integration/targets/postgresql/tasks/main.yml

@@ -820,6 +820,9 @@
 # Test postgresql_tablespace module
 - include: postgresql_tablespace.yml
 
+# Test postgresql_db module, specific options:
+- include: postgresql_db.yml
+
 # Test postgresql_privs
 - include: postgresql_privs.yml
 

test/integration/targets/postgresql/tasks/postgresql_db.yml

@@ -0,0 +1,169 @@
+# Copyright: (c) 2019, Andrew Klychkov (@Andersson007) <aaklychkov@mail.ru>
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+# The file for testing new options for postgresql_db module.
+
+- vars:
+    db_tablespace: bar
+    tblspc_location: /ssd
+    db_name: acme
+    block_parameters: &block_parameters
+      become_user: "{{ pg_user }}"
+      become: True
+    task_parameters: &task_parameters
+      register: result
+    pg_parameters: &pg_parameters
+      login_user: "{{ pg_user }}"
+
+  # Start tablespace option tests:
+  block:
+  # create tablespace for tests
+  - name: postgresql_db_tablespace - Create tablespace
+    <<: *task_parameters
+    postgresql_tablespace:
+      <<: *pg_parameters
+      login_db: postgres
+      name: "{{ db_tablespace }}"
+      location: "{{ tblspc_location }}"
+
+  # Check mode for DB creation with tablespace option:
+  - name: postgresql_db_tablespace - Create DB with tablespace option in check mode
+    <<: *task_parameters
+    check_mode: yes
+    postgresql_db:
+      <<: *pg_parameters
+      maintenance_db: postgres
+      name: "{{ db_name }}"
+      tablespace: "{{ db_tablespace }}"
+
+  - assert:
+      that:
+      - result.changed == true
+
+  - name: postgresql_db_tablespace - Check actual DB tablespace, rowcount must be 0 because actually nothing changed
+    <<: *task_parameters
+    postgresql_query:
+      <<: *pg_parameters
+      login_db: postgres
+      query: >
+        SELECT 1 FROM pg_database AS d JOIN pg_tablespace AS t
+        ON d.dattablespace = t.oid WHERE d.datname = '{{ db_name }}'
+        AND t.spcname = '{{ db_tablespace }}'
+
+  - assert:
+      that:
+      - result.rowcount == 0
+
+  # Actual mode for creation with tablespace option:
+  - name: postgresql_db_tablespace - Create DB with tablespace option
+    <<: *task_parameters
+    postgresql_db:
+      <<: *pg_parameters
+      maintenance_db: postgres
+      name: "{{ db_name }}"
+      tablespace: "{{ db_tablespace }}"
+
+  - assert:
+      that:
+      - result.changed == true
+
+  - name: postgresql_db_tablespace - Check actual DB tablespace, rowcount must be 1
+    <<: *task_parameters
+    postgresql_query:
+      <<: *pg_parameters
+      login_db: postgres
+      query: >
+        SELECT 1 FROM pg_database AS d JOIN pg_tablespace AS t
+        ON d.dattablespace = t.oid WHERE d.datname = '{{ db_name }}'
+        AND t.spcname = '{{ db_tablespace }}'
+
+  - assert:
+      that:
+      - result.rowcount == 1
+ 
+  # Try to change tablespace to the same:
+  - name: postgresql_db_tablespace - The same DB with tablespace option again
+    <<: *task_parameters
+    postgresql_db:
+      <<: *pg_parameters
+      maintenance_db: postgres
+      name: "{{ db_name }}"
+      tablespace: "{{ db_tablespace }}"
+
+  - assert:
+      that:
+      - result.changed == false
+
+  # Try to change tablespace in check_mode:
+  - name: postgresql_db_tablespace - Change tablespace in check_mode
+    <<: *task_parameters
+    check_mode: yes
+    postgresql_db:
+      <<: *pg_parameters
+      maintenance_db: postgres
+      name: "{{ db_name }}"
+      tablespace: pg_default
+
+  - assert:
+      that:
+      - result.changed == true
+
+  - name: postgresql_db_tablespace - Check actual DB tablespace, rowcount must be 1 because actually nothing changed
+    <<: *task_parameters
+    postgresql_query:
+      <<: *pg_parameters
+      login_db: postgres
+      query: >
+        SELECT 1 FROM pg_database AS d JOIN pg_tablespace AS t
+        ON d.dattablespace = t.oid WHERE d.datname = '{{ db_name }}'
+        AND t.spcname = '{{ db_tablespace }}'
+
+  - assert:
+      that:
+      - result.rowcount == 1
+
+  # Try to change tablespace to pg_default in actual mode:
+  - name: postgresql_db_tablespace - Change tablespace in actual mode
+    <<: *task_parameters
+    postgresql_db:
+      <<: *pg_parameters
+      maintenance_db: postgres
+      name: "{{ db_name }}"
+      tablespace: pg_default
+
+  - assert:
+      that:
+      - result.changed == true
+
+  - name: postgresql_db_tablespace - Check actual DB tablespace, rowcount must be 1
+    <<: *task_parameters
+    postgresql_query:
+      <<: *pg_parameters
+      login_db: postgres
+      query: >
+        SELECT 1 FROM pg_database AS d JOIN pg_tablespace AS t
+        ON d.dattablespace = t.oid WHERE d.datname = '{{ db_name }}'
+        AND t.spcname = 'pg_default'
+
+  - assert:
+      that:
+      - result.rowcount == 1
+
+  # Cleanup:
+  - name: postgresql_db_tablespace - Drop test DB
+    <<: *task_parameters
+    postgresql_db:
+      <<: *pg_parameters
+      maintenance_db: postgres
+      name: "{{ db_name }}"
+      state: absent
+
+  - name: postgresql_db_tablespace - Remove tablespace
+    <<: *task_parameters
+    postgresql_tablespace:
+      <<: *pg_parameters
+      login_db: postgres
+      name: "{{ db_tablespace }}"
+      state: absent
+
+  <<: *block_parameters
+  # End of tablespace block