From 9e1d525b206e0daf3108b629d23b34c0573a3ad8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=80=97=E5=AD=90?= <i@haozi.net>
Date: Sat, 23 Nov 2024 01:53:08 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=95=B0=E6=8D=AE=E5=BA=93=E6=9D=83?=
 =?UTF-8?q?=E9=99=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/db/mysql.go    | 24 ++++++++++++++++++++++++
 pkg/db/postgres.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/pkg/db/mysql.go b/pkg/db/mysql.go
index 6cf7a546..722bfee1 100644
--- a/pkg/db/mysql.go
+++ b/pkg/db/mysql.go
@@ -122,6 +122,30 @@ func (m *MySQL) PrivilegesGrant(user, database string) error {
 	return err
 }
 
+func (m *MySQL) UserPrivileges(user, host string) (map[string][]string, error) {
+	rows, err := m.Query(fmt.Sprintf("SHOW GRANTS FOR '%s'@'%s'", user, host))
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	privileges := make(map[string][]string)
+	for rows.Next() {
+		var grant string
+		if err := rows.Scan(&grant); err != nil {
+			continue
+		}
+
+		var db string
+		var privs []string
+		if _, err := fmt.Sscanf(grant, "GRANT %s ON %s TO", &privs, &db); err == nil {
+			privileges[db] = append(privileges[db], privs...)
+		}
+	}
+
+	return privileges, nil
+}
+
 func (m *MySQL) PrivilegesRevoke(user, database string) error {
 	_, err := m.Exec(fmt.Sprintf("REVOKE ALL PRIVILEGES ON %s.* FROM '%s'@'localhost'", database, user))
 	m.flushPrivileges()
diff --git a/pkg/db/postgres.go b/pkg/db/postgres.go
index 6031251e..a2e10a7d 100644
--- a/pkg/db/postgres.go
+++ b/pkg/db/postgres.go
@@ -3,6 +3,7 @@ package db
 import (
 	"database/sql"
 	"fmt"
+	"strings"
 
 	_ "github.com/lib/pq"
 
@@ -116,6 +117,40 @@ func (m *Postgres) UserPassword(user, password string) error {
 	return err
 }
 
+func (p *Postgres) UserPrivileges(user string) (map[string][]string, error) {
+	query := `
+		SELECT 
+			table_catalog as database_name,
+			string_agg(DISTINCT privilege_type, ',') as privileges
+		FROM information_schema.role_database_privileges 
+		WHERE grantee = $1
+		GROUP BY table_catalog`
+
+	rows, err := p.Query(query, user)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query database privileges: %w", err)
+	}
+	defer rows.Close()
+
+	privileges := make(map[string][]string)
+
+	for rows.Next() {
+		var dbName, privilegeStr string
+		if err := rows.Scan(&dbName, &privilegeStr); err != nil {
+			return nil, fmt.Errorf("failed to scan row: %w", err)
+		}
+
+		key := fmt.Sprintf("%s.*", dbName)
+		privileges[key] = strings.Split(privilegeStr, ",")
+	}
+
+	if err = rows.Err(); err != nil {
+		return nil, fmt.Errorf("error iterating rows: %w", err)
+	}
+
+	return privileges, nil
+}
+
 func (m *Postgres) PrivilegesGrant(user, database string) error {
 	if _, err := m.Exec(fmt.Sprintf("ALTER DATABASE %s OWNER TO %s", database, user)); err != nil {
 		return err