Replace gorm.DB with an interface #3008

dnephin · 2022-08-22T22:33:34Z

Summary

Branched from #3006, that will need to merge first.

This is a rather large PR because it was hard to split this change into smaller pieces. Introducing the interface should be the largest part of the change. After this everything should be easily broken up into much smaller changes.

This PR does the following:

introduces a GormTxn interface that is accepted by all the data query functions
adds methods to that interface to pass the Organization ID
updates the middleware to set the orgID on this transaction, and all functions to read from this value.

This approach makes our code much safer, because we are no longer defaulting to the default orgID. I'll call out some places where I had to explicitly do this in tests to prevent breaking many tests. I'll go through in a follow up to fix these places.

There are many commits, and a lot of them are very mechanism find/replace. I'll try to call out the most interesting bits with inline comments.

Related to #2415

dnephin · 2022-08-22T22:35:05Z

internal/access/access_key.go

+	s := []data.SelectorFunc{
+		data.Preload("IssuedForIdentity"),
+		data.ByOptionalIssuedFor(identityID),
+		data.ByOptionalName(name),
+	}
 	if !showExpired {
 		s = append(s, data.ByNotExpiredOrExtended())
 	}

-	return data.ListAccessKeys(db.Preload("IssuedForIdentity"), p, s...)
+	return data.ListAccessKeys(db, p, s...)


In a few places I had to convert Preload to a selector so that I could accept a data.GormTxn as the first argument to the query function, and only selectors would continue to use gorm.DB. This is one of the places, there should be a few others.

dnephin · 2022-08-22T22:37:30Z

internal/access/signup.go

+	db = data.NewTransaction(db.GormDB(), details.Org.ID)
+	c.Set("db", db)
+	rCtx := GetRequestContext(c)
+	rCtx.DBTxn = db
+	c.Set(RequestContextKey, rCtx)


Now that we're passing orgID as part of the transaction the data.CreateOrganization function can't set the org ID in the context. I think this is a good thing, as previously it was surprising that creating an org mutated the DB to make all queries use that org.

As a consequence of that change we have to set the tx here for the rest of the request, but this is the only place where we need to do this.

I've also restored the original name for CreateOrganization, now that it's not doing the SetContext part.

dnephin · 2022-08-22T22:42:26Z

internal/server/data/data.go

+	// FIXME: this is a hack to keep our tests passing. The db should not
+	// be scoped to an org ID.
+	return d.DefaultOrg.ID


Calling this out. This is not a regression, because it works this way on main, but it is a problem we should fix.

All of our test fixtures should be more explicit about which org they are using when creating things.

I don't understand this one. How does this do the right thing in production?

Exactly. This is why I prioritized this part of the data package work, because this is the behaviour we have now by setting the context when we create a database.

It "works" in production because requests use the transaction reference where we override the organization value.

The only code that doesn't do that right now is:

the config loading, where you fixed it to set this same default value

tests in the server package where we do the same thing by setting the default org.

I plan on removing this in the first follow up PR.

dnephin · 2022-08-22T22:44:41Z

internal/server/data/data_test.go

@@ -76,66 +75,8 @@ func TestSnowflakeIDSerialization(t *testing.T) {
 	})
 }

-func TestDatabaseSelectors(t *testing.T) {


Removed because this test is just a simulation of the selector behaviour, and we should be removing selector functions soon. Porting this test to the new interface would be a lot more work than other tests, and is not worth it.

dnephin · 2022-08-22T22:46:27Z

internal/server/data/identity.go

+	selectID := func(db *gorm.DB) *gorm.DB {
+		return db.Select("id")
+	}
+	selectors = append([]SelectorFunc{selectID}, selectors...)
+	toDelete, err := ListIdentities(tx, nil, selectors...)


Another place where I had to move a thing to the selectors. I think this is the only case that was not a Preload

dnephin · 2022-08-22T22:47:07Z

internal/server/data/migrations_test.go

 		assert.NilError(t, m.Migrate())

 		initialSchema = dumpSchema(t, os.Getenv("POSTGRESQL_CONNECTION"))

-		assert.NilError(t, db.Exec("DROP SCHEMA IF EXISTS testing CASCADE").Error)
+		_, err = db.Exec("DROP SCHEMA IF EXISTS testing CASCADE")


I converted the few remaining Exec calls to the new interface, which matches the stdlib. So you'll notice the .Error removed, and that it now returns two args.

dnephin · 2022-08-22T22:48:45Z

internal/server/data/organization.go

-	return org
-}
-
-func MustGetOrgFromContext(ctx context.Context) *models.Organization {


Now that we have a strongly typed interface we don't need to store the org in the context. This wasn't necessary to do all at once, but it was a very small change at the end (see the last commit).

We no longer panic, but an unset orgID will also be 0, which will never match a real org (not even the default), which will still error and should be caught by a test.

dnephin · 2022-08-22T22:50:09Z

internal/server/data/organization_test.go

-	pgsql := postgresDriver(t)
-	db := setupDB(t, pgsql)
+func TestCreateOrganization(t *testing.T) {
+	runDBTests(t, func(t *testing.T, db *DB) {


Updated this test to use runDBTests like all the others. We only really need to test against postgres right now, but I expect runDBTests to be useful in the future despite that.

One use case that comes to mind is testing postgres upgrades. We can run all of our data tests against multiple versions of postgres to catch problems.

dnephin · 2022-08-22T22:52:24Z

internal/server/middleware.go

-		withDBTxn(c.Request.Context(), srv.DB(), func(tx *gorm.DB) {
+		withDBTxn(c.Request.Context(), srv.DB().GormDB(), func(db *gorm.DB) {
+
+			tx := data.NewTransaction(db, 0)


As it was before, we now set the orgID in the transaction as part of these middleware. But until we have an org identified we query with orgID =0, which should ensure we never match an org.

BruceMacD · 2022-08-23T14:13:09Z

internal/access/signup_test.go

@@ -42,6 +43,10 @@ func TestSignup(t *testing.T) {
 		assert.Equal(t, identity.Name, user)
 		assert.Equal(t, identity.OrganizationID, org.ID)

+		// simulate a request


BruceMacD · 2022-08-23T14:19:28Z

internal/server/data/access_key.go

 	return list[models.AccessKey](db, p, selectors...)
 }

-func GetAccessKey(db *gorm.DB, selectors ...SelectorFunc) (*models.AccessKey, error) {
+func GetAccessKey(tx GormTxn, selectors ...SelectorFunc) (*models.AccessKey, error) {
+	db := tx.GormDB()


Why do some queries need tx.GormDB?

Anything that builds a query itself using gorm methods will need to call tx.GormDB. Mostly that is our helpers like add, save, get, but also a few others that do more interesting queries.

This will give us an easy way of finding all the queries that need to be replaced with sql.

I found this a little confusing too when reading through the PR. Maybe a note around GormDB() which explains when to use it?

We should be removing GormDB very soon, but I'll add a comment for the interim.

BruceMacD · 2022-08-23T14:20:20Z

internal/server/data/access_key.go

-		}
-		if err := SaveAccessKey(db, t); err != nil {
+
+		tx = NewTransaction(tx.GormDB(), t.OrganizationID)


Is this a nested transaction?

The naming is maybe a bit misleading. NewTransaction doesn't create a new DB transaction, it uses an existing reference to an sql.Tx, but sets a new orgID for the transaction.

This also took me a moment to realize what was going on, and how Exec() works inside the transaction.

I just noticed this is done in the wrong place. I'll add a comment and move it to outside of this block.

This NewTransaction wouldn't be necessary if we didn't have the panic thta was requested in these comments. It's kind of unfortunate that we're forcing unnecessary code because of safeguards that aren't correct in all cases.

I was able to remove this by moving where we panic:

https://github.com/infrahq/infra/compare/dnephin/data-interface-take2...dnephin/data-interface-org-id-panic?expand=1

I'll open that as a follow up

internal/server/data/data.go

BruceMacD · 2022-08-23T14:59:15Z

internal/server/middleware.go

@@ -157,6 +158,7 @@ func unauthenticatedMiddleware(srv *Server) gin.HandlerFunc {
 				sendAPIError(c, internal.ErrBadRequest)
 				return
 			}
+			authned.Organization = org


This should be set directly on the authned access key already, re-assigning feels dangerous. Or is this just loading the org for the access key's org ID?

This is the unauthenticated route, so it may not be set yet.

I wasn't sure where to store it. Maybe we should split up the Authenticated struct into something like this:

type RequestDetails struct { Authenticated Authenticated Organization *models.Organization }

That seems like a good idea

internal/server/passwordreset.go

ssoroka

Would like to see the return of MustGetOrg type semantics, where things can't quietly fail when missing. Even in cases where the calling code panics anyway, the panic message is more clear when explicit.

internal/server/data/data.go

ssoroka · 2022-08-23T15:46:40Z

internal/server/data/data.go

+	// FIXME: this is a hack to keep our tests passing. The db should not
+	// be scoped to an org ID.
+	return d.DefaultOrg.ID


I don't understand this one. How does this do the right thing in production?

internal/server/middleware.go

internal/server/data/organization.go

internal/access/organization.go

pdevine

Overall looks fine. We still need to somehow split this for privileged vs unprivileged requests, and I'm not 100% sure how to carve that up. We could hold two references inside of the transaction, although it would be preferable only to store the unprivileged connection just in case we mess something up and call something in the data layer which references the wrong connection.

Right now we need the privileged connection to:

Use the telemetry/metrics
Load the config (which needs to look up the default org, but we could drop this if we didn't have a default org at all)
DeleteIdentities() and DeleteGrants() also needs the work somehow in the config which gets rid of any users/grants which were set up previously. This gets really weird if the org changes.
Encryption keys (since they're not tied to an org -- the unprivileged connection can still look these up right now which is maybe not what we want)
Migrations
Sign-up (although this could also work if we don't turn on RLS for the organizations table)

pdevine · 2022-08-23T14:16:23Z

internal/server/authn/authn_method_test.go

@@ -15,7 +14,7 @@ import (
 	"github.com/infrahq/infra/internal/testing/patch"
 )

-func setupDB(t *testing.T) *gorm.DB {
+func setupDB(t *testing.T) *data.DB {
 	t.Helper()
 	driver := database.PostgresDriver(t, "_authn")
 	if driver == nil {


Should we be pulling out the SQLite driver now?

No reason to yet. That should be a separate change either way.

internal/server/data/data.go

pdevine · 2022-08-23T16:35:14Z

internal/server/data/access_key.go

 	return list[models.AccessKey](db, p, selectors...)
 }

-func GetAccessKey(db *gorm.DB, selectors ...SelectorFunc) (*models.AccessKey, error) {
+func GetAccessKey(tx GormTxn, selectors ...SelectorFunc) (*models.AccessKey, error) {
+	db := tx.GormDB()


I found this a little confusing too when reading through the PR. Maybe a note around GormDB() which explains when to use it?

pdevine · 2022-08-23T16:36:45Z

internal/server/data/access_key.go

-		}
-		if err := SaveAccessKey(db, t); err != nil {
+
+		tx = NewTransaction(tx.GormDB(), t.OrganizationID)


This also took me a moment to realize what was going on, and how Exec() works inside the transaction.

dnephin · 2022-08-23T17:31:25Z

We still need to somehow split this for privileged vs unprivileged requests, and I'm not 100% sure how to carve that up.

The data.DB struct should return two connection pools, one privileged, and one unprivileged.

When we choose which one to use when we create the transaction.

I don't think the changes were should impact that much. The biggest change might be in tests.

dnephin · 2022-08-23T17:41:18Z

I've restored the panic and added more godoc.

I don't like the panic, I think it's assuming too much about the caller, which maybe true for regular requests, but is not true for middleware and other callers of the data package. But I'm fine to add it for now , and we can revisit it in a smaller more focused PR.

There are some places where we have to create a transaction just to work around the panic. The data.Transaction wouldn't be necessary otherwise, because the org ID is already set on the object.

I've also added some godoc in a few places as requested.

Those changes are in the latest 2 commits. This should be ready for another review.

BruceMacD

This looks good to me pending any further comments from Steven and Patrick

We will be removing these selectors shortly.

So that all queries fit the same pattern. These special cases did not work with the new GormTxn interface approach.

with the OrgID

We now use the DB transaction to pass the org ID. Also renames CreateOrganization back now that the function doesn't automatically set the OrgID in the context.

dnephin requested review from ssoroka, BruceMacD, jmorganca and pdevine as code owners August 22, 2022 22:33

dnephin changed the title ~~Dnephin/data interface take2~~ Replace gorm.DB with an interface Aug 22, 2022

dnephin commented Aug 22, 2022

View reviewed changes

dnephin force-pushed the dnephin/data-interface-take2 branch from ccc3f42 to c898a37 Compare August 22, 2022 22:54

BruceMacD reviewed Aug 23, 2022

View reviewed changes

ssoroka reviewed Aug 23, 2022

View reviewed changes

pdevine reviewed Aug 23, 2022

View reviewed changes

dnephin force-pushed the dnephin/data-interface-take2 branch from 9c59375 to ee783cd Compare August 23, 2022 17:39

Base automatically changed from dnephin/data-remove-gorm-from-migrator to main August 23, 2022 21:23

BruceMacD approved these changes Aug 24, 2022

View reviewed changes

dnephin mentioned this pull request Aug 25, 2022

Remove db from gin context #3022

Merged

ssoroka approved these changes Aug 25, 2022

View reviewed changes

dnephin added 13 commits August 25, 2022 14:59

maintain: remove selector test

4704c12

We will be removing these selectors shortly.

improve: introduce the transaction struct

dc178f2

improve: change all write funcs to use the interface

87bb36e

improve: use runDBTests in org create test

a8eb2bf

improve: convert runDBTests signature

e041561

improve: convert remaining queries to the interface

314b6d1

maintain: create missing selectors

2aa5c97

So that all queries fit the same pattern. These special cases did not work with the new GormTxn interface approach.

maintain: temporary fix for org ID in tests

0e8f440

improve: update other packages for the new data interface

66dded0

fix: restore setting of org id after creating an org

be80209

fix: the server package for data interface changes

dbea8ee

improve: set transaction in middleware

b26c0e2

with the OrgID

maintain: remove organization from context

68fb79f

We now use the DB transaction to pass the org ID. Also renames CreateOrganization back now that the function doesn't automatically set the OrgID in the context.

dnephin added 2 commits August 25, 2022 15:02

improve: panic if org is not set

c5a355d

maintain: add more godoc around transactions

8245c6b

dnephin force-pushed the dnephin/data-interface-take2 branch from ee783cd to 8245c6b Compare August 25, 2022 19:02

dnephin mentioned this pull request Aug 25, 2022

Move where we panic about a missing org ID to reduce the calls to data.NewTransaction #3027

Merged

dnephin merged commit 6c3d610 into main Aug 25, 2022

dnephin deleted the dnephin/data-interface-take2 branch August 25, 2022 21:33

Replace gorm.DB with an interface #3008

Replace gorm.DB with an interface #3008

Conversation

dnephin commented Aug 22, 2022

Summary

Choose a reason for hiding this comment

dnephin Aug 22, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dnephin Aug 23, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ssoroka left a comment • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pdevine left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dnephin commented Aug 23, 2022

dnephin commented Aug 23, 2022 • edited

BruceMacD left a comment

Choose a reason for hiding this comment

dnephin Aug 22, 2022 •

edited

dnephin Aug 23, 2022 •

edited

ssoroka left a comment •

edited

dnephin commented Aug 23, 2022 •

edited