cognite-extractor-utils 7.2.2__tar.gz → 7.3.0__tar.gz

{cognite_extractor_utils-7.2.2 → cognite_extractor_utils-7.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cognite-extractor-utils
-Version: 7.2.2
+Version: 7.3.0
 Summary: Utilities for easier development of extractors for CDF
 Home-page: https://github.com/cognitedata/python-extractor-utils
 License: Apache-2.0
@@ -21,6 +21,7 @@ Requires-Dist: azure-keyvault-secrets (>=4.7.0,<5.0.0)
 Requires-Dist: cognite-sdk (>=7.43.3,<8.0.0)
 Requires-Dist: dacite (>=1.6.0,<2.0.0)
 Requires-Dist: decorator (>=5.1.1,<6.0.0)
+Requires-Dist: httpx (>=0.27.0,<0.28.0)
 Requires-Dist: more-itertools (>=10.0.0,<11.0.0)
 Requires-Dist: orjson (>=3.10.3,<4.0.0)
 Requires-Dist: prometheus-client (>0.7.0,<=1.0.0)

{cognite_extractor_utils-7.2.2 → cognite_extractor_utils-7.3.0}/cognite/extractorutils/__init__.py

@@ -16,5 +16,5 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "7.2.2"
+__version__ = "7.3.0"
 from .base import Extractor

{cognite_extractor_utils-7.2.2 → cognite_extractor_utils-7.3.0}/cognite/extractorutils/configtools/elements.py

@@ -20,7 +20,7 @@ from enum import Enum
 from logging.handlers import TimedRotatingFileHandler
 from time import sleep
 from typing import Any, Dict, List, Optional, Tuple, Union
-from urllib.parse import urljoin
+from urllib.parse import urljoin, urlparse
 
 import yaml
 from prometheus_client import REGISTRY, start_http_server
@@ -70,7 +70,7 @@ class AuthenticatorConfig:
 
     client_id: str
     scopes: List[str]
-    secret: Optional[str]
+    secret: Optional[str] = None
     tenant: Optional[str] = None
     token_url: Optional[str] = None
     resource: Optional[str] = None
@@ -264,6 +264,18 @@ class FileSizeConfig(yaml.YAMLObject):
         return self._expression
 
 
+path_elem_regex = re.compile(r"^([a-zA-Z0-9.\-_~!$&'()*+,;=:@]|%[A-F0-9]{2})*$")
+
+
+def _validate_https_url(value: str, name: str) -> None:
+    try:
+        url = urlparse(value)
+    except Exception as e:
+        raise InvalidConfigError(f"{name} ({value}) is not a valid URL") from e
+    if url.scheme != "https":
+        raise InvalidConfigError(f"{name} ({value}) must be HTTPS")
+
+
 @dataclass
 class CogniteConfig:
     """
@@ -272,10 +284,10 @@ class CogniteConfig:
 
     project: str
     idp_authentication: AuthenticatorConfig
-    data_set: Optional[EitherIdConfig]
-    data_set_id: Optional[int]
-    data_set_external_id: Optional[str]
-    extraction_pipeline: Optional[EitherIdConfig]
+    data_set: Optional[EitherIdConfig] = None
+    data_set_id: Optional[int] = None
+    data_set_external_id: Optional[str] = None
+    extraction_pipeline: Optional[EitherIdConfig] = None
     timeout: TimeIntervalConfig = TimeIntervalConfig("30s")
     connection: ConnectionConfig = field(default_factory=ConnectionConfig)
     security_categories: Optional[List[int]] = None
@@ -300,11 +312,20 @@ class CogniteConfig:
         global_config.disable_ssl = self.connection.disable_ssl
         global_config.proxies = self.connection.proxies
 
+        if not self.project:
+            raise InvalidConfigError("Project is not set")
+        if not path_elem_regex.match(self.project):
+            raise InvalidConfigError(f"Project ({self.project}) is not valid")
+
         credential_provider: CredentialProvider
         if self.idp_authentication.certificate:
             if self.idp_authentication.certificate.authority_url:
                 authority_url = self.idp_authentication.certificate.authority_url
+                _validate_https_url(self.idp_authentication.certificate.authority_url, "Authority URL")
             elif self.idp_authentication.tenant:
+                _validate_https_url(self.idp_authentication.authority, "Authority")
+                if not path_elem_regex.match(self.idp_authentication.tenant):
+                    raise InvalidConfigError(f"Tenant {self.idp_authentication.tenant} is not valid")
                 authority_url = urljoin(self.idp_authentication.authority, self.idp_authentication.tenant)
             else:
                 raise InvalidConfigError("Either authority-url or tenant is required for certificate authentication")
@@ -312,6 +333,8 @@ class CogniteConfig:
                 self.idp_authentication.certificate.path,
                 self.idp_authentication.certificate.password,
             )
+            if not self.idp_authentication.scopes:
+                _logger.warn("No scopes configured. Authenticating with CDF is unlikely to work correctly")
             credential_provider = OAuthClientCertificate(
                 authority_url=authority_url,
                 client_id=self.idp_authentication.client_id,
@@ -323,14 +346,20 @@ class CogniteConfig:
         elif self.idp_authentication.secret:
             kwargs: Dict[str, Any] = {}
             if self.idp_authentication.token_url:
+                _validate_https_url(self.idp_authentication.token_url, "Token URL")
                 kwargs["token_url"] = self.idp_authentication.token_url
             elif self.idp_authentication.tenant:
+                _validate_https_url(self.idp_authentication.authority, "Authority")
+                if not path_elem_regex.match(self.idp_authentication.tenant):
+                    raise InvalidConfigError(f"Tenant ({self.idp_authentication.tenant}) is not valid")
                 base_url = urljoin(self.idp_authentication.authority, self.idp_authentication.tenant)
                 kwargs["token_url"] = f"{base_url}/oauth2/v2.0/token"
             else:
                 raise InvalidConfigError("Either token-url or tenant is required for client credentials authentication")
             kwargs["client_id"] = self.idp_authentication.client_id
             kwargs["client_secret"] = self.idp_authentication.secret
+            if not self.idp_authentication.scopes:
+                _logger.warning("No scopes configured. Authenticating with CDF is unlikely to work correctly")
             kwargs["scopes"] = self.idp_authentication.scopes
             if token_custom_args is None:
                 token_custom_args = {}

{cognite_extractor_utils-7.2.2 → cognite_extractor_utils-7.3.0}/cognite/extractorutils/uploader/files.py

@@ -18,8 +18,9 @@ from io import BytesIO, RawIOBase
 from math import ceil
 from os import PathLike
 from types import TracebackType
-from typing import Any, BinaryIO, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, BinaryIO, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union
 
+from httpx import URL, Client, Headers, Request, StreamConsumed, SyncByteStream
 from requests.utils import super_len
 
 from cognite.client import CogniteClient
@@ -140,6 +141,22 @@ class ChunkedStream(RawIOBase, BinaryIO):
         return True
 
 
+class IOByteStream(SyncByteStream):
+    CHUNK_SIZE = 65_536
+
+    def __init__(self, stream: BinaryIO) -> None:
+        self._stream = stream
+        self._is_stream_consumed = False
+
+    def __iter__(self) -> Iterator[bytes]:
+        if self._is_stream_consumed:
+            raise StreamConsumed()
+        chunk = self._stream.read(self.CHUNK_SIZE)
+        while chunk:
+            yield chunk
+            chunk = self._stream.read(self.CHUNK_SIZE)
+
+
 class IOFileUploadQueue(AbstractUploadQueue):
     """
     Upload queue for files using BinaryIO
@@ -205,6 +222,8 @@ class IOFileUploadQueue(AbstractUploadQueue):
 
         self._full_queue = threading.Condition()
 
+        self._httpx_client = Client(follow_redirects=True)
+
         global _QUEUES, _QUEUES_LOCK
         with _QUEUES_LOCK:
             self._pool = ThreadPoolExecutor(
@@ -266,44 +285,32 @@ class IOFileUploadQueue(AbstractUploadQueue):
                         f"File {file_meta.external_id} is larger than 5GiB ({size})"
                         f", uploading in {chunks.chunk_count} chunks"
                     )
-                    with self.cdf_client.files.multipart_upload_session(
-                        file_meta.name if file_meta.name is not None else "",
-                        parts=chunks.chunk_count,
-                        overwrite=self.overwrite_existing,
-                        external_id=file_meta.external_id,
-                        source=file_meta.source,
-                        mime_type=file_meta.mime_type,
-                        metadata=file_meta.metadata,
-                        directory=file_meta.directory,
-                        asset_ids=file_meta.asset_ids,
-                        data_set_id=file_meta.data_set_id,
-                        labels=file_meta.labels,
-                        geo_location=file_meta.geo_location,
-                        source_created_time=file_meta.source_created_time,
-                        source_modified_time=file_meta.source_modified_time,
-                        security_categories=file_meta.security_categories,
-                    ) as session:
-                        while chunks.next_chunk():
-                            session.upload_part(chunks.current_chunk, chunks)
-                        file_meta = session.file_metadata
+
+                    res = self.cdf_client.files._post(
+                        url_path="/files/initmultipartupload",
+                        json=file_meta.dump(camel_case=True),
+                        params={"overwrite": self.overwrite_existing, "parts": chunks.chunk_count},
+                    )
+                    returned_file_metadata = res.json()
+                    upload_urls = returned_file_metadata["uploadUrls"]
+                    upload_id = returned_file_metadata["uploadId"]
+                    file_meta = FileMetadata.load(returned_file_metadata)
+
+                    for url in upload_urls:
+                        chunks.next_chunk()
+                        resp = self._httpx_client.send(self._get_file_upload_request(url, chunks, len(chunks)))
+                        resp.raise_for_status()
+
+                    self.cdf_client.files._post(
+                        url_path="/files/completemultipartupload", json={"id": file_meta.id, "uploadId": upload_id}
+                    )
+
                 else:
-                    file_meta = self.cdf_client.files.upload_bytes(
-                        file,
-                        file_meta.name if file_meta.name is not None else "",
-                        overwrite=self.overwrite_existing,
-                        external_id=file_meta.external_id,
-                        source=file_meta.source,
-                        mime_type=file_meta.mime_type,
-                        metadata=file_meta.metadata,
-                        directory=file_meta.directory,
-                        asset_ids=file_meta.asset_ids,
-                        data_set_id=file_meta.data_set_id,
-                        labels=file_meta.labels,
-                        geo_location=file_meta.geo_location,
-                        source_created_time=file_meta.source_created_time,
-                        source_modified_time=file_meta.source_modified_time,
-                        security_categories=file_meta.security_categories,
+                    file_meta, url = self.cdf_client.files.create(
+                        file_metadata=file_meta, overwrite=self.overwrite_existing
                     )
+                    resp = self._httpx_client.send(self._get_file_upload_request(url, file, size))
+                    resp.raise_for_status()
 
             if self.post_upload_function:
                 try:
@@ -338,6 +345,25 @@ class IOFileUploadQueue(AbstractUploadQueue):
             self.files_queued.inc()
             self.queue_size.set(self.upload_queue_size)
 
+    def _get_file_upload_request(self, url_str: str, stream: BinaryIO, size: int) -> Request:
+        url = URL(url_str)
+        headers = Headers(self._httpx_client.headers)
+        headers.update(
+            {
+                "Accept": "*/*",
+                "Content-Length": str(size),
+                "Host": url.netloc.decode("ascii"),
+                "x-cdp-app": self.cdf_client._config.client_name,
+            }
+        )
+
+        return Request(
+            method="PUT",
+            url=url,
+            stream=IOByteStream(stream),
+            headers=headers,
+        )
+
     def upload(self, fail_on_errors: bool = True, timeout: Optional[float] = None) -> None:
         """
         Wait for all uploads to finish

{cognite_extractor_utils-7.2.2 → cognite_extractor_utils-7.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.2.2"
+version = "7.3.0"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"
@@ -66,13 +66,14 @@ python-dotenv = "^1.0.0"
 azure-identity = "^1.14.0"
 azure-keyvault-secrets = "^4.7.0"
 orjson = "^3.10.3"
+httpx = "^0.27.0"
 
 [tool.poetry.extras]
 experimental = ["cognite-sdk-experimental"]
 
 [tool.poetry.group.dev.dependencies]
-mypy = "1.10.1"
-ruff = "^0.4.0"
+mypy = "1.11.1"
+ruff = "^0.5.0"
 pytest = "^8.0.0"
 pytest-cov = "^5.0.0"
 sphinx = "^7.0.0"
@@ -85,7 +86,7 @@ parameterized = "*"
 requests = "^2.31.0"
 types-requests = "^2.31.0.20240125"
 httpx = "^0.27.0"
-faker = "^25.2.0"
+faker = "^26.0.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]