feat(snowflake)!: transpile BASE64_DECODE_STRING/BINARY to DuckDB (#6837)

* transpilation of base64_decode*

* fixed tests, refactored replacement

Author: fivetran-kwoodbeck

sqlglot/dialects/duckdb.py

@@ -97,6 +97,57 @@
 _SEQ_RESTRICTED = (exp.Where, exp.Having, exp.AggFunc, exp.Order, exp.Select)
 
 
+def _apply_base64_alphabet_replacements(
+    result: exp.Expression,
+    alphabet: t.Optional[exp.Expression],
+    reverse: bool = False,
+) -> exp.Expression:
+    """
+    Apply base64 alphabet character replacements.
+
+    Base64 alphabet can be 1-3 chars: 1st = index 62 ('+'), 2nd = index 63 ('/'), 3rd = padding ('=').
+    zip truncates to the shorter string, so 1-char alphabet only replaces '+', 2-char replaces '+/', etc.
+
+    Args:
+        result: The expression to apply replacements to
+        alphabet: Custom alphabet literal (expected chars for +/=)
+        reverse: If False, replace default with custom (encode)
+                 If True, replace custom with default (decode)
+    """
+    if isinstance(alphabet, exp.Literal) and alphabet.is_string:
+        for default_char, new_char in zip("+/=", alphabet.this):
+            if new_char != default_char:
+                find, replace = (new_char, default_char) if reverse else (default_char, new_char)
+                result = exp.Replace(
+                    this=result,
+                    expression=exp.Literal.string(find),
+                    replacement=exp.Literal.string(replace),
+                )
+    return result
+
+
+def _base64_decode_sql(self: DuckDB.Generator, expression: exp.Expression, to_string: bool) -> str:
+    """
+    Transpile Snowflake BASE64_DECODE_STRING/BINARY to DuckDB.
+
+    DuckDB uses FROM_BASE64() which returns BLOB. For string output, wrap with DECODE().
+    Custom alphabets require REPLACE() calls to convert to standard base64.
+    """
+    input_expr = expression.this
+    alphabet = expression.args.get("alphabet")
+
+    # Handle custom alphabet by replacing non-standard chars with standard ones
+    input_expr = _apply_base64_alphabet_replacements(input_expr, alphabet, reverse=True)
+
+    # FROM_BASE64 returns BLOB
+    input_expr = exp.FromBase64(this=input_expr)
+
+    if to_string:
+        input_expr = exp.Decode(this=input_expr)
+
+    return self.sql(input_expr)
+
+
 def _last_day_sql(self: DuckDB.Generator, expression: exp.LastDay) -> str:
     """
     DuckDB's LAST_DAY only supports finding the last day of a month.
@@ -1536,6 +1587,8 @@ class Generator(generator.Generator):
             exp.ArrayUniqueAgg: lambda self, e: self.func(
                 "LIST", exp.Distinct(expressions=[e.this])
             ),
+            exp.Base64DecodeBinary: lambda self, e: _base64_decode_sql(self, e, to_string=False),
+            exp.Base64DecodeString: lambda self, e: _base64_decode_sql(self, e, to_string=True),
             exp.BitwiseAnd: lambda self, e: self._bitwise_op(e, "&"),
             exp.BitwiseAndAgg: _bitwise_agg_sql,
             exp.BitwiseLeftShift: _bitshift_sql,
@@ -2742,29 +2795,19 @@ def base64encode_sql(self, expression: exp.Base64Encode) -> str:
             # DuckDB TO_BASE64 requires BLOB input
             # Snowflake BASE64_ENCODE accepts both VARCHAR and BINARY - for VARCHAR it implicitly
             # encodes UTF-8 bytes. We add ENCODE unless the input is a binary type.
-            input_expr = expression.this
+            result = expression.this
 
             # Check if input is a string type - ENCODE only accepts VARCHAR
-            result = input_expr
-            if input_expr.is_type(*exp.DataType.TEXT_TYPES):
-                result = exp.Encode(this=input_expr)
+            if result.is_type(*exp.DataType.TEXT_TYPES):
+                result = exp.Encode(this=result)
 
             result = exp.ToBase64(this=result)
 
             max_line_length = expression.args.get("max_line_length")
             alphabet = expression.args.get("alphabet")
 
-            # Handle custom alphabet by replacing characters (applied before line breaks)
-            # Alphabet can be 1-3 chars: 1st = index 62 ('+'), 2nd = index 63 ('/'), 3rd = padding ('=')
-            # zip truncates to the shorter string, so 1-char alphabet only replaces '+', 2-char replaces '+/'
-            if isinstance(alphabet, exp.Literal) and alphabet.is_string:
-                for default_char, new_char in zip("+/=", alphabet.this):
-                    if new_char != default_char:
-                        result = exp.Replace(
-                            this=result,
-                            expression=exp.Literal.string(default_char),
-                            replacement=exp.Literal.string(new_char),
-                        )
+            # Handle custom alphabet by replacing standard chars with custom ones
+            result = _apply_base64_alphabet_replacements(result, alphabet)
 
             # Handle max_line_length by inserting newlines every N characters
             line_length = (

tests/dialects/test_snowflake.py

@@ -2562,12 +2562,6 @@ def test_snowflake(self):
             "SELECT ILIKE(col, 'pattern', '!')", "SELECT col ILIKE 'pattern' ESCAPE '!'"
         )
 
-        self.validate_identity("SELECT BASE64_DECODE_BINARY('SGVsbG8=')")
-        self.validate_identity("SELECT BASE64_DECODE_BINARY('SGVsbG8=')")
-
-        self.validate_identity("SELECT BASE64_DECODE_STRING('SGVsbG8gV29ybGQ=')")
-        self.validate_identity("SELECT BASE64_DECODE_STRING('SGVsbG8gV29ybGQ=', '+/=')")
-
         expr = self.validate_identity("SELECT BASE64_ENCODE('Hello World')")
         annotated = annotate_types(expr, dialect="snowflake")
         self.assertEqual(annotated.sql("duckdb"), "SELECT TO_BASE64(ENCODE('Hello World'))")
@@ -2593,11 +2587,34 @@ def test_snowflake(self):
             },
         )
 
-        self.validate_identity("SELECT TRY_BASE64_DECODE_BINARY('SGVsbG8=')")
-        self.validate_identity("SELECT TRY_BASE64_DECODE_BINARY('SGVsbG8=', '+/=')")
-
-        self.validate_identity("SELECT TRY_BASE64_DECODE_STRING('SGVsbG8gV29ybGQ=')")
-        self.validate_identity("SELECT TRY_BASE64_DECODE_STRING('SGVsbG8gV29ybGQ=', '+/=')")
+        self.validate_all(
+            "SELECT BASE64_DECODE_STRING('U25vd2ZsYWtl')",
+            write={
+                "snowflake": "SELECT BASE64_DECODE_STRING('U25vd2ZsYWtl')",
+                "duckdb": "SELECT DECODE(FROM_BASE64('U25vd2ZsYWtl'))",
+            },
+        )
+        self.validate_all(
+            "SELECT BASE64_DECODE_STRING('U25vd2ZsYWtl', '-_+')",
+            write={
+                "snowflake": "SELECT BASE64_DECODE_STRING('U25vd2ZsYWtl', '-_+')",
+                "duckdb": "SELECT DECODE(FROM_BASE64(REPLACE(REPLACE(REPLACE('U25vd2ZsYWtl', '-', '+'), '_', '/'), '+', '=')))",
+            },
+        )
+        self.validate_all(
+            "SELECT BASE64_DECODE_BINARY(x)",
+            write={
+                "snowflake": "SELECT BASE64_DECODE_BINARY(x)",
+                "duckdb": "SELECT FROM_BASE64(x)",
+            },
+        )
+        self.validate_all(
+            "SELECT BASE64_DECODE_BINARY(x, '-_+')",
+            write={
+                "snowflake": "SELECT BASE64_DECODE_BINARY(x, '-_+')",
+                "duckdb": "SELECT FROM_BASE64(REPLACE(REPLACE(REPLACE(x, '-', '+'), '_', '/'), '+', '='))",
+            },
+        )
 
         self.validate_identity("SELECT TRY_HEX_DECODE_BINARY('48656C6C6F')")