From a658f2d4a1f40f9ea51447560727538c95c3aa8a Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Mon, 5 Aug 2024 17:05:58 +0200
Subject: [PATCH 01/13] feat: add atlas and golang migrate

feat: install atlas from dist binary with nix

fix: switch from community version of atlas

feat: define initial migration

chore: make automigrate configurable and add cmd for local

refactor: get rid of make commands and use atlas config file

refactor: use golang-migrate format

chore: add golang migrate to env

refactor: write custom migrate util

refactor: expose migrate interface

feat: add url parser utility

fix: enable migrations under custom path

fix: dont return error on no new migrations

fix: fix lint and test

fix: update e2e config

feat: support multiple automigrate modes
---
 atlas.hcl                                     |  72 +++++++++++
 cmd/balance-worker/main.go                    |   7 +-
 cmd/notification-service/main.go              |   6 +-
 cmd/server/main.go                            |  67 +++++-----
 config.example.yaml                           |   1 +
 config/config.go                              |   1 +
 config/postgres.go                            |  41 +++++-
 e2e/config.yaml                               |   1 +
 flake.nix                                     |  24 ++++
 go.mod                                        |   1 +
 go.sum                                        |   8 +-
 openmeter/entitlement/metered/utils_test.go   |   8 +-
 .../productcatalog/adapter/feature_test.go    |  20 +--
 openmeter/registry/startup/db.go              |  26 ++++
 openmeter/testutils/pg_driver.go              |  39 ++++--
 pkg/framework/entutils/database.go            |  12 +-
 pkg/framework/entutils/paginate_test.go       |   2 +-
 pkg/framework/entutils/transaction_test.go    |   5 +-
 test/entitlement/regression/framework_test.go |   9 +-
 tools/migrate/migrate.go                      |  61 +++++++++
 .../migrations/20240807123504_init.down.sql   |  44 +++++++
 .../migrations/20240807123504_init.up.sql     | 122 ++++++++++++++++++
 tools/migrate/migrations/atlas.sum            |   3 +
 23 files changed, 503 insertions(+), 77 deletions(-)
 create mode 100644 atlas.hcl
 create mode 100644 openmeter/registry/startup/db.go
 create mode 100644 tools/migrate/migrate.go
 create mode 100644 tools/migrate/migrations/20240807123504_init.down.sql
 create mode 100644 tools/migrate/migrations/20240807123504_init.up.sql
 create mode 100644 tools/migrate/migrations/atlas.sum

diff --git a/atlas.hcl b/atlas.hcl
new file mode 100644
index 000000000..e951b7166
--- /dev/null
+++ b/atlas.hcl
@@ -0,0 +1,72 @@
+env "local" {
+  // Declare where the schema definition resides.
+  src = "ent://internal/ent/schema"
+
+  migration {
+    // Define the directory where the migrations are stored.
+    dir = "file://tools/migrate/migrations"
+    // We use golang-migrate
+    format = golang-migrate
+  }
+
+  format {
+    migrate {
+      diff = "{{ sql . \"  \" }}"
+    }
+  }
+
+  // Define the URL of the database which is managed in this environment.
+  url = "postgres://postgres:postgres@localhost:5432/postgres?search_path=public&sslmode=disable"
+
+  // Define the URL of the Dev Database for this environment
+  // See: https://atlasgo.io/concepts/dev-database
+  dev = "docker://postgres/15/dev?search_path=public"
+
+  lint {
+    // Lint the effects of the 100 latest migration files
+    latest = 100
+  }
+}
+
+// CAN be used for all remote deployments
+env "remote" {
+  // Declare where the schema definition resides.
+  src = "ent://internal/ent/schema"
+
+  migration {
+    // Define the directory where the migrations are stored.
+    dir = "file://tools/migrate/migrations"
+    // We use golang-migrate
+    format = golang-migrate
+    // Remote deployments already had auto deploy present
+    baseline = "20240807123504"
+  }
+
+  format {
+    migrate {
+      diff = "{{ sql . \"  \" }}"
+    }
+  }
+
+  // Define the URL of the Dev Database for this environment
+  // See: https://atlasgo.io/concepts/dev-database
+  dev = "docker://postgres/15/dev?search_path=public"
+}
+
+lint {
+    non_linear {
+        error = true
+    }
+
+    destructive {
+        error = false
+    }
+
+    data_depend {
+        error = true
+    }
+
+    incompatible {
+        error = true
+    }
+}
diff --git a/cmd/balance-worker/main.go b/cmd/balance-worker/main.go
index db5d687ab..ffd0f9bc4 100644
--- a/cmd/balance-worker/main.go
+++ b/cmd/balance-worker/main.go
@@ -33,6 +33,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/entitlement/balanceworker"
 	"github.com/openmeterio/openmeter/openmeter/meter"
 	registrybuilder "github.com/openmeterio/openmeter/openmeter/registry/builder"
+	"github.com/openmeterio/openmeter/openmeter/registry/startup"
 	"github.com/openmeterio/openmeter/openmeter/streaming/clickhouse_connector"
 	"github.com/openmeterio/openmeter/openmeter/watermill/driver/kafka"
 	watermillkafka "github.com/openmeterio/openmeter/openmeter/watermill/driver/kafka"
@@ -238,10 +239,8 @@ func main() {
 
 	entClient := entPostgresDriver.Client()
 
-	// Run database schema creation
-	err = entClient.Schema.Create(ctx)
-	if err != nil {
-		logger.Error("failed to create database schema", "error", err)
+	if err := startup.DB(conf.Postgres, entClient); err != nil {
+		logger.Error("failed to initialize database", "error", err)
 		os.Exit(1)
 	}
 
diff --git a/cmd/notification-service/main.go b/cmd/notification-service/main.go
index 8d05365a4..b4becb36d 100644
--- a/cmd/notification-service/main.go
+++ b/cmd/notification-service/main.go
@@ -35,6 +35,7 @@ import (
 	notificationservice "github.com/openmeterio/openmeter/openmeter/notification/service"
 	notificationwebhook "github.com/openmeterio/openmeter/openmeter/notification/webhook"
 	registrybuilder "github.com/openmeterio/openmeter/openmeter/registry/builder"
+	"github.com/openmeterio/openmeter/openmeter/registry/startup"
 	"github.com/openmeterio/openmeter/openmeter/streaming/clickhouse_connector"
 	"github.com/openmeterio/openmeter/openmeter/watermill/driver/kafka"
 	watermillkafka "github.com/openmeterio/openmeter/openmeter/watermill/driver/kafka"
@@ -241,9 +242,8 @@ func main() {
 	entClient := entPostgresDriver.Client()
 
 	// Run database schema creation
-	err = entClient.Schema.Create(ctx)
-	if err != nil {
-		logger.Error("failed to create database schema", "error", err)
+	if err := startup.DB(conf.Postgres, entClient); err != nil {
+		logger.Error("failed to initialize database", "error", err)
 		os.Exit(1)
 	}
 
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 8120d78e6..562bcad3d 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -35,7 +35,6 @@ import (
 
 	"github.com/openmeterio/openmeter/config"
 	"github.com/openmeterio/openmeter/openmeter/debug"
-	entdb "github.com/openmeterio/openmeter/openmeter/ent/db"
 	"github.com/openmeterio/openmeter/openmeter/ingest"
 	"github.com/openmeterio/openmeter/openmeter/ingest/ingestdriver"
 	"github.com/openmeterio/openmeter/openmeter/ingest/kafkaingest"
@@ -49,6 +48,7 @@ import (
 	notificationwebhook "github.com/openmeterio/openmeter/openmeter/notification/webhook"
 	"github.com/openmeterio/openmeter/openmeter/registry"
 	registrybuilder "github.com/openmeterio/openmeter/openmeter/registry/builder"
+	"github.com/openmeterio/openmeter/openmeter/registry/startup"
 	"github.com/openmeterio/openmeter/openmeter/server"
 	"github.com/openmeterio/openmeter/openmeter/server/authenticator"
 	"github.com/openmeterio/openmeter/openmeter/server/router"
@@ -322,46 +322,42 @@ func main() {
 	debugConnector := debug.NewDebugConnector(streamingConnector)
 	entitlementConnRegistry := &registry.Entitlement{}
 
-	var entClient *entdb.Client
-	if conf.Entitlements.Enabled {
-		// Initialize Postgres driver
-		var postgresDriver *pgdriver.Driver
-		postgresDriver, err = pgdriver.NewPostgresDriver(
-			ctx,
-			conf.Postgres.URL,
-			pgdriver.WithTracerProvider(otelTracerProvider),
-			pgdriver.WithMeterProvider(otelMeterProvider),
-		)
-		if err != nil {
-			logger.Error("failed to initialize postgres driver", "error", err)
-			os.Exit(1)
-		}
+	// Initialize Postgres driver
+	postgresDriver, err := pgdriver.NewPostgresDriver(
+		ctx,
+		conf.Postgres.URL,
+		pgdriver.WithTracerProvider(otelTracerProvider),
+		pgdriver.WithMeterProvider(otelMeterProvider),
+	)
+	if err != nil {
+		logger.Error("failed to initialize postgres driver", "error", err)
+		os.Exit(1)
+	}
 
-		defer func() {
-			if err = postgresDriver.Close(); err != nil {
-				logger.Error("failed to close postgres driver", "error", err)
-			}
-		}()
+	defer func() {
+		if err = postgresDriver.Close(); err != nil {
+			logger.Error("failed to close postgres driver", "error", err)
+		}
+	}()
 
-		// Initialize Ent driver
-		entPostgresDriver := entdriver.NewEntPostgresDriver(postgresDriver.DB())
-		defer func() {
-			if err = entPostgresDriver.Close(); err != nil {
-				logger.Error("failed to close ent driver", "error", err)
-			}
-		}()
+	// Initialize Ent driver
+	entPostgresDriver := entdriver.NewEntPostgresDriver(postgresDriver.DB())
+	defer func() {
+		if err = entPostgresDriver.Close(); err != nil {
+			logger.Error("failed to close ent driver", "error", err)
+		}
+	}()
 
-		entClient = entPostgresDriver.Client()
+	entClient := entPostgresDriver.Client()
 
-		// Run database schema creation
-		err = entClient.Schema.Create(ctx)
-		if err != nil {
-			logger.Error("failed to create schema in database", "error", err)
-			os.Exit(1)
-		}
+	if err := startup.DB(conf.Postgres, entClient); err != nil {
+		logger.Error("failed to initialize database", "error", err)
+		os.Exit(1)
+	}
 
-		logger.Info("Postgres client initialized")
+	logger.Info("Postgres client initialized")
 
+	if conf.Entitlements.Enabled {
 		entitlementConnRegistry = registrybuilder.GetEntitlementRegistry(registrybuilder.EntitlementOptions{
 			DatabaseClient:     entClient,
 			StreamingConnector: streamingConnector,
@@ -372,6 +368,7 @@ func main() {
 	}
 
 	var notificationService notification.Service
+
 	if conf.Notification.Enabled {
 		if !conf.Entitlements.Enabled {
 			logger.Error("failed to initialize notification service: entitlements must be enabled")
diff --git a/config.example.yaml b/config.example.yaml
index b212a0588..1fe95dca7 100644
--- a/config.example.yaml
+++ b/config.example.yaml
@@ -57,6 +57,7 @@ events:
 
 postgres:
   url: postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable
+  autoMigrate: ent # Runs migrations as part of the service startup, valid values are: ent, migration, false
 
 meters:
   # Sample meter to count API requests
diff --git a/config/config.go b/config/config.go
index 3e3a5b749..ac321bf9f 100644
--- a/config/config.go
+++ b/config/config.go
@@ -130,6 +130,7 @@ func SetViperDefaults(v *viper.Viper, flags *pflag.FlagSet) {
 
 	configureTelemetry(v, flags)
 
+	ConfigurePostgres(v)
 	ConfigureNamespace(v)
 	ConfigureIngest(v)
 	ConfigureAggregation(v)
diff --git a/config/postgres.go b/config/postgres.go
index 83534e0d1..16a9bf76e 100644
--- a/config/postgres.go
+++ b/config/postgres.go
@@ -1,10 +1,20 @@
 package config
 
-import "errors"
+import (
+	"errors"
+
+	"github.com/spf13/viper"
+)
 
 type PostgresConfig struct {
 	// URL is the PostgreSQL database connection URL.
 	URL string `yaml:"url"`
+	// AutoMigrate is a flag that indicates whether the database should be automatically migrated.
+	// Supported values are:
+	// - "false" to disable auto-migration at startup
+	// - "ent" to use ent Schema Upserts (the default value)
+	// - "migration" to use the migrations directory
+	AutoMigrate AutoMigrate `yaml:"autoMigrate"`
 }
 
 // Validate validates the configuration.
@@ -12,6 +22,35 @@ func (c PostgresConfig) Validate() error {
 	if c.URL == "" {
 		return errors.New("database URL is required")
 	}
+	if err := c.AutoMigrate.Validate(); err != nil {
+		return err
+	}
 
 	return nil
 }
+
+func ConfigurePostgres(v *viper.Viper) {
+	v.SetDefault("postgres.autoMigrate", "ent")
+}
+
+type AutoMigrate string
+
+const (
+	AutoMigrateEnt       AutoMigrate = "ent"
+	AutoMigrateMigration AutoMigrate = "migration"
+	AutoMigrateOff       AutoMigrate = "false"
+)
+
+func (a AutoMigrate) Enabled() bool {
+	// For all other values it's enabled
+	return !(a == "false")
+}
+
+func (a AutoMigrate) Validate() error {
+	switch a {
+	case AutoMigrateEnt, AutoMigrateMigration, AutoMigrateOff:
+		return nil
+	default:
+		return errors.New("invalid auto-migrate value")
+	}
+}
diff --git a/e2e/config.yaml b/e2e/config.yaml
index 48bcfc4db..45a8b57c1 100644
--- a/e2e/config.yaml
+++ b/e2e/config.yaml
@@ -30,6 +30,7 @@ entitlements:
 
 postgres:
   url: postgres://postgres:postgres@postgres:5432/postgres?sslmode=disable
+  autoMigrate: true
 
 meters:
   - slug: ingest
diff --git a/flake.nix b/flake.nix
index e3062d2f2..4c699832e 100644
--- a/flake.nix
+++ b/flake.nix
@@ -100,11 +100,15 @@
               # python
               poetry
 
+              self'.packages.atlasx
+
               just
               semver-tool
 
               dagger
               licensei
+
+              go-migrate
             ];
 
             env = {
@@ -146,6 +150,26 @@
               "-X main.version=v${version}"
             ];
           };
+
+          atlasx = pkgs.stdenv.mkDerivation rec {
+            pname = "atlasx";
+            version = "0.25.0";
+            src = pkgs.fetchurl {
+              # License: https://ariga.io/legal/atlas/eula/eula-20240804.pdf
+              url = "https://release.ariga.io/atlas/atlas-darwin-arm64-v${version}";
+              hash = "sha256-bYJtNDE13UhJWL4ALLKI0sHMZrDS//kFWzguGX63EAo=";
+            };
+
+            unpackPhase = ''
+              cp $src atlas
+            '';
+
+            installPhase = ''
+              mkdir -p $out/bin
+              cp atlas $out/bin/atlas
+            '';
+
+          };
         };
       };
     };
diff --git a/go.mod b/go.mod
index 6e91859b7..239283ef8 100644
--- a/go.mod
+++ b/go.mod
@@ -27,6 +27,7 @@ require (
 	github.com/go-viper/mapstructure/v2 v2.1.0
 	github.com/golang-cz/devslog v0.0.8
 	github.com/golang-jwt/jwt/v5 v5.2.1
+	github.com/golang-migrate/migrate/v4 v4.17.1
 	github.com/google/uuid v1.6.0
 	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/huandu/go-sqlbuilder v1.28.1
diff --git a/go.sum b/go.sum
index 0e05f026d..c0cc60737 100644
--- a/go.sum
+++ b/go.sum
@@ -92,8 +92,8 @@ github.com/Azure/azure-sdk-for-go/sdk/storage/azqueue v1.0.0/go.mod h1:GfT0aGew8
 github.com/Azure/azure-storage-blob-go v0.14.0/go.mod h1:SMqIBi+SuiQH32bvyjngEewEeXoPfKMgWlBDaYf6fck=
 github.com/Azure/go-amqp v1.0.4 h1:GX5OFOs706UjuFRD5PDKm3aOuLQ92F7DMbua+DKAYCc=
 github.com/Azure/go-amqp v1.0.4/go.mod h1:vZAogwdrkbyK3Mla8m/CxSc/aKdnTZ4IbPxl51Y5WZE=
-github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 h1:UQHMgLO+TxOElx5B5HZ4hJQsoJ/PvUvKRhJHDQXO8P8=
-github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
+github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 h1:L/gRVlceqvL25UVaW/CKtUDjefjrs0SPonmDGUVOYP0=
+github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
 github.com/Azure/go-autorest v14.2.0+incompatible/go.mod h1:r+4oMnoxhatjLLJ6zxSWATqVooLgysK6ZNox3g/xq24=
 github.com/Azure/go-autorest/autorest/adal v0.9.13/go.mod h1:W/MM4U6nLxnIskrw4UwWzlHfGjwUS50aOsc/I3yuU8M=
 github.com/Azure/go-autorest/autorest/date v0.3.0/go.mod h1:BI0uouVdmngYNUzGWeSYnokU+TrmwEsOqdt8Y6sso74=
@@ -411,6 +411,8 @@ github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 h1:tdlZCpZ/P9DhczC
 github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
+github.com/dhui/dktest v0.4.1 h1:/w+IWuDXVymg3IrRJCHHOkMK10m9aNVMOyD0X12YVTg=
+github.com/dhui/dktest v0.4.1/go.mod h1:DdOqcUpL7vgyP4GlF3X3w7HbSlz8cEQzwewPveYEQbA=
 github.com/dimfeld/httptreemux v5.0.1+incompatible h1:Qj3gVcDNoOthBAqftuD596rm4wg/adLLz5xh5CmpiCA=
 github.com/dimfeld/httptreemux v5.0.1+incompatible/go.mod h1:rbUlSV+CCpv/SuqUTP/8Bk2O3LyUV436/yaRGkhP6Z0=
 github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
@@ -592,6 +594,8 @@ github.com/golang-jwt/jwt/v4 v4.5.0 h1:7cYmW1XlMY7h7ii7UhUyChSgS5wUJEnm9uZVTGqOW
 github.com/golang-jwt/jwt/v4 v4.5.0/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
 github.com/golang-jwt/jwt/v5 v5.2.1 h1:OuVbFODueb089Lh128TAcimifWaLhJwVflnrgM17wHk=
 github.com/golang-jwt/jwt/v5 v5.2.1/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
+github.com/golang-migrate/migrate/v4 v4.17.1 h1:4zQ6iqL6t6AiItphxJctQb3cFqWiSpMnX7wLTPnnYO4=
+github.com/golang-migrate/migrate/v4 v4.17.1/go.mod h1:m8hinFyWBn0SA4QKHuKh175Pm9wjmxj3S2Mia7dbXzM=
 github.com/golang-sql/civil v0.0.0-20190719163853-cb61b32ac6fe/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0=
 github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA=
 github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0=
diff --git a/openmeter/entitlement/metered/utils_test.go b/openmeter/entitlement/metered/utils_test.go
index 3d5d19c1c..1c8915e5e 100644
--- a/openmeter/entitlement/metered/utils_test.go
+++ b/openmeter/entitlement/metered/utils_test.go
@@ -1,7 +1,6 @@
 package meteredentitlement_test
 
 import (
-	"context"
 	"sync"
 	"testing"
 	"time"
@@ -23,6 +22,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/testutils"
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
 	"github.com/openmeterio/openmeter/pkg/models"
+	"github.com/openmeterio/openmeter/tools/migrate"
 )
 
 type dependencies struct {
@@ -60,7 +60,7 @@ func setupConnector(t *testing.T) (meteredentitlement.Connector, *dependencies)
 	driver := testutils.InitPostgresDB(t)
 
 	// build db client & adapters
-	dbClient := db.NewClient(db.Driver(driver))
+	dbClient := db.NewClient(db.Driver(driver.EntDriver))
 
 	featureRepo := productcatalog_postgresadapter.NewPostgresFeatureRepo(dbClient, testLogger)
 	entitlementRepo := entitlement_postgresadapter.NewPostgresEntitlementRepo(dbClient)
@@ -71,8 +71,8 @@ func setupConnector(t *testing.T) (meteredentitlement.Connector, *dependencies)
 	m.Lock()
 	defer m.Unlock()
 	// migrate db
-	if err := dbClient.Schema.Create(context.Background()); err != nil {
-		t.Fatalf("failed to migrate database %s", err)
+	if err := migrate.Up(driver.URL); err != nil {
+		t.Fatalf("failed to migrate db: %s", err.Error())
 	}
 
 	mockPublisher := eventbus.NewMock(t)
diff --git a/openmeter/productcatalog/adapter/feature_test.go b/openmeter/productcatalog/adapter/feature_test.go
index 76077ff48..d8b6f9786 100644
--- a/openmeter/productcatalog/adapter/feature_test.go
+++ b/openmeter/productcatalog/adapter/feature_test.go
@@ -16,6 +16,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/testutils"
 	"github.com/openmeterio/openmeter/pkg/models"
 	"github.com/openmeterio/openmeter/pkg/pagination"
+	"github.com/openmeterio/openmeter/tools/migrate"
 )
 
 func TestCreateFeature(t *testing.T) {
@@ -216,13 +217,16 @@ func TestCreateFeature(t *testing.T) {
 			defer m.Unlock()
 
 			driver := testutils.InitPostgresDB(t)
-			dbClient := db.NewClient(db.Driver(driver))
+			defer driver.SQLDriver.Close()
+			dbClient := db.NewClient(db.Driver(driver.EntDriver))
+			defer dbClient.Close()
 
-			if err := dbClient.Schema.Create(context.Background()); err != nil {
-				t.Fatalf("failed to migrate database %s", err)
-			}
+			m.Lock()
+			defer m.Unlock()
 
-			defer dbClient.Close()
+			if err := migrate.Up(driver.URL); err != nil {
+				t.Fatalf("failed to migrate db: %s", err.Error())
+			}
 
 			dbConnector := adapter.NewPostgresFeatureRepo(dbClient, testutils.NewLogger(t))
 			tc.run(t, dbConnector)
@@ -235,11 +239,11 @@ func TestCreateFeature(t *testing.T) {
 		defer m.Unlock()
 
 		driver := testutils.InitPostgresDB(t)
-		dbClient := db.NewClient(db.Driver(driver))
+		dbClient := db.NewClient(db.Driver(driver.EntDriver))
 		defer dbClient.Close()
 
-		if err := dbClient.Schema.Create(context.Background()); err != nil {
-			t.Fatalf("failed to migrate database %s", err)
+		if err := migrate.Up(driver.URL); err != nil {
+			t.Fatalf("failed to migrate db: %s", err.Error())
 		}
 
 		dbConnector := adapter.NewPostgresFeatureRepo(dbClient, testutils.NewLogger(t))
diff --git a/openmeter/registry/startup/db.go b/openmeter/registry/startup/db.go
new file mode 100644
index 000000000..f5a929f68
--- /dev/null
+++ b/openmeter/registry/startup/db.go
@@ -0,0 +1,26 @@
+package startup
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/openmeterio/openmeter/config"
+	"github.com/openmeterio/openmeter/openmeter/ent/db"
+	"github.com/openmeterio/openmeter/tools/migrate"
+)
+
+func DB(cfg config.PostgresConfig, db *db.Client) error {
+	if cfg.AutoMigrate.Enabled() {
+		switch cfg.AutoMigrate {
+		case config.AutoMigrateEnt:
+			if err := db.Schema.Create(context.Background()); err != nil {
+				return fmt.Errorf("failed to migrate db: %w", err)
+			}
+		case config.AutoMigrateMigration:
+			if err := migrate.Up(cfg.URL); err != nil {
+				return fmt.Errorf("failed to migrate db: %w", err)
+			}
+		}
+	}
+	return nil
+}
diff --git a/openmeter/testutils/pg_driver.go b/openmeter/testutils/pg_driver.go
index 3e99c6222..1db0ba0b0 100644
--- a/openmeter/testutils/pg_driver.go
+++ b/openmeter/testutils/pg_driver.go
@@ -10,13 +10,15 @@ import (
 	entsql "entgo.io/ent/dialect/sql"
 	_ "github.com/jackc/pgx/v5/stdlib" // pgx database driver
 	"github.com/peterldowns/pgtestdb"
+
+	"github.com/openmeterio/openmeter/pkg/framework/entutils"
 )
 
-// EntMigrator is a migrator for pgtestdb.
-type EntMigrator struct{}
+// NoopMigrator is a migrator for pgtestdb.
+type NoopMigrator struct{}
 
 // Hash returns the md5 hash of the schema file.
-func (m *EntMigrator) Hash() (string, error) {
+func (m *NoopMigrator) Hash() (string, error) {
 	return "", nil
 }
 
@@ -24,7 +26,7 @@ func (m *EntMigrator) Hash() (string, error) {
 // database.
 //
 //	atlas schema apply --auto-approve --url $DB --to file://$schemaFilePath
-func (m *EntMigrator) Migrate(
+func (m *NoopMigrator) Migrate(
 	ctx context.Context,
 	db *sql.DB,
 	templateConf pgtestdb.Config,
@@ -33,7 +35,7 @@ func (m *EntMigrator) Migrate(
 }
 
 // Prepare is a no-op method.
-func (*EntMigrator) Prepare(
+func (*NoopMigrator) Prepare(
 	_ context.Context,
 	_ *sql.DB,
 	_ pgtestdb.Config,
@@ -42,7 +44,7 @@ func (*EntMigrator) Prepare(
 }
 
 // Verify is a no-op method.
-func (*EntMigrator) Verify(
+func (*NoopMigrator) Verify(
 	_ context.Context,
 	_ *sql.DB,
 	_ pgtestdb.Config,
@@ -50,7 +52,13 @@ func (*EntMigrator) Verify(
 	return nil
 }
 
-func InitPostgresDB(t *testing.T) *entsql.Driver {
+type TestDB struct {
+	EntDriver *entsql.Driver
+	SQLDriver *sql.DB
+	URL       string
+}
+
+func InitPostgresDB(t *testing.T) *TestDB {
 	t.Helper()
 
 	// Dagger will set the POSTGRES_HOST environment variable for `make test`.
@@ -62,12 +70,25 @@ func InitPostgresDB(t *testing.T) *entsql.Driver {
 	}
 
 	// TODO: fix migrations
-	return entsql.OpenDB(dialect.Postgres, pgtestdb.New(t, pgtestdb.Config{
+	dbConf := pgtestdb.Custom(t, pgtestdb.Config{
 		DriverName: "pgx",
 		User:       "postgres",
 		Password:   "postgres",
 		Host:       host,
 		Port:       "5432",
 		Options:    "sslmode=disable",
-	}, &EntMigrator{}))
+	}, &NoopMigrator{})
+
+	sqlDriver, err := entutils.GetSQLDriver(dbConf.URL())
+	if err != nil {
+		t.Fatalf("failed to get pg driver: %s", err)
+	}
+
+	entDriver := entsql.OpenDB(dialect.Postgres, sqlDriver)
+
+	return &TestDB{
+		EntDriver: entDriver,
+		SQLDriver: sqlDriver,
+		URL:       dbConf.URL(),
+	}
 }
diff --git a/pkg/framework/entutils/database.go b/pkg/framework/entutils/database.go
index 3701537f4..e3688d56b 100644
--- a/pkg/framework/entutils/database.go
+++ b/pkg/framework/entutils/database.go
@@ -1,6 +1,8 @@
 package entutils
 
 import (
+	"database/sql"
+
 	"entgo.io/ent/dialect"
 	entDialectSQL "entgo.io/ent/dialect/sql"
 	"github.com/XSAM/otelsql"
@@ -8,11 +10,15 @@ import (
 )
 
 // Deprecated: use NewEntPostgresDriver instead
-func GetPGDriver(databaseURL string) (*entDialectSQL.Driver, error) {
-	// TODO: inject trace and metrics provider
-	database, err := otelsql.Open("pgx", databaseURL, otelsql.WithAttributes(
+func GetSQLDriver(databaseURL string) (*sql.DB, error) {
+	return otelsql.Open("pgx", databaseURL, otelsql.WithAttributes(
 		semconv.DBSystemPostgreSQL,
 	))
+}
+
+func GetEntDriver(databaseURL string) (*entDialectSQL.Driver, error) {
+	// TODO: inject trace and metrics provider
+	database, err := GetSQLDriver(databaseURL)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/framework/entutils/paginate_test.go b/pkg/framework/entutils/paginate_test.go
index bd4749815..7eadc1835 100644
--- a/pkg/framework/entutils/paginate_test.go
+++ b/pkg/framework/entutils/paginate_test.go
@@ -21,7 +21,7 @@ func TestPaginate(t *testing.T) {
 	driver := testutils.InitPostgresDB(t)
 
 	// build db clients
-	dbClient := db.NewClient(db.Driver(driver))
+	dbClient := db.NewClient(db.Driver(driver.EntDriver))
 	defer dbClient.Close()
 
 	if err := dbClient.Schema.Create(context.Background()); err != nil {
diff --git a/pkg/framework/entutils/transaction_test.go b/pkg/framework/entutils/transaction_test.go
index 2500919fc..cf2d2d5a9 100644
--- a/pkg/framework/entutils/transaction_test.go
+++ b/pkg/framework/entutils/transaction_test.go
@@ -291,10 +291,11 @@ func TestTransaction(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			// create isolated pg db for tests
 			driver := testutils.InitPostgresDB(t)
+			defer driver.EntDriver.Close()
 
 			// build db clients
-			db1Client := db1.NewClient(db1.Driver(driver))
-			db2Client := db2.NewClient(db2.Driver(driver))
+			db1Client := db1.NewClient(db1.Driver(driver.EntDriver))
+			db2Client := db2.NewClient(db2.Driver(driver.EntDriver))
 
 			if err := db1Client.Schema.Create(context.Background()); err != nil {
 				t.Fatalf("failed to migrate database %s", err)
diff --git a/test/entitlement/regression/framework_test.go b/test/entitlement/regression/framework_test.go
index 6f33242a7..74b7f1cb5 100644
--- a/test/entitlement/regression/framework_test.go
+++ b/test/entitlement/regression/framework_test.go
@@ -1,7 +1,6 @@
 package framework_test
 
 import (
-	"context"
 	"log/slog"
 	"testing"
 	"time"
@@ -23,6 +22,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/testutils"
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
 	"github.com/openmeterio/openmeter/pkg/models"
+	"github.com/openmeterio/openmeter/tools/migrate"
 )
 
 type Dependencies struct {
@@ -53,13 +53,12 @@ func (d *Dependencies) Close() {
 
 func setupDependencies(t *testing.T) Dependencies {
 	log := slog.Default()
-	ctx := context.Background()
 	driver := testutils.InitPostgresDB(t)
 
 	// init db
-	dbClient := db.NewClient(db.Driver(driver))
-	if err := dbClient.Schema.Create(ctx); err != nil {
-		t.Fatalf("failed to migrate database %s", err)
+	dbClient := db.NewClient(db.Driver(driver.EntDriver))
+	if err := migrate.Up(driver.URL); err != nil {
+		t.Fatalf("failed to migrate db: %s", err.Error())
 	}
 
 	// Init product catalog
diff --git a/tools/migrate/migrate.go b/tools/migrate/migrate.go
new file mode 100644
index 000000000..da94664cc
--- /dev/null
+++ b/tools/migrate/migrate.go
@@ -0,0 +1,61 @@
+// A very lightweigh migration tool to replace `ent.Schema.Create` calls.
+package migrate
+
+import (
+	"embed"
+	"io/fs"
+	"net/url"
+
+	"github.com/golang-migrate/migrate/v4"
+	_ "github.com/golang-migrate/migrate/v4/database/postgres"
+	"github.com/golang-migrate/migrate/v4/source/iofs"
+)
+
+const (
+	MigrationsTable = "schema_om"
+)
+
+type Migrate = migrate.Migrate
+
+//go:embed migrations
+var omMigrations embed.FS
+
+// NewMigrate creates a new migrate instance.
+func NewMigrate(conn string, fs fs.FS, fsPath string) (*Migrate, error) {
+	d, err := iofs.New(fs, fsPath)
+	if err != nil {
+		return nil, err
+	}
+	return migrate.NewWithSourceInstance("iofs", d, conn)
+}
+
+func Up(conn string) error {
+	conn, err := SetMigrationTableName(conn, MigrationsTable)
+	if err != nil {
+		return err
+	}
+	m, err := NewMigrate(conn, omMigrations, "migrations")
+	if err != nil {
+		return err
+	}
+
+	defer m.Close()
+	err = m.Up()
+	if err != nil && err != migrate.ErrNoChange {
+		return err
+	}
+	return nil
+}
+
+func SetMigrationTableName(conn, tableName string) (string, error) {
+	parsedURL, err := url.Parse(conn)
+	if err != nil {
+		return "", err
+	}
+
+	values := parsedURL.Query()
+	values.Set("x-migrations-table", tableName)
+	parsedURL.RawQuery = values.Encode()
+
+	return parsedURL.String(), nil
+}
diff --git a/tools/migrate/migrations/20240807123504_init.down.sql b/tools/migrate/migrations/20240807123504_init.down.sql
new file mode 100644
index 000000000..519f19707
--- /dev/null
+++ b/tools/migrate/migrations/20240807123504_init.down.sql
@@ -0,0 +1,44 @@
+-- reverse: create index "usagereset_namespace_entitlement_id_reset_time" to table: "usage_resets"
+DROP INDEX "usagereset_namespace_entitlement_id_reset_time";
+-- reverse: create index "usagereset_namespace_entitlement_id" to table: "usage_resets"
+DROP INDEX "usagereset_namespace_entitlement_id";
+-- reverse: create index "usagereset_id" to table: "usage_resets"
+DROP INDEX "usagereset_id";
+-- reverse: create "usage_resets" table
+DROP TABLE "usage_resets";
+-- reverse: create index "grant_namespace_owner_id" to table: "grants"
+DROP INDEX "grant_namespace_owner_id";
+-- reverse: create index "grant_id" to table: "grants"
+DROP INDEX "grant_id";
+-- reverse: create index "grant_effective_at_expires_at" to table: "grants"
+DROP INDEX "grant_effective_at_expires_at";
+-- reverse: create "grants" table
+DROP TABLE "grants";
+-- reverse: create index "balancesnapshot_namespace_balance_at" to table: "balance_snapshots"
+DROP INDEX "balancesnapshot_namespace_balance_at";
+-- reverse: create index "balancesnapshot_namespace_balance" to table: "balance_snapshots"
+DROP INDEX "balancesnapshot_namespace_balance";
+-- reverse: create index "balancesnapshot_namespace_at" to table: "balance_snapshots"
+DROP INDEX "balancesnapshot_namespace_at";
+-- reverse: create "balance_snapshots" table
+DROP TABLE "balance_snapshots";
+-- reverse: create index "entitlement_namespace_subject_key" to table: "entitlements"
+DROP INDEX "entitlement_namespace_subject_key";
+-- reverse: create index "entitlement_namespace_id_subject_key" to table: "entitlements"
+DROP INDEX "entitlement_namespace_id_subject_key";
+-- reverse: create index "entitlement_namespace_id" to table: "entitlements"
+DROP INDEX "entitlement_namespace_id";
+-- reverse: create index "entitlement_namespace_feature_id_id" to table: "entitlements"
+DROP INDEX "entitlement_namespace_feature_id_id";
+-- reverse: create index "entitlement_namespace_current_usage_period_end" to table: "entitlements"
+DROP INDEX "entitlement_namespace_current_usage_period_end";
+-- reverse: create index "entitlement_id" to table: "entitlements"
+DROP INDEX "entitlement_id";
+-- reverse: create "entitlements" table
+DROP TABLE "entitlements";
+-- reverse: create index "feature_namespace_id" to table: "features"
+DROP INDEX "feature_namespace_id";
+-- reverse: create index "feature_id" to table: "features"
+DROP INDEX "feature_id";
+-- reverse: create "features" table
+DROP TABLE "features";
diff --git a/tools/migrate/migrations/20240807123504_init.up.sql b/tools/migrate/migrations/20240807123504_init.up.sql
new file mode 100644
index 000000000..eac451fa7
--- /dev/null
+++ b/tools/migrate/migrations/20240807123504_init.up.sql
@@ -0,0 +1,122 @@
+-- create "features" table
+CREATE TABLE "features" (
+  "id" character(26) NOT NULL,
+  "created_at" timestamptz NOT NULL,
+  "updated_at" timestamptz NOT NULL,
+  "deleted_at" timestamptz NULL,
+  "metadata" jsonb NULL,
+  "namespace" character varying NOT NULL,
+  "name" character varying NOT NULL,
+  "key" character varying NOT NULL,
+  "meter_slug" character varying NULL,
+  "meter_group_by_filters" jsonb NULL,
+  "archived_at" timestamptz NULL,
+  PRIMARY KEY ("id")
+);
+-- create index "feature_id" to table: "features"
+CREATE INDEX "feature_id" ON "features" ("id");
+-- create index "feature_namespace_id" to table: "features"
+CREATE INDEX "feature_namespace_id" ON "features" ("namespace", "id");
+-- create "entitlements" table
+CREATE TABLE "entitlements" (
+  "id" character(26) NOT NULL,
+  "namespace" character varying NOT NULL,
+  "metadata" jsonb NULL,
+  "created_at" timestamptz NOT NULL,
+  "updated_at" timestamptz NOT NULL,
+  "deleted_at" timestamptz NULL,
+  "entitlement_type" character varying NOT NULL,
+  "feature_key" character varying NOT NULL,
+  "subject_key" character varying NOT NULL,
+  "measure_usage_from" timestamptz NULL,
+  "issue_after_reset" double precision NULL,
+  "issue_after_reset_priority" smallint NULL,
+  "is_soft_limit" boolean NULL,
+  "config" jsonb NULL,
+  "usage_period_interval" character varying NULL,
+  "usage_period_anchor" timestamptz NULL,
+  "current_usage_period_start" timestamptz NULL,
+  "current_usage_period_end" timestamptz NULL,
+  "feature_id" character(26) NOT NULL,
+  PRIMARY KEY ("id"),
+  CONSTRAINT "entitlements_features_entitlement" FOREIGN KEY ("feature_id") REFERENCES "features" ("id") ON UPDATE NO ACTION ON DELETE NO ACTION
+);
+-- create index "entitlement_id" to table: "entitlements"
+CREATE INDEX "entitlement_id" ON "entitlements" ("id");
+-- create index "entitlement_namespace_current_usage_period_end" to table: "entitlements"
+CREATE INDEX "entitlement_namespace_current_usage_period_end" ON "entitlements" ("namespace", "current_usage_period_end");
+-- create index "entitlement_namespace_feature_id_id" to table: "entitlements"
+CREATE INDEX "entitlement_namespace_feature_id_id" ON "entitlements" ("namespace", "feature_id", "id");
+-- create index "entitlement_namespace_id" to table: "entitlements"
+CREATE INDEX "entitlement_namespace_id" ON "entitlements" ("namespace", "id");
+-- create index "entitlement_namespace_id_subject_key" to table: "entitlements"
+CREATE INDEX "entitlement_namespace_id_subject_key" ON "entitlements" ("namespace", "id", "subject_key");
+-- create index "entitlement_namespace_subject_key" to table: "entitlements"
+CREATE INDEX "entitlement_namespace_subject_key" ON "entitlements" ("namespace", "subject_key");
+-- create "balance_snapshots" table
+CREATE TABLE "balance_snapshots" (
+  "id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY,
+  "namespace" character varying NOT NULL,
+  "created_at" timestamptz NOT NULL,
+  "updated_at" timestamptz NOT NULL,
+  "deleted_at" timestamptz NULL,
+  "grant_balances" jsonb NOT NULL,
+  "balance" numeric NOT NULL,
+  "overage" numeric NOT NULL,
+  "at" timestamptz NOT NULL,
+  "owner_id" character(26) NOT NULL,
+  PRIMARY KEY ("id"),
+  CONSTRAINT "balance_snapshots_entitlements_balance_snapshot" FOREIGN KEY ("owner_id") REFERENCES "entitlements" ("id") ON UPDATE NO ACTION ON DELETE NO ACTION
+);
+-- create index "balancesnapshot_namespace_at" to table: "balance_snapshots"
+CREATE INDEX "balancesnapshot_namespace_at" ON "balance_snapshots" ("namespace", "at");
+-- create index "balancesnapshot_namespace_balance" to table: "balance_snapshots"
+CREATE INDEX "balancesnapshot_namespace_balance" ON "balance_snapshots" ("namespace", "balance");
+-- create index "balancesnapshot_namespace_balance_at" to table: "balance_snapshots"
+CREATE INDEX "balancesnapshot_namespace_balance_at" ON "balance_snapshots" ("namespace", "balance", "at");
+-- create "grants" table
+CREATE TABLE "grants" (
+  "id" character(26) NOT NULL,
+  "namespace" character varying NOT NULL,
+  "metadata" jsonb NULL,
+  "created_at" timestamptz NOT NULL,
+  "updated_at" timestamptz NOT NULL,
+  "deleted_at" timestamptz NULL,
+  "amount" numeric NOT NULL,
+  "priority" smallint NOT NULL DEFAULT 0,
+  "effective_at" timestamptz NOT NULL,
+  "expiration" jsonb NOT NULL,
+  "expires_at" timestamptz NOT NULL,
+  "voided_at" timestamptz NULL,
+  "reset_max_rollover" numeric NOT NULL,
+  "reset_min_rollover" numeric NOT NULL,
+  "recurrence_period" character varying NULL,
+  "recurrence_anchor" timestamptz NULL,
+  "owner_id" character(26) NOT NULL,
+  PRIMARY KEY ("id"),
+  CONSTRAINT "grants_entitlements_grant" FOREIGN KEY ("owner_id") REFERENCES "entitlements" ("id") ON UPDATE NO ACTION ON DELETE NO ACTION
+);
+-- create index "grant_effective_at_expires_at" to table: "grants"
+CREATE INDEX "grant_effective_at_expires_at" ON "grants" ("effective_at", "expires_at");
+-- create index "grant_id" to table: "grants"
+CREATE INDEX "grant_id" ON "grants" ("id");
+-- create index "grant_namespace_owner_id" to table: "grants"
+CREATE INDEX "grant_namespace_owner_id" ON "grants" ("namespace", "owner_id");
+-- create "usage_resets" table
+CREATE TABLE "usage_resets" (
+  "id" character(26) NOT NULL,
+  "namespace" character varying NOT NULL,
+  "created_at" timestamptz NOT NULL,
+  "updated_at" timestamptz NOT NULL,
+  "deleted_at" timestamptz NULL,
+  "reset_time" timestamptz NOT NULL,
+  "entitlement_id" character(26) NOT NULL,
+  PRIMARY KEY ("id"),
+  CONSTRAINT "usage_resets_entitlements_usage_reset" FOREIGN KEY ("entitlement_id") REFERENCES "entitlements" ("id") ON UPDATE NO ACTION ON DELETE NO ACTION
+);
+-- create index "usagereset_id" to table: "usage_resets"
+CREATE INDEX "usagereset_id" ON "usage_resets" ("id");
+-- create index "usagereset_namespace_entitlement_id" to table: "usage_resets"
+CREATE INDEX "usagereset_namespace_entitlement_id" ON "usage_resets" ("namespace", "entitlement_id");
+-- create index "usagereset_namespace_entitlement_id_reset_time" to table: "usage_resets"
+CREATE INDEX "usagereset_namespace_entitlement_id_reset_time" ON "usage_resets" ("namespace", "entitlement_id", "reset_time");
diff --git a/tools/migrate/migrations/atlas.sum b/tools/migrate/migrations/atlas.sum
new file mode 100644
index 000000000..cc9508f3f
--- /dev/null
+++ b/tools/migrate/migrations/atlas.sum
@@ -0,0 +1,3 @@
+h1:TXkJcI7ny+2HOMoQqMn1Ifjerd+UbEEhOVlrsn3uQ7w=
+20240807123504_init.down.sql h1:KOqgWiIg5UA5ozQ7iRzzLbIcElzkNEojFdsS3bWEYHk=
+20240807123504_init.up.sql h1:GQiJZKVuVE/Asmqc4barw6C14q/kuShAJMci8HTkO6Q=

From d3056e8e854d6088ca4593860bfe57d75b6eab80 Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Mon, 12 Aug 2024 14:56:14 +0200
Subject: [PATCH 02/13] docs(migrate): add migration docs

fix: fix e2e config

chore: regenerate init migration and add config test

test: add up down up test
---
 config/config_test.go                         |  3 ++
 docs/migration.md                             | 25 ++++++++++++
 e2e/config.yaml                               |  1 -
 tools/migrate/migrate.go                      |  4 +-
 tools/migrate/migrate_test.go                 | 38 +++++++++++++++++++
 ....down.sql => 20240812143658_init.down.sql} |  0
 ...init.up.sql => 20240812143658_init.up.sql} |  1 +
 tools/migrate/migrations/atlas.sum            |  6 +--
 8 files changed, 72 insertions(+), 6 deletions(-)
 create mode 100644 docs/migration.md
 create mode 100644 tools/migrate/migrate_test.go
 rename tools/migrate/migrations/{20240807123504_init.down.sql => 20240812143658_init.down.sql} (100%)
 rename tools/migrate/migrations/{20240807123504_init.up.sql => 20240812143658_init.up.sql} (99%)

diff --git a/config/config_test.go b/config/config_test.go
index 852b6b96f..ec2eac40c 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -40,6 +40,9 @@ func TestComplete(t *testing.T) {
 	}
 
 	expected := Configuration{
+		Postgres: PostgresConfig{
+			AutoMigrate: AutoMigrateEnt,
+		},
 		Address:     "127.0.0.1:8888",
 		Environment: "local",
 		Telemetry: TelemetryConfig{
diff --git a/docs/migration.md b/docs/migration.md
new file mode 100644
index 000000000..4d77fe95a
--- /dev/null
+++ b/docs/migration.md
@@ -0,0 +1,25 @@
+# Data Migrations in OpenMeter
+
+OpenMeter uses [ent](https://entgo.io) for its data storage and schema management. Database state is synced from the ent schema definitions under `internal/ent/schema` via either `ent` schema upsertions or migrations.
+
+## AutoMigrate
+
+OpenMeter can automatically sync the database schema even in multi-instance deployments. This behavior can be configured via `postgres.autoMigrate` in the configuration.
+- Choosing the calue of `ent` will internally call `ent.Schema.Create` which runs a schema upsertion. This is the default behavior, intended for development and testing.
+- Choosing the value of `migration` will automatically execute the scripts in the `migrations` directory.
+- Choosing the value of `false` will disable the automatic schema sync and lets the user manage the schema manually, for example `migrate -database "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable&x-migrations-table=schema_om" -path ./services/common/migrations/om up`
+
+## Generating Migrations
+
+OpenMeter uses [atlas](https://atlasgo.io/) to generate versioned migrations from changes in the ent schema.
+
+After changing the schema and running `go generate` you can create a new migration diff via `atlas migrate --env local diff <migration-name>`, the generated migration files will be placed in the `migrations` directory.
+
+## Running Migrations
+
+The recommended way to run migrations is to use the `golang-migrate` CLI tool. You can run the migrations via
+```bash
+migrate -database "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable&x-migrations-table=schema_om" -path ./services/common/migrations/om up
+```
+
+OpenMeter `autoMigrate` uses the `schema_om` lock table.
diff --git a/e2e/config.yaml b/e2e/config.yaml
index 45a8b57c1..48bcfc4db 100644
--- a/e2e/config.yaml
+++ b/e2e/config.yaml
@@ -30,7 +30,6 @@ entitlements:
 
 postgres:
   url: postgres://postgres:postgres@postgres:5432/postgres?sslmode=disable
-  autoMigrate: true
 
 meters:
   - slug: ingest
diff --git a/tools/migrate/migrate.go b/tools/migrate/migrate.go
index da94664cc..96f6ed7bd 100644
--- a/tools/migrate/migrate.go
+++ b/tools/migrate/migrate.go
@@ -18,7 +18,7 @@ const (
 type Migrate = migrate.Migrate
 
 //go:embed migrations
-var omMigrations embed.FS
+var OMMigrations embed.FS
 
 // NewMigrate creates a new migrate instance.
 func NewMigrate(conn string, fs fs.FS, fsPath string) (*Migrate, error) {
@@ -34,7 +34,7 @@ func Up(conn string) error {
 	if err != nil {
 		return err
 	}
-	m, err := NewMigrate(conn, omMigrations, "migrations")
+	m, err := NewMigrate(conn, OMMigrations, "migrations")
 	if err != nil {
 		return err
 	}
diff --git a/tools/migrate/migrate_test.go b/tools/migrate/migrate_test.go
new file mode 100644
index 000000000..ec8f357cc
--- /dev/null
+++ b/tools/migrate/migrate_test.go
@@ -0,0 +1,38 @@
+package migrate_test
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/openmeterio/openmeter/internal/testutils"
+	"github.com/openmeterio/openmeter/tools/migrate"
+)
+
+func TestUpDownUp(t *testing.T) {
+	testDB := testutils.InitPostgresDB(t)
+
+	migrator, err := migrate.NewMigrate(testDB.URL, migrate.OMMigrations, "migrations")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		err1, err2 := migrator.Close()
+		err := errors.Join(err1, err2)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	if err := migrator.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := migrator.Down(); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := migrator.Up(); err != nil {
+		t.Fatal(err)
+	}
+}
diff --git a/tools/migrate/migrations/20240807123504_init.down.sql b/tools/migrate/migrations/20240812143658_init.down.sql
similarity index 100%
rename from tools/migrate/migrations/20240807123504_init.down.sql
rename to tools/migrate/migrations/20240812143658_init.down.sql
diff --git a/tools/migrate/migrations/20240807123504_init.up.sql b/tools/migrate/migrations/20240812143658_init.up.sql
similarity index 99%
rename from tools/migrate/migrations/20240807123504_init.up.sql
rename to tools/migrate/migrations/20240812143658_init.up.sql
index eac451fa7..1d68d0f3b 100644
--- a/tools/migrate/migrations/20240807123504_init.up.sql
+++ b/tools/migrate/migrations/20240812143658_init.up.sql
@@ -32,6 +32,7 @@ CREATE TABLE "entitlements" (
   "issue_after_reset" double precision NULL,
   "issue_after_reset_priority" smallint NULL,
   "is_soft_limit" boolean NULL,
+  "preserve_overage_at_reset" boolean NULL,
   "config" jsonb NULL,
   "usage_period_interval" character varying NULL,
   "usage_period_anchor" timestamptz NULL,
diff --git a/tools/migrate/migrations/atlas.sum b/tools/migrate/migrations/atlas.sum
index cc9508f3f..84d8c8181 100644
--- a/tools/migrate/migrations/atlas.sum
+++ b/tools/migrate/migrations/atlas.sum
@@ -1,3 +1,3 @@
-h1:TXkJcI7ny+2HOMoQqMn1Ifjerd+UbEEhOVlrsn3uQ7w=
-20240807123504_init.down.sql h1:KOqgWiIg5UA5ozQ7iRzzLbIcElzkNEojFdsS3bWEYHk=
-20240807123504_init.up.sql h1:GQiJZKVuVE/Asmqc4barw6C14q/kuShAJMci8HTkO6Q=
+h1:fW1f9OMqGTAwpbvQXgX+2jGJFPN8plpswoqfJhKs/xE=
+20240812143658_init.down.sql h1:hoeC6yyHw+V6ihO9WkCKElh/OIVtb4Kttaqy98DGhBU=
+20240812143658_init.up.sql h1:OsuTQd+TZQVfz0NNkJSyiLfUqmWEfwoUQdhIKmd06GY=

From 80bafcfd037be02953d1a3c0cc67504e54fea356 Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Tue, 13 Aug 2024 17:47:22 +0200
Subject: [PATCH 03/13] feat(ci): add ci checks

---
 .github/workflows/ci.yaml | 36 ++++++++++++++++++++
 Makefile                  |  5 +++
 atlas.hcl                 | 50 +++++++++++++++++++++-------
 ci/e2e.go                 | 14 ++++++--
 ci/migrate.go             | 57 ++++++++++++++++++++++++++++++++
 dagger.json               |  4 +++
 flake.nix                 | 69 ++++++++++++++++++++++++++++-----------
 7 files changed, 201 insertions(+), 34 deletions(-)
 create mode 100644 ci/migrate.go

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 2092c555b..daf8295ad 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -72,6 +72,42 @@ jobs:
             engine.stderr.log
           retention-days: 14
 
+  migrations:
+    name: Migration Checks
+    runs-on: depot-ubuntu-latest-8
+
+    steps:
+      # Required as a workaround for Dagger to properly detect Git metadata
+      - name: Checkout repository
+        uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+        with:
+          fetch-depth: 0 # Needed to compare against base branch
+
+      - name: Run pipeline
+        uses: dagger/dagger-for-github@29a88e72255e732147ba18a670978b90bcc59efd # v6.4.0
+        with:
+          verb: call
+          module: github.com/${{ github.repository }}@${{ github.ref }}
+          args: --ref ${{ github.ref }} migrate check
+          cloud-token: ${{ secrets.DAGGER_CLOUD_TOKEN }}
+          version: ${{ env.DAGGER_VERSION }}
+
+      - name: Export Dagger Engine logs
+        id: export-dagger-engine-logs
+        run: docker logs $(docker container list --all --filter 'name=^dagger-engine-*' --format '{{.Names}}') > engine.stdout.log 2> engine.stderr.log
+        if: always()
+        continue-on-error: true
+
+      - name: Upload Dagger Engine logs as artifact
+        uses: actions/upload-artifact@89ef406dd8d7e03cfd12d9e0a4a378f454709029 # v4.3.5
+        if: always() && steps.export-dagger-engine-logs.outcome == 'success'
+        with:
+          name: "[${{ github.job }}] Dagger Engine logs"
+          path: |
+            engine.stdout.log
+            engine.stderr.log
+          retention-days: 14
+
   lint:
     name: Lint
     runs-on: depot-ubuntu-latest-8
diff --git a/Makefile b/Makefile
index 76f5a3fe0..29d03454c 100644
--- a/Makefile
+++ b/Makefile
@@ -18,6 +18,11 @@ gen-api: ## Generate API and SDKs
 	dagger call --source .:default generate web-sdk -o api/client/web
 	dagger call --source .:default generate python-sdk -o api/client/python
 
+.PHONY: migrate-check
+migrate-check: ## Validate migrations
+	$(call print-target)
+	dagger call --source .:default migrate check
+
 .PHONY: generate
 generate: ## Generate code
 	$(call print-target)
diff --git a/atlas.hcl b/atlas.hcl
index e951b7166..4a52f21ea 100644
--- a/atlas.hcl
+++ b/atlas.hcl
@@ -1,12 +1,9 @@
 env "local" {
-  // Declare where the schema definition resides.
-  src = "ent://internal/ent/schema"
+  src = "${local.schema_src}"
 
   migration {
-    // Define the directory where the migrations are stored.
-    dir = "file://tools/migrate/migrations"
-    // We use golang-migrate
-    format = golang-migrate
+    dir = "${local.migrations_dir}"
+    format = "${local.migrations_format}"
   }
 
   format {
@@ -15,8 +12,7 @@ env "local" {
     }
   }
 
-  // Define the URL of the database which is managed in this environment.
-  url = "postgres://postgres:postgres@localhost:5432/postgres?search_path=public&sslmode=disable"
+  url = "${local.local_url}"
 
   // Define the URL of the Dev Database for this environment
   // See: https://atlasgo.io/concepts/dev-database
@@ -28,18 +24,34 @@ env "local" {
   }
 }
 
+env "ci" {
+  src = "${local.schema_src}"
+
+  migration {
+    dir = "${local.migrations_dir}"
+    format = "${local.migrations_format}"
+  }
+
+  format {
+    migrate {
+      diff = "{{ sql . \"  \" }}"
+    }
+  }
+
+  dev = "${local.ci_url}"
+}
+
 // CAN be used for all remote deployments
 env "remote" {
-  // Declare where the schema definition resides.
-  src = "ent://internal/ent/schema"
+  src = "${local.schema_src}"
 
   migration {
     // Define the directory where the migrations are stored.
     dir = "file://tools/migrate/migrations"
     // We use golang-migrate
-    format = golang-migrate
+    format = "${local.migrations_format}"
     // Remote deployments already had auto deploy present
-    baseline = "20240807123504"
+    baseline = "${local.init_migration_ts}"
   }
 
   format {
@@ -53,6 +65,20 @@ env "remote" {
   dev = "docker://postgres/15/dev?search_path=public"
 }
 
+locals {
+    // Define the directory where the schema definition resides.
+    schema_src = "ent://internal/ent/schema"
+    // Define the initial migration timestamp
+    init_migration_ts = "20240807123504"
+    // Define the directory where the migrations are stored.
+    migrations_dir = "file://tools/migrate/migrations"
+    // We use golang-migrate
+    migrations_format = "golang-migrate"
+    // Define common connection URLs
+    local_url = "postgres://postgres:postgres@localhost:5432/postgres?search_path=public&sslmode=disable"
+    ci_url = "postgres://postgres:postgres@postgres:5432/postgres?search_path=public&sslmode=disable"
+}
+
 lint {
     non_linear {
         error = true
diff --git a/ci/e2e.go b/ci/e2e.go
index 47bf8e2b4..3e06b463b 100644
--- a/ci/e2e.go
+++ b/ci/e2e.go
@@ -67,14 +67,22 @@ func redis() *dagger.Service {
 		AsService()
 }
 
-func postgres() *dagger.Service {
+func pg() *dagger.Container {
 	return dag.Container().
 		From(fmt.Sprintf("postgres:%s", postgresVersion)).
 		WithEnvVariable("POSTGRES_USER", "postgres").
 		WithEnvVariable("POSTGRES_PASSWORD", "postgres").
 		WithEnvVariable("POSTGRES_DB", "postgres").
-		WithExposedPort(5432).
-		AsService()
+		WithExposedPort(5432)
+}
+
+func postgres() *dagger.Service {
+	return pg().AsService()
+}
+
+// Creates a postgres service unique by name
+func postgresNamed(name string) *dagger.Service {
+	return pg().WithLabel("uniq-name", name).AsService()
 }
 
 const (
diff --git a/ci/migrate.go b/ci/migrate.go
new file mode 100644
index 000000000..fd6eef110
--- /dev/null
+++ b/ci/migrate.go
@@ -0,0 +1,57 @@
+package main
+
+import (
+	"context"
+
+	"github.com/openmeterio/openmeter/ci/internal/dagger"
+	"github.com/sourcegraph/conc/pool"
+)
+
+func (m *Ci) Migrate() *Migrate {
+	return &Migrate{
+		Source: m.Source,
+	}
+}
+
+type Migrate struct {
+	Source *dagger.Directory
+}
+
+func (m *Migrate) Check(
+	ctx context.Context,
+) error {
+	nix := nix(m.Source)
+	p := pool.New().WithErrors().WithContext(ctx)
+
+	// Always validate schema is generated
+	p.Go(syncFunc(nix.
+		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c go generate ./internal/ent/..."}).
+		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c sh -c 'if [[ -n $(git status --porcelain ./internal/ent) ]]; then echo \"Ent schema wasnt generated\"; exit 1; fi' "}),
+	))
+	// Always validate migrations are in sync with schema
+	p.Go(syncFunc(nix.
+		WithServiceBinding("postgres", postgresNamed("no-diff")).
+		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c atlas migrate --env ci diff test"}).
+		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c sh -c 'if [[ -n $(git status --porcelain ./tools/migrate/migrations) ]]; then echo \"Migrations directory is dirty\"; exit 1; fi' "}),
+	))
+	// Always lint last 10 migrations
+	p.Go(syncFunc(nix.
+		WithServiceBinding("postgres", postgresNamed("last-10")).
+		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c atlas migrate --env ci lint --latest 10"}),
+	))
+	// Validate checksum is intact
+	p.Go(syncFunc(nix.
+		WithServiceBinding("postgres", postgresNamed("validate")).
+		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c atlas migrate --env ci validate"}),
+	))
+
+	return p.Wait()
+}
+
+func nix(src *dagger.Directory) *dagger.Container {
+	return dag.Nix().SetupNix(dagger.NixSetupNixOpts{
+		Src: src,
+	}).
+		// Note: we have to use `sh -c` as otherwise devenv nix assertion fails: https://github.com/cachix/devenv/blob/b285601679c7686f623791ad93a8e0debc322633/src/modules/top-level.nix#L229
+		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger"}) // Prepare environment
+}
diff --git a/dagger.json b/dagger.json
index 22191526b..3a5d7c87d 100644
--- a/dagger.json
+++ b/dagger.json
@@ -42,6 +42,10 @@
       "name": "kafka",
       "source": "github.com/sagikazarmark/daggerverse/kafka@00deb5158b5d3a71dae8926fd8a4119ac2f45cca"
     },
+    {
+      "name": "nix",
+      "source": "github.com/tsirysndr/daggerverse/nix@631932b459d218e641dec8047085f0cc87cf5f1c"
+    },
     {
       "name": "python",
       "source": "github.com/sagikazarmark/daggerverse/python@00deb5158b5d3a71dae8926fd8a4119ac2f45cca"
diff --git a/flake.nix b/flake.nix
index 4c699832e..8922df813 100644
--- a/flake.nix
+++ b/flake.nix
@@ -13,9 +13,9 @@
         inputs.devenv.flakeModule
       ];
 
-      systems = [ "x86_64-linux" "x86_64-darwin" "aarch64-darwin" ];
+      systems = [ "x86_64-linux" "x86_64-darwin" "aarch64-darwin" "aarch64-linux" ];
 
-      perSystem = { config, self', inputs', pkgs, system, ... }: rec {
+      perSystem = { config, self', inputs', pkgs, lib, system, ... }: rec {
         _module.args.pkgs = import inputs.nixpkgs {
           inherit system;
 
@@ -23,6 +23,7 @@
             (final: prev: {
               dagger = inputs'.dagger.packages.dagger;
               licensei = self'.packages.licensei;
+              atlasx = self'.packages.atlasx;
             })
           ];
         };
@@ -100,7 +101,7 @@
               # python
               poetry
 
-              self'.packages.atlasx
+              atlasx
 
               just
               semver-tool
@@ -126,6 +127,19 @@
           };
 
           ci = devenv.shells.default;
+
+          # Lighteweight target to use inside dagger
+          dagger = {
+            languages = {
+              go = devenv.shells.default.languages.go;
+            };
+            packages = with pkgs; [
+              gnumake
+              git
+              atlasx
+            ];
+            containers = devenv.shells.default.containers;
+          };
         };
 
         packages = {
@@ -151,25 +165,42 @@
             ];
           };
 
-          atlasx = pkgs.stdenv.mkDerivation rec {
-            pname = "atlasx";
-            version = "0.25.0";
-            src = pkgs.fetchurl {
-              # License: https://ariga.io/legal/atlas/eula/eula-20240804.pdf
-              url = "https://release.ariga.io/atlas/atlas-darwin-arm64-v${version}";
-              hash = "sha256-bYJtNDE13UhJWL4ALLKI0sHMZrDS//kFWzguGX63EAo=";
-            };
+          atlasx =
+            let
+              systemMappings = {
+                x86_64-linux = "linux-amd64";
+                x86_64-darwin = "darwin-amd64";
+                aarch64-darwin = "darwin-arm64";
+                aarch64-linux = "linux-arm64";
+              };
+              hashMappings = {
+                x86_64-linux = "sha256-6270kOQ0uqiv/ljHtAi41uCzb+bkf+99rnmsc87/n6w=";
+                x86_64-darwin = "sha256-KeOp6LeHIY59Y2DJVAhMcr9xyb3KItFqEs6y9+uA7rM=";
+                aarch64-darwin = "sha256-bYJtNDE13UhJWL4ALLKI0sHMZrDS//kFWzguGX63EAo=";
+                aarch64-linux = "sha256-pRLZo7bwFJ1Xxlw2Afi/tAT6HSEvQ/B83ZzHGzCKXT8=";
+              };
+            in
+            pkgs.stdenv.mkDerivation rec {
+              pname = "atlasx";
+              version = "0.25.0";
+
+              src = pkgs.fetchurl {
+                # License: https://ariga.io/legal/atlas/eula/eula-20240804.pdf
+                url = "https://release.ariga.io/atlas/atlas-${systemMappings."${system}"}-v${version}";
+                hash = hashMappings."${system}";
+              };
 
-            unpackPhase = ''
-              cp $src atlas
-            '';
+              unpackPhase = ''
+                cp $src atlas
+              '';
 
-            installPhase = ''
-              mkdir -p $out/bin
-              cp atlas $out/bin/atlas
-            '';
+              installPhase = ''
+                mkdir -p $out/bin
+                cp atlas $out/bin/atlas
+                chmod +x $out/bin/atlas
+              '';
 
-          };
+            };
         };
       };
     };

From 028441e523295fd1e96ff4e289a377f76f3ae21c Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Wed, 21 Aug 2024 16:08:09 +0200
Subject: [PATCH 04/13] refactor: use purpose built images and get rid of nix
 in dagger

---
 ci/migrate.go | 85 ++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 63 insertions(+), 22 deletions(-)

diff --git a/ci/migrate.go b/ci/migrate.go
index fd6eef110..ec15a25cb 100644
--- a/ci/migrate.go
+++ b/ci/migrate.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/openmeterio/openmeter/ci/internal/dagger"
 	"github.com/sourcegraph/conc/pool"
@@ -20,38 +21,78 @@ type Migrate struct {
 func (m *Migrate) Check(
 	ctx context.Context,
 ) error {
-	nix := nix(m.Source)
+	app := goModuleCross("").
+		WithSource(m.Source).
+		Container().
+		WithEnvVariable("GOFLAGS", "-tags=musl")
+
+	bin := dag.Container(dagger.ContainerOpts{
+		Platform: "linux/amd64",
+	}).From("arigaio/atlas:0.25.0").File("atlas")
+
+	atlas := app.
+		WithFile("/bin/atlas", bin).
+		WithDirectory("internal/ent", m.Source.Directory("internal/ent")).
+		WithDirectory("tools/migrate/migrations", m.Source.Directory("tools/migrate/migrations")).
+		WithFile("atlas.hcl", m.Source.File("atlas.hcl"))
+
 	p := pool.New().WithErrors().WithContext(ctx)
 
 	// Always validate schema is generated
-	p.Go(syncFunc(nix.
-		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c go generate ./internal/ent/..."}).
-		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c sh -c 'if [[ -n $(git status --porcelain ./internal/ent) ]]; then echo \"Ent schema wasnt generated\"; exit 1; fi' "}),
-	))
+	p.Go(func(ctx context.Context) error {
+		result := app.
+			WithExec([]string{"go", "generate", "-x", "-tags=musl", "-ldflags", "linkmode=external", "./internal/ent/..."}).
+			Directory("internal/ent")
+
+		source := m.Source.Directory("internal/ent")
+
+		err := diff(ctx, source, result)
+		if err != nil {
+			return fmt.Errorf("schema is not in sync with generated code")
+		}
+		return nil
+	})
+
 	// Always validate migrations are in sync with schema
-	p.Go(syncFunc(nix.
-		WithServiceBinding("postgres", postgresNamed("no-diff")).
-		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c atlas migrate --env ci diff test"}).
-		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c sh -c 'if [[ -n $(git status --porcelain ./tools/migrate/migrations) ]]; then echo \"Migrations directory is dirty\"; exit 1; fi' "}),
-	))
+	p.Go(func(ctx context.Context) error {
+		result := atlas.
+			WithServiceBinding("postgres", postgresNamed("no-diff")).
+			WithExec([]string{"atlas", "migrate", "--env", "ci", "diff", "test"}).
+			Directory("tools/migrate/migrations")
+
+		source := m.Source.Directory("tools/migrate/migrations")
+		err := diff(ctx, source, result)
+		if err != nil {
+			return fmt.Errorf("migrations are not in sync with schema")
+		}
+
+		return nil
+	})
+
 	// Always lint last 10 migrations
-	p.Go(syncFunc(nix.
-		WithServiceBinding("postgres", postgresNamed("last-10")).
-		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c atlas migrate --env ci lint --latest 10"}),
+	p.Go(syncFunc(
+		atlas.
+			WithServiceBinding("postgres", postgresNamed("last-10")).
+			WithExec([]string{"atlas", "migrate", "--env", "ci", "lint", "--latest", "10"}),
 	))
+
 	// Validate checksum is intact
-	p.Go(syncFunc(nix.
-		WithServiceBinding("postgres", postgresNamed("validate")).
-		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger -c atlas migrate --env ci validate"}),
+	p.Go(syncFunc(
+		atlas.
+			WithServiceBinding("postgres", postgresNamed("validate")).
+			WithExec([]string{"atlas", "migrate", "--env", "ci", "validate"}),
 	))
 
 	return p.Wait()
 }
 
-func nix(src *dagger.Directory) *dagger.Container {
-	return dag.Nix().SetupNix(dagger.NixSetupNixOpts{
-		Src: src,
-	}).
-		// Note: we have to use `sh -c` as otherwise devenv nix assertion fails: https://github.com/cachix/devenv/blob/b285601679c7686f623791ad93a8e0debc322633/src/modules/top-level.nix#L229
-		WithExec([]string{"sh", "-c", "nix develop --impure .#dagger"}) // Prepare environment
+func diff(ctx context.Context, d1, d2 *dagger.Directory) error {
+	_, err := dag.Container(dagger.ContainerOpts{Platform: ""}).
+		From(alpineBaseImage).
+		WithExec([]string{"apk", "add", "--update", "--no-cache", "ca-certificates", "tzdata", "bash"}).
+		WithDirectory("src", d1).
+		WithDirectory("res", d2).
+		WithExec([]string{"diff", "-u", "-r", "src", "res"}).
+		Sync(ctx)
+	return err
 }

From 1e0f361858bbedf09b5aa200b2e6ec64e66a2934 Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Thu, 22 Aug 2024 11:27:02 +0200
Subject: [PATCH 05/13] refactor: package naming

---
 cmd/balance-worker/main.go                      |  2 +-
 cmd/jobs/entitlement/init.go                    |  2 +-
 cmd/notification-service/main.go                |  2 +-
 cmd/server/main.go                              |  2 +-
 openmeter/testutils/pg_driver.go                | 17 +++++++++++++----
 .../entutils/{driver => entdriver}/driver.go    |  2 +-
 test/notification/testenv.go                    |  2 +-
 7 files changed, 19 insertions(+), 10 deletions(-)
 rename pkg/framework/entutils/{driver => entdriver}/driver.go (98%)

diff --git a/cmd/balance-worker/main.go b/cmd/balance-worker/main.go
index ffd0f9bc4..699624e75 100644
--- a/cmd/balance-worker/main.go
+++ b/cmd/balance-worker/main.go
@@ -40,7 +40,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
 	"github.com/openmeterio/openmeter/openmeter/watermill/router"
 	"github.com/openmeterio/openmeter/pkg/contextx"
-	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/driver"
+	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
 	"github.com/openmeterio/openmeter/pkg/framework/operation"
 	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 	"github.com/openmeterio/openmeter/pkg/gosundheit"
diff --git a/cmd/jobs/entitlement/init.go b/cmd/jobs/entitlement/init.go
index 09822bb77..14aaac1ea 100644
--- a/cmd/jobs/entitlement/init.go
+++ b/cmd/jobs/entitlement/init.go
@@ -15,7 +15,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/streaming/clickhouse_connector"
 	watermillkafka "github.com/openmeterio/openmeter/openmeter/watermill/driver/kafka"
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
-	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/driver"
+	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
 	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 	"github.com/openmeterio/openmeter/pkg/models"
 	"github.com/openmeterio/openmeter/pkg/slicesx"
diff --git a/cmd/notification-service/main.go b/cmd/notification-service/main.go
index b4becb36d..756edb7e7 100644
--- a/cmd/notification-service/main.go
+++ b/cmd/notification-service/main.go
@@ -42,7 +42,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
 	"github.com/openmeterio/openmeter/openmeter/watermill/router"
 	"github.com/openmeterio/openmeter/pkg/contextx"
-	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/driver"
+	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
 	"github.com/openmeterio/openmeter/pkg/framework/operation"
 	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 	"github.com/openmeterio/openmeter/pkg/gosundheit"
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 562bcad3d..5dc5722a1 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -59,7 +59,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
 	"github.com/openmeterio/openmeter/pkg/contextx"
 	"github.com/openmeterio/openmeter/pkg/errorsx"
-	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/driver"
+	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
 	"github.com/openmeterio/openmeter/pkg/framework/operation"
 	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 	"github.com/openmeterio/openmeter/pkg/gosundheit"
diff --git a/openmeter/testutils/pg_driver.go b/openmeter/testutils/pg_driver.go
index 1db0ba0b0..48bdfb296 100644
--- a/openmeter/testutils/pg_driver.go
+++ b/openmeter/testutils/pg_driver.go
@@ -11,7 +11,7 @@ import (
 	_ "github.com/jackc/pgx/v5/stdlib" // pgx database driver
 	"github.com/peterldowns/pgtestdb"
 
-	"github.com/openmeterio/openmeter/pkg/framework/entutils"
+	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 )
 
 // NoopMigrator is a migrator for pgtestdb.
@@ -79,16 +79,25 @@ func InitPostgresDB(t *testing.T) *TestDB {
 		Options:    "sslmode=disable",
 	}, &NoopMigrator{})
 
-	sqlDriver, err := entutils.GetSQLDriver(dbConf.URL())
+	postgresDriver, err := pgdriver.NewPostgresDriver(
+		context.TODO(),
+		dbConf.URL(),
+	)
 	if err != nil {
 		t.Fatalf("failed to get pg driver: %s", err)
 	}
 
-	entDriver := entsql.OpenDB(dialect.Postgres, sqlDriver)
+	defer func() {
+		if err = postgresDriver.Close(); err != nil {
+			t.Error("failed to close ent driver", "error", err)
+		}
+	}()
+
+	entDriver := entsql.OpenDB(dialect.Postgres, postgresDriver.DB())
 
 	return &TestDB{
 		EntDriver: entDriver,
-		SQLDriver: sqlDriver,
+		SQLDriver: postgresDriver.DB(),
 		URL:       dbConf.URL(),
 	}
 }
diff --git a/pkg/framework/entutils/driver/driver.go b/pkg/framework/entutils/entdriver/driver.go
similarity index 98%
rename from pkg/framework/entutils/driver/driver.go
rename to pkg/framework/entutils/entdriver/driver.go
index bb3f9344a..0730062d9 100644
--- a/pkg/framework/entutils/driver/driver.go
+++ b/pkg/framework/entutils/entdriver/driver.go
@@ -1,4 +1,4 @@
-package entutils
+package entdriver
 
 import (
 	"database/sql"
diff --git a/test/notification/testenv.go b/test/notification/testenv.go
index 7c9cc3102..698974399 100644
--- a/test/notification/testenv.go
+++ b/test/notification/testenv.go
@@ -16,7 +16,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/productcatalog"
 	productcatalogadapter "github.com/openmeterio/openmeter/openmeter/productcatalog/adapter"
 	"github.com/openmeterio/openmeter/pkg/defaultx"
-	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/driver"
+	entdriver "github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
 	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 )
 

From 55e58b5a853e6f901a3f62f662ac0aaa56eb2df9 Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Thu, 22 Aug 2024 12:04:24 +0200
Subject: [PATCH 06/13] chore: update to atlas 0.26.0

---
 ci/migrate.go |  2 +-
 flake.nix     | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/ci/migrate.go b/ci/migrate.go
index ec15a25cb..52574d4d1 100644
--- a/ci/migrate.go
+++ b/ci/migrate.go
@@ -28,7 +28,7 @@ func (m *Migrate) Check(
 
 	bin := dag.Container(dagger.ContainerOpts{
 		Platform: "linux/amd64",
-	}).From("arigaio/atlas:0.25.0").File("atlas")
+	}).From("arigaio/atlas:0.26.0").File("atlas")
 
 	atlas := app.
 		WithFile("/bin/atlas", bin).
diff --git a/flake.nix b/flake.nix
index 8922df813..04cec45b3 100644
--- a/flake.nix
+++ b/flake.nix
@@ -174,15 +174,15 @@
                 aarch64-linux = "linux-arm64";
               };
               hashMappings = {
-                x86_64-linux = "sha256-6270kOQ0uqiv/ljHtAi41uCzb+bkf+99rnmsc87/n6w=";
-                x86_64-darwin = "sha256-KeOp6LeHIY59Y2DJVAhMcr9xyb3KItFqEs6y9+uA7rM=";
-                aarch64-darwin = "sha256-bYJtNDE13UhJWL4ALLKI0sHMZrDS//kFWzguGX63EAo=";
-                aarch64-linux = "sha256-pRLZo7bwFJ1Xxlw2Afi/tAT6HSEvQ/B83ZzHGzCKXT8=";
+                x86_64-linux = "sha256-C/my+oOVX1DzdwuA2otud1Zer+0WOPjR749E8c03TOE=";
+                x86_64-darwin = "sha256-YWrLeH3d0x4I5lr3uArJqsQW3uY3XIy22PiNEmHrSLc=";
+                aarch64-darwin = "sha256-73G/kvh7gqXJpVidR7zB1e+0bkDegvFFavCsHYxSI1Q=";
+                aarch64-linux = "sha256-bQCCJF2tYkAAgwbBc1+nkdNtczmU+eRJlXRppulFvvc=";
               };
             in
             pkgs.stdenv.mkDerivation rec {
               pname = "atlasx";
-              version = "0.25.0";
+              version = "0.26.0";
 
               src = pkgs.fetchurl {
                 # License: https://ariga.io/legal/atlas/eula/eula-20240804.pdf

From ec1f033a333c55e62c7a21a636e53571253cb3f0 Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Thu, 22 Aug 2024 12:56:23 +0200
Subject: [PATCH 07/13] fix: pg connection closing in tests

---
 openmeter/entitlement/metered/utils_test.go   | 18 +++++++++++++-----
 .../productcatalog/adapter/feature_test.go    | 19 ++++++++-----------
 openmeter/testutils/pg_driver.go              | 16 +++++-----------
 pkg/framework/entutils/paginate_test.go       |  3 ++-
 pkg/framework/entutils/transaction_test.go    |  5 +++--
 test/entitlement/regression/framework_test.go | 15 +++++++++++----
 tools/migrate/migrate_test.go                 |  3 ++-
 7 files changed, 44 insertions(+), 35 deletions(-)

diff --git a/openmeter/entitlement/metered/utils_test.go b/openmeter/entitlement/metered/utils_test.go
index 1c8915e5e..ef74b3fae 100644
--- a/openmeter/entitlement/metered/utils_test.go
+++ b/openmeter/entitlement/metered/utils_test.go
@@ -21,12 +21,16 @@ import (
 	streaming_testutils "github.com/openmeterio/openmeter/openmeter/streaming/testutils"
 	"github.com/openmeterio/openmeter/openmeter/testutils"
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
+	"github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
+	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 	"github.com/openmeterio/openmeter/pkg/models"
 	"github.com/openmeterio/openmeter/tools/migrate"
 )
 
 type dependencies struct {
 	dbClient            *db.Client
+	pgDriver            *pgdriver.Driver
+	entDriver           *entdriver.EntPostgresDriver
 	featureRepo         productcatalog.FeatureRepo
 	entitlementRepo     entitlement.EntitlementRepo
 	usageResetRepo      meteredentitlement.UsageResetRepo
@@ -39,6 +43,8 @@ type dependencies struct {
 // Teardown cleans up the dependencies
 func (d *dependencies) Teardown() {
 	d.dbClient.Close()
+	d.entDriver.Close()
+	d.pgDriver.Close()
 }
 
 // When migrating in parallel with entgo it causes concurrent writes error
@@ -57,10 +63,10 @@ func setupConnector(t *testing.T) (meteredentitlement.Connector, *dependencies)
 	}})
 
 	// create isolated pg db for tests
-	driver := testutils.InitPostgresDB(t)
-
-	// build db client & adapters
-	dbClient := db.NewClient(db.Driver(driver.EntDriver))
+	testdb := testutils.InitPostgresDB(t)
+	dbClient := testdb.EntDriver.Client()
+	pgDriver := testdb.PGDriver
+	entDriver := testdb.EntDriver
 
 	featureRepo := productcatalog_postgresadapter.NewPostgresFeatureRepo(dbClient, testLogger)
 	entitlementRepo := entitlement_postgresadapter.NewPostgresEntitlementRepo(dbClient)
@@ -71,7 +77,7 @@ func setupConnector(t *testing.T) (meteredentitlement.Connector, *dependencies)
 	m.Lock()
 	defer m.Unlock()
 	// migrate db
-	if err := migrate.Up(driver.URL); err != nil {
+	if err := migrate.Up(testdb.URL); err != nil {
 		t.Fatalf("failed to migrate db: %s", err.Error())
 	}
 
@@ -108,6 +114,8 @@ func setupConnector(t *testing.T) (meteredentitlement.Connector, *dependencies)
 
 	return connector, &dependencies{
 		dbClient,
+		pgDriver,
+		entDriver,
 		featureRepo,
 		entitlementRepo,
 		usageResetRepo,
diff --git a/openmeter/productcatalog/adapter/feature_test.go b/openmeter/productcatalog/adapter/feature_test.go
index d8b6f9786..12d3959cd 100644
--- a/openmeter/productcatalog/adapter/feature_test.go
+++ b/openmeter/productcatalog/adapter/feature_test.go
@@ -9,7 +9,6 @@ import (
 	"github.com/oklog/ulid/v2"
 	"github.com/stretchr/testify/assert"
 
-	"github.com/openmeterio/openmeter/openmeter/ent/db"
 	db_feature "github.com/openmeterio/openmeter/openmeter/ent/db/feature"
 	"github.com/openmeterio/openmeter/openmeter/productcatalog"
 	"github.com/openmeterio/openmeter/openmeter/productcatalog/adapter"
@@ -216,15 +215,12 @@ func TestCreateFeature(t *testing.T) {
 			m.Lock()
 			defer m.Unlock()
 
-			driver := testutils.InitPostgresDB(t)
-			defer driver.SQLDriver.Close()
-			dbClient := db.NewClient(db.Driver(driver.EntDriver))
+			testdb := testutils.InitPostgresDB(t)
+			defer testdb.PGDriver.Close()
+			dbClient := testdb.EntDriver.Client()
 			defer dbClient.Close()
 
-			m.Lock()
-			defer m.Unlock()
-
-			if err := migrate.Up(driver.URL); err != nil {
+			if err := migrate.Up(testdb.URL); err != nil {
 				t.Fatalf("failed to migrate db: %s", err.Error())
 			}
 
@@ -238,11 +234,12 @@ func TestCreateFeature(t *testing.T) {
 		m.Lock()
 		defer m.Unlock()
 
-		driver := testutils.InitPostgresDB(t)
-		dbClient := db.NewClient(db.Driver(driver.EntDriver))
+		testdb := testutils.InitPostgresDB(t)
+		defer testdb.PGDriver.Close()
+		dbClient := testdb.EntDriver.Client()
 		defer dbClient.Close()
 
-		if err := migrate.Up(driver.URL); err != nil {
+		if err := migrate.Up(testdb.URL); err != nil {
 			t.Fatalf("failed to migrate db: %s", err.Error())
 		}
 
diff --git a/openmeter/testutils/pg_driver.go b/openmeter/testutils/pg_driver.go
index 48bdfb296..a3d74855b 100644
--- a/openmeter/testutils/pg_driver.go
+++ b/openmeter/testutils/pg_driver.go
@@ -6,11 +6,10 @@ import (
 	"os"
 	"testing"
 
-	"entgo.io/ent/dialect"
-	entsql "entgo.io/ent/dialect/sql"
 	_ "github.com/jackc/pgx/v5/stdlib" // pgx database driver
 	"github.com/peterldowns/pgtestdb"
 
+	"github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
 	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 )
 
@@ -53,7 +52,8 @@ func (*NoopMigrator) Verify(
 }
 
 type TestDB struct {
-	EntDriver *entsql.Driver
+	EntDriver *entdriver.EntPostgresDriver
+	PGDriver  *pgdriver.Driver
 	SQLDriver *sql.DB
 	URL       string
 }
@@ -87,17 +87,11 @@ func InitPostgresDB(t *testing.T) *TestDB {
 		t.Fatalf("failed to get pg driver: %s", err)
 	}
 
-	defer func() {
-		if err = postgresDriver.Close(); err != nil {
-			t.Error("failed to close ent driver", "error", err)
-		}
-	}()
-
-	entDriver := entsql.OpenDB(dialect.Postgres, postgresDriver.DB())
+	entDriver := entdriver.NewEntPostgresDriver(postgresDriver.DB())
 
 	return &TestDB{
+		PGDriver:  postgresDriver,
 		EntDriver: entDriver,
-		SQLDriver: postgresDriver.DB(),
 		URL:       dbConf.URL(),
 	}
 }
diff --git a/pkg/framework/entutils/paginate_test.go b/pkg/framework/entutils/paginate_test.go
index 7eadc1835..88f74690a 100644
--- a/pkg/framework/entutils/paginate_test.go
+++ b/pkg/framework/entutils/paginate_test.go
@@ -19,9 +19,10 @@ func TestPaginate(t *testing.T) {
 
 	// create isolated pg db for tests
 	driver := testutils.InitPostgresDB(t)
+	defer driver.PGDriver.Close()
 
 	// build db clients
-	dbClient := db.NewClient(db.Driver(driver.EntDriver))
+	dbClient := db.NewClient(db.Driver(driver.EntDriver.Driver()))
 	defer dbClient.Close()
 
 	if err := dbClient.Schema.Create(context.Background()); err != nil {
diff --git a/pkg/framework/entutils/transaction_test.go b/pkg/framework/entutils/transaction_test.go
index cf2d2d5a9..3aaf4a766 100644
--- a/pkg/framework/entutils/transaction_test.go
+++ b/pkg/framework/entutils/transaction_test.go
@@ -291,11 +291,12 @@ func TestTransaction(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			// create isolated pg db for tests
 			driver := testutils.InitPostgresDB(t)
+			defer driver.PGDriver.Close()
 			defer driver.EntDriver.Close()
 
 			// build db clients
-			db1Client := db1.NewClient(db1.Driver(driver.EntDriver))
-			db2Client := db2.NewClient(db2.Driver(driver.EntDriver))
+			db1Client := db1.NewClient(db1.Driver(driver.EntDriver.Driver()))
+			db2Client := db2.NewClient(db2.Driver(driver.EntDriver.Driver()))
 
 			if err := db1Client.Schema.Create(context.Background()); err != nil {
 				t.Fatalf("failed to migrate database %s", err)
diff --git a/test/entitlement/regression/framework_test.go b/test/entitlement/regression/framework_test.go
index 74b7f1cb5..7ddd65f3a 100644
--- a/test/entitlement/regression/framework_test.go
+++ b/test/entitlement/regression/framework_test.go
@@ -21,12 +21,16 @@ import (
 	streamingtestutils "github.com/openmeterio/openmeter/openmeter/streaming/testutils"
 	"github.com/openmeterio/openmeter/openmeter/testutils"
 	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
+	"github.com/openmeterio/openmeter/pkg/framework/entutils/entdriver"
+	"github.com/openmeterio/openmeter/pkg/framework/pgdriver"
 	"github.com/openmeterio/openmeter/pkg/models"
 	"github.com/openmeterio/openmeter/tools/migrate"
 )
 
 type Dependencies struct {
-	DBClient *db.Client
+	DBClient  *db.Client
+	PGDriver  *pgdriver.Driver
+	EntDriver *entdriver.EntPostgresDriver
 
 	GrantRepo           grant.Repo
 	BalanceSnapshotRepo balance.SnapshotRepo
@@ -49,14 +53,15 @@ type Dependencies struct {
 
 func (d *Dependencies) Close() {
 	d.DBClient.Close()
+	d.EntDriver.Close()
+	d.PGDriver.Close()
 }
 
 func setupDependencies(t *testing.T) Dependencies {
 	log := slog.Default()
 	driver := testutils.InitPostgresDB(t)
-
 	// init db
-	dbClient := db.NewClient(db.Driver(driver.EntDriver))
+	dbClient := db.NewClient(db.Driver(driver.EntDriver.Driver()))
 	if err := migrate.Up(driver.URL); err != nil {
 		t.Fatalf("failed to migrate db: %s", err.Error())
 	}
@@ -132,7 +137,9 @@ func setupDependencies(t *testing.T) Dependencies {
 	)
 
 	return Dependencies{
-		DBClient: dbClient,
+		DBClient:  dbClient,
+		PGDriver:  driver.PGDriver,
+		EntDriver: driver.EntDriver,
 
 		GrantRepo:      grantRepo,
 		GrantConnector: creditConnector,
diff --git a/tools/migrate/migrate_test.go b/tools/migrate/migrate_test.go
index ec8f357cc..1a3775ef8 100644
--- a/tools/migrate/migrate_test.go
+++ b/tools/migrate/migrate_test.go
@@ -4,12 +4,13 @@ import (
 	"errors"
 	"testing"
 
-	"github.com/openmeterio/openmeter/internal/testutils"
+	"github.com/openmeterio/openmeter/openmeter/testutils"
 	"github.com/openmeterio/openmeter/tools/migrate"
 )
 
 func TestUpDownUp(t *testing.T) {
 	testDB := testutils.InitPostgresDB(t)
+	defer testDB.PGDriver.Close()
 
 	migrator, err := migrate.NewMigrate(testDB.URL, migrate.OMMigrations, "migrations")
 	if err != nil {

From 56387745a1799eba090660ef5e12192ac82f83b5 Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Mon, 26 Aug 2024 14:10:17 +0200
Subject: [PATCH 08/13] chore: update atlas version and regenerate init
 migrations

---
 atlas.hcl                                     |  4 +-
 flake.nix                                     | 10 +-
 go.mod                                        |  6 +-
 go.sum                                        | 12 +--
 .../migrations/20240812143658_init.down.sql   | 44 ---------
 .../migrations/20240826120919_init.down.sql   | 84 +++++++++++++++++
 ...init.up.sql => 20240826120919_init.up.sql} | 94 +++++++++++++++++++
 tools/migrate/migrations/atlas.sum            |  6 +-
 8 files changed, 197 insertions(+), 63 deletions(-)
 delete mode 100644 tools/migrate/migrations/20240812143658_init.down.sql
 create mode 100644 tools/migrate/migrations/20240826120919_init.down.sql
 rename tools/migrate/migrations/{20240812143658_init.up.sql => 20240826120919_init.up.sql} (51%)

diff --git a/atlas.hcl b/atlas.hcl
index 4a52f21ea..3b340a865 100644
--- a/atlas.hcl
+++ b/atlas.hcl
@@ -67,9 +67,9 @@ env "remote" {
 
 locals {
     // Define the directory where the schema definition resides.
-    schema_src = "ent://internal/ent/schema"
+    schema_src = "ent://openmeter/ent/schema"
     // Define the initial migration timestamp
-    init_migration_ts = "20240807123504"
+    init_migration_ts = "20240826120919"
     // Define the directory where the migrations are stored.
     migrations_dir = "file://tools/migrate/migrations"
     // We use golang-migrate
diff --git a/flake.nix b/flake.nix
index 04cec45b3..19b781843 100644
--- a/flake.nix
+++ b/flake.nix
@@ -174,15 +174,15 @@
                 aarch64-linux = "linux-arm64";
               };
               hashMappings = {
-                x86_64-linux = "sha256-C/my+oOVX1DzdwuA2otud1Zer+0WOPjR749E8c03TOE=";
-                x86_64-darwin = "sha256-YWrLeH3d0x4I5lr3uArJqsQW3uY3XIy22PiNEmHrSLc=";
-                aarch64-darwin = "sha256-73G/kvh7gqXJpVidR7zB1e+0bkDegvFFavCsHYxSI1Q=";
-                aarch64-linux = "sha256-bQCCJF2tYkAAgwbBc1+nkdNtczmU+eRJlXRppulFvvc=";
+                x86_64-linux = "sha256-G+ZmZYJ4ep4Je49/V8qYKd/R7h1/KpYxl0wsjfB1LXw=";
+                x86_64-darwin = "sha256-kHvKZjNiQg5CJY46re0R+1FLzKN7VC4hmJGg8RB8n6A=";
+                aarch64-darwin = "sha256-OOaMsgvu4MCVCk0ayN4XSFewpcTdKF8ilqU1Be5Aii8=";
+                aarch64-linux = "sha256-EVg7QBd3L1NXGIhF1odtADSY6/paEuXCP2mY3ripFmA=";
               };
             in
             pkgs.stdenv.mkDerivation rec {
               pname = "atlasx";
-              version = "0.26.0";
+              version = "0.26.1";
 
               src = pkgs.fetchurl {
                 # License: https://ariga.io/legal/atlas/eula/eula-20240804.pdf
diff --git a/go.mod b/go.mod
index 239283ef8..4f89433fa 100644
--- a/go.mod
+++ b/go.mod
@@ -5,7 +5,7 @@ go 1.23.0
 replace github.com/getkin/kin-openapi => github.com/getkin/kin-openapi v0.123.0
 
 require (
-	entgo.io/ent v0.14.0
+	entgo.io/ent v0.14.1
 	github.com/AppsFlyer/go-sundheit v0.6.0
 	github.com/ClickHouse/clickhouse-go/v2 v2.28.0
 	github.com/IBM/sarama v1.43.2
@@ -393,14 +393,14 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.27.0 // indirect
 	golang.org/x/crypto v0.26.0 // indirect
-	golang.org/x/mod v0.17.0 // indirect
+	golang.org/x/mod v0.20.0 // indirect
 	golang.org/x/oauth2 v0.21.0 // indirect
 	golang.org/x/sync v0.8.0 // indirect
 	golang.org/x/sys v0.24.0 // indirect
 	golang.org/x/term v0.23.0 // indirect
 	golang.org/x/text v0.17.0 // indirect
 	golang.org/x/time v0.5.0 // indirect
-	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
+	golang.org/x/tools v0.24.0 // indirect
 	golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
 	google.golang.org/api v0.169.0 // indirect
 	google.golang.org/genproto v0.0.0-20240325203815-454cdb8f5daa // indirect
diff --git a/go.sum b/go.sum
index c0cc60737..f16fb5417 100644
--- a/go.sum
+++ b/go.sum
@@ -52,8 +52,8 @@ cuelang.org/go v0.7.0/go.mod h1:ix+3dM/bSpdG9xg6qpCgnJnpeLtciZu+O/rDbywoMII=
 dario.cat/mergo v1.0.0 h1:AGCNq9Evsj31mOgNPcLyXc+4PNABt905YmuqPYYpBWk=
 dario.cat/mergo v1.0.0/go.mod h1:uNxQE+84aUszobStD9th8a29P2fMDhsBdgRYvZOxGmk=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
-entgo.io/ent v0.14.0 h1:EO3Z9aZ5bXJatJeGqu/EVdnNr6K4mRq3rWe5owt0MC4=
-entgo.io/ent v0.14.0/go.mod h1:qCEmo+biw3ccBn9OyL4ZK5dfpwg++l1Gxwac5B1206A=
+entgo.io/ent v0.14.1 h1:fUERL506Pqr92EPHJqr8EYxbPioflJo6PudkrEA8a/s=
+entgo.io/ent v0.14.1/go.mod h1:MH6XLG0KXpkcDQhKiHfANZSzR55TJyPL5IGNpI8wpco=
 gioui.org v0.0.0-20210308172011-57750fc8a0a6/go.mod h1:RSH6KIUZ0p2xy5zHDxgAM4zumjgTw83q2ge/PI+yyw8=
 github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 h1:/vQbFIOMbk2FiG/kXiLl8BRyzTWDw7gX/Hz7Dd5eDMs=
 github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4/go.mod h1:hN7oaIRCjzsZ2dE+yG5k+rsdt3qcwykqK6HVGcKwsw4=
@@ -1485,8 +1485,8 @@ golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
-golang.org/x/mod v0.17.0 h1:zY54UmvipHiNd+pm+m0x9KhZ9hl1/7QNMyxXbc6ICqA=
-golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/mod v0.20.0 h1:utOm6MM3R3dnawAiJgn0y+xvuYRsm1RKM/4giyfDgV0=
+golang.org/x/mod v0.20.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
 golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20190108225652-1e06a53dbb7e/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
@@ -1674,8 +1674,8 @@ golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4f
 golang.org/x/tools v0.1.4/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
 golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
-golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d h1:vU5i/LfpvrRCpgM/VPfJLg5KjxD3E+hfT1SH+d9zLwg=
-golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
+golang.org/x/tools v0.24.0 h1:J1shsA93PJUEVaUSaay7UXAyE8aimq3GW0pjlolpa24=
+golang.org/x/tools v0.24.0/go.mod h1:YhNqVBIfWHdzvTLs0d8LCuMhkKUgSUKldakyV7W/WDQ=
 golang.org/x/xerrors v0.0.0-20190410155217-1f06c39b4373/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20190513163551-3ee3066db522/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
diff --git a/tools/migrate/migrations/20240812143658_init.down.sql b/tools/migrate/migrations/20240812143658_init.down.sql
deleted file mode 100644
index 519f19707..000000000
--- a/tools/migrate/migrations/20240812143658_init.down.sql
+++ /dev/null
@@ -1,44 +0,0 @@
--- reverse: create index "usagereset_namespace_entitlement_id_reset_time" to table: "usage_resets"
-DROP INDEX "usagereset_namespace_entitlement_id_reset_time";
--- reverse: create index "usagereset_namespace_entitlement_id" to table: "usage_resets"
-DROP INDEX "usagereset_namespace_entitlement_id";
--- reverse: create index "usagereset_id" to table: "usage_resets"
-DROP INDEX "usagereset_id";
--- reverse: create "usage_resets" table
-DROP TABLE "usage_resets";
--- reverse: create index "grant_namespace_owner_id" to table: "grants"
-DROP INDEX "grant_namespace_owner_id";
--- reverse: create index "grant_id" to table: "grants"
-DROP INDEX "grant_id";
--- reverse: create index "grant_effective_at_expires_at" to table: "grants"
-DROP INDEX "grant_effective_at_expires_at";
--- reverse: create "grants" table
-DROP TABLE "grants";
--- reverse: create index "balancesnapshot_namespace_balance_at" to table: "balance_snapshots"
-DROP INDEX "balancesnapshot_namespace_balance_at";
--- reverse: create index "balancesnapshot_namespace_balance" to table: "balance_snapshots"
-DROP INDEX "balancesnapshot_namespace_balance";
--- reverse: create index "balancesnapshot_namespace_at" to table: "balance_snapshots"
-DROP INDEX "balancesnapshot_namespace_at";
--- reverse: create "balance_snapshots" table
-DROP TABLE "balance_snapshots";
--- reverse: create index "entitlement_namespace_subject_key" to table: "entitlements"
-DROP INDEX "entitlement_namespace_subject_key";
--- reverse: create index "entitlement_namespace_id_subject_key" to table: "entitlements"
-DROP INDEX "entitlement_namespace_id_subject_key";
--- reverse: create index "entitlement_namespace_id" to table: "entitlements"
-DROP INDEX "entitlement_namespace_id";
--- reverse: create index "entitlement_namespace_feature_id_id" to table: "entitlements"
-DROP INDEX "entitlement_namespace_feature_id_id";
--- reverse: create index "entitlement_namespace_current_usage_period_end" to table: "entitlements"
-DROP INDEX "entitlement_namespace_current_usage_period_end";
--- reverse: create index "entitlement_id" to table: "entitlements"
-DROP INDEX "entitlement_id";
--- reverse: create "entitlements" table
-DROP TABLE "entitlements";
--- reverse: create index "feature_namespace_id" to table: "features"
-DROP INDEX "feature_namespace_id";
--- reverse: create index "feature_id" to table: "features"
-DROP INDEX "feature_id";
--- reverse: create "features" table
-DROP TABLE "features";
diff --git a/tools/migrate/migrations/20240826120919_init.down.sql b/tools/migrate/migrations/20240826120919_init.down.sql
new file mode 100644
index 000000000..00c642792
--- /dev/null
+++ b/tools/migrate/migrations/20240826120919_init.down.sql
@@ -0,0 +1,84 @@
+-- reverse: create index "usagereset_namespace_entitlement_id_reset_time" to table: "usage_resets"
+DROP INDEX "usagereset_namespace_entitlement_id_reset_time";
+-- reverse: create index "usagereset_namespace_entitlement_id" to table: "usage_resets"
+DROP INDEX "usagereset_namespace_entitlement_id";
+-- reverse: create index "usagereset_id" to table: "usage_resets"
+DROP INDEX "usagereset_id";
+-- reverse: create "usage_resets" table
+DROP TABLE "usage_resets";
+-- reverse: create "notification_event_delivery_status_events" table
+DROP TABLE "notification_event_delivery_status_events";
+-- reverse: create index "notificationevent_namespace_type" to table: "notification_events"
+DROP INDEX "notificationevent_namespace_type";
+-- reverse: create index "notificationevent_namespace_id" to table: "notification_events"
+DROP INDEX "notificationevent_namespace_id";
+-- reverse: create index "notificationevent_id" to table: "notification_events"
+DROP INDEX "notificationevent_id";
+-- reverse: create index "notificationevent_annotations" to table: "notification_events"
+DROP INDEX "notificationevent_annotations";
+-- reverse: create "notification_events" table
+DROP TABLE "notification_events";
+-- reverse: create index "notificationeventdeliverystatus_namespace_state" to table: "notification_event_delivery_status"
+DROP INDEX "notificationeventdeliverystatus_namespace_state";
+-- reverse: create index "notificationeventdeliverystatus_namespace_id" to table: "notification_event_delivery_status"
+DROP INDEX "notificationeventdeliverystatus_namespace_id";
+-- reverse: create index "notificationeventdeliverystatus_namespace_event_id_channel_id" to table: "notification_event_delivery_status"
+DROP INDEX "notificationeventdeliverystatus_namespace_event_id_channel_id";
+-- reverse: create index "notificationeventdeliverystatus_id" to table: "notification_event_delivery_status"
+DROP INDEX "notificationeventdeliverystatus_id";
+-- reverse: create "notification_event_delivery_status" table
+DROP TABLE "notification_event_delivery_status";
+-- reverse: create "notification_channel_rules" table
+DROP TABLE "notification_channel_rules";
+-- reverse: create index "notificationrule_namespace_type" to table: "notification_rules"
+DROP INDEX "notificationrule_namespace_type";
+-- reverse: create index "notificationrule_namespace_id" to table: "notification_rules"
+DROP INDEX "notificationrule_namespace_id";
+-- reverse: create index "notificationrule_id" to table: "notification_rules"
+DROP INDEX "notificationrule_id";
+-- reverse: create "notification_rules" table
+DROP TABLE "notification_rules";
+-- reverse: create index "notificationchannel_namespace_type" to table: "notification_channels"
+DROP INDEX "notificationchannel_namespace_type";
+-- reverse: create index "notificationchannel_namespace_id" to table: "notification_channels"
+DROP INDEX "notificationchannel_namespace_id";
+-- reverse: create index "notificationchannel_id" to table: "notification_channels"
+DROP INDEX "notificationchannel_id";
+-- reverse: create "notification_channels" table
+DROP TABLE "notification_channels";
+-- reverse: create index "grant_namespace_owner_id" to table: "grants"
+DROP INDEX "grant_namespace_owner_id";
+-- reverse: create index "grant_id" to table: "grants"
+DROP INDEX "grant_id";
+-- reverse: create index "grant_effective_at_expires_at" to table: "grants"
+DROP INDEX "grant_effective_at_expires_at";
+-- reverse: create "grants" table
+DROP TABLE "grants";
+-- reverse: create index "balancesnapshot_namespace_balance_at" to table: "balance_snapshots"
+DROP INDEX "balancesnapshot_namespace_balance_at";
+-- reverse: create index "balancesnapshot_namespace_balance" to table: "balance_snapshots"
+DROP INDEX "balancesnapshot_namespace_balance";
+-- reverse: create index "balancesnapshot_namespace_at" to table: "balance_snapshots"
+DROP INDEX "balancesnapshot_namespace_at";
+-- reverse: create "balance_snapshots" table
+DROP TABLE "balance_snapshots";
+-- reverse: create index "entitlement_namespace_subject_key" to table: "entitlements"
+DROP INDEX "entitlement_namespace_subject_key";
+-- reverse: create index "entitlement_namespace_id_subject_key" to table: "entitlements"
+DROP INDEX "entitlement_namespace_id_subject_key";
+-- reverse: create index "entitlement_namespace_id" to table: "entitlements"
+DROP INDEX "entitlement_namespace_id";
+-- reverse: create index "entitlement_namespace_feature_id_id" to table: "entitlements"
+DROP INDEX "entitlement_namespace_feature_id_id";
+-- reverse: create index "entitlement_namespace_current_usage_period_end" to table: "entitlements"
+DROP INDEX "entitlement_namespace_current_usage_period_end";
+-- reverse: create index "entitlement_id" to table: "entitlements"
+DROP INDEX "entitlement_id";
+-- reverse: create "entitlements" table
+DROP TABLE "entitlements";
+-- reverse: create index "feature_namespace_id" to table: "features"
+DROP INDEX "feature_namespace_id";
+-- reverse: create index "feature_id" to table: "features"
+DROP INDEX "feature_id";
+-- reverse: create "features" table
+DROP TABLE "features";
diff --git a/tools/migrate/migrations/20240812143658_init.up.sql b/tools/migrate/migrations/20240826120919_init.up.sql
similarity index 51%
rename from tools/migrate/migrations/20240812143658_init.up.sql
rename to tools/migrate/migrations/20240826120919_init.up.sql
index 1d68d0f3b..bf35fef85 100644
--- a/tools/migrate/migrations/20240812143658_init.up.sql
+++ b/tools/migrate/migrations/20240826120919_init.up.sql
@@ -103,6 +103,100 @@ CREATE INDEX "grant_effective_at_expires_at" ON "grants" ("effective_at", "expir
 CREATE INDEX "grant_id" ON "grants" ("id");
 -- create index "grant_namespace_owner_id" to table: "grants"
 CREATE INDEX "grant_namespace_owner_id" ON "grants" ("namespace", "owner_id");
+-- create "notification_channels" table
+CREATE TABLE "notification_channels" (
+  "id" character(26) NOT NULL,
+  "namespace" character varying NOT NULL,
+  "created_at" timestamptz NOT NULL,
+  "updated_at" timestamptz NOT NULL,
+  "deleted_at" timestamptz NULL,
+  "type" character varying NOT NULL,
+  "name" character varying NOT NULL,
+  "disabled" boolean NULL DEFAULT false,
+  "config" jsonb NOT NULL,
+  PRIMARY KEY ("id")
+);
+-- create index "notificationchannel_id" to table: "notification_channels"
+CREATE INDEX "notificationchannel_id" ON "notification_channels" ("id");
+-- create index "notificationchannel_namespace_id" to table: "notification_channels"
+CREATE INDEX "notificationchannel_namespace_id" ON "notification_channels" ("namespace", "id");
+-- create index "notificationchannel_namespace_type" to table: "notification_channels"
+CREATE INDEX "notificationchannel_namespace_type" ON "notification_channels" ("namespace", "type");
+-- create "notification_rules" table
+CREATE TABLE "notification_rules" (
+  "id" character(26) NOT NULL,
+  "namespace" character varying NOT NULL,
+  "created_at" timestamptz NOT NULL,
+  "updated_at" timestamptz NOT NULL,
+  "deleted_at" timestamptz NULL,
+  "type" character varying NOT NULL,
+  "name" character varying NOT NULL,
+  "disabled" boolean NULL DEFAULT false,
+  "config" jsonb NOT NULL,
+  PRIMARY KEY ("id")
+);
+-- create index "notificationrule_id" to table: "notification_rules"
+CREATE INDEX "notificationrule_id" ON "notification_rules" ("id");
+-- create index "notificationrule_namespace_id" to table: "notification_rules"
+CREATE INDEX "notificationrule_namespace_id" ON "notification_rules" ("namespace", "id");
+-- create index "notificationrule_namespace_type" to table: "notification_rules"
+CREATE INDEX "notificationrule_namespace_type" ON "notification_rules" ("namespace", "type");
+-- create "notification_channel_rules" table
+CREATE TABLE "notification_channel_rules" (
+  "notification_channel_id" character(26) NOT NULL,
+  "notification_rule_id" character(26) NOT NULL,
+  PRIMARY KEY ("notification_channel_id", "notification_rule_id"),
+  CONSTRAINT "notification_channel_rules_notification_channel_id" FOREIGN KEY ("notification_channel_id") REFERENCES "notification_channels" ("id") ON UPDATE NO ACTION ON DELETE CASCADE,
+  CONSTRAINT "notification_channel_rules_notification_rule_id" FOREIGN KEY ("notification_rule_id") REFERENCES "notification_rules" ("id") ON UPDATE NO ACTION ON DELETE CASCADE
+);
+-- create "notification_event_delivery_status" table
+CREATE TABLE "notification_event_delivery_status" (
+  "id" character(26) NOT NULL,
+  "namespace" character varying NOT NULL,
+  "created_at" timestamptz NOT NULL,
+  "updated_at" timestamptz NOT NULL,
+  "event_id" character varying NOT NULL,
+  "channel_id" character varying NOT NULL,
+  "state" character varying NOT NULL DEFAULT 'PENDING',
+  "reason" character varying NULL,
+  PRIMARY KEY ("id")
+);
+-- create index "notificationeventdeliverystatus_id" to table: "notification_event_delivery_status"
+CREATE INDEX "notificationeventdeliverystatus_id" ON "notification_event_delivery_status" ("id");
+-- create index "notificationeventdeliverystatus_namespace_event_id_channel_id" to table: "notification_event_delivery_status"
+CREATE INDEX "notificationeventdeliverystatus_namespace_event_id_channel_id" ON "notification_event_delivery_status" ("namespace", "event_id", "channel_id");
+-- create index "notificationeventdeliverystatus_namespace_id" to table: "notification_event_delivery_status"
+CREATE INDEX "notificationeventdeliverystatus_namespace_id" ON "notification_event_delivery_status" ("namespace", "id");
+-- create index "notificationeventdeliverystatus_namespace_state" to table: "notification_event_delivery_status"
+CREATE INDEX "notificationeventdeliverystatus_namespace_state" ON "notification_event_delivery_status" ("namespace", "state");
+-- create "notification_events" table
+CREATE TABLE "notification_events" (
+  "id" character(26) NOT NULL,
+  "namespace" character varying NOT NULL,
+  "created_at" timestamptz NOT NULL,
+  "type" character varying NOT NULL,
+  "payload" jsonb NOT NULL,
+  "annotations" jsonb NULL,
+  "rule_id" character(26) NOT NULL,
+  PRIMARY KEY ("id"),
+  CONSTRAINT "notification_events_notification_rules_events" FOREIGN KEY ("rule_id") REFERENCES "notification_rules" ("id") ON UPDATE NO ACTION ON DELETE NO ACTION
+);
+-- create index "notificationevent_annotations" to table: "notification_events"
+CREATE INDEX "notificationevent_annotations" ON "notification_events" USING gin ("annotations");
+-- create index "notificationevent_id" to table: "notification_events"
+CREATE INDEX "notificationevent_id" ON "notification_events" ("id");
+-- create index "notificationevent_namespace_id" to table: "notification_events"
+CREATE INDEX "notificationevent_namespace_id" ON "notification_events" ("namespace", "id");
+-- create index "notificationevent_namespace_type" to table: "notification_events"
+CREATE INDEX "notificationevent_namespace_type" ON "notification_events" ("namespace", "type");
+-- create "notification_event_delivery_status_events" table
+CREATE TABLE "notification_event_delivery_status_events" (
+  "notification_event_delivery_status_id" character(26) NOT NULL,
+  "notification_event_id" character(26) NOT NULL,
+  PRIMARY KEY ("notification_event_delivery_status_id", "notification_event_id"),
+  CONSTRAINT "notification_event_delivery_status_events_notification_event_de" FOREIGN KEY ("notification_event_delivery_status_id") REFERENCES "notification_event_delivery_status" ("id") ON UPDATE NO ACTION ON DELETE CASCADE,
+  CONSTRAINT "notification_event_delivery_status_events_notification_event_id" FOREIGN KEY ("notification_event_id") REFERENCES "notification_events" ("id") ON UPDATE NO ACTION ON DELETE CASCADE
+);
 -- create "usage_resets" table
 CREATE TABLE "usage_resets" (
   "id" character(26) NOT NULL,
diff --git a/tools/migrate/migrations/atlas.sum b/tools/migrate/migrations/atlas.sum
index 84d8c8181..bc4494a7a 100644
--- a/tools/migrate/migrations/atlas.sum
+++ b/tools/migrate/migrations/atlas.sum
@@ -1,3 +1,3 @@
-h1:fW1f9OMqGTAwpbvQXgX+2jGJFPN8plpswoqfJhKs/xE=
-20240812143658_init.down.sql h1:hoeC6yyHw+V6ihO9WkCKElh/OIVtb4Kttaqy98DGhBU=
-20240812143658_init.up.sql h1:OsuTQd+TZQVfz0NNkJSyiLfUqmWEfwoUQdhIKmd06GY=
+h1:EPap+DaNoXsLQC61DQvaBbCLBcYaGP1DHSglnzexChY=
+20240826120919_init.down.sql h1:AIbgwwngjkJEYa3yRZsIXQyBa2+qoZttwMXHxXEbHLI=
+20240826120919_init.up.sql h1:/hYHWF3Z3dab8SMKnw99ixVktCuJe2bAw5wstCZIEN8=

From a1e8e91c0d7299da40214374dd974b84c37262ed Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Tue, 27 Aug 2024 08:53:48 +0200
Subject: [PATCH 09/13] fix: update and fix CI

fix: update CI

fix: fix pg url

fix: run go generate
---
 .github/workflows/ci.yaml                     | 16 ++++++++++++++
 ci/e2e.go                                     |  4 ++--
 ci/migrate.go                                 | 21 +++++--------------
 ci/versions_pinned.go                         |  2 ++
 openmeter/ent/db/runtime/runtime.go           |  4 ++--
 .../testutils/ent1/db/runtime/runtime.go      |  4 ++--
 .../testutils/ent2/db/runtime/runtime.go      |  4 ++--
 quickstart/config.yaml                        |  4 ++++
 quickstart/docker-compose.yaml                | 20 ++++++++++++++++++
 9 files changed, 55 insertions(+), 24 deletions(-)

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index daf8295ad..2a2a3deac 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -253,6 +253,22 @@ jobs:
           cloud-token: ${{ secrets.DAGGER_CLOUD_TOKEN }}
           version: ${{ env.DAGGER_VERSION }}
 
+      - name: Export Container Logs
+        id: export-container-logs
+        run: docker logs $(docker container list --all --filter 'name=^*-openmeter-*' --format '{{.Names}}') > container.stdout.log 2> container.stderr.log
+        if: always()
+        continue-on-error: true
+
+      - name: Upload Container logs as artifact
+        uses: actions/upload-artifact@834a144ee995460fba8ed112a2fc961b36a5ec5a # v4.3.6
+        if: always() && steps.export-container-logs.outcome == 'success'
+        with:
+          name: "[${{ github.job }}] Container logs"
+          path: |
+            container.stdout.log
+            container.stderr.log
+          retention-days: 14
+
       - name: Cleanup Docker Compose
         run: docker compose -f docker-compose.yaml -f docker-compose.ci.yaml down -v
         working-directory: quickstart
diff --git a/ci/e2e.go b/ci/e2e.go
index 3e06b463b..3a0c76e0b 100644
--- a/ci/e2e.go
+++ b/ci/e2e.go
@@ -29,7 +29,7 @@ func (m *Ci) Etoe(
 		WithExec([]string{"openmeter-sink-worker", "--config", "/etc/openmeter/config.yaml"}).
 		AsService()
 
-	args := []string{"go", "test", "-count=1", "-v"}
+	args := []string{"go", "test", "-tags", "musl", "-count=1", "-v"}
 
 	if test != "" {
 		args = append(args, "-run", fmt.Sprintf("Test%s", test))
@@ -37,7 +37,7 @@ func (m *Ci) Etoe(
 
 	args = append(args, "./e2e/...")
 
-	return goModule().
+	return goModuleCross("").
 		WithModuleCache(cacheVolume("go-mod-e2e")).
 		WithBuildCache(cacheVolume("go-build-e2e")).
 		WithSource(m.Source).
diff --git a/ci/migrate.go b/ci/migrate.go
index 52574d4d1..a1186200c 100644
--- a/ci/migrate.go
+++ b/ci/migrate.go
@@ -28,11 +28,11 @@ func (m *Migrate) Check(
 
 	bin := dag.Container(dagger.ContainerOpts{
 		Platform: "linux/amd64",
-	}).From("arigaio/atlas:0.26.0").File("atlas")
+	}).From(atlasImage).File("atlas")
 
 	atlas := app.
 		WithFile("/bin/atlas", bin).
-		WithDirectory("internal/ent", m.Source.Directory("internal/ent")).
+		WithDirectory("openmeter/ent", m.Source.Directory("openmeter/ent")).
 		WithDirectory("tools/migrate/migrations", m.Source.Directory("tools/migrate/migrations")).
 		WithFile("atlas.hcl", m.Source.File("atlas.hcl"))
 
@@ -41,10 +41,10 @@ func (m *Migrate) Check(
 	// Always validate schema is generated
 	p.Go(func(ctx context.Context) error {
 		result := app.
-			WithExec([]string{"go", "generate", "-x", "-tags=musl", "-ldflags", "linkmode=external", "./internal/ent/..."}).
-			Directory("internal/ent")
+			WithExec([]string{"go", "generate", "-x", "./openmeter/ent/..."}).
+			Directory("openmeter/ent")
 
-		source := m.Source.Directory("internal/ent")
+		source := m.Source.Directory("openmeter/ent")
 
 		err := diff(ctx, source, result)
 		if err != nil {
@@ -85,14 +85,3 @@ func (m *Migrate) Check(
 
 	return p.Wait()
 }
-
-func diff(ctx context.Context, d1, d2 *dagger.Directory) error {
-	_, err := dag.Container(dagger.ContainerOpts{Platform: ""}).
-		From(alpineBaseImage).
-		WithExec([]string{"apk", "add", "--update", "--no-cache", "ca-certificates", "tzdata", "bash"}).
-		WithDirectory("src", d1).
-		WithDirectory("res", d2).
-		WithExec([]string{"diff", "-u", "-r", "src", "res"}).
-		Sync(ctx)
-	return err
-}
diff --git a/ci/versions_pinned.go b/ci/versions_pinned.go
index 8181c5b13..1ba1166db 100644
--- a/ci/versions_pinned.go
+++ b/ci/versions_pinned.go
@@ -17,4 +17,6 @@ const (
 	xxBaseImage    = "tonistiigi/xx:1.5.0@sha256:0c6a569797744e45955f39d4f7538ac344bfb7ebf0a54006a0a4297b153ccf0f"
 
 	alpineBaseImage = "alpine:3.20.2@sha256:0a4eaa0eecf5f8c050e5bba433f58c052be7587ee8af3e8b3910ef9ab5fbe9f5"
+
+	atlasImage = "arigaio/atlas:0.26.1@sha256:f77152f5458255410d2e59ad80a0fc661524067a90d297a0b62e6bda23437893"
 )
diff --git a/openmeter/ent/db/runtime/runtime.go b/openmeter/ent/db/runtime/runtime.go
index 503c0d47d..35e8ceacc 100644
--- a/openmeter/ent/db/runtime/runtime.go
+++ b/openmeter/ent/db/runtime/runtime.go
@@ -5,6 +5,6 @@ package runtime
 // The schema-stitching logic is generated in github.com/openmeterio/openmeter/openmeter/ent/db/runtime.go
 
 const (
-	Version = "v0.14.0"                                         // Version of ent codegen.
-	Sum     = "h1:EO3Z9aZ5bXJatJeGqu/EVdnNr6K4mRq3rWe5owt0MC4=" // Sum of ent codegen.
+	Version = "v0.14.1"                                         // Version of ent codegen.
+	Sum     = "h1:fUERL506Pqr92EPHJqr8EYxbPioflJo6PudkrEA8a/s=" // Sum of ent codegen.
 )
diff --git a/pkg/framework/entutils/testutils/ent1/db/runtime/runtime.go b/pkg/framework/entutils/testutils/ent1/db/runtime/runtime.go
index 65b76a15d..2990e6525 100644
--- a/pkg/framework/entutils/testutils/ent1/db/runtime/runtime.go
+++ b/pkg/framework/entutils/testutils/ent1/db/runtime/runtime.go
@@ -5,6 +5,6 @@ package runtime
 // The schema-stitching logic is generated in github.com/openmeterio/openmeter/pkg/framework/entutils/testutils/ent1/db/runtime.go
 
 const (
-	Version = "v0.14.0"                                         // Version of ent codegen.
-	Sum     = "h1:EO3Z9aZ5bXJatJeGqu/EVdnNr6K4mRq3rWe5owt0MC4=" // Sum of ent codegen.
+	Version = "v0.14.1"                                         // Version of ent codegen.
+	Sum     = "h1:fUERL506Pqr92EPHJqr8EYxbPioflJo6PudkrEA8a/s=" // Sum of ent codegen.
 )
diff --git a/pkg/framework/entutils/testutils/ent2/db/runtime/runtime.go b/pkg/framework/entutils/testutils/ent2/db/runtime/runtime.go
index 9e52a7f36..002cf9406 100644
--- a/pkg/framework/entutils/testutils/ent2/db/runtime/runtime.go
+++ b/pkg/framework/entutils/testutils/ent2/db/runtime/runtime.go
@@ -5,6 +5,6 @@ package runtime
 // The schema-stitching logic is generated in github.com/openmeterio/openmeter/pkg/framework/entutils/testutils/ent2/db/runtime.go
 
 const (
-	Version = "v0.14.0"                                         // Version of ent codegen.
-	Sum     = "h1:EO3Z9aZ5bXJatJeGqu/EVdnNr6K4mRq3rWe5owt0MC4=" // Sum of ent codegen.
+	Version = "v0.14.1"                                         // Version of ent codegen.
+	Sum     = "h1:fUERL506Pqr92EPHJqr8EYxbPioflJo6PudkrEA8a/s=" // Sum of ent codegen.
 )
diff --git a/quickstart/config.yaml b/quickstart/config.yaml
index db80c949e..c8a4a6e3a 100644
--- a/quickstart/config.yaml
+++ b/quickstart/config.yaml
@@ -10,6 +10,10 @@ sink:
   minCommitCount: 1
   namespaceRefetch: 1s
 
+postgres:
+  url: postgres://postgres:postgres@postgres:5432/postgres?sslmode=disable
+  autoMigrate: migration # Runs migrations as part of the service startup, valid values are: ent, migration, false
+
 meters:
   # Sample meter to count API requests
   - slug: api_requests_total        # Unique identifier for the meter
diff --git a/quickstart/docker-compose.yaml b/quickstart/docker-compose.yaml
index d0530fec7..1d4e4c785 100644
--- a/quickstart/docker-compose.yaml
+++ b/quickstart/docker-compose.yaml
@@ -11,6 +11,8 @@ services:
         condition: service_healthy
       clickhouse:
         condition: service_healthy
+      postgres:
+        condition: service_healthy
     ports:
       - 127.0.0.1:8888:8888
     volumes:
@@ -82,3 +84,21 @@ services:
     extends:
       file: ../docker-compose.yaml
       service: redis
+
+  postgres:
+    image: postgres:14.9
+    ports:
+      - "5432:5432"
+    environment:
+      - POSTGRES_USER=postgres
+      - POSTGRES_DB=postgres
+      - POSTGRES_PASSWORD=postgres
+    command:
+      - "postgres"
+      - "-c"
+      - "wal_level=logical"
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready", "-d", "db_prod"]
+      interval: 5s
+      timeout: 60s
+      retries: 5

From a04068200d036dcf04ccc403cdefddaecb59ac14 Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Tue, 27 Aug 2024 14:42:34 +0200
Subject: [PATCH 10/13] chore: create separate db for svix

---
 docker-compose.yaml        | 4 +++-
 etc/postgres-init/svix.sql | 7 +++++++
 2 files changed, 10 insertions(+), 1 deletion(-)
 create mode 100644 etc/postgres-init/svix.sql

diff --git a/docker-compose.yaml b/docker-compose.yaml
index 9577576b8..e5e983842 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -95,6 +95,8 @@ services:
       - POSTGRES_USER=postgres
       - POSTGRES_DB=postgres
       - POSTGRES_PASSWORD=postgres
+    volumes:
+      - ./etc/postgres-init:/docker-entrypoint-initdb.d/
     command:
       - "postgres"
       - "-c"
@@ -107,7 +109,7 @@ services:
     environment:
       WAIT_FOR: "true"  # We want to wait for the default services
       SVIX_REDIS_DSN: "redis://redis:6379"
-      SVIX_DB_DSN: "postgresql://postgres:postgres@postgres/postgres"
+      SVIX_DB_DSN: "postgresql://svix:svix@postgres/svix"
       SVIX_CACHE_TYPE: "redis"
       SVIX_JWT_SECRET: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpYXQiOjE3MjI5NzYyNzMsImV4cCI6MjAzODMzNjI3MywibmJmIjoxNzIyOTc2MjczLCJpc3MiOiJzdml4LXNlcnZlciIsInN1YiI6Im9yZ18yM3JiOFlkR3FNVDBxSXpwZ0d3ZFhmSGlyTXUifQ.PomP6JWRI62W5N4GtNdJm2h635Q5F54eij0J3BU-_Ds"
       SVIX_LOG_LEVEL: "info"
diff --git a/etc/postgres-init/svix.sql b/etc/postgres-init/svix.sql
new file mode 100644
index 000000000..f0993185f
--- /dev/null
+++ b/etc/postgres-init/svix.sql
@@ -0,0 +1,7 @@
+CREATE USER svix WITH PASSWORD 'svix';
+
+CREATE DATABASE svix;
+
+GRANT ALL PRIVILEGES ON DATABASE svix TO svix;
+
+ALTER DATABASE svix OWNER TO svix;
\ No newline at end of file

From 21795c21c70e42f116302c51fba13ac90f13b040 Mon Sep 17 00:00:00 2001
From: Alex Goth <64845621+GAlexIHU@users.noreply.github.com>
Date: Wed, 28 Aug 2024 10:48:32 +0200
Subject: [PATCH 11/13] chore: address pr comments

Co-authored-by: Krisztian Gacsal <chrisgacsal@users.noreply.github.com>
Signed-off-by: Alex Goth <64845621+GAlexIHU@users.noreply.github.com>

chore: remove redundancies of rebasing

fix: db startup invocation
---
 cmd/balance-worker/main.go         |  2 +-
 cmd/notification-service/main.go   |  2 +-
 cmd/server/main.go                 |  2 +-
 openmeter/registry/startup/db.go   | 13 ++++++++-----
 pkg/framework/entutils/database.go | 27 ---------------------------
 5 files changed, 11 insertions(+), 35 deletions(-)
 delete mode 100644 pkg/framework/entutils/database.go

diff --git a/cmd/balance-worker/main.go b/cmd/balance-worker/main.go
index 699624e75..8261f6c4c 100644
--- a/cmd/balance-worker/main.go
+++ b/cmd/balance-worker/main.go
@@ -239,7 +239,7 @@ func main() {
 
 	entClient := entPostgresDriver.Client()
 
-	if err := startup.DB(conf.Postgres, entClient); err != nil {
+	if err := startup.DB(ctx, conf.Postgres, entClient); err != nil {
 		logger.Error("failed to initialize database", "error", err)
 		os.Exit(1)
 	}
diff --git a/cmd/notification-service/main.go b/cmd/notification-service/main.go
index 756edb7e7..bc1cffa20 100644
--- a/cmd/notification-service/main.go
+++ b/cmd/notification-service/main.go
@@ -242,7 +242,7 @@ func main() {
 	entClient := entPostgresDriver.Client()
 
 	// Run database schema creation
-	if err := startup.DB(conf.Postgres, entClient); err != nil {
+	if err := startup.DB(ctx, conf.Postgres, entClient); err != nil {
 		logger.Error("failed to initialize database", "error", err)
 		os.Exit(1)
 	}
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 5dc5722a1..0e9a5b1e1 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -350,7 +350,7 @@ func main() {
 
 	entClient := entPostgresDriver.Client()
 
-	if err := startup.DB(conf.Postgres, entClient); err != nil {
+	if err := startup.DB(ctx, conf.Postgres, entClient); err != nil {
 		logger.Error("failed to initialize database", "error", err)
 		os.Exit(1)
 	}
diff --git a/openmeter/registry/startup/db.go b/openmeter/registry/startup/db.go
index f5a929f68..5d529a8bd 100644
--- a/openmeter/registry/startup/db.go
+++ b/openmeter/registry/startup/db.go
@@ -9,18 +9,21 @@ import (
 	"github.com/openmeterio/openmeter/tools/migrate"
 )
 
-func DB(cfg config.PostgresConfig, db *db.Client) error {
-	if cfg.AutoMigrate.Enabled() {
-		switch cfg.AutoMigrate {
+func DB(ctx context.Context, cfg config.PostgresConfig, db *db.Client) error {
+	if !cfg.AutoMigrate.Enabled() {
+		return nil
+	}
+	
+	switch cfg.AutoMigrate {
 		case config.AutoMigrateEnt:
-			if err := db.Schema.Create(context.Background()); err != nil {
+			if err := db.Schema.Create(ctx); err != nil {
 				return fmt.Errorf("failed to migrate db: %w", err)
 			}
 		case config.AutoMigrateMigration:
 			if err := migrate.Up(cfg.URL); err != nil {
 				return fmt.Errorf("failed to migrate db: %w", err)
 			}
-		}
 	}
+		
 	return nil
 }
diff --git a/pkg/framework/entutils/database.go b/pkg/framework/entutils/database.go
deleted file mode 100644
index e3688d56b..000000000
--- a/pkg/framework/entutils/database.go
+++ /dev/null
@@ -1,27 +0,0 @@
-package entutils
-
-import (
-	"database/sql"
-
-	"entgo.io/ent/dialect"
-	entDialectSQL "entgo.io/ent/dialect/sql"
-	"github.com/XSAM/otelsql"
-	semconv "go.opentelemetry.io/otel/semconv/v1.20.0"
-)
-
-// Deprecated: use NewEntPostgresDriver instead
-func GetSQLDriver(databaseURL string) (*sql.DB, error) {
-	return otelsql.Open("pgx", databaseURL, otelsql.WithAttributes(
-		semconv.DBSystemPostgreSQL,
-	))
-}
-
-func GetEntDriver(databaseURL string) (*entDialectSQL.Driver, error) {
-	// TODO: inject trace and metrics provider
-	database, err := GetSQLDriver(databaseURL)
-	if err != nil {
-		return nil, err
-	}
-
-	return entDialectSQL.OpenDB(dialect.Postgres, database), nil
-}

From db58146cd4feb47cad12c31ca3889d68978cb70b Mon Sep 17 00:00:00 2001
From: Alex Goth <64845621+GAlexIHU@users.noreply.github.com>
Date: Wed, 28 Aug 2024 10:59:43 +0200
Subject: [PATCH 12/13] chore: update ci/migrate.go

Co-authored-by: Krisztian Gacsal <chrisgacsal@users.noreply.github.com>
Signed-off-by: Alex Goth <64845621+GAlexIHU@users.noreply.github.com>
---
 ci/migrate.go | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/ci/migrate.go b/ci/migrate.go
index a1186200c..2750a0edc 100644
--- a/ci/migrate.go
+++ b/ci/migrate.go
@@ -18,9 +18,7 @@ type Migrate struct {
 	Source *dagger.Directory
 }
 
-func (m *Migrate) Check(
-	ctx context.Context,
-) error {
+func (m *Migrate) Check(ctx context.Context) error {
 	app := goModuleCross("").
 		WithSource(m.Source).
 		Container().

From 46e603d445121bd97f4611a75c68128752d26a84 Mon Sep 17 00:00:00 2001
From: Alex Goth <alex.goth@openmeter.io>
Date: Wed, 28 Aug 2024 11:13:13 +0200
Subject: [PATCH 13/13] chore: use golangci-lint for vscode

---
 .vscode/settings.json            |  4 +++-
 openmeter/registry/startup/db.go | 20 ++++++++++----------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 2730ac895..20c7dea09 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -9,5 +9,7 @@
   },
   "gopls": {
     "formatting.gofumpt": true
-  }
+  },
+  "go.lintTool": "golangci-lint",
+  "go.lintFlags": ["--fast", "--fix", "-c", ".golangci.yaml"]
 }
diff --git a/openmeter/registry/startup/db.go b/openmeter/registry/startup/db.go
index 5d529a8bd..0d800333e 100644
--- a/openmeter/registry/startup/db.go
+++ b/openmeter/registry/startup/db.go
@@ -13,17 +13,17 @@ func DB(ctx context.Context, cfg config.PostgresConfig, db *db.Client) error {
 	if !cfg.AutoMigrate.Enabled() {
 		return nil
 	}
-	
+
 	switch cfg.AutoMigrate {
-		case config.AutoMigrateEnt:
-			if err := db.Schema.Create(ctx); err != nil {
-				return fmt.Errorf("failed to migrate db: %w", err)
-			}
-		case config.AutoMigrateMigration:
-			if err := migrate.Up(cfg.URL); err != nil {
-				return fmt.Errorf("failed to migrate db: %w", err)
-			}
+	case config.AutoMigrateEnt:
+		if err := db.Schema.Create(ctx); err != nil {
+			return fmt.Errorf("failed to migrate db: %w", err)
+		}
+	case config.AutoMigrateMigration:
+		if err := migrate.Up(cfg.URL); err != nil {
+			return fmt.Errorf("failed to migrate db: %w", err)
+		}
 	}
-		
+
 	return nil
 }