feat: Somewhat Useful Exceptions for Unsupported Objects (#7)

* Support for Null Objects

* Raise Exceptions on Unsupported Objects

* Formatting

* Support For Glue Type Override

* Type Override Tests

* Typing & Documentation, Fix Checks

* Attempt to Fix CI Caching Issue

* Stop Running CI On Fork Branch

Author: robodair

.github/workflows/ci.yml

@@ -23,6 +23,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
+      id: setup-python
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}

README.md

@@ -17,7 +17,7 @@ When using `AWS Kinesis Firehose` in a configuration that receives JSONs and wri
 one needs to define a `AWS Glue` table so Firehose knows what schema to use when creating the parquet files.
 
 AWS Glue lets you define a schema using `Avro` or `JSON Schema` and then to create a table from that schema,
-but as of *May 2022`
+but as of *May 2022*
 there are limitations on AWS that tables that are created that way can't be used with Kinesis Firehose.
 
 <https://stackoverflow.com/questions/68125501/invalid-schema-error-in-aws-glue-created-via-terraform>
@@ -100,6 +100,62 @@ Alternatively you can run CLI with `-o` flag to set output file location:
 pydantic-glue -f example.py -c Foo -o example.json -l
 ```
 
+## Override the type for the AWS Glue Schema
+
+Wherever there is a `type` key in the input JSON Schema, an additional key `glue_type` may be
+defined to override the type that is used in the AWS Glue Schema. This is, for example, useful for
+a pydantic model that has a field of type `int` that is unix epoch time, while the column type you
+would like in Glue is of type `timestamp`.
+
+Additional JSON Schema keys to a pydantic model can be added by using the
+[`Field` function](https://docs.pydantic.dev/latest/api/fields/#pydantic.fields.Field)
+with the argument `json_schema_extra` like so:
+
+```python
+from pydantic import BaseModel, Field
+
+class A(BaseModel):
+    epoch_time: int = Field(
+        ...,
+        json_schema_extra={
+            "glue_type": "timestamp",
+        },
+    )
+```
+
+The resulting JSON Schema will be:
+
+```json
+{
+    "properties": {
+        "epoch_time": {
+            "glue_type": "timestamp",
+            "title": "Epoch Time",
+            "type": "integer"
+        }
+    },
+    "required": [
+        "epoch_time"
+    ],
+    "title": "A",
+    "type": "object"
+}
+```
+
+And the result after processing with pydantic-glue:
+
+```json
+{
+  "//": "Generated by pydantic-glue at 2022-05-25 12:35:55.333570. DO NOT EDIT",
+  "columns": {
+    "epoch_time": "timestamp",
+  }
+}
+```
+
+Recursing through object properties terminates when you supply a `glue_type` to use. If the type is
+complex, you must supply the full complex type yourself.
+
 ## How it works?
 
 * `pydantic` gets converted to JSON Schema

pydantic_glue/handler.py

@@ -4,14 +4,18 @@
 
 
 def dispatch(v: dict[str, Any]) -> str:
+
+    glue_type = v.get("glue_type", None)
+
+    if glue_type is not None:
+        return str(glue_type)
+
     if "anyOf" in v:
         return handle_union(v)
 
     t = v["type"]
 
     if t == "object":
-        if "additionalProperties" in v:
-            return handle_map(v)
         return handle_object(v)
 
     if t == "array":
@@ -55,6 +59,17 @@ def map_dispatch(o: dict[str, Any]) -> list[tuple[str, str]]:
 
 
 def handle_object(o: dict[str, Any]) -> str:
+    if "additionalProperties" in o:
+        if o["additionalProperties"] is True:
+            raise Exception("Glue Cannot Support a Map Without Types")
+        elif o["additionalProperties"]:
+            if "properties" in o:
+                raise NotImplementedError("Merging types of properties and additionalProperties")
+            return handle_map(o)
+
+    if "properties" not in o:
+        raise Exception("Object without properties or additionalProperties can't be represented")
+
     res = [f"{k}:{v}" for (k, v) in map_dispatch(o)]
     return f"struct<{','.join(res)}>"
 

tests/unit/test_convert.py

@@ -2,7 +2,8 @@
 import json
 from typing import Optional, Union
 
-from pydantic import BaseModel
+import pytest
+from pydantic import BaseModel, Field, field_serializer
 from pydantic_glue import convert
 
 
@@ -34,6 +35,31 @@ class A(BaseModel):
     assert convert(json.dumps(A.model_json_schema())) == expected
 
 
+def test_single_dict_as_string_column():
+    class A(BaseModel):
+        as_string: dict
+
+        @field_serializer("as_string")
+        def dump_json(self, value: dict) -> str:
+            return json.dumps(value)
+
+    expected = [("as_string", "string")]
+    assert convert(json.dumps(A.model_json_schema(mode="serialization"))) == expected
+
+
+def test_single_type_override_column():
+    class A(BaseModel):
+        special: int = Field(
+            ...,
+            json_schema_extra={
+                "glue_type": "gluetype",
+            },
+        )
+
+    expected = [("special", "gluetype")]
+    assert convert(json.dumps(A.model_json_schema())) == expected
+
+
 def test_single_date_column():
     class A(BaseModel):
         modifiedOn: datetime.date
@@ -155,7 +181,62 @@ class A(BaseModel):
         ("other", "string"),
     ]
 
-    assert convert(json.dumps(A.model_json_schema())) == expected
+    assert (
+        convert(
+            json.dumps(A.model_json_schema()),
+        )
+        == expected
+    )
+
+
+def test_custom_type():
+    class A(BaseModel):
+        unixtime: int = Field(
+            ...,
+            json_schema_extra={
+                "glue_type": "timestamp",
+            },
+        )
+        optional_unixtime: Optional[int] = Field(
+            ...,
+            json_schema_extra={
+                "glue_type": "timestamp",
+            },
+        )
+        clobber_union_unixtime: Optional[Union[int, str]] = Field(
+            ...,
+            json_schema_extra={
+                "glue_type": "timestamp",
+            },
+        )
+        correct_union_unixtime: Optional[Union[int, str]] = Field(
+            ...,
+            json_schema_extra={
+                "glue_type": "union<timestamp,string>",
+            },
+        )
+
+    expected = [
+        ("unixtime", "timestamp"),
+        ("optional_unixtime", "timestamp"),
+        ("clobber_union_unixtime", "timestamp"),
+        ("correct_union_unixtime", "union<timestamp,string>"),
+    ]
+
+    assert (
+        convert(
+            json.dumps(A.model_json_schema(mode="serialization")),
+        )
+        == expected
+    )
+
+
+def test_invalid_object_raises():
+    class A(BaseModel):
+        map_serialized_as_object: dict
+
+    with pytest.raises(Exception):
+        convert(json.dumps(A.model_json_schema()))
 
 
 def test_union_of_string_and_int():