From f98edb42a91f15c5b8a641d0b4a5264c6d5bc4bc Mon Sep 17 00:00:00 2001 From: Aditya Thebe Date: Fri, 6 Dec 2024 12:15:43 +0545 Subject: [PATCH] feat: benchmarks --- Makefile | 4 + bench/bench_test.go | 129 +++++++++++++++ bench/benchmark.md | 111 +++++++++++++ bench/utils_test.go | 186 ++++++++++++++++++++++ migrate/migrate.go | 6 +- models/permission.go | 8 +- models/permission_test.go | 4 +- schema/apply.go | 2 +- shutdown/shutdown.go | 7 +- tests/config_generator_test.go | 13 +- tests/config_relationship_test.go | 13 +- tests/config_type_summary_test.go | 15 +- tests/{ => generator}/config_generator.go | 2 +- tests/setup/common.go | 114 ++++++++----- views/034_rls_enable.sql | 57 ++----- views/035_rls_disable.sql | 4 + 16 files changed, 560 insertions(+), 115 deletions(-) create mode 100644 bench/bench_test.go create mode 100644 bench/benchmark.md create mode 100644 bench/utils_test.go rename tests/{ => generator}/config_generator.go (99%) diff --git a/Makefile b/Makefile index e68b43c4..d4aec425 100644 --- a/Makefile +++ b/Makefile @@ -6,6 +6,10 @@ ginkgo: test: ginkgo ginkgo -r -v +.PHONY: bench +bench: + go test -bench=. -benchtime=10s -timeout 30m github.com/flanksource/duty/bench + fmt: go fmt ./... diff --git a/bench/bench_test.go b/bench/bench_test.go new file mode 100644 index 00000000..4a7e9a28 --- /dev/null +++ b/bench/bench_test.go @@ -0,0 +1,129 @@ +package bench_test + +import ( + "fmt" + "os" + "testing" + + "github.com/flanksource/commons/logger" + "github.com/flanksource/duty/context" + "github.com/flanksource/duty/shutdown" + "github.com/flanksource/duty/tests/setup" +) + +type DistinctBenchConfig struct { + // view/table name + relation string + + // optional column to fetch. + // when left empty all columns are fetched (this is left empty for views with single column) + column string +} + +var benchConfigs = []DistinctBenchConfig{ + {"catalog_changes", "change_type"}, + {"config_changes", "change_type"}, + {"config_detail", "type"}, + {"config_names", "type"}, + {"config_summary", "type"}, + {"configs", "type"}, + + // These are single column views + {"analysis_types", ""}, + {"analyzer_types", ""}, + {"change_types", ""}, + {"config_classes", ""}, + {"config_types", ""}, +} + +var ( + testCtx context.Context + connUrl string + + // number of total configs in the database + testSizes = []int{10_000, 25_000, 50_000, 100_000} +) + +func setupTestDB(dbPath string) error { + logger.Infof("using %q as the pg data dir", dbPath) + os.Setenv(setup.DUTY_DB_DATA_DIR, dbPath) + + shutdown.AddHookWithPriority("delete data dir", shutdown.PriorityCritical+1, func() { + if err := os.RemoveAll(dbPath); err != nil { + logger.Errorf("failed to delete data dir: %v", err) + } + }) + + var err error + testCtx, err = setup.SetupDB("test", + setup.WithoutDummyData, // we generate the dummy data + setup.WithoutRLS, // start without RLS + ) + if err != nil { + return fmt.Errorf("failed to setup db: %v", err) + } + connUrl = testCtx.Value("db_url").(string) + return nil +} + +func TestMain(m *testing.M) { + shutdown.WaitForSignal() + + // Create a fixed postgres data directory + dbDataPath, err := os.CreateTemp("", "bench-pg-dir-*") + if err != nil { + shutdown.ShutdownAndExit(1, "failed to create temp dir for db") + } + + if err := setupTestDB(dbDataPath.Name()); err != nil { + shutdown.ShutdownAndExit(1, err.Error()) + } + + result := m.Run() + shutdown.ShutdownAndExit(result, "exiting ...") +} + +func BenchmarkMain(b *testing.B) { + for _, size := range testSizes { + resetPG(b, false) + _, err := setupConfigsForSize(testCtx, size) + if err != nil { + b.Fatalf("failed to setup configs for size %d: %v", size, err) + } + + b.Run(fmt.Sprintf("Sample-%d", size), func(b *testing.B) { + for _, config := range benchConfigs { + runBenchmark(b, config) + } + }) + } +} + +func runBenchmark(b *testing.B, config DistinctBenchConfig) { + b.Run(config.relation, func(b *testing.B) { + for _, rls := range []bool{false, true} { + resetPG(b, rls) + name := "Without RLS" + if rls { + name = "With RLS" + } + + b.Run(name, func(b *testing.B) { + for i := 0; i < b.N; i++ { + if rls { + b.StopTimer() + tags := sampleTags[i%len(sampleTags)] + if err := setupRLSPayload(testCtx, tags); err != nil { + b.Fatalf("failed to setup rls payload with tag(%v): %v", tags, err) + } + b.StartTimer() + } + + if err := fetchView(testCtx, config.relation, config.column); err != nil { + b.Fatalf("%v", err) + } + } + }) + } + }) +} diff --git a/bench/benchmark.md b/bench/benchmark.md new file mode 100644 index 00000000..b59a9a42 --- /dev/null +++ b/bench/benchmark.md @@ -0,0 +1,111 @@ +# PostgreSQL RLS Benchmark + +## Running Benchmarks + +query duration to fetch 10k, 25k, 50k and 100k config items in random are recorded. + +```bash +make bench +``` + +## Results + +Took ~23 minutes + +``` +goos: linux +goarch: amd64 +pkg: github.com/flanksource/duty/bench +cpu: Intel(R) Core(TM) i9-14900K +BenchmarkMain/Sample-10000/catalog_changes/Without_RLS-32 7210 1618398 ns/op +BenchmarkMain/Sample-10000/catalog_changes/With_RLS-32 468 25696464 ns/op +BenchmarkMain/Sample-10000/config_changes/Without_RLS-32 7225 1662744 ns/op +BenchmarkMain/Sample-10000/config_changes/With_RLS-32 472 25400088 ns/op +BenchmarkMain/Sample-10000/config_detail/Without_RLS-32 9204 1258744 ns/op +BenchmarkMain/Sample-10000/config_detail/With_RLS-32 1110 10922483 ns/op +BenchmarkMain/Sample-10000/config_names/Without_RLS-32 1630 7167192 ns/op +BenchmarkMain/Sample-10000/config_names/With_RLS-32 1068 11294074 ns/op +BenchmarkMain/Sample-10000/config_summary/Without_RLS-32 691 17307363 ns/op +BenchmarkMain/Sample-10000/config_summary/With_RLS-32 134 88561638 ns/op +BenchmarkMain/Sample-10000/configs/Without_RLS-32 5647 2071318 ns/op +BenchmarkMain/Sample-10000/configs/With_RLS-32 1111 10701505 ns/op +BenchmarkMain/Sample-10000/analysis_types/Without_RLS-32 9230 1267578 ns/op +BenchmarkMain/Sample-10000/analysis_types/With_RLS-32 9554 1249278 ns/op +BenchmarkMain/Sample-10000/analyzer_types/Without_RLS-32 9843 1186921 ns/op +BenchmarkMain/Sample-10000/analyzer_types/With_RLS-32 9657 1208368 ns/op +BenchmarkMain/Sample-10000/change_types/Without_RLS-32 7399 1602853 ns/op +BenchmarkMain/Sample-10000/change_types/With_RLS-32 7578 1618275 ns/op +BenchmarkMain/Sample-10000/config_classes/Without_RLS-32 9602 1053916 ns/op +BenchmarkMain/Sample-10000/config_classes/With_RLS-32 1102 10601675 ns/op +BenchmarkMain/Sample-10000/config_types/Without_RLS-32 9938 1220556 ns/op +BenchmarkMain/Sample-10000/config_types/With_RLS-32 1132 10687988 ns/op + +BenchmarkMain/Sample-25000/catalog_changes/Without_RLS-32 3189 3777448 ns/op +BenchmarkMain/Sample-25000/catalog_changes/With_RLS-32 199 59728301 ns/op +BenchmarkMain/Sample-25000/config_changes/Without_RLS-32 3106 3796288 ns/op +BenchmarkMain/Sample-25000/config_changes/With_RLS-32 202 59201120 ns/op +BenchmarkMain/Sample-25000/config_detail/Without_RLS-32 3986 2825246 ns/op +BenchmarkMain/Sample-25000/config_detail/With_RLS-32 472 25408187 ns/op +BenchmarkMain/Sample-25000/config_names/Without_RLS-32 712 16344633 ns/op +BenchmarkMain/Sample-25000/config_names/With_RLS-32 447 26696849 ns/op +BenchmarkMain/Sample-25000/config_summary/Without_RLS-32 274 43747108 ns/op +BenchmarkMain/Sample-25000/config_summary/With_RLS-32 45 242723303 ns/op +BenchmarkMain/Sample-25000/configs/Without_RLS-32 2466 4885648 ns/op +BenchmarkMain/Sample-25000/configs/With_RLS-32 470 25306394 ns/op +BenchmarkMain/Sample-25000/analysis_types/Without_RLS-32 4042 2932464 ns/op +BenchmarkMain/Sample-25000/analysis_types/With_RLS-32 4096 2936163 ns/op +BenchmarkMain/Sample-25000/analyzer_types/Without_RLS-32 4293 2776851 ns/op +BenchmarkMain/Sample-25000/analyzer_types/With_RLS-32 4180 2812037 ns/op +BenchmarkMain/Sample-25000/change_types/Without_RLS-32 3093 3864532 ns/op +BenchmarkMain/Sample-25000/change_types/With_RLS-32 3123 3806187 ns/op +BenchmarkMain/Sample-25000/config_classes/Without_RLS-32 4693 2435089 ns/op +BenchmarkMain/Sample-25000/config_classes/With_RLS-32 476 25211551 ns/op +BenchmarkMain/Sample-25000/config_types/Without_RLS-32 4164 2861676 ns/op +BenchmarkMain/Sample-25000/config_types/With_RLS-32 476 25352067 ns/op + +BenchmarkMain/Sample-50000/catalog_changes/Without_RLS-32 1560 7545395 ns/op +BenchmarkMain/Sample-50000/catalog_changes/With_RLS-32 100 117274979 ns/op +BenchmarkMain/Sample-50000/config_changes/Without_RLS-32 1573 7551748 ns/op +BenchmarkMain/Sample-50000/config_changes/With_RLS-32 99 117770448 ns/op +BenchmarkMain/Sample-50000/config_detail/Without_RLS-32 2101 5593338 ns/op +BenchmarkMain/Sample-50000/config_detail/With_RLS-32 242 49418844 ns/op +BenchmarkMain/Sample-50000/config_names/Without_RLS-32 378 31770900 ns/op +BenchmarkMain/Sample-50000/config_names/With_RLS-32 226 52552379 ns/op +BenchmarkMain/Sample-50000/config_summary/Without_RLS-32 128 90894472 ns/op +BenchmarkMain/Sample-50000/config_summary/With_RLS-32 25 473002784 ns/op +BenchmarkMain/Sample-50000/configs/Without_RLS-32 1251 9464835 ns/op +BenchmarkMain/Sample-50000/configs/With_RLS-32 238 49838197 ns/op +BenchmarkMain/Sample-50000/analysis_types/Without_RLS-32 2052 5801409 ns/op +BenchmarkMain/Sample-50000/analysis_types/With_RLS-32 2121 5712487 ns/op +BenchmarkMain/Sample-50000/analyzer_types/Without_RLS-32 2216 5442149 ns/op +BenchmarkMain/Sample-50000/analyzer_types/With_RLS-32 2169 5515249 ns/op +BenchmarkMain/Sample-50000/change_types/Without_RLS-32 1592 7552502 ns/op +BenchmarkMain/Sample-50000/change_types/With_RLS-32 1521 7634041 ns/op +BenchmarkMain/Sample-50000/config_classes/Without_RLS-32 2442 4780004 ns/op +BenchmarkMain/Sample-50000/config_classes/With_RLS-32 241 49653432 ns/op +BenchmarkMain/Sample-50000/config_types/Without_RLS-32 2145 5558880 ns/op +BenchmarkMain/Sample-50000/config_types/With_RLS-32 241 49518770 ns/op + +BenchmarkMain/Sample-100000/catalog_changes/Without_RLS-32 668 15792969 ns/op +BenchmarkMain/Sample-100000/catalog_changes/With_RLS-32 50 236585972 ns/op +BenchmarkMain/Sample-100000/config_changes/Without_RLS-32 670 15857288 ns/op +BenchmarkMain/Sample-100000/config_changes/With_RLS-32 49 237727030 ns/op +BenchmarkMain/Sample-100000/config_detail/Without_RLS-32 1060 11282955 ns/op +BenchmarkMain/Sample-100000/config_detail/With_RLS-32 121 98802558 ns/op +BenchmarkMain/Sample-100000/config_names/Without_RLS-32 175 68280940 ns/op +BenchmarkMain/Sample-100000/config_names/With_RLS-32 100 105502052 ns/op +BenchmarkMain/Sample-100000/config_summary/Without_RLS-32 67 169628955 ns/op +BenchmarkMain/Sample-100000/config_summary/With_RLS-32 12 984132710 ns/op +BenchmarkMain/Sample-100000/configs/Without_RLS-32 609 19589287 ns/op +BenchmarkMain/Sample-100000/configs/With_RLS-32 120 99833450 ns/op +BenchmarkMain/Sample-100000/analysis_types/Without_RLS-32 1039 11434234 ns/op +BenchmarkMain/Sample-100000/analysis_types/With_RLS-32 1064 11451964 ns/op +BenchmarkMain/Sample-100000/analyzer_types/Without_RLS-32 1110 10675073 ns/op +BenchmarkMain/Sample-100000/analyzer_types/With_RLS-32 1114 10854744 ns/op +BenchmarkMain/Sample-100000/change_types/Without_RLS-32 669 15856671 ns/op +BenchmarkMain/Sample-100000/change_types/With_RLS-32 668 16162332 ns/op +BenchmarkMain/Sample-100000/config_classes/Without_RLS-32 1261 9487116 ns/op +BenchmarkMain/Sample-100000/config_classes/With_RLS-32 121 98950319 ns/op +BenchmarkMain/Sample-100000/config_types/Without_RLS-32 1060 11280585 ns/op +BenchmarkMain/Sample-100000/config_types/With_RLS-32 121 99579524 ns/op +``` diff --git a/bench/utils_test.go b/bench/utils_test.go new file mode 100644 index 00000000..3a022dc7 --- /dev/null +++ b/bench/utils_test.go @@ -0,0 +1,186 @@ +package bench_test + +import ( + "crypto/rand" + "database/sql" + "encoding/json" + "errors" + "fmt" + "math/big" + "testing" + + "github.com/google/uuid" + + "github.com/flanksource/duty" + "github.com/flanksource/duty/api" + "github.com/flanksource/duty/context" + "github.com/flanksource/duty/models" + "github.com/flanksource/duty/shutdown" + pkgGenerator "github.com/flanksource/duty/tests/generator" + "github.com/flanksource/duty/tests/setup" +) + +var sampleTags = []map[string]string{ + {"cluster": "homelab"}, + {"cluster": "azure"}, + {"cluster": "aws"}, + {"cluster": "gcp"}, + {"cluster": "demo"}, + {"region": "eu-west-1"}, + {"region": "eu-west-2"}, + {"region": "us-east-1"}, + {"region": "us-east-2"}, +} + +func getRandomTag() map[string]string { + max := big.NewInt(int64(len(sampleTags))) + n, err := rand.Int(rand.Reader, max) + if err != nil { + panic(err) + } + + return sampleTags[n.Int64()] +} + +func generateConfigItems(ctx context.Context, count int) error { + for { + var totalConfigs int64 + if err := ctx.DB().Table("config_items").Count(&totalConfigs).Error; err != nil { + return err + } + + if totalConfigs > int64(count) { + break + } + + generator := pkgGenerator.ConfigGenerator{ + Nodes: pkgGenerator.ConfigTypeRequirements{ + Count: 3, + }, + Namespaces: pkgGenerator.ConfigTypeRequirements{ + Count: 10, + }, + DeploymentPerNamespace: pkgGenerator.ConfigTypeRequirements{ + Count: 10, + }, + ReplicaSetPerDeployment: pkgGenerator.ConfigTypeRequirements{ + Count: 2, + Deleted: 1, + }, + PodsPerReplicaSet: pkgGenerator.ConfigTypeRequirements{ + Count: 2, + NumChangesPerConfig: 1, + NumInsightsPerConfig: 2, + }, + Tags: getRandomTag(), + } + + generator.GenerateKubernetes() + if err := generator.Save(ctx.DB()); err != nil { + return err + } + } + + return nil +} + +func fetchView(ctx context.Context, view, column string) error { + selectColumns := "*" + if column != "" { + selectColumns = fmt.Sprintf("DISTINCT %s", column) // use distinct so we don't fetch a lot of rows + } + + var result []string + if err := ctx.DB().Select(selectColumns).Table(view).Scan(&result).Error; err != nil { + return fmt.Errorf("failed to fetch distinct types for %s: %w", view, err) + } + + return nil +} + +func setupRLSPayload(ctx context.Context, tags map[string]string) error { + bb, err := json.Marshal(tags) + if err != nil { + return err + } + + sql := fmt.Sprintf(`SET request.jwt.claims = '{"tags": [%s]}'`, string(bb)) + if err := ctx.DB().Exec(sql).Error; err != nil { + return err + } + + var jwt string + if err := ctx.DB().Raw(`SELECT current_setting('request.jwt.claims', TRUE)`).Scan(&jwt).Error; err != nil { + return err + } + + if err := ctx.DB().Exec(`SET role = 'postgrest_api'`).Error; err != nil { + return err + } + + if err := verifyRLSPayload(ctx); err != nil { + return err + } + + return nil +} + +func verifyRLSPayload(ctx context.Context) error { + var jwt sql.NullString + if err := ctx.DB().Raw(`SELECT current_setting('request.jwt.claims', TRUE)`).Scan(&jwt).Error; err != nil { + return err + } + + if !jwt.Valid { + return errors.New("jwt parameter not set") + } + + var role string + if err := ctx.DB().Raw(`SELECT CURRENT_USER`).Scan(&role).Error; err != nil { + return err + } + + if role != "postgrest_api" { + return errors.New("role is not set") + } + + return nil +} + +func setupConfigsForSize(ctx context.Context, size int) ([]uuid.UUID, error) { + if err := generateConfigItems(ctx, size); err != nil { + return nil, fmt.Errorf("failed to generate configs: %w", err) + } + + var configIDs []uuid.UUID + if err := ctx.DB().Select("id").Model(&models.ConfigItem{}).Find(&configIDs).Error; err != nil { + return nil, err + } + + return configIDs, nil +} + +func resetPG(b *testing.B, rlsEnable bool) { + if err := setup.RestartEmbeddedPG(); err != nil { + b.Fatalf("failed to restart embedded pg") + } + + if rlsEnable { + if err := duty.Migrate(duty.EnableRLS(duty.RunMigrations(api.NewConfig(connUrl)))); err != nil { + b.Fatalf("failed to enable rls: %v", err) + } + + if err := setupRLSPayload(testCtx, getRandomTag()); err != nil { + b.Fatalf("failed to setup tags after restart: %v", err) + } + } else { + if err := duty.Migrate(duty.RunMigrations(api.NewConfig(connUrl))); err != nil { + b.Fatalf("failed to enable rls: %v", err) + } + } + + // This is required due to a bug in how we handle rls_enable / disable scripts. + if err := testCtx.DB().Exec("DELETE FROM migration_logs").Error; err != nil { + shutdown.ShutdownAndExit(1, fmt.Sprintf("failed to delete migration logs: %v", err)) + } +} diff --git a/migrate/migrate.go b/migrate/migrate.go index 540466bf..ef0a5ed4 100644 --- a/migrate/migrate.go +++ b/migrate/migrate.go @@ -46,7 +46,7 @@ func RunMigrations(pool *sql.DB, config api.Config) error { if err := row.Scan(&name); err != nil { return fmt.Errorf("failed to get current database: %w", err) } - l.Infof("Migrating database %s", name) + l.V(1).Infof("Migrating database %s", name) if err := createMigrationLogTable(pool); err != nil { return fmt.Errorf("failed to create migration log table: %w", err) @@ -57,7 +57,7 @@ func RunMigrations(pool *sql.DB, config api.Config) error { return fmt.Errorf("failed to get executable scripts: %w", err) } - l.V(3).Infof("Running scripts") + l.V(3).Infof("Running %d scripts (functions)", len(allFunctions)) if err := runScripts(pool, allFunctions, config.SkipMigrationFiles); err != nil { return fmt.Errorf("failed to run scripts: %w", err) } @@ -73,7 +73,7 @@ func RunMigrations(pool *sql.DB, config api.Config) error { return fmt.Errorf("failed to apply schema migrations: %w", err) } - l.V(3).Infof("Running scripts for views") + l.V(3).Infof("Running %d scripts (views)", len(allViews)) if err := runScripts(pool, allViews, config.SkipMigrationFiles); err != nil { return fmt.Errorf("failed to run scripts for views: %w", err) } diff --git a/models/permission.go b/models/permission.go index c5cb5155..e04ede3c 100644 --- a/models/permission.go +++ b/models/permission.go @@ -59,19 +59,19 @@ func (t *Permission) Condition() string { var rule []string if t.ComponentID != nil { - rule = append(rule, fmt.Sprintf("r.obj.component != undefined && r.obj.component.id == %q", t.ComponentID.String())) + rule = append(rule, fmt.Sprintf("r.obj.component.id == %q", t.ComponentID.String())) } if t.ConfigID != nil { - rule = append(rule, fmt.Sprintf("r.obj.config != undefined && r.obj.config.id == %q", t.ConfigID.String())) + rule = append(rule, fmt.Sprintf("r.obj.config.id == %q", t.ConfigID.String())) } if t.CanaryID != nil { - rule = append(rule, fmt.Sprintf("r.obj.canary != undefined && r.obj.canary.id == %q", t.CanaryID.String())) + rule = append(rule, fmt.Sprintf("r.obj.canary.id == %q", t.CanaryID.String())) } if t.PlaybookID != nil { - rule = append(rule, fmt.Sprintf("r.obj.playbook != undefined && r.obj.playbook.id == %q", t.PlaybookID.String())) + rule = append(rule, fmt.Sprintf("r.obj.playbook.id == %q", t.PlaybookID.String())) } if len(t.Agents) > 0 || len(t.Tags) > 0 { diff --git a/models/permission_test.go b/models/permission_test.go index 6829d0cc..0b894726 100644 --- a/models/permission_test.go +++ b/models/permission_test.go @@ -19,7 +19,7 @@ func TestPermission_Condition(t *testing.T) { perm: Permission{ PlaybookID: lo.ToPtr(uuid.MustParse("33333333-3333-3333-3333-333333333333")), }, - expected: `r.obj.playbook != undefined && r.obj.playbook.id == "33333333-3333-3333-3333-333333333333"`, + expected: `r.obj.playbook.id == "33333333-3333-3333-3333-333333333333"`, }, { name: "Multiple fields II", @@ -27,7 +27,7 @@ func TestPermission_Condition(t *testing.T) { ConfigID: lo.ToPtr(uuid.MustParse("88888888-8888-8888-8888-888888888888")), PlaybookID: lo.ToPtr(uuid.MustParse("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")), }, - expected: `r.obj.config != undefined && r.obj.config.id == "88888888-8888-8888-8888-888888888888" && r.obj.playbook != undefined && r.obj.playbook.id == "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"`, + expected: `r.obj.config.id == "88888888-8888-8888-8888-888888888888" && r.obj.playbook.id == "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"`, }, { name: "No fields set", diff --git a/schema/apply.go b/schema/apply.go index 83adb16c..dcdb96c8 100644 --- a/schema/apply.go +++ b/schema/apply.go @@ -87,7 +87,7 @@ func Apply(ctx context.Context, connection string) error { return fmt.Errorf("applied %d changes and then failed: %w", len(changes), err) } - log.Infof("Applied %d changes", len(changes)) + log.V(1).Infof("Applied %d changes", len(changes)) return nil } diff --git a/shutdown/shutdown.go b/shutdown/shutdown.go index 872802a7..ad91b51d 100644 --- a/shutdown/shutdown.go +++ b/shutdown/shutdown.go @@ -48,7 +48,12 @@ var Shutdown = sync.OnceFunc(func() { func ShutdownAndExit(code int, msg string) { Shutdown() - logger.StandardLogger().WithSkipReportLevel(1).Errorf(msg) + if code == 0 { + logger.StandardLogger().WithSkipReportLevel(1).Infof(msg) + } else { + logger.StandardLogger().WithSkipReportLevel(1).Errorf(msg) + } + os.Exit(code) } diff --git a/tests/config_generator_test.go b/tests/config_generator_test.go index f57fb5fb..6c86e57f 100644 --- a/tests/config_generator_test.go +++ b/tests/config_generator_test.go @@ -2,26 +2,27 @@ package tests import ( "github.com/flanksource/duty/models" + pkgGenerator "github.com/flanksource/duty/tests/generator" "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) var _ = ginkgo.Describe("Config Generator", ginkgo.Ordered, func() { - generator := ConfigGenerator{ - Nodes: ConfigTypeRequirements{ + generator := pkgGenerator.ConfigGenerator{ + Nodes: pkgGenerator.ConfigTypeRequirements{ Count: 3, }, - Namespaces: ConfigTypeRequirements{ + Namespaces: pkgGenerator.ConfigTypeRequirements{ Count: 2, }, - DeploymentPerNamespace: ConfigTypeRequirements{ + DeploymentPerNamespace: pkgGenerator.ConfigTypeRequirements{ Count: 1, }, - ReplicaSetPerDeployment: ConfigTypeRequirements{ + ReplicaSetPerDeployment: pkgGenerator.ConfigTypeRequirements{ Count: 4, Deleted: 3, }, - PodsPerReplicaSet: ConfigTypeRequirements{ + PodsPerReplicaSet: pkgGenerator.ConfigTypeRequirements{ Count: 1, NumChangesPerConfig: 5, NumInsightsPerConfig: 2, diff --git a/tests/config_relationship_test.go b/tests/config_relationship_test.go index 43fcf336..62262f66 100644 --- a/tests/config_relationship_test.go +++ b/tests/config_relationship_test.go @@ -7,6 +7,7 @@ import ( "github.com/flanksource/duty/models" "github.com/flanksource/duty/query" "github.com/flanksource/duty/tests/fixtures/dummy" + pkgGenerator "github.com/flanksource/duty/tests/generator" "github.com/flanksource/duty/types" "github.com/flanksource/duty/upstream" "github.com/google/uuid" @@ -464,12 +465,12 @@ var _ = ginkgo.Describe("config relationship deletion test", func() { }) var _ = ginkgo.Describe("config relationship depth", ginkgo.Ordered, func() { - generator := ConfigGenerator{ - Nodes: ConfigTypeRequirements{Count: 3}, - Namespaces: ConfigTypeRequirements{Count: 2}, - DeploymentPerNamespace: ConfigTypeRequirements{Count: 2}, - ReplicaSetPerDeployment: ConfigTypeRequirements{Count: 4, Deleted: 3}, - PodsPerReplicaSet: ConfigTypeRequirements{Count: 1, NumChangesPerConfig: 5, NumInsightsPerConfig: 2}, + generator := pkgGenerator.ConfigGenerator{ + Nodes: pkgGenerator.ConfigTypeRequirements{Count: 3}, + Namespaces: pkgGenerator.ConfigTypeRequirements{Count: 2}, + DeploymentPerNamespace: pkgGenerator.ConfigTypeRequirements{Count: 2}, + ReplicaSetPerDeployment: pkgGenerator.ConfigTypeRequirements{Count: 4, Deleted: 3}, + PodsPerReplicaSet: pkgGenerator.ConfigTypeRequirements{Count: 1, NumChangesPerConfig: 5, NumInsightsPerConfig: 2}, Tags: map[string]string{ "test": "true", }, diff --git a/tests/config_type_summary_test.go b/tests/config_type_summary_test.go index cc294533..23bf9ff3 100644 --- a/tests/config_type_summary_test.go +++ b/tests/config_type_summary_test.go @@ -8,6 +8,7 @@ import ( "github.com/flanksource/duty/job" "github.com/flanksource/duty/models" "github.com/flanksource/duty/query" + "github.com/flanksource/duty/tests/generator" "github.com/flanksource/duty/types" ginkgo "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -96,13 +97,13 @@ var _ = ginkgo.Describe("Check config_class_summary view", ginkgo.Ordered, func( }) ginkgo.It("Should query config summary by type", func() { - gen := ConfigGenerator{} - gen.GenerateConfigItem("Test::type-A", "healthy", nil, nil, ConfigTypeRequirements{NumChangesPerConfig: 4, NumInsightsPerConfig: 3}) - gen.GenerateConfigItem("Test::type-A", "healthy", nil, nil, ConfigTypeRequirements{NumChangesPerConfig: 1, NumInsightsPerConfig: 2}) - gen.GenerateConfigItem("Test::type-A", "unhealthy", nil, nil, ConfigTypeRequirements{NumChangesPerConfig: 1, NumInsightsPerConfig: 2}) - gen.GenerateConfigItem("Test::type-B", "healthy", nil, nil, ConfigTypeRequirements{NumChangesPerConfig: 5, NumInsightsPerConfig: 1}) - gen.GenerateConfigItem("Test::type-B", "healthy", nil, nil, ConfigTypeRequirements{NumChangesPerConfig: 1, NumInsightsPerConfig: 3}) - gen.GenerateConfigItem("Test::type-C", "unhealthy", nil, nil, ConfigTypeRequirements{NumChangesPerConfig: 0, NumInsightsPerConfig: 0}) + gen := generator.ConfigGenerator{} + gen.GenerateConfigItem("Test::type-A", "healthy", nil, nil, generator.ConfigTypeRequirements{NumChangesPerConfig: 4, NumInsightsPerConfig: 3}) + gen.GenerateConfigItem("Test::type-A", "healthy", nil, nil, generator.ConfigTypeRequirements{NumChangesPerConfig: 1, NumInsightsPerConfig: 2}) + gen.GenerateConfigItem("Test::type-A", "unhealthy", nil, nil, generator.ConfigTypeRequirements{NumChangesPerConfig: 1, NumInsightsPerConfig: 2}) + gen.GenerateConfigItem("Test::type-B", "healthy", nil, nil, generator.ConfigTypeRequirements{NumChangesPerConfig: 5, NumInsightsPerConfig: 1}) + gen.GenerateConfigItem("Test::type-B", "healthy", nil, nil, generator.ConfigTypeRequirements{NumChangesPerConfig: 1, NumInsightsPerConfig: 3}) + gen.GenerateConfigItem("Test::type-C", "unhealthy", nil, nil, generator.ConfigTypeRequirements{NumChangesPerConfig: 0, NumInsightsPerConfig: 0}) for _, item := range gen.Generated.Configs { DefaultContext.DB().Create(&item) } diff --git a/tests/config_generator.go b/tests/generator/config_generator.go similarity index 99% rename from tests/config_generator.go rename to tests/generator/config_generator.go index 05e72e8f..835bc30c 100644 --- a/tests/config_generator.go +++ b/tests/generator/config_generator.go @@ -1,4 +1,4 @@ -package tests +package generator import ( "fmt" diff --git a/tests/setup/common.go b/tests/setup/common.go index 9c756adf..9d031fdb 100644 --- a/tests/setup/common.go +++ b/tests/setup/common.go @@ -1,7 +1,6 @@ package setup import ( - gocontext "context" "database/sql" "fmt" "net/http" @@ -18,6 +17,7 @@ import ( "github.com/flanksource/duty/context" "github.com/flanksource/duty/job" "github.com/flanksource/duty/models" + "github.com/flanksource/duty/shutdown" "github.com/flanksource/duty/telemetry" "github.com/flanksource/duty/tests/fixtures/dummy" "github.com/labstack/echo/v4" @@ -35,6 +35,14 @@ import ( _ "github.com/spf13/cobra" ) +// Env vars for embedded db +const ( + DUTY_DB_CREATE = "DUTY_DB_CREATE" + DUTY_DB_DATA_DIR = "DUTY_DB_DATA_DIR" + DUTY_DB_URL = "DUTY_DB_URL" + TEST_DB_PORT = "TEST_DB_PORT" +) + var DefaultContext context.Context var postgresServer *embeddedPG.EmbeddedPostgres @@ -42,7 +50,14 @@ var dummyData dummy.DummyData var PgUrl string var postgresDBUrl string -var dbName = "test" + +func RestartEmbeddedPG() error { + if err := postgresServer.Stop(); err != nil { + return err + } + + return postgresServer.Start() +} func init() { logger.UseSlog() @@ -80,14 +95,12 @@ func MustDB() *sql.DB { return db } +var WithoutRLS = "rls_disabled" var WithoutDummyData = "without_dummy_data" var WithExistingDatabase = "with_existing_database" -var recreateDatabase = os.Getenv("DUTY_DB_CREATE") != "false" - -var ShutdownHooks []func(ctx gocontext.Context) error +var recreateDatabase = os.Getenv(DUTY_DB_CREATE) != "false" func findFileInPath(filename string, depth int) string { - if !path.IsAbs(filename) { cwd, _ := os.Getwd() filename = path.Join(cwd, filename) @@ -107,15 +120,25 @@ func findFileInPath(filename string, depth int) string { } func BeforeSuiteFn(args ...interface{}) context.Context { + ctx, err := SetupDB("test", args...) + if err != nil { + shutdown.ShutdownAndExit(1, fmt.Sprintf("failed to setup db: %v", err)) + } + + DefaultContext = ctx + return ctx +} + +func SetupDB(dbName string, args ...interface{}) (context.Context, error) { if err := properties.LoadFile(findFileInPath("test.properties", 2)); err != nil { logger.Errorf("Failed to load test properties: %v", err) } - ShutdownHooks = append(ShutdownHooks, telemetry.InitTracer()) + defer telemetry.InitTracer() - var err error importDummyData := true - + disableRLS := false + dbOptions := []duty.StartOption{duty.DisablePostgrest, duty.RunMigrations} for _, arg := range args { if arg == WithoutDummyData { importDummyData = false @@ -123,17 +146,20 @@ func BeforeSuiteFn(args ...interface{}) context.Context { if arg == WithExistingDatabase { recreateDatabase = false } + if arg == WithoutRLS { + disableRLS = true + } } - if postgresServer != nil { - return DefaultContext + if !disableRLS { + dbOptions = append(dbOptions, duty.EnableRLS) } var port int - if val, ok := os.LookupEnv("TEST_DB_PORT"); ok { + if val, ok := os.LookupEnv(TEST_DB_PORT); ok { parsed, err := strconv.ParseInt(val, 10, 32) if err != nil { - panic(err) + return context.Context{}, err } port = int(parsed) @@ -142,7 +168,7 @@ func BeforeSuiteFn(args ...interface{}) context.Context { } PgUrl = fmt.Sprintf("postgres://postgres:postgres@localhost:%d/%s?sslmode=disable", port, dbName) - url := os.Getenv("DUTY_DB_URL") + url := os.Getenv(DUTY_DB_URL) if url != "" && !recreateDatabase { PgUrl = url } else if url != "" && recreateDatabase { @@ -151,52 +177,61 @@ func BeforeSuiteFn(args ...interface{}) context.Context { PgUrl = strings.Replace(url, "/postgres", "/"+dbName, 1) _ = execPostgres(postgresDBUrl, "DROP DATABASE "+dbName) if err := execPostgres(postgresDBUrl, "CREATE DATABASE "+dbName); err != nil { - panic(fmt.Sprintf("Cannot create %s: %v", dbName, err)) + return context.Context{}, fmt.Errorf("cannot create %s: %v", dbName, err) } - ShutdownHooks = append(ShutdownHooks, func(ctx gocontext.Context) error { - return execPostgres(postgresDBUrl, fmt.Sprintf("DROP DATABASE %s (FORCE)", dbName)) + shutdown.AddHookWithPriority("remote postgres", shutdown.PriorityCritical, func() { + if err := execPostgres(postgresDBUrl, fmt.Sprintf("DROP DATABASE %s (FORCE)", dbName)); err != nil { + logger.Errorf("execPostgres: %v", err) + } }) - } else if url == "" { + } else if url == "" && postgresServer == nil { config, _ := GetEmbeddedPGConfig(dbName, port) + + // allow data dir override + if v, ok := os.LookupEnv(DUTY_DB_DATA_DIR); ok { + config = config.DataPath(v) + } + postgresServer = embeddedPG.NewDatabase(config) logger.Infof("starting embedded postgres on port %d", port) - if err = postgresServer.Start(); err != nil { - panic(err.Error()) + if err := postgresServer.Start(); err != nil { + return context.Context{}, err } logger.Infof("Started postgres on port %d", port) - ShutdownHooks = append(ShutdownHooks, func(ctx gocontext.Context) error { - logger.Infof("Stopping postgres") - return postgresServer.Stop() + shutdown.AddHookWithPriority("embedded pg", shutdown.PriorityCritical, func() { + if err := postgresServer.Stop(); err != nil { + logger.Errorf("postgresServer.Stop: %v", err) + } }) } - ctx, _, err := duty.Start("test", duty.DisablePostgrest, duty.EnableRLS, duty.RunMigrations, duty.WithUrl(PgUrl)) + dbOptions = append(dbOptions, duty.WithUrl(PgUrl)) + ctx, _, err := duty.Start(dbName, dbOptions...) if err != nil { - panic(err.Error()) + return context.Context{}, err } - DefaultContext = ctx - if err := DefaultContext.DB().Exec("SET TIME ZONE 'UTC'").Error; err != nil { - panic(err.Error()) + if err := ctx.DB().Exec("SET TIME ZONE 'UTC'").Error; err != nil { + return context.Context{}, err } - DefaultContext = DefaultContext.WithValue("db_name", dbName).WithValue("db_url", PgUrl) + ctx = ctx.WithValue("db_name", dbName).WithValue("db_url", PgUrl) if importDummyData { - dummyData = dummy.GetStaticDummyData(DefaultContext.DB()) - if err := dummyData.Delete(DefaultContext.DB()); err != nil { + dummyData = dummy.GetStaticDummyData(ctx.DB()) + if err := dummyData.Delete(ctx.DB()); err != nil { logger.Errorf(err.Error()) } - err = dummyData.Populate(DefaultContext.DB()) + err = dummyData.Populate(ctx.DB()) if err != nil { - panic(err.Error()) + return context.Context{}, err } logger.Infof("Created dummy data %v", len(dummyData.Checks)) } - DefaultContext := DefaultContext.WithKubernetes(fake.NewSimpleClientset(&v1.ConfigMap{ + ctx = ctx.WithKubernetes(fake.NewSimpleClientset(&v1.ConfigMap{ ObjectMeta: metav1.ObjectMeta{ Name: "test-cm", Namespace: "default", @@ -213,18 +248,11 @@ func BeforeSuiteFn(args ...interface{}) context.Context { "foo": []byte("secret"), }}), nil) - return DefaultContext + return ctx, nil } func AfterSuiteFn() { - for _, fn := range ShutdownHooks { - if err := fn(DefaultContext); err != nil { - logger.Errorf(err.Error()) - } - } - // clear out hooks so they don't run again - ShutdownHooks = []func(ctx gocontext.Context) error{} - + shutdown.ShutdownAndExit(0, "") } // NewDB creates a new database from an existing context, and diff --git a/views/034_rls_enable.sql b/views/034_rls_enable.sql index 21062ff6..47d1f361 100644 --- a/views/034_rls_enable.sql +++ b/views/034_rls_enable.sql @@ -1,3 +1,12 @@ +CREATE +OR REPLACE FUNCTION is_rls_disabled () RETURNS BOOLEAN AS $$ +BEGIN + RETURN (current_setting('request.jwt.claims', TRUE) IS NULL + OR current_setting('request.jwt.claims', TRUE) = '' + OR current_setting('request.jwt.claims', TRUE)::jsonb ->> 'disable_rls' IS NOT NULL); +END; +$$ LANGUAGE plpgsql SECURITY DEFINER; + -- Policy config items ALTER TABLE config_items ENABLE ROW LEVEL SECURITY; @@ -6,12 +15,7 @@ DROP POLICY IF EXISTS config_items_auth ON config_items; CREATE POLICY config_items_auth ON config_items FOR ALL TO postgrest_api, postgrest_anon USING ( - CASE WHEN ( - current_setting('request.jwt.claims', TRUE) IS NULL - OR current_setting('request.jwt.claims', TRUE) = '' -- when the parameter is set, it cannot be deleted. It's value is set to empty string. - OR current_setting('request.jwt.claims', TRUE)::jsonb ->> 'disable_rls' IS NOT NULL - ) - THEN TRUE + CASE WHEN is_rls_disabled() THEN TRUE ELSE ( agent_id = ANY (ARRAY (SELECT (jsonb_array_elements_text(current_setting('request.jwt.claims')::jsonb -> 'agents'))::uuid)) OR @@ -31,12 +35,7 @@ DROP POLICY IF EXISTS config_changes_auth ON config_changes; CREATE POLICY config_changes_auth ON config_changes FOR ALL TO postgrest_api, postgrest_anon USING ( - CASE WHEN ( - current_setting('request.jwt.claims', TRUE) IS NULL - OR current_setting('request.jwt.claims', TRUE) = '' -- when the parameter is set, it cannot be deleted. It's value is set to empty string. - OR current_setting('request.jwt.claims', TRUE)::jsonb ->> 'disable_rls' IS NOT NULL - ) - THEN TRUE + CASE WHEN is_rls_disabled() THEN TRUE ELSE EXISTS ( -- just leverage the RLS on config_items SELECT 1 @@ -54,12 +53,7 @@ DROP POLICY IF EXISTS config_analysis_auth ON config_analysis; CREATE POLICY config_analysis_auth ON config_analysis FOR ALL TO postgrest_api, postgrest_anon USING ( - CASE WHEN ( - current_setting('request.jwt.claims', TRUE) IS NULL - OR current_setting('request.jwt.claims', TRUE) = '' -- when the parameter is set, it cannot be deleted. It's value is set to empty string. - OR current_setting('request.jwt.claims', TRUE)::jsonb ->> 'disable_rls' IS NOT NULL - ) - THEN TRUE + CASE WHEN is_rls_disabled() THEN TRUE ELSE EXISTS ( -- just leverage the RLS on config_items SELECT 1 @@ -77,12 +71,7 @@ DROP POLICY IF EXISTS config_relationships_auth ON config_relationships; CREATE POLICY config_relationships_auth ON config_relationships FOR ALL TO postgrest_api, postgrest_anon USING ( - CASE WHEN ( - current_setting('request.jwt.claims', TRUE) IS NULL - OR current_setting('request.jwt.claims', TRUE) = '' -- when the parameter is set, it cannot be deleted. It's value is set to empty string. - OR current_setting('request.jwt.claims', TRUE)::jsonb ->> 'disable_rls' IS NOT NULL - ) - THEN TRUE + CASE WHEN is_rls_disabled() THEN TRUE ELSE EXISTS ( -- just leverage the RLS on config_items SELECT 1 @@ -100,12 +89,7 @@ DROP POLICY IF EXISTS config_component_relationships_auth ON config_component_re CREATE POLICY config_component_relationships_auth ON config_component_relationships FOR ALL TO postgrest_api, postgrest_anon USING ( - CASE WHEN ( - current_setting('request.jwt.claims', TRUE) IS NULL - OR current_setting('request.jwt.claims', TRUE) = '' -- when the parameter is set, it cannot be deleted. It's value is set to empty string. - OR current_setting('request.jwt.claims', TRUE)::jsonb ->> 'disable_rls' IS NOT NULL - ) - THEN TRUE + CASE WHEN is_rls_disabled() THEN TRUE ELSE EXISTS ( -- just leverage the RLS on config_items SELECT 1 @@ -123,12 +107,7 @@ DROP POLICY IF EXISTS components_auth ON components; CREATE POLICY components_auth ON components FOR ALL TO postgrest_api, postgrest_anon USING ( - CASE WHEN ( - current_setting('request.jwt.claims', TRUE) IS NULL - OR current_setting('request.jwt.claims', TRUE) = '' -- when the parameter is set, it cannot be deleted. It's value is set to empty string. - OR current_setting('request.jwt.claims', TRUE)::jsonb ->> 'disable_rls' IS NOT NULL - ) - THEN TRUE + CASE WHEN is_rls_disabled() THEN TRUE ELSE ( agent_id = ANY (ARRAY (SELECT (jsonb_array_elements_text(current_setting('request.jwt.claims')::jsonb -> 'agents'))::uuid)) ) @@ -155,8 +134,4 @@ ALTER VIEW config_summary SET (security_invoker = true); ALTER VIEW config_tags SET (security_invoker = true); ALTER VIEW config_types SET (security_invoker = true); ALTER VIEW configs SET (security_invoker = true); -ALTER VIEW incidents_by_config SET (security_invoker = true); - --- ALTER MATERIALIZED VIEW config_item_summary_3d SET (security_invoker = true); --- ALTER MATERIALIZED VIEW config_item_summary_7d SET (security_invoker = true); --- ALTER MATERIALIZED VIEW config_item_summary_30d SET (security_invoker = true); +ALTER VIEW incidents_by_config SET (security_invoker = true); \ No newline at end of file diff --git a/views/035_rls_disable.sql b/views/035_rls_disable.sql index ab4d9ffb..56daf929 100644 --- a/views/035_rls_disable.sql +++ b/views/035_rls_disable.sql @@ -7,6 +7,10 @@ ALTER TABLE config_analysis DISABLE ROW LEVEL SECURITY; ALTER TABLE components DISABLE ROW LEVEL SECURITY; +ALTER TABLE config_component_relationships DISABLE ROW LEVEL SECURITY; + +ALTER TABLE config_relationships DISABLE ROW LEVEL SECURITY; + -- POLICIES DROP POLICY IF EXISTS config_items_auth ON config_items;