feat(cli): helper for cleaning up stale records (ory#2406)

Closes ory#952

Co-authored-by: Patrik <zepatrik@users.noreply.github.com>

cmd/cleanup/root.go

@@ -0,0 +1,22 @@
+package cleanup
+
+import (
+	"github.com/spf13/cobra"
+
+	"github.com/ory/x/configx"
+)
+
+func NewCleanupCmd() *cobra.Command {
+	c := &cobra.Command{
+		Use:   "cleanup",
+		Short: "Various cleanup helpers",
+	}
+	configx.RegisterFlags(c.PersistentFlags())
+	return c
+}
+
+func RegisterCommandRecursive(parent *cobra.Command) {
+	c := NewCleanupCmd()
+	parent.AddCommand(c)
+	c.AddCommand(NewCleanupSQLCmd())
+}

cmd/cleanup/sql.go

@@ -0,0 +1,59 @@
+/*
+Copyright © 2019 NAME HERE <EMAIL ADDRESS>
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package cleanup
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/ory/x/cmdx"
+
+	"github.com/spf13/cobra"
+
+	"github.com/ory/kratos/driver/config"
+
+	"github.com/ory/kratos/cmd/cliclient"
+	"github.com/ory/x/configx"
+)
+
+// cleanupSqlCmd represents the sql command
+func NewCleanupSQLCmd() *cobra.Command {
+	c := &cobra.Command{
+		Use:   "sql <database-url>",
+		Short: "Cleanup sql database from expired flows and sessions",
+		Long: `Run this command as frequently as you need.
+It is recommended to run this command close to the SQL instance (e.g. same subnet) instead of over the public internet.
+This decreases risk of failure and decreases time required.
+You can read in the database URL using the -e flag, for example:
+	export DSN=...
+	kratos cleanup sql -e
+### WARNING ###
+Before running this command on an existing database, create a back up!
+`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			err := cliclient.NewCleanupHandler().CleanupSQL(cmd, args)
+			if err != nil {
+				fmt.Fprintln(cmd.OutOrStdout(), err)
+				return cmdx.FailSilently(cmd)
+			}
+			return nil
+		},
+	}
+
+	configx.RegisterFlags(c.PersistentFlags())
+	c.Flags().BoolP("read-from-env", "e", true, "If set, reads the database connection string from the environment variable DSN or config file key dsn.")
+	c.Flags().Duration(config.ViperKeyDatabaseCleanupSleepTables, time.Minute, "How long to wait between each table cleanup")
+	c.Flags().IntP(config.ViperKeyDatabaseCleanupBatchSize, "b", 100, "Set the number of records to be cleaned per run")
+	c.Flags().Duration("keep-last", 0, "Don't remove records younger than")
+	return c
+}

cmd/cleanup/sql_test.go

@@ -0,0 +1,24 @@
+package cleanup
+
+import (
+	"bytes"
+	"io/ioutil"
+	"strings"
+	"testing"
+)
+
+func Test_ExecuteCleanupFailedDSN(t *testing.T) {
+	cmd := NewCleanupSQLCmd()
+	b := bytes.NewBufferString("")
+	cmd.SetOut(b)
+	cmd.SetArgs([]string{"--read-from-env=false"})
+	cmd.Execute()
+	out, err := ioutil.ReadAll(b)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(string(out), "expected to get the DSN as an argument") {
+		t.Fatalf("expected \"%s\" got \"%s\"", "expected to get the DSN as an argument", string(out))
+	}
+	cmd.Execute()
+}

cmd/cliclient/cleanup.go

@@ -0,0 +1,60 @@
+package cliclient
+
+import (
+	"github.com/pkg/errors"
+
+	"github.com/ory/x/configx"
+
+	"github.com/spf13/cobra"
+
+	"github.com/ory/kratos/driver"
+	"github.com/ory/kratos/driver/config"
+	"github.com/ory/x/flagx"
+)
+
+type CleanupHandler struct{}
+
+func NewCleanupHandler() *CleanupHandler {
+	return &CleanupHandler{}
+}
+
+func (h *CleanupHandler) CleanupSQL(cmd *cobra.Command, args []string) error {
+	opts := []configx.OptionModifier{
+		configx.WithFlags(cmd.Flags()),
+		configx.SkipValidation(),
+	}
+
+	if !flagx.MustGetBool(cmd, "read-from-env") {
+		if len(args) != 1 {
+			return errors.New(`expected to get the DSN as an argument, or the "read-from-env" flag`)
+		}
+		opts = append(opts, configx.WithValue(config.ViperKeyDSN, args[0]))
+	}
+
+	d := driver.NewWithoutInit(
+		cmd.Context(),
+		cmd.ErrOrStderr(),
+		opts...,
+	)
+	if len(d.Config(cmd.Context()).DSN()) == 0 {
+		return errors.New(`required config value "dsn" was not set`)
+	}
+
+	err := d.Init(cmd.Context(), driver.SkipNetworkInit)
+	if err != nil {
+		return errors.Wrap(err, "An error occurred initializing cleanup")
+	}
+
+	keepLast := flagx.MustGetDuration(cmd, "keep-last")
+
+	err = d.Persister().CleanupDatabase(
+		cmd.Context(),
+		d.Config(cmd.Context()).DatabaseCleanupSleepTables(),
+		keepLast,
+		d.Config(cmd.Context()).DatabaseCleanupBatchSize())
+	if err != nil {
+		return errors.Wrap(err, "An error occurred while cleaning up expired data")
+	}
+
+	return nil
+}

cmd/root.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/ory/kratos/cmd/cleanup"
+
 	"github.com/ory/kratos/driver/config"
 
 	"github.com/ory/kratos/cmd/courier"
@@ -36,6 +38,7 @@ func NewRootCmd() (cmd *cobra.Command) {
 	cmd.AddCommand(identities.NewListCmd(cmd))
 	migrate.RegisterCommandRecursive(cmd)
 	serve.RegisterCommandRecursive(cmd)
+	cleanup.RegisterCommandRecursive(cmd)
 	remote.RegisterCommandRecursive(cmd)
 	cmd.AddCommand(identities.NewValidateCmd())
 	cmd.AddCommand(cmdx.Version(&config.Version, &config.Commit, &config.Date))

continuity/persistence.go

@@ -2,6 +2,7 @@ package continuity
 
 import (
 	"context"
+	"time"
 
 	"github.com/gofrs/uuid"
 )
@@ -14,4 +15,5 @@ type Persister interface {
 	SaveContinuitySession(ctx context.Context, c *Container) error
 	GetContinuitySession(ctx context.Context, id uuid.UUID) (*Container, error)
 	DeleteContinuitySession(ctx context.Context, id uuid.UUID) error
+	DeleteExpiredContinuitySessions(ctx context.Context, deleteOlder time.Time, pageSize int) error
 }

continuity/test/persistence.go

@@ -94,5 +94,28 @@ func TestPersister(ctx context.Context, p interface {
 				require.ErrorIs(t, err, sqlcon.ErrNoRows)
 			})
 		})
+
+		t.Run("case=cleanup", func(t *testing.T) {
+			id := x.NewUUID()
+			yesterday := time.Now().Add(-24 * time.Hour).UTC().Truncate(time.Second)
+			m := sqlxx.NullJSONRawMessage(`{"foo": "bar"}`)
+			expected := continuity.Container{Name: "foo", IdentityID: x.PointToUUID(createIdentity(t).ID),
+				ExpiresAt: yesterday,
+				Payload:   m,
+			}
+			expected.ID = id
+
+			t.Run("can cleanup", func(t *testing.T) {
+				require.NoError(t, p.SaveContinuitySession(ctx, &expected))
+
+				assert.EqualValues(t, id, expected.ID)
+				assert.EqualValues(t, nid, expected.NID)
+
+				require.NoError(t, p.DeleteExpiredContinuitySessions(ctx, time.Now(), 5))
+
+				_, err := p.GetContinuitySession(ctx, id)
+				require.Error(t, err)
+			})
+		})
 	}
 }

driver/config/config.go

@@ -153,6 +153,8 @@ const (
 	ViperKeyHasherArgon2ConfigDedicatedMemory                = "hashers.argon2.dedicated_memory"
 	ViperKeyHasherBcryptCost                                 = "hashers.bcrypt.cost"
 	ViperKeyCipherAlgorithm                                  = "ciphers.algorithm"
+	ViperKeyDatabaseCleanupSleepTables                       = "database.cleanup.sleep.tables"
+	ViperKeyDatabaseCleanupBatchSize                         = "database.cleanup.batch_size"
 	ViperKeyLinkLifespan                                     = "selfservice.methods.link.config.lifespan"
 	ViperKeyLinkBaseURL                                      = "selfservice.methods.link.config.base_url"
 	ViperKeyPasswordHaveIBeenPwnedHost                       = "selfservice.methods.password.config.haveibeenpwned_host"
@@ -1074,6 +1076,14 @@ func (p *Config) SelfServiceLinkMethodBaseURL() *url.URL {
 	return p.p.RequestURIF(ViperKeyLinkBaseURL, p.SelfPublicURL())
 }
 
+func (p *Config) DatabaseCleanupSleepTables() time.Duration {
+	return p.p.Duration(ViperKeyDatabaseCleanupSleepTables)
+}
+
+func (p *Config) DatabaseCleanupBatchSize() int {
+	return p.p.Int(ViperKeyDatabaseCleanupBatchSize)
+}
+
 func (p *Config) SelfServiceFlowRecoveryAfterHooks(strategy string) []SelfServiceHook {
 	return p.selfServiceHooks(HookStrategyKey(ViperKeySelfServiceRecoveryAfter, strategy))
 }

driver/config/config_test.go

@@ -48,7 +48,7 @@ func TestViperProvider(t *testing.T) {
 		p := config.MustNew(t, logrusx.New("", ""), os.Stderr,
 			configx.WithConfigFiles("stub/.kratos.yaml"))
 
-		t.Run("gourp=client config", func(t *testing.T) {
+		t.Run("group=client config", func(t *testing.T) {
 			assert.False(t, p.ClientHTTPNoPrivateIPRanges(), "Should not have private IP ranges disabled per default")
 			p.MustSet(config.ViperKeyClientHTTPNoPrivateIPRanges, true)
 			assert.True(t, p.ClientHTTPNoPrivateIPRanges(), "Should disallow private IP ranges if set")
@@ -1152,3 +1152,17 @@ func TestCourierTemplatesConfig(t *testing.T) {
 		assert.Equal(t, courierTemplateConfig, c.CourierTemplatesHelper(config.ViperKeyCourierTemplatesRecoveryValidEmail))
 	})
 }
+
+func TestCleanup(t *testing.T) {
+	p := config.MustNew(t, logrusx.New("", ""), os.Stderr,
+		configx.WithConfigFiles("stub/.kratos.yaml"))
+
+	t.Run("group=cleanup config", func(t *testing.T) {
+		assert.Equal(t, p.DatabaseCleanupSleepTables(), 1*time.Minute)
+		p.MustSet(config.ViperKeyDatabaseCleanupSleepTables, "1s")
+		assert.Equal(t, p.DatabaseCleanupSleepTables(), time.Second)
+		assert.Equal(t, p.DatabaseCleanupBatchSize(), 100)
+		p.MustSet(config.ViperKeyDatabaseCleanupBatchSize, "1")
+		assert.Equal(t, p.DatabaseCleanupBatchSize(), 1)
+	})
+}

embedx/config.schema.json

@@ -1435,6 +1435,49 @@
         }
       }
     },
+    "database":  {
+      "type": "object",
+      "title": "Database related configuration",
+      "description": "Miscellaneous settings used in database related tasks (cleanup, etc.)",
+      "properties": {
+        "cleanup": {
+          "type": "object",
+          "title": "Database cleanup settings",
+          "description": "Settings that controls how the database cleanup process is configured (delays, batch size, etc.)",
+          "properties": {
+            "batch_size" : {
+              "type": "integer",
+              "title": "Number of records to clean in one iteration",
+              "description": "Controls how many records should be purged from one table during database cleanup task",
+              "minimum": 1,
+              "default": 100
+            },
+            "sleep": {
+              "type": "object",
+              "title": "Delays between various database cleanup phases",
+              "description": "Configures delays between each step of the cleanup process. It is useful to tune the process so it will be efficient and performant.",
+              "properties": {
+                "tables": {
+                  "type": "string",
+                  "title": "Delay between each table cleanups",
+                  "description": "Controls the delay time between cleaning each table in one cleanup iteration",
+                  "pattern": "^[0-9]+(ns|us|ms|s|m|h)$",
+                  "default": "1m"
+                }
+              }
+            },
+            "older_than": {
+              "type": "string",
+              "title": "Remove records older than",
+              "description": "Controls how old records do we want to leave",
+              "pattern": "^[0-9]+(ns|us|ms|s|m|h)$",
+              "default": "0s"
+            }
+          }
+        }
+      },
+      "additionalProperties": false
+    },
     "dsn": {
       "type": "string",
       "title": "Data Source Name",

persistence/reference.go

@@ -2,6 +2,7 @@ package persistence
 
 import (
 	"context"
+	"time"
 
 	"github.com/ory/x/networkx"
 
@@ -43,6 +44,7 @@ type Persister interface {
 	link.RecoveryTokenPersister
 	link.VerificationTokenPersister
 
+	CleanupDatabase(context.Context, time.Duration, time.Duration, int) error
 	Close(context.Context) error
 	Ping() error
 	MigrationStatus(c context.Context) (popx.MigrationStatuses, error)