cmd/tools: remove nullify-bad-user-agents tool
This tool is being removed because it has served its purpose and was blocking another removal from being verified. Change-Id: Ie888aa7ae1b153a34210af3a5d5a3682b381ba82
This commit is contained in:
Wilfred Asomani
parent
2f04e20627
commit
73ffa0827f
@ -1,308 +0,0 @@
|
||||
// Copyright (C) 2021 Storj Labs, Inc.
|
||||
// See LICENSE for copying information.
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/jackc/pgtype"
|
||||
pgx "github.com/jackc/pgx/v4"
|
||||
"github.com/zeebo/errs"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"storj.io/private/dbutil/pgutil"
|
||||
)
|
||||
|
||||
// MigrateTablesLimited runs the migration for each table with update count limits.
|
||||
func MigrateTablesLimited(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
err = MigrateUsersLimited(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating users: %w", err)
|
||||
}
|
||||
err = MigrateProjectsLimited(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating projects: %w", err)
|
||||
}
|
||||
err = MigrateAPIKeysLimited(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating api_keys: %w", err)
|
||||
}
|
||||
err = MigrateBucketMetainfosLimited(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating bucket_metainfos: %w", err)
|
||||
}
|
||||
err = MigrateValueAttributionsLimited(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating value_attributions: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateUsersLimited updates the user_agent column to corresponding Partners.Names or partner_id if applicable for a limited number
|
||||
// of rows.
|
||||
func MigrateUsersLimited(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
log.Info("beginning users migration", zap.Int("max updates", config.MaxUpdates))
|
||||
|
||||
// wrap select in anonymous function for deferred rows.Close()
|
||||
selected, err := func() (ids [][]byte, err error) {
|
||||
rows, err := conn.Query(ctx, `
|
||||
SELECT id FROM users
|
||||
WHERE user_agent = partner_id
|
||||
LIMIT $1
|
||||
`, config.MaxUpdates)
|
||||
if err != nil {
|
||||
return nil, errs.New("selecting ids for update: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var id pgtype.GenericBinary
|
||||
err = rows.Scan(&id)
|
||||
if err != nil {
|
||||
return nil, errs.New("scanning row: %w", err)
|
||||
}
|
||||
ids = append(ids, id.Bytes)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
row := conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE users
|
||||
SET user_agent = NULL
|
||||
WHERE users.id IN (SELECT unnest($1::bytea[]))
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated
|
||||
`, pgutil.ByteaArray(selected),
|
||||
)
|
||||
var updated int
|
||||
err = row.Scan(&updated)
|
||||
if err != nil {
|
||||
return errs.New("error scanning results: %w", err)
|
||||
}
|
||||
log.Info("updated rows", zap.Int("count", updated))
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateProjectsLimited updates the user_agent column to corresponding Partners.Names or partner_id if applicable for a limited number
|
||||
// of rows.
|
||||
func MigrateProjectsLimited(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
log.Info("beginning projects migration", zap.Int("max updates", config.MaxUpdates))
|
||||
|
||||
// wrap select in anonymous function for deferred rows.Close()
|
||||
selected, err := func() (ids [][]byte, err error) {
|
||||
rows, err := conn.Query(ctx, `
|
||||
SELECT id FROM projects
|
||||
WHERE user_agent = partner_id
|
||||
LIMIT $1
|
||||
`, config.MaxUpdates)
|
||||
if err != nil {
|
||||
return nil, errs.New("selecting ids for update: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var id pgtype.GenericBinary
|
||||
err = rows.Scan(&id)
|
||||
if err != nil {
|
||||
return nil, errs.New("scanning row: %w", err)
|
||||
}
|
||||
ids = append(ids, id.Bytes)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
row := conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE projects
|
||||
SET user_agent = NULL
|
||||
WHERE projects.id IN (SELECT unnest($1::bytea[]))
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated
|
||||
`, pgutil.ByteaArray(selected),
|
||||
)
|
||||
var updated int
|
||||
err = row.Scan(&updated)
|
||||
if err != nil {
|
||||
return errs.New("error scanning results: %w", err)
|
||||
}
|
||||
log.Info("updated rows", zap.Int("count", updated))
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateAPIKeysLimited updates the user_agent column to corresponding Partners.Names or partner_id if applicable for a limited number
|
||||
// of rows.
|
||||
func MigrateAPIKeysLimited(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
log.Info("beginning api_keys migration", zap.Int("max updates", config.MaxUpdates))
|
||||
|
||||
// wrap select in anonymous function for deferred rows.Close()
|
||||
selected, err := func() (ids [][]byte, err error) {
|
||||
rows, err := conn.Query(ctx, `
|
||||
SELECT id FROM api_keys
|
||||
WHERE user_agent = partner_id
|
||||
LIMIT $1
|
||||
`, config.MaxUpdates)
|
||||
if err != nil {
|
||||
return nil, errs.New("selecting ids for update: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var id pgtype.GenericBinary
|
||||
err = rows.Scan(&id)
|
||||
if err != nil {
|
||||
return nil, errs.New("scanning row: %w", err)
|
||||
}
|
||||
ids = append(ids, id.Bytes)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
row := conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE api_keys
|
||||
SET user_agent = NULL
|
||||
WHERE api_keys.id IN (SELECT unnest($1::bytea[]))
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated
|
||||
`, pgutil.ByteaArray(selected),
|
||||
)
|
||||
var updated int
|
||||
err = row.Scan(&updated)
|
||||
if err != nil {
|
||||
return errs.New("error scanning results: %w", err)
|
||||
}
|
||||
log.Info("updated rows", zap.Int("count", updated))
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateBucketMetainfosLimited updates the user_agent column to corresponding Partners.Names or partner_id if applicable for a limited number
|
||||
// of rows.
|
||||
func MigrateBucketMetainfosLimited(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
log.Info("beginning bucket_metainfos migration", zap.Int("max updates", config.MaxUpdates))
|
||||
|
||||
// wrap select in anonymous function for deferred rows.Close()
|
||||
selected, err := func() (ids [][]byte, err error) {
|
||||
rows, err := conn.Query(ctx, `
|
||||
SELECT id FROM bucket_metainfos
|
||||
WHERE user_agent = partner_id
|
||||
LIMIT $1
|
||||
`, config.MaxUpdates)
|
||||
if err != nil {
|
||||
return nil, errs.New("selecting ids for update: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var id pgtype.GenericBinary
|
||||
err = rows.Scan(&id)
|
||||
if err != nil {
|
||||
return nil, errs.New("scanning row: %w", err)
|
||||
}
|
||||
ids = append(ids, id.Bytes)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
row := conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE bucket_metainfos
|
||||
SET user_agent = NULL
|
||||
WHERE bucket_metainfos.id IN (SELECT unnest($1::bytea[]))
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated
|
||||
`, pgutil.ByteaArray(selected),
|
||||
)
|
||||
var updated int
|
||||
err = row.Scan(&updated)
|
||||
if err != nil {
|
||||
return errs.New("error scanning results: %w", err)
|
||||
}
|
||||
log.Info("updated rows", zap.Int("count", updated))
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateValueAttributionsLimited updates the user_agent column to corresponding Partners.Names or partner_id if applicable for a limited number
|
||||
// of rows.
|
||||
func MigrateValueAttributionsLimited(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
log.Info("beginning value_attributions migration", zap.Int("max updates", config.MaxUpdates))
|
||||
|
||||
// wrap select in anonymous function for deferred rows.Close()
|
||||
projects, buckets, err := func() (projectIDs [][]byte, buckets [][]byte, err error) {
|
||||
rows, err := conn.Query(ctx, `
|
||||
SELECT project_id, bucket_name FROM value_attributions
|
||||
WHERE user_agent = partner_id
|
||||
LIMIT $1
|
||||
`, config.MaxUpdates)
|
||||
if err != nil {
|
||||
return nil, nil, errs.New("selecting rows for update: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var projectID, bucketName pgtype.GenericBinary
|
||||
err = rows.Scan(&projectID, &bucketName)
|
||||
if err != nil {
|
||||
return nil, nil, errs.New("scanning row: %w", err)
|
||||
}
|
||||
projectIDs = append(projectIDs, projectID.Bytes)
|
||||
buckets = append(buckets, bucketName.Bytes)
|
||||
}
|
||||
return projectIDs, buckets, rows.Err()
|
||||
}()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
row := conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE value_attributions
|
||||
SET user_agent = NULL
|
||||
WHERE value_attributions.project_id IN (SELECT unnest($1::bytea[]))
|
||||
AND value_attributions.bucket_name IN (SELECT unnest($2::bytea[]))
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated
|
||||
`, pgutil.ByteaArray(projects), pgutil.ByteaArray(buckets),
|
||||
)
|
||||
var updated int
|
||||
err = row.Scan(&updated)
|
||||
if err != nil {
|
||||
return errs.New("error scanning results: %w", err)
|
||||
}
|
||||
log.Info("updated rows", zap.Int("count", updated))
|
||||
return nil
|
||||
}
|
||||
@ -1,113 +0,0 @@
|
||||
// Copyright (C) 2021 Storj Labs, Inc.
|
||||
// See LICENSE for copying information.
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
pgx "github.com/jackc/pgx/v4"
|
||||
"github.com/spacemonkeygo/monkit/v3"
|
||||
"github.com/spf13/cobra"
|
||||
flag "github.com/spf13/pflag"
|
||||
"github.com/zeebo/errs"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"storj.io/private/process"
|
||||
)
|
||||
|
||||
var mon = monkit.Package()
|
||||
|
||||
var (
|
||||
rootCmd = &cobra.Command{
|
||||
Use: "nullify-bad-user-agents",
|
||||
Short: "nullify-bad-user-agents",
|
||||
}
|
||||
|
||||
runCmd = &cobra.Command{
|
||||
Use: "run",
|
||||
Short: "run nullify-bad-user-agents",
|
||||
RunE: run,
|
||||
}
|
||||
|
||||
config Config
|
||||
)
|
||||
|
||||
func init() {
|
||||
rootCmd.AddCommand(runCmd)
|
||||
|
||||
config.BindFlags(runCmd.Flags())
|
||||
}
|
||||
|
||||
// Config defines configuration for migration.
|
||||
type Config struct {
|
||||
SatelliteDB string
|
||||
Limit int
|
||||
MaxUpdates int
|
||||
}
|
||||
|
||||
// BindFlags adds bench flags to the flagset.
|
||||
func (config *Config) BindFlags(flag *flag.FlagSet) {
|
||||
flag.StringVar(&config.SatelliteDB, "satellitedb", "", "connection URL for satelliteDB")
|
||||
flag.IntVar(&config.Limit, "limit", 1000, "number of updates to perform at once")
|
||||
flag.IntVar(&config.MaxUpdates, "max-updates", 0, "max number of updates to perform on each table")
|
||||
}
|
||||
|
||||
// VerifyFlags verifies whether the values provided are valid.
|
||||
func (config *Config) VerifyFlags() error {
|
||||
var errlist errs.Group
|
||||
if config.SatelliteDB == "" {
|
||||
errlist.Add(errors.New("flag '--satellitedb' is not set"))
|
||||
}
|
||||
return errlist.Err()
|
||||
}
|
||||
|
||||
func run(cmd *cobra.Command, args []string) error {
|
||||
if err := config.VerifyFlags(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
ctx, _ := process.Ctx(cmd)
|
||||
log := zap.L()
|
||||
return Migrate(ctx, log, config)
|
||||
}
|
||||
|
||||
func main() {
|
||||
process.Exec(rootCmd)
|
||||
}
|
||||
|
||||
// Migrate updates the user_agent column if user_agent = '\x00000000000000000000000000000000' and sets
|
||||
// it to NULL.
|
||||
// Affected tables:
|
||||
//
|
||||
// users
|
||||
// projects
|
||||
// api_keys
|
||||
// bucket_metainfos
|
||||
// value_attributions.
|
||||
func Migrate(ctx context.Context, log *zap.Logger, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
conn, err := pgx.Connect(ctx, config.SatelliteDB)
|
||||
if err != nil {
|
||||
return errs.New("unable to connect %q: %w", config.SatelliteDB, err)
|
||||
}
|
||||
defer func() {
|
||||
err = errs.Combine(err, conn.Close(ctx))
|
||||
}()
|
||||
|
||||
// The original migrations are already somewhat complex and in my opinion,
|
||||
// trying to edit them to be able to handle conditionally limiting updates increased the
|
||||
// complexity. While I think splitting out the limited update migrations isn't the
|
||||
// most ideal solution, since this code is temporary we don't need to worry about
|
||||
// maintenance concerns with having multiple queries.
|
||||
if config.MaxUpdates > 1000 {
|
||||
return errs.New("When running limited migration, set --max-updates to something less than 1000")
|
||||
}
|
||||
if config.MaxUpdates > 0 {
|
||||
return MigrateTablesLimited(ctx, log, conn, config)
|
||||
}
|
||||
|
||||
return MigrateTables(ctx, log, conn, config)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@ -1,511 +0,0 @@
|
||||
// Copyright (C) 2021 Storj Labs, Inc.
|
||||
// See LICENSE for copying information.
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
pgx "github.com/jackc/pgx/v4"
|
||||
"github.com/zeebo/errs"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"storj.io/private/dbutil/cockroachutil"
|
||||
)
|
||||
|
||||
// MigrateTables runs the migration for each table.
|
||||
func MigrateTables(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
err = MigrateUsers(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating users: %w", err)
|
||||
}
|
||||
err = MigrateProjects(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating projects: %w", err)
|
||||
}
|
||||
err = MigrateAPIKeys(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating api_keys: %w", err)
|
||||
}
|
||||
err = MigrateBucketMetainfos(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating bucket_metainfos: %w", err)
|
||||
}
|
||||
err = MigrateValueAttributions(ctx, log, conn, config)
|
||||
if err != nil {
|
||||
return errs.New("error migrating value_attributions: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateUsers updates the user_agent column to corresponding Partners.Names or partner_id if applicable.
|
||||
func MigrateUsers(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
// We select the next id then use limit as an offset which actually gives us limit+1 rows.
|
||||
offset := config.Limit - 1
|
||||
|
||||
startID := []byte{}
|
||||
nextID := []byte{}
|
||||
var total int
|
||||
more := true
|
||||
|
||||
// select the next ID after startID and offset the result. We will update relevant rows from the range of startID to nextID.
|
||||
_, err = conn.Prepare(ctx, "select-for-update-users", "SELECT id FROM users WHERE id > $1 ORDER BY id OFFSET $2 LIMIT 1")
|
||||
if err != nil {
|
||||
return errs.New("could not prepare select query")
|
||||
}
|
||||
|
||||
// update range from startID to nextID. Return count of updated rows to log progress.
|
||||
_, err = conn.Prepare(ctx, "update-limited-users", `
|
||||
WITH updated as (
|
||||
UPDATE users
|
||||
SET user_agent = NULL
|
||||
WHERE users.id > $1 AND users.id <= $2
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`)
|
||||
if err != nil {
|
||||
return errs.New("could not prepare update statement: %w", err)
|
||||
}
|
||||
|
||||
for {
|
||||
row := conn.QueryRow(ctx, "select-for-update-users", startID, offset)
|
||||
err = row.Scan(&nextID)
|
||||
if err != nil {
|
||||
if errs.Is(err, pgx.ErrNoRows) {
|
||||
more = false
|
||||
} else {
|
||||
return errs.New("unable to select row for update from users: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var updated int
|
||||
for {
|
||||
var row pgx.Row
|
||||
if more {
|
||||
row = conn.QueryRow(ctx, "update-limited-users", startID, nextID)
|
||||
} else {
|
||||
// if !more then the select statement reached the end of the table. Update to the end of the table.
|
||||
row = conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE users
|
||||
SET user_agent = NULL
|
||||
WHERE users.id > $1
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`, startID,
|
||||
)
|
||||
}
|
||||
err := row.Scan(&updated)
|
||||
if err != nil {
|
||||
if cockroachutil.NeedsRetry(err) {
|
||||
continue
|
||||
} else if errs.Is(err, pgx.ErrNoRows) {
|
||||
break
|
||||
}
|
||||
return errs.New("updating users %w", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
total += updated
|
||||
if !more {
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated))
|
||||
break
|
||||
}
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated), zap.Binary("last id", nextID))
|
||||
startID = nextID
|
||||
}
|
||||
log.Info("users migration complete", zap.Int("total rows updated", total))
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateProjects updates the user_agent column to corresponding PartnerInfo.Names or partner_id if applicable.
|
||||
func MigrateProjects(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
// We select the next id then use limit as an offset which actually gives us limit+1 rows.
|
||||
offset := config.Limit - 1
|
||||
|
||||
startID := []byte{}
|
||||
nextID := []byte{}
|
||||
var total int
|
||||
more := true
|
||||
|
||||
// select the next ID after startID and offset the result. We will update relevant rows from the range of startID to nextID.
|
||||
_, err = conn.Prepare(ctx, "select-for-update-projects", "SELECT id FROM projects WHERE id > $1 ORDER BY id OFFSET $2 LIMIT 1")
|
||||
if err != nil {
|
||||
return errs.New("could not prepare select query: %w", err)
|
||||
}
|
||||
|
||||
// update range from startID to nextID. Return count of updated rows to log progress.
|
||||
_, err = conn.Prepare(ctx, "update-limited-projects", `
|
||||
WITH updated as (
|
||||
UPDATE projects
|
||||
SET user_agent = NULL
|
||||
WHERE projects.id > $1 AND projects.id <= $2
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`)
|
||||
if err != nil {
|
||||
return errs.New("could not prepare update statement")
|
||||
}
|
||||
|
||||
for {
|
||||
row := conn.QueryRow(ctx, "select-for-update-projects", startID, offset)
|
||||
err = row.Scan(&nextID)
|
||||
if err != nil {
|
||||
if errs.Is(err, pgx.ErrNoRows) {
|
||||
more = false
|
||||
} else {
|
||||
return errs.New("unable to select row for update from projects: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
var updated int
|
||||
for {
|
||||
var row pgx.Row
|
||||
if more {
|
||||
row = conn.QueryRow(ctx, "update-limited-projects", startID, nextID)
|
||||
} else {
|
||||
// if !more then the select statement reached the end of the table. Update to the end of the table.
|
||||
row = conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE projects
|
||||
SET user_agent = NULL
|
||||
WHERE projects.id > $1
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`, startID,
|
||||
)
|
||||
}
|
||||
|
||||
err := row.Scan(&updated)
|
||||
if err != nil {
|
||||
if cockroachutil.NeedsRetry(err) {
|
||||
continue
|
||||
} else if errs.Is(err, pgx.ErrNoRows) {
|
||||
break
|
||||
}
|
||||
return errs.New("updating projects %w", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
total += updated
|
||||
if !more {
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated))
|
||||
break
|
||||
}
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated), zap.Binary("last id", nextID))
|
||||
startID = nextID
|
||||
}
|
||||
log.Info("projects migration complete", zap.Int("total rows updated", total))
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateAPIKeys updates the user_agent column to corresponding PartnerInfo.Names or partner_id if applicable.
|
||||
func MigrateAPIKeys(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
// We select the next id then use limit as an offset which actually gives us limit+1 rows.
|
||||
offset := config.Limit - 1
|
||||
|
||||
startID := []byte{}
|
||||
nextID := []byte{}
|
||||
var total int
|
||||
more := true
|
||||
|
||||
// select the next ID after startID and offset the result. We will update relevant rows from the range of startID to nextID.
|
||||
_, err = conn.Prepare(ctx, "select-for-update-api-keys", "SELECT id FROM api_keys WHERE id > $1 ORDER BY id OFFSET $2 LIMIT 1")
|
||||
if err != nil {
|
||||
return errs.New("could not prepare select query: %w", err)
|
||||
}
|
||||
|
||||
// update range from startID to nextID. Return count of updated rows to log progress.
|
||||
_, err = conn.Prepare(ctx, "update-limited-api-keys", `
|
||||
WITH updated as (
|
||||
UPDATE api_keys
|
||||
SET user_agent = NULL
|
||||
WHERE api_keys.id > $1 AND api_keys.id <= $2
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`)
|
||||
if err != nil {
|
||||
return errs.New("could not prepare update statement")
|
||||
}
|
||||
|
||||
for {
|
||||
row := conn.QueryRow(ctx, "select-for-update-api-keys", startID, offset)
|
||||
err = row.Scan(&nextID)
|
||||
if err != nil {
|
||||
if errs.Is(err, pgx.ErrNoRows) {
|
||||
more = false
|
||||
} else {
|
||||
return errs.New("unable to select row for update from api_keys: %w", err)
|
||||
}
|
||||
}
|
||||
var updated int
|
||||
for {
|
||||
var row pgx.Row
|
||||
if more {
|
||||
row = conn.QueryRow(ctx, "update-limited-api-keys", startID, nextID)
|
||||
} else {
|
||||
// if !more then the select statement reached the end of the table. Update to the end of the table.
|
||||
row = conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE api_keys
|
||||
SET user_agent = NULL
|
||||
WHERE api_keys.id > $1
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`, startID,
|
||||
)
|
||||
}
|
||||
err := row.Scan(&updated)
|
||||
if err != nil {
|
||||
if cockroachutil.NeedsRetry(err) {
|
||||
continue
|
||||
} else if errs.Is(err, pgx.ErrNoRows) {
|
||||
break
|
||||
}
|
||||
return errs.New("updating api_keys %w", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
total += updated
|
||||
if !more {
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated))
|
||||
break
|
||||
}
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated), zap.Binary("last id", nextID))
|
||||
startID = nextID
|
||||
}
|
||||
log.Info("api_keys migration complete", zap.Int("total rows updated", total))
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateBucketMetainfos updates the user_agent column to corresponding Partners.Names or partner_id if applicable.
|
||||
func MigrateBucketMetainfos(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
// We select the next id then use limit as an offset which actually gives us limit+1 rows.
|
||||
offset := config.Limit - 1
|
||||
|
||||
startID := []byte{}
|
||||
nextID := []byte{}
|
||||
var total int
|
||||
more := true
|
||||
|
||||
// select the next ID after startID and offset the result. We will update relevant rows from the range of startID to nextID.
|
||||
_, err = conn.Prepare(ctx, "select-for-update-bucket-metainfos", "SELECT id FROM bucket_metainfos WHERE id > $1 ORDER BY id OFFSET $2 LIMIT 1")
|
||||
if err != nil {
|
||||
return errs.New("could not prepare select query: %w", err)
|
||||
}
|
||||
|
||||
// update range from startID to nextID. Return count of updated rows to log progress.
|
||||
_, err = conn.Prepare(ctx, "update-limited-bucket-metainfos", `
|
||||
WITH updated as (
|
||||
UPDATE bucket_metainfos
|
||||
SET user_agent = NULL
|
||||
WHERE bucket_metainfos.id > $1 AND bucket_metainfos.id <= $2
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`)
|
||||
if err != nil {
|
||||
return errs.New("could not prepare update statement")
|
||||
}
|
||||
|
||||
for {
|
||||
row := conn.QueryRow(ctx, "select-for-update-bucket-metainfos", startID, offset)
|
||||
err = row.Scan(&nextID)
|
||||
if err != nil {
|
||||
if errs.Is(err, pgx.ErrNoRows) {
|
||||
more = false
|
||||
} else {
|
||||
return errs.New("unable to select row for update from bucket_metainfos: %w", err)
|
||||
}
|
||||
}
|
||||
var updated int
|
||||
for {
|
||||
var row pgx.Row
|
||||
if more {
|
||||
row = conn.QueryRow(ctx, "update-limited-bucket-metainfos", startID, nextID)
|
||||
} else {
|
||||
// if !more then the select statement reached the end of the table. Update to the end of the table.
|
||||
row = conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE bucket_metainfos
|
||||
SET user_agent = NULL
|
||||
WHERE bucket_metainfos.id > $1
|
||||
AND user_agent = partner_id
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`, startID,
|
||||
)
|
||||
}
|
||||
err := row.Scan(&updated)
|
||||
if err != nil {
|
||||
if cockroachutil.NeedsRetry(err) {
|
||||
continue
|
||||
} else if errs.Is(err, pgx.ErrNoRows) {
|
||||
break
|
||||
}
|
||||
return errs.New("updating bucket_metainfos %w", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
total += updated
|
||||
if !more {
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated))
|
||||
break
|
||||
}
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated), zap.Binary("last id", nextID))
|
||||
startID = nextID
|
||||
}
|
||||
log.Info("bucket_metainfos migration complete", zap.Int("total rows updated", total))
|
||||
return nil
|
||||
}
|
||||
|
||||
// MigrateValueAttributions updates the user_agent column to corresponding Partners.Names or partner_id if applicable.
|
||||
func MigrateValueAttributions(ctx context.Context, log *zap.Logger, conn *pgx.Conn, config Config) (err error) {
|
||||
defer mon.Task()(&ctx)(&err)
|
||||
|
||||
// We select the next id then use limit as an offset which actually gives us limit+1 rows.
|
||||
offset := config.Limit - 1
|
||||
|
||||
startProjectID := []byte{}
|
||||
startBucket := []byte{}
|
||||
nextProjectID := []byte{}
|
||||
nextBucket := []byte{}
|
||||
var total int
|
||||
more := true
|
||||
|
||||
// select the next ID after startID and offset the result. We will update relevant rows from the range of startID to nextID.
|
||||
_, err = conn.Prepare(ctx, "select-for-update-value-attributions", `
|
||||
SELECT project_id, bucket_name
|
||||
FROM value_attributions
|
||||
WHERE project_id > $1
|
||||
OR (
|
||||
project_id = $1 AND bucket_name > $2
|
||||
)
|
||||
ORDER BY project_id, bucket_name
|
||||
OFFSET $3
|
||||
LIMIT 1
|
||||
`)
|
||||
if err != nil {
|
||||
return errs.New("could not prepare select query: %w", err)
|
||||
}
|
||||
|
||||
// update range from startID to nextID. Return count of updated rows to log progress.
|
||||
_, err = conn.Prepare(ctx, "update-limited-value-attributions", `
|
||||
WITH updated as (
|
||||
UPDATE value_attributions
|
||||
SET user_agent = NULL
|
||||
WHERE user_agent = partner_id
|
||||
AND (
|
||||
value_attributions.project_id > $1
|
||||
OR (
|
||||
value_attributions.project_id = $1 AND value_attributions.bucket_name > $3
|
||||
)
|
||||
)
|
||||
AND (
|
||||
value_attributions.project_id < $2
|
||||
OR (
|
||||
value_attributions.project_id = $2 AND value_attributions.bucket_name <= $4
|
||||
)
|
||||
)
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`)
|
||||
if err != nil {
|
||||
return errs.New("could not prepare update statement")
|
||||
}
|
||||
|
||||
for {
|
||||
row := conn.QueryRow(ctx, "select-for-update-value-attributions", startProjectID, startBucket, offset)
|
||||
err = row.Scan(&nextProjectID, &nextBucket)
|
||||
if err != nil {
|
||||
if errs.Is(err, pgx.ErrNoRows) {
|
||||
more = false
|
||||
} else {
|
||||
return errs.New("unable to select row for update from value_attributions: %w", err)
|
||||
}
|
||||
}
|
||||
var updated int
|
||||
for {
|
||||
var row pgx.Row
|
||||
if more {
|
||||
row = conn.QueryRow(ctx, "update-limited-value-attributions", startProjectID, nextProjectID, startBucket, nextBucket)
|
||||
} else {
|
||||
// if !more then the select statement reached the end of the table. Update to the end of the table.
|
||||
row = conn.QueryRow(ctx, `
|
||||
WITH updated as (
|
||||
UPDATE value_attributions
|
||||
SET user_agent = NULL
|
||||
WHERE user_agent = partner_id
|
||||
AND (
|
||||
value_attributions.project_id > $1
|
||||
OR (
|
||||
value_attributions.project_id = $1 AND value_attributions.bucket_name > $2
|
||||
)
|
||||
)
|
||||
RETURNING 1
|
||||
)
|
||||
SELECT count(*)
|
||||
FROM updated;
|
||||
`, startProjectID, startBucket,
|
||||
)
|
||||
}
|
||||
err := row.Scan(&updated)
|
||||
if err != nil {
|
||||
if cockroachutil.NeedsRetry(err) {
|
||||
continue
|
||||
} else if errs.Is(err, pgx.ErrNoRows) {
|
||||
break
|
||||
}
|
||||
return errs.New("updating value_attributions %w", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
total += updated
|
||||
if !more {
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated))
|
||||
break
|
||||
}
|
||||
log.Info("batch update complete", zap.Int("rows updated", updated), zap.Binary("last project id", nextProjectID), zap.String("last bucket name", string(nextBucket)))
|
||||
startProjectID = nextProjectID
|
||||
startBucket = nextBucket
|
||||
}
|
||||
log.Info("value_attributions migration complete", zap.Int("total rows updated", total))
|
||||
return nil
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user