在 Go 中构建可扩展的多租户应用程序【译】

package main

type Tenant struct {
	ID   uint   `gorm:"primaryKey" json:"id"`
	Name string `json:"name"`
}

type Customer struct {
	ID       uint   `gorm:"primaryKey" json:"id"`
	Name     string `json:"name"`
	TenantID uint   `json:"tenant_id"`
}

func (s *Server) customersHandler(w http.ResponseWriter, r *http.Request) {
	tid, err := s.getTenantID(r)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte("invalid tenant"))
		return
	}
	var customers []Customer
	if err := s.db.
		Where("tenant_id = ?", tid).
		Find(&customers).Error; err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	json.NewEncoder(w).Encode(customers)
}

type OrderSumResult struct {
	CustomerID uint   `json:"customer_id"`
	Customer   string `json:"customer"`
	OrderCount int    `json:"order_count"`
}

func (s *Server) orderSumsHandler(w http.ResponseWriter, r *http.Request) {
	tid, err := s.getTenantID(r)
	if err != nil {
		w.WriteHeader(http.StatusUnauthorized)
		w.Write([]byte("unauthenticated"))
		return
	}
	var results []OrderSumResult
	err = s.db.Raw(`
		SELECT c.id as customer_id, c.name as customer, COUNT(o.id) as order_count
		FROM customers c
		LEFT JOIN orders o ON c.id = o.customer_id AND o.tenant_id = c.tenant_id
		WHERE c.tenant_id = ?
		GROUP BY c.id, c.name
	`, tid).Scan(&results).Error
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("query error"))
		return
	}
	json.NewEncoder(w).Encode(results)
}

package viewer

type middleware struct {
	*ent.Client
}

func Middleware(c *ent.Client) func(http.Handler) http.Handler {
	m := &middleware{c}
	return m.handle
}

func (m *middleware) handle(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		tenant, err := m.tenant(req)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		req = req.WithContext(
			NewContext(
				req.Context(),
				UserViewer{T: tenant},
			),
		)
		next.ServeHTTP(w, req)
	})
}

func (m *middleware) tenant(r *http.Request) (*ent.Tenant, error) {
	tid := r.Header.Get(Header)
	if tid == "" {
		return nil, http.ErrNoLocation
	}
	id, err := strconv.Atoi(tid)
	if err != nil {
		return nil, err
	}
	return m.Tenant.Get(r.Context(), id)
}

// NewServer creates a new HTTP server with ent client and returns http.Handler
func NewServer(client *ent.Client) http.Handler {
	s := &Server{client: client}
	r := chi.NewRouter()
	r.Use(middleware.Recoverer)
	r.Use(viewer.Middleware(client))

	r.Get("/customers", s.GetCustomers)
	r.Get("/products", s.GetProducts)
	r.Get("/orders", s.GetOrders)
	r.Get("/order-sums", s.orderSumsHandler)

	return r
}

package schema

import (
	"entgo.io/ent"
	"entgo.io/ent/schema/edge"
	"entgo.io/ent/schema/field"
)

// Tenant holds the schema definition for the Tenant entity.
type Tenant struct {
	ent.Schema
}

// Fields of the Tenant.
func (Tenant) Fields() []ent.Field {
	return []ent.Field{
		field.String("name").Unique(),
	}
}

// Edges of the Tenant.
func (Tenant) Edges() []ent.Edge {
	return []ent.Edge{
		edge.To("customers", Customer.Type),
	}
}

package schema

import (
	"entgo.io/ent"
	"entgo.io/ent/schema/edge"
	"entgo.io/ent/schema/field"
)

// Customer holds the schema definition for the Customer entity.
type Customer struct {
	ent.Schema
}

func (Customer) Fields() []ent.Field {
	return []ent.Field{
		field.String("name"),
	}
}

func (Customer) Edges() []ent.Edge {
	return []ent.Edge{
		edge.To("orders", Order.Type),
	}
}

func (Customer) Mixin() []ent.Mixin {
	return []ent.Mixin{TenantMixin{}}
}

package schema

import (
	"entgo.io/ent"
	"entgo.io/ent/schema/edge"
	"entgo.io/ent/schema/field"
	"entgo.io/ent/schema/mixin"
	"gophercon/entpriv/ent/rule"
)

// TenantMixin for embedding the tenant info in different schemas.
type TenantMixin struct {
	mixin.Schema
}

// Fields for all schemas that embed TenantMixin.
func (TenantMixin) Fields() []ent.Field {
	return []ent.Field{
		field.Int("tenant_id").
			Immutable(),
	}
}

// Edges for all schemas that embed TenantMixin.
func (TenantMixin) Edges() []ent.Edge {
	return []ent.Edge{
		edge.To("tenant", Tenant.Type).
			Field("tenant_id").
			Unique().
			Required().
			Immutable(),
	}
}

func (TenantMixin) Policy() ent.Policy {
	return rule.FilterTenantRule()
}

package rule

import (
	"context"
	"entgo.io/ent/entql"
	"gophercon/entpriv/ent/privacy"
	"gophercon/entpriv/viewer"
)

// FilterTenantRule is a query/mutation rule that filters out entities that are not in the tenant.
func FilterTenantRule() privacy.QueryMutationRule {
	// TenantsFilter is an interface to wrap WhereHasTenantWith()
	// predicate that is used by both `Group` and `User` schemas.
	type TenantsFilter interface {
		WhereTenantID(entql.IntP)
	}
	return privacy.FilterFunc(func(ctx context.Context, f privacy.Filter) error {
		view := viewer.FromContext(ctx)
		tid, ok := view.Tenant()
		if !ok {
			return privacy.Denyf("missing tenant information in viewer")
		}
		tf, ok := f.(TenantsFilter)
		if !ok {
			return privacy.Denyf("unexpected filter type %T", f)
		}
		// Make sure that a tenant reads only entities that have an edge to it.
		tf.WhereTenantID(entql.IntEQ(tid))
		// Skip to the next privacy rule (equivalent to return nil).
		return privacy.Skip
	})
}

// GetCustomers handles GET /customers request
func (s *Server) GetCustomers(w http.ResponseWriter, r *http.Request) {
	customers, err := s.client.Customer.Query().All(r.Context())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(customers)
}

--- Enable row-level security on the users table.
ALTER TABLE "customers" ENABLE ROW LEVEL SECURITY;

-- Create a policy that restricts access to
-- rows in the users table based on the current tenant.
CREATE POLICY tenant_isolation ON "users"
USING (
  "tenant_id" = current_setting('app.current_tenant')::integer
);

ctx := sql.WithIntVar(
    ctx, "app.current_tenant",
    tenant1.ID,
)

// Get only tenant1 customers
users := client.Customer.Query().AllX(ctx)

type middleware struct {
	url     *url.URL
	dialect string

	clients map[string]*ent.Client
	sync.Mutex
}

func (m *middleware) handle(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		client, err := m.tenantClient(req)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		ctx := ent.NewContext(req.Context(), client)
		next.ServeHTTP(w, req.WithContext(ctx))
	})
}

func (m *middleware) tenantClient(r *http.Request) (*ent.Client, error) {
	m.Lock()
	defer m.Unlock()
	name := r.Header.Get(Header)
	if name == "" {
		return nil, fmt.Errorf("missing tenant ID")
	}
	client, ok := m.clients[name]
	if !ok {
		return ent.Open(m.dialect, m.urlFor(name))
	}
	return client, nil
}

func (m *middleware) urlFor(t string) string {
	u := *m.url
	q := u.Query()
	q.Set("search_path", t)
	u.RawQuery = q.Encode()
	return u.String()
}

data "sql" "tenants" {
  url = var.url
  query = <<EOS
SELECT `schema_name`
  FROM `information_schema`.`schemata`
  WHERE `schema_name` LIKE ?
EOS
  args = [local.pattern]
}

env "prod" {
  for_each = toset(data.sql.tenants.values)
  url      = urlsetpath(var.url, each.value)
}