Repository: Azure/pykusto

Issue:
Allow `use_global_cache=True` with any authentication method
Use case: researcher was using pykusto from a Data Bricks notebook, and the default Azure CLI authentication was failing with the exception below.
When they provided their own `KustoClient` instance, they could no longer use `use_global_cache`.

```
FileNotFoundError Traceback (most recent call last) /databricks/python/lib/python3.7/site-packages/azure/cli/core/_session.py in load(self, filename, max_age) 47 self.save() ---> 48 with codecs_open(self.filename, 'r', encoding=self._encoding) as f: 49 self.data = json.load(f) /local_disk0/.ephemeral_nfs/envs/pythonEnv-55cda8d1-9409-4809-af73-a9a8df6e827b/lib/python3.7/codecs.py in open(filename, mode, encoding, errors, buffering) 903 mode = mode + 'b' --> 904 file = builtins.open(filename, mode, buffering) 905 if encoding is None: FileNotFoundError: [Errno 2] No such file or directory: '/root/.azure/azureProfile.json' During handling of the above exception, another exception occurred: FileNotFoundError Traceback (most recent call last) <command-1366941365999876> in <module> 8 9 # Use the client when creating the PyKustoClient ---> 10 client = PyKustoClient(cluster,use_global_cache=True) /databricks/python/lib/python3.7/site-packages/pykusto/_src/client.py in __init__(self, client_or_cluster, fetch_by_default, use_global_cache, retry_config) 119 else: 120 --> 121 self.__client = (self._cached_get_client_for_cluster if use_global_cache else self._get_client_for_cluster)(client_or_cluster) 122 self.__cluster_name = client_or_cluster 123 self._refresh_if_needed() /databricks/python/lib/python3.7/site-packages/pykusto/_src/client.py in _cached_get_client_for_cluster(cluster) 178 Provided for convenience during development, not recommended for general use. 179 """ --> 180 return PyKustoClient._get_client_for_cluster(cluster) 181 182 def _internal_get_items(self) -> Dict[str, 'Database']: /databricks/python/lib/python3.7/site-packages/pykusto/_src/client.py in _get_client_for_cluster(cluster) 164 # As a workaround, we first attempt to manually get the Azure CLI token, and see if it works. 165 # Get rid of this workaround once this is resolved: https://github.com/Azure/azure-kusto-python/issues/240 --> 166 stored_token = _get_azure_cli_auth_token() 167 if stored_token is None: 168 _logger.info("Failed to get Azure CLI token, falling back to AAD device authentication") /databricks/python/lib/python3.7/site-packages/azure/kusto/data/security.py in _get_azure_cli_auth_token() 30 31 azure_folder = get_config_dir() ---> 32 ACCOUNT.load(os.path.join(azure_folder, "azureProfile.json")) 33 profile = Profile(storage=ACCOUNT) 34 token_data = profile.get_raw_token()[0][2] /databricks/python/lib/python3.7/site-packages/azure/cli/core/_session.py in load(self, filename, max_age) 59 "Failed to load or parse file %s. It will be overridden by default settings.", 60 self.filename) ---> 61 self.save() 62 63 def save(self): /databricks/python/lib/python3.7/site-packages/azure/cli/core/_session.py in save(self) 63 def save(self): 64 if self.filename: ---> 65 with codecs_open(self.filename, 'w', encoding=self._encoding) as f: 66 json.dump(self.data, f) 67 /local_disk0/.ephemeral_nfs/envs/pythonEnv-55cda8d1-9409-4809-af73-a9a8df6e827b/lib/python3.7/codecs.py in open(filename, mode, encoding, errors, buffering) 902 # Force opening of the file in binary mode 903 mode = mode + 'b' --> 904 file = builtins.open(filename, mode, buffering) 905 if encoding is None: 906 return file FileNotFoundError: [Errno 2] No such file or directory: '/root/.azure/azureProfile.json'
```

Discussion:

Patch:
diff --git a/.idea/runConfigurations/Lint.xml b/.idea/runConfigurations/Lint.xml
index 5135e6a..aa679af 100644
--- a/.idea/runConfigurations/Lint.xml
+++ b/.idea/runConfigurations/Lint.xml
@@ -13,7 +13,7 @@
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
     <option name="SCRIPT_NAME" value="" />
-    <option name="PARAMETERS" value="-m flake8 pykusto test --count --select=E9,F63,F7,F82 --max-complexity=10 --max-line-length=127 --show-source --statistics" />
+    <option name="PARAMETERS" value="-m flake8 setup.py pykusto test --count --max-complexity=10 --max-line-length=180 --show-source --statistics" />
     <option name="SHOW_COMMAND_LINE" value="false" />
     <option name="EMULATE_TERMINAL" value="false" />
     <option name="MODULE_MODE" value="false" />
diff --git a/pykusto/__init__.py b/pykusto/__init__.py
index 1b9302d..e851ccb 100644
--- a/pykusto/__init__.py
+++ b/pykusto/__init__.py
@@ -1,3 +1,4 @@
+# flake8: noqa
 # Allows importing externally-facing classes without specifying submodules.
 # e.g. "from pykusto import PyKustoClient" instead of "from pykusto.client import PyKustoClient"
 # Also allows for a convenient list of all externally-facing classes as the autocomplete of "from pykusto import "
diff --git a/pykusto/_src/client.py b/pykusto/_src/client.py
index 76b9337..66c39f6 100644
--- a/pykusto/_src/client.py
+++ b/pykusto/_src/client.py
@@ -1,6 +1,5 @@
 from collections import defaultdict
 from fnmatch import fnmatch
-from functools import lru_cache
 from threading import Lock
 from typing import Union, List, Tuple, Dict, Generator, Optional, Set, Type, Callable
 from urllib.parse import urlparse
@@ -12,8 +11,6 @@ from azure.kusto.data._models import KustoResultRow as _KustoResultRow
 from azure.kusto.data.exceptions import KustoServiceError
 from azure.kusto.data.helpers import dataframe_from_result_table
 from azure.kusto.data.response import KustoResponseDataSet
-# noinspection PyProtectedMember
-from azure.kusto.data.security import _get_azure_cli_auth_token
 from redo import retrier
 
 from .expressions import BaseColumn, _AnyTypeColumn
@@ -47,9 +44,12 @@ class RetryConfig:
                 for exception_to_check in self.retry_exceptions:
                     if isinstance(e, exception_to_check):
                         if attempt == self.attempts:
-                            _logger.warn(f"Reached maximum number of attempts ({self.attempts}), raising exception")
+                            _logger.warning(f"Reached maximum number of attempts ({self.attempts}), raising exception")
                             raise
-                        _logger.info(f"Attempt number {attempt} out of {self.attempts} failed, previous sleep time was {sleep_time} seconds. Exception: {e}")
+                        _logger.info(
+                            f"Attempt number {attempt} out of {self.attempts} failed, "
+                            f"previous sleep time was {sleep_time} seconds. Exception: {e.__class__.__name__}('{str(e)}')"
+                        )
                         break
                 else:
                     raise
@@ -95,10 +95,15 @@ class PyKustoClient(_ItemFetcher):
     __first_execution: bool
     __first_execution_lock: Lock
     __retry_config: RetryConfig
+    __auth_method: Callable[[str], KustoConnectionStringBuilder]
+
+    __global_client_cache: Dict[str, KustoClient] = {}
+    __global_cache_lock: Lock = Lock()
 
     def __init__(
             self, client_or_cluster: Union[str, KustoClient], fetch_by_default: bool = True, use_global_cache: bool = False,
-            retry_config: RetryConfig = NO_RETRIES
+            retry_config: RetryConfig = NO_RETRIES,
+            auth_method: Callable[[str], KustoConnectionStringBuilder] = KustoConnectionStringBuilder.with_az_cli_authentication,
     ) -> None:
         """
         Create a new handle to Kusto cluster. The value of "fetch_by_default" is used for current instance, and also passed on to database instances.
@@ -111,15 +116,15 @@ class PyKustoClient(_ItemFetcher):
         self.__first_execution = True
         self.__first_execution_lock = Lock()
         self.__retry_config = retry_config
+        self.__auth_method = auth_method
         if isinstance(client_or_cluster, KustoClient):
             self.__client = client_or_cluster
             # noinspection PyProtectedMember
             self.__cluster_name = urlparse(client_or_cluster._query_endpoint).netloc
             assert not use_global_cache, "Global cache not supported when providing your own client instance"
         else:
-
-            self.__client = (self._cached_get_client_for_cluster if use_global_cache else self._get_client_for_cluster)(client_or_cluster)
             self.__cluster_name = client_or_cluster
+            self.__client = (self._cached_get_client_for_cluster if use_global_cache else self._get_client_for_cluster)()
         self._refresh_if_needed()
 
     def __repr__(self) -> str:
@@ -158,26 +163,21 @@ class PyKustoClient(_ItemFetcher):
     def get_cluster_name(self) -> str:
         return self.__cluster_name
 
-    @staticmethod
-    def _get_client_for_cluster(cluster: str) -> KustoClient:
-        # If we call 'with_az_cli_authentication' directly, in case of failure we will get an un-informative exception.
-        # As a workaround, we first attempt to manually get the Azure CLI token, and see if it works.
-        # Get rid of this workaround once this is resolved: https://github.com/Azure/azure-kusto-python/issues/240
-        stored_token = _get_azure_cli_auth_token()
-        if stored_token is None:
-            _logger.info("Failed to get Azure CLI token, falling back to AAD device authentication")
-            connection_string_builder = KustoConnectionStringBuilder.with_aad_device_authentication(cluster)
-        else:
-            connection_string_builder = KustoConnectionStringBuilder.with_az_cli_authentication(cluster)
-        return KustoClient(connection_string_builder)
+    def _get_client_for_cluster(self) -> KustoClient:
+        return KustoClient(self.__auth_method(self.__cluster_name))
 
-    @staticmethod
-    @lru_cache(maxsize=128)
-    def _cached_get_client_for_cluster(cluster: str) -> KustoClient:
+    def _cached_get_client_for_cluster(self) -> KustoClient:
         """
         Provided for convenience during development, not recommended for general use.
         """
-        return PyKustoClient._get_client_for_cluster(cluster)
+        with PyKustoClient.__global_cache_lock:
+            client = PyKustoClient.__global_client_cache.get(self.__cluster_name)
+            if client is None:
+                client = self._get_client_for_cluster()
+                PyKustoClient.__global_client_cache[self.__cluster_name] = client
+                assert len(PyKustoClient.__global_client_cache) <= 1024, "Global client cache cannot exceed size of 1024"
+
+        return client
 
     def _internal_get_items(self) -> Dict[str, 'Database']:
         # Retrieves database names, table names, column names and types for all databases. A database name is required
diff --git a/pykusto/_src/functions.py b/pykusto/_src/functions.py
index fe348d1..2885200 100644
--- a/pykusto/_src/functions.py
+++ b/pykusto/_src/functions.py
@@ -534,7 +534,7 @@ class Functions:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/make-timespanfunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     # def max_of(self): return
     #
@@ -553,7 +553,7 @@ class Functions:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/newguidfunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     @staticmethod
     def now(offset: TimespanType = None) -> _DatetimeExpression:
@@ -628,14 +628,14 @@ class Functions:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/percentile-tdigestfunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     @staticmethod
     def percentrank_tdigest() -> AnyExpression:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/percentrank-tdigestfunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     @staticmethod
     def pow(expr1: NumberType, expr2: NumberType) -> _NumberExpression:
@@ -1021,7 +1021,7 @@ class Functions:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/toguidfunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     @staticmethod
     def to_hex(expr1: NumberType, expr2: NumberType = None) -> _StringExpression:
@@ -1104,21 +1104,21 @@ class Functions:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/urldecodefunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     @staticmethod
     def url_encode() -> _StringExpression:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/urlencodefunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     @staticmethod
     def week_of_year() -> _NumberExpression:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/weekofyearfunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     # def welch_test(self): return
 
@@ -1127,7 +1127,7 @@ class Functions:
         """
         https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/zipfunction
         """
-        raise NotImplemented  # pragma: no cover
+        raise NotImplementedError()  # pragma: no cover
 
     # ----------------------------------------------------
     # Aggregation functions
diff --git a/pykusto/_src/item_fetcher.py b/pykusto/_src/item_fetcher.py
index c93368f..d4e5725 100644
--- a/pykusto/_src/item_fetcher.py
+++ b/pykusto/_src/item_fetcher.py
@@ -136,7 +136,7 @@ class _ItemFetcher(metaclass=ABCMeta):
                 # If an exception is raised in the future thread, it is re-raised here
                 self.__future.result(timeout=timeout_seconds)
             except TimeoutError:
-                _logger.warn("Timeout while waiting for response from Kusto")
+                _logger.warning("Timeout while waiting for response from Kusto")
 
     def blocking_refresh(self, timeout_seconds: Union[None, float] = None) -> None:
         self.refresh()
diff --git a/setup.py b/setup.py
index 0fc071b..b69db1d 100644
--- a/setup.py
+++ b/setup.py
@@ -23,20 +23,24 @@ setup(
     long_description_content_type="text/markdown",
     keywords="kusto azure-data-explorer client library query",
     install_requires=[
-        'azure-kusto-data>=0.0.43,<=0.1.0',  # In 0.0.43 some packages were renamed
-        'pandas>=0.24.1,<=1.2.3',  # azure-kusto-data requires 0.24.1
+        # Release notes: https://github.com/Azure/azure-kusto-python/releases
+        'azure-kusto-data==2.1.1',  # Earlier versions not supported because of: https://github.com/Azure/azure-kusto-python/issues/312
+
+        # Release notes: https://pandas.pydata.org/docs/whatsnew/index.html
+        # Note that starting from 1.2.0 support for Python 3.6 was dropped
+        'pandas>=0.25.0,<=1.2.4',  # Tests use DataFrame constructor options introduced in 0.25.0
+
         'redo==2.0.4',
     ],
     tests_require=[
         'pytest',
         'pytest-cov',
         'flake8',
-        'pandas>=0.25.0',  # Tests use DataFrame constructor options introduced in 0.25.0
     ],
     classifiers=[
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Developers",
-        "Topic :: Software Development",        
+        "Topic :: Software Development",
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",


Test patch:
diff --git a/.github/workflows/runtests.yml b/.github/workflows/runtests.yml
index 3759f53..4ede851 100644
--- a/.github/workflows/runtests.yml
+++ b/.github/workflows/runtests.yml
@@ -17,20 +17,23 @@ jobs:
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python 3.8
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v2
       with:
+        # TODO: test all relevant versions
         python-version: 3.8
     - name: Install dependencies
-      run: |        
+      # According to the internet using 'python -m pip' instead of 'pip' can prevent some issues
+      run: |
+        python -m pip install --upgrade pip
+        # For some reason transitive dependencies are not installed properly, unless we first explictly install 'azure-kusto-data'. Hence this ugly hack.
+        python -m pip install `grep -oP "azure-kusto-data==[^']+" setup.py`
         python setup.py install
+        python -m pip freeze
     - name: Lint with flake8
       run: |
-        pip install flake8
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+        python -m pip install flake8
+        python -m flake8 setup.py pykusto test --count --max-complexity=10 --max-line-length=180 --show-source --statistics
     - name: Test with pytest
       run: |
-        pip install pytest pytest-cov
-        pytest --cov=pykusto --cov-report term-missing --cov-fail-under=100
+        python -m pip install pytest pytest-cov
+        python -m pytest --cov=pykusto --cov-report term-missing --cov-fail-under=100
diff --git a/test/test_base.py b/test/test_base.py
index 7b19cc5..a46e84b 100644
--- a/test/test_base.py
+++ b/test/test_base.py
@@ -89,7 +89,8 @@ class CustomAssertLogsContext(_AssertLogsContext):
 # noinspection PyMissingConstructor
 class MockKustoResultTable(KustoResultTable):
     def __init__(self, rows: Tuple[Any, ...], columns: Tuple[str, ...]):
-        self.rows = tuple(KustoResultRow(columns, row) for row in rows)
+        self.kusto_result_rows = tuple(KustoResultRow(columns, row) for row in rows)
+        self.raw_rows = self.kusto_result_rows
         self.columns = tuple(type('Column', (object,), {'column_name': col, 'column_type': ''}) for col in columns)
 
 
@@ -230,5 +231,4 @@ class MockKustoClient(KustoClient):
         return response()
 
 
-
 test_logger = logging.getLogger("pykusto_test")
diff --git a/test/test_client.py b/test/test_client.py
index 0aca717..5fc05c2 100644
--- a/test/test_client.py
+++ b/test/test_client.py
@@ -2,12 +2,13 @@ import logging
 from typing import Type
 from unittest.mock import patch
 
-from azure.kusto.data import KustoClient
+from azure.kusto.data import KustoClient, KustoConnectionStringBuilder
 from azure.kusto.data.exceptions import KustoError
 
 from pykusto import PyKustoClient, column_generator as col, Query, KustoServiceError, RetryConfig, NO_RETRIES
 # noinspection PyProtectedMember
 from pykusto._src.logger import _logger
+# noinspection PyProtectedMember
 from test.test_base import TestBase, MockKustoClient, RecordedQuery, mock_response
 
 
@@ -75,7 +76,7 @@ class TestClient(TestBase):
 
     def test_default_authentication(self):
         mock_kusto_client = MockKustoClient()
-        with patch('pykusto._src.client.PyKustoClient._get_client_for_cluster', lambda s, cluster: mock_kusto_client):
+        with patch('pykusto._src.client.PyKustoClient._get_client_for_cluster', lambda s: mock_kusto_client):
             table = PyKustoClient('https://help.kusto.windows.net/')['test_db']['mock_table']
             Query().take(5).execute(table)
         self.assertIs(
@@ -155,21 +156,18 @@ class TestClient(TestBase):
         )
 
     def test_client_for_cluster_with_azure_cli_auth(self):
-        with patch('pykusto._src.client._get_azure_cli_auth_token', lambda: "MOCK_TOKEN"), self.assertLogs(_logger, logging.INFO) as cm:
+        with self.assertLogs(_logger, logging.INFO) as cm:
             client = PyKustoClient('https://help.kusto.windows.net', fetch_by_default=False)
             self.assertIsInstance(client._PyKustoClient__client, KustoClient)
             self.assertEqual('https://help.kusto.windows.net', client.get_cluster_name())
         self.assertEqual([], cm.output)
 
-    def test_client_for_cluster_fallback_to_aad_device_auth(self):
-        with patch('pykusto._src.client._get_azure_cli_auth_token', lambda: None), self.assertLogs(_logger, logging.INFO) as cm:
-            client = PyKustoClient('https://help.kusto.windows.net', fetch_by_default=False)
+    def test_client_for_cluster_with_aad_device_auth(self):
+        with self.assertLogs(_logger, logging.INFO) as cm:
+            client = PyKustoClient('https://help.kusto.windows.net', fetch_by_default=False, auth_method=KustoConnectionStringBuilder.with_aad_device_authentication)
             self.assertIsInstance(client._PyKustoClient__client, KustoClient)
             self.assertEqual('https://help.kusto.windows.net', client.get_cluster_name())
-        self.assertEqual(
-            ['INFO:pykusto:Failed to get Azure CLI token, falling back to AAD device authentication'],
-            cm.output
-        )
+        self.assertEqual([], cm.output)
 
     @staticmethod
     def unreliable_mock_kusto_client(number_of_failures: int, exception_type: Type[Exception] = KustoServiceError):
@@ -195,7 +193,7 @@ class TestClient(TestBase):
         self.assertEqual(
             [
                 f"INFO:pykusto:Attempt number {i} out of {number_of_retries} failed, "
-                f"previous sleep time was 0.1 seconds. Exception: (KustoServiceError(...), 'Mock exception for test')"
+                f"previous sleep time was 0.1 seconds. Exception: KustoServiceError('Mock exception for test')"
                 for i in range(1, number_of_retries)
             ],
             cm.output