From 1e1c9b8377f298ba29d0260a0b707de41668007f Mon Sep 17 00:00:00 2001
From: nszoni <namson98@gmail.com>
Date: Wed, 14 Feb 2024 16:10:51 +0100
Subject: [PATCH] alter generate series to fit synapse limitations and add
 tests

---
 .../synapse/macros/utils/generate_series.sql  | 61 +++++++++++++++++++
 .../adapter/test_generate_series.py           | 48 +++++++++++++++
 2 files changed, 109 insertions(+)
 create mode 100644 dbt/include/synapse/macros/utils/generate_series.sql
 create mode 100644 tests/functional/adapter/test_generate_series.py

diff --git a/dbt/include/synapse/macros/utils/generate_series.sql b/dbt/include/synapse/macros/utils/generate_series.sql
new file mode 100644
index 00000000..88fac113
--- /dev/null
+++ b/dbt/include/synapse/macros/utils/generate_series.sql
@@ -0,0 +1,61 @@
+{#
+Specifying more than one WITH clause in a CTE isn't allowed. For example, if a CTE query definition contains a subquery,
+that subquery can't contain a nested WITH clause that defines another CTE.
+#}
+
+{% macro get_powers_of_two(upper_bound) %}
+    {{ return(adapter.dispatch('get_powers_of_two', 'dbt')(upper_bound)) }}
+{% endmacro %}
+
+{% macro default__get_powers_of_two(upper_bound) %}
+
+    {% if upper_bound <= 0 %}
+    {{ exceptions.raise_compiler_error("upper bound must be positive") }}
+    {% endif %}
+
+    {% for _ in range(1, 100) %}
+       {% if upper_bound <= 2 ** loop.index %}{{ return(loop.index) }}{% endif %}
+    {% endfor %}
+
+{% endmacro %}
+
+
+{% macro generate_series(upper_bound) %}
+    {{ return(adapter.dispatch('generate_series', 'dbt')(upper_bound)) }}
+{% endmacro %}
+
+{% macro default__generate_series(upper_bound) %}
+
+    {% set n = dbt.get_powers_of_two(upper_bound) %}
+
+    with p as (
+        select 0 as generated_number union all select 1
+    ), unioned as (
+
+    select
+
+    {% for i in range(n) %}
+    p{{i}}.generated_number * power(2, {{i}})
+    {% if not loop.last %} + {% endif %}
+    {% endfor %}
+    + 1
+    as generated_number
+
+    from
+
+    {% for i in range(n) %}
+    p as p{{i}}
+    {% if not loop.last %} cross join {% endif %}
+    {% endfor %}
+
+    ),
+
+    generate_series as (
+    select *
+    from unioned
+    where generated_number <= {{upper_bound}}
+    )
+
+    select * from generate_series
+
+{% endmacro %}
diff --git a/tests/functional/adapter/test_generate_series.py b/tests/functional/adapter/test_generate_series.py
new file mode 100644
index 00000000..2e726dfe
--- /dev/null
+++ b/tests/functional/adapter/test_generate_series.py
@@ -0,0 +1,48 @@
+import pytest
+from dbt.tests.adapter.utils.base_utils import BaseUtils
+from dbt.tests.adapter.utils.fixture_generate_series import models__test_generate_series_yml
+
+# Cause of overriding fixture: Specifying more than one WITH clause in a CTE isn't allowed.
+# For example, if a CTE query definition contains a subquery,
+# that subquery can't contain a nested WITH clause that defines another CTE.
+
+models__test_generate_series_sql = """
+    {{ dbt.generate_series(10) }}
+    left join (
+        select 1 as expected
+        union all
+        select 2 as expected
+        union all
+        select 3 as expected
+        union all
+        select 4 as expected
+        union all
+        select 5 as expected
+        union all
+        select 6 as expected
+        union all
+        select 7 as expected
+        union all
+        select 8 as expected
+        union all
+        select 9 as expected
+        union all
+        select 10 as expected
+    ) as expected_numbers
+    on generate_series.generated_number = expected_numbers.expected
+"""
+
+
+class BaseGenerateSeries(BaseUtils):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "test_generate_series.yml": models__test_generate_series_yml,
+            "test_generate_series.sql": self.interpolate_macro_namespace(
+                models__test_generate_series_sql, "generate_series"
+            ),
+        }
+
+
+class TestGenerateSeries(BaseGenerateSeries):
+    pass