New module postgresql_query (#52555)

* New module postgresql_query

* New module postgresql_query: added tests

* New module postgresql_query: added path_to_script

* New module postgresql_query: fix doc

* New module postgresql_query: fix autocommit

* New module postgresql_query: added exception for centos6 tests

* New module postgresql_query: fixes

* New module postgresql_query: add psycopg2 check

* New module postgresql_query: add psycopg2 check, fix

* New module postgresql_query: add psycopg2 check, fix

* New module postgresql_query: add psycopg2 check, fix

* New module postgresql_query: fix a type

3 files changed, 613 insertions, 0 deletions

diff --git a/lib/ansible/modules/database/postgresql/postgresql_query.py b/lib/ansible/modules/database/postgresql/postgresql_query.py
new file mode 100644
index 0000000000..206eff6459
--- /dev/null
+++ b/lib/ansible/modules/database/postgresql/postgresql_query.py
@@ -0,0 +1,360 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+
+# Copyright: (c) 2017, Felix Archambault
+# 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)
+
+from __future__ import (absolute_import, division, print_function)
+__metaclass__ = type
+
+ANSIBLE_METADATA = {
+    'metadata_version': '1.1',
+    'supported_by': 'community',
+    'status': ['preview']
+}
+
+DOCUMENTATION = r'''
+---
+module: postgresql_query
+short_description: Run PostgreSQL queries
+description:
+- Runs arbitraty PostgreSQL queries.
+- Can run queries from SQL script files.
+version_added: "2.8"
+options:
+  query:
+    description:
+    - SQL query to run. Variables can be escaped with psycopg2 syntax U(http://initd.org/psycopg/docs/usage.html).
+    type: str
+  positional_args:
+    description:
+    - List of values to be passed as positional arguments to the query.
+    - Mutually exclusive with I(named_args).
+    type: list
+  named_args:
+    description:
+    - Dictionary of key-value arguments to pass to the query.
+    - Mutually exclusive with I(positional_args).
+    type: dict
+  path_to_script:
+    description:
+    - Path to SQL script on the remote host.
+    - Returns result of the last query in the script.
+    - Mutually exclusive with I(query).
+    type: path
+  session_role:
+    description:
+    - Switch to session_role after connecting. The specified session_role must
+      be a role that the current login_user is a member of.
+    - Permissions checking for SQL commands is carried out as though
+      the session_role were the one that had logged in originally.
+    type: str
+  db:
+    description:
+    - Name of database to connect to and run queries against.
+    type: str
+  port:
+    description:
+    - Database port to connect.
+    type: int
+    default: 5432
+  login_user:
+    description:
+    - User (role) used to authenticate with PostgreSQL.
+    type: str
+    default: postgres
+  login_password:
+    description:
+    - Password used to authenticate with PostgreSQL.
+    type: str
+  login_host:
+    description:
+    - Host running PostgreSQL.
+    type: str
+  login_unix_socket:
+    description:
+    - Path to a Unix domain socket for local connections.
+    type: str
+  ssl_mode:
+    description:
+    - Determines whether or with what priority a secure SSL TCP/IP connection
+      will be negotiated with the server.
+    - See U(https://www.postgresql.org/docs/current/static/libpq-ssl.html) for
+      more information on the modes.
+    - Default of C(prefer) matches libpq default.
+    type: str
+    default: prefer
+    choices: [ allow, disable, prefer, require, verify-ca, verify-full ]
+  ssl_rootcert:
+    description:
+    - Specifies the name of a file containing SSL certificate authority (CA)
+      certificate(s).
+    - If the file exists, the server's certificate will be
+      verified to be signed by one of these authorities.
+    type: str
+notes:
+- The default authentication assumes that you are either logging in as or
+  sudo'ing to the postgres account on the host.
+- To avoid "Peer authentication failed for user postgres" error,
+  use postgres user as a I(become_user).
+- This module uses psycopg2, a Python PostgreSQL database adapter. You must
+  ensure that psycopg2 is installed on the host before using this module. If
+  the remote host is the PostgreSQL server (which is the default case), then
+  PostgreSQL must also be installed on the remote host. For Ubuntu-based
+  systems, install the postgresql, libpq-dev, and python-psycopg2 packages
+  on the remote host before using this module.
+requirements: [ psycopg2 ]
+author:
+- Felix Archambault (@archf)
+- Andrew Klychkov (@Andersson007)
+- Will Rouesnel (@wrouesnel)
+'''
+
+EXAMPLES = r'''
+- name: Simple select query to acme db
+  postgresql_query:
+    db: acme
+    query: SELECT version()
+
+- name: Select query to db acme with positional arguments and non-default credentials
+  postgresql_query:
+    db: acme
+    login_user: django
+    login_password: mysecretpass
+    query: SELECT * FROM acme WHERE id = %s AND story = %s
+    positional_args:
+    - 1
+    - test
+
+- name: Select query to test_db with named_args
+  postgresql_query:
+    db: test_db
+    query: SELECT * FROM test WHERE id = %(id_val)s AND story = %(story_val)s
+    named_args:
+      id_val: 1
+      story_val: test
+
+- name: Insert query to db test_db
+  postgresql_query:
+    db: test_db
+    query: INSERT INTO test_db (id, story) VALUES (2, 'my_long_story')
+
+- name: Run queries from SQL script
+  postgresql_query:
+    db: test_db
+    path_to_script: /var/lib/pgsql/test.sql
+    positional_args:
+    - 1
+'''
+
+RETURN = r'''
+query:
+    description: Query that was tried to be executed.
+    returned: always
+    type: str
+    sample: 'SELECT * FROM bar'
+statusmessage:
+    description: Attribute containing the message returned by the command.
+    returned: always
+    type: str
+    sample: 'INSERT 0 1'
+query_result:
+    description:
+    - List of dictionaries in column:value form representing returned rows.
+    returned: changed
+    type: list
+    sample: [{"Column": "Value1"},{"Column": "Value2"}]
+rowcount:
+    description: Number of affected rows.
+    returned: changed
+    type: int
+    sample: 5
+'''
+
+import os
+
+try:
+    import psycopg2
+    HAS_PSYCOPG2 = True
+except ImportError:
+    HAS_PSYCOPG2 = False
+
+import ansible.module_utils.postgres as pgutils
+from ansible.module_utils.basic import AnsibleModule, missing_required_lib
+from ansible.module_utils.database import SQLParseError
+from ansible.module_utils.postgres import postgres_common_argument_spec
+from ansible.module_utils._text import to_native
+from ansible.module_utils.six import iteritems
+
+
+def connect_to_db(module, kw, autocommit=False):
+    try:
+        db_connection = psycopg2.connect(**kw)
+        if autocommit:
+            if psycopg2.__version__ >= '2.4.2':
+                db_connection.set_session(autocommit=True)
+            else:
+                db_connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+
+    except TypeError as e:
+        if 'sslrootcert' in e.args[0]:
+            module.fail_json(msg='Postgresql server must be at least '
+                                 'version 8.4 to support sslrootcert')
+
+        module.fail_json(msg="unable to connect to database: %s" % to_native(e))
+
+    except Exception as e:
+        module.fail_json(msg="unable to connect to database: %s" % to_native(e))
+
+    return db_connection
+
+# ===========================================
+# Module execution.
+#
+
+
+def main():
+    argument_spec = postgres_common_argument_spec()
+    argument_spec.update(
+        query=dict(type='str'),
+        db=dict(type='str'),
+        ssl_mode=dict(type='str', default='prefer', choices=['allow', 'disable', 'prefer', 'require', 'verify-ca', 'verify-full']),
+        ssl_rootcert=dict(type='str'),
+        positional_args=dict(type='list'),
+        named_args=dict(type='dict'),
+        session_role=dict(type='str'),
+        path_to_script=dict(type='path'),
+    )
+
+    module = AnsibleModule(
+        argument_spec=argument_spec,
+        mutually_exclusive=(('positional_args', 'named_args'),),
+        supports_check_mode=True,
+    )
+
+    if not HAS_PSYCOPG2:
+        module.fail_json(msg=missing_required_lib('psycopg2'))
+
+    query = module.params["query"]
+    positional_args = module.params["positional_args"]
+    named_args = module.params["named_args"]
+    sslrootcert = module.params["ssl_rootcert"]
+    session_role = module.params["session_role"]
+    path_to_script = module.params["path_to_script"]
+
+    if positional_args and named_args:
+        module.fail_json(msg="positional_args and named_args params are mutually exclusive")
+
+    if path_to_script and query:
+        module.fail_json(msg="path_to_script is mutually exclusive with query")
+
+    if path_to_script:
+        try:
+            query = open(path_to_script, 'r').read()
+        except Exception as e:
+            module.fail_json(msg="Cannot read file '%s' : %s" % (path_to_script, to_native(e)))
+
+    # To use defaults values, keyword arguments must be absent, so
+    # check which values are empty and don't include in the **kw
+    # dictionary
+    params_map = {
+        "login_host": "host",
+        "login_user": "user",
+        "login_password": "password",
+        "port": "port",
+        "db": "database",
+        "ssl_mode": "sslmode",
+        "ssl_rootcert": "sslrootcert"
+    }
+    kw = dict((params_map[k], v) for (k, v) in iteritems(module.params)
+              if k in params_map and v != '' and v is not None)
+
+    # If a login_unix_socket is specified, incorporate it here.
+    is_localhost = "host" not in kw or kw["host"] is None or kw["host"] == "localhost"
+    if is_localhost and module.params["login_unix_socket"] != "":
+        kw["host"] = module.params["login_unix_socket"]
+
+    if psycopg2.__version__ < '2.4.3' and sslrootcert:
+        module.fail_json(msg='psycopg2 must be at least 2.4.3 '
+                             'in order to user the ssl_rootcert parameter')
+
+    db_connection = connect_to_db(module, kw)
+    cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
+
+    # Switch role, if specified:
+    if session_role:
+        try:
+            cursor.execute('SET ROLE %s' % session_role)
+        except Exception as e:
+            module.fail_json(msg="Could not switch role: %s" % to_native(e))
+
+    # Prepare args:
+    if module.params["positional_args"]:
+        arguments = module.params["positional_args"]
+    elif module.params["named_args"]:
+        arguments = module.params["named_args"]
+    else:
+        arguments = None
+
+    # Set defaults:
+    changed = False
+
+    # Execute query:
+    try:
+        cursor.execute(query, arguments)
+    except Exception as e:
+        cursor.close()
+        db_connection.close()
+        module.fail_json(msg="Cannot execute SQL '%s' %s: %s" % (query, arguments, to_native(e)))
+
+    statusmessage = cursor.statusmessage
+    rowcount = cursor.rowcount
+
+    try:
+        query_result = [dict(row) for row in cursor.fetchall()]
+    except psycopg2.ProgrammingError as e:
+        if to_native(e) == 'no results to fetch':
+            query_result = {}
+
+    except Exception as e:
+        module.fail_json(msg="Cannot fetch rows from cursor: %s" % to_native(e))
+
+    if 'SELECT' not in statusmessage:
+        if 'UPDATE' in statusmessage or 'INSERT' in statusmessage or 'DELETE' in statusmessage:
+            s = statusmessage.split()
+            if len(s) == 3:
+                if statusmessage.split()[2] != '0':
+                    changed = True
+
+            elif len(s) == 2:
+                if statusmessage.split()[1] != '0':
+                    changed = True
+
+            else:
+                changed = True
+
+        else:
+            changed = True
+
+    if module.check_mode:
+        db_connection.rollback()
+    else:
+        db_connection.commit()
+
+    kw = dict(
+        changed=changed,
+        query=cursor.query,
+        statusmessage=statusmessage,
+        query_result=query_result,
+        rowcount=rowcount if rowcount >= 0 else 0,
+    )
+
+    cursor.close()
+    db_connection.close()
+
+    module.exit_json(**kw)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/test/integration/targets/postgresql/tasks/main.yml b/test/integration/targets/postgresql/tasks/main.yml
index cd82b33a52..3f35d44ee8 100644
--- a/test/integration/targets/postgresql/tasks/main.yml
+++ b/test/integration/targets/postgresql/tasks/main.yml
@@ -768,6 +768,9 @@
 # Test postgresql_idx module
 - include: postgresql_idx.yml
 
+# Test postgresql_query module
+- include: postgresql_query.yml
+
 # dump/restore tests per format
 # ============================================================
 - include: state_dump_restore.yml test_fixture=user file=dbdata.sql
diff --git a/test/integration/targets/postgresql/tasks/postgresql_query.yml b/test/integration/targets/postgresql/tasks/postgresql_query.yml
new file mode 100644
index 0000000000..5a8cb82f68
--- /dev/null
+++ b/test/integration/targets/postgresql/tasks/postgresql_query.yml
@@ -0,0 +1,250 @@
+# 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)
+
+# Prepare for tests:
+- name: postgresql_query - drop test table if exists
+  become_user: "{{ pg_user }}"
+  become: yes
+  shell: psql postgres -U "{{ pg_user }}" -t -c "DROP TABLE IF EXISTS test_table;"
+  ignore_errors: yes
+
+# Create test_table:
+- name: postgresql_query - create test table called test_table
+  become_user: "{{ pg_user }}"
+  become: yes
+  shell: psql postgres -U "{{ pg_user }}" -t -c "CREATE TABLE test_table (id int, story text);"
+  ignore_errors: yes
+
+- name: postgresql_query - insert some data into test_table
+  become_user: "{{ pg_user }}"
+  become: yes
+  shell: psql postgres -U "{{ pg_user }}" -t -c "INSERT INTO test_table (id, story) VALUES (1, 'first'), (2, 'second'), (3, 'third');"
+  ignore_errors: yes
+
+# Prepare SQL script:
+- name: postgresql_query - remove SQL script if exists
+  become_user: "{{ pg_user }}"
+  become: yes
+  file:
+    path: '~{{ pg_user}}/test.sql'
+    state: absent
+  ignore_errors: yes
+
+- name: postgresql_query - prepare SQL script
+  become_user: "{{ pg_user }}"
+  become: yes
+  shell: 'echo "{{ item }}" >> ~{{ pg_user}}/test.sql'
+  ignore_errors: yes
+  with_items:
+  - SELECT version();
+  - SELECT story FROM test_table
+  - WHERE id = %s;
+
+##############
+# Start tests:
+#
+
+# Run ANALYZE command:
+- name: postgresql_query - analyze test_table
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: ANALYZE test_table
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == true
+    - result.query == 'ANALYZE test_table'
+    - result.rowcount == 0
+    - result.statusmessage == 'ANALYZE'
+    - result.query_result == {}
+
+# Run queries from SQL script:
+- name: postgresql_query - run queries from SQL script
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    path_to_script: '~{{ pg_user }}/test.sql'
+    positional_args:
+    - 1
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == false
+    - result.query == 'SELECT version();\nSELECT story FROM test_table\nWHERE id = 1;\n'
+    - result.rowcount == 1
+    - result.statusmessage == 'SELECT 1' or result.statusmessage == 'SELECT'
+    - result.query_result[0].story == 'first'
+
+# Simple select query:
+- name: postgresql_query - simple select query to test_table
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: SELECT * FROM test_table
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == false
+    - result.query == 'SELECT * FROM test_table'
+    - result.rowcount == 3
+    - result.statusmessage == 'SELECT 3' or result.statusmessage == 'SELECT'
+    - result.query_result[0].id == 1
+    - result.query_result[1].id == 2
+    - result.query_result[2].id == 3
+    - result.query_result[0].story == 'first'
+    - result.query_result[1].story == 'second'
+    - result.query_result[2].story == 'third'
+
+# Select query with named_args:
+- name: postgresql_query - select query with named args
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: SELECT id FROM test_table WHERE id = %(id_val)s AND story = %(story_val)s
+    named_args:
+      id_val: 1
+      story_val: first
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == false
+    - result.query == "SELECT id FROM test_table WHERE id = 1 AND story = 'first'" or result.query == "SELECT id FROM test_table WHERE id = 1 AND story = E'first'"
+    - result.rowcount == 1
+    - result.statusmessage == 'SELECT 1' or result.statusmessage == 'SELECT'
+    - result.query_result[0].id == 1
+
+# Select query with positional arguments:
+- name: postgresql_query - select query with positional arguments
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: SELECT story FROM test_table WHERE id = %s AND story = %s
+    positional_args:
+    - 2
+    - second
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == false
+    - result.query == "SELECT story FROM test_table WHERE id = 2 AND story = 'second'" or result.query == "SELECT story FROM test_table WHERE id = 2 AND story = E'second'"
+    - result.rowcount == 1
+    - result.statusmessage == 'SELECT 1' or result.statusmessage == 'SELECT'
+    - result.query_result[0].story == 'second'
+
+# Simple update query (positional_args and named args were checked by the previous tests):
+- name: postgresql_query - simple update query
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: UPDATE test_table SET story = 'new' WHERE id = 3
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == true
+    - result.query == "UPDATE test_table SET story = 'new' WHERE id = 3"
+    - result.rowcount == 1
+    - result.statusmessage == 'UPDATE 1'
+    - result.query_result == {}
+
+# Try to update not existing row:
+- name: postgresql_query - try to update not existing row
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: UPDATE test_table SET story = 'new' WHERE id = 100
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == false
+    - result.query == "UPDATE test_table SET story = 'new' WHERE id = 100"
+    - result.rowcount == 0
+    - result.statusmessage == 'UPDATE 0'
+    - result.query_result == {}
+
+# Simple insert query positional_args:
+- name: postgresql_query - insert query
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: INSERT INTO test_table (id, story) VALUES (%s, %s)
+    positional_args:
+    - 4
+    - fourth
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == true
+    - result.query == "INSERT INTO test_table (id, story) VALUES (4, 'fourth')" or result.query == "INSERT INTO test_table (id, story) VALUES (4, E'fourth')"
+    - result.rowcount == 1
+    - result.statusmessage == 'INSERT 0 1'
+    - result.query_result == {}
+
+# Truncate table:
+- name: postgresql_query - truncate test_table
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: TRUNCATE test_table
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == true
+    - result.query == "TRUNCATE test_table"
+    - result.rowcount == 0
+    - result.statusmessage == 'TRUNCATE TABLE'
+    - result.query_result == {}
+
+# Try DDL query:
+- name: postgresql_query - alter test_table
+  become_user: "{{ pg_user }}"
+  become: yes
+  postgresql_query:
+    login_user: "{{ pg_user }}"
+    db: postgres
+    query: ALTER TABLE test_table ADD COLUMN foo int
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result.changed == true
+    - result.query == "ALTER TABLE test_table ADD COLUMN foo int"
+    - result.rowcount == 0
+    - result.statusmessage == 'ALTER TABLE'