| 385 | |
| 386 | # Represents a connection to Impala using the Beeswax API. |
| 387 | class BeeswaxConnection(ImpalaConnection): |
| 388 | |
| 389 | # This is based on ClientRequestState::BeeswaxQueryState(). |
| 390 | __QUERY_STATE_TO_EXEC_STATE = { |
| 391 | QueryState.CREATED: INITIALIZED, |
| 392 | QueryState.COMPILED: PENDING, |
| 393 | QueryState.RUNNING: RUNNING, |
| 394 | QueryState.FINISHED: FINISHED, |
| 395 | QueryState.EXCEPTION: ERROR, |
| 396 | # These are not official ExecState, but added to complete mapping. |
| 397 | QueryState.INITIALIZED: 'UNIMPLEMENTED_INITIALIZED', |
| 398 | } |
| 399 | |
| 400 | def __init__(self, host_port, use_kerberos=False, user=None, password=None, |
| 401 | use_ssl=False): |
| 402 | self.__beeswax_client = ImpalaBeeswaxClient(host_port, use_kerberos, user=user, |
| 403 | password=password, use_ssl=use_ssl) |
| 404 | self.__host_port = host_port |
| 405 | self.QUERY_STATES = self.__beeswax_client.query_states |
| 406 | |
| 407 | def get_test_protocol(self): |
| 408 | return BEESWAX |
| 409 | |
| 410 | def get_host_port(self): |
| 411 | return self.__host_port |
| 412 | |
| 413 | def set_configuration_option(self, name, value, is_log_sql=True): |
| 414 | # Only set the option if it's not already set to the same value. |
| 415 | name = name.lower() |
| 416 | value = str(value) |
| 417 | if self.__beeswax_client.get_query_option(name) != value: |
| 418 | self.__beeswax_client.set_query_option(name, value) |
| 419 | if is_log_sql: |
| 420 | self.log_client("\n\nset {0}={1};\n".format(name, value)) |
| 421 | return True |
| 422 | return False |
| 423 | |
| 424 | def clear_configuration(self): |
| 425 | self.__beeswax_client.clear_query_options() |
| 426 | # A hook in conftest sets tests.common.current_node. |
| 427 | if hasattr(tests.common, "current_node"): |
| 428 | self.set_configuration_option("client_identifier", tests.common.current_node) |
| 429 | |
| 430 | def connect(self): |
| 431 | try: |
| 432 | self.__beeswax_client.connect() |
| 433 | self.log_client("connected to %s with beeswax" % self.__host_port) |
| 434 | except Exception as e: |
| 435 | self.log_client("failed connecting to %s with beeswax" % self.__host_port) |
| 436 | raise e |
| 437 | |
| 438 | # TODO: rename to close_connection |
| 439 | def close(self): |
| 440 | self.log_client("closing beeswax connection to: %s" % self.__host_port) |
| 441 | self.__beeswax_client.close_connection() |
| 442 | |
| 443 | def close_query(self, operation_handle, fetch_profile_after_close=False): |
| 444 | self.log_handle(operation_handle, 'closing query for operation') |
no outgoing calls
no test coverage detected