aqora 0.31.3.dev0__tar.gz → 0.32.0.dev0__tar.gz

{aqora-0.31.3.dev0 → aqora-0.32.0.dev0}/Cargo.lock

@@ -139,7 +139,7 @@ dependencies = [
 
 [[package]]
 name = "aqora"
-version = "0.31.3-dev.0"
+version = "0.32.0-dev.0"
 dependencies = [
  "aqora-archiver",
  "aqora-client",
@@ -211,7 +211,7 @@ dependencies = [
 
 [[package]]
 name = "aqora-archiver"
-version = "0.31.3-dev.0"
+version = "0.32.0-dev.0"
 dependencies = [
  "crc32fast",
  "flate2",
@@ -232,7 +232,7 @@ dependencies = [
 
 [[package]]
 name = "aqora-client"
-version = "0.31.3-dev.0"
+version = "0.32.0-dev.0"
 dependencies = [
  "async-trait",
  "aws-config",
@@ -264,7 +264,7 @@ dependencies = [
 
 [[package]]
 name = "aqora-config"
-version = "0.31.3-dev.0"
+version = "0.32.0-dev.0"
 dependencies = [
  "pep440_rs",
  "pep508_rs",
@@ -276,7 +276,7 @@ dependencies = [
 
 [[package]]
 name = "aqora-data-utils"
-version = "0.31.3-dev.0"
+version = "0.32.0-dev.0"
 dependencies = [
  "aqora-client",
  "arrow",
@@ -325,7 +325,7 @@ dependencies = [
 
 [[package]]
 name = "aqora-runner"
-version = "0.31.3-dev.0"
+version = "0.32.0-dev.0"
 dependencies = [
  "aqora-config",
  "clap",
@@ -342,7 +342,7 @@ dependencies = [
 
 [[package]]
 name = "aqora-template"
-version = "0.31.3-dev.0"
+version = "0.32.0-dev.0"
 dependencies = [
  "derive_builder",
  "handlebars",

{aqora-0.31.3.dev0 → aqora-0.32.0.dev0}/Cargo.toml

@@ -3,7 +3,7 @@ resolver = "2"
 members = ["archiver", "config", "runner", "template", "data-utils", "client"]
 
 [workspace.package]
-version = "0.31.3-dev.0"
+version = "0.32.0-dev.0"
 license = "MIT"
 
 [package]

{aqora-0.31.3.dev0 → aqora-0.32.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aqora
-Version: 0.31.3.dev0
+Version: 0.32.0.dev0
 Requires-Dist: typing-extensions>=4.13.2
 Requires-Dist: uv>=0.3.2,<1.0.0
 Requires-Dist: cirq-core>=1.4 ; extra == 'cirq'

{aqora-0.31.3.dev0 → aqora-0.32.0.dev0}/python/aqora/_aqora.pyi

@@ -34,6 +34,7 @@ class ClientError(Exception): ...
 
 class Client:
     authenticated: bool
+    granted_scopes: list[str] | None
     def __init__(
         self, url: str | None = None, *, allow_insecure_host: bool | None = None
     ) -> None: ...
@@ -57,7 +58,7 @@ class Client:
     def with_token(self, token: str) -> Client: ...
     async def authorize_viewer(
         self, scope: Sequence[str] | None = None
-    ) -> ViewerAuthorization: ...
+    ) -> ViewerAuthorization | None: ...
 
 class ViewerAuthorization:
     url: str

aqora-0.32.0.dev0/python/aqora/auth.py

@@ -0,0 +1,109 @@
+import html
+import warnings
+from typing import Any, Sequence
+
+from ._aqora import Client
+
+# The grant a viewer made is remembered on marimo's runtime context, under this
+# attribute, so that re-running a cell does not ask the same person again. It is
+# stored on the context rather than in a module of our own because marimo installs
+# a context per session: the grant then lives and dies with the session and cannot
+# be reached from another one, even if marimo ever ran several sessions in one
+# process. Anything process-wide would hand one viewer's token to another.
+_GRANT_ATTR = "_aqora_viewer_grant"
+
+
+def _covers(granted: Sequence[str] | None, requested: Sequence[str] | None) -> bool:
+    return set(requested or ["default"]) <= set(granted or [])
+
+
+def _session_context() -> Any:
+    """marimo's runtime state for this session, or ``None`` if out of reach.
+
+    ``get_context`` is a marimo internal, so the import and the call are both
+    guarded; no context simply means no caching.
+    """
+    try:
+        from marimo._runtime.context import get_context
+
+        return get_context()
+    except Exception:
+        return None
+
+
+async def viewer_login(
+    scope: Sequence[str] | None = None,
+    *,
+    client: Client | None = None,
+    timeout: float = 600.0,
+    label: str = "Sign in with aqora",
+) -> Client:
+    """Ask the person viewing this notebook to grant it access to their aqora account.
+
+    Renders a sign-in link in the current cell output, waits until they approve,
+    and returns a client authenticated as *them*. The grant is reused for the rest
+    of their session, so re-running a cell does not ask again; a fresh visit does.
+    The default ``Client()`` inside a workspace runner acts as the workspace owner;
+    use the returned client for anything done on the viewer's behalf.
+
+    Outside a workspace runner, in a notebook on your own machine, there is no
+    viewer to ask: the client returned acts as whoever ran ``aqora login``, with
+    that account's full access rather than ``scope``. That is meant for
+    development; the consent flow only runs on aqora.
+    """
+    import marimo as mo  # imported lazily; aqora does not depend on marimo
+
+    session = _session_context()
+    if session is not None:
+        cached, granted = getattr(session, _GRANT_ATTR, (None, None))
+        if cached is not None and _covers(granted, scope):
+            return cached
+
+    client = client or Client()
+    auth = await client.authorize_viewer(scope)
+    if auth is None:
+        return await _local_login(client)
+    mo.output.append(
+        mo.Html(
+            f'<a href="{html.escape(auth.url, quote=True)}"'
+            ' target="_blank" rel="noopener noreferrer">'
+            f"{html.escape(label)}</a>"
+        )
+    )
+    viewer = await auth.wait(timeout=timeout)
+    mo.output.clear()
+    if session is not None:
+        try:
+            setattr(session, _GRANT_ATTR, (viewer, viewer.granted_scopes))
+        except Exception:
+            pass  # a context that refuses attributes just means no caching
+    return viewer
+
+
+async def _local_login(client: Client) -> Client:
+    """The developer's own ``aqora login``, standing in for a viewer."""
+    # `aqora/__init__` imports this module before it defines ClientError
+    from . import ClientError
+
+    if not client.authenticated:
+        await client.authenticate()
+    try:
+        data = await client.send("query { viewer { username } }")
+    except ClientError as error:
+        codes = {
+            (graphql_error.get("extensions") or {}).get("code")
+            for graphql_error in error.graphql_errors or []
+        }
+        if "NOT_AUTHORIZED" not in codes:
+            raise
+        raise ClientError(
+            "viewer_login: not inside an aqora workspace runner and no aqora"
+            " credentials found; run `aqora login` first"
+        ) from error
+    warnings.warn(
+        "viewer_login: not inside an aqora workspace runner; acting as"
+        f" {data['viewer']['username']} from `aqora login` with that account's"
+        " full access",
+        stacklevel=3,
+    )
+    return client

{aqora-0.31.3.dev0 → aqora-0.32.0.dev0}/src/graphql/oauth2_refresh.graphql

@@ -16,6 +16,7 @@ mutation Oauth2RefreshMutation(
       expiresIn
       accessToken
       refreshToken
+      scope
     }
   }
 }

{aqora-0.31.3.dev0 → aqora-0.32.0.dev0}/src/graphql/oauth2_token.graphql

@@ -8,6 +8,7 @@ mutation Oauth2TokenMutation($code: String!, $client_id: String!, $redirect_uri:
       expiresIn
       accessToken
       refreshToken
+      scope
     }
   }
 }

{aqora-0.31.3.dev0 → aqora-0.32.0.dev0}/src/oauth2.rs

@@ -148,14 +148,18 @@ pub(crate) struct IssuedTokens {
     pub access_token: String,
     pub refresh_token: String,
     pub expires_at: DateTime<Utc>,
+    /// The space delimited scopes the authorization server granted.
+    #[cfg_attr(not(feature = "extension-module"), allow(dead_code))]
+    pub scope: String,
 }
 
 impl IssuedTokens {
-    fn new(access_token: String, refresh_token: String, expires_in: i64) -> Self {
+    fn new(access_token: String, refresh_token: String, expires_in: i64, scope: String) -> Self {
         Self {
             access_token,
             refresh_token,
             expires_at: expires_at(expires_in),
+            scope,
         }
     }
 
@@ -202,7 +206,12 @@ pub(crate) async fn exchange_authorization_code(
         token.unauthorized,
         token.client_error,
         token.issued.map(|issued| {
-            IssuedTokens::new(issued.access_token, issued.refresh_token, issued.expires_in)
+            IssuedTokens::new(
+                issued.access_token,
+                issued.refresh_token,
+                issued.expires_in,
+                issued.scope,
+            )
         }),
     )
 }
@@ -232,7 +241,12 @@ pub(crate) async fn refresh_tokens(
         refresh.unauthorized,
         refresh.client_error,
         refresh.issued.map(|issued| {
-            IssuedTokens::new(issued.access_token, issued.refresh_token, issued.expires_in)
+            IssuedTokens::new(
+                issued.access_token,
+                issued.refresh_token,
+                issued.expires_in,
+                issued.scope,
+            )
         }),
     )
 }
@@ -408,7 +422,7 @@ mod tests {
         client
     }
 
-    const REFRESHED: &str = r#"{"data":{"oauth2Refresh":{"clientError":false,"unauthorized":false,"issued":{"expiresIn":3600,"accessToken":"refreshed-access","refreshToken":"refreshed-refresh"}}}}"#;
+    const REFRESHED: &str = r#"{"data":{"oauth2Refresh":{"clientError":false,"unauthorized":false,"issued":{"expiresIn":3600,"accessToken":"refreshed-access","refreshToken":"refreshed-refresh","scope":"default"}}}}"#;
 
     #[tokio::test]
     async fn viewer_credentials_returns_the_current_token() {
@@ -420,6 +434,7 @@ mod tests {
                 access_token: "access".to_string(),
                 refresh_token: "refresh".to_string(),
                 expires_at: Utc::now() + Duration::try_hours(1).unwrap(),
+                scope: "default".to_string(),
             },
         );
 
@@ -440,6 +455,7 @@ mod tests {
                 access_token: "access".to_string(),
                 refresh_token: "refresh".to_string(),
                 expires_at: Utc::now() + Duration::try_seconds(30).unwrap(),
+                scope: "default".to_string(),
             },
         );
 
@@ -460,6 +476,7 @@ mod tests {
                 access_token: "access".to_string(),
                 refresh_token: "refresh".to_string(),
                 expires_at: Utc::now() - Duration::try_hours(1).unwrap(),
+                scope: "default".to_string(),
             },
         );
 
@@ -509,6 +526,7 @@ mod tests {
                 access_token: "access".to_string(),
                 refresh_token: "refresh".to_string(),
                 expires_at: Utc::now() + Duration::try_seconds(30).unwrap(),
+                scope: "default".to_string(),
             },
         )));
 

{aqora-0.31.3.dev0 → aqora-0.32.0.dev0}/src/python_module.rs

@@ -24,7 +24,8 @@ use crate::{
     graphql_client::{authenticate_client, unauthenticated_client},
     oauth2::{
         exchange_code, new_authorization_request, oauth2_workspace_client_query, spawn_first_code,
-        subscribe_code, AuthorizationRequest, Oauth2WorkspaceClientQuery, ViewerCredentials,
+        subscribe_code, AuthorizationRequest, IssuedTokens, Oauth2WorkspaceClientQuery,
+        ViewerCredentials,
     },
     workspace::download_workspace_notebook,
 };
@@ -112,6 +113,9 @@ struct PyClient {
     url: Url,
     options: ClientOptions,
     inner: Arc<RwLock<PyClientInner>>,
+    /// The scopes the viewer granted this client, for clients built from a
+    /// viewer authorization.
+    granted_scopes: Option<Vec<String>>,
 }
 
 struct PyClientInner {
@@ -131,6 +135,48 @@ fn get_environ<'py>(py: Python<'py>, key: &str) -> PyResult<Option<String>> {
     }
 }
 
+/// The OAuth2 client this workspace runner authorizes viewers for, or `None`
+/// when the request was not made with a runner api key.
+async fn workspace_oauth2_client(
+    base: &Client,
+) -> PyResult<Option<oauth2_workspace_client_query::Oauth2WorkspaceClientQueryOauth2WorkspaceClient>>
+{
+    Ok(base
+        .send::<Oauth2WorkspaceClientQuery>(oauth2_workspace_client_query::Variables)
+        .await
+        .map_err(|error| client_error(error.into()))?
+        .oauth2_workspace_client)
+}
+
+/// A client authenticating as the viewer `tokens` were issued for, refreshing
+/// them through `base`.
+fn viewer_client(
+    url: Url,
+    options: ClientOptions,
+    base: Client,
+    client_id: String,
+    tokens: IssuedTokens,
+) -> PyResult<PyClient> {
+    let granted_scopes = tokens
+        .scope
+        .split_whitespace()
+        .map(ToString::to_string)
+        .collect();
+    let mut client = unauthenticated_client(url.clone(), options.clone()).map_err(client_error)?;
+    client.graphql_layer(CredentialsLayer::new(ViewerCredentials::new(
+        base, client_id, tokens,
+    )));
+    Ok(PyClient {
+        url,
+        options,
+        inner: Arc::new(RwLock::new(PyClientInner {
+            client,
+            authenticated: true,
+        })),
+        granted_scopes: Some(granted_scopes),
+    })
+}
+
 #[pymethods]
 impl PyClient {
     #[new]
@@ -175,6 +221,7 @@ impl PyClient {
                 client,
                 authenticated: false,
             })),
+            granted_scopes: None,
         })
     }
 
@@ -190,13 +237,15 @@ impl PyClient {
                 client,
                 authenticated: true,
             })),
+            granted_scopes: None,
         })
     }
 
     /// Start an OAuth2 authorization for the viewer of this workspace app.
     ///
-    /// Only works inside an aqora workspace runner, where requests without an
-    /// `Authorization` header are authenticated as the workspace.
+    /// Resolves to `None` outside an aqora workspace runner: only there are
+    /// requests without an `Authorization` header authenticated as the
+    /// workspace.
     #[pyo3(signature = (scope=None))]
     fn authorize_viewer<'py>(
         &self,
@@ -208,16 +257,9 @@ impl PyClient {
         let options = self.options.clone();
         future_into_py(py, async move {
             let base = inner.read().await.client.clone();
-            let workspace_client = base
-                .send::<Oauth2WorkspaceClientQuery>(oauth2_workspace_client_query::Variables)
-                .await
-                .map_err(|error| client_error(error.into()))?
-                .oauth2_workspace_client
-                .ok_or_else(|| {
-                    ClientError::new_err((
-                        "viewer login is only available inside an aqora workspace runner",
-                    ))
-                })?;
+            let Some(workspace_client) = workspace_oauth2_client(&base).await? else {
+                return Ok(None);
+            };
             let request = new_authorization_request(
                 &workspace_client.authorize_url,
                 &workspace_client.client_id,
@@ -231,13 +273,13 @@ impl PyClient {
             let stream = subscribe_code(&base, &request)
                 .await
                 .map_err(client_error)?;
-            Ok(PyViewerAuthorization {
+            Ok(Some(PyViewerAuthorization {
                 request: Arc::new(request),
                 first_code: Arc::new(Mutex::new(Some(spawn_first_code(stream)))),
                 base,
                 url,
                 options,
-            })
+            }))
         })
     }
 
@@ -398,6 +440,11 @@ impl PyClient {
     fn authenticated(&self) -> bool {
         self.inner.blocking_read().authenticated
     }
+
+    #[getter]
+    fn granted_scopes(&self) -> Option<Vec<String>> {
+        self.granted_scopes.clone()
+    }
 }
 
 /// A pending viewer authorization: send the viewer to `url` and `wait` for
@@ -467,21 +514,7 @@ impl PyViewerAuthorization {
             let tokens = exchange_code(&base, &request, &code)
                 .await
                 .map_err(client_error)?;
-            let mut client =
-                unauthenticated_client(url.clone(), options.clone()).map_err(client_error)?;
-            client.graphql_layer(CredentialsLayer::new(ViewerCredentials::new(
-                base,
-                request.client_id.clone(),
-                tokens,
-            )));
-            Ok(PyClient {
-                url,
-                options,
-                inner: Arc::new(RwLock::new(PyClientInner {
-                    client,
-                    authenticated: true,
-                })),
-            })
+            viewer_client(url, options, base, request.client_id.clone(), tokens)
         })
     }
 }