cognite-extractor-utils 5.0.1__py3-none-any.whl → 5.2.0__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.

@@ -16,5 +16,5 @@
16
16
  Cognite extractor utils is a Python package that simplifies the development of new extractors.
17
17
  """
18
18
 
19
- __version__ = "5.0.1"
19
+ __version__ = "5.2.0"
20
20
  from .base import Extractor
@@ -18,6 +18,7 @@ A module containing utilities meant for use inside the extractor-utils package
18
18
 
19
19
  import json
20
20
  from decimal import Decimal
21
+ from typing import Any, Dict, Union
21
22
 
22
23
 
23
24
  def _resolve_log_level(level: str) -> int:
@@ -25,17 +26,17 @@ def _resolve_log_level(level: str) -> int:
25
26
 
26
27
 
27
28
  class _DecimalEncoder(json.JSONEncoder):
28
- def default(self, obj):
29
+ def default(self, obj: Any) -> Dict[str, str]:
29
30
  if isinstance(obj, Decimal):
30
31
  return {"type": "decimal_encoded", "value": str(obj)}
31
32
  return super(_DecimalEncoder, self).default(obj)
32
33
 
33
34
 
34
35
  class _DecimalDecoder(json.JSONDecoder):
35
- def __init__(self, *args, **kwargs):
36
+ def __init__(self, *args: Any, **kwargs: Any) -> None:
36
37
  json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
37
38
 
38
- def object_hook(self, obj_dict):
39
+ def object_hook(self, obj_dict: Dict[str, str]) -> Union[Dict[str, str], Decimal]:
39
40
  if obj_dict.get("type") == "decimal_encoded":
40
41
  return Decimal(obj_dict["value"])
41
42
  return obj_dict
@@ -21,10 +21,10 @@ from threading import Event, Thread
21
21
  from types import TracebackType
22
22
  from typing import Any, Callable, Dict, Generic, Optional, Type, TypeVar
23
23
 
24
- from cognite.client import CogniteClient
25
- from cognite.client.data_classes import ExtractionPipeline, ExtractionPipelineRun
26
24
  from dotenv import find_dotenv, load_dotenv
27
25
 
26
+ from cognite.client import CogniteClient
27
+ from cognite.client.data_classes import ExtractionPipeline, ExtractionPipelineRun
28
28
  from cognite.extractorutils.configtools import BaseConfig, ConfigResolver, StateStoreConfig
29
29
  from cognite.extractorutils.exceptions import InvalidConfigError
30
30
  from cognite.extractorutils.metrics import BaseMetrics
@@ -94,7 +94,7 @@ class Extractor(Generic[CustomConfigClass]):
94
94
  self.run_handle = run_handle
95
95
  self.config_class = config_class
96
96
  self.use_default_state_store = use_default_state_store
97
- self.version = version
97
+ self.version = version or "unknown"
98
98
  self.cancellation_token = cancellation_token
99
99
  self.config_file_path = config_file_path
100
100
  self.continuous_extractor = continuous_extractor
@@ -117,7 +117,7 @@ class Extractor(Generic[CustomConfigClass]):
117
117
  if metrics:
118
118
  self.metrics = metrics
119
119
  else:
120
- self.metrics = BaseMetrics(extractor_name=name, extractor_version=version)
120
+ self.metrics = BaseMetrics(extractor_name=name, extractor_version=self.version)
121
121
 
122
122
  def _initial_load_config(self, override_path: Optional[str] = None) -> None:
123
123
  """
@@ -133,9 +133,9 @@ class Extractor(Generic[CustomConfigClass]):
133
133
  self.config_resolver = ConfigResolver.from_cli(self.name, self.description, self.version, self.config_class)
134
134
 
135
135
  self.config = self.config_resolver.config
136
- Extractor._config_singleton = self.config
136
+ Extractor._config_singleton = self.config # type: ignore
137
137
 
138
- def config_refresher():
138
+ def config_refresher() -> None:
139
139
  while not self.cancellation_token.is_set():
140
140
  self.cancellation_token.wait(self.reload_config_interval)
141
141
  if self.config_resolver.has_changed:
@@ -144,7 +144,7 @@ class Extractor(Generic[CustomConfigClass]):
144
144
  if self.reload_config_interval and self.reload_config_action != ReloadConfigAction.DO_NOTHING:
145
145
  Thread(target=config_refresher, name="ConfigReloader", daemon=True).start()
146
146
 
147
- def reload_config_callback(self):
147
+ def reload_config_callback(self) -> None:
148
148
  self.logger.error("Method for reloading configs has not been overridden in subclass")
149
149
 
150
150
  def _reload_config(self) -> None:
@@ -154,7 +154,7 @@ class Extractor(Generic[CustomConfigClass]):
154
154
  self.logger.info("Loading in new config file")
155
155
  self.config_resolver.accept_new_config()
156
156
  self.config = self.config_resolver.config
157
- Extractor._config_singleton = self.config
157
+ Extractor._config_singleton = self.config # type: ignore
158
158
 
159
159
  elif self.reload_config_action == ReloadConfigAction.SHUTDOWN:
160
160
  self.logger.info("Shutting down, expecting to be restarted")
@@ -193,7 +193,7 @@ class Extractor(Generic[CustomConfigClass]):
193
193
 
194
194
  try:
195
195
  self.state_store.initialize()
196
- except ValueError as e:
196
+ except ValueError:
197
197
  self.logger.exception("Could not load state store, using an empty state store as default")
198
198
 
199
199
  Extractor._statestore_singleton = self.state_store
@@ -212,18 +212,16 @@ class Extractor(Generic[CustomConfigClass]):
212
212
  )
213
213
  )
214
214
 
215
- def _report_error(self, exc_type: Type[BaseException], exc_val: BaseException, exc_tb: TracebackType) -> None:
215
+ def _report_error(self, exception: BaseException) -> None:
216
216
  """
217
217
  Called on an unsuccessful exit of the extractor
218
218
 
219
219
  Args:
220
- exc_type: Type of exception that caused the extractor to fail
221
- exc_val: Exception object that caused the extractor to fail
222
- exc_tb: Stack trace of where the extractor failed
220
+ exception: Exception object that caused the extractor to fail
223
221
  """
224
- self.logger.error("Unexpected error during extraction", exc_info=exc_val)
222
+ self.logger.error("Unexpected error during extraction", exc_info=exception)
225
223
  if self.extraction_pipeline:
226
- message = f"{exc_type.__name__}: {str(exc_val)}"[:1000]
224
+ message = f"{type(exception).__name__}: {str(exception)}"[:1000]
227
225
 
228
226
  self.logger.info(f"Reporting new failed run: {message}")
229
227
  self.cognite_client.extraction_pipelines.runs.create(
@@ -254,8 +252,8 @@ class Extractor(Generic[CustomConfigClass]):
254
252
  try:
255
253
  self._initial_load_config(override_path=self.config_file_path)
256
254
  except InvalidConfigError as e:
257
- print("Critical error: Could not read config file", file=sys.stderr)
258
- print(str(e), file=sys.stderr)
255
+ print("Critical error: Could not read config file", file=sys.stderr) # noqa: T201
256
+ print(str(e), file=sys.stderr) # noqa: T201
259
257
  sys.exit(1)
260
258
 
261
259
  if not self.configured_logger:
@@ -276,11 +274,11 @@ class Extractor(Generic[CustomConfigClass]):
276
274
  self.extraction_pipeline = self.config.cognite.get_extraction_pipeline(self.cognite_client)
277
275
 
278
276
  try:
279
- self.config.metrics.start_pushers(self.cognite_client)
277
+ self.config.metrics.start_pushers(self.cognite_client) # type: ignore
280
278
  except AttributeError:
281
279
  pass
282
280
 
283
- def heartbeat_loop():
281
+ def heartbeat_loop() -> None:
284
282
  while not self.cancellation_token.is_set():
285
283
  self.cancellation_token.wait(self.heartbeat_waiting_time)
286
284
 
@@ -289,7 +287,7 @@ class Extractor(Generic[CustomConfigClass]):
289
287
  try:
290
288
  self.cognite_client.extraction_pipelines.runs.create(
291
289
  ExtractionPipelineRun(
292
- extpipe_external_id=self.extraction_pipeline.external_id, status="seen"
290
+ extpipe_external_id=self.extraction_pipeline.external_id, status="seen" # type: ignore
293
291
  )
294
292
  )
295
293
  except Exception:
@@ -336,12 +334,12 @@ class Extractor(Generic[CustomConfigClass]):
336
334
  self.state_store.synchronize()
337
335
 
338
336
  try:
339
- self.config.metrics.stop_pushers()
337
+ self.config.metrics.stop_pushers() # type: ignore
340
338
  except AttributeError:
341
339
  pass
342
340
 
343
341
  if exc_val:
344
- self._report_error(exc_type, exc_val, exc_tb)
342
+ self._report_error(exc_val)
345
343
  else:
346
344
  self._report_success()
347
345
 
@@ -364,9 +362,9 @@ class Extractor(Generic[CustomConfigClass]):
364
362
 
365
363
  @classmethod
366
364
  def get_current_config(cls) -> CustomConfigClass:
367
- if Extractor._config_singleton is None:
365
+ if Extractor._config_singleton is None: # type: ignore
368
366
  raise ValueError("No config singleton created. Have a config file been loaded?")
369
- return Extractor._config_singleton
367
+ return Extractor._config_singleton # type: ignore
370
368
 
371
369
  @classmethod
372
370
  def get_current_statestore(cls) -> AbstractStateStore:
@@ -16,9 +16,9 @@
16
16
  Module containing tools for loading and verifying config files, and a YAML loader to automatically serialize these
17
17
  dataclasses from a config file.
18
18
 
19
- Configs are described as ``dataclass``\es, and use the ``BaseConfig`` class as a superclass to get a few things built-in:
20
- config version, Cognite project and logging. Use type hints to specify types, use the ``Optional`` type to specify that
21
- a config parameter is optional, and give the attribute a value to give it a default.
19
+ Configs are described as ``dataclass``\es, and use the ``BaseConfig`` class as a superclass to get a few things
20
+ built-in: config version, Cognite project and logging. Use type hints to specify types, use the ``Optional`` type to
21
+ specify that a config parameter is optional, and give the attribute a value to give it a default.
22
22
 
23
23
  For example, a config class for an extractor may look like the following:
24
24
 
@@ -14,7 +14,7 @@
14
14
  import base64
15
15
  import re
16
16
  from pathlib import Path
17
- from typing import Any, Dict, Optional, Tuple
17
+ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
18
18
 
19
19
  from cryptography.hazmat.primitives import hashes
20
20
  from cryptography.hazmat.primitives import serialization as serialization
@@ -37,11 +37,11 @@ def _to_snake_case(dictionary: Dict[str, Any], case_style: str) -> Dict[str, Any
37
37
  An updated dictionary with keys in the given convention.
38
38
  """
39
39
 
40
- def fix_list(list_, key_translator):
40
+ def fix_list(list_: List[Any], key_translator: Callable[[str], str]) -> List[Any]:
41
41
  if list_ is None:
42
42
  return []
43
43
 
44
- new_list = [None] * len(list_)
44
+ new_list: List[Any] = [None] * len(list_)
45
45
  for i, element in enumerate(list_):
46
46
  if isinstance(element, dict):
47
47
  new_list[i] = fix_dict(element, key_translator)
@@ -51,11 +51,11 @@ def _to_snake_case(dictionary: Dict[str, Any], case_style: str) -> Dict[str, Any
51
51
  new_list[i] = element
52
52
  return new_list
53
53
 
54
- def fix_dict(dict_, key_translator):
54
+ def fix_dict(dict_: Dict[str, Any], key_translator: Callable[[str], str]) -> Dict[str, Any]:
55
55
  if dict_ is None:
56
56
  return {}
57
57
 
58
- new_dict = {}
58
+ new_dict: Dict[str, Any] = {}
59
59
  for key in dict_:
60
60
  if isinstance(dict_[key], dict):
61
61
  new_dict[key_translator(key)] = fix_dict(dict_[key], key_translator)
@@ -65,10 +65,10 @@ def _to_snake_case(dictionary: Dict[str, Any], case_style: str) -> Dict[str, Any
65
65
  new_dict[key_translator(key)] = dict_[key]
66
66
  return new_dict
67
67
 
68
- def translate_hyphen(key):
68
+ def translate_hyphen(key: str) -> str:
69
69
  return key.replace("-", "_")
70
70
 
71
- def translate_camel(key):
71
+ def translate_camel(key: str) -> str:
72
72
  return re.sub(r"([A-Z]+)", r"_\1", key).strip("_").lower()
73
73
 
74
74
  if case_style == "snake" or case_style == "underscore":
@@ -81,7 +81,7 @@ def _to_snake_case(dictionary: Dict[str, Any], case_style: str) -> Dict[str, Any
81
81
  raise ValueError(f"Invalid case style: {case_style}")
82
82
 
83
83
 
84
- def _load_certificate_data(cert_path: str, password: Optional[str]) -> Tuple[str, str]:
84
+ def _load_certificate_data(cert_path: str, password: Optional[str]) -> Union[Tuple[str, str], Tuple[bytes, bytes]]:
85
85
  path = Path(cert_path)
86
86
  cert_data = Path(path).read_bytes()
87
87
 
@@ -93,16 +93,22 @@ def _load_certificate_data(cert_path: str, password: Optional[str]) -> Tuple[str
93
93
  format=serialization.PrivateFormat.TraditionalOpenSSL,
94
94
  encryption_algorithm=serialization.NoEncryption(),
95
95
  )
96
- return (base64.b16encode(cert.fingerprint(hashes.SHA1())), private_key_str)
96
+ return base64.b16encode(cert.fingerprint(hashes.SHA1())), private_key_str
97
97
  elif path.suffix == ".pfx":
98
- (private_key, cert, _) = pkcs12.load_key_and_certificates(
98
+ (private_key_pfx, cert_pfx, _) = pkcs12.load_key_and_certificates(
99
99
  cert_data, password=password.encode() if password else None
100
100
  )
101
- private_key_str = private_key.private_bytes(
101
+
102
+ if private_key_pfx is None:
103
+ raise InvalidConfigError(f"Can't load private key from {cert_path}")
104
+ if cert_pfx is None:
105
+ raise InvalidConfigError(f"Can't load certificate from {cert_path}")
106
+
107
+ private_key_str = private_key_pfx.private_bytes(
102
108
  encoding=serialization.Encoding.PEM,
103
109
  format=serialization.PrivateFormat.TraditionalOpenSSL,
104
110
  encryption_algorithm=serialization.NoEncryption(),
105
111
  )
106
- return (base64.b16encode(cert.fingerprint(hashes.SHA1())), private_key_str)
112
+ return base64.b16encode(cert_pfx.fingerprint(hashes.SHA1())), private_key_str
107
113
  else:
108
114
  raise InvalidConfigError(f"Unknown certificate format '{path.suffix}'. Allowed formats are 'pem' and 'pfx'")
@@ -20,15 +20,15 @@ from enum import Enum
20
20
  from logging.handlers import TimedRotatingFileHandler
21
21
  from threading import Event
22
22
  from time import sleep
23
- from typing import Dict, List, Optional, Tuple, Union
23
+ from typing import Any, Dict, List, Optional, Tuple, Union
24
24
  from urllib.parse import urljoin
25
25
 
26
26
  import yaml
27
- from cognite.client import ClientConfig, CogniteClient
28
- from cognite.client.credentials import OAuthClientCertificate, OAuthClientCredentials
29
- from cognite.client.data_classes import Asset, DataSet, ExtractionPipeline
30
27
  from prometheus_client import REGISTRY, start_http_server
31
28
 
29
+ from cognite.client import ClientConfig, CogniteClient
30
+ from cognite.client.credentials import CredentialProvider, OAuthClientCertificate, OAuthClientCredentials
31
+ from cognite.client.data_classes import Asset, DataSet, ExtractionPipeline
32
32
  from cognite.extractorutils.configtools._util import _load_certificate_data
33
33
  from cognite.extractorutils.exceptions import InvalidConfigError
34
34
  from cognite.extractorutils.metrics import AbstractMetricsPusher, CognitePusher, PrometheusPusher
@@ -63,6 +63,22 @@ class AuthenticatorConfig:
63
63
  certificate: Optional[CertificateConfig] = None
64
64
 
65
65
 
66
+ @dataclass
67
+ class ConnectionConfig:
68
+ """
69
+ Configuration parameters for the global_config python SDK settings
70
+ """
71
+
72
+ disable_gzip: bool = False
73
+ status_forcelist: List[int] = field(default_factory=lambda: [429, 502, 503, 504])
74
+ max_retries: int = 10
75
+ max_retries_connect: int = 3
76
+ max_retry_backoff: int = 30
77
+ max_connection_pool_size: int = 50
78
+ disable_ssl: bool = False
79
+ proxies: Dict[str, str] = field(default_factory=dict)
80
+
81
+
66
82
  @dataclass
67
83
  class EitherIdConfig:
68
84
  id: Optional[int]
@@ -74,7 +90,7 @@ class EitherIdConfig:
74
90
 
75
91
 
76
92
  class TimeIntervalConfig(yaml.YAMLObject):
77
- def __init__(self, expression):
93
+ def __init__(self, expression: str) -> None:
78
94
  self._interval, self._expression = TimeIntervalConfig._parse_expression(expression)
79
95
 
80
96
  @classmethod
@@ -130,7 +146,7 @@ class TimeIntervalConfig(yaml.YAMLObject):
130
146
 
131
147
 
132
148
  class FileSizeConfig(yaml.YAMLObject):
133
- def __init__(self, expression):
149
+ def __init__(self, expression: str) -> None:
134
150
  self._bytes, self._expression = FileSizeConfig._parse_expression(expression)
135
151
 
136
152
  @classmethod
@@ -220,16 +236,26 @@ class CogniteConfig:
220
236
  data_set_external_id: Optional[str]
221
237
  extraction_pipeline: Optional[EitherIdConfig]
222
238
  timeout: TimeIntervalConfig = TimeIntervalConfig("30s")
239
+ connection: ConnectionConfig = field(default_factory=ConnectionConfig)
223
240
  external_id_prefix: str = ""
224
241
  host: str = "https://api.cognitedata.com"
225
242
 
226
243
  def get_cognite_client(
227
- self, client_name: str, token_custom_args: Optional[Dict[str, str]] = None, use_experimental_sdk=False
244
+ self, client_name: str, token_custom_args: Optional[Dict[str, str]] = None, use_experimental_sdk: bool = False
228
245
  ) -> CogniteClient:
229
246
  from cognite.client.config import global_config
230
247
 
231
248
  global_config.disable_pypi_version_check = True
232
-
249
+ global_config.disable_gzip = self.connection.disable_gzip
250
+ global_config.status_forcelist = set(self.connection.status_forcelist)
251
+ global_config.max_retries = self.connection.max_retries
252
+ global_config.max_retries_connect = self.connection.max_retries_connect
253
+ global_config.max_retry_backoff = self.connection.max_retry_backoff
254
+ global_config.max_connection_pool_size = self.connection.max_connection_pool_size
255
+ global_config.disable_ssl = self.connection.disable_ssl
256
+ global_config.proxies = self.connection.proxies
257
+
258
+ credential_provider: CredentialProvider
233
259
  if self.idp_authentication.certificate:
234
260
  if self.idp_authentication.certificate.authority_url:
235
261
  authority_url = self.idp_authentication.certificate.authority_url
@@ -243,13 +269,13 @@ class CogniteConfig:
243
269
  credential_provider = OAuthClientCertificate(
244
270
  authority_url=authority_url,
245
271
  client_id=self.idp_authentication.client_id,
246
- cert_thumbprint=thumprint,
247
- certificate=key,
272
+ cert_thumbprint=str(thumprint),
273
+ certificate=str(key),
248
274
  scopes=self.idp_authentication.scopes,
249
275
  )
250
276
 
251
277
  elif self.idp_authentication.secret:
252
- kwargs = {}
278
+ kwargs: Dict[str, Any] = {}
253
279
  if self.idp_authentication.token_url:
254
280
  kwargs["token_url"] = self.idp_authentication.token_url
255
281
  elif self.idp_authentication.tenant:
@@ -264,7 +290,7 @@ class CogniteConfig:
264
290
  token_custom_args["resource"] = self.idp_authentication.resource
265
291
  if self.idp_authentication.audience:
266
292
  token_custom_args["audience"] = self.idp_authentication.audience
267
- credential_provider = OAuthClientCredentials(**kwargs, **token_custom_args)
293
+ credential_provider = OAuthClientCredentials(**kwargs, **token_custom_args) # type: ignore
268
294
 
269
295
  else:
270
296
  raise InvalidConfigError("No client certificate or secret provided")
@@ -278,7 +304,7 @@ class CogniteConfig:
278
304
  )
279
305
 
280
306
  if use_experimental_sdk:
281
- from cognite.experimental import CogniteClient as ExperimentalCogniteClient
307
+ from cognite.experimental import CogniteClient as ExperimentalCogniteClient # type: ignore
282
308
 
283
309
  return ExperimentalCogniteClient(client_config)
284
310
 
@@ -299,7 +325,7 @@ class CogniteConfig:
299
325
  return None
300
326
 
301
327
  return cdf_client.data_sets.retrieve(
302
- id=self.data_set.either_id.internal_id, external_id=self.data_set.either_id.external_id
328
+ id=self.data_set.either_id.internal_id, external_id=self.data_set.either_id.external_id # type: ignore
303
329
  )
304
330
 
305
331
  def get_extraction_pipeline(self, cdf_client: CogniteClient) -> Optional[ExtractionPipeline]:
@@ -308,8 +334,8 @@ class CogniteConfig:
308
334
 
309
335
  either_id = self.extraction_pipeline.either_id
310
336
  extraction_pipeline = cdf_client.extraction_pipelines.retrieve(
311
- id=either_id.internal_id,
312
- external_id=either_id.external_id,
337
+ id=either_id.internal_id, # type: ignore
338
+ external_id=either_id.external_id, # type: ignore
313
339
  )
314
340
  if extraction_pipeline is None:
315
341
  raise ValueError(f"Extraction pipeline with {either_id.type()} {either_id.content()} not found")
@@ -341,7 +367,7 @@ class LoggingConfig:
341
367
  # Prometheus and/or Cognite
342
368
  metrics: Optional[bool] = False
343
369
 
344
- def setup_logging(self, suppress_console=False) -> None:
370
+ def setup_logging(self, suppress_console: bool = False) -> None:
345
371
  """
346
372
  Sets up the default logger in the logging package to be configured as defined in this config object
347
373
 
@@ -430,6 +456,7 @@ class MetricsConfig:
430
456
 
431
457
  push_gateways = self.push_gateways or []
432
458
 
459
+ pusher: AbstractMetricsPusher
433
460
  for counter, push_gateway in enumerate(push_gateways):
434
461
  pusher = PrometheusPusher(
435
462
  job_name=push_gateway.job_name,
@@ -449,7 +476,7 @@ class MetricsConfig:
449
476
  if self.cognite:
450
477
  asset = None
451
478
 
452
- if self.cognite.asset_name is not None:
479
+ if self.cognite.asset_name is not None and self.cognite.asset_external_id:
453
480
  asset = Asset(name=self.cognite.asset_name, external_id=self.cognite.asset_external_id)
454
481
 
455
482
  pusher = CognitePusher(
@@ -40,10 +40,10 @@ def _load_yaml(
40
40
  source: Union[TextIO, str],
41
41
  config_type: Type[CustomConfigClass],
42
42
  case_style: str = "hyphen",
43
- expand_envvars=True,
43
+ expand_envvars: bool = True,
44
44
  dict_manipulator: Callable[[Dict[str, Any]], Dict[str, Any]] = lambda x: x,
45
45
  ) -> CustomConfigClass:
46
- def env_constructor(_: yaml.SafeLoader, node):
46
+ def env_constructor(_: yaml.SafeLoader, node: yaml.Node) -> bool:
47
47
  bool_values = {
48
48
  "true": True,
49
49
  "false": False,
@@ -61,7 +61,7 @@ def _load_yaml(
61
61
 
62
62
  # Safe to use load instead of safe_load since both loader classes are based on SafeLoader
63
63
  try:
64
- config_dict = yaml.load(source, Loader=loader)
64
+ config_dict = yaml.load(source, Loader=loader) # noqa: S506
65
65
  except ScannerError as e:
66
66
  location = e.problem_mark or e.context_mark
67
67
  formatted_location = f" at line {location.line+1}, column {location.column+1}" if location is not None else ""
@@ -80,7 +80,10 @@ def _load_yaml(
80
80
  raise InvalidConfigError(f"Unknown config parameter{'s' if len(unknowns) > 1 else ''} {', '.join(unknowns)}")
81
81
 
82
82
  except (dacite.WrongTypeError, dacite.MissingValueError, dacite.UnionMatchError) as e:
83
- path = e.field_path.replace("_", "-") if case_style == "hyphen" else e.field_path
83
+ if e.field_path:
84
+ path = e.field_path.replace("_", "-") if case_style == "hyphen" else e.field_path
85
+ else:
86
+ path = None
84
87
 
85
88
  def name(type_: Type) -> str:
86
89
  return type_.__name__ if hasattr(type_, "__name__") else str(type_)
@@ -106,7 +109,10 @@ def _load_yaml(
106
109
 
107
110
 
108
111
  def load_yaml(
109
- source: Union[TextIO, str], config_type: Type[CustomConfigClass], case_style: str = "hyphen", expand_envvars=True
112
+ source: Union[TextIO, str],
113
+ config_type: Type[CustomConfigClass],
114
+ case_style: str = "hyphen",
115
+ expand_envvars: bool = True,
110
116
  ) -> CustomConfigClass:
111
117
  """
112
118
  Read a YAML file, and create a config object based on its contents.
@@ -135,7 +141,7 @@ class ConfigResolver(Generic[CustomConfigClass]):
135
141
  self._config: Optional[CustomConfigClass] = None
136
142
  self._next_config: Optional[CustomConfigClass] = None
137
143
 
138
- def _reload_file(self):
144
+ def _reload_file(self) -> None:
139
145
  with open(self.config_path, "r") as stream:
140
146
  self._config_text = stream.read()
141
147
 
@@ -152,17 +158,17 @@ class ConfigResolver(Generic[CustomConfigClass]):
152
158
  def has_changed(self) -> bool:
153
159
  try:
154
160
  self._resolve_config()
155
- except Exception as e:
161
+ except Exception:
156
162
  _logger.exception("Failed to reload configuration file")
157
163
  return False
158
- return self._config._file_hash != self._next_config._file_hash
164
+ return self._config._file_hash != self._next_config._file_hash if self._config else True # type: ignore
159
165
 
160
166
  @property
161
167
  def config(self) -> CustomConfigClass:
162
168
  if self._config is None:
163
169
  self._resolve_config()
164
170
  self.accept_new_config()
165
- return self._config
171
+ return self._config # type: ignore
166
172
 
167
173
  def accept_new_config(self) -> None:
168
174
  self._config = self._next_config
@@ -196,11 +202,13 @@ class ConfigResolver(Generic[CustomConfigClass]):
196
202
  if local_part.cognite.host is not None:
197
203
  remote_part["cognite"]["host"] = local_part.cognite.host
198
204
  remote_part["cognite"]["project"] = local_part.cognite.project
205
+
206
+ # Ignoring None type, extraction pipelines is required at this point
199
207
  remote_part["cognite"]["extraction-pipeline"] = {}
200
- remote_part["cognite"]["extraction-pipeline"]["id"] = local_part.cognite.extraction_pipeline.id
208
+ remote_part["cognite"]["extraction-pipeline"]["id"] = local_part.cognite.extraction_pipeline.id # type: ignore
201
209
  remote_part["cognite"]["extraction-pipeline"][
202
210
  "external_id"
203
- ] = local_part.cognite.extraction_pipeline.external_id
211
+ ] = local_part.cognite.extraction_pipeline.external_id # type: ignore
204
212
 
205
213
  return remote_part
206
214
 
@@ -209,10 +217,10 @@ class ConfigResolver(Generic[CustomConfigClass]):
209
217
 
210
218
  if self.is_remote:
211
219
  _logger.debug("Loading remote config file")
212
- tmp_config: _BaseConfig = load_yaml(self._config_text, _BaseConfig)
220
+ tmp_config: _BaseConfig = load_yaml(self._config_text, _BaseConfig) # type: ignore
213
221
  client = tmp_config.cognite.get_cognite_client("config_resolver")
214
222
  response = client.extraction_pipelines.config.retrieve(
215
- tmp_config.cognite.get_extraction_pipeline(client).external_id
223
+ tmp_config.cognite.get_extraction_pipeline(client).external_id # type: ignore # ignoring extpipe None
216
224
  )
217
225
 
218
226
  self._next_config = _load_yaml(
@@ -43,23 +43,28 @@ import threading
43
43
  from abc import ABC, abstractmethod
44
44
  from threading import Event
45
45
  from time import sleep
46
- from typing import Any, Callable, Dict, List, Optional, T, Tuple, Type, Union
46
+ from types import TracebackType
47
+ from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
47
48
 
48
49
  import arrow
49
50
  import psutil
50
- from cognite.client import CogniteClient
51
- from cognite.client.data_classes import Asset, TimeSeries
52
- from cognite.client.exceptions import CogniteDuplicatedError
53
51
  from prometheus_client import Gauge, Info, Metric
54
52
  from prometheus_client.core import REGISTRY
55
53
  from prometheus_client.exposition import basic_auth_handler, delete_from_gateway, pushadd_to_gateway
56
54
 
55
+ from cognite.client import CogniteClient
56
+ from cognite.client.data_classes import Asset, Datapoints, DatapointsArray, TimeSeries
57
+ from cognite.client.exceptions import CogniteDuplicatedError
58
+
57
59
  from .util import ensure_time_series
58
60
 
59
61
  _metrics_singularities = {}
60
62
 
61
63
 
62
- def safe_get(cls: Type[T], *args, **kwargs) -> T:
64
+ T = TypeVar("T")
65
+
66
+
67
+ def safe_get(cls: Type[T], *args: Any, **kwargs: Any) -> T:
63
68
  """
64
69
  A factory for instances of metrics collections.
65
70
 
@@ -226,7 +231,9 @@ class AbstractMetricsPusher(ABC):
226
231
  self.start()
227
232
  return self
228
233
 
229
- def __exit__(self, exc_type, exc_val, exc_tb) -> None:
234
+ def __exit__(
235
+ self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
236
+ ) -> None:
230
237
  """
231
238
  Wraps around stop method, for use as context manager
232
239
 
@@ -270,7 +277,7 @@ class PrometheusPusher(AbstractMetricsPusher):
270
277
 
271
278
  self.url = url
272
279
 
273
- def _auth_handler(self, url: str, method: str, timeout: int, headers: Dict[str, str], data: Any) -> Callable:
280
+ def _auth_handler(self, url: str, method: str, timeout: int, headers: List[Tuple[str, str]], data: Any) -> Callable:
274
281
  """
275
282
  Returns a authentication handler against the Prometheus Pushgateway to use in the pushadd_to_gateway method.
276
283
 
@@ -298,7 +305,7 @@ class PrometheusPusher(AbstractMetricsPusher):
298
305
 
299
306
  except OSError as exp:
300
307
  self.logger.warning("Failed to push metrics to %s: %s", self.url, str(exp))
301
- except:
308
+ except Exception:
302
309
  self.logger.exception("Failed to push metrics to %s", self.url)
303
310
 
304
311
  self.logger.debug("Pushed metrics to %s", self.url)
@@ -354,6 +361,7 @@ class CognitePusher(AbstractMetricsPusher):
354
361
 
355
362
  if self.asset is not None:
356
363
  # Ensure that asset exist, and retrieve internal ID
364
+ asset: Optional[Asset]
357
365
  try:
358
366
  asset = self.cdf_client.assets.create(self.asset)
359
367
  except CogniteDuplicatedError:
@@ -374,7 +382,7 @@ class CognitePusher(AbstractMetricsPusher):
374
382
  name=metric.name,
375
383
  legacy_name=external_id,
376
384
  description=metric.documentation,
377
- asset_id=asset_id,
385
+ asset_id=asset_id, # type: ignore # this is optional. Type hint in SDK is wrong
378
386
  )
379
387
  )
380
388
 
@@ -386,7 +394,7 @@ class CognitePusher(AbstractMetricsPusher):
386
394
  """
387
395
  timestamp = int(arrow.get().float_timestamp * 1000)
388
396
 
389
- datapoints: List[Dict[str, Union[str, List[Tuple[float, float]]]]] = []
397
+ datapoints: List[Dict[str, Union[str, int, List[Any], Datapoints, DatapointsArray]]] = []
390
398
 
391
399
  for metric in REGISTRY.collect():
392
400
  if type(metric) == Metric and metric.type in ["gauge", "counter"]: