跳到主要内容

在 Ent Schema 中使用行级安全

PostgreSQL 中的行级安全（RLS）功能允许表根据用户角色实施策略，限制对行的访问或修改，从而增强 GRANT 提供的基础 SQL 标准权限。

一旦激活，对表的每次标准访问都必须遵守这些策略。如果表上未定义任何策略，则默认采用拒绝所有规则，即无法查看或更改任何行。这些策略可针对特定命令、角色或两者进行定制，从而实现对数据访问或更改的精细管理。

本指南介绍如何将行级安全（RLS）策略附加到您的 Ent 类型（对象），并配置 schema 迁移，以便使用 Atlas 将 RLS 和 Ent schema 作为单个迁移单元进行管理。

Atlas Pro 功能

本指南中使用的行级安全策略的 Atlas 支持仅对 Pro 用户可用。要使用此功能，请运行：

atlas login

安装 Atlas​

To install the latest release of Atlas, simply run one of the following commands in your terminal, or check out the Atlas website:

curl -sSf https://atlasgo.sh | sh

登录 Atlas​

$ atlas login a8m
您现已连接到 Atlas Cloud 上的 "a8m"。

复合 Schema​

ent/schema 包主要用于定义 Ent 类型（对象）、它们的字段、边和逻辑。表策略或任何其他数据库原生对象在 Ent 模型中没有表示形式。

为了扩展我们的 PostgreSQL schema 以包含我们的 Ent 类型及其策略，我们配置 Atlas 从复合 Schema 数据源读取 schema 状态。请按照以下步骤为您的项目配置此功能：

1. 让我们定义一个包含两种类型（表）的简单 schema：users 和 tenants：

ent/schema/tenant.go
// Tenant 保存 Tenant 实体的 schema 定义。
type Tenant struct {
    ent.Schema
}

// Tenant 的字段。
func (Tenant) Fields() []ent.Field {
    return []ent.Field{
        field.String("name"),
    }
}

// User 保存 User 实体的 schema 定义。
type User struct {
    ent.Schema
}

// User 的字段。
func (User) Fields() []ent.Field {
    return []ent.Field{
        field.String("name"),
        field.Int("tenant_id"),
    }
}

2. 现在，假设我们想基于 tenant_id 字段限制对 users 表的访问。我们可以通过在 users 表上定义行级安全（RLS）策略来实现这一点。以下是定义 RLS 策略的 SQL 代码：

schema.sql
--- 在 users 表上启用行级安全。
ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;

-- 创建一个策略，根据当前租户限制对 users 表中行的访问。
CREATE POLICY tenant_isolation ON "users"
    USING ("tenant_id" = current_setting('app.current_tenant')::integer);

3. 最后，我们创建一个简单的 atlas.hcl 配置文件，其中包含一个 composite_schema，该 schema 包含我们的 Ent schema 和在 schema.sql 中定义的自定义安全策略：

atlas.hcl
data "composite_schema" "app" {
  # 首先加载包含所有表的 ent schema。
  schema "public" {
    url = "ent://ent/schema"
  }
  # 然后，加载 RLS schema。
  schema "public" {
    url = "file://schema.sql"
  }
}

env "local" {
  src = data.composite_schema.app.url
  dev = "docker://postgres/15/dev?search_path=public"
}

使用​

设置好复合 schema 后，我们可以使用 atlas schema inspect 命令获取其表示形式，为其生成 schema 迁移，将其应用到数据库等。以下是一些帮助您开始使用 Atlas 的命令：

检查 Schema​

atlas schema inspect 命令通常用于检查数据库。但是，我们也可以使用它来检查我们的 composite_schema 并打印其 SQL 表示形式：

atlas schema inspect \
  --env local \
  --url env://src \
  --format '{{ sql . }}'

上述命令打印以下 SQL。注意，tenant_isolation 策略在 users 表之后的 schema 中定义：

-- 创建 "users" 表
CREATE TABLE "users" ("id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY, "name" character varying NOT NULL, "tenant_id" bigint NOT NULL, PRIMARY KEY ("id"));
-- 为 "users" 表启用行级安全
ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;
-- 创建策略 "tenant_isolation"
CREATE POLICY "tenant_isolation" ON "users" AS PERMISSIVE FOR ALL TO PUBLIC USING (tenant_id = (current_setting('app.current_tenant'::text))::integer);
-- 创建 "tenants" 表
CREATE TABLE "tenants" ("id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY, "name" character varying NOT NULL, PRIMARY KEY ("id"));

为 Schema 生成迁移​

要为 schema 生成迁移，请运行以下命令：

atlas migrate diff \
  --env local

注意，将创建一个新的迁移文件，其内容如下：

migrations/20240712090543.sql
-- 创建 "users" 表
CREATE TABLE "users" ("id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY, "name" character varying NOT NULL, "tenant_id" bigint NOT NULL, PRIMARY KEY ("id"));
-- 为 "users" 表启用行级安全
ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;
-- 创建策略 "tenant_isolation"
CREATE POLICY "tenant_isolation" ON "users" AS PERMISSIVE FOR ALL TO PUBLIC USING (tenant_id = (current_setting('app.current_tenant'::text))::integer);
-- 创建 "tenants" 表
CREATE TABLE "tenants" ("id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY, "name" character varying NOT NULL, PRIMARY KEY ("id"));

应用迁移​

要将上面生成的迁移应用到数据库，请运行以下命令：

atlas migrate apply \
  --env local \
  --url "postgres://postgres:pass@localhost:5432/database?search_path=public&sslmode=disable"
直接在数据库上应用 Schema

有时，需要在不生成迁移文件的情况下直接将 schema 应用到数据库。例如，在试验 schema 更改、启动数据库进行测试等情况下。在这种情况下，您可以使用以下命令将 schema 直接应用到数据库：

atlas schema apply \
  --env local \
  --url "postgres://postgres:pass@localhost:5432/database?search_path=public&sslmode=disable"

或者，使用 Atlas Go SDK：

ac, err := atlasexec.NewClient(".", "atlas")
if err != nil {
    log.Fatalf("failed to initialize client: %w", err)
}
// 自动使用所需 schema 更新数据库。
// 另一个选项是手动使用 'migrate apply' 或 'schema apply'。
if _, err := ac.SchemaApply(ctx, &atlasexec.SchemaApplyParams{
    Env: "local",
    URL: "postgres://postgres:pass@localhost:5432/database?search_path=public&sslmode=disable",
    AutoApprove: true,
}); err != nil {
    log.Fatalf("failed to apply schema changes: %w", err)
}

代码示例​

设置好 Ent schema 和 RLS 策略后，我们可以打开一个 Ent 客户端，并传递我们正在处理的不同变更和查询的相关租户 ID。这确保数据库维护我们的 RLS 策略：

ctx1, ctx2 := sql.WithIntVar(ctx, "app.current_tenant", a8m.ID), sql.WithIntVar(ctx, "app.current_tenant", r3m.ID)
users1 := client.User.Query().AllX(ctx1)
// Users1 只能看到来自租户 a8m 的用户。
users2 := client.User.Query().AllX(ctx2)
// Users2 只能看到来自租户 r3m 的用户。
实际示例

在实际应用中，用户可以使用钩子和拦截器根据用户上下文设置 app.current_tenant 变量。

本指南的代码可以在 GitHub 上找到。