Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion dataframely/_compat.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright (c) QuantCo 2025-2025
# Copyright (c) QuantCo 2025-2026
# SPDX-License-Identifier: BSD-3-Clause


Expand Down Expand Up @@ -29,13 +29,15 @@ class DeltaTable: # type: ignore # noqa: N801
try:
import sqlalchemy as sa
import sqlalchemy.dialects.mssql as sa_mssql
import sqlalchemy.dialects.postgresql as sa_postgresql
from sqlalchemy import Dialect
from sqlalchemy.dialects.mssql.pyodbc import MSDialect_pyodbc
from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
from sqlalchemy.sql.type_api import TypeEngine as sa_TypeEngine
except ImportError:
sa = _DummyModule("sqlalchemy") # type: ignore
sa_mssql = _DummyModule("sqlalchemy") # type: ignore
sa_postgresql = _DummyModule("sqlalchemy") # type: ignore

class sa_TypeEngine: # type: ignore # noqa: N801
pass
Expand Down Expand Up @@ -81,6 +83,7 @@ class Dialect: # type: ignore # noqa: N801
"pydantic_core_schema",
"pydantic",
"sa_mssql",
"sa_postgresql",
"sa_TypeEngine",
"sa",
]
11 changes: 7 additions & 4 deletions dataframely/columns/struct.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright (c) QuantCo 2025-2025
# Copyright (c) QuantCo 2025-2026
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations
Expand All @@ -8,7 +8,7 @@

import polars as pl

from dataframely._compat import pa, sa, sa_TypeEngine
from dataframely._compat import pa, sa, sa_postgresql, sa_TypeEngine
from dataframely._polars import PolarsDataType
from dataframely.random import Generator

Expand Down Expand Up @@ -107,8 +107,11 @@ def validation_rules(self, expr: pl.Expr) -> dict[str, pl.Expr]:
}

def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
# NOTE: We might want to add support for PostgreSQL's JSON in the future.
raise NotImplementedError("SQL column cannot have 'Struct' type.")
match dialect.name:
case "postgresql":
return sa_postgresql.JSONB()
case _:
raise NotImplementedError("SQL column cannot have 'Struct' type.")

@property
def pyarrow_dtype(self) -> pa.DataType:
Expand Down
5 changes: 3 additions & 2 deletions tests/columns/test_sqlalchemy_columns.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright (c) QuantCo 2025-2025
# Copyright (c) QuantCo 2025-2026
# SPDX-License-Identifier: BSD-3-Clause

import pytest
Expand Down Expand Up @@ -95,6 +95,7 @@ def test_mssql_datatype(column: Column, datatype: str) -> None:
(dy.String(regex="^[abc]{1,3}d$"), "VARCHAR(4)"),
(dy.Enum(["foo", "bar"]), "CHAR(3)"),
(dy.Enum(["a", "abc"]), "VARCHAR(3)"),
(dy.Struct({"a": dy.String(nullable=True)}), "JSONB"),
],
)
def test_postgres_datatype(column: Column, datatype: str) -> None:
Expand Down Expand Up @@ -152,7 +153,7 @@ def test_raise_for_array_column(dialect: Dialect) -> None:
dy.Array(dy.String(nullable=True), 1).sqlalchemy_dtype(dialect)


@pytest.mark.parametrize("dialect", [MSDialect_pyodbc(), PGDialect_psycopg2()])
@pytest.mark.parametrize("dialect", [MSDialect_pyodbc()])
def test_raise_for_struct_column(dialect: Dialect) -> None:
with pytest.raises(
NotImplementedError, match="SQL column cannot have 'Struct' type."
Expand Down
Loading