From 7f46e46bc3de21cbe1f5059f91e4148d8bf3354a Mon Sep 17 00:00:00 2001 From: Pere Miquel Brull Date: Fri, 8 Apr 2022 15:59:01 +0200 Subject: [PATCH] Fix #3943 - Airflow auth providers fix (#3947) * Fix airflow providers loading * Format * Format --- .../lineage/config/loader.py | 9 +++++++- .../lineage/config/providers.py | 23 +++++-------------- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/ingestion/src/airflow_provider_openmetadata/lineage/config/loader.py b/ingestion/src/airflow_provider_openmetadata/lineage/config/loader.py index 9dcf1b9d712..f14aef09414 100644 --- a/ingestion/src/airflow_provider_openmetadata/lineage/config/loader.py +++ b/ingestion/src/airflow_provider_openmetadata/lineage/config/loader.py @@ -20,6 +20,7 @@ from pydantic import BaseModel from airflow_provider_openmetadata.lineage.config.commons import LINEAGE from airflow_provider_openmetadata.lineage.config.providers import ( + InvalidAirflowProviderException, provider_config_registry, ) from metadata.generated.schema.metadataIngestion.workflow import ( @@ -45,7 +46,13 @@ def parse_airflow_config(airflow_service_name: str) -> AirflowLineageConfig: if auth_provider_type == AuthProvider.no_auth.value: security_config = None else: - load_security_config_fn = provider_config_registry.registry(auth_provider_type) + load_security_config_fn = provider_config_registry.registry.get( + auth_provider_type + ) + if not load_security_config_fn: + raise InvalidAirflowProviderException( + f"Cannot find {auth_provider_type} in airflow providers registry." + ) security_config = load_security_config_fn() return AirflowLineageConfig( diff --git a/ingestion/src/airflow_provider_openmetadata/lineage/config/providers.py b/ingestion/src/airflow_provider_openmetadata/lineage/config/providers.py index f7289565ac9..153d03c31d3 100644 --- a/ingestion/src/airflow_provider_openmetadata/lineage/config/providers.py +++ b/ingestion/src/airflow_provider_openmetadata/lineage/config/providers.py @@ -13,7 +13,6 @@ OpenMetadata Airflow Lineage Backend security providers config """ -from collections import namedtuple from airflow.configuration import conf @@ -24,26 +23,16 @@ from metadata.generated.schema.metadataIngestion.workflow import ( GoogleSSOConfig, OktaSSOConfig, ) +from metadata.utils.dispatch import enum_register + +provider_config_registry = enum_register() -def register_provider_config(): +class InvalidAirflowProviderException(Exception): """ - Helps us register custom functions to parse provider config + Raised when we cannot find the provider + in Airflow config """ - registry = dict() - - def add(provider: str): - def inner(fn): - registry[provider] = fn - return fn - - return inner - - Register = namedtuple("Register", ["add", "registry"]) - return Register(add, registry) - - -provider_config_registry = register_provider_config() @provider_config_registry.add(AuthProvider.google.value)