Skip to content

Commit aad36cf

Browse files
ethanndicksonethanndickson
committed
add pubsub to dbmem
1 parent f5a64b2 commit aad36cf

File tree

6 files changed

+65
-58
lines changed

6 files changed

+65
-58
lines changed

‎cli/server.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -580,14 +580,15 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
580580
return xerrors.Errorf("parse ssh config options %q: %w", vals.SSHConfig.SSHConfigOptions.String(), err)
581581
}
582582

583+
db, ps := dbmem.NewWithPubsub()
583584
options := &coderd.Options{
584585
AccessURL: vals.AccessURL.Value(),
585586
AppHostname: appHostname,
586587
AppHostnameRegex: appHostnameRegex,
587588
Logger: logger.Named("coderd"),
588-
Database: dbmem.New(),
589+
Database: db,
589590
BaseDERPMap: derpMap,
590-
Pubsub: pubsub.NewInMemory(),
591+
Pubsub: ps,
591592
CacheDir: cacheDir,
592593
GoogleTokenValidator: googleTokenValidator,
593594
ExternalAuthConfigs: externalAuthConfigs,
@@ -691,8 +692,9 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
691692
var pubsubWatchdogTimeout <-chan struct{}
692693
if vals.InMemoryDatabase {
693694
// This is only used for testing.
694-
options.Database = dbmem.New()
695-
options.Pubsub = pubsub.NewInMemory()
695+
db, ps := dbmem.NewWithPubsub()
696+
options.Pubsub = ps
697+
options.Database = db
696698
} else {
697699
sqlDB, dbURL, err := getPostgresDB(ctx, logger, vals.PostgresURL.String(), codersdk.PostgresAuth(vals.PostgresAuth), sqlDriver)
698700
if err != nil {

‎coderd/database/dbmem/dbmem.go

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626

2727
"github.com/coder/coder/v2/coderd/database"
2828
"github.com/coder/coder/v2/coderd/database/dbtime"
29+
"github.com/coder/coder/v2/coderd/database/pubsub"
2930
"github.com/coder/coder/v2/coderd/rbac"
3031
"github.com/coder/coder/v2/coderd/rbac/regosql"
3132
"github.com/coder/coder/v2/coderd/util/slice"
@@ -51,6 +52,17 @@ var (
5152

5253
// New returns an in-memory fake of the database.
5354
func New() database.Store {
55+
return newDbMem()
56+
}
57+
58+
func NewWithPubsub() (database.Store, pubsub.Pubsub) {
59+
q := newDbMem()
60+
ps := pubsub.NewInMemory()
61+
q.ps = ps
62+
return q, ps
63+
}
64+
65+
func newDbMem() *FakeQuerier {
5466
q := &FakeQuerier{
5567
mutex: &sync.RWMutex{},
5668
data: &data{
@@ -167,6 +179,7 @@ func (inTxMutex) RUnlock() {}
167179
// can do type checks.
168180
type FakeQuerier struct {
169181
mutex rwMutex
182+
ps pubsub.Pubsub
170183
*data
171184
}
172185

@@ -7776,7 +7789,7 @@ func (q *FakeQuerier) InsertWorkspace(_ context.Context, arg database.InsertWork
77767789
return workspace, nil
77777790
}
77787791

7779-
func (q *FakeQuerier) InsertWorkspaceAgent(_ context.Context, arg database.InsertWorkspaceAgentParams) (database.WorkspaceAgent, error) {
7792+
func (q *FakeQuerier) InsertWorkspaceAgent(ctx context.Context, arg database.InsertWorkspaceAgentParams) (database.WorkspaceAgent, error) {
77807793
if err := validateDatabaseType(arg); err != nil {
77817794
return database.WorkspaceAgent{}, err
77827795
}
@@ -7807,6 +7820,22 @@ func (q *FakeQuerier) InsertWorkspaceAgent(_ context.Context, arg database.Inser
78077820
}
78087821

78097822
q.workspaceAgents = append(q.workspaceAgents, agent)
7823+
if q.ps != nil {
7824+
workspace, err := q.getWorkspaceByAgentIDNoLock(ctx, agent.ID)
7825+
// Agents might not belong to a workspace (template imports)
7826+
if err != nil {
7827+
return agent, nil
7828+
}
7829+
msg, err := json.Marshal(codersdk.WorkspaceEvent{
7830+
Kind: codersdk.WorkspaceEventKindNewAgent,
7831+
WorkspaceID: workspace.ID,
7832+
AgentID: &agent.ID,
7833+
})
7834+
if err != nil {
7835+
return database.WorkspaceAgent{}, xerrors.Errorf("failed to marshal workspace event: %w", err)
7836+
}
7837+
_ = q.ps.Publish(codersdk.WorkspaceEventChannel(workspace.OwnerID), msg)
7838+
}
78107839
return agent, nil
78117840
}
78127841

‎coderd/database/dbtestutil/db.go

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,7 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
9595
opt(&o)
9696
}
9797

98-
db := dbmem.New()
99-
ps := pubsub.NewInMemory()
98+
db, ps := dbmem.NewWithPubsub()
10099
if WillUsePostgres() {
101100
connectionURL := os.Getenv("CODER_PG_CONNECTION_URL")
102101
if connectionURL == "" && o.url != "" {

‎coderd/database/dump.sql

Lines changed: 14 additions & 22 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/migrations/000260_workspace_updates.down.sql

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,4 @@
11
DROP TYPE agent_id_name_pair;
22

3-
DROP TRIGGER IF EXISTS new_workspace_notify ON workspaces;
4-
DROP FUNCTION IF EXISTS new_workspace_notify;
5-
63
DROP TRIGGER IF EXISTS new_agent_notify ON workspace_agents;
74
DROP FUNCTION IF EXISTS new_agent_notify;

‎coderd/database/migrations/000260_workspace_updates.up.sql

Lines changed: 14 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -3,35 +3,15 @@ CREATE TYPE agent_id_name_pair AS (
33
name text
44
);
55

6-
CREATE FUNCTION new_workspace_notify() RETURNS trigger
7-
LANGUAGE plpgsql
8-
AS $$
9-
DECLARE
10-
BEGIN
11-
-- Notify for new workspaces & ownership transfers
12-
IF TG_OP = 'INSERT' OR (TG_OP = 'UPDATE' AND NEW.owner_id <> OLD.owner_id) THEN
13-
-- Write to the notification channel `new_workspace:owner_id`
14-
-- with the workspace id as the payload.
15-
PERFORM pg_notify('new_workspace:' || NEW.owner_id, NEW.id::text);
16-
END IF;
17-
RETURN NEW;
18-
END;
19-
$$;
20-
21-
CREATE TRIGGER new_workspace_notify
22-
AFTER INSERT OR UPDATE ON workspaces
23-
FOR EACH ROW
24-
EXECUTE FUNCTION new_workspace_notify();
25-
26-
276
CREATE FUNCTION new_agent_notify() RETURNS trigger
287
LANGUAGE plpgsql
298
AS $$
309
DECLARE
31-
workspace_owner_id uuid;
10+
v_workspace_id uuid;
11+
v_workspace_owner_id uuid;
3212
BEGIN
33-
SELECT workspaces.owner_id
34-
INTO workspace_owner_id
13+
SELECT workspaces.owner_id, workspaces.id
14+
INTO v_workspace_owner_id, v_workspace_id
3515
FROM
3616
workspaces
3717
WHERE
@@ -58,9 +38,17 @@ BEGIN
5838
)
5939
);
6040
-- Agents might not belong to a workspace (template imports)
61-
IF workspace_owner_id IS NOT NULL THEN
41+
IF v_workspace_owner_id IS NOT NULL THEN
6242
-- Write to the notification channel `new_agent:workspace_owner_id`
63-
PERFORM pg_notify('new_agent:' || workspace_owner_id, '');
43+
PERFORM pg_notify(
44+
'workspace_owner:' || v_workspace_owner_id,
45+
json_build_object(
46+
'kind', 'new_agent',
47+
'workspace_id', v_workspace_id,
48+
'agent_id', NEW.id
49+
'agent_name', NEW.name
50+
)::text
51+
);
6452
END IF;
6553
RETURN NEW;
6654
END;

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy