postgresql modules: move connect_to_db and get_pg_version to module_utils/postgres.py (#55514)

* postgresql modules: move connect_to_db, get_pg_version to module_utils

* postgresql modules: move connect_to_db, get_pg_version to module_utils 2

* postgresql modules: move connect_to_db, get_pg_version to module_utils 3

lib/ansible/module_utils/postgres.py

@@ -34,6 +34,8 @@ try:
 except ImportError:
     HAS_PSYCOPG2 = False
 
+from ansible.module_utils._text import to_native
+
 
 class LibraryError(Exception):
     pass
@@ -59,3 +61,30 @@ def postgres_common_argument_spec():
         ssl_mode=dict(default='prefer', choices=['allow', 'disable', 'prefer', 'require', 'verify-ca', 'verify-full']),
         ca_cert=dict(aliases=['ssl_rootcert']),
     )
+
+
+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
+
+
+def get_pg_version(cursor):
+    cursor.execute("select current_setting('server_version_num')")
+    return int(cursor.fetchone()[0])

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

@@ -145,7 +145,7 @@ except ImportError:
     HAS_PSYCOPG2 = False
 
 from ansible.module_utils.basic import AnsibleModule, missing_required_lib
-from ansible.module_utils.postgres import postgres_common_argument_spec
+from ansible.module_utils.postgres import connect_to_db, postgres_common_argument_spec
 from ansible.module_utils.six import iteritems
 from ansible.module_utils._text import to_native
 from ansible.module_utils.database import pg_quote_identifier
@@ -248,26 +248,9 @@ def main():
     if psycopg2.__version__ < '2.4.3' and sslrootcert is not None:
         module.fail_json(msg='psycopg2 must be at least 2.4.3 in order to user the ca_cert parameter')
 
-    try:
+    db_connection = connect_to_db(module, kw, autocommit=True)
-        db_connection = psycopg2.connect(**kw)
-        # Enable autocommit so we can create databases
-        if psycopg2.__version__ >= '2.4.2':
-            db_connection.autocommit = True
-        else:
-            db_connection.set_isolation_level(psycopg2
-                                              .extensions
-                                              .ISOLATION_LEVEL_AUTOCOMMIT)
     cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
 
-    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), exception=traceback.format_exc())
-
-    except Exception as e:
-        module.fail_json(msg="unable to connect to database: %s" % to_native(e), exception=traceback.format_exc())
-
     if session_role:
         try:
             cursor.execute('SET ROLE %s' % pg_quote_identifier(session_role, 'role'))

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

@@ -241,7 +241,7 @@ except ImportError:
 
 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.postgres import connect_to_db, postgres_common_argument_spec
 from ansible.module_utils._text import to_native
 from ansible.module_utils.six import iteritems
 
@@ -485,22 +485,8 @@ def main():
     if psycopg2.__version__ < '2.4.3' and sslrootcert is not None:
         module.fail_json(msg='psycopg2 must be at least 2.4.3 in order to user the ca_cert parameter')
 
-    try:
+    db_connection = connect_to_db(module, kw, autocommit=True)
-        db_connection = psycopg2.connect(**kw)
-        if concurrent:
-            if psycopg2.__version__ >= '2.4.2':
-                db_connection.set_session(autocommit=True)
-            else:
-                db_connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
-
     cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
-    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))
 
     if session_role:
         try:

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

@@ -180,7 +180,7 @@ except ImportError:
     HAS_PSYCOPG2 = False
 
 from ansible.module_utils.basic import AnsibleModule, missing_required_lib
-from ansible.module_utils.postgres import postgres_common_argument_spec
+from ansible.module_utils.postgres import connect_to_db, postgres_common_argument_spec
 from ansible.module_utils.six import iteritems
 from ansible.module_utils._text import to_native
 from ansible.module_utils.database import pg_quote_identifier
@@ -293,18 +293,9 @@ def main():
     if psycopg2.__version__ < '2.4.3' and sslrootcert is not None:
         module.fail_json(msg='psycopg2 must be at least 2.4.3 in order to user the ca_cert parameter')
 
-    try:
+    db_connection = connect_to_db(module, kw, autocommit=False)
-        db_connection = psycopg2.connect(**kw)
     cursor = db_connection.cursor()
 
-    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), exception=traceback.format_exc())
-
-    except Exception as e:
-        module.fail_json(msg="unable to connect to database: %s" % to_native(e), exception=traceback.format_exc())
-
     if session_role:
         try:
             cursor.execute('SET ROLE %s' % pg_quote_identifier(session_role, 'role'))

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

@@ -143,33 +143,11 @@ except ImportError:
 
 from ansible.module_utils.basic import AnsibleModule, missing_required_lib
 from ansible.module_utils.database import SQLParseError, pg_quote_identifier
-from ansible.module_utils.postgres import postgres_common_argument_spec
+from ansible.module_utils.postgres import connect_to_db, 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
-
-
 class PgMembership(object):
     def __init__(self, module, cursor, groups, target_roles, fail_on_role):
         self.module = module

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

@@ -157,33 +157,11 @@ except ImportError:
 
 from ansible.module_utils.basic import AnsibleModule, missing_required_lib
 from ansible.module_utils.database import SQLParseError, pg_quote_identifier
-from ansible.module_utils.postgres import postgres_common_argument_spec
+from ansible.module_utils.postgres import connect_to_db, 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
-
-
 class PgOwnership(object):
     def __init__(self, module, cursor, role):
         self.module = module

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

@@ -147,32 +147,11 @@ except ImportError:
 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.postgres import connect_to_db, 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.
 #

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

@@ -131,7 +131,7 @@ except ImportError:
     HAS_PSYCOPG2 = False
 
 from ansible.module_utils.basic import AnsibleModule, missing_required_lib
-from ansible.module_utils.postgres import postgres_common_argument_spec
+from ansible.module_utils.postgres import connect_to_db, postgres_common_argument_spec
 from ansible.module_utils.database import SQLParseError, pg_quote_identifier
 from ansible.module_utils.six import iteritems
 from ansible.module_utils._text import to_native
@@ -266,26 +266,8 @@ def main():
         module.fail_json(
             msg='psycopg2 must be at least 2.4.3 in order to user the ca_cert parameter')
 
-    try:
+    db_connection = connect_to_db(module, kw, autocommit=True)
-        db_connection = psycopg2.connect(**kw)
+    cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
-        # Enable autocommit so we can create databases
-        if psycopg2.__version__ >= '2.4.2':
-            db_connection.autocommit = True
-        else:
-            db_connection.set_isolation_level(psycopg2
-                                              .extensions
-                                              .ISOLATION_LEVEL_AUTOCOMMIT)
-        cursor = db_connection.cursor(
-            cursor_factory=psycopg2.extras.DictCursor)
-
-    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), exception=traceback.format_exc())
-
-    except Exception as e:
-        module.fail_json(msg="unable to connect to database: %s" % to_native(e), exception=traceback.format_exc())
 
     if session_role:
         try:

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

@@ -167,7 +167,7 @@ except ImportError:
 
 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.postgres import connect_to_db, get_pg_version, postgres_common_argument_spec
 from ansible.module_utils.six import iteritems
 from ansible.module_utils._text import to_native
 
@@ -271,21 +271,6 @@ def param_set(cursor, module, name, value, context):
     return True
 
 
-def connect_to_db(module, kw, autocommit=False):
-    try:
-        db_connection = psycopg2.connect(**kw)
-        if autocommit:
-            db_connection.set_session(autocommit=True)
-
-    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.
 #
@@ -357,8 +342,7 @@ def main():
     cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
 
     # Check server version (needs 9.4 or later):
-    cursor.execute("select current_setting('server_version_num')")
+    ver = get_pg_version(cursor)
-    ver = int(cursor.fetchone()[0])
     if ver < PG_REQ_VER:
         module.warn("PostgreSQL is %s version but %s or later is required" % (ver, PG_REQ_VER))
         kw = dict(

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

@@ -149,37 +149,10 @@ except ImportError:
 
 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.postgres import connect_to_db, get_pg_version, postgres_common_argument_spec
 from ansible.module_utils._text import to_native
 
 
-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
-
-
-def get_pg_version(cursor):
-    cursor.execute("select current_setting('server_version_num')")
-    return int(cursor.fetchone()[0])
-
-
 # ===========================================
 # PostgreSQL module specific support methods.
 #

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

@@ -218,7 +218,7 @@ except ImportError:
 
 from ansible.module_utils.basic import AnsibleModule, missing_required_lib
 from ansible.module_utils.database import SQLParseError, pg_quote_identifier
-from ansible.module_utils.postgres import postgres_common_argument_spec
+from ansible.module_utils.postgres import connect_to_db, postgres_common_argument_spec
 from ansible.module_utils._text import to_native
 from ansible.module_utils.six import iteritems
 
@@ -525,15 +525,8 @@ def main():
     if psycopg2.__version__ < '2.4.3' and sslrootcert is not None:
         module.fail_json(msg='psycopg2 must be at least 2.4.3 in order to user the ca_cert parameter')
 
-    try:
+    db_connection = connect_to_db(module, kw, autocommit=False)
-        db_connection = psycopg2.connect(**kw)
     cursor = db_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
-    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))
 
     if session_role:
         try:

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

@@ -177,33 +177,11 @@ except ImportError:
 
 from ansible.module_utils.basic import AnsibleModule, missing_required_lib
 from ansible.module_utils.database import SQLParseError, pg_quote_identifier
-from ansible.module_utils.postgres import postgres_common_argument_spec
+from ansible.module_utils.postgres import connect_to_db, 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
-
-
 class PgTablespace(object):
     def __init__(self, module, cursor, name):
         self.module = module