diff --git a/pydantic_glue/handler.py b/pydantic_glue/handler.py index cd3a40c..16b1895 100644 --- a/pydantic_glue/handler.py +++ b/pydantic_glue/handler.py @@ -18,6 +18,10 @@ def dispatch(v: dict[str, Any]) -> str: return handle_array(v) if t == "string": + if v.get("format") == "date-time": + return "timestamp" + if v.get("format") == "date": + return "date" return "string" if t == "boolean": diff --git a/pyproject.toml b/pyproject.toml index 751f8c7..b8ce2c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [tool.poetry] name = "pydantic-glue" keywords = ["pydantic", "glue", "athena", "types", "convert"] -version = "0.2.1" +version = "0.3.0" description = "Convert pydantic model to aws glue schema for terraform" authors = ["Serhii Dimchenko "] readme = "README.md" diff --git a/tests/unit/test_convert.py b/tests/unit/test_convert.py index 0a17b99..3dfee22 100644 --- a/tests/unit/test_convert.py +++ b/tests/unit/test_convert.py @@ -1,3 +1,4 @@ +import datetime import json from typing import Optional, Union @@ -33,6 +34,22 @@ class A(BaseModel): assert convert(json.dumps(A.model_json_schema())) == expected +def test_single_date_column(): + class A(BaseModel): + modifiedOn: datetime.date + + expected = [("modifiedOn", "date")] + assert convert(json.dumps(A.model_json_schema())) == expected + + +def test_single_datetime_column(): + class A(BaseModel): + modifiedOn: datetime.datetime + + expected = [("modifiedOn", "timestamp")] + assert convert(json.dumps(A.model_json_schema())) == expected + + def test_multiple_string_column(): class A(BaseModel): hey: str