gtsocial-umbx

query_table_create.go (9145B)

---

 package bun
 
 import (
 	"context"
 	"database/sql"
 	"fmt"
 	"sort"
 	"strconv"
 	"strings"
 
 	"github.com/uptrace/bun/dialect/feature"
 	"github.com/uptrace/bun/dialect/sqltype"
 	"github.com/uptrace/bun/internal"
 	"github.com/uptrace/bun/schema"
 )
 
 type CreateTableQuery struct {
 	baseQuery
 
 	temp        bool
 	ifNotExists bool
 
 	// varchar changes the default length for VARCHAR columns.
 	// Because some dialects require that length is always specified for VARCHAR type,
 	// we will use the exact user-defined type if length is set explicitly, as in `bun:",type:varchar(5)"`,
 	// but assume the new default length when it's omitted, e.g. `bun:",type:varchar"`.
 	varchar int
 
 	fks         []schema.QueryWithArgs
 	partitionBy schema.QueryWithArgs
 	tablespace  schema.QueryWithArgs
 }
 
 var _ Query = (*CreateTableQuery)(nil)
 
 func NewCreateTableQuery(db *DB) *CreateTableQuery {
 	q := &CreateTableQuery{
 		baseQuery: baseQuery{
 			db:   db,
 			conn: db.DB,
 		},
 		varchar: db.Dialect().DefaultVarcharLen(),
 	}
 	return q
 }
 
 func (q *CreateTableQuery) Conn(db IConn) *CreateTableQuery {
 	q.setConn(db)
 	return q
 }
 
 func (q *CreateTableQuery) Model(model interface{}) *CreateTableQuery {
 	q.setModel(model)
 	return q
 }
 
 func (q *CreateTableQuery) Err(err error) *CreateTableQuery {
 	q.setErr(err)
 	return q
 }
 
 // ------------------------------------------------------------------------------
 
 func (q *CreateTableQuery) Table(tables ...string) *CreateTableQuery {
 	for _, table := range tables {
 		q.addTable(schema.UnsafeIdent(table))
 	}
 	return q
 }
 
 func (q *CreateTableQuery) TableExpr(query string, args ...interface{}) *CreateTableQuery {
 	q.addTable(schema.SafeQuery(query, args))
 	return q
 }
 
 func (q *CreateTableQuery) ModelTableExpr(query string, args ...interface{}) *CreateTableQuery {
 	q.modelTableName = schema.SafeQuery(query, args)
 	return q
 }
 
 func (q *CreateTableQuery) ColumnExpr(query string, args ...interface{}) *CreateTableQuery {
 	q.addColumn(schema.SafeQuery(query, args))
 	return q
 }
 
 // ------------------------------------------------------------------------------
 
 func (q *CreateTableQuery) Temp() *CreateTableQuery {
 	q.temp = true
 	return q
 }
 
 func (q *CreateTableQuery) IfNotExists() *CreateTableQuery {
 	q.ifNotExists = true
 	return q
 }
 
 // Varchar sets default length for VARCHAR columns.
 func (q *CreateTableQuery) Varchar(n int) *CreateTableQuery {
 	if n <= 0 {
 		q.setErr(fmt.Errorf("bun: illegal VARCHAR length: %d", n))
 		return q
 	}
 	q.varchar = n
 	return q
 }
 
 func (q *CreateTableQuery) ForeignKey(query string, args ...interface{}) *CreateTableQuery {
 	q.fks = append(q.fks, schema.SafeQuery(query, args))
 	return q
 }
 
 func (q *CreateTableQuery) PartitionBy(query string, args ...interface{}) *CreateTableQuery {
 	q.partitionBy = schema.SafeQuery(query, args)
 	return q
 }
 
 func (q *CreateTableQuery) TableSpace(tablespace string) *CreateTableQuery {
 	q.tablespace = schema.UnsafeIdent(tablespace)
 	return q
 }
 
 func (q *CreateTableQuery) WithForeignKeys() *CreateTableQuery {
 	for _, relation := range q.tableModel.Table().Relations {
 		if relation.Type == schema.ManyToManyRelation ||
 			relation.Type == schema.HasManyRelation {
 			continue
 		}
 
 		q = q.ForeignKey("(?) REFERENCES ? (?) ? ?",
 			Safe(appendColumns(nil, "", relation.BaseFields)),
 			relation.JoinTable.SQLName,
 			Safe(appendColumns(nil, "", relation.JoinFields)),
 			Safe(relation.OnUpdate),
 			Safe(relation.OnDelete),
 		)
 	}
 	return q
 }
 
 // ------------------------------------------------------------------------------
 
 func (q *CreateTableQuery) Operation() string {
 	return "CREATE TABLE"
 }
 
 func (q *CreateTableQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err
 	}
 	if q.table == nil {
 		return nil, errNilModel
 	}
 
 	b = append(b, "CREATE "...)
 	if q.temp {
 		b = append(b, "TEMP "...)
 	}
 	b = append(b, "TABLE "...)
 	if q.ifNotExists && fmter.Dialect().Features().Has(feature.TableNotExists) {
 		b = append(b, "IF NOT EXISTS "...)
 	}
 	b, err = q.appendFirstTable(fmter, b)
 	if err != nil {
 		return nil, err
 	}
 
 	b = append(b, " ("...)
 
 	for i, field := range q.table.Fields {
 		if i > 0 {
 			b = append(b, ", "...)
 		}
 
 		b = append(b, field.SQLName...)
 		b = append(b, " "...)
 		b = q.appendSQLType(b, field)
 		if field.NotNull {
 			b = append(b, " NOT NULL"...)
 		}
 		if field.AutoIncrement {
 			switch {
 			case fmter.Dialect().Features().Has(feature.AutoIncrement):
 				b = append(b, " AUTO_INCREMENT"...)
 			case fmter.Dialect().Features().Has(feature.Identity):
 				b = append(b, " IDENTITY"...)
 			}
 		}
 		if field.Identity {
 			if fmter.Dialect().Features().Has(feature.GeneratedIdentity) {
 				b = append(b, " GENERATED BY DEFAULT AS IDENTITY"...)
 			}
 		}
 		if field.SQLDefault != "" {
 			b = append(b, " DEFAULT "...)
 			b = append(b, field.SQLDefault...)
 		}
 	}
 
 	for i, col := range q.columns {
 		// Only pre-pend the comma if we are on subsequent iterations, or if there were fields/columns appended before
 		// this. This way if we are only appending custom column expressions we will not produce a syntax error with a
 		// leading comma.
 		if i > 0 || len(q.table.Fields) > 0 {
 			b = append(b, ", "...)
 		}
 		b, err = col.AppendQuery(fmter, b)
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	b = q.appendPKConstraint(b, q.table.PKs)
 	b = q.appendUniqueConstraints(fmter, b)
 	b, err = q.appendFKConstraints(fmter, b)
 	if err != nil {
 		return nil, err
 	}
 
 	b = append(b, ")"...)
 
 	if !q.partitionBy.IsZero() {
 		b = append(b, " PARTITION BY "...)
 		b, err = q.partitionBy.AppendQuery(fmter, b)
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	if !q.tablespace.IsZero() {
 		b = append(b, " TABLESPACE "...)
 		b, err = q.tablespace.AppendQuery(fmter, b)
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	return b, nil
 }
 
 func (q *CreateTableQuery) appendSQLType(b []byte, field *schema.Field) []byte {
 	// Most of the time these two will match, but for the cases where DiscoveredSQLType is dialect-specific,
 	// e.g. pgdialect would change sqltype.SmallInt to pgTypeSmallSerial for columns that have `bun:",autoincrement"`
 	if !strings.EqualFold(field.CreateTableSQLType, field.DiscoveredSQLType) {
 		return append(b, field.CreateTableSQLType...)
 	}
 
 	// For all common SQL types except VARCHAR, both UserDefinedSQLType and DiscoveredSQLType specify the correct type,
 	// and we needn't modify it. For VARCHAR columns, we will stop to check if a valid length has been set in .Varchar(int).
 	if !strings.EqualFold(field.CreateTableSQLType, sqltype.VarChar) || q.varchar <= 0 {
 		return append(b, field.CreateTableSQLType...)
 	}
 
 	b = append(b, sqltype.VarChar...)
 	b = append(b, "("...)
 	b = strconv.AppendInt(b, int64(q.varchar), 10)
 	b = append(b, ")"...)
 	return b
 }
 
 func (q *CreateTableQuery) appendUniqueConstraints(fmter schema.Formatter, b []byte) []byte {
 	unique := q.table.Unique
 
 	keys := make([]string, 0, len(unique))
 	for key := range unique {
 		keys = append(keys, key)
 	}
 	sort.Strings(keys)
 
 	for _, key := range keys {
 		if key == "" {
 			for _, field := range unique[key] {
 				b = q.appendUniqueConstraint(fmter, b, key, field)
 			}
 			continue
 		}
 		b = q.appendUniqueConstraint(fmter, b, key, unique[key]...)
 	}
 
 	return b
 }
 
 func (q *CreateTableQuery) appendUniqueConstraint(
 	fmter schema.Formatter, b []byte, name string, fields ...*schema.Field,
 ) []byte {
 	if name != "" {
 		b = append(b, ", CONSTRAINT "...)
 		b = fmter.AppendIdent(b, name)
 	} else {
 		b = append(b, ","...)
 	}
 	b = append(b, " UNIQUE ("...)
 	b = appendColumns(b, "", fields)
 	b = append(b, ")"...)
 	return b
 }
 
 func (q *CreateTableQuery) appendFKConstraints(
 	fmter schema.Formatter, b []byte,
 ) (_ []byte, err error) {
 	for _, fk := range q.fks {
 		b = append(b, ", FOREIGN KEY "...)
 		b, err = fk.AppendQuery(fmter, b)
 		if err != nil {
 			return nil, err
 		}
 	}
 	return b, nil
 }
 
 func (q *CreateTableQuery) appendPKConstraint(b []byte, pks []*schema.Field) []byte {
 	if len(pks) == 0 {
 		return b
 	}
 
 	b = append(b, ", PRIMARY KEY ("...)
 	b = appendColumns(b, "", pks)
 	b = append(b, ")"...)
 	return b
 }
 
 // ------------------------------------------------------------------------------
 
 func (q *CreateTableQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result, error) {
 	if err := q.beforeCreateTableHook(ctx); err != nil {
 		return nil, err
 	}
 
 	queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
 	if err != nil {
 		return nil, err
 	}
 
 	query := internal.String(queryBytes)
 
 	res, err := q.exec(ctx, q, query)
 	if err != nil {
 		return nil, err
 	}
 
 	if q.table != nil {
 		if err := q.afterCreateTableHook(ctx); err != nil {
 			return nil, err
 		}
 	}
 
 	return res, nil
 }
 
 func (q *CreateTableQuery) beforeCreateTableHook(ctx context.Context) error {
 	if hook, ok := q.table.ZeroIface.(BeforeCreateTableHook); ok {
 		if err := hook.BeforeCreateTable(ctx, q); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
 func (q *CreateTableQuery) afterCreateTableHook(ctx context.Context) error {
 	if hook, ok := q.table.ZeroIface.(AfterCreateTableHook); ok {
 		if err := hook.AfterCreateTable(ctx, q); err != nil {
 			return err
 		}
 	}
 	return nil
 }