Function selectTiDBTableSample

dumpling/export/dump.go:968–1013

(tctx *tcontext.Context, conn *BaseConn, meta TableMeta)

}

func selectTiDBTableSample(tctx *tcontext.Context, conn *BaseConn, meta TableMeta) (pkFields []string, pkVals [][]string, err error) {
	pkFields, pkColTypes, err := selectTiDBRowKeyFields(tctx, conn, meta, nil)
	if err != nil {
		return nil, nil, errors.Trace(err)
	}

	query := buildTiDBTableSampleQuery(pkFields, meta.DatabaseName(), meta.TableName())
	pkValNum := len(pkFields)
	var iter SQLRowIter
	rowRec := MakeRowReceiver(pkColTypes)
	buf := new(bytes.Buffer)

	err = conn.QuerySQL(tctx, func(rows *sql.Rows) error {
		if iter == nil {
			iter = &rowIter{
				rows: rows,
				args: make([]any, pkValNum),
			}
		}
		err = iter.Decode(rowRec)
		if err != nil {
			return errors.Trace(err)
		}
		pkValRow := make([]string, 0, pkValNum)
		for _, rec := range rowRec.receivers {
			rec.WriteToBuffer(buf, true)
			pkValRow = append(pkValRow, buf.String())
			buf.Reset()
		}
		pkVals = append(pkVals, pkValRow)
		return nil
	}, func() {
		if iter != nil {
			_ = iter.Close()
			iter = nil
		}
		rowRec = MakeRowReceiver(pkColTypes)
		pkVals = pkVals[:0]
		buf.Reset()
	}, query)
	if err == nil && iter != nil && iter.Error() != nil {
		err = iter.Error()
	}

	return pkFields, pkVals, err
}

func buildTiDBTableSampleQuery(pkFields []string, dbName, tblName string) string {
	template := "SELECT %s FROM `%s`.`%s` TABLESAMPLE REGIONS() ORDER BY %s"