Fix #3943 - Airflow auth providers fix (#3947)

* Fix airflow providers loading

* Format

* Format
This commit is contained in:
Pere Miquel Brull 2022-04-08 15:59:01 +02:00 committed by GitHub
parent 06a3e4c989
commit 7f46e46bc3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 14 additions and 18 deletions

View File

@ -20,6 +20,7 @@ from pydantic import BaseModel
from airflow_provider_openmetadata.lineage.config.commons import LINEAGE from airflow_provider_openmetadata.lineage.config.commons import LINEAGE
from airflow_provider_openmetadata.lineage.config.providers import ( from airflow_provider_openmetadata.lineage.config.providers import (
InvalidAirflowProviderException,
provider_config_registry, provider_config_registry,
) )
from metadata.generated.schema.metadataIngestion.workflow import ( 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: if auth_provider_type == AuthProvider.no_auth.value:
security_config = None security_config = None
else: 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() security_config = load_security_config_fn()
return AirflowLineageConfig( return AirflowLineageConfig(

View File

@ -13,7 +13,6 @@
OpenMetadata Airflow Lineage Backend security providers config OpenMetadata Airflow Lineage Backend security providers config
""" """
from collections import namedtuple
from airflow.configuration import conf from airflow.configuration import conf
@ -24,26 +23,16 @@ from metadata.generated.schema.metadataIngestion.workflow import (
GoogleSSOConfig, GoogleSSOConfig,
OktaSSOConfig, 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) @provider_config_registry.add(AuthProvider.google.value)