コード例 #1
0
ファイル: conn.go プロジェクト: jmptrader/tidb
// handleLoadData does the additional work after processing the 'load data' query.
// It sends client a file path, then reads the file content from client, inserts data into database.
func (cc *clientConn) handleLoadData(loadDataInfo *executor.LoadDataInfo) error {
	// If the server handles the load data request, the client has to set the ClientLocalFiles capability.
	if cc.capability&mysql.ClientLocalFiles == 0 {
		return errNotAllowedCommand
	}

	var err error
	defer func() {
		cc.ctx.SetValue(executor.LoadDataVarKey, nil)
		if err == nil {
			err = cc.ctx.CommitTxn()
		}
		if err == nil {
			return
		}
		err1 := cc.ctx.RollbackTxn()
		if err1 != nil {
			log.Errorf("load data rollback failed: %v", err1)
		}
	}()

	if loadDataInfo == nil {
		err = errors.New("load data info is empty")
		return errors.Trace(err)
	}
	err = cc.writeReq(loadDataInfo.Path)
	if err != nil {
		return errors.Trace(err)
	}

	var prevData []byte
	var curData []byte
	var shouldBreak bool
	for {
		curData, err = cc.readPacket()
		if err != nil {
			if terror.ErrorNotEqual(err, io.EOF) {
				log.Error(errors.ErrorStack(err))
				return errors.Trace(err)
			}
		}
		if len(curData) == 0 {
			shouldBreak = true
			if len(prevData) == 0 {
				break
			}
		}
		prevData, err = loadDataInfo.InsertData(prevData, curData)
		if err != nil {
			return errors.Trace(err)
		}
		if shouldBreak {
			break
		}
	}

	return nil
}
コード例 #2
0
ファイル: executor_test.go プロジェクト: jmptrader/tidb
func checkCases(cases []testCase, ld *executor.LoadDataInfo,
	c *C, tk *testkit.TestKit, ctx context.Context, selectSQL, deleteSQL string) {
	for _, ca := range cases {
		data, err1 := ld.InsertData(ca.data1, ca.data2)
		c.Assert(err1, IsNil)
		if ca.restData == nil {
			c.Assert(data, HasLen, 0,
				Commentf("data1:%v, data2:%v, data:%v", string(ca.data1), string(ca.data2), string(data)))
		} else {
			c.Assert(data, DeepEquals, ca.restData,
				Commentf("data1:%v, data2:%v, data:%v", string(ca.data1), string(ca.data2), string(data)))
		}
		err1 = ctx.CommitTxn()
		c.Assert(err1, IsNil)
		r := tk.MustQuery(selectSQL)
		r.Check(testkit.Rows(ca.expected...))
		tk.MustExec(deleteSQL)
	}
}