def test_get_response(self, mock_resp, mock_urlopen): client = PrestoClient('any_host', 'any_user') mock_urlopen.return_value = mock_resp mock_resp.read.return_value = '{"message": "ok!"}' client.get_response_from('any_uri') self.assertEqual(client.response_from_server, {"message": "ok!"})
def test_retrieve_rows(self, mock_uri, mock_get_from_uri, mock_port): mock_port.return_value = 8080 client = PrestoClient('any_host', 'any_user') dir = os.path.abspath(os.path.dirname(__file__)) with open(dir + '/resources/valid_rest_response_level1.txt') \ as json_file: client.response_from_server = json.load(json_file) mock_get_from_uri.return_value = True mock_uri.side_effect = [ "http://localhost:8080/v1/statement/2015_harih/2", "" ] self.assertEqual(client.get_rows(), []) self.assertEqual(client.next_uri, "http://localhost:8080/v1/statement/2015_harih/2") with open(dir + '/resources/valid_rest_response_level2.txt') \ as json_file: client.response_from_server = json.load(json_file) mock_uri.side_effect = [ "http://localhost:8080/v1/statement/2015_harih/2", "" ] expected_row = [["uuid1", "http://localhost:8080", "presto-main:0.97", True], ["uuid2", "http://worker:8080", "presto-main:0.97", False]] self.assertEqual(client.get_rows(), expected_row) self.assertEqual(client.next_uri, "")
def test_connection_failed(self, mock_conn, mock_port): mock_port.return_value = 8080 client = PrestoClient('any_host', 'any_user') client.execute_query("any_sql") self.assertTrue(mock_conn().close.called) self.assertFalse(client.execute_query("any_sql"))
def test_http_call_failed(self, mock_conn, mock_presto_config): client = PrestoClient('any_host', 'any_user') mock_conn.side_effect = HTTPException("Error") self.assertFalse(client.run_sql("any_sql")) mock_conn.side_effect = socket.error("Error") self.assertFalse(client.run_sql("any_sql"))
def test_http_answer_valid(self, mock_response, mock_request, mock_port): mock_port.return_value = 8080 client = PrestoClient('any_host', 'any_user') mock_response.return_value.read.return_value = '{}' type(mock_response.return_value).status = \ PropertyMock(return_value=200) self.assertTrue(client.execute_query('any_sql'))
def test_http_call_failed(self, mock_conn): client = PrestoClient('any_host', 'any_user', 8080) mock_conn.side_effect = HTTPException("Error") self.assertFalse(client.execute_query("any_sql")) mock_conn.side_effect = socket.error("Error") self.assertFalse(client.execute_query("any_sql"))
def test_http_answer_valid(self, mock_response, mock_request, mock_presto_config): client = PrestoClient('any_host', 'any_user') mock_response.return_value.read.return_value = '{}' type(mock_response.return_value).status = \ PropertyMock(return_value=200) self.assertEquals(client.run_sql('any_sql'), [])
def test_limit_rows(self, mock_uri, mock_get_from_uri): client = PrestoClient('any_host', 'any_user', 8080) dir = os.path.abspath(os.path.dirname(__file__)) with open(dir + '/files/valid_rest_response_level2.txt') as json_file: client.response_from_server = json.load(json_file) mock_get_from_uri.return_value = True mock_uri.side_effect = ["any_next_uri", ""] self.assertEqual(client.get_rows(0), [])
def test_default_request_called(self, mock_conn, mock_presto_config): client = PrestoClient('any_host', 'any_user') headers = {"X-Presto-Catalog": "hive", "X-Presto-Schema": "default", "X-Presto-User": '******', "X-Presto-Source": "presto-admin"} client.run_sql("any_sql") mock_conn.assert_called_with('any_host', 8080, False, URL_TIMEOUT_MS) mock_conn().request.assert_called_with("POST", "/v1/statement", "any_sql", headers) self.assertTrue(mock_conn().getresponse.called)
def test_limit_rows(self, mock_uri, mock_get_from_uri, mock_presto_config): client = PrestoClient('any_host', 'any_user') dir = os.path.abspath(os.path.dirname(__file__)) with open(dir + '/resources/valid_rest_response_level2.txt') \ as json_file: client.response_from_server = json.load(json_file) mock_get_from_uri.return_value = True mock_uri.side_effect = ["any_next_uri", ""] self.assertEqual(client.get_rows(0), [])
def test_default_request_called(self, mock_conn): client = PrestoClient('any_host', 'any_user', 8080) headers = {"X-Presto-Catalog": "hive", "X-Presto-Schema": "default", "X-Presto-User": '******'} client.execute_query("any_sql") mock_conn.assert_called_with('any_host', 8080, False, URL_TIMEOUT_MS) mock_conn().request.assert_called_with("POST", "/v1/statement", "any_sql", headers) self.assertTrue(mock_conn().getresponse.called)
def test_default_request_called(self, mock_conn, mock_presto_config): client = PrestoClient('any_host', 'any_user') headers = { "X-Presto-Catalog": "hive", "X-Presto-Schema": "default", "X-Presto-User": '******', "X-Presto-Source": "presto-admin" } client.run_sql("any_sql") mock_conn.assert_called_with('any_host', 8080, False, URL_TIMEOUT_MS) mock_conn().request.assert_called_with("POST", "/v1/statement", "any_sql", headers) self.assertTrue(mock_conn().getresponse.called)
def test_create_authorization_headers(self, mock_presto_config): auth_headers = PrestoClient._create_auth_headers( "Aladdin", "open sesame") expected_auth_headers = { "Authorization": "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ==" } self.assertEqual(auth_headers, expected_auth_headers)
def test_http_answer_not_json(self, mock_response, mock_request): client = PrestoClient('any_host', 'any_user', 8080) mock_response.return_value.read.return_value = 'NOT JSON!' type(mock_response.return_value).status =\ PropertyMock(return_value=200) self.assertRaisesRegexp(ValueError, 'No JSON object could be decoded', client.execute_query, 'any_sql')
def test_append_rows(self, mock_uri, mock_get_from_uri): client = PrestoClient('any_host', 'any_user', 8080) dir = os.path.abspath(os.path.dirname(__file__)) with open(dir + '/files/valid_rest_response_level2.txt') as json_file: client.response_from_server = json.load(json_file) mock_get_from_uri.return_value = True mock_uri.side_effect = ["any_next_uri", "any_next_next_uri", "", ""] expected_row = [["uuid1", "http://localhost:8080", "presto-main:0.97", True], ["uuid2", "http://worker:8080", "presto-main:0.97", False], ["uuid1", "http://localhost:8080", "presto-main:0.97", True], ["uuid2", "http://worker:8080", "presto-main:0.97", False]] self.assertEqual(client.get_rows(), expected_row)
def test_no_sql(self): client = PrestoClient('any_host', 'any_user') self.assertRaisesRegexp( InvalidArgumentError, "SQL query missing", client.execute_query, "", )
def test_append_rows(self, mock_uri, mock_get_from_uri, mock_presto_config): client = PrestoClient('any_host', 'any_user') dir = os.path.abspath(os.path.dirname(__file__)) with open(dir + '/resources/valid_rest_response_level2.txt') \ as json_file: client.response_from_server = json.load(json_file) mock_get_from_uri.return_value = True mock_uri.side_effect = ["any_next_uri", "any_next_next_uri", "", ""] expected_row = [ ["uuid1", "http://localhost:8080", "presto-main:0.97", True], ["uuid2", "http://worker:8080", "presto-main:0.97", False], ["uuid1", "http://localhost:8080", "presto-main:0.97", True], ["uuid2", "http://worker:8080", "presto-main:0.97", False] ] self.assertEqual(client.get_rows(), expected_row)
def test_no_sql(self, mock_presto_config): client = PrestoClient('any_host', 'any_user') self.assertRaisesRegexp( InvalidArgumentError, "SQL query missing", client.run_sql, "", )
def check_status_for_control_commands(): client = PrestoClient(env.host, env.port) print('Waiting to make sure we can connect to the Presto server on %s, ' 'please wait. This check will time out after %d minutes if the ' 'server does not respond.' % (env.host, (RETRY_TIMEOUT / 60))) if check_server_status(client): print('Server started successfully on: ' + env.host) else: warn('Server failed to start on: ' + env.host + '\nPlease check ' + REMOTE_PRESTO_LOG_DIR + '/server.log')
def check_status_for_control_commands(): client = PrestoClient(env.host, env.user) print('Waiting to make sure we can connect to the Presto server on %s, ' 'please wait. This check will time out after %d minutes if the ' 'server does not respond.' % (env.host, (RETRY_TIMEOUT / 60))) if check_server_status(client): print('Server started successfully on: ' + env.host) else: warn('Server failed to start on: ' + env.host + '\nPlease check ' + lookup_server_log_file(env.host) + ' and ' + lookup_launcher_log_file(env.host))
def create_presto_client(self, host=None): ips = self.cluster.get_ip_address_dict() config_path = os.path.join('~', LOCAL_CONF_DIR, COORDINATOR_DIR_NAME, CONFIG_PROPERTIES) config = self.cluster.exec_cmd_on_host(self.cluster.master, 'cat ' + config_path) user = '******' if host is None: host = self.cluster.master return PrestoClient( ips[host], user, PrestoConfig.from_file(StringIO(config), config_path, host))
def test_retrieve_rows(self, mock_uri, mock_get_from_uri, mock_port): mock_port.return_value = 8080 client = PrestoClient('any_host', 'any_user') dir = os.path.abspath(os.path.dirname(__file__)) with open(dir + '/resources/valid_rest_response_level1.txt') \ as json_file: client.response_from_server = json.load(json_file) mock_get_from_uri.return_value = True mock_uri.side_effect = [ "http://localhost:8080/v1/statement/2015_harih/2", "" ] self.assertEqual(client.get_rows(), []) self.assertEqual(client.next_uri, "http://localhost:8080/v1/statement/2015_harih/2") with open(dir + '/resources/valid_rest_response_level2.txt') \ as json_file: client.response_from_server = json.load(json_file) mock_uri.side_effect = [ "http://localhost:8080/v1/statement/2015_harih/2", "" ] expected_row = [[ "uuid1", "http://localhost:8080", "presto-main:0.97", True ], ["uuid2", "http://worker:8080", "presto-main:0.97", False]] self.assertEqual(client.get_rows(), expected_row) self.assertEqual(client.next_uri, "")
def get_query_id(self, sql, host=None): ips = self.cluster.get_ip_address_dict() if host is None: host = self.cluster.master client = PrestoClient(ips[host], 'root') run_sql(client, sql) query_runtime_info = run_sql(client, 'SELECT query_id FROM ' 'system.runtime.queries ' 'WHERE query = \'' + sql + '\'') if not query_runtime_info: raise PrestoError('Presto not started up yet.') for row in query_runtime_info: return row[0]
def test_execute_query_get_port(self, run_mock, conn_mock): client = PrestoClient('any_host', 'any_user') client.rows = ['hello'] client.next_uri = 'hello' client.response_from_server = {'hello': 'hello'} run_mock.return_value = _AttributeString('http-server.http.port=8080') run_mock.return_value.failed = False client.execute_query('select * from nation') self.assertEqual(client.port, 8080) self.assertEqual(client.rows, []) self.assertEqual(client.next_uri, '') self.assertEqual(client.response_from_server, {})
def testrun_sql_get_port(self, sudo_mock, conn_mock, mock_presto_config): client = PrestoClient('any_host', 'any_user') client.rows = ['hello'] client.next_uri = 'hello' client.response_from_server = {'hello': 'hello'} sudo_mock.return_value = _AttributeString('http-server.http.port=8080') sudo_mock.return_value.failed = False sudo_mock.return_value.return_code = 0 client.run_sql('select * from nation') self.assertEqual(client.port, 8080) self.assertEqual(client.rows, []) self.assertEqual(client.next_uri, '') self.assertEqual(client.response_from_server, {})
def get_status_from_coordinator(): with closing(PrestoClient(get_coordinator_role()[0], env.user)) as client: try: coordinator_status = run_sql(client, SYSTEM_RUNTIME_NODES) catalog_status = get_catalog_info_from(client) except BaseException as e: # Just log errors that come from a missing port or anything else; if # we can't connect to the coordinator, we just want to print out a # minimal status anyway. _LOGGER.warn(e.message) coordinator_status = [] catalog_status = [] with settings(hide('running')): node_information = execute(collect_node_information, hosts=get_host_list()) for host in get_host_list(): if isinstance(node_information[host], Exception): external_ip = 'Unknown' is_running = False error_message = node_information[host].message else: (external_ip, is_running, error_message) = node_information[host] print_status_header(external_ip, is_running, host) if error_message: print('\t' + error_message) elif not coordinator_status: print( '\tNo information available: unable to query coordinator') elif not is_running: print('\tNo information available') else: version_string = get_presto_version() version = strip_tag(split_version(version_string)) query, processor = NODE_INFO_PER_URI_SQL.for_version(version) # just get the node_info row for the host if server is up node_info_row = run_sql(client, query % external_ip) node_status = processor(node_info_row) if node_status: print_node_info(node_status, catalog_status) else: print( '\tNo information available: the coordinator has not yet' ' discovered this node')
def collect_node_information(): with closing(PrestoClient(get_coordinator_role()[0], env.user)) as client: with settings(hide('warnings')): error_message = check_presto_version() if error_message: external_ip = 'Unknown' is_running = False else: with settings(hide('warnings', 'aborts', 'stdout')): try: external_ip = get_ext_ip_of_node(client) except: external_ip = 'Unknown' try: is_running = service('status') except: is_running = False return external_ip, is_running, error_message
def system_info(): """ Gather system information like nodes in the system, presto version, presto-admin version, os version etc. """ if env.host not in fabricapi.get_coordinator_role(): return err_msg = 'Unable to access node information. ' \ 'Please check that server is up with command: server status' req = get_request(request_url(NODES_REQUEST_EXT), err_msg) if not os.path.exists(TMP_PRESTO_DEBUG): os.mkdir(TMP_PRESTO_DEBUG) downloaded_sys_info_loc = os.path.join(TMP_PRESTO_DEBUG, "sysinfo") node_info_file_name = os.path.join(downloaded_sys_info_loc, 'node_info.json') if not os.path.exists(downloaded_sys_info_loc): os.mkdir(downloaded_sys_info_loc) with open(node_info_file_name, 'w') as out_file: out_file.write(json.dumps(req.json(), indent=4)) _LOGGER.debug('Gathered node information in file: ' + node_info_file_name) conn_file_name = os.path.join(downloaded_sys_info_loc, 'connector_info.txt') client = PrestoClient(env.host, env.user) conn_info = get_connector_info_from(client) with open(conn_file_name, 'w') as out_file: out_file.write(conn_info + '\n') _LOGGER.debug('Gathered connector information in file: ' + conn_file_name) execute(get_system_info, downloaded_sys_info_loc, roles=env.roles) make_tarfile(OUTPUT_FILENAME_FOR_SYS_INFO, downloaded_sys_info_loc) print 'System info archive created: ' + OUTPUT_FILENAME_FOR_SYS_INFO
def check_server_status(): """ Checks if server is running for env.host. Retries connecting to server until server is up or till RETRY_TIMEOUT is reached Parameters: client - client that executes the query Returns: True or False """ if len(get_coordinator_role()) < 1: warn('No coordinator defined. Cannot verify server status.') client = PrestoClient(get_coordinator_role()[0], env.user) node_id = lookup_string_config( 'node.id', os.path.join(constants.REMOTE_CONF_DIR, 'node.properties'), env.host) try: return query_server_for_status(client, node_id) except RetryError: return False
def test_connection_failed(self, mock_conn, mock_presto_config): client = PrestoClient('any_host', 'any_user') client.run_sql("any_sql") self.assertTrue(mock_conn().close.called) self.assertFalse(client.run_sql("any_sql"))
def test_connection_failed(self, mock_conn): client = PrestoClient('any_host', 'any_user', 8080) client.execute_query("any_sql") self.assertTrue(mock_conn().close.called) self.assertFalse(client.execute_query("any_sql"))
def test_create_authorization_headers(self, mock_presto_config): auth_headers = PrestoClient._create_auth_headers("Aladdin", "open sesame") expected_auth_headers = {"Authorization": "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="} self.assertEqual(auth_headers, expected_auth_headers)
def test_no_user(self, mock_presto_config): client = PrestoClient('any_host', "") self.assertRaisesRegexp(InvalidArgumentError, "Username missing", client.run_sql, "any_sql")
def test_no_server(self, mock_presto_config): client = PrestoClient("", 'any_user') self.assertRaisesRegexp(InvalidArgumentError, "Server IP missing", client.run_sql, "any_sql")
def test_create_authorization_headers_fails_with_colon_in_user( self, mock_error, mock_presto_config): PrestoClient._create_auth_headers("Aladdin:1", "open sesame") error_message = "LDAP user cannot contain ':': Aladdin:1" mock_error.assert_called_once_with(error_message)
def test_create_authorization_headers_fails_with_empty_password( self, mock_error, mock_presto_config): PrestoClient._create_auth_headers("Aladdin", "") error_message = 'LDAP password (taken from internal-communication.authentication.ldap.password in ' \ '/etc/presto/config.properties on the coordinator) cannot be null or empty' mock_error.assert_called_once_with(error_message)
def test_create_authorization_headers_fails_with_empty_password(self, mock_error, mock_presto_config): PrestoClient._create_auth_headers("Aladdin", "") error_message = 'LDAP password (taken from internal-communication.authentication.ldap.password in ' \ '/etc/presto/config.properties on the coordinator) cannot be null or empty' mock_error.assert_called_once_with(error_message)
def test_create_authorization_headers_fails_with_null_user(self, mock_error, mock_presto_config): PrestoClient._create_auth_headers(None, "open sesame") error_message = 'LDAP user (taken from internal-communication.authentication.ldap.user in ' \ '/etc/presto/config.properties on the coordinator) cannot be null or empty' mock_error.assert_called_once_with(error_message)
def test_create_authorization_headers_fails_with_colon_in_user(self, mock_error, mock_presto_config): PrestoClient._create_auth_headers("Aladdin:1", "open sesame") error_message = "LDAP user cannot contain ':': Aladdin:1" mock_error.assert_called_once_with(error_message)