ariadne/internal/platform/database/database.go

package database

import (
	"context"
	"errors"
	"fmt"
	"sync/atomic"

	"git.toki-labs.com/toki/ariadne/internal/platform/config"
	"github.com/jackc/pgx/v5/pgxpool"
	"go.uber.org/fx"
	"go.uber.org/zap"
)

type Database struct {
	runtimePool  *pgxpool.Pool
	identityPool *pgxpool.Pool
	compatible   atomic.Bool
}

func New(lifecycle fx.Lifecycle, cfg config.Config, logger *zap.Logger) (*Database, error) {
	database := &Database{}
	if cfg.DatabaseURL == "" || cfg.DatabaseIdentityURL == "" {
		logger.Warn(
			"runtime and identity databases are not fully configured; readiness checks will fail",
		)
		return database, nil
	}

	runtimePool, err := newPool(cfg.DatabaseURL, "ariadne")
	if err != nil {
		return nil, fmt.Errorf("configure runtime database: %w", err)
	}
	identityPool, err := newPool(cfg.DatabaseIdentityURL, "ariadne-identity")
	if err != nil {
		runtimePool.Close()
		return nil, fmt.Errorf("configure identity database: %w", err)
	}
	database.runtimePool = runtimePool
	database.identityPool = identityPool

	lifecycle.Append(fx.Hook{
		OnStart: func(ctx context.Context) error {
			if err := runtimePool.Ping(ctx); err != nil {
				return fmt.Errorf("ping runtime database: %w", err)
			}
			if err := identityPool.Ping(ctx); err != nil {
				return fmt.Errorf("ping identity database: %w", err)
			}
			if err := validateSchema(ctx, runtimePool); err != nil {
				return err
			}
			if err := validateTenantTableCatalog(ctx, runtimePool); err != nil {
				return err
			}
			if err := validateRuntimeConnection(
				ctx,
				runtimePool,
				cfg.DatabaseRuntimeRole,
			); err != nil {
				return err
			}
			if err := validateIdentityConnection(
				ctx,
				identityPool,
				cfg.DatabaseIdentityRole,
			); err != nil {
				return err
			}
			database.compatible.Store(true)
			logger.Info("runtime and identity databases connected")
			return nil
		},
		OnStop: func(context.Context) error {
			database.compatible.Store(false)
			identityPool.Close()
			runtimePool.Close()
			return nil
		},
	})

	return database, nil
}

func newPool(databaseURL, applicationName string) (*pgxpool.Pool, error) {
	poolConfig, err := pgxpool.ParseConfig(databaseURL)
	if err != nil {
		return nil, err
	}
	poolConfig.ConnConfig.RuntimeParams["application_name"] = applicationName
	return pgxpool.NewWithConfig(context.Background(), poolConfig)
}

func validateRuntimeConnection(
	ctx context.Context,
	pool *pgxpool.Pool,
	expectedRole string,
) error {
	roleName, err := validateConnectionRole(ctx, pool, expectedRole, "runtime")
	if err != nil {
		return err
	}
	if err := validateNoTenantTableOwnership(ctx, pool, roleName); err != nil {
		return err
	}
	if err := validateRuntimePrivileges(ctx, pool, roleName); err != nil {
		return err
	}
	return nil
}

func validateIdentityConnection(
	ctx context.Context,
	pool *pgxpool.Pool,
	expectedRole string,
) error {
	roleName, err := validateConnectionRole(ctx, pool, expectedRole, "identity")
	if err != nil {
		return err
	}
	if err := validateNoTenantTableOwnership(ctx, pool, roleName); err != nil {
		return err
	}
	return validateIdentityPrivileges(ctx, pool, roleName)
}

func validateConnectionRole(
	ctx context.Context,
	pool *pgxpool.Pool,
	expectedRole string,
	purpose string,
) (string, error) {
	if expectedRole == "" {
		return "", fmt.Errorf("%s database role is not configured", purpose)
	}

	var (
		roleName      string
		isSuperuser   bool
		canBypassRLS  bool
		canCreateDB   bool
		canCreateRole bool
		canReplicate  bool
		canInherit    bool
	)
	if err := pool.QueryRow(ctx, `
		SELECT
			rolname,
			rolsuper,
			rolbypassrls,
			rolcreatedb,
			rolcreaterole,
			rolreplication,
			rolinherit
		FROM pg_roles
		WHERE rolname = current_user
	`).Scan(
		&roleName,
		&isSuperuser,
		&canBypassRLS,
		&canCreateDB,
		&canCreateRole,
		&canReplicate,
		&canInherit,
	); err != nil {
		return "", fmt.Errorf("inspect %s database role: %w", purpose, err)
	}
	if roleName != expectedRole {
		return "", fmt.Errorf(
			"%s database connected as %q, expected %q",
			purpose,
			roleName,
			expectedRole,
		)
	}
	if isSuperuser ||
		canBypassRLS ||
		canCreateDB ||
		canCreateRole ||
		canReplicate ||
		canInherit {
		return "", fmt.Errorf(
			"%s database role %q must be NOSUPERUSER, NOBYPASSRLS, "+
				"NOCREATEDB, NOCREATEROLE, NOREPLICATION and NOINHERIT",
			purpose,
			roleName,
		)
	}
	var (
		hasSchemaUsage  bool
		hasSchemaCreate bool
	)
	if err := pool.QueryRow(ctx, `
		SELECT
			has_schema_privilege(current_user, 'public', 'USAGE'),
			has_schema_privilege(current_user, 'public', 'CREATE')
	`).Scan(&hasSchemaUsage, &hasSchemaCreate); err != nil {
		return "", fmt.Errorf("inspect %s schema privileges: %w", purpose, err)
	}
	if !hasSchemaUsage || hasSchemaCreate {
		return "", fmt.Errorf(
			"%s database role %q must have USAGE but not CREATE on public schema",
			purpose,
			roleName,
		)
	}
	return roleName, nil
}

func validateNoTenantTableOwnership(
	ctx context.Context,
	pool *pgxpool.Pool,
	roleName string,
) error {
	var ownsTenantTable bool
	if err := pool.QueryRow(ctx, `
		SELECT EXISTS (
			SELECT 1
			FROM pg_class table_class
			JOIN pg_namespace namespace
			  ON namespace.oid = table_class.relnamespace
			WHERE namespace.nspname = 'public'
			  AND table_class.relkind IN ('r', 'p')
			  AND table_class.relowner = (
			      SELECT oid FROM pg_roles WHERE rolname = $1
			  )
			  AND table_class.relname = ANY($2)
		)
	`, roleName, orderedRuntimeTables()).Scan(&ownsTenantTable); err != nil {
		return fmt.Errorf("check %s table ownership: %w", roleName, err)
	}
	if ownsTenantTable {
		return fmt.Errorf("database role %q must not own tenant tables", roleName)
	}
	return nil
}

func (database *Database) Check(ctx context.Context) error {
	if database == nil ||
		database.runtimePool == nil ||
		database.identityPool == nil {
		return errors.New("database is not configured")
	}
	if !database.compatible.Load() {
		return errors.New("database schema or access policy is not compatible")
	}
	if err := database.runtimePool.Ping(ctx); err != nil {
		return fmt.Errorf("ping runtime database: %w", err)
	}
	if err := database.identityPool.Ping(ctx); err != nil {
		return fmt.Errorf("ping identity database: %w", err)
	}
	return nil
}