streamformatter.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159
  1. // Package streamformatter provides helper functions to format a stream.
  2. package streamformatter
  3. import (
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "github.com/docker/docker/pkg/jsonmessage"
  8. "github.com/docker/docker/pkg/progress"
  9. )
  10. const streamNewline = "\r\n"
  11. type jsonProgressFormatter struct{}
  12. func appendNewline(source []byte) []byte {
  13. return append(source, []byte(streamNewline)...)
  14. }
  15. // FormatStatus formats the specified objects according to the specified format (and id).
  16. func FormatStatus(id, format string, a ...interface{}) []byte {
  17. str := fmt.Sprintf(format, a...)
  18. b, err := json.Marshal(&jsonmessage.JSONMessage{ID: id, Status: str})
  19. if err != nil {
  20. return FormatError(err)
  21. }
  22. return appendNewline(b)
  23. }
  24. // FormatError formats the error as a JSON object
  25. func FormatError(err error) []byte {
  26. jsonError, ok := err.(*jsonmessage.JSONError)
  27. if !ok {
  28. jsonError = &jsonmessage.JSONError{Message: err.Error()}
  29. }
  30. if b, err := json.Marshal(&jsonmessage.JSONMessage{Error: jsonError, ErrorMessage: err.Error()}); err == nil {
  31. return appendNewline(b)
  32. }
  33. return []byte(`{"error":"format error"}` + streamNewline)
  34. }
  35. func (sf *jsonProgressFormatter) formatStatus(id, format string, a ...interface{}) []byte {
  36. return FormatStatus(id, format, a...)
  37. }
  38. // formatProgress formats the progress information for a specified action.
  39. func (sf *jsonProgressFormatter) formatProgress(id, action string, progress *jsonmessage.JSONProgress, aux interface{}) []byte {
  40. if progress == nil {
  41. progress = &jsonmessage.JSONProgress{}
  42. }
  43. var auxJSON *json.RawMessage
  44. if aux != nil {
  45. auxJSONBytes, err := json.Marshal(aux)
  46. if err != nil {
  47. return nil
  48. }
  49. auxJSON = new(json.RawMessage)
  50. *auxJSON = auxJSONBytes
  51. }
  52. b, err := json.Marshal(&jsonmessage.JSONMessage{
  53. Status: action,
  54. ProgressMessage: progress.String(),
  55. Progress: progress,
  56. ID: id,
  57. Aux: auxJSON,
  58. })
  59. if err != nil {
  60. return nil
  61. }
  62. return appendNewline(b)
  63. }
  64. type rawProgressFormatter struct{}
  65. func (sf *rawProgressFormatter) formatStatus(id, format string, a ...interface{}) []byte {
  66. return []byte(fmt.Sprintf(format, a...) + streamNewline)
  67. }
  68. func (sf *rawProgressFormatter) formatProgress(id, action string, progress *jsonmessage.JSONProgress, aux interface{}) []byte {
  69. if progress == nil {
  70. progress = &jsonmessage.JSONProgress{}
  71. }
  72. endl := "\r"
  73. if progress.String() == "" {
  74. endl += "\n"
  75. }
  76. return []byte(action + " " + progress.String() + endl)
  77. }
  78. // NewProgressOutput returns a progress.Output object that can be passed to
  79. // progress.NewProgressReader.
  80. func NewProgressOutput(out io.Writer) progress.Output {
  81. return &progressOutput{sf: &rawProgressFormatter{}, out: out, newLines: true}
  82. }
  83. // NewJSONProgressOutput returns a progress.Output that that formats output
  84. // using JSON objects
  85. func NewJSONProgressOutput(out io.Writer, newLines bool) progress.Output {
  86. return &progressOutput{sf: &jsonProgressFormatter{}, out: out, newLines: newLines}
  87. }
  88. type formatProgress interface {
  89. formatStatus(id, format string, a ...interface{}) []byte
  90. formatProgress(id, action string, progress *jsonmessage.JSONProgress, aux interface{}) []byte
  91. }
  92. type progressOutput struct {
  93. sf formatProgress
  94. out io.Writer
  95. newLines bool
  96. }
  97. // WriteProgress formats progress information from a ProgressReader.
  98. func (out *progressOutput) WriteProgress(prog progress.Progress) error {
  99. var formatted []byte
  100. if prog.Message != "" {
  101. formatted = out.sf.formatStatus(prog.ID, prog.Message)
  102. } else {
  103. jsonProgress := jsonmessage.JSONProgress{Current: prog.Current, Total: prog.Total, HideCounts: prog.HideCounts, Units: prog.Units}
  104. formatted = out.sf.formatProgress(prog.ID, prog.Action, &jsonProgress, prog.Aux)
  105. }
  106. _, err := out.out.Write(formatted)
  107. if err != nil {
  108. return err
  109. }
  110. if out.newLines && prog.LastUpdate {
  111. _, err = out.out.Write(out.sf.formatStatus("", ""))
  112. return err
  113. }
  114. return nil
  115. }
  116. // AuxFormatter is a streamFormatter that writes aux progress messages
  117. type AuxFormatter struct {
  118. io.Writer
  119. }
  120. // Emit emits the given interface as an aux progress message
  121. func (sf *AuxFormatter) Emit(aux interface{}) error {
  122. auxJSONBytes, err := json.Marshal(aux)
  123. if err != nil {
  124. return err
  125. }
  126. auxJSON := new(json.RawMessage)
  127. *auxJSON = auxJSONBytes
  128. msgJSON, err := json.Marshal(&jsonmessage.JSONMessage{Aux: auxJSON})
  129. if err != nil {
  130. return err
  131. }
  132. msgJSON = appendNewline(msgJSON)
  133. n, err := sf.Writer.Write(msgJSON)
  134. if n != len(msgJSON) {
  135. return io.ErrShortWrite
  136. }
  137. return err
  138. }