snapshot of https://github.com/ent/ent.git (master, 2026-09-15; blobs > 131072B omitted)

dialect/dialect.go

@@ -0,0 +1,208 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package dialect
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"fmt"
+	"log"
+
+	"github.com/google/uuid"
+)
+
+// Dialect names for external usage.
+const (
+	MySQL    = "mysql"
+	SQLite   = "sqlite3"
+	Postgres = "postgres"
+	Gremlin  = "gremlin"
+)
+
+// ExecQuerier wraps the 2 database operations.
+type ExecQuerier interface {
+	// Exec executes a query that does not return records. For example, in SQL, INSERT or UPDATE.
+	// It scans the result into the pointer v. For SQL drivers, it is dialect/sql.Result.
+	Exec(ctx context.Context, query string, args, v any) error
+	// Query executes a query that returns rows, typically a SELECT in SQL.
+	// It scans the result into the pointer v. For SQL drivers, it is *dialect/sql.Rows.
+	Query(ctx context.Context, query string, args, v any) error
+}
+
+// Driver is the interface that wraps all necessary operations for ent clients.
+type Driver interface {
+	ExecQuerier
+	// Tx starts and returns a new transaction.
+	// The provided context is used until the transaction is committed or rolled back.
+	Tx(context.Context) (Tx, error)
+	// Close closes the underlying connection.
+	Close() error
+	// Dialect returns the dialect name of the driver.
+	Dialect() string
+}
+
+// Tx wraps the Exec and Query operations in transaction.
+type Tx interface {
+	ExecQuerier
+	driver.Tx
+}
+
+type nopTx struct {
+	Driver
+}
+
+func (nopTx) Commit() error   { return nil }
+func (nopTx) Rollback() error { return nil }
+
+// NopTx returns a Tx with a no-op Commit / Rollback methods wrapping
+// the provided Driver d.
+func NopTx(d Driver) Tx {
+	return nopTx{d}
+}
+
+// DebugDriver is a driver that logs all driver operations.
+type DebugDriver struct {
+	Driver                               // underlying driver.
+	log    func(context.Context, ...any) // log function. defaults to log.Println.
+}
+
+// Debug gets a driver and an optional logging function, and returns
+// a new debugged-driver that prints all outgoing operations.
+func Debug(d Driver, logger ...func(...any)) Driver {
+	logf := log.Println
+	if len(logger) == 1 {
+		logf = logger[0]
+	}
+	drv := &DebugDriver{d, func(_ context.Context, v ...any) { logf(v...) }}
+	return drv
+}
+
+// DebugWithContext gets a driver and a logging function, and returns
+// a new debugged-driver that prints all outgoing operations with context.
+func DebugWithContext(d Driver, logger func(context.Context, ...any)) Driver {
+	drv := &DebugDriver{d, logger}
+	return drv
+}
+
+// Exec logs its params and calls the underlying driver Exec method.
+func (d *DebugDriver) Exec(ctx context.Context, query string, args, v any) error {
+	d.log(ctx, fmt.Sprintf("driver.Exec: query=%v args=%v", query, args))
+	return d.Driver.Exec(ctx, query, args, v)
+}
+
+// ExecContext logs its params and calls the underlying driver ExecContext method if it is supported.
+func (d *DebugDriver) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
+	drv, ok := d.Driver.(interface {
+		ExecContext(context.Context, string, ...any) (sql.Result, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Driver.ExecContext is not supported")
+	}
+	d.log(ctx, fmt.Sprintf("driver.ExecContext: query=%v args=%v", query, args))
+	return drv.ExecContext(ctx, query, args...)
+}
+
+// Query logs its params and calls the underlying driver Query method.
+func (d *DebugDriver) Query(ctx context.Context, query string, args, v any) error {
+	d.log(ctx, fmt.Sprintf("driver.Query: query=%v args=%v", query, args))
+	return d.Driver.Query(ctx, query, args, v)
+}
+
+// QueryContext logs its params and calls the underlying driver QueryContext method if it is supported.
+func (d *DebugDriver) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
+	drv, ok := d.Driver.(interface {
+		QueryContext(context.Context, string, ...any) (*sql.Rows, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Driver.QueryContext is not supported")
+	}
+	d.log(ctx, fmt.Sprintf("driver.QueryContext: query=%v args=%v", query, args))
+	return drv.QueryContext(ctx, query, args...)
+}
+
+// Tx adds an log-id for the transaction and calls the underlying driver Tx command.
+func (d *DebugDriver) Tx(ctx context.Context) (Tx, error) {
+	tx, err := d.Driver.Tx(ctx)
+	if err != nil {
+		return nil, err
+	}
+	id := uuid.New().String()
+	d.log(ctx, fmt.Sprintf("driver.Tx(%s): started", id))
+	return &DebugTx{tx, id, d.log, ctx}, nil
+}
+
+// BeginTx adds an log-id for the transaction and calls the underlying driver BeginTx command if it is supported.
+func (d *DebugDriver) BeginTx(ctx context.Context, opts *sql.TxOptions) (Tx, error) {
+	drv, ok := d.Driver.(interface {
+		BeginTx(context.Context, *sql.TxOptions) (Tx, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Driver.BeginTx is not supported")
+	}
+	tx, err := drv.BeginTx(ctx, opts)
+	if err != nil {
+		return nil, err
+	}
+	id := uuid.New().String()
+	d.log(ctx, fmt.Sprintf("driver.BeginTx(%s): started", id))
+	return &DebugTx{tx, id, d.log, ctx}, nil
+}
+
+// DebugTx is a transaction implementation that logs all transaction operations.
+type DebugTx struct {
+	Tx                                // underlying transaction.
+	id  string                        // transaction logging id.
+	log func(context.Context, ...any) // log function. defaults to fmt.Println.
+	ctx context.Context               // underlying transaction context.
+}
+
+// Exec logs its params and calls the underlying transaction Exec method.
+func (d *DebugTx) Exec(ctx context.Context, query string, args, v any) error {
+	d.log(ctx, fmt.Sprintf("Tx(%s).Exec: query=%v args=%v", d.id, query, args))
+	return d.Tx.Exec(ctx, query, args, v)
+}
+
+// ExecContext logs its params and calls the underlying transaction ExecContext method if it is supported.
+func (d *DebugTx) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
+	drv, ok := d.Tx.(interface {
+		ExecContext(context.Context, string, ...any) (sql.Result, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Tx.ExecContext is not supported")
+	}
+	d.log(ctx, fmt.Sprintf("Tx(%s).ExecContext: query=%v args=%v", d.id, query, args))
+	return drv.ExecContext(ctx, query, args...)
+}
+
+// Query logs its params and calls the underlying transaction Query method.
+func (d *DebugTx) Query(ctx context.Context, query string, args, v any) error {
+	d.log(ctx, fmt.Sprintf("Tx(%s).Query: query=%v args=%v", d.id, query, args))
+	return d.Tx.Query(ctx, query, args, v)
+}
+
+// QueryContext logs its params and calls the underlying transaction QueryContext method if it is supported.
+func (d *DebugTx) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
+	drv, ok := d.Tx.(interface {
+		QueryContext(context.Context, string, ...any) (*sql.Rows, error)
+	})
+	if !ok {
+		return nil, fmt.Errorf("Tx.QueryContext is not supported")
+	}
+	d.log(ctx, fmt.Sprintf("Tx(%s).QueryContext: query=%v args=%v", d.id, query, args))
+	return drv.QueryContext(ctx, query, args...)
+}
+
+// Commit logs this step and calls the underlying transaction Commit method.
+func (d *DebugTx) Commit() error {
+	d.log(d.ctx, fmt.Sprintf("Tx(%s): committed", d.id))
+	return d.Tx.Commit()
+}
+
+// Rollback logs this step and calls the underlying transaction Rollback method.
+func (d *DebugTx) Rollback() error {
+	d.log(d.ctx, fmt.Sprintf("Tx(%s): rollbacked", d.id))
+	return d.Tx.Rollback()
+}

dialect/entsql/annotation.go

@@ -0,0 +1,857 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package entsql
+
+import (
+	"errors"
+	"fmt"
+
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/schema"
+)
+
+// Annotation is a builtin schema annotation for attaching
+// SQL metadata to schema objects for both codegen and runtime.
+type Annotation struct {
+	// The Schema option allows setting the schema which the table belongs to.
+	// Note, this option is no-op for Ent default migration engine. However, schema
+	// extensions (like Atlas) can accept this option and implement it accordingly.
+	//
+	//	entsql.Annotation{
+	//		Schema: "public",
+	//	}
+	//
+	Schema string `json:"schema,omitempty"`
+
+	// The Table option allows overriding the default table
+	// name that is generated by ent. For example:
+	//
+	//	entsql.Annotation{
+	//		Table: "Users",
+	//	}
+	//
+	Table string `json:"table,omitempty"`
+
+	// Charset defines the character-set of the table. For example:
+	//
+	//	entsql.Annotation{
+	//		Charset: "utf8mb4",
+	//	}
+	//
+	Charset string `json:"charset,omitempty"`
+
+	// Collation defines the collation of the table (a set of rules for comparing
+	// characters in a character set). For example:
+	//
+	//	entsql.Annotation{
+	//		Collation: "utf8mb4_bin",
+	//	}
+	//
+	Collation string `json:"collation,omitempty"`
+
+	// Default specifies a literal default value of a column. Note that using
+	// this option overrides the default behavior of the code-generation.
+	//
+	//	entsql.Annotation{
+	//		Default: `{"key":"value"}`,
+	//	}
+	//
+	Default string `json:"default,omitempty"`
+
+	// DefaultExpr specifies an expression default value of a column. Using this option,
+	// users can define custom expressions to be set as database default values. Note that
+	// using this option overrides the default behavior of the code-generation.
+	//
+	//	entsql.Annotation{
+	//		DefaultExpr: "CURRENT_TIMESTAMP",
+	//	}
+	//
+	//	entsql.Annotation{
+	//		DefaultExpr: "uuid_generate_v4()",
+	//	}
+	//
+	//	entsql.Annotation{
+	//		DefaultExpr: "(a + b)",
+	//	}
+	//
+	DefaultExpr string `json:"default_expr,omitempty"`
+
+	// DefaultExpr specifies an expression default value of a column per dialect.
+	// See, DefaultExpr for full doc.
+	//
+	//	entsql.Annotation{
+	//		DefaultExprs: map[string]string{
+	//			dialect.MySQL:    "uuid()",
+	//			dialect.Postgres: "uuid_generate_v4",
+	//		}
+	//
+	DefaultExprs map[string]string `json:"default_exprs,omitempty"`
+
+	// Options defines the additional table options. For example:
+	//
+	//	entsql.Annotation{
+	//		Options: "ENGINE = INNODB",
+	//	}
+	//
+	Options string `json:"options,omitempty"`
+
+	// Size defines the column size in the generated schema. For example:
+	//
+	//	entsql.Annotation{
+	//		Size: 128,
+	//	}
+	//
+	Size int64 `json:"size,omitempty"`
+
+	// WithComments specifies whether fields' comments should
+	// be stored in the database schema as column comments.
+	//
+	//  withCommentsEnabled := true
+	//	entsql.WithComments{
+	//		WithComments: &withCommentsEnabled,
+	//	}
+	//
+	WithComments *bool `json:"with_comments,omitempty"`
+
+	// Incremental defines the auto-incremental behavior of a column. For example:
+	//
+	//  incrementalEnabled := true
+	//  entsql.Annotation{
+	//      Incremental: &incrementalEnabled,
+	//  }
+	//
+	// By default, this value is nil defaulting to whatever best fits each scenario.
+	//
+	Incremental *bool `json:"incremental,omitempty"`
+
+	// IncrementStart defines the auto-incremental start value of a column. For example:
+	//
+	//  incrementStart := 100
+	//  entsql.Annotation{
+	//      IncrementStart: &incrementStart,
+	//  }
+	//
+	// By default, this value is nil defaulting to whatever the database settings are.
+	//
+	IncrementStart *int `json:"increment_start,omitempty"`
+
+	// OnDelete specifies a custom referential action for DELETE operations on parent
+	// table that has matching rows in the child table.
+	//
+	// For example, in order to delete rows from the parent table and automatically delete
+	// their matching rows in the child table, pass the following annotation:
+	//
+	//	entsql.Annotation{
+	//		OnDelete: entsql.Cascade,
+	//	}
+	//
+	OnDelete ReferenceOption `json:"on_delete,omitempty"`
+
+	// Check allows injecting custom "DDL" for setting an unnamed "CHECK" clause in "CREATE TABLE".
+	//
+	//	entsql.Annotation{
+	//		Check: "age < 10",
+	//	}
+	//
+	Check string `json:"check,omitempty"`
+
+	// Checks allows injecting custom "DDL" for setting named "CHECK" clauses in "CREATE TABLE".
+	//
+	//	entsql.Annotation{
+	//		Checks: map[string]string{
+	//			"valid_discount": "price > discount_price",
+	//		},
+	//	}
+	//
+	Checks map[string]string `json:"checks,omitempty"`
+
+	// Skip indicates that the field or the schema is skipped/ignored during
+	// migration (e.g., defined externally).
+	//
+	//	entsql.Annotation{
+	//		Skip: true,
+	//	}
+	//
+	Skip bool `json:"skip,omitempty"`
+
+	// ViewAs allows defining a view for the schema. For example:
+	//
+	//	entsql.Annotation{
+	//		View: "SELECT name FROM users",
+	//	}
+	ViewAs string `json:"view_as,omitempty"`
+
+	// ViewFor allows defining a view for the schema per dialect. For example:
+	//
+	//	entsql.Annotation{
+	//		ViewFor: map[string]string{
+	//			dialect.MySQL:    "...",
+	//			dialect.Postgres: "...",
+	//		},
+	//	}
+	ViewFor map[string]string `json:"view_for,omitempty"`
+
+	// error occurs during annotation build. This field is not
+	// serialized to JSON and used only by the codegen loader.
+	err error
+}
+
+// Name describes the annotation name.
+func (Annotation) Name() string {
+	return "EntSQL"
+}
+
+// The Schema option allows setting the schema which the table belongs to.
+// Note, this option is no-op for Ent default migration engine. However, schema
+// extensions (like Atlas) can accept this option and implement it accordingly.
+//
+//	func (T) Annotations() []schema.Annotation {
+//		return []schema.Annotation{
+//			entsql.Schema("public"),
+//		}
+//	}
+func Schema(s string) *Annotation {
+	return &Annotation{
+		Schema: s,
+	}
+}
+
+// The Table option allows overriding the default table
+// name that is generated by ent. For example:
+//
+//	func (T) Annotations() []schema.Annotation {
+//		return []schema.Annotation{
+//			entsql.Table("Users"),
+//		}
+//	}
+func Table(t string) *Annotation {
+	return &Annotation{
+		Table: t,
+	}
+}
+
+// SchemaTable allows setting both schema and table name in one annotation.
+func SchemaTable(s, t string) *Annotation {
+	return &Annotation{
+		Schema: s,
+		Table:  t,
+	}
+}
+
+// Check allows injecting custom "DDL" for setting an unnamed "CHECK" clause in "CREATE TABLE".
+//
+//	entsql.Annotation{
+//		Check: "(`age` < 10)",
+//	}
+func Check(c string) *Annotation {
+	return &Annotation{
+		Check: c,
+	}
+}
+
+// Checks allows injecting custom "DDL" for setting named "CHECK" clauses in "CREATE TABLE".
+//
+//	entsql.Annotation{
+//		Checks: map[string]string{
+//			"valid_discount": "price > discount_price",
+//		},
+//	}
+func Checks(c map[string]string) *Annotation {
+	return &Annotation{
+		Checks: c,
+	}
+}
+
+// Skip indicates that the field or the schema is skipped/ignored during
+// migration (e.g., defined externally).
+func Skip() *Annotation {
+	return &Annotation{Skip: true}
+}
+
+// View specifies the definition of a view.
+func View(as string) *Annotation {
+	return &Annotation{ViewAs: as}
+}
+
+// ViewFor specifies the definition of a view.
+func ViewFor(dialect string, as func(*sql.Selector)) *Annotation {
+	b := sql.Dialect(dialect).Select()
+	as(b)
+	switch q, args := b.Query(); {
+	case len(args) > 0:
+		return &Annotation{
+			err: fmt.Errorf("entsql: view query should not contain arguments. got: %d", len(args)),
+		}
+	case q == "":
+		return &Annotation{
+			err: errors.New("entsql: view query is empty"),
+		}
+	case b.Err() != nil:
+		return &Annotation{
+			err: b.Err(),
+		}
+	default:
+		return &Annotation{
+			ViewFor: map[string]string{dialect: q},
+		}
+	}
+}
+
+// Default specifies a literal default value of a column. Note that using
+// this option overrides the default behavior of the code-generation.
+//
+//	entsql.Annotation{
+//		Default: `{"key":"value"}`,
+//	}
+func Default(literal string) *Annotation {
+	return &Annotation{
+		Default: literal,
+	}
+}
+
+// DefaultExpr specifies an expression default value for the annotated column.
+// Using this option, users can define custom expressions to be set as database
+// default values.Note that using this option overrides the default behavior of
+// the code-generation.
+//
+//	field.UUID("id", uuid.Nil).
+//		Default(uuid.New).
+//		Annotations(
+//			entsql.DefaultExpr("uuid_generate_v4()"),
+//		)
+func DefaultExpr(expr string) *Annotation {
+	return &Annotation{
+		DefaultExpr: expr,
+	}
+}
+
+// DefaultExprs specifies an expression default value for the annotated
+// column per dialect. See, DefaultExpr for full doc.
+//
+//	field.UUID("id", uuid.Nil).
+//		Default(uuid.New).
+//		Annotations(
+//			entsql.DefaultExprs(map[string]string{
+//				dialect.MySQL:    "uuid()",
+//				dialect.Postgres: "uuid_generate_v4()",
+//			}),
+//		)
+func DefaultExprs(exprs map[string]string) *Annotation {
+	return &Annotation{
+		DefaultExprs: exprs,
+	}
+}
+
+// WithComments specifies whether fields' comments should
+// be stored in the database schema as column comments.
+//
+//	func (T) Annotations() []schema.Annotation {
+//		return []schema.Annotation{
+//			entsql.WithComments(true),
+//		}
+//	}
+func WithComments(b bool) *Annotation {
+	return &Annotation{
+		WithComments: &b,
+	}
+}
+
+// OnDelete specifies a custom referential action for DELETE operations on parent
+// table that has matching rows in the child table.
+//
+// For example, in order to delete rows from the parent table and automatically delete
+// their matching rows in the child table, pass the following annotation:
+//
+//	func (T) Annotations() []schema.Annotation {
+//		return []schema.Annotation{
+//			entsql.OnDelete(entsql.Cascade),
+//		}
+//	}
+func OnDelete(opt ReferenceOption) *Annotation {
+	return &Annotation{
+		OnDelete: opt,
+	}
+}
+
+// IncrementStart specifies the starting value for auto-increment columns.
+//
+// For example, in order to define the starting value for auto-increment to be 100:
+//
+//	func (T) Annotations() []schema.Annotation {
+//		return []schema.Annotation{
+//			entsql.IncrementStart(100),
+//		}
+//	}
+func IncrementStart(i int) *Annotation {
+	return &Annotation{
+		IncrementStart: &i,
+	}
+}
+
+// Merge implements the schema.Merger interface.
+func (a Annotation) Merge(other schema.Annotation) schema.Annotation {
+	var ant Annotation
+	switch other := other.(type) {
+	case Annotation:
+		ant = other
+	case *Annotation:
+		if other != nil {
+			ant = *other
+		}
+	default:
+		return a
+	}
+	if s := ant.Schema; s != "" {
+		a.Schema = s
+	}
+	if t := ant.Table; t != "" {
+		a.Table = t
+	}
+	if c := ant.Charset; c != "" {
+		a.Charset = c
+	}
+	if c := ant.Collation; c != "" {
+		a.Collation = c
+	}
+	if d := ant.Default; d != "" {
+		a.Default = d
+	}
+	if d := ant.DefaultExpr; d != "" {
+		a.DefaultExpr = d
+	}
+	if d := ant.DefaultExprs; d != nil {
+		if a.DefaultExprs == nil {
+			a.DefaultExprs = make(map[string]string)
+		}
+		for dialect, x := range d {
+			a.DefaultExprs[dialect] = x
+		}
+	}
+	if o := ant.Options; o != "" {
+		a.Options = o
+	}
+	if s := ant.Size; s != 0 {
+		a.Size = s
+	}
+	if b := ant.WithComments; b != nil {
+		a.WithComments = b
+	}
+	if i := ant.Incremental; i != nil {
+		a.Incremental = i
+	}
+	if i := ant.IncrementStart; i != nil {
+		a.IncrementStart = i
+	}
+	if od := ant.OnDelete; od != "" {
+		a.OnDelete = od
+	}
+	if c := ant.Check; c != "" {
+		a.Check = c
+	}
+	if checks := ant.Checks; len(checks) > 0 {
+		if a.Checks == nil {
+			a.Checks = make(map[string]string)
+		}
+		for name, check := range checks {
+			a.Checks[name] = check
+		}
+	}
+	if ant.Skip {
+		a.Skip = true
+	}
+	if v := ant.ViewAs; v != "" {
+		a.ViewAs = v
+	}
+	if vf := ant.ViewFor; len(vf) > 0 {
+		if a.ViewFor == nil {
+			a.ViewFor = make(map[string]string)
+		}
+		for dialect, view := range vf {
+			a.ViewFor[dialect] = view
+		}
+	}
+	if ant.err != nil {
+		a.err = errors.Join(a.err, ant.err)
+	}
+	return a
+}
+
+// Err returns the error that occurred during annotation build, if any.
+func (a Annotation) Err() error {
+	return a.err
+}
+
+var _ interface {
+	schema.Annotation
+	schema.Merger
+} = (*Annotation)(nil)
+
+// ReferenceOption for constraint actions.
+type ReferenceOption string
+
+// Reference options (actions) specified by ON UPDATE and ON DELETE
+// subclauses of the FOREIGN KEY clause.
+const (
+	NoAction   ReferenceOption = "NO ACTION"
+	Restrict   ReferenceOption = "RESTRICT"
+	Cascade    ReferenceOption = "CASCADE"
+	SetNull    ReferenceOption = "SET NULL"
+	SetDefault ReferenceOption = "SET DEFAULT"
+)
+
+// IndexAnnotation is a builtin schema annotation for attaching
+// SQL metadata to schema indexes for both codegen and runtime.
+type IndexAnnotation struct {
+	// Prefix defines a column prefix for a single string column index.
+	// In MySQL, the following annotation maps to:
+	//
+	//	index.Fields("column").
+	//		Annotation(entsql.Prefix(100))
+	//
+	//	CREATE INDEX `table_column` ON `table`(`column`(100))
+	//
+	Prefix uint
+
+	// PrefixColumns defines column prefixes for a multi-column index.
+	// In MySQL, the following annotation maps to:
+	//
+	//	index.Fields("c1", "c2", "c3").
+	//		Annotation(
+	//			entsql.PrefixColumn("c1", 100),
+	//			entsql.PrefixColumn("c2", 200),
+	//		)
+	//
+	//	CREATE INDEX `table_c1_c2_c3` ON `table`(`c1`(100), `c2`(200), `c3`)
+	//
+	PrefixColumns map[string]uint
+
+	// Desc defines the DESC clause for a single column index.
+	// In MySQL, the following annotation maps to:
+	//
+	//	index.Fields("column").
+	//		Annotation(entsql.Desc())
+	//
+	//	CREATE INDEX `table_column` ON `table`(`column` DESC)
+	//
+	Desc bool
+
+	// DescColumns defines the DESC clause for columns in multi-column index.
+	// In MySQL, the following annotation maps to:
+	//
+	//	index.Fields("c1", "c2", "c3").
+	//		Annotation(
+	//			entsql.DescColumns("c1", "c2"),
+	//		)
+	//
+	//	CREATE INDEX `table_c1_c2_c3` ON `table`(`c1` DESC, `c2` DESC, `c3`)
+	//
+	DescColumns map[string]bool
+
+	// IncludeColumns defines the INCLUDE clause for the index.
+	// Works only in Postgres and its definition is as follows:
+	//
+	//	index.Fields("c1").
+	//		Annotation(
+	//			entsql.IncludeColumns("c2"),
+	//		)
+	//
+	//	CREATE INDEX "table_column" ON "table"("c1") INCLUDE ("c2")
+	//
+	IncludeColumns []string
+
+	// Type defines the type of the index.
+	// In MySQL, the following annotation maps to:
+	//
+	//	index.Fields("c1").
+	//		Annotation(
+	//			entsql.IndexType("FULLTEXT"),
+	//		)
+	//
+	//	CREATE FULLTEXT INDEX `table_c1` ON `table`(`c1`)
+	//
+	Type string
+
+	// Types is like the Type option but allows mapping an index-type per dialect.
+	//
+	//	index.Fields("c1").
+	//		Annotation(
+	//			entsql.IndexTypes(map[string]string{
+	//				dialect.MySQL:		"FULLTEXT",
+	//				dialect.Postgres:	"GIN",
+	//			}),
+	//		)
+	//
+	Types map[string]string
+
+	// OpClass defines the operator class for a single string column index.
+	// In PostgreSQL, the following annotation maps to:
+	//
+	//	index.Fields("column").
+	//		Annotation(
+	//			entsql.IndexType("BRIN"),
+	//			entsql.OpClass("int8_bloom_ops"),
+	//		)
+	//
+	//	CREATE INDEX "table_column" ON "table" USING BRIN ("column" int8_bloom_ops)
+	//
+	OpClass string
+
+	// OpClassColumns defines operator-classes for a multi-column index.
+	// In PostgreSQL, the following annotation maps to:
+	//
+	//	index.Fields("c1", "c2", "c3").
+	//		Annotation(
+	//			entsql.IndexType("BRIN"),
+	//			entsql.OpClassColumn("c1", "int8_bloom_ops"),
+	//			entsql.OpClassColumn("c2", "int8_minmax_multi_ops(values_per_range=8)"),
+	//		)
+	//
+	//	CREATE INDEX "table_column" ON "table" USING BRIN ("c1" int8_bloom_ops, "c2" int8_minmax_multi_ops(values_per_range=8), "c3")
+	//
+	OpClassColumns map[string]string
+
+	// IndexWhere allows configuring partial indexes in SQLite and PostgreSQL.
+	// Read more: https://postgresql.org/docs/current/indexes-partial.html.
+	//
+	// Note that the `WHERE` clause should be defined exactly like it is
+	// stored in the database (i.e. normal form). Read more about this on
+	// the Atlas website: https://atlasgo.io/concepts/dev-database#diffing.
+	//
+	//	index.Fields("a").
+	//		Annotations(
+	//			entsql.IndexWhere("b AND c > 0"),
+	//		)
+	//	CREATE INDEX "table_a" ON "table"("a") WHERE (b AND c > 0)
+	Where string
+}
+
+// Prefix returns a new index annotation with a single string column index.
+// In MySQL, the following annotation maps to:
+//
+//	index.Fields("column").
+//		Annotation(entsql.Prefix(100))
+//
+//	CREATE INDEX `table_column` ON `table`(`column`(100))
+func Prefix(prefix uint) *IndexAnnotation {
+	return &IndexAnnotation{
+		Prefix: prefix,
+	}
+}
+
+// PrefixColumn returns a new index annotation with column prefix for
+// multi-column indexes. In MySQL, the following annotation maps to:
+//
+//	index.Fields("c1", "c2", "c3").
+//		Annotation(
+//			entsql.PrefixColumn("c1", 100),
+//			entsql.PrefixColumn("c2", 200),
+//		)
+//
+//	CREATE INDEX `table_c1_c2_c3` ON `table`(`c1`(100), `c2`(200), `c3`)
+func PrefixColumn(name string, prefix uint) *IndexAnnotation {
+	return &IndexAnnotation{
+		PrefixColumns: map[string]uint{
+			name: prefix,
+		},
+	}
+}
+
+// OpClass defines the operator class for a single string column index.
+// In PostgreSQL, the following annotation maps to:
+//
+//	index.Fields("column").
+//		Annotation(
+//			entsql.IndexType("BRIN"),
+//			entsql.OpClass("int8_bloom_ops"),
+//		)
+//
+//	CREATE INDEX "table_column" ON "table" USING BRIN ("column" int8_bloom_ops)
+func OpClass(op string) *IndexAnnotation {
+	return &IndexAnnotation{
+		OpClass: op,
+	}
+}
+
+// OpClassColumn returns a new index annotation with column operator
+// class for multi-column indexes. In PostgreSQL, the following annotation maps to:
+//
+//	index.Fields("c1", "c2", "c3").
+//		Annotation(
+//			entsql.IndexType("BRIN"),
+//			entsql.OpClassColumn("c1", "int8_bloom_ops"),
+//			entsql.OpClassColumn("c2", "int8_minmax_multi_ops(values_per_range=8)"),
+//		)
+//
+//	CREATE INDEX "table_column" ON "table" USING BRIN ("c1" int8_bloom_ops, "c2" int8_minmax_multi_ops(values_per_range=8), "c3")
+func OpClassColumn(name, op string) *IndexAnnotation {
+	return &IndexAnnotation{
+		OpClassColumns: map[string]string{
+			name: op,
+		},
+	}
+}
+
+// Desc returns a new index annotation with the DESC clause for a
+// single column index. In MySQL, the following annotation maps to:
+//
+//	index.Fields("column").
+//		Annotation(entsql.Desc())
+//
+//	CREATE INDEX `table_column` ON `table`(`column` DESC)
+func Desc() *IndexAnnotation {
+	return &IndexAnnotation{
+		Desc: true,
+	}
+}
+
+// DescColumns returns a new index annotation with the DESC clause attached to
+// the columns in the index. In MySQL, the following annotation maps to:
+//
+//	index.Fields("c1", "c2", "c3").
+//		Annotation(
+//			entsql.DescColumns("c1", "c2"),
+//		)
+//
+//	CREATE INDEX `table_c1_c2_c3` ON `table`(`c1` DESC, `c2` DESC, `c3`)
+func DescColumns(names ...string) *IndexAnnotation {
+	ant := &IndexAnnotation{
+		DescColumns: make(map[string]bool, len(names)),
+	}
+	for i := range names {
+		ant.DescColumns[names[i]] = true
+	}
+	return ant
+}
+
+// IncludeColumns defines the INCLUDE clause for the index.
+// Works only in Postgres and its definition is as follows:
+//
+//	index.Fields("c1").
+//		Annotation(
+//			entsql.IncludeColumns("c2"),
+//		)
+//
+//	CREATE INDEX "table_column" ON "table"("c1") INCLUDE ("c2")
+func IncludeColumns(names ...string) *IndexAnnotation {
+	return &IndexAnnotation{IncludeColumns: names}
+}
+
+// IndexType defines the type of the index.
+// In MySQL, the following annotation maps to:
+//
+//	index.Fields("c1").
+//		Annotation(
+//			entsql.IndexType("FULLTEXT"),
+//		)
+//
+//	CREATE FULLTEXT INDEX `table_c1` ON `table`(`c1`)
+func IndexType(t string) *IndexAnnotation {
+	return &IndexAnnotation{Type: t}
+}
+
+// IndexTypes is like the Type option but allows mapping an index-type per dialect.
+//
+//	index.Fields("c1").
+//		Annotations(
+//			entsql.IndexTypes(map[string]string{
+//				dialect.MySQL:    "FULLTEXT",
+//				dialect.Postgres: "GIN",
+//			}),
+//		)
+func IndexTypes(types map[string]string) *IndexAnnotation {
+	return &IndexAnnotation{Types: types}
+}
+
+// IndexWhere allows configuring partial indexes in SQLite and PostgreSQL.
+// Read more: https://postgresql.org/docs/current/indexes-partial.html.
+//
+// Note that the `WHERE` clause should be defined exactly like it is
+// stored in the database (i.e. normal form). Read more about this on the
+// Atlas website: https://atlasgo.io/concepts/dev-database#diffing.
+//
+//	index.Fields("a").
+//		Annotations(
+//			entsql.IndexWhere("b AND c > 0"),
+//		)
+//	CREATE INDEX "table_a" ON "table"("a") WHERE (b AND c > 0)
+func IndexWhere(pred string) *IndexAnnotation {
+	return &IndexAnnotation{Where: pred}
+}
+
+// Name describes the annotation name.
+func (IndexAnnotation) Name() string {
+	return "EntSQLIndexes"
+}
+
+// Merge implements the schema.Merger interface.
+func (a IndexAnnotation) Merge(other schema.Annotation) schema.Annotation {
+	var ant IndexAnnotation
+	switch other := other.(type) {
+	case IndexAnnotation:
+		ant = other
+	case *IndexAnnotation:
+		if other != nil {
+			ant = *other
+		}
+	default:
+		return a
+	}
+	if ant.Prefix != 0 {
+		a.Prefix = ant.Prefix
+	}
+	if ant.PrefixColumns != nil {
+		if a.PrefixColumns == nil {
+			a.PrefixColumns = make(map[string]uint)
+		}
+		for column, prefix := range ant.PrefixColumns {
+			a.PrefixColumns[column] = prefix
+		}
+	}
+	if ant.OpClass != "" {
+		a.OpClass = ant.OpClass
+	}
+	if ant.OpClassColumns != nil {
+		if a.OpClassColumns == nil {
+			a.OpClassColumns = make(map[string]string)
+		}
+		for column, op := range ant.OpClassColumns {
+			a.OpClassColumns[column] = op
+		}
+	}
+	if ant.Desc {
+		a.Desc = ant.Desc
+	}
+	if ant.DescColumns != nil {
+		if a.DescColumns == nil {
+			a.DescColumns = make(map[string]bool)
+		}
+		for column, desc := range ant.DescColumns {
+			a.DescColumns[column] = desc
+		}
+	}
+	if ant.IncludeColumns != nil {
+		a.IncludeColumns = append(a.IncludeColumns, ant.IncludeColumns...)
+	}
+	if ant.Type != "" {
+		a.Type = ant.Type
+	}
+	if ant.Types != nil {
+		if a.Types == nil {
+			a.Types = make(map[string]string)
+		}
+		for dialect, t := range ant.Types {
+			a.Types[dialect] = t
+		}
+	}
+	if ant.Where != "" {
+		a.Where = ant.Where
+	}
+	return a
+}
+
+var _ interface {
+	schema.Annotation
+	schema.Merger
+} = (*IndexAnnotation)(nil)

dialect/gremlin/client.go

@@ -0,0 +1,77 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+)
+
+// RoundTripper is an interface representing the ability to execute a
+// single gremlin transaction, obtaining the Response for a given Request.
+type RoundTripper interface {
+	RoundTrip(context.Context, *Request) (*Response, error)
+}
+
+// The RoundTripperFunc type is an adapter to allow the use of ordinary functions as Gremlin RoundTripper.
+type RoundTripperFunc func(context.Context, *Request) (*Response, error)
+
+// RoundTrip calls f(ctx, r).
+func (f RoundTripperFunc) RoundTrip(ctx context.Context, r *Request) (*Response, error) {
+	return f(ctx, r)
+}
+
+// Interceptor provides a hook to intercept the execution of a Gremlin Request.
+type Interceptor func(RoundTripper) RoundTripper
+
+// A Client is a gremlin client.
+type Client struct {
+	// Transport specifies the mechanism by which individual
+	// Gremlin requests are made.
+	Transport RoundTripper
+}
+
+// MaxResponseSize defines the maximum response size allowed.
+const MaxResponseSize = 2 << 20
+
+// NewClient creates a gremlin client from config and options.
+func NewClient(cfg Config, opt ...Option) (*Client, error) {
+	return cfg.Build(opt...)
+}
+
+// NewHTTPClient creates an http based gremlin client.
+func NewHTTPClient(url string, client *http.Client) (*Client, error) {
+	transport, err := NewHTTPTransport(url, client)
+	if err != nil {
+		return nil, err
+	}
+	return &Client{transport}, nil
+}
+
+// Do sends a gremlin request and returns a gremlin response.
+func (c Client) Do(ctx context.Context, req *Request) (*Response, error) {
+	rsp, err := c.Transport.RoundTrip(ctx, req)
+	if err == nil {
+		err = rsp.Err()
+	}
+
+	// If we got an error, and the context has been canceled,
+	// the context's error is probably more useful.
+	if err != nil && ctx.Err() != nil {
+		err = ctx.Err()
+	}
+	return rsp, err
+}
+
+// Query issues an eval request via the Do function.
+func (c Client) Query(ctx context.Context, query string) (*Response, error) {
+	return c.Do(ctx, NewEvalRequest(query))
+}
+
+// Queryf formats a query string and invokes Query.
+func (c Client) Queryf(ctx context.Context, format string, args ...any) (*Response, error) {
+	return c.Query(ctx, fmt.Sprintf(format, args...))
+}

dialect/gremlin/client_test.go

@@ -0,0 +1,93 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"io"
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestNewClient(t *testing.T) {
+	var cfg Config
+	cfg.Endpoint.URL, _ = url.Parse("http://gremlin-server/gremlin")
+	c, err := NewClient(cfg)
+	assert.NotNil(t, c)
+	assert.NoError(t, err)
+}
+
+type mockRoundTripper struct{ mock.Mock }
+
+func (m *mockRoundTripper) RoundTrip(ctx context.Context, req *Request) (*Response, error) {
+	args := m.Called(ctx, req)
+	return args.Get(0).(*Response), args.Error(1)
+}
+
+func TestClientRequest(t *testing.T) {
+	ctx := context.Background()
+	req, rsp := &Request{}, &Response{}
+
+	var m mockRoundTripper
+	m.On("RoundTrip", ctx, req).
+		Run(func(mock.Arguments) { rsp.Status.Code = StatusSuccess }).
+		Return(rsp, nil).
+		Once()
+	defer m.AssertExpectations(t)
+
+	response, err := Client{&m}.Do(context.Background(), req)
+	assert.NoError(t, err)
+	assert.Equal(t, rsp, response)
+}
+
+func TestClientResponseError(t *testing.T) {
+	rsp := &Response{}
+
+	var m mockRoundTripper
+	m.On("RoundTrip", mock.Anything, mock.Anything).
+		Run(func(mock.Arguments) { rsp.Status.Code = StatusServerError }).
+		Return(rsp, nil).
+		Once()
+	defer m.AssertExpectations(t)
+
+	_, err := Client{&m}.Do(context.Background(), nil)
+	assert.Error(t, err)
+}
+
+func TestClientCanceledContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	var m mockRoundTripper
+	m.On("RoundTrip", ctx, mock.Anything).
+		Run(func(mock.Arguments) { cancel() }).
+		Return(&Response{}, io.ErrUnexpectedEOF).
+		Once()
+	defer m.AssertExpectations(t)
+
+	_, err := Client{&m}.Query(ctx, "g.E()")
+	assert.EqualError(t, err, context.Canceled.Error())
+}
+
+func TestClientQuery(t *testing.T) {
+	rsp := &Response{}
+	rsp.Status.Code = StatusNoContent
+
+	var m mockRoundTripper
+	m.On("RoundTrip", mock.Anything, mock.Anything).
+		Run(func(args mock.Arguments) {
+			req := args.Get(1).(*Request)
+			assert.Equal(t, "g.V(1)", req.Arguments[ArgsGremlin])
+		}).
+		Return(rsp, nil).
+		Once()
+	defer m.AssertExpectations(t)
+
+	rsp, err := Client{&m}.Queryf(context.Background(), "g.V(%d)", 1)
+	assert.NotNil(t, rsp)
+	assert.NoError(t, err)
+}

dialect/gremlin/config.go

@@ -0,0 +1,82 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"fmt"
+	"net/http"
+	"net/url"
+)
+
+type (
+	// Config offers a declarative way to construct a client.
+	Config struct {
+		Endpoint         Endpoint `env:"ENDPOINT" long:"endpoint" default:"" description:"gremlin endpoint to connect to"`
+		DisableExpansion bool     `env:"DISABLE_EXPANSION" long:"disable-expansion" description:"disable bindings expansion"`
+	}
+
+	// An Option configured client.
+	Option func(*options)
+
+	options struct {
+		interceptors []Interceptor
+		httpClient   *http.Client
+	}
+
+	// Endpoint wraps a url to add flag unmarshalling.
+	Endpoint struct {
+		*url.URL
+	}
+)
+
+// WithInterceptor adds interceptors to the client's transport.
+func WithInterceptor(interceptors ...Interceptor) Option {
+	return func(opts *options) {
+		opts.interceptors = append(opts.interceptors, interceptors...)
+	}
+}
+
+// WithHTTPClient assigns underlying http client to be used by http transport.
+func WithHTTPClient(client *http.Client) Option {
+	return func(opts *options) {
+		opts.httpClient = client
+	}
+}
+
+// Build constructs a client from Config.
+func (cfg Config) Build(opt ...Option) (c *Client, err error) {
+	opts := cfg.buildOptions(opt)
+	switch cfg.Endpoint.Scheme {
+	case "http", "https":
+		c, err = NewHTTPClient(cfg.Endpoint.String(), opts.httpClient)
+	default:
+		err = fmt.Errorf("unsupported endpoint scheme: %s", cfg.Endpoint.Scheme)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	for i := len(opts.interceptors) - 1; i >= 0; i-- {
+		c.Transport = opts.interceptors[i](c.Transport)
+	}
+	if !cfg.DisableExpansion {
+		c.Transport = ExpandBindings(c.Transport)
+	}
+	return c, nil
+}
+
+func (Config) buildOptions(opts []Option) options {
+	var o options
+	for _, opt := range opts {
+		opt(&o)
+	}
+	return o
+}
+
+// UnmarshalFlag implements flag.Unmarshaler interface.
+func (ep *Endpoint) UnmarshalFlag(value string) (err error) {
+	ep.URL, err = url.Parse(value)
+	return
+}

dialect/gremlin/config_test.go

@@ -0,0 +1,154 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/url"
+	"testing"
+
+	"github.com/jessevdk/go-flags"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestConfigParsing(t *testing.T) {
+	var cfg Config
+	_, err := flags.ParseArgs(&cfg, []string{
+		"--disable-expansion",
+		"--endpoint", "http://localhost:8182/gremlin",
+	})
+	assert.NoError(t, err)
+	assert.True(t, cfg.DisableExpansion)
+	assert.Equal(t, "http", cfg.Endpoint.Scheme)
+	assert.Equal(t, "http://localhost:8182/gremlin", cfg.Endpoint.String())
+
+	cfg = Config{}
+	_, err = flags.ParseArgs(&cfg, nil)
+	assert.NoError(t, err)
+	assert.NotNil(t, cfg.Endpoint.URL)
+}
+
+func TestConfigBuild(t *testing.T) {
+	tests := []struct {
+		name    string
+		cfg     Config
+		opts    []Option
+		wantErr bool
+	}{
+		{
+			name: "HTTP",
+			cfg: Config{
+				Endpoint: Endpoint{
+					URL: func() *url.URL {
+						u, _ := url.Parse("http://gremlin-server/gremlin")
+						return u
+					}(),
+				},
+			},
+		},
+		{
+			name: "NoScheme",
+			cfg: Config{
+				Endpoint: Endpoint{
+					URL: &url.URL{},
+				},
+			},
+			wantErr: true,
+		},
+		{
+			name: "BadScheme",
+			cfg: Config{
+				Endpoint: Endpoint{
+					URL: &url.URL{
+						Scheme: "bad",
+					},
+				},
+			},
+			wantErr: true,
+		},
+		{
+			name: "WithOptions",
+			cfg: Config{
+				Endpoint: Endpoint{
+					URL: func() *url.URL {
+						u, _ := url.Parse("http://gremlin-server/gremlin")
+						return u
+					}(),
+				},
+				DisableExpansion: true,
+			},
+			opts: []Option{WithHTTPClient(&http.Client{})},
+		},
+		{
+			name: "NoExpansion",
+			cfg: Config{
+				Endpoint: Endpoint{
+					URL: &url.URL{
+						Scheme: "bad",
+					},
+				},
+				DisableExpansion: true,
+			},
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			client, err := tc.cfg.Build(tc.opts...)
+			if !tc.wantErr {
+				assert.NotNil(t, client)
+				assert.NoError(t, err)
+			} else {
+				assert.Error(t, err)
+			}
+		})
+	}
+}
+
+type testRoundTripper struct{ mock.Mock }
+
+func (rt *testRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	args := rt.Called(req)
+	rsp, _ := args.Get(0).(*http.Response)
+	return rsp, args.Error(1)
+}
+
+func TestBuildWithHTTPClient(t *testing.T) {
+	var transport testRoundTripper
+	transport.On("RoundTrip", mock.Anything).
+		Return(nil, errors.New("noop")).
+		Once()
+	defer transport.AssertExpectations(t)
+
+	u, err := url.Parse("http://gremlin-server:8182/gremlin")
+	require.NoError(t, err)
+
+	client, err := Config{Endpoint: Endpoint{u}}.
+		Build(WithHTTPClient(&http.Client{Transport: &transport}))
+	require.NoError(t, err)
+	_, _ = client.Do(context.Background(), NewEvalRequest("g.V()"))
+}
+
+func TestExpandOrdering(t *testing.T) {
+	var cfg Config
+	cfg.Endpoint.URL, _ = url.Parse("http://gremlin-server/gremlin")
+	interceptor := func(RoundTripper) RoundTripper {
+		return RoundTripperFunc(func(ctx context.Context, req *Request) (*Response, error) {
+			assert.Equal(t, `g.V().hasLabel("user")`, req.Arguments[ArgsGremlin])
+			assert.Nil(t, req.Arguments[ArgsBindings])
+			return nil, errors.New("noop")
+		})
+	}
+	c, err := cfg.Build(WithInterceptor(interceptor))
+	require.NoError(t, err)
+	req := NewEvalRequest("g.V().hasLabel($1)", WithBindings(map[string]any{"$1": "user"}))
+	_, _ = c.Do(context.Background(), req)
+}

dialect/gremlin/driver.go

@@ -0,0 +1,58 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"fmt"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/gremlin/graph/dsl"
+)
+
+// Driver is a dialect.Driver implementation for TinkerPop gremlin.
+type Driver struct {
+	*Client
+}
+
+// NewDriver returns a new dialect.Driver implementation for gremlin.
+func NewDriver(c *Client) *Driver {
+	c.Transport = ExpandBindings(c.Transport)
+	return &Driver{c}
+}
+
+// Dialect implements the dialect.Dialect method.
+func (Driver) Dialect() string { return dialect.Gremlin }
+
+// Exec implements the dialect.Exec method.
+func (c *Driver) Exec(ctx context.Context, query string, args, v any) error {
+	vr, ok := v.(*Response)
+	if !ok {
+		return fmt.Errorf("dialect/gremlin: invalid type %T. expect *gremlin.Response", v)
+	}
+	bindings, ok := args.(dsl.Bindings)
+	if !ok {
+		return fmt.Errorf("dialect/gremlin: invalid type %T. expect map[string]any for bindings", args)
+	}
+	res, err := c.Do(ctx, NewEvalRequest(query, WithBindings(bindings)))
+	if err != nil {
+		return err
+	}
+	*vr = *res
+	return nil
+}
+
+// Query implements the dialect.Query method.
+func (c *Driver) Query(ctx context.Context, query string, args, v any) error {
+	return c.Exec(ctx, query, args, v)
+}
+
+// Close is a nop close call. It should close the connection in case of WS client.
+func (Driver) Close() error { return nil }
+
+// Tx returns a nop transaction.
+func (c *Driver) Tx(context.Context) (dialect.Tx, error) { return dialect.NopTx(c), nil }
+
+var _ dialect.Driver = (*Driver)(nil)

dialect/gremlin/encoding/graphson/bench_test.go

@@ -0,0 +1,99 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"testing"
+
+	jsoniter "github.com/json-iterator/go"
+)
+
+type book struct {
+	ID       string   `json:"id" graphson:"g:UUID"`
+	Title    string   `json:"title"`
+	Author   string   `json:"author"`
+	Pages    int      `json:"num_pages"`
+	Chapters []string `json:"chapters"`
+}
+
+func generateObject() *book {
+	return &book{
+		ID:       "21d5dcbf-1fd4-493e-9b74-d6c429f9e4a5",
+		Title:    "The Art of Computer Programming, Vol. 2",
+		Author:   "Donald E. Knuth",
+		Pages:    784,
+		Chapters: []string{"Random numbers", "Arithmetic"},
+	}
+}
+
+func BenchmarkMarshalObject(b *testing.B) {
+	obj := generateObject()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		_, err := Marshal(obj)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkUnmarshalObject(b *testing.B) {
+	b.ReportAllocs()
+	out, err := Marshal(generateObject())
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	obj := &book{}
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		err = Unmarshal(out, obj)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkMarshalInterface(b *testing.B) {
+	b.ReportAllocs()
+	data, err := jsoniter.Marshal(generateObject())
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	var obj any
+	if err = jsoniter.Unmarshal(data, &obj); err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		_, err = Marshal(obj)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkUnmarshalInterface(b *testing.B) {
+	b.ReportAllocs()
+	data, err := Marshal(generateObject())
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	var obj any
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		err = Unmarshal(data, &obj)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}

dialect/gremlin/encoding/graphson/common_test.go

@@ -0,0 +1,54 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/stretchr/testify/mock"
+)
+
+type mocker struct {
+	mock.Mock
+}
+
+// Encode belongs to jsoniter.ValEncoder interface.
+func (m *mocker) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	m.Called(ptr, stream)
+}
+
+// IsEmpty belongs to jsoniter.ValEncoder interface.
+func (m *mocker) IsEmpty(ptr unsafe.Pointer) bool {
+	args := m.Called(ptr)
+	return args.Bool(0)
+}
+
+// Decode implements jsoniter.ValDecoder interface.
+func (m *mocker) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	m.Called(ptr, iter)
+}
+
+// CheckType implements typeChecker interface.
+func (m *mocker) CheckType(typ Type) error {
+	args := m.Called(typ)
+	return args.Error(0)
+}
+
+// MarshalGraphson implements Marshaler interface.
+func (m *mocker) MarshalGraphson() ([]byte, error) {
+	args := m.Called()
+	data, err := args.Get(0), args.Error(1)
+	if data == nil {
+		return nil, err
+	}
+	return data.([]byte), err
+}
+
+// UnmarshalGraphson implements Unmarshaler interface.
+func (m *mocker) UnmarshalGraphson(data []byte) error {
+	args := m.Called(data)
+	return args.Error(0)
+}

dialect/gremlin/encoding/graphson/decode.go

@@ -0,0 +1,109 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"io"
+	"reflect"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+type decodeExtension struct {
+	jsoniter.DummyExtension
+}
+
+// Unmarshal parses the graphson encoded data and stores the result
+// in the value pointed to by v.
+func Unmarshal(data []byte, v any) error {
+	return config.Unmarshal(data, v)
+}
+
+// UnmarshalFromString parses the graphson encoded str and stores the result
+// in the value pointed to by v.
+func UnmarshalFromString(str string, v any) error {
+	return config.UnmarshalFromString(str, v)
+}
+
+// Decoder defines a graphson decoder.
+type Decoder interface {
+	Decode(any) error
+}
+
+// NewDecoder create a graphson decoder.
+func NewDecoder(r io.Reader) Decoder {
+	return config.NewDecoder(r)
+}
+
+// Unmarshaler is the interface implemented by types
+// that can unmarshal a graphson description of themselves.
+type Unmarshaler interface {
+	UnmarshalGraphson([]byte) error
+}
+
+// UpdateStructDescriptor decorates struct field encoders for graphson tags.
+func (ext decodeExtension) UpdateStructDescriptor(desc *jsoniter.StructDescriptor) {
+	for _, binding := range desc.Fields {
+		if tag, ok := binding.Field.Tag().Lookup("graphson"); ok && tag != "-" {
+			if dec := ext.DecoratorOfStructField(binding.Decoder, tag); dec != nil {
+				binding.Decoder = dec
+			}
+		}
+	}
+}
+
+// CreateDecoder returns a value decoder for type.
+func (ext decodeExtension) CreateDecoder(typ reflect2.Type) jsoniter.ValDecoder {
+	if dec := ext.DecoderOfRegistered(typ); dec != nil {
+		return dec
+	}
+	if dec := ext.DecoderOfUnmarshaler(typ); dec != nil {
+		return dec
+	}
+	if dec := ext.DecoderOfNative(typ); dec != nil {
+		return dec
+	}
+	switch typ.Kind() {
+	case reflect.Array:
+		return ext.DecoderOfArray(typ)
+	case reflect.Slice:
+		return ext.DecoderOfSlice(typ)
+	case reflect.Map:
+		return ext.DecoderOfMap(typ)
+	default:
+		return nil
+	}
+}
+
+// DecorateDecoder decorates an passed in value decoder for type.
+func (ext decodeExtension) DecorateDecoder(typ reflect2.Type, dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	if dec := ext.DecoratorOfRegistered(dec); dec != nil {
+		return dec
+	}
+	if dec := ext.DecoratorOfUnmarshaler(typ, dec); dec != nil {
+		return dec
+	}
+	if dec := ext.DecoratorOfTyper(typ, dec); dec != nil {
+		return dec
+	}
+	if dec := ext.DecoratorOfNative(typ, dec); dec != nil {
+		return dec
+	}
+	switch typ.Kind() {
+	case reflect.Ptr, reflect.Struct:
+		return dec
+	case reflect.Interface:
+		return ext.DecoratorOfInterface(typ, dec)
+	case reflect.Slice:
+		return ext.DecoratorOfSlice(typ, dec)
+	case reflect.Array:
+		return ext.DecoratorOfArray(dec)
+	case reflect.Map:
+		return ext.DecoratorOfMap(dec)
+	default:
+		return ext.DecoderOfError("graphson: unsupported type: %s", typ.String())
+	}
+}

dialect/gremlin/encoding/graphson/decode_test.go

@@ -0,0 +1,5 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson

dialect/gremlin/encoding/graphson/encode.go

@@ -0,0 +1,98 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"io"
+	"reflect"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+type encodeExtension struct {
+	jsoniter.DummyExtension
+}
+
+// Marshal returns the graphson encoding of v.
+func Marshal(v any) ([]byte, error) {
+	return config.Marshal(v)
+}
+
+// MarshalToString returns the graphson encoding of v as string.
+func MarshalToString(v any) (string, error) {
+	return config.MarshalToString(v)
+}
+
+// Encoder defines a graphson encoder.
+type Encoder interface {
+	Encode(any) error
+}
+
+// NewEncoder create a graphson encoder.
+func NewEncoder(w io.Writer) Encoder {
+	return config.NewEncoder(w)
+}
+
+// Marshaler is the interface implemented by types that
+// can marshal themselves as graphson.
+type Marshaler interface {
+	MarshalGraphson() ([]byte, error)
+}
+
+// UpdateStructDescriptor decorates struct field encoders for graphson tags.
+func (ext encodeExtension) UpdateStructDescriptor(desc *jsoniter.StructDescriptor) {
+	for _, binding := range desc.Fields {
+		if tag, ok := binding.Field.Tag().Lookup("graphson"); ok && tag != "-" {
+			if enc := ext.DecoratorOfStructField(binding.Encoder, tag); enc != nil {
+				binding.Encoder = enc
+			}
+		}
+	}
+}
+
+// CreateEncoder returns a value encoder for type.
+func (ext encodeExtension) CreateEncoder(typ reflect2.Type) jsoniter.ValEncoder {
+	if enc := ext.EncoderOfRegistered(typ); enc != nil {
+		return enc
+	}
+	if enc := ext.EncoderOfNative(typ); enc != nil {
+		return enc
+	}
+	switch typ.Kind() {
+	case reflect.Map:
+		return ext.EncoderOfMap(typ)
+	default:
+		return nil
+	}
+}
+
+// DecorateEncoder decorates an passed in value encoder for type.
+func (ext encodeExtension) DecorateEncoder(typ reflect2.Type, enc jsoniter.ValEncoder) jsoniter.ValEncoder {
+	if enc := ext.DecoratorOfRegistered(enc); enc != nil {
+		return enc
+	}
+	if enc := ext.DecoratorOfMarshaler(typ, enc); enc != nil {
+		return enc
+	}
+	if enc := ext.DecoratorOfTyper(typ, enc); enc != nil {
+		return enc
+	}
+	if enc := ext.DecoratorOfNative(typ, enc); enc != nil {
+		return enc
+	}
+	switch typ.Kind() {
+	case reflect.Ptr, reflect.Interface, reflect.Struct:
+		return enc
+	case reflect.Array:
+		return ext.DecoratorOfArray(enc)
+	case reflect.Slice:
+		return ext.DecoratorOfSlice(typ, enc)
+	case reflect.Map:
+		return ext.DecoratorOfMap(enc)
+	default:
+		return ext.EncoderOfError("graphson: unsupported type: %s", typ.String())
+	}
+}

dialect/gremlin/encoding/graphson/encode_test.go

@@ -0,0 +1,16 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEncodeUnsupportedType(t *testing.T) {
+	_, err := Marshal(func() {})
+	assert.Error(t, err)
+}

dialect/gremlin/encoding/graphson/error.go

@@ -0,0 +1,45 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"fmt"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+)
+
+// EncoderOfError returns a value encoder which always fails to encode.
+func (encodeExtension) EncoderOfError(format string, args ...any) jsoniter.ValEncoder {
+	return decoratorOfError(format, args...)
+}
+
+// DecoderOfError returns a value decoder which always fails to decode.
+func (decodeExtension) DecoderOfError(format string, args ...any) jsoniter.ValDecoder {
+	return decoratorOfError(format, args...)
+}
+
+func decoratorOfError(format string, args ...any) errorCodec {
+	err := fmt.Errorf(format, args...)
+	return errorCodec{err}
+}
+
+type errorCodec struct{ error }
+
+func (ec errorCodec) Encode(_ unsafe.Pointer, stream *jsoniter.Stream) {
+	if stream.Error == nil {
+		stream.Error = ec.error
+	}
+}
+
+func (errorCodec) IsEmpty(unsafe.Pointer) bool {
+	return false
+}
+
+func (ec errorCodec) Decode(_ unsafe.Pointer, iter *jsoniter.Iterator) {
+	if iter.Error == nil {
+		iter.Error = ec.error
+	}
+}

dialect/gremlin/encoding/graphson/error_test.go

@@ -0,0 +1,30 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"bytes"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestErrorCodec(t *testing.T) {
+	codec := errorCodec{errors.New("codec error")}
+	assert.False(t, codec.IsEmpty(nil))
+
+	var buf bytes.Buffer
+	stream := config.BorrowStream(&buf)
+	defer config.ReturnStream(stream)
+	codec.Encode(nil, stream)
+	assert.Empty(t, buf.Bytes())
+	assert.EqualError(t, stream.Error, codec.Error())
+
+	iter := config.BorrowIterator([]byte{})
+	defer config.ReturnIterator(iter)
+	codec.Decode(nil, iter)
+	assert.EqualError(t, iter.Error, codec.Error())
+}

dialect/gremlin/encoding/graphson/extension.go

@@ -0,0 +1,88 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"reflect"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+var (
+	typeEncoders = map[string]jsoniter.ValEncoder{}
+	typeDecoders = map[string]jsoniter.ValDecoder{}
+)
+
+// RegisterTypeEncoder register type encoder for typ.
+func RegisterTypeEncoder(typ string, enc jsoniter.ValEncoder) {
+	typeEncoders[typ] = enc
+}
+
+// RegisterTypeDecoder register type decoder for typ.
+func RegisterTypeDecoder(typ string, dec jsoniter.ValDecoder) {
+	typeDecoders[typ] = dec
+}
+
+type registeredEncoder struct{ jsoniter.ValEncoder }
+
+// EncoderOfRegistered returns a value encoder of a registered type.
+func (encodeExtension) EncoderOfRegistered(typ reflect2.Type) jsoniter.ValEncoder {
+	enc := typeEncoders[typ.String()]
+	if enc != nil {
+		return registeredEncoder{enc}
+	}
+	if typ.Kind() == reflect.Ptr {
+		ptrType := typ.(reflect2.PtrType)
+		enc := typeEncoders[ptrType.Elem().String()]
+		if enc != nil {
+			return registeredEncoder{
+				ValEncoder: &jsoniter.OptionalEncoder{
+					ValueEncoder: enc,
+				},
+			}
+		}
+	}
+	return nil
+}
+
+// DecoratorOfRegistered decorates a value encoder of a registered type.
+func (encodeExtension) DecoratorOfRegistered(enc jsoniter.ValEncoder) jsoniter.ValEncoder {
+	if _, ok := enc.(registeredEncoder); ok {
+		return enc
+	}
+	return nil
+}
+
+type registeredDecoder struct{ jsoniter.ValDecoder }
+
+// DecoderOfRegistered returns a value decoder of a registered type.
+func (decodeExtension) DecoderOfRegistered(typ reflect2.Type) jsoniter.ValDecoder {
+	dec := typeDecoders[typ.String()]
+	if dec != nil {
+		return registeredDecoder{dec}
+	}
+	if typ.Kind() == reflect.Ptr {
+		ptrType := typ.(reflect2.PtrType)
+		dec := typeDecoders[ptrType.Elem().String()]
+		if dec != nil {
+			return registeredDecoder{
+				ValDecoder: &jsoniter.OptionalDecoder{
+					ValueType:    ptrType.Elem(),
+					ValueDecoder: dec,
+				},
+			}
+		}
+	}
+	return nil
+}
+
+// DecoratorOfRegistered decorates a value decoder of a registered type.
+func (decodeExtension) DecoratorOfRegistered(dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	if _, ok := dec.(registeredDecoder); ok {
+		return dec
+	}
+	return nil
+}

dialect/gremlin/encoding/graphson/init.go

@@ -0,0 +1,16 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	jsoniter "github.com/json-iterator/go"
+)
+
+var config = jsoniter.Config{}.Froze()
+
+func init() {
+	config.RegisterExtension(&encodeExtension{})
+	config.RegisterExtension(&decodeExtension{})
+}

dialect/gremlin/encoding/graphson/interface.go

@@ -0,0 +1,153 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+	"reflect"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+// DecoratorOfInterface decorates a value decoder of an interface type.
+func (decodeExtension) DecoratorOfInterface(typ reflect2.Type, dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	if _, ok := typ.(*reflect2.UnsafeEFaceType); ok {
+		return efaceDecoder{typ, dec}
+	}
+	return dec
+}
+
+type efaceDecoder struct {
+	typ reflect2.Type
+	jsoniter.ValDecoder
+}
+
+func (dec efaceDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	switch next := iter.WhatIsNext(); next {
+	case jsoniter.StringValue, jsoniter.BoolValue, jsoniter.NilValue:
+		dec.ValDecoder.Decode(ptr, iter)
+	case jsoniter.ObjectValue:
+		dec.decode(ptr, iter)
+	default:
+		iter.ReportError("decode empty interface", fmt.Sprintf("unexpected value type: %d", next))
+	}
+}
+
+func (dec efaceDecoder) decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	data := iter.SkipAndReturnBytes()
+	if iter.Error != nil && iter.Error != io.EOF {
+		return
+	}
+
+	rtype, err := dec.reflectBytes(data)
+	if err != nil {
+		iter.ReportError("decode empty interface", err.Error())
+		return
+	}
+
+	it := config.BorrowIterator(data)
+	defer config.ReturnIterator(it)
+
+	var val any
+	if rtype != nil {
+		val = rtype.New()
+		it.ReadVal(val)
+		val = rtype.Indirect(val)
+	} else {
+		if jsoniter.Get(data, TypeKey).LastError() == nil {
+			vk := jsoniter.Get(data, ValueKey)
+			if vk.LastError() == nil {
+				val = vk.GetInterface()
+			}
+		}
+		if val == nil {
+			val = it.Read()
+		}
+	}
+
+	if it.Error != nil && it.Error != io.EOF {
+		iter.ReportError("decode empty interface", it.Error.Error())
+		return
+	}
+
+	// nolint: gas
+	dec.typ.UnsafeSet(ptr, unsafe.Pointer(&val))
+}
+
+func (dec efaceDecoder) reflectBytes(data []byte) (reflect2.Type, error) {
+	typ := Type(jsoniter.Get(data, TypeKey).ToString())
+	rtype := dec.reflectType(typ)
+	if rtype != nil {
+		return rtype, nil
+	}
+
+	switch typ {
+	case listType:
+		return dec.reflectSlice(data)
+	case mapType:
+		return dec.reflectMap(data)
+	default:
+		return nil, nil
+	}
+}
+
+func (efaceDecoder) reflectType(typ Type) reflect2.Type {
+	switch typ {
+	case doubleType:
+		return reflect2.TypeOf(float64(0))
+	case floatType:
+		return reflect2.TypeOf(float32(0))
+	case byteType:
+		return reflect2.TypeOf(uint8(0))
+	case int16Type:
+		return reflect2.TypeOf(int16(0))
+	case int32Type:
+		return reflect2.TypeOf(int32(0))
+	case int64Type, bigIntegerType:
+		return reflect2.TypeOf(int64(0))
+	case byteBufferType:
+		return reflect2.TypeOf([]byte{})
+	default:
+		return nil
+	}
+}
+
+func (efaceDecoder) reflectSlice(data []byte) (reflect2.Type, error) {
+	var elem any
+	if err := Unmarshal(data, &[...]*any{&elem}); err != nil {
+		return nil, fmt.Errorf("cannot read first list element: %w", err)
+	}
+
+	if elem == nil {
+		return reflect2.TypeOf([]any{}), nil
+	}
+
+	sliceType := reflect.SliceOf(reflect.TypeOf(elem))
+	return reflect2.Type2(sliceType), nil
+}
+
+func (efaceDecoder) reflectMap(data []byte) (reflect2.Type, error) {
+	var key, elem any
+	if err := Unmarshal(
+		bytes.Replace(data, []byte(mapType), []byte(listType), 1),
+		&[...]*any{&key, &elem},
+	); err != nil {
+		return nil, fmt.Errorf("cannot unmarshal first map item: %w", err)
+	}
+
+	if key == nil {
+		return reflect2.TypeOf(map[any]any{}), nil
+	} else if elem == nil {
+		return nil, errors.New("expect map element, but found only key")
+	}
+
+	mapType := reflect.MapOf(reflect.TypeOf(key), reflect.TypeOf(elem))
+	return reflect2.Type2(mapType), nil
+}

dialect/gremlin/encoding/graphson/interface_test.go

@@ -0,0 +1,352 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDecodeInterface(t *testing.T) {
+	tests := []struct {
+		name    string
+		in      string
+		want    any
+		wantErr bool
+	}{
+		{
+			name: "Boolean",
+			in:   "false",
+			want: false,
+		},
+		{
+			name: "String",
+			in:   `"str"`,
+			want: "str",
+		},
+		{
+			name: "Double",
+			in: `{
+				"@type": "g:Double",
+				"@value": 3.14
+			}`,
+			want: float64(3.14),
+		},
+		{
+			name: "Float",
+			in: `{
+				"@type": "g:Float",
+				"@value": -22.567
+			}`,
+			want: float32(-22.567),
+		},
+		{
+			name: "Int32",
+			in: `{
+				"@type": "g:Int32",
+				"@value": 9000
+			}`,
+			want: int32(9000),
+		},
+		{
+			name: "Int64",
+			in: `{
+				"@type": "g:Int64",
+				"@value": 188786
+			}`,
+			want: int64(188786),
+		},
+		{
+			name: "BigInteger",
+			in: `{
+				"@type": "gx:BigInteger",
+				"@value": 352353463712
+			}`,
+			want: int64(352353463712),
+		},
+		{
+			name: "Byte",
+			in: `{
+				"@type": "gx:Byte",
+				"@value": 100
+			}`,
+			want: uint8(100),
+		},
+		{
+			name: "Int16",
+			in: `{
+				"@type": "gx:Int16",
+				"@value": 2000
+			}`,
+			want: int16(2000),
+		},
+		{
+			name: "UnknownType",
+			in: `{
+				"@type": "g:T",
+				"@value": "label"
+			}`,
+			want: "label",
+		},
+		{
+			name:    "UntypedArray",
+			in:      "[]",
+			wantErr: true,
+		},
+		{
+			name: "NoType",
+			in: `{
+				"@typ": "g:Int32",
+				"@value": 345
+			}`,
+			wantErr: true,
+		},
+		{
+			name: "BadObject",
+			in: `{
+				"@type": "g:Int32",
+				"@value": 345
+			`,
+			wantErr: true,
+		},
+		{
+			name: "BadList",
+			in: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@val": 123457990
+					}
+				]
+			}`,
+			wantErr: true,
+		},
+		{
+			name: "BadMap",
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@val": 123457990
+					},
+					"First"
+				]
+			}`,
+			wantErr: true,
+		},
+		{
+			name: "KeyOnlyMap",
+			in: `{
+				"@type": "g:Map",
+				"@value": ["Key"]
+			}`,
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			var got any
+			err := UnmarshalFromString(tc.in, &got)
+			if !tc.wantErr {
+				require.NoError(t, err)
+				assert.Equal(t, tc.want, got)
+			} else {
+				assert.Error(t, err)
+			}
+		})
+	}
+}
+
+func TestDecodeInterfaceSlice(t *testing.T) {
+	tests := []struct {
+		in   string
+		want any
+	}{
+		{
+			in: `{
+				"@type": "g:List",
+				"@value": []
+			}`,
+			want: []any{},
+		},
+		{
+			in: `{
+				"@type": "g:List",
+				"@value": ["x", "y", "z"]
+			}`,
+			want: []string{"x", "y", "z"},
+		},
+		{
+			in: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@value": 123457990
+					},
+					{
+						"@type": "g:Int64",
+						"@value": 23456111
+					},
+					{
+						"@type": "g:Int64",
+						"@value": -687450
+					}
+				]
+			}`,
+			want: []int64{123457990, 23456111, -687450},
+		},
+		{
+			in: `{
+				"@type": "gx:ByteBuffer",
+				"@value": "AQIDBAU="
+			}`,
+			want: []byte{1, 2, 3, 4, 5},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(fmt.Sprintf("%T", tc.want), func(t *testing.T) {
+			t.Parallel()
+			var got any
+			err := UnmarshalFromString(tc.in, &got)
+			require.NoError(t, err)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}
+
+func TestDecodeInterfaceMap(t *testing.T) {
+	tests := []struct {
+		in   string
+		want any
+	}{
+		{
+			in: `{
+				"@type": "g:Map",
+				"@value": []
+			}`,
+			want: map[any]any{},
+		},
+		{
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					"Sep",
+					{
+						"@type": "g:Int32",
+						"@value": 9
+					},
+					"Oct",
+					{
+						"@type": "g:Int32",
+						"@value": 10
+					},
+					"Nov",
+					{
+						"@type": "g:Int32",
+						"@value": 11
+					}
+				]
+			}`,
+			want: map[string]int32{
+				"Sep": int32(9),
+				"Oct": int32(10),
+				"Nov": int32(11),
+			},
+		},
+		{
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					"One",
+					{
+						"@type": "g:List",
+						"@value": [
+							{
+								"@type": "g:Int32",
+								"@value": 1
+							}
+						]
+					},
+					"Two",
+					{
+						"@type": "g:List",
+						"@value": [
+							{
+								"@type": "g:Int32",
+								"@value": 2
+							}
+						]
+					},
+					"Three",
+					{
+						"@type": "g:List",
+						"@value": [
+							{
+								"@type": "g:Int32",
+								"@value": 3
+							}
+						]
+					}
+				]
+			}`,
+			want: map[string][]int32{
+				"One":   {1},
+				"Two":   {2},
+				"Three": {3},
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(fmt.Sprintf("%T", tc.want), func(t *testing.T) {
+			t.Parallel()
+			var got any
+			err := UnmarshalFromString(tc.in, &got)
+			require.NoError(t, err)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}
+
+func TestDecodeInterfaceObject(t *testing.T) {
+	book := struct {
+		ID       string   `json:"id" graphson:"g:UUID"`
+		Title    string   `json:"title"`
+		Author   string   `json:"author"`
+		Pages    int      `json:"num_pages"`
+		Chapters []string `json:"chapters"`
+	}{
+		ID:       "21d5dcbf-1fd4-493e-9b74-d6c429f9e4a5",
+		Title:    "The Art of Computer Programming, Vol. 2",
+		Author:   "Donald E. Knuth",
+		Pages:    784,
+		Chapters: []string{"Random numbers", "Arithmetic"},
+	}
+	data, err := Marshal(book)
+	require.NoError(t, err)
+
+	var v any
+	err = Unmarshal(data, &v)
+	require.NoError(t, err)
+
+	obj := v.(map[string]any)
+	assert.Equal(t, book.ID, obj["id"])
+	assert.Equal(t, book.Title, obj["title"])
+	assert.Equal(t, book.Author, obj["author"])
+	assert.EqualValues(t, book.Pages, obj["num_pages"])
+	assert.ElementsMatch(t, book.Chapters, obj["chapters"])
+}

dialect/gremlin/encoding/graphson/lazy.go

@@ -0,0 +1,74 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"fmt"
+	"sync"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+// LazyEncoderOf returns a lazy encoder for type.
+func (encodeExtension) LazyEncoderOf(typ reflect2.Type) jsoniter.ValEncoder {
+	return &lazyEncoder{resolve: func() jsoniter.ValEncoder {
+		return config.EncoderOf(typ)
+	}}
+}
+
+// LazyDecoderOf returns a lazy unique decoder for type.
+func (decodeExtension) LazyDecoderOf(typ reflect2.Type) jsoniter.ValDecoder {
+	return &lazyDecoder{resolve: func() jsoniter.ValDecoder {
+		dec := config.DecoderOf(reflect2.PtrTo(typ))
+		if td, ok := dec.(typeDecoder); ok {
+			td.typeChecker = &uniqueType{elemChecker: td.typeChecker}
+			dec = td
+		}
+		return dec
+	}}
+}
+
+type lazyEncoder struct {
+	jsoniter.ValEncoder
+	resolve func() jsoniter.ValEncoder
+	once    sync.Once
+}
+
+func (enc *lazyEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	enc.once.Do(func() { enc.ValEncoder = enc.resolve() })
+	enc.ValEncoder.Encode(ptr, stream)
+}
+
+func (enc *lazyEncoder) IsEmpty(ptr unsafe.Pointer) bool {
+	enc.once.Do(func() { enc.ValEncoder = enc.resolve() })
+	return enc.ValEncoder.IsEmpty(ptr)
+}
+
+type lazyDecoder struct {
+	jsoniter.ValDecoder
+	resolve func() jsoniter.ValDecoder
+	once    sync.Once
+}
+
+func (dec *lazyDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	dec.once.Do(func() { dec.ValDecoder = dec.resolve() })
+	dec.ValDecoder.Decode(ptr, iter)
+}
+
+type uniqueType struct {
+	typ         Type
+	once        sync.Once
+	elemChecker typeChecker
+}
+
+func (u *uniqueType) CheckType(other Type) error {
+	u.once.Do(func() { u.typ = other })
+	if u.typ != other {
+		return fmt.Errorf("expect type %s, but found %s", u.typ, other)
+	}
+	return u.elemChecker.CheckType(u.typ)
+}

dialect/gremlin/encoding/graphson/lazy_test.go

@@ -0,0 +1,46 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"sync/atomic"
+	"testing"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestLazyEncode(t *testing.T) {
+	var m mocker
+	m.On("IsEmpty", mock.Anything).Return(false).Once()
+	m.On("Encode", mock.Anything, mock.Anything).Once()
+	defer m.AssertExpectations(t)
+
+	var cnt uint32
+	var enc jsoniter.ValEncoder = &lazyEncoder{resolve: func() jsoniter.ValEncoder {
+		assert.Equal(t, uint32(1), atomic.AddUint32(&cnt, 1))
+		return &m
+	}}
+
+	enc.IsEmpty(nil)
+	enc.Encode(nil, nil)
+}
+
+func TestLazyDecode(t *testing.T) {
+	var m mocker
+	m.On("Decode", mock.Anything, mock.Anything).Times(3)
+	defer m.AssertExpectations(t)
+
+	var cnt uint32
+	var dec jsoniter.ValDecoder = &lazyDecoder{resolve: func() jsoniter.ValDecoder {
+		assert.Equal(t, uint32(1), atomic.AddUint32(&cnt, 1))
+		return &m
+	}}
+
+	for i := 0; i < 3; i++ {
+		dec.Decode(nil, nil)
+	}
+}

dialect/gremlin/encoding/graphson/map.go

@@ -0,0 +1,116 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+// EncoderOfMap returns a value encoder of a map type.
+func (ext encodeExtension) EncoderOfMap(typ reflect2.Type) jsoniter.ValEncoder {
+	mapType := typ.(reflect2.MapType)
+	return &mapEncoder{
+		mapType: mapType,
+		keyEnc:  ext.LazyEncoderOf(mapType.Key()),
+		elemEnc: ext.LazyEncoderOf(mapType.Elem()),
+	}
+}
+
+// DecoratorOfMap decorates a value encoder of a map type.
+func (encodeExtension) DecoratorOfMap(enc jsoniter.ValEncoder) jsoniter.ValEncoder {
+	return typeEncoder{enc, mapType}
+}
+
+type mapEncoder struct {
+	mapType reflect2.MapType
+	keyEnc  jsoniter.ValEncoder
+	elemEnc jsoniter.ValEncoder
+}
+
+func (enc *mapEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	iter := enc.mapType.UnsafeIterate(ptr)
+	if !iter.HasNext() {
+		stream.WriteEmptyArray()
+		return
+	}
+
+	stream.WriteArrayStart()
+	for {
+		key, elem := iter.UnsafeNext()
+		enc.keyEnc.Encode(key, stream)
+		stream.WriteMore()
+		enc.elemEnc.Encode(elem, stream)
+		if !iter.HasNext() {
+			break
+		}
+		stream.WriteMore()
+	}
+	stream.WriteArrayEnd()
+}
+
+func (enc *mapEncoder) IsEmpty(ptr unsafe.Pointer) bool {
+	return !enc.mapType.UnsafeIterate(ptr).HasNext()
+}
+
+// DecoderOfMap returns a value decoder of a map type.
+func (ext decodeExtension) DecoderOfMap(typ reflect2.Type) jsoniter.ValDecoder {
+	mapType := typ.(reflect2.MapType)
+	keyType, elemType := mapType.Key(), mapType.Elem()
+	return &mapDecoder{
+		mapType:  mapType,
+		keyType:  keyType,
+		elemType: elemType,
+		keyDec:   ext.LazyDecoderOf(keyType),
+		elemDec:  ext.LazyDecoderOf(elemType),
+	}
+}
+
+// DecoratorOfMap decorates a value decoder of a map type.
+func (decodeExtension) DecoratorOfMap(dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	return typeDecoder{dec, mapType}
+}
+
+type mapDecoder struct {
+	mapType  reflect2.MapType
+	keyType  reflect2.Type
+	elemType reflect2.Type
+	keyDec   jsoniter.ValDecoder
+	elemDec  jsoniter.ValDecoder
+}
+
+func (dec *mapDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	mapType := dec.mapType
+	if mapType.UnsafeIsNil(ptr) {
+		mapType.UnsafeSet(ptr, mapType.UnsafeMakeMap(0))
+	}
+
+	var key unsafe.Pointer
+	if !iter.ReadArrayCB(func(iter *jsoniter.Iterator) bool {
+		if key == nil {
+			key = dec.keyType.UnsafeNew()
+			dec.keyDec.Decode(key, iter)
+			return iter.Error == nil
+		}
+
+		elem := dec.elemType.UnsafeNew()
+		dec.elemDec.Decode(elem, iter)
+		if iter.Error != nil {
+			return false
+		}
+
+		mapType.UnsafeSetIndex(ptr, key, elem)
+		key = nil
+		return true
+	}) {
+		return
+	}
+
+	if key != nil {
+		iter.ReportError("decode map", "odd number of map items")
+	}
+}

dialect/gremlin/encoding/graphson/map_test.go

@@ -0,0 +1,355 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"strings"
+	"testing"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEncodeMap(t *testing.T) {
+	tests := []struct {
+		name string
+		in   any
+		want string
+	}{
+		{
+			name: "simple",
+			in: map[int32]string{
+				3: "Mar",
+				1: "Jan",
+				2: "Feb",
+			},
+			want: `[
+				{
+					"@type": "g:Int32",
+					"@value": 1
+				},
+				"Jan",
+				{
+					"@type": "g:Int32",
+					"@value": 2
+				},
+				"Feb",
+				{
+					"@type": "g:Int32",
+					"@value": 3
+				},
+				"Mar"
+			]`,
+		},
+		{
+			name: "mixed",
+			in: map[string]any{
+				"byte":   byte('a'),
+				"string": "str",
+				"slice":  []int{1, 2, 3},
+				"map":    map[string]int{},
+			},
+			want: `[
+				"byte",
+				{
+					"@type": "gx:Byte",
+					"@value": 97
+				},
+				"string",
+				"str",
+				"slice",
+				{
+					"@type": "g:List",
+					"@value": [
+						{
+							"@type": "g:Int64",
+							"@value": 1
+						},
+						{
+							"@type": "g:Int64",
+							"@value": 2
+						},
+						{
+							"@type": "g:Int64",
+							"@value": 3
+						}
+					]
+				},
+				"map",
+				{
+					"@type": "g:Map",
+					"@value": []
+				}
+			]`,
+		},
+		{
+			name: "struct-key",
+			in: map[struct {
+				K string `json:"key"`
+			}]int32{
+
+				{"result"}: 42,
+			},
+			want: `[
+				{
+					"key": "result"
+				},
+				{
+					"@type": "g:Int32",
+					"@value": 42
+				}
+			]`,
+		},
+		{
+			name: "nil",
+			in:   map[string]uint8(nil),
+			want: "null",
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			data, err := Marshal(tc.in)
+			require.NoError(t, err)
+
+			assert.Equal(t, "g:Map", jsoniter.Get(data, "@type").ToString())
+			var want []any
+			err = jsoniter.UnmarshalFromString(tc.want, &want)
+			require.NoError(t, err)
+
+			got, ok := jsoniter.Get(data, "@value").GetInterface().([]any)
+			require.True(t, ok)
+			assert.ElementsMatch(t, want, got)
+		})
+	}
+}
+
+func TestDecodeMap(t *testing.T) {
+	tests := []struct {
+		name string
+		in   string
+		want any
+	}{
+		{
+			name: "empty",
+			in: `{
+				"@type": "g:Map",
+				"@value": []
+			}`,
+			want: map[int]int{},
+		},
+		{
+			name: "simple",
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					{
+						"@type": "g:Int32",
+						"@value": 6
+					},
+					"Jun",
+					{
+						"@type": "g:Int32",
+						"@value": 7
+					},
+					"Jul",
+					{
+						"@type": "g:Int32",
+						"@value": 8
+					},
+					"Aug"
+				]
+			}`,
+			want: map[int]string{
+				6: "Jun",
+				7: "Jul",
+				8: "Aug",
+			},
+		},
+		{
+			name: "duplicate",
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					"Sep",
+					{
+						"@type": "g:Int32",
+						"@value": 9
+					},
+					"Oct",
+					{
+						"@type": "g:Int32",
+						"@value": 65
+					},
+					"Oct",
+					{
+						"@type": "g:Int32",
+						"@value": 10
+					},
+					"Nov",
+					null
+				]
+			}`,
+			want: map[string]*int{
+				"Sep": func() *int { v := 9; return &v }(),
+				"Oct": func() *int { v := 10; return &v }(),
+				"Nov": nil,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			typ := reflect2.TypeOf(tc.want).(reflect2.MapType)
+			got := typ.MakeMap(0)
+			err := UnmarshalFromString(tc.in, got)
+			require.NoError(t, err)
+			assert.Equal(t, tc.want, typ.Indirect(got))
+		})
+	}
+}
+
+func TestDecodeMapIntoNil(t *testing.T) {
+	var got map[int64]int32
+	err := UnmarshalFromString(`{
+		"@type": "g:Map",
+		"@value": [
+			{
+				"@type": "g:Int64",
+				"@value": 9
+			},
+			{
+				"@type": "g:Int32",
+				"@value": -9
+			},
+			{
+				"@type": "g:Int64",
+				"@value": 99
+			},
+			{
+				"@type": "g:Int32",
+				"@value": -99
+			},
+			{
+				"@type": "g:Int64",
+				"@value": 999
+			},
+			{
+				"@type": "g:Int32",
+				"@value": -999
+			}
+		]
+	}`, &got)
+	require.NoError(t, err)
+	assert.Equal(t, map[int64]int32{9: -9, 99: -99, 999: -999}, got)
+}
+
+func TestDecodeBadMap(t *testing.T) {
+	tests := []struct {
+		name string
+		in   string
+	}{
+		{
+			name: "BadValue",
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@value": 9
+					},
+					{
+						"@type": "g:Int32",
+						"@value": "55"
+					}
+				]
+			}`,
+		},
+		{
+			name: "NoValue",
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@value": 9
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 9
+					},
+					{
+						"@type": "g:Int64",
+						"@value": 42
+					}
+				]
+			}`,
+		},
+		{
+			name: "AlterKeyType",
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@value": 9
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 9
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 42
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 42
+					}
+				]
+			}`,
+		},
+		{
+			name: "AlterValType",
+			in: `{
+				"@type": "g:Map",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@value": 9
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 9
+					},
+					{
+						"@type": "g:Int64",
+						"@value": 42
+					},
+					{
+						"@type": "g:Int64",
+						"@value": 42
+					}
+				]
+			}`,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			var v map[int]int
+			err := NewDecoder(strings.NewReader(tc.in)).Decode(&v)
+			assert.Error(t, err)
+		})
+	}
+}

dialect/gremlin/encoding/graphson/marshaler.go

@@ -0,0 +1,133 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"fmt"
+	"io"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+// DecoratorOfMarshaler decorates a value encoder of a Marshaler interface.
+func (ext encodeExtension) DecoratorOfMarshaler(typ reflect2.Type, enc jsoniter.ValEncoder) jsoniter.ValEncoder {
+	if typ == marshalerType {
+		enc := marshalerEncoder{enc, typ}
+		return directMarshalerEncoder{enc}
+	}
+	if typ.Implements(marshalerType) {
+		return marshalerEncoder{enc, typ}
+	}
+	ptrType := reflect2.PtrTo(typ)
+	if ptrType.Implements(marshalerType) {
+		ptrEnc := ext.LazyEncoderOf(ptrType)
+		enc := marshalerEncoder{ptrEnc, ptrType}
+		return referenceEncoder{enc}
+	}
+	return nil
+}
+
+// DecoderOfUnmarshaler returns a value decoder of an Unmarshaler interface.
+func (decodeExtension) DecoderOfUnmarshaler(typ reflect2.Type) jsoniter.ValDecoder {
+	ptrType := reflect2.PtrTo(typ)
+	if ptrType.Implements(unmarshalerType) {
+		return referenceDecoder{
+			unmarshalerDecoder{ptrType},
+		}
+	}
+	return nil
+}
+
+// DecoratorOfUnmarshaler decorates a value encoder of an Unmarshaler interface.
+func (decodeExtension) DecoratorOfUnmarshaler(typ reflect2.Type, dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	if reflect2.PtrTo(typ).Implements(unmarshalerType) {
+		return dec
+	}
+	return nil
+}
+
+var (
+	marshalerType   = reflect2.TypeOfPtr((*Marshaler)(nil)).Elem()
+	unmarshalerType = reflect2.TypeOfPtr((*Unmarshaler)(nil)).Elem()
+)
+
+type marshalerEncoder struct {
+	jsoniter.ValEncoder
+	reflect2.Type
+}
+
+func (enc marshalerEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	marshaler := enc.Type.UnsafeIndirect(ptr).(Marshaler)
+	enc.encode(marshaler, stream)
+}
+
+func (enc marshalerEncoder) encode(marshaler Marshaler, stream *jsoniter.Stream) {
+	data, err := marshaler.MarshalGraphson()
+	if err != nil {
+		stream.Error = fmt.Errorf("graphson: error calling MarshalGraphson for type %s: %w", enc.Type, err)
+		return
+	}
+	if !config.Valid(data) {
+		stream.Error = fmt.Errorf("graphson: syntax error when marshaling type %s", enc.Type)
+		return
+	}
+	_, stream.Error = stream.Write(data)
+}
+
+type directMarshalerEncoder struct {
+	marshalerEncoder
+}
+
+func (enc directMarshalerEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	marshaler := *(*Marshaler)(ptr)
+	enc.encode(marshaler, stream)
+}
+
+type referenceEncoder struct {
+	jsoniter.ValEncoder
+}
+
+func (enc referenceEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	// nolint: gas
+	enc.ValEncoder.Encode(unsafe.Pointer(&ptr), stream)
+}
+
+func (enc referenceEncoder) IsEmpty(ptr unsafe.Pointer) bool {
+	// nolint: gas
+	return enc.ValEncoder.IsEmpty(unsafe.Pointer(&ptr))
+}
+
+type unmarshalerDecoder struct {
+	reflect2.Type
+}
+
+func (dec unmarshalerDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	bytes := iter.SkipAndReturnBytes()
+	if iter.Error != nil && iter.Error != io.EOF {
+		return
+	}
+
+	unmarshaler := dec.UnsafeIndirect(ptr).(Unmarshaler)
+	if err := unmarshaler.UnmarshalGraphson(bytes); err != nil {
+		iter.ReportError(
+			"unmarshal graphson",
+			fmt.Sprintf(
+				"graphson: error calling UnmarshalGraphson for type %s: %s",
+				dec.Type, err,
+			),
+		)
+	}
+}
+
+type referenceDecoder struct {
+	jsoniter.ValDecoder
+}
+
+func (dec referenceDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	// nolint: gas
+	dec.ValDecoder.Decode(unsafe.Pointer(&ptr), iter)
+}

dialect/gremlin/encoding/graphson/marshaler_test.go

@@ -0,0 +1,101 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"errors"
+	"fmt"
+	"reflect"
+	"testing"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMarshalerEncode(t *testing.T) {
+	want := []byte(`{"@type": "g:Int32", "@value": 42}`)
+	m := &mocker{}
+	call := m.On("MarshalGraphson").Return(want, nil)
+	defer m.AssertExpectations(t)
+
+	tests := []any{m, &m, func() *Marshaler { marshaler := Marshaler(m); return &marshaler }(), Marshaler(nil)}
+	call.Times(len(tests) - 1)
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(fmt.Sprintf("%T", tc), func(t *testing.T) {
+			got, err := Marshal(tc)
+			assert.NoError(t, err)
+			if !reflect2.IsNil(tc) {
+				assert.Equal(t, want, got)
+			} else {
+				assert.Equal(t, []byte("null"), got)
+			}
+		})
+	}
+}
+
+func TestMarshalerError(t *testing.T) {
+	errStr := "marshaler error"
+	m := &mocker{}
+	m.On("MarshalGraphson").Return(nil, errors.New(errStr)).Once()
+	defer m.AssertExpectations(t)
+
+	_, err := Marshal(m)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), errStr)
+}
+
+func TestBadMarshaler(t *testing.T) {
+	m := &mocker{}
+	m.On("MarshalGraphson").Return([]byte(`{"@type": "g:Int32", "@value":`), nil).Once()
+	defer m.AssertExpectations(t)
+	_, err := Marshal(m)
+	assert.Error(t, err)
+}
+
+func TestUnmarshalerDecode(t *testing.T) {
+	data := `{"@type": "g:UUID", "@value": "cb682578-9d92-4499-9ebc-5c6aa73c5397"}`
+	var value string
+
+	m := &mocker{}
+	m.On("UnmarshalGraphson", mock.Anything).
+		Run(func(args mock.Arguments) {
+			data := args.Get(0).([]byte)
+			value = jsoniter.Get(data, "@value").ToString()
+		}).
+		Return(nil).
+		Once()
+	defer m.AssertExpectations(t)
+
+	err := UnmarshalFromString(data, m)
+	require.NoError(t, err)
+	assert.Equal(t, "cb682578-9d92-4499-9ebc-5c6aa73c5397", value)
+}
+
+func TestUnmarshalerError(t *testing.T) {
+	errStr := "unmarshaler error"
+	m := &mocker{}
+	m.On("UnmarshalGraphson", mock.Anything).Return(errors.New(errStr)).Once()
+	defer m.AssertExpectations(t)
+
+	err := Unmarshal([]byte(`{}`), m)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(),
+		fmt.Sprintf("graphson: error calling UnmarshalGraphson for type %s: %s",
+			reflect.TypeOf(m), errStr,
+		),
+	)
+}
+
+func TestUnmarshalBadInput(t *testing.T) {
+	m := &mocker{}
+	defer m.AssertExpectations(t)
+	err := UnmarshalFromString(`{"@type"}`, m)
+	assert.Error(t, err)
+}

dialect/gremlin/encoding/graphson/native.go

@@ -0,0 +1,137 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"fmt"
+	"io"
+	"math"
+	"reflect"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+// EncoderOfNative returns a value encoder of a native type.
+func (encodeExtension) EncoderOfNative(typ reflect2.Type) jsoniter.ValEncoder {
+	switch typ.Kind() {
+	case reflect.Float64:
+		return float64Encoder{typ}
+	default:
+		return nil
+	}
+}
+
+// DecoratorOfNative decorates a value encoder of a native type.
+func (encodeExtension) DecoratorOfNative(typ reflect2.Type, enc jsoniter.ValEncoder) jsoniter.ValEncoder {
+	switch typ.Kind() {
+	case reflect.Bool, reflect.String:
+		return enc
+	case reflect.Int64, reflect.Int, reflect.Uint32:
+		return typeEncoder{enc, int64Type}
+	case reflect.Int32, reflect.Int8, reflect.Uint16:
+		return typeEncoder{enc, int32Type}
+	case reflect.Int16:
+		return typeEncoder{enc, int16Type}
+	case reflect.Uint64, reflect.Uint:
+		return typeEncoder{enc, bigIntegerType}
+	case reflect.Uint8:
+		return typeEncoder{enc, byteType}
+	case reflect.Float32:
+		return typeEncoder{enc, floatType}
+	case reflect.Float64:
+		return typeEncoder{enc, doubleType}
+	default:
+		return nil
+	}
+}
+
+// DecoderOfNative returns a value decoder of a native type.
+func (decodeExtension) DecoderOfNative(typ reflect2.Type) jsoniter.ValDecoder {
+	switch typ.Kind() {
+	case reflect.Float64:
+		return float64Decoder{typ}
+	default:
+		return nil
+	}
+}
+
+// DecoratorOfNative decorates a value decoder of a native type.
+func (decodeExtension) DecoratorOfNative(typ reflect2.Type, dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	switch typ.Kind() {
+	case reflect.Bool:
+		return dec
+	case reflect.String:
+		return typeDecoder{dec, typeCheckerFunc(func(Type) error { return nil })}
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
+		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+		return typeDecoder{dec, integerTypes}
+	case reflect.Float32:
+		return typeDecoder{dec, floatTypes}
+	case reflect.Float64:
+		return typeDecoder{dec, doubleTypes}
+	default:
+		return nil
+	}
+}
+
+type float64Encoder struct {
+	reflect2.Type
+}
+
+func (enc float64Encoder) IsEmpty(ptr unsafe.Pointer) bool {
+	return enc.UnsafeIndirect(ptr).(float64) == 0
+}
+
+func (enc float64Encoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	f := enc.UnsafeIndirect(ptr).(float64)
+	switch {
+	case math.IsNaN(f):
+		stream.WriteString("NaN")
+	case math.IsInf(f, 1):
+		stream.WriteString("Infinity")
+	case math.IsInf(f, -1):
+		stream.WriteString("-Infinity")
+	default:
+		stream.WriteFloat64(f)
+	}
+}
+
+type float64Decoder struct {
+	reflect2.Type
+}
+
+var (
+	integerTypes = Types{byteType, int16Type, int32Type, int64Type, bigIntegerType}
+	floatTypes   = append(integerTypes, floatType, bigDecimal)
+	doubleTypes  = append(floatTypes, doubleType)
+)
+
+func (dec float64Decoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	var val float64
+	switch next := iter.WhatIsNext(); next {
+	case jsoniter.NumberValue:
+		val = iter.ReadFloat64()
+	case jsoniter.StringValue:
+		switch str := iter.ReadString(); str {
+		case "NaN":
+			val = math.NaN()
+		case "Infinity":
+			val = math.Inf(1)
+		case "-Infinity":
+			val = math.Inf(-1)
+		default:
+			iter.ReportError("decode float64", "invalid value "+str)
+		}
+	default:
+		iter.ReportError("decode float64", fmt.Sprintf("unexpected value type: %d", next))
+	}
+
+	if iter.Error == nil || iter.Error == io.EOF {
+		// nolint: gas
+		dec.UnsafeSet(ptr, unsafe.Pointer(&val))
+	}
+}

dialect/gremlin/encoding/graphson/native_test.go

@@ -0,0 +1,299 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"fmt"
+	"math"
+	"testing"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEncodeNative(t *testing.T) {
+	tests := []struct {
+		in      any
+		want    string
+		wantErr bool
+	}{
+		{
+			in:   true,
+			want: "true",
+		},
+		{
+			in:   "hello",
+			want: `"hello"`,
+		},
+		{
+			in: int8(120),
+			want: `{
+				"@type": "g:Int32",
+				"@value": 120
+			}`,
+		},
+		{
+			in: int16(-16),
+			want: `{
+				"@type": "gx:Int16",
+				"@value": -16
+			}`,
+		},
+		{
+			in: int32(3232),
+			want: `{
+				"@type": "g:Int32",
+				"@value": 3232
+			}`,
+		},
+		{
+			in: int64(646464),
+			want: `{
+				"@type": "g:Int64",
+				"@value": 646464
+			}`,
+		},
+		{
+			in: int(127001),
+			want: `{
+				"@type": "g:Int64",
+				"@value": 127001
+			}`,
+		},
+		{
+			in: uint8(81),
+			want: `{
+				"@type": "gx:Byte",
+				"@value": 81
+			}`,
+		},
+		{
+			in: uint16(12345),
+			want: `{
+				"@type": "g:Int32",
+				"@value": 12345
+			}`,
+		},
+		{
+			in: uint32(123454321),
+			want: `{
+				"@type": "g:Int64",
+				"@value": 123454321
+			}`,
+		},
+		{
+			in: uint64(1234567890),
+			want: `{
+				"@type": "gx:BigInteger",
+				"@value": 1234567890
+			}`,
+		},
+		{
+			in: uint(9876543210),
+			want: `{
+				"@type" :"gx:BigInteger",
+				"@value": 9876543210
+			}`,
+		},
+		{
+			in: float32(math.Pi),
+			want: `{
+				"@type": "g:Float",
+				"@value": 3.1415927
+			}`,
+		},
+		{
+			in: float64(math.E),
+			want: `{
+				"@type": "g:Double",
+				"@value": 2.718281828459045
+			}`,
+		},
+		{
+			in: math.NaN(),
+			want: `{
+				"@type": "g:Double",
+				"@value": "NaN"
+			}`,
+		},
+		{
+			in: math.Inf(1),
+			want: `{
+				"@type": "g:Double",
+				"@value": "Infinity"
+			}`,
+		},
+		{
+			in: math.Inf(-1),
+			want: `{
+				"@type": "g:Double",
+				"@value": "-Infinity"
+			}`,
+		},
+		{
+			in: func() *int { v := 7142; return &v }(),
+			want: `{
+				"@type": "g:Int64",
+				"@value": 7142
+			}`,
+		},
+		{
+			in: func() any { v := int16(6116); return &v }(),
+			want: `{
+				"@type": "gx:Int16",
+				"@value": 6116
+			}`,
+		},
+		{
+			in:   nil,
+			want: "null",
+		},
+		{
+			in:      make(chan int),
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(fmt.Sprintf("%T", tc.in), func(t *testing.T) {
+			t.Parallel()
+			got, err := MarshalToString(tc.in)
+			if !tc.wantErr {
+				assert.NoError(t, err)
+				assert.JSONEq(t, tc.want, got)
+			} else {
+				assert.Error(t, err)
+				assert.Empty(t, got)
+			}
+		})
+	}
+}
+
+func TestDecodeNative(t *testing.T) {
+	tests := []struct {
+		in   string
+		want any
+	}{
+		{
+			in:   `{"@type": "g:Float", "@value": 3.14}`,
+			want: float32(3.14),
+		},
+		{
+			in: `{"@type": "g:Float", "@value": "Float"}`,
+		},
+		{
+			in:   `{"@type": "g:Double", "@value": 2.71}`,
+			want: float64(2.71),
+		},
+		{
+			in:   `{"@type": "gx:BigDecimal", "@value": 3.142}`,
+			want: float32(3.142),
+		},
+		{
+			in:   `{"@type": "gx:BigDecimal", "@value": 55512.5176}`,
+			want: float64(55512.5176),
+		},
+		{
+			in:   `{"@type": "g:T", "@value": "world"}`,
+			want: "world",
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(fmt.Sprintf("%T", tc.want), func(t *testing.T) {
+			t.Parallel()
+			if tc.want != nil {
+				typ := reflect2.TypeOf(tc.want)
+				got := typ.New()
+				err := UnmarshalFromString(tc.in, got)
+				require.NoError(t, err)
+				assert.Equal(t, tc.want, typ.Indirect(got))
+			} else {
+				var msg jsoniter.RawMessage
+				err := UnmarshalFromString(tc.in, &msg)
+				assert.Error(t, err)
+			}
+		})
+	}
+}
+
+func TestDecodeTypeMismatch(t *testing.T) {
+	t.Run("FloatToInt", func(t *testing.T) {
+		var v int
+		err := UnmarshalFromString(`{"@type": "g:Float", "@value": 3.14}`, &v)
+		assert.Error(t, err)
+	})
+	t.Run("DoubleToFloat", func(t *testing.T) {
+		var v float32
+		err := UnmarshalFromString(`{"@type": "g:Double", "@value": 5.51}`, &v)
+		assert.Error(t, err)
+	})
+	t.Run("BigDecimalToUint64", func(t *testing.T) {
+		var v uint64
+		err := UnmarshalFromString(`{"@type": "gx:BigDecimal", "@value": 5645.51834}`, &v)
+		assert.Error(t, err)
+	})
+}
+
+func TestDecodeNaNInfinity(t *testing.T) {
+	tests := []struct {
+		data   []byte
+		expect func(*testing.T, float64, error)
+	}{
+		{
+			data: []byte(`{"@type": "g:Double", "@value": "NaN"}`),
+			expect: func(t *testing.T, f float64, err error) {
+				assert.NoError(t, err)
+				assert.True(t, math.IsNaN(f))
+			},
+		},
+		{
+			data: []byte(`{"@type": "g:Double", "@value": "Infinity"}`),
+			expect: func(t *testing.T, f float64, err error) {
+				assert.NoError(t, err)
+				assert.True(t, math.IsInf(f, 1))
+			},
+		},
+		{
+			data: []byte(`{"@type": "g:Double", "@value": "-Infinity"}`),
+			expect: func(t *testing.T, f float64, err error) {
+				assert.NoError(t, err)
+				assert.True(t, math.IsInf(f, -1))
+			},
+		},
+		{
+			data: []byte(`{"@type": "g:Double", "@value": "Junk"}`),
+			expect: func(t *testing.T, _ float64, err error) {
+				assert.Error(t, err)
+			},
+		},
+		{
+			data: []byte(`{"@type": "g:Double", "@value": [42]}`),
+			expect: func(t *testing.T, _ float64, err error) {
+				assert.Error(t, err)
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		var f float64
+		err := Unmarshal(tc.data, &f)
+		tc.expect(t, f, err)
+	}
+}
+
+func TestDecodeTypeDefinition(t *testing.T) {
+	type Status int
+	const StatusOk Status = 42
+
+	var status Status
+	err := UnmarshalFromString(`{"@type": "g:Int64", "@value": 42}`, &status)
+	assert.NoError(t, err)
+	assert.Equal(t, StatusOk, status)
+}

dialect/gremlin/encoding/graphson/raw.go

@@ -0,0 +1,35 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"errors"
+)
+
+// RawMessage is a raw encoded graphson value.
+type RawMessage []byte
+
+// RawMessage must implement Marshaler/Unmarshaler interfaces.
+var (
+	_ Marshaler   = (*RawMessage)(nil)
+	_ Unmarshaler = (*RawMessage)(nil)
+)
+
+// MarshalGraphson returns m as the graphson encoding of m.
+func (m RawMessage) MarshalGraphson() ([]byte, error) {
+	if m == nil {
+		return []byte("null"), nil
+	}
+	return m, nil
+}
+
+// UnmarshalGraphson sets *m to a copy of data.
+func (m *RawMessage) UnmarshalGraphson(data []byte) error {
+	if m == nil {
+		return errors.New("graphson.RawMessage: UnmarshalGraphson on nil pointer")
+	}
+	*m = append((*m)[0:0], data...)
+	return nil
+}

dialect/gremlin/encoding/graphson/raw_test.go

@@ -0,0 +1,32 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRawMessageEncoding(t *testing.T) {
+	var s struct{ M RawMessage }
+	got, err := MarshalToString(s)
+	require.NoError(t, err)
+	assert.Equal(t, `{"M":null}`, got)
+
+	s.M = []byte(`"155a"`)
+	got, err = MarshalToString(s)
+	require.NoError(t, err)
+	assert.JSONEq(t, `{"M": "155a"}`, got)
+
+	err = (*RawMessage)(nil).UnmarshalGraphson(s.M)
+	assert.Error(t, err)
+
+	s.M = nil
+	err = UnmarshalFromString(got, &s)
+	require.NoError(t, err)
+	assert.Equal(t, `"155a"`, string(s.M))
+}

dialect/gremlin/encoding/graphson/slice.go

@@ -0,0 +1,142 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"fmt"
+	"io"
+	"reflect"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+// DecoratorOfSlice decorates a value encoder of a slice type.
+func (encodeExtension) DecoratorOfSlice(typ reflect2.Type, enc jsoniter.ValEncoder) jsoniter.ValEncoder {
+	encoder := typeEncoder{ValEncoder: enc}
+	sliceType := typ.(reflect2.SliceType)
+	if sliceType.Elem().Kind() == reflect.Uint8 {
+		encoder.Type = byteBufferType
+	} else {
+		encoder.Type = listType
+	}
+	return sliceEncoder{sliceType, encoder}
+}
+
+// DecoratorOfArray decorates a value encoder of an array type.
+func (encodeExtension) DecoratorOfArray(enc jsoniter.ValEncoder) jsoniter.ValEncoder {
+	return typeEncoder{enc, listType}
+}
+
+// DecoderOfSlice returns a value decoder of a slice type.
+func (ext decodeExtension) DecoderOfSlice(typ reflect2.Type) jsoniter.ValDecoder {
+	sliceType := typ.(reflect2.SliceType)
+	elemType := sliceType.Elem()
+	if elemType.Kind() == reflect.Uint8 {
+		return nil
+	}
+	return sliceDecoder{
+		sliceType: sliceType,
+		elemDec:   ext.LazyDecoderOf(elemType),
+	}
+}
+
+// DecoderOfArray returns a value decoder of an array type.
+func (ext decodeExtension) DecoderOfArray(typ reflect2.Type) jsoniter.ValDecoder {
+	arrayType := typ.(reflect2.ArrayType)
+	return arrayDecoder{
+		arrayType: arrayType,
+		elemDec:   ext.LazyDecoderOf(arrayType.Elem()),
+	}
+}
+
+// DecoratorOfSlice decorates a value decoder of a slice type.
+func (ext decodeExtension) DecoratorOfSlice(typ reflect2.Type, dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	if typ.(reflect2.SliceType).Elem().Kind() == reflect.Uint8 {
+		return typeDecoder{dec, byteBufferType}
+	}
+	return typeDecoder{dec, listType}
+}
+
+// DecoratorOfArray decorates a value decoder of an array type.
+func (ext decodeExtension) DecoratorOfArray(dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	return typeDecoder{dec, listType}
+}
+
+type sliceEncoder struct {
+	sliceType reflect2.SliceType
+	jsoniter.ValEncoder
+}
+
+func (enc sliceEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	if enc.sliceType.UnsafeIsNil(ptr) {
+		stream.WriteNil()
+	} else {
+		enc.ValEncoder.Encode(ptr, stream)
+	}
+}
+
+type sliceDecoder struct {
+	sliceType reflect2.SliceType
+	elemDec   jsoniter.ValDecoder
+}
+
+func (dec sliceDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	dec.decode(ptr, iter)
+	if iter.Error != nil && iter.Error != io.EOF {
+		iter.Error = fmt.Errorf("decoding slice %s: %w", dec.sliceType, iter.Error)
+	}
+}
+
+func (dec sliceDecoder) decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	sliceType := dec.sliceType
+	if iter.ReadNil() {
+		sliceType.UnsafeSetNil(ptr)
+		return
+	}
+
+	sliceType.UnsafeSet(ptr, sliceType.UnsafeMakeSlice(0, 0))
+	var length int
+
+	iter.ReadArrayCB(func(iter *jsoniter.Iterator) bool {
+		idx := length
+		length++
+		sliceType.UnsafeGrow(ptr, length)
+		elem := sliceType.UnsafeGetIndex(ptr, idx)
+		dec.elemDec.Decode(elem, iter)
+		return iter.Error == nil
+	})
+}
+
+type arrayDecoder struct {
+	arrayType reflect2.ArrayType
+	elemDec   jsoniter.ValDecoder
+}
+
+func (dec arrayDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	dec.decode(ptr, iter)
+	if iter.Error != nil && iter.Error != io.EOF {
+		iter.Error = fmt.Errorf("decoding array %s: %w", dec.arrayType, iter.Error)
+	}
+}
+
+func (dec arrayDecoder) decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	var (
+		arrayType = dec.arrayType
+		length    int
+	)
+	iter.ReadArrayCB(func(iter *jsoniter.Iterator) bool {
+		if length < arrayType.Len() {
+			idx := length
+			length++
+			elem := arrayType.UnsafeGetIndex(ptr, idx)
+			dec.elemDec.Decode(elem, iter)
+		} else {
+			iter.Skip()
+		}
+		return iter.Error == nil
+	})
+}

dialect/gremlin/encoding/graphson/slice_test.go

@@ -0,0 +1,219 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/modern-go/reflect2"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEncodeArray(t *testing.T) {
+	t.Parallel()
+	got, err := MarshalToString([...]string{"a", "b", "c"})
+	require.NoError(t, err)
+	want := `{ "@type": "g:List", "@value": ["a", "b", "c"]}`
+	assert.JSONEq(t, want, got)
+}
+
+func TestEncodeSlice(t *testing.T) {
+	tests := []struct {
+		in   any
+		want string
+	}{
+		{
+			in: []int32{5, 6, 7, 8},
+			want: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Int32",
+						"@value": 5
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 6
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 7
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 8
+					}
+				]
+			}`,
+		},
+		{
+			in: []byte{1, 2, 3, 4, 5},
+			want: `{
+				"@type": "gx:ByteBuffer",
+				"@value": "AQIDBAU="
+			}`,
+		},
+		{
+			in: [...]byte{4, 5},
+			want: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "gx:Byte",
+						"@value": 4
+					},
+					{
+						"@type": "gx:Byte",
+						"@value": 5
+					}
+				]
+			}`,
+		},
+		{
+			in:   []uint64(nil),
+			want: "null",
+		},
+	}
+	for _, tc := range tests {
+		tc := tc
+		t.Run(fmt.Sprintf("%T", tc.in), func(t *testing.T) {
+			t.Parallel()
+			var got bytes.Buffer
+			err := NewEncoder(&got).Encode(tc.in)
+			assert.NoError(t, err)
+			assert.JSONEq(t, tc.want, got.String())
+		})
+	}
+}
+
+func TestDecodeSlice(t *testing.T) {
+	tests := []struct {
+		in   string
+		want any
+	}{
+		{
+			in: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Int32",
+						"@value": 3
+					},
+					{
+						"@type": "g:Int32",
+						"@value": -2
+					},
+					{
+						"@type": "g:Int32",
+						"@value": 1
+					}
+				]
+			}`,
+			want: []int32{3, -2, 1},
+		},
+		{
+			in: `{
+				"@type": "g:List",
+				"@value": ["a", "b", "c"]
+			}`,
+			want: []string{"a", "b", "c"},
+		},
+		{
+			in: `{
+				"@type": "gx:ByteBuffer",
+				"@value": "AQIDBAU="
+			}`,
+			want: []byte{1, 2, 3, 4, 5},
+		},
+		{
+			in: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "gx:Byte",
+						"@value": 42
+					},
+					{
+						"@type": "gx:Byte",
+						"@value": 55
+					},
+					{
+						"@type": "gx:Byte",
+						"@value": 94
+					}
+				]
+			}`,
+			want: [...]byte{42, 55},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(fmt.Sprintf("%T", tc.want), func(t *testing.T) {
+			t.Parallel()
+			typ := reflect2.TypeOf(tc.want)
+			got := typ.New()
+			err := NewDecoder(strings.NewReader(tc.in)).Decode(got)
+			require.NoError(t, err)
+			assert.Equal(t, tc.want, typ.Indirect(got))
+		})
+	}
+}
+
+func TestDecodeBadSlice(t *testing.T) {
+	tests := []struct {
+		name string
+		in   string
+		new  func() any
+	}{
+		{
+			name: "TypeMismatch",
+			in: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Int32",
+						"@value": 3
+					},
+					{
+						"@type": "g:Int64",
+						"@value": 2
+					}
+				]
+			}`,
+			new: func() any { return &[]int{} },
+		},
+		{
+			name: "BadValue",
+			in: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Int32",
+						"@value": 3
+					},
+					{
+						"@type": "g:Int32",
+						"@value": "2"
+					}
+				]
+			}`,
+			new: func() any { return &[2]int{} },
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			err := NewDecoder(strings.NewReader(tc.in)).Decode(tc.new())
+			assert.Error(t, err)
+		})
+	}
+}

dialect/gremlin/encoding/graphson/struct.go

@@ -0,0 +1,39 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import jsoniter "github.com/json-iterator/go"
+
+// DecoratorOfStructField decorates a struct field value encoder.
+func (encodeExtension) DecoratorOfStructField(enc jsoniter.ValEncoder, tag string) jsoniter.ValEncoder {
+	typ, _ := parseTag(tag)
+	if typ == "" {
+		return nil
+	}
+
+	encoder, ok := enc.(typeEncoder)
+	if !ok {
+		encoder = typeEncoder{ValEncoder: enc}
+	}
+	encoder.Type = Type(typ)
+
+	return encoder
+}
+
+// DecoratorOfStructField decorates a struct field value decoder.
+func (decodeExtension) DecoratorOfStructField(dec jsoniter.ValDecoder, tag string) jsoniter.ValDecoder {
+	typ, _ := parseTag(tag)
+	if typ == "" {
+		return nil
+	}
+
+	decoder, ok := dec.(typeDecoder)
+	if !ok {
+		decoder = typeDecoder{ValDecoder: dec}
+	}
+	decoder.typeChecker = Type(typ)
+
+	return decoder
+}

dialect/gremlin/encoding/graphson/struct_test.go

@@ -0,0 +1,205 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"testing"
+
+	"github.com/modern-go/reflect2"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEncodeStruct(t *testing.T) {
+	tests := []struct {
+		name string
+		in   any
+		want string
+	}{
+		{
+			name: "Simple",
+			in: struct {
+				S string
+				I int
+			}{
+				S: "string",
+				I: 1000,
+			},
+			want: `{
+				"S":"string",
+				"I": {
+					"@type": "g:Int64",
+					"@value": 1000
+				}
+			}`,
+		},
+		{
+			name: "Tagged",
+			in: struct {
+				ID   string            `json:"requestId" graphson:"g:UUID"`
+				Seq  int               `json:"seq" graphson:"g:Int32"`
+				Op   string            `json:"op" graphson:","`
+				Args map[string]string `json:"args"`
+			}{
+				ID:  "cb682578-9d92-4499-9ebc-5c6aa73c5397",
+				Seq: 42,
+				Op:  "authentication",
+				Args: map[string]string{
+					"sasl": "AHN0ZXBocGhlbgBwYXNzd29yZA==",
+				},
+			},
+			want: `{
+				"requestId": {
+					"@type": "g:UUID",
+					"@value": "cb682578-9d92-4499-9ebc-5c6aa73c5397"
+				},
+				"seq": {
+					"@type": "g:Int32",
+					"@value": 42
+				},
+				"op": "authentication",
+				"args": {
+					"@type": "g:Map",
+					"@value": ["sasl", "AHN0ZXBocGhlbgBwYXNzd29yZA=="]
+				}
+			}`,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			data, err := MarshalToString(tc.in)
+			require.NoError(t, err)
+			assert.JSONEq(t, tc.want, data)
+		})
+	}
+}
+
+func TestEncodeNestedStruct(t *testing.T) {
+	type S struct {
+		Parent *S  `json:"parent,omitempty"`
+		ID     int `json:"id" graphson:"g:Int32"`
+	}
+
+	v := S{Parent: &S{ID: 1}, ID: 2}
+	want := `{
+		"id": {
+			"@type": "g:Int32",
+			"@value": 2
+		},
+		"parent": {
+			"id": {
+				"@type": "g:Int32",
+				"@value": 1
+			}
+		}
+	}`
+
+	got, err := MarshalToString(&v)
+	require.NoError(t, err)
+	assert.JSONEq(t, want, got)
+}
+
+func TestDecodeStruct(t *testing.T) {
+	tests := []struct {
+		name string
+		in   string
+		want any
+	}{
+		{
+			name: "Simple",
+			in: `{
+				"S":"str",
+				"I": {
+					"@type": "g:Int32",
+					"@value": 9999
+				}
+			}`,
+			want: struct {
+				S string
+				I int32
+			}{
+				S: "str",
+				I: 9999,
+			},
+		},
+		{
+			name: "Tagged",
+			in: `{
+				"requestId": {
+					"@type": "g:UUID",
+					"@value": "cb682578-9d92-4499-9ebc-5c6aa73c5397"
+				},
+				"seq": {
+					"@type": "g:Int32",
+					"@value": 42
+				},
+				"op": "authentication",
+				"args": {
+					"@type": "g:Map",
+					"@value": ["sasl", "AHN0ZXBocGhlbgBwYXNzd29yZA=="]
+				}
+			}`,
+			want: struct {
+				ID   string            `json:"requestId" graphson:"g:UUID"`
+				Seq  int               `json:"seq" graphson:"g:Int32"`
+				Op   string            `json:"op" graphson:","`
+				Args map[string]string `json:"args"`
+			}{
+				ID:  "cb682578-9d92-4499-9ebc-5c6aa73c5397",
+				Seq: 42,
+				Op:  "authentication",
+				Args: map[string]string{
+					"sasl": "AHN0ZXBocGhlbgBwYXNzd29yZA==",
+				},
+			},
+		},
+		{
+			name: "Empty",
+			in:   `{}`,
+			want: struct{}{},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			typ := reflect2.TypeOf(tc.want)
+			got := typ.New()
+			err := UnmarshalFromString(tc.in, got)
+			require.NoError(t, err)
+			assert.Equal(t, tc.want, typ.Indirect(got))
+		})
+	}
+}
+
+func TestDecodeNestedStruct(t *testing.T) {
+	type S struct {
+		Parent *S  `json:"parent,omitempty"`
+		ID     int `json:"id" graphson:"g:Int32"`
+	}
+
+	in := `{
+		"id": {
+			"@type": "g:Int32",
+			"@value": 37
+		},
+		"parent": {
+			"id": {
+				"@type": "g:Int32",
+				"@value": 65
+			}
+		}
+	}`
+	var got S
+	err := UnmarshalFromString(in, &got)
+	require.NoError(t, err)
+
+	want := S{Parent: &S{ID: 65}, ID: 37}
+	assert.Equal(t, want, got)
+}

dialect/gremlin/encoding/graphson/tags.go

@@ -0,0 +1,39 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"strings"
+)
+
+type tagOptions string
+
+// parseTag splits a struct field's graphson tag into its type and
+// comma-separated options.
+func parseTag(tag string) (string, tagOptions) {
+	if idx := strings.Index(tag, ","); idx != -1 {
+		return tag[:idx], tagOptions(tag[idx+1:])
+	}
+	return tag, ""
+}
+
+// Contains reports whether a comma-separated list of options
+// contains a particular substr flag. substr must be surrounded by a
+// string boundary or commas.
+func (opts tagOptions) Contains(opt string) bool {
+	s := string(opts)
+	for s != "" {
+		var next string
+		i := strings.Index(s, ",")
+		if i >= 0 {
+			s, next = s[:i], s[i+1:]
+		}
+		if s == opt {
+			return true
+		}
+		s = next
+	}
+	return false
+}

dialect/gremlin/encoding/graphson/tags_test.go

@@ -0,0 +1,60 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestParseTag(t *testing.T) {
+	tests := []struct {
+		name string
+		tag  string
+		typ  string
+		opts tagOptions
+	}{
+		{
+			name: "Empty",
+		},
+		{
+			name: "TypeOnly",
+			tag:  "g:Int32",
+			typ:  "g:Int32",
+		},
+		{
+			name: "OptsOnly",
+			tag:  ",opt1,opt2",
+			opts: "opt1,opt2",
+		},
+		{
+			name: "TypeAndOpts",
+			tag:  "g:UUID,opt3,opt4",
+			typ:  "g:UUID",
+			opts: "opt3,opt4",
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			typ, opts := parseTag(tc.tag)
+			assert.Equal(t, tc.typ, typ)
+			assert.Equal(t, tc.opts, opts)
+		})
+	}
+}
+
+func TestTagOptionsContains(t *testing.T) {
+	_, opts := parseTag(",opt1,opt2,opt3")
+	assert.True(t, opts.Contains("opt1"))
+	assert.True(t, opts.Contains("opt2"))
+	assert.True(t, opts.Contains("opt3"))
+	assert.False(t, opts.Contains("opt4"))
+	assert.False(t, opts.Contains("opt11"))
+	assert.False(t, opts.Contains(""))
+}

dialect/gremlin/encoding/graphson/time.go

@@ -0,0 +1,34 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"time"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+)
+
+func init() {
+	RegisterTypeEncoder("time.Time", typeEncoder{timeCodec{}, Timestamp})
+	RegisterTypeDecoder("time.Time", typeDecoder{timeCodec{}, Types{Timestamp, Date}})
+}
+
+type timeCodec struct{}
+
+func (timeCodec) IsEmpty(ptr unsafe.Pointer) bool {
+	ts := *((*time.Time)(ptr))
+	return ts.IsZero()
+}
+
+func (timeCodec) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	ts := *((*time.Time)(ptr))
+	stream.WriteInt64(ts.UnixNano() / time.Millisecond.Nanoseconds())
+}
+
+func (timeCodec) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	ns := iter.ReadInt64() * time.Millisecond.Nanoseconds()
+	*((*time.Time)(ptr)) = time.Unix(0, ns)
+}

dialect/gremlin/encoding/graphson/time_test.go

@@ -0,0 +1,36 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTimeEncoding(t *testing.T) {
+	const ms = 1481750076295
+	ts := time.Unix(0, ms*time.Millisecond.Nanoseconds())
+
+	for _, v := range []any{ts, &ts} {
+		got, err := MarshalToString(v)
+		require.NoError(t, err)
+		assert.JSONEq(t, `{ "@type": "g:Timestamp", "@value": 1481750076295 }`, got)
+	}
+
+	strs := []string{
+		`{ "@type": "g:Timestamp", "@value": 1481750076295 }`,
+		`{ "@type": "g:Date", "@value": 1481750076295 }`,
+	}
+
+	for _, str := range strs {
+		var v time.Time
+		err := UnmarshalFromString(str, &v)
+		assert.NoError(t, err)
+		assert.True(t, ts.Equal(v))
+	}
+}

dialect/gremlin/encoding/graphson/type.go

@@ -0,0 +1,155 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/modern-go/reflect2"
+)
+
+// A Type is a graphson type.
+type Type string
+
+// graphson typed value types.
+const (
+	// core
+	doubleType Type = "g:Double"
+	floatType  Type = "g:Float"
+	int32Type  Type = "g:Int32"
+	int64Type  Type = "g:Int64"
+	listType   Type = "g:List"
+	mapType    Type = "g:Map"
+	Timestamp  Type = "g:Timestamp"
+	Date       Type = "g:Date"
+
+	// extended
+	bigIntegerType Type = "gx:BigInteger"
+	bigDecimal     Type = "gx:BigDecimal"
+	byteType       Type = "gx:Byte"
+	byteBufferType Type = "gx:ByteBuffer"
+	int16Type      Type = "gx:Int16"
+)
+
+// String implements fmt.Stringer interface.
+func (typ Type) String() string {
+	return string(typ)
+}
+
+// CheckType implements typeChecker interface.
+func (typ Type) CheckType(other Type) error {
+	if typ != other {
+		return fmt.Errorf("expect type %s, but found %s", typ, other)
+	}
+	return nil
+}
+
+// Types is a slice of Type.
+type Types []Type
+
+// Contains reports whether a slice of types contains a particular type.
+func (types Types) Contains(typ Type) bool {
+	for i := range types {
+		if types[i] == typ {
+			return true
+		}
+	}
+	return false
+}
+
+// String implements fmt.Stringer interface.
+func (types Types) String() string {
+	var builder strings.Builder
+	builder.WriteByte('[')
+	for i := range types {
+		if i > 0 {
+			builder.WriteByte(',')
+		}
+		builder.WriteString(types[i].String())
+	}
+	builder.WriteByte(']')
+	return builder.String()
+}
+
+// CheckType implements typeChecker interface.
+func (types Types) CheckType(typ Type) error {
+	if !types.Contains(typ) {
+		return fmt.Errorf("expect any of %s, but found %s", types, typ)
+	}
+	return nil
+}
+
+// Typer is the interface implemented by types that
+// define an underlying graphson type.
+type Typer interface {
+	GraphsonType() Type
+}
+
+var typerType = reflect2.TypeOfPtr((*Typer)(nil)).Elem()
+
+// DecoratorOfTyper decorates a value encoder of a Typer interface.
+func (ext encodeExtension) DecoratorOfTyper(typ reflect2.Type, enc jsoniter.ValEncoder) jsoniter.ValEncoder {
+	if typ.Kind() != reflect.Struct {
+		return nil
+	}
+	if typ.Implements(typerType) {
+		return typerEncoder{
+			typeEncoder: typeEncoder{ValEncoder: enc},
+			typerOf: func(ptr unsafe.Pointer) Typer {
+				return typ.UnsafeIndirect(ptr).(Typer)
+			},
+		}
+	}
+	ptrType := reflect2.PtrTo(typ)
+	if ptrType.Implements(typerType) {
+		return typerEncoder{
+			typeEncoder: typeEncoder{ValEncoder: enc},
+			typerOf: func(ptr unsafe.Pointer) Typer {
+				// nolint: gas
+				return ptrType.UnsafeIndirect(unsafe.Pointer(&ptr)).(Typer)
+			},
+		}
+	}
+	return nil
+}
+
+// DecoratorOfTyper decorates a value decoder of a Typer interface.
+func (ext decodeExtension) DecoratorOfTyper(typ reflect2.Type, dec jsoniter.ValDecoder) jsoniter.ValDecoder {
+	ptrType := reflect2.PtrTo(typ)
+	if ptrType.Implements(typerType) {
+		return typerDecoder{
+			typeDecoder: typeDecoder{ValDecoder: dec},
+			typerOf: func(ptr unsafe.Pointer) Typer {
+				// nolint: gas
+				return ptrType.UnsafeIndirect(unsafe.Pointer(&ptr)).(Typer)
+			},
+		}
+	}
+	return nil
+}
+
+type typerEncoder struct {
+	typeEncoder
+	typerOf func(unsafe.Pointer) Typer
+}
+
+func (enc typerEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	enc.typeEncoder.Type = enc.typerOf(ptr).GraphsonType()
+	enc.typeEncoder.Encode(ptr, stream)
+}
+
+type typerDecoder struct {
+	typeDecoder
+	typerOf func(unsafe.Pointer) Typer
+}
+
+func (dec typerDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	dec.typeDecoder.typeChecker = dec.typerOf(ptr).GraphsonType()
+	dec.typeDecoder.Decode(ptr, iter)
+}

dialect/gremlin/encoding/graphson/type_test.go

@@ -0,0 +1,95 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestTypeCheckType(t *testing.T) {
+	assert.NoError(t, int32Type.CheckType(int32Type))
+	assert.Error(t, int32Type.CheckType(int64Type))
+}
+
+func TestTypesCheckType(t *testing.T) {
+	assert.NoError(t, Types{int16Type, int32Type, int64Type}.CheckType(int32Type))
+	assert.Error(t, Types{floatType, doubleType}.CheckType(bigIntegerType))
+}
+
+func TestTypesString(t *testing.T) {
+	assert.Equal(t, "[]", Types{}.String())
+	assert.Equal(t, "[gx:Byte]", Types{byteType}.String())
+	assert.Equal(t, "[gx:Int16,g:Int32,g:Int64]", Types{int16Type, int32Type, int64Type}.String())
+}
+
+type vertex struct {
+	ID    int    `json:"id"`
+	Label string `json:"label"`
+}
+
+func (vertex) GraphsonType() Type {
+	return Type("g:Vertex")
+}
+
+type mockVertex struct {
+	mock.Mock `json:"-"`
+	ID        int    `json:"id"`
+	Label     string `json:"label"`
+}
+
+func (m *mockVertex) GraphsonType() Type {
+	return m.Called().Get(0).(Type)
+}
+
+func TestEncodeTyper(t *testing.T) {
+	m := &mockVertex{ID: 42, Label: "person"}
+	m.On("GraphsonType").Return(Type("g:Vertex")).Twice()
+	defer m.AssertExpectations(t)
+
+	v := vertex{ID: m.ID, Label: m.Label}
+	var vv Typer = v
+
+	want := `{
+	   "@type": "g:Vertex",
+	   "@value": {
+		   "id": {
+			   "@type": "g:Int64",
+			   "@value": 42
+		   },
+		   "label": "person"
+	   }
+	}`
+
+	for _, tc := range []any{m, &m, v, vv, &vv} {
+		got, err := MarshalToString(tc)
+		assert.NoError(t, err)
+		assert.JSONEq(t, want, got)
+	}
+}
+
+func TestDecodeTyper(t *testing.T) {
+	var m mockVertex
+	m.On("GraphsonType").Return(Type("g:Vertex")).Once()
+	defer m.AssertExpectations(t)
+
+	in := `{
+	   "@type": "g:Vertex",
+	   "@value": {
+		   "id": {
+			   "@type": "g:Int64",
+			   "@value": 55
+		   },
+		   "label": "user"
+	   }
+	}`
+
+	err := UnmarshalFromString(in, &m)
+	assert.NoError(t, err)
+	assert.Equal(t, 55, m.ID)
+	assert.Equal(t, "user", m.Label)
+}

dialect/gremlin/encoding/graphson/util.go

@@ -0,0 +1,114 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"errors"
+	"io"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+)
+
+// graphson encoding type / value keys
+const (
+	TypeKey  = "@type"
+	ValueKey = "@value"
+)
+
+// typeEncoder adds graphson type information to a value encoder.
+type typeEncoder struct {
+	jsoniter.ValEncoder
+	Type Type
+}
+
+// Encode belongs to jsoniter.ValEncoder interface.
+func (enc typeEncoder) Encode(ptr unsafe.Pointer, stream *jsoniter.Stream) {
+	stream.WriteObjectStart()
+	stream.WriteObjectField(TypeKey)
+	stream.WriteString(enc.Type.String())
+	stream.WriteMore()
+	stream.WriteObjectField(ValueKey)
+	enc.ValEncoder.Encode(ptr, stream)
+	stream.WriteObjectEnd()
+}
+
+type (
+	// typeDecoder decorates a value decoder and adds graphson type verification.
+	typeDecoder struct {
+		jsoniter.ValDecoder
+		typeChecker
+	}
+
+	// typeChecker defines an interface for graphson type verification.
+	typeChecker interface {
+		CheckType(Type) error
+	}
+
+	// typeCheckerFunc allows the use of functions as type checkers.
+	typeCheckerFunc func(Type) error
+
+	// typeValue defines a graphson type / value pair.
+	typeValue struct {
+		Type  Type
+		Value jsoniter.RawMessage
+	}
+)
+
+// Decode belongs to jsoniter.ValDecoder interface.
+func (dec typeDecoder) Decode(ptr unsafe.Pointer, iter *jsoniter.Iterator) {
+	if iter.WhatIsNext() != jsoniter.ObjectValue {
+		dec.ValDecoder.Decode(ptr, iter)
+		return
+	}
+
+	data := iter.SkipAndReturnBytes()
+	if iter.Error != nil && iter.Error != io.EOF {
+		return
+	}
+
+	var tv typeValue
+	if err := jsoniter.Unmarshal(data, &tv); err != nil {
+		iter.ReportError("unmarshal type value", err.Error())
+		return
+	}
+
+	if err := dec.CheckType(tv.Type); err != nil {
+		iter.ReportError("check type", err.Error())
+		return
+	}
+
+	it := config.BorrowIterator(tv.Value)
+	defer config.ReturnIterator(it)
+
+	dec.ValDecoder.Decode(ptr, it)
+	if it.Error != nil && it.Error != io.EOF {
+		iter.ReportError("decode value", it.Error.Error())
+	}
+}
+
+// UnmarshalJSON implements json.Unmarshaler interface.
+func (tv *typeValue) UnmarshalJSON(data []byte) error {
+	var v struct {
+		Type  *Type               `json:"@type"`
+		Value jsoniter.RawMessage `json:"@value"`
+	}
+
+	if err := jsoniter.Unmarshal(data, &v); err != nil {
+		return err
+	}
+	if v.Type == nil || v.Value == nil {
+		return errors.New("missing type or value")
+	}
+
+	tv.Type = *v.Type
+	tv.Value = v.Value
+	return nil
+}
+
+// CheckType implements typeChecker interface.
+func (f typeCheckerFunc) CheckType(typ Type) error {
+	return f(typ)
+}

dialect/gremlin/encoding/graphson/util_test.go

@@ -0,0 +1,129 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graphson
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"testing"
+	"unsafe"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTypeEncode(t *testing.T) {
+	var got bytes.Buffer
+	stream := config.BorrowStream(&got)
+	defer config.ReturnStream(stream)
+
+	typ, val := int32Type, 42
+	ptr := unsafe.Pointer(&val)
+
+	var m mocker
+	m.On("Encode", ptr, stream).
+		Run(func(args mock.Arguments) {
+			stream := args.Get(1).(*jsoniter.Stream)
+			stream.WriteInt(val)
+		}).
+		Once()
+	defer m.AssertExpectations(t)
+
+	typeEncoder{&m, typ}.Encode(ptr, stream)
+	require.NoError(t, stream.Flush())
+	want := fmt.Sprintf(`{"@type": "%s", "@value": %d}`, typ, val)
+	assert.JSONEq(t, want, got.String())
+}
+
+func TestTypeDecode(t *testing.T) {
+	typ, val := int64Type, 84
+	ptr := unsafe.Pointer(&val)
+
+	data := fmt.Sprintf(`{"@value": %d, "@type": "%s"}`, val, typ)
+	iter := config.BorrowIterator([]byte(data))
+	defer config.ReturnIterator(iter)
+
+	m := &mocker{}
+	m.On("CheckType", typ).
+		Return(nil).
+		Once()
+	m.On("Decode", ptr, mock.Anything).
+		Run(func(args mock.Arguments) {
+			iter := args.Get(1).(*jsoniter.Iterator)
+			assert.Equal(t, val, iter.ReadInt())
+		}).
+		Once()
+	defer m.AssertExpectations(t)
+
+	typeDecoder{m, m}.Decode(ptr, iter)
+	assert.NoError(t, iter.Error)
+}
+
+func TestTypeDecodeBadType(t *testing.T) {
+	typ, val := int64Type, 55
+	ptr := unsafe.Pointer(&val)
+
+	m := &mocker{}
+	m.On("CheckType", typ).Return(errors.New("bad type")).Once()
+	defer m.AssertExpectations(t)
+
+	data := fmt.Sprintf(`{"@type": "%s", "@value": %d}`, typ, val)
+	iter := config.BorrowIterator([]byte(data))
+	defer config.ReturnIterator(iter)
+
+	typeDecoder{m, m}.Decode(ptr, iter)
+	require.Error(t, iter.Error)
+	assert.Contains(t, iter.Error.Error(), "bad type")
+}
+
+func TestTypeDecodeDuplicateField(t *testing.T) {
+	data := `{"@type": "gx:Byte", "@value": 33, "@type": "g:Int32"}`
+	iter := config.BorrowIterator([]byte(data))
+	defer config.ReturnIterator(iter)
+	var ptr unsafe.Pointer
+
+	m := &mocker{}
+	m.On("CheckType", mock.MatchedBy(func(typ Type) bool { return typ == int32Type })).
+		Return(nil).
+		Once()
+	m.On("Decode", ptr, mock.Anything).
+		Run(func(args mock.Arguments) {
+			args.Get(1).(*jsoniter.Iterator).Skip()
+			require.NoError(t, iter.Error)
+		}).
+		Once()
+	defer m.AssertExpectations(t)
+
+	typeDecoder{m, m}.Decode(ptr, iter)
+	assert.NoError(t, iter.Error)
+}
+
+func TestTypeDecodeMissingField(t *testing.T) {
+	data := `{"@type": "g:Int32"}`
+	iter := config.BorrowIterator([]byte(data))
+	defer config.ReturnIterator(iter)
+
+	m := &mocker{}
+	defer m.AssertExpectations(t)
+
+	typeDecoder{m, m}.Decode(nil, iter)
+	require.Error(t, iter.Error)
+	assert.Contains(t, iter.Error.Error(), "missing type or value")
+}
+
+func TestTypeDecodeSyntaxError(t *testing.T) {
+	data := `{"@type": "gx:Int16", "@value", 65000}`
+	iter := config.BorrowIterator([]byte(data))
+	defer config.ReturnIterator(iter)
+
+	m := &mocker{}
+	defer m.AssertExpectations(t)
+
+	typeDecoder{m, m}.Decode(nil, iter)
+	assert.Error(t, iter.Error)
+}

dialect/gremlin/encoding/mime.go

@@ -0,0 +1,30 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package encoding
+
+import (
+	"bytes"
+)
+
+// Mime defines a gremlin mime type.
+type Mime []byte
+
+// GraphSON3Mime mime headers.
+var (
+	GraphSON3Mime = NewMime("application/vnd.gremlin-v3.0+json")
+)
+
+// NewMime creates a wire format mime header.
+func NewMime(s string) Mime {
+	var buf bytes.Buffer
+	buf.WriteByte(byte(len(s)))
+	buf.WriteString(s)
+	return buf.Bytes()
+}
+
+// String implements fmt.Stringer interface.
+func (m Mime) String() string {
+	return string(m[1:])
+}

dialect/gremlin/encoding/mime_test.go

@@ -0,0 +1,21 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package encoding
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewMime(t *testing.T) {
+	str := "application/vnd.gremlin-v2.0+json"
+	mime := NewMime(str)
+	require.Len(t, mime, len(str)+1)
+	assert.EqualValues(t, len(str), mime[0])
+	assert.EqualValues(t, str, mime[1:])
+	assert.Equal(t, str, mime.String())
+}

dialect/gremlin/example_test.go

@@ -0,0 +1,46 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"flag"
+	"log"
+	"os"
+	"time"
+)
+
+func ExampleClient_Query() {
+	addr := flag.String("gremlin-server", os.Getenv("GREMLIN_SERVER"), "gremlin server address")
+	flag.Parse()
+
+	if *addr == "" {
+		log.Fatal("missing gremlin server address")
+	}
+
+	client, err := NewHTTPClient(*addr, nil)
+	if err != nil {
+		log.Fatalf("creating client: %v", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+
+	rsp, err := client.Query(ctx, "g.E()")
+	if err != nil {
+		log.Fatalf("executing query: %v", err)
+	}
+
+	edges, err := rsp.ReadEdges()
+	if err != nil {
+		log.Fatalf("unmashal edges")
+	}
+
+	defer cancel()
+
+	for _, e := range edges {
+		log.Println(e.String())
+	}
+	// - Output:
+}

dialect/gremlin/expand.go

@@ -0,0 +1,47 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"fmt"
+	"sort"
+	"strings"
+
+	jsoniter "github.com/json-iterator/go"
+)
+
+// ExpandBindings expands the given RoundTripper and expands the request bindings into the Gremlin traversal.
+func ExpandBindings(rt RoundTripper) RoundTripper {
+	return RoundTripperFunc(func(ctx context.Context, r *Request) (*Response, error) {
+		bindings, ok := r.Arguments[ArgsBindings]
+		if !ok {
+			return rt.RoundTrip(ctx, r)
+		}
+		query, ok := r.Arguments[ArgsGremlin]
+		if !ok {
+			return rt.RoundTrip(ctx, r)
+		}
+		{
+			query, bindings := query.(string), bindings.(map[string]any)
+			keys := make(sort.StringSlice, 0, len(bindings))
+			for k := range bindings {
+				keys = append(keys, k)
+			}
+			sort.Sort(sort.Reverse(keys))
+			kv := make([]string, 0, len(bindings)*2)
+			for _, k := range keys {
+				s, err := jsoniter.MarshalToString(bindings[k])
+				if err != nil {
+					return nil, fmt.Errorf("marshal bindings value for key %s: %w", k, err)
+				}
+				kv = append(kv, k, s)
+			}
+			delete(r.Arguments, ArgsBindings)
+			r.Arguments[ArgsGremlin] = strings.NewReplacer(kv...).Replace(query)
+		}
+		return rt.RoundTrip(ctx, r)
+	})
+}

dialect/gremlin/expand_test.go

@@ -0,0 +1,71 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestExpandBindings(t *testing.T) {
+	tests := []struct {
+		req       *Request
+		wantErr   bool
+		wantQuery string
+	}{
+		{
+			req:       NewEvalRequest("no bindings"),
+			wantQuery: "no bindings",
+		},
+		{
+			req:       NewEvalRequest("g.V($0)", WithBindings(map[string]any{"$0": 1})),
+			wantQuery: "g.V(1)",
+		},
+		{
+			req:       NewEvalRequest("g.V().has($1, $2)", WithBindings(map[string]any{"$1": "name", "$2": "a8m"})),
+			wantQuery: "g.V().has(\"name\", \"a8m\")",
+		},
+		{
+			req:       NewEvalRequest("g.V().limit(n)", WithBindings(map[string]any{"n": 10})),
+			wantQuery: "g.V().limit(10)",
+		},
+		{
+			req:     NewEvalRequest("g.V()", WithBindings(map[string]any{"$0": func() {}})),
+			wantErr: true,
+		},
+		{
+			req:       NewEvalRequest("g.V().has($0, $1)", WithBindings(map[string]any{"$0": "active", "$1": true})),
+			wantQuery: "g.V().has(\"active\", true)",
+		},
+		{
+			req:       NewEvalRequest("g.V().has($1, $11)", WithBindings(map[string]any{"$1": "active", "$11": true})),
+			wantQuery: "g.V().has(\"active\", true)",
+		},
+	}
+	for i, tt := range tests {
+		tt := tt
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			rt := ExpandBindings(RoundTripperFunc(func(ctx context.Context, r *Request) (*Response, error) {
+				assert.Equal(t, tt.wantQuery, r.Arguments[ArgsGremlin])
+				return nil, nil
+			}))
+			_, err := rt.RoundTrip(context.Background(), tt.req)
+			assert.Equal(t, tt.wantErr, err != nil)
+		})
+	}
+}
+
+func TestExpandBindingsNoQuery(t *testing.T) {
+	rt := ExpandBindings(RoundTripperFunc(func(ctx context.Context, r *Request) (*Response, error) {
+		return nil, nil
+	}))
+	_, err := rt.RoundTrip(context.Background(), &Request{Arguments: map[string]any{
+		ArgsBindings: map[string]any{},
+	}})
+	assert.NoError(t, err)
+}

dialect/gremlin/graph/dsl/__/dsl.go

@@ -0,0 +1,75 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package __
+
+import "entgo.io/ent/dialect/gremlin/graph/dsl"
+
+// As is the api for calling __.As().
+func As(args ...any) *dsl.Traversal { return New().As(args...) }
+
+// Is is the api for calling __.Is().
+func Is(args ...any) *dsl.Traversal { return New().Is(args...) }
+
+// Not is the api for calling __.Not().
+func Not(args ...any) *dsl.Traversal { return New().Not(args...) }
+
+// Has is the api for calling __.Has().
+func Has(args ...any) *dsl.Traversal { return New().Has(args...) }
+
+// HasNot is the api for calling __.HasNot().
+func HasNot(args ...any) *dsl.Traversal { return New().HasNot(args...) }
+
+// Or is the api for calling __.Or().
+func Or(args ...any) *dsl.Traversal { return New().Or(args...) }
+
+// And is the api for calling __.And().
+func And(args ...any) *dsl.Traversal { return New().And(args...) }
+
+// In is the api for calling __.In().
+func In(args ...any) *dsl.Traversal { return New().In(args...) }
+
+// Out is the api for calling __.Out().
+func Out(args ...any) *dsl.Traversal { return New().Out(args...) }
+
+// OutE is the api for calling __.OutE().
+func OutE(args ...any) *dsl.Traversal { return New().OutE(args...) }
+
+// InE is the api for calling __.InE().
+func InE(args ...any) *dsl.Traversal { return New().InE(args...) }
+
+// InV is the api for calling __.InV().
+func InV(args ...any) *dsl.Traversal { return New().InV(args...) }
+
+// V is the api for calling __.V().
+func V(args ...any) *dsl.Traversal { return New().V(args...) }
+
+// OutV is the api for calling __.OutV().
+func OutV(args ...any) *dsl.Traversal { return New().OutV(args...) }
+
+// Values is the api for calling __.Values().
+func Values(args ...string) *dsl.Traversal { return New().Values(args...) }
+
+// Union is the api for calling __.Union().
+func Union(args ...any) *dsl.Traversal { return New().Union(args...) }
+
+// Constant is the api for calling __.Constant().
+func Constant(args ...any) *dsl.Traversal { return New().Constant(args...) }
+
+// Properties is the api for calling __.Properties().
+func Properties(args ...any) *dsl.Traversal { return New().Properties(args...) }
+
+// OtherV is the api for calling __.OtherV().
+func OtherV() *dsl.Traversal { return New().OtherV() }
+
+// Count is the api for calling __.Count().
+func Count() *dsl.Traversal { return New().Count() }
+
+// Drop is the api for calling __.Drop().
+func Drop() *dsl.Traversal { return New().Drop() }
+
+// Fold is the api for calling __.Fold().
+func Fold() *dsl.Traversal { return New().Fold() }
+
+func New() *dsl.Traversal { return new(dsl.Traversal).Add(dsl.Token("__")) }

dialect/gremlin/graph/dsl/dsl.go

@@ -0,0 +1,214 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+// Package dsl provide an API for writing gremlin dsl queries almost as-is
+// in Go without using strings in the code.
+//
+// Note that, the API is not type-safe and assume the provided query and
+// its arguments are valid.
+package dsl
+
+import (
+	"fmt"
+	"strings"
+	"time"
+)
+
+// Node represents a DSL step in the traversal.
+type Node interface {
+	// Code returns the code representation of the element and its bindings (if any).
+	Code() (string, []any)
+}
+
+type (
+	// Token holds a simple token, like assignment.
+	Token string
+	// List represents a list of elements.
+	List struct {
+		Elements []any
+	}
+	// Func represents a function call.
+	Func struct {
+		Name string
+		Args []any
+	}
+	// Block represents a block/group of nodes.
+	Block struct {
+		Nodes []any
+	}
+	// Var represents a variable assignment and usage.
+	Var struct {
+		Name string
+		Elem any
+	}
+)
+
+// Code stringified the token.
+func (t Token) Code() (string, []any) { return string(t), nil }
+
+// Code returns the code representation of a list.
+func (l List) Code() (string, []any) {
+	c, args := codeList(", ", l.Elements...)
+	return fmt.Sprintf("[%s]", c), args
+}
+
+// Code returns the code representation of a function call.
+func (f Func) Code() (string, []any) {
+	c, args := codeList(", ", f.Args...)
+	return fmt.Sprintf("%s(%s)", f.Name, c), args
+}
+
+// Code returns the code representation of group/block of nodes.
+func (b Block) Code() (string, []any) {
+	return codeList("; ", b.Nodes...)
+}
+
+// Code returns the code representation of variable declaration or its identifier.
+func (v Var) Code() (string, []any) {
+	c, args := code(v.Elem)
+	if v.Name == "" {
+		return c, args
+	}
+	return fmt.Sprintf("%s = %s", v.Name, c), args
+}
+
+// predefined nodes.
+var (
+	G   = Token("g")
+	Dot = Token(".")
+)
+
+// NewFunc returns a new function node.
+func NewFunc(name string, args ...any) *Func {
+	return &Func{Name: name, Args: args}
+}
+
+// NewList returns a new list node.
+func NewList(args ...any) *List {
+	return &List{Elements: args}
+}
+
+// Querier is the interface that wraps the Query method.
+type Querier interface {
+	// Query returns the query-string (similar to the Gremlin byte-code) and its bindings.
+	Query() (string, Bindings)
+}
+
+// Bindings are used to associate a variable with a value.
+type Bindings map[string]any
+
+// Add adds new value to the bindings map, formats it if needed, and returns its generated name.
+func (b Bindings) Add(v any) string {
+	k := fmt.Sprintf("$%x", len(b))
+	switch v := v.(type) {
+	case time.Time:
+		b[k] = v.UnixNano()
+	default:
+		b[k] = v
+	}
+	return k
+}
+
+// Cardinality of vertex properties.
+type Cardinality string
+
+// Cardinality options.
+const (
+	Set    Cardinality = "set"
+	Single Cardinality = "single"
+)
+
+// Code implements the Node interface.
+func (c Cardinality) Code() (string, []any) { return string(c), nil }
+
+// Keyword defines a Gremlin keyword.
+type Keyword string
+
+// Keyword options.
+const (
+	ID Keyword = "id"
+)
+
+// Code implements the Node interface.
+func (k Keyword) Code() (string, []any) { return string(k), nil }
+
+// Order of vertex properties.
+type Order string
+
+// Order options.
+const (
+	Incr    Order = "incr"
+	Decr    Order = "decr"
+	Shuffle Order = "shuffle"
+)
+
+// Code implements the Node interface.
+func (o Order) Code() (string, []any) { return string(o), nil }
+
+// Column references a particular type of column in a complex data structure such as a Map, a Map.Entry, or a Path.
+type Column string
+
+// Column options.
+const (
+	Keys   Column = "keys"
+	Values Column = "values"
+)
+
+// Code implements the Node interface.
+func (o Column) Code() (string, []any) { return string(o), nil }
+
+// Scope used for steps that have a variable scope which alter the manner in which the step will behave in relation to how the traverses are processed.
+type Scope string
+
+// Scope options.
+const (
+	Local  Scope = "local"
+	Global Scope = "global"
+)
+
+// Code implements the Node interface.
+func (s Scope) Code() (string, []any) { return string(s), nil }
+
+func codeList(sep string, vs ...any) (string, []any) {
+	var (
+		br   strings.Builder
+		args []any
+	)
+	for i, node := range vs {
+		if i > 0 {
+			br.WriteString(sep)
+		}
+		c, nargs := code(node)
+		br.WriteString(c)
+		args = append(args, nargs...)
+	}
+	return br.String(), args
+}
+
+func code(v any) (string, []any) {
+	switch n := v.(type) {
+	case Node:
+		return n.Code()
+	case *Traversal:
+		var (
+			b    strings.Builder
+			args []any
+		)
+		for i := range n.nodes {
+			code, nargs := n.nodes[i].Code()
+			b.WriteString(code)
+			args = append(args, nargs...)
+		}
+		return b.String(), args
+	default:
+		return "%s", []any{v}
+	}
+}
+
+func sface(args []string) (v []any) {
+	for _, s := range args {
+		v = append(v, s)
+	}
+	return
+}

dialect/gremlin/graph/dsl/dsl_test.go

@@ -0,0 +1,238 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package dsl_test
+
+import (
+	"strconv"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin/graph/dsl"
+	"entgo.io/ent/dialect/gremlin/graph/dsl/__"
+	"entgo.io/ent/dialect/gremlin/graph/dsl/g"
+	"entgo.io/ent/dialect/gremlin/graph/dsl/p"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTraverse(t *testing.T) {
+	tests := []struct {
+		input     dsl.Querier
+		wantQuery string
+		wantBinds dsl.Bindings
+	}{
+		{
+			input:     g.V(5),
+			wantQuery: "g.V($0)",
+			wantBinds: dsl.Bindings{"$0": 5},
+		},
+		{
+			input:     g.V(2).Both("knows"),
+			wantQuery: "g.V($0).both($1)",
+			wantBinds: dsl.Bindings{"$0": 2, "$1": "knows"},
+		},
+		{
+			input:     g.V(49).BothE("knows").OtherV().ValueMap(),
+			wantQuery: "g.V($0).bothE($1).otherV().valueMap()",
+			wantBinds: dsl.Bindings{"$0": 49, "$1": "knows"},
+		},
+		{
+			input:     g.AddV("person").Property("name", "a8m").Next(),
+			wantQuery: "g.addV($0).property($1, $2).next()",
+			wantBinds: dsl.Bindings{"$0": "person", "$1": "name", "$2": "a8m"},
+		},
+		{
+			input: dsl.Each([]any{1, 2, 3}, func(it *dsl.Traversal) *dsl.Traversal {
+				return g.V(it)
+			}),
+			wantQuery: "[$0, $1, $2].each { g.V(it) }",
+			wantBinds: dsl.Bindings{"$0": 1, "$1": 2, "$2": 3},
+		},
+		{
+			input: dsl.Each([]any{g.V(1).Next()}, func(it *dsl.Traversal) *dsl.Traversal {
+				return it.ID()
+			}),
+			wantQuery: "[g.V($0).next()].each { it.id() }",
+			wantBinds: dsl.Bindings{"$0": 1},
+		},
+		{
+			input:     g.AddV("person").AddE("knows").To(g.V(2)),
+			wantQuery: "g.addV($0).addE($1).to(g.V($2))",
+			wantBinds: dsl.Bindings{"$0": "person", "$1": "knows", "$2": 2},
+		},
+		{
+			input: func() *dsl.Traversal {
+				v1 := g.V(2).Next()
+				v2 := g.AddV("person").Property("name", "a8m")
+				e1 := g.V(v1).AddE("knows").To(v2)
+				return dsl.Group(v1, v2, e1)
+			}(),
+			wantQuery: "t0 = g.V($0).next(); t1 = g.addV($1).property($2, $3); t2 = g.V(t0).addE($4).to(t1); t2",
+			wantBinds: dsl.Bindings{"$0": 2, "$1": "person", "$2": "name", "$3": "a8m", "$4": "knows"},
+		},
+		{
+			input: func() *dsl.Traversal {
+				v1 := g.AddV("person")
+				each := dsl.Each([]any{1, 2, 3}, func(it *dsl.Traversal) *dsl.Traversal {
+					return g.V(v1).AddE("knows").To(g.V(it)).Next()
+				})
+				return dsl.Group(v1, each)
+			}(),
+			wantQuery: "t0 = g.addV($0); t1 = [$1, $2, $3].each { g.V(t0).addE($4).to(g.V(it)).next() }; t1",
+			wantBinds: dsl.Bindings{"$0": "person", "$1": 1, "$2": 2, "$3": 3, "$4": "knows"},
+		},
+		{
+			input: g.V().HasLabel("person").
+				Choose(__.Values("age").Is(p.LTE(20))),
+			wantQuery: "g.V().hasLabel($0).choose(__.values($1).is(lte($2)))",
+			wantBinds: dsl.Bindings{"$0": "person", "$1": "age", "$2": 20},
+		},
+		{
+			input:     g.AddV("person").Property("name", "a8m").Properties(),
+			wantQuery: "g.addV($0).property($1, $2).properties()",
+			wantBinds: dsl.Bindings{"$0": "person", "$1": "name", "$2": "a8m"},
+		},
+		{
+			input: func() *dsl.Traversal {
+				v1 := g.AddV("person").Next()
+				e1 := g.V(v1).AddE("knows").To(g.V(2).Next())
+				return dsl.Group(v1, e1, g.V(v1).ValueMap(true))
+			}(),
+			wantQuery: "t0 = g.addV($0).next(); t1 = g.V(t0).addE($1).to(g.V($2).next()); t2 = g.V(t0).valueMap($3); t2",
+			wantBinds: dsl.Bindings{"$0": "person", "$1": "knows", "$2": 2, "$3": true},
+		},
+		{
+			input: func() *dsl.Traversal {
+				vs := g.V().HasLabel("person").ToList()
+				edge := g.V(vs).AddE("assoc").To(g.V(1)).Iterate()
+				each := dsl.Each(vs, func(it *dsl.Traversal) *dsl.Traversal {
+					return g.V(1).AddE("inverse").To(it).Next()
+				})
+				return dsl.Group(vs, edge, each)
+			}(),
+			wantQuery: "t0 = g.V().hasLabel($0).toList(); t1 = g.V(t0).addE($1).to(g.V($2)).iterate(); t2 = t0.each { g.V($3).addE($4).to(it).next() }; t2",
+			wantBinds: dsl.Bindings{"$0": "person", "$1": "assoc", "$2": 1, "$3": 1, "$4": "inverse"},
+		},
+		{
+			input:     g.V().Where(__.Or(__.Has("age", 29), __.Has("age", 30))),
+			wantQuery: "g.V().where(__.or(__.has($0, $1), __.has($2, $3)))",
+			wantBinds: dsl.Bindings{"$0": "age", "$1": 29, "$2": "age", "$3": 30},
+		},
+		{
+			input:     g.V().Has("name", p.Containing("le")).Has("name", p.StartingWith("A")),
+			wantQuery: `g.V().has($0, containing($1)).has($2, startingWith($3))`,
+			wantBinds: dsl.Bindings{"$0": "name", "$1": "le", "$2": "name", "$3": "A"},
+		},
+		{
+			input:     g.AddV().Property(dsl.Single, "age", 32).ValueMap(),
+			wantQuery: "g.addV().property(single, $0, $1).valueMap()",
+			wantBinds: dsl.Bindings{"$0": "age", "$1": 32},
+		},
+		{
+			input:     g.V().Count(),
+			wantQuery: "g.V().count()",
+			wantBinds: dsl.Bindings{},
+		},
+		{
+			input:     g.V().HasNot("age"),
+			wantQuery: "g.V().hasNot($0)",
+			wantBinds: dsl.Bindings{"$0": "age"},
+		},
+		{
+			input: func() *dsl.Traversal {
+				v := g.V().HasID(1)
+				u := v.Clone().InE().Drop()
+				return dsl.Join(v, u)
+			}(),
+			wantQuery: "g.V().hasId($0); g.V().hasId($1).inE().drop()",
+			wantBinds: dsl.Bindings{"$0": 1, "$1": 1},
+		},
+		{
+			input: func() *dsl.Traversal {
+				v := g.V().HasID(1)
+				u := v.Clone().InE().Drop()
+				w := u.Clone()
+				return dsl.Join(v, u, w)
+			}(),
+			wantQuery: "g.V().hasId($0); g.V().hasId($1).inE().drop(); g.V().hasId($2).inE().drop()",
+			wantBinds: dsl.Bindings{"$0": 1, "$1": 1, "$2": 1},
+		},
+		{
+			input:     g.V().OutE("knows").Where(__.InV().Has("name", "a8m")).OutV(),
+			wantQuery: "g.V().outE($0).where(__.inV().has($1, $2)).outV()",
+			wantBinds: dsl.Bindings{"$0": "knows", "$1": "name", "$2": "a8m"},
+		},
+		{
+			input:     g.V().Has("name", p.Within("a8m", "alex")),
+			wantQuery: "g.V().has($0, within($1, $2))",
+			wantBinds: dsl.Bindings{"$0": "name", "$1": "a8m", "$2": "alex"},
+		},
+		{
+			input:     g.V().HasID(p.Within(1, 2)),
+			wantQuery: "g.V().hasId(within($0, $1))",
+			wantBinds: dsl.Bindings{"$0": 1, "$1": 2},
+		},
+		{
+			input:     g.V().HasID(p.Without(1, 2)),
+			wantQuery: "g.V().hasId(without($0, $1))",
+			wantBinds: dsl.Bindings{"$0": 1, "$1": 2},
+		},
+		{
+			input:     g.V().Order().By("name"),
+			wantQuery: "g.V().order().by($0)",
+			wantBinds: dsl.Bindings{"$0": "name"},
+		},
+		{
+			input:     g.V().Order().By("name", dsl.Incr),
+			wantQuery: "g.V().order().by($0, incr)",
+			wantBinds: dsl.Bindings{"$0": "name"},
+		},
+		{
+			input:     g.V().Order().By("name", dsl.Incr).Undo(),
+			wantQuery: "g.V().order()",
+			wantBinds: dsl.Bindings{},
+		},
+		{
+			input:     g.V().OutE("knows").Where(__.InV().Has("name", "a8m")).Undo(),
+			wantQuery: "g.V().outE($0)",
+			wantBinds: dsl.Bindings{"$0": "knows"},
+		},
+		{
+			input:     g.V().Has("name").Group().By("name").By("age").Select(dsl.Values),
+			wantQuery: "g.V().has($0).group().by($1).by($2).select(values)",
+			wantBinds: dsl.Bindings{"$0": "name", "$1": "name", "$2": "age"},
+		},
+		{
+			input:     g.V().Fold().Unfold(),
+			wantQuery: "g.V().fold().unfold()",
+			wantBinds: dsl.Bindings{},
+		},
+		{
+			input: g.V().Has("person", "name", "a8m").Count().Coalesce(
+				__.Is(p.NEQ(0)).Constant("unique constraint failed"),
+				g.AddV("person").Property("name", "a8m").ValueMap(true),
+			),
+			wantQuery: "g.V().has($0, $1, $2).count().coalesce(__.is(neq($3)).constant($4), g.addV($5).property($6, $7).valueMap($8))",
+			wantBinds: dsl.Bindings{"$0": "person", "$1": "name", "$2": "a8m", "$3": 0, "$4": "unique constraint failed", "$5": "person", "$6": "name", "$7": "a8m", "$8": true},
+		},
+		{
+			input:     g.V().Has("age").Property("age", __.Union(__.Values("age"), __.Constant(10)).Sum()).ValueMap(),
+			wantQuery: "g.V().has($0).property($1, __.union(__.values($2), __.constant($3)).sum()).valueMap()",
+			wantBinds: dsl.Bindings{"$0": "age", "$1": "age", "$2": "age", "$3": 10},
+		},
+		{
+			input:     g.V().Has("age").SideEffect(__.Properties("name").Drop()).ValueMap(),
+			wantQuery: "g.V().has($0).sideEffect(__.properties($1).drop()).valueMap()",
+			wantBinds: dsl.Bindings{"$0": "age", "$1": "name"},
+		},
+	}
+	for i, tt := range tests {
+		tt := tt
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			query, bindings := tt.input.Query()
+			require.Equal(t, tt.wantQuery, query)
+			require.Equal(t, tt.wantBinds, bindings)
+		})
+	}
+}

dialect/gremlin/graph/dsl/g/g.go

@@ -0,0 +1,19 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package g
+
+import "entgo.io/ent/dialect/gremlin/graph/dsl"
+
+// V is the api for calling g.V().
+func V(args ...any) *dsl.Traversal { return dsl.NewTraversal().V(args...) }
+
+// E is the api for calling g.E().
+func E(args ...any) *dsl.Traversal { return dsl.NewTraversal().E(args...) }
+
+// AddV is the api for calling g.AddV().
+func AddV(args ...any) *dsl.Traversal { return dsl.NewTraversal().AddV(args...) }
+
+// AddE is the api for calling g.AddE().
+func AddE(args ...any) *dsl.Traversal { return dsl.NewTraversal().AddE(args...) }

dialect/gremlin/graph/dsl/p/p.go

@@ -0,0 +1,93 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package p
+
+import (
+	"entgo.io/ent/dialect/gremlin/graph/dsl"
+)
+
+// EQ is the equal predicate.
+func EQ(v any) *dsl.Traversal {
+	return op("eq", v)
+}
+
+// NEQ is the not-equal predicate.
+func NEQ(v any) *dsl.Traversal {
+	return op("neq", v)
+}
+
+// GT is the greater than predicate.
+func GT(v any) *dsl.Traversal {
+	return op("gt", v)
+}
+
+// GTE is the greater than or equal predicate.
+func GTE(v any) *dsl.Traversal {
+	return op("gte", v)
+}
+
+// LT is the less than predicate.
+func LT(v any) *dsl.Traversal {
+	return op("lt", v)
+}
+
+// LTE is the less than or equal predicate.
+func LTE(v any) *dsl.Traversal {
+	return op("lte", v)
+}
+
+// Between is the between/contains predicate.
+func Between(v, u any) *dsl.Traversal {
+	return op("between", v, u)
+}
+
+// StartingWith is the prefix test predicate.
+func StartingWith(prefix string) *dsl.Traversal {
+	return op("startingWith", prefix)
+}
+
+// EndingWith is the suffix test predicate.
+func EndingWith(suffix string) *dsl.Traversal {
+	return op("endingWith", suffix)
+}
+
+// Containing is the sub string test predicate.
+func Containing(substr string) *dsl.Traversal {
+	return op("containing", substr)
+}
+
+// NotStartingWith is the negation of StartingWith.
+func NotStartingWith(prefix string) *dsl.Traversal {
+	return op("notStartingWith", prefix)
+}
+
+// NotEndingWith is the negation of EndingWith.
+func NotEndingWith(suffix string) *dsl.Traversal {
+	return op("notEndingWith", suffix)
+}
+
+// NotContaining is the negation of Containing.
+func NotContaining(substr string) *dsl.Traversal {
+	return op("notContaining", substr)
+}
+
+// Within Determines if a value is within the specified list of values.
+func Within[T any](args ...T) *dsl.Traversal {
+	return op("within", args...)
+}
+
+// Without determines if a value is not within the specified list of values.
+func Without[T any](args ...T) *dsl.Traversal {
+	return op("without", args...)
+}
+
+func op[T any](name string, args ...T) *dsl.Traversal {
+	t := &dsl.Traversal{}
+	vs := make([]any, len(args))
+	for i, arg := range args {
+		vs[i] = arg
+	}
+	return t.Add(dsl.NewFunc(name, vs...))
+}

dialect/gremlin/graph/dsl/traversal.go

@@ -0,0 +1,426 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package dsl
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+)
+
+// Traversal mimics the TinkerPop graph traversal.
+type Traversal struct {
+	// nodes holds the dsl nodes. first element is the reference name
+	// of the TinkerGraph. defaults to "g".
+	nodes []Node
+	errs  []error
+}
+
+// NewTraversal returns a new default traversal with "g" as a reference name to the Graph.
+func NewTraversal() *Traversal {
+	return &Traversal{nodes: []Node{G}}
+}
+
+// Group groups a list of traversals into one. all traversals are assigned into a temporary
+// variables named by their index. The last variable functions as a return value of the query.
+// Note that, this "temporary hack" is not perfect and may not work in some cases because of
+// the limitation of evaluation order.
+func Group(trs ...*Traversal) *Traversal {
+	var (
+		b     = Block{}
+		names = make(map[*Traversal]Token)
+	)
+	for i, tr := range trs {
+		if _, ok := names[tr]; ok {
+			continue
+		}
+		v := &Var{Name: fmt.Sprintf("t%d", i), Elem: &Traversal{nodes: tr.nodes}}
+		b.Nodes = append(b.Nodes, v)
+		names[tr] = Token(v.Name)
+	}
+	for _, tr := range trs {
+		tr.nodes = []Node{names[tr]}
+	}
+	b.Nodes = append(b.Nodes, names[trs[len(trs)-1]])
+	return &Traversal{nodes: []Node{b}}
+}
+
+// Join joins a list of traversals with a semicolon separator.
+func Join(trs ...*Traversal) *Traversal {
+	b := Block{}
+	for _, tr := range trs {
+		b.Nodes = append(b.Nodes, &Traversal{nodes: tr.nodes})
+	}
+	return &Traversal{nodes: []Node{b}}
+}
+
+// AddError adds an error to the traversal.
+func (t *Traversal) AddError(err error) *Traversal {
+	t.errs = append(t.errs, err)
+	return t
+}
+
+// Err returns a concatenated error of all errors encountered during
+// the query-building, or were added manually by calling AddError.
+func (t *Traversal) Err() error {
+	if len(t.errs) == 0 {
+		return nil
+	}
+	br := strings.Builder{}
+	for i := range t.errs {
+		if i > 0 {
+			br.WriteString("; ")
+		}
+		br.WriteString(t.errs[i].Error())
+	}
+	return errors.New(br.String())
+}
+
+// V step is usually used to start a traversal but it may also be used mid-traversal.
+func (t *Traversal) V(args ...any) *Traversal {
+	t.Add(Dot, NewFunc("V", args...))
+	return t
+}
+
+// OtherV maps the Edge to the incident vertex that was not just traversed from in the path history.
+func (t *Traversal) OtherV() *Traversal {
+	t.Add(Dot, NewFunc("otherV"))
+	return t
+}
+
+// E step is usually used to start a traversal but it may also be used mid-traversal.
+func (t *Traversal) E(args ...any) *Traversal {
+	t.Add(Dot, NewFunc("E", args...))
+	return t
+}
+
+// AddV adds a vertex.
+func (t *Traversal) AddV(args ...any) *Traversal {
+	t.Add(Dot, NewFunc("addV", args...))
+	return t
+}
+
+// AddE adds an edge.
+func (t *Traversal) AddE(args ...any) *Traversal {
+	t.Add(Dot, NewFunc("addE", args...))
+	return t
+}
+
+// Next gets the next n-number of results from the traversal.
+func (t *Traversal) Next() *Traversal {
+	return t.Add(Dot, NewFunc("next"))
+}
+
+// Drop removes elements and properties from the graph.
+func (t *Traversal) Drop() *Traversal {
+	return t.Add(Dot, NewFunc("drop"))
+}
+
+// Property sets a Property value and related meta properties if supplied,
+// if supported by the Graph and if the Element is a VertexProperty.
+func (t *Traversal) Property(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("property", args...))
+}
+
+// Both maps the Vertex to its adjacent vertices given the edge labels.
+func (t *Traversal) Both(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("both", args...))
+}
+
+// BothE maps the Vertex to its incident edges given the edge labels.
+func (t *Traversal) BothE(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("bothE", args...))
+}
+
+// Has filters vertices, edges and vertex properties based on their properties.
+// See: http://tinkerpop.apache.org/docs/current/reference/#has-step.
+func (t *Traversal) Has(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("has", args...))
+}
+
+// HasNot filters vertices, edges and vertex properties based on the non-existence of properties.
+// See: http://tinkerpop.apache.org/docs/current/reference/#has-step.
+func (t *Traversal) HasNot(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("hasNot", args...))
+}
+
+// HasID filters vertices, edges and vertex properties based on their identifier.
+func (t *Traversal) HasID(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("hasId", args...))
+}
+
+// HasLabel filters vertices, edges and vertex properties based on their label.
+func (t *Traversal) HasLabel(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("hasLabel", args...))
+}
+
+// HasNext returns true if the iteration has more elements.
+func (t *Traversal) HasNext() *Traversal {
+	return t.Add(Dot, NewFunc("hasNext"))
+}
+
+// Match maps the Traverser to a Map of bindings as specified by the provided match traversals.
+func (t *Traversal) Match(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("match", args...))
+}
+
+// Choose routes the current traverser to a particular traversal branch option which allows the creation of if-then-else like semantics within a traversal.
+func (t *Traversal) Choose(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("choose", args...))
+}
+
+// Select arbitrary values from the traversal.
+func (t *Traversal) Select(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("select", args...))
+}
+
+// Group organizes objects in the stream into a Map.Calls to group() are typically accompanied with by() modulators which help specify how the grouping should occur.
+func (t *Traversal) Group() *Traversal {
+	return t.Add(Dot, NewFunc("group"))
+}
+
+// Values maps the Element to the values of the associated properties given the provide property keys.
+func (t *Traversal) Values(args ...string) *Traversal {
+	return t.Add(Dot, NewFunc("values", sface(args)...))
+}
+
+// ValueMap maps the Element to a Map of the property values key'd according to their Property.key().
+func (t *Traversal) ValueMap(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("valueMap", args...))
+}
+
+// Properties maps the Element to its associated properties given the provide property keys.
+func (t *Traversal) Properties(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("properties", args...))
+}
+
+// Range filters the objects in the traversal by the number of them to pass through the stream.
+func (t *Traversal) Range(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("range", args...))
+}
+
+// Limit filters the objects in the traversal by the number of them to pass through the stream, where only the first n objects are allowed as defined by the limit argument.
+func (t *Traversal) Limit(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("limit", args...))
+}
+
+// ID maps the Element to its Element.id().
+func (t *Traversal) ID() *Traversal {
+	return t.Add(Dot, NewFunc("id"))
+}
+
+// Label maps the Element to its Element.label().
+func (t *Traversal) Label() *Traversal {
+	return t.Add(Dot, NewFunc("label"))
+}
+
+// From provides from()-modulation to respective steps.
+func (t *Traversal) From(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("from", args...))
+}
+
+// To used as a modifier to addE(String) this method specifies the traversal to use for selecting the incoming vertex of the newly added Edge.
+func (t *Traversal) To(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("to", args...))
+}
+
+// As provides a label to the step that can be accessed later in the traversal by other steps.
+func (t *Traversal) As(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("as", args...))
+}
+
+// Or ensures that at least one of the provided traversals yield a result.
+func (t *Traversal) Or(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("or", args...))
+}
+
+// And ensures that all of the provided traversals yield a result.
+func (t *Traversal) And(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("and", args...))
+}
+
+// Is filters the E object if it is not P.eq(V) to the provided value.
+func (t *Traversal) Is(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("is", args...))
+}
+
+// Not removes objects from the traversal stream when the traversal provided as an argument does not return any objects.
+func (t *Traversal) Not(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("not", args...))
+}
+
+// In maps the Vertex to its incoming adjacent vertices given the edge labels.
+func (t *Traversal) In(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("in", args...))
+}
+
+// Where filters the current object based on the object itself or the path history.
+func (t *Traversal) Where(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("where", args...))
+}
+
+// Out maps the Vertex to its outgoing adjacent vertices given the edge labels.
+func (t *Traversal) Out(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("out", args...))
+}
+
+// OutE maps the Vertex to its outgoing incident edges given the edge labels.
+func (t *Traversal) OutE(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("outE", args...))
+}
+
+// InE maps the Vertex to its incoming incident edges given the edge labels.
+func (t *Traversal) InE(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("inE", args...))
+}
+
+// OutV maps the Edge to its outgoing/tail incident Vertex.
+func (t *Traversal) OutV(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("outV", args...))
+}
+
+// InV maps the Edge to its incoming/head incident Vertex.
+func (t *Traversal) InV(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("inV", args...))
+}
+
+// ToList puts all the results into a Groovy list.
+func (t *Traversal) ToList() *Traversal {
+	return t.Add(Dot, NewFunc("toList"))
+}
+
+// Iterate iterates the traversal presumably for the generation of side-effects.
+func (t *Traversal) Iterate() *Traversal {
+	return t.Add(Dot, NewFunc("iterate"))
+}
+
+// Count maps the traversal stream to its reduction as a sum of the Traverser.bulk() values
+// (i.e. count the number of traversers up to this point).
+func (t *Traversal) Count(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("count", args...))
+}
+
+// Order all the objects in the traversal up to this point and then emit them one-by-one in their ordered sequence.
+func (t *Traversal) Order(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("order", args...))
+}
+
+// By can be applied to a number of different step to alter their behaviors.
+// This form is essentially an identity() modulation.
+func (t *Traversal) By(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("by", args...))
+}
+
+// Fold rolls up objects in the stream into an aggregate list..
+func (t *Traversal) Fold() *Traversal {
+	return t.Add(Dot, NewFunc("fold"))
+}
+
+// Unfold unrolls a Iterator, Iterable or Map into a linear form or simply emits the object if it is not one of those types.
+func (t *Traversal) Unfold() *Traversal {
+	return t.Add(Dot, NewFunc("unfold"))
+}
+
+// Sum maps the traversal stream to its reduction as a sum of the Traverser.get() values multiplied by their Traverser.bulk().
+func (t *Traversal) Sum(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("sum", args...))
+}
+
+// Mean determines the mean value in the stream.
+func (t *Traversal) Mean(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("mean", args...))
+}
+
+// Min determines the smallest value in the stream.
+func (t *Traversal) Min(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("min", args...))
+}
+
+// Max determines the greatest value in the stream.
+func (t *Traversal) Max(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("max", args...))
+}
+
+// Coalesce evaluates the provided traversals and returns the result of the first traversal to emit at least one object.
+func (t *Traversal) Coalesce(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("coalesce", args...))
+}
+
+// Dedup removes all duplicates in the traversal stream up to this point.
+func (t *Traversal) Dedup(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("dedup", args...))
+}
+
+// Constant maps any object to a fixed E value.
+func (t *Traversal) Constant(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("constant", args...))
+}
+
+// Union merges the results of an arbitrary number of traversals.
+func (t *Traversal) Union(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("union", args...))
+}
+
+// SideEffect allows the traverser to proceed unchanged, but yield some computational
+// sideEffect in the process.
+func (t *Traversal) SideEffect(args ...any) *Traversal {
+	return t.Add(Dot, NewFunc("sideEffect", args...))
+}
+
+// Each is a Groovy each-loop function.
+func Each(v any, cb func(it *Traversal) *Traversal) *Traversal {
+	t := &Traversal{}
+	switch v := v.(type) {
+	case *Traversal:
+		t.Add(&Var{Elem: v})
+	case []any:
+		t.Add(NewList(v...))
+	default:
+		t.Add(Token("undefined"))
+	}
+	t.Add(Dot, Token("each"), Token(" { "))
+	t.Add(cb(&Traversal{nodes: []Node{Token("it")}}).nodes...)
+	t.Add(Token(" }"))
+	return t
+}
+
+// Add is the public API for adding new nodes to the traversal by its sub packages.
+func (t *Traversal) Add(n ...Node) *Traversal {
+	t.nodes = append(t.nodes, n...)
+	return t
+}
+
+// Query returns the query-representation and its binding of this traversal object.
+func (t *Traversal) Query() (string, Bindings) {
+	var (
+		names    []any
+		query    strings.Builder
+		bindings = Bindings{}
+	)
+	for _, n := range t.nodes {
+		code, args := n.Code()
+		query.WriteString(code)
+		for _, arg := range args {
+			names = append(names, bindings.Add(arg))
+		}
+	}
+	return fmt.Sprintf(query.String(), names...), bindings
+}
+
+// Clone creates a deep copy of an existing traversal.
+func (t *Traversal) Clone() *Traversal {
+	if t == nil {
+		return nil
+	}
+	return &Traversal{nodes: append(make([]Node, 0, len(t.nodes)), t.nodes...)}
+}
+
+// Undo reverts the last-step of the traversal.
+func (t *Traversal) Undo() *Traversal {
+	if n := len(t.nodes); n > 2 {
+		t.nodes = t.nodes[:n-2]
+	}
+	return t
+}

dialect/gremlin/graph/edge.go

@@ -0,0 +1,94 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graph
+
+import (
+	"fmt"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+)
+
+type (
+	// An Edge between two vertices.
+	Edge struct {
+		Element
+		OutV, InV Vertex
+	}
+
+	// graphson edge repr.
+	edge struct {
+		Element
+		OutV      any    `json:"outV"`
+		OutVLabel string `json:"outVLabel"`
+		InV       any    `json:"inV"`
+		InVLabel  string `json:"inVLabel"`
+	}
+)
+
+// NewEdge create a new graph edge.
+func NewEdge(id any, label string, outV, inV Vertex) Edge {
+	return Edge{
+		Element: NewElement(id, label),
+		OutV:    outV,
+		InV:     inV,
+	}
+}
+
+// String implements fmt.Stringer interface.
+func (e Edge) String() string {
+	return fmt.Sprintf("e[%v][%v-%s->%v]", e.ID, e.OutV.ID, e.Label, e.InV.ID)
+}
+
+// MarshalGraphson implements graphson.Marshaler interface.
+func (e Edge) MarshalGraphson() ([]byte, error) {
+	return graphson.Marshal(edge{
+		Element:   e.Element,
+		OutV:      e.OutV.ID,
+		OutVLabel: e.OutV.Label,
+		InV:       e.InV.ID,
+		InVLabel:  e.InV.Label,
+	})
+}
+
+// UnmarshalGraphson implements graphson.Unmarshaler interface.
+func (e *Edge) UnmarshalGraphson(data []byte) error {
+	var edge edge
+	if err := graphson.Unmarshal(data, &edge); err != nil {
+		return fmt.Errorf("unmarshalling edge: %w", err)
+	}
+
+	*e = NewEdge(
+		edge.ID, edge.Label,
+		NewVertex(edge.OutV, edge.OutVLabel),
+		NewVertex(edge.InV, edge.InVLabel),
+	)
+	return nil
+}
+
+// GraphsonType implements graphson.Typer interface.
+func (edge) GraphsonType() graphson.Type {
+	return "g:Edge"
+}
+
+// Property denotes a key/value pair associated with an edge.
+type Property struct {
+	Key   string `json:"key"`
+	Value any    `json:"value"`
+}
+
+// NewProperty create a new graph edge property.
+func NewProperty(key string, value any) Property {
+	return Property{key, value}
+}
+
+// GraphsonType implements graphson.Typer interface.
+func (Property) GraphsonType() graphson.Type {
+	return "g:Property"
+}
+
+// String implements fmt.Stringer interface.
+func (p Property) String() string {
+	return fmt.Sprintf("p[%s->%v]", p.Key, p.Value)
+}

dialect/gremlin/graph/edge_test.go

@@ -0,0 +1,108 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graph
+
+import (
+	"fmt"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEdgeString(t *testing.T) {
+	e := NewEdge(
+		13, "develops",
+		NewVertex(1, ""),
+		NewVertex(10, ""),
+	)
+	assert.Equal(t, "e[13][1-develops->10]", fmt.Sprint(e))
+}
+
+func TestEdgeEncoding(t *testing.T) {
+	t.Parallel()
+
+	e := NewEdge(13, "develops",
+		NewVertex(1, "person"),
+		NewVertex(10, "software"),
+	)
+	got, err := graphson.MarshalToString(e)
+	require.NoError(t, err)
+
+	want := `{
+		"@type" : "g:Edge",
+		"@value" : {
+			"id" : {
+			  "@type" : "g:Int64",
+			  "@value" : 13
+			},
+			"label" : "develops",
+			"inVLabel" : "software",
+			"outVLabel" : "person",
+			"inV" : {
+			  "@type" : "g:Int64",
+			  "@value" : 10
+			},
+			"outV" : {
+			  "@type" : "g:Int64",
+			  "@value" : 1
+			}
+		}
+	}`
+	assert.JSONEq(t, want, got)
+
+	e = Edge{}
+	err = graphson.UnmarshalFromString(got, &e)
+	require.NoError(t, err)
+
+	assert.Equal(t, NewElement(int64(13), "develops"), e.Element)
+	assert.Equal(t, NewVertex(int64(1), "person"), e.OutV)
+	assert.Equal(t, NewVertex(int64(10), "software"), e.InV)
+}
+
+func TestPropertyEncoding(t *testing.T) {
+	t.Parallel()
+
+	props := []Property{
+		NewProperty("from", int32(2017)),
+		NewProperty("to", int32(2019)),
+	}
+	got, err := graphson.MarshalToString(props)
+	require.NoError(t, err)
+
+	want := `{
+		"@type" : "g:List",
+		"@value" : [
+			{
+				"@type" : "g:Property",
+				"@value" : {
+					"key" : "from",
+					"value" : {
+						"@type" : "g:Int32",
+						"@value" : 2017
+					}
+				}
+			},
+			{
+				"@type" : "g:Property",
+				"@value" : {
+					"key" : "to",
+					"value" : {
+						"@type" : "g:Int32",
+						"@value" : 2019
+					}
+				}
+			}
+		]
+	}`
+	assert.JSONEq(t, want, got)
+}
+
+func TestPropertyString(t *testing.T) {
+	p := NewProperty("since", 2019)
+	assert.Equal(t, "p[since->2019]", fmt.Sprint(p))
+}

dialect/gremlin/graph/element.go

@@ -0,0 +1,16 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graph
+
+// Element defines a base struct for graph elements.
+type Element struct {
+	ID    any    `json:"id"`
+	Label string `json:"label"`
+}
+
+// NewElement create a new graph element.
+func NewElement(id any, label string) Element {
+	return Element{id, label}
+}

dialect/gremlin/graph/valuemap.go

@@ -0,0 +1,57 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graph
+
+import (
+	"errors"
+	"fmt"
+	"reflect"
+
+	"github.com/mitchellh/mapstructure"
+)
+
+// ValueMap models a .valueMap() gremlin response.
+type ValueMap []map[string]any
+
+// Decode decodes a value map into v.
+func (m ValueMap) Decode(v any) error {
+	rv := reflect.ValueOf(v)
+	if rv.Kind() != reflect.Ptr {
+		return errors.New("cannot unmarshal into a non pointer")
+	}
+	if rv.IsNil() {
+		return errors.New("cannot unmarshal into a nil pointer")
+	}
+
+	if rv.Elem().Kind() != reflect.Slice {
+		v = &[]any{v}
+	}
+	return m.decode(v)
+}
+
+func (m ValueMap) decode(v any) error {
+	cfg := mapstructure.DecoderConfig{
+		DecodeHook: func(f, t reflect.Kind, data any) (any, error) {
+			if f == reflect.Slice && t != reflect.Slice {
+				rv := reflect.ValueOf(data)
+				if rv.Len() == 1 {
+					data = rv.Index(0).Interface()
+				}
+			}
+			return data, nil
+		},
+		Result:  v,
+		TagName: "json",
+	}
+
+	dec, err := mapstructure.NewDecoder(&cfg)
+	if err != nil {
+		return fmt.Errorf("creating structure decoder: %w", err)
+	}
+	if err := dec.Decode(m); err != nil {
+		return fmt.Errorf("decoding value map: %w", err)
+	}
+	return nil
+}

dialect/gremlin/graph/valuemap_test.go

@@ -0,0 +1,74 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graph
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestValueMapDecodeOne(t *testing.T) {
+	vm := ValueMap{map[string]any{
+		"id":    int64(1),
+		"label": "person",
+		"name":  []any{"marko"},
+		"age":   []any{int32(29)},
+	}}
+
+	var ent struct {
+		ID    uint64 `json:"id"`
+		Label string `json:"label"`
+		Name  string `json:"name"`
+		Age   uint8  `json:"age"`
+	}
+	err := vm.Decode(&ent)
+	require.NoError(t, err)
+
+	assert.Equal(t, uint64(1), ent.ID)
+	assert.Equal(t, "person", ent.Label)
+	assert.Equal(t, "marko", ent.Name)
+	assert.Equal(t, uint8(29), ent.Age)
+}
+
+func TestValueMapDecodeMany(t *testing.T) {
+	vm := ValueMap{
+		map[string]any{
+			"id":    int64(1),
+			"label": "person",
+			"name":  []any{"chico"},
+		},
+		map[string]any{
+			"id":    int64(2),
+			"label": "person",
+			"name":  []any{"dico"},
+		},
+	}
+
+	ents := []struct {
+		ID    int    `json:"id"`
+		Label string `json:"label"`
+		Name  string `json:"name"`
+	}{}
+	err := vm.Decode(&ents)
+	require.NoError(t, err)
+
+	require.Len(t, ents, 2)
+	assert.Equal(t, 1, ents[0].ID)
+	assert.Equal(t, "person", ents[0].Label)
+	assert.Equal(t, "chico", ents[0].Name)
+	assert.Equal(t, 2, ents[1].ID)
+	assert.Equal(t, "person", ents[1].Label)
+	assert.Equal(t, "dico", ents[1].Name)
+}
+
+func TestValueMapDecodeBadInput(t *testing.T) {
+	type s struct{ Name string }
+	err := ValueMap{}.Decode(s{})
+	assert.Error(t, err)
+	err = ValueMap{}.Decode((*s)(nil))
+	assert.Error(t, err)
+}

dialect/gremlin/graph/vertex.go

@@ -0,0 +1,62 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graph
+
+import (
+	"fmt"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+)
+
+// Vertex represents a graph vertex.
+type Vertex struct {
+	Element
+}
+
+// NewVertex create a new graph vertex.
+func NewVertex(id any, label string) Vertex {
+	if label == "" {
+		label = "vertex"
+	}
+	return Vertex{
+		Element: NewElement(id, label),
+	}
+}
+
+// GraphsonType implements graphson.Typer interface.
+func (Vertex) GraphsonType() graphson.Type {
+	return "g:Vertex"
+}
+
+// String implements fmt.Stringer interface.
+func (v Vertex) String() string {
+	return fmt.Sprintf("v[%v]", v.ID)
+}
+
+// VertexProperty denotes a key/value pair associated with a vertex.
+type VertexProperty struct {
+	ID    any    `json:"id"`
+	Key   string `json:"label"`
+	Value any    `json:"value"`
+}
+
+// NewVertexProperty create a new graph vertex property.
+func NewVertexProperty(id any, key string, value any) VertexProperty {
+	return VertexProperty{
+		ID:    id,
+		Key:   key,
+		Value: value,
+	}
+}
+
+// GraphsonType implements graphson.Typer interface.
+func (VertexProperty) GraphsonType() graphson.Type {
+	return "g:VertexProperty"
+}
+
+// String implements fmt.Stringer interface.
+func (vp VertexProperty) String() string {
+	return fmt.Sprintf("vp[%s->%v]", vp.Key, vp.Value)
+}

dialect/gremlin/graph/vertex_test.go

@@ -0,0 +1,86 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package graph
+
+import (
+	"fmt"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestVertexCreation(t *testing.T) {
+	v := NewVertex(45, "person")
+	assert.Equal(t, 45, v.ID)
+	assert.Equal(t, "person", v.Label)
+	v = NewVertex(46, "")
+	assert.Equal(t, "vertex", v.Label)
+}
+
+func TestVertexString(t *testing.T) {
+	v := NewVertex(42, "")
+	assert.Equal(t, "v[42]", fmt.Sprint(v))
+}
+
+func TestVertexEncoding(t *testing.T) {
+	t.Parallel()
+
+	v := NewVertex(1, "user")
+	got, err := graphson.MarshalToString(v)
+	require.NoError(t, err)
+
+	want := `{
+		"@type" : "g:Vertex",
+		"@value" : {
+			"id" : {
+				"@type" : "g:Int64",
+				"@value" : 1
+			},
+			"label" : "user"
+		}
+	}`
+	assert.JSONEq(t, want, got)
+
+	v = Vertex{}
+	err = graphson.UnmarshalFromString(got, &v)
+	require.NoError(t, err)
+
+	assert.Equal(t, int64(1), v.ID)
+	assert.Equal(t, "user", v.Label)
+}
+
+func TestVertexPropertyEncoding(t *testing.T) {
+	t.Parallel()
+
+	vp := NewVertexProperty("46ab60c2-918c-4cc4-a13b-350510e8908a", "name", "alex")
+	got, err := graphson.MarshalToString(vp)
+	require.NoError(t, err)
+
+	want := `{
+		"@type" : "g:VertexProperty",
+		"@value" : {
+			"id" : "46ab60c2-918c-4cc4-a13b-350510e8908a",
+			"label": "name",
+			"value": "alex"
+		}
+	}`
+	assert.JSONEq(t, want, got)
+
+	vp = VertexProperty{}
+	err = graphson.UnmarshalFromString(got, &vp)
+	require.NoError(t, err)
+
+	assert.Equal(t, "46ab60c2-918c-4cc4-a13b-350510e8908a", vp.ID)
+	assert.Equal(t, "name", vp.Key)
+	assert.Equal(t, "alex", vp.Value)
+}
+
+func TestVertexPropertyString(t *testing.T) {
+	vp := NewVertexProperty(55, "country", "israel")
+	assert.Equal(t, "vp[country->israel]", fmt.Sprint(vp))
+}

dialect/gremlin/http.go

@@ -0,0 +1,85 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	jsoniter "github.com/json-iterator/go"
+)
+
+type httpTransport struct {
+	client *http.Client
+	url    string
+}
+
+// NewHTTPTransport returns a new http transport.
+func NewHTTPTransport(urlStr string, client *http.Client) (RoundTripper, error) {
+	u, err := url.Parse(urlStr)
+	if err != nil {
+		return nil, fmt.Errorf("gremlin/http: parsing url: %w", err)
+	}
+	if client == nil {
+		client = http.DefaultClient
+	}
+	return &httpTransport{client, u.String()}, nil
+}
+
+// RoundTrip implements RouterTripper interface.
+func (t *httpTransport) RoundTrip(ctx context.Context, req *Request) (*Response, error) {
+	if req.Operation != OpsEval {
+		return nil, fmt.Errorf("gremlin/http: unsupported operation: %q", req.Operation)
+	}
+	if _, ok := req.Arguments[ArgsGremlin]; !ok {
+		return nil, errors.New("gremlin/http: missing query expression")
+	}
+
+	pr, pw := io.Pipe()
+	defer pr.Close()
+	go func() {
+		err := jsoniter.NewEncoder(pw).Encode(req.Arguments)
+		if err != nil {
+			err = fmt.Errorf("gremlin/http: encoding request: %w", err)
+		}
+		_ = pw.CloseWithError(err)
+	}()
+
+	var br io.Reader
+	{
+		req, err := http.NewRequest(http.MethodPost, t.url, pr)
+		if err != nil {
+			return nil, fmt.Errorf("gremlin/http: creating http request: %w", err)
+		}
+		req.Header.Set("Content-Type", "application/json")
+
+		rsp, err := t.client.Do(req.WithContext(ctx))
+		if err != nil {
+			return nil, fmt.Errorf("gremlin/http: posting http request: %w", err)
+		}
+		defer rsp.Body.Close()
+
+		if rsp.StatusCode < http.StatusOK || rsp.StatusCode > http.StatusPartialContent {
+			body, _ := io.ReadAll(rsp.Body)
+			return nil, fmt.Errorf("gremlin/http: status=%q, body=%q", rsp.Status, body)
+		}
+		if rsp.ContentLength > MaxResponseSize {
+			return nil, errors.New("gremlin/http: context length exceeds limit")
+		}
+		br = rsp.Body
+	}
+
+	var rsp Response
+	if err := graphson.NewDecoder(io.LimitReader(br, MaxResponseSize)).Decode(&rsp); err != nil {
+		return nil, fmt.Errorf("gremlin/http: decoding response: %w", err)
+	}
+	return &rsp, nil
+}

dialect/gremlin/http_test.go

@@ -0,0 +1,130 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHTTPTransportRoundTripper(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, r.Header.Get("Content-Type"), "application/json")
+		got, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		assert.JSONEq(t, `{"gremlin": "g.V(1)", "language": "gremlin-groovy"}`, string(got))
+
+		_, err = io.WriteString(w, `{
+			"requestId": "f679127f-8701-425c-af55-049a44720db6",
+			"result": {
+				"data": {
+					"@type": "g:List",
+					"@value": [
+						{
+							"@type": "g:Vertex",
+							"@value": {
+								"id": {
+									"@type": "g:Int64",
+									"@value": 1
+								},
+								"label": "person"
+							}
+						}
+					]
+				},
+				"meta": {
+					"@type": "g:Map",
+					"@value": []
+				}
+			},
+			"status": {
+				"attributes": {
+					"@type": "g:Map",
+					"@value": []
+				},
+				"code": 200,
+				"message": ""
+			}
+		}`)
+		require.NoError(t, err)
+	}))
+	defer srv.Close()
+
+	transport, err := NewHTTPTransport(srv.URL, nil)
+	require.NoError(t, err)
+
+	rsp, err := transport.RoundTrip(context.Background(), NewEvalRequest("g.V(1)"))
+	require.NoError(t, err)
+
+	assert.Equal(t, "f679127f-8701-425c-af55-049a44720db6", rsp.RequestID)
+	assert.Equal(t, 200, rsp.Status.Code)
+	assert.Empty(t, rsp.Status.Message)
+
+	v := jsoniter.Get(rsp.Result.Data, graphson.ValueKey, 0, graphson.ValueKey)
+	require.NoError(t, v.LastError())
+	assert.Equal(t, 1, v.Get("id", graphson.ValueKey).ToInt())
+	assert.Equal(t, "person", v.Get("label").ToString())
+}
+
+func TestNewHTTPTransportBadURL(t *testing.T) {
+	transport, err := NewHTTPTransport(":", nil)
+	assert.Nil(t, transport)
+	assert.Error(t, err)
+}
+
+func TestHTTPTransportBadRequest(t *testing.T) {
+	transport, err := NewHTTPTransport("example.com", nil)
+	require.NoError(t, err)
+
+	req := NewEvalRequest("g.V()")
+	req.Operation = ""
+	rsp, err := transport.RoundTrip(context.Background(), req)
+	assert.EqualError(t, err, `gremlin/http: unsupported operation: ""`)
+	assert.Nil(t, rsp)
+
+	req = NewEvalRequest("g.V()")
+	delete(req.Arguments, ArgsGremlin)
+	rsp, err = transport.RoundTrip(context.Background(), req)
+	assert.EqualError(t, err, "gremlin/http: missing query expression")
+	assert.Nil(t, rsp)
+}
+
+func TestHTTPTransportBadResponseStatus(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+	}))
+	defer srv.Close()
+
+	transport, err := NewHTTPTransport(srv.URL, nil)
+	require.NoError(t, err)
+
+	_, err = transport.RoundTrip(context.Background(), NewEvalRequest("g.E()."))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), http.StatusText(http.StatusInternalServerError))
+}
+
+func TestHTTPTransportBadResponseBody(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, err := io.WriteString(w, "{{{")
+		require.NoError(t, err)
+	}))
+	defer srv.Close()
+
+	transport, err := NewHTTPTransport(srv.URL, nil)
+	require.NoError(t, err)
+
+	_, err = transport.RoundTrip(context.Background(), NewEvalRequest("g.E()."))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "decoding response")
+}

dialect/gremlin/internal/ws/conn.go

@@ -0,0 +1,315 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ws
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"sync"
+	"time"
+
+	"entgo.io/ent/dialect/gremlin"
+	"entgo.io/ent/dialect/gremlin/encoding"
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/gorilla/websocket"
+	"golang.org/x/sync/errgroup"
+)
+
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 5 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 10 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+)
+
+type (
+	// A Dialer contains options for connecting to Gremlin server.
+	Dialer struct {
+		// Underlying websocket dialer.
+		websocket.Dialer
+
+		// Gremlin server basic auth credentials.
+		user, pass string
+	}
+
+	// Conn performs operations on a gremlin server.
+	Conn struct {
+		// Underlying websocket connection.
+		conn *websocket.Conn
+
+		// Credentials for basic authentication.
+		user, pass string
+
+		// Goroutine tracking.
+		ctx context.Context
+		grp *errgroup.Group
+
+		// Channel of outbound requests.
+		send chan io.Reader
+
+		// Map of in flight requests.
+		inflight sync.Map
+	}
+
+	// inflight tracks request state.
+	inflight struct {
+		// partially received data
+		frags []graphson.RawMessage
+
+		// response channel
+		result chan<- result
+	}
+
+	// represents an execution result.
+	result struct {
+		rsp *gremlin.Response
+		err error
+	}
+)
+
+var (
+	// DefaultDialer is a dialer with all fields set to the default values.
+	DefaultDialer = &Dialer{
+		Dialer: websocket.Dialer{
+			Proxy:            http.ProxyFromEnvironment,
+			HandshakeTimeout: 5 * time.Second,
+			WriteBufferSize:  8192,
+			ReadBufferSize:   8192,
+		},
+	}
+
+	// ErrConnClosed is returned by the Conn's Execute method when
+	// the underlying gremlin server connection is closed.
+	ErrConnClosed = errors.New("gremlin: server connection closed")
+
+	// ErrDuplicateRequest is returned by the Conns Execute method on
+	// request identifier key collision.
+	ErrDuplicateRequest = errors.New("gremlin: duplicate request")
+)
+
+// Dial creates a new connection by calling DialContext with a background context.
+func (d *Dialer) Dial(uri string) (*Conn, error) {
+	return d.DialContext(context.Background(), uri)
+}
+
+// DialContext creates a new Gremlin connection.
+func (d *Dialer) DialContext(ctx context.Context, uri string) (*Conn, error) {
+	c, rsp, err := d.Dialer.DialContext(ctx, uri, nil)
+	if err != nil {
+		return nil, fmt.Errorf("gremlin: dialing uri %s: %w", uri, err)
+	}
+	defer rsp.Body.Close()
+
+	conn := &Conn{
+		conn: c,
+		user: d.user,
+		pass: d.pass,
+		send: make(chan io.Reader),
+	}
+	conn.grp, conn.ctx = errgroup.WithContext(context.Background())
+
+	conn.grp.Go(conn.sender)
+	conn.grp.Go(conn.receiver)
+
+	return conn, nil
+}
+
+// Execute executes a request against a Gremlin server.
+func (c *Conn) Execute(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	// buffered result channel prevents receiver block on context cancellation
+	result := make(chan result, 1)
+
+	// request id must be unique across inflight request
+	if _, loaded := c.inflight.LoadOrStore(req.RequestID, &inflight{result: result}); loaded {
+		return nil, ErrDuplicateRequest
+	}
+
+	pr, pw := io.Pipe()
+	defer pr.Close()
+
+	// stream graphson encoding into request
+	c.grp.Go(func() error {
+		err := graphson.NewEncoder(pw).Encode(req)
+		if err != nil {
+			err = fmt.Errorf("encoding request: %w", err)
+		}
+		pw.CloseWithError(err)
+		return err
+	})
+
+	// local copy for single write
+	send := c.send
+
+	for {
+		select {
+		case <-c.ctx.Done():
+			c.inflight.Delete(req.RequestID)
+			return nil, ErrConnClosed
+		case <-ctx.Done():
+			c.inflight.Delete(req.RequestID)
+			return nil, ctx.Err()
+		case send <- pr:
+			send = nil
+		case result := <-result:
+			return result.rsp, result.err
+		}
+	}
+}
+
+// Close connection with a Gremlin server.
+func (c *Conn) Close() error {
+	c.grp.Go(func() error { return ErrConnClosed })
+	_ = c.grp.Wait()
+	return nil
+}
+
+func (c *Conn) sender() error {
+	pinger := time.NewTicker(pingPeriod)
+	defer pinger.Stop()
+
+	// closing connection terminates receiver
+	defer c.conn.Close()
+
+	for {
+		select {
+		case r := <-c.send:
+			// ensure write completes within a window
+			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
+
+			// fetch next message writer
+			w, err := c.conn.NextWriter(websocket.BinaryMessage)
+			if err != nil {
+				return fmt.Errorf("getting message writer: %w", err)
+			}
+
+			// write mime header
+			if _, err := w.Write(encoding.GraphSON3Mime); err != nil {
+				return fmt.Errorf("writing mime header: %w", err)
+			}
+
+			// write request body
+			if _, err := io.Copy(w, r); err != nil {
+				return fmt.Errorf("writing request: %w", err)
+			}
+
+			// finish message write
+			if err := w.Close(); err != nil {
+				return fmt.Errorf("closing message writer: %w", err)
+			}
+		case <-c.ctx.Done():
+			// connection closing
+			return c.conn.WriteControl(
+				websocket.CloseMessage,
+				websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
+				time.Time{},
+			)
+		case <-pinger.C:
+			// periodic connection keepalive
+			if err := c.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeWait)); err != nil {
+				return fmt.Errorf("writing ping message: %w", err)
+			}
+		}
+	}
+}
+
+func (c *Conn) receiver() error {
+	// handle keepalive responses
+	c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	c.conn.SetPongHandler(func(string) error {
+		return c.conn.SetReadDeadline(time.Now().Add(pongWait))
+	})
+
+	// complete all in flight requests on termination
+	defer c.inflight.Range(func(id, ifr any) bool {
+		ifr.(*inflight).result <- result{err: ErrConnClosed}
+		c.inflight.Delete(id)
+		return true
+	})
+
+	for {
+		// rely on sender connection close during termination
+		_, r, err := c.conn.NextReader()
+		if err != nil {
+			return fmt.Errorf("writing ping message: %w", err)
+		}
+
+		// decode received response
+		var rsp gremlin.Response
+		if err := graphson.NewDecoder(r).Decode(&rsp); err != nil {
+			return fmt.Errorf("reading response: %w", err)
+		}
+
+		ifr, ok := c.inflight.Load(rsp.RequestID)
+		if !ok {
+			// context cancellation aborts inflight requests
+			continue
+		}
+
+		// handle incoming response
+		if done := c.receive(ifr.(*inflight), &rsp); done {
+			// stop tracking finished requests
+			c.inflight.Delete(rsp.RequestID)
+		}
+	}
+}
+
+func (c *Conn) receive(ifr *inflight, rsp *gremlin.Response) bool {
+	result := result{rsp: rsp}
+	switch rsp.Status.Code {
+	case gremlin.StatusSuccess:
+		// quickly handle non fragmented responses
+		if ifr.frags == nil {
+			break
+		}
+		// handle fragment
+		fallthrough
+	case gremlin.StatusPartialContent:
+		// append received fragment
+		var frag []graphson.RawMessage
+		if err := graphson.Unmarshal(rsp.Result.Data, &frag); err != nil {
+			result.err = fmt.Errorf("decoding response fragment: %w", err)
+			break
+		}
+		ifr.frags = append(ifr.frags, frag...)
+
+		// partial response requires additional fragments
+		if rsp.Status.Code == gremlin.StatusPartialContent {
+			return false
+		}
+
+		// reassemble fragmented response
+		if rsp.Result.Data, result.err = graphson.Marshal(ifr.frags); result.err != nil {
+			result.err = fmt.Errorf("assembling fragmented response: %w", result.err)
+		}
+	case gremlin.StatusAuthenticate:
+		// receiver should never block
+		c.grp.Go(func() error {
+			var buf bytes.Buffer
+			if err := graphson.NewEncoder(&buf).Encode(
+				gremlin.NewAuthRequest(rsp.RequestID, c.user, c.pass),
+			); err != nil {
+				return fmt.Errorf("encoding auth request: %w", err)
+			}
+			select {
+			case c.send <- &buf:
+			case <-c.ctx.Done():
+			}
+			return c.ctx.Err()
+		})
+		return false
+	}
+
+	ifr.result <- result
+	return true
+}

dialect/gremlin/internal/ws/conn_test.go

@@ -0,0 +1,390 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ws
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"sync"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin"
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/gorilla/websocket"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type conn struct{ *websocket.Conn }
+
+func (c conn) ReadRequest() (*gremlin.Request, error) {
+	_, data, err := c.ReadMessage()
+	if err != nil {
+		return nil, err
+	}
+	var req gremlin.Request
+	if err := graphson.Unmarshal(data[data[0]+1:], &req); err != nil {
+		return nil, err
+	}
+	return &req, nil
+}
+
+func (c conn) WriteResponse(rsp *gremlin.Response) error {
+	data, err := graphson.Marshal(rsp)
+	if err != nil {
+		return err
+	}
+	return c.WriteMessage(websocket.BinaryMessage, data)
+}
+
+func serve(handler func(conn)) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		upgrader := websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}
+		c, _ := upgrader.Upgrade(w, r, nil)
+		defer c.Close()
+		handler(conn{c})
+		for {
+			_, _, err := c.ReadMessage()
+			if err != nil {
+				break
+			}
+		}
+	}))
+}
+
+func TestConnectClosure(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	defer wg.Wait()
+
+	srv := serve(func(conn conn) {
+		defer wg.Done()
+		_, _, err := conn.ReadMessage()
+		assert.True(t, websocket.IsCloseError(err, websocket.CloseNormalClosure))
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+
+	err = conn.Close()
+	assert.NoError(t, err)
+
+	_, err = conn.Execute(context.Background(), gremlin.NewEvalRequest("g.V()"))
+	assert.EqualError(t, err, ErrConnClosed.Error())
+}
+
+func TestSimpleQuery(t *testing.T) {
+	srv := serve(func(conn conn) {
+		typ, data, err := conn.ReadMessage()
+		require.NoError(t, err)
+		assert.Equal(t, websocket.BinaryMessage, typ)
+
+		var req gremlin.Request
+		err = graphson.Unmarshal(data[data[0]+1:], &req)
+		require.NoError(t, err)
+		assert.Equal(t, "g.V()", req.Arguments["gremlin"])
+
+		rsp := gremlin.Response{RequestID: req.RequestID}
+		rsp.Status.Code = gremlin.StatusNoContent
+		err = conn.WriteResponse(&rsp)
+		require.NoError(t, err)
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer assert.Condition(t, func() bool { return assert.NoError(t, conn.Close()) })
+
+	rsp, err := conn.Execute(context.Background(), gremlin.NewEvalRequest("g.V()"))
+	assert.NoError(t, err)
+	require.NotNil(t, rsp)
+	assert.Equal(t, gremlin.StatusNoContent, rsp.Status.Code)
+}
+
+func TestDuplicateRequest(t *testing.T) {
+	// skip until flakiness will be fixed.
+	t.SkipNow()
+	srv := serve(func(conn conn) {
+		req, err := conn.ReadRequest()
+		require.NoError(t, err)
+
+		rsp := gremlin.Response{RequestID: req.RequestID}
+		rsp.Status.Code = gremlin.StatusNoContent
+		err = conn.WriteResponse(&rsp)
+		require.NoError(t, err)
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	var errors [2]error
+	req := gremlin.NewEvalRequest("g.V()")
+
+	var wg sync.WaitGroup
+	wg.Add(len(errors))
+
+	for i := range errors {
+		go func(i int) {
+			_, errors[i] = conn.Execute(context.Background(), req)
+			wg.Done()
+		}(i)
+	}
+	wg.Wait()
+
+	err = errors[0]
+	if err == nil {
+		err = errors[1]
+	}
+	assert.EqualError(t, err, ErrDuplicateRequest.Error())
+}
+
+func TestConnectCancellation(t *testing.T) {
+	srv := serve(func(conn) {})
+	defer srv.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	conn, err := DefaultDialer.DialContext(ctx, "ws://"+srv.Listener.Addr().String())
+	assert.Error(t, err)
+	assert.Nil(t, conn)
+}
+
+func TestQueryCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	srv := serve(func(conn conn) {
+		if _, _, err := conn.ReadMessage(); err == nil {
+			cancel()
+		}
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	_, err = conn.Execute(ctx, gremlin.NewEvalRequest("g.E()"))
+	assert.EqualError(t, err, context.Canceled.Error())
+}
+
+func TestBadResponse(t *testing.T) {
+	tests := []struct {
+		name   string
+		mangle func(*gremlin.Response) *gremlin.Response
+	}{
+		{
+			name: "NoStatus",
+			mangle: func(rsp *gremlin.Response) *gremlin.Response {
+				return rsp
+			},
+		},
+		{
+			name: "Malformed",
+			mangle: func(rsp *gremlin.Response) *gremlin.Response {
+				rsp.Status.Code = gremlin.StatusMalformedRequest
+				rsp.Status.Message = "bad request"
+				return rsp
+			},
+		},
+		{
+			name: "Unknown",
+			mangle: func(rsp *gremlin.Response) *gremlin.Response {
+				rsp.Status.Code = 424242
+				return rsp
+			},
+		},
+	}
+
+	srv := serve(func(conn conn) {
+		for {
+			req, err := conn.ReadRequest()
+			if err != nil {
+				break
+			}
+
+			idx, err := strconv.ParseInt(req.Arguments["gremlin"].(string), 10, 0)
+			require.NoError(t, err)
+
+			err = conn.WriteResponse(tests[idx].mangle(&gremlin.Response{RequestID: req.RequestID}))
+			require.NoError(t, err)
+		}
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	var wg sync.WaitGroup
+	wg.Add(len(tests))
+
+	ctx := context.Background()
+	for i, tc := range tests {
+		i, tc := i, tc
+		t.Run(tc.name, func(t *testing.T) {
+			defer wg.Done()
+			rsp, err := conn.Execute(ctx, gremlin.NewEvalRequest(strconv.FormatInt(int64(i), 10)))
+			assert.NoError(t, err)
+			assert.True(t, rsp.IsErr())
+		})
+	}
+	wg.Wait()
+}
+
+func TestServerHangup(t *testing.T) {
+	// skip until flakiness will be fixed.
+	t.SkipNow()
+	srv := serve(func(conn conn) { _ = conn.Close() })
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	_, err = conn.Execute(context.Background(), gremlin.NewEvalRequest("g.V()"))
+	assert.EqualError(t, err, ErrConnClosed.Error())
+	assert.Error(t, conn.ctx.Err())
+}
+
+func TestCanceledLongRequest(t *testing.T) {
+	// skip until flakiness will be fixed.
+	t.SkipNow()
+	ctx, cancel := context.WithCancel(context.Background())
+	srv := serve(func(conn conn) {
+		var responses [3]*gremlin.Response
+		for i := 0; i < len(responses); i++ {
+			req, err := conn.ReadRequest()
+			require.NoError(t, err)
+
+			rsp := gremlin.Response{RequestID: req.RequestID}
+			rsp.Status.Code = gremlin.StatusSuccess
+			rsp.Result.Data = graphson.RawMessage(`"ok"`)
+			responses[i] = &rsp
+		}
+
+		cancel()
+
+		responses[0], responses[2] = responses[2], responses[0]
+		for i := 0; i < len(responses); i++ {
+			err := conn.WriteResponse(responses[i])
+			require.NoError(t, err)
+		}
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	var wg sync.WaitGroup
+	wg.Add(3)
+	defer wg.Wait()
+
+	for i := 0; i < 3; i++ {
+		go func(ctx context.Context, idx int) {
+			defer wg.Done()
+			rsp, err := conn.Execute(ctx, gremlin.NewEvalRequest("g.V()"))
+			if idx > 0 {
+				assert.NoError(t, err)
+				assert.EqualValues(t, []byte(`"ok"`), rsp.Result.Data)
+			} else {
+				assert.EqualError(t, err, context.Canceled.Error())
+			}
+		}(ctx, i)
+		ctx = context.Background()
+	}
+}
+
+func TestPartialResponse(t *testing.T) {
+	type kv struct {
+		Key   string
+		Value int
+	}
+	kvs := []kv{
+		{"one", 1},
+		{"two", 2},
+		{"three", 3},
+	}
+	srv := serve(func(conn conn) {
+		req, err := conn.ReadRequest()
+		require.NoError(t, err)
+
+		for i := range kvs {
+			data, err := graphson.Marshal([]kv{kvs[i]})
+			require.NoError(t, err)
+
+			rsp := gremlin.Response{RequestID: req.RequestID}
+			rsp.Result.Data = graphson.RawMessage(data)
+
+			if i != len(kvs)-1 {
+				rsp.Status.Code = gremlin.StatusPartialContent
+			} else {
+				rsp.Status.Code = gremlin.StatusSuccess
+			}
+
+			err = conn.WriteResponse(&rsp)
+			require.NoError(t, err)
+		}
+	})
+	defer srv.Close()
+
+	conn, err := DefaultDialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer conn.Close()
+
+	rsp, err := conn.Execute(context.Background(), gremlin.NewEvalRequest("g.E()"))
+	assert.NoError(t, err)
+
+	var result []kv
+	err = graphson.Unmarshal(rsp.Result.Data, &result)
+	require.NoError(t, err)
+	assert.Equal(t, kvs, result)
+}
+
+func TestAuthentication(t *testing.T) {
+	user, pass := "username", "password"
+	srv := serve(func(conn conn) {
+		req, err := conn.ReadRequest()
+		require.NoError(t, err)
+
+		rsp := gremlin.Response{RequestID: req.RequestID}
+		rsp.Status.Code = gremlin.StatusAuthenticate
+		err = conn.WriteResponse(&rsp)
+		require.NoError(t, err)
+
+		areq, err := conn.ReadRequest()
+		require.NoError(t, err)
+
+		var acreds gremlin.Credentials
+		err = acreds.UnmarshalText([]byte(areq.Arguments["sasl"].(string)))
+		assert.NoError(t, err)
+		areq.Arguments["sasl"] = acreds
+		assert.Equal(t, gremlin.NewAuthRequest(req.RequestID, user, pass), areq)
+
+		rsp = gremlin.Response{RequestID: req.RequestID}
+		rsp.Status.Code = gremlin.StatusNoContent
+		err = conn.WriteResponse(&rsp)
+		require.NoError(t, err)
+	})
+	defer srv.Close()
+
+	dialer := *DefaultDialer
+	dialer.user = user
+	dialer.pass = pass
+
+	client, err := dialer.Dial("ws://" + srv.Listener.Addr().String())
+	require.NoError(t, err)
+	defer client.Close()
+
+	_, err = client.Execute(context.Background(), gremlin.NewEvalRequest("g.E().drop()"))
+	assert.NoError(t, err)
+}

dialect/gremlin/ocgremlin/client.go

@@ -0,0 +1,64 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"go.opencensus.io/trace"
+)
+
+// Transport is an gremlin.RoundTripper that instruments all outgoing requests with
+// OpenCensus stats and tracing.
+type Transport struct {
+	// Base is a wrapped gremlin.RoundTripper that does the actual requests.
+	Base gremlin.RoundTripper
+
+	// StartOptions are applied to the span started by this Transport around each
+	// request.
+	//
+	// StartOptions.SpanKind will always be set to trace.SpanKindClient
+	// for spans started by this transport.
+	StartOptions trace.StartOptions
+
+	// GetStartOptions allows to set start options per request. If set,
+	// StartOptions is going to be ignored.
+	GetStartOptions func(context.Context, *gremlin.Request) trace.StartOptions
+
+	// NameFromRequest holds the function to use for generating the span name
+	// from the information found in the outgoing Gremlin Request. By default the
+	// name equals the URL Path.
+	FormatSpanName func(context.Context, *gremlin.Request) string
+
+	// WithQuery, if set to true, will enable recording of gremlin queries in spans.
+	// Only allow this if it is safe to have queries recorded with respect to
+	// security.
+	WithQuery bool
+}
+
+// RoundTrip implements gremlin.RoundTripper, delegating to Base and recording stats and traces for the request.
+func (t *Transport) RoundTrip(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	spanNameFormatter := t.FormatSpanName
+	if spanNameFormatter == nil {
+		spanNameFormatter = func(context.Context, *gremlin.Request) string {
+			return "gremlin:traversal"
+		}
+	}
+	startOpts := t.StartOptions
+	if t.GetStartOptions != nil {
+		startOpts = t.GetStartOptions(ctx, req)
+	}
+
+	var rt gremlin.RoundTripper = &traceTransport{
+		base:           t.Base,
+		formatSpanName: spanNameFormatter,
+		startOptions:   startOpts,
+		withQuery:      t.WithQuery,
+	}
+	rt = statsTransport{rt}
+	return rt.RoundTrip(ctx, req)
+}

dialect/gremlin/ocgremlin/client_test.go

@@ -0,0 +1,74 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"github.com/stretchr/testify/mock"
+	"go.opencensus.io/trace"
+)
+
+type mockExporter struct {
+	mock.Mock
+}
+
+func (e *mockExporter) ExportSpan(s *trace.SpanData) {
+	e.Called(s)
+}
+
+func TestTransportOptions(t *testing.T) {
+	tests := []struct {
+		name     string
+		spanName string
+		wantName string
+	}{
+		{
+			name:     "Default formatter",
+			wantName: "gremlin:traversal",
+		},
+		{
+			name:     "Custom formatter",
+			spanName: "tester",
+			wantName: "tester",
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			var exporter mockExporter
+			exporter.On(
+				"ExportSpan",
+				mock.MatchedBy(func(s *trace.SpanData) bool { return s.Name == tt.wantName })).
+				Once()
+			defer exporter.AssertExpectations(t)
+			trace.RegisterExporter(&exporter)
+			defer trace.UnregisterExporter(&exporter)
+
+			transport := &mockTransport{}
+			transport.On("RoundTrip", mock.Anything, mock.Anything).
+				Return(nil, errors.New("noop")).
+				Once()
+			defer transport.AssertExpectations(t)
+
+			rt := &Transport{
+				Base: transport,
+				GetStartOptions: func(context.Context, *gremlin.Request) trace.StartOptions {
+					return trace.StartOptions{Sampler: trace.AlwaysSample()}
+				},
+			}
+			if tt.spanName != "" {
+				rt.FormatSpanName = func(context.Context, *gremlin.Request) string {
+					return tt.spanName
+				}
+			}
+			_, _ = rt.RoundTrip(context.Background(), gremlin.NewEvalRequest("g.E()"))
+		})
+	}
+}

dialect/gremlin/ocgremlin/stats.go

@@ -0,0 +1,117 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"go.opencensus.io/stats"
+	"go.opencensus.io/stats/view"
+	"go.opencensus.io/tag"
+)
+
+// The following measures are supported for use in custom views.
+var (
+	RequestCount = stats.Int64(
+		"gremlin/request_count",
+		"Number of Gremlin requests started",
+		stats.UnitDimensionless,
+	)
+	ResponseBytes = stats.Int64(
+		"gremlin/response_bytes",
+		"Total number of bytes in response data",
+		stats.UnitBytes,
+	)
+	RoundTripLatency = stats.Float64(
+		"gremlin/roundtrip_latency",
+		"End-to-end latency",
+		stats.UnitMilliseconds,
+	)
+)
+
+// The following tags are applied to stats recorded by this package.
+var (
+	// StatusCode is the numeric Gremlin response status code,
+	// or "error" if a transport error occurred and no status code was read.
+	StatusCode, _ = tag.NewKey("gremlin_status_code")
+)
+
+// Default distributions used by views in this package.
+var (
+	DefaultSizeDistribution    = view.Distribution(32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144, 524288, 1048576)
+	DefaultLatencyDistribution = view.Distribution(1, 2, 3, 4, 5, 6, 8, 10, 13, 16, 20, 25, 30, 40, 50, 65, 80, 100, 130, 160, 200, 250, 300, 400, 500, 650, 800, 1000, 2000, 5000, 10000, 20000, 50000, 100000)
+)
+
+// Package ocgremlin provides some convenience views for measures.
+// You still need to register these views for data to actually be collected.
+var (
+	RequestCountView = &view.View{
+		Name:        "gremlin/request_count",
+		Measure:     RequestCount,
+		Aggregation: view.Count(),
+		Description: "Count of Gremlin requests started",
+	}
+
+	ResponseCountView = &view.View{
+		Name:        "gremlin/response_count",
+		Measure:     RoundTripLatency,
+		Aggregation: view.Count(),
+		Description: "Count of responses received, by response status",
+		TagKeys:     []tag.Key{StatusCode},
+	}
+
+	ResponseBytesView = &view.View{
+		Name:        "gremlin/response_bytes",
+		Measure:     ResponseBytes,
+		Aggregation: DefaultSizeDistribution,
+		Description: "Total number of bytes in response data",
+	}
+
+	RoundTripLatencyView = &view.View{
+		Name:        "gremlin/roundtrip_latency",
+		Measure:     RoundTripLatency,
+		Aggregation: DefaultLatencyDistribution,
+		Description: "End-to-end latency, by response code",
+		TagKeys:     []tag.Key{StatusCode},
+	}
+)
+
+// Views are the default views provided by this package.
+func Views() []*view.View {
+	return []*view.View{
+		RequestCountView,
+		ResponseCountView,
+		ResponseBytesView,
+		RoundTripLatencyView,
+	}
+}
+
+// statsTransport is an gremlin.RoundTripper that collects stats for the outgoing requests.
+type statsTransport struct {
+	base gremlin.RoundTripper
+}
+
+func (t statsTransport) RoundTrip(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	stats.Record(ctx, RequestCount.M(1))
+	start := time.Now()
+	rsp, err := t.base.RoundTrip(ctx, req)
+	latency := float64(time.Since(start)) / float64(time.Millisecond)
+	var (
+		tags = make([]tag.Mutator, 1)
+		ms   = []stats.Measurement{RoundTripLatency.M(latency)}
+	)
+	if err == nil {
+		tags[0] = tag.Upsert(StatusCode, strconv.Itoa(rsp.Status.Code))
+		ms = append(ms, ResponseBytes.M(int64(len(rsp.Result.Data))))
+	} else {
+		tags[0] = tag.Upsert(StatusCode, "error")
+	}
+	_ = stats.RecordWithTags(ctx, tags, ms...)
+	return rsp, err
+}

dialect/gremlin/ocgremlin/stats_test.go

@@ -0,0 +1,94 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin"
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"go.opencensus.io/stats/view"
+)
+
+func TestStatsCollection(t *testing.T) {
+	err := view.Register(
+		RequestCountView,
+		ResponseCountView,
+		ResponseBytesView,
+		RoundTripLatencyView,
+	)
+	require.NoError(t, err)
+
+	req := gremlin.NewEvalRequest("g.E()")
+	rsp := &gremlin.Response{RequestID: req.RequestID}
+	rsp.Status.Code = gremlin.StatusSuccess
+	rsp.Result.Data = graphson.RawMessage(
+		`{"@type": "g:List", "@value": [{"@type": "g:Int32", "@value": 42}]}`,
+	)
+
+	transport := &mockTransport{}
+	transport.On("RoundTrip", mock.Anything, mock.Anything).
+		Return(rsp, nil).
+		Once()
+	defer transport.AssertExpectations(t)
+
+	rt := &statsTransport{transport}
+	_, _ = rt.RoundTrip(context.Background(), req)
+
+	tests := []struct {
+		name   string
+		expect func(*testing.T, *view.Row)
+	}{
+		{
+			name: "gremlin/request_count",
+			expect: func(t *testing.T, row *view.Row) {
+				count, ok := row.Data.(*view.CountData)
+				require.True(t, ok)
+				assert.Equal(t, int64(1), count.Value)
+			},
+		},
+		{
+			name: "gremlin/response_count",
+			expect: func(t *testing.T, row *view.Row) {
+				count, ok := row.Data.(*view.CountData)
+				require.True(t, ok)
+				assert.Equal(t, int64(1), count.Value)
+			},
+		},
+		{
+			name: "gremlin/response_bytes",
+			expect: func(t *testing.T, row *view.Row) {
+				data, ok := row.Data.(*view.DistributionData)
+				require.True(t, ok)
+				assert.EqualValues(t, len(rsp.Result.Data), data.Sum())
+			},
+		},
+		{
+			name: "gremlin/roundtrip_latency",
+			expect: func(t *testing.T, row *view.Row) {
+				data, ok := row.Data.(*view.DistributionData)
+				require.True(t, ok)
+				assert.NotZero(t, data.Sum())
+			},
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name[strings.Index(tt.name, "/")+1:], func(t *testing.T) {
+			v := view.Find(tt.name)
+			assert.NotNil(t, v)
+			rows, err := view.RetrieveData(tt.name)
+			require.NoError(t, err)
+			require.Len(t, rows, 1)
+			tt.expect(t, rows[0])
+		})
+	}
+}

dialect/gremlin/ocgremlin/trace.go

@@ -0,0 +1,133 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+	"fmt"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"go.opencensus.io/trace"
+)
+
+// Attributes recorded on the span for the requests.
+const (
+	RequestIDAttribute = "gremlin.request_id"
+	OperationAttribute = "gremlin.operation"
+	QueryAttribute     = "gremlin.query"
+	BindingAttribute   = "gremlin.binding"
+	CodeAttribute      = "gremlin.code"
+	MessageAttribute   = "gremlin.message"
+)
+
+type traceTransport struct {
+	base           gremlin.RoundTripper
+	startOptions   trace.StartOptions
+	formatSpanName func(context.Context, *gremlin.Request) string
+	withQuery      bool
+}
+
+func (t *traceTransport) RoundTrip(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	ctx, span := trace.StartSpan(ctx,
+		t.formatSpanName(ctx, req),
+		trace.WithSampler(t.startOptions.Sampler),
+		trace.WithSpanKind(trace.SpanKindClient),
+	)
+	defer span.End()
+
+	span.AddAttributes(requestAttrs(req, t.withQuery)...)
+	rsp, err := t.base.RoundTrip(ctx, req)
+	if err != nil {
+		span.SetStatus(trace.Status{Code: trace.StatusCodeUnknown, Message: err.Error()})
+		return rsp, err
+	}
+
+	span.AddAttributes(responseAttrs(rsp)...)
+	span.SetStatus(TraceStatus(rsp.Status.Code))
+	return rsp, err
+}
+
+func requestAttrs(req *gremlin.Request, withQuery bool) []trace.Attribute {
+	attrs := []trace.Attribute{
+		trace.StringAttribute(RequestIDAttribute, req.RequestID),
+		trace.StringAttribute(OperationAttribute, req.Operation),
+	}
+	if withQuery {
+		query, _ := req.Arguments[gremlin.ArgsGremlin].(string)
+		attrs = append(attrs, trace.StringAttribute(QueryAttribute, query))
+		if bindings, ok := req.Arguments[gremlin.ArgsBindings].(map[string]any); ok {
+			attrs = append(attrs, bindingsAttrs(bindings)...)
+		}
+	}
+	return attrs
+}
+
+func bindingsAttrs(bindings map[string]any) []trace.Attribute {
+	attrs := make([]trace.Attribute, 0, len(bindings))
+	for key, val := range bindings {
+		key = BindingAttribute + "." + key
+		attrs = append(attrs, bindingToAttr(key, val))
+	}
+	return attrs
+}
+
+func bindingToAttr(key string, val any) trace.Attribute {
+	switch v := val.(type) {
+	case nil:
+		return trace.StringAttribute(key, "")
+	case int64:
+		return trace.Int64Attribute(key, v)
+	case float64:
+		return trace.Float64Attribute(key, v)
+	case string:
+		return trace.StringAttribute(key, v)
+	case bool:
+		return trace.BoolAttribute(key, v)
+	default:
+		s := fmt.Sprintf("%v", v)
+		if len(s) > 256 {
+			s = s[:256]
+		}
+		return trace.StringAttribute(key, s)
+	}
+}
+
+func responseAttrs(rsp *gremlin.Response) []trace.Attribute {
+	attrs := []trace.Attribute{
+		trace.Int64Attribute(CodeAttribute, int64(rsp.Status.Code)),
+	}
+	if rsp.Status.Message != "" {
+		attrs = append(attrs, trace.StringAttribute(MessageAttribute, rsp.Status.Message))
+	}
+	return attrs
+}
+
+// TraceStatus is a utility to convert the gremlin status code to a trace.Status.
+func TraceStatus(status int) trace.Status {
+	var code int32
+	switch status {
+	case gremlin.StatusSuccess,
+		gremlin.StatusNoContent,
+		gremlin.StatusPartialContent:
+		code = trace.StatusCodeOK
+	case gremlin.StatusUnauthorized:
+		code = trace.StatusCodePermissionDenied
+	case gremlin.StatusAuthenticate:
+		code = trace.StatusCodeUnauthenticated
+	case gremlin.StatusMalformedRequest,
+		gremlin.StatusInvalidRequestArguments,
+		gremlin.StatusScriptEvaluationError:
+		code = trace.StatusCodeInvalidArgument
+	case gremlin.StatusServerError,
+		gremlin.StatusServerSerializationError:
+		code = trace.StatusCodeInternal
+	case gremlin.StatusServerTimeout:
+		code = trace.StatusCodeDeadlineExceeded
+	default:
+		code = trace.StatusCodeUnknown
+	}
+	return trace.Status{Code: code, Message: gremlin.StatusText(status)}
+}

dialect/gremlin/ocgremlin/trace_test.go

@@ -0,0 +1,254 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"go.opencensus.io/trace"
+)
+
+type mockTransport struct {
+	mock.Mock
+}
+
+func (t *mockTransport) RoundTrip(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	args := t.Called(ctx, req)
+	rsp, _ := args.Get(0).(*gremlin.Response)
+	return rsp, args.Error(1)
+}
+
+func TestTraceTransportRoundTrip(t *testing.T) {
+	_, parent := trace.StartSpan(context.Background(), "parent")
+	tests := []struct {
+		name   string
+		parent *trace.Span
+	}{
+		{
+			name:   "no parent",
+			parent: nil,
+		},
+		{
+			name:   "parent",
+			parent: parent,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			transport := &mockTransport{}
+			transport.On("RoundTrip", mock.Anything, mock.Anything).
+				Run(func(args mock.Arguments) {
+					span := trace.FromContext(args.Get(0).(context.Context))
+					require.NotNil(t, span)
+					if tt.parent != nil {
+						assert.Equal(t, tt.parent.SpanContext().TraceID, span.SpanContext().TraceID)
+					}
+				}).
+				Return(nil, errors.New("noop")).
+				Once()
+			defer transport.AssertExpectations(t)
+
+			ctx, req := context.Background(), gremlin.NewEvalRequest("g.V()")
+			if tt.parent != nil {
+				ctx = trace.NewContext(ctx, tt.parent)
+			}
+			rt := &Transport{Base: transport}
+			_, _ = rt.RoundTrip(ctx, req)
+		})
+	}
+}
+
+type testExporter struct {
+	spans []*trace.SpanData
+}
+
+func (t *testExporter) ExportSpan(s *trace.SpanData) {
+	t.spans = append(t.spans, s)
+}
+
+func TestEndToEnd(t *testing.T) {
+	trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
+	var exporter testExporter
+	trace.RegisterExporter(&exporter)
+	defer trace.UnregisterExporter(&exporter)
+
+	req := gremlin.NewEvalRequest("g.V()")
+	rsp := &gremlin.Response{
+		RequestID: req.RequestID,
+	}
+	rsp.Status.Code = 200
+	rsp.Status.Message = "OK"
+
+	var transport mockTransport
+	transport.On("RoundTrip", mock.Anything, mock.Anything).
+		Return(rsp, nil).
+		Once()
+	defer transport.AssertExpectations(t)
+
+	rt := &Transport{Base: &transport, WithQuery: true}
+	_, err := rt.RoundTrip(context.Background(), req)
+	require.NoError(t, err)
+
+	require.Len(t, exporter.spans, 1)
+	attrs := exporter.spans[0].Attributes
+	assert.Len(t, attrs, 5)
+	assert.Equal(t, req.RequestID, attrs["gremlin.request_id"])
+	assert.Equal(t, req.Operation, attrs["gremlin.operation"])
+	assert.Equal(t, req.Arguments[gremlin.ArgsGremlin], attrs["gremlin.query"])
+	assert.Equal(t, int64(200), attrs["gremlin.code"])
+	assert.Equal(t, "OK", attrs["gremlin.message"])
+}
+
+func TestRequestAttributes(t *testing.T) {
+	tests := []struct {
+		name      string
+		makeReq   func() *gremlin.Request
+		wantAttrs []trace.Attribute
+	}{
+		{
+			name: "Query without bindings",
+			makeReq: func() *gremlin.Request {
+				req := gremlin.NewEvalRequest("g.E().count()")
+				req.RequestID = "a8b5c664-03ca-4175-a9e7-569b46f3551c"
+				return req
+			},
+			wantAttrs: []trace.Attribute{
+				trace.StringAttribute("gremlin.request_id", "a8b5c664-03ca-4175-a9e7-569b46f3551c"),
+				trace.StringAttribute("gremlin.operation", "eval"),
+				trace.StringAttribute("gremlin.query", "g.E().count()"),
+			},
+		},
+		{
+			name: "Query with bindings",
+			makeReq: func() *gremlin.Request {
+				bindings := map[string]any{
+					"$1": "user", "$2": int64(42),
+					"$3": 3.14, "$4": bytes.Repeat([]byte{0xff}, 257),
+					"$5": true, "$6": nil,
+				}
+				req := gremlin.NewEvalRequest(
+					`g.V().hasLabel($1).has("age",$2).has("v",$3).limit($4).valueMap($5)`,
+					gremlin.WithBindings(bindings),
+				)
+				req.RequestID = "d3d986fa-bd22-41bd-b2f7-ef2f1f639260"
+				return req
+			},
+			wantAttrs: []trace.Attribute{
+				trace.StringAttribute("gremlin.request_id", "d3d986fa-bd22-41bd-b2f7-ef2f1f639260"),
+				trace.StringAttribute("gremlin.operation", "eval"),
+				trace.StringAttribute("gremlin.query", `g.V().hasLabel($1).has("age",$2).has("v",$3).limit($4).valueMap($5)`),
+				trace.StringAttribute("gremlin.binding.$1", "user"),
+				trace.Int64Attribute("gremlin.binding.$2", 42),
+				trace.Float64Attribute("gremlin.binding.$3", 3.14),
+				trace.StringAttribute("gremlin.binding.$4", func() string {
+					str := fmt.Sprintf("%v", bytes.Repeat([]byte{0xff}, 256))
+					return str[:256]
+				}()),
+				trace.BoolAttribute("gremlin.binding.$5", true),
+				trace.StringAttribute("gremlin.binding.$6", ""),
+			},
+		},
+		{
+			name: "Authentication",
+			makeReq: func() *gremlin.Request {
+				return gremlin.NewAuthRequest(
+					"d239d950-59a1-41a7-a103-908f976ebd89",
+					"user", "pass",
+				)
+			},
+			wantAttrs: []trace.Attribute{
+				trace.StringAttribute("gremlin.request_id", "d239d950-59a1-41a7-a103-908f976ebd89"),
+				trace.StringAttribute("gremlin.operation", "authentication"),
+				trace.StringAttribute("gremlin.query", ""),
+			},
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			req := tt.makeReq()
+			attrs := requestAttrs(req, true)
+			for _, attr := range attrs {
+				assert.Contains(t, tt.wantAttrs, attr)
+			}
+			assert.Len(t, attrs, len(tt.wantAttrs))
+		})
+	}
+}
+
+func TestResponseAttributes(t *testing.T) {
+	tests := []struct {
+		name      string
+		makeRsp   func() *gremlin.Response
+		wantAttrs []trace.Attribute
+	}{
+		{
+			name: "Success no message",
+			makeRsp: func() *gremlin.Response {
+				var rsp gremlin.Response
+				rsp.Status.Code = 204
+				return &rsp
+			},
+			wantAttrs: []trace.Attribute{
+				trace.Int64Attribute("gremlin.code", 204),
+			},
+		},
+		{
+			name: "Authenticate with message",
+			makeRsp: func() *gremlin.Response {
+				var rsp gremlin.Response
+				rsp.Status.Code = 407
+				rsp.Status.Message = "login required"
+				return &rsp
+			},
+			wantAttrs: []trace.Attribute{
+				trace.Int64Attribute("gremlin.code", 407),
+				trace.StringAttribute("gremlin.message", "login required"),
+			},
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			rsp := tt.makeRsp()
+			attrs := responseAttrs(rsp)
+			assert.Equal(t, tt.wantAttrs, attrs)
+		})
+	}
+}
+
+func TestTraceStatus(t *testing.T) {
+	tests := []struct {
+		in   int
+		want trace.Status
+	}{
+		{200, trace.Status{Code: trace.StatusCodeOK, Message: "Success"}},
+		{204, trace.Status{Code: trace.StatusCodeOK, Message: "No Content"}},
+		{206, trace.Status{Code: trace.StatusCodeOK, Message: "Partial Content"}},
+		{401, trace.Status{Code: trace.StatusCodePermissionDenied, Message: "Unauthorized"}},
+		{407, trace.Status{Code: trace.StatusCodeUnauthenticated, Message: "Authenticate"}},
+		{498, trace.Status{Code: trace.StatusCodeInvalidArgument, Message: "Malformed Request"}},
+		{499, trace.Status{Code: trace.StatusCodeInvalidArgument, Message: "Invalid Request Arguments"}},
+		{500, trace.Status{Code: trace.StatusCodeInternal, Message: "Server Error"}},
+		{597, trace.Status{Code: trace.StatusCodeInvalidArgument, Message: "Script Evaluation Error"}},
+		{598, trace.Status{Code: trace.StatusCodeDeadlineExceeded, Message: "Server Timeout"}},
+		{599, trace.Status{Code: trace.StatusCodeInternal, Message: "Server Serialization Error"}},
+		{600, trace.Status{Code: trace.StatusCodeUnknown, Message: ""}},
+	}
+	for _, tt := range tests {
+		assert.Equal(t, tt.want, TraceStatus(tt.in))
+	}
+}

dialect/gremlin/request.go

@@ -0,0 +1,110 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"bytes"
+	"encoding/base64"
+	"errors"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+type (
+	// A Request models a request message sent to the server.
+	Request struct {
+		RequestID string         `json:"requestId" graphson:"g:UUID"`
+		Operation string         `json:"op"`
+		Processor string         `json:"processor"`
+		Arguments map[string]any `json:"args"`
+	}
+
+	// RequestOption enables request customization.
+	RequestOption func(*Request)
+
+	// Credentials holds request plain auth credentials.
+	Credentials struct{ Username, Password string }
+)
+
+// NewEvalRequest returns a new evaluation request request.
+func NewEvalRequest(query string, opts ...RequestOption) *Request {
+	r := &Request{
+		RequestID: uuid.New().String(),
+		Operation: OpsEval,
+		Arguments: map[string]any{
+			ArgsGremlin:  query,
+			ArgsLanguage: "gremlin-groovy",
+		},
+	}
+	for i := range opts {
+		opts[i](r)
+	}
+	return r
+}
+
+// NewAuthRequest returns a new auth request.
+func NewAuthRequest(requestID, username, password string) *Request {
+	return &Request{
+		RequestID: requestID,
+		Operation: OpsAuthentication,
+		Arguments: map[string]any{
+			ArgsSasl: Credentials{
+				Username: username,
+				Password: password,
+			},
+			ArgsSaslMechanism: "PLAIN",
+		},
+	}
+}
+
+// WithBindings sets request bindings.
+func WithBindings(bindings map[string]any) RequestOption {
+	return func(r *Request) {
+		r.Arguments[ArgsBindings] = bindings
+	}
+}
+
+// WithEvalTimeout sets script evaluation timeout.
+func WithEvalTimeout(timeout time.Duration) RequestOption {
+	return func(r *Request) {
+		r.Arguments[ArgsEvalTimeout] = int64(timeout / time.Millisecond)
+	}
+}
+
+// MarshalText implements encoding.TextMarshaler interface.
+func (c Credentials) MarshalText() ([]byte, error) {
+	var buf bytes.Buffer
+	buf.WriteByte(0)
+	buf.WriteString(c.Username)
+	buf.WriteByte(0)
+	buf.WriteString(c.Password)
+
+	enc := base64.StdEncoding
+	text := make([]byte, enc.EncodedLen(buf.Len()))
+	enc.Encode(text, buf.Bytes())
+	return text, nil
+}
+
+// UnmarshalText implements encoding.TextUnmarshaler interface.
+func (c *Credentials) UnmarshalText(text []byte) error {
+	enc := base64.StdEncoding
+	data := make([]byte, enc.DecodedLen(len(text)))
+
+	n, err := enc.Decode(data, text)
+	if err != nil {
+		return err
+	}
+	data = data[:n]
+
+	parts := bytes.SplitN(data, []byte{0}, 3)
+	if len(parts) != 3 {
+		return errors.New("bad credentials data")
+	}
+
+	c.Username = string(parts[1])
+	c.Password = string(parts[2])
+	return nil
+}

dialect/gremlin/request_test.go

@@ -0,0 +1,135 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEvaluateRequestEncode(t *testing.T) {
+	req := NewEvalRequest("g.V(x)",
+		WithBindings(map[string]any{"x": 1}),
+		WithEvalTimeout(time.Second),
+	)
+	data, err := graphson.Marshal(req)
+	require.NoError(t, err)
+
+	var got map[string]any
+	err = json.Unmarshal(data, &got)
+	require.NoError(t, err)
+
+	assert.Equal(t, map[string]any{
+		"@type":  "g:UUID",
+		"@value": req.RequestID,
+	}, got["requestId"])
+	assert.Equal(t, req.Operation, got["op"])
+	assert.Equal(t, req.Processor, got["processor"])
+
+	args := got["args"].(map[string]any)
+	assert.Equal(t, "g:Map", args["@type"])
+	assert.ElementsMatch(t, args["@value"], []any{
+		"gremlin", "g.V(x)", "language", "gremlin-groovy",
+		"scriptEvaluationTimeout", map[string]any{
+			"@type":  "g:Int64",
+			"@value": float64(1000),
+		},
+		"bindings", map[string]any{
+			"@type": "g:Map",
+			"@value": []any{
+				"x",
+				map[string]any{
+					"@type":  "g:Int64",
+					"@value": float64(1),
+				},
+			},
+		},
+	})
+}
+
+func TestEvaluateRequestWithoutBindingsEncode(t *testing.T) {
+	req := NewEvalRequest("g.E()")
+	got, err := graphson.MarshalToString(req)
+	require.NoError(t, err)
+	assert.NotContains(t, got, "bindings")
+}
+
+func TestAuthenticateRequestEncode(t *testing.T) {
+	req := NewAuthRequest("41d2e28a-20a4-4ab0-b379-d810dede3786", "user", "pass")
+	data, err := graphson.Marshal(req)
+	require.NoError(t, err)
+
+	var got map[string]any
+	err = json.Unmarshal(data, &got)
+	require.NoError(t, err)
+
+	assert.Equal(t, map[string]any{
+		"@type":  "g:UUID",
+		"@value": req.RequestID,
+	}, got["requestId"])
+	assert.Equal(t, req.Operation, got["op"])
+	assert.Equal(t, req.Processor, got["processor"])
+
+	args := got["args"].(map[string]any)
+	assert.Equal(t, "g:Map", args["@type"])
+	assert.ElementsMatch(t, args["@value"], []any{
+		"sasl", "AHVzZXIAcGFzcw==", "saslMechanism", "PLAIN",
+	})
+}
+
+func TestCredentialsMarshaling(t *testing.T) {
+	want := Credentials{
+		Username: "username",
+		Password: "password",
+	}
+
+	text, err := want.MarshalText()
+	assert.NoError(t, err)
+	assert.Equal(t, "AHVzZXJuYW1lAHBhc3N3b3Jk", string(text))
+
+	var got Credentials
+	err = got.UnmarshalText(text)
+	assert.NoError(t, err)
+	assert.Equal(t, want, got)
+}
+
+func TestCredentialsBadEncodingMarshaling(t *testing.T) {
+	tests := []struct {
+		name string
+		text []byte
+	}{
+		{
+			name: "BadBase64",
+			text: []byte{0x12},
+		},
+		{
+			name: "Empty",
+			text: []byte{},
+		},
+		{
+			name: "BadPrefix",
+			text: []byte("Kg=="),
+		},
+		{
+			name: "NoSeparator",
+			text: []byte("AHVzZXI="),
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			var creds Credentials
+			err := creds.UnmarshalText(tc.text)
+			assert.Error(t, err)
+		})
+	}
+}

dialect/gremlin/response.go

@@ -0,0 +1,127 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"errors"
+	"fmt"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+	"entgo.io/ent/dialect/gremlin/graph"
+)
+
+// A Response models a response message received from the server.
+type Response struct {
+	RequestID string `json:"requestId" graphson:"g:UUID"`
+	Status    struct {
+		Code       int            `json:"code"`
+		Attributes map[string]any `json:"attributes"`
+		Message    string         `json:"message"`
+	} `json:"status"`
+	Result struct {
+		Data graphson.RawMessage `json:"data"`
+		Meta map[string]any      `json:"meta"`
+	} `json:"result"`
+}
+
+// IsErr returns whether response indicates an error.
+func (rsp *Response) IsErr() bool {
+	switch rsp.Status.Code {
+	case StatusSuccess, StatusNoContent, StatusPartialContent:
+		return false
+	default:
+		return true
+	}
+}
+
+// Err returns an error representing response status.
+func (rsp *Response) Err() error {
+	if rsp.IsErr() {
+		return fmt.Errorf("gremlin: code=%d, message=%q", rsp.Status.Code, rsp.Status.Message)
+	}
+	return nil
+}
+
+// ReadVal reads gremlin response data into v.
+func (rsp *Response) ReadVal(v any) error {
+	if err := rsp.Err(); err != nil {
+		return err
+	}
+	if err := graphson.Unmarshal(rsp.Result.Data, v); err != nil {
+		return fmt.Errorf("gremlin: unmarshal response data: type=%T: %w", v, err)
+	}
+	return nil
+}
+
+// ReadVertices returns response data as slice of vertices.
+func (rsp *Response) ReadVertices() ([]graph.Vertex, error) {
+	var v []graph.Vertex
+	err := rsp.ReadVal(&v)
+	return v, err
+}
+
+// ReadVertexProperties returns response data as slice of vertex properties.
+func (rsp *Response) ReadVertexProperties() ([]graph.VertexProperty, error) {
+	var vp []graph.VertexProperty
+	err := rsp.ReadVal(&vp)
+	return vp, err
+}
+
+// ReadEdges returns response data as slice of edges.
+func (rsp *Response) ReadEdges() ([]graph.Edge, error) {
+	var e []graph.Edge
+	err := rsp.ReadVal(&e)
+	return e, err
+}
+
+// ReadProperties returns response data as slice of properties.
+func (rsp *Response) ReadProperties() ([]graph.Property, error) {
+	var p []graph.Property
+	err := rsp.ReadVal(&p)
+	return p, err
+}
+
+// ReadValueMap returns response data as a value map.
+func (rsp *Response) ReadValueMap() (graph.ValueMap, error) {
+	var m graph.ValueMap
+	err := rsp.ReadVal(&m)
+	return m, err
+}
+
+// ReadBool returns response data as a bool.
+func (rsp *Response) ReadBool() (bool, error) {
+	var b [1]*bool
+	if err := rsp.ReadVal(&b); err != nil {
+		return false, err
+	}
+	if b[0] == nil {
+		return false, errors.New("gremlin: no boolean value")
+	}
+	return *b[0], nil
+}
+
+// ReadInt returns response data as an int.
+func (rsp *Response) ReadInt() (int, error) {
+	var v [1]*int
+	if err := rsp.ReadVal(&v); err != nil {
+		return 0, err
+	}
+	if v[0] == nil {
+		return 0, errors.New("gremlin: no integer value")
+	}
+	return *v[0], nil
+}
+
+// ReadString returns response data as a string.
+func (rsp *Response) ReadString() (string, error) {
+	var v [1]*string
+	if err := rsp.ReadVal(&v); err != nil {
+		return "", err
+	}
+	if v[0] == nil {
+		return "", errors.New("gremlin: no string value")
+	}
+	return *v[0], nil
+}

dialect/gremlin/response_test.go

@@ -0,0 +1,564 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"reflect"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+	"entgo.io/ent/dialect/gremlin/graph"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDecodeResponse(t *testing.T) {
+	in := `{
+		"requestId": "a65f2d39-1efa-45d2-a06a-c736476500fc",
+		"result": {
+			"data": {
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Map",
+						"@value": [
+							{
+								"@type": "g:T",
+								"@value": "id"
+							},
+							{
+								"@type": "g:Int64",
+								"@value": 1
+							},
+							{
+								"@type": "g:T",
+								"@value": "label"
+							},
+							"person",
+							"name",
+							{
+								"@type": "g:List",
+								"@value": [
+									"marko"
+								]
+							},
+							"age",
+							{
+								"@type": "g:List",
+								"@value": [
+									{
+										"@type": "g:Int32",
+										"@value": 29
+									}
+								]
+							}
+						]
+					},
+					{
+						"@type": "g:Map",
+						"@value": [
+							{
+								"@type": "g:T",
+								"@value": "id"
+							},
+							{
+								"@type": "g:Int64",
+								"@value": 6
+							},
+							{
+								"@type": "g:T",
+								"@value": "label"
+							},
+							"person",
+							"name",
+							{
+								"@type": "g:List",
+								"@value": [
+									"peter"
+								]
+							},
+							"age",
+							{
+								"@type": "g:List",
+								"@value": [
+									{
+										"@type": "g:Int32",
+										"@value": 35
+									}
+								]
+							}
+						]
+					}
+				]
+			},
+			"meta": {
+				"@type": "g:Map",
+				"@value": []
+			}
+		},
+		"status": {
+			"attributes": {
+				"@type": "g:Map",
+				"@value": []
+			},
+			"code": 200,
+			"message": ""
+		}
+	}`
+
+	var rsp Response
+	err := graphson.UnmarshalFromString(in, &rsp)
+	require.NoError(t, err)
+
+	assert.Equal(t, "a65f2d39-1efa-45d2-a06a-c736476500fc", rsp.RequestID)
+	assert.Equal(t, 200, rsp.Status.Code)
+	assert.Empty(t, rsp.Status.Message)
+	assert.Empty(t, rsp.Status.Attributes)
+	assert.Empty(t, rsp.Result.Meta)
+
+	var vm graph.ValueMap
+	err = graphson.Unmarshal(rsp.Result.Data, &vm)
+	require.NoError(t, err)
+	require.Len(t, vm, 2)
+
+	type person struct {
+		ID   int64  `json:"id"`
+		Name string `json:"name"`
+		Age  int    `json:"age"`
+	}
+	var people []person
+	err = vm.Decode(&people)
+	require.NoError(t, err)
+	assert.Equal(t, []person{
+		{1, "marko", 29},
+		{6, "peter", 35},
+	}, people)
+}
+
+func TestDecodeResponseWithError(t *testing.T) {
+	in := `{
+		"requestId": "41d2e28a-20a4-4ab0-b379-d810dede3786",
+		"result": {
+			"data": null,
+			"meta": {
+				"@type": "g:Map",
+				"@value": []
+			}
+		},
+		"status": {
+			"attributes": {
+				"@type": "g:Map",
+				"@value": []
+			},
+			"code": 500,
+			"message": "Database Down"
+		}
+	}`
+	var rsp Response
+	err := graphson.UnmarshalFromString(in, &rsp)
+	require.NoError(t, err)
+
+	err = rsp.Err()
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "Database Down")
+
+	rsp = Response{}
+	err = graphson.UnmarshalFromString(`{"status": null}`, &rsp)
+	require.NoError(t, err)
+	assert.Error(t, rsp.Err())
+}
+
+func TestResponseReadVal(t *testing.T) {
+	var rsp Response
+	rsp.Status.Code = StatusSuccess
+	rsp.Result.Data = []byte(`{"@type": "g:Int32", "@value": 15}`)
+
+	var v int32
+	err := rsp.ReadVal(&v)
+	assert.NoError(t, err)
+	assert.Equal(t, int32(15), v)
+
+	var s string
+	err = rsp.ReadVal(&s)
+	assert.Error(t, err)
+
+	rsp.Status.Code = StatusServerError
+	err = rsp.ReadVal(&v)
+	assert.Error(t, err)
+}
+
+func TestResponseReadGraphElements(t *testing.T) {
+	tests := []struct {
+		method string
+		data   string
+		want   any
+	}{
+		{
+			method: "ReadVertices",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Vertex",
+						"@value": {
+							"id": {
+								"@type": "g:Int64",
+								"@value": 1
+							},
+							"label": "person"
+						}
+					},
+					{
+						"@type": "g:Vertex",
+						"@value": {
+							"id": {
+								"@type": "g:Int64",
+								"@value": 6
+							},
+							"label": "person"
+						}
+					}
+				]
+			}`,
+			want: []graph.Vertex{
+				graph.NewVertex(int64(1), "person"),
+				graph.NewVertex(int64(6), "person"),
+			},
+		},
+		{
+			method: "ReadVertexProperties",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:VertexProperty",
+						"@value": {
+							"id": {
+								"@type": "g:Int64",
+								"@value": 0
+							},
+							"label": "name",
+							"value": "marko"
+						}
+					},
+					{
+						"@type": "g:VertexProperty",
+						"@value": {
+							"id": {
+								"@type": "g:Int64",
+								"@value": 2
+							},
+							"label": "age",
+							"value": {
+								"@type": "g:Int32",
+								"@value": 29
+							}
+						}
+					}
+				]
+			}`,
+			want: []graph.VertexProperty{
+				graph.NewVertexProperty(int64(0), "name", "marko"),
+				graph.NewVertexProperty(int64(2), "age", int32(29)),
+			},
+		},
+		{
+			method: "ReadEdges",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Edge",
+						"@value": {
+							"id": {
+								"@type": "g:Int32",
+								"@value": 12
+							},
+							"inV": {
+								"@type": "g:Int64",
+								"@value": 3
+							},
+							"inVLabel": "software",
+							"label": "created",
+							"outV": {
+								"@type": "g:Int64",
+								"@value": 6
+							},
+							"outVLabel": "person"
+						}
+					}
+				]
+			}`,
+			want: []graph.Edge{
+				graph.NewEdge(int32(12), "created",
+					graph.NewVertex(int64(6), "person"),
+					graph.NewVertex(int64(3), "software"),
+				),
+			},
+		},
+		{
+			method: "ReadProperties",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Property",
+						"@value": {
+							"key": "weight",
+							"value": {
+								"@type": "g:Double",
+								"@value": 0.2
+							}
+						}
+					}
+				]
+			}`,
+			want: []graph.Property{
+				graph.NewProperty("weight", float64(0.2)),
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.method, func(t *testing.T) {
+			t.Parallel()
+			var rsp Response
+			rsp.Status.Code = StatusSuccess
+			rsp.Result.Data = []byte(tc.data)
+			vals := reflect.ValueOf(&rsp).MethodByName(tc.method).Call(nil)
+			require.Len(t, vals, 2)
+			require.True(t, vals[1].IsNil())
+			assert.Equal(t, tc.want, vals[0].Interface())
+		})
+	}
+}
+
+func TestResponseReadValueMap(t *testing.T) {
+	t.Parallel()
+	var rsp Response
+	rsp.Status.Code = StatusSuccess
+	rsp.Result.Data = []byte(`{
+		"@type": "g:List",
+		"@value": [
+			{
+				"@type": "g:Map",
+				"@value": [
+					"name",
+					{
+						"@type": "g:List",
+						"@value": [
+							"alex"
+						]
+					}
+				]
+			}
+		]
+	}`)
+	m, err := rsp.ReadValueMap()
+	require.NoError(t, err)
+
+	var name string
+	err = m.Decode(&struct {
+		Name *string `json:"name"`
+	}{&name})
+	require.NoError(t, err)
+	assert.Equal(t, "alex", name)
+}
+
+func TestResponseReadBool(t *testing.T) {
+	tests := []struct {
+		name    string
+		data    string
+		want    bool
+		wantErr bool
+	}{
+		{
+			name: "Simple",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					true
+				]
+			}`,
+			want: true,
+		},
+		{
+			name: "Multi",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					false,
+					true
+				]
+			}`,
+			want: false,
+		},
+		{
+			name: "Empty",
+			data: `{
+				"@type": "g:List",
+				"@value": []
+			}`,
+			wantErr: true,
+		},
+		{
+			name: "BadType",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					"user"
+				]
+			}`,
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			var rsp Response
+			rsp.Status.Code = StatusSuccess
+			rsp.Result.Data = []byte(tc.data)
+			got, err := rsp.ReadBool()
+			if tc.wantErr {
+				assert.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tc.want, got)
+			}
+		})
+	}
+}
+
+func TestResponseReadInt(t *testing.T) {
+	tests := []struct {
+		name    string
+		data    string
+		want    int
+		wantErr bool
+	}{
+		{
+			name: "Simple",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@value": 42
+					}
+				]
+			}`,
+			want: 42,
+		},
+		{
+			name: "Multi",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					{
+						"@type": "g:Int64",
+						"@value": 55
+					},
+					{
+						"@type": "g:Int64",
+						"@value": 13
+					}
+				]
+			}`,
+			want: 55,
+		},
+		{
+			name: "Empty",
+			data: `{
+				"@type": "g:List",
+				"@value": []
+			}`,
+			wantErr: true,
+		},
+		{
+			name: "BadType",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					true
+				]
+			}`,
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			var rsp Response
+			rsp.Status.Code = StatusSuccess
+			rsp.Result.Data = []byte(tc.data)
+			got, err := rsp.ReadInt()
+			if tc.wantErr {
+				assert.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tc.want, got)
+			}
+		})
+	}
+}
+
+func TestResponseReadString(t *testing.T) {
+	tests := []struct {
+		name    string
+		data    string
+		want    string
+		wantErr bool
+	}{
+		{
+			name: "Simple",
+			data: `{
+				"@type": "g:List",
+				"@value": ["foo"]
+			}`,
+			want: "foo",
+		},
+		{
+			name: "Empty",
+			data: `{
+				"@type": "g:List",
+				"@value": []
+			}`,
+			wantErr: true,
+		},
+		{
+			name: "BadType",
+			data: `{
+				"@type": "g:List",
+				"@value": [
+					true
+				]
+			}`,
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			var rsp Response
+			rsp.Status.Code = StatusSuccess
+			rsp.Result.Data = []byte(tc.data)
+			got, err := rsp.ReadString()
+			if tc.wantErr {
+				assert.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tc.want, got)
+			}
+		})
+	}
+}

dialect/gremlin/status.go

@@ -0,0 +1,69 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+const (
+	// StatusSuccess is returned on success.
+	StatusSuccess = 200
+
+	// StatusNoContent means the server processed the request but there is no result to return.
+	StatusNoContent = 204
+
+	// StatusPartialContent indicates the server successfully returned some content, but there
+	// is more in the stream to arrive wait for a success code to signify the end.
+	StatusPartialContent = 206
+
+	// StatusUnauthorized means the request attempted to access resources that
+	// the requesting user did not have access to.
+	StatusUnauthorized = 401
+
+	// StatusAuthenticate denotes a challenge from the server for the client to authenticate its request.
+	StatusAuthenticate = 407
+
+	// StatusMalformedRequest means the request message was not properly formatted which means it could not be parsed at
+	// all or the "op" code  was not recognized such that Gremlin Server could properly route it for processing.
+	// Check the message format and retry the request.
+	StatusMalformedRequest = 498
+
+	// StatusInvalidRequestArguments means the request message was parsable, but the arguments supplied in the message
+	// were in conflict or incomplete. Check the message format and retry the request.
+	StatusInvalidRequestArguments = 499
+
+	// StatusServerError indicates a general server error occurred that prevented the request from being processed.
+	StatusServerError = 500
+
+	// StatusScriptEvaluationError is returned when the script submitted for processing evaluated in the ScriptEngine
+	// with errors and could not be processed. Check the script submitted for syntax errors or other problems
+	// and then resubmit.
+	StatusScriptEvaluationError = 597
+
+	// StatusServerTimeout means the server exceeded one of the timeout settings for the request and could therefore
+	// only partially responded or did not respond at all.
+	StatusServerTimeout = 598
+
+	// StatusServerSerializationError means the server was not capable of serializing an object that was returned from the
+	// script supplied on the request. Either transform the object into something Gremlin Server can process within
+	// the script or install mapper serialization classes to Gremlin Server.
+	StatusServerSerializationError = 599
+)
+
+var statusText = map[int]string{
+	StatusSuccess:                  "Success",
+	StatusNoContent:                "No Content",
+	StatusPartialContent:           "Partial Content",
+	StatusUnauthorized:             "Unauthorized",
+	StatusAuthenticate:             "Authenticate",
+	StatusMalformedRequest:         "Malformed Request",
+	StatusInvalidRequestArguments:  "Invalid Request Arguments",
+	StatusServerError:              "Server Error",
+	StatusScriptEvaluationError:    "Script Evaluation Error",
+	StatusServerTimeout:            "Server Timeout",
+	StatusServerSerializationError: "Server Serialization Error",
+}
+
+// StatusText returns status text of code.
+func StatusText(code int) string {
+	return statusText[code]
+}

dialect/gremlin/status_test.go

@@ -0,0 +1,16 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestStatusText(t *testing.T) {
+	assert.NotEmpty(t, StatusText(StatusSuccess))
+	assert.Empty(t, StatusText(4242))
+}

dialect/gremlin/tokens.go

@@ -0,0 +1,69 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package gremlin
+
+// Gremlin server operations.
+const (
+	// OpsAuthentication used by the client to authenticate itself.
+	OpsAuthentication = "authentication"
+
+	// OpsBytecode used for a request that contains the Bytecode representation of a Traversal.
+	OpsBytecode = "bytecode"
+
+	// OpsEval used to evaluate a Gremlin script provided as a string.
+	OpsEval = "eval"
+
+	// OpsGather used to get a particular side-effect as produced by a previously executed Traversal.
+	OpsGather = "gather"
+
+	// OpsKeys used to get all the keys of all side-effects as produced by a previously executed Traversal.
+	OpsKeys = "keys"
+
+	// OpsClose used to get all the keys of all side-effects as produced by a previously executed Traversal.
+	OpsClose = "close"
+)
+
+// Gremlin server operation processors.
+const (
+	// ProcessorTraversal is the default operation processor.
+	ProcessorTraversal = "traversal"
+)
+
+const (
+	// ArgsBatchSize allows to defines the number of iterations each ResponseMessage should contain
+	ArgsBatchSize = "batchSize"
+
+	// ArgsBindings allows to provide a map of key/value pairs to apply
+	// as variables in the context of the Gremlin script.
+	ArgsBindings = "bindings"
+
+	// ArgsAliases allows to define aliases that represent globally bound Graph and TraversalSource objects.
+	ArgsAliases = "aliases"
+
+	// ArgsGremlin corresponds to the Traversal to evaluate.
+	ArgsGremlin = "gremlin"
+
+	// ArgsSideEffect allows to specify the unique identifier for the request.
+	ArgsSideEffect = "sideEffect"
+
+	// ArgsSideEffectKey allows to specify the key for a specific side-effect.
+	ArgsSideEffectKey = "sideEffectKey"
+
+	// ArgsAggregateTo describes how side-effect data should be treated.
+	ArgsAggregateTo = "aggregateTo"
+
+	// ArgsLanguage allows to change the flavor of Gremlin used (e.g. gremlin-groovy).
+	ArgsLanguage = "language"
+
+	// ArgsEvalTimeout allows to override the server setting that determines
+	// the maximum time to wait for a script to execute on the server.
+	ArgsEvalTimeout = "scriptEvaluationTimeout"
+
+	// ArgsSasl defines the response to the server authentication challenge.
+	ArgsSasl = "sasl"
+
+	// ArgsSaslMechanism defines the SASL mechanism (e.g. PLAIN).
+	ArgsSaslMechanism = "saslMechanism"
+)

dialect/sql/bench_test.go

@@ -0,0 +1,37 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"testing"
+
+	"entgo.io/ent/dialect"
+)
+
+func BenchmarkInsertBuilder_Default(b *testing.B) {
+	for _, d := range []string{dialect.SQLite, dialect.MySQL, dialect.Postgres} {
+		b.Run(d, func(b *testing.B) {
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				Dialect(d).Insert("users").Default().Returning("id").Query()
+			}
+		})
+	}
+}
+
+func BenchmarkInsertBuilder_Small(b *testing.B) {
+	for _, d := range []string{dialect.SQLite, dialect.MySQL, dialect.Postgres} {
+		b.Run(d, func(b *testing.B) {
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				Dialect(d).Insert("users").
+					Columns("id", "age", "first_name", "last_name", "nickname", "spouse_id", "created_at", "updated_at").
+					Values(1, 30, "Ariel", "Mashraki", "a8m", 2, "2009-11-10 23:00:00", "2009-11-10 23:00:00").
+					Returning("id").
+					Query()
+			}
+		})
+	}
+}

dialect/sql/driver.go

@@ -0,0 +1,305 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"entgo.io/ent/dialect"
+)
+
+// Driver is a dialect.Driver implementation for SQL based databases.
+type Driver struct {
+	Conn
+	dialect string
+}
+
+// NewDriver creates a new Driver with the given Conn and dialect.
+func NewDriver(dialect string, c Conn) *Driver {
+	return &Driver{dialect: dialect, Conn: c}
+}
+
+// Open wraps the database/sql.Open method and returns a dialect.Driver that implements the an ent/dialect.Driver interface.
+func Open(dialect, source string) (*Driver, error) {
+	db, err := sql.Open(dialect, source)
+	if err != nil {
+		return nil, err
+	}
+	return NewDriver(dialect, Conn{db, dialect}), nil
+}
+
+// OpenDB wraps the given database/sql.DB method with a Driver.
+func OpenDB(dialect string, db *sql.DB) *Driver {
+	return NewDriver(dialect, Conn{db, dialect})
+}
+
+// DB returns the underlying *sql.DB instance.
+func (d Driver) DB() *sql.DB {
+	return d.ExecQuerier.(*sql.DB)
+}
+
+// Dialect implements the dialect.Dialect method.
+func (d Driver) Dialect() string {
+	// If the underlying driver is wrapped with a telemetry driver.
+	for _, name := range []string{dialect.MySQL, dialect.SQLite, dialect.Postgres} {
+		if strings.HasPrefix(d.dialect, name) {
+			return name
+		}
+	}
+	return d.dialect
+}
+
+// Tx starts and returns a transaction.
+func (d *Driver) Tx(ctx context.Context) (dialect.Tx, error) {
+	return d.BeginTx(ctx, nil)
+}
+
+// BeginTx starts a transaction with options.
+func (d *Driver) BeginTx(ctx context.Context, opts *TxOptions) (dialect.Tx, error) {
+	tx, err := d.DB().BeginTx(ctx, opts)
+	if err != nil {
+		return nil, err
+	}
+	return &Tx{
+		Conn: Conn{tx, d.dialect},
+		Tx:   tx,
+	}, nil
+}
+
+// Close closes the underlying connection.
+func (d *Driver) Close() error { return d.DB().Close() }
+
+// Tx implements dialect.Tx interface.
+type Tx struct {
+	Conn
+	driver.Tx
+}
+
+// ctyVarsKey is the key used for attaching and reading the context variables.
+type ctxVarsKey struct{}
+
+// sessionVars holds sessions/transactions variables to set before every statement.
+type sessionVars struct {
+	vars []struct{ k, v string }
+}
+
+// WithVar returns a new context that holds the session variable to be executed before every query.
+func WithVar(ctx context.Context, name, value string) context.Context {
+	sv, _ := ctx.Value(ctxVarsKey{}).(sessionVars)
+	sv.vars = append(sv.vars, struct {
+		k, v string
+	}{
+		k: name,
+		v: value,
+	})
+	return context.WithValue(ctx, ctxVarsKey{}, sv)
+}
+
+// VarFromContext returns the session variable value from the context.
+func VarFromContext(ctx context.Context, name string) (string, bool) {
+	sv, _ := ctx.Value(ctxVarsKey{}).(sessionVars)
+	for _, s := range sv.vars {
+		if s.k == name {
+			return s.v, true
+		}
+	}
+	return "", false
+}
+
+// WithIntVar calls WithVar with the string representation of the value.
+func WithIntVar(ctx context.Context, name string, value int) context.Context {
+	return WithVar(ctx, name, strconv.Itoa(value))
+}
+
+// ExecQuerier wraps the standard Exec and Query methods.
+type ExecQuerier interface {
+	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
+	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
+}
+
+// Conn implements dialect.ExecQuerier given ExecQuerier.
+type Conn struct {
+	ExecQuerier
+	dialect string
+}
+
+// Exec implements the dialect.Exec method.
+func (c Conn) Exec(ctx context.Context, query string, args, v any) (rerr error) {
+	argv, ok := args.([]any)
+	if !ok {
+		return fmt.Errorf("dialect/sql: invalid type %T. expect []any for args", v)
+	}
+	ex, cf, err := c.maySetVars(ctx)
+	if err != nil {
+		return err
+	}
+	if cf != nil {
+		defer func() { rerr = errors.Join(rerr, cf()) }()
+	}
+	switch v := v.(type) {
+	case nil:
+		if _, err := ex.ExecContext(ctx, query, argv...); err != nil {
+			return err
+		}
+	case *sql.Result:
+		res, err := ex.ExecContext(ctx, query, argv...)
+		if err != nil {
+			return err
+		}
+		*v = res
+	default:
+		return fmt.Errorf("dialect/sql: invalid type %T. expect *sql.Result", v)
+	}
+	return nil
+}
+
+// Query implements the dialect.Query method.
+func (c Conn) Query(ctx context.Context, query string, args, v any) error {
+	vr, ok := v.(*Rows)
+	if !ok {
+		return fmt.Errorf("dialect/sql: invalid type %T. expect *sql.Rows", v)
+	}
+	argv, ok := args.([]any)
+	if !ok {
+		return fmt.Errorf("dialect/sql: invalid type %T. expect []any for args", args)
+	}
+	ex, cf, err := c.maySetVars(ctx)
+	if err != nil {
+		return err
+	}
+	rows, err := ex.QueryContext(ctx, query, argv...)
+	if err != nil {
+		if cf != nil {
+			err = errors.Join(err, cf())
+		}
+		return err
+	}
+	*vr = Rows{rows}
+	if cf != nil {
+		vr.ColumnScanner = rowsWithCloser{rows, cf}
+	}
+	return nil
+}
+
+// maySetVars sets the session variables before executing a query.
+func (c Conn) maySetVars(ctx context.Context) (ExecQuerier, func() error, error) {
+	sv, _ := ctx.Value(ctxVarsKey{}).(sessionVars)
+	if len(sv.vars) == 0 {
+		return c, nil, nil
+	}
+	var (
+		ex    ExecQuerier  // Underlying ExecQuerier.
+		cf    func() error // Close function.
+		reset []string     // Reset variables.
+		seen  = make(map[string]struct{}, len(sv.vars))
+	)
+	switch e := c.ExecQuerier.(type) {
+	case *sql.Tx:
+		ex = e
+	case *sql.DB:
+		conn, err := e.Conn(ctx)
+		if err != nil {
+			return nil, nil, err
+		}
+		ex, cf = conn, conn.Close
+	}
+	for _, s := range sv.vars {
+		if _, ok := seen[s.k]; !ok {
+			switch c.dialect {
+			case dialect.Postgres:
+				reset = append(reset, fmt.Sprintf("RESET %s", s.k))
+			case dialect.MySQL:
+				reset = append(reset, fmt.Sprintf("SET %s = NULL", s.k))
+			}
+			seen[s.k] = struct{}{}
+		}
+		if _, err := ex.ExecContext(ctx, fmt.Sprintf("SET %s = '%s'", s.k, s.v)); err != nil {
+			if cf != nil {
+				err = errors.Join(err, cf())
+			}
+			return nil, nil, err
+		}
+	}
+	// If there are variables to reset, and we need to return the
+	// connection to the pool, we need to clean up the variables.
+	if cls := cf; cf != nil && len(reset) > 0 {
+		cf = func() error {
+			for _, q := range reset {
+				if _, err := ex.ExecContext(ctx, q); err != nil {
+					return errors.Join(err, cls())
+				}
+			}
+			return cls()
+		}
+	}
+	return ex, cf, nil
+}
+
+var _ dialect.Driver = (*Driver)(nil)
+
+type (
+	// Rows wraps the sql.Rows to avoid locks copy.
+	Rows struct{ ColumnScanner }
+	// Result is an alias to sql.Result.
+	Result = sql.Result
+	// NullBool is an alias to sql.NullBool.
+	NullBool = sql.NullBool
+	// NullInt64 is an alias to sql.NullInt64.
+	NullInt64 = sql.NullInt64
+	// NullString is an alias to sql.NullString.
+	NullString = sql.NullString
+	// NullFloat64 is an alias to sql.NullFloat64.
+	NullFloat64 = sql.NullFloat64
+	// NullTime represents a time.Time that may be null.
+	NullTime = sql.NullTime
+	// TxOptions holds the transaction options to be used in DB.BeginTx.
+	TxOptions = sql.TxOptions
+)
+
+// NullScanner implements the sql.Scanner interface such that it
+// can be used as a scan destination, similar to the types above.
+type NullScanner struct {
+	S     sql.Scanner
+	Valid bool // Valid is true if the Scan value is not NULL.
+}
+
+// Scan implements the Scanner interface.
+func (n *NullScanner) Scan(value any) error {
+	n.Valid = value != nil
+	if n.Valid {
+		return n.S.Scan(value)
+	}
+	return nil
+}
+
+// ColumnScanner is the interface that wraps the standard
+// sql.Rows methods used for scanning database rows.
+type ColumnScanner interface {
+	Close() error
+	ColumnTypes() ([]*sql.ColumnType, error)
+	Columns() ([]string, error)
+	Err() error
+	Next() bool
+	NextResultSet() bool
+	Scan(dest ...any) error
+}
+
+// rowsWithCloser wraps the ColumnScanner interface with a custom Close hook.
+type rowsWithCloser struct {
+	ColumnScanner
+	closer func() error
+}
+
+// Close closes the underlying ColumnScanner and calls the custom closer.
+func (r rowsWithCloser) Close() error {
+	err := r.ColumnScanner.Close()
+	return errors.Join(err, r.closer())
+}

dialect/sql/driver_test.go

@@ -0,0 +1,93 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"context"
+	"testing"
+
+	"entgo.io/ent/dialect"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithVars(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	db.SetMaxOpenConns(1)
+	drv := OpenDB(dialect.Postgres, db)
+	mock.ExpectExec("SET foo = 'bar'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectQuery("SELECT 1").WillReturnRows(sqlmock.NewRows([]string{"1"}).AddRow(1))
+	mock.ExpectExec("RESET foo").WillReturnResult(sqlmock.NewResult(0, 0))
+	rows := &Rows{}
+	err = drv.Query(
+		WithVar(context.Background(), "foo", "bar"),
+		"SELECT 1",
+		[]any{},
+		rows,
+	)
+	require.NoError(t, err)
+	require.NoError(t, rows.Close(), "rows should be closed to release the connection")
+	require.NoError(t, mock.ExpectationsWereMet())
+
+	mock.ExpectExec("SET foo = 'bar'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("SET foo = 'baz'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectQuery("SELECT 1").WillReturnRows(sqlmock.NewRows([]string{"1"}).AddRow(1))
+	mock.ExpectExec("RESET foo").WillReturnResult(sqlmock.NewResult(0, 0))
+	err = drv.Query(
+		WithVar(WithVar(context.Background(), "foo", "bar"), "foo", "baz"),
+		"SELECT 1",
+		[]any{},
+		rows,
+	)
+	require.NoError(t, err)
+	require.NoError(t, rows.Close(), "rows should be closed to release the connection")
+	require.NoError(t, mock.ExpectationsWereMet())
+
+	mock.ExpectBegin()
+	mock.ExpectExec("SET foo = 'bar'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectQuery("SELECT 1").WillReturnRows(sqlmock.NewRows([]string{"1"}).AddRow(1))
+	mock.ExpectCommit()
+	tx, err := drv.Tx(context.Background())
+	require.NoError(t, err)
+	err = tx.Query(
+		WithVar(context.Background(), "foo", "bar"),
+		"SELECT 1",
+		[]any{},
+		rows,
+	)
+	require.NoError(t, err)
+	require.NoError(t, tx.Commit())
+	require.NoError(t, mock.ExpectationsWereMet())
+	// Rows should not be closed to release the session,
+	// as a transaction is always scoped to a single connection.
+
+	mock.ExpectExec("SET foo = 'qux'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("INSERT INTO users DEFAULT VALUES").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("RESET foo").WillReturnResult(sqlmock.NewResult(0, 0))
+	err = drv.Exec(
+		WithVar(context.Background(), "foo", "qux"),
+		"INSERT INTO users DEFAULT VALUES",
+		[]any{},
+		nil,
+	)
+	require.NoError(t, err)
+	require.NoError(t, mock.ExpectationsWereMet())
+	// No rows are returned, so no need to close them.
+
+	mock.ExpectExec("SET foo = 'foo'").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("INSERT INTO users DEFAULT VALUES").WillReturnResult(sqlmock.NewResult(0, 0))
+	mock.ExpectExec("RESET foo").WillReturnResult(sqlmock.NewResult(0, 0))
+	err = drv.Exec(
+		WithVar(context.Background(), "foo", "foo"),
+		"INSERT INTO users DEFAULT VALUES",
+		[]any{},
+		nil,
+	)
+	require.NoError(t, err)
+	require.NoError(t, mock.ExpectationsWereMet())
+	// No rows are returned, so no need to close them.
+}

dialect/sql/scan.go

@@ -0,0 +1,425 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"database/sql"
+	"database/sql/driver"
+	"encoding/json"
+	"fmt"
+	"reflect"
+	"strings"
+	"time"
+)
+
+// ScanOne scans one row to the given value. It fails if the rows holds more than 1 row.
+func ScanOne(rows ColumnScanner, v any) error {
+	columns, err := rows.Columns()
+	if err != nil {
+		return fmt.Errorf("sql/scan: failed getting column names: %w", err)
+	}
+	if n := len(columns); n != 1 {
+		return fmt.Errorf("sql/scan: unexpected number of columns: %d", n)
+	}
+	if !rows.Next() {
+		if err := rows.Err(); err != nil {
+			return err
+		}
+		return sql.ErrNoRows
+	}
+	if err := rows.Scan(v); err != nil {
+		return err
+	}
+	if rows.Next() {
+		return fmt.Errorf("sql/scan: expect exactly one row in result set")
+	}
+	return rows.Err()
+}
+
+// ScanInt64 scans and returns an int64 from the rows.
+func ScanInt64(rows ColumnScanner) (int64, error) {
+	var n int64
+	if err := ScanOne(rows, &n); err != nil {
+		return 0, err
+	}
+	return n, nil
+}
+
+// ScanInt scans and returns an int from the rows.
+func ScanInt(rows ColumnScanner) (int, error) {
+	n, err := ScanInt64(rows)
+	if err != nil {
+		return 0, err
+	}
+	return int(n), nil
+}
+
+// ScanBool scans and returns a boolean from the rows.
+func ScanBool(rows ColumnScanner) (bool, error) {
+	var b bool
+	if err := ScanOne(rows, &b); err != nil {
+		return false, err
+	}
+	return b, nil
+}
+
+// ScanString scans and returns a string from the rows.
+func ScanString(rows ColumnScanner) (string, error) {
+	var s string
+	if err := ScanOne(rows, &s); err != nil {
+		return "", err
+	}
+	return s, nil
+}
+
+// ScanValue scans and returns a driver.Value from the rows.
+func ScanValue(rows ColumnScanner) (driver.Value, error) {
+	var v driver.Value
+	if err := ScanOne(rows, &v); err != nil {
+		return "", err
+	}
+	return v, nil
+}
+
+// ScanSlice scans the given ColumnScanner (basically, sql.Row or sql.Rows) into the given slice.
+func ScanSlice(rows ColumnScanner, v any) error {
+	columns, err := rows.Columns()
+	if err != nil {
+		return fmt.Errorf("sql/scan: failed getting column names: %w", err)
+	}
+	rv := reflect.ValueOf(v)
+	switch {
+	case rv.Kind() != reflect.Ptr:
+		if t := reflect.TypeOf(v); t != nil {
+			return fmt.Errorf("sql/scan: ScanSlice(non-pointer %s)", t)
+		}
+		fallthrough
+	case rv.IsNil():
+		return fmt.Errorf("sql/scan: ScanSlice(nil)")
+	}
+	rv = reflect.Indirect(rv)
+	if k := rv.Kind(); k != reflect.Slice {
+		return fmt.Errorf("sql/scan: invalid type %s. expected slice as an argument", k)
+	}
+	scan, err := scanType(rv.Type().Elem(), columns)
+	if err != nil {
+		return err
+	}
+	if n, m := len(columns), len(scan.columns); n > m {
+		return fmt.Errorf("sql/scan: columns do not match (%d > %d)", n, m)
+	}
+	for rows.Next() {
+		values := scan.values()
+		if err := rows.Scan(values...); err != nil {
+			return fmt.Errorf("sql/scan: failed scanning rows: %w", err)
+		}
+		vv, err := scan.value(values...)
+		if err != nil {
+			return err
+		}
+		rv.Set(reflect.Append(rv, vv))
+	}
+	return rows.Err()
+}
+
+// rowScan is the configuration for scanning one sql.Row.
+type rowScan struct {
+	// column types of a row.
+	columns []reflect.Type
+	// value functions that converts the row columns (result) to a reflect.Value.
+	value func(v ...any) (reflect.Value, error)
+}
+
+// values returns a []any from the configured column types.
+func (r *rowScan) values() []any {
+	values := make([]any, len(r.columns))
+	for i := range r.columns {
+		values[i] = reflect.New(r.columns[i]).Interface()
+	}
+	return values
+}
+
+// scanType returns rowScan for the given reflect.Type.
+func scanType(typ reflect.Type, columns []string) (*rowScan, error) {
+	switch k := typ.Kind(); {
+	case assignable(typ):
+		return &rowScan{
+			columns: []reflect.Type{typ},
+			value: func(v ...any) (reflect.Value, error) {
+				return reflect.Indirect(reflect.ValueOf(v[0])), nil
+			},
+		}, nil
+	case k == reflect.Ptr:
+		return scanPtr(typ, columns)
+	case k == reflect.Struct:
+		return scanStruct(typ, columns)
+	default:
+		return nil, fmt.Errorf("sql/scan: unsupported type ([]%s)", k)
+	}
+}
+
+var (
+	timeType     = reflect.TypeOf(time.Time{})
+	scannerType  = reflect.TypeOf((*sql.Scanner)(nil)).Elem()
+	nullJSONType = reflect.TypeOf((*nullJSON)(nil)).Elem()
+)
+
+// nullJSON represents a json.RawMessage that may be NULL.
+type nullJSON json.RawMessage
+
+// Scan implements the sql.Scanner interface.
+func (j *nullJSON) Scan(v interface{}) error {
+	if v == nil {
+		return nil
+	}
+	*j = v.([]byte)
+	return nil
+}
+
+// assignable reports if the given type can be assigned directly by `Rows.Scan`.
+func assignable(typ reflect.Type) bool {
+	switch k := typ.Kind(); {
+	case typ.Implements(scannerType):
+	case k == reflect.Interface && typ.NumMethod() == 0:
+	case k == reflect.String || k >= reflect.Bool && k <= reflect.Float64:
+	case (k == reflect.Slice || k == reflect.Array) && typ.Elem().Kind() == reflect.Uint8:
+	default:
+		return false
+	}
+	return true
+}
+
+// scanStruct returns the configuration for scanning a sql.Row into a struct.
+func scanStruct(typ reflect.Type, columns []string) (*rowScan, error) {
+	var (
+		scan  = &rowScan{}
+		idxs  = make([][]int, 0, typ.NumField())
+		names = make(map[string][]int, typ.NumField())
+	)
+	for i := 0; i < typ.NumField(); i++ {
+		f := typ.Field(i)
+		// Skip unexported fields.
+		if f.PkgPath != "" {
+			continue
+		}
+		// Support 1-level embedding to accept types as `type T struct {ent.T; V int}`.
+		if typ := f.Type; f.Anonymous && typ.Kind() == reflect.Struct {
+			for j := 0; j < typ.NumField(); j++ {
+				names[columnName(typ.Field(j))] = []int{i, j}
+			}
+			continue
+		}
+		names[columnName(f)] = []int{i}
+	}
+	for _, c := range columns {
+		var idx []int
+		// Normalize columns if necessary,
+		// for example: COUNT(*) => count.
+		switch name := strings.Split(c, "(")[0]; {
+		case names[name] != nil:
+			idx = names[name]
+		case names[strings.ToLower(name)] != nil:
+			idx = names[strings.ToLower(name)]
+		default:
+			return nil, fmt.Errorf("sql/scan: missing struct field for column: %s (%s)", c, name)
+		}
+		idxs = append(idxs, idx)
+		rtype := typ.Field(idx[0]).Type
+		if len(idx) > 1 {
+			rtype = rtype.Field(idx[1]).Type
+		}
+		switch {
+		// If the field is not support by the standard
+		// convertAssign, assume it is a JSON field.
+		case !supportsScan(rtype):
+			rtype = nullJSONType
+		// Create a pointer to the actual reflect
+		// types to accept optional struct fields.
+		case !nillable(rtype):
+			rtype = reflect.PtrTo(rtype)
+		}
+		scan.columns = append(scan.columns, rtype)
+	}
+	scan.value = func(vs ...any) (reflect.Value, error) {
+		st := reflect.New(typ).Elem()
+		for i, v := range vs {
+			rv := reflect.Indirect(reflect.ValueOf(v))
+			if rv.IsNil() {
+				continue
+			}
+			idx := idxs[i]
+			rvalue, ft := st.Field(idx[0]), st.Type().Field(idx[0])
+			if len(idx) > 1 {
+				// Embedded field.
+				rvalue, ft = rvalue.Field(idx[1]), ft.Type.Field(idx[1])
+			}
+			switch {
+			case rv.Type() == nullJSONType:
+				if rv = reflect.Indirect(rv); rv.IsNil() {
+					continue
+				}
+				if err := json.Unmarshal(rv.Bytes(), rvalue.Addr().Interface()); err != nil {
+					return reflect.Value{}, fmt.Errorf("unmarshal field %q: %w", ft.Name, err)
+				}
+			case !nillable(rvalue.Type()):
+				rv = reflect.Indirect(rv)
+				fallthrough
+			default:
+				rvalue.Set(rv)
+			}
+		}
+		return st, nil
+	}
+	return scan, nil
+}
+
+// columnName returns the column name of a struct-field.
+func columnName(f reflect.StructField) string {
+	name := strings.ToLower(f.Name)
+	if tag, ok := f.Tag.Lookup("sql"); ok {
+		name = tag
+	} else if tag, ok := f.Tag.Lookup("json"); ok {
+		name = strings.Split(tag, ",")[0]
+	}
+	return name
+}
+
+// nillable reports if the reflect-type can have nil value.
+func nillable(t reflect.Type) bool {
+	switch t.Kind() {
+	case reflect.Interface, reflect.Slice, reflect.Map, reflect.Ptr, reflect.UnsafePointer:
+		return true
+	}
+	return false
+}
+
+// scanPtr wraps the underlying type with rowScan.
+func scanPtr(typ reflect.Type, columns []string) (*rowScan, error) {
+	typ = typ.Elem()
+	scan, err := scanType(typ, columns)
+	if err != nil {
+		return nil, err
+	}
+	wrap := scan.value
+	scan.value = func(vs ...any) (reflect.Value, error) {
+		v, err := wrap(vs...)
+		if err != nil {
+			return reflect.Value{}, err
+		}
+		pt := reflect.PtrTo(v.Type())
+		pv := reflect.New(pt.Elem())
+		pv.Elem().Set(v)
+		return pv, nil
+	}
+	return scan, nil
+}
+
+func supportsScan(t reflect.Type) bool {
+	if t.Implements(scannerType) || reflect.PtrTo(t).Implements(scannerType) {
+		return true
+	}
+	if t.Kind() == reflect.Ptr {
+		t = t.Elem()
+	}
+	switch t.Kind() {
+	case reflect.Bool, reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
+		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64,
+		reflect.Float32, reflect.Float64, reflect.Pointer, reflect.String:
+		return true
+	case reflect.Slice:
+		return t == reflect.TypeOf(sql.RawBytes(nil)) || t == reflect.TypeOf([]byte(nil))
+	case reflect.Interface:
+		return t == reflect.TypeOf((*any)(nil)).Elem()
+	default:
+		return t == reflect.TypeOf(time.Time{}) || t.Implements(scannerType)
+	}
+}
+
+// UnknownType is a named type to any indicates the info
+// needs to be extracted from the underlying rows.
+type UnknownType any
+
+// ScanTypeOf returns the type used for scanning column i from the database.
+func ScanTypeOf(rows *Rows, i int) any {
+	unknown := new(any)
+	ct, err := rows.ColumnTypes()
+	if err != nil || len(ct) <= i {
+		return unknown
+	}
+	rt := ct[i].ScanType()
+	if rt.Kind() == reflect.Pointer {
+		rt = rt.Elem()
+	}
+	// Handle NULL values.
+	switch k := rt.Kind(); k {
+	case reflect.Bool:
+		rt = reflect.TypeOf(sql.NullBool{})
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
+		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+		rt = reflect.TypeOf(sql.NullInt64{})
+	case reflect.Float32, reflect.Float64:
+		rt = reflect.TypeOf(sql.NullFloat64{})
+	case reflect.String:
+		rt = reflect.TypeOf(sql.NullString{})
+	default:
+		if k == reflect.Struct && rt == timeType {
+			rt = reflect.TypeOf(sql.NullTime{})
+		}
+	}
+	return reflect.New(rt).Interface()
+}
+
+// SelectValues maps a selected column to its value.
+// Used by the generated code for storing runtime selected columns/expressions.
+type SelectValues map[string]any
+
+// Set sets the value of the given column.
+func (s *SelectValues) Set(name string, v any) {
+	if *s == nil {
+		*s = make(SelectValues)
+	}
+	if pv, ok := v.(*any); ok && pv != nil {
+		v = *pv
+	}
+	(*s)[name] = v
+}
+
+// Get returns the value of the given column.
+func (s SelectValues) Get(name string) (any, error) {
+	v, ok := s[name]
+	if !ok {
+		return nil, fmt.Errorf("%s value was not selected", name)
+	}
+	if v == nil {
+		return nil, nil
+	}
+	switch rv := reflect.Indirect(reflect.ValueOf(v)).Interface().(type) {
+	case NullString:
+		if rv.Valid {
+			return rv.String, nil
+		}
+	case NullInt64:
+		if rv.Valid {
+			return rv.Int64, nil
+		}
+	case NullFloat64:
+		if rv.Valid {
+			return rv.Float64, nil
+		}
+	case NullBool:
+		if rv.Valid {
+			return rv.Bool, nil
+		}
+	case NullTime:
+		if rv.Valid {
+			return rv.Time, nil
+		}
+	case sql.RawBytes:
+		return []byte(rv), nil
+	default:
+		return rv, nil
+	}
+	return nil, nil
+}

dialect/sql/scan_test.go

@@ -0,0 +1,367 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"database/sql"
+	"database/sql/driver"
+	"encoding/json"
+	"testing"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+func TestScanSlice(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"name"}).
+		AddRow("foo").
+		AddRow("bar")
+	var v0 []string
+	require.NoError(t, ScanSlice(toRows(mock), &v0))
+	require.Equal(t, []string{"foo", "bar"}, v0)
+
+	mock = sqlmock.NewRows([]string{"age"}).
+		AddRow(1).
+		AddRow(2)
+	var v1 []int
+	require.NoError(t, ScanSlice(toRows(mock), &v1))
+	require.Equal(t, []int{1, 2}, v1)
+
+	mock = sqlmock.NewRows([]string{"name", "COUNT(*)"}).
+		AddRow("foo", 1).
+		AddRow("bar", 2)
+	var v2 []struct {
+		Name  string
+		Count int
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v2))
+	require.Equal(t, "foo", v2[0].Name)
+	require.Equal(t, "bar", v2[1].Name)
+	require.Equal(t, 1, v2[0].Count)
+	require.Equal(t, 2, v2[1].Count)
+
+	mock = sqlmock.NewRows([]string{"nick_name", "COUNT(*)"}).
+		AddRow("foo", 1).
+		AddRow("bar", 2)
+	var v3 []struct {
+		Count int
+		Name  string `json:"nick_name"`
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v3))
+	require.Equal(t, "foo", v3[0].Name)
+	require.Equal(t, "bar", v3[1].Name)
+	require.Equal(t, 1, v3[0].Count)
+	require.Equal(t, 2, v3[1].Count)
+
+	mock = sqlmock.NewRows([]string{"nick_name", "COUNT(*)"}).
+		AddRow("foo", 1).
+		AddRow("bar", 2)
+	var v4 []*struct {
+		Count   int
+		Name    string `json:"nick_name"`
+		Ignored string `json:"string"`
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v4))
+	require.Equal(t, "foo", v4[0].Name)
+	require.Equal(t, "bar", v4[1].Name)
+	require.Equal(t, 1, v4[0].Count)
+	require.Equal(t, 2, v4[1].Count)
+
+	mock = sqlmock.NewRows([]string{"nick_name", "COUNT(*)"}).
+		AddRow("foo", 1).
+		AddRow("bar", 2)
+	var v5 []*struct {
+		Count int
+		Name  string `json:"name" sql:"nick_name"`
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v5))
+	require.Equal(t, "foo", v5[0].Name)
+	require.Equal(t, "bar", v5[1].Name)
+	require.Equal(t, 1, v5[0].Count)
+	require.Equal(t, 2, v5[1].Count)
+
+	mock = sqlmock.NewRows([]string{"age", "name"}).
+		AddRow(1, nil).
+		AddRow(nil, "a8m")
+	var v6 []struct {
+		Age  NullInt64
+		Name NullString
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v6))
+	require.EqualValues(t, 1, v6[0].Age.Int64)
+	require.False(t, v6[0].Name.Valid)
+	require.False(t, v6[1].Age.Valid)
+	require.Equal(t, "a8m", v6[1].Name.String)
+
+	u1, u2 := uuid.New().String(), uuid.New().String()
+	mock = sqlmock.NewRows([]string{"ids"}).
+		AddRow([]byte(u1)).
+		AddRow([]byte(u2))
+	var ids []uuid.UUID
+	require.NoError(t, ScanSlice(toRows(mock), &ids))
+	require.Equal(t, u1, ids[0].String())
+	require.Equal(t, u2, ids[1].String())
+
+	mock = sqlmock.NewRows([]string{"pids"}).
+		AddRow([]byte(u1)).
+		AddRow([]byte(u2))
+	var pids []*uuid.UUID
+	require.NoError(t, ScanSlice(toRows(mock), &pids))
+	require.Equal(t, u1, pids[0].String())
+	require.Equal(t, u2, pids[1].String())
+
+	mock = sqlmock.NewRows([]string{"id", "first", "last"}).
+		AddRow(10, "Ariel", "Mashraki")
+	err := ScanSlice(toRows(mock), nil)
+	require.EqualError(t, err, "sql/scan: ScanSlice(nil)")
+	type P struct {
+		_     int
+		ID    int
+		First string
+		Last  string
+	}
+	var p []P
+	err = ScanSlice(toRows(mock), p)
+	require.EqualError(t, err, "sql/scan: ScanSlice(non-pointer []sql.P)")
+
+	require.NoError(t, ScanSlice(toRows(mock), &p))
+	require.Equal(t, 10, p[0].ID)
+	require.Equal(t, "Ariel", p[0].First)
+	require.Equal(t, "Mashraki", p[0].Last)
+
+	var pp []struct{ _, id int }
+	mock = sqlmock.NewRows([]string{"id"}).
+		AddRow(10)
+	err = ScanSlice(toRows(mock), &pp)
+	require.EqualError(t, err, "sql/scan: missing struct field for column: id (id)")
+	require.Empty(t, pp)
+}
+
+func TestScanSlice_CamelTags(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"nickName"}).
+		AddRow("foo").
+		AddRow("bar")
+	var v []*struct {
+		NickName string `json:"nickName"`
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v))
+	require.Equal(t, "foo", v[0].NickName)
+}
+
+func TestScanJSON(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"v", "p"}).
+		AddRow([]byte(`{"i": 1, "s":"a8m"}`), []byte(`{"i": 1, "s":"a8m"}`)).
+		AddRow([]byte(`{"i": 2, "s":"tmr"}`), []byte(`{"i": 2, "s":"tmr"}`)).
+		AddRow([]byte(nil), []byte(`null`)).
+		AddRow(nil, nil)
+	var v1 []*struct {
+		V struct {
+			I int    `json:"i"`
+			S string `json:"s"`
+		} `json:"v"`
+		P *struct {
+			I int    `json:"i"`
+			S string `json:"s"`
+		} `json:"p"`
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v1))
+	require.Equal(t, 1, v1[0].V.I)
+	require.Equal(t, "a8m", v1[0].V.S)
+	require.Equal(t, v1[0].V, *v1[0].P)
+	require.Equal(t, 2, v1[1].V.I)
+	require.Equal(t, "tmr", v1[1].V.S)
+	require.Equal(t, v1[1].V, *v1[1].P)
+	require.Equal(t, 0, v1[2].V.I)
+	require.Equal(t, "", v1[2].V.S)
+	require.Nil(t, v1[2].P)
+	require.Equal(t, 0, v1[3].V.I)
+	require.Equal(t, "", v1[3].V.S)
+	require.Nil(t, v1[3].P)
+
+	mock = sqlmock.NewRows([]string{"v", "p"}).
+		AddRow([]byte(`[1]`), []byte(`[1]`)).
+		AddRow([]byte(`[2]`), []byte(`[2]`))
+	var v2 []*struct {
+		V []int  `json:"v"`
+		P *[]int `json:"p"`
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v2))
+	require.Equal(t, []int{1}, v2[0].V)
+	require.Equal(t, v2[0].V, *v2[0].P)
+	require.Equal(t, []int{2}, v2[1].V)
+	require.Equal(t, v2[1].V, *v2[1].P)
+
+	mock = sqlmock.NewRows([]string{"v", "p"}).
+		AddRow([]byte(`null`), []byte(`{}`)).
+		AddRow(nil, nil)
+	var v3 []*struct {
+		V json.RawMessage  `json:"v"`
+		P *json.RawMessage `json:"p"`
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v3))
+	require.Equal(t, json.RawMessage("null"), v3[0].V)
+	require.Equal(t, json.RawMessage("{}"), *v3[0].P)
+	require.Equal(t, json.RawMessage(nil), v3[1].V)
+	require.Nil(t, v3[1].P)
+
+	// Unmarshal errors.
+	mock = sqlmock.NewRows([]string{"v", "p"}).
+		AddRow([]byte(`{invalid}`), []byte(`{}`))
+	require.EqualError(t, ScanSlice(toRows(mock), &v1), `unmarshal field "V": invalid character 'i' looking for beginning of object key string`)
+	mock = sqlmock.NewRows([]string{"v", "p"}).
+		AddRow([]byte(``), []byte(``))
+	require.EqualError(t, ScanSlice(toRows(mock), &v1), `unmarshal field "V": unexpected end of JSON input`)
+}
+
+func TestScanNestedStruct(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"name", "age"}).
+		AddRow("foo", 1).
+		AddRow("bar", 2).
+		AddRow("baz", nil)
+	type T struct{ Name string }
+	var v []struct {
+		T
+		Age int
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v))
+	require.Equal(t, "foo", v[0].Name)
+	require.Equal(t, 1, v[0].Age)
+	require.Equal(t, "bar", v[1].Name)
+	require.Equal(t, 2, v[1].Age)
+	require.Equal(t, "baz", v[2].Name)
+	require.Equal(t, 0, v[2].Age)
+
+	mock = sqlmock.NewRows([]string{"name", "age"}).
+		AddRow("foo", 1).
+		AddRow("bar", nil)
+	type T1 struct{ Name **string }
+	var v1 []struct {
+		T1
+		Age *int
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v1))
+	require.Equal(t, "foo", **v1[0].Name)
+	require.Equal(t, "bar", **v1[1].Name)
+	require.Equal(t, 1, *v1[0].Age)
+	require.Nil(t, v1[1].Age)
+}
+
+func TestScanSlicePtr(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"name"}).
+		AddRow("foo").
+		AddRow("bar")
+	var v0 []*string
+	require.NoError(t, ScanSlice(toRows(mock), &v0))
+	require.Equal(t, "foo", *v0[0])
+	require.Equal(t, "bar", *v0[1])
+
+	mock = sqlmock.NewRows([]string{"age"}).
+		AddRow(1).
+		AddRow(2)
+	var v1 []**int
+	require.NoError(t, ScanSlice(toRows(mock), &v1))
+	require.Equal(t, 1, **v1[0])
+	require.Equal(t, 2, **v1[1])
+
+	mock = sqlmock.NewRows([]string{"age", "name"}).
+		AddRow(1, "a8m").
+		AddRow(2, "nati")
+	var v2 []*struct {
+		Age  *int
+		Name **string
+	}
+	require.NoError(t, ScanSlice(toRows(mock), &v2))
+	require.Equal(t, 1, *v2[0].Age)
+	require.Equal(t, "a8m", **v2[0].Name)
+	require.Equal(t, 2, *v2[1].Age)
+	require.Equal(t, "nati", **v2[1].Name)
+}
+
+func TestScanInt64(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"age"}).
+		AddRow("10").
+		AddRow("20")
+	n, err := ScanInt64(toRows(mock))
+	require.Error(t, err)
+	require.Zero(t, n)
+
+	mock = sqlmock.NewRows([]string{"age", "count"}).
+		AddRow("10", "1")
+	n, err = ScanInt64(toRows(mock))
+	require.Error(t, err)
+	require.Zero(t, n)
+
+	mock = sqlmock.NewRows([]string{"count"}).
+		AddRow(10)
+	n, err = ScanInt64(toRows(mock))
+	require.NoError(t, err)
+	require.EqualValues(t, 10, n)
+
+	mock = sqlmock.NewRows([]string{"count"}).
+		AddRow("10")
+	n, err = ScanInt64(toRows(mock))
+	require.NoError(t, err)
+	require.EqualValues(t, 10, n)
+}
+
+func TestScanValue(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"count"}).
+		AddRow(10)
+	n, err := ScanValue(toRows(mock))
+	require.NoError(t, err)
+	require.EqualValues(t, 10, n)
+}
+
+func TestScanOne(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"name"}).
+		AddRow("10").
+		AddRow("20")
+	err := ScanOne(toRows(mock), new(string))
+	require.Error(t, err, "multiple lines")
+
+	mock = sqlmock.NewRows([]string{"name"}).
+		AddRow("10")
+	err = ScanOne(toRows(mock), "")
+	require.Error(t, err, "not a pointer")
+
+	mock = sqlmock.NewRows([]string{"name"}).
+		AddRow("10")
+	var s string
+	err = ScanOne(toRows(mock), &s)
+	require.NoError(t, err)
+	require.Equal(t, "10", s)
+}
+
+func TestInterface(t *testing.T) {
+	mock := sqlmock.NewRows([]string{"age"}).
+		AddRow("10").
+		AddRow("20")
+	var values []driver.Value
+	err := ScanSlice(toRows(mock), &values)
+	require.NoError(t, err)
+	require.Equal(t, []driver.Value{"10", "20"}, values)
+
+	mock = sqlmock.NewRows([]string{"age"}).
+		AddRow(10).
+		AddRow(20)
+	values = values[:0:0]
+	err = ScanSlice(toRows(mock), &values)
+	require.NoError(t, err)
+	require.Equal(t, []driver.Value{int64(10), int64(20)}, values)
+}
+
+func TestScanTypeOf(t *testing.T) {
+	mock := &Rows{ColumnScanner: toRows(sqlmock.NewRows([]string{"age"}).AddRow(10))}
+	tv := ScanTypeOf(mock, 0)
+	require.IsType(t, (*any)(nil), tv)
+}
+
+func toRows(mrows *sqlmock.Rows) *sql.Rows {
+	db, mock, _ := sqlmock.New()
+	mock.ExpectQuery("").WillReturnRows(mrows)
+	rows, _ := db.Query("")
+	return rows
+}

dialect/sql/schema/migrate.go

@@ -0,0 +1,160 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"context"
+	"fmt"
+	"math"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/schema/field"
+)
+
+const (
+	// TypeTable defines the table name holding the type information.
+	TypeTable = "ent_types"
+
+	// MaxTypes defines the max number of types can be created when
+	// defining universal ids. The left 16-bits are reserved.
+	MaxTypes = math.MaxUint16
+)
+
+// NewTypesTable returns a new table for holding the global-id information.
+func NewTypesTable() *Table {
+	return NewTable(TypeTable).
+		AddPrimary(&Column{Name: "id", Type: field.TypeUint, Increment: true}).
+		AddColumn(&Column{Name: "type", Type: field.TypeString, Unique: true})
+}
+
+// MigrateOption allows configuring Atlas using functional arguments.
+type MigrateOption func(*Atlas)
+
+// WithGlobalUniqueID sets the universal ids options to the migration.
+// Defaults to false.
+func WithGlobalUniqueID(b bool) MigrateOption {
+	return func(a *Atlas) {
+		a.universalID = b
+	}
+}
+
+// WithIndent sets Atlas to generate SQL statements with indentation.
+// An empty string indicates no indentation.
+func WithIndent(indent string) MigrateOption {
+	return func(a *Atlas) {
+		a.indent = indent
+	}
+}
+
+// WithErrNoPlan sets Atlas to returns a migrate.ErrNoPlan in case
+// the migration plan is empty. Defaults to false.
+func WithErrNoPlan(b bool) MigrateOption {
+	return func(a *Atlas) {
+		a.errNoPlan = b
+	}
+}
+
+// WithSchemaName sets the database schema for the migration.
+// If not set, the CURRENT_SCHEMA() is used.
+func WithSchemaName(ns string) MigrateOption {
+	return func(a *Atlas) {
+		a.schema = ns
+	}
+}
+
+// WithDropColumn sets the columns dropping option to the migration.
+// Defaults to false.
+func WithDropColumn(b bool) MigrateOption {
+	return func(a *Atlas) {
+		a.dropColumns = b
+	}
+}
+
+// WithDropIndex sets the indexes dropping option to the migration.
+// Defaults to false.
+func WithDropIndex(b bool) MigrateOption {
+	return func(a *Atlas) {
+		a.dropIndexes = b
+	}
+}
+
+// WithForeignKeys enables creating foreign-key in ddl. Defaults to true.
+func WithForeignKeys(b bool) MigrateOption {
+	return func(a *Atlas) {
+		a.withForeignKeys = b
+	}
+}
+
+// WithHooks adds a list of hooks to the schema migration.
+func WithHooks(hooks ...Hook) MigrateOption {
+	return func(a *Atlas) {
+		a.hooks = append(a.hooks, hooks...)
+	}
+}
+
+type (
+	// Creator is the interface that wraps the Create method.
+	Creator interface {
+		// Create creates the given tables in the database. See Migrate.Create for more details.
+		Create(context.Context, ...*Table) error
+	}
+
+	// The CreateFunc type is an adapter to allow the use of ordinary function as Creator.
+	// If f is a function with the appropriate signature, CreateFunc(f) is a Creator that calls f.
+	CreateFunc func(context.Context, ...*Table) error
+
+	// Hook defines the "create middleware". A function that gets a Creator and returns a Creator.
+	// For example:
+	//
+	//	hook := func(next schema.Creator) schema.Creator {
+	//		return schema.CreateFunc(func(ctx context.Context, tables ...*schema.Table) error {
+	//			fmt.Println("Tables:", tables)
+	//			return next.Create(ctx, tables...)
+	//		})
+	//	}
+	//
+	Hook func(Creator) Creator
+)
+
+// Create calls f(ctx, tables...).
+func (f CreateFunc) Create(ctx context.Context, tables ...*Table) error {
+	return f(ctx, tables...)
+}
+
+// exist checks if the given COUNT query returns a value >= 1.
+func exist(ctx context.Context, conn dialect.ExecQuerier, query string, args ...any) (bool, error) {
+	rows := &sql.Rows{}
+	if err := conn.Query(ctx, query, args, rows); err != nil {
+		return false, fmt.Errorf("reading schema information %w", err)
+	}
+	defer rows.Close()
+	n, err := sql.ScanInt(rows)
+	if err != nil {
+		return false, err
+	}
+	return n > 0, nil
+}
+
+func indexOf(a []string, s string) int {
+	for i := range a {
+		if a[i] == s {
+			return i
+		}
+	}
+	return -1
+}
+
+type sqlDialect interface {
+	atBuilder
+	dialect.Driver
+	init(context.Context) error
+	tableExist(context.Context, dialect.ExecQuerier, string) (bool, error)
+}
+
+// verifyRanger wraps the method for verifying global-id range correctness.
+type verifyRanger interface {
+	verifyRange(context.Context, dialect.ExecQuerier, *Table, int64) error
+}

dialect/sql/schema/migrate_test.go

@@ -0,0 +1,464 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strings"
+	"sync"
+	"testing"
+	"text/template"
+	"time"
+
+	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/schema"
+	"ariga.io/atlas/sql/sqlite"
+	"ariga.io/atlas/sql/sqltool"
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/entsql"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/schema/field"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	_ "github.com/mattn/go-sqlite3"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMigrate_SchemaName(t *testing.T) {
+	db, mk, err := sqlmock.New()
+	require.NoError(t, err)
+	mk.ExpectQuery(escape("SHOW server_version_num")).
+		WillReturnRows(sqlmock.NewRows([]string{"server_version_num"}).AddRow("130000"))
+	mk.ExpectQuery(escape("SELECT current_setting('server_version_num'), current_setting('default_table_access_method', true), current_setting('crdb_version', true)")).
+		WillReturnRows(sqlmock.NewRows([]string{"current_setting", "current_setting", "current_setting"}).AddRow("130000", "heap", ""))
+	mk.ExpectQuery("SELECT nspname AS schema_name,.+").
+		WithArgs("public"). // Schema "public" param is used.
+		WillReturnRows(sqlmock.NewRows([]string{"schema_name", "comment"}).AddRow("public", "default schema"))
+	mk.ExpectQuery("SELECT t3.oid, t1.table_schema,.+").
+		WillReturnRows(sqlmock.NewRows([]string{}))
+	m, err := NewMigrate(sql.OpenDB("postgres", db), WithSchemaName("public"), WithDiffHook(func(next Differ) Differ {
+		return DiffFunc(func(current, desired *schema.Schema) ([]schema.Change, error) {
+			return nil, nil // Noop.
+		})
+	}))
+	require.NoError(t, err)
+	require.NoError(t, m.Create(context.Background()))
+	require.NoError(t, mk.ExpectationsWereMet())
+
+	// Without schema name the CURRENT_SCHEMA is used.
+	mk.ExpectQuery(escape("SHOW server_version_num")).
+		WillReturnRows(sqlmock.NewRows([]string{"server_version_num"}).AddRow("130000"))
+	mk.ExpectQuery(escape("SELECT current_setting('server_version_num'), current_setting('default_table_access_method', true), current_setting('crdb_version', true)")).
+		WillReturnRows(sqlmock.NewRows([]string{"current_setting", "current_setting", "current_setting"}).AddRow("130000", "heap", ""))
+	mk.ExpectQuery("SELECT nspname AS schema_name,.+CURRENT_SCHEMA().+").
+		WillReturnRows(sqlmock.NewRows([]string{"schema_name", "comment"}).AddRow("public", "default schema"))
+	mk.ExpectQuery("SELECT t3.oid, t1.table_schema,.+").
+		WillReturnRows(sqlmock.NewRows([]string{}))
+	m, err = NewMigrate(sql.OpenDB("postgres", db), WithDiffHook(func(next Differ) Differ {
+		return DiffFunc(func(current, desired *schema.Schema) ([]schema.Change, error) {
+			return nil, nil // Noop.
+		})
+	}))
+	require.NoError(t, err)
+	require.NoError(t, m.Create(context.Background()))
+}
+
+func escape(query string) string {
+	rows := strings.Split(query, "\n")
+	for i := range rows {
+		rows[i] = strings.TrimPrefix(rows[i], " ")
+	}
+	query = strings.Join(rows, " ")
+	return strings.TrimSpace(regexp.QuoteMeta(query)) + "$"
+}
+
+func TestMigrate_Formatter(t *testing.T) {
+	db, _, err := sqlmock.New()
+	require.NoError(t, err)
+
+	// If no formatter is given it will be set according to the given migration directory implementation.
+	for _, tt := range []struct {
+		dir migrate.Dir
+		fmt migrate.Formatter
+	}{
+		{&migrate.LocalDir{}, sqltool.GolangMigrateFormatter},
+		{&sqltool.GolangMigrateDir{}, sqltool.GolangMigrateFormatter},
+		{&sqltool.GooseDir{}, sqltool.GooseFormatter},
+		{&sqltool.DBMateDir{}, sqltool.DBMateFormatter},
+		{&sqltool.FlywayDir{}, sqltool.FlywayFormatter},
+		{&sqltool.LiquibaseDir{}, sqltool.LiquibaseFormatter},
+		{struct{ migrate.Dir }{}, sqltool.GolangMigrateFormatter}, // default one if migration dir is unknown
+	} {
+		m, err := NewMigrate(sql.OpenDB("", db), WithDir(tt.dir))
+		require.NoError(t, err)
+		require.Equal(t, tt.fmt, m.fmt)
+	}
+
+	// If a formatter is given, it is not overridden.
+	m, err := NewMigrate(sql.OpenDB("", db), WithDir(&migrate.LocalDir{}), WithFormatter(migrate.DefaultFormatter))
+	require.NoError(t, err)
+	require.Equal(t, migrate.DefaultFormatter, m.fmt)
+}
+
+func TestMigrate_DiffJoinTableAllocationBC(t *testing.T) {
+	// Due to a bug in previous versions, if the universal ID option was enabled and the schema did contain an M2M
+	// relation, the join table would have had an entry for the join table in the types table. This test ensures,
+	// that the PK range allocated for the join table stays in place, since it's removal would break existing projects
+	// due to shifted ranges.
+
+	db, err := sql.Open(dialect.SQLite, "file:test?mode=memory&_fk=1")
+	require.NoError(t, err)
+
+	// Mock an existing database with an allocation for a join table.
+	for _, stmt := range []string{
+		"CREATE TABLE `groups` (`id` integer NOT NULL PRIMARY KEY AUTOINCREMENT, `name` text NOT NULL);",
+		"CREATE INDEX `short` ON `groups` (`id`);",
+		"CREATE INDEX `long____________________________1cb2e7e47a309191385af4ad320875b1` ON `groups` (`id`);",
+		"CREATE TABLE `users` (`id` integer NOT NULL PRIMARY KEY AUTOINCREMENT, `name` text NOT NULL);",
+		"INSERT INTO sqlite_sequence (name, seq) VALUES (\"users\", 4294967296);",
+		"CREATE TABLE `user_groups` (`user_id` integer NOT NULL, `group_id` integer NOT NULL, PRIMARY KEY (`user_id`, `group_id`), CONSTRAINT `user_groups_user_id` FOREIGN KEY (`user_id`) REFERENCES `users` (`id`) ON DELETE CASCADE, CONSTRAINT `user_groups_group_id` FOREIGN KEY (`group_id`) REFERENCES `groups` (`id`) ON DELETE CASCADE);",
+		"INSERT INTO sqlite_sequence (name, seq) VALUES (\"user_groups\", 8589934592);",
+		"CREATE TABLE `ent_types` (`id` integer NOT NULL PRIMARY KEY AUTOINCREMENT, `type` text NOT NULL);",
+		"CREATE UNIQUE INDEX `ent_types_type_key` ON `ent_types` (`type`);",
+		"INSERT INTO `ent_types` (`type`) VALUES ('groups'), ('users'), ('user_groups');",
+		"INSERT INTO `groups` (`name`) VALUES ('seniors'), ('juniors')",
+		"INSERT INTO `users` (`name`) VALUES ('masseelch'), ('a8m'), ('rotemtam')",
+		"INSERT INTO `user_groups` (`user_id`, `group_id`) VALUES (4294967297, 1), (4294967298, 1), (4294967299, 2)",
+	} {
+		_, err := db.ExecContext(context.Background(), stmt)
+		require.NoError(t, err)
+	}
+
+	// Expect to have no changes when migration runs with fix.
+	m, err := NewMigrate(db, WithGlobalUniqueID(true), WithDiffHook(func(next Differ) Differ {
+		return DiffFunc(func(current, desired *schema.Schema) ([]schema.Change, error) {
+			changes, err := next.Diff(current, desired)
+			if err != nil {
+				return nil, err
+			}
+			require.Len(t, changes, 0)
+			return changes, nil
+		})
+	}))
+	require.NoError(t, err)
+	require.NoError(t, m.Create(context.Background(), tables...))
+
+	// Expect to have no changes to the allocation when the join table is dropped.
+	m, err = NewMigrate(db, WithGlobalUniqueID(true))
+	require.NoError(t, err)
+	require.NoError(t, m.Create(context.Background(), groupsTable, usersTable))
+
+	rows, err := db.QueryContext(context.Background(), "SELECT `type` from `ent_types` ORDER BY `id` ASC")
+	require.NoError(t, err)
+	var types []string
+	for rows.Next() {
+		var typ string
+		require.NoError(t, rows.Scan(&typ))
+		types = append(types, typ)
+	}
+	require.NoError(t, rows.Err())
+	require.Equal(t, []string{"groups", "users", "user_groups"}, types)
+}
+
+var (
+	groupsColumns = []*Column{
+		{Name: "id", Type: field.TypeInt, Increment: true},
+		{Name: "name", Type: field.TypeString},
+	}
+	groupsTable = &Table{
+		Name:       "groups",
+		Columns:    groupsColumns,
+		PrimaryKey: []*Column{groupsColumns[0]},
+		Indexes: []*Index{
+			{
+				Name:    "short",
+				Columns: []*Column{groupsColumns[0]}},
+			{
+				Name:    "long_" + strings.Repeat("_", 60),
+				Columns: []*Column{groupsColumns[0]},
+			},
+		},
+	}
+	usersColumns = []*Column{
+		{Name: "id", Type: field.TypeInt, Increment: true},
+		{Name: "name", Type: field.TypeString},
+	}
+	usersTable = &Table{
+		Name:       "users",
+		Columns:    usersColumns,
+		PrimaryKey: []*Column{usersColumns[0]},
+	}
+	userGroupsColumns = []*Column{
+		{Name: "user_id", Type: field.TypeInt},
+		{Name: "group_id", Type: field.TypeInt},
+	}
+	userGroupsTable = &Table{
+		Name:       "user_groups",
+		Columns:    userGroupsColumns,
+		PrimaryKey: []*Column{userGroupsColumns[0], userGroupsColumns[1]},
+		ForeignKeys: []*ForeignKey{
+			{
+				Symbol:     "user_groups_user_id",
+				Columns:    []*Column{userGroupsColumns[0]},
+				RefColumns: []*Column{usersColumns[0]},
+				OnDelete:   Cascade,
+			},
+			{
+				Symbol:     "user_groups_group_id",
+				Columns:    []*Column{userGroupsColumns[1]},
+				RefColumns: []*Column{groupsColumns[0]},
+				OnDelete:   Cascade,
+			},
+		},
+	}
+	tables = []*Table{
+		groupsTable,
+		usersTable,
+		userGroupsTable,
+	}
+	petColumns = []*Column{
+		{Name: "id", Type: field.TypeInt, Increment: true},
+	}
+	petsTable = &Table{
+		Name:       "pets",
+		Columns:    petColumns,
+		PrimaryKey: petColumns,
+	}
+)
+
+func init() {
+	userGroupsTable.ForeignKeys[0].RefTable = usersTable
+	userGroupsTable.ForeignKeys[1].RefTable = groupsTable
+}
+
+func TestMigrate_Diff(t *testing.T) {
+	ctx := context.Background()
+
+	db, err := sql.Open(dialect.SQLite, "file:test?mode=memory&_fk=1")
+	require.NoError(t, err)
+
+	p := t.TempDir()
+	d, err := migrate.NewLocalDir(p)
+	require.NoError(t, err)
+
+	m, err := NewMigrate(db, WithDir(d))
+	require.NoError(t, err)
+	require.NoError(t, m.Diff(ctx, &Table{Name: "users"}))
+	v := time.Now().UTC().Format("20060102150405")
+	requireFileEqual(t, filepath.Join(p, v+"_changes.up.sql"), "-- create \"users\" table\nCREATE TABLE `users` ();\n")
+	requireFileEqual(t, filepath.Join(p, v+"_changes.down.sql"), "-- reverse: create \"users\" table\nDROP TABLE `users`;\n")
+	require.FileExists(t, filepath.Join(p, migrate.HashFileName))
+
+	// Test integrity file.
+	p = t.TempDir()
+	d, err = migrate.NewLocalDir(p)
+	require.NoError(t, err)
+	m, err = NewMigrate(db, WithDir(d))
+	require.NoError(t, err)
+	require.NoError(t, m.Diff(ctx, &Table{Name: "users"}))
+	requireFileEqual(t, filepath.Join(p, v+"_changes.up.sql"), "-- create \"users\" table\nCREATE TABLE `users` ();\n")
+	requireFileEqual(t, filepath.Join(p, v+"_changes.down.sql"), "-- reverse: create \"users\" table\nDROP TABLE `users`;\n")
+	require.FileExists(t, filepath.Join(p, migrate.HashFileName))
+	require.NoError(t, d.WriteFile("tmp.sql", nil))
+	require.ErrorIs(t, m.Diff(ctx, &Table{Name: "users"}), migrate.ErrChecksumMismatch)
+
+	p = t.TempDir()
+	d, err = migrate.NewLocalDir(p)
+	require.NoError(t, err)
+	f, err := migrate.NewTemplateFormatter(
+		template.Must(template.New("").Parse("{{ .Name }}.sql")),
+		template.Must(template.New("").Parse(
+			`{{ range .Changes }}{{ printf "%s;\n" .Cmd }}{{ end }}`,
+		)),
+	)
+	require.NoError(t, err)
+
+	// Join tables (mapping between user and group) will not result in an entry to the types table.
+	m, err = NewMigrate(db, WithFormatter(f), WithDir(d), WithGlobalUniqueID(true))
+	require.NoError(t, err)
+	require.NoError(t, m.Diff(ctx, tables...))
+	changesSQL := strings.Join([]string{
+		"CREATE TABLE `groups` (`id` integer NOT NULL PRIMARY KEY AUTOINCREMENT, `name` text NOT NULL);",
+		"CREATE INDEX `short` ON `groups` (`id`);",
+		"CREATE INDEX `long____________________________1cb2e7e47a309191385af4ad320875b1` ON `groups` (`id`);",
+		"CREATE TABLE `users` (`id` integer NOT NULL PRIMARY KEY AUTOINCREMENT, `name` text NOT NULL);",
+		fmt.Sprintf("INSERT INTO sqlite_sequence (name, seq) VALUES (\"users\", %d);", 1<<32),
+		"CREATE TABLE `user_groups` (`user_id` integer NOT NULL, `group_id` integer NOT NULL, PRIMARY KEY (`user_id`, `group_id`), CONSTRAINT `user_groups_user_id` FOREIGN KEY (`user_id`) REFERENCES `users` (`id`) ON DELETE CASCADE, CONSTRAINT `user_groups_group_id` FOREIGN KEY (`group_id`) REFERENCES `groups` (`id`) ON DELETE CASCADE);",
+		"CREATE TABLE `ent_types` (`id` integer NOT NULL PRIMARY KEY AUTOINCREMENT, `type` text NOT NULL);",
+		"CREATE UNIQUE INDEX `ent_types_type_key` ON `ent_types` (`type`);",
+		"INSERT INTO `ent_types` (`type`) VALUES ('groups'), ('users');",
+		"",
+	}, "\n")
+	requireFileEqual(t, filepath.Join(p, "changes.sql"), changesSQL)
+
+	// Skipping table creation should write only the ent_type insertion.
+	m, err = NewMigrate(db, WithFormatter(f), WithDir(d), WithGlobalUniqueID(true), WithDiffOptions(schema.DiffSkipChanges(&schema.AddTable{})))
+	require.NoError(t, err)
+	require.NoError(t, m.Diff(ctx, tables...))
+	requireFileEqual(t, filepath.Join(p, "changes.sql"), "INSERT INTO `ent_types` (`type`) VALUES ('groups'), ('users');\n")
+
+	// Enable indentations.
+	m, err = NewMigrate(db, WithFormatter(f), WithDir(d), WithGlobalUniqueID(true), WithIndent("  "))
+	require.NoError(t, err)
+	// Adding another node will result in a new entry to the TypeTable (without actually creating it).
+	_, err = db.ExecContext(ctx, changesSQL, nil, nil)
+	require.NoError(t, err)
+	require.NoError(t, m.NamedDiff(ctx, "changes_2", petsTable))
+	requireFileEqual(t,
+		filepath.Join(p, "changes_2.sql"), strings.Join([]string{
+			"CREATE TABLE `pets` (\n  `id` integer NOT NULL PRIMARY KEY AUTOINCREMENT\n);",
+			fmt.Sprintf("INSERT INTO sqlite_sequence (name, seq) VALUES (\"pets\", %d);", 2<<32),
+			"INSERT INTO `ent_types` (`type`) VALUES ('pets');", "",
+		}, "\n"))
+
+	// Checksum will be updated as well.
+	require.NoError(t, migrate.Validate(d))
+
+	require.NoError(t, m.NamedDiff(ctx, "no_changes"), "should not error if WithErrNoPlan is not set")
+	// Enable WithErrNoPlan.
+	m, err = NewMigrate(db, WithFormatter(f), WithDir(d), WithGlobalUniqueID(true), WithErrNoPlan(true))
+	require.NoError(t, err)
+	err = m.NamedDiff(ctx, "no_changes")
+	require.ErrorIs(t, err, migrate.ErrNoPlan)
+}
+
+func requireFileEqual(t *testing.T, name, contents string) {
+	c, err := os.ReadFile(name)
+	require.NoError(t, err)
+	require.Equal(t, contents, string(c))
+}
+
+func TestMigrateWithoutForeignKeys(t *testing.T) {
+	tbl := &schema.Table{
+		Name: "tbl",
+		Columns: []*schema.Column{
+			{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}},
+		},
+	}
+	fk := &schema.ForeignKey{
+		Symbol:     "fk",
+		Table:      tbl,
+		Columns:    tbl.Columns[1:],
+		RefTable:   tbl,
+		RefColumns: tbl.Columns[:1],
+		OnUpdate:   schema.NoAction,
+		OnDelete:   schema.Cascade,
+	}
+	tbl.ForeignKeys = append(tbl.ForeignKeys, fk)
+	t.Run("AddTable", func(t *testing.T) {
+		mdiff := DiffFunc(func(_, _ *schema.Schema) ([]schema.Change, error) {
+			return []schema.Change{
+				&schema.AddTable{
+					T: tbl,
+				},
+			}, nil
+		})
+		df, err := withoutForeignKeys(mdiff).Diff(nil, nil)
+		require.NoError(t, err)
+		require.Len(t, df, 1)
+		actual, ok := df[0].(*schema.AddTable)
+		require.True(t, ok)
+		require.Nil(t, actual.T.ForeignKeys)
+	})
+	t.Run("ModifyTable", func(t *testing.T) {
+		mdiff := DiffFunc(func(_, _ *schema.Schema) ([]schema.Change, error) {
+			return []schema.Change{
+				&schema.ModifyTable{
+					T: tbl,
+					Changes: []schema.Change{
+						&schema.AddIndex{
+							I: &schema.Index{
+								Name: "id_key",
+								Parts: []*schema.IndexPart{
+									{C: tbl.Columns[0]},
+								},
+							},
+						},
+						&schema.DropForeignKey{
+							F: fk,
+						},
+						&schema.AddForeignKey{
+							F: fk,
+						},
+						&schema.ModifyForeignKey{
+							From:   fk,
+							To:     fk,
+							Change: schema.ChangeRefColumn,
+						},
+						&schema.AddColumn{
+							C: &schema.Column{Name: "name", Type: &schema.ColumnType{Type: &schema.StringType{T: "varchar(255)"}}},
+						},
+					},
+				},
+			}, nil
+		})
+		df, err := withoutForeignKeys(mdiff).Diff(nil, nil)
+		require.NoError(t, err)
+		require.Len(t, df, 1)
+		actual, ok := df[0].(*schema.ModifyTable)
+		require.True(t, ok)
+		require.Len(t, actual.Changes, 2)
+		addIndex, ok := actual.Changes[0].(*schema.AddIndex)
+		require.True(t, ok)
+		require.EqualValues(t, "id_key", addIndex.I.Name)
+		addColumn, ok := actual.Changes[1].(*schema.AddColumn)
+		require.True(t, ok)
+		require.EqualValues(t, "name", addColumn.C.Name)
+	})
+}
+
+func TestAtlas_StateReader(t *testing.T) {
+	db, err := sql.Open(dialect.SQLite, "file:test?mode=memory&_fk=1")
+	require.NoError(t, err)
+	m, err := NewMigrate(db)
+	require.NoError(t, err)
+	realm, err := m.StateReader(&Table{
+		Name: "users",
+		Columns: []*Column{
+			{Name: "id", Type: field.TypeInt64, Increment: true},
+			{Name: "name", Type: field.TypeString},
+			{Name: "active", Type: field.TypeBool},
+		},
+		Annotation: &entsql.Annotation{
+			IncrementStart: func(i int) *int { return &i }(100),
+		},
+	}).ReadState(context.Background())
+	require.NoError(t, err)
+	require.NotNil(t, realm)
+	require.Len(t, realm.Schemas, 1)
+	require.Len(t, realm.Schemas[0].Tables, 1)
+	require.Equal(t, "users", realm.Schemas[0].Tables[0].Name)
+	require.Equal(t, []schema.Attr{&sqlite.AutoIncrement{Seq: 100}}, realm.Schemas[0].Tables[0].Attrs)
+	require.Equal(t,
+		realm.Schemas[0].Tables[0].Columns,
+		[]*schema.Column{
+			schema.NewIntColumn("id", "integer").
+				AddAttrs(&sqlite.AutoIncrement{}),
+			schema.NewStringColumn("name", "text"),
+			schema.NewBoolColumn("active", "bool"),
+		},
+	)
+}
+
+func TestAtlas_ParallelCreate(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(10)
+	for i := 0; i < 10; i++ {
+		db, err := sql.Open(dialect.SQLite, fmt.Sprintf("file:test-%d?mode=memory&_fk=1", i))
+		require.NoError(t, err)
+		m, err := NewMigrate(db)
+		require.NoError(t, err)
+		go func() {
+			defer wg.Done()
+			require.NoError(t, m.Create(context.Background(), petsTable))
+			require.NoError(t, db.Close())
+		}()
+	}
+	wg.Wait()
+}

dialect/sql/schema/mysql.go

@@ -0,0 +1,333 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"reflect"
+	"strconv"
+	"strings"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/schema/field"
+
+	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/mysql"
+	"ariga.io/atlas/sql/schema"
+)
+
+// MySQL adapter for Atlas migration engine.
+type MySQL struct {
+	dialect.Driver
+	schema  string
+	version string
+}
+
+// init loads the MySQL version from the database for later use in the migration process.
+func (d *MySQL) init(ctx context.Context) error {
+	if d.version != "" {
+		return nil // already initialized.
+	}
+	rows := &sql.Rows{}
+	if err := d.Query(ctx, "SHOW VARIABLES LIKE 'version'", []any{}, rows); err != nil {
+		return fmt.Errorf("mysql: querying mysql version %w", err)
+	}
+	defer rows.Close()
+	if !rows.Next() {
+		if err := rows.Err(); err != nil {
+			return err
+		}
+		return fmt.Errorf("mysql: version variable was not found")
+	}
+	version := make([]string, 2)
+	if err := rows.Scan(&version[0], &version[1]); err != nil {
+		return fmt.Errorf("mysql: scanning mysql version: %w", err)
+	}
+	d.version = version[1]
+	return nil
+}
+
+func (d *MySQL) tableExist(ctx context.Context, conn dialect.ExecQuerier, name string) (bool, error) {
+	query, args := sql.Select(sql.Count("*")).From(sql.Table("TABLES").Schema("INFORMATION_SCHEMA")).
+		Where(sql.And(
+			d.matchSchema(),
+			sql.EQ("TABLE_NAME", name),
+		)).Query()
+	return exist(ctx, conn, query, args...)
+}
+
+// matchSchema returns the predicate for matching table schema.
+func (d *MySQL) matchSchema(columns ...string) *sql.Predicate {
+	column := "TABLE_SCHEMA"
+	if len(columns) > 0 {
+		column = columns[0]
+	}
+	if d.schema != "" {
+		return sql.EQ(column, d.schema)
+	}
+	return sql.EQ(column, sql.Raw("(SELECT DATABASE())"))
+}
+
+func (d *MySQL) atOpen(conn dialect.ExecQuerier) (migrate.Driver, error) {
+	return mysql.Open(&db{ExecQuerier: conn})
+}
+
+func (d *MySQL) atTable(t1 *Table, t2 *schema.Table) {
+	t2.SetCharset("utf8mb4").SetCollation("utf8mb4_bin")
+	if t1.Annotation == nil {
+		return
+	}
+	if charset := t1.Annotation.Charset; charset != "" {
+		t2.SetCharset(charset)
+	}
+	if collate := t1.Annotation.Collation; collate != "" {
+		t2.SetCollation(collate)
+	}
+	if opts := t1.Annotation.Options; opts != "" {
+		t2.AddAttrs(&mysql.CreateOptions{
+			V: opts,
+		})
+	}
+	// Check if the connected database supports the CHECK clause.
+	// For MySQL, is >= "8.0.16" and for MariaDB it is "10.2.1".
+	v1, v2 := d.version, "8.0.16"
+	if v, ok := d.mariadb(); ok {
+		v1, v2 = v, "10.2.1"
+	}
+	if compareVersions(v1, v2) >= 0 {
+		setAtChecks(t1, t2)
+	}
+}
+
+func (d *MySQL) supportsDefault(c *Column) bool {
+	_, maria := d.mariadb()
+	switch c.Default.(type) {
+	case Expr, map[string]Expr:
+		if maria {
+			return compareVersions(d.version, "10.2.0") >= 0
+		}
+		return c.supportDefault() && compareVersions(d.version, "8.0.0") >= 0
+	default:
+		return c.supportDefault() || maria
+	}
+}
+
+func (d *MySQL) supportsUUID() bool {
+	_, maria := d.mariadb()
+	return maria && compareVersions(d.version, "10.7.0") >= 0
+}
+
+func (d *MySQL) atTypeC(c1 *Column, c2 *schema.Column) error {
+	if c1.SchemaType != nil && c1.SchemaType[dialect.MySQL] != "" {
+		t, err := mysql.ParseType(strings.ToLower(c1.SchemaType[dialect.MySQL]))
+		if err != nil {
+			return err
+		}
+		c2.Type.Type = t
+		return nil
+	}
+	var t schema.Type
+	switch c1.Type {
+	case field.TypeBool:
+		t = &schema.BoolType{T: "boolean"}
+	case field.TypeBlob:
+		return fmt.Errorf("blob fields are not stored in the database")
+	case field.TypeInt8:
+		t = &schema.IntegerType{T: mysql.TypeTinyInt}
+	case field.TypeUint8:
+		t = &schema.IntegerType{T: mysql.TypeTinyInt, Unsigned: true}
+	case field.TypeInt16:
+		t = &schema.IntegerType{T: mysql.TypeSmallInt}
+	case field.TypeUint16:
+		t = &schema.IntegerType{T: mysql.TypeSmallInt, Unsigned: true}
+	case field.TypeInt32:
+		t = &schema.IntegerType{T: mysql.TypeInt}
+	case field.TypeUint32:
+		t = &schema.IntegerType{T: mysql.TypeInt, Unsigned: true}
+	case field.TypeInt, field.TypeInt64:
+		t = &schema.IntegerType{T: mysql.TypeBigInt}
+	case field.TypeUint, field.TypeUint64:
+		t = &schema.IntegerType{T: mysql.TypeBigInt, Unsigned: true}
+	case field.TypeBytes:
+		size := int64(math.MaxUint16)
+		if c1.Size > 0 {
+			size = c1.Size
+		}
+		switch {
+		case size <= math.MaxUint8:
+			t = &schema.BinaryType{T: mysql.TypeTinyBlob}
+		case size <= math.MaxUint16:
+			t = &schema.BinaryType{T: mysql.TypeBlob}
+		case size < 1<<24:
+			t = &schema.BinaryType{T: mysql.TypeMediumBlob}
+		case size <= math.MaxUint32:
+			t = &schema.BinaryType{T: mysql.TypeLongBlob}
+		}
+	case field.TypeJSON:
+		t = &schema.JSONType{T: mysql.TypeJSON}
+		if compareVersions(d.version, "5.7.8") == -1 {
+			t = &schema.BinaryType{T: mysql.TypeLongBlob}
+		}
+	case field.TypeString:
+		size := c1.Size
+		if size == 0 {
+			size = d.defaultSize(c1)
+		}
+		switch {
+		case c1.typ == "tinytext", c1.typ == "text":
+			t = &schema.StringType{T: c1.typ}
+		case size <= math.MaxUint16:
+			t = &schema.StringType{T: mysql.TypeVarchar, Size: int(size)}
+		case size == 1<<24-1:
+			t = &schema.StringType{T: mysql.TypeMediumText}
+		default:
+			t = &schema.StringType{T: mysql.TypeLongText}
+		}
+	case field.TypeFloat32, field.TypeFloat64:
+		t = &schema.FloatType{T: c1.scanTypeOr(mysql.TypeDouble)}
+	case field.TypeTime:
+		t = &schema.TimeType{T: c1.scanTypeOr(mysql.TypeTimestamp)}
+		// In MariaDB or in MySQL < v8.0.2, the TIMESTAMP column has both `DEFAULT CURRENT_TIMESTAMP`
+		// and `ON UPDATE CURRENT_TIMESTAMP` if neither is specified explicitly. this behavior is
+		// suppressed if the column is defined with a `DEFAULT` clause or with the `NULL` attribute.
+		if _, maria := d.mariadb(); maria || compareVersions(d.version, "8.0.2") == -1 && c1.Default == nil {
+			c2.SetNull(c1.Attr == "")
+		}
+	case field.TypeEnum:
+		t = &schema.EnumType{T: mysql.TypeEnum, Values: c1.Enums}
+	case field.TypeUUID:
+		if d.supportsUUID() {
+			// Native support for the uuid type
+			t = &schema.UUIDType{T: mysql.TypeUUID}
+		} else {
+			// "CHAR(X) BINARY" is treated as "CHAR(X) COLLATE latin1_bin", and in MySQL < 8,
+			// and "COLLATE utf8mb4_bin" in MySQL >= 8. However we already set the table to
+			t = &schema.StringType{T: mysql.TypeChar, Size: 36}
+			c2.SetCollation("utf8mb4_bin")
+		}
+	default:
+		t, err := mysql.ParseType(strings.ToLower(c1.typ))
+		if err != nil {
+			return err
+		}
+		c2.Type.Type = t
+	}
+	c2.Type.Type = t
+	return nil
+}
+
+func (d *MySQL) atUniqueC(t1 *Table, c1 *Column, t2 *schema.Table, c2 *schema.Column) {
+	// For UNIQUE columns, MySQL create an implicit index
+	// named as the column with an extra index in case the
+	// name is already taken (<e.g. c>, <c_2>, <c_3>, ...).
+	for _, idx := range t1.Indexes {
+		// Index also defined explicitly, and will be add in atIndexes.
+		if idx.Unique && d.atImplicitIndexName(idx, c1) {
+			return
+		}
+	}
+	t2.AddIndexes(schema.NewUniqueIndex(c1.Name).AddColumns(c2))
+}
+
+func (d *MySQL) atIncrementC(t *schema.Table, c *schema.Column) {
+	if c.Default != nil {
+		t.Attrs = removeAttr(t.Attrs, reflect.TypeOf(&mysql.AutoIncrement{}))
+	} else {
+		c.AddAttrs(&mysql.AutoIncrement{})
+	}
+}
+
+func (d *MySQL) atIncrementT(t *schema.Table, v int64) {
+	if v >= 0 {
+		t.AddAttrs(mysql.NewAutoIncrement(uint64(v)))
+	}
+}
+
+func (d *MySQL) atImplicitIndexName(idx *Index, c1 *Column) bool {
+	if idx.Name == c1.Name {
+		return true
+	}
+	if !strings.HasPrefix(idx.Name, c1.Name+"_") {
+		return false
+	}
+	i, err := strconv.ParseInt(strings.TrimLeft(idx.Name, c1.Name+"_"), 10, 64)
+	return err == nil && i > 1
+}
+
+func (d *MySQL) atIndex(idx1 *Index, t2 *schema.Table, idx2 *schema.Index) error {
+	prefix := indexParts(idx1)
+	for _, c1 := range idx1.Columns {
+		c2, ok := t2.Column(c1.Name)
+		if !ok {
+			return fmt.Errorf("unexpected index %q column: %q", idx1.Name, c1.Name)
+		}
+		part := &schema.IndexPart{C: c2}
+		if v, ok := prefix[c1.Name]; ok {
+			part.AddAttrs(&mysql.SubPart{Len: int(v)})
+		}
+		idx2.AddParts(part)
+	}
+	if t, ok := indexType(idx1, dialect.MySQL); ok {
+		idx2.AddAttrs(&mysql.IndexType{T: t})
+	}
+	return nil
+}
+
+func (*MySQL) atTypeRangeSQL(ts ...string) string {
+	for i := range ts {
+		ts[i] = fmt.Sprintf("('%s')", ts[i])
+	}
+	return fmt.Sprintf("INSERT INTO `%s` (`type`) VALUES %s", TypeTable, strings.Join(ts, ", "))
+}
+
+// mariadb reports if the migration runs on MariaDB and returns the semver string.
+func (d *MySQL) mariadb() (string, bool) {
+	idx := strings.Index(d.version, "MariaDB")
+	if idx == -1 {
+		return "", false
+	}
+	return d.version[:idx-1], true
+}
+
+// defaultSize returns the default size for MySQL/MariaDB varchar type
+// based on column size, charset and table indexes, in order to avoid
+// index prefix key limit (767) for older versions of MySQL/MariaDB.
+func (d *MySQL) defaultSize(c *Column) int64 {
+	size := DefaultStringLen
+	version, checked := d.version, "5.7.0"
+	if v, ok := d.mariadb(); ok {
+		version, checked = v, "10.2.2"
+	}
+	switch {
+	// Version is >= 5.7 for MySQL, or >= 10.2.2 for MariaDB.
+	case compareVersions(version, checked) != -1:
+	// Column is non-unique, or not part of any index (reaching
+	// the error 1071).
+	case !c.Unique && len(c.indexes) == 0 && !c.PrimaryKey():
+	default:
+		size = 191
+	}
+	return size
+}
+
+// indexParts returns a map holding the sub_part mapping if exists.
+func indexParts(idx *Index) map[string]uint {
+	parts := make(map[string]uint)
+	if idx.Annotation == nil {
+		return parts
+	}
+	// If prefix (without a name) was defined on the
+	// annotation, map it to the single column index.
+	if idx.Annotation.Prefix > 0 && len(idx.Columns) == 1 {
+		parts[idx.Columns[0].Name] = idx.Annotation.Prefix
+	}
+	for column, part := range idx.Annotation.PrefixColumns {
+		parts[column] = part
+	}
+	return parts
+}

dialect/sql/schema/postgres.go

@@ -0,0 +1,260 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"strconv"
+	"strings"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/schema/field"
+
+	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/postgres"
+	"ariga.io/atlas/sql/schema"
+)
+
+// Postgres adapter for Atlas migration engine.
+type Postgres struct {
+	dialect.Driver
+	schema  string
+	version string
+}
+
+// init loads the Postgres version from the database for later use in the migration process.
+// It returns an error if the server version is lower than v10.
+func (d *Postgres) init(ctx context.Context) error {
+	if d.version != "" {
+		return nil // already initialized.
+	}
+	rows := &sql.Rows{}
+	if err := d.Query(ctx, "SHOW server_version_num", []any{}, rows); err != nil {
+		return fmt.Errorf("querying server version %w", err)
+	}
+	defer rows.Close()
+	if !rows.Next() {
+		if err := rows.Err(); err != nil {
+			return err
+		}
+		return fmt.Errorf("server_version_num variable was not found")
+	}
+	var version string
+	if err := rows.Scan(&version); err != nil {
+		return fmt.Errorf("scanning version: %w", err)
+	}
+	if len(version) < 6 {
+		return fmt.Errorf("malformed version: %s", version)
+	}
+	d.version = fmt.Sprintf("%s.%s.%s", version[:2], version[2:4], version[4:])
+	if compareVersions(d.version, "10.0.0") == -1 {
+		return fmt.Errorf("unsupported postgres version: %s", d.version)
+	}
+	return nil
+}
+
+// tableExist checks if a table exists in the database and current schema.
+func (d *Postgres) tableExist(ctx context.Context, conn dialect.ExecQuerier, name string) (bool, error) {
+	query, args := sql.Dialect(dialect.Postgres).
+		Select(sql.Count("*")).From(sql.Table("tables").Schema("information_schema")).
+		Where(sql.And(
+			d.matchSchema(),
+			sql.EQ("table_name", name),
+		)).Query()
+	return exist(ctx, conn, query, args...)
+}
+
+// matchSchema returns the predicate for matching table schema.
+func (d *Postgres) matchSchema(columns ...string) *sql.Predicate {
+	column := "table_schema"
+	if len(columns) > 0 {
+		column = columns[0]
+	}
+	if d.schema != "" {
+		return sql.EQ(column, d.schema)
+	}
+	return sql.EQ(column, sql.Raw("CURRENT_SCHEMA()"))
+}
+
+// maxCharSize defines the maximum size of limited character types in Postgres (10 MB).
+const maxCharSize = 10 << 20
+
+func (d *Postgres) atOpen(conn dialect.ExecQuerier) (migrate.Driver, error) {
+	return postgres.Open(&db{ExecQuerier: conn})
+}
+
+func (d *Postgres) atTable(t1 *Table, t2 *schema.Table) {
+	if t1.Annotation != nil {
+		setAtChecks(t1, t2)
+	}
+}
+
+func (d *Postgres) supportsDefault(*Column) bool {
+	// PostgreSQL supports default values for all standard types.
+	return true
+}
+
+func (d *Postgres) atTypeC(c1 *Column, c2 *schema.Column) error {
+	if c1.SchemaType != nil && c1.SchemaType[dialect.Postgres] != "" {
+		t, err := postgres.ParseType(strings.ToLower(c1.SchemaType[dialect.Postgres]))
+		if err != nil {
+			return err
+		}
+		c2.Type.Type = t
+		if s, ok := t.(*postgres.SerialType); c1.foreign != nil && ok {
+			c2.Type.Type = s.IntegerType()
+		}
+		return nil
+	}
+	var t schema.Type
+	switch c1.Type {
+	case field.TypeBlob:
+		return fmt.Errorf("blob fields are not stored in the database")
+	case field.TypeBool:
+		t = &schema.BoolType{T: postgres.TypeBoolean}
+	case field.TypeUint8, field.TypeInt8, field.TypeInt16:
+		t = &schema.IntegerType{T: postgres.TypeSmallInt}
+	case field.TypeUint16, field.TypeInt32:
+		t = &schema.IntegerType{T: postgres.TypeInt}
+	case field.TypeUint32, field.TypeInt, field.TypeUint, field.TypeInt64, field.TypeUint64:
+		t = &schema.IntegerType{T: postgres.TypeBigInt}
+	case field.TypeFloat32:
+		t = &schema.FloatType{T: c1.scanTypeOr(postgres.TypeReal)}
+	case field.TypeFloat64:
+		t = &schema.FloatType{T: c1.scanTypeOr(postgres.TypeDouble)}
+	case field.TypeBytes:
+		t = &schema.BinaryType{T: postgres.TypeBytea}
+	case field.TypeUUID:
+		t = &postgres.UUIDType{T: postgres.TypeUUID}
+	case field.TypeJSON:
+		t = &schema.JSONType{T: postgres.TypeJSONB}
+	case field.TypeString:
+		t = &schema.StringType{T: postgres.TypeVarChar}
+		if c1.Size > maxCharSize {
+			t = &schema.StringType{T: postgres.TypeText}
+		}
+	case field.TypeTime:
+		t = &schema.TimeType{T: c1.scanTypeOr(postgres.TypeTimestampWTZ)}
+	case field.TypeEnum:
+		// Although atlas supports enum types, we keep backwards compatibility
+		// with previous versions of ent and use varchar (see cType).
+		t = &schema.StringType{T: postgres.TypeVarChar}
+	case field.TypeOther:
+		t = &schema.UnsupportedType{T: c1.typ}
+	default:
+		t, err := postgres.ParseType(strings.ToLower(c1.typ))
+		if err != nil {
+			return err
+		}
+		c2.Type.Type = t
+	}
+	c2.Type.Type = t
+	return nil
+}
+
+func (d *Postgres) atUniqueC(t1 *Table, c1 *Column, t2 *schema.Table, c2 *schema.Column) {
+	// For UNIQUE columns, PostgreSQL creates an implicit index named
+	// "<table>_<column>_key<i>".
+	for _, idx := range t1.Indexes {
+		// Index also defined explicitly, and will be added in atIndexes.
+		if idx.Unique && d.atImplicitIndexName(idx, t1, c1) {
+			return
+		}
+	}
+	t2.AddIndexes(schema.NewUniqueIndex(fmt.Sprintf("%s_%s_key", t1.Name, c1.Name)).AddColumns(c2))
+}
+
+func (d *Postgres) atImplicitIndexName(idx *Index, t1 *Table, c1 *Column) bool {
+	p := fmt.Sprintf("%s_%s_key", t1.Name, c1.Name)
+	if idx.Name == p {
+		return true
+	}
+	i, err := strconv.ParseInt(strings.TrimPrefix(idx.Name, p), 10, 64)
+	return err == nil && i > 0
+}
+
+func (d *Postgres) atIncrementC(t *schema.Table, c *schema.Column) {
+	// Skip marking this column as an identity in case it is
+	// serial type or a default was already defined for it.
+	if _, ok := c.Type.Type.(*postgres.SerialType); ok || c.Default != nil {
+		t.Attrs = removeAttr(t.Attrs, reflect.TypeOf(&postgres.Identity{}))
+		return
+	}
+	id := &postgres.Identity{}
+	for _, a := range t.Attrs {
+		if a, ok := a.(*postgres.Identity); ok {
+			id = a
+		}
+	}
+	c.AddAttrs(id)
+}
+
+func (d *Postgres) atIncrementT(t *schema.Table, v int64) {
+	t.AddAttrs(&postgres.Identity{Sequence: &postgres.Sequence{Start: v}})
+}
+
+// indexOpClass returns a map holding the operator-class mapping if exists.
+func indexOpClass(idx *Index) map[string]string {
+	opc := make(map[string]string)
+	if idx.Annotation == nil {
+		return opc
+	}
+	// If operator-class (without a name) was defined on
+	// the annotation, map it to the single column index.
+	if idx.Annotation.OpClass != "" && len(idx.Columns) == 1 {
+		opc[idx.Columns[0].Name] = idx.Annotation.OpClass
+	}
+	for column, op := range idx.Annotation.OpClassColumns {
+		opc[column] = op
+	}
+	return opc
+}
+
+func (d *Postgres) atIndex(idx1 *Index, t2 *schema.Table, idx2 *schema.Index) error {
+	opc := indexOpClass(idx1)
+	for _, c1 := range idx1.Columns {
+		c2, ok := t2.Column(c1.Name)
+		if !ok {
+			return fmt.Errorf("unexpected index %q column: %q", idx1.Name, c1.Name)
+		}
+		part := &schema.IndexPart{C: c2}
+		if v, ok := opc[c1.Name]; ok {
+			var op postgres.IndexOpClass
+			if err := op.UnmarshalText([]byte(v)); err != nil {
+				return fmt.Errorf("unmarshalling operator-class %q for column %q: %v", v, c1.Name, err)
+			}
+			part.Attrs = append(part.Attrs, &op)
+		}
+		idx2.AddParts(part)
+	}
+	if t, ok := indexType(idx1, dialect.Postgres); ok {
+		idx2.AddAttrs(&postgres.IndexType{T: t})
+	}
+	if ant, supportsInclude := idx1.Annotation, compareVersions(d.version, "11.0.0") >= 0; ant != nil && len(ant.IncludeColumns) > 0 && supportsInclude {
+		columns := make([]*schema.Column, len(ant.IncludeColumns))
+		for i, ic := range ant.IncludeColumns {
+			c, ok := t2.Column(ic)
+			if !ok {
+				return fmt.Errorf("include column %q was not found for index %q", ic, idx1.Name)
+			}
+			columns[i] = c
+		}
+		idx2.AddAttrs(&postgres.IndexInclude{Columns: columns})
+	}
+	if idx1.Annotation != nil && idx1.Annotation.Where != "" {
+		idx2.AddAttrs(&postgres.IndexPredicate{P: idx1.Annotation.Where})
+	}
+	return nil
+}
+
+func (*Postgres) atTypeRangeSQL(ts ...string) string {
+	for i := range ts {
+		ts[i] = fmt.Sprintf("('%s')", ts[i])
+	}
+	return fmt.Sprintf(`INSERT INTO "%s" ("type") VALUES %s`, TypeTable, strings.Join(ts, ", "))
+}

dialect/sql/schema/schema.go

@@ -0,0 +1,695 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+// Package schema contains all schema migration logic for SQL dialects.
+package schema
+
+import (
+	"context"
+	"fmt"
+	"slices"
+	"strconv"
+	"strings"
+	"sync"
+
+	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/mysql"
+	"ariga.io/atlas/sql/postgres"
+	"ariga.io/atlas/sql/schema"
+	"ariga.io/atlas/sql/sqlite"
+	entdialect "entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/entsql"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/schema/field"
+)
+
+const (
+	// DefaultStringLen describes the default length for string/varchar types.
+	DefaultStringLen int64 = 255
+	// Null is the string representation of NULL in SQL.
+	Null = "NULL"
+	// PrimaryKey is the string representation of PKs in SQL.
+	PrimaryKey = "PRI"
+	// UniqueKey is the string representation of PKs in SQL.
+	UniqueKey = "UNI"
+)
+
+// Table schema definition for SQL dialects.
+type Table struct {
+	mu          sync.Mutex
+	Name        string
+	Schema      string
+	Columns     []*Column
+	columns     map[string]*Column
+	Indexes     []*Index
+	PrimaryKey  []*Column
+	ForeignKeys []*ForeignKey
+	Annotation  *entsql.Annotation
+	Comment     string
+	View        bool   // Indicate the table is a view.
+	Pos         string // filename:line of the ent schema definition.
+}
+
+// NewTable returns a new table with the given name.
+func NewTable(name string) *Table {
+	return &Table{
+		Name:    name,
+		columns: make(map[string]*Column),
+	}
+}
+
+// NewView returns a new view with the given name.
+func NewView(name string) *Table {
+	t := NewTable(name)
+	t.View = true
+	return t
+}
+
+// SetComment sets the table comment.
+func (t *Table) SetComment(c string) *Table {
+	t.Comment = c
+	return t
+}
+
+// SetSchema sets the table schema.
+func (t *Table) SetSchema(s string) *Table {
+	t.Schema = s
+	return t
+}
+
+// SetPos sets the table position.
+func (t *Table) SetPos(p string) *Table {
+	t.Pos = p
+	return t
+}
+
+// AddPrimary adds a new primary key to the table.
+func (t *Table) AddPrimary(c *Column) *Table {
+	c.Key = PrimaryKey
+	t.AddColumn(c)
+	t.PrimaryKey = append(t.PrimaryKey, c)
+	return t
+}
+
+// AddForeignKey adds a foreign key to the table.
+func (t *Table) AddForeignKey(fk *ForeignKey) *Table {
+	t.ForeignKeys = append(t.ForeignKeys, fk)
+	return t
+}
+
+// AddColumn adds a new column to the table.
+func (t *Table) AddColumn(c *Column) *Table {
+	t.columns[c.Name] = c
+	t.Columns = append(t.Columns, c)
+	return t
+}
+
+// HasColumn reports if the table contains a column with the given name.
+func (t *Table) HasColumn(name string) bool {
+	_, ok := t.Column(name)
+	return ok
+}
+
+// Column returns the column with the given name. If exists.
+func (t *Table) Column(name string) (*Column, bool) {
+	if c, ok := t.columns[name]; ok {
+		return c, true
+	}
+	// In case the column was added
+	// directly to the Columns field.
+	for _, c := range t.Columns {
+		if c.Name == name {
+			return c, true
+		}
+	}
+	return nil, false
+}
+
+// SetAnnotation the entsql.Annotation on the table.
+func (t *Table) SetAnnotation(ant *entsql.Annotation) *Table {
+	t.Annotation = ant
+	return t
+}
+
+// AddIndex creates and adds a new index to the table from the given options.
+func (t *Table) AddIndex(name string, unique bool, columns []string) *Table {
+	return t.addIndex(&Index{
+		Name:    name,
+		Unique:  unique,
+		columns: columns,
+		Columns: make([]*Column, 0, len(columns)),
+	})
+}
+
+// AddIndex creates and adds a new index to the table from the given options.
+func (t *Table) addIndex(idx *Index) *Table {
+	for _, name := range idx.columns {
+		c, ok := t.columns[name]
+		if ok {
+			c.indexes.append(idx)
+			idx.Columns = append(idx.Columns, c)
+		}
+	}
+	t.Indexes = append(t.Indexes, idx)
+	return t
+}
+
+// column returns a table column by its name.
+// faster than map lookup for most cases.
+func (t *Table) column(name string) (*Column, bool) {
+	for _, c := range t.Columns {
+		if c.Name == name {
+			return c, true
+		}
+	}
+	return nil, false
+}
+
+// Index returns a table index by its exact name.
+func (t *Table) Index(name string) (*Index, bool) {
+	idx, ok := t.index(name)
+	if ok && idx.Name == name {
+		return idx, ok
+	}
+	return nil, false
+}
+
+// index returns a table index by its name.
+func (t *Table) index(name string) (*Index, bool) {
+	for _, idx := range t.Indexes {
+		if name == idx.Name || name == idx.realname {
+			return idx, true
+		}
+		// Same as below, there are cases where the index name
+		// is unknown (created automatically on column constraint).
+		if len(idx.Columns) == 1 && idx.Columns[0].Name == name {
+			return idx, true
+		}
+	}
+	// If it is an "implicit index" (unique constraint on
+	// table creation) and it wasn't loaded in table scanning.
+	c, ok := t.column(name)
+	if !ok {
+		// Postgres naming convention for unique constraint (<table>_<column>_key).
+		name = strings.TrimPrefix(name, t.Name+"_")
+		name = strings.TrimSuffix(name, "_key")
+		c, ok = t.column(name)
+	}
+	if ok && c.Unique {
+		return &Index{Name: name, Unique: c.Unique, Columns: []*Column{c}, columns: []string{c.Name}}, true
+	}
+	return nil, false
+}
+
+// CopyTables returns a deep-copy of the given tables. This utility function is
+// useful for copying the generated schema tables (i.e. migrate.Tables) before
+// running schema migration when there is a need for execute multiple migrations
+// concurrently. e.g. running parallel unit-tests using the generated enttest package.
+func CopyTables(tables []*Table) ([]*Table, error) {
+	var (
+		copyT  = make([]*Table, len(tables))
+		byName = make(map[string]*Table)
+	)
+	for i, t := range tables {
+		copyT[i] = &Table{
+			Name:        t.Name,
+			Columns:     make([]*Column, len(t.Columns)),
+			Indexes:     make([]*Index, len(t.Indexes)),
+			ForeignKeys: make([]*ForeignKey, len(t.ForeignKeys)),
+		}
+		for j, c := range t.Columns {
+			cc := *c
+			// SchemaType and Enums are read-only fields.
+			cc.indexes = nil
+			cc.foreign = nil
+			copyT[i].Columns[j] = &cc
+		}
+		if at := t.Annotation; at != nil {
+			cat := *at
+			copyT[i].Annotation = &cat
+		}
+		byName[t.Name] = copyT[i]
+	}
+	for i, t := range tables {
+		ct := copyT[i]
+		for _, c := range t.PrimaryKey {
+			cc, ok := ct.column(c.Name)
+			if !ok {
+				return nil, fmt.Errorf("sql/schema: missing primary key column %q", c.Name)
+			}
+			ct.PrimaryKey = append(ct.PrimaryKey, cc)
+		}
+		for j, idx := range t.Indexes {
+			cidx := &Index{
+				Name:    idx.Name,
+				Unique:  idx.Unique,
+				Columns: make([]*Column, len(idx.Columns)),
+			}
+			if at := idx.Annotation; at != nil {
+				cat := *at
+				cidx.Annotation = &cat
+			}
+			for k, c := range idx.Columns {
+				cc, ok := ct.column(c.Name)
+				if !ok {
+					return nil, fmt.Errorf("sql/schema: missing index column %q", c.Name)
+				}
+				cidx.Columns[k] = cc
+			}
+			ct.Indexes[j] = cidx
+		}
+		for j, fk := range t.ForeignKeys {
+			cfk := &ForeignKey{
+				Symbol:     fk.Symbol,
+				OnUpdate:   fk.OnUpdate,
+				OnDelete:   fk.OnDelete,
+				Columns:    make([]*Column, len(fk.Columns)),
+				RefColumns: make([]*Column, len(fk.RefColumns)),
+			}
+			for k, c := range fk.Columns {
+				cc, ok := ct.column(c.Name)
+				if !ok {
+					return nil, fmt.Errorf("sql/schema: missing foreign-key column %q", c.Name)
+				}
+				cfk.Columns[k] = cc
+			}
+			cref, ok := byName[fk.RefTable.Name]
+			if !ok {
+				return nil, fmt.Errorf("sql/schema: missing foreign-key ref-table %q", fk.RefTable.Name)
+			}
+			cfk.RefTable = cref
+			for k, c := range fk.RefColumns {
+				cc, ok := cref.column(c.Name)
+				if !ok {
+					return nil, fmt.Errorf("sql/schema: missing foreign-key ref-column %q", c.Name)
+				}
+				cfk.RefColumns[k] = cc
+			}
+			ct.ForeignKeys[j] = cfk
+		}
+	}
+	return copyT, nil
+}
+
+// Column schema definition for SQL dialects.
+type Column struct {
+	Name       string            // column name.
+	Type       field.Type        // column type.
+	SchemaType map[string]string // optional schema type per dialect.
+	Attr       string            // extra attributes.
+	Size       int64             // max size parameter for string, blob, etc.
+	Key        string            // key definition (PRI, UNI or MUL).
+	Unique     bool              // column with unique constraint.
+	Increment  bool              // auto increment attribute.
+	Nullable   bool              // null or not null attribute.
+	Default    any               // default value.
+	Enums      []string          // enum values.
+	Collation  string            // collation type (utf8mb4_unicode_ci, utf8mb4_general_ci)
+	typ        string            // row column type (used for Rows.Scan).
+	indexes    Indexes           // linked indexes.
+	foreign    *ForeignKey       // linked foreign-key.
+	Comment    string            // optional column comment.
+}
+
+// Expr represents a raw expression. It is used to distinguish between
+// literal values and raw expressions when defining default values.
+type Expr string
+
+// UniqueKey returns boolean indicates if this column is a unique key.
+// Used by the migration tool when parsing the `DESCRIBE TABLE` output Go objects.
+func (c *Column) UniqueKey() bool { return c.Key == UniqueKey }
+
+// PrimaryKey returns boolean indicates if this column is on of the primary key columns.
+// Used by the migration tool when parsing the `DESCRIBE TABLE` output Go objects.
+func (c *Column) PrimaryKey() bool { return c.Key == PrimaryKey }
+
+// ConvertibleTo reports whether a column can be converted to the new column without altering its data.
+func (c *Column) ConvertibleTo(d *Column) bool {
+	switch {
+	case c.Type == d.Type:
+		if c.Size != 0 && d.Size != 0 {
+			// Types match and have a size constraint.
+			return c.Size <= d.Size
+		}
+		return true
+	case c.IntType() && d.IntType() || c.UintType() && d.UintType():
+		return c.Type <= d.Type
+	case c.UintType() && d.IntType():
+		// uintX can not be converted to intY, when X > Y.
+		return c.Type-field.TypeUint8 <= d.Type-field.TypeInt8
+	case c.Type == field.TypeString && d.Type == field.TypeEnum ||
+		c.Type == field.TypeEnum && d.Type == field.TypeString:
+		return true
+	case c.Type.Integer() && d.Type == field.TypeString:
+		return true
+	}
+	return c.FloatType() && d.FloatType()
+}
+
+// IntType reports whether the column is an int type (int8 ... int64).
+func (c Column) IntType() bool { return c.Type >= field.TypeInt8 && c.Type <= field.TypeInt64 }
+
+// UintType reports of the given type is a uint type (int8 ... int64).
+func (c Column) UintType() bool { return c.Type >= field.TypeUint8 && c.Type <= field.TypeUint64 }
+
+// FloatType reports of the given type is a float type (float32, float64).
+func (c Column) FloatType() bool { return c.Type == field.TypeFloat32 || c.Type == field.TypeFloat64 }
+
+// ScanDefault scans the default value string to its interface type.
+func (c *Column) ScanDefault(value string) error {
+	switch {
+	case strings.ToUpper(value) == Null: // ignore.
+	case c.IntType():
+		v := &sql.NullInt64{}
+		if err := v.Scan(value); err != nil {
+			return fmt.Errorf("scanning int value for column %q: %w", c.Name, err)
+		}
+		c.Default = v.Int64
+	case c.UintType():
+		v := &sql.NullInt64{}
+		if err := v.Scan(value); err != nil {
+			return fmt.Errorf("scanning uint value for column %q: %w", c.Name, err)
+		}
+		c.Default = uint64(v.Int64)
+	case c.FloatType():
+		v := &sql.NullFloat64{}
+		if err := v.Scan(value); err != nil {
+			return fmt.Errorf("scanning float value for column %q: %w", c.Name, err)
+		}
+		c.Default = v.Float64
+	case c.Type == field.TypeBool:
+		v := &sql.NullBool{}
+		if err := v.Scan(value); err != nil {
+			return fmt.Errorf("scanning bool value for column %q: %w", c.Name, err)
+		}
+		c.Default = v.Bool
+	case c.Type == field.TypeString || c.Type == field.TypeEnum:
+		v := &sql.NullString{}
+		if err := v.Scan(value); err != nil {
+			return fmt.Errorf("scanning string value for column %q: %w", c.Name, err)
+		}
+		c.Default = v.String
+	case c.Type == field.TypeJSON:
+		v := &sql.NullString{}
+		if err := v.Scan(value); err != nil {
+			return fmt.Errorf("scanning json value for column %q: %w", c.Name, err)
+		}
+		c.Default = v.String
+	case c.Type == field.TypeBytes:
+		c.Default = []byte(value)
+	case c.Type == field.TypeUUID:
+		// skip function
+		if !strings.Contains(value, "()") {
+			c.Default = value
+		}
+	default:
+		return fmt.Errorf("unsupported default type: %v default to %q", c.Type, value)
+	}
+	return nil
+}
+
+// supportDefault reports if the column type supports default value.
+func (c Column) supportDefault() bool {
+	switch t := c.Type; t {
+	case field.TypeString, field.TypeEnum:
+		return c.Size < 1<<16 // not a text.
+	case field.TypeBool, field.TypeTime, field.TypeUUID:
+		return true
+	default:
+		return t.Numeric()
+	}
+}
+
+// scanTypeOr returns the scanning type or the given value.
+func (c *Column) scanTypeOr(t string) string {
+	if c.typ != "" {
+		return strings.ToLower(c.typ)
+	}
+	return t
+}
+
+// ForeignKey definition for creation.
+type ForeignKey struct {
+	Symbol     string          // foreign-key name. Generated if empty.
+	Columns    []*Column       // table column
+	RefTable   *Table          // referenced table.
+	RefColumns []*Column       // referenced columns.
+	OnUpdate   ReferenceOption // action on update.
+	OnDelete   ReferenceOption // action on delete.
+}
+
+// ReferenceOption for constraint actions.
+type ReferenceOption string
+
+// Reference options.
+const (
+	NoAction   ReferenceOption = "NO ACTION"
+	Restrict   ReferenceOption = "RESTRICT"
+	Cascade    ReferenceOption = "CASCADE"
+	SetNull    ReferenceOption = "SET NULL"
+	SetDefault ReferenceOption = "SET DEFAULT"
+)
+
+// ConstName returns the constant name of a reference option. It's used by entc for printing the constant name in templates.
+func (r ReferenceOption) ConstName() string {
+	return strings.ReplaceAll(strings.Title(strings.ToLower(string(r))), " ", "")
+}
+
+// Index definition for table index.
+type Index struct {
+	Name       string                  // index name.
+	Unique     bool                    // uniqueness.
+	Columns    []*Column               // actual table columns.
+	Annotation *entsql.IndexAnnotation // index annotation.
+	columns    []string                // columns loaded from query scan.
+	realname   string                  // real name in the database (Postgres only).
+}
+
+// Indexes used for scanning all sql.Rows into a list of indexes, because
+// multiple sql rows can represent the same index (multi-columns indexes).
+type Indexes []*Index
+
+// append wraps the basic `append` function by filtering duplicates indexes.
+func (i *Indexes) append(idx1 *Index) {
+	for _, idx2 := range *i {
+		if idx2.Name == idx1.Name {
+			return
+		}
+	}
+	*i = append(*i, idx1)
+}
+
+// compareVersions returns an integer comparing the 2 versions.
+func compareVersions(v1, v2 string) int {
+	pv1, ok1 := parseVersion(v1)
+	pv2, ok2 := parseVersion(v2)
+	if !ok1 && !ok2 {
+		return 0
+	}
+	if !ok1 {
+		return -1
+	}
+	if !ok2 {
+		return 1
+	}
+	if v := compare(pv1.major, pv2.major); v != 0 {
+		return v
+	}
+	if v := compare(pv1.minor, pv2.minor); v != 0 {
+		return v
+	}
+	return compare(pv1.patch, pv2.patch)
+}
+
+// version represents a parsed MySQL version.
+type version struct {
+	major int
+	minor int
+	patch int
+}
+
+// parseVersion returns an integer comparing the 2 versions.
+func parseVersion(v string) (*version, bool) {
+	parts := strings.Split(v, ".")
+	if len(parts) == 0 {
+		return nil, false
+	}
+	var (
+		err error
+		ver = &version{}
+	)
+	for i, e := range []*int{&ver.major, &ver.minor, &ver.patch} {
+		if i == len(parts) {
+			break
+		}
+		if *e, err = strconv.Atoi(strings.Split(parts[i], "-")[0]); err != nil {
+			return nil, false
+		}
+	}
+	return ver, true
+}
+
+func compare(v1, v2 int) int {
+	if v1 == v2 {
+		return 0
+	}
+	if v1 < v2 {
+		return -1
+	}
+	return 1
+}
+
+func indexType(idx *Index, d string) (string, bool) {
+	ant := idx.Annotation
+	if ant == nil {
+		return "", false
+	}
+	if ant.Types != nil && ant.Types[d] != "" {
+		return ant.Types[d], true
+	}
+	if ant.Type != "" {
+		return ant.Type, true
+	}
+	return "", false
+}
+
+type driver struct {
+	sqlDialect
+	schema.Differ
+	migrate.PlanApplier
+}
+
+var drivers = func(v string) map[string]driver {
+	return map[string]driver{
+		entdialect.SQLite: {
+			&SQLite{
+				WithForeignKeys: true,
+				Driver:          nopDriver{dialect: entdialect.SQLite},
+			},
+			sqlite.DefaultDiff,
+			sqlite.DefaultPlan,
+		},
+		entdialect.MySQL: {
+			&MySQL{
+				version: v,
+				Driver:  nopDriver{dialect: entdialect.MySQL},
+			},
+			mysql.DefaultDiff,
+			mysql.DefaultPlan,
+		},
+		entdialect.Postgres: {
+			&Postgres{
+				version: v,
+				Driver:  nopDriver{dialect: entdialect.Postgres},
+			},
+			postgres.DefaultDiff,
+			postgres.DefaultPlan,
+		},
+	}
+}
+
+type DDLArgs struct {
+	// Dialect and Version of the target database.
+	Dialect, Version string
+	// HashSymbols indicates whether to hash long symbols in the DDL.
+	HashSymbols bool
+	// Tables to dump.
+	Tables []*Table
+	// Options to pass to the migration plan engine.
+	Options []migrate.PlanOption
+}
+
+// Dump the schema DDL for the given tables.
+//
+// Deprecated: use DDL instead.
+func Dump(ctx context.Context, dialect, version string, tables []*Table, opts ...migrate.PlanOption) (string, error) {
+	return DDL(ctx, DDLArgs{
+		Dialect: dialect,
+		Version: version,
+		Tables:  tables,
+		Options: opts,
+	})
+}
+
+// DDL the schema DDL for the given tables.
+func DDL(ctx context.Context, args DDLArgs) (string, error) {
+	args.Options = append([]migrate.PlanOption{func(o *migrate.PlanOptions) {
+		o.Mode = migrate.PlanModeDump
+		o.Indent = "  "
+	}}, args.Options...)
+	d, ok := drivers(args.Version)[args.Dialect]
+	if !ok {
+		return "", fmt.Errorf("unsupported dialect %q", args.Dialect)
+	}
+	a := &Atlas{
+		sqlDialect:  d,
+		dialect:     args.Dialect,
+		hashSymbols: args.HashSymbols,
+	}
+	r, err := a.StateReader(args.Tables...).ReadState(ctx)
+	if err != nil {
+		return "", err
+	}
+	// Since the Atlas version bundled with Ent does not support view management,
+	// simply spit out the definition instead of letting Atlas plan them.
+	var vs []*schema.View
+	for _, s := range r.Schemas {
+		vs = append(vs, s.Views...)
+		s.Views = nil
+	}
+	var c schema.Changes
+	if slices.ContainsFunc(args.Tables, func(t *Table) bool { return t.Schema != "" }) {
+		c, err = d.RealmDiff(&schema.Realm{}, r)
+	} else {
+		c, err = d.SchemaDiff(&schema.Schema{}, r.Schemas[0])
+	}
+	if err != nil {
+		return "", err
+	}
+	p, err := d.PlanChanges(ctx, "dump", c, args.Options...)
+	if err != nil {
+		return "", err
+	}
+	for _, v := range vs {
+		q, _ := sql.Dialect(args.Dialect).
+			CreateView(v.Name).
+			Schema(v.Schema.Name).
+			Columns(func(cols []*schema.Column) (bs []*sql.ColumnBuilder) {
+				for _, c := range cols {
+					bs = append(bs, sql.Dialect(args.Dialect).Column(c.Name).Type(c.Type.Raw))
+				}
+				return
+			}(v.Columns)...).
+			As(sql.Raw(v.Def)).
+			Query()
+		p.Changes = append(p.Changes, &migrate.Change{
+			Cmd:     q,
+			Comment: fmt.Sprintf("Add %q view", v.Name),
+		})
+	}
+	for _, t := range args.Tables {
+		p.Directives = append(p.Directives, fmt.Sprintf(
+			"-- atlas:pos %s%s[type=%s] %s",
+			func() string {
+				if t.Schema != "" {
+					return t.Schema + "[type=schema]."
+				}
+				return ""
+			}(),
+			t.Name,
+			func() string {
+				if t.View {
+					return "view"
+				}
+				return "table"
+			}(),
+			t.Pos,
+		))
+	}
+	f, err := migrate.DefaultFormatter.FormatFile(p)
+	if err != nil {
+		return "", err
+	}
+	return string(f.Bytes()), nil
+}

dialect/sql/schema/schema_test.go

@@ -0,0 +1,390 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"testing"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/entsql"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/schema/field"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestColumn_ConvertibleTo(t *testing.T) {
+	c1 := &Column{Type: field.TypeString, Size: 10}
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeString, Size: 10}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeString, Size: 255}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeString, Size: 9}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeFloat32}))
+
+	c1 = &Column{Type: field.TypeFloat32}
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeFloat32}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeFloat64}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeString}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint}))
+
+	c1 = &Column{Type: field.TypeFloat64}
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeFloat32}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeFloat64}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeString}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint}))
+
+	c1 = &Column{Type: field.TypeUint}
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeUint}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeInt}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeInt64}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeUint64}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeInt8}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint8}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint16}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint32}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeString}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeString, Size: 1}))
+
+	c1 = &Column{Type: field.TypeInt}
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeInt}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeInt64}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeInt8}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeInt32}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint8}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint16}))
+	require.False(t, c1.ConvertibleTo(&Column{Type: field.TypeUint32}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeString}))
+	require.True(t, c1.ConvertibleTo(&Column{Type: field.TypeString, Size: 1}))
+}
+
+func TestColumn_ScanDefault(t *testing.T) {
+	c1 := &Column{Type: field.TypeString, Size: 10}
+	require.NoError(t, c1.ScanDefault("Hello World"))
+	require.Equal(t, "Hello World", c1.Default)
+	require.NoError(t, c1.ScanDefault("1"))
+	require.Equal(t, "1", c1.Default)
+
+	c1 = &Column{Type: field.TypeInt64}
+	require.NoError(t, c1.ScanDefault("128"))
+	require.Equal(t, int64(128), c1.Default)
+	require.NoError(t, c1.ScanDefault("1"))
+	require.Equal(t, int64(1), c1.Default)
+	require.Error(t, c1.ScanDefault("foo"))
+
+	c1 = &Column{Type: field.TypeUint64}
+	require.NoError(t, c1.ScanDefault("128"))
+	require.Equal(t, uint64(128), c1.Default)
+	require.NoError(t, c1.ScanDefault("1"))
+	require.Equal(t, uint64(1), c1.Default)
+	require.Error(t, c1.ScanDefault("foo"))
+
+	c1 = &Column{Type: field.TypeFloat64}
+	require.NoError(t, c1.ScanDefault("128.1"))
+	require.Equal(t, 128.1, c1.Default)
+	require.NoError(t, c1.ScanDefault("1"))
+	require.Equal(t, float64(1), c1.Default)
+	require.Error(t, c1.ScanDefault("foo"))
+
+	c1 = &Column{Type: field.TypeBool}
+	require.NoError(t, c1.ScanDefault("1"))
+	require.Equal(t, true, c1.Default)
+	require.NoError(t, c1.ScanDefault("true"))
+	require.Equal(t, true, c1.Default)
+	require.NoError(t, c1.ScanDefault("0"))
+	require.Equal(t, false, c1.Default)
+	require.NoError(t, c1.ScanDefault("false"))
+	require.Equal(t, false, c1.Default)
+	require.Error(t, c1.ScanDefault("foo"))
+
+	c1 = &Column{Type: field.TypeUUID}
+	require.NoError(t, c1.ScanDefault("gen_random_uuid()"))
+	require.Equal(t, nil, c1.Default)
+	require.NoError(t, c1.ScanDefault("00000000-0000-0000-0000-000000000000"))
+	require.Equal(t, "00000000-0000-0000-0000-000000000000", c1.Default)
+}
+
+func TestCopyTables(t *testing.T) {
+	users := &Table{
+		Name: "users",
+		Columns: []*Column{
+			{Name: "id", Type: field.TypeInt},
+			{Name: "name", Type: field.TypeString},
+			{Name: "spouse_id", Type: field.TypeInt},
+		},
+	}
+	users.PrimaryKey = users.Columns[:1]
+	users.Indexes = append(users.Indexes, &Index{
+		Name:    "name",
+		Columns: users.Columns[1:2],
+	})
+	users.AddForeignKey(&ForeignKey{
+		Columns:    users.Columns[2:],
+		RefTable:   users,
+		RefColumns: users.Columns[:1],
+		OnUpdate:   SetNull,
+	})
+	users.SetAnnotation(&entsql.Annotation{Table: "Users"})
+	pets := &Table{
+		Name: "pets",
+		Columns: []*Column{
+			{Name: "id", Type: field.TypeInt},
+			{Name: "name", Type: field.TypeString},
+			{Name: "owner_id", Type: field.TypeInt},
+		},
+	}
+	pets.Indexes = append(pets.Indexes, &Index{
+		Name:       "name",
+		Unique:     true,
+		Columns:    pets.Columns[1:2],
+		Annotation: entsql.Desc(),
+	})
+	pets.AddForeignKey(&ForeignKey{
+		Columns:    pets.Columns[2:],
+		RefTable:   users,
+		RefColumns: users.Columns[:1],
+		OnDelete:   SetDefault,
+	})
+	tables := []*Table{users, pets}
+	copyT, err := CopyTables(tables)
+	require.NoError(t, err)
+	require.Equal(t, tables, copyT)
+}
+
+func TestDDL(t *testing.T) {
+	const (
+		hash   = "249590215c5bc8be0106146e65d7fa7f"
+		idx    = "super_duper_mega_ultra_hyper_giga_colossal_unbelievably_long_index_name"
+		hashMY = "super_duper_mega_ultra_hyper_gi_249590215c5bc8be0106146e65d7fa7f"
+		hashPG = "super_duper_mega_ultra_hyper_g_249590215c5bc8be0106146e65d7fa7f"
+	)
+	tbls := func() []*Table {
+		users := &Table{
+			Name: "users",
+			Pos:  "users.go:15",
+			Columns: []*Column{
+				{Name: "id", Type: field.TypeInt},
+				{Name: "name", Type: field.TypeString},
+				{Name: "spouse_id", Type: field.TypeInt},
+			},
+		}
+		users.PrimaryKey = users.Columns[:1]
+		users.Indexes = append(users.Indexes, &Index{
+			Name:    "name",
+			Columns: users.Columns[1:2],
+		})
+		users.AddForeignKey(&ForeignKey{
+			Columns:    users.Columns[2:],
+			RefTable:   users,
+			RefColumns: users.Columns[:1],
+			OnUpdate:   SetDefault,
+		})
+		users.SetAnnotation(&entsql.Annotation{Table: "Users"})
+		pets := &Table{
+			Name: "pets",
+			Pos:  "pets.go:15",
+			Columns: []*Column{
+				{Name: "id", Type: field.TypeInt},
+				{Name: "name", Type: field.TypeString},
+				{Name: "fur_color", Type: field.TypeEnum, Enums: []string{"black", "white"}},
+				{Name: "owner_id", Type: field.TypeInt},
+			},
+		}
+		pets.Indexes = append(pets.Indexes, &Index{
+			Name:       idx,
+			Unique:     true,
+			Columns:    pets.Columns[1:2],
+			Annotation: entsql.Desc(),
+		})
+		pets.AddForeignKey(&ForeignKey{
+			Columns:    pets.Columns[3:],
+			RefTable:   users,
+			RefColumns: users.Columns[:1],
+			OnDelete:   SetDefault,
+		})
+		petsWithoutFur := &Table{
+			Name:       "pets_without_fur",
+			Pos:        "pets.go:30",
+			View:       true,
+			Columns:    append(pets.Columns[:2], pets.Columns[3]),
+			Annotation: entsql.View("SELECT id, name, owner_id FROM pets"),
+		}
+		petNames := &Table{
+			Name:    "pet_names",
+			Pos:     "pets.go:45",
+			View:    true,
+			Columns: pets.Columns[1:1],
+			Annotation: entsql.ViewFor(dialect.Postgres, func(s *sql.Selector) {
+				s.Select("name").From(sql.Table("pets"))
+			}),
+		}
+		return []*Table{users, pets, petsWithoutFur, petNames}
+	}
+	my := func(length int, idx string) string {
+		return fmt.Sprintf(strings.ReplaceAll(`-- Add new schema named "s1"
+CREATE DATABASE $s1$;
+-- Add new schema named "s2"
+CREATE DATABASE $s2$;
+-- Add new schema named "s3"
+CREATE DATABASE $s3$;
+-- Create "users" table
+CREATE TABLE $s1$.$users$ (
+  $id$ bigint NOT NULL,
+  $name$ varchar(%d) NOT NULL,
+  $spouse_id$ bigint NOT NULL,
+  PRIMARY KEY ($id$),
+  INDEX $name$ ($name$),
+  FOREIGN KEY ($spouse_id$) REFERENCES $s1$.$users$ ($id$) ON UPDATE SET DEFAULT
+) CHARSET utf8mb4 COLLATE utf8mb4_bin;
+-- Create "pets" table
+CREATE TABLE $s2$.$pets$ (
+  $id$ bigint NOT NULL,
+  $name$ varchar(%d) NOT NULL,
+  $owner_id$ bigint NOT NULL,
+  $owner_id$ bigint NOT NULL,
+  UNIQUE INDEX $%s$ ($name$ DESC),
+  FOREIGN KEY ($owner_id$) REFERENCES $s1$.$users$ ($id$) ON DELETE SET DEFAULT
+) CHARSET utf8mb4 COLLATE utf8mb4_bin;
+-- Add "pets_without_fur" view
+CREATE VIEW $s3$.$pets_without_fur$ ($id$, $name$, $owner_id$) AS SELECT id, name, owner_id FROM pets;
+`, "$", "`"), length, length, idx)
+	}
+
+	pg := func(idx string) string {
+		return fmt.Sprintf(`-- Add new schema named "s1"
+CREATE SCHEMA "s1";
+-- Add new schema named "s2"
+CREATE SCHEMA "s2";
+-- Add new schema named "s3"
+CREATE SCHEMA "s3";
+-- Create "users" table
+CREATE TABLE "s1"."users" (
+  "id" bigint NOT NULL,
+  "name" character varying NOT NULL,
+  "spouse_id" bigint NOT NULL,
+  PRIMARY KEY ("id"),
+  FOREIGN KEY ("spouse_id") REFERENCES "s1"."users" ("id") ON UPDATE SET DEFAULT
+);
+-- Create index "name" to table: "users"
+CREATE INDEX "name" ON "s1"."users" ("name");
+-- Create "pets" table
+CREATE TABLE "s2"."pets" (
+  "id" bigint NOT NULL,
+  "name" character varying NOT NULL,
+  "owner_id" bigint NOT NULL,
+  "owner_id" bigint NOT NULL,
+  FOREIGN KEY ("owner_id") REFERENCES "s1"."users" ("id") ON DELETE SET DEFAULT
+);
+-- Create index "%s" to table: "pets"
+CREATE UNIQUE INDEX "%s" ON "s2"."pets" ("name" DESC);
+-- Add "pets_without_fur" view
+CREATE VIEW "s3"."pets_without_fur" ("id", "name", "owner_id") AS SELECT id, name, owner_id FROM pets;
+-- Add "pet_names" view
+CREATE VIEW "s3"."pet_names" AS SELECT "name" FROM "pets";
+`, idx, idx)
+	}
+
+	for _, tt := range []struct {
+		dialect, version, expected string
+		hash                       bool
+	}{
+		{
+			dialect.SQLite, "",
+			fmt.Sprintf(strings.ReplaceAll(`-- Create "users" table
+CREATE TABLE $users$ (
+  $id$ integer NOT NULL,
+  $name$ text NOT NULL,
+  $spouse_id$ integer NOT NULL,
+  PRIMARY KEY ($id$),
+  FOREIGN KEY ($spouse_id$) REFERENCES $users$ ($id$) ON UPDATE SET DEFAULT
+);
+-- Create index "name" to table: "users"
+CREATE INDEX $name$ ON $users$ ($name$);
+-- Create "pets" table
+CREATE TABLE $pets$ (
+  $id$ integer NOT NULL,
+  $name$ text NOT NULL,
+  $owner_id$ integer NOT NULL,
+  $owner_id$ integer NOT NULL,
+  FOREIGN KEY ($owner_id$) REFERENCES $users$ ($id$) ON DELETE SET DEFAULT
+);
+-- Create index "%s" to table: "pets"
+CREATE UNIQUE INDEX $%s$ ON $pets$ ($name$ DESC);
+-- Add "pets_without_fur" view
+CREATE VIEW $pets_without_fur$ ($id$, $name$, $owner_id$) AS SELECT id, name, owner_id FROM pets;
+`, "$", "`"), idx, idx), false,
+		},
+		{dialect.MySQL, "5.6", my(255, idx), false},
+		{dialect.MySQL, "5.6", my(191, hashMY), true},
+		{dialect.MySQL, "5.7", my(255, idx), false},
+		{dialect.MySQL, "8", my(255, hashMY), true},
+		{dialect.Postgres, "12", pg(idx), false},
+		{dialect.Postgres, "13", pg(idx), false},
+		{dialect.Postgres, "14", pg(hashPG), true},
+		{dialect.Postgres, "15", pg(hashPG), true},
+	} {
+		n := tt.dialect
+		if tt.version != "" {
+			n += ":" + tt.version
+		}
+		pos := `-- atlas:pos users[type=table] users.go:15
+-- atlas:pos pets[type=table] pets.go:15
+-- atlas:pos pets_without_fur[type=view] pets.go:30
+-- atlas:pos pet_names[type=view] pets.go:45
+
+`
+		tables := tbls()
+		if tt.dialect != dialect.SQLite {
+			tables[0].Schema = "s1"
+			tables[1].Schema = "s2"
+			tables[2].Schema = "s3"
+			tables[3].Schema = "s3"
+			pos = `-- atlas:pos s1[type=schema].users[type=table] users.go:15
+-- atlas:pos s2[type=schema].pets[type=table] pets.go:15
+-- atlas:pos s3[type=schema].pets_without_fur[type=view] pets.go:30
+-- atlas:pos s3[type=schema].pet_names[type=view] pets.go:45
+
+`
+		}
+		t.Run(n, func(t *testing.T) {
+			ac, err := DDL(context.Background(), DDLArgs{
+				Dialect:     tt.dialect,
+				Version:     tt.version,
+				Tables:      tables,
+				HashSymbols: tt.hash,
+			})
+			require.NoError(t, err)
+			require.Equal(t, pos+tt.expected, ac)
+		})
+		t.Run(n+" single schema", func(t *testing.T) {
+			ac, err := DDL(context.Background(), DDLArgs{
+				Dialect:     tt.dialect,
+				Version:     tt.version,
+				Tables:      tables[0:1],
+				HashSymbols: tt.hash,
+			})
+			require.NoError(t, err)
+			if tt.dialect != dialect.SQLite {
+				require.Contains(t, ac, "s1[type=schema].")
+				require.NotContains(t, ac, "s2[type=schema].")
+				require.Contains(t, ac, "-- Add new schema named \"s1\"")
+			}
+		})
+		t.Run(n+" no schema", func(t *testing.T) {
+			tables[0].Schema = ""
+			ac, err := DDL(context.Background(), DDLArgs{
+				Dialect:     tt.dialect,
+				Version:     tt.version,
+				Tables:      tables[0:1],
+				HashSymbols: tt.hash,
+			})
+			require.NoError(t, err)
+			require.NotContains(t, ac, "[type=schema].")
+			require.Contains(t, ac, "[type=table]")
+		})
+	}
+}

dialect/sql/schema/sqlite.go

@@ -0,0 +1,227 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"context"
+	stdsql "database/sql"
+	"fmt"
+	"reflect"
+	"strconv"
+	"strings"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/schema/field"
+
+	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/schema"
+	"ariga.io/atlas/sql/sqlite"
+)
+
+type (
+	// SQLite adapter for Atlas migration engine.
+	SQLite struct {
+		dialect.Driver
+		WithForeignKeys bool
+	}
+	// SQLiteTx implements dialect.Tx.
+	SQLiteTx struct {
+		dialect.Tx
+		commit   func() error // Override Commit to toggle foreign keys back on after Commit.
+		rollback func() error // Override Rollback to toggle foreign keys back on after Rollback.
+	}
+)
+
+// Tx implements opens a transaction.
+func (d *SQLite) Tx(ctx context.Context) (dialect.Tx, error) {
+	db := &db{d}
+	if _, err := db.ExecContext(ctx, "PRAGMA foreign_keys = off"); err != nil {
+		return nil, fmt.Errorf("sqlite: set 'foreign_keys = off': %w", err)
+	}
+	t, err := d.Driver.Tx(ctx)
+	if err != nil {
+		return nil, err
+	}
+	tx := &tx{t}
+	cm, err := sqlite.CommitFunc(ctx, db, tx, true)
+	if err != nil {
+		return nil, err
+	}
+	return &SQLiteTx{Tx: t, commit: cm, rollback: sqlite.RollbackFunc(ctx, db, tx, true)}, nil
+}
+
+// Commit ensures foreign keys are toggled back on after commit.
+func (tx *SQLiteTx) Commit() error {
+	return tx.commit()
+}
+
+// Rollback ensures foreign keys are toggled back on after rollback.
+func (tx *SQLiteTx) Rollback() error {
+	return tx.rollback()
+}
+
+// init makes sure that foreign_keys support is enabled.
+func (d *SQLite) init(ctx context.Context) error {
+	on, err := exist(ctx, d, "PRAGMA foreign_keys")
+	if err != nil {
+		return fmt.Errorf("sqlite: check foreign_keys pragma: %w", err)
+	}
+	if !on {
+		// foreign_keys pragma is off, either enable it by execute "PRAGMA foreign_keys=ON"
+		// or add the following parameter in the connection string "_fk=1".
+		return fmt.Errorf("sqlite: foreign_keys pragma is off: missing %q in the connection string", "_fk=1")
+	}
+	return nil
+}
+
+func (d *SQLite) tableExist(ctx context.Context, conn dialect.ExecQuerier, name string) (bool, error) {
+	query, args := sql.Select().Count().
+		From(sql.Table("sqlite_master")).
+		Where(sql.And(
+			sql.EQ("type", "table"),
+			sql.EQ("name", name),
+		)).
+		Query()
+	return exist(ctx, conn, query, args...)
+}
+
+func (d *SQLite) atOpen(conn dialect.ExecQuerier) (migrate.Driver, error) {
+	return sqlite.Open(&db{ExecQuerier: conn})
+}
+
+func (d *SQLite) atTable(t1 *Table, t2 *schema.Table) {
+	if t1.Annotation != nil {
+		setAtChecks(t1, t2)
+	}
+}
+
+func (d *SQLite) supportsDefault(*Column) bool {
+	// SQLite supports default values for all standard types.
+	return true
+}
+
+func (d *SQLite) atTypeC(c1 *Column, c2 *schema.Column) error {
+	if c1.SchemaType != nil && c1.SchemaType[dialect.SQLite] != "" {
+		t, err := sqlite.ParseType(strings.ToLower(c1.SchemaType[dialect.SQLite]))
+		if err != nil {
+			return err
+		}
+		c2.Type.Type = t
+		return nil
+	}
+	var t schema.Type
+	switch c1.Type {
+	case field.TypeBlob:
+		return fmt.Errorf("blob fields are not stored in the database")
+	case field.TypeBool:
+		t = &schema.BoolType{T: "bool"}
+	case field.TypeInt8, field.TypeUint8, field.TypeInt16, field.TypeUint16, field.TypeInt32,
+		field.TypeUint32, field.TypeUint, field.TypeInt, field.TypeInt64, field.TypeUint64:
+		t = &schema.IntegerType{T: sqlite.TypeInteger}
+	case field.TypeBytes:
+		t = &schema.BinaryType{T: sqlite.TypeBlob}
+	case field.TypeString, field.TypeEnum:
+		// SQLite does not impose any length restrictions on
+		// the length of strings, BLOBs or numeric values.
+		t = &schema.StringType{T: sqlite.TypeText}
+	case field.TypeFloat32, field.TypeFloat64:
+		t = &schema.FloatType{T: sqlite.TypeReal}
+	case field.TypeTime:
+		t = &schema.TimeType{T: "datetime"}
+	case field.TypeJSON:
+		t = &schema.JSONType{T: "json"}
+	case field.TypeUUID:
+		t = &sqlite.UUIDType{T: "uuid"}
+	case field.TypeOther:
+		t = &schema.UnsupportedType{T: c1.typ}
+	default:
+		t, err := sqlite.ParseType(strings.ToLower(c1.typ))
+		if err != nil {
+			return err
+		}
+		c2.Type.Type = t
+	}
+	c2.Type.Type = t
+	return nil
+}
+
+func (d *SQLite) atUniqueC(t1 *Table, c1 *Column, t2 *schema.Table, c2 *schema.Column) {
+	// For UNIQUE columns, SQLite create an implicit index named
+	// "sqlite_autoindex_<table>_<i>". Ent uses the PostgreSQL approach
+	// in its migration, and name these indexes as "<table>_<column>_key".
+	for _, idx := range t1.Indexes {
+		// Index also defined explicitly, and will be add in atIndexes.
+		if idx.Unique && d.atImplicitIndexName(idx, t1, c1) {
+			return
+		}
+	}
+	t2.AddIndexes(schema.NewUniqueIndex(fmt.Sprintf("%s_%s_key", t2.Name, c1.Name)).AddColumns(c2))
+}
+
+func (d *SQLite) atImplicitIndexName(idx *Index, t1 *Table, c1 *Column) bool {
+	if idx.Name == c1.Name {
+		return true
+	}
+	p := fmt.Sprintf("sqlite_autoindex_%s_", t1.Name)
+	if !strings.HasPrefix(idx.Name, p) {
+		return false
+	}
+	i, err := strconv.ParseInt(strings.TrimPrefix(idx.Name, p), 10, 64)
+	return err == nil && i > 0
+}
+
+func (d *SQLite) atIncrementC(t *schema.Table, c *schema.Column) {
+	if c.Default != nil {
+		t.Attrs = removeAttr(t.Attrs, reflect.TypeOf(&sqlite.AutoIncrement{}))
+	} else {
+		c.AddAttrs(&sqlite.AutoIncrement{})
+	}
+}
+
+func (d *SQLite) atIncrementT(t *schema.Table, v int64) {
+	t.AddAttrs(&sqlite.AutoIncrement{Seq: v})
+}
+
+func (d *SQLite) atIndex(idx1 *Index, t2 *schema.Table, idx2 *schema.Index) error {
+	for _, c1 := range idx1.Columns {
+		c2, ok := t2.Column(c1.Name)
+		if !ok {
+			return fmt.Errorf("unexpected index %q column: %q", idx1.Name, c1.Name)
+		}
+		idx2.AddParts(&schema.IndexPart{C: c2})
+	}
+	if idx1.Annotation != nil && idx1.Annotation.Where != "" {
+		idx2.AddAttrs(&sqlite.IndexPredicate{P: idx1.Annotation.Where})
+	}
+	return nil
+}
+
+func (*SQLite) atTypeRangeSQL(ts ...string) string {
+	for i := range ts {
+		ts[i] = fmt.Sprintf("('%s')", ts[i])
+	}
+	return fmt.Sprintf("INSERT INTO `%s` (`type`) VALUES %s", TypeTable, strings.Join(ts, ", "))
+}
+
+type tx struct {
+	dialect.Tx
+}
+
+func (tx *tx) QueryContext(ctx context.Context, query string, args ...any) (*stdsql.Rows, error) {
+	rows := &sql.Rows{}
+	if err := tx.Query(ctx, query, args, rows); err != nil {
+		return nil, err
+	}
+	return rows.ColumnScanner.(*stdsql.Rows), nil
+}
+
+func (tx *tx) ExecContext(ctx context.Context, query string, args ...any) (stdsql.Result, error) {
+	var r stdsql.Result
+	if err := tx.Exec(ctx, query, args, &r); err != nil {
+		return nil, err
+	}
+	return r, nil
+}

dialect/sql/schema/writer.go

@@ -0,0 +1,365 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"regexp"
+	"strconv"
+	"strings"
+	"time"
+	"unicode"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+
+	"ariga.io/atlas/sql/migrate"
+)
+
+type (
+	// WriteDriver is a driver that writes all driver exec operations to its writer.
+	// Note that this driver is used only for printing or writing statements to SQL
+	// files, and may require manual changes to the generated SQL statements.
+	WriteDriver struct {
+		dialect.Driver // optional driver for query calls.
+		io.Writer      // target for exec statements.
+		FormatFunc     func(string) (string, error)
+	}
+	// DirWriter implements the io.Writer interface
+	// for writing to an Atlas managed directory.
+	DirWriter struct {
+		Dir       migrate.Dir       // target directory.
+		Formatter migrate.Formatter // optional formatter.
+		b         bytes.Buffer      // working buffer.
+		changes   []*migrate.Change // changes to flush.
+	}
+)
+
+// Write implements the io.Writer interface.
+func (d *DirWriter) Write(p []byte) (int, error) {
+	return d.b.Write(trimReturning(p))
+}
+
+// Change converts all written statement so far into a migration
+// change with the given comment.
+func (d *DirWriter) Change(comment string) {
+	// Trim semicolon and new line, because formatter adds it.
+	d.changes = append(d.changes, &migrate.Change{Comment: comment, Cmd: strings.TrimRight(d.b.String(), ";\n")})
+	d.b.Reset()
+}
+
+// Flush flushes the written statements to the directory.
+func (d *DirWriter) Flush(name string) error {
+	switch {
+	case d.b.Len() != 0:
+		return fmt.Errorf("writer has undocumented change. Use Change or FlushChange instead")
+	case len(d.changes) == 0:
+		return errors.New("writer has no changes to flush")
+	default:
+		return migrate.NewPlanner(nil, d.Dir, migrate.PlanFormat(d.Formatter)).
+			WritePlan(&migrate.Plan{
+				Name:    name,
+				Changes: d.changes,
+			})
+	}
+}
+
+// FlushChange combines Change and Flush.
+func (d *DirWriter) FlushChange(name, comment string) error {
+	d.Change(comment)
+	return d.Flush(name)
+}
+
+// NewWriteDriver creates a dialect.Driver that writes all driver exec statement to its writer.
+func NewWriteDriver(dialect string, w io.Writer) *WriteDriver {
+	return &WriteDriver{
+		Writer: w,
+		Driver: nopDriver{dialect: dialect},
+	}
+}
+
+// Exec implements the dialect.Driver.Exec method.
+func (w *WriteDriver) Exec(_ context.Context, query string, args, res any) error {
+	if rr, ok := res.(*sql.Result); ok {
+		*rr = noResult{}
+	}
+	if !strings.HasSuffix(query, ";") {
+		query += ";"
+	}
+	if args != nil {
+		args, ok := args.([]any)
+		if !ok {
+			return fmt.Errorf("unexpected args type: %T", args)
+		}
+		query = w.expandArgs(query, args)
+	}
+	_, err := io.WriteString(w, query+"\n")
+	return err
+}
+
+// Query implements the dialect.Driver.Query method.
+func (w *WriteDriver) Query(ctx context.Context, query string, args, res any) error {
+	if strings.HasPrefix(query, "INSERT") || strings.HasPrefix(query, "UPDATE") {
+		if err := w.Exec(ctx, query, args, nil); err != nil {
+			return err
+		}
+		if rr, ok := res.(*sql.Rows); ok {
+			cols := func() []string {
+				// If the query has a RETURNING clause, mock the result.
+				var clause string
+			outer:
+				for i := 0; i < len(query); i++ {
+					switch q := query[i]; {
+					case q == '\'', q == '"', q == '`': // string or identifier
+						_, skip := skipQuoted(query, i)
+						if skip == -1 {
+							return nil // malformed SQL
+						}
+						i = skip
+						continue
+					case reReturning.MatchString(query[i:]):
+						var j int
+					inner:
+						// Forward until next unquoted ';' appears, or we reach the end of the query.
+						for j = i; j < len(query); j++ {
+							switch query[j] {
+							case '\'', '"', '`': // string or identifier
+								_, skip := skipQuoted(query, j)
+								if skip == -1 {
+									return nil // malformed RETURNING clause
+								}
+								j = skip
+							case ';':
+								break inner
+							}
+						}
+						clause = query[i:j]
+						break outer
+					}
+				}
+				cols := strings.Split(reReturning.ReplaceAllString(clause, ""), ",")
+				for i := range cols {
+					cols[i] = strings.TrimSpace(cols[i])
+				}
+				return cols
+			}()
+			*rr = sql.Rows{ColumnScanner: &noRows{cols: cols}}
+		}
+		return nil
+	}
+	switch w.Driver.(type) {
+	case nil, nopDriver:
+		return errors.New("query is not supported by the WriteDriver")
+	default:
+		return w.Driver.Query(ctx, query, args, res)
+	}
+}
+
+// expandArgs combines to arguments and statement into a single statement to
+// print or write into a file (before editing).
+// Note, the output may be incorrect or unsafe SQL and require manual changes.
+func (w *WriteDriver) expandArgs(query string, args []any) string {
+	var (
+		b    strings.Builder
+		p    = w.placeholder()
+		scan = w.scanPlaceholder()
+	)
+	for i := 0; i < len(query); i++ {
+	Top:
+		switch query[i] {
+		case p:
+			idx, size := scan(query[i+1:])
+			// Unrecognized placeholder.
+			if idx < 0 || idx >= len(args) {
+				return query
+			}
+			i += size
+			v, err := w.formatArg(args[idx])
+			if err != nil {
+				// Unexpected formatting error.
+				return query
+			}
+			b.WriteString(v)
+		// String or identifier.
+		case '\'', '"', '`':
+			for j := i + 1; j < len(query); j++ {
+				switch query[j] {
+				case '\\':
+					j++
+				case query[i]:
+					b.WriteString(query[i : j+1])
+					i = j
+					break Top
+				}
+			}
+			// Unexpected EOS.
+			return query
+		default:
+			b.WriteByte(query[i])
+		}
+	}
+	return b.String()
+}
+
+func (w *WriteDriver) scanPlaceholder() func(string) (int, int) {
+	switch w.Dialect() {
+	case dialect.Postgres:
+		return func(s string) (int, int) {
+			var i int
+			for i < len(s) && unicode.IsDigit(rune(s[i])) {
+				i++
+			}
+			idx, err := strconv.ParseInt(s[:i], 10, 64)
+			if err != nil {
+				return -1, 0
+			}
+			// Placeholders are 1-based.
+			return int(idx) - 1, i
+		}
+	default:
+		idx := -1
+		return func(string) (int, int) {
+			idx++
+			return idx, 0
+		}
+	}
+}
+
+func (w *WriteDriver) placeholder() byte {
+	if w.Dialect() == dialect.Postgres {
+		return '$'
+	}
+	return '?'
+}
+
+func (w *WriteDriver) formatArg(v any) (string, error) {
+	if w.FormatFunc != nil {
+		return w.FormatFunc(fmt.Sprint(v))
+	}
+	switch v := v.(type) {
+	case nil:
+		return "NULL", nil
+	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
+		return fmt.Sprintf("%d", v), nil
+	case float32, float64:
+		return fmt.Sprintf("%g", v), nil
+	case bool:
+		if v {
+			return "1", nil
+		} else {
+			return "0", nil
+		}
+	case string:
+		return "'" + strings.ReplaceAll(v, "'", "''") + "'", nil
+	case json.RawMessage:
+		return "'" + strings.ReplaceAll(string(v), "'", "''") + "'", nil
+	case []byte:
+		return "{{ BINARY_VALUE }}", nil
+	case time.Time:
+		return "{{ TIME_VALUE }}", nil
+	case fmt.Stringer:
+		return "'" + strings.ReplaceAll(v.String(), "'", "''") + "'", nil
+	default:
+		return "{{ VALUE }}", nil
+	}
+}
+
+var reReturning = regexp.MustCompile(`(?i)^\s?RETURNING`)
+
+// trimReturning trims any RETURNING suffix from INSERT/UPDATE queries.
+// Note, that the output may be incorrect or unsafe SQL and require manual changes.
+func trimReturning(query []byte) []byte {
+	var b bytes.Buffer
+loop:
+	for i := 0; i < len(query); i++ {
+		switch q := query[i]; {
+		case q == '\'', q == '"', q == '`': // string or identifier
+			s, skip := skipQuoted(query, i)
+			if skip == -1 {
+				return query
+			}
+			b.Write(s)
+			i = skip
+			continue
+		case reReturning.Match(query[i:]):
+			// Forward until next unquoted ';' appears.
+			for j := i; j < len(query); j++ { // skip "RETURNING"
+				switch query[j] {
+				case '\'', '"', '`': // string or identifier
+					_, skip := skipQuoted(query, j)
+					if skip == -1 {
+						return query
+					}
+					j = skip
+				case ';':
+					b.WriteString(";")
+					i += j
+					continue loop
+				}
+			}
+		}
+		b.WriteByte(query[i])
+	}
+	return b.Bytes()
+}
+
+func skipQuoted[T []byte | string](query T, idx int) (T, int) {
+	for j := idx + 1; j < len(query); j++ {
+		switch query[j] {
+		case '\\':
+			j++
+		case query[idx]:
+			return query[idx : j+1], j
+		}
+	}
+	// Unexpected EOS.
+	return query, -1
+}
+
+// Tx writes the transaction start.
+func (w *WriteDriver) Tx(context.Context) (dialect.Tx, error) {
+	return dialect.NopTx(w), nil
+}
+
+// noResult represents a zero result.
+type noResult struct{}
+
+func (noResult) LastInsertId() (int64, error) { return 0, nil }
+func (noResult) RowsAffected() (int64, error) { return 0, nil }
+
+// noRows represents no rows.
+type noRows struct {
+	sql.ColumnScanner
+	cols []string
+	done bool
+}
+
+func (*noRows) Close() error { return nil }
+func (*noRows) Err() error   { return nil }
+func (r *noRows) Next() bool {
+	if !r.done {
+		r.done = true
+		return true
+	}
+	return false
+}
+func (r *noRows) Columns() ([]string, error) { return r.cols, nil }
+func (*noRows) Scan(...any) error            { return nil }
+
+type nopDriver struct {
+	dialect.Driver
+	dialect string
+}
+
+func (d nopDriver) Dialect() string { return d.dialect }
+
+func (nopDriver) Query(context.Context, string, any, any) error {
+	return nil
+}

dialect/sql/schema/writer_test.go

@@ -0,0 +1,189 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"bytes"
+	"context"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/dialect/sql/sqljson"
+
+	"ariga.io/atlas/sql/migrate"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWriteDriver(t *testing.T) {
+	b := &bytes.Buffer{}
+	w := NewWriteDriver(dialect.MySQL, b)
+	ctx := context.Background()
+	tx, err := w.Tx(ctx)
+	require.NoError(t, err)
+	err = tx.Query(ctx, "SELECT `name` FROM `users`", nil, nil)
+	require.EqualError(t, err, "query is not supported by the WriteDriver")
+	err = tx.Exec(ctx, "ALTER TABLE `users` ADD COLUMN `age` int", nil, nil)
+	require.NoError(t, err)
+	err = tx.Exec(ctx, "ALTER TABLE `users` ADD COLUMN `NAME` varchar(100);", nil, nil)
+	require.NoError(t, err)
+	require.NoError(t, tx.Commit())
+	lines := strings.Split(b.String(), "\n")
+	require.Len(t, lines, 3)
+	require.Equal(t, "ALTER TABLE `users` ADD COLUMN `age` int;", lines[0])
+	require.Equal(t, "ALTER TABLE `users` ADD COLUMN `NAME` varchar(100);", lines[1])
+	require.Empty(t, lines[2], "file ends with blank line")
+
+	b.Reset()
+	query, args := sql.Update("users").Schema("test").Set("a", 1).Set("b", "a").Set("c", "'c'").Set("d", true).Where(sql.EQ("p", 0.2)).Query()
+	err = w.Exec(ctx, query, args, nil)
+	require.NoError(t, err)
+	require.Equal(t, "UPDATE `test`.`users` SET `a` = 1, `b` = 'a', `c` = '''c''', `d` = 1 WHERE `p` = 0.2;\n", b.String())
+
+	b.Reset()
+	query, args = sql.Dialect(dialect.MySQL).Update("users").Schema("test").Set("a", "{}").Where(sqljson.ValueIsNull("a")).Query()
+	err = w.Exec(ctx, query, args, nil)
+	require.NoError(t, err)
+	require.Equal(t, "UPDATE `test`.`users` SET `a` = '{}' WHERE JSON_CONTAINS(`a`, 'null', '$');\n", b.String())
+
+	b.Reset()
+	w = NewWriteDriver(dialect.Postgres, b)
+	query, args = sql.Dialect(dialect.Postgres).Update("users").Set("id", uuid.Nil).Set("a", 1).Set("b", time.Now()).Query()
+	err = w.Exec(ctx, query, args, nil)
+	require.NoError(t, err)
+	require.Equal(t, `UPDATE "users" SET "id" = '00000000-0000-0000-0000-000000000000', "a" = 1, "b" = {{ TIME_VALUE }};`+"\n", b.String())
+
+	b.Reset()
+	err = w.Exec(ctx, `INSERT INTO "users" (name) VALUES("a8m") RETURNING id`, nil, nil)
+	require.NoError(t, err)
+	require.Equal(t, `INSERT INTO "users" (name) VALUES("a8m") RETURNING id;`+"\n", b.String())
+
+	// batchCreator uses tx.Query when doing an insert
+	b.Reset()
+	err = w.Query(ctx, `INSERT INTO "users" (name) VALUES("a8m") RETURNING id`, nil, nil)
+	require.NoError(t, err)
+	require.Equal(t, `INSERT INTO "users" (name) VALUES("a8m") RETURNING id;`+"\n", b.String())
+
+	// correct columns are extracted from a returning clause and returned by sql.ColumnScanner.
+	for q, cols := range map[string][]string{
+		`INSERT INTO "users" (name) VALUES("a8m") RETURNING id`:                          {"id"},
+		`INSERT INTO "users" (name) VALUES("a8m") RETURNING id, "name"`:                  {"id", `"name"`},
+		`INSERT INTO "users" (name) VALUES("a8m") RETURNING "id", "name"`:                {`"id"`, `"name"`},
+		`INSERT INTO "users" (name) VALUES("a8m") RETURNING "id", "name"; DROP "groups"`: {`"id"`, `"name"`},
+	} {
+		var rows sql.Rows
+		err = w.Query(ctx, q, nil, &rows)
+		require.NoError(t, err)
+		require.True(t, rows.Next())
+		c, err := rows.Columns()
+		require.NoError(t, err)
+		require.Equal(t, cols, c)
+		require.NoError(t, rows.Scan())
+	}
+	b.Reset()
+}
+
+func TestDirWriter(t *testing.T) {
+	for _, tt := range []struct {
+		dialect  string
+		exec     []string
+		comments []string
+		args     [][]any
+		want     string
+	}{
+		{
+			dialect.MySQL,
+			[]string{
+				"UPDATE `test`.`users` SET `a` = ?",
+				"UPDATE `test`.`users` SET `b` = ?",
+			},
+			[]string{
+				"Comment 1.",
+				"Comment 2.",
+			},
+			[][]any{
+				{1},
+				{2},
+			},
+			"-- Comment 1.\nUPDATE `test`.`users` SET `a` = 1;\n-- Comment 2.\nUPDATE `test`.`users` SET `b` = 2;\n",
+		},
+		{
+			dialect.Postgres,
+			[]string{
+				"INSERT INTO \"users\" (\"name\", \"email\") VALUES ($1, $2) RETURNING \"id\"",
+				"INSERT INTO \"groups\" (\"name\") VALUES ($1) RETURNING \"id\"",
+			},
+			[]string{
+				"Seed users table",
+				"Seed groups table",
+			},
+			[][]any{
+				{"masseelch", "j@ariga.io"},
+				{"admins"},
+			},
+			strings.Join([]string{
+				"-- Seed users table\nINSERT INTO \"users\" (\"name\", \"email\") VALUES ('masseelch', 'j@ariga.io');\n",
+				"-- Seed groups table\nINSERT INTO \"groups\" (\"name\") VALUES ('admins');\n",
+			}, ""),
+		},
+		{
+			dialect.SQLite,
+			[]string{
+				"INSERT INTO `users` (`name`, `email`) VALUES (?, ?) RETURNING `id`",
+				"INSERT INTO `groups` (`name`) VALUES (?) RETURNING `id`",
+			},
+			[]string{
+				"Seed users table",
+				"Seed groups table",
+			},
+			[][]any{
+				{"masseelch", "j@ariga.io"},
+				{"admins"},
+			},
+			strings.Join([]string{
+				"-- Seed users table\nINSERT INTO `users` (`name`, `email`) VALUES ('masseelch', 'j@ariga.io');\n",
+				"-- Seed groups table\nINSERT INTO `groups` (`name`) VALUES ('admins');\n",
+			}, ""),
+		},
+		{
+			dialect.SQLite + " no space",
+			[]string{"INSERT INTO `users` (`name`) VALUES (?)RETURNING `id`"},
+			[]string{"Seed users table"},
+			[][]any{{"masseelch"}},
+			"-- Seed users table\nINSERT INTO `users` (`name`) VALUES ('masseelch');\n",
+		},
+	} {
+		t.Run(tt.dialect, func(t *testing.T) {
+			var (
+				p   = t.TempDir()
+				dir = func() migrate.Dir {
+					d, err := migrate.NewLocalDir(p)
+					require.NoError(t, err)
+					return d
+				}()
+				w   = &DirWriter{Dir: dir}
+				drv = NewWriteDriver(tt.dialect, w)
+			)
+			for i := range tt.exec {
+				require.NoError(t, drv.Exec(context.Background(), tt.exec[i], tt.args[i], nil))
+				w.Change(tt.comments[i])
+			}
+			require.NoError(t, w.Flush("migration_file"))
+			files, err := os.ReadDir(p)
+			require.NoError(t, err)
+			require.Len(t, files, 2)
+			require.Contains(t, files[0].Name(), "_migration_file.sql")
+			buf, err := os.ReadFile(filepath.Join(p, files[0].Name()))
+			require.NoError(t, err)
+			require.Equal(t, tt.want, string(buf))
+			require.Equal(t, "atlas.sum", files[1].Name())
+		})
+	}
+}

dialect/sql/sql.go

@@ -0,0 +1,428 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"fmt"
+	"strings"
+
+	"entgo.io/ent/dialect"
+)
+
+// The following helpers exist to simplify the way raw predicates
+// are defined and used in both ent/schema and generated code. For
+// full predicates API, check out the sql.P in builder.go.
+
+// FieldIsNull returns a raw predicate to check if the given field is NULL.
+func FieldIsNull(name string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(IsNull(s.C(name)))
+	}
+}
+
+// FieldNotNull returns a raw predicate to check if the given field is not NULL.
+func FieldNotNull(name string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(NotNull(s.C(name)))
+	}
+}
+
+// FieldEQ returns a raw predicate to check if the given field equals to the given value.
+func FieldEQ(name string, v any) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(EQ(s.C(name), v))
+	}
+}
+
+// FieldsEQ returns a raw predicate to check if the given fields (columns) are equal.
+func FieldsEQ(field1, field2 string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(ColumnsEQ(s.C(field1), s.C(field2)))
+	}
+}
+
+// FieldNEQ returns a raw predicate to check if the given field does not equal to the given value.
+func FieldNEQ(name string, v any) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(NEQ(s.C(name), v))
+	}
+}
+
+// FieldsNEQ returns a raw predicate to check if the given fields (columns) are not equal.
+func FieldsNEQ(field1, field2 string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(ColumnsNEQ(s.C(field1), s.C(field2)))
+	}
+}
+
+// FieldGT returns a raw predicate to check if the given field is greater than the given value.
+func FieldGT(name string, v any) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(GT(s.C(name), v))
+	}
+}
+
+// FieldsGT returns a raw predicate to check if field1 is greater than field2.
+func FieldsGT(field1, field2 string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(ColumnsGT(s.C(field1), s.C(field2)))
+	}
+}
+
+// FieldGTE returns a raw predicate to check if the given field is greater than or equal the given value.
+func FieldGTE(name string, v any) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(GTE(s.C(name), v))
+	}
+}
+
+// FieldsGTE returns a raw predicate to check if field1 is greater than or equal field2.
+func FieldsGTE(field1, field2 string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(ColumnsGTE(s.C(field1), s.C(field2)))
+	}
+}
+
+// FieldLT returns a raw predicate to check if the value of the field is less than the given value.
+func FieldLT(name string, v any) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(LT(s.C(name), v))
+	}
+}
+
+// FieldsLT returns a raw predicate to check if field1 is lower than field2.
+func FieldsLT(field1, field2 string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(ColumnsLT(s.C(field1), s.C(field2)))
+	}
+}
+
+// FieldLTE returns a raw predicate to check if the value of the field is less than the given value.
+func FieldLTE(name string, v any) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(LTE(s.C(name), v))
+	}
+}
+
+// FieldsLTE returns a raw predicate to check if field1 is lower than or equal field2.
+func FieldsLTE(field1, field2 string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(ColumnsLTE(s.C(field1), s.C(field2)))
+	}
+}
+
+// FieldsHasPrefix returns a raw predicate to checks if field1 begins with the value of field2.
+func FieldsHasPrefix(field1, field2 string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(ColumnsHasPrefix(s.C(field1), s.C(field2)))
+	}
+}
+
+// FieldIn returns a raw predicate to check if the value of the field is IN the given values.
+func FieldIn[T any](name string, vs ...T) func(*Selector) {
+	return func(s *Selector) {
+		v := make([]any, len(vs))
+		for i := range v {
+			v[i] = vs[i]
+		}
+		s.Where(In(s.C(name), v...))
+	}
+}
+
+// FieldNotIn returns a raw predicate to check if the value of the field is NOT IN the given values.
+func FieldNotIn[T any](name string, vs ...T) func(*Selector) {
+	return func(s *Selector) {
+		v := make([]any, len(vs))
+		for i := range v {
+			v[i] = vs[i]
+		}
+		s.Where(NotIn(s.C(name), v...))
+	}
+}
+
+// FieldEqualFold returns a raw predicate to check if the field is equal to the given string under case-folding.
+func FieldEqualFold(name string, substr string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(EqualFold(s.C(name), substr))
+	}
+}
+
+// FieldHasPrefix returns a raw predicate to check if the field has the given prefix.
+func FieldHasPrefix(name string, prefix string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(HasPrefix(s.C(name), prefix))
+	}
+}
+
+// FieldHasPrefixFold returns a raw predicate to check if the field has the given prefix with case-folding
+func FieldHasPrefixFold(name string, prefix string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(HasPrefixFold(s.C(name), prefix))
+	}
+}
+
+// FieldHasSuffix returns a raw predicate to check if the field has the given suffix.
+func FieldHasSuffix(name string, suffix string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(HasSuffix(s.C(name), suffix))
+	}
+}
+
+// FieldHasSuffixFold returns a raw predicate to check if the field has the given suffix with case-folding
+func FieldHasSuffixFold(name string, suffix string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(HasSuffixFold(s.C(name), suffix))
+	}
+}
+
+// FieldContains returns a raw predicate to check if the field contains the given substring.
+func FieldContains(name string, substr string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(Contains(s.C(name), substr))
+	}
+}
+
+// FieldContainsFold returns a raw predicate to check if the field contains the given substring with case-folding.
+func FieldContainsFold(name string, substr string) func(*Selector) {
+	return func(s *Selector) {
+		s.Where(ContainsFold(s.C(name), substr))
+	}
+}
+
+// AndPredicates returns a new predicate for joining multiple generated predicates with AND between them.
+func AndPredicates[P ~func(*Selector)](predicates ...P) func(*Selector) {
+	return func(s *Selector) {
+		s.CollectPredicates()
+		for _, p := range predicates {
+			p(s)
+		}
+		collected := s.CollectedPredicates()
+		s.UncollectedPredicates()
+		switch len(collected) {
+		case 0:
+		case 1:
+			s.Where(collected[0])
+		default:
+			s.Where(And(collected...))
+		}
+	}
+}
+
+// OrPredicates returns a new predicate for joining multiple generated predicates with OR between them.
+func OrPredicates[P ~func(*Selector)](predicates ...P) func(*Selector) {
+	return func(s *Selector) {
+		s.CollectPredicates()
+		for _, p := range predicates {
+			p(s)
+		}
+		collected := s.CollectedPredicates()
+		s.UncollectedPredicates()
+		switch len(collected) {
+		case 0:
+		case 1:
+			s.Where(collected[0])
+		default:
+			s.Where(Or(collected...))
+		}
+	}
+}
+
+// NotPredicates wraps the generated predicates with NOT. For example, NOT(P), NOT((P1 AND P2)).
+func NotPredicates[P ~func(*Selector)](predicates ...P) func(*Selector) {
+	return func(s *Selector) {
+		s.CollectPredicates()
+		for _, p := range predicates {
+			p(s)
+		}
+		collected := s.CollectedPredicates()
+		s.UncollectedPredicates()
+		switch len(collected) {
+		case 0:
+		case 1:
+			s.Where(Not(collected[0]))
+		default:
+			s.Where(Not(And(collected...)))
+		}
+	}
+}
+
+// ColumnCheck is a function that verifies whether the
+// specified column exists within the given table.
+type ColumnCheck func(table, column string) error
+
+// NewColumnCheck returns a function that verifies whether the specified column exists
+// within the given table. This function is utilized by the generated code to validate
+// column names in ordering functions.
+func NewColumnCheck(checks map[string]func(string) bool) ColumnCheck {
+	return func(table, column string) error {
+		check, ok := checks[table]
+		if !ok {
+			return fmt.Errorf("unknown table %q", table)
+		}
+		if !check(column) {
+			return fmt.Errorf("unknown column %q for table %q", column, table)
+		}
+		return nil
+	}
+}
+
+type (
+	// OrderFieldTerm represents an ordering by a field.
+	OrderFieldTerm struct {
+		OrderTermOptions
+		Field string // Field name.
+	}
+	// OrderExprTerm represents an ordering by an expression.
+	OrderExprTerm struct {
+		OrderTermOptions
+		Expr func(*Selector) Querier // Expression.
+	}
+	// OrderTerm represents an ordering by a term.
+	OrderTerm interface {
+		term()
+	}
+	// OrderTermOptions represents options for ordering by a term.
+	OrderTermOptions struct {
+		Desc       bool   // Whether to sort in descending order.
+		As         string // Optional alias.
+		Selected   bool   // Whether the term should be selected.
+		NullsFirst bool   // Whether to sort nulls first.
+		NullsLast  bool   // Whether to sort nulls last.
+	}
+	// OrderTermOption is an option for ordering by a term.
+	OrderTermOption func(*OrderTermOptions)
+)
+
+// OrderDesc returns an option to sort in descending order.
+func OrderDesc() OrderTermOption {
+	return func(o *OrderTermOptions) {
+		o.Desc = true
+	}
+}
+
+// OrderAsc returns an option to sort in ascending order.
+func OrderAsc() OrderTermOption {
+	return func(o *OrderTermOptions) {
+		o.Desc = false
+	}
+}
+
+// OrderAs returns an option to set the alias for the ordering.
+func OrderAs(as string) OrderTermOption {
+	return func(o *OrderTermOptions) {
+		o.As = as
+	}
+}
+
+// OrderSelected returns an option to select the ordering term.
+func OrderSelected() OrderTermOption {
+	return func(o *OrderTermOptions) {
+		o.Selected = true
+	}
+}
+
+// OrderSelectAs returns an option to set and select the alias for the ordering.
+func OrderSelectAs(as string) OrderTermOption {
+	return func(o *OrderTermOptions) {
+		o.As = as
+		o.Selected = true
+	}
+}
+
+// OrderNullsFirst returns an option to sort nulls first.
+func OrderNullsFirst() OrderTermOption {
+	return func(o *OrderTermOptions) {
+		o.NullsFirst = true
+	}
+}
+
+// OrderNullsLast returns an option to sort nulls last.
+func OrderNullsLast() OrderTermOption {
+	return func(o *OrderTermOptions) {
+		o.NullsLast = true
+	}
+}
+
+// NewOrderTermOptions returns a new OrderTermOptions from the given options.
+func NewOrderTermOptions(opts ...OrderTermOption) *OrderTermOptions {
+	o := &OrderTermOptions{}
+	for _, opt := range opts {
+		opt(o)
+	}
+	return o
+}
+
+// OrderByField returns an ordering by the given field.
+func OrderByField(field string, opts ...OrderTermOption) *OrderFieldTerm {
+	return &OrderFieldTerm{Field: field, OrderTermOptions: *NewOrderTermOptions(opts...)}
+}
+
+// OrderBySum returns an ordering by the sum of the given field.
+func OrderBySum(field string, opts ...OrderTermOption) *OrderExprTerm {
+	return orderByAgg("SUM", field, opts...)
+}
+
+// OrderByCount returns an ordering by the count of the given field.
+func OrderByCount(field string, opts ...OrderTermOption) *OrderExprTerm {
+	return orderByAgg("COUNT", field, opts...)
+}
+
+// orderByAgg returns an ordering by the aggregation of the given field.
+func orderByAgg(fn, field string, opts ...OrderTermOption) *OrderExprTerm {
+	return &OrderExprTerm{
+		OrderTermOptions: *NewOrderTermOptions(
+			append(
+				// Default alias is "<func>_<field>".
+				[]OrderTermOption{OrderAs(fmt.Sprintf("%s_%s", strings.ToLower(fn), field))},
+				opts...,
+			)...,
+		),
+		Expr: func(s *Selector) Querier {
+			var c string
+			switch {
+			case field == "*", isFunc(field):
+				c = field
+			default:
+				c = s.C(field)
+			}
+			return Raw(fmt.Sprintf("%s(%s)", fn, c))
+		},
+	}
+}
+
+// OrderByRand returns a term to natively order by a random value.
+func OrderByRand() func(*Selector) {
+	return func(s *Selector) {
+		s.OrderExprFunc(func(b *Builder) {
+			switch s.Dialect() {
+			case dialect.MySQL:
+				b.WriteString("RAND()")
+			default:
+				b.WriteString("RANDOM()")
+			}
+		})
+	}
+}
+
+// ToFunc returns a function that sets the ordering on the given selector.
+// This is used by the generated code.
+func (f *OrderFieldTerm) ToFunc() func(*Selector) {
+	return func(s *Selector) {
+		s.OrderExprFunc(func(b *Builder) {
+			b.WriteString(s.C(f.Field))
+			if f.Desc {
+				b.WriteString(" DESC")
+			}
+			if f.NullsFirst {
+				b.WriteString(" NULLS FIRST")
+			} else if f.NullsLast {
+				b.WriteString(" NULLS LAST")
+			}
+		})
+	}
+}
+
+func (OrderFieldTerm) term() {}
+func (OrderExprTerm) term()  {}

dialect/sql/sql_test.go

@@ -0,0 +1,471 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sql
+
+import (
+	"testing"
+
+	"entgo.io/ent/dialect"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFieldIsNull(t *testing.T) {
+	p := FieldIsNull("name")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` IS NULL", query)
+		require.Empty(t, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" IS NULL`, query)
+		require.Empty(t, args)
+	})
+}
+
+func TestFieldNotNull(t *testing.T) {
+	p := FieldNotNull("name")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` IS NOT NULL", query)
+		require.Empty(t, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" IS NOT NULL`, query)
+		require.Empty(t, args)
+	})
+}
+
+func TestFieldEQ(t *testing.T) {
+	p := FieldEQ("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` = ?", query)
+		require.Equal(t, []any{"a8m"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" = $1`, query)
+		require.Equal(t, []any{"a8m"}, args)
+	})
+}
+
+func TestFieldsEQ(t *testing.T) {
+	p := FieldsEQ("create_time", "update_time")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`create_time` = `users`.`update_time`", query)
+		require.Empty(t, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."create_time" = "users"."update_time"`, query)
+		require.Empty(t, args)
+	})
+}
+
+func TestFieldsNEQ(t *testing.T) {
+	p := FieldsNEQ("create_time", "update_time")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`create_time` <> `users`.`update_time`", query)
+		require.Empty(t, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."create_time" <> "users"."update_time"`, query)
+		require.Empty(t, args)
+	})
+}
+
+func TestFieldNEQ(t *testing.T) {
+	p := FieldNEQ("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` <> ?", query)
+		require.Equal(t, []any{"a8m"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" <> $1`, query)
+		require.Equal(t, []any{"a8m"}, args)
+	})
+}
+
+func TestFieldGT(t *testing.T) {
+	p := FieldGT("stars", 1000)
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`stars` > ?", query)
+		require.Equal(t, []any{1000}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."stars" > $1`, query)
+		require.Equal(t, []any{1000}, args)
+	})
+}
+
+func TestFieldsGT(t *testing.T) {
+	p := FieldsGT("a", "b")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`a` > `users`.`b`", query)
+		require.Empty(t, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."a" > "users"."b"`, query)
+		require.Empty(t, args)
+	})
+}
+
+func TestFieldGTE(t *testing.T) {
+	p := FieldGTE("stars", 1000)
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`stars` >= ?", query)
+		require.Equal(t, []any{1000}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."stars" >= $1`, query)
+		require.Equal(t, []any{1000}, args)
+	})
+}
+
+func TestFieldsGTE(t *testing.T) {
+	p := FieldsGTE("a", "b")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`a` >= `users`.`b`", query)
+		require.Empty(t, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."a" >= "users"."b"`, query)
+		require.Empty(t, args)
+	})
+}
+
+func TestFieldLT(t *testing.T) {
+	p := FieldLT("stars", 1000)
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`stars` < ?", query)
+		require.Equal(t, []any{1000}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."stars" < $1`, query)
+		require.Equal(t, []any{1000}, args)
+	})
+}
+
+func TestFieldsLT(t *testing.T) {
+	p := FieldsLT("a", "b")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`a` < `users`.`b`", query)
+		require.Empty(t, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."a" < "users"."b"`, query)
+		require.Empty(t, args)
+	})
+}
+
+func TestFieldLTE(t *testing.T) {
+	p := FieldLTE("stars", 1000)
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`stars` <= ?", query)
+		require.Equal(t, []any{1000}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."stars" <= $1`, query)
+		require.Equal(t, []any{1000}, args)
+	})
+}
+
+func TestFieldsLTE(t *testing.T) {
+	p := FieldsLTE("a", "b")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`a` <= `users`.`b`", query)
+		require.Empty(t, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."a" <= "users"."b"`, query)
+		require.Empty(t, args)
+	})
+}
+
+func TestFieldIn(t *testing.T) {
+	p := FieldIn("name", "a8m", "foo", "bar")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` IN (?, ?, ?)", query)
+		require.Equal(t, []any{"a8m", "foo", "bar"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" IN ($1, $2, $3)`, query)
+		require.Equal(t, []any{"a8m", "foo", "bar"}, args)
+	})
+}
+
+func TestFieldNotIn(t *testing.T) {
+	p := FieldNotIn("id", 1, 2, 3)
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`id` NOT IN (?, ?, ?)", query)
+		require.Equal(t, []any{1, 2, 3}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."id" NOT IN ($1, $2, $3)`, query)
+		require.Equal(t, []any{1, 2, 3}, args)
+	})
+}
+
+func TestFieldEqualFold(t *testing.T) {
+	p := FieldEqualFold("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` COLLATE utf8mb4_general_ci = ?", query)
+		require.Equal(t, []any{"a8m"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" ILIKE $1`, query)
+		require.Equal(t, []any{"a8m"}, args)
+	})
+}
+
+func TestFieldHasPrefix(t *testing.T) {
+	p := FieldHasPrefix("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` LIKE ?", query)
+		require.Equal(t, []any{"a8m%"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" LIKE $1`, query)
+		require.Equal(t, []any{"a8m%"}, args)
+	})
+}
+
+func TestFieldHasPrefixFold(t *testing.T) {
+	p := FieldHasPrefixFold("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` COLLATE utf8mb4_general_ci LIKE ?", query)
+		require.Equal(t, []any{"a8m%"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" ILIKE $1`, query)
+		require.Equal(t, []any{"a8m%"}, args)
+	})
+}
+
+func TestFieldHasSuffix(t *testing.T) {
+	p := FieldHasSuffix("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` LIKE ?", query)
+		require.Equal(t, []any{"%a8m"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" LIKE $1`, query)
+		require.Equal(t, []any{"%a8m"}, args)
+	})
+}
+
+func TestFieldHasSuffixFold(t *testing.T) {
+	p := FieldHasSuffixFold("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` COLLATE utf8mb4_general_ci LIKE ?", query)
+		require.Equal(t, []any{"%a8m"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" ILIKE $1`, query)
+		require.Equal(t, []any{"%a8m"}, args)
+	})
+}
+
+func TestFieldContains(t *testing.T) {
+	p := FieldContains("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` LIKE ?", query)
+		require.Equal(t, []any{"%a8m%"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" LIKE $1`, query)
+		require.Equal(t, []any{"%a8m%"}, args)
+	})
+}
+
+func TestFieldContainsFold(t *testing.T) {
+	p := FieldContainsFold("name", "a8m")
+	t.Run("MySQL", func(t *testing.T) {
+		s := Dialect(dialect.MySQL).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, "SELECT * FROM `users` WHERE `users`.`name` COLLATE utf8mb4_general_ci LIKE ?", query)
+		require.Equal(t, []any{"%a8m%"}, args)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		s := Dialect(dialect.Postgres).Select("*").From(Table("users"))
+		p(s)
+		query, args := s.Query()
+		require.Equal(t, `SELECT * FROM "users" WHERE "users"."name" ILIKE $1`, query)
+		require.Equal(t, []any{"%a8m%"}, args)
+	})
+}
+
+func TestAndPredicates(t *testing.T) {
+	s := Select("*").From(Table("users")).Where(EQ("name", "a8m"))
+	p := AndPredicates(
+		FieldEQ("a", "foo"),
+		FieldEQ("b", 1),
+		func(s *Selector) {
+			petT := Table("pets").As("p")
+			s.Join(petT).On(petT.C("owner_id"), s.C("id"))
+		},
+	)
+	p(s)
+	query, args := s.Query()
+	require.Equal(t, "SELECT * FROM `users` JOIN `pets` AS `p` ON `p`.`owner_id` = `users`.`id` WHERE `name` = ? AND (`users`.`a` = ? AND `users`.`b` = ?)", query)
+	require.Equal(t, []any{"a8m", "foo", 1}, args)
+}
+
+func TestOrPredicates(t *testing.T) {
+	s := Select("*").From(Table("users")).Where(EQ("name", "a8m"))
+	p := OrPredicates(
+		AndPredicates(
+			FieldEQ("a", "foo"),
+			FieldEQ("b", 1),
+		),
+		func(s *Selector) {
+			petT := Table("pets").As("p")
+			s.Join(petT).On(petT.C("owner_id"), s.C("id"))
+			s.Where(EQ(petT.C("name"), "c"))
+		},
+	)
+	p(s)
+	query, args := s.Query()
+	require.Equal(t, "SELECT * FROM `users` JOIN `pets` AS `p` ON `p`.`owner_id` = `users`.`id` WHERE `name` = ? AND ((`users`.`a` = ? AND `users`.`b` = ?) OR `p`.`name` = ?)", query)
+	require.Equal(t, []any{"a8m", "foo", 1, "c"}, args)
+}
+
+func TestNotPredicates(t *testing.T) {
+	s := Select("*").From(Table("users")).Where(EQ("name", "a8m"))
+	NotPredicates(FieldEQ("a", "a"), FieldEQ("b", "b"))(s)
+	NotPredicates(FieldEQ("c", "c"))(s)
+	query, args := s.Query()
+	require.Equal(t, "SELECT * FROM `users` WHERE (`name` = ? AND (NOT (`users`.`a` = ? AND `users`.`b` = ?))) AND (NOT (`users`.`c` = ?))", query)
+	require.Equal(t, []any{"a8m", "a", "b", "c"}, args)
+}

dialect/sql/sqlgraph/blob.go

@@ -0,0 +1,155 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sqlgraph
+
+import (
+	"context"
+
+	"entgo.io/ent"
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+)
+
+// BlobSpec configures SQL-level blob key queries and implements [ent.BlobQuerier].
+type BlobSpec struct {
+	Driver  dialect.Driver
+	Table   string
+	Columns map[string]string // field name -> key column name
+	// CheckRefs holds the fields whose keys more than one row can hold, the ones
+	// declaring CheckRefs on the schema. Only these are looked up before a delete;
+	// a key generated fresh per write has no other holder to find. An empty
+	// CheckRefs makes [BlobSpec.ReferencedBlobKeys] a no-op.
+	CheckRefs map[string]bool
+	Predicate func(*sql.Selector)
+}
+
+// maxBlobKeysPerQuery bounds the keys per reference lookup, keeping the
+// statement clear of driver placeholder limits on large deletes.
+const maxBlobKeysPerQuery = 500
+
+// ReferencedBlobKeys implements [ent.BlobRefChecker]. It returns the subset of keys
+// some row holds in its field's key column. [BlobSpec.Predicate] is deliberately not
+// applied: the question is which keys are in use anywhere, and callers time the lookup
+// so the mutation's own rows already read the way they will settle.
+//
+// Only fields listed in [BlobSpec.CheckRefs] are looked up. Keys of any other field are
+// reported as held by no row, which is what an unshared key strategy guarantees anyway.
+//
+// The lookup only asks whether a key is used, never how often, so an index on each
+// blob key column lets it stop at the first match.
+func (s *BlobSpec) ReferencedBlobKeys(ctx context.Context, keys []ent.BlobKey) ([]ent.BlobKey, error) {
+	if len(keys) == 0 || len(s.CheckRefs) == 0 {
+		return nil, nil
+	}
+	// Group the distinct keys by the column they would be stored in.
+	lookup := make(map[string][]any)
+	seen := make(map[ent.BlobKey]bool, len(keys))
+	for _, k := range keys {
+		if k.Key == "" || seen[k] || !s.CheckRefs[k.Field] {
+			continue
+		}
+		seen[k] = true
+		if col, ok := s.Columns[k.Field]; ok {
+			lookup[col] = append(lookup[col], k.Key)
+		}
+	}
+	var referenced []ent.BlobKey
+	for field := range s.CheckRefs {
+		col, ok := s.Columns[field]
+		if !ok {
+			continue
+		}
+		vals := lookup[col]
+		for len(vals) > 0 {
+			n := min(len(vals), maxBlobKeysPerQuery)
+			found, err := s.selectReferenced(ctx, col, vals[:n])
+			if err != nil {
+				return nil, err
+			}
+			for _, key := range found {
+				referenced = append(referenced, ent.BlobKey{Field: field, Key: key})
+			}
+			vals = vals[n:]
+		}
+	}
+	return referenced, nil
+}
+
+// selectReferenced returns which of vals appear in col.
+func (s *BlobSpec) selectReferenced(ctx context.Context, col string, vals []any) ([]string, error) {
+	query, args := sql.Dialect(s.Driver.Dialect()).
+		Select(col).
+		Distinct().
+		From(sql.Table(s.Table)).
+		Where(sql.In(col, vals...)).
+		Query()
+	rows := &sql.Rows{}
+	if err := s.Driver.Query(ctx, query, args, rows); err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var found []string
+	for rows.Next() {
+		var key string
+		if err := rows.Scan(&key); err != nil {
+			return nil, err
+		}
+		found = append(found, key)
+	}
+	return found, rows.Err()
+}
+
+// QueryBlobKeys implements [ent.BlobQuerier].
+// If fields is nil, all columns are queried (for deletes);
+// otherwise only the named fields are queried.
+func (s *BlobSpec) QueryBlobKeys(ctx context.Context, fields []string) ([]ent.BlobKey, error) {
+	cols := s.Columns
+	if len(fields) > 0 {
+		cols = make(map[string]string, len(fields))
+		for _, f := range fields {
+			if c, ok := s.Columns[f]; ok {
+				cols[f] = c
+			}
+		}
+	}
+	if len(cols) == 0 {
+		return nil, nil
+	}
+	names := make([]string, 0, len(cols))
+	colNames := make([]string, 0, len(cols))
+	for field, col := range cols {
+		names = append(names, field)
+		colNames = append(colNames, col)
+	}
+	selector := sql.Dialect(s.Driver.Dialect()).
+		Select(colNames...).
+		From(sql.Table(s.Table))
+	if s.Predicate != nil {
+		s.Predicate(selector)
+	}
+	query, args := selector.Query()
+	rows := &sql.Rows{}
+	if err := s.Driver.Query(ctx, query, args, rows); err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var keys []ent.BlobKey
+	for rows.Next() {
+		vals := make([]*string, len(colNames))
+		ptrs := make([]any, len(colNames))
+		for i := range vals {
+			ptrs[i] = &vals[i]
+		}
+		if err := rows.Scan(ptrs...); err != nil {
+			return nil, err
+		}
+		for i, v := range vals {
+			if v != nil && *v != "" {
+				keys = append(keys, ent.BlobKey{Field: names[i], Key: *v})
+			}
+		}
+	}
+	return keys, rows.Err()
+}

dialect/sql/sqlgraph/blob_test.go

@@ -0,0 +1,182 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sqlgraph
+
+import (
+	"context"
+	"database/sql/driver"
+	"fmt"
+	"testing"
+
+	"entgo.io/ent"
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBlobSpecReferencedBlobKeys(t *testing.T) {
+	tests := []struct {
+		name    string
+		dialect string
+		keys    []ent.BlobKey
+		expect  func(sqlmock.Sqlmock)
+		want    []ent.BlobKey
+	}{
+		{
+			name:    "keys held by a row",
+			dialect: dialect.MySQL,
+			keys: []ent.BlobKey{
+				{Field: "content", Key: "k1"},
+				{Field: "content", Key: "k2"},
+			},
+			expect: func(m sqlmock.Sqlmock) {
+				m.ExpectQuery(escape("SELECT DISTINCT `content_key` FROM `documents` WHERE `content_key` IN (?, ?)")).
+					WithArgs("k1", "k2").
+					WillReturnRows(sqlmock.NewRows([]string{"content_key"}).AddRow("k1"))
+			},
+			want: []ent.BlobKey{{Field: "content", Key: "k1"}},
+		},
+		{
+			name:    "keys held by no row",
+			dialect: dialect.SQLite,
+			keys:    []ent.BlobKey{{Field: "content", Key: "gone"}},
+			expect: func(m sqlmock.Sqlmock) {
+				m.ExpectQuery(escape("SELECT DISTINCT `content_key` FROM `documents` WHERE `content_key` IN (?)")).
+					WithArgs("gone").
+					WillReturnRows(sqlmock.NewRows([]string{"content_key"}))
+			},
+			want: nil,
+		},
+		{
+			// The mutation's own rows are not filtered out here -- callers time the
+			// lookup so those rows already read the way they will settle.
+			name:    "the mutation's predicate is not applied",
+			dialect: dialect.Postgres,
+			keys:    []ent.BlobKey{{Field: "content", Key: "k1"}},
+			expect: func(m sqlmock.Sqlmock) {
+				m.ExpectQuery(escape(`SELECT DISTINCT "content_key" FROM "documents" WHERE "content_key" IN ($1)`)).
+					WithArgs("k1").
+					WillReturnRows(sqlmock.NewRows([]string{"content_key"}).AddRow("k1"))
+			},
+			want: []ent.BlobKey{{Field: "content", Key: "k1"}},
+		},
+		{
+			name:    "duplicate keys are looked up once",
+			dialect: dialect.MySQL,
+			keys: []ent.BlobKey{
+				{Field: "content", Key: "k1"},
+				{Field: "content", Key: "k1"},
+			},
+			expect: func(m sqlmock.Sqlmock) {
+				m.ExpectQuery(escape("SELECT DISTINCT `content_key` FROM `documents` WHERE `content_key` IN (?)")).
+					WithArgs("k1").
+					WillReturnRows(sqlmock.NewRows([]string{"content_key"}).AddRow("k1"))
+			},
+			want: []ent.BlobKey{{Field: "content", Key: "k1"}},
+		},
+		{
+			name:    "unknown fields and empty keys are skipped",
+			dialect: dialect.MySQL,
+			keys: []ent.BlobKey{
+				{Field: "nosuchfield", Key: "k1"},
+				{Field: "content", Key: ""},
+			},
+			expect: func(sqlmock.Sqlmock) {},
+			want:   nil,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			db, mock, err := sqlmock.New()
+			require.NoError(t, err)
+			tt.expect(mock)
+			spec := &BlobSpec{
+				Driver: sql.OpenDB(tt.dialect, db),
+				Table:  "documents",
+				// Set to prove ReferencedBlobKeys ignores it.
+				Predicate: func(s *sql.Selector) { s.Where(sql.EQ("id", 7)) },
+				Columns:   map[string]string{"content": "content_key", "meta": "meta_key"},
+				CheckRefs: map[string]bool{"content": true},
+			}
+			referenced, err := spec.ReferencedBlobKeys(context.Background(), tt.keys)
+			require.NoError(t, err)
+			require.Equal(t, tt.want, referenced)
+			require.NoError(t, mock.ExpectationsWereMet())
+		})
+	}
+}
+
+// A key set larger than maxBlobKeysPerQuery is split across statements so the
+// query stays clear of driver placeholder limits.
+func TestBlobSpecReferencedBlobKeysChunks(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	var (
+		keys  = make([]ent.BlobKey, 0, maxBlobKeysPerQuery+1)
+		first = make([]driver.Value, 0, maxBlobKeysPerQuery)
+	)
+	for i := range maxBlobKeysPerQuery + 1 {
+		key := fmt.Sprintf("k%d", i)
+		keys = append(keys, ent.BlobKey{Field: "content", Key: key})
+		if i < maxBlobKeysPerQuery {
+			first = append(first, key)
+		}
+	}
+	// The first statement carries a full chunk, the leftover key its own. Matched
+	// loosely — escape anchors the pattern, and the argument list asserts the size.
+	mock.ExpectQuery("SELECT DISTINCT").
+		WithArgs(first...).
+		WillReturnRows(sqlmock.NewRows([]string{"content_key"}).AddRow(keys[0].Key))
+	mock.ExpectQuery(escape("SELECT DISTINCT `content_key` FROM `documents` WHERE `content_key` IN (?)")).
+		WithArgs(keys[maxBlobKeysPerQuery].Key).
+		WillReturnRows(sqlmock.NewRows([]string{"content_key"}).AddRow(keys[maxBlobKeysPerQuery].Key))
+	spec := &BlobSpec{
+		Driver:    sql.OpenDB(dialect.MySQL, db),
+		Table:     "documents",
+		Columns:   map[string]string{"content": "content_key"},
+		CheckRefs: map[string]bool{"content": true},
+	}
+	referenced, err := spec.ReferencedBlobKeys(context.Background(), keys)
+	require.NoError(t, err)
+	require.Equal(t, []ent.BlobKey{keys[0], keys[maxBlobKeysPerQuery]}, referenced)
+	require.NoError(t, mock.ExpectationsWereMet())
+}
+
+// Empty input, and fields that did not opt in, must not reach the database.
+func TestBlobSpecReferencedBlobKeysNoop(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	spec := &BlobSpec{
+		Driver:    sql.OpenDB(dialect.MySQL, db),
+		Table:     "documents",
+		Columns:   map[string]string{"content": "content_key"},
+		CheckRefs: map[string]bool{"content": true},
+	}
+	keys := []ent.BlobKey{{Field: "content", Key: "k"}}
+	referenced, err := spec.ReferencedBlobKeys(context.Background(), nil)
+	require.NoError(t, err)
+	require.Empty(t, referenced)
+
+	// No field opted in: the whole lookup is skipped.
+	spec.CheckRefs = nil
+	referenced, err = spec.ReferencedBlobKeys(context.Background(), keys)
+	require.NoError(t, err)
+	require.Empty(t, referenced)
+
+	// Some other field opted in, but not this one.
+	spec.CheckRefs = map[string]bool{"thumbnail": true}
+	referenced, err = spec.ReferencedBlobKeys(context.Background(), keys)
+	require.NoError(t, err)
+	require.Empty(t, referenced)
+
+	spec.CheckRefs = map[string]bool{"content": true}
+	spec.Columns = nil
+	referenced, err = spec.ReferencedBlobKeys(context.Background(), keys)
+	require.NoError(t, err)
+	require.Empty(t, referenced)
+	require.NoError(t, mock.ExpectationsWereMet())
+}

dialect/sql/sqlgraph/entql.go

@@ -0,0 +1,334 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sqlgraph
+
+import (
+	"fmt"
+
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/entql"
+)
+
+type (
+	// A Schema holds a representation of ent/schema at runtime. Each Node
+	// represents a single schema-type and its relations in the graph (storage).
+	//
+	// It is used for translating common graph traversal operations to the
+	// underlying SQL storage. For example, an operation like `has_edge(E)`,
+	// will be translated to an SQL lookup based on the relation type and the
+	// FK configuration.
+	Schema struct {
+		Nodes []*Node
+	}
+
+	// A Node in the graph holds the SQL information for an ent/schema.
+	Node struct {
+		NodeSpec
+
+		// Type holds the node type (schema name).
+		Type string
+
+		// Fields maps from field names to their spec.
+		Fields map[string]*FieldSpec
+
+		// Edges maps from edge names to their spec.
+		Edges map[string]struct {
+			To   *Node
+			Spec *EdgeSpec
+		}
+	}
+)
+
+// AddE adds an edge to the graph. It fails, if one of the node
+// types is missing.
+//
+//	g.AddE("pets", spec, "user", "pet")
+//	g.AddE("friends", spec, "user", "user")
+func (g *Schema) AddE(name string, spec *EdgeSpec, from, to string) error {
+	var fromT, toT *Node
+	for i := range g.Nodes {
+		t := g.Nodes[i].Type
+		if t == from {
+			fromT = g.Nodes[i]
+		}
+		if t == to {
+			toT = g.Nodes[i]
+		}
+	}
+	if fromT == nil || toT == nil {
+		return fmt.Errorf("from/to type was not found")
+	}
+	if fromT.Edges == nil {
+		fromT.Edges = make(map[string]struct {
+			To   *Node
+			Spec *EdgeSpec
+		})
+	}
+	fromT.Edges[name] = struct {
+		To   *Node
+		Spec *EdgeSpec
+	}{
+		To:   toT,
+		Spec: spec,
+	}
+	return nil
+}
+
+// MustAddE is like AddE but panics if the edge can be added to the graph.
+func (g *Schema) MustAddE(name string, spec *EdgeSpec, from, to string) {
+	if err := g.AddE(name, spec, from, to); err != nil {
+		panic(err)
+	}
+}
+
+// EvalP evaluates the entql predicate on the given selector (query builder).
+func (g *Schema) EvalP(nodeType string, p entql.P, selector *sql.Selector) error {
+	var node *Node
+	for i := range g.Nodes {
+		if g.Nodes[i].Type == nodeType {
+			node = g.Nodes[i]
+			break
+		}
+	}
+	if node == nil {
+		return fmt.Errorf("node %s was not found in the graph schema", nodeType)
+	}
+	pr, err := evalExpr(node, selector, p)
+	if err != nil {
+		return err
+	}
+	selector.Where(pr)
+	return nil
+}
+
+// FuncSelector represents a selector function to be used as an entql foreign-function.
+const FuncSelector entql.Func = "func_selector"
+
+// wrappedFunc wraps the selector-function to an ent-expression.
+type wrappedFunc struct {
+	entql.Expr
+	Func func(*sql.Selector)
+}
+
+// WrapFunc wraps a selector-func with an entql call expression.
+func WrapFunc(s func(*sql.Selector)) *entql.CallExpr {
+	return &entql.CallExpr{
+		Func: FuncSelector,
+		Args: []entql.Expr{wrappedFunc{Func: s}},
+	}
+}
+
+var (
+	binary = [...]sql.Op{
+		entql.OpEQ:    sql.OpEQ,
+		entql.OpNEQ:   sql.OpNEQ,
+		entql.OpGT:    sql.OpGT,
+		entql.OpGTE:   sql.OpGTE,
+		entql.OpLT:    sql.OpLT,
+		entql.OpLTE:   sql.OpLTE,
+		entql.OpIn:    sql.OpIn,
+		entql.OpNotIn: sql.OpNotIn,
+	}
+	nary = [...]func(...*sql.Predicate) *sql.Predicate{
+		entql.OpAnd: sql.And,
+		entql.OpOr:  sql.Or,
+	}
+	strFunc = map[entql.Func]func(string, string) *sql.Predicate{
+		entql.FuncContains:     sql.Contains,
+		entql.FuncContainsFold: sql.ContainsFold,
+		entql.FuncEqualFold:    sql.EqualFold,
+		entql.FuncHasPrefix:    sql.HasPrefix,
+		entql.FuncHasSuffix:    sql.HasSuffix,
+	}
+	nullFunc = [...]func(string) *sql.Predicate{
+		entql.OpEQ:  sql.IsNull,
+		entql.OpNEQ: sql.NotNull,
+	}
+)
+
+// state represents the state of a predicate evaluation.
+// Note that, the evaluation output is a predicate to be
+// applied on the database.
+type state struct {
+	sql.Builder
+	context  *Node
+	selector *sql.Selector
+}
+
+// evalExpr evaluates the entql expression and returns a new SQL predicate to be applied on the database.
+func evalExpr(context *Node, selector *sql.Selector, expr entql.Expr) (p *sql.Predicate, err error) {
+	ex := &state{
+		context:  context,
+		selector: selector,
+	}
+	defer catch(&err)
+	p = ex.evalExpr(expr)
+	return
+}
+
+// evalExpr evaluates any expression.
+func (e *state) evalExpr(expr entql.Expr) *sql.Predicate {
+	switch expr := expr.(type) {
+	case *entql.BinaryExpr:
+		return e.evalBinary(expr)
+	case *entql.UnaryExpr:
+		return sql.Not(e.evalExpr(expr.X))
+	case *entql.NaryExpr:
+		ps := make([]*sql.Predicate, len(expr.Xs))
+		for i, x := range expr.Xs {
+			ps[i] = e.evalExpr(x)
+		}
+		return nary[expr.Op](ps...)
+	case *entql.CallExpr:
+		switch expr.Func {
+		case entql.FuncHasPrefix, entql.FuncHasSuffix, entql.FuncContains, entql.FuncEqualFold, entql.FuncContainsFold:
+			expect(len(expr.Args) == 2, "invalid number of arguments for %s", expr.Func)
+			f, ok := expr.Args[0].(*entql.Field)
+			expect(ok, "*entql.Field, got %T", expr.Args[0])
+			v, ok := expr.Args[1].(*entql.Value)
+			expect(ok, "*entql.Value, got %T", expr.Args[1])
+			s, ok := v.V.(string)
+			expect(ok, "string value, got %T", v.V)
+			return strFunc[expr.Func](e.field(f), s)
+		case entql.FuncHasEdge:
+			expect(len(expr.Args) > 0, "invalid number of arguments for %s", expr.Func)
+			edge, ok := expr.Args[0].(*entql.Edge)
+			expect(ok, "*entql.Edge, got %T", expr.Args[0])
+			return e.evalEdge(edge.Name, expr.Args[1:]...)
+		}
+	}
+	panic("invalid")
+}
+
+// evalBinary evaluates binary expressions.
+func (e *state) evalBinary(expr *entql.BinaryExpr) *sql.Predicate {
+	switch expr.Op {
+	case entql.OpOr:
+		return sql.Or(e.evalExpr(expr.X), e.evalExpr(expr.Y))
+	case entql.OpAnd:
+		return sql.And(e.evalExpr(expr.X), e.evalExpr(expr.Y))
+	case entql.OpEQ, entql.OpNEQ:
+		if expr.Y == (*entql.Value)(nil) {
+			f, ok := expr.X.(*entql.Field)
+			expect(ok, "*entql.Field, got %T", expr.Y)
+			return nullFunc[expr.Op](e.field(f))
+		}
+		fallthrough
+	default:
+		field, ok := expr.X.(*entql.Field)
+		expect(ok, "expr.X to be *entql.Field (got %T)", expr.X)
+		_, ok = expr.Y.(*entql.Field)
+		if !ok {
+			_, ok = expr.Y.(*entql.Value)
+		}
+		expect(ok, "expr.Y to be *entql.Field or *entql.Value (got %T)", expr.X)
+		switch x := expr.Y.(type) {
+		case *entql.Field:
+			return sql.ColumnsOp(e.field(field), e.field(x), binary[expr.Op])
+		case *entql.Value:
+			c := e.field(field)
+			return sql.P(func(b *sql.Builder) {
+				b.Ident(c).WriteOp(binary[expr.Op])
+				args(b, x)
+			})
+		default:
+			panic("unreachable")
+		}
+	}
+}
+
+// evalEdge evaluates has-edge and has-edge-with calls.
+func (e *state) evalEdge(name string, exprs ...entql.Expr) *sql.Predicate {
+	edge, ok := e.context.Edges[name]
+	expect(ok, "edge %q was not found for node %q", name, e.context.Type)
+	var fromC, toC string
+	switch {
+	case edge.To.ID != nil:
+		toC = edge.To.ID.Column
+	// Edge-owner points to its edge schema.
+	case edge.To.CompositeID != nil && !edge.Spec.Inverse:
+		toC = edge.To.CompositeID[0].Column
+	// Edge-backref points to its edge schema.
+	case edge.To.CompositeID != nil && edge.Spec.Inverse:
+		toC = edge.To.CompositeID[1].Column
+	default:
+		panic(evalError{fmt.Sprintf("expect id definition for edge %q", name)})
+	}
+	switch {
+	case e.context.ID != nil:
+		fromC = e.context.ID.Column
+	case e.context.CompositeID != nil && (edge.Spec.Rel == M2O || (edge.Spec.Rel == O2O && edge.Spec.Inverse)):
+		// An edge-schema with a composite id can query
+		// only edges that it owns (holds the foreign-key).
+	default:
+		panic(evalError{fmt.Sprintf("unexpected edge-query from an edge-schema %q", e.context.Type)})
+	}
+	step := NewStep(
+		From(e.context.Table, fromC),
+		To(edge.To.Table, toC),
+		Edge(edge.Spec.Rel, edge.Spec.Inverse, edge.Spec.Table, edge.Spec.Columns...),
+	)
+	selector := e.selector.Clone().SetP(nil)
+	selector.SetTotal(e.Total())
+	if len(exprs) == 0 {
+		HasNeighbors(selector, step)
+		return selector.P()
+	}
+	HasNeighborsWith(selector, step, func(s *sql.Selector) {
+		for _, expr := range exprs {
+			if cx, ok := expr.(*entql.CallExpr); ok && cx.Func == FuncSelector {
+				expect(len(cx.Args) == 1, "invalid number of arguments for %s", FuncSelector)
+				wrapped, ok := cx.Args[0].(wrappedFunc)
+				expect(ok, "invalid argument for %s: %T", FuncSelector, cx.Args[0])
+				wrapped.Func(s)
+			} else {
+				p, err := evalExpr(edge.To, s, expr)
+				expect(err == nil, "edge evaluation failed for %s->%s: %s", e.context.Type, name, err)
+				s.Where(p)
+			}
+		}
+	})
+	return selector.P()
+}
+
+func (e *state) field(f *entql.Field) string {
+	_, ok := e.context.Fields[f.Name]
+	expect(ok || e.context.ID.Column == f.Name, "field %q was not found for node %q", f.Name, e.context.Type)
+	return e.selector.C(f.Name)
+}
+
+func args(b *sql.Builder, v *entql.Value) {
+	vs, ok := v.V.([]any)
+	if !ok {
+		b.Arg(v.V)
+		return
+	}
+	b.WriteByte('(').Args(vs...).WriteByte(')')
+}
+
+// expect panics if the condition is false.
+func expect(cond bool, msg string, args ...any) {
+	if !cond {
+		panic(evalError{fmt.Sprintf("expect "+msg, args...)})
+	}
+}
+
+type evalError struct {
+	msg string
+}
+
+func (p evalError) Error() string {
+	return fmt.Sprintf("sqlgraph: %s", p.msg)
+}
+
+func catch(err *error) {
+	if e := recover(); e != nil {
+		xerr, ok := e.(evalError)
+		if !ok {
+			panic(e)
+		}
+		*err = xerr
+	}
+}

dialect/sql/sqlgraph/entql_test.go

@@ -0,0 +1,171 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sqlgraph
+
+import (
+	"strconv"
+	"testing"
+
+	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/sql"
+	"entgo.io/ent/entql"
+	"entgo.io/ent/schema/field"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGraph_AddE(t *testing.T) {
+	g := &Schema{
+		Nodes: []*Node{{Type: "user"}, {Type: "pet"}},
+	}
+	err := g.AddE("pets", &EdgeSpec{Rel: O2M}, "user", "pet")
+	assert.NoError(t, err)
+	err = g.AddE("owner", &EdgeSpec{Rel: O2M}, "pet", "user")
+	assert.NoError(t, err)
+	err = g.AddE("groups", &EdgeSpec{Rel: M2M}, "pet", "groups")
+	assert.Error(t, err)
+}
+
+func TestGraph_EvalP(t *testing.T) {
+	g := &Schema{
+		Nodes: []*Node{
+			{
+				Type: "user",
+				NodeSpec: NodeSpec{
+					Table: "users",
+					ID:    &FieldSpec{Column: "uid"},
+				},
+				Fields: map[string]*FieldSpec{
+					"name": {Column: "name", Type: field.TypeString},
+					"last": {Column: "last", Type: field.TypeString},
+				},
+			},
+			{
+				Type: "pet",
+				NodeSpec: NodeSpec{
+					Table: "pets",
+					ID:    &FieldSpec{Column: "pid"},
+				},
+				Fields: map[string]*FieldSpec{
+					"name": {Column: "name", Type: field.TypeString},
+				},
+			},
+			{
+				Type: "group",
+				NodeSpec: NodeSpec{
+					Table: "groups",
+					ID:    &FieldSpec{Column: "gid"},
+				},
+				Fields: map[string]*FieldSpec{
+					"name": {Column: "name", Type: field.TypeString},
+				},
+			},
+		},
+	}
+	err := g.AddE("pets", &EdgeSpec{Rel: O2M, Table: "pets", Columns: []string{"owner_id"}}, "user", "pet")
+	require.NoError(t, err)
+	err = g.AddE("owner", &EdgeSpec{Rel: M2O, Inverse: true, Table: "pets", Columns: []string{"owner_id"}}, "pet", "user")
+	require.NoError(t, err)
+	err = g.AddE("groups", &EdgeSpec{Rel: M2M, Table: "user_groups", Columns: []string{"user_id", "group_id"}}, "user", "group")
+	require.NoError(t, err)
+	err = g.AddE("users", &EdgeSpec{Rel: M2M, Inverse: true, Table: "user_groups", Columns: []string{"user_id", "group_id"}}, "group", "user")
+	require.NoError(t, err)
+
+	tests := []struct {
+		s         *sql.Selector
+		p         entql.P
+		wantQuery string
+		wantArgs  []any
+		wantErr   bool
+	}{
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")),
+			p:         entql.FieldHasPrefix("name", "a"),
+			wantQuery: `SELECT * FROM "users" WHERE "users"."name" LIKE $1`,
+			wantArgs:  []any{"a%"},
+		},
+		{
+			s: sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")).
+				Where(sql.EQ("age", 1)),
+			p:         entql.FieldHasPrefix("name", "a"),
+			wantQuery: `SELECT * FROM "users" WHERE "age" = $1 AND "users"."name" LIKE $2`,
+			wantArgs:  []any{1, "a%"},
+		},
+		{
+			s: sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")).
+				Where(sql.EQ("age", 1)),
+			p:         entql.FieldHasPrefix("name", "a"),
+			wantQuery: `SELECT * FROM "users" WHERE "age" = $1 AND "users"."name" LIKE $2`,
+			wantArgs:  []any{1, "a%"},
+		},
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")),
+			p:         entql.EQ(entql.F("name"), entql.F("last")),
+			wantQuery: `SELECT * FROM "users" WHERE "users"."name" = "users"."last"`,
+		},
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")),
+			p:         entql.EQ(entql.F("name"), entql.F("last")),
+			wantQuery: `SELECT * FROM "users" WHERE "users"."name" = "users"."last"`,
+		},
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")),
+			p:         entql.And(entql.FieldNil("name"), entql.FieldNotNil("last")),
+			wantQuery: `SELECT * FROM "users" WHERE "users"."name" IS NULL AND "users"."last" IS NOT NULL`,
+		},
+		{
+			s: sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")).
+				Where(sql.EQ("foo", "bar")),
+			p:         entql.Or(entql.FieldEQ("name", "foo"), entql.FieldEQ("name", "baz")),
+			wantQuery: `SELECT * FROM "users" WHERE "foo" = $1 AND ("users"."name" = $2 OR "users"."name" = $3)`,
+			wantArgs:  []any{"bar", "foo", "baz"},
+		},
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")),
+			p:         entql.HasEdge("pets"),
+			wantQuery: `SELECT * FROM "users" WHERE EXISTS (SELECT "pets"."owner_id" FROM "pets" WHERE "users"."uid" = "pets"."owner_id")`,
+		},
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")),
+			p:         entql.HasEdge("groups"),
+			wantQuery: `SELECT * FROM "users" WHERE "users"."uid" IN (SELECT "user_groups"."user_id" FROM "user_groups")`,
+		},
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")),
+			p:         entql.HasEdgeWith("pets", entql.Or(entql.FieldEQ("name", "pedro"), entql.FieldEQ("name", "xabi"))),
+			wantQuery: `SELECT * FROM "users" WHERE EXISTS (SELECT "pets"."owner_id" FROM "pets" WHERE "users"."uid" = "pets"."owner_id" AND ("pets"."name" = $1 OR "pets"."name" = $2))`,
+			wantArgs:  []any{"pedro", "xabi"},
+		},
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")).Where(sql.EQ("active", true)),
+			p:         entql.HasEdgeWith("groups", entql.Or(entql.FieldEQ("name", "GitHub"), entql.FieldEQ("name", "GitLab"))),
+			wantQuery: `SELECT * FROM "users" WHERE "active" AND "users"."uid" IN (SELECT "user_groups"."user_id" FROM "user_groups" JOIN "groups" AS "t1" ON "user_groups"."group_id" = "t1"."gid" WHERE "t1"."name" = $1 OR "t1"."name" = $2)`,
+			wantArgs:  []any{"GitHub", "GitLab"},
+		},
+		{
+			s:         sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")).Where(sql.EQ("active", true)),
+			p:         entql.And(entql.HasEdge("pets"), entql.HasEdge("groups"), entql.EQ(entql.F("name"), entql.F("uid"))),
+			wantQuery: `SELECT * FROM "users" WHERE "active" AND (EXISTS (SELECT "pets"."owner_id" FROM "pets" WHERE "users"."uid" = "pets"."owner_id") AND "users"."uid" IN (SELECT "user_groups"."user_id" FROM "user_groups") AND "users"."name" = "users"."uid")`,
+		},
+		{
+			s: sql.Dialect(dialect.Postgres).Select().From(sql.Table("users")).Where(sql.EQ("active", true)),
+			p: entql.HasEdgeWith("pets", entql.FieldEQ("name", "pedro"), WrapFunc(func(s *sql.Selector) {
+				s.Where(sql.EQ("owner_id", 10))
+			})),
+			wantQuery: `SELECT * FROM "users" WHERE "active" AND EXISTS (SELECT "pets"."owner_id" FROM "pets" WHERE ("users"."uid" = "pets"."owner_id" AND "pets"."name" = $1) AND "owner_id" = $2)`,
+			wantArgs:  []any{"pedro", 10},
+		},
+	}
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			err = g.EvalP("user", tt.p, tt.s)
+			require.Equal(t, tt.wantErr, err != nil, err)
+			query, args := tt.s.Query()
+			require.Equal(t, tt.wantQuery, query)
+			require.Equal(t, tt.wantArgs, args)
+		})
+	}
+}

dialect/sql/sqlgraph/errors.go

@@ -0,0 +1,74 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sqlgraph
+
+import (
+	"errors"
+	"strings"
+)
+
+// IsConstraintError returns true if the error resulted from a database constraint violation.
+func IsConstraintError(err error) bool {
+	var e *ConstraintError
+	return errors.As(err, &e) ||
+		IsUniqueConstraintError(err) ||
+		IsForeignKeyConstraintError(err) ||
+		IsCheckConstraintError(err)
+}
+
+// IsUniqueConstraintError reports if the error resulted from a DB uniqueness constraint violation.
+// e.g. duplicate value in unique index.
+func IsUniqueConstraintError(err error) bool {
+	if err == nil {
+		return false
+	}
+	for _, s := range []string{
+		"Error 1062",                 // MySQL
+		"violates unique constraint", // Postgres
+		"UNIQUE constraint failed",   // SQLite
+	} {
+		if strings.Contains(err.Error(), s) {
+			return true
+		}
+	}
+	return false
+}
+
+// IsForeignKeyConstraintError reports if the error resulted from a database foreign-key constraint violation.
+// e.g. parent row does not exist.
+func IsForeignKeyConstraintError(err error) bool {
+	if err == nil {
+		return false
+	}
+	for _, s := range []string{
+		"Error 1451",                      // MySQL (Cannot delete or update a parent row).
+		"Error 1452",                      // MySQL (Cannot add or update a child row).
+		"violates foreign key constraint", // Postgres
+		"FOREIGN KEY constraint failed",   // SQLite
+	} {
+		if strings.Contains(err.Error(), s) {
+			return true
+		}
+	}
+	return false
+}
+
+// IsCheckConstraintError reports if the error resulted from a database check constraint violation.
+// e.g. a value does not satisfy a check condition.
+func IsCheckConstraintError(err error) bool {
+	if err == nil {
+		return false
+	}
+	for _, s := range []string{
+		"Error 3819",                // MySQL
+		"violates check constraint", // Postgres
+		"CHECK constraint failed",   // SQLite
+	} {
+		if strings.Contains(err.Error(), s) {
+			return true
+		}
+	}
+	return false
+}