Merge pull request #2446 from b6d/postgresql_db-module-extension

postgresql_db: add args lc_collate, lc_ctype

library/postgresql_db

@@ -64,6 +64,16 @@ options:
       - Encoding of the database
     required: false
     default: null
+  lc_collate:
+    description:
+      - Collation order (LC_COLLATE) to use in the database. Must match collation order of template database unless C(template0) is used as template.
+    required: false
+    default: null
+  lc_ctype:
+    description:
+      - Character classification (LC_CTYPE) to use in the database (e.g. lower, upper, ...) Must match LC_CTYPE of template database unless C(template0) is used as template.
+    required: false
+    default: null
   state:
     description:
       - The database state
@@ -73,6 +83,8 @@ options:
 examples:
    - code: "postgresql_db: db=acme"
      description: Create a new database with name C(acme)
+   - code: "postgresql_db: db=acme encoding='UTF-8' lc_collate='de_DE.UTF-8' lc_ctype='de_DE.UTF-8' template='template0'"
+     description: Create a new database with name C(acme) and specific encoding and locale settings. If a template different from C(template0) is specified, encoding and locale settings must match those of the template.
 notes:
    - The default authentication assumes that you are either logging in as or sudo'ing to the C(postgres) account on the host.
    - This module uses I(psycopg2), a Python PostgreSQL database adapter. You must ensure that psycopg2 is installed on
@@ -83,11 +95,16 @@ author: Lorin Hochstein
 
 try:
     import psycopg2
+    import psycopg2.extras
 except ImportError:
     postgresqldb_found = False
 else:
     postgresqldb_found = True
 
+class NotSupportedError(Exception):
+    pass
+
+
 # ===========================================
 # PostgreSQL module specific support methods.
 #
@@ -97,11 +114,21 @@ def set_owner(cursor, db, owner):
     cursor.execute(query)
     return True
 
-def db_owned_by(cursor, db, user):
-    query = """SELECT * FROM pg_database JOIN pg_user ON datdba = usesysid
-    WHERE usename = %(user)s and datname = %(db)s"""
-    cursor.execute(query, {'db':db, 'user':user})
-    return cursor.rowcount == 1
+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 usename AS owner,
+    pg_encoding_to_char(encoding) AS encoding, encoding AS encoding_id,
+    datcollate AS lc_collate, datctype AS lc_ctype
+    FROM pg_database JOIN pg_user ON pg_user.usesysid = pg_database.datdba
+    WHERE datname = %(db)s
+    """
+    cursor.execute(query, {'db':db})
+    return cursor.fetchone()
 
 def db_exists(cursor, db):
     query = "SELECT * FROM pg_database WHERE datname=%(db)s"
@@ -116,7 +143,7 @@ def db_delete(cursor, db):
     else:
         return False
 
-def db_create(cursor, db, owner, template, encoding):
+def db_create(cursor, db, owner, template, encoding, lc_collate, lc_ctype):
     if not db_exists(cursor, db):
         if owner:
             owner = " OWNER \"%s\"" % owner
@@ -124,13 +151,37 @@ def db_create(cursor, db, owner, template, encoding):
             template = " TEMPLATE \"%s\"" % template
         if encoding:
             encoding = " ENCODING '%s'" % encoding
-        query = "CREATE DATABASE \"%s\"%s%s%s" % (db, owner, template, encoding)
+        if lc_collate:
+            lc_collate = " LC_COLLATE '%s'" % lc_collate
+        if lc_ctype:
+            lc_ctype = " LC_CTYPE '%s'" % lc_ctype
+        query = 'CREATE DATABASE "%s"%s%s%s%s%s' % (db, owner,
+                                                    template, encoding,
+                                                    lc_collate, lc_ctype)
         cursor.execute(query)
         return True
-    elif owner and not db_owned_by(cursor, db, owner):
-        return set_owner(cursor, db, owner)
     else:
-        return False
+        db_info = get_db_info(cursor, db)
+        if (encoding and 
+            get_encoding_id(cursor, encoding) != db_info['encoding_id']):
+            raise NotSupportedError(
+                'Changing database encoding is not supported. '
+                'Current encoding: %s' % db_info['encoding']
+            )
+        elif lc_collate and lc_collate != db_info['lc_collate']:
+            raise NotSupportedError(
+                'Changing LC_COLLATE is not supported. '
+                'Current LC_COLLATE: %s' % db_info['lc_collate']
+            )
+        elif lc_ctype and lc_ctype != db_info['lc_ctype']:
+            raise NotSupportedError(
+                'Changing LC_CTYPE is not supported.'
+                'Current LC_CTYPE: %s' % db_info['lc_ctype']
+            )
+        elif owner and owner != db_info['owner']:
+            return set_owner(cursor, db, owner)
+        else:
+            return False
 
 # ===========================================
 # Module execution.
@@ -147,6 +198,8 @@ def main():
             owner=dict(default=""),
             template=dict(default=""),
             encoding=dict(default=""),
+            lc_collate=dict(default=""),
+            lc_ctype=dict(default=""),
             state=dict(default="present", choices=["absent", "present"]),
         ),
         supports_check_mode = True
@@ -160,6 +213,8 @@ def main():
     owner = module.params["owner"]
     template = module.params["template"]
     encoding = module.params["encoding"]
+    lc_collate = module.params["lc_collate"]
+    lc_ctype = module.params["lc_ctype"]
     state = module.params["state"]
     changed = False
 
@@ -183,7 +238,8 @@ def main():
             db_connection.set_isolation_level(psycopg2
                                               .extensions
                                               .ISOLATION_LEVEL_AUTOCOMMIT)
-        cursor = db_connection.cursor()
+        cursor = db_connection.cursor(
+                cursor_factory=psycopg2.extras.DictCursor)
     except Exception, e:
         module.fail_json(msg="unable to connect to database: %s" % e)
 
@@ -195,8 +251,10 @@ def main():
             changed = db_delete(cursor, db)
 
         elif state == "present":
-            changed = db_create(cursor, db, owner, template, encoding)
-
+            changed = db_create(cursor, db, owner, template, encoding,
+                                lc_collate, lc_ctype)
+    except NotSupportedError, e:
+        module.fail_json(msg=str(e))
     except Exception, e:
         module.fail_json(msg="Database query failed: %s" % e)