From dcdb8dcff2f0e3e9165c30f2adb3b9ce5a68381c Mon Sep 17 00:00:00 2001 From: Alexander Kazarin Date: Tue, 21 Apr 2020 16:55:41 +0300 Subject: [PATCH] add Kerberos support to PrestoHook --- airflow/providers/presto/hooks/presto.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/airflow/providers/presto/hooks/presto.py b/airflow/providers/presto/hooks/presto.py index b316311301556..d877171d2cbc1 100644 --- a/airflow/providers/presto/hooks/presto.py +++ b/airflow/providers/presto/hooks/presto.py @@ -19,6 +19,8 @@ from prestodb.exceptions import DatabaseError from prestodb.transaction import IsolationLevel +from airflow.configuration import conf +from airflow.security import utils from airflow.hooks.dbapi_hook import DbApiHook @@ -44,7 +46,21 @@ class PrestoHook(DbApiHook): def get_conn(self): """Returns a connection object""" db = self.get_connection(self.presto_conn_id) # pylint: disable=no-member - auth = prestodb.auth.BasicAuthentication(db.login, db.password) if db.password else None + + auth = None + if db.extra_dejson.get('kerberos'): + principal = db.extra_dejson.get('kerberos_principal', conf.get('kerberos', 'principal')) + if "_HOST" in principal: + principal = utils.replace_hostname_pattern( + utils.get_components(principal)) + auth = prestodb.auth.KerberosAuthentication( + config=db.extra_dejson.get('kerberos_config', '/etc/krb5.conf'), + service_name=db.extra_dejson.get('kerberos_service', 'presto'), + principal=principal, + ca_bundle=db.extra_dejson.get('kerberos_ca_bundle'), + ) + elif db.password: + auth = prestodb.auth.BasicAuthentication(db.login, db.password) return prestodb.dbapi.connect( host=db.host,