SQLAlchemy Enum type breaks

[sebastian-correa]

I'm trying to use sqlmodel to handle a BigQuery dataset. The model I'm using has an Enum field, declared from a Python Enum. Table creation fails.

Environment details

• MacOS 12.6.3 on an Intel Macbook Pro 13.
• Python 3.10.8.
• pip 22.3 (using poetry 1.2.2).
• sqlalchemy-bigquery version:

    Name: sqlalchemy-bigquery
    Version: 1.6.1
    Summary: SQLAlchemy dialect for BigQuery
    Home-page: https://github.com/googleapis/python-bigquery-sqlalchemy
    Author: The Sqlalchemy-Bigquery Authors
    Author-email: googleapis-packages@google.com
    License: 
    Location: /Users/me/Dev/myproj/.venv/lib/python3.10/site-packages
    Requires: future, google-api-core, google-auth, google-cloud-bigquery, google-cloud-bigquery-storage, packaging, pyarrow, 
    sqlalchemy
    Required-by:
  

Steps to reproduce

1. Have a SQLAlchemy model with an Enum column declared from a python Enum.
2. Run create_all or try to add and commit to a hand-made table in BigQuery (the code Enum column mapped to a column with the same name, of type string with length = max Enum member length).

Code example

from enum import Enum

from sqlmodel import Column
from sqlmodel import Enum as DbEnum
from sqlmodel import Field, SQLModel, create_engine


class ValidMembers(str, Enum):
    member_1 = "member 1"
    member_2 = "member 2"
    member_3 = "member 3"


class TestTable(SQLModel, table=True):
    __tablename__ = "test_table"

    id: int | None = Field(default=None, primary_key=True)
    member: ValidMembers = Field(
        sa_column=Column(
            DbEnum(ValidMembers),
            nullable=False,
        )
    )


ENGINE = create_engine(url="bigquery://<project_id>/<default_dataset>", credentials_path="path/to/credentials-sa.json")


def create_db_and_tables():
    SQLModel.metadata.create_all(ENGINE)


if __name__ == "__main__":
    create_db_and_tables()

Stack trace

Traceback (most recent call last):
  File "/Users/me/Dev/myproj/test.py", line 34, in <module>
    create_db_and_tables()
  File "/Users/me/Dev/myproj/test.py", line 30, in create_db_and_tables
    SQLModel.metadata.create_all(ENGINE)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/schema.py", line 4917, in create_all
    bind._run_ddl_visitor(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/future/engine.py", line 342, in _run_ddl_visitor
    conn._run_ddl_visitor(visitorcallable, element, **kwargs)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 2211, in _run_ddl_visitor
    visitorcallable(self.dialect, self, **kwargs).traverse_single(element)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/visitors.py", line 524, in traverse_single
    return meth(obj, **kw)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/ddl.py", line 838, in visit_metadata
    metadata.dispatch.before_create(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/event/attr.py", line 334, in __call__
    fn(*args, **kw)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py", line 875, in __call__
    return getattr(self.target, self.name)(*arg, **kw)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/sqltypes.py", line 1203, in _on_metadata_create
    t = self.dialect_impl(bind.dialect)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 649, in dialect_impl
    return self._dialect_info(dialect)["impl"]
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 726, in _dialect_info
    impl = self._gen_dialect_impl(dialect)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 735, in _gen_dialect_impl
    return dialect.type_descriptor(self)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/default.py", line 578, in type_descriptor
    return sqltypes.adapt_type(typeobj, self.colspecs)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 1974, in adapt_type
    return typeobj.adapt(impltype)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/sqltypes.py", line 1711, in adapt
    return super(Enum, self).adapt(impltype, **kw)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 1134, in adapt
    return super(Emulated, self).adapt(impltype, **kw)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 759, in adapt
    return util.constructor_copy(self, cls, **kw)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py", line 1325, in constructor_copy
    return cls(*args, **kw)
TypeError: String.__init__() got an unexpected keyword argument '_enums'

[sebastian-correa]

Hey, @shollyman or @meredithslota! I see you usually respond to issues in this repo. Did you have a chance to look into this?

[benvdh-incentro]

@sebastian-correa As far as I can see in the BigQuery Data types docs, it does not have an Enum data type, at least not in the sense that a server like PostgreSQL does:

https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#string_type

You might have more luck implementing it more along these lines (as BigQuery does support integers):
fastapi/sqlmodel#96 (comment)

[sebastian-correa]

Hey @benvdh-incentro, thanks for your tips. I tried them and the issue got solved, though some other unrelated issues popped up; see the end of the message for more details.

I would have expected this dialect to handle the case with Enum columns, as members can be mapped to strings (always) using their names. If you debug the package's code, you'll see that the column gets mapped correctly to a string, but the code fails later because the string constructor doesn't take an _enums field.

In our particular case, we want to support connecting to BQ or to a PSQL instance by an environment variable, that's why we can't have too much custom code on the model side.

I think that if this dialect mapped an Enum column to a string column with the length of the maximum Enum member name, and it didn't forward the _enum kwarg, everything would work. This obviously won't have any constraints, but that's fine in our use-case.

---

Enum issue aside, if we accept to work with the ChoiceType, my full model also has an updated_at column which is a datetime with a python default of None:

class TestTable(SQLModel, table=True):
    __tablename__ = "test_table"

    id: int | None = Field(default=None, primary_key=True)
    member: ValidMembers = Field(
        sa_column=Column(
            ChoiceType(ValidMembers, impl=String(length=8)),
            nullable=False,
        )
    )
    updated_at: datetime | None = Field(
        default=None,
        sa_column=Column(
            "updated_at",
            DateTime(timezone=False),
            onupdate=utc_now,
            default=utc_now,
            nullable=False,
        ),
    )

with Session(ENGINE) as session:
    session.add(TestTable(member=ValidMembers.member_1))
    session.commit()

The column's value should get set when the item is inserted into the DB or updated, but be None in the Python object. However, I get this error:

Traceback (most recent call last):
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/dbapi/cursor.py", line 203, in _execute
    self._query_job.result()
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/job/query.py", line 1499, in result
    do_get_result()
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/api_core/retry.py", line 349, in retry_wrapped_func
    return retry_target(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/api_core/retry.py", line 191, in retry_target
    return target()
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/job/query.py", line 1489, in do_get_result
    super(QueryJob, self).result(retry=retry, timeout=timeout)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/job/base.py", line 728, in result
    return super(_AsyncJob, self).result(timeout=timeout, **kwargs)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/api_core/future/polling.py", line 261, in result
    raise self._exception
google.api_core.exceptions.BadRequest: 400 Required field updated_at cannot be null

Location: US
Job ID: <some_uuid>


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1900, in _execute_context
    self.dialect.do_execute(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/default.py", line 736, in do_execute
    cursor.execute(statement, parameters)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/dbapi/_helpers.py", line 494, in with_closed_check
    return method(self, *args, **kwargs)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/dbapi/cursor.py", line 166, in execute
    self._execute(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/dbapi/cursor.py", line 205, in _execute
    raise exceptions.DatabaseError(exc)
google.cloud.bigquery.dbapi.exceptions.DatabaseError: 400 Required field updated_at cannot be null

Location: US
Job ID: <some_uuid>


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/me/Dev/myproj/test.py", line 47, in <module>
    session.commit()
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 1451, in commit
    self._transaction.commit(_to_root=self.future)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 829, in commit
    self._prepare_impl()
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 808, in _prepare_impl
    self.session.flush()
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 3386, in flush
    self._flush(objects)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 3525, in _flush
    with util.safe_reraise():
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py", line 70, in __exit__
    compat.raise_(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/util/compat.py", line 208, in raise_
    raise exception
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/session.py", line 3486, in _flush
    flush_context.execute()
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/unitofwork.py", line 456, in execute
    rec.execute(self)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/unitofwork.py", line 630, in execute
    util.preloaded.orm_persistence.save_obj(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/persistence.py", line 245, in save_obj
    _emit_insert_statements(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/orm/persistence.py", line 1238, in _emit_insert_statements
    result = connection._execute_20(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1705, in _execute_20
    return meth(self, args_10style, kwargs_10style, execution_options)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/sql/elements.py", line 333, in _execute_on_connection
    return connection._execute_clauseelement(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1572, in _execute_clauseelement
    ret = self._execute_context(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1943, in _execute_context
    self._handle_dbapi_exception(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 2124, in _handle_dbapi_exception
    util.raise_(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/util/compat.py", line 208, in raise_
    raise exception
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1900, in _execute_context
    self.dialect.do_execute(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/sqlalchemy/engine/default.py", line 736, in do_execute
    cursor.execute(statement, parameters)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/dbapi/_helpers.py", line 494, in with_closed_check
    return method(self, *args, **kwargs)
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/dbapi/cursor.py", line 166, in execute
    self._execute(
  File "/Users/me/Dev/myproj/.venv/lib/python3.10/site-packages/google/cloud/bigquery/dbapi/cursor.py", line 205, in _execute
    raise exceptions.DatabaseError(exc)
sqlalchemy.exc.DatabaseError: (google.cloud.bigquery.dbapi.exceptions.DatabaseError) 400 Required field updated_at cannot be null

Location: US
Job ID: <some_uuid>

[SQL: INSERT INTO `test_table` (`member`, `updated_at`) VALUES (%(member:STRING)s, %(updated_at:DATETIME)s)]
[parameters: {'member': 'member 1', 'updated_at': None}]
(Background on this error at: https://sqlalche.me/e/14/4xp6)

[meredithslota]

Thank you for the detailed report!

[chalmerlowe]

This appears to be a very specific issue with limited overall impact. Due to staffing and team priorities I am going to mark this as will not fix and close.

[chalmerlowe]

This PR should take care of the primary symptom raised by this Issue.
#1111