orca-sdk 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

orca_sdk/conftest.py

@@ -4,11 +4,13 @@ from typing import Generator
 from uuid import uuid4
 
 import pytest
+import pytest_asyncio
 from datasets import ClassLabel, Dataset, Features, Value
 
 from ._utils.auth import _create_api_key, _delete_org
+from .async_client import OrcaAsyncClient
 from .classification_model import ClassificationModel
-from .client import orca_api
+from .client import OrcaClient
 from .credentials import OrcaCredentials
 from .datasource import Datasource
 from .embedding_model import PretrainedEmbeddingModel
@@ -44,51 +46,51 @@ def _create_org_id():
     return "10e50000-0000-4000-a000-" + str(uuid4())[24:]
 
 
-@pytest.fixture()
-def api_url_reset():
-    original_base_url = orca_api.base_url
-    yield
-    orca_api.base_url = original_base_url
+@pytest.fixture(scope="session")
+def org_id():
+    return _create_org_id()
 
 
 @pytest.fixture(scope="session")
-def org_id():
+def other_org_id():
     return _create_org_id()
 
 
 @pytest.fixture(autouse=True, scope="session")
 def api_key(org_id) -> Generator[str, None, None]:
     api_key = _create_api_key(org_id=org_id, name="orca_sdk_test")
-    OrcaCredentials.set_api_key(api_key, check_validity=True)
-    yield api_key
+    with OrcaClient(api_key=api_key).use():
+        yield api_key
     _delete_org(org_id)
 
 
+# We cannot use a session scoped fixture because async pytest tears down the client after each test
 @pytest.fixture(autouse=True)
-def authenticated(api_key):
-    OrcaCredentials.set_api_key(api_key, check_validity=False)
+def authenticate_async_client(api_key) -> Generator[None, None, None]:
+    with OrcaAsyncClient(api_key=api_key).use():
+        yield
 
 
-@pytest.fixture()
-def unauthenticated(api_key):
-    OrcaCredentials.set_api_key(str(uuid4()), check_validity=False)
-    yield
-    # Need to reset the api key to the original api key so following tests don't fail
-    OrcaCredentials.set_api_key(api_key, check_validity=False)
+@pytest.fixture(scope="session")
+def unauthenticated_client() -> OrcaClient:
+    return OrcaClient(api_key=str(uuid4()))
 
 
-@pytest.fixture()
-def other_org_id():
-    return _create_org_id()
+@pytest_asyncio.fixture()
+def unauthenticated_async_client() -> OrcaAsyncClient:
+    return OrcaAsyncClient(api_key=str(uuid4()))
 
 
-@pytest.fixture()
-def unauthorized(api_key, other_org_id):
+@pytest.fixture(scope="session")
+def unauthorized_client(other_org_id):
     different_api_key = _create_api_key(org_id=other_org_id, name="orca_sdk_test_other_org")
-    OrcaCredentials.set_api_key(different_api_key, check_validity=False)
-    yield
-    OrcaCredentials.set_api_key(api_key, check_validity=False)
-    _delete_org(other_org_id)
+    return OrcaClient(api_key=different_api_key)
+
+
+@pytest.fixture(scope="session")
+def predict_only_client() -> OrcaClient:
+    predict_api_key = OrcaCredentials.create_api_key("orca_sdk_test_predict", scopes={"PREDICT"})
+    return OrcaClient(api_key=predict_api_key)
 
 
 @pytest.fixture(scope="session")
@@ -209,18 +211,17 @@ def writable_memoryset(datasource: Datasource, api_key: str) -> Generator[Labele
         yield memoryset
     finally:
         # Restore the memoryset to a clean state for the next test.
-        OrcaCredentials.set_api_key(api_key, check_validity=False)
-
-        if LabeledMemoryset.exists("test_writable_memoryset"):
-            memoryset.refresh()
+        with OrcaClient(api_key=api_key).use():
+            if LabeledMemoryset.exists("test_writable_memoryset"):
+                memoryset.refresh()
 
-            memory_ids = [memoryset[i].memory_id for i in range(len(memoryset))]
+                memory_ids = [memoryset[i].memory_id for i in range(len(memoryset))]
 
-            if memory_ids:
-                memoryset.delete(memory_ids)
-            memoryset.refresh()
-            assert len(memoryset) == 0
-            memoryset.insert(SAMPLE_DATA)
+                if memory_ids:
+                    memoryset.delete(memory_ids)
+                memoryset.refresh()
+                assert len(memoryset) == 0
+                memoryset.insert(SAMPLE_DATA)
         # If the test dropped the memoryset, do nothing — it will be recreated on the next use.
 
 

orca_sdk/credentials.py

@@ -1,10 +1,13 @@
+import os
 from datetime import datetime
 from typing import Literal, NamedTuple
 
 import httpx
-from httpx import ConnectError, Headers
+from httpx import ConnectError, Headers, HTTPTransport
+from typing_extensions import deprecated
 
-from .client import orca_api
+from .async_client import OrcaAsyncClient
+from .client import OrcaClient
 
 Scope = Literal["ADMINISTER", "PREDICT"]
 """
@@ -15,19 +18,31 @@ The scopes of an API key.
 """
 
 
-class ApiKeyInfo(NamedTuple):
+class ApiKeyInfo:
     """
-    Named tuple containing information about an API key
+    Information about an API key
+
+    Note:
+        The value of the API key is only available at creation time.
 
     Attributes:
         name: Unique name of the API key
         created_at: When the API key was created
+        scopes: The scopes of the API key
     """
 
     name: str
     created_at: datetime
     scopes: set[Scope]
 
+    def __init__(self, name: str, created_at: datetime, scopes: set[Scope]):
+        self.name = name
+        self.created_at = created_at
+        self.scopes = scopes
+
+    def __repr__(self) -> str:
+        return "ApiKey({ " + f"name: '{self.name}', scopes: <{'|'.join(self.scopes)}>" + "})"
+
 
 class OrcaCredentials:
     """
@@ -42,8 +57,9 @@ class OrcaCredentials:
         Returns:
             True if you are authenticated, False otherwise
         """
+        client = OrcaClient._resolve_client()
         try:
-            return orca_api.GET("/auth")
+            return client.GET("/auth")
         except ValueError as e:
             if "Invalid API key" in str(e):
                 return False
@@ -57,8 +73,10 @@ class OrcaCredentials:
         Returns:
             True if the API is healthy, False otherwise
         """
+        client = OrcaClient._resolve_client()
         try:
-            orca_api.GET("/check/healthy")
+            # we don't want a retry transport here, so we use httpx directly
+            httpx.get(f"{client.base_url}/check/healthy")
         except Exception:
             return False
         return True
@@ -71,13 +89,14 @@ class OrcaCredentials:
         Returns:
             A list of named tuples, with the name and creation date time of the API key
         """
+        client = OrcaClient._resolve_client()
         return [
             ApiKeyInfo(
                 name=api_key["name"],
                 created_at=datetime.fromisoformat(api_key["created_at"]),
                 scopes=set(api_key["scope"]),
             )
-            for api_key in orca_api.GET("/auth/api_key")
+            for api_key in client.GET("/auth/api_key")
         ]
 
     @staticmethod
@@ -92,7 +111,8 @@ class OrcaCredentials:
         Returns:
             The secret value of the API key. Make sure to save this value as it will not be shown again.
         """
-        res = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        res = client.POST(
             "/auth/api_key",
             json={"name": name, "scope": list(scopes)},
         )
@@ -109,8 +129,12 @@ class OrcaCredentials:
         Raises:
             ValueError: if the API key is not found
         """
-        orca_api.DELETE("/auth/api_key/{name_or_id}", params={"name_or_id": name})
+        client = OrcaClient._resolve_client()
+        client.DELETE("/auth/api_key/{name_or_id}", params={"name_or_id": name})
+
+    # TODO: remove deprecated methods after 2026-01-01
 
+    @deprecated("Use `OrcaClient.api_key` instead")
     @staticmethod
     def set_api_key(api_key: str, check_validity: bool = True):
         """
@@ -126,17 +150,24 @@ class OrcaCredentials:
         Raises:
             ValueError: if the API key is invalid and `check_validity` is True
         """
-        OrcaCredentials.set_api_headers({"Api-Key": api_key})
+        sync_client = OrcaClient._resolve_client()
+        sync_client.api_key = api_key
         if check_validity:
-            orca_api.GET("/auth")
+            sync_client.GET("/auth")
 
+        async_client = OrcaAsyncClient._resolve_client()
+        async_client.api_key = api_key
+
+    @deprecated("Use `OrcaClient.base_url` instead")
     @staticmethod
     def get_api_url() -> str:
         """
         Get the base URL of the Orca API that is currently being used
         """
-        return str(orca_api.base_url)
+        client = OrcaClient._resolve_client()
+        return str(client.base_url)
 
+    @deprecated("Use `OrcaClient.base_url` instead")
     @staticmethod
     def set_api_url(url: str, check_validity: bool = True):
         """
@@ -156,12 +187,17 @@ class OrcaCredentials:
             except ConnectError as e:
                 raise ValueError(f"No API found at {url}") from e
 
-        orca_api.base_url = url
+        sync_client = OrcaClient._resolve_client()
+        sync_client.base_url = url
+
+        async_client = OrcaAsyncClient._resolve_client()
+        async_client.base_url = url
 
         # check if the api passes the health check
         if check_validity:
             OrcaCredentials.is_healthy()
 
+    @deprecated("Use `OrcaClient.headers` instead")
     @staticmethod
     def set_api_headers(headers: dict[str, str]):
         """
@@ -174,4 +210,8 @@ class OrcaCredentials:
             New keys are merged into the existing headers, this will overwrite headers with the
             same name, but leave other headers untouched.
         """
-        orca_api.headers.update(Headers(headers))
+        sync_client = OrcaClient._resolve_client()
+        sync_client.headers.update(Headers(headers))
+
+        async_client = OrcaAsyncClient._resolve_client()
+        async_client.headers.update(Headers(headers))

orca_sdk/credentials_test.py

@@ -2,56 +2,120 @@ from uuid import uuid4
 
 import pytest
 
-from .client import orca_api
+from .client import OrcaClient
 from .credentials import OrcaCredentials
 
 
+def test_is_authenticated():
+    assert OrcaCredentials.is_authenticated()
+
+
+def test_is_authenticated_false(unauthenticated_client):
+    with unauthenticated_client.use():
+        assert not OrcaCredentials.is_authenticated()
+
+
+def test_is_healthy():
+    assert OrcaCredentials.is_healthy()
+
+
+def test_is_healthy_false(api_key):
+    with OrcaClient(api_key=api_key, base_url="http://localhost:1582").use():
+        assert not OrcaCredentials.is_healthy()
+
+
 def test_list_api_keys():
     api_keys = OrcaCredentials.list_api_keys()
     assert len(api_keys) >= 1
     assert "orca_sdk_test" in [api_key.name for api_key in api_keys]
 
 
-def test_list_api_keys_unauthenticated(unauthenticated):
-    with pytest.raises(ValueError, match="Invalid API key"):
-        OrcaCredentials.list_api_keys()
+def test_list_api_keys_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            OrcaCredentials.list_api_keys()
 
 
-def test_is_authenticated():
-    assert OrcaCredentials.is_authenticated()
+def test_manage_api_key():
+    api_key_name = f"orca_sdk_test_{uuid4().hex[:8]}"
+    api_key = OrcaCredentials.create_api_key(api_key_name)
+    assert api_key is not None
+    assert len(api_key) > 0
+    assert api_key_name in [aki.name for aki in OrcaCredentials.list_api_keys()]
+    OrcaCredentials.revoke_api_key(api_key_name)
+    assert api_key_name not in [aki.name for aki in OrcaCredentials.list_api_keys()]
 
 
-def test_is_authenticated_false(unauthenticated):
-    assert not OrcaCredentials.is_authenticated()
+def test_create_api_key_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            OrcaCredentials.create_api_key(f"orca_sdk_test_{uuid4().hex[:8]}")
 
 
-def test_set_api_key(api_key, unauthenticated):
-    assert not OrcaCredentials.is_authenticated()
-    OrcaCredentials.set_api_key(api_key)
-    assert OrcaCredentials.is_authenticated()
+def test_create_api_key_unauthorized(predict_only_client):
+    with predict_only_client.use():
+        with pytest.raises(PermissionError):
+            OrcaCredentials.create_api_key(f"orca_sdk_test_{uuid4().hex[:8]}")
 
 
-def test_set_invalid_api_key(api_key):
-    assert OrcaCredentials.is_authenticated()
-    with pytest.raises(ValueError, match="Invalid API key"):
-        OrcaCredentials.set_api_key(str(uuid4()))
-    assert not OrcaCredentials.is_authenticated()
+def test_revoke_api_key_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            OrcaCredentials.revoke_api_key(f"orca_sdk_test_{uuid4().hex[:8]}")
 
 
-def test_set_api_url(api_url_reset):
-    OrcaCredentials.set_api_url("http://api.orcadb.ai")
-    assert str(orca_api.base_url) == "http://api.orcadb.ai"
+def test_revoke_api_key_unauthorized(predict_only_client):
+    with predict_only_client.use():
+        with pytest.raises(PermissionError):
+            OrcaCredentials.revoke_api_key(f"orca_sdk_test_{uuid4().hex[:8]}")
 
 
-def test_set_invalid_base_url():
-    with pytest.raises(ValueError, match="No API found at http://localhost:1582"):
-        OrcaCredentials.set_api_url("http://localhost:1582")
+def test_create_api_key_already_exists():
+    with pytest.raises(ValueError, match="API key with this name already exists"):
+        OrcaCredentials.create_api_key("orca_sdk_test")
 
 
-def test_is_healthy():
-    assert OrcaCredentials.is_healthy()
+def test_set_api_key(api_key):
+    client = OrcaClient(api_key=str(uuid4()))
+    with client.use():
+        assert not OrcaCredentials.is_authenticated()
+        client.api_key = api_key
+        assert client.api_key == api_key
+        assert OrcaCredentials.is_authenticated()
+
+
+def test_set_base_url(api_key):
+    client = OrcaClient(base_url="http://localhost:1582")
+    assert client.base_url == "http://localhost:1582"
+    client.base_url = "http://localhost:1583"
+    assert client.base_url == "http://localhost:1583"
+
+
+# deprecated methods:
+
+
+def test_deprecated_set_api_key(api_key):
+    with OrcaClient(api_key=str(uuid4())).use():
+        assert not OrcaCredentials.is_authenticated()
+        OrcaCredentials.set_api_key(api_key)
+        assert OrcaCredentials.is_authenticated()
+
+
+def test_deprecated_set_invalid_api_key(api_key):
+    with OrcaClient(api_key=api_key).use():
+        assert OrcaCredentials.is_authenticated()
+        with pytest.raises(ValueError, match="Invalid API key"):
+            OrcaCredentials.set_api_key(str(uuid4()))
+        assert not OrcaCredentials.is_authenticated()
+
+
+def test_deprecated_set_api_url(api_key):
+    with OrcaClient(api_key=api_key).use():
+        OrcaCredentials.set_api_url("http://api.orcadb.ai")
+        assert str(OrcaClient._resolve_client().base_url) == "http://api.orcadb.ai"
 
 
-def test_is_healthy_false(api_url_reset):
-    OrcaCredentials.set_api_url("http://localhost:1582", check_validity=False)
-    assert not OrcaCredentials.is_healthy()
+def test_deprecated_set_invalid_api_url(api_key):
+    with OrcaClient(api_key=api_key).use():
+        with pytest.raises(ValueError, match="No API found at http://localhost:1582"):
+            OrcaCredentials.set_api_url("http://localhost:1582")

orca_sdk/datasource.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
 import logging
-import os
 import tempfile
 import zipfile
 from datetime import datetime
 from io import BytesIO
 from os import PathLike
 from pathlib import Path
-from typing import Literal, Union, cast
+from typing import Any, Literal, Union, cast
 
 import pandas as pd
 import pyarrow as pa
@@ -22,7 +21,7 @@ from tqdm.auto import tqdm
 from ._utils.common import CreateMode, DropMode
 from ._utils.data_parsing import hf_dataset_from_torch
 from ._utils.tqdm_file_reader import TqdmFileReader
-from .client import DatasourceMetadata, orca_api
+from .client import DatasourceMetadata, OrcaClient
 
 
 def _upload_files_to_datasource(
@@ -55,7 +54,8 @@ def _upload_files_to_datasource(
             files.append(("files", (file_path.name, cast(bytes, tqdm_reader))))
 
         # Use manual HTTP request for file uploads
-        metadata = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        metadata = client.POST(
             "/datasource/upload",
             files=files,
             data={"name": name, "description": description},
@@ -268,7 +268,8 @@ class Datasource:
         if existing is not None:
             return existing
 
-        metadata = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        metadata = client.POST(
             "/datasource",
             json={"name": name, "description": description, "content": data},
         )
@@ -302,7 +303,8 @@ class Datasource:
         if existing is not None:
             return existing
 
-        metadata = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        metadata = client.POST(
             "/datasource",
             json={"name": name, "description": description, "content": data},
         )
@@ -361,7 +363,8 @@ class Datasource:
         parquet.write_table(pyarrow_table, buffer)
         parquet_bytes = buffer.getvalue()
 
-        metadata = orca_api.POST(
+        client = OrcaClient._resolve_client()
+        metadata = client.POST(
             "/datasource/upload",
             files=[("files", ("data.parquet", parquet_bytes))],
             data={"name": name, "description": description},
@@ -429,7 +432,8 @@ class Datasource:
         Raises:
             LookupError: If the datasource does not exist
         """
-        return cls(orca_api.GET("/datasource/{name_or_id}", params={"name_or_id": name_or_id}))
+        client = OrcaClient._resolve_client()
+        return cls(client.GET("/datasource/{name_or_id}", params={"name_or_id": name_or_id}))
 
     @classmethod
     def exists(cls, name_or_id: str) -> bool:
@@ -456,7 +460,8 @@ class Datasource:
         Returns:
             A list of all datasource handles in the OrcaCloud
         """
-        return [cls(metadata) for metadata in orca_api.GET("/datasource")]
+        client = OrcaClient._resolve_client()
+        return [cls(metadata) for metadata in client.GET("/datasource")]
 
     @classmethod
     def drop(cls, name_or_id: str, if_not_exists: DropMode = "error") -> None:
@@ -472,7 +477,8 @@ class Datasource:
             LookupError: If the datasource does not exist and if_not_exists is `"error"`
         """
         try:
-            orca_api.DELETE("/datasource/{name_or_id}", params={"name_or_id": name_or_id})
+            client = OrcaClient._resolve_client()
+            client.DELETE("/datasource/{name_or_id}", params={"name_or_id": name_or_id})
             logging.info(f"Deleted datasource {name_or_id}")
         except LookupError:
             if if_not_exists == "error":
@@ -481,6 +487,50 @@ class Datasource:
     def __len__(self) -> int:
         return self.length
 
+    def query(
+        self,
+        offset: int = 0,
+        limit: int = 100,
+        shuffle: bool = False,
+        shuffle_seed: int | None = None,
+        filters: list[tuple[str, Literal["==", "!=", ">", ">=", "<", "<=", "in", "not in", "like"], Any]] = [],
+    ) -> list[dict[str, Any]]:
+        """
+        Query the datasource for rows with pagination and filtering support.
+
+        Params:
+            offset: Number of rows to skip
+            limit: Maximum number of rows to return
+            shuffle: Whether to shuffle the dataset before pagination
+            shuffle_seed: Seed for shuffling (for reproducible results)
+            filters: List of filter tuples. Each tuple contains:
+                - field (str): Column name to filter on
+                - op (str): Operator ("==", "!=", ">", ">=", "<", "<=", "in", "not in", "like")
+                - value: Value to compare against
+
+        Returns:
+            List of rows from the datasource
+
+        Examples:
+            >>> datasource.query(filters=[("age", ">", 25)])
+            >>> datasource.query(filters=[("city", "in", ["NYC", "LA"])])
+            >>> datasource.query(filters=[("name", "like", "John")])
+        """
+
+        client = OrcaClient._resolve_client()
+        response = client.POST(
+            "/datasource/{name_or_id}/rows",
+            params={"name_or_id": self.id},
+            json={
+                "limit": limit,
+                "offset": offset,
+                "shuffle": shuffle,
+                "shuffle_seed": shuffle_seed,
+                "filters": [{"field": field, "op": op, "value": value} for field, op, value in filters],
+            },
+        )
+        return response
+
     def download(
         self, output_dir: str | PathLike, file_type: Literal["hf_dataset", "json", "csv"] = "hf_dataset"
     ) -> None:
@@ -497,7 +547,8 @@ class Datasource:
         extension = "zip" if file_type == "hf_dataset" else file_type
         output_path = Path(output_dir) / f"{self.name}.{extension}"
         with open(output_path, "wb") as download_file:
-            with orca_api.stream("GET", f"/datasource/{self.id}/download", params={"file_type": file_type}) as response:
+            client = OrcaClient._resolve_client()
+            with client.stream("GET", f"/datasource/{self.id}/download", params={"file_type": file_type}) as response:
                 total_chunks = int(response.headers["X-Total-Chunks"]) if "X-Total-Chunks" in response.headers else None
                 with tqdm(desc="Downloading", total=total_chunks, disable=total_chunks is None) as progress:
                     for chunk in response.iter_bytes():
@@ -521,4 +572,5 @@ class Datasource:
         Returns:
             A list of dictionaries representation of the datasource.
         """
-        return orca_api.GET("/datasource/{name_or_id}/download", params={"name_or_id": self.id, "file_type": "json"})
+        client = OrcaClient._resolve_client()
+        return client.GET("/datasource/{name_or_id}/download", params={"name_or_id": self.id, "file_type": "json"})