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]

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