FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Merge pull request #27 from leynier/fork-10 · synodriver/sqlmodel@7f934d6 · GitHub

Repository navigation

Commit 7f934d6

Browse files
authored
Merge pull request fastapi#27 from leynier/fork-10
Fork 10
2 parents 633b230 + 705b4d8 commit 7f934d6

2 files changed

Lines changed: 76 additions & 13 deletions

File tree

‎sqlmodel/main.py‎

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -30,18 +30,9 @@
3030
from pydantic.main import ModelMetaclass, validate_model
3131
from pydantic.typing import ForwardRef, NoArgAnyCallable, resolve_annotations
3232
from pydantic.utils import ROOT_KEY, Representation
33-
from sqlalchemy import (
34-
Boolean,
35-
Column,
36-
Date,
37-
DateTime,
38-
Float,
39-
ForeignKey,
40-
Integer,
41-
Interval,
42-
Numeric,
43-
inspect,
44-
)
33+
from sqlalchemy import Boolean, Column, Date, DateTime
34+
from sqlalchemy import Enum as sa_Enum
35+
from sqlalchemy import Float, ForeignKey, Integer, Interval, Numeric, inspect
4536
from sqlalchemy.orm import RelationshipProperty, declared_attr, registry, relationship
4637
from sqlalchemy.orm.attributes import set_attribute
4738
from sqlalchemy.orm.decl_api import DeclarativeMeta
@@ -395,7 +386,7 @@ def get_sqlachemy_type(field: ModelField) -> Any:
395386
if issubclass(field.type_, time):
396387
return Time
397388
if issubclass(field.type_, Enum):
398-
return Enum
389+
return sa_Enum(field.type_)
399390
if issubclass(field.type_, bytes):
400391
return LargeBinary
401392
if issubclass(field.type_, Decimal):

‎tests/test_enums.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import enum
2+
import uuid
3+
4+
from sqlalchemy import create_mock_engine
5+
from sqlalchemy.sql.type_api import TypeEngine
6+
from sqlmodel import Field, SQLModel
7+
8+
"""
9+
Tests related to Enums
10+
11+
Associated issues:
12+
* https://github.com/tiangolo/sqlmodel/issues/96
13+
* https://github.com/tiangolo/sqlmodel/issues/164
14+
"""
15+
16+
17+
class MyEnum1(enum.Enum):
18+
A = "A"
19+
B = "B"
20+
21+
22+
class MyEnum2(enum.Enum):
23+
C = "C"
24+
D = "D"
25+
26+
27+
class BaseModel(SQLModel):
28+
id: uuid.UUID = Field(primary_key=True)
29+
enum_field: MyEnum2
30+
31+
32+
class FlatModel(SQLModel, table=True):
33+
id: uuid.UUID = Field(primary_key=True)
34+
enum_field: MyEnum1
35+
36+
37+
class InheritModel(BaseModel, table=True):
38+
pass
39+
40+
41+
def pg_dump(sql: TypeEngine, *args, **kwargs):
42+
dialect = sql.compile(dialect=postgres_engine.dialect)
43+
sql_str = str(dialect).rstrip()
44+
if sql_str:
45+
print(sql_str + ";")
46+
47+
48+
def sqlite_dump(sql: TypeEngine, *args, **kwargs):
49+
dialect = sql.compile(dialect=sqlite_engine.dialect)
50+
sql_str = str(dialect).rstrip()
51+
if sql_str:
52+
print(sql_str + ";")
53+
54+
55+
postgres_engine = create_mock_engine("postgresql://", pg_dump)
56+
sqlite_engine = create_mock_engine("sqlite://", sqlite_dump)
57+
58+
59+
def test_postgres_ddl_sql(capsys):
60+
SQLModel.metadata.create_all(bind=postgres_engine, checkfirst=False)
61+
62+
captured = capsys.readouterr()
63+
assert "CREATE TYPE myenum1 AS ENUM ('A', 'B');" in captured.out
64+
assert "CREATE TYPE myenum2 AS ENUM ('C', 'D');" in captured.out
65+
66+
67+
def test_sqlite_ddl_sql(capsys):
68+
SQLModel.metadata.create_all(bind=sqlite_engine, checkfirst=False)
69+
70+
captured = capsys.readouterr()
71+
assert "enum_field VARCHAR(1) NOT NULL" in captured.out
72+
assert "CREATE TYPE" not in captured.out

0 commit comments

Comments
 (0)

Back | FazBrowse Home | New Git URL