postgres modules: move params mapping from main to connect_to_db (#55549)

* params mapping to connect_to_db

lib/ansible/module_utils/postgres.py

@@ -29,23 +29,26 @@
 
 try:
     import psycopg2
-    import psycopg2.extras
+    from psycopg2.extras import DictCursor
     HAS_PSYCOPG2 = True
 except ImportError:
     HAS_PSYCOPG2 = False
 
+from ansible.module_utils.basic import missing_required_lib
 from ansible.module_utils._text import to_native
+from ansible.module_utils.six import iteritems
 
 
 class LibraryError(Exception):
     pass
 
 
-def ensure_libs(sslrootcert=None):
+def ensure_libs(module):
     if not HAS_PSYCOPG2:
-        raise LibraryError('psycopg2 is not installed. we need psycopg2.')
-    if sslrootcert and psycopg2.__version__ < '2.4.3':
-        raise LibraryError('psycopg2 must be at least 2.4.3 in order to use the ca_cert parameter')
+        module.fail_json(msg=missing_required_lib('psycopg2'))
+
+    if module.params.get('ca_cert') and psycopg2.__version__ < '2.4.3':
+        module.fail_json(msg='psycopg2 must be at least 2.4.3 in order to user the ca_cert parameter')
 
     # no problems
     return None
@@ -63,7 +66,42 @@ def postgres_common_argument_spec():
     )
 
 
-def connect_to_db(module, kw, autocommit=False):
+def connect_to_db(module, autocommit=False, fail_on_conn=True, warn_db_default=True):
+
+    ensure_libs(module)
+
+    # 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",
+        "ssl_mode": "sslmode",
+        "ca_cert": "sslrootcert"
+    }
+
+    # Might be different in the modules:
+    if module.params.get('db'):
+        params_map['db'] = 'database'
+    elif module.params.get('database'):
+        params_map['database'] = 'database'
+    elif module.params.get('login_db'):
+        params_map['login_db'] = 'database'
+    else:
+        if warn_db_default:
+            module.warn('Database name has not been passed, '
+                        'used default database to connect to.')
+
+    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"]
+
     try:
         db_connection = psycopg2.connect(**kw)
         if autocommit:
@@ -72,19 +110,31 @@ def connect_to_db(module, kw, autocommit=False):
             else:
                 db_connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
 
+        # Switch role, if specified:
+        cursor = db_connection.cursor(cursor_factory=DictCursor)
+        if module.params.get('session_role'):
+            try:
+                cursor.execute('SET ROLE %s' % module.params['session_role'])
+            except Exception as e:
+                module.fail_json(msg="Could not switch role: %s" % to_native(e))
+        cursor.close()
+
     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))
+        if fail_on_conn:
+            module.fail_json(msg="unable to connect to database: %s" % to_native(e))
+        else:
+            module.warn("PostgreSQL server is unavailable: %s" % to_native(e))
+            db_connection = None
 
     except Exception as e:
-        module.fail_json(msg="unable to connect to database: %s" % to_native(e))
+        if fail_on_conn:
+            module.fail_json(msg="unable to connect to database: %s" % to_native(e))
+        else:
+            module.warn("PostgreSQL server is unavailable: %s" % to_native(e))
+            db_connection = None
 
     return db_connection
-
-
-def get_pg_version(cursor):
-    cursor.execute("select current_setting('server_version_num')")
-    return int(cursor.fetchone()[0])