Skip to content

Commit

Permalink
crossengine support dremio software
Browse files Browse the repository at this point in the history
  • Loading branch information
XN137 committed Sep 20, 2024
1 parent c7af055 commit a4ab011
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 23 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

import java.util.List;
import org.apache.spark.sql.SparkSession;
import org.assertj.core.api.ThrowableAssert;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
Expand Down Expand Up @@ -48,10 +49,13 @@ public class ITCrossEngineDremioPermissions {

private static final String NAMESPACE = System.getProperty("nesqueit.namespace", "db");
private static final String READ_ONLY_TABLE = "dremio_readonly"; // pre-created on the outside
private static final String EXPECTED_PERMISSION_ERROR =
String.format(
"Update entity is not allowed on reference 'main' on content '%s.%s'",
NAMESPACE, READ_ONLY_TABLE);

private static void assertPermissionError(ThrowableAssert.ThrowingCallable callable) {
assertThatThrownBy(callable)
.hasStackTraceContaining("Failed Access Check")
.hasStackTraceContaining("is not allowed")
.hasStackTraceContaining("on content '" + NAMESPACE + ".dremio_readonly'");
}

@BeforeAll
public static void setupNamepspace(
Expand All @@ -73,8 +77,7 @@ public void testReadOnlyTableSpark(
new SparkTestEngine("SPARK-READONLY", readOnlySpark, NAMESPACE);
List<Object> newRow = toRow(777);

assertThatThrownBy(() -> readOnlySparkEngine.insertRow(READ_ONLY_TABLE, newRow))
.hasStackTraceContaining(EXPECTED_PERMISSION_ERROR);
assertPermissionError(() -> readOnlySparkEngine.insertRow(READ_ONLY_TABLE, newRow));
assertThat(readOnlySparkEngine.selectRowsOrderedById(READ_ONLY_TABLE)).doesNotContain(newRow);

DremioTestEngine dremioEngine = new DremioTestEngine("DREMIO", dremio, NAMESPACE);
Expand All @@ -92,8 +95,7 @@ public void testReadOnlyTableFlink(
new FlinkTestEngine("FLINK-READONLY", readOnlyFlink, NAMESPACE);
List<Object> newRow = toRow(888);

assertThatThrownBy(() -> readOnlyFlinkEngine.insertRow(READ_ONLY_TABLE, newRow))
.hasStackTraceContaining(EXPECTED_PERMISSION_ERROR);
assertPermissionError(() -> readOnlyFlinkEngine.insertRow(READ_ONLY_TABLE, newRow));
assertThat(readOnlyFlinkEngine.selectRowsOrderedById(READ_ONLY_TABLE)).doesNotContain(newRow);

DremioTestEngine dremioEngine = new DremioTestEngine("DREMIO", dremio, NAMESPACE);
Expand All @@ -110,8 +112,7 @@ public void testReadOnlyTableDremio(
new DremioTestEngine("DREMIO-READONLY", readOnlyDremio, NAMESPACE);
List<Object> newRow = toRow(999);

assertThatThrownBy(() -> readOnlyDremioEngine.insertRow(READ_ONLY_TABLE, newRow))
.hasStackTraceContaining(EXPECTED_PERMISSION_ERROR);
assertPermissionError(() -> readOnlyDremioEngine.insertRow(READ_ONLY_TABLE, newRow));
assertThat(readOnlyDremioEngine.selectRowsOrderedById(READ_ONLY_TABLE)).doesNotContain(newRow);

DremioTestEngine dremioEngine = new DremioTestEngine("DREMIO", dremio, NAMESPACE);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,9 @@
import org.intellij.lang.annotations.Language;

public class DremioHelper {
private String token;
private String projectUrl;
private String catalogName;
private final String token;
private final String apiBaseUrl;
private final String catalogName;

private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

Expand All @@ -57,9 +57,9 @@ private static JsonNode parseJson(String json, String url) {
}
}

DremioHelper(String token, String apiBaseUrl, String projectId, String catalogName) {
DremioHelper(String token, String apiBaseUrl, String catalogName) {
this.token = token;
this.projectUrl = apiBaseUrl + "/v0/projects/" + projectId;
this.apiBaseUrl = apiBaseUrl;
this.catalogName = catalogName;
}

Expand All @@ -68,7 +68,7 @@ public String getCatalogName() {
}

public String getJobUrl(String jobId) {
return projectUrl + "/job/" + jobId;
return apiBaseUrl + "/job/" + jobId;
}

private String createPayload(String query) throws JsonProcessingException {
Expand Down Expand Up @@ -123,7 +123,8 @@ private static void sleep(Duration duration) {
}

private void waitForJobCompletion(String jobId, String query) throws IOException {
// See docs: https://docs.dremio.com/cloud/reference/api/job/
// docs cloud: https://docs.dremio.com/cloud/reference/api/job/
// docs sw: https://docs.dremio.com/current/reference/api/job/
String url = getJobUrl(jobId);
Set<String> finalJobStates = new HashSet<>(asList("COMPLETED", "FAILED", "CANCELED"));
// Default Timeout for engine-startup is 5min
Expand Down Expand Up @@ -165,7 +166,8 @@ private void waitForJobCompletion(String jobId, String query) throws IOException
}

private List<List<Object>> fetchQueryResult(String jobId) throws IOException {
// See docs: https://docs.dremio.com/cloud/reference/api/job/job-results
// docs cloud: https://docs.dremio.com/cloud/reference/api/job/job-results
// docs sw: https://docs.dremio.com/current/reference/api/job/job-results
String url = getJobUrl(jobId) + "/results";
String result = performHttpRequest(url, null);

Expand All @@ -181,9 +183,10 @@ private List<List<Object>> fetchQueryResult(String jobId) throws IOException {
}

private String submitQueryAndGetJobId(String query) throws IOException {
// See docs: https://docs.dremio.com/cloud/api/sql/
// docs cloud: https://docs.dremio.com/cloud/api/sql/
// docs sw: https://docs.dremio.com/current/reference/api/sql/
String payload = createPayload(query);
String url = projectUrl + "/sql";
String url = apiBaseUrl + "/sql";
String result = performHttpRequest(url, payload);
JsonNode node = parseJson(result, url);
JsonNode idNode = node.get("id");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -79,11 +79,21 @@ private DremioHelper buildDremioHelperFromSystemProperties(String systemProperty
if (systemPropertyPrefix == null || !systemPropertyPrefix.endsWith(".")) {
throw new IllegalArgumentException("Invalid systemPropertyPrefix: " + systemPropertyPrefix);
}
String apiBaseUrl = dremioUrl(); // same for all DremioHelper test params
String baseUrl = dremioUrl(); // same for all DremioHelper test params
String token = readRequiredSystemProperty(systemPropertyPrefix + "token");
String projectId = readRequiredSystemProperty(systemPropertyPrefix + "project-id");
String catalogName = readRequiredSystemProperty(systemPropertyPrefix + "catalog-name");

return new DremioHelper(token, apiBaseUrl, projectId, catalogName);
String projectId = System.getProperty(systemPropertyPrefix + "project-id");
String apiBaseUrl;
if (projectId != null) {
// dremio cloud
// https://docs.dremio.com/cloud/reference/api/
apiBaseUrl = baseUrl + "/v0/projects/" + projectId;
} else {
// dremio software
// https://docs.dremio.com/current/reference/api/
apiBaseUrl = baseUrl + "/api/v3";
}
return new DremioHelper(token, apiBaseUrl, catalogName);
}
}

0 comments on commit a4ab011

Please sign in to comment.