Add external MCP servers as mcp connectors with per-user OAuth

admin/database/database.go

@@ -345,6 +345,26 @@ type DB interface {
 	DeleteBillingIssue(ctx context.Context, id string) error
 	DeleteBillingIssueByTypeForOrg(ctx context.Context, orgID string, errorType BillingIssueType) error
 
+	// FindConnectorOAuthClient returns the client for a connector's OAuth configuration, identified by its hash.
+	FindConnectorOAuthClient(ctx context.Context, projectID, connector, configHash string) (*ConnectorOAuthClient, error)
+	FindConnectorOAuthClientByID(ctx context.Context, id string) (*ConnectorOAuthClient, error)
+	// FindConnectorOAuthClients returns the clients for all OAuth configurations a connector has had, such as before a scope change or on other branches.
+	FindConnectorOAuthClients(ctx context.Context, projectID, connector string) ([]*ConnectorOAuthClient, error)
+	// FindOrInsertConnectorOAuthClient inserts a client, unless a client for the same connector and configuration hash already exists
+	// (for example, because another request inserted one concurrently), in which case it returns the existing client instead.
+	FindOrInsertConnectorOAuthClient(ctx context.Context, opts *InsertConnectorOAuthClientOptions) (*ConnectorOAuthClient, error)
+	// DeleteConnectorOAuthClient deletes a client and the tokens obtained through it.
+	DeleteConnectorOAuthClient(ctx context.Context, id string) error
+	// DeleteConnectorOAuthClientsWithoutTokens deletes a connector's clients that hold no tokens, except the client for keepConfigHash.
+	DeleteConnectorOAuthClientsWithoutTokens(ctx context.Context, projectID, connector, keepConfigHash string) error
+	// AcquireConnectorOAuthTokenLock takes an exclusive lock on a user's token for a client until the current transaction ends.
+	// It must be called within a transaction started with NewTx. Reads and writes of the token (refresh, disconnect, a completed authorization)
+	// take it first, so they are serialized across all admin processes.
+	AcquireConnectorOAuthTokenLock(ctx context.Context, oauthClientID, userID string) error
+	FindConnectorOAuthToken(ctx context.Context, oauthClientID, userID string) (*ConnectorOAuthToken, error)
+	UpsertConnectorOAuthToken(ctx context.Context, opts *UpsertConnectorOAuthTokenOptions) (*ConnectorOAuthToken, error)
+	DeleteConnectorOAuthToken(ctx context.Context, oauthClientID, userID string) error
+
 	FindProjectVariables(ctx context.Context, projectID string, environment *string) ([]*ProjectVariable, error)
 	UpsertProjectVariable(ctx context.Context, projectID, environment string, vars map[string]string, userID string) ([]*ProjectVariable, error)
 	DeleteProjectVariables(ctx context.Context, projectID, environment string, vars []string) error
@@ -1271,6 +1291,48 @@ type Asset struct {
 	CreatedOn      time.Time `db:"created_on"`
 }
 
+// ConnectorOAuthClient is an OAuth client Rill uses with a connector's authorization server, shared by the users of a project's connector.
+// ConfigHash identifies the connector's OAuth configuration it was created for; a configuration change leads to a new client.
+// Clients are never modified, because tokens and pending authorizations are bound to them.
+// Client is the JSON encoding of runtime/pkg/oauthclient.Client; it is decrypted when read.
+type ConnectorOAuthClient struct {
+	ID                    string    `db:"id"`
+	ProjectID             string    `db:"project_id"`
+	Connector             string    `db:"connector"`
+	ConfigHash            string    `db:"config_hash"`
+	Client                []byte    `db:"client"`
+	ClientEncryptionKeyID string    `db:"client_encryption_key_id"`
+	CreatedOn             time.Time `db:"created_on"`
+}
+
+// InsertConnectorOAuthClientOptions defines options for creating a ConnectorOAuthClient.
+type InsertConnectorOAuthClientOptions struct {
+	ProjectID  string `validate:"required"`
+	Connector  string `validate:"required"`
+	ConfigHash string `validate:"required"`
+	Client     []byte `validate:"required"`
+}
+
+// ConnectorOAuthToken is a user's OAuth token obtained through a ConnectorOAuthClient. Tokens are never shared between users.
+// The token is only used and refreshed through the client that obtained it.
+// Token is the JSON encoding of golang.org/x/oauth2.Token; it is decrypted when read.
+type ConnectorOAuthToken struct {
+	ID                   string    `db:"id"`
+	OAuthClientID        string    `db:"oauth_client_id"`
+	UserID               string    `db:"user_id"`
+	Token                []byte    `db:"token"`
+	TokenEncryptionKeyID string    `db:"token_encryption_key_id"`
+	CreatedOn            time.Time `db:"created_on"`
+	UpdatedOn            time.Time `db:"updated_on"`
+}
+
+// UpsertConnectorOAuthTokenOptions defines options for creating or replacing a ConnectorOAuthToken.
+type UpsertConnectorOAuthTokenOptions struct {
+	OAuthClientID string `validate:"required"`
+	UserID        string `validate:"required"`
+	Token         []byte `validate:"required"`
+}
+
 // ProjectVariable represents a key-value variable for a project, possible for a specific environment (e.g. production or development).
 type ProjectVariable struct {
 	ID                   string    `db:"id"`

admin/database/postgres/migrations/0098.sql

@@ -0,0 +1,35 @@
+-- OAuth state for connectors that authorize users against an external OAuth server (e.g. driver "mcp" with auth "oauth").
+
+-- An OAuth client Rill uses with a connector's authorization server: discovered endpoints plus the registered (or pre-registered) client,
+-- shared by the users of a project's connector.
+-- A client is bound to a hash of the connector's OAuth configuration (as computed by the runtime),
+-- so changing the configuration (URL, pre-registered client, scopes) leads to a new client and never reuses authorizations granted for the old one.
+-- Clients are immutable.
+CREATE TABLE connector_oauth_clients (
+    id UUID NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
+    project_id UUID NOT NULL REFERENCES projects (id) ON DELETE CASCADE,
+    connector TEXT NOT NULL,
+    config_hash TEXT NOT NULL,
+    -- JSON of the client (see runtime/pkg/oauthclient.Client), encrypted when a key ID is set
+    client TEXT NOT NULL,
+    client_encryption_key_id TEXT NOT NULL DEFAULT '',
+    created_on TIMESTAMPTZ NOT NULL DEFAULT now()
+);
+
+CREATE UNIQUE INDEX connector_oauth_clients_project_id_connector_config_hash_idx ON connector_oauth_clients (project_id, connector, config_hash);
+
+-- A user's OAuth token, obtained through a client. Tokens are never shared between users.
+-- A token is only used (and refreshed) through the client that obtained it.
+CREATE TABLE connector_oauth_tokens (
+    id UUID NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
+    oauth_client_id UUID NOT NULL REFERENCES connector_oauth_clients (id) ON DELETE CASCADE,
+    user_id UUID NOT NULL REFERENCES users (id) ON DELETE CASCADE,
+    -- JSON of the token (access token, refresh token, expiry), encrypted when a key ID is set
+    token TEXT NOT NULL,
+    token_encryption_key_id TEXT NOT NULL DEFAULT '',
+    created_on TIMESTAMPTZ NOT NULL DEFAULT now(),
+    updated_on TIMESTAMPTZ NOT NULL DEFAULT now()
+);
+
+CREATE UNIQUE INDEX connector_oauth_tokens_oauth_client_id_user_id_idx ON connector_oauth_tokens (oauth_client_id, user_id);
+CREATE INDEX connector_oauth_tokens_user_id_idx ON connector_oauth_tokens (user_id);

admin/database/postgres/postgres.go

@@ -4288,3 +4288,191 @@ func dedupeAndSortResourceNames(resources []database.ResourceName) []database.Re
 
 	return deduped
 }
+
+func (c *connection) FindConnectorOAuthClient(ctx context.Context, projectID, connector, configHash string) (*database.ConnectorOAuthClient, error) {
+	res := &database.ConnectorOAuthClient{}
+	err := c.getDB(ctx).QueryRowxContext(ctx, "SELECT * FROM connector_oauth_clients WHERE project_id=$1 AND connector=$2 AND config_hash=$3", projectID, connector, configHash).StructScan(res)
+	if err != nil {
+		return nil, parseErr("connector oauth client", err)
+	}
+	err = c.decryptConnectorOAuthClient(res)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
+func (c *connection) FindConnectorOAuthClientByID(ctx context.Context, id string) (*database.ConnectorOAuthClient, error) {
+	res := &database.ConnectorOAuthClient{}
+	err := c.getDB(ctx).QueryRowxContext(ctx, "SELECT * FROM connector_oauth_clients WHERE id=$1", id).StructScan(res)
+	if err != nil {
+		return nil, parseErr("connector oauth client", err)
+	}
+	err = c.decryptConnectorOAuthClient(res)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
+func (c *connection) FindOrInsertConnectorOAuthClient(ctx context.Context, opts *database.InsertConnectorOAuthClientOptions) (*database.ConnectorOAuthClient, error) {
+	if err := database.Validate(opts); err != nil {
+		return nil, err
+	}
+	value, keyID, err := c.encryptText(opts.Client)
+	if err != nil {
+		return nil, err
+	}
+	res := &database.ConnectorOAuthClient{}
+	err = c.getDB(ctx).QueryRowxContext(ctx, `
+		INSERT INTO connector_oauth_clients (project_id, connector, config_hash, client, client_encryption_key_id)
+		VALUES ($1, $2, $3, $4, $5)
+		ON CONFLICT (project_id, connector, config_hash) DO UPDATE SET config_hash = EXCLUDED.config_hash
+		RETURNING *`,
+		opts.ProjectID, opts.Connector, opts.ConfigHash, value, keyID,
+	).StructScan(res)
+	if err != nil {
+		return nil, parseErr("connector oauth client", err)
+	}
+	err = c.decryptConnectorOAuthClient(res)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
+func (c *connection) FindConnectorOAuthClients(ctx context.Context, projectID, connector string) ([]*database.ConnectorOAuthClient, error) {
+	var res []*database.ConnectorOAuthClient
+	err := c.getDB(ctx).SelectContext(ctx, &res, "SELECT * FROM connector_oauth_clients WHERE project_id=$1 AND connector=$2", projectID, connector)
+	if err != nil {
+		return nil, parseErr("connector oauth clients", err)
+	}
+	for _, r := range res {
+		err = c.decryptConnectorOAuthClient(r)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return res, nil
+}
+
+func (c *connection) DeleteConnectorOAuthClientsWithoutTokens(ctx context.Context, projectID, connector, keepConfigHash string) error {
+	_, err := c.getDB(ctx).ExecContext(ctx, `
+		DELETE FROM connector_oauth_clients c
+		WHERE c.project_id=$1 AND c.connector=$2 AND c.config_hash<>$3
+		AND NOT EXISTS (SELECT 1 FROM connector_oauth_tokens t WHERE t.oauth_client_id=c.id)
+	`, projectID, connector, keepConfigHash)
+	if err != nil {
+		return parseErr("connector oauth clients", err)
+	}
+	return nil
+}
+
+func (c *connection) DeleteConnectorOAuthClient(ctx context.Context, id string) error {
+	res, err := c.getDB(ctx).ExecContext(ctx, "DELETE FROM connector_oauth_clients WHERE id=$1", id)
+	return checkDeleteRow("connector oauth client", res, err)
+}
+
+func (c *connection) decryptConnectorOAuthClient(res *database.ConnectorOAuthClient) error {
+	client, err := c.decryptText(res.Client, res.ClientEncryptionKeyID)
+	if err != nil {
+		return err
+	}
+	res.Client = client
+	res.ClientEncryptionKeyID = ""
+	return nil
+}
+
+func (c *connection) AcquireConnectorOAuthTokenLock(ctx context.Context, oauthClientID, userID string) error {
+	if txFromContext(ctx) == nil {
+		return errors.New("postgres: connector OAuth locks must be acquired within a transaction")
+	}
+	key := "connector_oauth_tokens/" + oauthClientID + "/" + userID
+	_, err := c.getDB(ctx).ExecContext(ctx, "SELECT pg_advisory_xact_lock(hashtext($1))", key)
+	if err != nil {
+		return parseErr("connector oauth lock", err)
+	}
+	return nil
+}
+
+func (c *connection) FindConnectorOAuthToken(ctx context.Context, oauthClientID, userID string) (*database.ConnectorOAuthToken, error) {
+	res := &database.ConnectorOAuthToken{}
+	err := c.getDB(ctx).QueryRowxContext(ctx, "SELECT * FROM connector_oauth_tokens WHERE oauth_client_id=$1 AND user_id=$2", oauthClientID, userID).StructScan(res)
+	if err != nil {
+		return nil, parseErr("connector oauth token", err)
+	}
+	err = c.decryptConnectorOAuthToken(res)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
+func (c *connection) UpsertConnectorOAuthToken(ctx context.Context, opts *database.UpsertConnectorOAuthTokenOptions) (*database.ConnectorOAuthToken, error) {
+	if err := database.Validate(opts); err != nil {
+		return nil, err
+	}
+	value, keyID, err := c.encryptText(opts.Token)
+	if err != nil {
+		return nil, err
+	}
+	res := &database.ConnectorOAuthToken{}
+	err = c.getDB(ctx).QueryRowxContext(ctx, `
+		INSERT INTO connector_oauth_tokens (oauth_client_id, user_id, token, token_encryption_key_id)
+		VALUES ($1, $2, $3, $4)
+		ON CONFLICT (oauth_client_id, user_id) DO UPDATE SET
+			token = EXCLUDED.token,
+			token_encryption_key_id = EXCLUDED.token_encryption_key_id,
+			updated_on = now()
+		RETURNING *`,
+		opts.OAuthClientID, opts.UserID, value, keyID,
+	).StructScan(res)
+	if err != nil {
+		return nil, parseErr("connector oauth token", err)
+	}
+	err = c.decryptConnectorOAuthToken(res)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
+func (c *connection) DeleteConnectorOAuthToken(ctx context.Context, oauthClientID, userID string) error {
+	res, err := c.getDB(ctx).ExecContext(ctx, "DELETE FROM connector_oauth_tokens WHERE oauth_client_id=$1 AND user_id=$2", oauthClientID, userID)
+	return checkDeleteRow("connector oauth token", res, err)
+}
+
+func (c *connection) decryptConnectorOAuthToken(res *database.ConnectorOAuthToken) error {
+	token, err := c.decryptText(res.Token, res.TokenEncryptionKeyID)
+	if err != nil {
+		return err
+	}
+	res.Token = token
+	res.TokenEncryptionKeyID = ""
+	return nil
+}
+
+// encryptText encrypts a value for storage in a TEXT column.
+// It returns the value to store (base64 of the ciphertext when encryption is configured, the plain text otherwise) and the key ID used.
+func (c *connection) encryptText(plain []byte) (string, string, error) {
+	encrypted, keyID, err := c.encrypt(plain)
+	if err != nil {
+		return "", "", err
+	}
+	if keyID == "" {
+		return string(plain), "", nil
+	}
+	return base64.StdEncoding.EncodeToString(encrypted), keyID, nil
+}
+
+// decryptText reverses encryptText.
+func (c *connection) decryptText(stored []byte, keyID string) ([]byte, error) {
+	if keyID == "" {
+		return stored, nil
+	}
+	decoded, err := base64.StdEncoding.DecodeString(string(stored))
+	if err != nil {
+		return nil, err
+	}
+	return c.decrypt(decoded, keyID)
+}

admin/database/postgres/postgres_test.go

@@ -43,6 +43,7 @@ func TestPostgres(t *testing.T) {
 	t.Run("TestProjectsForUsersWithPagination", func(t *testing.T) { testProjectsForUserWithPagination(t, db) })
 	t.Run("TestMembersWithPagination", func(t *testing.T) { testOrgsMembersPagination(t, db) })
 	t.Run("TestUpsertProjectVariable", func(t *testing.T) { testUpsertProjectVariable(t, db) })
+	t.Run("TestConnectorOAuth", func(t *testing.T) { testConnectorOAuth(t, db) })
 	t.Run("TestManagedGitRepos", func(t *testing.T) { testManagedGitRepos(t, db) })
 	t.Run("TestOrganizationMemberUserAttributes", func(t *testing.T) { testOrganizationMemberUserAttributes(t, db) })
 	t.Run("TestOrganizationInviteAttributes", func(t *testing.T) { testOrganizationInviteAttributes(t, db) })
@@ -525,6 +526,79 @@ func testOrgsMembersPagination(t *testing.T, db database.DB) {
 	require.NoError(t, db.DeleteOrganization(ctx, "alpha"))
 }
 
+func testConnectorOAuth(t *testing.T, db database.DB) {
+	ctx := context.Background()
+	_, projectID, userID := seed(t, db)
+
+	// Clients are unique per connector and configuration hash; a duplicate insert returns the existing client
+	c1, err := db.FindOrInsertConnectorOAuthClient(ctx, &database.InsertConnectorOAuthClientOptions{
+		ProjectID:  projectID,
+		Connector:  "notion",
+		ConfigHash: "hash1",
+		Client:     []byte(`{"client_id":"a"}`),
+	})
+	require.NoError(t, err)
+	require.Equal(t, `{"client_id":"a"}`, string(c1.Client))
+	dup, err := db.FindOrInsertConnectorOAuthClient(ctx, &database.InsertConnectorOAuthClientOptions{
+		ProjectID:  projectID,
+		Connector:  "notion",
+		ConfigHash: "hash1",
+		Client:     []byte(`{"client_id":"b"}`),
+	})
+	require.NoError(t, err)
+	require.Equal(t, c1.ID, dup.ID)
+	require.Equal(t, `{"client_id":"a"}`, string(dup.Client))
+	c2, err := db.FindOrInsertConnectorOAuthClient(ctx, &database.InsertConnectorOAuthClientOptions{
+		ProjectID:  projectID,
+		Connector:  "notion",
+		ConfigHash: "hash2",
+		Client:     []byte(`{"client_id":"c"}`),
+	})
+	require.NoError(t, err)
+
+	found, err := db.FindConnectorOAuthClient(ctx, projectID, "notion", "hash1")
+	require.NoError(t, err)
+	require.Equal(t, c1.ID, found.ID)
+	require.Equal(t, `{"client_id":"a"}`, string(found.Client))
+	_, err = db.FindConnectorOAuthClient(ctx, projectID, "linear", "hash1")
+	require.ErrorIs(t, err, database.ErrNotFound)
+
+	// Tokens are unique per client and user, and are replaced on upsert
+	_, err = db.UpsertConnectorOAuthToken(ctx, &database.UpsertConnectorOAuthTokenOptions{OAuthClientID: c1.ID, UserID: userID, Token: []byte(`{"access_token":"t1"}`)})
+	require.NoError(t, err)
+	_, err = db.UpsertConnectorOAuthToken(ctx, &database.UpsertConnectorOAuthTokenOptions{OAuthClientID: c1.ID, UserID: userID, Token: []byte(`{"access_token":"t2"}`)})
+	require.NoError(t, err)
+	_, err = db.UpsertConnectorOAuthToken(ctx, &database.UpsertConnectorOAuthTokenOptions{OAuthClientID: c2.ID, UserID: userID, Token: []byte(`{"access_token":"t3"}`)})
+	require.NoError(t, err)
+	tok, err := db.FindConnectorOAuthToken(ctx, c1.ID, userID)
+	require.NoError(t, err)
+	require.Equal(t, `{"access_token":"t2"}`, string(tok.Token))
+
+	require.NoError(t, db.DeleteConnectorOAuthToken(ctx, c1.ID, userID))
+	_, err = db.FindConnectorOAuthToken(ctx, c1.ID, userID)
+	require.ErrorIs(t, err, database.ErrNotFound)
+	require.ErrorIs(t, db.DeleteConnectorOAuthToken(ctx, c1.ID, userID), database.ErrNotFound)
+
+	// A connector has a client per configuration; clients without tokens are deleted, except the current one
+	clients, err := db.FindConnectorOAuthClients(ctx, projectID, "notion")
+	require.NoError(t, err)
+	require.Len(t, clients, 2)
+	require.NoError(t, db.DeleteConnectorOAuthClientsWithoutTokens(ctx, projectID, "notion", "hash1"))
+	clients, err = db.FindConnectorOAuthClients(ctx, projectID, "notion")
+	require.NoError(t, err)
+	require.Len(t, clients, 2, "c1 is current and c2 holds a token")
+	require.NoError(t, db.DeleteConnectorOAuthClientsWithoutTokens(ctx, projectID, "notion", "hash2"))
+	_, err = db.FindConnectorOAuthClientByID(ctx, c1.ID)
+	require.ErrorIs(t, err, database.ErrNotFound)
+	_, err = db.FindConnectorOAuthClientByID(ctx, c2.ID)
+	require.NoError(t, err)
+
+	// Deleting a client deletes its tokens
+	require.NoError(t, db.DeleteConnectorOAuthClient(ctx, c2.ID))
+	_, err = db.FindConnectorOAuthToken(ctx, c2.ID, userID)
+	require.ErrorIs(t, err, database.ErrNotFound)
+}
+
 func testUpsertProjectVariable(t *testing.T, db database.DB) {
 	_, projectID, userID := seed(t, db)