From 10ebae293c6a6f25d237d462ae4a60fb6e434660 Mon Sep 17 00:00:00 2001
From: Marco Baumgartl <marco.baumgartl@stock3.com>
Date: Wed, 31 Jul 2024 15:46:09 +0200
Subject: [PATCH] fix: Handle JOIN table dependencies

---
 lib/sql-query-optimizer.js            | 33 ++++++++----
 test/unit/sql-query-optimizer.spec.js | 73 +++++++++++++++++++++++++++
 2 files changed, 95 insertions(+), 11 deletions(-)

diff --git a/lib/sql-query-optimizer.js b/lib/sql-query-optimizer.js
index 65696ae..043a056 100644
--- a/lib/sql-query-optimizer.js
+++ b/lib/sql-query-optimizer.js
@@ -33,23 +33,34 @@ const getAttributeExpressions = filterTree(
  * @private
  */
 function getRequiredTables(ast) {
-    const tableExpressions = ['columns', 'where', 'groupby', 'orderby']
+    const tblExpressions = ['columns', 'where', 'groupby', 'orderby']
         .filter((clause) => Object.hasOwn(ast, clause) && ast[clause] !== null)
         .flatMap((clause) => getTableExpressions(ast[clause]));
 
     if (ast.from.length > 1) {
-        ast.from
-            // extract joined tables that are not marked as required
-            .filter((expr) => {
-                const table = expr.as !== null ? expr.as : expr.table;
-                return expr.join !== undefined && tableExpressions.some((tblExpr) => tblExpr.table === table);
-            })
-            // get tables used in joins
-            .forEach((join) => tableExpressions.push(...getTableExpressions(join.on)));
+        const joins = ast.from.filter((expr) => expr.on !== undefined);
+        const joinTblDeps = joins.reduce((tblDeps, expr) => {
+            const depTblExprs = getTableExpressions(expr).filter((tblExpr) => expr.table !== tblExpr.table);
+            return {
+                ...tblDeps,
+                ...(depTblExprs.length ? { [expr.table]: depTblExprs } : {})
+            };
+        }, {});
+
+        tblExpressions.push(
+            // add joined tables not marked as required yet
+            ...joins
+                .filter((expr) => {
+                    const table = expr.as !== null ? expr.as : expr.table;
+                    return tblExpressions.some((tblExpr) => tblExpr.table === table);
+                })
+                .flatMap((expr) => getTableExpressions(expr.on))
+        );
+
+        tblExpressions.forEach((tblExpr) => tblExpressions.push(...(joinTblDeps[tblExpr.table] ?? [])));
     }
 
-    const requiredTables = tableExpressions.map(({ table }) => table);
-
+    const requiredTables = tblExpressions.map(({ table }) => table);
     return unique(requiredTables);
 }
 
diff --git a/test/unit/sql-query-optimizer.spec.js b/test/unit/sql-query-optimizer.spec.js
index 9be8df9..8454932 100644
--- a/test/unit/sql-query-optimizer.spec.js
+++ b/test/unit/sql-query-optimizer.spec.js
@@ -109,6 +109,79 @@ describe('SQL query optimizer', () => {
         assert.deepEqual(optimizedAst.from, [{ db: null, table: 't', as: null }]);
     });
 
+    it('should not remove JOIN dependencies from AST', () => {
+        /* should not remove t1 LEFT JOIN because it's required for JOINs on t2/t3
+         * SELECT t.id
+         * FROM t
+         * LEFT JOIN t1 ON t.id = t1.id
+         * LEFT JOIN t2 ON t1.id = t2.id
+         * LEFT JOIN t3 ON t2.id = t3.id
+         * WHERE t3.id = 1
+         */
+        const from = [
+            { db: null, table: 't', as: null },
+            {
+                db: null,
+                table: 't1',
+                as: null,
+                join: 'LEFT JOIN',
+                on: {
+                    type: 'binary_expr',
+                    operator: '=',
+                    left: { type: 'column_ref', table: 't', column: 'id' },
+                    right: { type: 'column_ref', table: 't1', column: 'id' }
+                }
+            },
+            {
+                db: null,
+                table: 't2',
+                as: null,
+                join: 'LEFT JOIN',
+                on: {
+                    type: 'binary_expr',
+                    operator: '=',
+                    left: { type: 'column_ref', table: 't1', column: 'id' },
+                    right: { type: 'column_ref', table: 't2', column: 'id' }
+                }
+            },
+            {
+                db: null,
+                table: 't3',
+                as: null,
+                join: 'LEFT JOIN',
+                on: {
+                    type: 'binary_expr',
+                    operator: '=',
+                    left: { type: 'column_ref', table: 't2', column: 'id' },
+                    right: { type: 'column_ref', table: 't3', column: 'id' }
+                }
+            }
+        ];
+        const optimizedAst = optimize(
+            {
+                with: null,
+                type: 'select',
+                options: null,
+                distinct: null,
+                columns: [{ expr: { type: 'column_ref', table: 't', column: 'id' }, as: null }],
+                from,
+                where: {
+                    type: 'binary_expr',
+                    operator: '=',
+                    left: { type: 'column_ref', table: 't3', column: 'id' },
+                    right: { type: 'number', value: 1 }
+                },
+                groupby: null,
+                having: null,
+                orderby: null,
+                limit: null
+            },
+            ['id']
+        );
+
+        assert.deepEqual(optimizedAst.from, from);
+    });
+
     it('should pay attention to table aliases', () => {
         /**
          * t3 must not be removed because it's used by it's alias