Function getStatementWithResultLimitInline

backend/plugin/db/tidb/query.go:98–152

(statement string, limit int)

}

func getStatementWithResultLimitInline(statement string, limit int) (string, error) {
	stmtList, err := tidbparser.ParseTiDB(statement, "", "")
	if err != nil {
		return "", errors.Wrapf(err, "failed to parse tidb statement: %s", statement)
	}
	if len(stmtList) != 1 {
		return "", errors.Errorf("expect one single statement in the query, %s", statement)
	}
	restoreFlags := format.DefaultRestoreFlags | format.RestoreStringWithoutDefaultCharset
	stmt := stmtList[0]
	switch stmt := stmt.(type) {
	case *tidbast.SelectStmt:
		if stmt.Limit != nil && stmt.Limit.Count != nil {
			if v, ok := stmt.Limit.Count.(*tidbdriver.ValueExpr); ok {
				userLimit := int(v.GetInt64())
				if limit < userLimit {
					userLimit = limit
				}
				stmt.Limit.Count = tidbast.NewValueExpr(int64(userLimit), "", "")
			}
		} else {
			stmt.Limit = &tidbast.Limit{
				Count: tidbast.NewValueExpr(int64(limit), "", ""),
			}
		}
		var buffer strings.Builder
		ctx := format.NewRestoreCtx(restoreFlags, &buffer)
		if err := stmt.Restore(ctx); err != nil {
			return "", err
		}
		return buffer.String(), nil
	case *tidbast.SetOprStmt:
		if stmt.Limit != nil && stmt.Limit.Count != nil {
			if v, ok := stmt.Limit.Count.(*tidbdriver.ValueExpr); ok {
				userLimit := int(v.GetInt64())
				if limit < userLimit {
					userLimit = limit
				}
				stmt.Limit.Count = tidbast.NewValueExpr(int64(userLimit), "", "")
			}
		} else {
			stmt.Limit = &tidbast.Limit{
				Count: tidbast.NewValueExpr(int64(limit), "", ""),
			}
		}
		var buffer strings.Builder
		ctx := format.NewRestoreCtx(restoreFlags, &buffer)
		if err := stmt.Restore(ctx); err != nil {
			return "", err
		}
		return buffer.String(), nil
	default:
	}
	return statement, nil
}