Skip to content

Commit cc5549e

Browse files
dmitrys-odysseusdmitrys-odysseus
authored
Merge pull request #400 from OHDSI/bigquery
Reinstate support for BigQuery
2 parents 30eedd9 + 63b9203 commit cc5549e

11 files changed

+199
-73
lines changed

‎engine/pom.xml

+6
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,7 @@
359359
<google.gax.version>2.13.0</google.gax.version>
360360
<proto.google.cloud.bigquerystorage.v1.version>2.12.0</proto.google.cloud.bigquerystorage.v1.version>
361361
<opencensus.contrib.http.util.version>0.31.0</opencensus.contrib.http.util.version>
362+
<io.grpc.grpc-core.version>1.45.0</io.grpc.grpc-core.version>
362363
<!-- BigQuery JDBC driver path -->
363364
<bigquery.classpath>/bigquery</bigquery.classpath>
364365
</properties>
@@ -428,6 +429,11 @@
428429
<artifactId>opencensus-contrib-http-util</artifactId>
429430
<version>${opencensus.contrib.http.util.version}</version>
430431
</dependency>
432+
<dependency>
433+
<groupId>io.grpc</groupId>
434+
<artifactId>grpc-core</artifactId>
435+
<version>${io.grpc.grpc-core.version}</version>
436+
</dependency>
431437
</dependencies>
432438
<build>
433439
<plugins>

‎engine/src/main/docker/Dockerfile

+1-4
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
FROM eclipse-temurin:8-jdk
1+
FROM eclipse-temurin:8-jre
22
LABEL maintainer="Odysseus Data Services, Inc: [email protected]"
33
VOLUME /tmp
44
VOLUME /libs-r
@@ -10,9 +10,6 @@ RUN ln -fs /usr/share/zoneinfo/Etc/GMT /etc/localtime && dpkg-reconfigure -f non
1010
RUN update-rc.d cron defaults && systemctl enable rsyslog
1111
ENV R_INSTALL_SCRIPT=libs.r
1212
ADD execution-engine-exec.jar /execution-engine.jar
13-
RUN mkdir /impala
14-
ADD Cloudera_ImpalaJDBC4_2.5.41.zip /impala/impala-jdbc.zip
15-
RUN cd /impala && unzip impala-jdbc.zip && rm -f impala-jdbc.zip
1613
RUN mkdir /runtimes
1714
COPY descriptor_base.json /runtimes/descriptor_base.json
1815
RUN ls -la /runtimes/*

‎engine/src/main/java/com/odysseusinc/arachne/executionengine/auth/AuthEffects.java

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package com.odysseusinc.arachne.executionengine.auth;
2+
3+
import java.util.Map;
4+
5+
public interface AuthEffects {
6+
interface AddEnvironmentVariables extends AuthEffects {
7+
Map<String, String> getEnvVars();
8+
}
9+
10+
interface ModifyUrl extends AuthEffects {
11+
String getNewUrl();
12+
}
13+
14+
interface Cleanup extends AuthEffects {
15+
void cleanup();
16+
}
17+
}

‎engine/src/main/java/com/odysseusinc/arachne/executionengine/auth/BigQueryCredentialsProvider.java

+64
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package com.odysseusinc.arachne.executionengine.auth;
2+
3+
import com.google.common.collect.ImmutableMap;
4+
import com.odysseusinc.arachne.commons.types.DBMSType;
5+
import com.odysseusinc.arachne.execution_engine_common.api.v1.dto.DataSourceUnsecuredDTO;
6+
import lombok.AllArgsConstructor;
7+
import lombok.extern.slf4j.Slf4j;
8+
import org.apache.commons.io.FileUtils;
9+
import org.apache.commons.io.IOUtils;
10+
import org.springframework.stereotype.Service;
11+
12+
import java.io.IOException;
13+
import java.io.OutputStream;
14+
import java.nio.file.Files;
15+
import java.nio.file.Path;
16+
import java.nio.file.Paths;
17+
import java.util.Map;
18+
19+
@Service
20+
@Slf4j
21+
public class BigQueryCredentialsProvider implements CredentialsProvider {
22+
private static final String RUNTIME_BQ_KEYFILE = "BQ_KEYFILE";
23+
24+
@Override
25+
public AuthEffects apply(DataSourceUnsecuredDTO dataSource, Path analysisDir, String analysisDirInContainer) {
26+
if (dataSource.getType() == DBMSType.BIGQUERY) {
27+
try {
28+
Path keyFile = Files.createTempFile(analysisDir, "", ".json").toAbsolutePath();
29+
try (OutputStream out = Files.newOutputStream(keyFile)) {
30+
IOUtils.write(dataSource.getKeyfile(), out);
31+
}
32+
Path innerPath = Paths.get(analysisDirInContainer).resolve(keyFile.getFileName());
33+
return new Effects(dataSource.getConnectionString(), keyFile, innerPath.toString());
34+
} catch (IOException e) {
35+
log.error("Failed to resolve BigQuery authentication for Source: [{}]", dataSource.getName(), e);
36+
throw new RuntimeException(e);
37+
}
38+
} else {
39+
return null;
40+
}
41+
}
42+
43+
@AllArgsConstructor
44+
private static class Effects implements AuthEffects, AuthEffects.AddEnvironmentVariables, AuthEffects.ModifyUrl, AuthEffects.Cleanup {
45+
private final String newUrl;
46+
private final Path keyFile;
47+
private final String keyFileInContainer;
48+
49+
@Override
50+
public Map<String, String> getEnvVars() {
51+
return ImmutableMap.of(RUNTIME_BQ_KEYFILE, keyFileInContainer);
52+
}
53+
54+
@Override
55+
public String getNewUrl() {
56+
return newUrl + ";OAuthPvtKeyPath=" + keyFileInContainer + ";";
57+
}
58+
59+
@Override
60+
public void cleanup() {
61+
FileUtils.deleteQuietly(keyFile.toFile());
62+
}
63+
}
64+
}

‎engine/src/main/java/com/odysseusinc/arachne/executionengine/auth/CredentialsProvider.java

+9
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
package com.odysseusinc.arachne.executionengine.auth;
2+
3+
import com.odysseusinc.arachne.execution_engine_common.api.v1.dto.DataSourceUnsecuredDTO;
4+
5+
import java.nio.file.Path;
6+
7+
public interface CredentialsProvider {
8+
AuthEffects apply(DataSourceUnsecuredDTO dataSource, Path analysisDir, String analysisDirInContainer);
9+
}

‎engine/src/main/java/com/odysseusinc/arachne/executionengine/auth/KerberosCredentialsProvider.java

+60
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
package com.odysseusinc.arachne.executionengine.auth;
2+
3+
import com.odysseusinc.arachne.execution_engine_common.api.v1.dto.DataSourceUnsecuredDTO;
4+
import com.odysseusinc.datasourcemanager.krblogin.KerberosService;
5+
import com.odysseusinc.datasourcemanager.krblogin.KrbConfig;
6+
import com.odysseusinc.datasourcemanager.krblogin.RuntimeServiceMode;
7+
import lombok.AllArgsConstructor;
8+
import lombok.extern.slf4j.Slf4j;
9+
import org.apache.commons.io.FileUtils;
10+
import org.springframework.beans.factory.annotation.Autowired;
11+
import org.springframework.stereotype.Service;
12+
13+
import java.io.File;
14+
import java.io.IOException;
15+
import java.nio.file.Path;
16+
import java.util.Map;
17+
18+
@Service
19+
@Slf4j
20+
public class KerberosCredentialsProvider implements CredentialsProvider {
21+
@Autowired
22+
private KerberosService kerberosService;
23+
24+
@Override
25+
public AuthEffects apply(DataSourceUnsecuredDTO ds, Path analysisDir, String analysisDirInContainer) {
26+
Path keys = analysisDir.resolve("keys");
27+
File keystoreDir = keys.toFile();
28+
29+
if (ds.getUseKerberos()) {
30+
keystoreDir.mkdirs();
31+
try {
32+
KrbConfig config = kerberosService.runKinit(ds, RuntimeServiceMode.SINGLE, keystoreDir);
33+
return new Effects(config, keys);
34+
} catch (IOException e) {
35+
log.error("Failed to resolve Kerberos auth for Datasource: {}", ds.getName(), e);
36+
throw new RuntimeException(e);
37+
}
38+
} else {
39+
return null;
40+
}
41+
}
42+
43+
@AllArgsConstructor
44+
private static class Effects implements AuthEffects, AuthEffects.AddEnvironmentVariables, AuthEffects.Cleanup {
45+
private final KrbConfig config;
46+
private final Path keyFile;
47+
48+
@Override
49+
public Map<String, String> getEnvVars() {
50+
return config.getIsolatedRuntimeEnvs();
51+
}
52+
53+
@Override
54+
public void cleanup() {
55+
FileUtils.deleteQuietly(keyFile.toFile());
56+
FileUtils.deleteQuietly(config.getComponents().getKeytabPath().toFile());
57+
FileUtils.deleteQuietly(config.getConfPath().toFile());
58+
}
59+
}
60+
}

‎engine/src/main/java/com/odysseusinc/arachne/executionengine/config/AnalisysConfig.java

+6
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
package com.odysseusinc.arachne.executionengine.config;
2424

25+
import com.odysseusinc.datasourcemanager.jdbc.auth.BigQueryAuthResolver;
2526
import com.odysseusinc.datasourcemanager.krblogin.KerberosService;
2627
import com.odysseusinc.datasourcemanager.krblogin.KerberosServiceImpl;
2728
import org.springframework.beans.factory.annotation.Value;
@@ -69,4 +70,9 @@ public KerberosService kerberosService() {
6970
return new KerberosServiceImpl(timeout, kinitPath, configPath);
7071
}
7172

73+
@Bean
74+
public BigQueryAuthResolver bigQueryAuthResolver() {
75+
return new BigQueryAuthResolver();
76+
};
77+
7278
}

‎engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/KerberosSupport.java

-35
This file was deleted.

‎engine/src/main/java/com/odysseusinc/arachne/executionengine/execution/r/DockerService.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,10 @@
1515
import com.odysseusinc.arachne.execution_engine_common.api.v1.dto.AnalysisSyncRequestDTO;
1616
import com.odysseusinc.arachne.execution_engine_common.api.v1.dto.Stage;
1717
import com.odysseusinc.arachne.execution_engine_common.descriptor.dto.DockerEnvironmentDTO;
18+
import com.odysseusinc.arachne.executionengine.auth.AuthEffects;
1819
import com.odysseusinc.arachne.executionengine.config.properties.DockerRegistryProperties;
1920
import com.odysseusinc.arachne.executionengine.execution.Overseer;
2021
import com.odysseusinc.arachne.executionengine.util.Streams;
21-
import com.odysseusinc.datasourcemanager.krblogin.KrbConfig;
2222
import lombok.extern.slf4j.Slf4j;
2323
import org.apache.commons.lang3.exception.ExceptionUtils;
2424
import org.springframework.beans.factory.annotation.Autowired;
@@ -95,9 +95,9 @@ public DockerService(DockerRegistryProperties properties) {
9595

9696
@Override
9797
protected Overseer analyze(
98-
AnalysisSyncRequestDTO analysis, File analysisDir, BiConsumer<String, String> sendCallback, Integer updateInterval, KrbConfig krbConfig
98+
AnalysisSyncRequestDTO analysis, File analysisDir, BiConsumer<String, String> sendCallback, Integer updateInterval, AuthEffects auth
9999
) {
100-
List<String> env = buildRuntimeEnvVariables(analysis.getDataSource(), krbConfig.getIsolatedRuntimeEnvs()).entrySet().stream().map(entry ->
100+
List<String> env = buildRuntimeEnvVariables(analysis.getDataSource(), auth).entrySet().stream().map(entry ->
101101
entry.getKey() + "=" + entry.getValue()).collect(Collectors.toList()
102102
);
103103
Instant started = Instant.now();

0 commit comments

Comments
 (0)