From 09d68678ee4c950db655bc15b358f3001e61b963 Mon Sep 17 00:00:00 2001
From: Andrew Klychkov <aaklychkov@mail.ru>
Date: Mon, 20 Jul 2020 10:08:02 +0300
Subject: [PATCH] postgresql_query: add search_path parameter (#653)

* postgresql_query: add search_path parameter

* add CI tests

* add ref to seealso

* add changelog fragment

* fix test syntax

* fix test syntax

* fix

* fix

* fix CI syntax

* cosmetic change

* improve CI test

* move CI tests to the right place

* improve CI
---
 ...postgresql_query_add_search_path_param.yml |  2 +
 .../database/postgresql/postgresql_query.py   | 34 +++++++++
 .../tasks/postgresql_query_initial.yml        | 71 +++++++++++++++++++
 3 files changed, 107 insertions(+)
 create mode 100644 changelogs/fragments/653-postgresql_query_add_search_path_param.yml

diff --git a/changelogs/fragments/653-postgresql_query_add_search_path_param.yml b/changelogs/fragments/653-postgresql_query_add_search_path_param.yml
new file mode 100644
index 0000000000..72ba15b8ab
--- /dev/null
+++ b/changelogs/fragments/653-postgresql_query_add_search_path_param.yml
@@ -0,0 +1,2 @@
+minor_changes:
+- postgresql_query - add search_path parameter (https://github.com/ansible-collections/community.general/issues/625).
diff --git a/plugins/modules/database/postgresql/postgresql_query.py b/plugins/modules/database/postgresql/postgresql_query.py
index a4c0a408d0..7b03621083 100644
--- a/plugins/modules/database/postgresql/postgresql_query.py
+++ b/plugins/modules/database/postgresql/postgresql_query.py
@@ -75,8 +75,17 @@ options:
     type: bool
     default: yes
     version_added: '0.2.0'
+  search_path:
+    description:
+    - List of schema names to look in.
+    type: list
+    elements: str
+    version_added: '1.0.0'
 seealso:
 - module: community.general.postgresql_db
+- name: PostgreSQL Schema reference
+  description: Complete reference of the PostgreSQL schema documentation.
+  link: https://www.postgresql.org/docs/current/ddl-schemas.html
 author:
 - Felix Archambault (@archf)
 - Andrew Klychkov (@Andersson007)
@@ -152,6 +161,16 @@ EXAMPLES = r'''
     positional_args:
     - '{{ my_list }}'
     - '{{ my_arr|string }}'
+
+# Select from test table looking into app1 schema first, then,
+# if the schema doesn't exist or the table hasn't been found there,
+# try to find it in the schema public
+- name: Select from test using search_path
+  community.general.postgresql_query:
+    query: SELECT * FROM test_array_table
+    search_path:
+    - app1
+    - public
 '''
 
 RETURN = r'''
@@ -242,6 +261,16 @@ def convert_elements_to_pg_arrays(obj):
     return obj
 
 
+def set_search_path(cursor, search_path):
+    """Set session's search_path.
+
+    Args:
+        cursor (Psycopg2 cursor): Database cursor object.
+        search_path (str): String containing comma-separated schema names.
+    """
+    cursor.execute('SET search_path TO %s' % search_path)
+
+
 def main():
     argument_spec = postgres_common_argument_spec()
     argument_spec.update(
@@ -254,6 +283,7 @@ def main():
         autocommit=dict(type='bool', default=False),
         encoding=dict(type='str'),
         trust_input=dict(type='bool', default=True),
+        search_path=dict(type='list', elements='str'),
     )
 
     module = AnsibleModule(
@@ -270,6 +300,7 @@ def main():
     encoding = module.params["encoding"]
     session_role = module.params["session_role"]
     trust_input = module.params["trust_input"]
+    search_path = module.params["search_path"]
 
     if not trust_input:
         # Check input for potentially dangerous elements:
@@ -300,6 +331,9 @@ def main():
         db_connection.set_client_encoding(encoding)
     cursor = db_connection.cursor(cursor_factory=DictCursor)
 
+    if search_path:
+        set_search_path(cursor, '%s' % ','.join([x.strip(' ') for x in search_path]))
+
     # Prepare args:
     if module.params.get("positional_args"):
         arguments = module.params["positional_args"]
diff --git a/tests/integration/targets/postgresql_query/tasks/postgresql_query_initial.yml b/tests/integration/targets/postgresql_query/tasks/postgresql_query_initial.yml
index 5b75e74ebb..68a22b4e26 100644
--- a/tests/integration/targets/postgresql_query/tasks/postgresql_query_initial.yml
+++ b/tests/integration/targets/postgresql_query/tasks/postgresql_query_initial.yml
@@ -457,3 +457,74 @@
     name: test_array_table
     state: absent
   when: postgres_version_resp.stdout is version('9.4', '>=')
+
+#############################
+# Check search_path parameter
+
+- name: postgresql_set - create test schemas
+  become_user: '{{ pg_user }}'
+  become: true
+  postgresql_schema:
+    login_user: '{{ pg_user }}'
+    login_db: postgres
+    name: '{{ item }}'
+  loop:
+  - query_test1
+  - query_test2
+
+- name: postgresql_set - create test tables
+  become_user: '{{ pg_user }}'
+  become: true
+  postgresql_table:
+    login_user: '{{ pg_user }}'
+    login_db: postgres
+    name: '{{ item }}'
+    columns:
+    - id int
+  loop:
+  - 'query_test1.test1'
+  - 'query_test2.test2'
+
+- name: postgresql_query - insert data
+  become_user: '{{ pg_user }}'
+  become: true
+  postgresql_query:
+    login_user: '{{ pg_user }}'
+    login_db: postgres
+    query: 'INSERT INTO {{ item }} (id) VALUES (1)'
+    search_path:
+    - query_test1
+    - query_test2
+  loop:
+  - test1
+  - test2
+
+- name: postgresql_query - get data
+  become_user: '{{ pg_user }}'
+  become: true
+  postgresql_query:
+    login_user: '{{ pg_user }}'
+    login_db: postgres
+    query: 'SELECT id FROM test1'
+    search_path:
+    - query_test1
+    - query_test2
+  register: result
+
+- assert:
+    that:
+    - result.rowcount == 1
+
+- name: postgresql_query - get data, must fail
+  become_user: '{{ pg_user }}'
+  become: true
+  postgresql_query:
+    login_user: '{{ pg_user }}'
+    login_db: postgres
+    query: 'SELECT id FROM test1'
+  register: result
+  ignore_errors: yes
+
+- assert:
+    that:
+    - result is failed