From 4193d4614a19cfc2c3efb00b7a09ee8ca209a45d Mon Sep 17 00:00:00 2001 From: noapinsler Date: Mon, 1 Apr 2024 14:03:51 +0300 Subject: [PATCH 1/7] add autodiscovery for dbs --- job.go | 45 +++++++++++++++++++++++++++++ postgresql.go | 78 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+) create mode 100644 postgresql.go diff --git a/job.go b/job.go index a1d8ac6f..2f8d7103 100644 --- a/job.go +++ b/job.go @@ -243,6 +243,51 @@ func (j *Job) updateConnections() { conn = strings.Replace(conn, "AUTHTOKEN", url.QueryEscape(token), 1) } + if strings.HasPrefix(conn, "postgres://") || strings.HasPrefix(conn, "pg://") { + u, err := url.Parse(conn) + var filteredDBs []string + if err != nil { + level.Error(j.log).Log("msg", "Failed to parse URL", "url", conn, "err", err) + continue + } + if strings.Contains(u.Path, "include") || strings.Contains(u.Path, "exclude") { + if strings.Contains(u.Path, "include") && strings.Contains(u.Path, "exclude") { + fmt.Printf("You cannot use exclude and include: %s, error: %v\n", conn, err) + return // Or handle the error appropriately + } else { + extractedPath := u.Path //save pattern + u.Path = "/postgres" + dsn := u.String() + databases, err := listDatabases(dsn) // Corrected: closing parenthesis + if err != nil { + fmt.Printf("Error listing databases: %v\n", err) + continue + } + filteredDBs, err = filterDatabases(databases, extractedPath) + if err != nil { + fmt.Printf("Error filtering databases: %v\n", err) + continue + } + + for _, db := range filteredDBs { + u.Path = "/" + db // Set the path to the filtered database name + newUserDSN := u.String() + // Append to your connections slice, assuming j.conns is defined correctly + j.conns = append(j.conns, &connection{ + conn: nil, // Assuming you populate this later + url: newUserDSN, + driver: u.Scheme, + host: u.Host, + database: db, + user: u.User.Username(), + }) + } + } + } + continue + + } + u, err := url.Parse(conn) if err != nil { level.Error(j.log).Log("msg", "Failed to parse URL", "url", conn, "err", err) diff --git a/postgresql.go b/postgresql.go new file mode 100644 index 00000000..e6ef6d0b --- /dev/null +++ b/postgresql.go @@ -0,0 +1,78 @@ +package main + +import ( + "database/sql" + "fmt" + "strings" + + _ "github.com/lib/pq" +) + +const ( + INCLUDE_DBS = "/include:" + EXCLUDE_DBS = "/exclude:" +) + +func listDatabases(connStr string) ([]string, error) { + + db, err := sql.Open("postgres", connStr) + if err != nil { + return nil, err + } + defer db.Close() + + rows, err := db.Query("SELECT datname FROM pg_database WHERE datistemplate = false;") + if err != nil { + return nil, err + } + defer rows.Close() + + var databases []string + for rows.Next() { + var dbname string + if err := rows.Scan(&dbname); err != nil { + return nil, err + } + databases = append(databases, dbname) + } + + return databases, nil +} + +func filterDatabases(databases []string, pattern string) ([]string, error) { + var filtered []string + mode, dbs := parsePattern(pattern) + + dbMap := make(map[string]bool) + for _, db := range strings.Split(dbs, ",") { + dbMap[db] = true + } + + if mode == INCLUDE_DBS { + for _, dbname := range databases { + if _, ok := dbMap[dbname]; ok { + filtered = append(filtered, dbname) + } + } + } else if mode == EXCLUDE_DBS { + for _, dbname := range databases { + if _, ok := dbMap[dbname]; !ok { + filtered = append(filtered, dbname) + } + } + } else { + // If mode is neither include nor exclude, return an error + return nil, fmt.Errorf("invalid pattern: %s", pattern) + } + + return filtered, nil +} + +func parsePattern(pattern string) (mode string, dbs string) { + if strings.HasPrefix(pattern, INCLUDE_DBS) { + return INCLUDE_DBS, pattern[len(INCLUDE_DBS):] + } else if strings.HasPrefix(pattern, EXCLUDE_DBS) { + return EXCLUDE_DBS, pattern[len(EXCLUDE_DBS):] + } + return "", "" +} From 6a79c657bcf1d81997bfff1b2a047df900745a6c Mon Sep 17 00:00:00 2001 From: noapinsler Date: Tue, 2 Apr 2024 10:27:52 +0300 Subject: [PATCH 2/7] add regex --- postgresql.go | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/postgresql.go b/postgresql.go index e6ef6d0b..54b331c5 100644 --- a/postgresql.go +++ b/postgresql.go @@ -3,6 +3,7 @@ package main import ( "database/sql" "fmt" + "regexp" "strings" _ "github.com/lib/pq" @@ -43,26 +44,27 @@ func filterDatabases(databases []string, pattern string) ([]string, error) { var filtered []string mode, dbs := parsePattern(pattern) - dbMap := make(map[string]bool) - for _, db := range strings.Split(dbs, ",") { - dbMap[db] = true + // Compile the pattern into a regex + dbRegex, err := regexp.Compile(dbs) + if err != nil { + return nil, fmt.Errorf("invalid regex pattern: %s", err) } if mode == INCLUDE_DBS { for _, dbname := range databases { - if _, ok := dbMap[dbname]; ok { + if dbRegex.MatchString(dbname) { filtered = append(filtered, dbname) } } } else if mode == EXCLUDE_DBS { for _, dbname := range databases { - if _, ok := dbMap[dbname]; !ok { + if !dbRegex.MatchString(dbname) { filtered = append(filtered, dbname) } } } else { // If mode is neither include nor exclude, return an error - return nil, fmt.Errorf("invalid pattern: %s", pattern) + return nil, fmt.Errorf("invalid pattern mode: %s", pattern) } return filtered, nil From 0971c812e5a5d913d6149ce8efeadc70b2d0c38e Mon Sep 17 00:00:00 2001 From: noapinsler Date: Tue, 2 Apr 2024 10:55:23 +0300 Subject: [PATCH 3/7] add regex --- postgresql.go | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/postgresql.go b/postgresql.go index 54b331c5..f116a7bc 100644 --- a/postgresql.go +++ b/postgresql.go @@ -44,27 +44,27 @@ func filterDatabases(databases []string, pattern string) ([]string, error) { var filtered []string mode, dbs := parsePattern(pattern) - // Compile the pattern into a regex - dbRegex, err := regexp.Compile(dbs) - if err != nil { - return nil, fmt.Errorf("invalid regex pattern: %s", err) - } + // Split the dbs string into individual patterns + dbPatterns := strings.Split(dbs, ",") + + // Process each database name against the patterns + for _, dbname := range databases { + include := false - if mode == INCLUDE_DBS { - for _, dbname := range databases { - if dbRegex.MatchString(dbname) { - filtered = append(filtered, dbname) + for _, dbPattern := range dbPatterns { + matched, err := regexp.MatchString(dbPattern, dbname) + if err != nil { + return nil, fmt.Errorf("invalid pattern: %s", dbPattern) } - } - } else if mode == EXCLUDE_DBS { - for _, dbname := range databases { - if !dbRegex.MatchString(dbname) { - filtered = append(filtered, dbname) + if matched { + include = true + break } } - } else { - // If mode is neither include nor exclude, return an error - return nil, fmt.Errorf("invalid pattern mode: %s", pattern) + + if (mode == INCLUDE_DBS && include) || (mode == EXCLUDE_DBS && !include) { + filtered = append(filtered, dbname) + } } return filtered, nil From 84d4bf28c2f34086c481dc61236c2d521629279e Mon Sep 17 00:00:00 2001 From: noapinsler Date: Tue, 16 Apr 2024 15:52:25 +0300 Subject: [PATCH 4/7] add regex --- job.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/job.go b/job.go index 2f8d7103..d7b49d68 100644 --- a/job.go +++ b/job.go @@ -282,10 +282,9 @@ func (j *Job) updateConnections() { user: u.User.Username(), }) } + continue } } - continue - } u, err := url.Parse(conn) @@ -351,6 +350,8 @@ func (j *Job) updateConnections() { } j.conns = append(j.conns, newConn) + //print j + fmt.Println(j) } } } From 604cc4450ed911738c34fa7a5f0a65811660af67 Mon Sep 17 00:00:00 2001 From: noapinsler Date: Tue, 16 Apr 2024 17:41:57 +0300 Subject: [PATCH 5/7] remove debug --- job.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/job.go b/job.go index d7b49d68..1e28b829 100644 --- a/job.go +++ b/job.go @@ -350,8 +350,6 @@ func (j *Job) updateConnections() { } j.conns = append(j.conns, newConn) - //print j - fmt.Println(j) } } } From 231a95df78c7a2b41d4856a8f0a681552f723b32 Mon Sep 17 00:00:00 2001 From: noapinsler <82938819+noapinsler@users.noreply.github.com> Date: Wed, 17 Apr 2024 09:44:12 +0300 Subject: [PATCH 6/7] Update job.go remove debug comments --- job.go | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/job.go b/job.go index 19e9579a..29b9acd5 100644 --- a/job.go +++ b/job.go @@ -253,12 +253,12 @@ func (j *Job) updateConnections() { if strings.Contains(u.Path, "include") || strings.Contains(u.Path, "exclude") { if strings.Contains(u.Path, "include") && strings.Contains(u.Path, "exclude") { fmt.Printf("You cannot use exclude and include: %s, error: %v\n", conn, err) - return // Or handle the error appropriately + return } else { extractedPath := u.Path //save pattern u.Path = "/postgres" dsn := u.String() - databases, err := listDatabases(dsn) // Corrected: closing parenthesis + databases, err := listDatabases(dsn) if err != nil { fmt.Printf("Error listing databases: %v\n", err) continue @@ -272,9 +272,8 @@ func (j *Job) updateConnections() { for _, db := range filteredDBs { u.Path = "/" + db // Set the path to the filtered database name newUserDSN := u.String() - // Append to your connections slice, assuming j.conns is defined correctly j.conns = append(j.conns, &connection{ - conn: nil, // Assuming you populate this later + conn: nil, url: newUserDSN, driver: u.Scheme, host: u.Host, From 7ad079a8445ddb9be56f740a5f345cc50e54fa56 Mon Sep 17 00:00:00 2001 From: noapinsler <82938819+noapinsler@users.noreply.github.com> Date: Wed, 3 Jul 2024 15:43:21 +0300 Subject: [PATCH 7/7] Update job.go change printf to logger --- job.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/job.go b/job.go index 29b9acd5..0239bfbc 100644 --- a/job.go +++ b/job.go @@ -252,7 +252,7 @@ func (j *Job) updateConnections() { } if strings.Contains(u.Path, "include") || strings.Contains(u.Path, "exclude") { if strings.Contains(u.Path, "include") && strings.Contains(u.Path, "exclude") { - fmt.Printf("You cannot use exclude and include: %s, error: %v\n", conn, err) + level.Error(j.log).Log("msg", "You cannot use exclude and include:", "url", conn, "err", err) return } else { extractedPath := u.Path //save pattern @@ -260,12 +260,12 @@ func (j *Job) updateConnections() { dsn := u.String() databases, err := listDatabases(dsn) if err != nil { - fmt.Printf("Error listing databases: %v\n", err) + level.Error(j.log).Log("msg", "Error listing databases", "url", conn, "err", err) continue } filteredDBs, err = filterDatabases(databases, extractedPath) if err != nil { - fmt.Printf("Error filtering databases: %v\n", err) + level.Error(j.log).Log("msg", "Error filtering databases", "url", conn, "err", err) continue }