feat: Support integer types other than BIGINT

target_postgres/connector.py

@@ -18,7 +18,14 @@
 import sqlalchemy as sa
 from singer_sdk import SQLConnector
 from singer_sdk.connectors.sql import JSONSchemaToSQL
-from sqlalchemy.dialects.postgresql import ARRAY, BIGINT, BYTEA, JSONB, UUID
+from sqlalchemy.dialects.postgresql import (
+    ARRAY,
+    BIGINT,
+    BYTEA,
+    JSONB,
+    SMALLINT,
+    UUID,
+)
 from sqlalchemy.engine import URL
 from sqlalchemy.engine.url import make_url
 from sqlalchemy.types import (
@@ -255,12 +262,22 @@ def _handle_array_type(self, jsonschema: dict) -> ARRAY | JSONB:
         # Case 3: tuples
         return ARRAY(JSONB()) if isinstance(items, list) else 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()
+
+        return BIGINT()
+
     @cached_property
     def jsonschema_to_sql(self) -> JSONSchemaToSQL:
         """Return a JSONSchemaToSQL instance with custom type handling."""
         to_sql = JSONSchemaToPostgres(content_encoding=self.interpret_content_encoding)
         to_sql.fallback_type = TEXT
-        to_sql.register_type_handler("integer", BIGINT)
+        to_sql.register_type_handler("integer", self._handle_integer_type)
         to_sql.register_type_handler("object", JSONB)
         to_sql.register_type_handler("array", self._handle_array_type)
         to_sql.register_format_handler("date-time", TIMESTAMP)

target_postgres/tests/test_types.py

@@ -2,8 +2,9 @@
 
 import pytest
 import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import BIGINT, SMALLINT
 
-from target_postgres.connector import NOTYPE, PostgresConnector
+from target_postgres.connector import NOTYPE, JSONSchemaToPostgres, PostgresConnector
 
 
 @pytest.fixture
@@ -36,3 +37,48 @@ def connector():
 def test_type_hierarchy(connector, types, expected):
     """Test that types are merged correctly."""
     assert type(connector.merge_sql_types(types)) is expected
+
+
+class TestJSONSchemaToPostgres:
+    """Test JSONSchemaToPostgres class."""
+
+    @pytest.fixture
+    def to_postgres(self, connector: PostgresConnector):
+        """Create a JSONSchemaToPostgres instance."""
+        return connector.jsonschema_to_sql
+
+    def test_datetime_string(self, to_postgres: JSONSchemaToPostgres):
+        """Test conversion of JSON schema string to Postgres datetime."""
+        result = to_postgres.to_sql_type({"type": "string", "format": "date-time"})
+        assert type(result) is sa.TIMESTAMP
+
+    @pytest.mark.parametrize(
+        ("jsonschema", "expected"),
+        [
+            pytest.param({"type": "integer"}, BIGINT, id="default"),
+            pytest.param(
+                {"type": "integer", "maximum": 2**15 - 1},
+                SMALLINT,
+                id="smallint",
+            ),
+            pytest.param(
+                {"type": "integer", "maximum": 2**31 - 1},
+                sa.INTEGER,
+                id="integer",
+            ),
+            pytest.param(
+                {"type": "integer", "maximum": 2**31 + 1},
+                BIGINT,
+                id="bigint",
+            ),
+        ],
+    )
+    def test_integers(
+        self,
+        to_postgres: JSONSchemaToPostgres,
+        jsonschema: dict,
+        expected: type[sa.types.TypeEngine],
+    ):
+        """Test conversion of JSON schema types to Postgres types."""
+        result = to_postgres.to_sql_type(jsonschema)
+        assert type(result) is expected