Mypy sqlalchemy issues with TaskInstance / context - Column[String] instead of str

[HTRafal]

Apache Airflow version

2.7.1

What happened

I'm trying to migrate from 2.5.0 to 2.7.1. This simple snippet fails mypy linting now:

from datetime import datetime
from airflow.decorators import dag

@dag(schedule=None, start_date=datetime(2021, 1, 1))
def alwayspass():
    def go():
        from airflow.operators.python import get_current_context
        context = get_current_context()
        # dag_test.py:28: error: Incompatible types in assignment (expression has type "Union[Column[String], str]", variable has type "str")  [assignment]
        task_id: str = context["task_instance"].task_id
    go()

job = alwayspass()

I digged a bit and I believe it's due to this update - https://github.com/apache/airflow/pull/31486/files#diff-63f51df3fc306cb89897939da64682b68738cf21cb267d289d48e15bbc97bce8R50
You're using SQLAlchemy 2.0 style while constraints.txt still uses SQLAlchemy==1.4.49 - https://raw.githubusercontent.com/apache/airflow/constraints-2.7.1/constraints-3.10.txt
Based on sqlalchemy documentation on mypy, you have to use declarative_base or the decorator for mypy plugin to work - https://docs.sqlalchemy.org/en/14/orm/extensions/mypy.html

What you think should happen instead

No mypy issues.

How to reproduce

pyenv virtualenv 3.10.9 airflow-2.7.1
pyenv shell airflow-2.7.1
pip install "apache-airflow[celery]==2.7.1" "sqlalchemy[mypy]" --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.7.1/constraints-3.10.txt"
mypy --check-untyped-defs dag_test.py

dag_test.py:

from datetime import datetime
from airflow.decorators import dag

@dag(schedule=None, start_date=datetime(2021, 1, 1))
def alwayspass():
    def go():
        from airflow.operators.python import get_current_context
        context = get_current_context()
        # dag_test.py:28: error: Incompatible types in assignment (expression has type "Union[Column[String], str]", variable has type "str")  [assignment]
        task_id: str = context["task_instance"].task_id
    go()

job = alwayspass()

Operating System

Ubuntu 22.04.3 LTS

Versions of Apache Airflow Providers

Default installation from pypi - https://airflow.apache.org/docs/apache-airflow/stable/installation/installing-from-pypi.html

Deployment

Virtualenv installation

Deployment details

Default installation from pypi - https://airflow.apache.org/docs/apache-airflow/stable/installation/installing-from-pypi.html

Anything else

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg[bot]]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[Taragolis]

From https://docs.sqlalchemy.org/en/20/orm/extensions/mypy.html

Warning

Deprecated since version 2.0: The SQLAlchemy Mypy Plugin is DEPRECATED, and will be removed possibly as early as the SQLAlchemy 2.1 release. We would urge users to please migrate away from it ASAP.

This plugin cannot be maintained across constantly changing releases of mypy and its stability going forward CANNOT be guaranteed.

Modern SQLAlchemy now offers fully pep-484 compliant mapping syntaxes; see the linked section for migration details.

You're using SQLAlchemy 2.0 style while constraints.txt still uses SQLAlchemy==1.4.49 -

Airflow doesn't support SA 2.0 yet #28723

[HTRafal]

Airflow doesn't support SA 2.0 yet

Not asking for SA 2 support. We don't modify SA version. It's set to 1.4.49 in airflow constraints. The problem is that you started using SA 2 style models and now mypy doesn't properly lint it (with mypy plugin for SA 1.4).

Maybe the request wasn't clear enough. To be able to use mypy with airflow 2.7.1 we have to put a bunch of # type:ignore in our code. I have however found the exact commit where you changed SA model base to a 2.0 style that broke mypy linting. My expectation was that maybe you could revert that back until you fully switch to SA 2.0? It seemed like that model change shouldn't be included until requirements specify sqlalchemy<2.0.

[Taragolis]

In general changes for migrate to SA2 would be added anyway, all changes which added in airflow codebase should be both compatible with SA 1.4 and SA 2.0 until we could say Airflow is SA2 compatible, for compatible I mean it should work on both versions. And this changes can not be done in the one go.

My expectation was that maybe you could revert that back until you fully switch to SA 2.0? It seemed like that model change shouldn't be included until requirements specify sqlalchemy<2.0.

I'm not sure about which changes are you talk about there are plenty number of changes made for make Airflow a bit more compatible with SA 2.0 by different Airflow contributors.

I'm not even sure that Airflow codebase was ever use sqlalchemy2-stubs, let alone support it.

I would propose to make PR and in this case you have an opportunity to show Airflow Community what you think would be better.
Airflow is a OSS and developed by the community, most of the community is also a users of Airflow

---

I just keep this links here

1. Contributor's Quick Start
2. Airflow Breeze
3. CONTRIBUTING.rst

[hterik]

all changes which added in airflow codebase should be both compatible with SA 1.4 and SA 2.0 until we could say Airflow is SA2 compatible, for compatible I mean it should work on both versions.

In current state of things it is not however. If you include mypy in the definition of compatible, it is not compatible with 1.4 anymore.
Based on discussion in #28723, 2.0 migration was not supposed to happen until after Airflow 2.7

The recent change https://github.com/apache/airflow/pull/31486/files#diff-63f51df3fc306cb89897939da64682b68738cf21cb267d289d48e15bbc97bce8R50 removes use of declarative_base, which breaks mypy on 1.4: https://docs.sqlalchemy.org/en/14/orm/extensions/mypy.html
(All this was described in the original ticket.)

I don't see an easy PR just simply fixing this, that would mean reverting changes above which i am not very keen on doing without maintainers blessing first. I would also like to see progress moving forward on upgrading to 2.0 instead of going backwards. The expectations on this issue would be to know what is the stance on such compatibility and maybe understanding more the roadmap on when 2.0 migration will be completed, maybe the latter is better tracked in #28723 alreday.

[hterik]

This is now fixed in Airflow 3.2 👍