From c98648f676ef1f8864e2415b83f4a03cdc67497c Mon Sep 17 00:00:00 2001
From: Paul Harrison <paul@harrison.sh>
Date: Tue, 2 Apr 2024 13:18:07 +0100
Subject: [PATCH] Do not require mocks for ARRAY JOIN clause arguments

The `ARRAY JOIN` clause in Clickhouse is currently recognised as a table
that requires a mock by the library. This is not the case as this clause
is used for expanding arrays in the source table into separate rows.
This commit adds an additional clause to the `get_source_tables` helper
function to ignore joins of kind `ARRAY`.

Closes #48
---
 CHANGELOG.md                   |  1 +
 src/sql_mock/helpers.py        |  4 ++++
 tests/sql_mock/test_helpers.py | 24 ++++++++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f2bc49..5a1f407 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * Add default target path for dbt
 * Improve replacement of tables (also taking into account missing alias)
+* Do not require mocks for ARRAY JOIN clause arguments
 
 ## [0.6.0]
 
diff --git a/src/sql_mock/helpers.py b/src/sql_mock/helpers.py
index 4d975f4..36bd88e 100644
--- a/src/sql_mock/helpers.py
+++ b/src/sql_mock/helpers.py
@@ -144,6 +144,10 @@ def get_source_tables(query, dialect) -> List[str]:
         #     then `source` will be a Table instance.
         for alias, (node, source) in scope.selected_sources.items()
         if isinstance(source, sqlglot.exp.Table)
+        # When using ARRAY joins sqlglot percieves the inputs as tables even though they are infact not.
+        # This fixes it but does not allow for multiple types of joins to be mixed with the ARRAY JOIN,
+        # For now we consider it a reasonable solution.
+        and not (node.parent.key == "join" and any(join.kind == "ARRAY" for join in node.parent_select.args["joins"]))
     }
 
     return list(tables)
diff --git a/tests/sql_mock/test_helpers.py b/tests/sql_mock/test_helpers.py
index d9dbc73..fdb24d9 100644
--- a/tests/sql_mock/test_helpers.py
+++ b/tests/sql_mock/test_helpers.py
@@ -321,3 +321,27 @@ def test_query_with_comment(self):
 
         expected = ["table_1"]
         assert res == expected
+
+    def test_query_with_array_joins(self):
+        """...then the array join fields should not be treated as tables to be mocked"""
+
+        query = """
+        SELECT
+            sum(1) AS impressions,
+            city,
+            browser
+        FROM
+        (
+            SELECT
+                ['Istanbul', 'Berlin', 'Bobruisk'] AS cities,
+                ['Firefox', 'Chrome', 'Chrome'] AS browsers
+        )
+        ARRAY JOIN
+            cities AS city,
+            browsers AS browser
+        """
+
+        res = get_source_tables(query, dialect="clickhouse")
+
+        expected = []
+        assert res == expected