func (sei *session) handleSet(stmt *sqlparser.Set) error {
// get first arg
expr := stmt.Exprs[0]
name := string(expr.Name.Name)
switch name {
case "dbnum":
buf := sqlparser.NewTrackedBuffer(nil)
expr.Expr.Format(buf)
_, ok := expr.Expr.(sqlparser.NumVal)
if ok {
i, err := strconv.Atoi(buf.String())
if err != nil {
return sei.writeError(mysql.NewDefaultError(mysql.ER_WRONG_VALUE_FOR_VAR, buf.String(), name))
}
sei.dbnum = i
glog.Infof("set dbnum=%v", sei.dbnum)
return sei.writeOK(nil)
}
return sei.writeError(mysql.NewDefaultError(mysql.ER_WRONG_VALUE_FOR_VAR, buf.String(), name))
case "partitionkey":
buf := sqlparser.NewTrackedBuffer(nil)
expr.Expr.Format(buf)
_, ok := expr.Expr.(sqlparser.StrVal)
if ok {
sei.partitionKey = buf.String()
glog.Infof("set partitionKey=%v", sei.partitionKey)
return sei.writeOK(nil)
}
return sei.writeError(mysql.NewDefaultError(mysql.ER_WRONG_VALUE_FOR_VAR, buf.String(), name))
default:
return sei.writeError(mysql.NewDefaultError(mysql.ER_UNKNOWN_SYSTEM_VARIABLE, name))
}
}
func (sei *session) handleExplain(stmt *sqlparser.Explain) error {
if sei.db == "" {
return sei.writeError(mysql.NewDefaultError(mysql.ER_NO_DB_ERROR))
}
// build plan
plan, err := sei.buildPlan(stmt.SQL)
if err != nil {
return sei.writeError(err)
}
// build result
names := []string{"ID", "Backends", "SQLs"}
var r *mysql.Resultset
if len(plan.SQLs) == 0 {
r, _ = sei.buildEmptySet(names, []interface{}{""})
} else {
var values [][]interface{}
for id, sql := range plan.SQLs {
values = append(values, []interface{}{
id,
sql.Backend.Name,
strings.Join(sql.SQL, ";"),
})
}
r, _ = sei.buildResultset(nil, names, values)
}
return sei.writeResultset(sei.status, r)
}
func (sei *session) handleFieldList(data []byte) error {
index := bytes.IndexByte(data, 0x00)
table := string(data[0:index])
wildcard := string(data[index+1:])
if sei.db == "" {
return mysql.NewDefaultError(mysql.ER_NO_DB_ERROR)
}
tableInfo, err := sei.getMeta().GetTable(sei.user, sei.db, table)
if err != nil {
return err
}
co, err := sei.getBpool().GetConn(tableInfo.Backends[0])
if err != nil {
return err
}
defer sei.getBpool().PushConn(tableInfo.Backends[0], co, err)
fs, err := co.FieldList(table, wildcard)
if err != nil {
return err
}
return sei.writeFieldList(sei.status, fs)
}
func (sei *session) useDB(db string) error {
isExist, err := sei.server.info.IsDBExist(sei.user, db)
if err != nil {
return err
}
if isExist {
sei.db = db
return sei.writeOK(nil)
}
return sei.writeError(mysql.NewDefaultError(mysql.ER_BAD_DB_ERROR, db))
}
func (sei *session) readHandshakeResponse() error {
data, err := sei.pkg.ReadPacket()
if err != nil {
return err
}
pos := 0
// capability
sei.capability = binary.LittleEndian.Uint32(data[:4])
pos += 4
// skip max packet size
pos += 4
// charset, skip, if you want to use another charset, use set names
// sei.collation = CollationId(data[pos])
pos++
// skip reserved 23[00]
pos += 23
// user name
sei.user = string(data[pos : pos+bytes.IndexByte(data[pos:], 0)])
pos += len(sei.user) + 1
// auth length and auth
authLen := int(data[pos])
pos++
auth := data[pos : pos+authLen]
pos += authLen
if sei.capability&mysql.CLIENT_CONNECT_WITH_DB > 0 {
if len(data[pos:]) == 0 {
return nil
}
sei.db = string(data[pos : pos+bytes.IndexByte(data[pos:], 0)])
pos += len(sei.db) + 1
glog.Info("DB is :", sei.db)
}
// user and password check.
if !sei.server.info.CheckUser(sei.user, auth, sei.salt, sei.db) {
glog.Infof("User(%v) password or user name error", sei.user)
return mysql.NewDefaultError(mysql.ER_ACCESS_DENIED_ERROR, sei.user, "", "YES")
}
return nil
}
func (sei *session) handleQuery(data []byte) error {
sql := strings.TrimRight(string(data), ";")
stmt, err := sqlparser.Parse(sql)
if err != nil {
glog.Infof("parse sql(%s) error:%v", sql, err)
return sei.writeError(mysql.NewDefaultError(mysql.ER_SYNTAX_ERROR))
}
switch v := stmt.(type) {
case *sqlparser.Explain:
return sei.handleExplain(v)
case *sqlparser.Select:
return sei.handleSelect(v)
case *sqlparser.Insert:
return sei.handleInsert(v)
case *sqlparser.Update:
return sei.handleUpdate(v)
case *sqlparser.Delete:
return sei.handleDelete(v)
case *sqlparser.Set:
// only support like `SET autocommit=1`
return sei.handleSet(v)
case *sqlparser.DDL:
return sei.handleDDL(sql)
case *sqlparser.Show:
r, err := sei.handleShow(v)
if err != nil {
glog.Infof("handle show stmt has error:%v", err)
sei.writeError(err)
// not throw the error
return nil
}
return sei.writeResultset(sei.status, r)
case *sqlparser.UseDB:
return sei.handleUseDB(v)
default:
return fmt.Errorf("statement %T not support now", stmt)
}
return nil
}
func (sei *session) handleNormalExecute(stmt sqlparser.Statement) error {
if sei.db == "" {
return sei.writeError(mysql.NewDefaultError(mysql.ER_NO_DB_ERROR))
}
plan, err := sei.buildPlan(stmt)
if err != nil {
return sei.writeError(err)
}
rs, err := sei.executePlan(plan)
if err != nil {
return sei.writeError(mysql.NewError(mysql.ER_UNKNOWN_ERROR, err.Error()))
}
// merge results
return sei.mergeExecResult(rs)
}
func (sei *session) handleSelect(stmt *sqlparser.Select) error {
if stmt.From == nil {
return sei.handleSimpleSelect(stmt)
}
if len(stmt.From) > 1 {
return sei.writeError(mysql.NewDefaultError(mysql.ER_SYNTAX_ERROR))
}
plan, err := sei.buildPlan(stmt)
if err != nil {
return sei.writeError(mysql.NewError(mysql.ER_UNKNOWN_ERROR, err.Error()))
}
rs, err := sei.executePlan(plan)
if err != nil {
return sei.writeError(mysql.NewError(mysql.ER_UNKNOWN_ERROR, err.Error()))
}
offset, count := plan.getOffetCount()
// merge select rs
return sei.mergeSelectResult(rs, stmt, offset, count)
}
func (sei *session) buildInsertPlan(stmt *sqlparser.Insert) (*Plan, error) {
plan := &Plan{}
isNoCol := false
isNeedInitCol := true
colLen := len(stmt.Columns)
// get table
table, err := sei.getMeta().GetTable(sei.user, sei.db, string(stmt.Table.Name))
if err != nil {
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, err.Error())
}
plan.Table = table
if colLen == 0 {
isNoCol = true
}
// get sqls
values, ok := stmt.Rows.(sqlparser.Values)
if !ok {
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, "unsupport")
}
for row, value := range values {
tuple, ok := value.(sqlparser.ValTuple)
if !ok {
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, "unsupport")
}
if (isNoCol && len(tuple) != table.Table.ColsLen) ||
!isNoCol && len(tuple) != colLen {
return nil, mysql.NewDefaultError(mysql.ER_WRONG_VALUE_COUNT_ON_ROW, row+1)
}
// rewrite auto key
if isNoCol {
// rewrite by index
for _, autokey := range table.Table.AutoKeys {
id, err := table.GetKey(autokey.Name, config.Config.Proxy.AutoKeyInterval, sei.server.info)
if err != nil {
glog.Warningf("Get autokey(%v) get error(%v)", autokey.Name, err)
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, err.Error())
}
if autokey.Type == meta.TypeKeyInt {
tuple[autokey.Index] = sqlparser.NumVal(id)
} else {
tuple[autokey.Index] = sqlparser.StrVal(id)
}
}
} else {
for _, autokey := range table.Table.AutoKeys {
isExist := false
id, err := table.GetKey(autokey.Name, config.Config.Proxy.AutoKeyInterval, sei.server.info)
if err != nil {
glog.Warningf("Get autokey(%v) get error(%v)", autokey.Name, err)
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, err.Error())
}
for n, column := range stmt.Columns {
col, ok := column.(*sqlparser.NonStarExpr)
if !ok {
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, "unsupport")
}
colName, ok := col.Expr.(*sqlparser.ColName)
if !ok {
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, "unsupport")
}
if autokey.Name == string(colName.Name) && n < len(tuple) {
isExist = true
if autokey.Type == meta.TypeKeyInt {
tuple[n] = sqlparser.NumVal(id)
} else {
tuple[n] = sqlparser.StrVal(id)
}
}
}
if !isExist {
if isNeedInitCol {
stmt.Columns = append(stmt.Columns, &sqlparser.NonStarExpr{
Expr: &sqlparser.ColName{
Name: sqlparser.SQLName(autokey.Name),
},
})
}
if autokey.Type == meta.TypeKeyInt {
tuple = append(tuple, sqlparser.NumVal(id))
} else {
tuple = append(tuple, sqlparser.StrVal(id))
}
}
}
}
values[row] = tuple
isNeedInitCol = false
}
stmt.Rows = values
sqls := make(map[int]sqlparser.InsertRows)
// split stmt
for row, value := range values {
tuple := value.(sqlparser.ValTuple)
var value string
if isNoCol {
// get hash key
switch v := tuple[table.Table.PartitionKey.Index].(type) {
case sqlparser.NumVal:
value = string(v)
case sqlparser.StrVal:
value = string(v)
default:
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, "unknow partition key type")
}
} else {
for col, column := range stmt.Columns {
name := string(column.(*sqlparser.NonStarExpr).Expr.(*sqlparser.ColName).Name)
if name == table.Table.PartitionKey.Name {
switch v := tuple[col].(type) {
case sqlparser.NumVal:
value = string(v)
case sqlparser.StrVal:
value = string(v)
default:
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, "unknow partition key type")
}
}
}
}
if value == "" {
glog.Warning("Get partition key error value is null")
return nil, mysql.NewError(mysql.ER_UNKNOWN_ERROR, "unknow partition key type")
}
index := FindForKey(value, len(table.Backends))
s, ok := sqls[index]
if !ok {
sqls[index] = sqlparser.Values{values[row]}
} else {
sqls[index] = append(s.(sqlparser.Values), values[row])
}
}
var lastIDs []uint64
if len(table.Table.AutoIns) > 0 {
index := 0
if isNoCol {
// range column
index = table.Table.AutoIns[0].Index
} else {
// get index
index = sort.Search(len(stmt.Columns), func(i int) bool {
return string(stmt.Columns[i].(*sqlparser.NonStarExpr).Expr.(*sqlparser.ColName).Name) ==
table.Table.AutoIns[0].Name
})
}
lastIDs = make([]uint64, len(sqls))
n := 0
for _, sql := range sqls {
switch v := sql.(sqlparser.Values)[0].(sqlparser.ValTuple)[index].(type) {
case sqlparser.NumVal:
glog.Info(n, len(lastIDs))
lastIDs[n], _ = strconv.ParseUint(string(v), 10, 64)
case sqlparser.StrVal:
lastIDs[n], _ = strconv.ParseUint(string(v), 10, 64)
}
n++
}
}
for k, v := range sqls {
stmt.Rows = v
plan.generateOneSQL(stmt, k)
}
return plan, nil
}
func (sei *session) handleDDL(sql string) error {
stmt, err := parser.ParseOneStmt(sql, "", "")
if err != nil {
glog.Infof("parse ddl sql(%s) error:%v", sql, err)
return sei.writeError(mysql.NewDefaultError(mysql.ER_SYNTAX_ERROR))
}
switch v := stmt.(type) {
case *ast.CreateDatabaseStmt:
dbname := v.Name
id, rows, err := sei.ddlManage().CreateDatabase(sei.user, dbname, sei.dbnum)
glog.Infof("DDL plan id(%v)", id)
if err != nil {
glog.Infof("CREATE DATABASE has an error(%v)", err)
err = sei.writeError(err)
} else {
// one time only creata a db
r := &mysql.Result{
AffectedRows: rows,
}
err = sei.writeOK(r)
}
return err
case *ast.CreateTableStmt:
if sei.db == "" {
return sei.writeError(mysql.NewDefaultError(mysql.ER_NO_DB_ERROR))
}
table := &meta.Table{
Scheme: "hash",
Name: v.Table.Name.String(),
PartitionKey: &meta.Key{
Name: sei.partitionKey,
},
ColsLen: len(v.Cols),
}
existMap := make(map[string]bool)
// get constraints
for _, constraint := range v.Constraints {
if constraint.Tp == ast.ConstraintPrimaryKey ||
constraint.Tp == ast.ConstraintUniq {
if len(constraint.Keys) > 1 {
err := mysql.NewError(mysql.ER_SYNTAX_ERROR,
"not support constraint keys' length > 1")
return sei.writeError(err)
}
// get type
name := constraint.Keys[0].Column.Name.String()
index, typ := sei.getFieldType(v.Cols, name)
if typ == meta.TypeKeyUnknow {
err := mysql.NewError(mysql.ER_SYNTAX_ERROR,
"unsupport key's type ")
return sei.writeError(err)
}
if constraint.Tp == ast.ConstraintPrimaryKey && sei.partitionKey == "" {
// set primary key for partition key
table.PartitionKey.Name = name
table.PartitionKey.Type = typ
table.PartitionKey.Index = index
}
table.AutoKeys = append(table.AutoKeys, &meta.Key{
Name: name,
Type: typ,
Index: index,
})
existMap[name] = true
}
}
// check auto increment
for _, col := range v.Cols {
for n, option := range col.Options {
t := sei.getOneFeildType(col)
if t == meta.TypeKeyUnknow {
err := mysql.NewError(mysql.ER_SYNTAX_ERROR,
"unsupport key's type ")
return sei.writeError(err)
}
switch option.Tp {
case ast.ColumnOptionAutoIncrement, ast.ColumnOptionPrimaryKey, ast.ColumnOptionUniq:
if ast.ColumnOptionPrimaryKey == option.Tp && table.PartitionKey.Name == "" {
table.PartitionKey.Name = col.Name.Name.String()
table.PartitionKey.Type = t
table.PartitionKey.Index = n
}
// check if exist not append
if existMap[col.Name.Name.String()] {
continue
}
table.AutoKeys = append(table.AutoKeys, &meta.Key{
Name: col.Name.Name.String(),
Type: t,
Index: n,
})
existMap[col.Name.Name.String()] = true
}
if option.Tp == ast.ColumnOptionAutoIncrement {
glog.Infof("record auto increment option index(%v)", n)
// record
table.AutoIns = append(table.AutoIns, &meta.Key{
Name: col.Name.Name.String(),
Type: t,
Index: n,
})
}
}
}
// check partition key
if table.PartitionKey.Name == "" {
err := mysql.NewError(mysql.ER_SYNTAX_ERROR,
"partitionKey is null")
return sei.writeError(err)
}
data, _ := json.MarshalIndent(table, "", "\t")
glog.Info("table is\n", string(data))
id, rows, err := sei.ddlManage().CreateTable(sei.user, sei.db, sql, table)
glog.Infof("DDL plan id(%v)", id)
if err != nil {
glog.Infof("CREATE TABLE has an error(%v)", err)
err = sei.writeError(err)
} else {
// one time only creata a db
r := &mysql.Result{
AffectedRows: rows,
}
err = sei.writeOK(r)
}
return err
default:
return fmt.Errorf("create statement %T not support now", stmt)
}
}
func (d *Manage) doTask(t *Task) *Result {
// TODO: impl
r := new(Result)
plan := t.Plan
// get status
status, err := d.GetTaskStatus(plan.UserName, plan.ID)
if err != nil {
r.err = err
return r
}
// prepare handle task
switch t.Type {
case CreateDB:
// check the db
isExist, err := d.info.IsDBExist(plan.UserName, plan.DBName)
if err != nil {
r.err = err
return r
}
if isExist {
glog.Infof("database(%v) is already exist in user(%v)", plan.DBName, plan.UserName)
r.err = mysql.NewDefaultError(mysql.ER_DB_CREATE_EXISTS, plan.DBName)
return r
}
case CreateTable:
// check table
isExist, err := d.info.IsTableExist(plan.UserName, plan.DBName, plan.Table.Name)
if err != nil {
r.err = err
return r
}
if isExist {
glog.Infof("table(%v) is already exist in user(%v) database(%v)", plan.Table, plan.UserName, plan.DBName)
r.err = mysql.NewDefaultError(mysql.ER_DB_CREATE_EXISTS, plan.DBName)
return r
}
default:
r.err = fmt.Errorf("not support task's type:%v", t.Type)
}
// execute task
for _, sp := range plan.SubPlans {
d.updateTaskStatus(plan.UserName, plan.ID, sp.SubDatabase.Name, Doing, "", status, sp)
if err := d.executeSubPlan(sp, t); err != nil {
d.updateTaskStatus(plan.UserName, plan.ID, sp.SubDatabase.Name, Fail, err.Error(), status, sp)
r.err = err
break
} else {
r.affectedRows++
d.updateTaskStatus(plan.UserName, plan.ID, sp.SubDatabase.Name, Done, "", status, sp)
}
// record
plan.FinishNodes = append(plan.FinishNodes, sp.SubDatabase)
if len(plan.SubPlans) > 1 {
plan.SubPlans = plan.SubPlans[1:]
data, _ := json.Marshal(t)
d.info.UpdateTask(t.Plan.UserName, t.Seq, string(data))
} else {
// register result
if err = d.registerTaskResult(t); err != nil {
r.err = err
break
}
}
}
return r
}