gh-ost/vendor/github.com/siddontang/go-mysql/schema/schema.go

// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package schema

import (
	"fmt"
	"strings"

	"github.com/juju/errors"
	"github.com/siddontang/go-mysql/mysql"
)

const (
	TYPE_NUMBER    = iota + 1 // tinyint, smallint, mediumint, int, bigint, year
	TYPE_FLOAT                // float, double
	TYPE_ENUM                 // enum
	TYPE_SET                  // set
	TYPE_STRING               // other
	TYPE_DATETIME             // datetime
	TYPE_TIMESTAMP            // timestamp
	TYPE_DATE                 // date
	TYPE_TIME                 // time
	TYPE_BIT                  // bit
	TYPE_JSON                 // json
)

type TableColumn struct {
	Name       string
	Type       int
	IsAuto     bool
	EnumValues []string
	SetValues  []string
}

type Index struct {
	Name        string
	Columns     []string
	Cardinality []uint64
}

type Table struct {
	Schema string
	Name   string

	Columns   []TableColumn
	Indexes   []*Index
	PKColumns []int
}

func (ta *Table) String() string {
	return fmt.Sprintf("%s.%s", ta.Schema, ta.Name)
}

func (ta *Table) AddColumn(name string, columnType string, extra string) {
	index := len(ta.Columns)
	ta.Columns = append(ta.Columns, TableColumn{Name: name})

	if strings.Contains(columnType, "int") || strings.HasPrefix(columnType, "year") {
		ta.Columns[index].Type = TYPE_NUMBER
	} else if strings.HasPrefix(columnType, "float") ||
		strings.HasPrefix(columnType, "double") ||
		strings.HasPrefix(columnType, "decimal") {
		ta.Columns[index].Type = TYPE_FLOAT
	} else if strings.HasPrefix(columnType, "enum") {
		ta.Columns[index].Type = TYPE_ENUM
		ta.Columns[index].EnumValues = strings.Split(strings.Replace(
			strings.TrimSuffix(
				strings.TrimPrefix(
					columnType, "enum("),
				")"),
			"'", "", -1),
			",")
	} else if strings.HasPrefix(columnType, "set") {
		ta.Columns[index].Type = TYPE_SET
		ta.Columns[index].SetValues = strings.Split(strings.Replace(
			strings.TrimSuffix(
				strings.TrimPrefix(
					columnType, "set("),
				")"),
			"'", "", -1),
			",")
	} else if strings.HasPrefix(columnType, "datetime") {
		ta.Columns[index].Type = TYPE_DATETIME
	} else if strings.HasPrefix(columnType, "timestamp") {
		ta.Columns[index].Type = TYPE_TIMESTAMP
	} else if strings.HasPrefix(columnType, "time") {
		ta.Columns[index].Type = TYPE_TIME
	} else if "date" == columnType {
		ta.Columns[index].Type = TYPE_DATE
	} else if strings.HasPrefix(columnType, "bit") {
		ta.Columns[index].Type = TYPE_BIT
	} else if strings.HasPrefix(columnType, "json") {
		ta.Columns[index].Type = TYPE_JSON
	} else {
		ta.Columns[index].Type = TYPE_STRING
	}

	if extra == "auto_increment" {
		ta.Columns[index].IsAuto = true
	}
}

func (ta *Table) FindColumn(name string) int {
	for i, col := range ta.Columns {
		if col.Name == name {
			return i
		}
	}
	return -1
}

func (ta *Table) GetPKColumn(index int) *TableColumn {
	return &ta.Columns[ta.PKColumns[index]]
}

func (ta *Table) AddIndex(name string) (index *Index) {
	index = NewIndex(name)
	ta.Indexes = append(ta.Indexes, index)
	return index
}

func NewIndex(name string) *Index {
	return &Index{name, make([]string, 0, 8), make([]uint64, 0, 8)}
}

func (idx *Index) AddColumn(name string, cardinality uint64) {
	idx.Columns = append(idx.Columns, name)
	if cardinality == 0 {
		cardinality = uint64(len(idx.Cardinality) + 1)
	}
	idx.Cardinality = append(idx.Cardinality, cardinality)
}

func (idx *Index) FindColumn(name string) int {
	for i, colName := range idx.Columns {
		if name == colName {
			return i
		}
	}
	return -1
}

func NewTable(conn mysql.Executer, schema string, name string) (*Table, error) {
	ta := &Table{
		Schema:  schema,
		Name:    name,
		Columns: make([]TableColumn, 0, 16),
		Indexes: make([]*Index, 0, 8),
	}

	if err := ta.fetchColumns(conn); err != nil {
		return nil, err
	}

	if err := ta.fetchIndexes(conn); err != nil {
		return nil, err
	}

	return ta, nil
}

func (ta *Table) fetchColumns(conn mysql.Executer) error {
	r, err := conn.Execute(fmt.Sprintf("describe `%s`.`%s`", ta.Schema, ta.Name))
	if err != nil {
		return errors.Trace(err)
	}

	for i := 0; i < r.RowNumber(); i++ {
		name, _ := r.GetString(i, 0)
		colType, _ := r.GetString(i, 1)
		extra, _ := r.GetString(i, 5)

		ta.AddColumn(name, colType, extra)
	}

	return nil
}

func (ta *Table) fetchIndexes(conn mysql.Executer) error {
	r, err := conn.Execute(fmt.Sprintf("show index from `%s`.`%s`", ta.Schema, ta.Name))
	if err != nil {
		return errors.Trace(err)
	}
	var currentIndex *Index
	currentName := ""

	for i := 0; i < r.RowNumber(); i++ {
		indexName, _ := r.GetString(i, 2)
		if currentName != indexName {
			currentIndex = ta.AddIndex(indexName)
			currentName = indexName
		}
		cardinality, _ := r.GetUint(i, 6)
		colName, _ := r.GetString(i, 4)
		currentIndex.AddColumn(colName, cardinality)
	}

	if len(ta.Indexes) == 0 {
		return nil
	}

	pkIndex := ta.Indexes[0]
	if pkIndex.Name != "PRIMARY" {
		return nil
	}

	ta.PKColumns = make([]int, len(pkIndex.Columns))
	for i, pkCol := range pkIndex.Columns {
		ta.PKColumns[i] = ta.FindColumn(pkCol)
	}

	return nil
}