From 06fd933861adda7e1362ff7c8a3642f62b2eeb8b Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Thu, 5 Jun 2025 09:19:01 +0200 Subject: [PATCH 01/11] Update connection config schema --- .../unstable/configuration/models.py | 72 ++++++++++++------- cognite/extractorutils/unstable/core/base.py | 4 +- .../extractorutils/unstable/core/runtime.py | 2 +- tests/test_unstable/conftest.py | 3 +- tests/test_unstable/test_base.py | 2 +- tests/test_unstable/test_configuration.py | 3 +- tests/test_unstable/test_errors.py | 4 +- tests/test_unstable/test_runtime.py | 4 +- 8 files changed, 58 insertions(+), 36 deletions(-) diff --git a/cognite/extractorutils/unstable/configuration/models.py b/cognite/extractorutils/unstable/configuration/models.py index c0dd8ebb..33dd802b 100644 --- a/cognite/extractorutils/unstable/configuration/models.py +++ b/cognite/extractorutils/unstable/configuration/models.py @@ -6,7 +6,7 @@ from typing import Annotated, Any, Literal from humps import kebabize -from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler +from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler, field_validator from pydantic_core import CoreSchema, core_schema from typing_extensions import assert_never @@ -45,23 +45,31 @@ class ConfigModel(BaseModel): ) -class _ClientCredentialsConfig(ConfigModel): - type: Literal["client-credentials"] +class BaseAuthConfig(ConfigModel): client_id: str + scopes: list[str] + + @field_validator("scopes", mode="before", json_schema_input_type=str | list[str]) + @classmethod + def cast_scopes(cls, scopes: str | list[str]) -> list[str]: + if isinstance(scopes, str): + return [scope.strip() for scope in scopes.split(",")] + return scopes + + +class _ClientCredentialsConfig(BaseAuthConfig): + type: Literal["client-credentials"] client_secret: str token_url: str - scopes: list[str] resource: str | None = None audience: str | None = None -class _ClientCertificateConfig(ConfigModel): +class _ClientCertificateConfig(BaseAuthConfig): type: Literal["client-certificate"] - client_id: str path: Path password: str | None = None authority_url: str - scopes: list[str] AuthenticationConfig = Annotated[_ClientCredentialsConfig | _ClientCertificateConfig, Field(discriminator="type")] @@ -139,23 +147,40 @@ def __repr__(self) -> str: return self._expression -class _ConnectionParameters(ConfigModel): - gzip_compression: bool = False - status_forcelist: list[int] = Field(default_factory=lambda: [429, 502, 503, 504]) +class RetriesConfig(ConfigModel): max_retries: int = 10 - max_retries_connect: int = 3 - max_retry_backoff: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s")) - max_connection_pool_size: int = 50 - ssl_verify: bool = True - proxies: dict[str, str] = Field(default_factory=dict) + max_backoff: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s")) timeout: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s")) +class SslCertificatesConfig(ConfigModel): + verify: bool = True + allowed_thumbprints: list[str] | None = None + + @field_validator("allowed_thumbprints", mode="before", json_schema_input_type=str | list[str] | None) + @classmethod + def cast_scopes(cls, scopes: str | list[str] | None) -> list[str] | None: + if scopes is None: + return None + if isinstance(scopes, str): + return [scope.strip() for scope in scopes.split(",")] + return scopes + + +class _ConnectionParameters(ConfigModel): + retries: RetriesConfig = Field(default_factory=RetriesConfig) + ssl_certificates: SslCertificatesConfig = Field(default_factory=SslCertificatesConfig) + + +class IntegrationConfig(ConfigModel): + external_id: str + + class ConnectionConfig(ConfigModel): project: str base_url: str - integration: str + integration: IntegrationConfig authentication: AuthenticationConfig @@ -165,14 +190,9 @@ def get_cognite_client(self, client_name: str) -> CogniteClient: from cognite.client.config import global_config global_config.disable_pypi_version_check = True - global_config.disable_gzip = not self.connection.gzip_compression - global_config.status_forcelist = set(self.connection.status_forcelist) - global_config.max_retries = self.connection.max_retries - global_config.max_retries_connect = self.connection.max_retries_connect - global_config.max_retry_backoff = self.connection.max_retry_backoff.seconds - global_config.max_connection_pool_size = self.connection.max_connection_pool_size - global_config.disable_ssl = not self.connection.ssl_verify - global_config.proxies = self.connection.proxies + global_config.max_retries = self.connection.retries.max_retries + global_config.max_retry_backoff = self.connection.retries.max_backoff.seconds + global_config.disable_ssl = not self.connection.ssl_certificates.verify credential_provider: CredentialProvider match self.authentication: @@ -210,7 +230,7 @@ def get_cognite_client(self, client_name: str) -> CogniteClient: project=self.project, base_url=self.base_url, client_name=client_name, - timeout=self.connection.timeout.seconds, + timeout=self.connection.retries.timeout.seconds, credentials=credential_provider, ) @@ -242,7 +262,7 @@ def from_environment(cls) -> "ConnectionConfig": return ConnectionConfig( project=os.environ["COGNITE_PROJECT"], base_url=os.environ["COGNITE_BASE_URL"], - integration=os.environ["COGNITE_INTEGRATION"], + integration=IntegrationConfig(external_id=os.environ["COGNITE_INTEGRATION"]), authentication=auth, ) diff --git a/cognite/extractorutils/unstable/core/base.py b/cognite/extractorutils/unstable/core/base.py index 40dd4b50..a1886754 100644 --- a/cognite/extractorutils/unstable/core/base.py +++ b/cognite/extractorutils/unstable/core/base.py @@ -155,7 +155,7 @@ def _checkin(self) -> None: res = self.cognite_client.post( f"/api/v1/projects/{self.cognite_client.config.project}/integrations/checkin", json={ - "externalId": self.connection_config.integration, + "externalId": self.connection_config.integration.external_id, "taskEvents": task_updates, "errors": error_updates, }, @@ -265,7 +265,7 @@ def _report_extractor_info(self) -> None: self.cognite_client.post( f"/api/v1/projects/{self.cognite_client.config.project}/integrations/extractorinfo", json={ - "externalId": self.connection_config.integration, + "externalId": self.connection_config.integration.external_id, "activeConfigRevision": self.current_config_revision, "extractor": { "version": self.VERSION, diff --git a/cognite/extractorutils/unstable/core/runtime.py b/cognite/extractorutils/unstable/core/runtime.py index 991e9a15..cf675f5b 100644 --- a/cognite/extractorutils/unstable/core/runtime.py +++ b/cognite/extractorutils/unstable/core/runtime.py @@ -162,7 +162,7 @@ def _try_get_application_config( application_config, current_config_revision = load_from_cdf( self._cognite_client, - connection_config.integration, + connection_config.integration.external_id, self._extractor_class.CONFIG_TYPE, ) diff --git a/tests/test_unstable/conftest.py b/tests/test_unstable/conftest.py index 41bfd0f2..a159ddb4 100644 --- a/tests/test_unstable/conftest.py +++ b/tests/test_unstable/conftest.py @@ -12,6 +12,7 @@ from cognite.extractorutils.unstable.configuration.models import ( ConnectionConfig, ExtractorConfig, + IntegrationConfig, _ClientCredentialsConfig, ) from cognite.extractorutils.unstable.core.base import Extractor @@ -75,7 +76,7 @@ def connection_config(extraction_pipeline: str) -> ConnectionConfig: return ConnectionConfig( project=os.environ["COGNITE_DEV_PROJECT"], base_url=os.environ["COGNITE_DEV_BASE_URL"], - integration=extraction_pipeline, + integration=IntegrationConfig(external_id=extraction_pipeline), authentication=_ClientCredentialsConfig( type="client-credentials", client_id=os.environ.get("COGNITE_DEV_CLIENT_ID", os.environ["COGNITE_CLIENT_ID"]), diff --git a/tests/test_unstable/test_base.py b/tests/test_unstable/test_base.py index c15420cb..92258a80 100644 --- a/tests/test_unstable/test_base.py +++ b/tests/test_unstable/test_base.py @@ -79,7 +79,7 @@ def test_simple_task_report( # Test that the task run is entered into the history for that task res = extractor.cognite_client.get( - f"/api/v1/projects/{extractor.cognite_client.config.project}/integrations/history?integration={connection_config.integration}&taskName=TestTask", + f"/api/v1/projects/{extractor.cognite_client.config.project}/integrations/history?integration={connection_config.integration.external_id}&taskName=TestTask", headers={"cdf-version": "alpha"}, ).json() diff --git a/tests/test_unstable/test_configuration.py b/tests/test_unstable/test_configuration.py index 139b2f58..5cf36ebf 100644 --- a/tests/test_unstable/test_configuration.py +++ b/tests/test_unstable/test_configuration.py @@ -9,7 +9,8 @@ project: test-project base-url: https://baseurl.com -integration: test-pipeline +integration: + external_id: test-pipeline authentication: type: client-credentials diff --git a/tests/test_unstable/test_errors.py b/tests/test_unstable/test_errors.py index e5b8e4f4..bce4258f 100644 --- a/tests/test_unstable/test_errors.py +++ b/tests/test_unstable/test_errors.py @@ -170,7 +170,7 @@ def test_reporting_errors( extractor._checkin() res = extractor.cognite_client.get( - f"/api/v1/projects/{extractor.cognite_client.config.project}/integrations/errors?integration={connection_config.integration}", + f"/api/v1/projects/{extractor.cognite_client.config.project}/integrations/errors?integration={connection_config.integration.external_id}", headers={"cdf-version": "alpha"}, ).json()["items"] assert len(res) == 1 @@ -187,7 +187,7 @@ def test_reporting_errors( extractor._checkin() res = extractor.cognite_client.get( - f"/api/v1/projects/{extractor.cognite_client.config.project}/integrations/errors?integration={connection_config.integration}", + f"/api/v1/projects/{extractor.cognite_client.config.project}/integrations/errors?integration={connection_config.integration.external_id}", headers={"cdf-version": "alpha"}, ).json()["items"] assert len(res) == 1 diff --git a/tests/test_unstable/test_runtime.py b/tests/test_unstable/test_runtime.py index af49a84b..40cc579c 100644 --- a/tests/test_unstable/test_runtime.py +++ b/tests/test_unstable/test_runtime.py @@ -46,7 +46,7 @@ def test_load_cdf_config(connection_config: ConnectionConfig) -> None: cognite_client.post( url=f"/api/v1/projects/{cognite_client.config.project}/odin/config", json={ - "externalId": connection_config.integration, + "externalId": connection_config.integration.external_id, "config": "parameter-one: 123\nparameter-two: abc\n", }, headers={"cdf-version": "alpha"}, @@ -81,7 +81,7 @@ def set_config_after_delay() -> None: cognite_client.post( url=f"/api/v1/projects/{cognite_client.config.project}/odin/config", json={ - "externalId": connection_config.integration, + "externalId": connection_config.integration.external_id, "config": "parameter-one: 123\nparameter-two: abc\n", }, headers={"cdf-version": "alpha"}, From e7a6855a587c69903626426cb140753be2adf255 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Thu, 5 Jun 2025 09:59:01 +0200 Subject: [PATCH 02/11] feat: cast credentials --- cognite/extractorutils/configtools/loaders.py | 3 ++- tests/test_unstable/test_runtime.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/cognite/extractorutils/configtools/loaders.py b/cognite/extractorutils/configtools/loaders.py index 8f9d8b14..405deee5 100644 --- a/cognite/extractorutils/configtools/loaders.py +++ b/cognite/extractorutils/configtools/loaders.py @@ -27,6 +27,7 @@ import dacite import yaml +from azure.core.credentials import TokenCredential from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ServiceRequestError from azure.identity import ClientSecretCredential, DefaultAzureCredential from azure.keyvault.secrets import SecretClient @@ -121,7 +122,7 @@ def _init_client(self) -> None: "Invalid KeyVault authentication method. Possible values : default or client-secret" ) - self.client = SecretClient(vault_url=vault_url, credential=self.credentials) + self.client = SecretClient(vault_url=vault_url, credential=cast(TokenCredential, self.credentials)) def __call__(self, _: yaml.SafeLoader, node: yaml.Node) -> str: self._init_client() diff --git a/tests/test_unstable/test_runtime.py b/tests/test_unstable/test_runtime.py index 40cc579c..2c170caa 100644 --- a/tests/test_unstable/test_runtime.py +++ b/tests/test_unstable/test_runtime.py @@ -114,7 +114,7 @@ def cancel_after_delay() -> None: # There should be one error reported from initially attempting to run without a config errors = cognite_client.get( url=f"/api/v1/projects/{cognite_client.config.project}/odin/errors", - params={"integration": connection_config.integration}, + params={"integration": connection_config.integration.external_id}, ).json() assert len(errors["items"]) == 1 From 66e341249f032af7d9450a83a37e3fe5ced54644 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Thu, 5 Jun 2025 10:08:50 +0200 Subject: [PATCH 03/11] fix: use external id --- cognite/extractorutils/unstable/core/runtime.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cognite/extractorutils/unstable/core/runtime.py b/cognite/extractorutils/unstable/core/runtime.py index cf675f5b..af0787cb 100644 --- a/cognite/extractorutils/unstable/core/runtime.py +++ b/cognite/extractorutils/unstable/core/runtime.py @@ -201,7 +201,7 @@ def _safe_get_application_config( self._cognite_client.post( f"/api/v1/projects/{self._cognite_client.config.project}/odin/checkin", json={ - "externalId": connection_config.integration, + "externalId": connection_config.integration.external_id, "errors": [error.model_dump()], }, headers={"cdf-version": "alpha"}, From 002bc0ac0c005a1faf50b6d8e26d550c427ca0b4 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Thu, 5 Jun 2025 10:19:41 +0200 Subject: [PATCH 04/11] feat: update method --- cognite/extractorutils/unstable/configuration/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cognite/extractorutils/unstable/configuration/models.py b/cognite/extractorutils/unstable/configuration/models.py index 33dd802b..d73453ec 100644 --- a/cognite/extractorutils/unstable/configuration/models.py +++ b/cognite/extractorutils/unstable/configuration/models.py @@ -159,7 +159,7 @@ class SslCertificatesConfig(ConfigModel): @field_validator("allowed_thumbprints", mode="before", json_schema_input_type=str | list[str] | None) @classmethod - def cast_scopes(cls, scopes: str | list[str] | None) -> list[str] | None: + def cast_thumbprints(cls, scopes: str | list[str] | None) -> list[str] | None: if scopes is None: return None if isinstance(scopes, str): From 7bf093943d3f38c40923aa759db0ba014a0ee8e2 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Thu, 5 Jun 2025 10:24:42 +0200 Subject: [PATCH 05/11] test: test casting --- tests/test_unstable/test_configuration.py | 24 +++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/tests/test_unstable/test_configuration.py b/tests/test_unstable/test_configuration.py index 5cf36ebf..0eb9693a 100644 --- a/tests/test_unstable/test_configuration.py +++ b/tests/test_unstable/test_configuration.py @@ -1,6 +1,8 @@ import os from io import StringIO +import pytest + from cognite.client.credentials import OAuthClientCredentials from cognite.extractorutils.unstable.configuration.loaders import ConfigFormat, load_io from cognite.extractorutils.unstable.configuration.models import ConnectionConfig, _ClientCredentialsConfig @@ -19,17 +21,35 @@ token-url: https://get-a-token.com/token scopes: - scopea + - scopeb +""" + +CONFIG_EXAMPLE_ONLY_REQUIRED2 = """ +project: test-project +base-url: https://baseurl.com + +integration: + external_id: test-pipeline + +authentication: + type: client-credentials + client-id: testid + client-secret: very_secret123 + token-url: https://get-a-token.com/token + scopes: scopea, scopeb """ -def test_load_from_io() -> None: - stream = StringIO(CONFIG_EXAMPLE_ONLY_REQUIRED) +@pytest.mark.parametrize("config_str", [CONFIG_EXAMPLE_ONLY_REQUIRED, CONFIG_EXAMPLE_ONLY_REQUIRED2]) +def test_load_from_io(config_str: str) -> None: + stream = StringIO(config_str) config = load_io(stream, ConfigFormat.YAML, ConnectionConfig) assert config.project == "test-project" assert config.base_url == "https://baseurl.com" assert config.authentication.type == "client-credentials" assert config.authentication.client_secret == "very_secret123" + assert config.authentication.scopes == ["scopea", "scopeb"] def test_from_env() -> None: From b8f55fd2876b03c54af935bc97209fc6fbd66795 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Thu, 5 Jun 2025 10:36:01 +0200 Subject: [PATCH 06/11] chore: change variable to match --- .../extractorutils/unstable/configuration/models.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cognite/extractorutils/unstable/configuration/models.py b/cognite/extractorutils/unstable/configuration/models.py index d73453ec..c2aa1409 100644 --- a/cognite/extractorutils/unstable/configuration/models.py +++ b/cognite/extractorutils/unstable/configuration/models.py @@ -159,12 +159,12 @@ class SslCertificatesConfig(ConfigModel): @field_validator("allowed_thumbprints", mode="before", json_schema_input_type=str | list[str] | None) @classmethod - def cast_thumbprints(cls, scopes: str | list[str] | None) -> list[str] | None: - if scopes is None: + def cast_thumbprints(cls, thumbprints: str | list[str] | None) -> list[str] | None: + if thumbprints is None: return None - if isinstance(scopes, str): - return [scope.strip() for scope in scopes.split(",")] - return scopes + if isinstance(thumbprints, str): + return [scope.strip() for scope in thumbprints.split(",")] + return thumbprints class _ConnectionParameters(ConfigModel): From 73abd5d66a464a90fcbd51bb195f2b2e4a74b229 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Fri, 6 Jun 2025 11:38:32 +0200 Subject: [PATCH 07/11] test: comprehensive test of full config-schema --- .../unstable/configuration/models.py | 14 +- .../extractorutils/unstable/core/runtime.py | 8 +- tests/test_unstable/test_configuration.py | 144 ++++++++++++++++++ 3 files changed, 155 insertions(+), 11 deletions(-) diff --git a/cognite/extractorutils/unstable/configuration/models.py b/cognite/extractorutils/unstable/configuration/models.py index c2aa1409..098389de 100644 --- a/cognite/extractorutils/unstable/configuration/models.py +++ b/cognite/extractorutils/unstable/configuration/models.py @@ -45,7 +45,7 @@ class ConfigModel(BaseModel): ) -class BaseAuthConfig(ConfigModel): +class BaseCredentialsConfig(ConfigModel): client_id: str scopes: list[str] @@ -57,7 +57,7 @@ def cast_scopes(cls, scopes: str | list[str]) -> list[str]: return scopes -class _ClientCredentialsConfig(BaseAuthConfig): +class _ClientCredentialsConfig(BaseCredentialsConfig): type: Literal["client-credentials"] client_secret: str token_url: str @@ -65,7 +65,7 @@ class _ClientCredentialsConfig(BaseAuthConfig): audience: str | None = None -class _ClientCertificateConfig(BaseAuthConfig): +class _ClientCertificateConfig(BaseCredentialsConfig): type: Literal["client-certificate"] path: Path password: str | None = None @@ -75,7 +75,7 @@ class _ClientCertificateConfig(BaseAuthConfig): AuthenticationConfig = Annotated[_ClientCredentialsConfig | _ClientCertificateConfig, Field(discriminator="type")] -class TimeIntervalConfig: +class TimeIntervalConfig(str): """ Configuration parameter for setting a time interval """ @@ -148,16 +148,16 @@ def __repr__(self) -> str: class RetriesConfig(ConfigModel): - max_retries: int = 10 + max_retries: int = Field(default=10, ge=-1) max_backoff: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s")) timeout: TimeIntervalConfig = Field(default_factory=lambda: TimeIntervalConfig("30s")) class SslCertificatesConfig(ConfigModel): verify: bool = True - allowed_thumbprints: list[str] | None = None + allow_list: list[str] | None = None - @field_validator("allowed_thumbprints", mode="before", json_schema_input_type=str | list[str] | None) + @field_validator("allow_list", mode="before", json_schema_input_type=str | list[str] | None) @classmethod def cast_thumbprints(cls, thumbprints: str | list[str] | None) -> list[str] | None: if thumbprints is None: diff --git a/cognite/extractorutils/unstable/core/runtime.py b/cognite/extractorutils/unstable/core/runtime.py index af0787cb..976ecba1 100644 --- a/cognite/extractorutils/unstable/core/runtime.py +++ b/cognite/extractorutils/unstable/core/runtime.py @@ -24,13 +24,13 @@ load_file, load_from_cdf, ) -from cognite.extractorutils.unstable.configuration.models import ConnectionConfig +from cognite.extractorutils.unstable.configuration.models import ConnectionConfig, ExtractorConfig from cognite.extractorutils.unstable.core._dto import Error from cognite.extractorutils.unstable.core.errors import ErrorLevel from cognite.extractorutils.util import now from ._messaging import RuntimeMessage -from .base import ConfigRevision, ConfigType, Extractor, FullConfig +from .base import ConfigRevision, Extractor, FullConfig __all__ = ["Runtime", "ExtractorType"] @@ -141,7 +141,7 @@ def _try_get_application_config( self, args: Namespace, connection_config: ConnectionConfig, - ) -> tuple[ConfigType, ConfigRevision]: + ) -> tuple[ExtractorConfig, ConfigRevision]: current_config_revision: ConfigRevision if args.local_override: @@ -172,7 +172,7 @@ def _safe_get_application_config( self, args: Namespace, connection_config: ConnectionConfig, - ) -> tuple[ConfigType, ConfigRevision] | None: + ) -> tuple[ExtractorConfig, ConfigRevision] | None: prev_error: str | None = None while not self._cancellation_token.is_cancelled: diff --git a/tests/test_unstable/test_configuration.py b/tests/test_unstable/test_configuration.py index 0eb9693a..8cd3dc0d 100644 --- a/tests/test_unstable/test_configuration.py +++ b/tests/test_unstable/test_configuration.py @@ -39,6 +39,94 @@ scopes: scopea, scopeb """ +FULL_CONFIG_EXAMPLE = """ +project: test-project +base-url: https://baseurl.com + +integration: + external_id: test-pipeline +""" + +CLIENT_CREDENTIALS_STRING = """ +authentication: + type: client-credentials + client-id: testid + client-secret: very_secret123 + token-url: https://get-a-token.com/token + scopes: scopea, scopeb + +connection: + retries: + max-retries: 1 + max-backoff: 2d + timeout: 3h + ssl-certificates: + verify: true + allow-list: thumbprint1, thumbprint2 +""" +CLIENT_CERTIFICATES_STRING = """ +authentication: + type: client-certificate + client-id: testid + path: /path/to/cert.pem + password: very-strong-password + authority-url: https://you-are-authorized.com + scopes: scopea, scopeb + +connection: + retries: + max-retries: 1 + max-backoff: 2d + timeout: 3h + ssl-certificates: + verify: true + allow-list: thumbprint1, thumbprint2 +""" + +CLIENT_CREDENTIALS_LIST = """ +authentication: + type: client-credentials + client-id: testid + client-secret: very_secret123 + token-url: https://get-a-token.com/token + scopes: + - scopea + - scopeb + +connection: + retries: + max-retries: 1 + max-backoff: 2d + timeout: 3h + ssl-certificates: + verify: true + allow-list: + - thumbprint1 + - thumbprint2 +""" +CLIENT_CERTIFICATES_LIST = """ +authentication: + type: client-certificate + client-id: testid + path: /path/to/cert.pem + password: very-strong-password + authority-url: https://you-are-authorized.com + scopes: + - scopea + - scopeb + +connection: + retries: + max-retries: 1 + max-backoff: 2d + timeout: 3h + ssl-certificates: + verify: true + allow-list: + - thumbprint1 + - thumbprint2 +""" + @pytest.mark.parametrize("config_str", [CONFIG_EXAMPLE_ONLY_REQUIRED, CONFIG_EXAMPLE_ONLY_REQUIRED2]) def test_load_from_io(config_str: str) -> None: @@ -47,11 +135,67 @@ def test_load_from_io(config_str: str) -> None: assert config.project == "test-project" assert config.base_url == "https://baseurl.com" + assert config.integration.external_id == "test-pipeline" assert config.authentication.type == "client-credentials" assert config.authentication.client_secret == "very_secret123" assert config.authentication.scopes == ["scopea", "scopeb"] +@pytest.mark.parametrize( + "config_str", [FULL_CONFIG_EXAMPLE + CLIENT_CREDENTIALS_STRING, FULL_CONFIG_EXAMPLE + CLIENT_CREDENTIALS_LIST] +) +def test_full_config_client_credentials(config_str: str) -> None: + stream = StringIO(config_str) + config = load_io(stream, ConfigFormat.YAML, ConnectionConfig) + + assert config.project == "test-project" + assert config.base_url == "https://baseurl.com" + assert config.integration.external_id == "test-pipeline" + + assert config.authentication.type == "client-credentials" + assert config.authentication.client_id == "testid" + assert config.authentication.client_secret == "very_secret123" + assert config.authentication.token_url == "https://get-a-token.com/token" + assert config.authentication.scopes == ["scopea", "scopeb"] + + assert config.connection.retries.max_retries == 1 + assert config.connection.retries.max_backoff.seconds == 2 * 24 * 60 * 60 + assert config.connection.retries.max_backoff == "2d" + assert config.connection.retries.timeout.seconds == 3 * 60 * 60 + assert config.connection.retries.timeout == "3h" + + assert config.connection.ssl_certificates.verify + assert config.connection.ssl_certificates.allow_list == ["thumbprint1", "thumbprint2"] + + +@pytest.mark.parametrize( + "config_str", [FULL_CONFIG_EXAMPLE + CLIENT_CERTIFICATES_STRING, FULL_CONFIG_EXAMPLE + CLIENT_CERTIFICATES_LIST] +) +def test_full_config_client_certificates(config_str: str) -> None: + stream = StringIO(config_str) + config = load_io(stream, ConfigFormat.YAML, ConnectionConfig) + + assert config.project == "test-project" + assert config.base_url == "https://baseurl.com" + assert config.integration.external_id == "test-pipeline" + + assert config.authentication.type == "client-certificate" + assert config.authentication.client_id == "testid" + assert config.authentication.password == "very-strong-password" + assert config.authentication.path.as_posix() == "/path/to/cert.pem" + assert config.authentication.authority_url == "https://you-are-authorized.com" + assert config.authentication.scopes == ["scopea", "scopeb"] + + assert config.connection.retries.max_retries == 1 + assert config.connection.retries.max_backoff.seconds == 2 * 24 * 60 * 60 + assert config.connection.retries.max_backoff == "2d" + assert config.connection.retries.timeout.seconds == 3 * 60 * 60 + assert config.connection.retries.timeout == "3h" + + assert config.connection.ssl_certificates.verify + assert config.connection.ssl_certificates.allow_list == ["thumbprint1", "thumbprint2"] + + def test_from_env() -> None: # Init a config from env config = ConnectionConfig.from_environment() From 947518e431af868982a90c8830c8e7fb9024fc49 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Wed, 11 Jun 2025 15:12:12 +0200 Subject: [PATCH 08/11] Fix type issue in `KeyValutLoader` (#447) `SecretClient` does not accept `None`s as the `credential`. This is a type issue, that `mypy` only _sometimes_ catches. There is no reason the `credentials` should be a class variable, it's only ever used inside the `_init_client` method, so it's fine to move it to a local variable. --- cognite/extractorutils/configtools/loaders.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cognite/extractorutils/configtools/loaders.py b/cognite/extractorutils/configtools/loaders.py index bf667709..36174d02 100644 --- a/cognite/extractorutils/configtools/loaders.py +++ b/cognite/extractorutils/configtools/loaders.py @@ -65,7 +65,6 @@ class KeyVaultLoader: def __init__(self, config: dict | None): self.config = config - self.credentials: DefaultAzureCredential | ClientSecretCredential | None = None self.client: SecretClient | None = None def _init_client(self) -> None: @@ -89,9 +88,10 @@ def _init_client(self) -> None: vault_url = f"https://{keyvault_name}.vault.azure.net" + credentials: TokenCredential if self.config["authentication-method"] == KeyVaultAuthenticationMethod.DEFAULT.value: _logger.info("Using Azure DefaultCredentials to access KeyVault") - self.credentials = DefaultAzureCredential() + credentials = DefaultAzureCredential() elif self.config["authentication-method"] == KeyVaultAuthenticationMethod.CLIENTSECRET.value: auth_parameters = ("client-id", "tenant-id", "secret") @@ -108,7 +108,7 @@ def _init_client(self) -> None: client_id = os.path.expandvars(self.config["client-id"]) secret = os.path.expandvars(self.config["secret"]) - self.credentials = ClientSecretCredential( + credentials = ClientSecretCredential( tenant_id=tenant_id, client_id=client_id, client_secret=secret, @@ -122,7 +122,7 @@ def _init_client(self) -> None: "Invalid KeyVault authentication method. Possible values : default or client-secret" ) - self.client = SecretClient(vault_url=vault_url, credential=cast(TokenCredential, self.credentials)) + self.client = SecretClient(vault_url=vault_url, credential=credentials) def __call__(self, _: yaml.SafeLoader, node: yaml.Node) -> str: self._init_client() From ee2bb86e9232105004328af5beb8fa580e2bc272 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Tue, 17 Jun 2025 14:37:20 +0200 Subject: [PATCH 09/11] feat: scopes config list --- .../unstable/configuration/models.py | 58 ++++++++++++------- tests/test_unstable/conftest.py | 5 +- tests/test_unstable/test_configuration.py | 46 ++++++++------- 3 files changed, 64 insertions(+), 45 deletions(-) diff --git a/cognite/extractorutils/unstable/configuration/models.py b/cognite/extractorutils/unstable/configuration/models.py index 514c3403..12d10224 100644 --- a/cognite/extractorutils/unstable/configuration/models.py +++ b/cognite/extractorutils/unstable/configuration/models.py @@ -1,12 +1,13 @@ import os import re +from collections.abc import Iterator from datetime import timedelta from enum import Enum from pathlib import Path from typing import Annotated, Any, Literal from humps import kebabize -from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler, field_validator +from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler from pydantic_core import CoreSchema, core_schema from typing_extensions import assert_never @@ -45,16 +46,34 @@ class ConfigModel(BaseModel): ) -class BaseCredentialsConfig(ConfigModel): - client_id: str - scopes: list[str] +class Scopes(str): + def __init__(self, scopes: str) -> None: + self._scopes = list(scopes.split(" ")) - @field_validator("scopes", mode="before", json_schema_input_type=str | list[str]) @classmethod - def cast_scopes(cls, scopes: str | list[str]) -> list[str]: - if isinstance(scopes, str): - return [scope.strip() for scope in scopes.split(",")] - return scopes + def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema: + return core_schema.no_info_after_validator_function(cls, handler(str)) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Scopes): + return NotImplemented + return self._scopes == other._scopes + + def __hash__(self) -> int: + return hash(self._scopes) + + def __iter__(self) -> Iterator[str]: + return iter(self._scopes) + + # def __next__(self) -> str: + # if not self._scopes: + # raise StopIteration + # return self._scopes.pop(0) + + +class BaseCredentialsConfig(ConfigModel): + client_id: str + scopes: Scopes class _ClientCredentialsConfig(BaseCredentialsConfig): @@ -75,7 +94,7 @@ class _ClientCertificateConfig(BaseCredentialsConfig): AuthenticationConfig = Annotated[_ClientCredentialsConfig | _ClientCertificateConfig, Field(discriminator="type")] -class TimeIntervalConfig(str): +class TimeIntervalConfig: """ Configuration parameter for setting a time interval """ @@ -157,15 +176,6 @@ class SslCertificatesConfig(ConfigModel): verify: bool = True allow_list: list[str] | None = None - @field_validator("allow_list", mode="before", json_schema_input_type=str | list[str] | None) - @classmethod - def cast_thumbprints(cls, thumbprints: str | list[str] | None) -> list[str] | None: - if thumbprints is None: - return None - if isinstance(thumbprints, str): - return [scope.strip() for scope in thumbprints.split(",")] - return thumbprints - class _ConnectionParameters(ConfigModel): retries: RetriesConfig = Field(default_factory=RetriesConfig) @@ -220,7 +230,7 @@ def get_cognite_client(self, client_name: str) -> CogniteClient: client_id=client_certificate.client_id, cert_thumbprint=str(thumbprint), certificate=str(key), - scopes=client_certificate.scopes, + scopes=list(client_certificate.scopes), ) case _: @@ -245,7 +255,9 @@ def from_environment(cls) -> "ConnectionConfig": client_id=os.environ["COGNITE_CLIENT_ID"], client_secret=os.environ["COGNITE_CLIENT_SECRET"], token_url=os.environ["COGNITE_TOKEN_URL"], - scopes=os.environ["COGNITE_TOKEN_SCOPES"].split(","), + scopes=Scopes( + os.environ["COGNITE_TOKEN_SCOPES"], + ), ) elif "COGNITE_CLIENT_CERTIFICATE_PATH" in os.environ: auth = _ClientCertificateConfig( @@ -254,7 +266,9 @@ def from_environment(cls) -> "ConnectionConfig": path=Path(os.environ["COGNITE_CLIENT_CERTIFICATE_PATH"]), password=os.environ.get("COGNITE_CLIENT_CERTIFICATE_PATH"), authority_url=os.environ["COGNITE_AUTHORITY_URL"], - scopes=os.environ["COGNITE_TOKEN_SCOPES"].split(","), + scopes=Scopes( + os.environ["COGNITE_TOKEN_SCOPES"], + ), ) else: raise KeyError("Missing auth, either COGNITE_CLIENT_SECRET or COGNITE_CLIENT_CERTIFICATE_PATH must be set") diff --git a/tests/test_unstable/conftest.py b/tests/test_unstable/conftest.py index a159ddb4..5de580a7 100644 --- a/tests/test_unstable/conftest.py +++ b/tests/test_unstable/conftest.py @@ -13,6 +13,7 @@ ConnectionConfig, ExtractorConfig, IntegrationConfig, + Scopes, _ClientCredentialsConfig, ) from cognite.extractorutils.unstable.core.base import Extractor @@ -81,7 +82,9 @@ def connection_config(extraction_pipeline: str) -> ConnectionConfig: type="client-credentials", client_id=os.environ.get("COGNITE_DEV_CLIENT_ID", os.environ["COGNITE_CLIENT_ID"]), client_secret=os.environ.get("COGNITE_DEV_CLIENT_SECRET", os.environ["COGNITE_CLIENT_SECRET"]), - scopes=os.environ["COGNITE_DEV_TOKEN_SCOPES"].split(","), + scopes=Scopes( + os.environ["COGNITE_DEV_TOKEN_SCOPES"], + ), token_url=os.environ.get("COGNITE_DEV_TOKEN_URL", os.environ["COGNITE_TOKEN_URL"]), ), ) diff --git a/tests/test_unstable/test_configuration.py b/tests/test_unstable/test_configuration.py index 8cd3dc0d..d7c3f050 100644 --- a/tests/test_unstable/test_configuration.py +++ b/tests/test_unstable/test_configuration.py @@ -5,7 +5,11 @@ from cognite.client.credentials import OAuthClientCredentials from cognite.extractorutils.unstable.configuration.loaders import ConfigFormat, load_io -from cognite.extractorutils.unstable.configuration.models import ConnectionConfig, _ClientCredentialsConfig +from cognite.extractorutils.unstable.configuration.models import ( + ConnectionConfig, + TimeIntervalConfig, + _ClientCredentialsConfig, +) CONFIG_EXAMPLE_ONLY_REQUIRED = """ project: test-project @@ -19,9 +23,7 @@ client-id: testid client-secret: very_secret123 token-url: https://get-a-token.com/token - scopes: - - scopea - - scopeb + scopes: scopea scopeb """ CONFIG_EXAMPLE_ONLY_REQUIRED2 = """ @@ -36,7 +38,7 @@ client-id: testid client-secret: very_secret123 token-url: https://get-a-token.com/token - scopes: scopea, scopeb + scopes: scopea scopeb """ FULL_CONFIG_EXAMPLE = """ @@ -53,7 +55,7 @@ client-id: testid client-secret: very_secret123 token-url: https://get-a-token.com/token - scopes: scopea, scopeb + scopes: scopea scopeb connection: retries: @@ -62,7 +64,9 @@ timeout: 3h ssl-certificates: verify: true - allow-list: thumbprint1, thumbprint2 + allow-list: + - thumbprint1 + - thumbprint2 """ CLIENT_CERTIFICATES_STRING = """ authentication: @@ -71,7 +75,7 @@ path: /path/to/cert.pem password: very-strong-password authority-url: https://you-are-authorized.com - scopes: scopea, scopeb + scopes: scopea scopeb connection: retries: @@ -80,7 +84,9 @@ timeout: 3h ssl-certificates: verify: true - allow-list: thumbprint1, thumbprint2 + allow-list: + - thumbprint1 + - thumbprint2 """ CLIENT_CREDENTIALS_LIST = """ @@ -89,9 +95,7 @@ client-id: testid client-secret: very_secret123 token-url: https://get-a-token.com/token - scopes: - - scopea - - scopeb + scopes: scopea scopeb connection: retries: @@ -111,9 +115,7 @@ path: /path/to/cert.pem password: very-strong-password authority-url: https://you-are-authorized.com - scopes: - - scopea - - scopeb + scopes: scopea scopeb connection: retries: @@ -138,7 +140,7 @@ def test_load_from_io(config_str: str) -> None: assert config.integration.external_id == "test-pipeline" assert config.authentication.type == "client-credentials" assert config.authentication.client_secret == "very_secret123" - assert config.authentication.scopes == ["scopea", "scopeb"] + assert list(config.authentication.scopes) == ["scopea", "scopeb"] @pytest.mark.parametrize( @@ -156,13 +158,13 @@ def test_full_config_client_credentials(config_str: str) -> None: assert config.authentication.client_id == "testid" assert config.authentication.client_secret == "very_secret123" assert config.authentication.token_url == "https://get-a-token.com/token" - assert config.authentication.scopes == ["scopea", "scopeb"] + assert list(config.authentication.scopes) == ["scopea", "scopeb"] assert config.connection.retries.max_retries == 1 assert config.connection.retries.max_backoff.seconds == 2 * 24 * 60 * 60 - assert config.connection.retries.max_backoff == "2d" + assert config.connection.retries.max_backoff == TimeIntervalConfig("2d") assert config.connection.retries.timeout.seconds == 3 * 60 * 60 - assert config.connection.retries.timeout == "3h" + assert config.connection.retries.timeout == TimeIntervalConfig("3h") assert config.connection.ssl_certificates.verify assert config.connection.ssl_certificates.allow_list == ["thumbprint1", "thumbprint2"] @@ -184,13 +186,13 @@ def test_full_config_client_certificates(config_str: str) -> None: assert config.authentication.password == "very-strong-password" assert config.authentication.path.as_posix() == "/path/to/cert.pem" assert config.authentication.authority_url == "https://you-are-authorized.com" - assert config.authentication.scopes == ["scopea", "scopeb"] + assert list(config.authentication.scopes) == ["scopea", "scopeb"] assert config.connection.retries.max_retries == 1 assert config.connection.retries.max_backoff.seconds == 2 * 24 * 60 * 60 - assert config.connection.retries.max_backoff == "2d" + assert config.connection.retries.max_backoff == TimeIntervalConfig("2d") assert config.connection.retries.timeout.seconds == 3 * 60 * 60 - assert config.connection.retries.timeout == "3h" + assert config.connection.retries.timeout == TimeIntervalConfig("3h") assert config.connection.ssl_certificates.verify assert config.connection.ssl_certificates.allow_list == ["thumbprint1", "thumbprint2"] From 0eeb72fed6d7b7d289082b8652575f56cd520d71 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Thu, 19 Jun 2025 09:41:31 +0200 Subject: [PATCH 10/11] feat: expose connection parameters --- cognite/extractorutils/unstable/configuration/models.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/cognite/extractorutils/unstable/configuration/models.py b/cognite/extractorutils/unstable/configuration/models.py index 12d10224..854f961a 100644 --- a/cognite/extractorutils/unstable/configuration/models.py +++ b/cognite/extractorutils/unstable/configuration/models.py @@ -65,11 +65,6 @@ def __hash__(self) -> int: def __iter__(self) -> Iterator[str]: return iter(self._scopes) - # def __next__(self) -> str: - # if not self._scopes: - # raise StopIteration - # return self._scopes.pop(0) - class BaseCredentialsConfig(ConfigModel): client_id: str @@ -177,7 +172,7 @@ class SslCertificatesConfig(ConfigModel): allow_list: list[str] | None = None -class _ConnectionParameters(ConfigModel): +class ConnectionParameters(ConfigModel): retries: RetriesConfig = Field(default_factory=RetriesConfig) ssl_certificates: SslCertificatesConfig = Field(default_factory=SslCertificatesConfig) @@ -194,7 +189,7 @@ class ConnectionConfig(ConfigModel): authentication: AuthenticationConfig - connection: _ConnectionParameters = Field(default_factory=_ConnectionParameters) + connection: ConnectionParameters = Field(default_factory=ConnectionParameters) def get_cognite_client(self, client_name: str) -> CogniteClient: from cognite.client.config import global_config From 3d8d15459e1942c6e035ba8eee3cb35bb0540663 Mon Sep 17 00:00:00 2001 From: Babatunde Aromire Date: Thu, 19 Jun 2025 11:09:13 +0200 Subject: [PATCH 11/11] fix: update cognite sdk version --- cognite/extractorutils/metrics.py | 3 ++- pyproject.toml | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/cognite/extractorutils/metrics.py b/cognite/extractorutils/metrics.py index bc73b296..c6fb5217 100644 --- a/cognite/extractorutils/metrics.py +++ b/cognite/extractorutils/metrics.py @@ -54,6 +54,7 @@ def __init__(self): from cognite.client import CogniteClient from cognite.client.data_classes import Asset, Datapoints, DatapointsArray, TimeSeries +from cognite.client.data_classes.data_modeling import NodeId from cognite.client.exceptions import CogniteDuplicatedError from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.util import EitherId @@ -411,7 +412,7 @@ def _push_to_server(self) -> None: """ timestamp = int(arrow.get().float_timestamp * 1000) - datapoints: list[dict[str, str | int | list[Any] | Datapoints | DatapointsArray]] = [] + datapoints: list[dict[str, str | int | list[Any] | Datapoints | DatapointsArray | NodeId]] = [] for metric in REGISTRY.collect(): if isinstance(metric, Metric) and metric.type in ["gauge", "counter"]: diff --git a/pyproject.toml b/pyproject.toml index c97a61a7..8c0ac583 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ classifiers = [ ] dependencies = [ - "cognite-sdk>=7.59.0", + "cognite-sdk>=7.75.2", "prometheus-client>=0.7.0,<=1.0.0", "arrow>=1.0.0", "pyyaml>=5.3.0,<7",