ingest

38class IngestionResult:
39    """
40    The result of an ingestion.
41    """
42
43    status: IngestionStatus
44    "Will be `Queued` if the ingestion is queued, or `Success` if the ingestion is streaming and successful."
45
46    database: str
47    "The name of the database where the ingestion was performed."
48
49    table: str
50    "The name of the table where the ingestion was performed."
51
52    source_id: uuid.UUID
53    "The source id of the ingestion."
54
55    blob_uri: Optional[str]
56    "The blob uri of the ingestion, if exists."
57
58    def __init__(self, status: IngestionStatus, database: str, table: str, source_id: uuid.UUID, blob_uri: Optional[str] = None):
59        self.status = status
60        self.database = database
61        self.table = table
62        self.source_id = source_id
63        self.blob_uri = blob_uri
64
65    def __repr__(self):
66        # Remove query parameters from blob_uri, if exists
67        obfuscated_path = None
68        if isinstance(self.blob_uri, str):
69            obfuscated_path = self.blob_uri.split("?")[0].split(";")[0]
70        blob_uri = f", obfuscated_blob_uri={obfuscated_path}" if obfuscated_path else ""
71        return f"IngestionResult(status={self.status}, database={self.database}, table={self.table}, source_id={self.source_id}{blob_uri})"

26class IngestionStatus(Enum):
27    """
28    The ingestion was queued.
29    """
30
31    QUEUED = "QUEUED"
32    """
33    The ingestion was successfully streamed
34    """
35    SUCCESS = "SUCCESS"

131class BlobDescriptor(DescriptorBase):
132    """BlobDescriptor is used to describe a blob that will be used as an ingestion source"""
133
134    _BLOB_URI = "blob_uri"
135
136    def __init__(self, path: str, size: Optional[int] = None, source_id: OptionalUUID = None):
137        """
138        :param path: blob uri.
139        :type path: str.
140        :param size: estimated size of file if known.
141        :type size: Optional[int].
142        :param source_id: a v4 uuid to serve as the source's id.
143        :type source_id: OptionalUUID
144        """
145        self.path: str = path
146        self.size: Optional[int] = size
147        self.source_id: uuid.UUID = ensure_uuid(source_id)
148
149    def get_tracing_attributes(self) -> dict:
150        # Remove query parameters from self.path, if exists
151        if self.path:
152            obfuscated_path = self.path.split("?")[0].split(";")[0]
153        return {self._BLOB_URI: obfuscated_path, self._SOURCE_ID: str(self.source_id)}
154
155    def fill_size(self):
156        if not self.size:
157            self.size = BlobClient.from_blob_url(self.path).get_blob_properties().size

 41class FileDescriptor(DescriptorBase):
 42    """FileDescriptor is used to describe a file that will be used as an ingestion source."""
 43
 44    # Gzip keeps the decompressed stream size as a UINT32 in the last 4 bytes of the stream, however this poses a limit to the expressed size which is 4GB
 45    # The standard says that when the size is bigger than 4GB, the UINT rolls over.
 46    # The below constant expresses the maximal size of a compressed stream that will not cause the UINT32 to rollover given a maximal compression ratio of 1:40
 47    GZIP_MAX_DISK_SIZE_FOR_DETECTION = int(4 * 1024 * 1024 * 1024 / 40)
 48    DEFAULT_COMPRESSION_RATIO = 11
 49    _FILE_PATH = "file_path"
 50
 51    def __init__(self, path: str, size: Optional[int] = None, source_id: OptionalUUID = None):
 52        """
 53        :param path: file path.
 54        :type path: str.
 55        :param size: estimated size of file if known. if None or 0 will try to guess.
 56        :type size: Optional[int].
 57        :param source_id: a v4 uuid to serve as the source's id.
 58        :type source_id: OptionalUUID
 59        """
 60        self.path: str = path
 61        self._size: Optional[int] = size
 62        self._detect_size_once: bool = not size
 63
 64        self.source_id: uuid.UUID = ensure_uuid(source_id)
 65        self.stream_name: str = os.path.basename(self.path)
 66
 67    @property
 68    def size(self) -> int:
 69        if self._detect_size_once:
 70            self._detect_size()
 71            self._detect_size_once = False
 72
 73        return self._size
 74
 75    @size.setter
 76    def size(self, size: int):
 77        if size:
 78            self._size = size
 79            self._detect_size_once = False
 80
 81    def _detect_size(self):
 82        uncompressed_size = 0
 83        if self.path.endswith(".gz"):
 84            # This logic follow after the C# implementation
 85            # See IngstionHelpers.cs for an explanation as to what stands behind it
 86            with open(self.path, "rb") as f:
 87                disk_size = f.seek(-4, SEEK_END)
 88                uncompressed_size = struct.unpack("I", f.read(4))[0]
 89                if (disk_size >= uncompressed_size) or (disk_size >= self.GZIP_MAX_DISK_SIZE_FOR_DETECTION):
 90                    uncompressed_size = disk_size * self.DEFAULT_COMPRESSION_RATIO
 91
 92        elif self.path.endswith(".zip"):
 93            with ZipFile(self.path) as zip_archive:
 94                for f in zip_archive.infolist():
 95                    uncompressed_size += f.file_size
 96
 97        else:
 98            uncompressed_size = os.path.getsize(self.path)
 99
100        self._size = uncompressed_size
101
102    @property
103    def is_compressed(self) -> bool:
104        return self.path.endswith(".gz") or self.path.endswith(".zip")
105
106    def open(self, should_compress: bool) -> BytesIO:
107        if should_compress:
108            file_stream = self.compress_stream()
109        else:
110            file_stream = open(self.path, "rb")
111        return file_stream
112
113    def compress_stream(self) -> BytesIO:
114        self.stream_name += ".gz"
115        file_stream = BytesIO()
116        with open(self.path, "rb") as f_in, GzipFile(filename="data", fileobj=file_stream, mode="wb") as f_out:
117            shutil.copyfileobj(f_in, f_out)
118        file_stream.seek(0)
119        return file_stream
120
121    def get_tracing_attributes(self) -> dict:
122        return {self._FILE_PATH: self.stream_name, self._SOURCE_ID: str(self.source_id)}
123
124    @classmethod
125    def get_instance(cls, file_descriptor: Union["FileDescriptor", str]) -> "FileDescriptor":
126        if not isinstance(file_descriptor, cls):
127            return cls(file_descriptor)
128        return file_descriptor

160class StreamDescriptor(DescriptorBase):
161    """StreamDescriptor is used to describe a stream that will be used as ingestion source"""
162
163    _STREAM_NAME = "stream_name"
164
165    # TODO: currently we always assume that streams are gz compressed (will get compressed before sending), should we expand that?
166    def __init__(
167        self, stream: IO[AnyStr], source_id: OptionalUUID = None, is_compressed: bool = False, stream_name: Optional[str] = None, size: Optional[int] = None
168    ):
169        """
170        :param stream: in-memory stream object.
171        :type stream: io.BaseIO
172        :param source_id: a v4 uuid to serve as the sources id.
173        :type source_id: OptionalUUID
174        :param is_compressed: specify if the provided stream is compressed
175        :type is_compressed: boolean
176        """
177        self.stream: IO[AnyStr] = stream
178        self.source_id: uuid.UUID = ensure_uuid(source_id)
179        self.is_compressed: bool = is_compressed
180        self.stream_name: str = stream_name
181        if self.stream_name is None:
182            self.stream_name = "stream"
183            if is_compressed:
184                self.stream_name += ".gz"
185        self.size: Optional[int] = size
186
187    def compress_stream(self) -> None:
188        stream = self.stream
189        zipped_stream = BytesIO()
190        stream_buffer = stream.read()
191        with GzipFile(filename="data", fileobj=zipped_stream, mode="wb") as f_out:
192            if isinstance(stream_buffer, str):
193                data = bytes(stream_buffer, "utf-8")
194                f_out.write(data)
195            else:
196                f_out.write(stream_buffer)
197        zipped_stream.seek(0)
198        self.is_compressed = True
199        self.stream_name += ".gz"
200        self.stream = zipped_stream
201
202    @staticmethod
203    def from_file_descriptor(file_descriptor: Union[FileDescriptor, str]) -> "StreamDescriptor":
204        """
205        Transforms FileDescriptor instance into StreamDescriptor instance. Note that stream is open when instance is returned
206        :param Union[FileDescriptor, str] file_descriptor: File Descriptor instance
207        :return new StreamDescriptor instance
208        """
209        descriptor = FileDescriptor.get_instance(file_descriptor)
210        stream = open(descriptor.path, "rb")
211        is_compressed = descriptor.path.endswith(".gz") or descriptor.path.endswith(".zip")
212        stream_descriptor = StreamDescriptor(stream, descriptor.source_id, is_compressed, descriptor.stream_name, descriptor.size)
213        return stream_descriptor
214
215    @classmethod
216    def get_instance(cls, stream_descriptor: Union["StreamDescriptor", IO[AnyStr]]) -> "StreamDescriptor":
217        if not isinstance(stream_descriptor, cls):
218            descriptor = cls(stream_descriptor)
219        else:
220            descriptor = copy(stream_descriptor)
221        return descriptor
222
223    def get_tracing_attributes(self) -> dict:
224        return {self._STREAM_NAME: self.stream_name, self._SOURCE_ID: str(self.source_id)}

24class KustoMissingMappingError(KustoClientError):
25    """
26    Raised when provided a mapping kind without a mapping reference or column mapping.
27    """

 7class KustoMappingError(KustoClientError):
 8    """
 9    Raised when the provided mapping arguments are invalid.
10    """

40class KustoQueueError(KustoClientError):
41    """Raised when not succeeding to upload message to queue in all retries"""
42
43    def __init__(self):
44        message = "Failed to upload message to queues in all reties."
45        super(KustoQueueError, self).__init__(message)

13class KustoDuplicateMappingError(KustoClientError):
14    """
15    Raised when ingestion properties include both
16    column mappings and a mapping reference
17    """
18
19    def __init__(self):
20        message = "Ingestion properties can't contain both an explicit mapping and a mapping reference."
21        super(KustoDuplicateMappingError, self).__init__(message)

30class KustoInvalidEndpointError(KustoClientError):
31    """Raised when trying to ingest to invalid cluster type."""
32
33    def __init__(self, expected_service_type, actual_service_type, suggested_endpoint_url=None):
34        message = f"You are using '{expected_service_type}' client type, but the provided endpoint is of ServiceType '{actual_service_type}'. Initialize the client with the appropriate endpoint URI"
35        if suggested_endpoint_url:
36            message = message + ": '" + suggested_endpoint_url + "'"
37        super(KustoInvalidEndpointError, self).__init__(message)

148class KustoClientError(KustoError):
149    """Raised when a Kusto client is unable to send or complete a request."""

 26class QueuedIngestClient(BaseIngestClient):
 27    """
 28    Queued ingest client provides methods to allow queued ingestion into kusto (ADX).
 29    To learn more about the different types of ingestions and when to use each, visit:
 30    https://docs.microsoft.com/en-us/azure/data-explorer/ingest-data-overview#ingestion-methods
 31    """
 32
 33    _INGEST_PREFIX = "ingest-"
 34    _SERVICE_CLIENT_TIMEOUT_SECONDS = 10 * 60
 35    _MAX_RETRIES = 3
 36
 37    def __init__(self, kcsb: Union[str, KustoConnectionStringBuilder], auto_correct_endpoint: bool = True):
 38        """Kusto Ingest Client constructor.
 39        :param kcsb: The connection string to initialize KustoClient.
 40        """
 41        super().__init__()
 42        if not isinstance(kcsb, KustoConnectionStringBuilder):
 43            kcsb = KustoConnectionStringBuilder(kcsb)
 44
 45        if auto_correct_endpoint:
 46            kcsb["Data Source"] = BaseIngestClient.get_ingestion_endpoint(kcsb.data_source)
 47
 48        self._proxy_dict: Optional[Dict[str, str]] = None
 49        self._connection_datasource = kcsb.data_source
 50        self._resource_manager = _ResourceManager(KustoClient(kcsb))
 51        self._endpoint_service_type = None
 52        self._suggested_endpoint_uri = None
 53        self.application_for_tracing = kcsb.client_details.application_for_tracing
 54        self.client_version_for_tracing = kcsb.client_details.version_for_tracing
 55
 56    def close(self) -> None:
 57        self._resource_manager.close()
 58        super().close()
 59
 60    def set_proxy(self, proxy_url: str):
 61        self._resource_manager.set_proxy(proxy_url)
 62        self._proxy_dict = {"http": proxy_url, "https": proxy_url}
 63
 64    @distributed_trace(name_of_span="QueuedIngestClient.ingest_from_file", kind=SpanKind.CLIENT)
 65    def ingest_from_file(self, file_descriptor: Union[FileDescriptor, str], ingestion_properties: IngestionProperties) -> IngestionResult:
 66        """Enqueue an ingest command from local files.
 67        To learn more about ingestion methods go to:
 68        https://docs.microsoft.com/en-us/azure/data-explorer/ingest-data-overview#ingestion-methods
 69        :param file_descriptor: a FileDescriptor to be ingested.
 70        :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties.
 71        """
 72        file_descriptor = FileDescriptor.get_instance(file_descriptor)
 73        IngestTracingAttributes.set_ingest_descriptor_attributes(file_descriptor, ingestion_properties)
 74
 75        super().ingest_from_file(file_descriptor, ingestion_properties)
 76
 77        containers = self._get_containers()
 78
 79        file_descriptor, should_compress = BaseIngestClient._prepare_file(file_descriptor, ingestion_properties)
 80        with file_descriptor.open(should_compress) as stream:
 81            blob_descriptor = self.upload_blob(
 82                containers,
 83                file_descriptor,
 84                ingestion_properties.database,
 85                ingestion_properties.table,
 86                stream,
 87                self._proxy_dict,
 88                self._SERVICE_CLIENT_TIMEOUT_SECONDS,
 89                self._MAX_RETRIES,
 90            )
 91        return self.ingest_from_blob(blob_descriptor, ingestion_properties=ingestion_properties)
 92
 93    @distributed_trace(name_of_span="QueuedIngestClient.ingest_from_stream", kind=SpanKind.CLIENT)
 94    def ingest_from_stream(self, stream_descriptor: Union[StreamDescriptor, IO[AnyStr]], ingestion_properties: IngestionProperties) -> IngestionResult:
 95        """Ingest from io streams.
 96        :param stream_descriptor: An object that contains a description of the stream to be ingested.
 97        :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties.
 98        """
 99        stream_descriptor = StreamDescriptor.get_instance(stream_descriptor)
100        IngestTracingAttributes.set_ingest_descriptor_attributes(stream_descriptor, ingestion_properties)
101
102        super().ingest_from_stream(stream_descriptor, ingestion_properties)
103
104        containers = self._get_containers()
105
106        stream_descriptor = BaseIngestClient._prepare_stream(stream_descriptor, ingestion_properties)
107        blob_descriptor = self.upload_blob(
108            containers,
109            stream_descriptor,
110            ingestion_properties.database,
111            ingestion_properties.table,
112            stream_descriptor.stream,
113            self._proxy_dict,
114            self._SERVICE_CLIENT_TIMEOUT_SECONDS,
115            self._MAX_RETRIES,
116        )
117        return self.ingest_from_blob(blob_descriptor, ingestion_properties=ingestion_properties)
118
119    @distributed_trace(name_of_span="QueuedIngestClient.ingest_from_blob", kind=SpanKind.CLIENT)
120    def ingest_from_blob(self, blob_descriptor: BlobDescriptor, ingestion_properties: IngestionProperties) -> IngestionResult:
121        """Enqueue an ingest command from azure blobs.
122        To learn more about ingestion methods go to:
123        https://docs.microsoft.com/en-us/azure/data-explorer/ingest-data-overview#ingestion-methods
124        :param azure.kusto.ingest.BlobDescriptor blob_descriptor: An object that contains a description of the blob to be ingested.
125        :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties.
126        """
127        IngestTracingAttributes.set_ingest_descriptor_attributes(blob_descriptor, ingestion_properties)
128
129        if self._is_closed:
130            raise KustoClosedError()
131
132        queues = self._resource_manager.get_ingestion_queues()
133
134        authorization_context = self._resource_manager.get_authorization_context()
135        ingestion_blob_info = IngestionBlobInfo(
136            blob_descriptor,
137            ingestion_properties=ingestion_properties,
138            auth_context=authorization_context,
139            application_for_tracing=self.application_for_tracing,
140            client_version_for_tracing=self.client_version_for_tracing,
141        )
142        ingestion_blob_info_json = ingestion_blob_info.to_json()
143        retries_left = min(self._MAX_RETRIES, len(queues))
144        for queue in queues:
145            try:
146                with QueueServiceClient(queue.account_uri, proxies=self._proxy_dict) as queue_service:
147                    with queue_service.get_queue_client(queue=queue.object_name, message_encode_policy=TextBase64EncodePolicy()) as queue_client:
148                        # trace enqueuing of blob for ingestion
149                        invoker = lambda: queue_client.send_message(content=ingestion_blob_info_json, timeout=self._SERVICE_CLIENT_TIMEOUT_SECONDS)
150                        enqueue_trace_attributes = IngestTracingAttributes.create_enqueue_request_attributes(queue_client.queue_name, blob_descriptor.source_id)
151                        MonitoredActivity.invoke(invoker, name_of_span="QueuedIngestClient.enqueue_request", tracing_attributes=enqueue_trace_attributes)
152
153                self._resource_manager.report_resource_usage_result(queue.storage_account_name, True)
154                return IngestionResult(
155                    IngestionStatus.QUEUED, ingestion_properties.database, ingestion_properties.table, blob_descriptor.source_id, blob_descriptor.path
156                )
157            except Exception as e:
158                retries_left = retries_left - 1
159                # TODO: log the retry once we have a proper logging system
160                self._resource_manager.report_resource_usage_result(queue.storage_account_name, False)
161                if retries_left == 0:
162                    raise KustoQueueError() from e
163
164    def _get_containers(self) -> List[_ResourceUri]:
165        return self._resource_manager.get_containers()
166
167    def upload_blob(
168        self,
169        containers: List[_ResourceUri],
170        descriptor: Union[FileDescriptor, "StreamDescriptor"],
171        database: str,
172        table: str,
173        stream: IO[AnyStr],
174        proxy_dict: Optional[Dict[str, str]],
175        timeout: int,
176        max_retries: int,
177    ) -> "BlobDescriptor":
178        """
179        Uploads and transforms FileDescriptor or StreamDescriptor into a BlobDescriptor instance
180        :param List[_ResourceUri] containers: blob containers
181        :param Union[FileDescriptor, "StreamDescriptor"] descriptor:
182        :param string database: database to be ingested to
183        :param string table: table to be ingested to
184        :param IO[AnyStr] stream: stream to be ingested from
185        :param Optional[Dict[str, str]] proxy_dict: proxy urls
186        :param int timeout: Azure service call timeout in seconds
187        :return new BlobDescriptor instance
188        """
189        blob_name = "{db}__{table}__{guid}__{file}".format(db=database, table=table, guid=descriptor.source_id, file=descriptor.stream_name)
190
191        retries_left = min(max_retries, len(containers))
192        for container in containers:
193            try:
194                blob_service = BlobServiceClient(container.account_uri, proxies=proxy_dict)
195                blob_client = blob_service.get_blob_client(container=container.object_name, blob=blob_name)
196                blob_client.upload_blob(data=stream, timeout=timeout)
197                self._resource_manager.report_resource_usage_result(container.storage_account_name, True)
198                return BlobDescriptor(blob_client.url, descriptor.size, descriptor.source_id)
199            except Exception as e:
200                retries_left = retries_left - 1
201                # TODO: log the retry once we have a proper logging system
202                self._resource_manager.report_resource_usage_result(container.storage_account_name, False)
203                if retries_left == 0:
204                    raise KustoBlobError(e)

26class ValidationPolicy:
27    """Validation policy to ingest command."""
28
29    def __init__(self, validation_options=ValidationOptions.DoNotValidate, validation_implications=ValidationImplications.BestEffort):
30        self.ValidationOptions = validation_options
31        self.ValidationImplications = validation_implications

19class ValidationImplications(IntEnum):
20    """Validation implications to ingest command."""
21
22    Fail = 0
23    BestEffort = 1

11class ValidationOptions(IntEnum):
12    """Validation options to ingest command."""
13
14    DoNotValidate = 0
15    ValidateCsvInputConstantColumns = 1
16    ValidateCsvInputColumnLevelOnly = 2

34class ReportLevel(IntEnum):
35    """Report level to ingest command."""
36
37    FailuresOnly = 0
38    DoNotReport = 1
39    FailuresAndSuccesses = 2