Skip to content

Commit e0392b0

Browse files
committed
check token
1 parent 9a32437 commit e0392b0

2 files changed

Lines changed: 57 additions & 22 deletions

File tree

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJsonStreamParser.java

Lines changed: 54 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -111,12 +111,15 @@ public List<Object[]> parseStream(InputStream inputStream) throws IOException {
111111
}
112112

113113
try (JsonParser parser = JSON_FACTORY.createParser(inputStream)) {
114-
while (!parser.isClosed()) {
115-
JsonToken token = parser.nextToken();
114+
JsonToken token;
115+
while ((token = parser.nextToken()) != null && !parser.isClosed()) {
116116
if (token == JsonToken.FIELD_NAME && "rows".equals(parser.currentName())) {
117117
// Found "rows" array
118118
if (parser.nextToken() == JsonToken.START_ARRAY) {
119-
while (parser.nextToken() != JsonToken.END_ARRAY && !parser.isClosed()) {
119+
JsonToken rowToken;
120+
while ((rowToken = parser.nextToken()) != null
121+
&& rowToken != JsonToken.END_ARRAY
122+
&& !parser.isClosed()) {
120123
Object[] row = parseSingleRow(parser);
121124
if (row != null) {
122125
rows.add(row);
@@ -139,12 +142,18 @@ private Object[] parseSingleRow(JsonParser parser) throws IOException {
139142
int colCount = fieldList.size();
140143
Object[] row = new Object[colCount];
141144

142-
while (parser.nextToken() != JsonToken.END_OBJECT && !parser.isClosed()) {
145+
JsonToken token;
146+
while ((token = parser.nextToken()) != null
147+
&& token != JsonToken.END_OBJECT
148+
&& !parser.isClosed()) {
143149
String fieldName = parser.currentName();
144150
if ("f".equals(fieldName)) {
145151
if (parser.nextToken() == JsonToken.START_ARRAY) {
146152
int colIdx = 0;
147-
while (parser.nextToken() != JsonToken.END_ARRAY && !parser.isClosed()) {
153+
JsonToken cellToken;
154+
while ((cellToken = parser.nextToken()) != null
155+
&& cellToken != JsonToken.END_ARRAY
156+
&& !parser.isClosed()) {
148157
if (colIdx < colCount) {
149158
row[colIdx] = parseCell(parser, colIdx);
150159
} else {
@@ -167,7 +176,10 @@ private Object parseCell(JsonParser parser, int colIdx) throws IOException {
167176
}
168177

169178
Object cellValue = null;
170-
while (parser.nextToken() != JsonToken.END_OBJECT && !parser.isClosed()) {
179+
JsonToken token;
180+
while ((token = parser.nextToken()) != null
181+
&& token != JsonToken.END_OBJECT
182+
&& !parser.isClosed()) {
171183
String name = parser.currentName();
172184
if ("v".equals(name)) {
173185
parser.nextToken(); // move to value token
@@ -195,7 +207,10 @@ private FieldValue parseComplexFieldValue(JsonParser parser, Field field) throws
195207
List<FieldValue> elements = new ArrayList<>();
196208
if (parser.currentToken() == JsonToken.START_ARRAY) {
197209
Field elementField = field.toBuilder().setMode(Field.Mode.REQUIRED).build();
198-
while (parser.nextToken() != JsonToken.END_ARRAY && !parser.isClosed()) {
210+
JsonToken elemToken;
211+
while ((elemToken = parser.nextToken()) != null
212+
&& elemToken != JsonToken.END_ARRAY
213+
&& !parser.isClosed()) {
199214
// element is { "v": ... }
200215
elements.add(parseComplexFieldValue(parser, elementField));
201216
}
@@ -206,25 +221,45 @@ private FieldValue parseComplexFieldValue(JsonParser parser, Field field) throws
206221
if (parser.currentToken() == JsonToken.START_OBJECT) {
207222
FieldList subFields = field.getSubFields();
208223
int subIdx = 0;
209-
while (parser.nextToken() != JsonToken.END_OBJECT && !parser.isClosed()) {
210-
if ("f".equals(parser.currentName()) && parser.nextToken() == JsonToken.START_ARRAY) {
211-
while (parser.nextToken() != JsonToken.END_ARRAY && !parser.isClosed()) {
212-
Field subField =
213-
subFields != null && subIdx < subFields.size() ? subFields.get(subIdx) : null;
214-
fields.add(
215-
subField != null
216-
? parseComplexFieldValue(parser, subField)
217-
: FieldValue.of(Attribute.PRIMITIVE, null));
218-
subIdx++;
224+
int depth = 1;
225+
while (depth > 0 && !parser.isClosed()) {
226+
JsonToken token = parser.nextToken();
227+
if (token == null) {
228+
break;
229+
}
230+
if (token == JsonToken.START_OBJECT) {
231+
depth++;
232+
} else if (token == JsonToken.END_OBJECT) {
233+
depth--;
234+
} else if (token == JsonToken.FIELD_NAME) {
235+
if ("f".equals(parser.currentName()) && parser.nextToken() == JsonToken.START_ARRAY) {
236+
JsonToken subToken;
237+
while ((subToken = parser.nextToken()) != null
238+
&& subToken != JsonToken.END_ARRAY
239+
&& !parser.isClosed()) {
240+
Field subField =
241+
subFields != null && subIdx < subFields.size() ? subFields.get(subIdx) : null;
242+
fields.add(
243+
subField != null
244+
? parseComplexFieldValue(parser, subField)
245+
: FieldValue.of(Attribute.PRIMITIVE, null));
246+
subIdx++;
247+
}
219248
}
220249
}
221250
}
251+
if (subFields != null && fields.size() == subFields.size()) {
252+
return FieldValue.of(Attribute.RECORD, FieldValueList.of(fields, subFields));
253+
}
222254
}
223-
return FieldValue.of(Attribute.RECORD, FieldValueList.of(fields, field.getSubFields()));
255+
return FieldValue.of(Attribute.RECORD, FieldValueList.of(fields));
224256
} else {
225257
if (parser.currentToken() == JsonToken.START_OBJECT) {
226258
String val = null;
227-
while (parser.nextToken() != JsonToken.END_OBJECT && !parser.isClosed()) {
259+
JsonToken token;
260+
while ((token = parser.nextToken()) != null
261+
&& token != JsonToken.END_OBJECT
262+
&& !parser.isClosed()) {
228263
if ("v".equals(parser.currentName())) {
229264
parser.nextToken();
230265
val = parser.currentToken() == JsonToken.VALUE_NULL ? null : parser.getText();

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJsonStreamParserTest.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -182,11 +182,11 @@ public void testStreamParsingAndCoercionAllTypes() throws Exception {
182182

183183
@Test
184184
public void testMemoryAllocationReduction() throws Exception {
185-
ThreadMXBean baseBean = ManagementFactory.getThreadMXBean();
186-
if (!(baseBean instanceof com.sun.management.ThreadMXBean)) {
185+
java.lang.management.ThreadMXBean baseBean = ManagementFactory.getThreadMXBean();
186+
if (!(baseBean instanceof ThreadMXBean)) {
187187
return;
188188
}
189-
com.sun.management.ThreadMXBean threadBean = (com.sun.management.ThreadMXBean) baseBean;
189+
ThreadMXBean threadBean = (ThreadMXBean) baseBean;
190190
long threadId = Thread.currentThread().getId();
191191

192192
FieldList simpleFieldList =

0 commit comments

Comments
 (0)