Setup SQL server based on pgx frontend/backend protocol

admin/server/auth/middleware.go

@@ -75,6 +75,22 @@ func (a *Authenticator) HTTPMiddlewareLenient(next http.Handler) http.Handler {
 	return a.httpMiddleware(next, true)
 }
 
+type PostgresPassword struct{}
+
+func (a *Authenticator) PostgresAuthHandler() func(ctx context.Context, username, password string) (context.Context, bool, error) {
+	return func(ctx context.Context, username, password string) (context.Context, bool, error) {
+		// Clients do not pass Bearer to avoid encoding the password
+		// The default token type is Bearer. In case of different token type, the type will be passed in postgres additional properties
+		password = fmt.Sprintf("bearer %s", password)
+		ctx = context.WithValue(ctx, PostgresPassword{}, password)
+		newCtx, err := a.parseClaimsFromBearer(ctx, password)
+		if err != nil {
+			newCtx = context.WithValue(ctx, claimsContextKey{}, anonClaims{})
+		}
+		return newCtx, true, nil
+	}
+}
+
 // httpMiddleware is the actual implementation of HTTPMiddleware and HTTPMiddlewareLenient.
 func (a *Authenticator) httpMiddleware(next http.Handler, lenient bool) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

admin/server/postgres.go

@@ -0,0 +1,204 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	wire "github.com/jeroenrinzema/psql-wire"
+	"github.com/lib/pq/oid"
+	"github.com/rilldata/rill/admin/server/auth"
+	runtimeauth "github.com/rilldata/rill/runtime/server/auth"
+	"go.uber.org/zap"
+)
+
+func (s *Server) QueryHandler(ctx context.Context, query string) (wire.PreparedStatements, error) {
+	s.logger.Debug("query", zap.String("query", query))
+	if strings.Trim(query, " ") == "" {
+		return wire.Prepared(wire.NewStatement(func(ctx context.Context, writer wire.DataWriter, parameters []wire.Parameter) error {
+			return writer.Empty()
+		})), nil
+	}
+
+	upperQuery := strings.ToUpper(query)
+	if strings.HasPrefix(upperQuery, "SET") {
+		return wire.Prepared(wire.NewStatement(func(ctx context.Context, writer wire.DataWriter, parameters []wire.Parameter) error {
+			return writer.Complete("SET")
+		}, wire.WithColumns(nil))), nil
+	}
+
+	if strings.HasPrefix(upperQuery, "BEGIN") || strings.HasPrefix(upperQuery, "COMMIT") || strings.HasPrefix(upperQuery, "ROLLBACK") {
+		return wire.Prepared(wire.NewStatement(func(ctx context.Context, writer wire.DataWriter, parameters []wire.Parameter) error {
+			return writer.Complete(strings.Trim(upperQuery, ";"))
+		}, wire.WithColumns(nil))), nil
+	}
+
+	clientParams := wire.ClientParameters(ctx)
+	// database is org.password
+	tokens := strings.Split(clientParams[wire.ParamDatabase], ".")
+	if len(tokens) != 2 {
+		return nil, fmt.Errorf("invalid org or project")
+	}
+	org := tokens[0]
+	project := tokens[1]
+
+	// Find the production deployment for the project we're proxying to
+	proj, err := s.admin.DB.FindProjectByName(ctx, org, project)
+	if err != nil {
+		return nil, fmt.Errorf("invalid org or project")
+	}
+
+	if proj.ProdDeploymentID == nil {
+		return nil, fmt.Errorf("no prod deployment for project")
+	}
+	depl, err := s.admin.DB.FindDeployment(ctx, *proj.ProdDeploymentID)
+	if err != nil {
+		return nil, fmt.Errorf("no prod deployment for project")
+	}
+
+	var jwt string
+	claims := auth.GetClaims(ctx)
+	switch claims.OwnerType() {
+	case auth.OwnerTypeAnon:
+		// If the client is not authenticated with the admin service, we just proxy the contents of the password to the runtime (if any).
+		password := ctx.Value(auth.PostgresPassword{}).(string)
+		if len(password) >= 6 && strings.EqualFold(password[0:6], "bearer") {
+			jwt = strings.TrimSpace(password[6:])
+		}
+	case auth.OwnerTypeUser, auth.OwnerTypeService:
+		// If the client is authenticated with the admin service, we issue a new ephemeral runtime JWT.
+		// The JWT should have the same permissions/configuration as one they would get by calling AdminService.GetProject.
+		permissions := claims.ProjectPermissions(ctx, proj.OrganizationID, depl.ProjectID)
+		if !permissions.ReadProd {
+			return nil, fmt.Errorf("does not have permission to access the production deployment")
+		}
+
+		var attr map[string]any
+		if claims.OwnerType() == auth.OwnerTypeUser {
+			attr, err = s.jwtAttributesForUser(ctx, claims.OwnerID(), proj.OrganizationID, permissions)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		jwt, err = s.issuer.NewToken(runtimeauth.TokenOptions{
+			AudienceURL: depl.RuntimeAudience,
+			Subject:     claims.OwnerID(),
+			TTL:         runtimeAccessTokenDefaultTTL,
+			InstancePermissions: map[string][]runtimeauth.Permission{
+				depl.RuntimeInstanceID: {
+					// TODO: Remove ReadProfiling and ReadRepo (may require frontend changes)
+					runtimeauth.ReadObjects,
+					runtimeauth.ReadMetrics,
+					runtimeauth.ReadProfiling,
+					runtimeauth.ReadRepo,
+					runtimeauth.ReadAPI,
+				},
+			},
+			Attributes: attr,
+		})
+		if err != nil {
+			return nil, err
+		}
+	default:
+		return nil, fmt.Errorf("runtime proxy not available for owner type %q", claims.OwnerType())
+	}
+
+	// Track usage of the deployment
+	s.admin.Used.Deployment(depl.ID)
+
+	hostURL, err := url.Parse(depl.RuntimeHost)
+	if err != nil {
+		return nil, err
+	}
+	hostURL.Scheme = "postgres"
+	hostURL.Host = hostURL.Hostname() + ":" + strconv.FormatInt(int64(15432), 10)
+	hostURL.User = url.UserPassword("postgres", fmt.Sprintf("Bearer %s", jwt))
+	hostURL.Path = depl.RuntimeInstanceID
+	conn, err := connectionPool(ctx, hostURL.String())
+	if err != nil {
+		s.logger.Info("error in get connection pool", zap.Error(err))
+		return nil, err
+	}
+
+	rows, err := conn.Query(ctx, query) // query to underlying host
+	if err != nil {
+		s.logger.Info("error in query", zap.Error(err))
+		return nil, err
+	}
+	defer rows.Close()
+
+	// handle schema
+	var cols []wire.Column
+	fds := rows.FieldDescriptions()
+	for _, fd := range fds {
+		cols = append(cols, wire.Column{
+			Table: int32(fd.TableOID),
+			Name:  fd.Name,
+			Oid:   oid.Oid(fd.DataTypeOID),
+			Width: fd.DataTypeSize,
+			Attr:  int16(fd.TableAttributeNumber),
+		})
+	}
+
+	// handle data
+	// NOTE :: This creates a copy of data and stores this till client starts reading data. This is required so that we
+	// can close runtime connection and not wait for client to complete reading whole data which can leak connection.
+	// We can improve this logic in future.
+	var data [][]any
+	for rows.Next() {
+		d, err := rows.Values()
+		if err != nil {
+			s.logger.Info("error in fetching next row", zap.Error(err))
+			return nil, err
+		}
+		data = append(data, d)
+	}
+	if rows.Err() != nil {
+		s.logger.Info("error in fetching rows", zap.Error(err))
+		return nil, err
+	}
+
+	handle := func(ctx context.Context, writer wire.DataWriter, parameters []wire.Parameter) error {
+		for i := 0; i < len(data); i++ {
+			if err := writer.Row(data[i]); err != nil {
+				return err
+			}
+		}
+		return writer.Complete("OK")
+	}
+	return wire.Prepared(wire.NewStatement(handle, wire.WithColumns(cols))), nil
+}
+
+var (
+	runtimePool map[string]*pgxpool.Pool = make(map[string]*pgxpool.Pool)
+	mu          sync.Mutex
+)
+
+func connectionPool(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
+	mu.Lock()
+	defer mu.Unlock()
+
+	pool, ok := runtimePool[dsn]
+	if ok {
+		return pool, nil
+	}
+
+	config, err := pgxpool.ParseConfig(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse dsn: %w", err)
+	}
+
+	// Runtime JWts are valid for 30 minutes only
+	config.MaxConnLifetime = time.Minute * 29
+	// since runtimes get restarted more often than actual DB servers. Consider if this should be reduced to even less time
+	// also consider if we should add some health check on connection acquisition
+	config.HealthCheckPeriod = time.Minute
+
+	return pgxpool.NewWithConfig(ctx, config)
+}

admin/server/server.go

@@ -48,6 +48,7 @@ var (
 type Options struct {
 	HTTPPort               int
 	GRPCPort               int
+	PostgresPort           int
 	ExternalURL            string
 	FrontendURL            string
 	AllowedOrigins         []string
@@ -189,6 +190,12 @@ func (s *Server) ServeHTTP(ctx context.Context) error {
 	})
 }
 
+// ServePostgres Starts the postrges server.
+func (s *Server) ServePostgres(ctx context.Context) error {
+	authHandler := s.authenticator.PostgresAuthHandler()
+	return graceful.ServePostgres(ctx, s.QueryHandler, authHandler, s.opts.PostgresPort, false, s.logger)
+}
+
 // HTTPHandler HTTP handler serving REST gateway.
 func (s *Server) HTTPHandler(ctx context.Context) (http.Handler, error) {
 	// Create REST gateway

cli/cmd/admin/start.go

@@ -48,6 +48,7 @@ type Config struct {
 	TracesExporter           observability.Exporter `default:"" split_words:"true"`
 	HTTPPort                 int                    `default:"8080" split_words:"true"`
 	GRPCPort                 int                    `default:"9090" split_words:"true"`
+	PostgresPort             int                    `default:"25432" split_words:"true"`
 	DebugPort                int                    `split_words:"true"`
 	ExternalURL              string                 `default:"http://localhost:8080" split_words:"true"`
 	ExternalGRPCURL          string                 `envconfig:"external_grpc_url"`
@@ -288,6 +289,7 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 				srv, err := server.New(logger, adm, issuer, limiter, activityClient, &server.Options{
 					HTTPPort:               conf.HTTPPort,
 					GRPCPort:               conf.GRPCPort,
+					PostgresPort:           conf.PostgresPort,
 					ExternalURL:            conf.ExternalURL,
 					FrontendURL:            conf.FrontendURL,
 					AllowedOrigins:         conf.AllowedOrigins,
@@ -306,6 +308,7 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 				}
 				group.Go(func() error { return srv.ServeGRPC(cctx) })
 				group.Go(func() error { return srv.ServeHTTP(cctx) })
+				group.Go(func() error { return srv.ServePostgres(cctx) })
 				if conf.DebugPort != 0 {
 					group.Go(func() error { return debugserver.ServeHTTP(cctx, conf.DebugPort) })
 				}

cli/cmd/runtime/start.go

@@ -43,7 +43,6 @@ import (
 	_ "github.com/rilldata/rill/runtime/drivers/s3"
 	_ "github.com/rilldata/rill/runtime/drivers/salesforce"
 	_ "github.com/rilldata/rill/runtime/drivers/slack"
-	_ "github.com/rilldata/rill/runtime/drivers/snowflake"
 	_ "github.com/rilldata/rill/runtime/drivers/sqlite"
 	_ "github.com/rilldata/rill/runtime/reconcilers"
 	_ "github.com/rilldata/rill/runtime/resolvers"
@@ -62,6 +61,7 @@ type Config struct {
 	HTTPPort                int                    `default:"8080" split_words:"true"`
 	GRPCPort                int                    `default:"9090" split_words:"true"`
 	DebugPort               int                    `default:"6060" split_words:"true"`
+	PostgresPort            int                    `default:"15432" split_words:"true"`
 	AllowedOrigins          []string               `default:"*" split_words:"true"`
 	SessionKeyPairs         []string               `split_words:"true"`
 	AuthEnable              bool                   `default:"false" split_words:"true"`
@@ -237,12 +237,14 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 			srvOpts := &server.Options{
 				HTTPPort:        conf.HTTPPort,
 				GRPCPort:        conf.GRPCPort,
+				PostgresPort:    conf.PostgresPort,
 				AllowedOrigins:  conf.AllowedOrigins,
 				ServePrometheus: conf.MetricsExporter == observability.PrometheusExporter,
 				SessionKeyPairs: keyPairs,
 				AuthEnable:      conf.AuthEnable,
 				AuthIssuerURL:   conf.AuthIssuerURL,
 				AuthAudienceURL: conf.AuthAudienceURL,
+				DataDir:         conf.DataDir,
 			}
 			s, err := server.NewServer(ctx, srvOpts, rt, logger, limiter, activityClient)
 			if err != nil {
@@ -253,6 +255,7 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 			group, cctx := errgroup.WithContext(ctx)
 			group.Go(func() error { return s.ServeGRPC(cctx) })
 			group.Go(func() error { return s.ServeHTTP(cctx, nil) })
+			group.Go(func() error { return s.ServePostgres(cctx, true) })
 			if conf.DebugPort != 0 {
 				group.Go(func() error { return debugserver.ServeHTTP(cctx, conf.DebugPort) })
 			}

cli/cmd/start/start.go

@@ -24,6 +24,7 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 	var olapDSN string
 	var httpPort int
 	var grpcPort int
+	var postgresPort int
 	var verbose bool
 	var debug bool
 	var readonly bool
@@ -157,7 +158,7 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 
 			userID, _ := ch.CurrentUserID(cmd.Context())
 
-			err = app.Serve(httpPort, grpcPort, !noUI, !noOpen, readonly, userID, tlsCertPath, tlsKeyPath)
+			err = app.Serve(httpPort, grpcPort, postgresPort, !noUI, !noOpen, readonly, userID, tlsCertPath, tlsKeyPath)
 			if err != nil {
 				return fmt.Errorf("serve: %w", err)
 			}
@@ -170,6 +171,7 @@ func StartCmd(ch *cmdutil.Helper) *cobra.Command {
 	startCmd.Flags().BoolVar(&noOpen, "no-open", false, "Do not open browser")
 	startCmd.Flags().IntVar(&httpPort, "port", 9009, "Port for HTTP")
 	startCmd.Flags().IntVar(&grpcPort, "port-grpc", 49009, "Port for gRPC (internal)")
+	startCmd.Flags().IntVar(&postgresPort, "port-postgres", 0, "Port for postgres server")
 	startCmd.Flags().BoolVar(&readonly, "readonly", false, "Show only dashboards in UI")
 	startCmd.Flags().BoolVar(&noUI, "no-ui", false, "Serve only the backend")
 	startCmd.Flags().BoolVar(&verbose, "verbose", false, "Sets the log level to debug")

cli/pkg/local/app.go

@@ -334,7 +334,7 @@ func (a *App) Close() error {
 	return nil
 }
 
-func (a *App) Serve(httpPort, grpcPort int, enableUI, openBrowser, readonly bool, userID, tlsCertPath, tlsKeyPath string) error {
+func (a *App) Serve(httpPort, grpcPort, postgresPort int, enableUI, openBrowser, readonly bool, userID, tlsCertPath, tlsKeyPath string) error {
 	// Get analytics info
 	installID, enabled, err := dotrill.AnalyticsInfo()
 	if err != nil {
@@ -378,6 +378,7 @@ func (a *App) Serve(httpPort, grpcPort int, enableUI, openBrowser, readonly bool
 	opts := &runtimeserver.Options{
 		HTTPPort:        httpPort,
 		GRPCPort:        grpcPort,
+		PostgresPort:    postgresPort,
 		TLSCertPath:     tlsCertPath,
 		TLSKeyPath:      tlsKeyPath,
 		AllowedOrigins:  []string{"*"},
@@ -408,6 +409,11 @@ func (a *App) Serve(httpPort, grpcPort int, enableUI, openBrowser, readonly bool
 	if a.Debug {
 		group.Go(func() error { return debugserver.ServeHTTP(ctx, 6060) })
 	}
+	if postgresPort != 0 {
+		group.Go(func() error {
+			return runtimeServer.ServePostgres(ctx, false)
+		})
+	}
 
 	// Open the browser when health check succeeds
 	go a.pollServer(ctx, httpPort, enableUI && openBrowser, secure)