kaxil commented on issue #52378:
URL: https://github.com/apache/airflow/issues/52378#issuecomment-3015071634
More details, especially if you plan to use AI for this work (you can feed
this as a prompt):
---
## ⚠️ CRITICAL WARNING ⚠️
**DO NOT MODIFY THE TASK SDK (`airflow.sdk`)**
- Never add methods to `airflow.sdk.BaseOperator`
- Never modify any components in `airflow.sdk`
- Use version compatibility patterns to work around SDK limitations
- If the entire provider directory does not import `BaseOperator`, no
changes should be made to the Provider.
## Summary
Update all provider packages to use the new Task SDK `BaseOperator` import
for Airflow 3.0 compatibility while maintaining backward compatibility with
Airflow 2.x. This involves updating import statements, fixing test patterns,
and resolving xcom_push method signature issues.
## Migration Strategy
### 1. Version Compatibility Pattern
**For ALL providers**, use this centralized version compatibility approach:
```python
# In
providers/{provider_name}/src/airflow/providers/{provider_name}/version_compat.py
from __future__ import annotations
def get_base_airflow_version_tuple() -> tuple[int, int, int]:
from packaging.version import Version
from airflow import __version__
airflow_version = Version(__version__)
return airflow_version.major, airflow_version.minor,
airflow_version.micro
# Choose the right version flag based on your provider's needs:
# - Use AIRFLOW_V_3_0_PLUS for most providers (BaseOperator,
BaseSensorOperator, etc.)
# - Use AIRFLOW_V_3_1_PLUS only if you have decorators with xcom_push issues
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseOperator # Use BaseOperator for most
providers
# Only include other components if your provider uses them:
# from airflow.sdk import BaseOperatorLink, BaseSensorOperator
# from airflow.sdk.execution_time.xcom import XCom
else:
from airflow.models import BaseOperator
# Only include other components if your provider uses them:
# from airflow.models import BaseOperatorLink, XCom
# from airflow.sensors.base import BaseSensorOperator
__all__ = [
"AIRFLOW_V_3_0_PLUS",
"BaseOperator",
# Only include what your provider actually uses:
# "BaseSensorOperator",
# "BaseOperatorLink",
# "XCom",
]
```
**IMPORTANT**: The above is a **template showing all possible imports**.
Each provider should:
1. **Always include**: `AIRFLOW_V_3_1_PLUS` and `BaseOperator` (every
provider needs these)
2. **Only include if used**: Remove any SDK components your provider doesn't
actually use
3. **Keep it minimal**: Only import what you need to avoid unnecessary
dependencies
**Example for a simple provider** (like most providers with only operators,
no decorators):
```python
# Minimal version_compat.py for providers with only operators
from __future__ import annotations
def get_base_airflow_version_tuple() -> tuple[int, int, int]:
from packaging.version import Version
from airflow import __version__
airflow_version = Version(__version__)
return airflow_version.major, airflow_version.minor,
airflow_version.micro
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import BaseOperator
else:
from airflow.models import BaseOperator
__all__ = [
"AIRFLOW_V_3_0_PLUS",
"BaseOperator",
]
```
**Example for providers with decorators that have xcom_push issues** (use
AIRFLOW_V_3_1_PLUS):
```python
# For providers with decorators that trigger xcom_push bugs in 3.0.x
from __future__ import annotations
def get_base_airflow_version_tuple() -> tuple[int, int, int]:
from packaging.version import Version
from airflow import __version__
airflow_version = Version(__version__)
return airflow_version.major, airflow_version.minor,
airflow_version.micro
AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
if AIRFLOW_V_3_1_PLUS:
from airflow.sdk import BaseOperator
else:
from airflow.models import BaseOperator
__all__ = [
"AIRFLOW_V_3_1_PLUS",
"BaseOperator",
]
```
**Key Decision Rules**:
- Use `AIRFLOW_V_3_0_PLUS` for BaseOperator imports in most providers
- Use `AIRFLOW_V_3_1_PLUS` for BaseOperator ONLY if your provider has
decorators that use xcom_push (avoids SDK BaseOperator xcom_push bug in Airflow
3.0.x)
- Use `AIRFLOW_V_3_0_PLUS` for other SDK components (BaseOperatorLink,
BaseSensorOperator, XCom)
### 2. Update All Provider Files
**Replace direct imports** in all operator, sensor, hook, and link files:
```python
# OLD PATTERN
from airflow.models.baseoperator import BaseOperator
# NEW PATTERN
from airflow.providers.{provider_name}.version_compat import BaseOperator
```
**For TYPE_CHECKING blocks**, maintain compatibility:
```python
if TYPE_CHECKING:
from airflow.providers.{provider_name}.version_compat import
AIRFLOW_V_3_1_PLUS
if AIRFLOW_V_3_1_PLUS:
from airflow.sdk import BaseOperator
else:
from airflow.models import BaseOperator
```
### 3. Fix BaseOperatorLink.persist Method Signatures
**Critical Issue**: Many providers have `BaseOperatorLink.persist` methods
that unnecessarily pass `task_instance: BaseOperator` parameter and call
`task_instance.xcom_push()`, causing AttributeError with SDK BaseOperator.
**Fix Pattern**:
```python
# OLD PATTERN - causes AttributeError
@staticmethod
def persist(
context: Context,
task_instance: BaseOperator, # ❌ Remove this parameter
log_view_url: str,
):
task_instance.xcom_push(key="link_key", value=log_view_url) # ❌ Fails
# NEW PATTERN - works universally
@staticmethod
def persist(
context: Context,
log_view_url: str,
):
context["task_instance"].xcom_push(key="link_key", value=log_view_url)
# ✅ Works
```
**Common xcom_push Signature Errors**:
```python
# WRONG - extra context parameter
context["task_instance"].xcom_push(context, key="key", value="value")
# CORRECT - no context parameter
context["task_instance"].xcom_push(key="key", value="value")
```
### 4. Test Pattern Updates
**Mock Context Type Issues**:
```python
# OLD PATTERN - causes mypy errors
mock_context = {"ti": mock_ti}
result = task.execute(context=mock_context) # ❌ Type error
# NEW PATTERN - add type ignore for test simplicity
mock_context = {"ti": mock_ti, "task_instance": mock_ti}
result = task.execute(context=mock_context) # type: ignore[arg-type] # ✅
Works
```
**Decorator Test Patterns**:
```python
# OLD PATTERN - fails in Airflow 3.0
@task.virtualenv
def f():
return "hello"
ret = f()
ret.operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # ❌
AttributeError
# NEW PATTERN - use dag_maker.run_ti()
@task.virtualenv
def f():
return "hello"
ret = f()
dr = dag_maker.create_dagrun()
dag_maker.run_ti("f", dr) # ✅ Works - use function name directly
```
**Critical DAG Creation Pattern**:
```python
# WRONG - causes task instance issues
dr = self.create_dag_run()
dag_maker.run_ti("task_id", dr) # ❌ May fail
# CORRECT - ensures proper task instance creation
dr = self.dag_maker.create_dagrun()
dag_maker.run_ti("task_id", dr) # ✅ Works correctly
```
## Automation Checklist
For agentic AI to automate this migration:
### Phase 1: Version Compatibility Setup
1. ✅ Create/update `version_compat.py` in each provider's root package
2. ✅ Use the standardized version compatibility pattern above
3. ✅ Always include `BaseOperator` (required for all providers)
4. ✅ **Choose the right version flag**: Use `AIRFLOW_V_3_0_PLUS` for most
providers, `AIRFLOW_V_3_1_PLUS` only for providers with decorator xcom_push
issues
5. ✅ **ONLY include other SDK components if the provider actually uses them**
6. ✅ Remove unused imports and comments to keep the file minimal and clean
### Phase 2: Import Updates
1. ✅ Replace all direct `BaseOperator` imports with version_compat imports
2. ✅ Update TYPE_CHECKING blocks to use version_compat
3. ✅ Replace other SDK imports (BaseOperatorLink, BaseSensorOperator, XCom)
### Phase 3: Method Signature Fixes
1. ✅ Find all `BaseOperatorLink.persist` methods with `task_instance:
BaseOperator` parameters
2. ✅ Remove the unnecessary `task_instance` parameter
3. ✅ Replace `task_instance.xcom_push()` with
`context["task_instance"].xcom_push()`
4. ✅ Remove extra `context` parameters from xcom_push calls
### Phase 4: Test Updates
1. ✅ Replace `ret.operator.run()` patterns with `dag_maker.run_ti()`
2. ✅ Use function names directly for decorator tests
3. ✅ Replace `self.create_dag_run()` with `self.dag_maker.create_dagrun()`
4. ✅ Add `# type: ignore[arg-type]` for mock context type issues
### Phase 5: Validation
1. ✅ Run `breeze testing providers-tests --test-type
"Providers[{provider_name}]"`
2. ✅ Verify no `AttributeError: object has no attribute 'run'` errors
3. ✅ Verify no `AttributeError: object has no attribute 'xcom_push'` errors
4. ✅ Run code quality checks:
```bash
# Linting and formatting
pre-commit run ruff --all-files
pre-commit run ruff-format --all-files
# Type checking (in breeze shell)
mypy providers/{provider_name}/src/airflow/providers/{provider_name}/
providers/{provider_name}/tests/unit/
```
5. ✅ Test with both Airflow 2.x and 3.x versions
## Key Files to Update Per Provider
**Core Files**:
- `src/airflow/providers/{provider}/version_compat.py` (create/update)
- `src/airflow/providers/{provider}/operators/*.py` (BaseOperator imports)
- `src/airflow/providers/{provider}/sensors/*.py` (BaseSensorOperator
imports)
- `src/airflow/providers/{provider}/hooks/*.py` (type hint imports)
- `src/airflow/providers/{provider}/links/*.py` (BaseOperatorLink imports +
persist methods)
**Test Files**:
- `tests/unit/{provider}/**/*.py` (test patterns, mock contexts)
- `tests/system/{provider}/**/*.py` (if applicable)
## Completed Providers
- ✅ **Standard Provider**: Complete migration with all patterns
- ✅ **OpenLineage Provider**: BaseOperator imports updated
- ✅ **Airbyte Provider**: Version compatibility implemented
- ✅ **Alibaba Provider**: Including MaxComputeLogViewLink.persist fix
- ✅ **Google Provider**: Marketing platform operators + links fixed
- ✅ **Asana Provider**: Minimal version_compat.py with AIRFLOW_V_3_0_PLUS +
BaseOperator (clean example)
- ✅ **ArangoDB Provider**: AIRFLOW_V_3_0_PLUS + BaseOperator +
BaseSensorOperator, cleaned up existing version_compat.py
- ✅ **Apprise Provider**: No migration needed (only hooks/notifications, no
operators/sensors)
- ✅ **Apache Pig Provider**: Minimal version_compat.py with
AIRFLOW_V_3_0_PLUS + BaseOperator
- ✅ **DBT Cloud Provider**: Enhanced version_compat.py with
AIRFLOW_V_3_0_PLUS + BaseOperator + BaseSensorOperator + BaseOperatorLink + XCom
- ✅ **Atlassian JIRA Provider**: Enhanced version_compat.py with
AIRFLOW_V_3_0_PLUS + BaseOperator + BaseSensorOperator, fixed xcom_push pattern
in operator
## Common Error Patterns & Solutions
**Error**: `"xcom_push" gets multiple values for keyword argument "key"`
**Solution**: Remove extra `context` parameter from xcom_push call
**Error**: `Argument 1 to "xcom_push" has incompatible type "Context";
expected "str"`
**Solution**: Remove `context` as first argument to xcom_push
**Error**: `Argument "context" to "execute" has incompatible type "dict[str,
MagicMock]"; expected "Context"`
**Solution**: Add `# type: ignore[arg-type]` to test execute calls
**Error**: `AttributeError: object has no attribute 'run'`
**Solution**: Replace `operator.run()` with `dag_maker.run_ti(task_id, dr)`
**Error**: `Name "BaseSensorOperator" already defined`
**Solution**: Ensure single import from version_compat, remove duplicate
imports
## Testing in Breeze
```bash
# Test specific provider
breeze testing providers-tests --test-type "Providers[{provider_name}]"
# Test with different Airflow versions
breeze shell --airflow-version 2.10.5 # Test backward compatibility
breeze shell --airflow-version 3.0.2 # Test forward compatibility
breeze shell # Test with latest (3.1.0+)
# Run specific failing tests
breeze testing providers-tests
providers/{provider}/tests/unit/path/to/test.py::TestClass::test_method -v
# Code quality validation
pre-commit run ruff --all-files # Linting
pre-commit run ruff-format --all-files # Formatting
# Type checking (run in breeze shell)
breeze shell
mypy providers/{provider_name}/src/airflow/providers/{provider_name}/
providers/{provider_name}/tests/unit/
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
