diff --git a/next_cvat/client/client.py b/next_cvat/client/client.py index d2833d9..197407c 100644 --- a/next_cvat/client/client.py +++ b/next_cvat/client/client.py @@ -5,6 +5,7 @@ from cvat_sdk import Client as CVATClient from cvat_sdk import make_client +from cvat_sdk.api_client.exceptions import ApiException from pydantic import BaseModel, field_validator from next_cvat.access_token import AccessToken @@ -27,6 +28,11 @@ class Client(BaseModel): token: str | None = None host: str = DEFAULT_CVAT_HOST + #: Organization slug (e.g. ``nextml``) that requests are made in the context of. + #: ``None`` uses whatever organization the requested resource belongs to, ``""`` + #: the personal workspace. Override via ``CVAT_ORGANIZATION``. + organization: str | None = None + @field_validator("host", mode="before") @classmethod def _normalize_host(cls, value: str | None) -> str: @@ -76,23 +82,55 @@ def basic_cvat_client(self) -> Generator[CVATClient, None, None]: host=self.host, credentials=(self.username, self.password) ) as client: client.login((self.username, self.password)) + client.organization_slug = self.organization yield client @contextmanager def token_cvat_client(self) -> Generator[CVATClient, None, None]: with make_client(host=self.host) as client: - token = AccessToken.deserialize(self.token) - - # Only set Authorization header if we have a real API key (not session-based) - if token.api_key != "session-based-auth": - client.api_client.set_default_header( - "Authorization", f"Token {token.api_key}" - ) + try: + token = AccessToken.deserialize(self.token) + except ValueError: + # A personal access token created in the CVAT UI, rather than a + # session serialized by create_token(). + self.authenticate_with_access_token_(client, self.token) + else: + # Only set Authorization header if we have a real API key (not session-based) + if token.api_key != "session-based-auth": + client.api_client.set_default_header( + "Authorization", f"Token {token.api_key}" + ) + + client.api_client.cookies["sessionid"] = token.sessionid + client.api_client.cookies["csrftoken"] = token.csrftoken + + client.organization_slug = self.organization + yield client - client.api_client.cookies["sessionid"] = token.sessionid - client.api_client.cookies["csrftoken"] = token.csrftoken + @staticmethod + def authenticate_with_access_token_(client: CVATClient, token: str) -> str: + """Authenticate ``client`` with a personal access token, returning the scheme used. - yield client + Which scheme a token wants depends on how the server issued it: tokens from + ``POST /api/auth/login`` are ``Token``, while personal access tokens created in + the UI are ``Bearer``. They are not reliably distinguishable by shape, so try + both and keep whichever the server accepts. + """ + for scheme in ("Bearer", "Token"): + client.api_client.set_default_header("Authorization", f"{scheme} {token}") + try: + client.users.retrieve_current_user() + except ApiException as exception: + if exception.status not in (401, 403): + raise + else: + return scheme + + client.api_client.default_headers.pop("Authorization", None) + raise ValueError( + "The token was rejected by the server as both a Bearer and a Token " + "credential. It may have expired, or belong to a different CVAT host." + ) def create_token(self) -> AccessToken: with self.basic_cvat_client() as client: diff --git a/next_cvat/client/project.py b/next_cvat/client/project.py index be56b1d..d2a2cb5 100644 --- a/next_cvat/client/project.py +++ b/next_cvat/client/project.py @@ -88,8 +88,12 @@ def create_task_( # Get project details to get the organization ID project = client.projects.retrieve(self.id) - # Set organization header - client.api_client.set_default_header("X-Organization", "NextMLAB") + # Creating a task in an organization's project requires that organization's + # context. Unless one was configured explicitly, use the project's own. + if self.client.organization is None and project.organization is not None: + client.organization_slug = client.organizations.retrieve( + project.organization + ).slug # Create task in the project spec = models.TaskWriteRequest( diff --git a/next_cvat/settings.py b/next_cvat/settings.py index f1d9650..38e05f4 100644 --- a/next_cvat/settings.py +++ b/next_cvat/settings.py @@ -16,5 +16,6 @@ class Settings( password: str | None = None token: str | None = None host: str | None = None + organization: str | None = None return Settings()