Use [min, max) range to infer integer type

MeltanoLabs · Nov 28, 2024 · 5966233 · 5966233
1 parent 598a218
commit 5966233
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 8 deletions.
diff --git a/target_postgres/connector.py b/target_postgres/connector.py
@@ -5,6 +5,7 @@
 import atexit
 import io
 import itertools
+import math
 import signal
 import sys
 import typing as t
@@ -264,11 +265,12 @@ def _handle_array_type(self, jsonschema: dict) -> ARRAY | JSONB:
 
     def _handle_integer_type(self, jsonschema: dict) -> SMALLINT | INTEGER | BIGINT:
         """Handle integer type."""
-        if maximum := jsonschema.get("maximum"):
-            if maximum < 2**15:
-                return SMALLINT()
-            if maximum < 2**31:
-                return INTEGER()
+        minimum = jsonschema.get("minimum", -math.inf)
+        maximum = jsonschema.get("maximum", math.inf)
+        if minimum >= -(2**15) and maximum < 2**15:
+            return SMALLINT()
+        if minimum >= -(2**31) and maximum < 2**31:
+            return INTEGER()
 
         return BIGINT()
 

diff --git a/target_postgres/tests/test_types.py b/target_postgres/tests/test_types.py
@@ -58,17 +58,38 @@ def test_datetime_string(self, to_postgres: JSONSchemaToPostgres):
             pytest.param({"type": "integer"}, BIGINT, id="default"),
             pytest.param({"type": ["integer", "null"]}, BIGINT, id="default-nullable"),
             pytest.param(
-                {"type": "integer", "maximum": 2**15 - 1},
+                {
+                    "type": "integer",
+                    "minimum": 0,
+                    "maximum": 2**15 - 1,
+                },
                 SMALLINT,
                 id="smallint",
             ),
             pytest.param(
-                {"type": "integer", "maximum": 2**31 - 1},
+                {
+                    "type": "integer",
+                    "minimum": -5,
+                    "maximum": 5,
+                },
+                SMALLINT,
+                id="negative-smallint",
+            ),
+            pytest.param(
+                {
+                    "type": "integer",
+                    "minimum": 0,
+                    "maximum": 2**31 - 1,
+                },
                 sa.INTEGER,
                 id="integer",
             ),
             pytest.param(
-                {"type": "integer", "maximum": 2**31 + 1},
+                {
+                    "type": "integer",
+                    "minimum": 0,
+                    "maximum": 2**31 + 1,
+                },
                 BIGINT,
                 id="bigint",
             ),