Skip to content

Commit baca331

Browse files
Neenu1995Neenu1995
committed
add tests
1 parent fa36876 commit baca331

2 files changed

Lines changed: 235 additions & 10 deletions

File tree

‎java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcContextProxy.java‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -81,9 +81,10 @@ private static String extractConnectionId(Object target) {
8181

8282
@Override
8383
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
84+
String methodName = method.getName();
85+
8486
// Handle standard Object methods explicitly
8587
if (method.getDeclaringClass() == Object.class) {
86-
String methodName = method.getName();
8788
if (methodName.equals("equals")) {
8889
Object other = args[0];
8990
if (other == null) {
@@ -106,7 +107,7 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
106107
}
107108

108109
// Support standard JDBC Wrapper unwrap operations
109-
if (method.getName().equals("unwrap") && args != null && args.length == 1) {
110+
if (methodName.equals("unwrap") && args != null && args.length == 1) {
110111
Class<?> iface = (Class<?>) args[0];
111112
if (iface.isInstance(target)) {
112113
return target;
@@ -117,7 +118,7 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
117118
throw e.getCause();
118119
}
119120
}
120-
if (method.getName().equals("isWrapperFor") && args != null && args.length == 1) {
121+
if (methodName.equals("isWrapperFor") && args != null && args.length == 1) {
121122
Class<?> iface = (Class<?>) args[0];
122123
if (iface.isInstance(target)) {
123124
return true;
@@ -131,18 +132,21 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
131132

132133
// Wrap execution in the context of the active connection for all non-bypassed methods
133134
try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(connectionId)) {
134-
if (LOG.isLoggable(Level.FINER)) {
135-
LOG.logp(Level.FINER, target.getClass().getName(), method.getName(), "++entry++");
135+
boolean isFinerLoggable = LOG.isLoggable(Level.FINER);
136+
String targetClassName = null;
137+
if (isFinerLoggable) {
138+
targetClassName = target.getClass().getName();
139+
LOG.logp(Level.FINER, targetClassName, methodName, "++entry++");
136140
}
137141

138142
Object result = method.invoke(target, args);
139143

140144
// Suppress exit logging for Connection.close() since its file handler is unmounted during
141145
// execution
142-
if (LOG.isLoggable(Level.FINER)
146+
if (isFinerLoggable
143147
&& !(java.sql.Connection.class.isAssignableFrom(interfaceType)
144-
&& "close".equals(method.getName()))) {
145-
LOG.logp(Level.FINER, target.getClass().getName(), method.getName(), "++exit++");
148+
&& "close".equals(methodName))) {
149+
LOG.logp(Level.FINER, targetClassName, methodName, "++exit++");
146150
}
147151

148152
// Symmetrical Cascade: Dynamic ResultSet concrete classes are deliberately unproxied here.
@@ -170,7 +174,7 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
170174
// context
171175
try (BigQueryJdbcMdc.MdcCloseable mdc = BigQueryJdbcMdc.registerInstance(connectionId)) {
172176
String errMsg = cause.getMessage() != null ? cause.getMessage() : cause.toString();
173-
LOG.severe("Exception occurred during " + method.getName() + ": " + errMsg, cause);
177+
LOG.severe("Exception occurred during " + methodName + ": " + errMsg, cause);
174178
}
175179

176180
throw cause;

‎java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcCustomLoggerTest.java‎

Lines changed: 222 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,28 @@
2121
import static org.junit.jupiter.api.Assertions.assertSame;
2222
import static org.junit.jupiter.api.Assertions.assertTrue;
2323

24+
import com.google.cloud.bigquery.Field;
25+
import com.google.cloud.bigquery.FieldList;
26+
import com.google.cloud.bigquery.FieldValue;
27+
import com.google.cloud.bigquery.FieldValue.Attribute;
28+
import com.google.cloud.bigquery.FieldValueList;
29+
import com.google.cloud.bigquery.LegacySQLTypeName;
2430
import java.util.ArrayList;
31+
import java.util.Collections;
32+
import java.util.HashMap;
2533
import java.util.List;
34+
import java.util.Map;
35+
import java.util.Map.Entry;
2636
import java.util.logging.Handler;
2737
import java.util.logging.Level;
2838
import java.util.logging.LogRecord;
39+
import org.apache.arrow.vector.util.JsonStringArrayList;
40+
import org.apache.arrow.vector.util.JsonStringHashMap;
2941
import org.junit.jupiter.api.AfterEach;
3042
import org.junit.jupiter.api.BeforeEach;
3143
import org.junit.jupiter.api.Test;
3244

33-
public class BigQueryJdbcCustomLoggerTest {
45+
public class BigQueryJdbcCustomLoggerTest extends BigQueryJdbcLoggingBaseTest {
3446

3547
private BigQueryJdbcCustomLogger logger;
3648
private TestHandler testHandler;
@@ -287,4 +299,213 @@ public void testResultSetLoggerParameterPreservation() {
287299
assertEquals("conn-123", params[0]);
288300
assertEquals("job-999", params[1]);
289301
}
302+
303+
@Test
304+
public void testArrowDatatypesConnectionIdPropagation() throws Exception {
305+
String connectionId = "conn-arrow-999";
306+
BigQueryJdbcResultSetLogger logger =
307+
BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowArray.class, connectionId);
308+
309+
Level originalRootLevel = BigQueryJdbcRootLogger.getRootLogger().getLevel();
310+
Level originalLoggerLevel = logger.getLevel();
311+
Map<Handler, Level> originalHandlerLevels = new HashMap<>();
312+
for (Handler h : BigQueryJdbcRootLogger.getRootLogger().getHandlers()) {
313+
originalHandlerLevels.put(h, h.getLevel());
314+
h.setLevel(Level.ALL);
315+
}
316+
317+
BigQueryJdbcRootLogger.getRootLogger().setLevel(Level.ALL);
318+
logger.setLevel(Level.ALL);
319+
320+
try {
321+
// 1. Arrow Array
322+
Field arraySchema = Field.of("name", LegacySQLTypeName.INTEGER);
323+
JsonStringArrayList<Object> arrayValues = new JsonStringArrayList<>();
324+
arrayValues.add(123L);
325+
BigQueryArrowArray array = new BigQueryArrowArray(arraySchema, arrayValues, logger);
326+
327+
capturedLogs.clear();
328+
array.LOG.fine("Log from Arrow Array");
329+
330+
assertEquals(1, capturedLogs.size());
331+
assertEquals(connectionId, capturedLogs.get(0).getParameters()[0]);
332+
333+
// 2. Arrow Struct
334+
FieldList structSchema = FieldList.of(Field.of("col", LegacySQLTypeName.INTEGER));
335+
JsonStringHashMap<String, Object> structValues = new JsonStringHashMap<>();
336+
structValues.put("col", 456L);
337+
BigQueryJdbcResultSetLogger structLogger =
338+
BigQueryJdbcResultSetLogger.getLogger(BigQueryArrowStruct.class, connectionId);
339+
structLogger.setLevel(Level.ALL);
340+
BigQueryArrowStruct struct =
341+
new BigQueryArrowStruct(structSchema, structValues, structLogger);
342+
343+
capturedLogs.clear();
344+
struct.LOG.fine("Log from Arrow Struct");
345+
346+
assertEquals(1, capturedLogs.size());
347+
assertEquals(connectionId, capturedLogs.get(0).getParameters()[0]);
348+
349+
// 3. Nested Struct in Arrow Array
350+
FieldList subFields = FieldList.of(Field.of("sub_col", LegacySQLTypeName.INTEGER));
351+
Field arrayNestedSchema =
352+
Field.newBuilder("struct_arr", LegacySQLTypeName.RECORD, subFields)
353+
.setMode(Field.Mode.REPEATED)
354+
.build();
355+
JsonStringHashMap<String, Object> map = new JsonStringHashMap<>();
356+
map.put("sub_col", 456L);
357+
JsonStringArrayList<Object> nestedValues = new JsonStringArrayList<>();
358+
nestedValues.add(map);
359+
BigQueryArrowArray arrayWithNested =
360+
new BigQueryArrowArray(arrayNestedSchema, nestedValues, logger);
361+
362+
Object[] result = (Object[]) arrayWithNested.getArray();
363+
BigQueryArrowStruct nestedStruct = (BigQueryArrowStruct) result[0];
364+
365+
capturedLogs.clear();
366+
nestedStruct.LOG.fine("Log from Nested Struct");
367+
368+
assertEquals(1, capturedLogs.size());
369+
assertEquals(connectionId, capturedLogs.get(0).getParameters()[0]);
370+
371+
// 4. Nested Array in Arrow Struct
372+
Field arrayField =
373+
Field.newBuilder("array_col", LegacySQLTypeName.INTEGER)
374+
.setMode(Field.Mode.REPEATED)
375+
.build();
376+
FieldList nestedSchema = FieldList.of(arrayField);
377+
JsonStringArrayList<Object> arrayVal = new JsonStringArrayList<>();
378+
arrayVal.add(123L);
379+
JsonStringHashMap<String, Object> values = new JsonStringHashMap<>();
380+
values.put("array_col", arrayVal);
381+
382+
BigQueryArrowStruct structWithNestedArray =
383+
new BigQueryArrowStruct(nestedSchema, values, structLogger);
384+
Object[] attributes = structWithNestedArray.getAttributes();
385+
BigQueryArrowArray nestedArray = (BigQueryArrowArray) attributes[0];
386+
387+
capturedLogs.clear();
388+
nestedArray.LOG.fine("Log from Nested Array");
389+
390+
assertEquals(1, capturedLogs.size());
391+
assertEquals(connectionId, capturedLogs.get(0).getParameters()[0]);
392+
393+
} finally {
394+
BigQueryJdbcRootLogger.getRootLogger().setLevel(originalRootLevel);
395+
logger.setLevel(originalLoggerLevel);
396+
for (Entry<Handler, Level> entry : originalHandlerLevels.entrySet()) {
397+
entry.getKey().setLevel(entry.getValue());
398+
}
399+
}
400+
}
401+
402+
@Test
403+
public void testJsonDatatypesConnectionIdPropagation() throws Exception {
404+
String connectionId = "conn-json-888";
405+
BigQueryJdbcResultSetLogger logger =
406+
BigQueryJdbcResultSetLogger.getLogger(BigQueryJsonArray.class, connectionId);
407+
408+
Level originalRootLevel = BigQueryJdbcRootLogger.getRootLogger().getLevel();
409+
Level originalLoggerLevel = logger.getLevel();
410+
Map<Handler, Level> originalHandlerLevels = new HashMap<>();
411+
for (Handler h : BigQueryJdbcRootLogger.getRootLogger().getHandlers()) {
412+
originalHandlerLevels.put(h, h.getLevel());
413+
h.setLevel(Level.ALL);
414+
}
415+
416+
BigQueryJdbcRootLogger.getRootLogger().setLevel(Level.ALL);
417+
logger.setLevel(Level.ALL);
418+
419+
try {
420+
// 1. JSON Array
421+
Field arraySchema = Field.of("name", LegacySQLTypeName.INTEGER);
422+
FieldValue arrayValue =
423+
FieldValue.of(
424+
Attribute.REPEATED,
425+
FieldValueList.of(
426+
Collections.singletonList(FieldValue.of(Attribute.PRIMITIVE, "123"))));
427+
BigQueryJsonArray array = new BigQueryJsonArray(arraySchema, arrayValue, logger);
428+
429+
capturedLogs.clear();
430+
array.LOG.fine("Log from JSON Array");
431+
432+
assertEquals(1, capturedLogs.size());
433+
assertEquals(connectionId, capturedLogs.get(0).getParameters()[0]);
434+
435+
// 2. JSON Struct
436+
FieldList structSchema = FieldList.of(Field.of("col", LegacySQLTypeName.INTEGER));
437+
FieldValue structValue =
438+
FieldValue.of(
439+
Attribute.RECORD,
440+
FieldValueList.of(
441+
Collections.singletonList(FieldValue.of(Attribute.PRIMITIVE, "456"))));
442+
BigQueryJdbcResultSetLogger structLogger =
443+
BigQueryJdbcResultSetLogger.getLogger(BigQueryJsonStruct.class, connectionId);
444+
structLogger.setLevel(Level.ALL);
445+
BigQueryJsonStruct struct = new BigQueryJsonStruct(structSchema, structValue, structLogger);
446+
447+
capturedLogs.clear();
448+
struct.LOG.fine("Log from JSON Struct");
449+
450+
assertEquals(1, capturedLogs.size());
451+
assertEquals(connectionId, capturedLogs.get(0).getParameters()[0]);
452+
453+
// 3. Nested Struct in JSON Array
454+
FieldList subFields = FieldList.of(Field.of("sub_col", LegacySQLTypeName.INTEGER));
455+
Field arrayNestedSchema =
456+
Field.newBuilder("struct_arr", LegacySQLTypeName.RECORD, subFields)
457+
.setMode(Field.Mode.REPEATED)
458+
.build();
459+
FieldValue recordVal =
460+
FieldValue.of(
461+
Attribute.RECORD,
462+
FieldValueList.of(
463+
Collections.singletonList(FieldValue.of(Attribute.PRIMITIVE, "789"))));
464+
FieldValue listVal =
465+
FieldValue.of(
466+
Attribute.REPEATED, FieldValueList.of(Collections.singletonList(recordVal)));
467+
BigQueryJsonArray arrayWithNested = new BigQueryJsonArray(arrayNestedSchema, listVal, logger);
468+
469+
Object[] result = (Object[]) arrayWithNested.getArray();
470+
BigQueryJsonStruct nestedStruct = (BigQueryJsonStruct) result[0];
471+
472+
capturedLogs.clear();
473+
nestedStruct.LOG.fine("Log from Nested JSON Struct");
474+
475+
assertEquals(1, capturedLogs.size());
476+
assertEquals(connectionId, capturedLogs.get(0).getParameters()[0]);
477+
478+
// 4. Nested Array in JSON Struct
479+
Field arrayField =
480+
Field.newBuilder("array_col", LegacySQLTypeName.INTEGER)
481+
.setMode(Field.Mode.REPEATED)
482+
.build();
483+
FieldList nestedSchema = FieldList.of(arrayField);
484+
FieldValue arrayVal =
485+
FieldValue.of(
486+
Attribute.REPEATED,
487+
FieldValueList.of(
488+
Collections.singletonList(FieldValue.of(Attribute.PRIMITIVE, "123"))));
489+
FieldValue rootVal =
490+
FieldValue.of(Attribute.RECORD, FieldValueList.of(Collections.singletonList(arrayVal)));
491+
492+
BigQueryJsonStruct structWithNestedArray =
493+
new BigQueryJsonStruct(nestedSchema, rootVal, structLogger);
494+
Object[] attributes = structWithNestedArray.getAttributes();
495+
BigQueryJsonArray nestedArray = (BigQueryJsonArray) attributes[0];
496+
497+
capturedLogs.clear();
498+
nestedArray.LOG.fine("Log from Nested JSON Array");
499+
500+
assertEquals(1, capturedLogs.size());
501+
assertEquals(connectionId, capturedLogs.get(0).getParameters()[0]);
502+
503+
} finally {
504+
BigQueryJdbcRootLogger.getRootLogger().setLevel(originalRootLevel);
505+
logger.setLevel(originalLoggerLevel);
506+
for (Entry<Handler, Level> entry : originalHandlerLevels.entrySet()) {
507+
entry.getKey().setLevel(entry.getValue());
508+
}
509+
}
510+
}
290511
}

0 commit comments

Comments
 (0)