From b5f1fc850b841b2e8909b00af8762e1e25c8ecab Mon Sep 17 00:00:00 2001 From: Abhijeet Kasurde Date: Sun, 5 Nov 2017 16:14:40 +0530 Subject: [PATCH] Allow user to specify maintenance DB (#32510) This fix allows user to specify alternative maintenance DB required for initial connection in Postgresql_db module. Also, adds pep8 related fixes. Fixes: #30017 Signed-off-by: Abhijeet Kasurde --- .../database/postgresql/postgresql_db.py | 56 ++++++++++++------- test/sanity/pep8/legacy-files.txt | 1 - 2 files changed, 36 insertions(+), 21 deletions(-) diff --git a/lib/ansible/modules/database/postgresql/postgresql_db.py b/lib/ansible/modules/database/postgresql/postgresql_db.py index f9b7a675086..8dec1efde7f 100644 --- a/lib/ansible/modules/database/postgresql/postgresql_db.py +++ b/lib/ansible/modules/database/postgresql/postgresql_db.py @@ -72,6 +72,11 @@ options: version_added: "2.4" description: - Further arguments for pg_dump or pg_restore. Used when state is "dump" or "restore" + maintenance_db: + version_added: "2.5" + description: + - The value specifies the initial database (which is also called as maintenance DB) that Ansible connects to. + default: postgres author: "Ansible Core Team" extends_documentation_fragment: - postgres @@ -112,18 +117,16 @@ EXAMPLES = ''' target_opts: "-n public" ''' +import os +import pipes +import subprocess import traceback -HAS_PSYCOPG2 = False try: import psycopg2 import psycopg2.extras - import pipes - import subprocess - import os - except ImportError: - pass + HAS_PSYCOPG2 = False else: HAS_PSYCOPG2 = True @@ -149,11 +152,13 @@ def set_owner(cursor, db, owner): cursor.execute(query) return True + def get_encoding_id(cursor, encoding): query = "SELECT pg_char_to_encoding(%(encoding)s) AS encoding_id;" cursor.execute(query, {'encoding': encoding}) return cursor.fetchone()['encoding_id'] + def get_db_info(cursor, db): query = """ SELECT rolname AS owner, @@ -165,11 +170,13 @@ def get_db_info(cursor, db): cursor.execute(query, {'db': db}) return cursor.fetchone() + def db_exists(cursor, db): query = "SELECT * FROM pg_database WHERE datname=%(db)s" cursor.execute(query, {'db': db}) return cursor.rowcount == 1 + def db_delete(cursor, db): if db_exists(cursor, db): query = "DROP DATABASE %s" % pg_quote_identifier(db, 'database') @@ -178,6 +185,7 @@ def db_delete(cursor, db): else: return False + def db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype): params = dict(enc=encoding, collate=lc_collate, ctype=lc_ctype) if not db_exists(cursor, db): @@ -218,6 +226,7 @@ def db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype): else: return False + def db_matches(cursor, db, owner, template, encoding, lc_collate, lc_ctype): if not db_exists(cursor, db): return False @@ -235,6 +244,7 @@ def db_matches(cursor, db, owner, template, encoding, lc_collate, lc_ctype): else: return True + def db_dump(module, target, target_opts="", db=None, user=None, @@ -270,13 +280,14 @@ def db_dump(module, target, target_opts="", return do_with_password(module, cmd, password) + def db_restore(module, target, target_opts="", - db=None, - user=None, - password=None, - host=None, - port=None, - **kw): + db=None, + user=None, + password=None, + host=None, + port=None, + **kw): flags = login_flags(db, host, port, user) comp_prog_path = None @@ -321,6 +332,7 @@ def db_restore(module, target, target_opts="", return do_with_password(module, cmd, password) + def login_flags(db, host, port, user, db_prefix=True): """ returns a list of connection argument strings each prefixed @@ -344,6 +356,7 @@ def login_flags(db, host, port, user, db_prefix=True): flags.append(' --username={0}'.format(user)) return flags + def do_with_password(module, cmd, password): env = {} if password: @@ -355,6 +368,7 @@ def do_with_password(module, cmd, password): # Module execution. # + def main(): argument_spec = pgutils.postgres_common_argument_spec() argument_spec.update(dict( @@ -367,6 +381,7 @@ def main(): state=dict(default="present", choices=["absent", "present", "dump", "restore"]), target=dict(default="", type="path"), target_opts=dict(default=""), + maintenance_db=dict(default="postgres"), )) module = AnsibleModule( @@ -387,19 +402,20 @@ def main(): target_opts = module.params["target_opts"] state = module.params["state"] changed = False + maintenance_db = module.params['maintenance_db'] # 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", - "ssl_rootcert":"sslrootcert" + "login_host": "host", + "login_user": "user", + "login_password": "password", + "port": "port", + "ssl_mode": "sslmode", + "ssl_rootcert": "sslrootcert" } - kw = dict( (params_map[k], v) for (k, v) in iteritems(module.params) + 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. @@ -414,7 +430,7 @@ def main(): try: pgutils.ensure_libs(sslrootcert=module.params.get('ssl_rootcert')) - db_connection = psycopg2.connect(database="postgres", **kw) + db_connection = psycopg2.connect(database=maintenance_db, **kw) # Enable autocommit so we can create databases if psycopg2.__version__ >= '2.4.2': diff --git a/test/sanity/pep8/legacy-files.txt b/test/sanity/pep8/legacy-files.txt index 020a122e6d1..bceb65db4bb 100644 --- a/test/sanity/pep8/legacy-files.txt +++ b/test/sanity/pep8/legacy-files.txt @@ -128,7 +128,6 @@ lib/ansible/modules/database/misc/riak.py lib/ansible/modules/database/mongodb/mongodb_parameter.py lib/ansible/modules/database/mongodb/mongodb_user.py lib/ansible/modules/database/mssql/mssql_db.py -lib/ansible/modules/database/postgresql/postgresql_db.py lib/ansible/modules/database/postgresql/postgresql_ext.py lib/ansible/modules/database/postgresql/postgresql_lang.py lib/ansible/modules/database/postgresql/postgresql_schema.py