Fixed handling of GENERATED fields

Author: percona-csalguero

dbutils/dbutils.go

@@ -0,0 +1,67 @@
+package dbutils
+
+import (
+	"database/sql"
+	"fmt"
+	"net/url"
+
+	_ "github.com/go-sql-driver/mysql"
+	"github.com/pkg/errors"
+)
+
+/*
+EXPLAIN SELECT COUNT(*) FROM sakila.actor:
++----+-------------+-------+------------+-------+---------------+---------------------+---------+------+------+----------+-------------+
+| id | select_type | table | partitions | type  | possible_keys | key                 | key_len | ref  | rows | filtered | Extra       |
++----+-------------+-------+------------+-------+---------------+---------------------+---------+------+------+----------+-------------+
+|  1 | SIMPLE      | actor | NULL       | index | NULL          | idx_actor_last_name | 137     | NULL |  200 |   100.00 | Using index |
++----+-------------+-------+------------+-------+---------------+---------------------+---------+------+------+----------+-------------+
+*/
+type ExplainRow struct {
+	ID           int
+	SelectType   string
+	Table        string
+	Partitions   sql.NullString
+	Type         string
+	PossibleKeys sql.NullString
+	Key          sql.NullString
+	KeyLen       sql.NullInt64
+	Ref          sql.NullString
+	Rows         int64
+	Filtered     float64
+	Extra        sql.NullString
+}
+
+func GetApproxRowsCount(conn *sql.DB, schema, table string) (int64, error) {
+	query := fmt.Sprintf("EXPLAIN SELECT COUNT(*) FROM `%s`.`%s`",
+		url.QueryEscape(schema), url.QueryEscape(table))
+	var exp ExplainRow
+	err := conn.QueryRow(query).Scan(&exp.ID, &exp.SelectType, &exp.Table, &exp.Partitions, &exp.Type,
+		&exp.PossibleKeys, &exp.Key, &exp.KeyLen, &exp.Ref, &exp.Rows, &exp.Filtered, &exp.Extra)
+	if err != nil {
+		return 0, err
+	}
+	return exp.Rows, nil
+}
+
+// Explain will run MySQL EXPLAIN for a given query
+// If DB connection doesn't have a default schema, the query must have a fully qualified schema.table.
+func Explain(conn *sql.DB, query string) ([]ExplainRow, error) {
+	query = "EXPLAIN " + query
+	explain := []ExplainRow{}
+
+	rows, err := conn.Query(query)
+	if err != nil {
+		return nil, errors.Wrap(err, "cannot run explain")
+	}
+	for rows.Next() {
+		var exp ExplainRow
+		err := rows.Scan(&exp.ID, &exp.SelectType, &exp.Table, &exp.Partitions, &exp.Type,
+			&exp.PossibleKeys, &exp.Key, &exp.KeyLen, &exp.Ref, &exp.Rows, &exp.Filtered, &exp.Extra)
+		if err != nil {
+			return nil, errors.Wrap(err, "error reading explain rows")
+		}
+		explain = append(explain, exp)
+	}
+	return explain, nil
+}

docker-compose.yml

@@ -0,0 +1,13 @@
+version: '3'
+services:
+  mysql:
+    image: ${MYSQL_IMAGE:-mysql:5.7}
+    ports:
+      - ${MYSQL_HOST:-127.0.0.1}:${MYSQL_PORT:-3306}:3306
+    environment:
+      - MYSQL_ALLOW_EMPTY_PASSWORD=yes
+    # MariaDB >= 10.0.12 doesn't enable Performance Schema by default so we need to do it manually
+    # https://mariadb.com/kb/en/mariadb/performance-schema-overview/#activating-the-performance-schema
+    command: --performance-schema --secure-file-priv=""
+    volumes:
+      - ./testdata/schema/:/docker-entrypoint-initdb.d/:rw

internal/getters/getters.go

@@ -208,7 +208,7 @@ func (r *RandomTime) Value() interface{} {
 	h := rand.Int63n(24)
 	m := rand.Int63n(60)
 	s := rand.Int63n(60)
-	return fmt.Sprintf("%02d:%02d:%02d", h, m, s)
+	return fmt.Sprintf("'%02d:%02d:%02d'", h, m, s)
 }
 
 func (r *RandomTime) String() string {

main.go

@@ -273,6 +273,11 @@ func makeValueFuncs(conn *sql.DB, fields []tableparser.Field) insertValues {
 			strings.Contains(field.Extra, "auto_increment") {
 			continue
 		}
+
+		if field.GenerationExpression != "" {
+			continue
+		}
+
 		if field.Constraint != nil {
 			samples, err := getSamples(conn, field.Constraint.ReferencedTableSchema,
 				field.Constraint.ReferencedTableName,
@@ -285,6 +290,7 @@ func makeValueFuncs(conn *sql.DB, fields []tableparser.Field) insertValues {
 			values = append(values, getters.NewRandomSample(field.ColumnName, samples, field.IsNullable))
 			continue
 		}
+
 		maxValue := maxValues["bigint"]
 		if m, ok := maxValues[field.DataType]; ok {
 			maxValue = m
@@ -326,10 +332,16 @@ func getFieldNames(fields []tableparser.Field) []string {
 		if !isSupportedType(field.DataType) {
 			continue
 		}
+
 		if !field.IsNullable && field.ColumnKey == "PRI" &&
 			strings.Contains(field.Extra, "auto_increment") {
 			continue
 		}
+
+		if field.GenerationExpression != "" {
+			continue
+		}
+
 		fieldNames = append(fieldNames, backticks(field.ColumnName))
 	}
 	return fieldNames

tableparser/tableparser.go

@@ -13,14 +13,20 @@ type Table struct {
 	Schema      string
 	Name        string
 	Fields      []Field
-	Indexes     []Index
+	Indexes     map[string]Index
 	Constraints []Constraint
 	Triggers    []Trigger
 	//
 	conn *sql.DB
 }
 
 type Index struct {
+	Name   string
+	Unique bool
+	Fields []string
+}
+
+type IndexField struct {
 	NonUnique    bool
 	KeyName      string
 	SeqInIndex   int
@@ -119,28 +125,7 @@ func (t *Table) parse() error {
 	//                           |          |       +-------- extra info (unsigned, etc)
 	//                           |          |       |
 	re := regexp.MustCompile("^(.*?)(?:\\((.*?)\\)(.*))?$")
-	query := "SELECT `TABLE_CATALOG`," +
-		"`TABLE_SCHEMA`," +
-		"`TABLE_NAME`," +
-		"`COLUMN_NAME`," +
-		"`ORDINAL_POSITION`," +
-		"`COLUMN_DEFAULT`," +
-		"`IS_NULLABLE`," +
-		"`DATA_TYPE`," +
-		"`CHARACTER_MAXIMUM_LENGTH`," +
-		"`CHARACTER_OCTET_LENGTH`," +
-		"`NUMERIC_PRECISION`," +
-		"`NUMERIC_SCALE`," +
-		"`DATETIME_PRECISION`," +
-		"`CHARACTER_SET_NAME`," +
-		"`COLLATION_NAME`," +
-		"`COLUMN_TYPE`," +
-		"`COLUMN_KEY`," +
-		"`EXTRA`," +
-		"`PRIVILEGES`," +
-		"`COLUMN_COMMENT`," +
-		"`GENERATION_EXPRESSION`" +
-		" FROM `information_schema`.`COLUMNS`" +
+	query := "SELECT * FROM `information_schema`.`COLUMNS`" +
 		fmt.Sprintf(" WHERE TABLE_SCHEMA = '%s' AND TABLE_NAME = '%s'", t.Schema, t.Name)
 
 	constraints := constraintsAsMap(t.Constraints)
@@ -154,7 +139,8 @@ func (t *Table) parse() error {
 	for rows.Next() {
 		var f Field
 		var allowNull string
-		err := rows.Scan(&f.TableCatalog,
+		fields := []interface{}{
+			&f.TableCatalog,
 			&f.TableSchema,
 			&f.TableName,
 			&f.ColumnName,
@@ -174,10 +160,13 @@ func (t *Table) parse() error {
 			&f.Extra,
 			&f.Privileges,
 			&f.ColumnComment,
-			&f.GenerationExpression,
-		)
+		}
+
+		if cols, err := rows.Columns(); err == nil && len(cols) > 20 { //&& cols[20] == "GENERATION_EXPRESSION" {
+			fields = append(fields, &f.GenerationExpression)
+		}
+		err := rows.Scan(fields...)
 		if err != nil {
-			fmt.Println(err)
 			continue
 		}
 
@@ -207,26 +196,44 @@ func (t *Table) parse() error {
 	return nil
 }
 
-func getIndexes(db *sql.DB, schema, tableName string) ([]Index, error) {
+func (t *Table) FieldNames() []string {
+	fields := []string{}
+	for _, field := range t.Fields {
+		fields = append(fields, field.ColumnName)
+	}
+	return fields
+}
+
+func getIndexes(db *sql.DB, schema, tableName string) (map[string]Index, error) {
 	query := fmt.Sprintf("SHOW INDEXES FROM `%s`.`%s`", schema, tableName)
 	rows, err := db.Query(query)
 	if err != nil {
 		return nil, err
 	}
 	defer rows.Close()
 
-	indexes := []Index{}
+	indexes := make(map[string]Index)
 
 	for rows.Next() {
-		var i Index
+		var i IndexField
 		var table string
 		err := rows.Scan(&table, &i.NonUnique, &i.KeyName, &i.SeqInIndex,
 			&i.ColumnName, &i.Collation, &i.Cardinality, &i.SubPart,
 			&i.Packed, &i.Null, &i.IndexType, &i.Comment, &i.IndexComment)
 		if err != nil {
 			return nil, fmt.Errorf("cannot read constraints: %s", err)
 		}
-		indexes = append(indexes, i)
+		if index, ok := indexes[i.KeyName]; !ok {
+			indexes[i.KeyName] = Index{
+				Name:   i.KeyName,
+				Unique: !i.NonUnique,
+				Fields: []string{i.ColumnName},
+			}
+
+		} else {
+			index.Fields = append(index.Fields, i.ColumnName)
+			index.Unique = index.Unique || !i.NonUnique
+		}
 	}
 
 	return indexes, nil

tableparser/tableparser_test.go

@@ -4,12 +4,19 @@ import (
 	"reflect"
 	"testing"
 
-	tu "github.com/Percona-Lab/random_data_load/internal/testutils"
+	tu "github.com/Percona-Lab/random_data_load/testutils"
 	_ "github.com/go-sql-driver/mysql"
+	version "github.com/hashicorp/go-version"
 )
 
-func TestParse(t *testing.T) {
+func TestParse56(t *testing.T) {
 	db := tu.GetMySQLConnection(t)
+	v := tu.GetVersion(t, db)
+	v56, _ := version.NewVersion("5.6")
+
+	if v.GreaterThan(v56) {
+		t.Skipf("This test runs under MySQL < 5.7 and version is %s", v.String())
+	}
 
 	table, err := NewTable(db, "sakila", "film")
 	if err != nil {
@@ -23,6 +30,30 @@ func TestParse(t *testing.T) {
 	}
 }
 
+func TestParse57(t *testing.T) {
+	db := tu.GetMySQLConnection(t)
+	v := tu.GetVersion(t, db)
+	v57, _ := version.NewVersion("5.7")
+
+	if v.LessThan(v57) {
+		t.Skipf("This test runs under MySQL 5.7+ and version is %s", v.String())
+	}
+
+	table, err := NewTable(db, "sakila", "film")
+	if err != nil {
+		t.Error(err)
+	}
+	if tu.UpdateSamples() {
+		tu.WriteJson(t, "table002.json", table)
+	}
+	var want *Table
+	tu.LoadJson(t, "table002.json", &want)
+
+	if !reflect.DeepEqual(table, want) {
+		t.Error("Table film was not correctly parsed")
+	}
+}
+
 func TestGetIndexes(t *testing.T) {
 	db := tu.GetMySQLConnection(t)
 	want := []Index{}