airflow_dag_hashicorp_vault_secret_retrieval_with_vault_hook.py

1import os
2from datetime import datetime
3
4from airflow import DAG
5from airflow.decorators import task
6from airflow.providers.hashicorp.hooks.vault import VaultHook
7
8# Environment variables for configuration
9ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
10DAG_ID = "example_hashicorp_vault"
11
12with DAG(
13    dag_id=DAG_ID,
14    start_date=datetime(2021, 1, 1),
15    schedule=None,
16    catchup=False,
17    tags=["example"],
18) as dag:
19
20    @task(task_id="get_secret_from_vault")
21    def get_secret_from_vault():
22        """
23        Example task that uses VaultHook to fetch a secret.
24        """
25        # Note: 'vault_conn_id' refers to an Airflow Connection of type 'vault'
26        hook = VaultHook(vault_conn_id="vault_default")
27        
28        # Replace 'secret/data/my-secret' with your actual Vault path
29        secret_data = hook.get_secret(path="secret/data/my-secret")
30        
31        if secret_data:
32            print(f"Successfully retrieved secret keys: {list(secret_data.keys())}")
33        else:
34            print("No secret found or access denied.")
35
36    get_secret_from_vault_task = get_secret_from_vault()
37
38    from tests.system.utils import get_test_run  # noqa: E402
39
40    # Needed to run the example as a system test
41    test_run = get_test_run(dag)