Migrate AzureBatchHook/Operator to azure-batch 15.x (track 2 SDK)

[arieleli01212]

Description

The azure-batch package underwent a complete API overhaul in version 15.x, moving from the old Azure SDK track 1 (BatchServiceClient, batch_auth) to track 2 (BatchClient, azure-identity). This PR updates AzureBatchHook and AzureBatchOperator to work with the new SDK and removes the <15.0.0 upper bound cap that was explicitly blocking users from upgrading.

Changes

hooks/batch.py

• Replace BatchServiceClient with BatchClient(endpoint, credential)
• Replace batch_auth.SharedKeyCredentials with AzureNamedKeyCredential
• Catch ResourceExistsError (from azure.core.exceptions) instead of BatchErrorException
• Rename model types: PoolAddParameter → BatchPoolCreateOptions, JobAddParameter → BatchJobCreateOptions, TaskAddParameter → BatchTaskCreateOptions, ImageReference → BatchVmImageReference, PoolInformation → BatchPoolInfo
• Update enum references to use SCREAMING_SNAKE_CASE values: BatchNodeState.IDLE, BatchTaskState.COMPLETED, BatchTaskExecutionResult.FAILURE
• Replace the nested sub-client calls (client.pool.add, client.compute_node.list, etc.) with the new flat API (client.create_pool, client.list_nodes, etc.)
• Use LRO methods begin_delete_pool, begin_delete_job, begin_terminate_job
• Remove CloudServiceConfiguration support — Azure has retired cloud-service VM pools

operators/batch.py

• Update type hints for batch_job_manager_task, batch_job_preparation_task, etc. to use the new Batch-prefixed names
• Remove os_family and os_version parameters (no longer applicable without CloudServiceConfiguration)
• Update node-state set to use BatchNodeState.START_TASK_FAILED / UNUSABLE / IDLE
• Update on_kill() and clean_up() to use the LRO methods

pyproject.toml

• Lift the azure-batch constraint from >=8.0.0,<15.0.0 to >=15.0.0

Tests

• Rewrite hook and operator tests to use the new BatchClient, model types, and enum values

Relates

Fixes the blocking comment left in pyproject.toml about needing a hook rewrite before lifting the <15.0.0 cap.

closes: #66466

Testing

• I ran pytest providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_batch.py and all tests pass
• I ran pytest providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_batch.py and all tests pass

---

Important

🛠️ Maintainer triage note for @arieleli01212 · by @potiuk · 2026-06-17 14:51 UTC

Helpful heads-up from the maintainers — please address before this PR can be reviewed:

• Static / docs checks failing (CI image checks / Build documentation (--docs-only), CI image checks / Build documentation (--spellcheck-only), CI image checks / Static checks, MyPy providers checks). Run them locally with prek run --all-files (or pre-commit run --all-files) and push the fixes.
• Failing test jobs: Non-DB tests: providers / Non-DB-prov::3.10:-amazon,celer...standard, Special tests / Latest Boto test: providers / All-prov:LatestBoto-Postgres:14:3.10:-amazon,celer...standard, Special tests / Pendulum2 test: providers / All-prov:Pendulum2-Postgres:14:3.10:-amazon,celer...standard, provider distributions tests / Compat 2.11.1:P3.10: … (+3 more). Reproduce and fix locally, then push.
• See the Pull Request quality criteria.

The ball is in your court — you've been assigned to this PR. Fix the above, then mark it Ready for review.

_{Automated triage — may be imperfect; a maintainer takes the next look.}

[boring-cyborg]

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
Here are some useful points:

• Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
• In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example Dag that shows how users should use it.
• Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
• Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
• Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
• Be sure to read the Airflow Coding style.
• Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
  Apache Airflow is a community-driven project and together we are making it better 🚀.
  In case of doubts contact the developers at:
  Mailing List: dev@airflow.apache.org
  Slack: https://s.apache.org/airflow-slack

[potiuk]

@arieleli01212 A few things need addressing before review — see our Pull Request quality criteria.

• ❌ Other failing CI checks. See docs.
• ❌ Pre-commit / static checks. See docs.
• ❌ Provider tests. See docs.

No rush.

---

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

[jscheffl]

Oh, this is a full re-write which will generate breaking changes on users side. At least providers version need a major version bump. Can you please add a notice to changelog.rst for the provider similar like we had it recently in Google?

https://github.com/apache/airflow/blob/main/providers/google/docs/changelog.rst?plain=1#L33

Note: Put the waring and migration details above the most recent version number, do NOT create a new one, this will be made by release manager once the provider release is cut.

[arieleli01212]

Done — added a breaking change notice with the migration details to providers/microsoft/azure/docs/changelog.rst, placed above the 13.3.0 header as requested. It covers the client/credential changes, the removed os_family/os_version pool config, the renamed Batch model classes, and the now-required vm_publisher on the operator.

[jscheffl]

Couple of CI errors.

[github-actions]

uv.lock on main just moved via #69795 ("Add Common Data Quality provider skeleton"), commit 2e511e6 and this PR currently conflicts.

Quickest fix:

git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

[potiuk]

@arieleli01212 A few things need addressing before review — see our Pull Request quality criteria.

• ❌ Merge conflicts with main. See docs.

No rush.

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

[eladkal]

@arieleli01212 can you rebase and fix conflicts?
Also, did you test the changes against real Azure env?

[arieleli01212]

Done, rebased on latest main and the conflict (the removed generated/provider_dependencies.json) is sorted out now.

On the real Azure question — no, I haven't run this against a live Batch account, I don't have one set up. I validated it against the provider's unit tests, which I updated for the 15.x models (BatchTask/BatchTaskState/get_pool etc.). So the coverage is the same shape as before the SDK bump, just not an end-to-end run on real infra. If a live smoke test is a hard requirement before merge I can try to stand up a Batch account, but that'll take me a bit — happy to go that route if you'd prefer.

[eladkal]

but that'll take me a bit — happy to go that route if you'd prefer.

We always prefer to validate against real life scenario.
Azure doesn't have system tests in place thus we prefer to validate such changes.
Validation against real azure account increases the confidence the fix is right - it's not mandatory but preferred.
Let me know how it goes

[eladkal]

@arieleli01212 were you able to set it up?
If not we'll relay on the tests we have set

[nailo2c]

Hi, thanks for the PR, one question :)

[nailo2c]

I wonder why we need to ignore the type here.

According to the BatchClient doc, credential should be either AzureNamedKeyCredential or TokenCredential. If credential is an AzureIdentityCredentialAdapter, as shown on line 112, will it still work?

[eladkal]

cc @arieleli01212 can you address this review comment?

[arieleli01212]

Rebased again on latest main — conflicts are resolved (mostly the azure deps in pyproject/uv.lock plus the changelog).

I wasn't able to get a real Azure Batch account stood up on my end, so let's go with relying on the existing tests as you suggested. I did make sure the unit tests exercise the 15.x code paths I changed (the BatchClient/get_pool/BatchTask bits), so that's the coverage we've got. If anyone on the team does have a Batch account handy and wants to give it a real-world run before merging, that'd obviously be the strongest signal.