zozo123 opened a new pull request, #76157:
URL: https://github.com/apache/airflow/pull/76157
Related: #74194 (part 2: the resumable standalone operator with deferrable
polling)
This builds on #74419 and #74517, which add `DatabricksGenieHook` (part 1).
It is stacked on #74517 because part 1 has not merged yet. The first four
commits on this branch are the part 1 commits plus one fixup that removes a
`common.ai` extra the rebase duplicated. **Only the last commit (`Add a
resumable Databricks Genie operator`) is new in this PR.** I will rebase once
part 1 merges. Please treat it as a draft until then.
### What it adds
- `DatabricksGenieOperator` (`operators/genie.py`): asks a Genie space a
question, either as a new conversation or as a follow-up in an existing
`conversation_id`. It returns the same bounded result that the hook produces.
It can also attach to an existing message through `conversation_id` +
`message_id`.
- `DatabricksGenieMessageTrigger` (`triggers/genie.py`): polls the message
status from the triggerer when `deferrable=True`. It only reads the status. The
worker fetches and bounds the result when it resumes.
- A how-to page (`docs/operators/genie_operator.rst`), a system test Dag
(`example_databricks_genie.py`), the `provider.yaml` / `get_provider_info.py`
registrations, and unit tests for both classes.
### Design
- **Resuming.** The operator uses `ResumableJobMixin`. On Airflow 3.3+ with
`durable=True`, it stores `{space_id, conversation_id, message_id,
fingerprint}` in the task state store. A retry does one of three things:
- reconnects to a message that is still running;
- returns the answer of a message that has completed;
- asks again if the message failed, was cancelled or expired.
- **Question fingerprint.** The fingerprint is a hash of the space, the
conversation and the question. If the templated question changes between tries,
the stored message is ignored, so an old answer is never returned for a new
question.
- **Unconfirmed submissions.** Sending the question is not idempotent. The
operator writes a short-lived marker before the POST and removes it once the
message IDs are stored.
- A definite 4xx rejection removes the marker, and the task can retry
normally.
- A timeout, connection error or 5xx fails the task with
`AirflowFailException` and keeps the marker, because Genie may have accepted
the question.
- The next try refuses to ask again. The user can attach to the message
with `conversation_id`/`message_id`, or opt in with `resubmit_unconfirmed=True`.
- **Deferral.** Before it defers, the worker reads the message once through
the hook's HTTPS-enforcing request path, so a message that has already finished
never defers. The trigger checks the connection scheme the same way and allows
plain HTTP only to loopback addresses. `timeout` bounds both the sync wait and
the deferral.
- **Before Airflow 3.3.** There is no task state store, so the operator
behaves like a normal non-resumable task: each try asks again.
- **Cancellation.** There is no `on_kill`, because the Genie conversation
API has no endpoint to cancel a message.
### Dependency on part 1 and limitations
- This uses hook internals from #74517: `_request(..., retry_read=False)`,
`_summarize_result`, `_truncate_result` and `_PENDING_STATUSES`. If part 1
exposes public primitives (for example `start_message` / `get_message` /
`summarize_message`), this operator can switch to them and get smaller. I
implemented against the interface as it stands in #74517 at `dba5ef1`.
- `breeze` could not run locally because no Docker was available, so the
`update-providers-build-files`, `check-provider-yaml-valid` and
`check-template-fields-valid` hooks were skipped. I edited `provider.yaml` and
`get_provider_info.py` by hand and am relying on CI for these checks.
### Testing
All of these passed locally on Python 3.13 with Airflow 3.3+, so the durable
tests ran:
- `uv run --project providers/databricks pytest
providers/databricks/tests/unit/databricks/{operators,triggers,hooks}/test_genie.py`:
104 passed.
- `uv run --project providers/databricks pytest
providers/databricks/tests/unit -n 8`: 1294 passed, 13 skipped.
- `prek run --stage pre-commit --files <changed files>`: everything passed
except the three Docker-only hooks above.
- `ruff format` / `ruff check` (0.17.0): clean.
- `mypy` on the new files: clean.
- `scripts/in_container/run_template_fields_check.py` on the operator:
passed.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
Generated-by: Grok Bot following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
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]
