Propagate context in case of an alias also (#93)

* add wraps

* add test

* work around alias

* changeset

Author: vogievetsky

.changeset/tiny-ties-sell.md

@@ -0,0 +1,5 @@
+---
+'druid-query-toolkit': patch
+---
+
+Propogate context in case of an alias also

src/sql/sql-case/sql-case.spec.ts

@@ -26,6 +26,32 @@ describe('CaseExpression', () => {
     backAndForth(sql, SqlCase);
   });
 
+  it('.ifThenElse', () => {
+    // Test with condition, then, and else
+    const condition = SqlExpression.parse('x > 5');
+    const thenExpr = SqlExpression.parse('"result is true"');
+    const elseExpr = SqlExpression.parse('"result is false"');
+
+    expect(SqlCase.ifThenElse(condition, thenExpr, elseExpr).toString()).toEqual(
+      'CASE WHEN x > 5 THEN "result is true" ELSE "result is false" END',
+    );
+
+    // Test with literal values
+    expect(SqlCase.ifThenElse(condition, 'yes', 'no').toString()).toEqual(
+      `CASE WHEN x > 5 THEN 'yes' ELSE 'no' END`,
+    );
+
+    // Test numeric literals
+    expect(SqlCase.ifThenElse(condition, 1, 0).toString()).toEqual(
+      'CASE WHEN x > 5 THEN 1 ELSE 0 END',
+    );
+
+    // Test without else expression
+    expect(SqlCase.ifThenElse(condition, thenExpr).toString()).toEqual(
+      'CASE WHEN x > 5 THEN "result is true" END',
+    );
+  });
+
   it('caseless CASE Expression', () => {
     const sql = `CASE WHEN B THEN C END`;
 

src/sql/sql-case/sql-case.ts

@@ -15,6 +15,7 @@
 import type { SqlBaseValue, SqlTypeDesignator, Substitutor } from '../sql-base';
 import { SqlBase } from '../sql-base';
 import { SqlExpression } from '../sql-expression';
+import type { LiteralValue } from '../sql-literal/sql-literal';
 import { SeparatedArray, SPACE } from '../utils';
 
 import { SqlWhenThenPart } from './sql-when-then-part';
@@ -34,14 +35,15 @@ export class SqlCase extends SqlExpression {
 
   static ifThenElse(
     conditionExpression: SqlExpression,
-    thenExpression: SqlExpression,
-    elseExpression?: SqlExpression,
+    thenExpression: SqlExpression | LiteralValue,
+    elseExpression?: SqlExpression | LiteralValue,
   ) {
     return new SqlCase({
       whenThenParts: SeparatedArray.fromSingleValue(
         SqlWhenThenPart.create(conditionExpression, thenExpression),
       ),
-      elseExpression,
+      elseExpression:
+        typeof elseExpression !== 'undefined' ? SqlExpression.wrap(elseExpression) : undefined,
     });
   }
 

src/sql/sql-case/sql-when-then-part.ts

@@ -15,6 +15,7 @@
 import type { SqlBaseValue, SqlTypeDesignator, Substitutor } from '../sql-base';
 import { SqlBase } from '../sql-base';
 import { SqlExpression } from '../sql-expression';
+import type { LiteralValue } from '../sql-literal/sql-literal';
 import { SeparatedArray, Separator } from '../utils';
 
 export interface SqlWhenThenPartValue extends SqlBaseValue {
@@ -30,13 +31,13 @@ export class SqlWhenThenPart extends SqlBase {
 
   static create(
     whenExpressions: SeparatedArray<SqlExpression> | SqlExpression[] | SqlExpression,
-    thenExpression: SqlExpression,
+    thenExpression: SqlExpression | LiteralValue,
   ): SqlWhenThenPart {
     return new SqlWhenThenPart({
       whenExpressions: SeparatedArray.fromArray(
         whenExpressions instanceof SqlExpression ? [whenExpressions] : whenExpressions,
       ),
-      thenExpression: SqlExpression.verify(thenExpression),
+      thenExpression: SqlExpression.wrap(thenExpression),
     });
   }
 

src/sql/sql-query/sql-query.spec.ts

@@ -183,7 +183,7 @@ describe('SqlQuery', () => {
       `);
     });
 
-    it('works in advanced case', () => {
+    it('works with context', () => {
       const query = SqlQuery.selectStarFrom(
         SqlQuery.selectStarFrom(SqlTable.create('lol')).changeContext({ a: 1 }),
       )
@@ -209,6 +209,33 @@ describe('SqlQuery', () => {
         WHERE channel  =  '#en.wikipedia'
       `);
     });
+
+    it('works with context and alias', () => {
+      const query = SqlQuery.selectStarFrom(
+        SqlQuery.selectStarFrom(SqlTable.create('lol')).changeContext({ a: 1 }).as('t'),
+      )
+        .changeSelectExpressions([
+          SqlColumn.create('channel'),
+          SqlColumn.create('page'),
+          SqlColumn.create('user'),
+          SqlColumn.create('as'),
+        ])
+        .changeWhereExpression(SqlExpression.parse(`channel  =  '#en.wikipedia'`));
+
+      expect(String(query)).toEqual(sane`
+        SET a = 1;
+        SELECT
+          "channel",
+          "page",
+          "user",
+          "as"
+        FROM (
+          SELECT *
+          FROM "lol"
+        ) AS "t"
+        WHERE channel  =  '#en.wikipedia'
+      `);
+    });
   });
 
   describe('.from', () => {

src/sql/sql-query/sql-query.ts

@@ -102,10 +102,13 @@ export class SqlQuery extends SqlExpression {
   static from(from: string | SqlExpression | SqlFromClause): SqlQuery {
     // Extract context from the inner query if given
     let contextStatements: SeparatedArray<SqlSetStatement> | undefined;
-    if (from instanceof SqlQuery || from instanceof SqlWithQuery) {
-      contextStatements = from.contextStatements;
-      if (contextStatements) {
-        from = from.changeContextStatements(undefined);
+    if (from instanceof SqlExpression) {
+      const underlyingFrom = from.getUnderlyingExpression();
+      if (underlyingFrom instanceof SqlQuery || underlyingFrom instanceof SqlWithQuery) {
+        contextStatements = underlyingFrom.contextStatements;
+        if (contextStatements) {
+          from = from.changeUnderlyingExpression(underlyingFrom.changeContextStatements(undefined));
+        }
       }
     }