Skip to content

Commit

Permalink
cli: Add role support for db shell (#1486)
Browse files Browse the repository at this point in the history
ekerfelt authored Oct 15, 2024
1 parent fbe032f commit d61400d
Showing 7 changed files with 701 additions and 553 deletions.
39 changes: 35 additions & 4 deletions cli/cmd/encore/db.go
Original file line number Diff line number Diff line change
@@ -25,12 +25,28 @@ var dbCmd = &cobra.Command{
}

var (
resetAll bool
testDB bool
shadowDB bool
nsName string
resetAll bool
testDB bool
shadowDB bool
write bool
admin bool
superuser bool
nsName string
)

func getDBRole() daemonpb.DBRole {
switch {
case superuser:
return daemonpb.DBRole_DB_ROLE_SUPERUSER
case admin:
return daemonpb.DBRole_DB_ROLE_ADMIN
case write:
return daemonpb.DBRole_DB_ROLE_WRITE
default:
return daemonpb.DBRole_DB_ROLE_READ
}
}

var dbResetCmd = &cobra.Command{
Use: "reset <database-names...|--all>",
Short: "Resets the databases with the given names. Use --all to reset all databases.",
@@ -119,6 +135,7 @@ when using tools like Prisma.
EnvName: dbEnv,
ClusterType: dbClusterType(),
Namespace: nonZeroPtr(nsName),
Role: getDBRole(),
})
if err != nil {
fatalf("could not connect to the database for service %s: %v", dbName, err)
@@ -199,6 +216,7 @@ when using tools like Prisma.
Port: dbProxyPort,
ClusterType: dbClusterType(),
Namespace: nonZeroPtr(nsName),
Role: getDBRole(),
})
if err != nil {
log.Fatal().Err(err).Msg("could not setup db proxy")
@@ -256,6 +274,7 @@ when using tools like Prisma.
EnvName: dbEnv,
ClusterType: dbClusterType(),
Namespace: nonZeroPtr(nsName),
Role: getDBRole(),
})
if err != nil {
st, ok := status.FromError(err)
@@ -284,19 +303,31 @@ func init() {
dbShellCmd.Flags().StringVarP(&dbEnv, "env", "e", "local", "Environment name to connect to (such as \"prod\")")
dbShellCmd.Flags().BoolVarP(&testDB, "test", "t", false, "Connect to the integration test database (implies --env=local)")
dbShellCmd.Flags().BoolVar(&shadowDB, "shadow", false, "Connect to the shadow database (implies --env=local)")
dbShellCmd.Flags().BoolVar(&write, "write", false, "Connect with write privileges")
dbShellCmd.Flags().BoolVar(&admin, "admin", false, "Connect with admin privileges")
dbShellCmd.Flags().BoolVar(&superuser, "superuser", false, "Connect as a superuser")
dbShellCmd.MarkFlagsMutuallyExclusive("write", "admin", "superuser")
dbCmd.AddCommand(dbShellCmd)

dbProxyCmd.Flags().StringVarP(&nsName, "namespace", "n", "", "Namespace to use (defaults to active namespace)")
dbProxyCmd.Flags().StringVarP(&dbEnv, "env", "e", "local", "Environment name to connect to (such as \"prod\")")
dbProxyCmd.Flags().Int32VarP(&dbProxyPort, "port", "p", 0, "Port to listen on (defaults to a random port)")
dbProxyCmd.Flags().BoolVarP(&testDB, "test", "t", false, "Connect to the integration test database (implies --env=local)")
dbProxyCmd.Flags().BoolVar(&shadowDB, "shadow", false, "Connect to the shadow database (implies --env=local)")
dbProxyCmd.Flags().BoolVar(&write, "write", false, "Connect with write privileges")
dbProxyCmd.Flags().BoolVar(&admin, "admin", false, "Connect with admin privileges")
dbProxyCmd.Flags().BoolVar(&superuser, "superuser", false, "Connect as a superuser")
dbProxyCmd.MarkFlagsMutuallyExclusive("write", "admin", "superuser")
dbCmd.AddCommand(dbProxyCmd)

dbConnURICmd.Flags().StringVarP(&nsName, "namespace", "n", "", "Namespace to use (defaults to active namespace)")
dbConnURICmd.Flags().StringVarP(&dbEnv, "env", "e", "local", "Environment name to connect to (such as \"prod\")")
dbConnURICmd.Flags().BoolVarP(&testDB, "test", "t", false, "Connect to the integration test database (implies --env=local)")
dbConnURICmd.Flags().BoolVar(&shadowDB, "shadow", false, "Connect to the shadow database (implies --env=local)")
dbConnURICmd.Flags().BoolVar(&write, "write", false, "Connect with write privileges")
dbConnURICmd.Flags().BoolVar(&admin, "admin", false, "Connect with admin privileges")
dbConnURICmd.Flags().BoolVar(&superuser, "superuser", false, "Connect as a superuser")
dbConnURICmd.MarkFlagsMutuallyExclusive("write", "admin", "superuser")
dbCmd.AddCommand(dbConnURICmd)
}

18 changes: 16 additions & 2 deletions cli/daemon/db.go
Original file line number Diff line number Diff line change
@@ -21,6 +21,20 @@ import (
daemonpb "encr.dev/proto/encore/daemon"
)

func toRoleType(role daemonpb.DBRole) sqldb.RoleType {
switch role {
case daemonpb.DBRole_DB_ROLE_READ:
return sqldb.RoleRead
case daemonpb.DBRole_DB_ROLE_WRITE:
return sqldb.RoleWrite
case daemonpb.DBRole_DB_ROLE_ADMIN:
return sqldb.RoleAdmin
default:
return sqldb.RoleRead
}

}

// DBConnect starts the database and returns the DSN for connecting to it.
func (s *Server) DBConnect(ctx context.Context, req *daemonpb.DBConnectRequest) (*daemonpb.DBConnectResponse, error) {
if req.EnvName == "local" {
@@ -33,7 +47,7 @@ func (s *Server) DBConnect(ctx context.Context, req *daemonpb.DBConnectRequest)
} else if appID == "" {
return nil, errNotLinked
}
port, passwd, err := sqldb.OneshotProxy(appID, req.EnvName)
port, passwd, err := sqldb.OneshotProxy(appID, req.EnvName, toRoleType(req.Role))
if err != nil {
return nil, err
}
@@ -224,7 +238,7 @@ func (s *Server) DBProxy(params *daemonpb.DBProxyRequest, stream daemonpb.Daemon
FrontendTLS: nil,
DialBackend: func(ctx context.Context, startup *pgproxy.StartupData) (pgproxy.LogicalConn, error) {
startupData := startup.Raw.Encode(nil)
ws, err := platform.DBConnect(ctx, appID, params.EnvName, startup.Database, startupData)
ws, err := platform.DBConnect(ctx, appID, params.EnvName, startup.Database, toRoleType(params.Role).String(), startupData)
if err != nil {
return nil, err
}
2 changes: 2 additions & 0 deletions cli/daemon/sqldb/cluster.go
Original file line number Diff line number Diff line change
@@ -425,6 +425,8 @@ func (roles EncoreRoles) find(typ RoleType) (Role, bool) {

type RoleType string

func (r RoleType) String() string { return string(r) }

const (
RoleSuperuser RoleType = "superuser"
RoleAdmin RoleType = "admin"
8 changes: 4 additions & 4 deletions cli/daemon/sqldb/remote.go
Original file line number Diff line number Diff line change
@@ -19,7 +19,7 @@ import (
// OneshotProxy listens on a random port for a single connection, and proxies that connection to a remote db.
// It reports the one-time password and port to use.
// Once a connection has been established, it stops listening.
func OneshotProxy(appSlug, envSlug string) (port int, passwd string, err error) {
func OneshotProxy(appSlug, envSlug string, role RoleType) (port int, passwd string, err error) {
ln, err := net.Listen("tcp", "localhost:0")
if err != nil {
return 0, "", err
@@ -30,11 +30,11 @@ func OneshotProxy(appSlug, envSlug string) (port int, passwd string, err error)
}
passwd = base64.RawURLEncoding.EncodeToString(passwdBytes[:])

go oneshotServer(context.Background(), ln, passwd, appSlug, envSlug)
go oneshotServer(context.Background(), ln, passwd, appSlug, envSlug, role)
return ln.Addr().(*net.TCPAddr).Port, passwd, nil
}

func oneshotServer(ctx context.Context, ln net.Listener, passwd, appSlug, envSlug string) error {
func oneshotServer(ctx context.Context, ln net.Listener, passwd, appSlug, envSlug string, role RoleType) error {
proxy := &pgproxy.SingleBackendProxy{
RequirePassword: passwd != "",
FrontendTLS: nil,
@@ -43,7 +43,7 @@ func oneshotServer(ctx context.Context, ln net.Listener, passwd, appSlug, envSlu
return nil, fmt.Errorf("bad password")
}
startupData := startup.Raw.Encode(nil)
ws, err := platform.DBConnect(ctx, appSlug, envSlug, startup.Database, startupData)
ws, err := platform.DBConnect(ctx, appSlug, envSlug, startup.Database, role.String(), startupData)
if err != nil {
var e platform.Error
if errors.As(err, &e) && e.HTTPCode == 404 {
5 changes: 4 additions & 1 deletion cli/internal/platform/api.go
Original file line number Diff line number Diff line change
@@ -136,8 +136,11 @@ func GetEnvMeta(ctx context.Context, appSlug, envName string) (*metav1.Data, err
return &md, nil
}

func DBConnect(ctx context.Context, appSlug, envSlug, dbName string, startupData []byte) (*websocket.Conn, error) {
func DBConnect(ctx context.Context, appSlug, envSlug, dbName, role string, startupData []byte) (*websocket.Conn, error) {
path := escapef("/apps/%s/envs/%s/sqldb-connect/%s", appSlug, envSlug, dbName)
if role != "" {
path += "?role=" + url.QueryEscape(role)
}
return wsDial(ctx, path, true, map[string]string{
"X-Startup-Message": base64.StdEncoding.EncodeToString(startupData),
})
1,171 changes: 629 additions & 542 deletions proto/encore/daemon/daemon.pb.go

Large diffs are not rendered by default.

11 changes: 11 additions & 0 deletions proto/encore/daemon/daemon.proto
Original file line number Diff line number Diff line change
@@ -235,6 +235,16 @@ message DBConnectRequest {
// namespace is the infrastructure namespace to use.
// If empty the active namespace is used.
optional string namespace = 5;

DBRole role = 6;
}

enum DBRole {
DB_ROLE_UNSPECIFIED = 0;
DB_ROLE_SUPERUSER = 1;
DB_ROLE_ADMIN = 2;
DB_ROLE_WRITE = 3;
DB_ROLE_READ = 4;
}

enum DBClusterType {
@@ -257,6 +267,7 @@ message DBProxyRequest {
// namespace is the infrastructure namespace to use.
// If empty the active namespace is used.
optional string namespace = 5;
DBRole role = 6;
}

message DBResetRequest {

0 comments on commit d61400d

Please sign in to comment.