Add managed streaming ingest

azure-kusto-data/azure/kusto/data/_models.py

@@ -8,7 +8,7 @@
 from typing import Iterator, List, Any, Union, Optional, Dict
 
 from . import _converters
-from .exceptions import KustoServiceError, KustoStreamingQueryError
+from .exceptions import KustoMultiApiError, KustoStreamingQueryError
 
 
 class WellKnownDataSet(Enum):
@@ -158,7 +158,7 @@ def __init__(self, json_table: Dict[str, Any]):
         super().__init__(json_table)
         errors = [row for row in json_table["Rows"] if isinstance(row, dict)]
         if errors:
-            raise KustoServiceError(errors[0]["OneApiErrors"][0]["error"]["@message"], kusto_response=json_table)
+            raise KustoMultiApiError(errors)
 
     @property
     def rows(self) -> List[KustoResultRow]:

azure-kusto-data/azure/kusto/data/aio/streaming_response.py

@@ -5,7 +5,7 @@
 from ijson import IncompleteJSONError
 
 from azure.kusto.data._models import WellKnownDataSet
-from azure.kusto.data.exceptions import KustoTokenParsingError, KustoServiceError, KustoUnsupportedApiError
+from azure.kusto.data.exceptions import KustoTokenParsingError, KustoUnsupportedApiError, KustoApiError, KustoMultiApiError
 from azure.kusto.data.streaming_response import JsonTokenType, FrameType, JsonToken
 
 
@@ -163,9 +163,7 @@ async def row_iterator(self) -> Iterator[list]:
         while True:
             token = await self.reader.read_token_of_type(JsonTokenType.START_ARRAY, JsonTokenType.END_ARRAY, JsonTokenType.START_MAP)
             if token.token_type == JsonTokenType.START_MAP:
-                raise KustoServiceError(
-                    "Received error in data: " + str(self.parse_object(skip_start=True))
-                )  # Todo - replace this with a better error once managed streaming ingest is merged
+                raise KustoMultiApiError([await self.parse_object(skip_start=True)])
             if token.token_type == JsonTokenType.END_ARRAY:
                 return
             yield await self.parse_array(skip_start=True)

azure-kusto-data/azure/kusto/data/client.py

@@ -8,7 +8,7 @@
 from copy import copy
 from datetime import timedelta
 from enum import Enum, unique
-from typing import TYPE_CHECKING, Union, Callable, Optional, Any, Coroutine, List, Tuple
+from typing import TYPE_CHECKING, Union, Callable, Optional, Any, Coroutine, List, Tuple, AnyStr, IO
 
 import requests
 from requests import Response
@@ -17,7 +17,7 @@
 
 from ._version import VERSION
 from .data_format import DataFormat
-from .exceptions import KustoServiceError
+from .exceptions import KustoServiceError, KustoApiError
 from .response import KustoResponseDataSetV1, KustoResponseDataSetV2, KustoStreamingResponseDataSet, KustoResponseDataSet
 from .security import _AadHelper
 from .streaming_response import StreamingDataSetEnumerator, JsonTokenReader
@@ -92,6 +92,14 @@ def parse(cls, key: str) -> "KustoConnectionStringBuilder.ValidKeywords":
                 return cls.msi_auth
             if key in ["msi_type"]:
                 return cls.msi_params
+            if key in ["az cli"]:
+                return cls.az_cli
+            if key in ["interactive login"]:
+                return cls.interactive_login
+            if key in ["login hint"]:
+                return cls.login_hint
+            if key in ["domain hint"]:
+                return cls.domain_hint
             raise KeyError(key)
 
         def is_secret(self) -> bool:
@@ -367,6 +375,7 @@ def with_aad_managed_service_identity_authentication(
         if object_id is not None:
             # Until we upgrade azure-identity to version 1.4.1, only client_id is excepted as a hint for user managed service identity
             raise ValueError("User Managed Service Identity with object_id is temporarily not supported by azure identity 1.3.1. Please use client_id instead.")
+            # noinspection PyUnreachableCode
             params["object_id"] = object_id
             exclusive_pcount += 1
 
@@ -375,6 +384,7 @@ def with_aad_managed_service_identity_authentication(
             raise ValueError(
                 "User Managed Service Identity with msi_res_id is temporarily not supported by azure identity 1.3.1. Please use client_id instead."
             )
+            # noinspection PyUnreachableCode
             params["msi_res_id"] = msi_res_id
             exclusive_pcount += 1
 
@@ -580,6 +590,8 @@ class ClientRequestProperties:
     For more information please look at: https://docs.microsoft.com/en-us/azure/kusto/api/netfx/request-properties
     """
 
+    OPTION_CLIENT_REQUEST_ID = "ClientRequestId"
+
     results_defer_partial_query_failures_option_name = "deferpartialqueryfailures"
     request_timeout_option_name = "servertimeout"
     no_request_timeout_option_name = "norequesttimeout"
@@ -717,15 +729,17 @@ def _handle_http_error(
             if payload:
                 raise KustoServiceError("The ingestion endpoint does not exist. Please enable streaming ingestion on your cluster.", response) from exception
 
-            raise KustoServiceError("The requested endpoint '{}' does not exist.".format(endpoint), response) from exception
+            raise KustoServiceError(f"The requested endpoint '{endpoint}' does not exist.", response) from exception
 
         if payload:
-            raise KustoServiceError(
-                "An error occurred while trying to ingest: Status: {0.status_code}, Reason: {0.reason}, Text: {1}.".format(response, response_text), response
-            ) from exception
+            message = f"An error occurred while trying to ingest: Status: {status}, Reason: {response.reason}, Text: {response_text}."
+            if response_json:
+                raise KustoApiError(response_json, message, response) from exception
+
+            raise KustoServiceError(message, response) from exception
 
         if response_json:
-            raise KustoServiceError([response_json], response) from exception
+            raise KustoApiError(response_json, http_response=response) from exception
 
         if response_text:
             raise KustoServiceError(response_text, response) from exception
@@ -864,7 +878,7 @@ def execute_streaming_ingest(
         self,
         database: str,
         table: str,
-        stream: io.IOBase,
+        stream: IO[AnyStr],
         stream_format: Union[DataFormat, str],
         properties: Optional[ClientRequestProperties] = None,
         mapping_name: str = None,
@@ -915,7 +929,7 @@ def _execute(
         endpoint: str,
         database: str,
         query: Optional[str],
-        payload: Optional[io.IOBase],
+        payload: Optional[IO[AnyStr]],
         timeout: timedelta,
         properties: Optional[ClientRequestProperties] = None,
         stream_response: bool = False,
@@ -934,7 +948,7 @@ def _execute(
                 response.raise_for_status()
                 return response
             except Exception as e:
-                raise self._handle_http_error(e, self._query_endpoint, None, response, response.json(), response.text)
+                raise self._handle_http_error(e, self._query_endpoint, None, response, response.status_code, response.json(), response.text)
 
         response_json = None
         try:

azure-kusto-data/azure/kusto/data/exceptions.py

@@ -2,13 +2,14 @@
 # Licensed under the MIT License
 from typing import List, Union, TYPE_CHECKING, Optional, Dict, Any
 
-import requests
-
 if TYPE_CHECKING:
+    import requests
+
     try:
         from aiohttp import ClientResponse
     except ImportError:
         # No aio installed, ignore
+        ClientResponse = None
         pass
 
 
@@ -57,10 +58,78 @@ def get_partial_results(self) -> Optional[Dict[str, Any]]:
         return self.kusto_response
 
 
+class OneApiError:
+    def __init__(self, code: str, message: str, type: str, description: str, context: dict, permanent: bool) -> None:
+        self.code = code
+        self.message = message
+        self.type = type
+        self.description = description
+        self.context = context
+        self.permanent = permanent
+
+    @staticmethod
+    def from_dict(obj: dict) -> "OneApiError":
+        code = obj["code"]
+        message = obj["message"]
+        type = obj["@type"]
+        description = obj["@message"]
+        context = obj["@context"]
+        permanent = obj["@permanent"]
+        return OneApiError(code, message, type, description, context, permanent)
+
+
+class KustoMultiApiError(KustoServiceError):
+    """
+    Represents a collection of standard API errors from kusto. Use `get_api_errors()` to retrieve more details.
+    """
+
+    def __init__(self, errors: List[dict]):
+        self.errors = KustoMultiApiError.parse_errors(errors)
+        messages = [error.description for error in self.errors]
+        super().__init__(messages[0] if len(self.errors) == 1 else messages)
+
+    def get_api_errors(self) -> List[OneApiError]:
+        return self.errors
+
+    @staticmethod
+    def parse_errors(errors: List[dict]) -> List[OneApiError]:
+        parsed_errors = []
+        for error_block in errors:
+            one_api_errors = error_block.get("OneApiErrors", None)
+            if not one_api_errors:
+                continue
+            for inner_error in one_api_errors:
+                error_dict = inner_error.get("error", None)
+                if error_dict:
+                    parsed_errors.append(OneApiError.from_dict(error_dict))
+        return parsed_errors
+
+
+class KustoApiError(KustoServiceError):
+    """
+    Represents a standard API error from kusto. Use `get_api_error()` to retrieve more details.
+    """
+
+    def __init__(self, error_dict: dict, message: str = None, http_response: "Union[requests.Response, ClientResponse, None]" = None, kusto_response=None):
+        self.error = OneApiError.from_dict(error_dict["error"])
+        super().__init__(message or self.error.description, http_response, kusto_response)
+
+    def get_api_error(self) -> OneApiError:
+        return self.error
+
+
 class KustoClientError(KustoError):
     """Raised when a Kusto client is unable to send or complete a request."""
 
 
+class KustoBlobError(KustoClientError):
+    def __init__(self, inner: Exception):
+        self.inner = inner
+
+    def message(self) -> str:
+        return f"Failed to upload blob: {self.inner}"
+
+
 class KustoUnsupportedApiError(KustoError):
     """Raised when a Kusto client is unable to send or complete a request."""
 

azure-kusto-data/azure/kusto/data/streaming_response.py

@@ -5,7 +5,7 @@
 from ijson import IncompleteJSONError
 
 from azure.kusto.data._models import WellKnownDataSet
-from azure.kusto.data.exceptions import KustoServiceError, KustoTokenParsingError, KustoUnsupportedApiError
+from azure.kusto.data.exceptions import KustoServiceError, KustoTokenParsingError, KustoUnsupportedApiError, KustoApiError, KustoMultiApiError
 
 
 class JsonTokenType(Enum):
@@ -202,9 +202,7 @@ def row_iterator(self) -> Iterator[list]:
             if token.token_type == JsonTokenType.START_MAP:
                 # Todo - this method of error handling may be problematic, since after raising an error the iteration stops.
                 #  This means that if there are more data or even more errors, we can't read them
-                raise KustoServiceError(
-                    "Received error in data: " + str(self.parse_object(skip_start=True))
-                )  # Todo - replace this with a better error once managed streaming ingest is merged
+                raise KustoMultiApiError([self.parse_object(skip_start=True)])
             if token.token_type == JsonTokenType.END_ARRAY:
                 return
             yield self.parse_array(skip_start=True)