From 13eb147af38db8e508540764a72aec39693af988 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Bedn=C3=A1=C5=99?= Date: Mon, 12 Aug 2024 11:46:47 +0200 Subject: [PATCH] fix: import PolarsDataframeSerializer (#105) * fix: import PolarsDataframeSerializer * docs: update CHANGELOG.md * fix: import PolarsDataframeSerializer --- CHANGELOG.md | 6 +++- .../write_client/client/_base.py | 3 +- ...dataframe_serializer.py => test_polars.py} | 36 +++++++++++++++++-- 3 files changed, 41 insertions(+), 4 deletions(-) rename tests/{test_polars_dataframe_serializer.py => test_polars.py} (53%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3fa4345..57918e1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,10 +2,14 @@ ## 0.8.0 [unreleased] +### Features + +1. [#101](https://github.com/InfluxCommunity/influxdb3-python/pull/101): Add support for InfluxDB Edge (OSS) authentication + ### Bug Fixes 1. [#100](https://github.com/InfluxCommunity/influxdb3-python/pull/100): InfluxDB Edge (OSS) error handling -1. [#101](https://github.com/InfluxCommunity/influxdb3-python/pull/101): Add support for InfluxDB Edge authentication +1. [#105](https://github.com/InfluxCommunity/influxdb3-python/pull/105): Importing Polars serialization module ## 0.7.0 [2024-07-11] diff --git a/influxdb_client_3/write_client/client/_base.py b/influxdb_client_3/write_client/client/_base.py index d1b0408..32664d9 100644 --- a/influxdb_client_3/write_client/client/_base.py +++ b/influxdb_client_3/write_client/client/_base.py @@ -248,7 +248,8 @@ def _serialize(self, record, write_precision, payload, **kwargs): self._serialize(Point.from_dict(record, write_precision=write_precision, **kwargs), write_precision, payload, **kwargs) elif 'polars' in str(type(record)): - from influxdb_client_3.write_client.client.write.dataframe_serializer import PolarsDataframeSerializer + from influxdb_client_3.write_client.client.write.polars_dataframe_serializer import \ + PolarsDataframeSerializer serializer = PolarsDataframeSerializer(record, self._point_settings, write_precision, **kwargs) self._serialize(serializer.serialize(), write_precision, payload, **kwargs) diff --git a/tests/test_polars_dataframe_serializer.py b/tests/test_polars.py similarity index 53% rename from tests/test_polars_dataframe_serializer.py rename to tests/test_polars.py index 0157f05..dabde23 100644 --- a/tests/test_polars_dataframe_serializer.py +++ b/tests/test_polars.py @@ -1,7 +1,9 @@ -import unittest import importlib.util +import unittest +from unittest.mock import Mock -from influxdb_client_3 import PointSettings +from influxdb_client_3 import PointSettings, InfluxDBClient3 +from influxdb_client_3.write_client import WriteService from influxdb_client_3.write_client.client.write.polars_dataframe_serializer import polars_data_frame_to_list_of_points @@ -29,3 +31,33 @@ def test_to_list_of_points(self): 'iot-devices,building=5a name="iot-devices",temperature=72.2 1664798460000000000' ] self.assertEqual(expected, actual) + + +@unittest.skipIf(importlib.util.find_spec("polars") is None, 'Polars package not installed') +class TestWritePolars(unittest.TestCase): + def setUp(self): + self.client = InfluxDBClient3( + host="localhost", + org="my_org", + database="my_db", + token="my_token" + ) + + def test_write_polars(self): + import polars as pl + df = pl.DataFrame({ + "time": pl.Series(["2024-08-01 00:00:00", "2024-08-01 01:00:00"]).str.to_datetime(time_unit='ns'), + "temperature": [22.4, 21.8], + }) + self.client._write_api._write_service = Mock(spec=WriteService) + + self.client.write( + database="database", + record=df, + data_frame_measurement_name="measurement", + data_frame_timestamp_column="time", + ) + + actual = self.client._write_api._write_service.post_write.call_args[1]['body'] + self.assertEqual(b'measurement temperature=22.4 1722470400000000000\n' + b'measurement temperature=21.8 1722474000000000000', actual)