From 511ab8697c4af4eefef00d638c56502d66938245 Mon Sep 17 00:00:00 2001
From: Pepe Barbe <dev@antropoide.net>
Date: Tue, 21 Aug 2012 14:25:19 -0500
Subject: [PATCH] Query for all active privileges instead Use a different
 method to query for current privileges at the table and database level. This
 method is more robust if newer privileges are added in future versions and
 also supports the ALL wildcard.

---
 postgresql_user | 72 ++++++++++++++++++++++++++++++++++---------------
 1 file changed, 51 insertions(+), 21 deletions(-)

diff --git a/postgresql_user b/postgresql_user
index dbef34e0ec5..8bde46aedcb 100755
--- a/postgresql_user
+++ b/postgresql_user
@@ -16,6 +16,8 @@
 # You should have received a copy of the GNU General Public License
 # along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
 
+import re
+
 try:
     import psycopg2
 except ImportError:
@@ -75,42 +77,70 @@ def user_delete(cursor, user):
 
 def has_table_privilege(cursor, user, table, priv):
     query = 'SELECT has_table_privilege(%s, %s, %s)'
-    cursor.execute(query, user, table, priv)
+    cursor.execute(query, (user, table, priv))
     return cursor.fetchone()[0]
 
+def get_table_privileges(cursor, user, table):
+    if '.' in table:
+        schema, table = table.split('.', 1)
+    else:
+        schema = 'public'
+    query = '''SELECT privilege_type FROM information_schema.role_table_grants
+    WHERE grantee=%s AND table_name=%s AND table_schema=%s'''
+    cursor.execute(query, (user, table, schema))
+    return set([x[0] for x in cursor.fetchall()])
+    
+
 def grant_table_privilege(cursor, user, table, priv):
-    if has_table_privilege(cursor, user, table, priv):
-        return False
+    prev_priv = get_table_privileges(cursor, user, table)
     query = 'GRANT %s ON TABLE %s TO %s' % (priv, table, user)
     cursor.execute(query)
-    return True
+    curr_priv = get_table_privileges(cursor, user, table)
+    return len(curr_priv) > len(prev_priv)
 
 def revoke_table_privilege(cursor, user, table, priv):
-    if not has_table_privilege(cursor, user, table, priv):
-        return False
+    prev_priv = get_table_privileges(cursor, user, table)
     query = 'REVOKE %s ON TABLE %s FROM %s' % (priv, table, user)
     cursor.execute(query)
-    return True
+    curr_priv = get_table_privileges(cursor, user, table)
+    return len(curr_priv) < len(prev_priv)
+
 
+def get_database_privileges(cursor, user, db):
+    priv_map = {
+        'C':'CREATE',
+        'T':'TEMPORARY',
+        'c':'CONNECT',
+    }
+    query = 'SELECT datacl FROM pg_database WHERE datname = %s'
+    cursor.execute(query, (db,))
+    datacl = cursor.fetchone()[0]
+    r = re.search('%s=(C?T?c?)/[a-z]+\,?' % user, datacl)
+    if r is None:
+        return []
+    o = []
+    for v in r.group(1):
+        o.append(priv_map[v])
+    return o
 
 def has_database_privilege(cursor, user, db, priv):
     query = 'SELECT has_database_privilege(%s, %s, %s)'
-    cursor.execute(query, user, db, priv)
+    cursor.execute(query, (user, db, priv))
     return cursor.fetchone()[0]
 
 def grant_database_privilege(cursor, user, db, priv):
-    if has_database_privilege(cursor, user, db, priv):
-        return False
+    prev_priv = get_database_privileges(cursor, user, db)
     query = 'GRANT %s ON DATABASE %s TO %s' % (priv, db, user)
     cursor.execute(query)
-    return True
+    curr_priv = get_database_privileges(cursor, user, db)
+    return len(curr_priv) > len(prev_priv)
 
 def revoke_database_privilege(cursor, user, db, priv):
-    if not has_database_privilege(cursor, user, db, priv):
-        return False
+    prev_priv = get_database_privileges(cursor, user, db)
     query = 'REVOKE %s ON DATABASE %s FROM %s' % (priv, db, user)
     cursor.execute(query)
-    return True
+    curr_priv = get_database_privileges(cursor, user, db)
+    return len(curr_priv) < len(prev_priv)
 
 def revoke_privileges(cursor, user, privs):
     if privs is None:
@@ -122,7 +152,7 @@ def revoke_privileges(cursor, user, privs):
             'table':revoke_table_privilege,
             'database':revoke_database_privilege
         }[type_]
-        for name, privileges in privs[type_].iteritem():
+        for name, privileges in privs[type_].iteritems():
             for privilege in privileges:
                 changed = revoke_func(cursor, user, name, privilege)\
                         or changed
@@ -139,7 +169,7 @@ def grant_privileges(cursor, user, privs):
             'table':grant_table_privilege,
             'database':grant_database_privilege
         }[type_]
-        for name, privileges in privs[type_].iteritem():
+        for name, privileges in privs[type_].iteritems():
             for privilege in privileges:
                 changed = grant_func(cursor, user, name, privilege)\
                         or changed
@@ -161,7 +191,7 @@ def parse_privs(privs, db):
     if privs is None:
         return privs
 
-    privs = {
+    o_privs = {
         'database':{},
         'table':{}
     }
@@ -169,15 +199,15 @@ def parse_privs(privs, db):
         if ':' not in token:
             type_ = 'database'
             name = db
-            privileges = token
+            priv_set = set(x.strip() for x in token.split(','))
         else:
             type_ = 'table'
             name, privileges = token.split(':', 1)
-        privileges = privileges.split(',')
+            priv_set = set(x.strip() for x in privileges.split(','))
         
-        privs[type_][name] = privileges
+        o_privs[type_][name] = priv_set
     
-    return privs
+    return o_privs
 
 # ===========================================
 # Module execution.