Support Airflow operator extra link about astronomer-cosmos HOT 1 CLOSED

🤖

Hey @joppevos! Great to see you contributing to the Cosmos repository again. Your suggestions are always insightful.

Thank you for your detailed request and suggestions, joppevos. I agree that adding support for Airflow operator extra links in Cosmos could significantly improve task discoverability and debugging.

To implement this, we could extend the BaseOperator class and override the get_extra_links method. This method should return a dictionary where the keys are the names of the links and the values are the URLs for those links.

Here's a basic example of how you could implement this:

from airflow.models import BaseOperator

class DbtOperatorWithExtraLinks(BaseOperator):
    def get_extra_links(self, dttm, link_name):
        if link_name == 'dbt_run_results':
            # Here you would generate the URL for the dbt run_results.json file
            # This could involve reading the invocation_id from the file and using it in the URL
            return 'http://path/to/dbt/run_results.json'
        elif link_name == 'dbt_python_api_results':
            # Here you would generate the URL for the dbt Python API results
            # This could involve calling the dbt Python API and getting the results
            return 'http://path/to/dbt/python/api/results'
        else:
            return super().get_extra_links(dttm, link_name)

Then, in your calculate_operator_class function, you would return the path to this new operator class instead of the existing one:

def calculate_operator_class(
    execution_mode: ExecutionMode,
    dbt_class: str,
) -> str:
    """
    Given an execution mode and dbt class, return the operator class path to use.

    :param execution_mode: Cosmos execution mode (e.g. local, virtualenv, docker, kubernetes)
    :param dbt_class: The dbt command being used (e.g. DbtSnapshot, DbtRun, DbtSeed)
    :returns: path string to the correspondent Cosmos Airflow operator
    (e.g. cosmos.operators.localDbtSnapshotLocalOperator)
    """
    return f"cosmos.operators.{execution_mode.value}.DbtOperatorWithExtraLinks"

Please note that this is a simplified example and you would need to adapt it to your specific use case, including handling the generation of the URLs for the extra links.

Regarding your suggestions about leveraging the invocation_id written by dbt in run_results.json or using dbt's Python API, both approaches seem feasible. The choice between them would depend on the specific requirements and constraints of your use case.

I hope this helps! If you have any further questions or need more clarification, please let me know.

Sources

• cosmos/airflow/graph.py
• cosmos/operators/base.py

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

from astronomer-cosmos.