Databases: Grant and Revoke privileges on DBs (#20)

Author: phoebesimon

pkg/connector/database.go

@@ -248,14 +248,29 @@ func (r *databaseSyncer) Delete(ctx context.Context, resourceId *v2.ResourceId)
 }
 
 func (r *databaseSyncer) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) ([]*v2.Grant, annotations.Annotations, error) {
-	if principal.Id.ResourceType != databaseResourceType.Id {
+	if principal.Id.ResourceType != roleResourceType.Id {
 		return nil, nil, fmt.Errorf("baton-postgres: only users and roles can have roles granted")
 	}
 
-	// TODO: pass IDs into client.Grant() and look up the names there
-	dbName := entitlement.Resource.DisplayName
+	// Parse the Entitlement ID to get the database ID and privilege name
+	_, dbIdStr, privilegeName, isGrant, err := parseEntitlementID(entitlement.Id)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	dbID, err := strconv.ParseInt(dbIdStr, 10, 64)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// Look up the database by ID
+	pgDb, err := r.client.GetDatabase(ctx, dbID)
+	if err != nil {
+		return nil, nil, err
+	}
+
 	principalName := principal.DisplayName
-	err := r.client.GrantDatabase(ctx, dbName, principalName, entitlement.GetDisplayName())
+	err = r.client.GrantDatabase(ctx, pgDb.Name, principalName, privilegeName, isGrant)
 	return nil, nil, err
 }
 

pkg/postgres/databases.go

@@ -143,13 +143,24 @@ func (c *Client) DeleteDatabase(ctx context.Context, dbName string) error {
 	return err
 }
 
-func (c *Client) GrantDatabase(ctx context.Context, dbName string, principalName string, privilege string) error {
+func transformPrivilege(privilege string) string {
+	return strings.ReplaceAll(privilege, "-", "")
+}
+
+func (c *Client) GrantDatabase(ctx context.Context, dbName string, principalName string, privilege string, isGrant bool) error {
 	l := ctxzap.Extract(ctx)
 	l.Debug("granting database", zap.String("dbName", dbName), zap.String("principalName", principalName), zap.String("privilege", privilege))
 
 	sanitizedDbName := pgx.Identifier{dbName}.Sanitize()
 	sanitizedPrincipalName := pgx.Identifier{principalName}.Sanitize()
-	q := fmt.Sprintf("GRANT %s ON DATABASE %s TO %s", privilege, sanitizedDbName, sanitizedPrincipalName)
+	sanitizedPrivilege := pgx.Identifier{transformPrivilege(privilege)}.Sanitize()
+	var q string
+	if isGrant {
+		q = fmt.Sprintf("GRANT %s ON DATABASE %s TO %s WITH GRANT OPTION", sanitizedPrivilege, sanitizedDbName, sanitizedPrincipalName)
+	} else {
+		q = fmt.Sprintf("GRANT %s ON DATABASE %s TO %s", sanitizedPrivilege, sanitizedDbName, sanitizedPrincipalName)
+	}
+
 	_, err := c.db.Exec(ctx, q)
 	return err
 }
@@ -159,7 +170,13 @@ func (c *Client) RevokeDatabase(ctx context.Context, dbName string, target strin
 
 	sanitizedDbName := pgx.Identifier{dbName}.Sanitize()
 	sanitizedTarget := pgx.Identifier{target}.Sanitize()
-	q := fmt.Sprintf("REVOKE %s ON DATABASE %s FROM %s", privilege, sanitizedDbName, sanitizedTarget)
+	sanitizedPrivilege := pgx.Identifier{transformPrivilege(privilege)}.Sanitize()
+	var q string
+	if isGrant {
+		q = fmt.Sprintf("REVOKE GRANT OPTION for %s ON DATABASE %s FROM %s", sanitizedPrivilege, sanitizedDbName, sanitizedTarget)
+	} else {
+		q = fmt.Sprintf("REVOKE %s ON DATABASE %s FROM %s", sanitizedPrivilege, sanitizedDbName, sanitizedTarget)
+	}
 
 	l.Debug("revoking role from member", zap.String("query", q))
 	_, err := c.db.Exec(ctx, q)