Replaced Resultset with VOTableDocument

opencadc · Dec 10, 2024 · dccb870 · dccb870
1 parent 841d6b8
commit dccb870
Showing 1 changed file with 47 additions and 33 deletions.
diff --git a/cadc-tap-server/src/main/java/ca/nrc/cadc/tap/DefaultTableWriter.java b/cadc-tap-server/src/main/java/ca/nrc/cadc/tap/DefaultTableWriter.java
@@ -311,7 +311,8 @@ public void write(ResultSet rs, OutputStream out) throws IOException
     public void write(ResultSet rs, OutputStream out, Long maxrec)
             throws IOException {
         if (this.getContentType().equals("application/vnd.apache.parquet")) {
-            tableWriter.write(rs, out, maxrec == null ? Long.MAX_VALUE : maxrec);
+            VOTableDocument voTableDocument = prepareVOTableDocument(rs);
+            tableWriter.write(voTableDocument, out, maxrec == null ? Long.MAX_VALUE : maxrec);
         } else {
             Writer writer = new BufferedWriter(new OutputStreamWriter(out, "UTF-8"));
             this.write(rs, writer, maxrec);
@@ -324,26 +325,7 @@ public void write(ResultSet rs, Writer out) throws IOException
         this.write(rs, out, null);
     }
 
-    @Override
-    public void write(ResultSet rs, Writer out, Long maxrec) throws IOException
-    {
-        if (rs != null && log.isDebugEnabled())
-            try { log.debug("resultSet column count: " + rs.getMetaData().getColumnCount()); }
-            catch(Exception oops) { log.error("failed to check resultset column count", oops); }
-
-        if (rssTableWriter != null)
-        {
-            rssTableWriter.setJob(job);
-            rssTableWriter.setSelectList(selectList);
-            rssTableWriter.setFormatFactory(formatFactory);
-            rssTableWriter.setQueryInfo(queryInfo);
-            if (maxrec != null)
-                rssTableWriter.write(rs, out, maxrec);
-            else
-                rssTableWriter.write(rs, out);
-            return;
-        }
-
+    private VOTableDocument prepareVOTableDocument(ResultSet rs) throws IOException {
         VOTableDocument votableDocument = new VOTableDocument();
 
         VOTableResource resultsResource = new VOTableResource("results");
@@ -353,9 +335,27 @@ public void write(ResultSet rs, Writer out, Long maxrec) throws IOException
         List<Format<Object>> formats = formatFactory.getFormats(selectList);
 
         List<String> serviceIDs = new ArrayList<String>();
-        int listIndex = 0;
 
         // Add the metadata elements.
+        addVOTableMetadata(formats, resultsTable, serviceIDs);
+
+        resultsResource.setTable(resultsTable);
+        votableDocument.getResources().add(resultsResource);
+
+        // Add the "meta" resources to describe services for each columnID in list columnIDs that we recognize
+        addMetaResources(votableDocument, serviceIDs);
+
+        setVOTableInfo(resultsResource);
+
+        ResultSetTableData tableData = new ResultSetTableData(rs, formats);
+        resultsTable.setTableData(tableData);
+        this.rowcount = tableData.getRowCount(); // TODO: it was getting set after the write method was called
+
+        return votableDocument;
+    }
+
+    private void addVOTableMetadata(List<Format<Object>> formats, VOTableTable resultsTable, List<String> serviceIDs) {
+        int listIndex = 0;
         for (TapSelectItem resultCol : selectList)
         {
             VOTableField newField = createVOTableField(resultCol);
@@ -376,14 +376,9 @@ public void write(ResultSet rs, Writer out, Long maxrec) throws IOException
 
             listIndex++;
         }
+    }
 
-        resultsResource.setTable(resultsTable);
-        votableDocument.getResources().add(resultsResource);
-
-        // Add the "meta" resources to describe services for each columnID in
-        // list columnIDs that we recognize
-        addMetaResources(votableDocument, serviceIDs);
-
+    private void setVOTableInfo(VOTableResource resultsResource) {
         VOTableInfo info = new VOTableInfo("QUERY_STATUS", "OK");
         resultsResource.getInfos().add(info);
 
@@ -397,16 +392,35 @@ public void write(ResultSet rs, Writer out, Long maxrec) throws IOException
             info = new VOTableInfo("QUERY", queryInfo);
             resultsResource.getInfos().add(info);
         }
+    }
+
+    @Override
+    public void write(ResultSet rs, Writer out, Long maxrec) throws IOException
+    {
+        if (rs != null && log.isDebugEnabled())
+            try { log.debug("resultSet column count: " + rs.getMetaData().getColumnCount()); }
+            catch(Exception oops) { log.error("failed to check resultset column count", oops); }
+
+        if (rssTableWriter != null)
+        {
+            rssTableWriter.setJob(job);
+            rssTableWriter.setSelectList(selectList);
+            rssTableWriter.setFormatFactory(formatFactory);
+            rssTableWriter.setQueryInfo(queryInfo);
+            if (maxrec != null)
+                rssTableWriter.write(rs, out, maxrec);
+            else
+                rssTableWriter.write(rs, out);
+            return;
+        }
+
+        VOTableDocument votableDocument = prepareVOTableDocument(rs);
 
-        ResultSetTableData tableData = new ResultSetTableData(rs, formats);
-        resultsTable.setTableData(tableData);
-
         if (maxrec != null)
             tableWriter.write(votableDocument, out, maxrec);
         else
             tableWriter.write(votableDocument, out);
 
-        this.rowcount = tableData.getRowCount();
     }
 
     // HACK: need to allow an ObsCore.access_url formatter to access this info