Replace ARRAY_CONSTRUCT with Array (#19)

* f

* gitignore

* f

* vbump + lint

Author: DavidMias

.gitignore

@@ -1 +1,12 @@
 .vscode
+
+# Pycharm
+.idea
+.idea/
+*.DS_Store
+.python-version
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class

CHANGELOG.md

@@ -1,5 +1,7 @@
 # Yellowbox Snowglobe Changelog
-## NEXT
+## 0.2.6
+### Added
+* Added support for ARRAY_CONSTRUCT for PostgreSQL
 ### Internal
 * removed cahce action from cicd
 * removed autofix from lint script

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "yellowbox-snowglobe"
-version = "0.2.5"
+version = "0.2.6"
 description = ""
 authors = ["Biocatch LTD <serverteam@biocatch.com>"]
 license = "MIT"
@@ -9,7 +9,7 @@ license = "MIT"
 python = "^3.8"
 yellowbox = { version = ">=0.7.0", extras = ["postgresql", "webserver"] }
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
 pytest = "*"
 pytest-asyncio = "*"
 pytest-cov = "*"

pyproject.toml

@@ -18,7 +18,7 @@ def test_connect(docker_client):
         assert results == [(2, "two"), (3, "three")]
 
 
-@mark.asyncio()
+@mark.asyncio
 async def test_connect_async(docker_client):
     async with SnowGlobeService.arun(docker_client) as service:
         with connector.connect(**service.local_connection_kwargs()) as conn:
@@ -32,7 +32,7 @@ async def test_connect_async(docker_client):
             assert results == [(2, "two"), (3, "three")]
 
 
-@mark.asyncio()
+@mark.asyncio
 async def test_simultaneous_connections(docker_client):
     async with SnowGlobeService.arun(docker_client) as service:
         with connector.connect(**service.local_connection_kwargs()) as conn:

tests/blackbox/test_simple.py

@@ -19,7 +19,7 @@ def test_lateral_flatten(connection):
     )
     res = (
         connection.cursor()
-        .execute("select x, root.value from bar, lateral flatten(d) as root" " where root.value*root.value = x;")
+        .execute("select x, root.value from bar, lateral flatten(d) as root where root.value*root.value = x;")
         .fetchall()
     )
     assert res == [(4, 2), (9, 3)]
@@ -56,7 +56,7 @@ def test_bools(connection):
 def test_null_bools_and_dates(connection):
     connection.cursor().execute("create table bar (x timestamp, y boolean)")
     connection.cursor().execute(
-        "insert into bar values " "('2014-01-01 16:00:00', true)," " (null, false), " "('2023-01-08 17:00:00', null)"
+        "insert into bar values ('2014-01-01 16:00:00', true), (null, false), ('2023-01-08 17:00:00', null)"
     )
     res = connection.cursor().execute("select * from bar;").fetchall()
     assert res == [(datetime(2014, 1, 1, 16, 0), True), (None, False), (datetime(2023, 1, 8, 17, 0), None)]
@@ -82,3 +82,31 @@ def test_json(connection, db, query, expected):
     connection.cursor().execute("""insert into bar values (1, '{"a":"1", "b":"1"}'), (2, '{"a":"2", "b":"2"}')""")
     res = connection.cursor().execute(query).fetchall()
     assert res == expected
+
+
+@mark.parametrize(
+    ("query", "expected"),
+    [
+        (
+            "select ARRAY_CONSTRUCT(NULL, x) from bar;",
+            [
+                ([None, "hello"],),
+            ],
+        ),
+        (
+            "select ARRAY_CONSTRUCT(ARRAY_CONSTRUCT(NULL, x)) from bar;",
+            [
+                (
+                    [
+                        [None, "hello"],
+                    ],
+                ),
+            ],
+        ),
+    ],
+)
+def test_array(connection, db, query, expected):
+    connection.cursor().execute("create table bar (x text)")
+    connection.cursor().execute("""insert into bar (x) values ('hello')""")
+    res = connection.cursor().execute(query).fetchall()
+    assert res == expected

tests/blackbox/test_syntax.py

@@ -57,6 +57,57 @@ def split_sql_to_statements(query: str) -> Iterator[str]:
         yield last_bit
 
 
+def find_matching_paren(text: str, start_pos: int) -> int:
+    """
+    Find the position of the closing parenthesis ')' corresponding to the opening parenthesis '(' at start_pos.
+    Handle nested parentheses by counting open/closed parentheses.
+    """
+    if start_pos >= len(text) or text[start_pos] != "(":
+        return -1
+    depth = 1
+    pos = start_pos + 1
+    while pos < len(text) and depth > 0:
+        if text[pos] == "(":
+            depth += 1
+        elif text[pos] == ")":
+            depth -= 1
+        pos += 1
+    return pos - 1 if depth == 0 else -1
+
+
+def replace_array_construct(text: str) -> str:
+    """
+    Replaces all occurrences of ARRAY_CONSTRUCT(...) with Array[...] while correctly handling nested parentheses.
+    """
+    result = []
+    i = 0
+    while i < len(text):
+        # Search for "ARRAY_CONSTRUCT(" (case-insensitive)
+        array_match = re.search(r"(?i)\bARRAY_CONSTRUCT\(", text[i:])
+        if not array_match:
+            result.append(text[i:])
+            break
+
+        array_start = i + array_match.start()
+        paren_start = i + array_match.end() - 1  # Position of the '('
+
+        # Find the corresponding closing parenthesis
+        paren_end = find_matching_paren(text, paren_start)
+        if paren_end == -1:
+            # Parenthesis not closed, leave as is
+            result.append(text[i:])
+            break
+
+        # Add text before ARRAY_CONSTRUCT(
+        result.append(text[i:array_start])
+        # Add Array[ with the content between parentheses
+        content = text[paren_start + 1 : paren_end]
+        result.append(f"Array[{content}]")
+        i = paren_end + 1
+
+    return "".join(result)
+
+
 """
 A Rule is replacement rule that converts a snowflake-dialect query to a postgresql query.
 for example there's a rule that will turn "a..b" into "a.public.b"
@@ -146,6 +197,8 @@ class Rule:
 def repl_part(part: Union[str, TextLiteral], rules: Iterable[Rule]) -> str:
     if isinstance(part, TextLiteral):
         return part.value
+    # Replace ARRAY_CONSTRUCT() with Array[] before applying the other rules
+    part = replace_array_construct(part)
     ret_parts = []
     while part:
         best_match = None