зеркало из https://github.com/Azure/pykusto.git
Using print instead of log for displaying the device auth message (#170)
This commit is contained in:
Родитель
d682465b92
Коммит
8a627e4fe7
|
@ -6,7 +6,6 @@ import pandas as pd
|
|||
from azure.kusto.data import ClientRequestProperties, KustoClient
|
||||
|
||||
from pykusto import PyKustoClient, NO_RETRIES, KustoResponse, KQL, RetryConfig
|
||||
from .logger import _logger
|
||||
|
||||
|
||||
class DataframeBasedKustoResponse(KustoResponse):
|
||||
|
@ -30,6 +29,7 @@ class PySparkKustoClient(PyKustoClient):
|
|||
"""
|
||||
Handle to a Kusto cluster, to be used inside a PySpark notebook.
|
||||
"""
|
||||
|
||||
def __init__(self, cluster: str, linked_service: str = None, fetch_by_default: bool = True) -> None:
|
||||
"""
|
||||
Create a new handle to a Kusto cluster. The value of "fetch_by_default" is used for current instance, and also passed on to database instances.
|
||||
|
@ -71,7 +71,7 @@ class PySparkKustoClient(PyKustoClient):
|
|||
assert self.__linked_service is None, "Device authentication can be used only when a linked_service was not provided to the client constructor"
|
||||
# noinspection PyProtectedMember
|
||||
device_auth = self.__spark_context._jvm.com.microsoft.kusto.spark.authentication.DeviceAuthentication(self.__cluster_name, "common")
|
||||
_logger.info(device_auth.getDeviceCodeMessage())
|
||||
print(device_auth.getDeviceCodeMessage()) # Logging is better than printing, but the PySpark notebook does not display logs by default
|
||||
self.option('accessToken', device_auth.acquireToken)
|
||||
|
||||
# noinspection PyUnresolvedReferences,PyPackageRequirements
|
||||
|
|
|
@ -2,6 +2,8 @@ import json
|
|||
import logging
|
||||
import sys
|
||||
from concurrent.futures import Future
|
||||
from contextlib import contextmanager
|
||||
from io import StringIO
|
||||
from threading import Event
|
||||
from typing import Callable, Tuple, Any, List, Optional, Union, Type
|
||||
from unittest import TestCase
|
||||
|
@ -270,4 +272,15 @@ def nested_attribute_dict(attributes: str, value: Any) -> Any:
|
|||
return result
|
||||
|
||||
|
||||
@contextmanager
|
||||
def captured_stdout():
|
||||
new_out = StringIO()
|
||||
old_out = sys.stdout
|
||||
try:
|
||||
sys.stdout = new_out
|
||||
yield sys.stdout
|
||||
finally:
|
||||
sys.stdout = old_out
|
||||
|
||||
|
||||
test_logger = logging.getLogger("pykusto_test")
|
||||
|
|
|
@ -1,4 +1,3 @@
|
|||
import logging
|
||||
from unittest.mock import patch
|
||||
|
||||
import pandas as pd
|
||||
|
@ -7,8 +6,7 @@ from pykusto import Query, PySparkKustoClient
|
|||
# noinspection PyProtectedMember
|
||||
from pykusto._src.expressions import _StringColumn, _NumberColumn
|
||||
# noinspection PyProtectedMember
|
||||
from pykusto._src.logger import _logger
|
||||
from test.test_base import TestBase, nested_attribute_dict
|
||||
from test.test_base import TestBase, nested_attribute_dict, captured_stdout
|
||||
|
||||
|
||||
class MockDataFrameReader:
|
||||
|
@ -112,11 +110,11 @@ class TestClient(TestBase):
|
|||
mock_spark_session = MockSparkSession(expected_df)
|
||||
mock_spark_context = MockSparkContext('MOCK_TOKEN')
|
||||
|
||||
with patch('pykusto._src.pyspark_client.PySparkKustoClient._PySparkKustoClient__get_spark_session_and_context', lambda s: (mock_spark_session, mock_spark_context)),\
|
||||
self.assertLogs(_logger, logging.INFO) as cm:
|
||||
with patch('pykusto._src.pyspark_client.PySparkKustoClient._PySparkKustoClient__get_spark_session_and_context', lambda s: (mock_spark_session, mock_spark_context)), \
|
||||
captured_stdout() as out:
|
||||
client = PySparkKustoClient('https://help.kusto.windows.net/', fetch_by_default=False)
|
||||
|
||||
self.assertEqual(["INFO:pykusto:To sign in, use a lubricated goat to open the pod bay doors."], cm.output)
|
||||
self.assertEqual("To sign in, use a lubricated goat to open the pod bay doors.", out.getvalue().strip())
|
||||
table = client['test_db']['mock_table']
|
||||
actual_df = Query(table).take(5).to_dataframe()
|
||||
self.assertTrue(expected_df.equals(actual_df))
|
||||
|
|
Загрузка…
Ссылка в новой задаче