| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| # Superclass of all HS2 tests containing commonly used functions. |
| |
| from getpass import getuser |
| from TCLIService import TCLIService |
| from ImpalaService import ImpalaHiveServer2Service |
| from thrift.transport.TSocket import TSocket |
| from thrift.transport.TTransport import TBufferedTransport |
| from thrift.protocol import TBinaryProtocol |
| from tests.common.impala_test_suite import ImpalaTestSuite, IMPALAD_HS2_HOST_PORT |
| from time import sleep, time |
| |
| |
| def add_session_helper(self, protocol_version, conf_overlay, close_session, fn): |
| """Helper function used in the various needs_session decorators before to set up |
| a session, call fn(), then optionally tear down the session.""" |
| open_session_req = TCLIService.TOpenSessionReq() |
| open_session_req.username = getuser() |
| open_session_req.configuration = dict() |
| if conf_overlay is not None: |
| open_session_req.configuration = conf_overlay |
| open_session_req.client_protocol = protocol_version |
| resp = self.hs2_client.OpenSession(open_session_req) |
| HS2TestSuite.check_response(resp) |
| self.session_handle = resp.sessionHandle |
| assert protocol_version <= resp.serverProtocolVersion |
| try: |
| fn() |
| finally: |
| if close_session: |
| close_session_req = TCLIService.TCloseSessionReq() |
| close_session_req.sessionHandle = resp.sessionHandle |
| HS2TestSuite.check_response(self.hs2_client.CloseSession(close_session_req)) |
| self.session_handle = None |
| |
| def needs_session(protocol_version= |
| TCLIService.TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V6, |
| conf_overlay=None, |
| close_session=True, |
| cluster_properties=None): |
| def session_decorator(fn): |
| """Decorator that establishes a session and sets self.session_handle. When the test is |
| finished, the session is closed. |
| """ |
| def add_session(self): |
| add_session_helper(self, protocol_version, conf_overlay, close_session, |
| lambda: fn(self)) |
| return add_session |
| |
| return session_decorator |
| |
| |
| # same as needs_session but takes in a cluster_properties as a argument |
| # cluster_properties is defined as a fixture in conftest.py which allows us |
| # to pass it as an argument to a test. However, it does not work well with |
| # decorators without installing new modules. |
| # Ref: https://stackoverflow.com/questions/19614658 |
| def needs_session_cluster_properties(protocol_version= |
| TCLIService.TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V6, |
| conf_overlay=None, |
| close_session=True): |
| def session_decorator(fn): |
| """Decorator that establishes a session and sets self.session_handle. When the test is |
| finished, the session is closed. |
| """ |
| def add_session(self, cluster_properties, unique_database): |
| add_session_helper(self, protocol_version, conf_overlay, close_session, |
| lambda: fn(self, cluster_properties, unique_database)) |
| return add_session |
| |
| return session_decorator |
| |
| def operation_id_to_query_id(operation_id): |
| lo, hi = operation_id.guid[:8], operation_id.guid[8:] |
| lo = ''.join(['%0.2X' % ord(c) for c in lo[::-1]]) |
| hi = ''.join(['%0.2X' % ord(c) for c in hi[::-1]]) |
| return "%s:%s" % (lo, hi) |
| |
| |
| def create_session_handle_without_secret(session_handle): |
| """Create a HS2 session handle with the same session ID as 'session_handle' but a |
| bogus secret of the right length, i.e. 16 bytes.""" |
| return TCLIService.TSessionHandle(TCLIService.THandleIdentifier( |
| session_handle.sessionId.guid, r"xxxxxxxxxxxxxxxx")) |
| |
| |
| def create_op_handle_without_secret(op_handle): |
| """Create a HS2 operation handle with same parameters as 'op_handle' but with a bogus |
| secret of the right length, i.e. 16 bytes.""" |
| op_id = TCLIService.THandleIdentifier(op_handle.operationId.guid, r"xxxxxxxxxxxxxxxx") |
| return TCLIService.TOperationHandle( |
| op_id, op_handle.operationType, op_handle.hasResultSet) |
| |
| |
| class HS2TestSuite(ImpalaTestSuite): |
| HS2_V6_COLUMN_TYPES = ['boolVal', 'stringVal', 'byteVal', 'i16Val', 'i32Val', 'i64Val', |
| 'doubleVal', 'binaryVal'] |
| |
| def setup(self): |
| host, port = IMPALAD_HS2_HOST_PORT.split(":") |
| self.socket = TSocket(host, port) |
| self.transport = TBufferedTransport(self.socket) |
| self.transport.open() |
| self.protocol = TBinaryProtocol.TBinaryProtocol(self.transport) |
| self.hs2_client = ImpalaHiveServer2Service.Client(self.protocol) |
| |
| def teardown(self): |
| if self.socket: |
| self.socket.close() |
| |
| @staticmethod |
| def check_response(response, |
| expected_status_code = TCLIService.TStatusCode.SUCCESS_STATUS, |
| expected_error_prefix = None): |
| assert response.status.statusCode == expected_status_code |
| if expected_status_code != TCLIService.TStatusCode.SUCCESS_STATUS\ |
| and expected_error_prefix is not None: |
| assert response.status.errorMessage.startswith(expected_error_prefix) |
| |
| @staticmethod |
| def check_invalid_session(response): |
| """Checks that the HS2 API response is the correct response if the session is invalid, |
| i.e. the session doesn't exist or the secret is invalid.""" |
| HS2TestSuite.check_response(response, TCLIService.TStatusCode.ERROR_STATUS, |
| "Invalid session id:") |
| |
| @staticmethod |
| def check_invalid_query(response, expect_legacy_err=False): |
| """Checks that the HS2 API response is the correct response if the query is invalid, |
| i.e. the query doesn't exist, doesn't match the session provided, or the secret is |
| invalid. """ |
| if expect_legacy_err: |
| # Some operations return non-standard errors like "Query id ... not found". |
| expected_err = "Query id" |
| else: |
| # We should standardise on this error message. |
| expected_err = "Invalid query handle:" |
| HS2TestSuite.check_response(response, TCLIService.TStatusCode.ERROR_STATUS, |
| expected_err) |
| |
| @staticmethod |
| def check_profile_access_denied(response, user): |
| """Checks that the HS2 API response is the correct response if the user is not |
| authorised to access the query's profile.""" |
| HS2TestSuite.check_response(response, TCLIService.TStatusCode.ERROR_STATUS, |
| "User {0} is not authorized to access the runtime " |
| "profile or execution summary".format(user)) |
| |
| def close(self, op_handle): |
| close_op_req = TCLIService.TCloseOperationReq() |
| close_op_req.operationHandle = op_handle |
| close_op_resp = self.hs2_client.CloseOperation(close_op_req) |
| assert close_op_resp.status.statusCode == TCLIService.TStatusCode.SUCCESS_STATUS |
| |
| @staticmethod |
| def get_num_rows(result_set): |
| # rows will always be set, so the only way to tell if we should use it is to see if |
| # any columns are set |
| if result_set.columns is None or len(result_set.columns) == 0: |
| return len(result_set.rows) |
| |
| assert result_set.columns is not None |
| for col_type in HS2TestSuite.HS2_V6_COLUMN_TYPES: |
| typed_col = getattr(result_set.columns[0], col_type) |
| if typed_col != None: |
| return len(typed_col.values) |
| |
| assert False |
| |
| def fetch(self, fetch_results_req): |
| """Wrapper around ImpalaHiveServer2Service.FetchResults(fetch_results_req) that |
| issues the given fetch request until the TCLIService.TStatusCode transitions from |
| STILL_EXECUTING_STATUS to SUCCESS_STATUS. If a fetch response contains the |
| STILL_EXECUTING_STATUS then rows are not yet available for consumption (e.g. the |
| query is still running and has not produced any rows yet). This status may be |
| returned to the client if the FETCH_ROWS_TIMEOUT_MS is hit.""" |
| fetch_results_resp = None |
| while fetch_results_resp is None or \ |
| fetch_results_resp.status.statusCode == \ |
| TCLIService.TStatusCode.STILL_EXECUTING_STATUS: |
| fetch_results_resp = self.hs2_client.FetchResults(fetch_results_req) |
| HS2TestSuite.check_response(fetch_results_resp) |
| return fetch_results_resp |
| |
| def fetch_at_most(self, handle, orientation, size, expected_num_rows = None): |
| """Fetches at most size number of rows from the query identified by the given |
| operation handle. Uses the given fetch orientation. Asserts that the fetch returns a |
| success status, and that the number of rows returned is equal to given |
| expected_num_rows (if given). It is only safe for expected_num_rows to be 0 or 1: |
| Impala does not guarantee that a larger result set will be returned in one go. Use |
| fetch_until() for repeated fetches.""" |
| assert expected_num_rows is None or expected_num_rows in (0, 1) |
| fetch_results_req = TCLIService.TFetchResultsReq() |
| fetch_results_req.operationHandle = handle |
| fetch_results_req.orientation = orientation |
| fetch_results_req.maxRows = size |
| fetch_results_resp = self.fetch(fetch_results_req) |
| if expected_num_rows is not None: |
| assert self.get_num_rows(fetch_results_resp.results) == expected_num_rows |
| return fetch_results_resp |
| |
| def fetch_until(self, handle, orientation, size, expected_num_rows = None): |
| """Tries to fetch exactly 'size' rows from the given query handle, with the given |
| fetch orientation, by repeatedly issuing fetch(size - num rows already fetched) |
| calls. Returns fewer than 'size' rows if either a fetch() returns 0 rows (indicating |
| EOS) or 'expected_num_rows' rows are returned. If 'expected_num_rows' is set to None, |
| it defaults to 'size', so that the effect is to both ask for and expect the same |
| number of rows.""" |
| assert expected_num_rows is None or (size >= expected_num_rows) |
| fetch_results_req = TCLIService.TFetchResultsReq() |
| fetch_results_req.operationHandle = handle |
| fetch_results_req.orientation = orientation |
| fetch_results_req.maxRows = size |
| fetch_results_resp = self.fetch(fetch_results_req) |
| num_rows_fetched = self.get_num_rows(fetch_results_resp.results) |
| if expected_num_rows is None: expected_num_rows = size |
| while num_rows_fetched < expected_num_rows: |
| # Always try to fetch at most 'size' |
| fetch_results_req.maxRows = size - num_rows_fetched |
| fetch_results_req.orientation = TCLIService.TFetchOrientation.FETCH_NEXT |
| fetch_results_resp = self.fetch(fetch_results_req) |
| last_fetch_size = self.get_num_rows(fetch_results_resp.results) |
| assert last_fetch_size > 0 |
| num_rows_fetched += last_fetch_size |
| |
| assert num_rows_fetched == expected_num_rows |
| |
| def fetch_fail(self, handle, orientation, expected_error_prefix): |
| """Attempts to fetch rows from the query identified by the given operation handle. |
| Asserts that the fetch returns an error with an error message matching the given |
| expected_error_prefix.""" |
| fetch_results_req = TCLIService.TFetchResultsReq() |
| fetch_results_req.operationHandle = handle |
| fetch_results_req.orientation = orientation |
| fetch_results_req.maxRows = 100 |
| fetch_results_resp = self.hs2_client.FetchResults(fetch_results_req) |
| HS2TestSuite.check_response(fetch_results_resp, TCLIService.TStatusCode.ERROR_STATUS, |
| expected_error_prefix) |
| return fetch_results_resp |
| |
| def result_metadata(self, handle): |
| """ Gets the schema for the query identified by the handle """ |
| req = TCLIService.TGetResultSetMetadataReq() |
| req.operationHandle = handle |
| resp = self.hs2_client.GetResultSetMetadata(req) |
| HS2TestSuite.check_response(resp) |
| return resp |
| |
| def column_results_to_string(self, columns): |
| """Quick-and-dirty way to get a readable string to compare the output of a |
| columnar-oriented query to its expected output""" |
| formatted = "" |
| num_rows = 0 |
| # Determine the number of rows by finding the type of the first column |
| for col_type in HS2TestSuite.HS2_V6_COLUMN_TYPES: |
| typed_col = getattr(columns[0], col_type) |
| if typed_col != None: |
| num_rows = len(typed_col.values) |
| break |
| |
| for i in xrange(num_rows): |
| row = [] |
| for c in columns: |
| for col_type in HS2TestSuite.HS2_V6_COLUMN_TYPES: |
| typed_col = getattr(c, col_type) |
| if typed_col != None: |
| indicator = ord(typed_col.nulls[i / 8]) |
| if indicator & (1 << (i % 8)): |
| row.append("NULL") |
| else: |
| row.append(str(typed_col.values[i])) |
| break |
| formatted += (", ".join(row) + "\n") |
| return (num_rows, formatted) |
| |
| def get_operation_status(self, operation_handle): |
| """Executes GetOperationStatus with the given operation handle and returns the |
| TGetOperationStatusResp""" |
| get_operation_status_req = TCLIService.TGetOperationStatusReq() |
| get_operation_status_req.operationHandle = operation_handle |
| get_operation_status_resp = \ |
| self.hs2_client.GetOperationStatus(get_operation_status_req) |
| return get_operation_status_resp |
| |
| def wait_for_operation_state(self, operation_handle, expected_state, \ |
| timeout = 10, interval = 1): |
| """Waits for the operation to reach expected_state by polling GetOperationStatus every |
| interval seconds, returning the TGetOperationStatusResp, or raising an assertion after |
| timeout seconds.""" |
| start_time = time() |
| while (time() - start_time < timeout): |
| get_operation_status_resp = self.get_operation_status(operation_handle) |
| HS2TestSuite.check_response(get_operation_status_resp) |
| if get_operation_status_resp.operationState is expected_state: |
| return get_operation_status_resp |
| sleep(interval) |
| assert False, 'Did not reach expected operation state %s in time, actual state was ' \ |
| '%s' % (expected_state, get_operation_status_resp.operationState) |
| |
| def wait_for_admission_control(self, operation_handle, timeout = 10): |
| """Waits for the admission control processing of the query to complete by polling |
| GetOperationStatus every interval seconds, returning the TGetOperationStatusResp, |
| or raising an assertion after timeout seconds.""" |
| start_time = time() |
| while (time() - start_time < timeout): |
| get_operation_status_resp = self.get_operation_status(operation_handle) |
| HS2TestSuite.check_response(get_operation_status_resp) |
| if TCLIService.TOperationState.INITIALIZED_STATE < \ |
| get_operation_status_resp.operationState < \ |
| TCLIService.TOperationState.PENDING_STATE: |
| return get_operation_status_resp |
| sleep(0.05) |
| assert False, 'Did not complete admission control processing in time, current ' \ |
| 'operation state of query: %s' % (get_operation_status_resp.operationState) |
| |
| def execute_statement(self, statement, conf_overlay=None, |
| expected_status_code=TCLIService.TStatusCode.SUCCESS_STATUS, |
| expected_error_prefix=None): |
| """Executes statement and checks if the response meets the expectations. |
| If so, it returns the response.""" |
| execute_statement_req = TCLIService.TExecuteStatementReq() |
| execute_statement_req.sessionHandle = self.session_handle |
| execute_statement_req.statement = statement |
| if conf_overlay: |
| execute_statement_req.confOverlay = conf_overlay |
| execute_statement_resp = self.hs2_client.ExecuteStatement(execute_statement_req) |
| HS2TestSuite.check_response(execute_statement_resp, expected_status_code, |
| expected_error_prefix) |
| return execute_statement_resp |
