Speed up truncation of Spanner tables (used in integration tests) (#102278)

Speed up truncation of Spanner tables.
pull/102303/head
Peter Štibraný 4 months ago committed by GitHub
parent 3f270d1de0
commit 11552e5c57
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 20
      pkg/services/sqlstore/migrator/spanner_dialect.go

@ -14,6 +14,7 @@ import (
"cloud.google.com/go/spanner/admin/database/apiv1/databasepb" "cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
"github.com/googleapis/gax-go/v2" "github.com/googleapis/gax-go/v2"
spannerdriver "github.com/googleapis/go-sql-spanner" spannerdriver "github.com/googleapis/go-sql-spanner"
"github.com/grafana/dskit/concurrency"
"google.golang.org/api/option" "google.golang.org/api/option"
"google.golang.org/grpc" "google.golang.org/grpc"
"google.golang.org/grpc/codes" "google.golang.org/grpc/codes"
@ -132,13 +133,16 @@ func (s *SpannerDialect) ColStringNoPk(col *Column) string {
} }
func (s *SpannerDialect) TruncateDBTables(engine *xorm.Engine) error { func (s *SpannerDialect) TruncateDBTables(engine *xorm.Engine) error {
tables, err := engine.DBMetas() // Get tables names only, no columns or indexes.
tables, err := engine.Dialect().GetTables()
if err != nil { if err != nil {
return err return err
} }
sess := engine.NewSession() sess := engine.NewSession()
defer sess.Close() defer sess.Close()
var statements []string
for _, table := range tables { for _, table := range tables {
switch table.Name { switch table.Name {
case "": case "":
@ -147,17 +151,17 @@ func (s *SpannerDialect) TruncateDBTables(engine *xorm.Engine) error {
continue continue
case "dashboard_acl": case "dashboard_acl":
// keep default dashboard permissions // keep default dashboard permissions
if _, err := sess.Exec(fmt.Sprintf("DELETE FROM %v WHERE dashboard_id != -1 AND org_id != -1;", s.Quote(table.Name))); err != nil { statements = append(statements, fmt.Sprintf("DELETE FROM %v WHERE dashboard_id != -1 AND org_id != -1;", s.Quote(table.Name)))
return fmt.Errorf("failed to truncate table %q: %w", table.Name, err)
}
default: default:
if _, err := sess.Exec(fmt.Sprintf("DELETE FROM %v WHERE TRUE;", s.Quote(table.Name))); err != nil { statements = append(statements, fmt.Sprintf("DELETE FROM %v WHERE TRUE;", s.Quote(table.Name)))
return fmt.Errorf("failed to truncate table %q: %w", table.Name, err)
}
} }
} }
return nil // Run statements concurrently.
return concurrency.ForEachJob(context.Background(), len(statements), 10, func(ctx context.Context, idx int) error {
_, err := sess.Exec(statements[idx])
return err
})
} }
// CleanDB drops all existing tables and their indexes. // CleanDB drops all existing tables and their indexes.

Loading…
Cancel
Save