1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
|
// Copyright (C) MongoDB, Inc. 2014-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
package mongoexport
import (
"encoding/csv"
"fmt"
"github.com/mongodb/mongo-tools/common/bsonutil"
"github.com/mongodb/mongo-tools/common/json"
"gopkg.in/mgo.v2/bson"
"io"
"reflect"
"strconv"
"strings"
)
// type for reflect code
var marshalDType = reflect.TypeOf(bsonutil.MarshalD{})
// CSVExportOutput is an implementation of ExportOutput that writes documents to the output in CSV format.
type CSVExportOutput struct {
// Fields is a list of field names in the bson documents to be exported.
// A field can also use dot-delimited modifiers to address nested structures,
// for example "location.city" or "addresses.0".
Fields []string
// NumExported maintains a running total of the number of documents written.
NumExported int64
// NoHeaderLine, if set, will export CSV data without a list of field names at the first line
NoHeaderLine bool
csvWriter *csv.Writer
}
// NewCSVExportOutput returns a CSVExportOutput configured to write output to the
// given io.Writer, extracting the specified fields only.
func NewCSVExportOutput(fields []string, noHeaderLine bool, out io.Writer) *CSVExportOutput {
return &CSVExportOutput{
fields,
0,
noHeaderLine,
csv.NewWriter(out),
}
}
// WriteHeader writes a comma-delimited list of fields as the output header row.
func (csvExporter *CSVExportOutput) WriteHeader() error {
if !csvExporter.NoHeaderLine {
csvExporter.csvWriter.Write(csvExporter.Fields)
return csvExporter.csvWriter.Error()
}
return nil
}
// WriteFooter is a no-op for CSV export formats.
func (csvExporter *CSVExportOutput) WriteFooter() error {
// no CSV footer
return nil
}
// Flush writes any pending data to the underlying I/O stream.
func (csvExporter *CSVExportOutput) Flush() error {
csvExporter.csvWriter.Flush()
return csvExporter.csvWriter.Error()
}
// ExportDocument writes a line to output with the CSV representation of a document.
func (csvExporter *CSVExportOutput) ExportDocument(document bson.D) error {
rowOut := make([]string, 0, len(csvExporter.Fields))
extendedDoc, err := bsonutil.ConvertBSONValueToJSON(document)
if err != nil {
return err
}
for _, fieldName := range csvExporter.Fields {
fieldVal := extractFieldByName(fieldName, extendedDoc)
if fieldVal == nil {
rowOut = append(rowOut, "")
} else if reflect.TypeOf(fieldVal) == reflect.TypeOf(bson.M{}) ||
reflect.TypeOf(fieldVal) == reflect.TypeOf(bson.D{}) ||
reflect.TypeOf(fieldVal) == marshalDType ||
reflect.TypeOf(fieldVal) == reflect.TypeOf([]interface{}{}) {
buf, err := json.Marshal(fieldVal)
if err != nil {
rowOut = append(rowOut, "")
} else {
rowOut = append(rowOut, string(buf))
}
} else {
rowOut = append(rowOut, fmt.Sprintf("%v", fieldVal))
}
}
csvExporter.csvWriter.Write(rowOut)
csvExporter.NumExported++
return csvExporter.csvWriter.Error()
}
// extractFieldByName takes a field name and document, and returns a value representing
// the value of that field in the document in a format that can be printed as a string.
// It will also handle dot-delimited field names for nested arrays or documents.
func extractFieldByName(fieldName string, document interface{}) interface{} {
dotParts := strings.Split(fieldName, ".")
var subdoc interface{} = document
for _, path := range dotParts {
docValue := reflect.ValueOf(subdoc)
if !docValue.IsValid() {
return ""
}
docType := docValue.Type()
docKind := docType.Kind()
if docKind == reflect.Map {
subdocVal := docValue.MapIndex(reflect.ValueOf(path))
if subdocVal.Kind() == reflect.Invalid {
return ""
}
subdoc = subdocVal.Interface()
} else if docKind == reflect.Slice {
if docType == marshalDType {
// dive into a D as a document
asD := bson.D(subdoc.(bsonutil.MarshalD))
var err error
subdoc, err = bsonutil.FindValueByKey(path, &asD)
if err != nil {
return ""
}
} else {
// check that the path can be converted to int
arrayIndex, err := strconv.Atoi(path)
if err != nil {
return ""
}
// bounds check for slice
if arrayIndex < 0 || arrayIndex >= docValue.Len() {
return ""
}
subdocVal := docValue.Index(arrayIndex)
if subdocVal.Kind() == reflect.Invalid {
return ""
}
subdoc = subdocVal.Interface()
}
} else {
// trying to index into a non-compound type - just return blank.
return ""
}
}
return subdoc
}
|
