Skip to content

Commit

Permalink
Fix list query history
Browse files Browse the repository at this point in the history
  • Loading branch information
mgyucht committed Jul 25, 2023
1 parent b71d391 commit 01b1d9b
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -93,21 +93,11 @@ private <I> Request withQuery(Request in, I entity) {
}
try {
// deterministic query string: in the order of class fields
for (Field field : entity.getClass().getDeclaredFields()) {
QueryParam param = field.getAnnotation(QueryParam.class);
if (param == null) {
continue;
}
field.setAccessible(true);
Object value = field.get(entity);
field.setAccessible(false);
if (value == null) {
continue;
}
in.withQueryParam(param.value(), value.toString());
for (Map.Entry<String, Object> e : HeaderSerializer.serialize(entity).entrySet()) {
in.withQueryParam(e.getKey(), mapper.writeValueAsString(e.getValue()));
}
return in;
} catch (IllegalAccessException e) {
} catch (JsonProcessingException e) {
throw new DatabricksException("Cannot create query string: " + e.getMessage(), e);
}
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
package com.databricks.sdk.core;

import com.databricks.sdk.support.QueryParam;
import com.fasterxml.jackson.annotation.JsonProperty;

import java.lang.reflect.Field;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class HeaderSerializer {
public static Map<String, Object> serialize(Object o) {
Map<String, Object> flattened = flattenObject(o);
for (Field f : o.getClass().getDeclaredFields()) {
QueryParam queryParam = f.getAnnotation(QueryParam.class);
if (queryParam == null) {
flattened.remove(getFieldName(f));
}
}

return flattened;
}

private static final List<Class<?>> primitiveTypes = Arrays.asList(
boolean.class,
Boolean.class,
byte.class,
Byte.class,
char.class,
Character.class,
short.class,
Short.class,
int.class,
Integer.class,
long.class,
Long.class,
float.class,
Float.class,
double.class,
Double.class,
String.class
);

private static String getFieldName(Field f) {
JsonProperty jsonProperty = f.getAnnotation(JsonProperty.class);
if (jsonProperty != null) {
return jsonProperty.value();
} else {
return f.getName();
}
}

private static Map<String, Object> flattenObject(Object o) {
Map<String, Object> result = new HashMap<>();
Field[] fields = o.getClass().getDeclaredFields();
for (Field f : fields) {
f.setAccessible(true);
try {
String name = getFieldName(f);
Object value = f.get(o);
if (value == null) {
continue;
}
// check if object is a primitive type
if (primitiveTypes.contains(f.getType()) || Iterable.class.isAssignableFrom(f.getType())) {
result.put(name, value);
} else {
// recursively flatten the object
Map<String, Object> flattened = flattenObject(value);
for (Map.Entry<String, Object> entry : flattened.entrySet()) {
result.put(name + "." + entry.getKey(), entry.getValue());
}
}
} catch (IllegalAccessException e) {
throw new RuntimeException(e);
} finally {
f.setAccessible(false);
}
}
return result;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,9 @@ private Iterator<T> outerIterator() {
return new Iterator<T>() {
@Override
public boolean hasNext() {
if (currentPage == null) {
return false;
}
if (currentPage.hasNext()) {
return true;
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
package com.databricks.sdk.integration;

import com.databricks.sdk.WorkspaceClient;
import com.databricks.sdk.integration.framework.EnvContext;
import com.databricks.sdk.integration.framework.EnvTest;
import com.databricks.sdk.service.sql.ListQueryHistoryRequest;
import com.databricks.sdk.service.sql.QueryFilter;
import com.databricks.sdk.service.sql.QueryInfo;
import com.databricks.sdk.service.sql.TimeRange;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;

@EnvContext("workspace")
@ExtendWith(EnvTest.class)
public class SqlIT {
@Test
void listQueryHistory(WorkspaceClient w) {
TimeRange timeRange = new TimeRange()
.setStartTimeMs(1690243200000L)
.setEndTimeMs(1690329600000L);
ListQueryHistoryRequest request = new ListQueryHistoryRequest()
.setFilterBy(new QueryFilter().setQueryStartTimeRange(timeRange));
Iterable<QueryInfo> queries = w.queryHistory().list(request);
for (QueryInfo query : queries) {
System.out.println(query);
}
}
}

0 comments on commit 01b1d9b

Please sign in to comment.