add schema

Author: siddontang

schema/schema.go

@@ -0,0 +1,153 @@
+// 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"
+	"github.com/siddontang/go-mysql/client"
+)
+
+type TableColumn struct {
+	Name   string
+	Type   string
+	IsAuto bool
+}
+
+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) AddColumn(name string, columnType string, extra string) {
+	index := len(ta.Columns)
+	ta.Columns = append(ta.Columns, TableColumn{Name: name})
+
+	ta.Columns[index].Type = columnType
+	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 *client.Conn, 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 *client.Conn) error {
+	r, err := conn.Execute(fmt.Sprintf("describe %s.%s", ta.Schema, ta.Name))
+	if err != nil {
+		return 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 *client.Conn) error {
+	r, err := conn.Execute(fmt.Sprintf("show index from %s.%s", ta.Schema, ta.Name))
+	if err != nil {
+		return 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
+}

schema/schema_test.go

@@ -0,0 +1,63 @@
+package schema
+
+import (
+	"flag"
+	"fmt"
+	"github.com/siddontang/go-mysql/client"
+	. "gopkg.in/check.v1"
+	"testing"
+)
+
+// use docker mysql for test
+var host = flag.String("host", "127.0.0.1", "MySQL host")
+
+func Test(t *testing.T) {
+	TestingT(t)
+}
+
+type schemaTestSuite struct {
+	conn *client.Conn
+}
+
+var _ = Suite(&schemaTestSuite{})
+
+func (s *schemaTestSuite) SetUpSuite(c *C) {
+	var err error
+	s.conn, err = client.Connect(fmt.Sprintf("%s:%d", *host, 3306), "root", "", "test")
+	c.Assert(err, IsNil)
+}
+
+func (s *schemaTestSuite) TearDownSuite(c *C) {
+	if s.conn != nil {
+		s.conn.Close()
+	}
+}
+
+func (s *schemaTestSuite) TestSchema(c *C) {
+	str := `
+        CREATE TABLE IF NOT EXISTS schema_test (
+            id INT,
+            id1 INT,
+            id2 INT,
+            name VARCHAR(256),
+            e ENUM("a", "b", "c"),
+            f FLOAT,
+            PRIMARY KEY(id2, id),
+            UNIQUE (id1),
+            INDEX name_idx (name)
+        ) ENGINE = INNODB;
+    `
+
+	_, err := s.conn.Execute(str)
+	c.Assert(err, IsNil)
+
+	ta, err := NewTable(s.conn, "test", "schema_test")
+	c.Assert(err, IsNil)
+
+	c.Assert(ta.Columns, HasLen, 6)
+	c.Assert(ta.Indexes, HasLen, 3)
+	c.Assert(ta.PKColumns, DeepEquals, []int{2, 0})
+	c.Assert(ta.Indexes[0].Columns, HasLen, 2)
+	c.Assert(ta.Indexes[0].Name, Equals, "PRIMARY")
+	c.Assert(ta.Indexes[2].Name, Equals, "name_idx")
+}

vitess_license

@@ -0,0 +1,28 @@
+Copyright 2012, Google Inc.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+    * Redistributions of source code must retain the above copyright
+notice, this list of conditions and the following disclaimer.
+    * Redistributions in binary form must reproduce the above
+copyright notice, this list of conditions and the following disclaimer
+in the documentation and/or other materials provided with the
+distribution.
+    * Neither the name of Google Inc. nor the names of its
+contributors may be used to endorse or promote products derived from
+this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.