Skip to content

Conversation

Ishankoradia
Copy link
Contributor

@Ishankoradia Ishankoradia commented Jul 13, 2025

Summary by CodeRabbit

  • New Features
    • Added support for cancelling running Airbyte jobs directly from the interface.
    • Updated integration to improve handling of Airbyte job cancellations.

Copy link

coderabbitai bot commented Jul 13, 2025

Walkthrough

A new Prefect flow, run_airbyte_cancel_job, was introduced to handle the cancellation of Airbyte jobs by extracting relevant information from the payload and invoking the cancellation logic. This flow is now conditionally called within deployment_schedule_flow_v4 when the task slug is "airbyte-cancel". Necessary imports were updated accordingly. Additionally, the prefect-airbyte dependency version was updated from v0.5 to v0.6 in pyproject.toml.

Changes

File(s) Change Summary
proxy/prefect_flows.py Added run_airbyte_cancel_job flow, integrated it into deployment_schedule_flow_v4, and updated imports to include cancel_job.
pyproject.toml Updated prefect-airbyte Git revision tag from "v0.5" to "v0.6" in [tool.uv.sources].

Sequence Diagram(s)

sequenceDiagram
    participant Client
    participant PrefectFlow
    participant AirbyteServerBlock
    participant AirbyteConnectionBlock
    participant AirbyteAPI

    Client->>PrefectFlow: Trigger deployment_schedule_flow_v4(payload)
    PrefectFlow->>PrefectFlow: Check payload["task_slug"]
    alt task_slug == "airbyte-cancel"
        PrefectFlow->>AirbyteServerBlock: Load Airbyte server block
        PrefectFlow->>AirbyteConnectionBlock: Load Airbyte connection block
        PrefectFlow->>AirbyteAPI: Call cancel_job(connection, job_id)
        AirbyteAPI-->>PrefectFlow: Return cancellation result
        PrefectFlow->>Client: Log and return result
    else other task_slug
        PrefectFlow->>...: (Other task flows)
    end
Loading

Possibly related PRs


📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 1aa7a91 and d4d7d59.

⛔ Files ignored due to path filters (1)
  • uv.lock is excluded by !**/*.lock
📒 Files selected for processing (1)
  • pyproject.toml (1 hunks)
✅ Files skipped from review due to trivial changes (1)
  • pyproject.toml

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Explain this complex logic.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai explain this code block.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and explain its main purpose.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 3eea68d and bcde992.

📒 Files selected for processing (1)
  • proxy/prefect_flows.py (3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
proxy/prefect_flows.py (2)
tests/test_service.py (1)
  • load (75-76)
proxy/helpers.py (2)
  • info (32-45)
  • error (47-60)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
  • GitHub Check: check (3.10)
🔇 Additional comments (2)
proxy/prefect_flows.py (2)

18-18: LGTM! Import addition is correct.

The import of cancel_connection_job is properly placed and necessary for the new cancellation functionality.


342-343: LGTM! Integration follows established patterns.

The conditional branch for "airbyte-cancel" is correctly implemented and consistent with the existing code structure for other Airbyte operations.

Comment on lines 96 to 121
# task config for cancelling a running airbyte job
# {
# type AIRBYTECONNECTION,
# slug: "airbyte-cancel"
# airbyte_server_block: str
# connection_id: str
# job_id: str
# }
@flow
def run_airbyte_cancel_job(payload: dict):
"""cancel a running airbyte job"""
try:
airbyte_server_block = payload["airbyte_server_block"]
serverblock = AirbyteServer.load(airbyte_server_block)
connection_block = AirbyteConnection(
airbyte_server=serverblock,
connection_id=payload["connection_id"],
)
result = cancel_connection_job(connection_block, payload["job_id"])
logger.info("airbyte job cancel result=")
logger.info(result)
return result
except Exception as error: # skipcq PYL-W0703
logger.error(str(error)) # "Job cancel failed."
raise

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Add input validation for required payload fields.

The flow implementation follows the established pattern but lacks input validation for required fields. Consider adding validation to prevent KeyError exceptions if required fields are missing from the payload.

@flow
def run_airbyte_cancel_job(payload: dict):
    """cancel a running airbyte job"""
    try:
+        # Validate required fields
+        required_fields = ["airbyte_server_block", "connection_id", "job_id"]
+        for field in required_fields:
+            if field not in payload:
+                raise ValueError(f"Missing required field: {field}")
+        
        airbyte_server_block = payload["airbyte_server_block"]
        serverblock = AirbyteServer.load(airbyte_server_block)
        connection_block = AirbyteConnection(
            airbyte_server=serverblock,
            connection_id=payload["connection_id"],
        )
        result = cancel_connection_job(connection_block, payload["job_id"])
        logger.info("airbyte job cancel result=")
        logger.info(result)
        return result
    except Exception as error:  # skipcq PYL-W0703
        logger.error(str(error))  # "Job cancel failed."
        raise
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# task config for cancelling a running airbyte job
# {
# type AIRBYTECONNECTION,
# slug: "airbyte-cancel"
# airbyte_server_block: str
# connection_id: str
# job_id: str
# }
@flow
def run_airbyte_cancel_job(payload: dict):
"""cancel a running airbyte job"""
try:
airbyte_server_block = payload["airbyte_server_block"]
serverblock = AirbyteServer.load(airbyte_server_block)
connection_block = AirbyteConnection(
airbyte_server=serverblock,
connection_id=payload["connection_id"],
)
result = cancel_connection_job(connection_block, payload["job_id"])
logger.info("airbyte job cancel result=")
logger.info(result)
return result
except Exception as error: # skipcq PYL-W0703
logger.error(str(error)) # "Job cancel failed."
raise
@flow
def run_airbyte_cancel_job(payload: dict):
"""cancel a running airbyte job"""
try:
# Validate required fields
required_fields = ["airbyte_server_block", "connection_id", "job_id"]
for field in required_fields:
if field not in payload:
raise ValueError(f"Missing required field: {field}")
airbyte_server_block = payload["airbyte_server_block"]
serverblock = AirbyteServer.load(airbyte_server_block)
connection_block = AirbyteConnection(
airbyte_server=serverblock,
connection_id=payload["connection_id"],
)
result = cancel_connection_job(connection_block, payload["job_id"])
logger.info("airbyte job cancel result=")
logger.info(result)
return result
except Exception as error: # skipcq PYL-W0703
logger.error(str(error)) # "Job cancel failed."
raise
🤖 Prompt for AI Agents
In proxy/prefect_flows.py around lines 96 to 121, the run_airbyte_cancel_job
flow lacks validation for required keys in the payload dictionary, which can
cause KeyError exceptions. Add explicit checks at the start of the function to
verify that "airbyte_server_block", "connection_id", and "job_id" keys exist in
the payload. If any are missing, raise a clear exception or handle the error
gracefully before proceeding with the rest of the flow logic.

@codecov-commenter
Copy link

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 76.90%. Comparing base (3eea68d) to head (d4d7d59).

Additional details and impacted files
@@           Coverage Diff           @@
##             main     #212   +/-   ##
=======================================
  Coverage   76.90%   76.90%           
=======================================
  Files           5        5           
  Lines        1429     1429           
=======================================
  Hits         1099     1099           
  Misses        330      330           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants