feat(optimizer)!: annotate types for Snowflake TRIM function

Author: fivetran-BradfordPaskewitz

sqlglot/dialects/snowflake.py

@@ -1362,9 +1362,10 @@ class Generator(generator.Generator):
 
         TYPE_MAPPING = {
             **generator.Generator.TYPE_MAPPING,
+            exp.DataType.Type.BIGDECIMAL: "DOUBLE",
             exp.DataType.Type.NESTED: "OBJECT",
             exp.DataType.Type.STRUCT: "OBJECT",
-            exp.DataType.Type.BIGDECIMAL: "DOUBLE",
+            exp.DataType.Type.TEXT: "VARCHAR",
         }
 
         TOKEN_MAPPING = {

tests/dialects/test_dialect.py

@@ -209,7 +209,7 @@ def test_cast(self):
                 "postgres": "CAST(a AS TEXT)",
                 "presto": "CAST(a AS VARCHAR)",
                 "redshift": "CAST(a AS VARCHAR(MAX))",
-                "snowflake": "CAST(a AS TEXT)",
+                "snowflake": "CAST(a AS VARCHAR)",
                 "spark": "CAST(a AS STRING)",
                 "starrocks": "CAST(a AS STRING)",
                 "tsql": "CAST(a AS VARCHAR(MAX))",
@@ -293,7 +293,7 @@ def test_cast(self):
                 "postgres": "CAST(a AS TEXT)",
                 "presto": "CAST(a AS VARCHAR)",
                 "redshift": "CAST(a AS VARCHAR(MAX))",
-                "snowflake": "CAST(a AS TEXT)",
+                "snowflake": "CAST(a AS VARCHAR)",
                 "spark": "CAST(a AS STRING)",
                 "starrocks": "CAST(a AS STRING)",
                 "tsql": "CAST(a AS VARCHAR(MAX))",

tests/dialects/test_snowflake.py

@@ -85,7 +85,10 @@ def test_snowflake(self):
         self.validate_identity("SELECT CONVERT_TIMEZONE('UTC', 'America/Los_Angeles', col)")
         self.validate_identity("ALTER TABLE a SWAP WITH b")
         self.validate_identity("SELECT MATCH_CONDITION")
-        self.validate_identity("SELECT * REPLACE (CAST(col AS TEXT) AS scol) FROM t")
+        self.validate_identity(
+            "SELECT * REPLACE (CAST(col AS TEXT) AS scol) FROM t",
+            "SELECT * REPLACE (CAST(col AS VARCHAR) AS scol) FROM t",
+        )
         self.validate_identity("1 /* /* */")
         self.validate_identity("TO_TIMESTAMP(col, fmt)")
         self.validate_identity("SELECT TO_CHAR(CAST('12:05:05' AS TIME))")
@@ -178,7 +181,7 @@ def test_snowflake(self):
         )
         self.validate_identity(
             "SELECT a:from::STRING, a:from || ' test' ",
-            "SELECT CAST(GET_PATH(a, 'from') AS TEXT), GET_PATH(a, 'from') || ' test'",
+            "SELECT CAST(GET_PATH(a, 'from') AS VARCHAR), GET_PATH(a, 'from') || ' test'",
         )
         self.validate_identity(
             "SELECT a:select",
@@ -187,7 +190,7 @@ def test_snowflake(self):
         self.validate_identity("x:from", "GET_PATH(x, 'from')")
         self.validate_identity(
             "value:values::string::int",
-            "CAST(CAST(GET_PATH(value, 'values') AS TEXT) AS INT)",
+            "CAST(CAST(GET_PATH(value, 'values') AS VARCHAR) AS INT)",
         )
         self.validate_identity(
             """SELECT GET_PATH(PARSE_JSON('{"y": [{"z": 1}]}'), 'y[0]:z')""",
@@ -2638,25 +2641,26 @@ def test_swap(self):
         assert isinstance(ast.args["actions"][0], exp.SwapTable)
 
     def test_try_cast(self):
-        self.validate_identity("SELECT TRY_CAST(x AS DOUBLE)")
-        self.validate_identity("SELECT TRY_CAST(FOO() AS TEXT)")
-
-        self.validate_all("TRY_CAST('foo' AS TEXT)", read={"hive": "CAST('foo' AS STRING)"})
-        self.validate_all("CAST(5 + 5 AS TEXT)", read={"hive": "CAST(5 + 5 AS STRING)"})
+        self.validate_all("TRY_CAST('foo' AS VARCHAR)", read={"hive": "CAST('foo' AS STRING)"})
+        self.validate_all("CAST(5 + 5 AS VARCHAR)", read={"hive": "CAST(5 + 5 AS STRING)"})
         self.validate_all(
-            "CAST(TRY_CAST('2020-01-01' AS DATE) AS TEXT)",
+            "CAST(TRY_CAST('2020-01-01' AS DATE) AS VARCHAR)",
             read={
                 "hive": "CAST(CAST('2020-01-01' AS DATE) AS STRING)",
-                "snowflake": "CAST(TRY_CAST('2020-01-01' AS DATE) AS TEXT)",
+                "snowflake": "CAST(TRY_CAST('2020-01-01' AS DATE) AS VARCHAR)",
             },
         )
         self.validate_all(
-            "TRY_CAST('val' AS TEXT)",
+            "TRY_CAST('val' AS VARCHAR)",
             read={
                 "hive": "CAST('val' AS STRING)",
-                "snowflake": "TRY_CAST('val' AS TEXT)",
+                "snowflake": "TRY_CAST('val' AS VARCHAR)",
             },
         )
+        self.validate_identity("SELECT TRY_CAST(x AS DOUBLE)")
+        self.validate_identity(
+            "SELECT TRY_CAST(FOO() AS TEXT)", "SELECT TRY_CAST(FOO() AS VARCHAR)"
+        )
 
         from sqlglot.optimizer.annotate_types import annotate_types
 
@@ -2670,7 +2674,7 @@ def test_try_cast(self):
 
                 expression = annotate_types(expression, schema={"t": {"x": value_type}})
                 self.assertEqual(
-                    expression.sql(dialect="snowflake"), f"SELECT {func}(t.x AS TEXT) FROM t"
+                    expression.sql(dialect="snowflake"), f"SELECT {func}(t.x AS VARCHAR) FROM t"
                 )
 
     def test_copy(self):

tests/fixtures/optimizer/annotate_functions.sql

@@ -1531,10 +1531,6 @@ STRING;
 -- Snowflake
 --------------------------------------
 
-# dialect: snowflake
-LEAST(x::DECIMAL(18, 2));
-DECIMAL(18, 2);
-
 # dialect: snowflake
 CHARINDEX('world', 'hello world');
 INT;
@@ -1543,18 +1539,6 @@ INT;
 CHARINDEX('world', 'hello world', 1);
 INT;
 
-# dialect: snowflake
-REVERSE('Hello, world!');
-VARCHAR;
-
-# dialect: snowflake
-REVERSE(tbl.str_col);
-TEXT;
-
-# dialect: snowflake
-REVERSE(tbl.bin_col);
-BINARY;
-
 # dialect: snowflake
 CONCAT('Hello', 'World!');
 VARCHAR;
@@ -1591,10 +1575,6 @@ BOOLEAN;
 CONTAINS(tbl.bin_col, NULL);
 BOOLEAN;
 
-# dialect: snowflake
-REVERSE(NULL);
-VARCHAR;
-
 # dialect: snowflake
 ENDSWITH('hello world', 'world');
 BOOLEAN;
@@ -1612,53 +1592,61 @@ ENDSWITH(tbl.bin_col, NULL);
 BOOLEAN;
 
 # dialect: snowflake
-STARTSWITH('hello world', 'hello');
-BOOLEAN;
+LEAST(x::DECIMAL(18, 2));
+DECIMAL(18, 2);
 
 # dialect: snowflake
-STARTSWITH(tbl.str_col, 'test');
-BOOLEAN;
+LEFT('hello world', 5);
+VARCHAR;
 
 # dialect: snowflake
-STARTSWITH(tbl.bin_col, tbl.bin_col);
-BOOLEAN;
+LEFT(tbl.str_col, 3);
+STRING;
 
 # dialect: snowflake
-STARTSWITH(tbl.bin_col, NULL);
-BOOLEAN;
+LEFT(tbl.bin_col, 3);
+BINARY;
 
 # dialect: snowflake
-SUBSTR('hello world', 1, 5);
-VARCHAR;
+LEFT(tbl.bin_col, NULL);
+BINARY;
 
 # dialect: snowflake
-SUBSTR(tbl.str_col, 1, 3);
-STRING;
+LEN(tbl.str_col);
+INT;
 
 # dialect: snowflake
-SUBSTR(tbl.bin_col, 1, 3);
-BINARY;
+LEN(tbl.bin_col);
+INT;
 
 # dialect: snowflake
-SUBSTR(tbl.str_col, NULL);
-STRING;
+LENGTH(tbl.str_col);
+INT;
 
 # dialect: snowflake
-LEFT('hello world', 5);
+LENGTH(tbl.bin_col);
+INT;
+
+# dialect: snowflake
+LOWER(tbl.str_col);
 VARCHAR;
 
 # dialect: snowflake
-LEFT(tbl.str_col, 3);
-STRING;
+REVERSE('Hello, world!');
+VARCHAR;
 
 # dialect: snowflake
-LEFT(tbl.bin_col, 3);
-BINARY;
+REVERSE(tbl.str_col);
+VARCHAR;
 
 # dialect: snowflake
-LEFT(tbl.bin_col, NULL);
+REVERSE(tbl.bin_col);
 BINARY;
 
+# dialect: snowflake
+REVERSE(NULL);
+VARCHAR;
+
 # dialect: snowflake
 RIGHT('hello world', 5);
 VARCHAR;
@@ -1676,23 +1664,55 @@ RIGHT(tbl.str_col, NULL);
 STRING;
 
 # dialect: snowflake
-LENGTH(tbl.str_col);
-INT;
+STARTSWITH('hello world', 'hello');
+BOOLEAN;
 
 # dialect: snowflake
-LENGTH(tbl.bin_col);
-INT;
+STARTSWITH(tbl.str_col, 'test');
+BOOLEAN;
 
 # dialect: snowflake
-LEN(tbl.str_col);
-INT;
+STARTSWITH(tbl.bin_col, tbl.bin_col);
+BOOLEAN;
 
 # dialect: snowflake
-LEN(tbl.bin_col);
-INT;
+STARTSWITH(tbl.bin_col, NULL);
+BOOLEAN;
 
 # dialect: snowflake
-LOWER(tbl.str_col);
+SUBSTR('hello world', 1, 5);
+VARCHAR;
+
+# dialect: snowflake
+SUBSTR(tbl.str_col, 1, 3);
+STRING;
+
+# dialect: snowflake
+SUBSTR(tbl.bin_col, 1, 3);
+BINARY;
+
+# dialect: snowflake
+SUBSTR(tbl.str_col, NULL);
+STRING;
+
+# dialect: snowflake
+TRIM('hello world');
+VARCHAR;
+
+# dialect: snowflake
+TRIM('hello world', 'hello');
+VARCHAR;
+
+# dialect: snowflake
+TRIM(tbl.str_col);
+VARCHAR;
+
+# dialect: snowflake
+TRIM(tbl.str_col, tbl.str_col);
+VARCHAR;
+
+# dialect: snowflake
+TRIM(NULL);
 VARCHAR;
 
 --------------------------------------