Files
impala/tests/shell/test_shell_client.py
Joe McDonnell ea0969a772 IMPALA-11980 (part 2): Fix absolute import issues for impala_shell
Python 3 changed the behavior of imports with PEP328. Existing
imports become absolute unless they use the new relative import
syntax. This adapts the impala-shell code to use absolute
imports, fixing issues where it is imported from our test code.

There are several parts to this:
1. It moves impala shell code into shell/impala_shell.
   This matches the directory structure of the PyPi package.
2. It changes the imports in the shell code to be
   absolute paths (i.e. impala_shell.foo rather than foo).
   This fixes issues with Python 3 absolute imports.
   It also eliminates the need for ugly hacks in the PyPi
   package's __init__.py.
3. This changes Thrift generation to put it directly in
   $IMPALA_HOME/shell rather than $IMPALA_HOME/shell/gen-py.
   This means that the generated Thrift code is rooted in
   the same directory as the shell code.
4. This changes the PYTHONPATH to include $IMPALA_HOME/shell
   and not $IMPALA_HOME/shell/gen-py. This means that the
   test code is using the same import paths as the pypi
   package.

With all of these changes, the source code is very close
to the directory structure of the PyPi package. As long as
CMake has generated the thrift files and the Python version
file, only a few differences remain. This removes those
differences by moving the setup.py / MANIFEST.in and other
files from the packaging directory to the top-level
shell/ directory. This means that one can pip install
directly from the source code. i.e. pip install $IMPALA_HOME/shell

This also moves the shell tarball generation script to the
packaging directory and changes bin/impala-shell.sh to use
Python 3.

This sorts the imports using isort for the affected Python files.

Testing:
 - Ran a regular core job with Python 2
 - Ran a core job with Python 3 and verified that the absolute
   import issues are gone.

Change-Id: Ica75a24fa6bcb78999b9b6f4f4356951b81c3124
Reviewed-on: http://gerrit.cloudera.org:8080/22330
Reviewed-by: Riza Suminto <riza.suminto@cloudera.com>
Reviewed-by: Michael Smith <michael.smith@cloudera.com>
Tested-by: Riza Suminto <riza.suminto@cloudera.com>
2025-05-21 15:14:11 +00:00

107 lines
4.4 KiB
Python

#!/usr/bin/env impala-python
# -*- coding: utf-8 -*-
#
# 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.
from __future__ import absolute_import, division, print_function
from impala_shell.impala_client import ImpalaBeeswaxClient, ImpalaHS2Client
from tests.common.impala_test_suite import ImpalaTestSuite
from tests.common.test_dimensions import (
create_client_protocol_dimension,
create_client_protocol_no_strict_dimension,
create_single_exec_option_dimension,
create_uncompressed_text_dimension,
)
from tests.shell.util import get_impalad_host_port
class TestShellClient(ImpalaTestSuite):
"""Tests for the Impala Shell clients: ImpalaBeeswaxClient and ImpalaHS2Client."""
@classmethod
def add_test_dimensions(cls):
super(TestShellClient, cls).add_test_dimensions()
# Limit to uncompressed text with default exec options
cls.ImpalaTestMatrix.add_dimension(
create_uncompressed_text_dimension(cls.get_workload()))
cls.ImpalaTestMatrix.add_dimension(create_single_exec_option_dimension())
# Run with beeswax and HS2
cls.ImpalaTestMatrix.add_dimension(create_client_protocol_dimension())
cls.ImpalaTestMatrix.add_dimension(create_client_protocol_no_strict_dimension())
def test_fetch_size(self, vector):
"""Tests that when result spooling is disabled, setting a small batch_size causes
the shell to fetch a single batch at a time, even when the configured fetch size is
larger than the batch_size."""
handle = None
num_rows = 100
batch_size = 10
query_options = {'batch_size': str(batch_size), 'spool_query_results': 'false'}
client = self.__get_shell_client(vector)
try:
client.connect()
handle = client.execute_query(
"select * from functional.alltypes limit {0}".format(num_rows), query_options)
self.__fetch_rows(client.fetch(handle), batch_size, num_rows)
finally:
if handle is not None: client.close_query(handle)
client.close_connection()
def test_fetch_size_result_spooling(self, vector):
"""Tests that when result spooling is enabled, that the exact fetch_size is honored
even if a small batch_size is configured."""
handle = None
fetch_size = 20
num_rows = 100
query_options = {'batch_size': '10', 'spool_query_results': 'true'}
client = self.__get_shell_client(vector, fetch_size)
try:
client.connect()
handle = client.execute_query(
"select * from functional.alltypes limit {0}".format(num_rows), query_options)
self.__fetch_rows(client.fetch(handle), num_rows // fetch_size, num_rows)
finally:
if handle is not None: client.close_query(handle)
client.close_connection()
def __fetch_rows(self, fetch_batches, num_batches, num_rows):
"""Fetches all rows using the given fetch_batches generator. Asserts that num_batches
batches are produced by the generator and that num_rows are returned."""
num_batches_count = 0
rows_per_batch = num_rows // num_batches
for fetch_batch in fetch_batches:
assert len(fetch_batch) == rows_per_batch
num_batches_count += 1
if num_batches_count == num_batches: break
assert num_batches_count == num_batches
def __get_shell_client(self, vector, fetch_size=1024):
"""Returns the client specified by the protocol in the given vector."""
impalad = get_impalad_host_port(vector).split(":")
protocol = vector.get_value("protocol")
if protocol == 'hs2':
return ImpalaHS2Client(impalad, fetch_size, None)
elif protocol == 'hs2-http':
return ImpalaHS2Client(impalad, fetch_size, None,
use_http_base_transport=True, http_path='cliservice')
elif protocol == 'beeswax':
return ImpalaBeeswaxClient(impalad, fetch_size, None)