Astro Databricks

Affordable Databricks Workflows in Apache Airflow

Astro Databricks is an Apache Airflow provider created by Astronomer for an optimal Databricks experience. With the DatabricksTaskGroup, Astro Datatricks allows you to run from Databricks workflows without the need of running Jobs individually, which can result in 75% cost reduction.

Prerequisites

• Apache Airflow >= 2.2.4
• Python >= 2.7
• Databricks account
• Previously created Databricks Notebooks

Install

pip install astro-provider-databricks

Quickstart

1. Use pre-existing or create two simple Databricks Notebooks. Their identifiers will be used in step (5). The original example DAG uses:

   • Shared/Notebook_1
   • Shared/Notebook_2

2. Generate a Databricks Personal Token. This will be used in step (6).

3. Ensure that your Airflow environment is set up correctly by running the following commands:

   export AIRFLOW_HOME=`pwd`
   
   airflow db init
   

4. Create using your preferred way a Databricks Airflow connection (so Airflow can access Databricks using your credentials). This can be done by running the following command, replacing the login and password (with your access token):

airflow connections add 'databricks_conn' \
    --conn-json '{
        "conn_type": "databricks",
        "login": "some.email@yourcompany.com",
        "host": "https://dbc-c9390870-65ef.cloud.databricks.com/",
        "password": "personal-access-token"
    }'

5. Copy the following workflow into a file named example_databricks_workflow.py and add it to the dags directory of your Airflow project:

   https://github.com/astronomer/astro-provider-databricks/blob/45897543a5e34d446c84b3fbc4f6f7a3ed16cdf7/example_dags/example_databricks_workflow.py#L48-L101

   Alternatively, you can download example_databricks_workflow.py

    curl -O https://raw.githubusercontent.com/astronomer/astro-provider-databricks/main/example_dags/example_databricks_workflow.py
   

6. Run the example DAG:

   airflow dags test example_databricks_workflow `date -Iseconds`
   

This will create a Databricks Workflow with two Notebook jobs.

Available features

• DatabricksWorkflowTaskGroup: Airflow task group that allows users to create a Databricks Workflow.
• DatabricksNotebookOperator: Airflow operator which abstracts a pre-existing Databricks Notebook. Can be used independently to run the Notebook, or within a Databricks Workflow Task Group.
• AstroDatabricksPlugin: An Airflow plugin which is installed by the default. It allows users, by using the UI, to view a Databricks job and retry running it in case of failure.

Documentation

The documentation is a work in progress--we aim to follow the Diátaxis system:

• Reference Guide

Changelog

Astro Databricks follows semantic versioning for releases. Read changelog to understand more about the changes introduced to each version.

Contribution guidelines

All contributions, bug reports, bug fixes, documentation improvements, enhancements, and ideas are welcome.

Read the Contribution Guidelines for a detailed overview on how to contribute.

Contributors and maintainers should abide by the Contributor Code of Conduct.

License

Apache Licence 2.0