Merge pull request #702 from lorin/postgresql-module

Add postgresql_db and postgresql_user module.

library/postgresql_db

@@ -0,0 +1,94 @@
+#!/usr/bin/python
+
+# This file is part of Ansible
+#
+# Ansible is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Ansible is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
+
+try:
+    import psycopg2
+except ImportError:
+    postgresqldb_found = False
+else:
+    postgresqldb_found = True
+
+# ===========================================
+# PostgreSQL module specific support methods.
+#
+
+
+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):
+    query = "DROP DATABASE %s" % db
+    cursor.execute(query)
+    return True
+
+
+def db_create(cursor, db):
+    query = "CREATE DATABASE %s" % db
+    cursor.execute(query)
+    return True
+
+# ===========================================
+# Module execution.
+#
+
+
+def main():
+    module = AnsibleModule(
+        argument_spec=dict(
+            loginuser=dict(default="postgres"),
+            loginpass=dict(default=""),
+            loginhost=dict(default=""),
+            db=dict(required=True),
+            state=dict(default="present", choices=["absent", "present"]),
+        )
+    )
+
+    if not postgresqldb_found:
+        module.fail_json(msg="the python psycopg2 module is required")
+
+    db = module.params["db"]
+    state = module.params["state"]
+    changed = False
+    try:
+        db_connection = psycopg2.connect(host=module.params["loginhost"],
+                                         user=module.params["loginuser"],
+                                         password=module.params["loginpass"],
+                                         database="template1")
+        # Enable autocommit so we can create databases
+        db_connection.autocommit = True
+        cursor = db_connection.cursor()
+    except Exception as e:
+        module.fail_json(msg="unable to connect to database: %s" % e)
+
+    try:
+        if db_exists(cursor, db):
+            if state == "absent":
+                changed = db_delete(cursor, db)
+        else:
+            if state == "present":
+                changed = db_create(cursor, db)
+    except Exception as e:
+        module.fail_json(msg="Database query failed: %s" % e)
+
+    module.exit_json(changed=changed, db=db)
+
+# this is magic, see lib/ansible/module_common.py
+#<<INCLUDE_ANSIBLE_MODULE_COMMON>>
+main()

library/postgresql_user

@@ -0,0 +1,152 @@
+#!/usr/bin/python
+
+# This file is part of Ansible
+#
+# Ansible is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Ansible is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
+
+try:
+    import psycopg2
+except ImportError:
+    postgresqldb_found = False
+else:
+    postgresqldb_found = True
+
+# ===========================================
+# PostgreSQL module specific support methods.
+#
+
+
+def user_exists(cursor, user):
+    query = "SELECT rolname FROM pg_roles WHERE rolname=%(user)s"
+    cursor.execute(query, {'user': user})
+    return cursor.rowcount > 0
+
+
+def user_add(cursor, user, passwd, db):
+    """Create a new user with write access to the database"""
+    query = "CREATE USER %(user)s with PASSWORD '%(passwd)s'"
+    cursor.execute(query % {"user": user, "passwd": passwd})
+    grant_privileges(cursor, user, db)
+    return True
+
+
+def has_privileges(cursor, user, db):
+    """Check if the user has create privileges on the database"""
+    query = "SELECT has_database_privilege(%(user)s, %(db)s, 'CREATE')"
+    cursor.execute(query, {'user': user, 'db': db})
+    return cursor.fetchone()[0]
+
+
+def grant_privileges(cursor, user, db):
+    """Grant all privileges on the database"""
+    query = "GRANT ALL PRIVILEGES ON DATABASE %(db)s TO %(user)s"
+    cursor.execute(query % {'user': user, 'db': db})
+
+
+def revoke_privileges(cursor, user, db):
+    """Revoke all privileges on the database"""
+    query = "REVOKE ALL PRIVILEGES ON DATABASE %(db)s FROM %(user)s"
+    cursor.execute(query % {'user': user, 'db': db})
+
+
+def user_mod(cursor, user, passwd, db):
+    """Update password and permissions"""
+    changed = False
+
+    # Handle passwords.
+    if passwd is not None:
+        select = "SELECT rolpassword FROM pg_authid where rolname=%(user)s"
+        cursor.execute(select, {"user": user})
+        current_pass_hash = cursor.fetchone()[0]
+        # Not sure how to hash the new password, so we just initiate the
+        # change and check if the hash changed
+        alter = "ALTER USER %(user)s WITH PASSWORD '%(passwd)s'"
+        cursor.execute(alter % {"user": user, "passwd": passwd})
+        cursor.execute(select, {"user": user})
+        new_pass_hash = cursor.fetchone()[0]
+        if current_pass_hash != new_pass_hash:
+            changed = True
+
+    # Handle privileges.
+    # For now, we just check if the user has access to the database
+    if not has_privileges(cursor, user, db):
+        grant_privileges(cursor, user, db)
+        changed = True
+
+    return changed
+
+
+def user_delete(cursor, user, db):
+    """Delete a user, first revoking privileges"""
+    revoke_privileges(cursor, user, db)
+    cursor.execute("DROP USER %(user)s" % {'user': user})
+    return True
+
+
+
+# ===========================================
+# Module execution.
+#
+
+
+def main():
+    module = AnsibleModule(
+        argument_spec=dict(
+            loginuser=dict(default="postgres"),
+            loginpass=dict(default=""),
+            loginhost=dict(default=""),
+            user=dict(required=True),
+            passwd=dict(default=None),
+            state=dict(default="present", choices=["absent", "present"]),
+            db=dict(required=True),
+        )
+    )
+    user = module.params["user"]
+    passwd = module.params["passwd"]
+    state = module.params["state"]
+    db = module.params["db"]
+
+    if not postgresqldb_found:
+        module.fail_json(msg="the python psycopg2 module is required")
+
+    try:
+        db_connection = psycopg2.connect(host=module.params["loginhost"],
+                                         user=module.params["loginuser"],
+                                         password=module.params["loginpass"],
+                                         database=db)
+        cursor = db_connection.cursor()
+    except Exception as e:
+        module.fail_json(msg="unable to connect to database: %s" % e)
+
+    if state == "present":
+        if user_exists(cursor, user):
+            changed = user_mod(cursor, user, passwd, db)
+        else:
+            if passwd is None:
+                msg = "passwd parameter required when adding a user"
+                module.fail_json(msg=msg)
+            changed = user_add(cursor, user, passwd, db)
+
+    elif state == "absent":
+        if user_exists(cursor, user):
+            changed = user_delete(cursor, user, db)
+        else:
+            changed = False
+    # Commit the database changes
+    db_connection.commit()
+    module.exit_json(changed=changed, user=user)
+
+# this is magic, see lib/ansible/module_common.py
+#<<INCLUDE_ANSIBLE_MODULE_COMMON>>
+main()