Skip to content

Commit 6837fb3

Browse files
blktblkt
authored
Add base queries (#115)
This change adds an initial set of queries that should cover most use cases. All non-trivial listing queries are paginated, and no batch inserts were added at this point. I also took the chance to fix some typos in the initial migration.
1 parent afc6c34 commit 6837fb3

File tree

14 files changed

+1572
-44
lines changed

14 files changed

+1572
-44
lines changed

‎database/migrate.go‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
// Package database provides functions to migrate the database.
2+
// This is a temporary hack while we wait for the tooling of choice
3+
// to be on main.
4+
package database
5+
6+
import (
7+
"context"
8+
_ "embed"
9+
10+
"github.com/jackc/pgx/v5"
11+
)
12+
13+
//go:embed migrations/000001_init.up.sql
14+
var initMigrationUp string
15+
16+
//go:embed migrations/000001_init.down.sql
17+
var initMigrationDown string
18+
19+
// MigrateUp executes the database migrations
20+
func MigrateUp(ctx context.Context, db *pgx.Conn) error {
21+
_, err := db.Exec(ctx, initMigrationUp)
22+
return err
23+
}
24+
25+
// MigrateDown executes the database migrations in reverse order
26+
func MigrateDown(ctx context.Context, db *pgx.Conn) error {
27+
_, err := db.Exec(ctx, initMigrationDown)
28+
return err
29+
}

‎database/migrations/000001_init.up.sql‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@ CREATE TABLE registry (
1212
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
1313
name TEXT NOT NULL,
1414
reg_type registry_type NOT NULL DEFAULT 'LOCAL',
15-
created_at TIMESTAMPZ DEFAULT NOW(),
16-
updated_at TIMESTAMPZ DEFAULT NOW(),
15+
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
16+
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
1717
UNIQUE(name)
1818
);
1919

@@ -33,8 +33,8 @@ CREATE TABLE registry_sync (
3333
reg_id UUID REFERENCES registry(id) ON DELETE CASCADE,
3434
sync_status sync_status NOT NULL DEFAULT 'IN_PROGRESS',
3535
error_msg TEXT, -- Populated if sync_status = 'FAILED'
36-
started_at TIMESTAMPZ DEFAULT NOW(),
37-
ended_at TIMESTAMPZ
36+
started_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
37+
ended_at TIMESTAMP WITH TIME ZONE
3838
);
3939

4040
CREATE INDEX registry_sync_started_at_idx ON registry_sync(reg_id, started_at);
@@ -47,8 +47,8 @@ CREATE TABLE mcp_server (
4747
name TEXT NOT NULL,
4848
version TEXT NOT NULL,
4949
reg_id UUID REFERENCES registry(id) ON DELETE CASCADE,
50-
created_at TIMESTAMPZ DEFAULT NOW(),
51-
updated_at TIMESTAMPZ DEFAULT NOW(),
50+
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
51+
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
5252
description TEXT,
5353
title TEXT,
5454
website TEXT,

‎database/queries/registry.sql‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,27 @@
1-
-- name: InsertRegistry :exec
2-
INSERT INTO registry (name, reg_type) VALUES ($1, $2);
1+
-- name: ListRegistries :many
2+
SELECT id,
3+
name,
4+
reg_type,
5+
created_at,
6+
updated_at
7+
FROM registry
8+
WHERE (sqlc.narg(next)::timestamp with time zone IS NULL OR created_at > sqlc.narg(next))
9+
OR (sqlc.narg(prev)::timestamp with time zone IS NULL AND created_at < sqlc.narg(prev))
10+
ORDER BY
11+
-- next page sorting
12+
CASE WHEN sqlc.narg(next)::timestamp with time zone IS NULL THEN created_at END ASC,
13+
-- previous page sorting
14+
CASE WHEN sqlc.narg(prev)::timestamp with time zone IS NULL THEN created_at END DESC
15+
LIMIT sqlc.arg(size)::bigint;
16+
17+
-- name: GetRegistry :one
18+
SELECT id,
19+
name,
20+
reg_type,
21+
created_at,
22+
updated_at
23+
FROM registry
24+
WHERE id = sqlc.arg(id);
25+
26+
-- name: InsertRegistry :one
27+
INSERT INTO registry (name, reg_type) VALUES ($1, $2) RETURNING id;

‎database/queries/servers.sql‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
-- name: ListServers :many
2+
SELECT r.reg_type as registry_type,
3+
s.id,
4+
s.name,
5+
s.version,
6+
l.latest_server_id IS NOT NULL AS is_latest,
7+
s.created_at,
8+
s.updated_at,
9+
s.description,
10+
s.title,
11+
s.website,
12+
s.upstream_meta,
13+
s.server_meta,
14+
s.repository_url,
15+
s.repository_id,
16+
s.repository_subfolder,
17+
s.repository_type
18+
FROM mcp_server s
19+
JOIN registry r ON s.reg_id = r.id
20+
LEFT JOIN latest_server_version l ON s.id = l.latest_server_id
21+
WHERE (sqlc.narg(next)::timestamp with time zone IS NULL OR sqlc.narg(next) > s.created_at)
22+
ORDER BY
23+
-- next page sorting
24+
CASE WHEN sqlc.narg(next)::timestamp with time zone IS NULL THEN s.reg_type END ASC,
25+
CASE WHEN sqlc.narg(next)::timestamp with time zone IS NULL THEN s.name END ASC,
26+
CASE WHEN sqlc.narg(next)::timestamp with time zone IS NULL THEN s.created_at END ASC,
27+
CASE WHEN sqlc.narg(next)::timestamp with time zone IS NULL THEN s.version END ASC, -- acts as tie breaker
28+
-- previous page sorting
29+
CASE WHEN sqlc.narg(prev)::timestamp with time zone IS NULL THEN s.reg_type END DESC,
30+
CASE WHEN sqlc.narg(prev)::timestamp with time zone IS NULL THEN s.name END DESC,
31+
CASE WHEN sqlc.narg(prev)::timestamp with time zone IS NULL THEN s.created_at END DESC,
32+
CASE WHEN sqlc.narg(prev)::timestamp with time zone IS NULL THEN s.version END DESC -- acts as tie breaker
33+
LIMIT sqlc.arg(size)::bigint;
34+
35+
-- name: ListServerPackages :many
36+
SELECT p.server_id,
37+
p.registry_type,
38+
p.pkg_registry_url,
39+
p.pkg_identifier,
40+
p.pkg_version,
41+
p.runtime_hint,
42+
p.runtime_arguments,
43+
p.package_arguments,
44+
p.env_vars,
45+
p.sha256_hash,
46+
p.transport,
47+
p.transport_url,
48+
p.transport_headers
49+
FROM mcp_server_package p
50+
JOIN mcp_server s ON p.server_id = s.id
51+
WHERE s.id IN (sqlc.slice(server_ids)::UUID[])
52+
ORDER BY p.pkg_version DESC;
53+
54+
-- name: ListServerRemotes :many
55+
SELECT r.server_id,
56+
r.transport,
57+
r.transport_url,
58+
r.transport_headers
59+
FROM mcp_server_remote r
60+
WHERE r.server_id IN (sqlc.slice(server_ids)::UUID[])
61+
ORDER BY r.transport, r.transport_url;
62+
63+
-- name: ListServerVersions :many
64+
SELECT s.id,
65+
s.name,
66+
s.version,
67+
s.created_at,
68+
s.updated_at,
69+
s.description,
70+
s.title,
71+
s.website,
72+
s.upstream_meta,
73+
s.server_meta,
74+
s.repository_url,
75+
s.repository_id,
76+
s.repository_subfolder,
77+
s.repository_type
78+
FROM mcp_server s
79+
WHERE s.name = sqlc.arg(name)
80+
ORDER BY
81+
CASE WHEN sqlc.narg(next)::timestamp with time zone IS NULL THEN s.created_at END ASC,
82+
CASE WHEN sqlc.narg(next)::timestamp with time zone IS NULL THEN s.version END DESC -- acts as tie breaker
83+
LIMIT sqlc.arg(size)::bigint;
84+
85+
-- name: UpsertServerVersion :exec
86+
INSERT INTO mcp_server (
87+
name,
88+
version,
89+
reg_id,
90+
created_at,
91+
updated_at,
92+
description,
93+
title,
94+
website,
95+
upstream_meta,
96+
server_meta,
97+
repository_url,
98+
repository_id,
99+
repository_subfolder,
100+
repository_type
101+
) VALUES (
102+
sqlc.arg(name),
103+
sqlc.arg(version),
104+
sqlc.arg(reg_id),
105+
CURRENT_TIMESTAMP,
106+
CURRENT_TIMESTAMP,
107+
sqlc.narg(description),
108+
sqlc.narg(title),
109+
sqlc.narg(website),
110+
sqlc.narg(upstream_meta),
111+
sqlc.narg(server_meta),
112+
sqlc.narg(repository_url),
113+
sqlc.narg(repository_id),
114+
sqlc.narg(repository_subfolder),
115+
sqlc.narg(repository_type)
116+
) ON CONFLICT (reg_id, name, version)
117+
DO UPDATE SET
118+
updated_at = CURRENT_TIMESTAMP,
119+
description = sqlc.narg(description),
120+
title = sqlc.narg(title),
121+
website = sqlc.narg(website),
122+
upstream_meta = sqlc.narg(upstream_meta),
123+
server_meta = sqlc.narg(server_meta),
124+
repository_url = sqlc.narg(repository_url),
125+
repository_id = sqlc.narg(repository_id),
126+
repository_subfolder = sqlc.narg(repository_subfolder),
127+
repository_type = sqlc.narg(repository_type);
128+
129+
-- name: UpsertLatestServerVersion :exec
130+
INSERT INTO latest_server_version (
131+
reg_id,
132+
name,
133+
version,
134+
latest_server_id
135+
) VALUES (
136+
sqlc.arg(reg_id),
137+
sqlc.arg(name),
138+
sqlc.arg(version),
139+
sqlc.arg(server_id)
140+
) ON CONFLICT (reg_id, name)
141+
DO UPDATE SET
142+
version = sqlc.arg(version),
143+
latest_server_id = sqlc.arg(server_id);
144+
145+
-- name: UpsertServerPackage :exec
146+
INSERT INTO mcp_server_package (
147+
server_id,
148+
registry_type,
149+
pkg_registry_url,
150+
pkg_identifier,
151+
pkg_version,
152+
runtime_hint,
153+
runtime_arguments,
154+
package_arguments,
155+
env_vars,
156+
sha256_hash,
157+
transport,
158+
transport_url,
159+
transport_headers
160+
) VALUES (
161+
sqlc.arg(server_id),
162+
sqlc.arg(registry_type),
163+
sqlc.arg(pkg_registry_url),
164+
sqlc.arg(pkg_identifier),
165+
sqlc.arg(pkg_version),
166+
sqlc.narg(runtime_hint),
167+
sqlc.narg(runtime_arguments),
168+
sqlc.narg(package_arguments),
169+
sqlc.narg(env_vars),
170+
sqlc.narg(sha256_hash),
171+
sqlc.arg(transport),
172+
sqlc.narg(transport_url),
173+
sqlc.narg(transport_headers)
174+
);
175+
176+
-- name: UpsertServerRemote :exec
177+
INSERT INTO mcp_server_remote (
178+
server_id,
179+
transport,
180+
transport_url,
181+
transport_headers
182+
) VALUES (
183+
sqlc.arg(server_id),
184+
sqlc.arg(transport),
185+
sqlc.narg(transport_url),
186+
sqlc.narg(transport_headers)
187+
) ON CONFLICT (server_id, transport, transport_url)
188+
DO UPDATE SET
189+
transport_headers = sqlc.narg(transport_headers);
190+
191+
-- name: UpsertServerIcon :exec
192+
INSERT INTO mcp_server_icon (
193+
server_id,
194+
source_uri,
195+
mime_type,
196+
theme
197+
) VALUES (
198+
sqlc.arg(server_id),
199+
sqlc.arg(source_uri),
200+
sqlc.arg(mime_type),
201+
sqlc.arg(theme)
202+
) ON CONFLICT (server_id, source_uri, mime_type, theme)
203+
DO UPDATE SET
204+
theme = sqlc.arg(theme);

‎database/queries/sync.sql‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
-- name: GetRegistrySync :one
2+
SELECT id,
3+
reg_id,
4+
sync_status,
5+
error_msg,
6+
started_at,
7+
ended_at
8+
FROM registry_sync
9+
WHERE id = sqlc.arg(id);
10+
11+
-- name: InsertRegistrySync :exec
12+
INSERT INTO registry_sync (
13+
reg_id,
14+
sync_status,
15+
error_msg,
16+
started_at
17+
) VALUES (
18+
sqlc.arg(reg_id),
19+
sqlc.arg(sync_status),
20+
sqlc.narg(error_msg),
21+
CURRENT_TIMESTAMP
22+
);
23+
24+
-- name: UpdateRegistrySync :exec
25+
UPDATE registry_sync SET
26+
sync_status = sqlc.arg(sync_status),
27+
error_msg = sqlc.narg(error_msg),
28+
ended_at = sqlc.arg(ended_at)
29+
WHERE id = sqlc.arg(id);

‎database/testcontainers.go‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package database
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/jackc/pgx/v5"
8+
"github.com/stretchr/testify/require"
9+
tc "github.com/testcontainers/testcontainers-go"
10+
tclog "github.com/testcontainers/testcontainers-go/log"
11+
"github.com/testcontainers/testcontainers-go/modules/postgres"
12+
)
13+
14+
type nopLogger struct{}
15+
16+
func (*nopLogger) Printf(_ string, _ ...any) {}
17+
18+
var _ tclog.Logger = (*nopLogger)(nil)
19+
20+
var (
21+
dbName = "testdb"
22+
dbUser = "testuser"
23+
dbPass = "testpass"
24+
)
25+
26+
// SetupTestDB creates a Postgres container using testcontainers and runs migrations
27+
func SetupTestDB(t *testing.T) (*pgx.Conn, func()) {
28+
t.Helper()
29+
30+
ctx := context.Background()
31+
32+
// Start Postgres container
33+
postgresContainer, err := postgres.Run(
34+
ctx,
35+
"postgres:16-alpine",
36+
postgres.WithDatabase(dbName),
37+
postgres.WithUsername(dbUser),
38+
postgres.WithPassword(dbPass),
39+
postgres.BasicWaitStrategies(),
40+
tc.WithLogger(&nopLogger{}),
41+
)
42+
require.NoError(t, err)
43+
44+
// Get connection string
45+
connStr, err := postgresContainer.ConnectionString(ctx, "sslmode=disable")
46+
require.NoError(t, err)
47+
48+
// Connect to database with retry logic
49+
var db *pgx.Conn
50+
db, err = pgx.Connect(ctx, connStr)
51+
require.NoError(t, err)
52+
53+
// Run migrations
54+
err = MigrateUp(ctx, db)
55+
require.NoError(t, err)
56+
57+
err = MigrateDown(ctx, db)
58+
require.NoError(t, err)
59+
60+
err = MigrateUp(ctx, db)
61+
require.NoError(t, err)
62+
63+
cleanupFunc := func() {
64+
//nolint:gosec
65+
_ = db.Close(ctx)
66+
tc.CleanupContainer(t, postgresContainer)
67+
}
68+
69+
return db, cleanupFunc
70+
}

0 commit comments

Comments
 (0)