OpenLineage with Airflow on GCP Composer

apache-airflow-providers-openlineage

transport={"type": "http","url": "<host name of the ADOC tenant>","endpoint": "/torch-pipeline/api/v1/lineage","verify": false,"custom_headers": {"accessKey": "<access key>","secretKey": "<secret key>"}} namespace="<openlineage namespace - choose an appropriate name>"

transport={"type": "http","url": "<host name of the ADOC tenant>","endpoint": "/torch-pipeline/api/v1/lineage", "auth": {"type": "tokenproviders.access_key_secret_token_provider.AccessKeySecretKeyTokenProvider"}}

# Copyright 2018-2023 contributors to the OpenLineage project # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import json import logging import os from datetime import datetime, timedelta from typing import Any import requests from airflow.configuration import conf from airflow.hooks.base import BaseHook from airflow.models import Variable from dateutil import parser from openlineage.client.transport.http import TokenProvider log = logging.getLogger(__name__) # Set the environment variable to bypass proxy os.environ["no_proxy"] = "*" CLIENT_ID = "acceldata-app" TOKEN_PROVIDER_PATH = "/admin/api/onboarding/token-exchange?grant_type=api_keys" ACCELDATA_LINEAGE_URL = "acceldata_lineage_url" ACCELDATA_LINEAGE_ENDPOINT = "acceldata_lineage_endpoint" ACCELDATA_ACCESS_KEY = "acceldata_access_key" ACCELDATA_SECRET_KEY = "acceldata_secret_key" ACCELDATA_BEARER_TOKEN = "acceldata_bearer_token" ACCELDATA_EXPIRES_AT = "acceldata_expires_at" class AccessKeySecretKeyTokenProvider(TokenProvider): def __init__(self, config: dict[str, str]) -> None: super().__init__(config) log.info( "Initializing AccessKeySecretKeyTokenProvider. Starting credential resolution." ) self.access_key = config.get("access_key") self.secret_key = config.get("secret_key") if self.is_config_loaded(): log.info("Config is already loaded. Skipping loading from the config.") else: ol_config = self._load_openlineage_config() access_key, secret_key, credential_source, transport = ( self._resolve_credentials(ol_config) ) if not access_key or not secret_key: log.error( "Credential resolution failed. Neither Airflow Connection nor " "OpenLineage config provides valid access_key and secret_key." ) return self._finalize_and_persist( transport=transport, access_key=access_key, secret_key=secret_key, credential_source=credential_source, ) def _load_openlineage_config(self) -> dict[str, str]: config = self.load_config() return config def _resolve_credentials(self, config): transport, access_key_from_config, secret_key_from_config = ( self.extract_transport_and_auth(config) ) # Try Airflow Connection first try: log.info( "Attempting to load credentials from Airflow Connection: acceldata_connection" ) conn = BaseHook.get_connection("acceldata_connection") log.info( "Airflow Connection loaded successfully (host=%s)", conn.host ) if conn.login and conn.password: log.info( "Using credentials from Airflow Connection " "(access_key_present=True, secret_key_present=True)" ) return conn.login, conn.password, "airflow_connection", transport log.warning( "Airflow Connection credentials are incomplete " "(access_key_present=%s, secret_key_present=%s).", bool(conn.login), bool(conn.password), ) except Exception as e: log.info( "Airflow Connection 'acceldata_connection' not available. " "Falling back to OpenLineage config. Details=%s", e, ) # Fallback to OpenLineage config if access_key_from_config and secret_key_from_config: log.info( "Using credentials from OpenLineage config after fallback." ) return access_key_from_config, secret_key_from_config, "openlineage_config", transport return None, None, None, None def _finalize_and_persist( self, *, transport, access_key: str, secret_key: str, credential_source: str, ): self.access_key = access_key self.secret_key = secret_key log.info( "Persisting OpenLineage configuration and resolved credentials " "(credential_source=%s)", credential_source, ) self.persist_config_to_airflow_variables( transport, self.access_key, self.secret_key ) # Checking if the acceldata_lineage_backend_url is already available @staticmethod def is_config_loaded(): from airflow.models import Variable log.info( "Checking OpenLineage configuration state: " "verifying transport (url, endpoint) and credential availability " "from Airflow Variables or Airflow Connection" ) acceldata_lineage_url = Variable.get(ACCELDATA_LINEAGE_URL, None) acceldata_lineage_endpoint = Variable.get(ACCELDATA_LINEAGE_ENDPOINT, None) acceldata_secret_key = Variable.get(ACCELDATA_SECRET_KEY, None) acceldata_access_key = Variable.get(ACCELDATA_ACCESS_KEY, None) if acceldata_lineage_url is not None and acceldata_lineage_endpoint is not None and acceldata_secret_key is not None and acceldata_access_key is not None: return True return False @staticmethod def persist_url(transport): if "url" in transport: log.info("URL found in the transport") url = transport["url"] log.info( "Persisting OpenLineage backend URL to Airflow Variables (%s)", url ) Variable.set(ACCELDATA_LINEAGE_URL, url) else: log.error( "Missing required OpenLineage configuration: transport.url not found" ) @staticmethod def persist_endpoint(transport): if "endpoint" in transport: log.info("endpoint found in the transport") endpoint = transport["endpoint"] log.info( "Persisting OpenLineage endpoint to Airflow Variables (%s)", endpoint ) Variable.set(ACCELDATA_LINEAGE_ENDPOINT, endpoint) else: log.error( "Missing required OpenLineage configuration: transport.endpoint not found" ) @staticmethod def persist_access_key(access_key: str | None): if not access_key: log.info( "access_key not provided for persistence. Skipping Airflow Variable update." ) return log.info("Persisting access_key to Airflow Variables") Variable.set(ACCELDATA_ACCESS_KEY, access_key) @staticmethod def persist_secret_key(secret_key: str | None): if not secret_key: log.info( "secret_key not provided for persistence. Skipping Airflow Variable update." ) return log.info("Persisting secret_key to Airflow Variables") Variable.set(ACCELDATA_SECRET_KEY, secret_key) def load_config(self) -> Dict[str, Any]: log.info( "Attempting to load [openlineage] section from airflow.cfg" ) try: # Accessing the [openlineage] section directly using conf openlineage_config = conf.getsection('openlineage') if not isinstance(openlineage_config, dict): log.error( "Invalid OpenLineage configuration type. " "Expected dict but got %s", type(openlineage_config).__name__, ) return {} log.info( "Successfully loaded OpenLineage configuration: %s", openlineage_config, ) return openlineage_config except Exception as e: log.error( "Failed to read [openlineage] configuration from airflow.cfg. Error=%s", e, ) return {} @staticmethod def get_bearer_token(token): return f"Bearer {token}" @staticmethod def _update_token_to_cache(token: str, expires_in: int): expiration = datetime.now() + timedelta(seconds=expires_in) log.info( "Updating cached bearer token and expiry timestamp in Airflow Variables" ) Variable.set(ACCELDATA_BEARER_TOKEN, token) Variable.set(ACCELDATA_EXPIRES_AT, expiration) log.info("Bearer token cache updated successfully") @staticmethod def _get_token_from_cache(): log.info( "Fetching cached bearer token and expiry timestamp from Airflow Variables" ) cached_token = Variable.get(ACCELDATA_BEARER_TOKEN, None) expires_at = Variable.get(ACCELDATA_EXPIRES_AT, None) return cached_token, expires_at @staticmethod def validate_token(cached_token, expires_at): # Parse the date string into a datetime object if expires_at is None: log.warning( "Bearer token not found in cache. Token will be generated." ) return False else: expires_at_date = parser.parse(expires_at) if cached_token is not None and expires_at_date > datetime.now(): log.info("Cached bearer token is valid") return True else: log.warning("Cached bearer token has expired") return False def _fetch_token_from_admin_central(self) -> tuple[Any, Any] | None: try: headers = { 'Content-Type': 'application/json', # Set the content type to JSON } access_key = Variable.get(ACCELDATA_ACCESS_KEY, None) secret_key = Variable.get(ACCELDATA_SECRET_KEY, None) Variable.get(ACCELDATA_LINEAGE_URL, None) data = { 'clientId': CLIENT_ID, 'secretKey': secret_key, 'accessKey': access_key } token_provider_base_url = Variable.get(ACCELDATA_LINEAGE_URL, None) if token_provider_base_url is not None: token_provider_url = token_provider_base_url + TOKEN_PROVIDER_PATH response = requests.post(token_provider_url, json=data, headers=headers) log.info( "Requesting bearer token from Admin Central " "(url=%s, clientId=%s)", token_provider_url, CLIENT_ID, ) log.info(response) if response.status_code == 200: log.info("Bearer token successfully retrieved from Admin Central") token_data = response.json() token = token_data.get('access_token') expires_in = token_data.get('expires_in') expires_in_sixty_seconds_ago = expires_in - 60 return token, expires_in_sixty_seconds_ago else: log.error( "Failed to fetch bearer token from Admin Central " "(status_code=%s, response_body=%s)", response.status_code, response.text, ) return None, None else: log.warning( "OpenLineage backend URL (%s) is not configured in Airflow Variables. " "Token generation cannot proceed without this value.", ACCELDATA_LINEAGE_URL, ) except Exception as e: log.error( "Exception occurred while fetching bearer token from Admin Central. Error=%s", e, ) def refresh_token(self): log.info("Refreshing bearer token") token, expires_in = self._fetch_token_from_admin_central() if token: self._update_token_to_cache(token, expires_in) else: log.error("Bearer token refresh failed") @staticmethod def extract_transport_and_auth(config): """ Extract transport config along with access_key and secret_key (if present) from OpenLineage configuration. Returns: tuple[dict | None, str | None, str | None] -> (transport, access_key, secret_key) """ if "transport" not in config: log.error("Missing required OpenLineage configuration: transport") return None, None, None try: transport = json.loads(config.get("transport")) except Exception as e: log.error("Failed to parse OpenLineage transport config. Error=%s", e) return None, None, None auth = transport.get("auth", {}) access_key = auth.get("access_key") secret_key = auth.get("secret_key") if not access_key: log.info( "transport.auth.access_key not found in OpenLineage config" ) if not secret_key: log.info( "transport.auth.secret_key not found in OpenLineage config" ) return transport, access_key, secret_key # To persist the openlineage configuration to Airflow variables def persist_config_to_airflow_variables(self, transport, access_key, secret_key): log.info("Persisting OpenLineage configuration to Airflow Variables") self.persist_url(transport) self.persist_endpoint(transport) self.persist_access_key(access_key) self.persist_secret_key(secret_key) def get_bearer(self): log.info("Checking bearer token in the cache") cached_token, expires_at = self._get_token_from_cache() is_valid_token = self.validate_token(cached_token, expires_at) if is_valid_token: return self.get_bearer_token(cached_token) else: log.info( "Either the token doesn't exist or has expired. Refreshing token") self.refresh_token() cached_token, expires_at = self._get_token_from_cache() return self.get_bearer_token(cached_token)