datachain 0.14.2__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/remote/studio.py

@@ -1,47 +1,64 @@
-import base64
 import json
 import logging
 import os
 from collections.abc import AsyncIterator, Iterable, Iterator
 from datetime import datetime, timedelta, timezone
 from struct import unpack
-from typing import (
-    Any,
-    Generic,
-    Optional,
-    TypeVar,
-)
+from typing import Any, BinaryIO, Generic, TypeVar
 from urllib.parse import urlparse, urlunparse
 
 import websockets
 from requests.exceptions import HTTPError, Timeout
 
 from datachain.config import Config
+from datachain.dataset import DatasetRecord
 from datachain.error import DataChainError
 from datachain.utils import STUDIO_URL, retry_with_backoff
 
 T = TypeVar("T")
-LsData = Optional[list[dict[str, Any]]]
-DatasetInfoData = Optional[dict[str, Any]]
-DatasetRowsData = Optional[Iterable[dict[str, Any]]]
-DatasetJobVersionsData = Optional[dict[str, Any]]
-DatasetExportStatus = Optional[dict[str, Any]]
-DatasetExportSignedUrls = Optional[list[str]]
-FileUploadData = Optional[dict[str, Any]]
-JobData = Optional[dict[str, Any]]
+LsData = list[dict[str, Any]] | None
+DatasetInfoData = dict[str, Any] | None
+DatasetRowsData = Iterable[dict[str, Any]] | None
+DatasetJobVersionsData = dict[str, Any] | None
+DatasetExportStatus = dict[str, Any] | None
+DatasetExportSignedUrls = list[str] | None
+FileUploadData = dict[str, Any] | None
+JobData = dict[str, Any] | None
+JobListData = list[dict[str, Any]]
+ClusterListData = list[dict[str, Any]]
 
 logger = logging.getLogger("datachain")
 
 DATASET_ROWS_CHUNK_SIZE = 8192
 
 
+def get_studio_env_variable(name: str) -> Any:
+    """
+    Get the value of a DataChain Studio environment variable.
+    It first checks for the variable prefixed with 'DATACHAIN_STUDIO_',
+    then checks for the deprecated 'DVC_STUDIO_' prefix.
+    If neither is set, it returns the provided default value.
+    """
+    if (value := os.environ.get(f"DATACHAIN_STUDIO_{name}")) is not None:
+        return value
+    if (value := os.environ.get(f"DVC_STUDIO_{name}")) is not None:  # deprecated
+        logger.warning(
+            "Environment variable 'DVC_STUDIO_%s' is deprecated, "
+            "use 'DATACHAIN_STUDIO_%s' instead.",
+            name,
+            name,
+        )
+        return value
+    return None
+
+
 def _is_server_error(status_code: int) -> bool:
     return str(status_code).startswith("5")
 
 
 def is_token_set() -> bool:
     return (
-        bool(os.environ.get("DVC_STUDIO_TOKEN"))
+        bool(get_studio_env_variable("TOKEN"))
         or Config().read().get("studio", {}).get("token") is not None
     )
 
@@ -56,10 +73,11 @@ def _parse_dates(obj: dict, date_fields: list[str]):
 
 
 class Response(Generic[T]):
-    def __init__(self, data: T, ok: bool, message: str) -> None:
+    def __init__(self, data: T, ok: bool, message: str, status: int) -> None:
         self.data = data
         self.ok = ok
         self.message = message
+        self.status = status
 
     def __repr__(self):
         return (
@@ -69,7 +87,7 @@ class Response(Generic[T]):
 
 
 class StudioClient:
-    def __init__(self, timeout: float = 3600.0, team: Optional[str] = None) -> None:
+    def __init__(self, timeout: float = 3600.0, team: str | None = None) -> None:
         self._check_dependencies()
         self.timeout = timeout
         self._config = None
@@ -77,12 +95,12 @@ class StudioClient:
 
     @property
     def token(self) -> str:
-        token = os.environ.get("DVC_STUDIO_TOKEN") or self.config.get("token")
+        token = get_studio_env_variable("TOKEN") or self.config.get("token")
 
         if not token:
             raise DataChainError(
                 "Studio token is not set. Use `datachain auth login` "
-                "or environment variable `DVC_STUDIO_TOKEN` to set it."
+                "or environment variable `DATACHAIN_STUDIO_TOKEN` to set it."
             )
 
         return token
@@ -90,8 +108,8 @@ class StudioClient:
     @property
     def url(self) -> str:
         return (
-            os.environ.get("DVC_STUDIO_URL") or self.config.get("url") or STUDIO_URL
-        ) + "/api"
+            get_studio_env_variable("URL") or self.config.get("url") or STUDIO_URL
+        ).rstrip("/") + "/api"
 
     @property
     def config(self) -> dict:
@@ -106,13 +124,13 @@ class StudioClient:
         return self._team
 
     def _get_team(self) -> str:
-        team = os.environ.get("DVC_STUDIO_TEAM") or self.config.get("team")
+        team = get_studio_env_variable("TEAM") or self.config.get("team")
 
         if not team:
             raise DataChainError(
                 "Studio team is not set. "
                 "Use `datachain auth team <team_name>` "
-                "or environment variable `DVC_STUDIO_TEAM` to set it. "
+                "or environment variable `DATACHAIN_STUDIO_TEAM` to set it. "
                 "You can also set `studio.team` in the config file."
             )
 
@@ -130,7 +148,7 @@ class StudioClient:
             ) from None
 
     def _send_request_msgpack(
-        self, route: str, data: dict[str, Any], method: Optional[str] = "POST"
+        self, route: str, data: dict[str, Any], method: str | None = "POST"
     ) -> Response[Any]:
         import msgpack
         import requests
@@ -164,11 +182,11 @@ class StudioClient:
             message = "Indexing in progress"
         else:
             message = content.get("message", "")
-        return Response(response_data, ok, message)
+        return Response(response_data, ok, message, response.status_code)
 
     @retry_with_backoff(retries=3, errors=(HTTPError, Timeout))
     def _send_request(
-        self, route: str, data: dict[str, Any], method: Optional[str] = "POST"
+        self, route: str, data: dict[str, Any], method: str | None = "POST"
     ) -> Response[Any]:
         """
         Function that communicate Studio API.
@@ -214,7 +232,46 @@ class StudioClient:
         else:
             message = ""
 
-        return Response(data, ok, message)
+        return Response(data, ok, message, response.status_code)
+
+    def _send_multipart_request(
+        self, route: str, files: dict[str, Any], params: dict[str, Any] | None = None
+    ) -> Response[Any]:
+        """
+        Function that communicates with Studio API using multipart/form-data.
+        It will raise an exception, and try to retry, if 5xx status code is
+        returned, or if Timeout exceptions is thrown from the requests lib
+        """
+        import requests
+
+        # Add team_name to params
+        request_params = {**(params or {}), "team_name": self.team}
+
+        response = requests.post(
+            url=f"{self.url}/{route}",
+            files=files,
+            params=request_params,
+            headers={
+                "Authorization": f"token {self.token}",
+            },
+            timeout=self.timeout,
+        )
+
+        ok = response.ok
+        try:
+            data = json.loads(response.content.decode("utf-8"))
+        except json.decoder.JSONDecodeError:
+            data = {}
+
+        if not ok:
+            if response.status_code == 403:
+                message = f"Not authorized for the team {self.team}"
+            else:
+                message = data.get("message", "")
+        else:
+            message = ""
+
+        return Response(data, ok, message, response.status_code)
 
     @staticmethod
     def _unpacker_hook(code, data):
@@ -282,21 +339,27 @@ class StudioClient:
             response = self._send_request_msgpack("datachain/ls", {"source": path})
             yield path, response
 
-    def ls_datasets(self) -> Response[LsData]:
-        return self._send_request("datachain/datasets", {}, method="GET")
+    def ls_datasets(self, prefix: str | None = None) -> Response[LsData]:
+        return self._send_request(
+            "datachain/datasets", {"prefix": prefix}, method="GET"
+        )
 
     def edit_dataset(
         self,
         name: str,
-        new_name: Optional[str] = None,
-        description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        namespace: str,
+        project: str,
+        new_name: str | None = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
     ) -> Response[DatasetInfoData]:
         body = {
             "new_name": new_name,
-            "dataset_name": name,
+            "name": name,
+            "namespace": namespace,
+            "project": project,
             "description": description,
-            "labels": labels,
+            "attrs": attrs,
         }
 
         return self._send_request(
@@ -307,44 +370,44 @@ class StudioClient:
     def rm_dataset(
         self,
         name: str,
-        version: Optional[int] = None,
-        force: Optional[bool] = False,
+        namespace: str,
+        project: str,
+        version: str | None = None,
+        force: bool | None = False,
     ) -> Response[DatasetInfoData]:
         return self._send_request(
             "datachain/datasets",
             {
-                "dataset_name": name,
-                "dataset_version": version,
+                "name": name,
+                "namespace": namespace,
+                "project": project,
+                "version": version,
                 "force": force,
             },
             method="DELETE",
         )
 
-    def dataset_info(self, name: str) -> Response[DatasetInfoData]:
+    def dataset_info(
+        self, namespace: str, project: str, name: str
+    ) -> Response[DatasetInfoData]:
         def _parse_dataset_info(dataset_info):
             _parse_dates(dataset_info, ["created_at", "finished_at"])
             for version in dataset_info.get("versions"):
                 _parse_dates(version, ["created_at"])
+            _parse_dates(dataset_info.get("project"), ["created_at"])
+            _parse_dates(dataset_info.get("project").get("namespace"), ["created_at"])
 
             return dataset_info
 
         response = self._send_request(
-            "datachain/datasets/info", {"dataset_name": name}, method="GET"
+            "datachain/datasets/info",
+            {"namespace": namespace, "project": project, "name": name},
+            method="GET",
         )
         if response.ok:
             response.data = _parse_dataset_info(response.data)
         return response
 
-    def dataset_rows_chunk(
-        self, name: str, version: int, offset: int
-    ) -> Response[DatasetRowsData]:
-        req_data = {"dataset_name": name, "dataset_version": version}
-        return self._send_request_msgpack(
-            "datachain/datasets/rows",
-            {**req_data, "offset": offset, "limit": DATASET_ROWS_CHUNK_SIZE},
-            method="GET",
-        )
-
     def dataset_job_versions(self, job_id: str) -> Response[DatasetJobVersionsData]:
         return self._send_request(
             "datachain/datasets/dataset_job_versions",
@@ -353,40 +416,57 @@ class StudioClient:
         )
 
     def export_dataset_table(
-        self, name: str, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> Response[DatasetExportSignedUrls]:
         return self._send_request(
             "datachain/datasets/export",
-            {"dataset_name": name, "dataset_version": version},
+            {
+                "namespace": dataset.project.namespace.name,
+                "project": dataset.project.name,
+                "name": dataset.name,
+                "version": version,
+            },
             method="GET",
         )
 
     def dataset_export_status(
-        self, name: str, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> Response[DatasetExportStatus]:
         return self._send_request(
             "datachain/datasets/export-status",
-            {"dataset_name": name, "dataset_version": version},
+            {
+                "namespace": dataset.project.namespace.name,
+                "project": dataset.project.name,
+                "name": dataset.name,
+                "version": version,
+            },
             method="GET",
         )
 
-    def upload_file(self, content: bytes, file_name: str) -> Response[FileUploadData]:
-        data = {
-            "file_content": base64.b64encode(content).decode("utf-8"),
-            "file_name": file_name,
-        }
-        return self._send_request("datachain/upload-file", data)
+    def upload_file(
+        self, file_obj: BinaryIO, file_name: str
+    ) -> Response[FileUploadData]:
+        # Prepare multipart form data
+        files = {"file": (file_name, file_obj, "application/octet-stream")}
+
+        return self._send_multipart_request("datachain/jobs/files", files)
 
     def create_job(
         self,
         query: str,
         query_type: str,
-        environment: Optional[str] = None,
-        workers: Optional[int] = None,
-        query_name: Optional[str] = None,
-        files: Optional[list[str]] = None,
-        python_version: Optional[str] = None,
-        requirements: Optional[str] = None,
+        environment: str | None = None,
+        workers: int | None = None,
+        query_name: str | None = None,
+        files: list[str] | None = None,
+        python_version: str | None = None,
+        requirements: str | None = None,
+        repository: str | None = None,
+        priority: int | None = None,
+        cluster: str | None = None,
+        start_time: str | None = None,
+        cron: str | None = None,
+        credentials_name: str | None = None,
     ) -> Response[JobData]:
         data = {
             "query": query,
@@ -397,12 +477,34 @@ class StudioClient:
             "files": files,
             "python_version": python_version,
             "requirements": requirements,
+            "repository": repository,
+            "priority": priority,
+            "compute_cluster_name": cluster,
+            "start_after": start_time,
+            "cron_expression": cron,
+            "credentials_name": credentials_name,
         }
-        return self._send_request("datachain/job", data)
+        return self._send_request("datachain/jobs/", data)
+
+    def get_jobs(
+        self,
+        status: str | None = None,
+        limit: int = 20,
+        job_id: str | None = None,
+    ) -> Response[JobListData]:
+        params: dict[str, Any] = {"limit": limit}
+        if status is not None:
+            params["status"] = status
+        if job_id is not None:
+            params["job_id"] = job_id
+        return self._send_request("datachain/jobs/", params, method="GET")
 
     def cancel_job(
         self,
         job_id: str,
     ) -> Response[JobData]:
-        url = f"datachain/job/{job_id}/cancel"
+        url = f"datachain/jobs/{job_id}/cancel"
         return self._send_request(url, data={}, method="POST")
+
+    def get_clusters(self) -> Response[ClusterListData]:
+        return self._send_request("datachain/clusters/", {}, method="GET")

datachain/script_meta.py

@@ -1,6 +1,6 @@
 import re
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 try:
     import tomllib
@@ -59,23 +59,23 @@ class ScriptConfig:
 
     """
 
-    python_version: Optional[str]
+    python_version: str | None
     dependencies: list[str]
     attachments: dict[str, str]
     params: dict[str, Any]
     inputs: dict[str, Any]
     outputs: dict[str, Any]
-    num_workers: Optional[int] = None
+    num_workers: int | None = None
 
     def __init__(
         self,
-        python_version: Optional[str] = None,
-        dependencies: Optional[list[str]] = None,
-        attachments: Optional[dict[str, str]] = None,
-        params: Optional[dict[str, Any]] = None,
-        inputs: Optional[dict[str, Any]] = None,
-        outputs: Optional[dict[str, Any]] = None,
-        num_workers: Optional[int] = None,
+        python_version: str | None = None,
+        dependencies: list[str] | None = None,
+        attachments: dict[str, str] | None = None,
+        params: dict[str, Any] | None = None,
+        inputs: dict[str, Any] | None = None,
+        outputs: dict[str, Any] | None = None,
+        num_workers: int | None = None,
     ):
         self.python_version = python_version
         self.dependencies = dependencies or []
@@ -98,7 +98,7 @@ class ScriptConfig:
         return self.attachments.get(name, default)
 
     @staticmethod
-    def read(script: str) -> Optional[dict]:
+    def read(script: str) -> dict | None:
         """Converts inline script metadata to dict with all found data"""
         regex = (
             r"(?m)^# \/\/\/ (?P<type>[a-zA-Z0-9-]+)[ \t]*$[\r\n|\r|\n]"
@@ -119,7 +119,7 @@ class ScriptConfig:
         return None
 
     @staticmethod
-    def parse(script: str) -> Optional["ScriptConfig"]:
+    def parse(script: str) -> "ScriptConfig | None":
         """
         Method that is parsing inline script metadata from datachain script and
         instantiating ScriptConfig class with found data. If no inline metadata is

datachain/semver.py

@@ -0,0 +1,68 @@
+# Maximum version number for semver (major.minor.patch) is 999999.999999.999999
+# this number was chosen because value("999999.999999.999999") < 2**63 - 1
+MAX_VERSION_NUMBER = 999_999
+
+
+def parse(version: str) -> tuple[int, int, int]:
+    """Parsing semver into 3 integers: major, minor, patch"""
+    validate(version)
+    parts = version.split(".")
+    return int(parts[0]), int(parts[1]), int(parts[2])
+
+
+def validate(version: str) -> None:
+    """
+    Raises exception if version doesn't have valid semver format which is:
+    <major>.<minor>.<patch> or one of version parts is not positive integer
+    """
+    error_message = (
+        "Invalid version. It should be in format: <major>.<minor>.<patch> where"
+        " each version part is positive integer"
+    )
+    parts = version.split(".")
+    if len(parts) != 3:
+        raise ValueError(error_message)
+    for part in parts:
+        try:
+            val = int(part)
+            assert 0 <= val <= MAX_VERSION_NUMBER
+        except (ValueError, AssertionError):
+            raise ValueError(error_message) from None
+
+
+def create(major: int = 0, minor: int = 0, patch: int = 0) -> str:
+    """Creates new semver from 3 integers: major, minor and patch"""
+    if not (
+        0 <= major <= MAX_VERSION_NUMBER
+        and 0 <= minor <= MAX_VERSION_NUMBER
+        and 0 <= patch <= MAX_VERSION_NUMBER
+    ):
+        raise ValueError("Major, minor and patch must be greater or equal to zero")
+
+    return ".".join([str(major), str(minor), str(patch)])
+
+
+def value(version: str) -> int:
+    """
+    Calculate integer value of a version. This is useful when comparing two versions.
+    """
+    major, minor, patch = parse(version)
+    limit = MAX_VERSION_NUMBER + 1
+    return major * (limit**2) + minor * limit + patch
+
+
+def compare(v1: str, v2: str) -> int:
+    """
+    Compares 2 versions and returns:
+       -1 if v1 < v2
+        0 if v1 == v2
+        1 if v1 > v2
+    """
+    v1_val = value(v1)
+    v2_val = value(v2)
+
+    if v1_val < v2_val:
+        return -1
+    if v1_val > v2_val:
+        return 1
+    return 0

datachain/sql/__init__.py

@@ -1,6 +1,8 @@
 from sqlalchemy.sql.elements import literal
 from sqlalchemy.sql.expression import column
 
+# Import PostgreSQL dialect registration (registers PostgreSQL type converter)
+from . import postgresql_dialect  # noqa: F401
 from .default import setup as default_setup
 from .selectable import select, values
 

datachain/sql/functions/array.py

@@ -1,6 +1,6 @@
 from sqlalchemy.sql.functions import GenericFunction
 
-from datachain.sql.types import Boolean, Float, Int64
+from datachain.sql.types import Boolean, Float, Int64, String
 from datachain.sql.utils import compiler_not_implemented
 
 
@@ -48,6 +48,37 @@ class contains(GenericFunction):  # noqa: N801
     inherit_cache = True
 
 
+class slice(GenericFunction):  # noqa: N801
+    """
+    Returns a slice of the array.
+    """
+
+    package = "array"
+    name = "slice"
+    inherit_cache = True
+
+
+class join(GenericFunction):  # noqa: N801
+    """
+    Returns the concatenation of the array elements.
+    """
+
+    type = String()
+    package = "array"
+    name = "join"
+    inherit_cache = True
+
+
+class get_element(GenericFunction):  # noqa: N801
+    """
+    Returns the element at the given index in the array.
+    """
+
+    package = "array"
+    name = "get_element"
+    inherit_cache = True
+
+
 class sip_hash_64(GenericFunction):  # noqa: N801
     """
     Computes the SipHash-64 hash of the array.
@@ -63,4 +94,5 @@ compiler_not_implemented(cosine_distance)
 compiler_not_implemented(euclidean_distance)
 compiler_not_implemented(length)
 compiler_not_implemented(contains)
+compiler_not_implemented(get_element)
 compiler_not_implemented(sip_hash_64)

datachain/sql/postgresql_dialect.py

@@ -0,0 +1,9 @@
+"""
+PostgreSQL dialect registration for DataChain.
+"""
+
+from datachain.sql.postgresql_types import PostgreSQLTypeConverter
+from datachain.sql.types import register_backend_types
+
+# Register PostgreSQL type converter
+register_backend_types("postgresql", PostgreSQLTypeConverter())

datachain/sql/postgresql_types.py

@@ -0,0 +1,21 @@
+"""
+PostgreSQL-specific type converter for DataChain.
+
+Handles PostgreSQL-specific type mappings that differ from the default dialect.
+"""
+
+from sqlalchemy.dialects import postgresql
+
+from datachain.sql.types import TypeConverter
+
+
+class PostgreSQLTypeConverter(TypeConverter):
+    """PostgreSQL-specific type converter."""
+
+    def datetime(self):
+        """PostgreSQL uses TIMESTAMP WITH TIME ZONE to preserve timezone information."""
+        return postgresql.TIMESTAMP(timezone=True)
+
+    def json(self):
+        """PostgreSQL uses JSONB for better performance and query capabilities."""
+        return postgresql.JSONB()

datachain/sql/sqlite/__init__.py

@@ -1,4 +1,8 @@
-from .base import create_user_defined_sql_functions, setup, sqlite_dialect
+from .base import (
+    create_user_defined_sql_functions,
+    setup,
+    sqlite_dialect,
+)
 
 __all__ = [
     "create_user_defined_sql_functions",