From 87b3ac02e1ca937010c3b2a872dd8d73e1b65929 Mon Sep 17 00:00:00 2001
From: Matias Holte <matias.holte@cognite.com>
Date: Mon, 11 Dec 2023 13:09:36 +0100
Subject: [PATCH] make subscription data calls more efficient (#1538)

---
 CHANGELOG.md                                    | 4 ++++
 cognite/client/_api/datapoints_subscriptions.py | 5 +++--
 cognite/client/_version.py                      | 2 +-
 pyproject.toml                                  | 2 +-
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 31b0e5d8d4..092c59b90d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,10 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [7.5.5] - 2023-12-07
+### Added
+- Added `poll_timeout` parameter on `time_series.subscriptions.iterate_data`. Will keep the connection open and waiting,
+  until new data is available, up to `poll_timeout` seconds.
 
 ## [7.5.4] - 2023-12-06
 ### Changed
diff --git a/cognite/client/_api/datapoints_subscriptions.py b/cognite/client/_api/datapoints_subscriptions.py
index 4eb602e6ff..444b62a711 100644
--- a/cognite/client/_api/datapoints_subscriptions.py
+++ b/cognite/client/_api/datapoints_subscriptions.py
@@ -211,6 +211,7 @@ def iterate_data(
         start: str | None = None,
         limit: int = DEFAULT_LIMIT_READ,
         partition: int = 0,
+        poll_timeout: int = 5,
         cursor: str | None = None,
     ) -> Iterator[DatapointSubscriptionBatch]:
         """`Iterate over data from a given subscription. <https://pr-2221.specs.preview.cogniteapp.com/20230101-beta.json.html#tag/Data-point-subscriptions/operation/listSubscriptionData>`_
@@ -227,6 +228,7 @@ def iterate_data(
             start (str | None): When to start the iteration. If set to None, the iteration will start from the beginning. The format is "N[timeunit]-ago", where timeunit is w,d,h,m (week, day, hour, minute). For example, "12h-ago" will start the iteration from 12 hours ago. You can also set it to "now" to jump straight to the end. Defaults to None.
             limit (int): Approximate number of results to return across all partitions.
             partition (int): The partition to iterate over. Defaults to 0.
+            poll_timeout (int): How many seconds to wait for new data, until an empty response is sent. Defaults to 5.
             cursor (str | None): Optional cursor to start iterating from.
 
         Yields:
@@ -252,8 +254,6 @@ def iterate_data(
                 ...     print(f"Added {len(batch.subscription_changes.added)} timeseries")
                 ...     print(f"Removed {len(batch.subscription_changes.removed)} timeseries")
                 ...     print(f"Changed timeseries data in {len(batch.updates)} updates")
-                ...     if not batch.has_next:
-                ...         time.sleep(1)
         """
         self._warning.warn()
 
@@ -263,6 +263,7 @@ def iterate_data(
                 "externalId": external_id,
                 "partitions": [p.dump(camel_case=True) for p in current_partitions],
                 "limit": limit,
+                "pollTimeoutSeconds": poll_timeout,
             }
             if start is not None:
                 body["initializeCursors"] = start
diff --git a/cognite/client/_version.py b/cognite/client/_version.py
index f9beaaba47..b92368a200 100644
--- a/cognite/client/_version.py
+++ b/cognite/client/_version.py
@@ -1,4 +1,4 @@
 from __future__ import annotations
 
-__version__ = "7.5.4"
+__version__ = "7.5.5"
 __api_subversion__ = "V20220125"
diff --git a/pyproject.toml b/pyproject.toml
index c91fa5bdac..44f97657fd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "cognite-sdk"
 
-version = "7.5.4"
+version = "7.5.5"
 description = "Cognite Python SDK"
 readme = "README.md"
 documentation = "https://cognite-sdk-python.readthedocs-hosted.com"