From 5581758bf314a4b9df507d46a243dc8b46d52cde Mon Sep 17 00:00:00 2001
From: Josh Mandel <jmandel@alum.mit.edu>
Date: Tue, 17 Oct 2023 18:06:16 -0500
Subject: [PATCH] Ensure unions have aligned column names

---
 tests/content/union.json                    | 29 +++++++++++++++++++--
 tests/generate.test.js                      | 13 ++++++---
 tests/reference-implementation/extractor.js | 20 +++++++++-----
 3 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/tests/content/union.json b/tests/content/union.json
index f6e836f..24062db 100644
--- a/tests/content/union.json
+++ b/tests/content/union.json
@@ -51,11 +51,17 @@
             "union": [
               {
                 "forEach": "name",
-                "column": [{ "path": "family", "name": "a" }]
+                "column": [
+                  { "path": "family", "name": "a" },
+                  { "path": "{}", "name": "b" }
+                ]
               },
               {
                 "forEach": "contact.name",
-                "column": [{ "path": "family", "name": "b" }]
+                "column": [
+                  { "path": "family", "name": "b" },
+                  { "path": "{}", "name": "a" }
+                ]
               }
             ]
           }
@@ -67,6 +73,25 @@
         { "a": null, "b": "f3" }
       ]
     },
+    {
+      "title": "Union with mismatching branches",
+      "view": {
+        "resource": "Patient",
+        "select": [
+          {
+            "union": [
+              {
+                "column": [{ "path": "'a'", "name": "a" }]
+              },
+              {
+                "column": [{ "path": "'b'", "name": "b" }]
+              }
+            ]
+          }
+        ]
+      },
+      "expectError": true
+    },
     {
       "title": "Union with forEachOrNull",
       "view": {
diff --git a/tests/generate.test.js b/tests/generate.test.js
index 4913edd..7f840e9 100644
--- a/tests/generate.test.js
+++ b/tests/generate.test.js
@@ -1,4 +1,4 @@
-import { runTests } from './reference-implementation/index.js'
+import { runTests, getColumns } from './reference-implementation/index.js'
 import fhirpath from 'fhirpath'
 import Ajv from 'ajv'
 import fs from 'fs'
@@ -18,6 +18,7 @@ function validatePathToSubset(path) {
     'Identifier',
     'LiteralTerm',
     'BooleanLiteral',
+    'NullLiteral',
     'StringLiteral',
     'NumberLiteral',
     'MemberInvocation',
@@ -71,7 +72,10 @@ function validatePathToSubset(path) {
     }
   }
   const ast = fhirpath.parse(path)
-  return validateChildren(ast)
+  const problems = validateChildren(ast)
+  if (problems) {
+    throw problems
+  }
 }
 
 function buildFhirpathFormat(allowExtendedFhirpath) {
@@ -79,10 +83,13 @@ function buildFhirpathFormat(allowExtendedFhirpath) {
     type: 'string',
     validate: (v) => {
       try {
-        if (!allowExtendedFhirpath && validatePathToSubset(v)) return false
+        if (!allowExtendedFhirpath) {
+          validatePathToSubset(v)
+        }
         fhirpath.compile(v)
         return true
       } catch (err) {
+        console.log('Invalid fhirpath', v, err)
         return false
       }
     },
diff --git a/tests/reference-implementation/extractor.js b/tests/reference-implementation/extractor.js
index 171e6dd..7f7858b 100644
--- a/tests/reference-implementation/extractor.js
+++ b/tests/reference-implementation/extractor.js
@@ -85,6 +85,18 @@ function compileViewDefinition(viewDefinition) {
   for (let field of subViews(viewDefinition)) {
     compileViewDefinition(field)
   }
+
+  const cols = (viewDefinition.union ?? [])
+    .map((u) =>
+      getColumns({ select: [u] })
+        .reduce((acc, c) => acc.concat(c.name), [])
+        .sort(),
+    )
+    .map((u) => JSON.stringify(u))
+
+  if (cols.some((c) => c !== cols[0])) {
+    throw `Unions use different columns: ${cols}`
+  }
 }
 
 function cartesianProduct([first, ...rest]) {
@@ -122,15 +134,9 @@ function extractFields(obj, viewDefinition, context = {}) {
     }
 
     const unionBindings = []
-    const unionColumns = getColumns({ union })
-      .reduce((acc, c) => {
-      acc[c.name] = null
-      return acc
-    }, {})
-
     for (const u of union ?? []) {
       for (const r of extract(nestedObject, { select: [u] }, context)) {
-        unionBindings.push({ ...unionColumns, ...r })
+        unionBindings.push(r)
       }
     }