gitee.com/go-genie/sqlx@v1.0.3/connectors/postgresql/schemas.go

package postgresql

import (
	"fmt"
	"github.com/pkg/errors"
	"regexp"
	"strings"

	"gitee.com/go-genie/sqlx"
	"gitee.com/go-genie/sqlx/builder"
)

func toInterfaces(list ...string) []interface{} {
	s := make([]interface{}, len(list))
	for i, v := range list {
		s[i] = v
	}
	return s
}

var reUsing = regexp.MustCompile(`USING ([^ ]+)`)

const COMMAND_MIGRATE = "migrate"
const COMMAND_GENERATE = "generate"

func dbFromInformationSchema(db sqlx.DBExecutor, usage string) (*sqlx.Database, error) {
	d := db.D()

	dbName := d.Name
	dbSchema := d.Schema
	tableNames := d.Tables.TableNames()

	d = sqlx.NewDatabase(dbName).WithSchema(dbSchema)

	tableColumnSchema := SchemaDatabase.T(&ColumnSchema{}).WithSchema("information_schema")
	columnSchemaList := make([]ColumnSchema, 0)

	tableSchema := "public"
	if d.Schema != "" {
		tableSchema = d.Schema
	}

	tableCondition := builder.EmptyCond().And(tableColumnSchema.F("TABLE_SCHEMA").Eq(tableSchema))
	switch usage {
	case COMMAND_GENERATE:
		tableCondition = tableCondition.And(tableColumnSchema.F("TABLE_NAME").RightLike("t_")).And(tableColumnSchema.F("TABLE_NAME").Neq("t_sql_meta_enum"))
	case COMMAND_MIGRATE:
		tableCondition = tableCondition.And(tableColumnSchema.F("TABLE_NAME").In(toInterfaces(tableNames...)...))
	default:
		return nil, errors.New("unknown usage")
	}

	stmt := builder.Select(tableColumnSchema.Columns.Clone()).From(tableColumnSchema,
		builder.Where(
			tableCondition,
		),
	)

	err := db.QueryExprAndScan(stmt, &columnSchemaList)
	if err != nil {
		return nil, err
	}

	for i := range columnSchemaList {
		columnSchema := columnSchemaList[i]

		table := d.Table(columnSchema.TABLE_NAME)
		if table == nil {
			table = builder.T(columnSchema.TABLE_NAME)
			d.AddTable(table)
		}

		table.AddCol(colFromColumnSchema(&columnSchema))
	}

	if tableColumnSchema.Columns.Len() != 0 {
		tableIndexSchema := SchemaDatabase.T(&IndexSchema{})

		indexCondition := builder.EmptyCond().And(tableIndexSchema.F("TABLE_SCHEMA").Eq(tableSchema))
		switch usage {
		case COMMAND_GENERATE:
			indexCondition = indexCondition.And(tableIndexSchema.F("TABLE_NAME").RightLike("t_")).And(tableIndexSchema.F("TABLE_NAME").Neq("t_sql_meta_enum"))
		case COMMAND_MIGRATE:
			indexCondition = indexCondition.And(tableIndexSchema.F("TABLE_NAME").In(toInterfaces(tableNames...)...))
		default:
			return nil, errors.New("unknown usage")
		}

		indexList := make([]IndexSchema, 0)

		err = db.QueryExprAndScan(
			builder.Select(tableIndexSchema.Columns.Clone()).
				From(
					tableIndexSchema,
					builder.Where(
						indexCondition,
					),
				),
			&indexList,
		)

		if err != nil {
			return nil, err
		}

		for _, indexSchema := range indexList {
			table := d.Table(indexSchema.TABLE_NAME)

			key := &builder.Key{}
			switch usage {
			case COMMAND_GENERATE:
				key.Name = strings.ToLower(indexSchema.INDEX_NAME)
			case COMMAND_MIGRATE:
				key.Name = strings.ToLower(indexSchema.INDEX_NAME[len(table.Name)+1:])
			}

			key.Method = strings.ToUpper(reUsing.FindString(indexSchema.INDEX_DEF)[6:])
			key.IsUnique = strings.Contains(indexSchema.INDEX_DEF, "UNIQUE")

			key.Def.Expr = strings.Replace(strings.TrimSpace(reUsing.Split(indexSchema.INDEX_DEF, 2)[1]), ", ", ",", -1)

			table.AddKey(key)
		}
	}

	return d, nil
}

var SchemaDatabase = sqlx.NewDatabase("INFORMATION_SCHEMA")

func init() {
	SchemaDatabase.Register(&ColumnSchema{})
	SchemaDatabase.Register(&IndexSchema{})
}

func colFromColumnSchema(columnSchema *ColumnSchema) *builder.Column {
	col := builder.Col(columnSchema.COLUMN_NAME)

	defaultValue := columnSchema.COLUMN_DEFAULT

	if defaultValue != "" {
		col.AutoIncrement = strings.HasSuffix(columnSchema.COLUMN_DEFAULT, "_seq'::regclass)")

		if !col.AutoIncrement {
			if !strings.Contains(defaultValue, "'::") && '0' <= defaultValue[0] && defaultValue[0] <= '9' {
				defaultValue = fmt.Sprintf("'%s'::integer", defaultValue)
			}
			col.Default = &defaultValue
		}
	}

	dataType := columnSchema.DATA_TYPE

	if col.AutoIncrement {
		if strings.HasPrefix(dataType, "big") {
			dataType = "bigserial"
		} else {
			dataType = "serial"
		}
	}

	col.DataType = dataType

	// numeric type
	if columnSchema.NUMERIC_PRECISION > 0 {
		col.Length = columnSchema.NUMERIC_PRECISION
		col.Decimal = columnSchema.NUMERIC_SCALE
	} else {
		col.Length = columnSchema.CHARACTER_MAXIMUM_LENGTH
	}

	if columnSchema.IS_NULLABLE == "YES" {
		col.Null = true
	}

	return col
}

type ColumnSchema struct {
	TABLE_SCHEMA             string `db:"table_schema"`
	TABLE_NAME               string `db:"table_name"`
	COLUMN_NAME              string `db:"column_name"`
	DATA_TYPE                string `db:"data_type"`
	IS_NULLABLE              string `db:"is_nullable"`
	COLUMN_DEFAULT           string `db:"column_default"`
	CHARACTER_MAXIMUM_LENGTH uint64 `db:"character_maximum_length"`
	NUMERIC_PRECISION        uint64 `db:"numeric_precision"`
	NUMERIC_SCALE            uint64 `db:"numeric_scale"`
}

func (ColumnSchema) TableName() string {
	return "columns"
}

type IndexSchema struct {
	TABLE_SCHEMA string `db:"schemaname"`
	TABLE_NAME   string `db:"tablename"`
	INDEX_NAME   string `db:"indexname"`
	INDEX_DEF    string `db:"indexdef"`
}

func (IndexSchema) TableName() string {
	return "pg_indexes"
}