Skip to content

Commit 255551a

Browse files
authored
Add custom user agent string when using LDAP connection to tag dbt adapter connection. (#69)
Currently this user agent string is <b>dbt/cloudera-impala-v[ADAPTER_VERSION]</b> Testplan (WIP): a) Use the PR: cloudera/impyla#498 to first update impyla to enable this feature b) In a sample dbt project, use dbt debug c) Use Atlas to filter the above agent string and match the queries that were executed.
1 parent 5708c98 commit 255551a

File tree

2 files changed

+20
-14
lines changed

2 files changed

+20
-14
lines changed

dbt/adapters/impala/connections.py

Lines changed: 19 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
from dbt.events.functions import fire_event
3030
from dbt.events.types import ConnectionUsed, SQLQuery, SQLQueryStatus
3131

32-
from dbt.logger import GLOBAL_LOGGER as LOGGER
32+
from dbt.events import AdapterLogger
3333

3434
import impala.dbapi
3535
from impala.error import DatabaseError
@@ -41,10 +41,13 @@
4141

4242
import json
4343

44+
from dbt.adapters.impala.__version__ import version as ADAPTER_VERSION
45+
4446
DEFAULT_IMPALA_HOST = "localhost"
4547
DEFAULT_IMPALA_PORT = 21050
4648
DEFAULT_MAX_RETRIES = 3
4749

50+
logger = AdapterLogger("Impala")
4851

4952
@dataclass
5053
class ImpalaCredentials(Credentials):
@@ -119,22 +122,22 @@ def exception_handler(self, sql: str):
119122
try:
120123
yield
121124
except HttpError as httpError:
122-
LOGGER.debug("Authorization error: {}".format(httpError))
125+
logger.debug("Authorization error: {}".format(httpError))
123126
raise dbt.exceptions.RuntimeException ("HTTP Authorization error: " + str(httpError) + ", please check your credentials")
124127
except HiveServer2Error as servError:
125-
LOGGER.debug("Server connection error: {}".format(servError))
128+
logger.debug("Server connection error: {}".format(servError))
126129
raise dbt.exceptions.RuntimeException ("Unable to establish connection to Impala server: " + str(servError))
127130
except DatabaseError as dbError:
128-
LOGGER.debug("Database connection error: {}".format(str(dbError)))
131+
logger.debug("Database connection error: {}".format(str(dbError)))
129132
raise dbt.exceptions.DatabaseException("Database Connection error: " + str(dbError))
130133
except Exception as exc:
131-
LOGGER.debug("Error running SQL: {}".format(sql))
134+
logger.debug("Error running SQL: {}".format(sql))
132135
raise dbt.exceptions.RuntimeException(str(exc))
133136

134137
@classmethod
135138
def open(cls, connection):
136139
if connection.state == ConnectionState.OPEN:
137-
LOGGER.debug("Connection is already open, skipping open.")
140+
logger.debug("Connection is already open, skipping open.")
138141
return connection
139142

140143
credentials = connection.credentials
@@ -148,6 +151,8 @@ def open(cls, connection):
148151
if (
149152
credentials.auth_type == "LDAP" or credentials.auth_type == "ldap"
150153
): # ldap connection
154+
custom_user_agent = 'dbt/cloudera-impala-v' + ADAPTER_VERSION
155+
logger.debug("Using user agent: {}".format(custom_user_agent))
151156
handle = impala.dbapi.connect(
152157
host=credentials.host,
153158
port=credentials.port,
@@ -158,6 +163,7 @@ def open(cls, connection):
158163
use_ssl=credentials.use_ssl,
159164
http_path=credentials.http_path,
160165
retries=credentials.retries,
166+
user_agent=custom_user_agent
161167
)
162168
auth_type = "ldap"
163169
elif (
@@ -186,7 +192,7 @@ def open(cls, connection):
186192
connection.state = ConnectionState.OPEN
187193
connection.handle = handle
188194
except Exception as ex:
189-
LOGGER.debug("Connection error {}".format(ex))
195+
logger.debug("Connection error {}".format(ex))
190196
connection_ex = ex
191197
connection.state = ConnectionState.FAIL
192198
connection.handle = None
@@ -232,7 +238,7 @@ def close(cls, connection):
232238

233239
return connection
234240
except Exception as err:
235-
LOGGER.debug(f"Error closing connection {err}")
241+
logger.debug(f"Error closing connection {err}")
236242

237243
@classmethod
238244
def get_response(cls, cursor):
@@ -243,16 +249,16 @@ def cancel(self, connection):
243249
connection.handle.close()
244250

245251
def add_begin_query(self, *args, **kwargs):
246-
LOGGER.debug("NotImplemented: add_begin_query")
252+
logger.debug("NotImplemented: add_begin_query")
247253

248254
def add_commit_query(self, *args, **kwargs):
249-
LOGGER.debug("NotImplemented: add_commit_query")
255+
logger.debug("NotImplemented: add_commit_query")
250256

251257
def commit(self, *args, **kwargs):
252-
LOGGER.debug("NotImplemented: commit")
258+
logger.debug("NotImplemented: commit")
253259

254260
def rollback(self, *args, **kwargs):
255-
LOGGER.debug("NotImplemented: rollback")
261+
logger.debug("NotImplemented: rollback")
256262

257263
def add_query(
258264
self,
@@ -272,7 +278,7 @@ def add_query(
272278
additional_info = json.loads(self.query_header.comment.query_comment.strip())
273279
except Exception as ex: # silently ignore error for parsing
274280
additional_info = {}
275-
LOGGER.debug(f"Unable to get query header {ex}")
281+
logger.debug(f"Unable to get query header {ex}")
276282

277283
with self.exception_handler(sql):
278284
if abridge_sql_log:

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ def _get_dbt_core_version():
6565
include_package_data=True,
6666
install_requires=[
6767
'dbt-core~={}'.format(dbt_core_version),
68-
"impyla>=0.18a5",
68+
"impyla==0.18a7",
6969
"python-decouple>=3.6"
7070
],
7171
classifiers=[

0 commit comments

Comments
 (0)