AssertionError: TypeDecorator implementations require a class-level variable 'impl'

[dhirschfeld]

Things to check first

• I have searched the existing issues and didn't find my bug already reported there

• 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

[dhirschfeld]

No luck with the latest rc5.

This is where the AssertionError is bubbling up:

sqlacodegen/src/sqlacodegen/generators.py

Lines 675 to 679 in 2a60532

	try:
	new_coltype = coltype.adapt(supercls)
	except TypeError:
	# If the adaptation fails, don't try again
	break

[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 so I suspect there's something fishy about their (relatively new) sqlalchemy integration.

If I patch my implementation as:

                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:

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:

        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:

except (TypeError, AssertionError):

to also handle this specific case.

Thoughts?

[dhirschfeld]

If you wanted to be a bit more defensive you could instead do:

                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.