feat(postgres): Support DIV() func for integer division

tobymao · Jun 5, 2024 · 6d065b2 · 6d065b2
1 parent ff55ec1
commit 6d065b2
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 0 deletions.
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
@@ -8,6 +8,7 @@
     Dialect,
     JSON_EXTRACT_TYPE,
     any_value_to_max_sql,
+    binary_from_function,
     bool_xor_sql,
     datestrtodate_sql,
     build_formatted_time,
@@ -338,6 +339,7 @@ class Tokenizer(tokens.Tokenizer):
             "REGTYPE": TokenType.OBJECT_IDENTIFIER,
             "FLOAT": TokenType.DOUBLE,
         }
+        KEYWORDS.pop("DIV")
 
         SINGLE_TOKENS = {
             **tokens.Tokenizer.SINGLE_TOKENS,
@@ -356,6 +358,7 @@ class Parser(parser.Parser):
         FUNCTIONS = {
             **parser.Parser.FUNCTIONS,
             "DATE_TRUNC": build_timestamp_trunc,
+            "DIV": binary_from_function(exp.IntDiv),
             "GENERATE_SERIES": _build_generate_series,
             "JSON_EXTRACT_PATH": build_json_extract_path(exp.JSONExtract),
             "JSON_EXTRACT_PATH_TEXT": build_json_extract_path(exp.JSONExtractScalar),
@@ -504,6 +507,7 @@ class Generator(generator.Generator):
             exp.DateSub: _date_add_sql("-"),
             exp.Explode: rename_func("UNNEST"),
             exp.GroupConcat: _string_agg_sql,
+            exp.IntDiv: rename_func("DIV"),
             exp.JSONExtract: _json_extract_sql("JSON_EXTRACT_PATH", "->"),
             exp.JSONExtractScalar: _json_extract_sql("JSON_EXTRACT_PATH_TEXT", "->>"),
             exp.JSONBExtract: lambda self, e: self.binary(e, "#>"),

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -1083,9 +1083,14 @@ def test_bigquery(self):
         )
         self.validate_all(
             "DIV(x, y)",
+            read={
+                "duckdb": "x // y",
+                "postgres": "DIV(x, y)",
+            },
             write={
                 "bigquery": "DIV(x, y)",
                 "duckdb": "x // y",
+                "postgres": "DIV(x, y)",
             },
         )
         self.validate_all(

diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -723,6 +723,7 @@ def test_postgres(self):
         self.validate_identity("cast(a as FLOAT)", "CAST(a AS DOUBLE PRECISION)")
         self.validate_identity("cast(a as FLOAT8)", "CAST(a AS DOUBLE PRECISION)")
         self.validate_identity("cast(a as FLOAT4)", "CAST(a AS REAL)")
+        self.validate_identity("DIV(5, 4)").assert_is(exp.IntDiv)
 
     def test_ddl(self):
         # Checks that user-defined types are parsed into DataType instead of Identifier