diff --git a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/SampleList.java b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/SampleList.java index 1c3b07717d1..a18712f14b0 100644 --- a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/SampleList.java +++ b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/SampleList.java @@ -7,8 +7,6 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.util.*; -import java.util.stream.Collectors; -import org.apache.commons.lang.StringUtils; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.broadinstitute.hellbender.utils.bigquery.BigQueryUtils; @@ -20,9 +18,9 @@ public class SampleList { private Map sampleIdMap = new HashMap<>(); private Map sampleNameMap = new HashMap<>(); - public SampleList(String sampleTableName, File sampleFile, boolean printDebugInformation) { + public SampleList(String sampleTableName, File sampleFile, String executionProjectId, boolean printDebugInformation) { if (sampleTableName != null) { - initializeMaps(new TableReference(sampleTableName, SchemaUtils.SAMPLE_FIELDS), printDebugInformation); + initializeMaps(new TableReference(sampleTableName, SchemaUtils.SAMPLE_FIELDS), executionProjectId, printDebugInformation); } else if (sampleFile != null) { initializeMaps(sampleFile); } else { @@ -50,26 +48,8 @@ public Map getMap() { return sampleIdMap; } -// protected Map getSampleNameMap(TableReference sampleTable, List samples, boolean printDebugInformation) { -// Map results = new HashMap<>(); -// // create optional where clause -// String whereClause = ""; -// if (samples != null && samples.size() > 0) { -// whereClause = " WHERE " + SchemaUtils.SAMPLE_NAME_FIELD_NAME + " in (\'" + StringUtils.join(samples, "\',\'") + "\') "; -// } -// -// TableResult queryResults = querySampleTable(sampleTable.getFQTableName(), whereClause, printDebugInformation); -// -// // Add our samples to our map: -// for (final FieldValueList row : queryResults.iterateAll()) { -// results.put(row.get(1).getStringValue(), (int) row.get(0).getLongValue()); -// } -// return results; -// } - - - protected void initializeMaps(TableReference sampleTable, boolean printDebugInformation) { - TableResult queryResults = querySampleTable(sampleTable.getFQTableName(), "", printDebugInformation); + protected void initializeMaps(TableReference sampleTable, String executionProjectId, boolean printDebugInformation) { + TableResult queryResults = querySampleTable(sampleTable.getFQTableName(), "", executionProjectId, printDebugInformation); // Add our samples to our map: for (final FieldValueList row : queryResults.iterateAll()) { @@ -95,15 +75,14 @@ protected void initializeMaps(File cohortSampleFile) { } } - private TableResult querySampleTable(String fqSampleTableName, String whereClause, boolean printDebugInformation) { + private TableResult querySampleTable(String fqSampleTableName, String whereClause, String executionProjectId, boolean printDebugInformation) { // Get the query string: final String sampleListQueryString = "SELECT " + SchemaUtils.SAMPLE_ID_FIELD_NAME + ", " + SchemaUtils.SAMPLE_NAME_FIELD_NAME + - " FROM `" + fqSampleTableName + "`" + whereClause; + " FROM `" + fqSampleTableName + "`" + whereClause; - - // Execute the query: - final TableResult result = BigQueryUtils.executeQuery(sampleListQueryString); + // Execute the query: + final TableResult result = BigQueryUtils.executeQuery(BigQueryUtils.getBigQueryEndPoint(executionProjectId) , sampleListQueryString, false); // Show our pretty results: if (printDebugInformation) { diff --git a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/arrays/ArrayExtractCohort.java b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/arrays/ArrayExtractCohort.java index 8851c432315..2ea182cef93 100644 --- a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/arrays/ArrayExtractCohort.java +++ b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/arrays/ArrayExtractCohort.java @@ -198,7 +198,7 @@ protected void onStartup() { vcfWriter = createVCFWriter(IOUtils.getPath(outputVcfPathString)); - SampleList sampleIdMap = new SampleList(sampleTableName, cohortSampleFile, printDebugInformation); + SampleList sampleIdMap = new SampleList(sampleTableName, cohortSampleFile, readProjectID, printDebugInformation); // Map sampleIdMap; VCFHeader header = CommonCode.generateRawArrayVcfHeader(new HashSet<>(sampleIdMap.getSampleNames()), reference.getSequenceDictionary()); diff --git a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/arrays/CreateImputedIngestFiles.java b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/arrays/CreateImputedIngestFiles.java index 221e35626ef..6fd997b232c 100644 --- a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/arrays/CreateImputedIngestFiles.java +++ b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/arrays/CreateImputedIngestFiles.java @@ -64,7 +64,7 @@ public void onTraversalStart() { // Get sample name final VCFHeader inputVCFHeader = getHeaderForVariants(); TableReference sampleTable = new TableReference(sampleListFQTablename, SchemaUtils.SAMPLE_FIELDS); - sampleNameMap = new SampleList(sampleListFQTablename, null, printDebugInformation); + sampleNameMap = new SampleList(sampleListFQTablename, null, null, printDebugInformation); tableToCreatorMap = new HashMap<>(); Map> tableNumberToSampleList = new HashMap<>(); diff --git a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractCohort.java b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractCohort.java index a90362e374d..59583abd741 100644 --- a/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractCohort.java +++ b/src/main/java/org/broadinstitute/hellbender/tools/variantdb/nextgen/ExtractCohort.java @@ -58,7 +58,7 @@ public class ExtractCohort extends ExtractTool { protected void onStartup() { super.onStartup(); - SampleList sampleList = new SampleList(sampleTableName, sampleFileName, printDebugInformation); + SampleList sampleList = new SampleList(sampleTableName, sampleFileName, projectID, printDebugInformation); Set sampleNames = new HashSet<>(sampleList.getSampleNames()); VCFHeader header = CommonCode.generateVcfHeader(sampleNames, reference.getSequenceDictionary()); diff --git a/src/main/java/org/broadinstitute/hellbender/utils/bigquery/BigQueryUtils.java b/src/main/java/org/broadinstitute/hellbender/utils/bigquery/BigQueryUtils.java index c335bbf24db..845211290b8 100644 --- a/src/main/java/org/broadinstitute/hellbender/utils/bigquery/BigQueryUtils.java +++ b/src/main/java/org/broadinstitute/hellbender/utils/bigquery/BigQueryUtils.java @@ -40,6 +40,15 @@ public static BigQuery getBigQueryEndPoint() { return BigQueryOptions.getDefaultInstance().getService(); } + /** + * @param executionProjectId The google project that should be used to execute this query + * + * @return A {@link BigQuery} object that can be used to interact with a BigQuery data set. + */ + public static BigQuery getBigQueryEndPoint(String executionProjectId) { + return (executionProjectId != null) ? BigQueryOptions.newBuilder().setProjectId(executionProjectId).build().getService() : getBigQueryEndPoint(); + } + /** * Executes the given {@code queryString} on the default instance of {@link BigQuery} as created by {@link #getBigQueryEndPoint()}. * Will block until results are returned.