fix(postgres): sha256 support

tobymao · Jun 7, 2024 · d96459f · d96459f
1 parent 0e1a1fb
commit d96459f
Show file tree

Hide file tree

Showing 8 changed files with 20 additions and 9 deletions.
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
@@ -22,6 +22,7 @@
     build_date_delta_with_interval,
     regexp_replace_sql,
     rename_func,
+    sha256_sql,
     timestrtotime_sql,
     ts_or_ds_add_cast,
     unit_to_var,
@@ -637,9 +638,7 @@ class Generator(generator.Generator):
                 ]
             ),
             exp.SHA: rename_func("SHA1"),
-            exp.SHA2: lambda self, e: self.func(
-                "SHA256" if e.text("length") == "256" else "SHA512", e.this
-            ),
+            exp.SHA2: sha256_sql,
             exp.StabilityProperty: lambda self, e: (
                 "DETERMINISTIC" if e.name == "IMMUTABLE" else "NOT DETERMINISTIC"
             ),

diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
@@ -14,6 +14,7 @@
     no_pivot_sql,
     build_json_extract_path,
     rename_func,
+    sha256_sql,
     var_map_sql,
     timestamptrunc_sql,
 )
@@ -758,9 +759,7 @@ class Generator(generator.Generator):
             exp.MD5Digest: rename_func("MD5"),
             exp.MD5: lambda self, e: self.func("LOWER", self.func("HEX", self.func("MD5", e.this))),
             exp.SHA: rename_func("SHA1"),
-            exp.SHA2: lambda self, e: self.func(
-                "SHA256" if e.text("length") == "256" else "SHA512", e.this
-            ),
+            exp.SHA2: sha256_sql,
             exp.UnixToTime: _unix_to_time_sql,
             exp.TimestampTrunc: timestamptrunc_sql(zone=True),
             exp.Variance: rename_func("varSamp"),

diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
@@ -1181,3 +1181,7 @@ def _builder(dtype: exp.DataType) -> exp.DataType:
         return exp.DataType.build(f"DECIMAL({params})")
 
     return _builder
+
+
+def sha256_sql(self: Generator, expression: exp.SHA2) -> str:
+    return self.func(f"SHA{expression.text('length') or '256'}", expression.this)
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
@@ -26,6 +26,7 @@
     build_json_extract_path,
     build_timestamp_trunc,
     rename_func,
+    sha256_sql,
     str_position_sql,
     struct_extract_sql,
     timestamptrunc_sql,
@@ -362,6 +363,9 @@ class Parser(parser.Parser):
             "TO_CHAR": build_formatted_time(exp.TimeToStr, "postgres"),
             "TO_TIMESTAMP": _build_to_timestamp,
             "UNNEST": exp.Explode.from_arg_list,
+            "SHA256": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(256)),
+            "SHA384": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(384)),
+            "SHA512": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(512)),
         }
 
         FUNCTION_PARSERS = {
@@ -534,6 +538,7 @@ class Generator(generator.Generator):
                     transforms.eliminate_qualify,
                 ]
             ),
+            exp.SHA2: sha256_sql,
             exp.StrPosition: str_position_sql,
             exp.StrToDate: lambda self, e: self.func("TO_DATE", e.this, self.format_time(e)),
             exp.StrToTime: lambda self, e: self.func("TO_TIMESTAMP", e.this, self.format_time(e)),

diff --git a/sqlglot/dialects/presto.py b/sqlglot/dialects/presto.py
@@ -21,6 +21,7 @@
     regexp_extract_sql,
     rename_func,
     right_to_substring_sql,
+    sha256_sql,
     struct_extract_sql,
     str_position_sql,
     timestamptrunc_sql,
@@ -452,9 +453,7 @@ class Generator(generator.Generator):
             ),
             exp.MD5Digest: rename_func("MD5"),
             exp.SHA: rename_func("SHA1"),
-            exp.SHA2: lambda self, e: self.func(
-                "SHA256" if e.text("length") == "256" else "SHA512", e.this
-            ),
+            exp.SHA2: sha256_sql,
         }
 
         RESERVED_KEYWORDS = {

diff --git a/sqlglot/dialects/redshift.py b/sqlglot/dialects/redshift.py
@@ -211,6 +211,7 @@ class Generator(Postgres.Generator):
 
         # Redshift supports LAST_DAY(..)
         TRANSFORMS.pop(exp.LastDay)
+        TRANSFORMS.pop(exp.SHA2)
 
         RESERVED_KEYWORDS = {
             "aes128",

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -761,12 +761,15 @@ def test_bigquery(self):
                 "clickhouse": "SHA256(x)",
                 "presto": "SHA256(x)",
                 "trino": "SHA256(x)",
+                "postgres": "SHA256(x)",
             },
             write={
                 "bigquery": "SHA256(x)",
                 "spark2": "SHA2(x, 256)",
                 "clickhouse": "SHA256(x)",
+                "postgres": "SHA256(x)",
                 "presto": "SHA256(x)",
+                "redshift": "SHA2(x, 256)",
                 "trino": "SHA256(x)",
             },
         )

diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -8,6 +8,7 @@ class TestPostgres(Validator):
     dialect = "postgres"
 
     def test_postgres(self):
+        self.validate_identity("SHA384(x)")
         self.validate_identity(
             'CREATE TABLE x (a TEXT COLLATE "de_DE")', "CREATE TABLE x (a TEXT COLLATE de_DE)"
         )