coder · sreya · Aug 17, 2022 · Aug 9, 2022 · Aug 9, 2022 · Aug 9, 2022
diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -73,6 +73,7 @@ type data struct {
 	organizations       []database.Organization
 	organizationMembers []database.OrganizationMember
 	users               []database.User
+	userLinks           []database.UserLink
 
 	// New tables
 	auditLogs                      []database.AuditLog
@@ -1453,20 +1454,16 @@ func (q *fakeQuerier) InsertAPIKey(_ context.Context, arg database.InsertAPIKeyP
 
 	//nolint:gosimple
 	key := database.APIKey{
-		ID:                arg.ID,
-		LifetimeSeconds:   arg.LifetimeSeconds,
-		HashedSecret:      arg.HashedSecret,
-		IPAddress:         arg.IPAddress,
-		UserID:            arg.UserID,
-		ExpiresAt:         arg.ExpiresAt,
-		CreatedAt:         arg.CreatedAt,
-		UpdatedAt:         arg.UpdatedAt,
-		LastUsed:          arg.LastUsed,
-		LoginType:         arg.LoginType,
-		OAuthAccessToken:  arg.OAuthAccessToken,
-		OAuthRefreshToken: arg.OAuthRefreshToken,
-		OAuthIDToken:      arg.OAuthIDToken,
-		OAuthExpiry:       arg.OAuthExpiry,
+		ID:              arg.ID,
+		LifetimeSeconds: arg.LifetimeSeconds,
+		HashedSecret:    arg.HashedSecret,
+		IPAddress:       arg.IPAddress,
+		UserID:          arg.UserID,
+		ExpiresAt:       arg.ExpiresAt,
+		CreatedAt:       arg.CreatedAt,
+		UpdatedAt:       arg.UpdatedAt,
+		LastUsed:        arg.LastUsed,
+		LoginType:       arg.LoginType,
 	}
 	q.apiKeys = append(q.apiKeys, key)
 	return key, nil
@@ -1898,9 +1895,6 @@ func (q *fakeQuerier) UpdateAPIKeyByID(_ context.Context, arg database.UpdateAPI
 		apiKey.LastUsed = arg.LastUsed
 		apiKey.ExpiresAt = arg.ExpiresAt
 		apiKey.IPAddress = arg.IPAddress
-		apiKey.OAuthAccessToken = arg.OAuthAccessToken
-		apiKey.OAuthRefreshToken = arg.OAuthRefreshToken
-		apiKey.OAuthExpiry = arg.OAuthExpiry
 		q.apiKeys[index] = apiKey
 		return nil
 	}
@@ -2259,3 +2253,80 @@ func (q *fakeQuerier) GetDeploymentID(_ context.Context) (string, error) {
 
 	return q.deploymentID, nil
 }
+
+func (q *fakeQuerier) GetUserLinkByLinkedID(_ context.Context, id string) (database.UserLink, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	for _, link := range q.userLinks {
+		if link.LinkedID == id {
+			return link, nil
+		}
+	}
+	return database.UserLink{}, sql.ErrNoRows
+}
+
+func (q *fakeQuerier) GetUserLinkByUserIDLoginType(_ context.Context, params database.GetUserLinkByUserIDLoginTypeParams) (database.UserLink, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	for _, link := range q.userLinks {
+		if link.UserID == params.UserID && link.LoginType == params.LoginType {
+			return link, nil
+		}
+	}
+	return database.UserLink{}, sql.ErrNoRows
+}
+
+func (q *fakeQuerier) InsertUserLink(_ context.Context, args database.InsertUserLinkParams) (database.UserLink, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	//nolint:gosimple
+	link := database.UserLink{
+		UserID:            args.UserID,
+		LoginType:         args.LoginType,
+		LinkedID:          args.LinkedID,
+		OAuthAccessToken:  args.OAuthAccessToken,
+		OAuthRefreshToken: args.OAuthRefreshToken,
+		OAuthExpiry:       args.OAuthExpiry,
+	}
+
+	q.userLinks = append(q.userLinks, link)
+
+	return link, nil
+}
+
+func (q *fakeQuerier) UpdateUserLinkedID(_ context.Context, params database.UpdateUserLinkedIDParams) (database.UserLink, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	for i, link := range q.userLinks {
+		if link.UserID == params.UserID && link.LoginType == params.LoginType {
+			link.LinkedID = params.LinkedID
+
+			q.userLinks[i] = link
+			return link, nil
+		}
+	}
+
+	return database.UserLink{}, sql.ErrNoRows
+}
+
+func (q *fakeQuerier) UpdateUserLink(_ context.Context, params database.UpdateUserLinkParams) (database.UserLink, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	for i, link := range q.userLinks {
+		if link.UserID == params.UserID && link.LoginType == params.LoginType {
+			link.OAuthAccessToken = params.OAuthAccessToken
+			link.OAuthRefreshToken = params.OAuthRefreshToken
+			link.OAuthExpiry = params.OAuthExpiry
+
+			q.userLinks[i] = link
+			return link, nil
+		}
+	}
+
+	return database.UserLink{}, sql.ErrNoRows
+}
diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/generate.sh b/coderd/database/generate.sh
@@ -13,6 +13,8 @@ SCRIPT_DIR=$(dirname "${BASH_SOURCE[0]}")
 (
 	cd "$SCRIPT_DIR"
 
+	# Dump the updated schema.
+	go run dump/main.go
 	# The logic below depends on the exact version being correct :(
 	go run github.com/kyleconroy/sqlc/cmd/[email protected] generate
 

diff --git a/coderd/database/migrations/000034_linked_user_id.down.sql b/coderd/database/migrations/000034_linked_user_id.down.sql
@@ -0,0 +1,22 @@
+-- This migration makes no attempt to try to populate
+-- the oauth_access_token, oauth_refresh_token, and oauth_expiry
+-- columns of api_key rows with the values from the dropped user_links
+-- table.
+BEGIN;
+
+DROP TABLE IF EXISTS user_links;
+
+ALTER TABLE
+  api_keys
+ADD COLUMN oauth_access_token text DEFAULT ''::text NOT NULL;
+
+ALTER TABLE
+  api_keys
+ADD COLUMN oauth_refresh_token text DEFAULT ''::text NOT NULL;
+
+ALTER TABLE
+  api_keys
+ADD COLUMN oauth_expiry timestamp with time zone DEFAULT '0001-01-01 00:00:00+00'::timestamp with time zone NOT NULL;
+
+
+COMMIT;
diff --git a/coderd/database/migrations/000034_linked_user_id.up.sql b/coderd/database/migrations/000034_linked_user_id.up.sql
@@ -0,0 +1,53 @@
+BEGIN;
+
+CREATE TABLE IF NOT EXISTS user_links (
+	user_id uuid NOT NULL,
+	login_type login_type NOT NULL,
+	linked_id text DEFAULT ''::text NOT NULL,
+	oauth_access_token text DEFAULT ''::text NOT NULL,
+	oauth_refresh_token text DEFAULT ''::text NOT NULL,
+	oauth_expiry timestamp with time zone DEFAULT '0001-01-01 00:00:00+00'::timestamp with time zone NOT NULL,
+	UNIQUE(user_id, login_type)
+);
+
+-- This migrates columns on api_keys to the new user_links table.
+-- It does this by finding all the API keys for each user, choosing
+-- the most recently updated for each one and then assigning its relevant
+-- values to the user_links table.
+-- A user should at most have a row for an OIDC account and a Github account.
+-- 'password' login types are ignored.
+
+INSERT INTO user_links 
+	( 
+		user_id,
+		login_type,
+		linked_id,
+		oauth_access_token,
+		oauth_refresh_token,
+		oauth_expiry
+	)
+SELECT 
+	keys.user_id, 
+	keys.login_type,
+	'',
+	keys.oauth_access_token,
+	keys.oauth_refresh_token,
+	keys.oauth_expiry 
+FROM 
+	( 
+		SELECT 
+			row_number() OVER (partition by user_id, login_type ORDER BY updated_at DESC) AS x, 
+			api_keys.* FROM api_keys
+	) as keys
+ WHERE x=1 AND keys.login_type != 'password';
+
+-- Drop columns that have been migrated to user_links.
+-- It appears the 'oauth_id_token' was unused and so it has
+-- been dropped here as well to avoid future confusion.
+ALTER TABLE api_keys 
+	DROP COLUMN oauth_access_token,
+	DROP COLUMN oauth_refresh_token,
+	DROP COLUMN oauth_id_token,
+	DROP COLUMN oauth_expiry;
+
+COMMIT;
diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go