cognite-extractor-utils 7.7.0__py3-none-any.whl → 7.8.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of cognite-extractor-utils might be problematic. Click here for more details.

Files changed (37) hide show
  1. cognite/examples/unstable/extractors/simple_extractor/config/config.yaml +3 -0
  2. cognite/examples/unstable/extractors/simple_extractor/config/connection_config.yaml +10 -0
  3. cognite/examples/unstable/extractors/simple_extractor/main.py +81 -0
  4. cognite/extractorutils/__init__.py +1 -1
  5. cognite/extractorutils/_inner_util.py +2 -2
  6. cognite/extractorutils/base.py +1 -1
  7. cognite/extractorutils/configtools/elements.py +4 -2
  8. cognite/extractorutils/configtools/loaders.py +18 -4
  9. cognite/extractorutils/exceptions.py +1 -1
  10. cognite/extractorutils/metrics.py +8 -6
  11. cognite/extractorutils/statestore/watermark.py +6 -3
  12. cognite/extractorutils/threading.py +2 -2
  13. cognite/extractorutils/unstable/configuration/exceptions.py +28 -1
  14. cognite/extractorutils/unstable/configuration/models.py +157 -32
  15. cognite/extractorutils/unstable/core/_dto.py +80 -7
  16. cognite/extractorutils/unstable/core/base.py +171 -106
  17. cognite/extractorutils/unstable/core/checkin_worker.py +428 -0
  18. cognite/extractorutils/unstable/core/errors.py +2 -2
  19. cognite/extractorutils/unstable/core/logger.py +49 -0
  20. cognite/extractorutils/unstable/core/runtime.py +200 -31
  21. cognite/extractorutils/unstable/core/tasks.py +2 -2
  22. cognite/extractorutils/uploader/_base.py +1 -1
  23. cognite/extractorutils/uploader/assets.py +1 -1
  24. cognite/extractorutils/uploader/data_modeling.py +1 -1
  25. cognite/extractorutils/uploader/events.py +1 -1
  26. cognite/extractorutils/uploader/files.py +4 -4
  27. cognite/extractorutils/uploader/raw.py +1 -1
  28. cognite/extractorutils/uploader/time_series.py +4 -4
  29. cognite/extractorutils/uploader_extractor.py +2 -2
  30. cognite/extractorutils/uploader_types.py +3 -3
  31. cognite/extractorutils/util.py +8 -6
  32. {cognite_extractor_utils-7.7.0.dist-info → cognite_extractor_utils-7.8.1.dist-info}/METADATA +4 -3
  33. cognite_extractor_utils-7.8.1.dist-info/RECORD +55 -0
  34. cognite_extractor_utils-7.8.1.dist-info/entry_points.txt +2 -0
  35. cognite_extractor_utils-7.7.0.dist-info/RECORD +0 -50
  36. {cognite_extractor_utils-7.7.0.dist-info → cognite_extractor_utils-7.8.1.dist-info}/WHEEL +0 -0
  37. {cognite_extractor_utils-7.7.0.dist-info → cognite_extractor_utils-7.8.1.dist-info}/licenses/LICENSE +0 -0
@@ -28,9 +28,13 @@ import os
28
28
  import sys
29
29
  import time
30
30
  from argparse import ArgumentParser, Namespace
31
- from multiprocessing import Process, Queue
31
+ from dataclasses import dataclass
32
+ from logging.handlers import NTEventLogHandler as WindowsEventHandler
33
+ from multiprocessing import Event, Process, Queue
34
+ from multiprocessing.synchronize import Event as MpEvent
32
35
  from pathlib import Path
33
36
  from random import randint
37
+ from threading import Thread
34
38
  from typing import Any, Generic, TypeVar
35
39
  from uuid import uuid4
36
40
 
@@ -44,24 +48,80 @@ from cognite.client.exceptions import (
44
48
  CogniteConnectionError,
45
49
  )
46
50
  from cognite.extractorutils.threading import CancellationToken
47
- from cognite.extractorutils.unstable.configuration.exceptions import InvalidConfigError
51
+ from cognite.extractorutils.unstable.configuration.exceptions import InvalidArgumentError, InvalidConfigError
48
52
  from cognite.extractorutils.unstable.configuration.loaders import (
49
53
  load_file,
50
54
  load_from_cdf,
51
55
  )
52
- from cognite.extractorutils.unstable.configuration.models import ConnectionConfig
56
+ from cognite.extractorutils.unstable.configuration.models import ConnectionConfig, ExtractorConfig
53
57
  from cognite.extractorutils.unstable.core._dto import Error
58
+ from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker
54
59
  from cognite.extractorutils.unstable.core.errors import ErrorLevel
55
60
  from cognite.extractorutils.util import now
56
61
 
57
62
  from ._messaging import RuntimeMessage
58
- from .base import ConfigRevision, ConfigType, Extractor, FullConfig
63
+ from .base import ConfigRevision, Extractor, FullConfig
59
64
 
60
65
  __all__ = ["ExtractorType", "Runtime"]
61
66
 
62
67
  ExtractorType = TypeVar("ExtractorType", bound=Extractor)
63
68
 
64
69
 
70
+ @dataclass
71
+ class _RuntimeControls:
72
+ cancel_event: MpEvent
73
+ message_queue: Queue
74
+
75
+
76
+ def _extractor_process_entrypoint(
77
+ extractor_class: type[Extractor],
78
+ controls: _RuntimeControls,
79
+ config: FullConfig,
80
+ checkin_worker: CheckinWorker,
81
+ ) -> None:
82
+ logger = logging.getLogger(f"{extractor_class.EXTERNAL_ID}.runtime")
83
+ checkin_worker.active_revision = config.current_config_revision
84
+ checkin_worker.set_on_fatal_error_handler(lambda _: on_fatal_error(controls))
85
+ checkin_worker.set_on_revision_change_handler(lambda _: on_revision_changed(controls))
86
+ if config.application_config.retry_startup:
87
+ checkin_worker.set_retry_startup(config.application_config.retry_startup)
88
+ extractor = extractor_class._init_from_runtime(config, checkin_worker)
89
+ extractor._attach_runtime_controls(
90
+ cancel_event=controls.cancel_event,
91
+ message_queue=controls.message_queue,
92
+ )
93
+
94
+ try:
95
+ with extractor:
96
+ extractor.run()
97
+
98
+ except Exception:
99
+ logger.exception("Extractor crashed, will attempt restart")
100
+ controls.message_queue.put(RuntimeMessage.RESTART)
101
+
102
+
103
+ def on_revision_changed(controls: _RuntimeControls) -> None:
104
+ """
105
+ Handle a change in the configuration revision.
106
+
107
+ Args:
108
+ controls(_RuntimeControls): The runtime controls containing the message queue and cancellation event.
109
+ """
110
+ controls.message_queue.put(RuntimeMessage.RESTART)
111
+ controls.cancel_event.set()
112
+
113
+
114
+ def on_fatal_error(controls: _RuntimeControls) -> None:
115
+ """
116
+ Handle a fatal error in the extractor.
117
+
118
+ Args:
119
+ logger(logging.Logger): The logger to use for logging messages.
120
+ controls(_RuntimeControls): The runtime controls containing the message queue and cancellation event.
121
+ """
122
+ controls.cancel_event.set()
123
+
124
+
65
125
  class Runtime(Generic[ExtractorType]):
66
126
  """
67
127
  The runtime for an extractor.
@@ -82,6 +142,7 @@ class Runtime(Generic[ExtractorType]):
82
142
  self._message_queue: Queue[RuntimeMessage] = Queue()
83
143
  self.logger = logging.getLogger(f"{self._extractor_class.EXTERNAL_ID}.runtime")
84
144
  self._setup_logging()
145
+ self._cancel_event: MpEvent | None = None
85
146
 
86
147
  self._cognite_client: CogniteClient
87
148
 
@@ -113,11 +174,32 @@ class Runtime(Generic[ExtractorType]):
113
174
  default=None,
114
175
  help="Include to use a local application configuration instead of fetching it from CDF",
115
176
  )
177
+ argparser.add_argument(
178
+ "-l",
179
+ "--log-level",
180
+ choices=["debug", "info", "warning", "error", "critical"],
181
+ type=str,
182
+ required=False,
183
+ default=None,
184
+ help="Set the logging level for the runtime.",
185
+ )
116
186
  argparser.add_argument(
117
187
  "--skip-init-checks",
118
188
  action="store_true",
119
189
  help="Skip any checks during startup. Useful for debugging, not recommended for production deployments.",
120
190
  )
191
+ argparser.add_argument(
192
+ "--cwd",
193
+ nargs=1,
194
+ type=Path,
195
+ required=False,
196
+ help="Set the current working directory for the extractor.",
197
+ )
198
+ argparser.add_argument(
199
+ "--service",
200
+ action="store_true",
201
+ help="Run the extractor as a Windows service (only supported on Windows).",
202
+ )
121
203
 
122
204
  return argparser
123
205
 
@@ -138,31 +220,55 @@ class Runtime(Generic[ExtractorType]):
138
220
 
139
221
  root.addHandler(console_handler)
140
222
 
141
- def _inner_run(
142
- self,
143
- message_queue: Queue,
144
- config: FullConfig,
145
- ) -> None:
146
- # This code is run inside the new extractor process
147
- extractor = self._extractor_class._init_from_runtime(config)
148
- extractor._set_runtime_message_queue(message_queue)
223
+ if sys.platform == "win32":
224
+ try:
225
+ event_log_handler = WindowsEventHandler(self._extractor_class.NAME)
149
226
 
150
- try:
151
- with extractor:
152
- extractor.run()
227
+ event_log_handler.setLevel(logging.INFO)
228
+ root.addHandler(event_log_handler)
153
229
 
154
- except Exception:
155
- self.logger.exception("Extractor crashed, will attempt restart")
156
- message_queue.put(RuntimeMessage.RESTART)
230
+ self.logger.info("Windows Event Log handler enabled for startup.")
231
+ except ImportError:
232
+ self.logger.warning(
233
+ "Failed to import the 'pywin32' package. This should install automatically on windows. "
234
+ "Please try reinstalling to resolve this issue."
235
+ )
236
+ except Exception as e:
237
+ self.logger.warning(f"Failed to initialize Windows Event Log handler: {e}")
238
+
239
+ def _start_cancellation_watcher(self, mp_cancel_event: MpEvent) -> None:
240
+ """
241
+ Start the inter-process cancellation watcher thread.
242
+
243
+ This creates a daemon thread that waits for the runtime's CancellationToken
244
+ and sets the multiprocessing Event to signal the child process.
245
+ """
246
+
247
+ def cancellation_watcher() -> None:
248
+ """Waits for the runtime token and sets the shared event."""
249
+ self._cancellation_token.wait()
250
+ mp_cancel_event.set()
251
+
252
+ watcher_thread = Thread(target=cancellation_watcher, daemon=True, name="RuntimeCancelWatcher")
253
+ watcher_thread.start()
157
254
 
158
255
  def _spawn_extractor(
159
256
  self,
160
257
  config: FullConfig,
258
+ checkin_worker: CheckinWorker,
161
259
  ) -> Process:
162
- self._message_queue = Queue()
260
+ self._cancel_event = Event()
261
+
262
+ self._start_cancellation_watcher(self._cancel_event)
263
+
264
+ controls = _RuntimeControls(
265
+ cancel_event=self._cancel_event,
266
+ message_queue=self._message_queue,
267
+ )
268
+
163
269
  process = Process(
164
- target=self._inner_run,
165
- args=(self._message_queue, config),
270
+ target=_extractor_process_entrypoint,
271
+ args=(self._extractor_class, controls, config, checkin_worker),
166
272
  )
167
273
 
168
274
  process.start()
@@ -173,15 +279,15 @@ class Runtime(Generic[ExtractorType]):
173
279
  self,
174
280
  args: Namespace,
175
281
  connection_config: ConnectionConfig,
176
- ) -> tuple[ConfigType, ConfigRevision]:
282
+ ) -> tuple[ExtractorConfig, ConfigRevision]:
177
283
  current_config_revision: ConfigRevision
178
284
 
179
- if args.local_override:
285
+ if args.force_local_config:
180
286
  self.logger.info("Loading local application config")
181
287
 
182
288
  current_config_revision = "local"
183
289
  try:
184
- application_config = load_file(args.local_override[0], self._extractor_class.CONFIG_TYPE)
290
+ application_config = load_file(args.force_local_config[0], self._extractor_class.CONFIG_TYPE)
185
291
  except InvalidConfigError as e:
186
292
  self.logger.critical(str(e))
187
293
  raise e
@@ -194,17 +300,29 @@ class Runtime(Generic[ExtractorType]):
194
300
 
195
301
  application_config, current_config_revision = load_from_cdf(
196
302
  self._cognite_client,
197
- connection_config.integration,
303
+ connection_config.integration.external_id,
198
304
  self._extractor_class.CONFIG_TYPE,
199
305
  )
200
306
 
201
307
  return application_config, current_config_revision
202
308
 
309
+ def _try_set_cwd(self, args: Namespace) -> None:
310
+ if args.cwd is not None and len(args.cwd) > 0:
311
+ try:
312
+ resolved_path = Path(args.cwd[0]).resolve(strict=True)
313
+ os.chdir(resolved_path)
314
+ self.logger.info(f"Changed working directory to {resolved_path}")
315
+ except (FileNotFoundError, OSError) as e:
316
+ self.logger.critical(f"Could not change working directory to {args.cwd[0]}: {e}")
317
+ raise InvalidArgumentError(f"Could not change working directory to {args.cwd[0]}: {e}") from e
318
+
319
+ self.logger.info(f"Using {os.getcwd()} as working directory")
320
+
203
321
  def _safe_get_application_config(
204
322
  self,
205
323
  args: Namespace,
206
324
  connection_config: ConnectionConfig,
207
- ) -> tuple[ConfigType, ConfigRevision] | None:
325
+ ) -> tuple[ExtractorConfig, ConfigRevision] | None:
208
326
  prev_error: str | None = None
209
327
 
210
328
  while not self._cancellation_token.is_cancelled:
@@ -222,7 +340,7 @@ class Runtime(Generic[ExtractorType]):
222
340
  ts = now()
223
341
  error = Error(
224
342
  external_id=str(uuid4()),
225
- level=ErrorLevel.fatal.value,
343
+ level=ErrorLevel.fatal,
226
344
  start_time=ts,
227
345
  end_time=ts,
228
346
  description=error_message,
@@ -233,8 +351,8 @@ class Runtime(Generic[ExtractorType]):
233
351
  self._cognite_client.post(
234
352
  f"/api/v1/projects/{self._cognite_client.config.project}/odin/checkin",
235
353
  json={
236
- "externalId": connection_config.integration,
237
- "errors": [error.model_dump()],
354
+ "externalId": connection_config.integration.external_id,
355
+ "errors": [error.model_dump(mode="json")],
238
356
  },
239
357
  headers={"cdf-version": "alpha"},
240
358
  )
@@ -251,7 +369,7 @@ class Runtime(Generic[ExtractorType]):
251
369
  self._cognite_client.post(
252
370
  f"/api/v1/projects/{self._cognite_client.config.project}/odin/checkin",
253
371
  json={
254
- "externalId": connection_config.integration,
372
+ "externalId": connection_config.integration.external_id,
255
373
  },
256
374
  headers={"cdf-version": "alpha"},
257
375
  )
@@ -304,7 +422,47 @@ class Runtime(Generic[ExtractorType]):
304
422
 
305
423
  self.logger.info(f"Started runtime with PID {os.getpid()}")
306
424
 
425
+ if args.service and sys.platform == "win32":
426
+ # Import here to avoid dependency on non-Windows systems
427
+ try:
428
+ from simple_winservice import ( # type: ignore[import-not-found]
429
+ ServiceHandle,
430
+ register_service,
431
+ run_service,
432
+ )
433
+ except ImportError:
434
+ self.logger.critical("simple-winservice library is not installed.")
435
+ sys.exit(1)
436
+
437
+ # Cancellation function for the service
438
+ def cancel_service() -> None:
439
+ self.logger.info("Service cancellation requested.")
440
+ self._cancellation_token.cancel()
441
+
442
+ # Wrap the main runtime loop in a function for the service
443
+ def service_main(handle: ServiceHandle, service_args: list[str]) -> None:
444
+ handle.event_log_info("Extractor Windows service is starting.")
445
+ try:
446
+ self._main_runtime(args)
447
+ except Exception as exc:
448
+ handle.event_log_error(f"Service crashed: {exc}")
449
+ self.logger.critical(f"Service crashed: {exc}", exc_info=True)
450
+ sys.exit(1)
451
+ handle.event_log_info("Extractor Windows service is stopping.")
452
+
453
+ # Register and run the service
454
+ register_service(service_main, self._extractor_class.NAME, cancel_service)
455
+ run_service()
456
+ return
457
+ elif args.service and sys.platform != "win32":
458
+ self.logger.critical("--service is only supported on Windows.")
459
+ sys.exit(1)
460
+
461
+ self._main_runtime(args)
462
+
463
+ def _main_runtime(self, args: Namespace) -> None:
307
464
  try:
465
+ self._try_set_cwd(args)
308
466
  connection_config = load_file(args.connection_config[0], ConnectionConfig)
309
467
  except InvalidConfigError as e:
310
468
  self.logger.error(str(e))
@@ -318,6 +476,15 @@ class Runtime(Generic[ExtractorType]):
318
476
  # exist yet, and I have not found a way to represent it in a generic way that isn't just an Any in disguise.
319
477
  application_config: Any
320
478
  config: tuple[Any, ConfigRevision] | None
479
+ cognite_client = connection_config.get_cognite_client(
480
+ f"{self._extractor_class.EXTERNAL_ID}-{self._extractor_class.VERSION}"
481
+ )
482
+
483
+ checkin_worker = CheckinWorker(
484
+ cognite_client,
485
+ connection_config.integration.external_id,
486
+ self.logger,
487
+ )
321
488
 
322
489
  while not self._cancellation_token.is_cancelled:
323
490
  config = self._safe_get_application_config(args, connection_config)
@@ -334,7 +501,9 @@ class Runtime(Generic[ExtractorType]):
334
501
  connection_config=connection_config,
335
502
  application_config=application_config,
336
503
  current_config_revision=current_config_revision,
337
- )
504
+ log_level_override=args.log_level,
505
+ ),
506
+ checkin_worker,
338
507
  )
339
508
  process.join()
340
509
 
@@ -28,7 +28,7 @@ class TaskContext(CogniteLogger):
28
28
  This class is used to log errors and messages related to the task execution.
29
29
  """
30
30
 
31
- def __init__(self, task: "Task", extractor: "Extractor"):
31
+ def __init__(self, task: "Task", extractor: "Extractor") -> None:
32
32
  super().__init__()
33
33
  self._task = task
34
34
  self._extractor = extractor
@@ -87,7 +87,7 @@ class ScheduledTask(_Task):
87
87
  target: TaskTarget,
88
88
  description: str | None = None,
89
89
  schedule: ScheduleConfig,
90
- ):
90
+ ) -> None:
91
91
  super().__init__(name=name, target=target, description=description)
92
92
  self.schedule = schedule
93
93
 
@@ -50,7 +50,7 @@ class AbstractUploadQueue(ABC):
50
50
  trigger_log_level: str = "DEBUG",
51
51
  thread_name: str | None = None,
52
52
  cancellation_token: CancellationToken | None = None,
53
- ):
53
+ ) -> None:
54
54
  self.cdf_client = cdf_client
55
55
 
56
56
  self.threshold = max_queue_size if max_queue_size is not None else -1
@@ -64,7 +64,7 @@ class AssetUploadQueue(AbstractUploadQueue):
64
64
  trigger_log_level: str = "DEBUG",
65
65
  thread_name: str | None = None,
66
66
  cancellation_token: CancellationToken | None = None,
67
- ):
67
+ ) -> None:
68
68
  super().__init__(
69
69
  cdf_client,
70
70
  post_upload_function,
@@ -50,7 +50,7 @@ class InstanceUploadQueue(AbstractUploadQueue):
50
50
  auto_create_start_nodes: bool = True,
51
51
  auto_create_end_nodes: bool = True,
52
52
  auto_create_direct_relations: bool = True,
53
- ):
53
+ ) -> None:
54
54
  super().__init__(
55
55
  cdf_client,
56
56
  post_upload_function,
@@ -62,7 +62,7 @@ class EventUploadQueue(AbstractUploadQueue):
62
62
  trigger_log_level: str = "DEBUG",
63
63
  thread_name: str | None = None,
64
64
  cancellation_token: CancellationToken | None = None,
65
- ):
65
+ ) -> None:
66
66
  # Super sets post_upload and threshold
67
67
  super().__init__(
68
68
  cdf_client,
@@ -103,13 +103,13 @@ class ChunkedStream(RawIOBase, BinaryIO):
103
103
  # resolve the same way. These four useless methods with liberal use of Any are
104
104
  # required to satisfy mypy.
105
105
  # This may be solvable by changing the typing in the python SDK to use typing.Protocol.
106
- def writelines(self, __lines: Any) -> None:
106
+ def writelines(self, __lines: Any) -> None: # noqa: ANN401
107
107
  """
108
108
  Not supported for ChunkedStream.
109
109
  """
110
110
  raise NotImplementedError()
111
111
 
112
- def write(self, __b: Any) -> int:
112
+ def write(self, __b: Any) -> int: # noqa: ANN401
113
113
  """
114
114
  Not supported for ChunkedStream.
115
115
  """
@@ -250,7 +250,7 @@ class IOFileUploadQueue(AbstractUploadQueue):
250
250
  max_parallelism: int | None = None,
251
251
  failure_logging_path: None | str = None,
252
252
  ssl_verify: bool | str = True,
253
- ):
253
+ ) -> None:
254
254
  # Super sets post_upload and threshold
255
255
  super().__init__(
256
256
  cdf_client,
@@ -696,7 +696,7 @@ class FileUploadQueue(IOFileUploadQueue):
696
696
  overwrite_existing: bool = False,
697
697
  cancellation_token: CancellationToken | None = None,
698
698
  ssl_verify: bool | str = True,
699
- ):
699
+ ) -> None:
700
700
  # Super sets post_upload and threshold
701
701
  super().__init__(
702
702
  cdf_client=cdf_client,
@@ -67,7 +67,7 @@ class RawUploadQueue(AbstractUploadQueue):
67
67
  trigger_log_level: str = "DEBUG",
68
68
  thread_name: str | None = None,
69
69
  cancellation_token: CancellationToken | None = None,
70
- ):
70
+ ) -> None:
71
71
  # Super sets post_upload and thresholds
72
72
  super().__init__(
73
73
  cdf_client,
@@ -122,7 +122,7 @@ class BaseTimeSeriesUploadQueue(AbstractUploadQueue, Generic[IdType]):
122
122
  trigger_log_level: str = "DEBUG",
123
123
  thread_name: str | None = None,
124
124
  cancellation_token: CancellationToken | None = None,
125
- ):
125
+ ) -> None:
126
126
  # Super sets post_upload and threshold
127
127
  super().__init__(
128
128
  cdf_client,
@@ -248,7 +248,7 @@ class TimeSeriesUploadQueue(BaseTimeSeriesUploadQueue[EitherId]):
248
248
  create_missing: Callable[[str, DataPointList], TimeSeries] | bool = False,
249
249
  data_set_id: int | None = None,
250
250
  cancellation_token: CancellationToken | None = None,
251
- ):
251
+ ) -> None:
252
252
  # Super sets post_upload and threshold
253
253
  super().__init__(
254
254
  cdf_client,
@@ -429,7 +429,7 @@ class CDMTimeSeriesUploadQueue(BaseTimeSeriesUploadQueue[NodeId]):
429
429
  create_missing: Callable[[NodeId, DataPointList], CogniteExtractorTimeSeriesApply] | bool = False,
430
430
  cancellation_token: CancellationToken | None = None,
431
431
  source: DirectRelationReference | None = None,
432
- ):
432
+ ) -> None:
433
433
  super().__init__(
434
434
  cdf_client,
435
435
  post_upload_function,
@@ -636,7 +636,7 @@ class SequenceUploadQueue(AbstractUploadQueue):
636
636
  thread_name: str | None = None,
637
637
  create_missing: bool = False,
638
638
  cancellation_token: CancellationToken | None = None,
639
- ):
639
+ ) -> None:
640
640
  # Super sets post_upload and threshold
641
641
  super().__init__(
642
642
  cdf_client,
@@ -105,7 +105,7 @@ class UploaderExtractor(Extractor[UploaderExtractorConfigClass]):
105
105
  heartbeat_waiting_time: int = 600,
106
106
  handle_interrupts: bool = True,
107
107
  middleware: list[Callable[[dict], dict]] | None = None,
108
- ):
108
+ ) -> None:
109
109
  super().__init__(
110
110
  name=name,
111
111
  description=description,
@@ -165,7 +165,7 @@ class UploaderExtractor(Extractor[UploaderExtractorConfigClass]):
165
165
  else:
166
166
  raise ValueError(f"Unexpected type: {type(peek)}")
167
167
 
168
- def _apply_middleware(self, item: Any) -> Any:
168
+ def _apply_middleware(self, item: Any) -> Any: # noqa: ANN401
169
169
  for mw in self.middleware:
170
170
  item = mw(item)
171
171
  return item
@@ -18,7 +18,7 @@ class InsertDatapoints:
18
18
  A class representing a batch of datapoints to be inserted into a time series.
19
19
  """
20
20
 
21
- def __init__(self, *, id: int | None = None, external_id: str | None = None, datapoints: list[DataPoint]): # noqa: A002
21
+ def __init__(self, *, id: int | None = None, external_id: str | None = None, datapoints: list[DataPoint]) -> None: # noqa: A002
22
22
  self.id = id
23
23
  self.external_id = external_id
24
24
  self.datapoints = datapoints
@@ -29,7 +29,7 @@ class InsertCDMDatapoints:
29
29
  A class representing a batch of datapoints to be inserted into a cdm time series.
30
30
  """
31
31
 
32
- def __init__(self, *, instance_id: NodeId, datapoints: list[DataPoint]):
32
+ def __init__(self, *, instance_id: NodeId, datapoints: list[DataPoint]) -> None:
33
33
  self.instance_id = instance_id
34
34
  self.datapoints = datapoints
35
35
 
@@ -39,7 +39,7 @@ class RawRow:
39
39
  A class representing a row of data to be inserted into a RAW table.
40
40
  """
41
41
 
42
- def __init__(self, db_name: str, table_name: str, row: _Row | Iterable[_Row]):
42
+ def __init__(self, db_name: str, table_name: str, row: _Row | Iterable[_Row]) -> None:
43
43
  self.db_name = db_name
44
44
  self.table_name = table_name
45
45
  if isinstance(row, Iterable):
@@ -30,12 +30,14 @@ from typing import Any, TypeVar
30
30
  from decorator import decorator
31
31
 
32
32
  from cognite.client import CogniteClient
33
+ from cognite.client._api.assets import AssetsAPI
34
+ from cognite.client._api.time_series import TimeSeriesAPI
33
35
  from cognite.client.data_classes import Asset, ExtractionPipelineRun, TimeSeries
34
36
  from cognite.client.exceptions import CogniteAPIError, CogniteException, CogniteFileUploadError, CogniteNotFoundError
35
37
  from cognite.extractorutils.threading import CancellationToken
36
38
 
37
39
 
38
- def _ensure(endpoint: Any, items: Iterable[Any]) -> None:
40
+ def _ensure(endpoint: TimeSeriesAPI | AssetsAPI, items: Iterable[Any]) -> None:
39
41
  try:
40
42
  external_ids = [ts.external_id for ts in items]
41
43
 
@@ -90,7 +92,7 @@ class EitherId:
90
92
  TypeError: If none of both of id types are set.
91
93
  """
92
94
 
93
- def __init__(self, **kwargs: int | str | None):
95
+ def __init__(self, **kwargs: int | str | None) -> None:
94
96
  internal_id = kwargs.get("id")
95
97
  external_id = kwargs.get("externalId") or kwargs.get("external_id")
96
98
 
@@ -127,7 +129,7 @@ class EitherId:
127
129
  """
128
130
  return self.internal_id or self.external_id # type: ignore # checked to be not None in init
129
131
 
130
- def __eq__(self, other: Any) -> bool:
132
+ def __eq__(self, other: object) -> bool:
131
133
  """
132
134
  Compare with another object. Only returns true if other is an EitherId with the same type and content.
133
135
 
@@ -210,7 +212,7 @@ def add_extraction_pipeline(
210
212
 
211
213
  def decorator_ext_pip(input_function: Callable[..., _T1]) -> Callable[..., _T1]:
212
214
  @wraps(input_function)
213
- def wrapper_ext_pip(*args: Any, **kwargs: Any) -> _T1:
215
+ def wrapper_ext_pip(*args: Any, **kwargs: Any) -> _T1: # noqa: ANN401
214
216
  ##############################
215
217
  # Setup Extraction Pipelines #
216
218
  ##############################
@@ -397,7 +399,7 @@ def retry(
397
399
  """
398
400
 
399
401
  @decorator
400
- def retry_decorator(f: Callable[..., _T2], *fargs: Any, **fkwargs: Any) -> _T2:
402
+ def retry_decorator(f: Callable[..., _T2], *fargs: Any, **fkwargs: Any) -> _T2: # noqa: ANN401
401
403
  args = fargs if fargs else []
402
404
  kwargs = fkwargs if fkwargs else {}
403
405
 
@@ -657,7 +659,7 @@ def iterable_to_stream(
657
659
  def readable(self) -> bool:
658
660
  return True
659
661
 
660
- def readinto(self, buffer: Any) -> int | None:
662
+ def readinto(self, buffer: "WritableBuffer") -> int | None: # type: ignore[name-defined] # noqa: F821
661
663
  try:
662
664
  # Bytes to return
663
665
  ln = len(buffer)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cognite-extractor-utils
3
- Version: 7.7.0
3
+ Version: 7.8.1
4
4
  Summary: Utilities for easier development of extractors for CDF
5
5
  Project-URL: repository, https://github.com/cognitedata/python-extractor-utils
6
6
  Author-email: Mathias Lohne <mathias.lohne@cognite.com>
@@ -12,9 +12,9 @@ Requires-Python: >=3.10
12
12
  Requires-Dist: arrow>=1.0.0
13
13
  Requires-Dist: azure-identity>=1.14.0
14
14
  Requires-Dist: azure-keyvault-secrets>=4.7.0
15
- Requires-Dist: cognite-sdk>=7.59.0
15
+ Requires-Dist: cognite-sdk>=7.75.2
16
16
  Requires-Dist: croniter>=6.0.0
17
- Requires-Dist: dacite<1.9.0,>=1.6.0
17
+ Requires-Dist: dacite<1.10.0,>=1.9.2
18
18
  Requires-Dist: decorator>=5.1.1
19
19
  Requires-Dist: httpx<1,>=0.27.0
20
20
  Requires-Dist: jsonlines>=4.0.0
@@ -26,6 +26,7 @@ Requires-Dist: pydantic>=2.8.2
26
26
  Requires-Dist: pyhumps>=3.8.0
27
27
  Requires-Dist: python-dotenv>=1.0.0
28
28
  Requires-Dist: pyyaml<7,>=5.3.0
29
+ Requires-Dist: simple-winservice>=0.1.0; sys_platform == 'win32'
29
30
  Requires-Dist: typing-extensions<5,>=3.7.4
30
31
  Provides-Extra: experimental
31
32
  Requires-Dist: cognite-sdk-experimental; extra == 'experimental'