Repository: agronholm/sqlacodegen

Issue:
AssertionError: TypeDecorator implementations require a class-level variable 'impl'
### Things to check first

- [X] I have searched the existing issues and didn't find my bug already reported there

- [X] I have checked that my bug is still present in the latest release


### Sqlacodegen version

3.0.0rc3

### SQLAlchemy version

2.0.25

### RDBMS vendor

Other

### What happened?

I'm seeing the below error when trying to generate code for a `databricks` "database"
```
❯ sqlacodegen "${db_uri}"
Traceback (most recent call last):
  File "/opt/python/envs/dev310/bin/sqlacodegen", line 11, in <module>
    sys.exit(main())
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlacodegen/cli.py", line 92, in main
    outfile.write(generator.generate())
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlacodegen/generators.py", line 172, in generate
    self.fix_column_types(table)
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlacodegen/generators.py", line 649, in fix_column_types
    column.type = self.get_adapted_type(column.type)
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlacodegen/generators.py", line 681, in get_adapted_type
    new_coltype = coltype.adapt(supercls)
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 1033, in adapt
    return util.constructor_copy(
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py", line 1422, in constructor_copy
    return cls(*args, **kw)
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 1692, in __init__
    raise AssertionError(
AssertionError: TypeDecorator implementations require a class-level variable 'impl' which refers to the class of type being decorated
```


Just posting to remind me to investigate further and in case others are also seeing this...

### Database schema for reproducing the bug

_No response_

Hints from issue comments:
@dhirschfeld: No luck with the latest rc5.

This is where the `AssertionError` is bubbling up:
https://github.com/agronholm/sqlacodegen/blob/2a6053224d28da27c9b89aa611b9dd4c27637fe6/src/sqlacodegen/generators.py#L675-L679@dhirschfeld: It appears to work fine if I change it to `except Exception` so it seems to be a simple case of too narrow of an `except` clause.

Is `except Exception` an acceptable fix? It could be considered too broad, in which case `except (TypeError, AssertionError):` might be better?

Happy to put in a PR...@agronholm: Before suppressing any exceptions, I'd like to understand the root cause of this. Could you find out for me?@dhirschfeld: Yep, I'm digging into it now...@dhirschfeld: This is with [`databricks-sql-python`](https://github.com/databricks/databricks-sql-python) so I suspect there's something fishy about their (relatively new) `sqlalchemy` integration.

If I patch my implementation as:
```python
                try:
                    new_coltype = coltype.adapt(supercls)
                except TypeError:
                    # If the adaptation fails, don't try again
                    break
                except Exception:
                    print('#'*80)
                    print(coltype)
                    print(coltype.__class__)
                    print(coltype.__class__.mro())
                    print(supercls)
                    print('#'*80)
                    raise
```
...I get:
```
################################################################################
DATETIME
<class 'databricks.sqlalchemy._types.TIMESTAMP'>
[
    <class 'databricks.sqlalchemy._types.TIMESTAMP'>,
    <class 'sqlalchemy.sql.type_api.TypeDecorator'>,
    <class 'sqlalchemy.sql.base.SchemaEventTarget'>,
    <class 'sqlalchemy.event.registry.EventTarget'>,
    <class 'sqlalchemy.sql.type_api.ExternalType'>,
    <class 'sqlalchemy.sql.type_api.TypeEngineMixin'>,
    <class 'sqlalchemy.sql.type_api.TypeEngine'>,
    <class 'sqlalchemy.sql.visitors.Visitable'>,
    <class 'typing.Generic'>,
    <class 'object'>
]
<class 'sqlalchemy.sql.type_api.TypeDecorator'>
################################################################################
Traceback (most recent call last):
  File "/opt/python/envs/dev310/bin/sqlacodegen", line 10, in <module>
    sys.exit(main())
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlacodegen/cli.py", line 101, in main
    outfile.write(generator.generate())
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlacodegen/generators.py", line 171, in generate
    self.fix_column_types(table)
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlacodegen/generators.py", line 644, in fix_column_types
    column.type = self.get_adapted_type(column.type)
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlacodegen/generators.py", line 676, in get_adapted_type
    new_coltype = coltype.adapt(supercls)
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 1032, in adapt
    return util.constructor_copy(
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py", line 1414, in constructor_copy
    return cls(*args, **kw)
  File "/opt/python/envs/dev310/lib/python3.10/site-packages/sqlalchemy/sql/type_api.py", line 1686, in __init__
    raise AssertionError(
AssertionError: TypeDecorator implementations require a class-level variable 'impl' which refers to the class of type being decorated
```@agronholm: Yeah, `AssertionError` indicates that something is very, very wrong there, and I'd prefer not to mask such errors.@dhirschfeld: Agreed - it's highlighting a actual problem with their implementation of the `TIMESTAMP` type.

I'll close this and open an issue on their repo...@dhirschfeld: Thanks for taking the time to triage this issue! @dhirschfeld: It's curious because the `TIMESTAMP` class does have a class-level `impl` attribute:
```python
impl = sqlalchemy.types.DateTime
```
https://github.com/databricks/databricks-sql-python/blob/50489346440cdf9e547b597254643ee4ceb28c19/src/databricks/sqlalchemy/_types.py#L147@dhirschfeld: ...and seems to be following the recommended recipe to augment existing types:
https://docs.sqlalchemy.org/en/20/core/custom_types.html#augmenting-existing-types@dhirschfeld: It seems that `sqlqcodegen` is trying to directly instantiate a `TypeDecorator` class (as that's the correct super-class of `TIMESTAMP`) but that will never work as it's not designed to be instantiated directly and in this case the constructor can raise an `AssertionError`:
```python
        if not hasattr(self.__class__, "impl"):
            raise AssertionError(
                "TypeDecorator implementations "
                "require a class-level variable "
                "'impl' which refers to the class of "
                "type being decorated"
            )
```
https://github.com/sqlalchemy/sqlalchemy/blob/9bcc4da735891d09a4c850c5f29b3abeef13ce27/lib/sqlalchemy/sql/type_api.py#L1685-L1691@dhirschfeld: ...so I think perhaps the issue lies with a too-strict exception clause here.

In the specific instance that the `coltype` is a subclass of `TypeDecorator` the call to `coltype.adapt(supercls)` will fail with an `AssertionError` so the except clause should then be:
```python
except (TypeError, AssertionError):
```
to also handle this specific case.

Thoughts?@dhirschfeld: If you wanted to be a bit more defensive you could instead do:
```python
                try:
                    new_coltype = coltype.adapt(supercls)
                except TypeError:
                    # If the adaptation fails, don't try again
                    break
                except AssertionError:
                    if supercls is TypeDecorator:
                        # trying to instantiate a TypeDecorator class
                        # will fail with an AssertionError
                        break
                    else:
                        raise
```@agronholm: I think I would like to check if the coltype is a type decorator, thus avoiding `AssertionError` completely.

Solution patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 431359f..a1159ca 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -9,6 +9,8 @@ Version history
 - Fixed two rendering issues in ``ENUM`` columns when a non-default schema is used: an
   unwarranted positional argument and missing the ``schema`` argument
 - Fixed ``AttributeError`` when metadata contains user defined column types
+- Fixed ``AssertionError`` when metadata contains a column type that is a type decorator
+  with an all-uppercase name
 
 **3.0.0rc5**
 
diff --git a/src/sqlacodegen/generators.py b/src/sqlacodegen/generators.py
index b3fc777..f992624 100644
--- a/src/sqlacodegen/generators.py
+++ b/src/sqlacodegen/generators.py
@@ -36,6 +36,7 @@
     String,
     Table,
     Text,
+    TypeDecorator,
     UniqueConstraint,
 )
 from sqlalchemy.dialects.postgresql import JSONB
@@ -684,7 +685,7 @@ def get_adapted_type(self, coltype: Any) -> Any:
                 supercls, "__visit_name__"
             ):
                 # Don't try to adapt UserDefinedType as it's not a proper column type
-                if supercls is UserDefinedType:
+                if supercls is UserDefinedType or issubclass(supercls, TypeDecorator):
                     return coltype
 
                 # Hack to fix adaptation of the Enum class which is broken since


Test patch:
diff --git a/tests/test_generator_tables.py b/tests/test_generator_tables.py
index fc99b7a..0bbd322 100644
--- a/tests/test_generator_tables.py
+++ b/tests/test_generator_tables.py
@@ -4,6 +4,7 @@
 
 import pytest
 from _pytest.fixtures import FixtureRequest
+from sqlalchemy import TypeDecorator
 from sqlalchemy.dialects import mysql, postgresql
 from sqlalchemy.engine import Engine
 from sqlalchemy.schema import (
@@ -18,7 +19,7 @@
     UniqueConstraint,
 )
 from sqlalchemy.sql.expression import text
-from sqlalchemy.sql.sqltypes import NullType
+from sqlalchemy.sql.sqltypes import DateTime, NullType
 from sqlalchemy.types import INTEGER, NUMERIC, SMALLINT, VARCHAR, Text
 
 from sqlacodegen.generators import CodeGenerator, TablesGenerator
@@ -26,6 +27,11 @@
 from .conftest import validate_code
 
 
+# This needs to be uppercased to trigger #315
+class TIMESTAMP_DECORATOR(TypeDecorator[DateTime]):
+    impl = DateTime
+
+
 @pytest.fixture
 def generator(
     request: FixtureRequest, metadata: MetaData, engine: Engine
@@ -45,12 +51,15 @@ def test_fancy_coltypes(generator: CodeGenerator) -> None:
         Column("bool", postgresql.BOOLEAN),
         Column("vector", VECTOR(3)),
         Column("number", NUMERIC(10, asdecimal=False)),
+        Column("timestamp", TIMESTAMP_DECORATOR()),
         schema="someschema",
     )
 
     validate_code(
         generator.generate(),
         """\
+        from tests.test_generator_tables import TIMESTAMP_DECORATOR
+
         from pgvector.sqlalchemy.vector import VECTOR
         from sqlalchemy import Boolean, Column, Enum, MetaData, Numeric, Table
 
@@ -63,6 +72,7 @@ def test_fancy_coltypes(generator: CodeGenerator) -> None:
             Column('bool', Boolean),
             Column('vector', VECTOR(3)),
             Column('number', Numeric(10, asdecimal=False)),
+            Column('timestamp', TIMESTAMP_DECORATOR),
             schema='someschema'
         )
         """,