class KustoStreamingIngestClient(BaseIngestClient): """Kusto streaming ingest client for Python. KustoStreamingIngestClient works with both 2.x and 3.x flavors of Python. All primitive types are supported. Tests are run using pytest. """ def __init__(self, kcsb: Union[KustoConnectionStringBuilder, str]): """Kusto Streaming Ingest Client constructor. :param KustoConnectionStringBuilder kcsb: The connection string to initialize KustoClient. """ self._kusto_client = KustoClient(kcsb) def set_proxy(self, proxy_url: str): self._kusto_client.set_proxy(proxy_url) def ingest_from_file(self, file_descriptor: Union[FileDescriptor, str], ingestion_properties: IngestionProperties) -> IngestionResult: """Ingest from local files. :param file_descriptor: a FileDescriptor to be ingested. :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties. """ stream_descriptor = StreamDescriptor.from_file_descriptor(file_descriptor) with stream_descriptor.stream: return self.ingest_from_stream(stream_descriptor, ingestion_properties) def ingest_from_stream(self, stream_descriptor: Union[StreamDescriptor, IO[AnyStr]], ingestion_properties: IngestionProperties) -> IngestionResult: """Ingest from io streams. :param azure.kusto.ingest.StreamDescriptor stream_descriptor: An object that contains a description of the stream to be ingested. :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties. """ return self._ingest_from_stream_with_client_request_id(stream_descriptor, ingestion_properties, None) def _ingest_from_stream_with_client_request_id( self, stream_descriptor: Union[StreamDescriptor, IO[AnyStr]], ingestion_properties: IngestionProperties, client_request_id: Optional[str] ) -> IngestionResult: stream_descriptor = BaseIngestClient._prepare_stream(stream_descriptor, ingestion_properties) additional_properties = None if client_request_id: additional_properties = ClientRequestProperties() additional_properties.client_request_id = client_request_id self._kusto_client.execute_streaming_ingest( ingestion_properties.database, ingestion_properties.table, stream_descriptor.stream, ingestion_properties.format.name, additional_properties, mapping_name=ingestion_properties.ingestion_mapping_reference, ) return IngestionResult(IngestionStatus.SUCCESS, ingestion_properties.database, ingestion_properties.table, stream_descriptor.source_id)
class KustoStreamingIngestClient: """Kusto streaming ingest client for Python. KustoStreamingIngestClient works with both 2.x and 3.x flavors of Python. All primitive types are supported. Tests are run using pytest. """ _mapping_required_formats = { DataFormat.JSON, DataFormat.SINGLEJSON, DataFormat.AVRO, DataFormat.MULTIJSON } def __init__(self, kcsb: KustoConnectionStringBuilder): """Kusto Streaming Ingest Client constructor. :param KustoConnectionStringBuilder kcsb: The connection string to initialize KustoClient. """ self._kusto_client = KustoClient(kcsb) def ingest_from_dataframe(self, df: "pandas.DataFrame", ingestion_properties: IngestionProperties): """Ingest from pandas DataFrame. :param pandas.DataFrame df: input dataframe to ingest. :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties. """ from pandas import DataFrame if not isinstance(df, DataFrame): raise ValueError("Expected DataFrame instance, found {}".format( type(df))) file_name = "df_{timestamp}_{pid}.csv.gz".format(timestamp=int( time.time()), pid=os.getpid()) temp_file_path = os.path.join(tempfile.gettempdir(), file_name) df.to_csv(temp_file_path, index=False, encoding="utf-8", header=False, compression="gzip") ingestion_properties.format = DataFormat.CSV self.ingest_from_file(temp_file_path, ingestion_properties) os.unlink(temp_file_path) def ingest_from_file(self, file_descriptor: Union[FileDescriptor, str], ingestion_properties: IngestionProperties): """Ingest from local files. :param file_descriptor: a FileDescriptor to be ingested. :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties. """ if isinstance(file_descriptor, FileDescriptor): descriptor = file_descriptor else: descriptor = FileDescriptor(file_descriptor) stream = open(descriptor.path, "rb") is_compressed = descriptor.path.endswith( ".gz") or descriptor.path.endswith(".zip") stream_descriptor = StreamDescriptor(stream, descriptor.source_id, is_compressed) self.ingest_from_stream(stream_descriptor, ingestion_properties) if stream is not None: stream.close() def ingest_from_stream(self, stream_descriptor: Union[IO[AnyStr], StreamDescriptor], ingestion_properties: IngestionProperties): """Ingest from io streams. :param azure.kusto.ingest.StreamDescriptor stream_descriptor: An object that contains a description of the stream to be ingested. :param azure.kusto.ingest.IngestionProperties ingestion_properties: Ingestion properties. """ if not isinstance(stream_descriptor, StreamDescriptor): stream_descriptor = StreamDescriptor(stream_descriptor) if isinstance(stream_descriptor.stream, TextIOWrapper): stream = stream_descriptor.stream.buffer else: stream = stream_descriptor.stream if ingestion_properties.format in self._mapping_required_formats and ingestion_properties.ingestion_mapping_reference is None: raise KustoMissingMappingReferenceError() if not stream_descriptor.is_compressed: zipped_stream = BytesIO() buffer = stream.read() with GzipFile(filename="data", fileobj=zipped_stream, mode="wb") as f_out: if isinstance(buffer, str): data = bytes(buffer, "utf-8") f_out.write(data) else: f_out.write(buffer) zipped_stream.seek(0) stream = zipped_stream self._kusto_client.execute_streaming_ingest( ingestion_properties.database, ingestion_properties.table, stream, ingestion_properties.format.name, mapping_name=ingestion_properties.ingestion_mapping_reference, )