Skip to content

Commit d9b9e0e

Browse files
sweatybridgesweatybridge
authored
feat: support custom query timeout for type gen (#4082)
1 parent 7a28eab commit d9b9e0e

File tree

6 files changed

+54
-35
lines changed

6 files changed

+54
-35
lines changed

‎cmd/gen.go‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package cmd
33
import (
44
"os"
55
"os/signal"
6+
"time"
67

78
env "github.com/Netflix/go-env"
89
"github.com/go-errors/errors"
@@ -56,6 +57,7 @@ var (
5657
},
5758
Value: types.LangTypescript,
5859
}
60+
queryTimeout time.Duration
5961
postgrestV9Compat bool
6062
swiftAccessControl = utils.EnumFlag{
6163
Allowed: []string{
@@ -88,7 +90,7 @@ var (
8890
return err
8991
}
9092
}
91-
return types.Run(ctx, flags.ProjectRef, flags.DbConfig, lang.Value, schema, postgrestV9Compat, swiftAccessControl.Value, afero.NewOsFs())
93+
return types.Run(ctx, flags.ProjectRef, flags.DbConfig, lang.Value, schema, postgrestV9Compat, swiftAccessControl.Value, queryTimeout, afero.NewOsFs())
9294
},
9395
Example: ` supabase gen types --local
9496
supabase gen types --linked --lang=go
@@ -126,8 +128,13 @@ func init() {
126128
genTypesCmd.MarkFlagsMutuallyExclusive("local", "linked", "project-id", "db-url")
127129
typeFlags.Var(&lang, "lang", "Output language of the generated types.")
128130
typeFlags.StringSliceVarP(&schema, "schema", "s", []string{}, "Comma separated list of schema to include.")
131+
// Direct connection only flags
129132
typeFlags.Var(&swiftAccessControl, "swift-access-control", "Access control for Swift generated types.")
130-
typeFlags.BoolVar(&postgrestV9Compat, "postgrest-v9-compat", false, "Generate types compatible with PostgREST v9 and below. Only use together with --db-url.")
133+
genTypesCmd.MarkFlagsMutuallyExclusive("linked", "project-id", "swift-access-control")
134+
typeFlags.BoolVar(&postgrestV9Compat, "postgrest-v9-compat", false, "Generate types compatible with PostgREST v9 and below.")
135+
genTypesCmd.MarkFlagsMutuallyExclusive("linked", "project-id", "postgrest-v9-compat")
136+
typeFlags.DurationVar(&queryTimeout, "query-timeout", time.Second*15, "Maximum timeout allowed for the database query.")
137+
genTypesCmd.MarkFlagsMutuallyExclusive("linked", "project-id", "query-timeout")
131138
genCmd.AddCommand(genTypesCmd)
132139
keyFlags := genKeysCmd.Flags()
133140
keyFlags.StringVar(&flags.ProjectRef, "project-ref", "", "Project ref of the Supabase project.")

‎internal/db/diff/migra.go‎

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,6 @@ var (
2323
//go:embed templates/migra.ts
2424
diffSchemaTypeScript string
2525

26-
//go:embed templates/staging-ca-2021.crt
27-
caStaging string
28-
//go:embed templates/prod-ca-2021.crt
29-
caProd string
30-
3126
managedSchemas = []string{
3227
// Local development
3328
"_analytics",
@@ -107,12 +102,10 @@ func loadSchema(ctx context.Context, dbURL string, options ...func(*pgx.ConnConf
107102

108103
func DiffSchemaMigra(ctx context.Context, source, target string, schema []string, options ...func(*pgx.ConnConfig)) (string, error) {
109104
env := []string{"SOURCE=" + source, "TARGET=" + target}
110-
// node-postgres does not support sslmode=prefer
111-
if require, err := types.IsRequireSSL(ctx, target, options...); err != nil {
105+
if ca, err := types.GetRootCA(ctx, target, options...); err != nil {
112106
return "", err
113-
} else if require {
114-
rootCA := caStaging + caProd
115-
env = append(env, "SSL_CA="+rootCA)
107+
} else if len(ca) > 0 {
108+
env = append(env, "SSL_CA="+ca)
116109
}
117110
if len(schema) > 0 {
118111
env = append(env, "INCLUDED_SCHEMAS="+strings.Join(schema, ","))

‎internal/db/diff/templates/prod-ca-2021.crt‎ renamed to ‎internal/gen/types/templates/prod-ca-2021.crt‎

File renamed without changes.

‎internal/db/diff/templates/staging-ca-2021.crt‎ renamed to ‎internal/gen/types/templates/staging-ca-2021.crt‎

File renamed without changes.

‎internal/gen/types/types.go‎

Lines changed: 33 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,11 @@ package types
22

33
import (
44
"context"
5+
_ "embed"
56
"fmt"
67
"os"
78
"strings"
9+
"time"
810

911
"github.com/docker/docker/api/types/container"
1012
"github.com/docker/docker/api/types/network"
@@ -28,7 +30,7 @@ const (
2830
SwiftInternalAccessControl = "internal"
2931
)
3032

31-
func Run(ctx context.Context, projectId string, dbConfig pgconn.Config, lang string, schemas []string, postgrestV9Compat bool, swiftAccessControl string, fsys afero.Fs, options ...func(*pgx.ConnConfig)) error {
33+
func Run(ctx context.Context, projectId string, dbConfig pgconn.Config, lang string, schemas []string, postgrestV9Compat bool, swiftAccessControl string, queryTimeout time.Duration, fsys afero.Fs, options ...func(*pgx.ConnConfig)) error {
3234
originalURL := utils.ToPostgresURL(dbConfig)
3335
// Add default schemas if --schema flag is not specified
3436
if len(schemas) == 0 {
@@ -77,26 +79,27 @@ func Run(ctx context.Context, projectId string, dbConfig pgconn.Config, lang str
7779
}
7880

7981
fmt.Fprintln(os.Stderr, "Connecting to", dbConfig.Host, dbConfig.Port)
80-
escaped := utils.ToPostgresURL(dbConfig)
81-
if require, err := IsRequireSSL(ctx, originalURL, options...); err != nil {
82+
env := []string{
83+
"PG_META_DB_URL=" + utils.ToPostgresURL(dbConfig),
84+
fmt.Sprintf("PG_CONN_TIMEOUT_SECS=%.0f", queryTimeout.Seconds()),
85+
fmt.Sprintf("PG_QUERY_TIMEOUT_SECS=%.0f", queryTimeout.Seconds()),
86+
"PG_META_GENERATE_TYPES=" + lang,
87+
"PG_META_GENERATE_TYPES_INCLUDED_SCHEMAS=" + included,
88+
"PG_META_GENERATE_TYPES_SWIFT_ACCESS_CONTROL=" + swiftAccessControl,
89+
fmt.Sprintf("PG_META_GENERATE_TYPES_DETECT_ONE_TO_ONE_RELATIONSHIPS=%v", !postgrestV9Compat),
90+
}
91+
if ca, err := GetRootCA(ctx, originalURL, options...); err != nil {
8292
return err
83-
} else if require {
84-
// node-postgres does not support sslmode=prefer
85-
escaped += "&sslmode=require"
93+
} else if len(ca) > 0 {
94+
env = append(env, "PG_META_DB_SSL_ROOT_CERT="+ca)
8695
}
8796

8897
return utils.DockerRunOnceWithConfig(
8998
ctx,
9099
container.Config{
91100
Image: utils.Config.Studio.PgmetaImage,
92-
Env: []string{
93-
"PG_META_DB_URL=" + escaped,
94-
"PG_META_GENERATE_TYPES=" + lang,
95-
"PG_META_GENERATE_TYPES_INCLUDED_SCHEMAS=" + included,
96-
"PG_META_GENERATE_TYPES_SWIFT_ACCESS_CONTROL=" + swiftAccessControl,
97-
fmt.Sprintf("PG_META_GENERATE_TYPES_DETECT_ONE_TO_ONE_RELATIONSHIPS=%v", !postgrestV9Compat),
98-
},
99-
Cmd: []string{"node", "dist/server/server.js"},
101+
Env: env,
102+
Cmd: []string{"node", "dist/server/server.js"},
100103
},
101104
hostConfig,
102105
network.NetworkingConfig{},
@@ -106,7 +109,22 @@ func Run(ctx context.Context, projectId string, dbConfig pgconn.Config, lang str
106109
)
107110
}
108111

109-
func IsRequireSSL(ctx context.Context, dbUrl string, options ...func(*pgx.ConnConfig)) (bool, error) {
112+
var (
113+
//go:embed templates/staging-ca-2021.crt
114+
caStaging string
115+
//go:embed templates/prod-ca-2021.crt
116+
caProd string
117+
)
118+
119+
func GetRootCA(ctx context.Context, dbURL string, options ...func(*pgx.ConnConfig)) (string, error) {
120+
// node-postgres does not support sslmode=prefer
121+
if require, err := isRequireSSL(ctx, dbURL, options...); !require {
122+
return "", err
123+
}
124+
return caStaging + caProd, nil
125+
}
126+
127+
func isRequireSSL(ctx context.Context, dbUrl string, options ...func(*pgx.ConnConfig)) (bool, error) {
110128
conn, err := utils.ConnectByUrl(ctx, dbUrl+"&sslmode=require", options...)
111129
if err != nil {
112130
if strings.HasSuffix(err.Error(), "(server refused TLS connection)") {

‎internal/gen/types/types_test.go‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"errors"
66
"net/http"
77
"testing"
8+
"time"
89

910
"github.com/docker/docker/api/types/container"
1011
"github.com/h2non/gock"
@@ -48,7 +49,7 @@ func TestGenLocalCommand(t *testing.T) {
4849
conn := pgtest.NewConn()
4950
defer conn.Close(t)
5051
// Run test
51-
assert.NoError(t, Run(context.Background(), "", dbConfig, LangTypescript, []string{}, true, "", fsys, conn.Intercept))
52+
assert.NoError(t, Run(context.Background(), "", dbConfig, LangTypescript, []string{}, true, "", time.Second, fsys, conn.Intercept))
5253
// Validate api
5354
assert.Empty(t, apitest.ListUnmatchedRequests())
5455
})
@@ -63,7 +64,7 @@ func TestGenLocalCommand(t *testing.T) {
6364
Get("/v" + utils.Docker.ClientVersion() + "/containers/" + utils.DbId).
6465
Reply(http.StatusServiceUnavailable)
6566
// Run test
66-
assert.Error(t, Run(context.Background(), "", dbConfig, LangTypescript, []string{}, true, "", fsys))
67+
assert.Error(t, Run(context.Background(), "", dbConfig, LangTypescript, []string{}, true, "", time.Second, fsys))
6768
// Validate api
6869
assert.Empty(t, apitest.ListUnmatchedRequests())
6970
})
@@ -83,7 +84,7 @@ func TestGenLocalCommand(t *testing.T) {
8384
Get("/v" + utils.Docker.ClientVersion() + "/images").
8485
Reply(http.StatusServiceUnavailable)
8586
// Run test
86-
assert.Error(t, Run(context.Background(), "", dbConfig, LangTypescript, []string{}, true, "", fsys))
87+
assert.Error(t, Run(context.Background(), "", dbConfig, LangTypescript, []string{}, true, "", time.Second, fsys))
8788
// Validate api
8889
assert.Empty(t, apitest.ListUnmatchedRequests())
8990
})
@@ -106,7 +107,7 @@ func TestGenLocalCommand(t *testing.T) {
106107
conn := pgtest.NewConn()
107108
defer conn.Close(t)
108109
// Run test
109-
assert.NoError(t, Run(context.Background(), "", dbConfig, LangSwift, []string{}, true, SwiftInternalAccessControl, fsys, conn.Intercept))
110+
assert.NoError(t, Run(context.Background(), "", dbConfig, LangSwift, []string{}, true, SwiftInternalAccessControl, time.Second, fsys, conn.Intercept))
110111
// Validate api
111112
assert.Empty(t, apitest.ListUnmatchedRequests())
112113
})
@@ -129,7 +130,7 @@ func TestGenLinkedCommand(t *testing.T) {
129130
Reply(200).
130131
JSON(api.TypescriptResponse{Types: ""})
131132
// Run test
132-
assert.NoError(t, Run(context.Background(), projectId, pgconn.Config{}, LangTypescript, []string{}, true, "", fsys))
133+
assert.NoError(t, Run(context.Background(), projectId, pgconn.Config{}, LangTypescript, []string{}, true, "", time.Second, fsys))
133134
// Validate api
134135
assert.Empty(t, apitest.ListUnmatchedRequests())
135136
})
@@ -144,7 +145,7 @@ func TestGenLinkedCommand(t *testing.T) {
144145
Get("/v1/projects/" + projectId + "/types/typescript").
145146
ReplyError(errNetwork)
146147
// Run test
147-
err := Run(context.Background(), projectId, pgconn.Config{}, LangTypescript, []string{}, true, "", fsys)
148+
err := Run(context.Background(), projectId, pgconn.Config{}, LangTypescript, []string{}, true, "", time.Second, fsys)
148149
// Validate api
149150
assert.ErrorIs(t, err, errNetwork)
150151
assert.Empty(t, apitest.ListUnmatchedRequests())
@@ -159,7 +160,7 @@ func TestGenLinkedCommand(t *testing.T) {
159160
Get("/v1/projects/" + projectId + "/types/typescript").
160161
Reply(http.StatusServiceUnavailable)
161162
// Run test
162-
assert.Error(t, Run(context.Background(), projectId, pgconn.Config{}, LangTypescript, []string{}, true, "", fsys))
163+
assert.Error(t, Run(context.Background(), projectId, pgconn.Config{}, LangTypescript, []string{}, true, "", time.Second, fsys))
163164
})
164165
}
165166

@@ -184,7 +185,7 @@ func TestGenRemoteCommand(t *testing.T) {
184185
conn := pgtest.NewConn()
185186
defer conn.Close(t)
186187
// Run test
187-
assert.NoError(t, Run(context.Background(), "", dbConfig, LangTypescript, []string{"public"}, true, "", afero.NewMemMapFs(), conn.Intercept))
188+
assert.NoError(t, Run(context.Background(), "", dbConfig, LangTypescript, []string{"public"}, true, "", time.Second, afero.NewMemMapFs(), conn.Intercept))
188189
// Validate api
189190
assert.Empty(t, apitest.ListUnmatchedRequests())
190191
})

0 commit comments

Comments
 (0)