Emit asset events in async execution mode

[pankajastro]

closes: #2141

This PR implements asset event emission for BigQuery async execution mode, addressing issue #2141. The changes enable Airflow to track dataset lineage when using async operators by registering output datasets after task execution.

Key changes:

• Added asset/dataset event registration in execute() and execute_complete() methods
• Implemented helper methods _register_event() to generate BigQuery asset URIs and register them

[netlify]

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

Name	Link
🔨 Latest commit	76a8253
🔍 Latest deploy log	https://app.netlify.com/projects/sunny-pastelito-5ecb04/deploys/693926c10eca2100083dbd22

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.86%. Comparing base (37e796a) to head (572774b).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2184   +/-   ##
=======================================
  Coverage   97.85%   97.86%           
=======================================
  Files          94       94           
  Lines        6066     6083   +17     
=======================================
+ Hits         5936     5953   +17     
  Misses        130      130           

☔ 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.

[Copilot]

Pull request overview

This PR implements asset event emission for BigQuery async execution mode, addressing issue #2141. The changes enable Airflow to track dataset lineage when using async operators by registering output datasets after task execution.

Key changes:

• Added asset/dataset event registration in execute() and execute_complete() methods
• Implemented helper methods _get_asset_uri() and _register_event() to generate BigQuery asset URIs and register them
• Added Airflow 2/3 compatibility for Asset/Dataset imports

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
cosmos/operators/_asynchronous/bigquery.py	Adds asset event emission with new _get_asset_uri() and _register_event() methods, including Airflow 2/3 compatible Asset imports
tests/operators/_asynchronous/test_bigquery.py	Disables dataset emission in existing tests to prevent failures due to uninitialized gcp_project/dataset attributes

Comments suppressed due to low confidence (4)

cosmos/operators/_asynchronous/bigquery.py:265

• The register_dataset() call passes an empty list for inlets (input datasets). While this may be intentional because OpenLineage events are not available in async execution mode (when enable_setup_async_task is True), adding a comment would clarify this design decision:

def _register_event(self, context: Context) -> None:
    output = [Asset(uri=self._get_asset_uri())]
    # Inlets (input datasets) are not available in async execution mode
    # because OpenLineage events are not calculated when build_and_run_cmd is not called
    self.register_dataset([], output, context)

cosmos/operators/_asynchronous/bigquery.py:213

• The _register_event() method is called in execute() when emit_datasets is True, but when settings.enable_setup_async_task is False, the _store_template_fields() method returns early (line 216-218) without setting self.gcp_project and self.dataset. This means _get_asset_uri() will use the empty string values initialized in __init__ (lines 129-130), resulting in invalid asset URIs like bigquery:///unknown_model or bigquery://..unknown_model.

Consider checking if gcp_project and dataset are set before calling _register_event(), or skip asset registration when enable_setup_async_task is False:

if self.emit_datasets and settings.enable_setup_async_task:
    self._register_event(context)

        if not settings.enable_setup_async_task:
            self.log.info("SQL cannot be made available, skipping registration of compiled_sql template field")

cosmos/operators/_asynchronous/bigquery.py:12

• The logger is imported at the top of the file but is never used. Since logging is being done via self.log (inherited from the base class), consider removing the unused logger import and declaration if it's not needed, or use it for module-level logging if applicable.

from cosmos.operators.base import _sanitize_xcom_key

try:

cosmos/operators/_asynchronous/bigquery.py:265

• The new helper methods _get_asset_uri() and _register_event() lack docstrings. Adding documentation would improve code maintainability and help other developers understand their purpose and behavior.

Consider adding docstrings:

def _get_asset_uri(self) -> str:
    """
    Generate the BigQuery asset URI for dataset registration.
    
    Returns a URI in the format:
    - Airflow 3: bigquery://{project}/{dataset}/{unique_id}
    - Airflow 2: bigquery://{project}.{dataset}.{unique_id}
    """
    ...

def _register_event(self, context: Context) -> None:
    """
    Register the asset/dataset event for this task execution.
    
    Args:
        context: The Airflow task execution context.
    """
    ...

        if AIRFLOW_VERSION.major >= 3:
            return f"bigquery://{self.gcp_project}/{self.dataset}/{unique_id}"
        else:
            return f"bigquery://{self.gcp_project}.{self.dataset}.{unique_id}"

    def _register_event(self, context: Context) -> None:
        output = [Asset(uri=self._get_asset_uri())]
        self.register_dataset([], output, context)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.