From 22e6457605a7158e99acf869945466e67c11ce52 Mon Sep 17 00:00:00 2001
From: Christopher Swenson <cswenson@meta.com>
Date: Wed, 11 Dec 2024 12:43:25 -0600
Subject: [PATCH] Make `select: *` work with composite sources

---
 .../src/lang/ast/field-space/query-spaces.ts  | 51 ++++++++++---------
 .../databases/all/composite_sources.spec.ts   | 10 ++++
 2 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/packages/malloy/src/lang/ast/field-space/query-spaces.ts b/packages/malloy/src/lang/ast/field-space/query-spaces.ts
index e21abb8d7..fb12006b4 100644
--- a/packages/malloy/src/lang/ast/field-space/query-spaces.ts
+++ b/packages/malloy/src/lang/ast/field-space/query-spaces.ts
@@ -359,33 +359,34 @@ export abstract class QuerySpace extends QueryOperationSpace {
         const wildPath = this.expandedWild[name];
         if (wildPath) {
           fields.push({type: 'fieldref', path: wildPath.path});
-          continue;
-        }
-        // TODO handle wildcards for composite sources
-        const fieldQueryDef = field.getQueryFieldDef(this.exprSpace);
-        if (fieldQueryDef) {
-          const typeDesc = field.typeDesc();
-          nextCompositeFieldUsage = typeDesc.compositeFieldUsage;
-          // Filter out fields whose type is 'error', which means that a totally bad field
-          // isn't sent to the compiler, where it will wig out.
-          // TODO Figure out how to make errors generated by `canContain` go in the right place,
-          // maybe by adding a logable element to SpaceFields.
-          if (
-            typeDesc &&
-            typeDesc.type !== 'error' &&
-            this.canContain(typeDesc) &&
-            !isEmptyNest(fieldQueryDef)
-          ) {
-            fields.push(fieldQueryDef);
+          nextCompositeFieldUsage =
+            wildPath.entry.typeDesc().compositeFieldUsage;
+        } else {
+          const fieldQueryDef = field.getQueryFieldDef(this.exprSpace);
+          if (fieldQueryDef) {
+            const typeDesc = field.typeDesc();
+            nextCompositeFieldUsage = typeDesc.compositeFieldUsage;
+            // Filter out fields whose type is 'error', which means that a totally bad field
+            // isn't sent to the compiler, where it will wig out.
+            // TODO Figure out how to make errors generated by `canContain` go in the right place,
+            // maybe by adding a logable element to SpaceFields.
+            if (
+              typeDesc &&
+              typeDesc.type !== 'error' &&
+              this.canContain(typeDesc) &&
+              !isEmptyNest(fieldQueryDef)
+            ) {
+              fields.push(fieldQueryDef);
+            }
           }
+          // TODO I removed the error here because during calculation of the refinement space,
+          // (see creation of a QuerySpace) we add references to all the fields from
+          // the refinement, but they don't have definitions. So in the case where we
+          // don't have a field def, we "know" that that field is already in the query,
+          // and we don't need to worry about actually adding it. Previously, this was also true for
+          // project statements, where we added "*" as a field and also all the individual
+          // fields, but the individual fields didn't have field defs.
         }
-        // TODO I removed the error here because during calculation of the refinement space,
-        // (see creation of a QuerySpace) we add references to all the fields from
-        // the refinement, but they don't have definitions. So in the case where we
-        // don't have a field def, we "know" that that field is already in the query,
-        // and we don't need to worry about actually adding it. Previously, this was also true for
-        // project statements, where we added "*" as a field and also all the individual
-        // fields, but the individual fields didn't have field defs.
       }
       const next = this.applyNextCompositeFieldUsage(
         source,
diff --git a/test/src/databases/all/composite_sources.spec.ts b/test/src/databases/all/composite_sources.spec.ts
index 816100fc8..9bc70ee97 100644
--- a/test/src/databases/all/composite_sources.spec.ts
+++ b/test/src/databases/all/composite_sources.spec.ts
@@ -243,4 +243,14 @@ describe.each(runtimes.runtimeList)('%s', (databaseName, runtime) => {
       }
     `).malloyResultMatches(runtime, {'foo.x': 1});
   });
+  it('composite with select *', async () => {
+    await expect(`
+      ##! experimental.composite_sources
+      source: state_facts is ${databaseName}.table('malloytest.state_facts')
+      source: x is compose(state_facts, state_facts extend { dimension: foo is 1 }) extend {
+        accept: foo
+      }
+      run: x -> { select: * }
+    `).malloyResultMatches(runtime, {foo: 1});
+  });
 });