Merge pull request #63 from atomscale-ai/enhancement/project_id_to_stream

Add `project_id` support to RHEED streamer

Author: munrojm

src/atomscale/streaming/rheed_stream.pyi

@@ -25,6 +25,7 @@ class RHEEDStreamer:
         chunk_size: int,
         stream_name: str | None = None,
         physical_sample: str | None = None,
+        project_id: str | None = None,
     ) -> str: ...
     def run(
         self,

src/atomscale/streaming/src/initialize.rs

@@ -1,6 +1,7 @@
 use anyhow::{Context, Result};
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
+use serde_json::Value;
 
 #[derive(Serialize, Debug)]
 #[serde(rename_all = "snake_case")] // Ensures JSON fields are snake_case (e.g., data_id)
@@ -9,6 +10,20 @@ pub struct RHEEDStreamSettings {
     pub rotational_period: f64,
     pub rotations_per_min: f64,
     pub fps_capture_rate: f64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub project_id: Option<String>,
+}
+
+// Project-related structs for GET /projects/ response
+#[derive(Deserialize, Debug)]
+struct DataConfiguration {
+    api_configuration: Option<Value>,
+}
+
+#[derive(Deserialize, Debug)]
+struct ProjectSummary {
+    id: String,
+    configuration: Option<DataConfiguration>,
 }
 
 /// POST request to initialize a RHEED stream
@@ -58,10 +73,10 @@ pub async fn ensure_physical_sample_link(
     api_key: &str,
     data_id: &str,
     sample_name: &str,
-) -> Result<()> {
+) -> Result<String> {
     let sample_name = sample_name.trim();
     if sample_name.is_empty() {
-        return Ok(());
+        anyhow::bail!("sample_name cannot be empty");
     }
 
     let list_url = format!("{base_endpoint}/physical_samples/");
@@ -102,7 +117,7 @@ pub async fn ensure_physical_sample_link(
     let link_url = format!("{base_endpoint}/data_entries/physical_sample");
     let link_body = LinkPhysicalSampleRequest {
         data_ids: vec![data_id.to_string()],
-        physical_sample_id: sample_id,
+        physical_sample_id: sample_id.clone(),
     };
 
     client
@@ -115,5 +130,63 @@ pub async fn ensure_physical_sample_link(
         .error_for_status()
         .context("physical sample link returned error status")?;
 
+    Ok(sample_id)
+}
+
+/// Updates the project's tracking_physical_sample_id in its configuration.
+/// Fetches current configuration, updates the tracking sample, and POSTs back.
+pub async fn update_project_tracking_sample(
+    client: &Client,
+    base_endpoint: &str,
+    api_key: &str,
+    project_id: &str,
+    physical_sample_id: &str,
+) -> Result<()> {
+    // GET /projects/ to find the project and its current configuration
+    let projects_url = format!("{base_endpoint}/projects/");
+    let projects: Vec<ProjectSummary> = client
+        .get(&projects_url)
+        .header("X-API-KEY", api_key)
+        .send()
+        .await
+        .context("failed to request projects")?
+        .error_for_status()
+        .context("projects list returned error status")?
+        .json()
+        .await
+        .context("failed to deserialize projects list")?;
+
+    // Find the project by ID
+    let project = projects
+        .into_iter()
+        .find(|p| p.id == project_id)
+        .ok_or_else(|| anyhow::anyhow!("project with id {} not found", project_id))?;
+
+    // Build updated configuration, preserving existing fields
+    let mut config = match project.configuration {
+        Some(data_config) => data_config.api_configuration.unwrap_or_else(|| Value::Object(Default::default())),
+        None => Value::Object(Default::default()),
+    };
+
+    // Update tracking_physical_sample_id in the configuration
+    if let Value::Object(ref mut map) = config {
+        map.insert(
+            "tracking_physical_sample_id".to_string(),
+            Value::String(physical_sample_id.to_string()),
+        );
+    }
+
+    // POST /projects/{project_id}/configuration with the updated config
+    let config_url = format!("{base_endpoint}/projects/{project_id}/configuration");
+    client
+        .post(&config_url)
+        .header("X-API-KEY", api_key)
+        .json(&config)
+        .send()
+        .await
+        .context("failed to update project configuration")?
+        .error_for_status()
+        .context("project configuration update returned error status")?;
+
     Ok(())
 }

src/atomscale/streaming/src/lib.rs

@@ -13,7 +13,10 @@ mod utils;
 use utils::{generic_post, init_tracing_once};
 
 mod initialize;
-use initialize::{ensure_physical_sample_link, post_for_initialization, RHEEDStreamSettings};
+use initialize::{
+    ensure_physical_sample_link, post_for_initialization, update_project_tracking_sample,
+    RHEEDStreamSettings,
+};
 
 mod upload;
 use upload::{
@@ -109,7 +112,7 @@ impl RHEEDStreamer {
     }
 
     ////Initialize stream
-    /// initialize(self, stream_name: Optional[str] = None, fps: float, rotations_per_min: float, chunk_size: int, physical_sample: Optional[str] = None) -> str
+    /// initialize(self, fps: float, rotations_per_min: float, chunk_size: int, stream_name: Optional[str] = None, physical_sample: Optional[str] = None, project_id: Optional[str] = None) -> str
     ///
     /// Creates a new **remote data item** for this stream and returns its `data_id`.
     /// Also captures runtime configuration used for subsequent chunk uploads.
@@ -121,21 +124,24 @@ impl RHEEDStreamer {
     /// After streaming via `run(...)` or `push(...)`, call `finalize(data_id)` to mark the stream as complete.
     ///
     /// Args:
-    ///     stream_name (Optional[str]): Human-readable name shown in the platform. If `None` or an empty string,
-    ///         a default like `"RHEED Stream @ 1:23PM"` is used.
     ///     fps (float): Capture rate in frames per second.
     ///     rotations_per_min (float): Wafer/crystal rotations per minute; use `0.0` for stationary operation.
     ///     chunk_size (int): The **intended** number of frames per chunk you will send with `run(...)` or `push(...)`.
+    ///     stream_name (Optional[str]): Human-readable name shown in the platform. If `None` or an empty string,
+    ///         a default like `"RHEED Stream @ 1:23PM"` is used.
     ///     physical_sample (Optional[str]): Name of a physical sample to associate with the data item; matched case-insensitively or created if missing.
+    ///     project_id (Optional[str]): UUID of a project to associate with the stream. When provided along with
+    ///         `physical_sample`, the project's `tracking_physical_sample_id` configuration is automatically updated
+    ///         to link the physical sample to the project for growth monitoring.
     ///
     /// Returns:
     ///     str: The created `data_id` for this stream.
     ///
     /// Raises:
     ///     RuntimeError: If the initialization POST fails.
-    #[pyo3(signature = (fps, rotations_per_min, chunk_size, stream_name=None, physical_sample=None))]
+    #[pyo3(signature = (fps, rotations_per_min, chunk_size, stream_name=None, physical_sample=None, project_id=None))]
     #[pyo3(
-        text_signature = "(fps, rotations_per_min, chunk_size, stream_name=None, physical_sample=None)"
+        text_signature = "(fps, rotations_per_min, chunk_size, stream_name=None, physical_sample=None, project_id=None)"
     )]
     fn initialize(
         &mut self,
@@ -144,6 +150,7 @@ impl RHEEDStreamer {
         chunk_size: usize,
         stream_name: Option<String>,
         physical_sample: Option<String>,
+        project_id: Option<String>,
     ) -> PyResult<String> {
         // Guard: chunk_size must be >= ceil(2 * fps)
         let min_chunk = (2.0 * fps).ceil() as usize;
@@ -165,6 +172,10 @@ impl RHEEDStreamer {
             .map(|s| s.trim().to_string())
             .filter(|s| !s.is_empty());
 
+        let project_id = project_id
+            .map(|s| s.trim().to_string())
+            .filter(|s| !s.is_empty());
+
         let fpr = (fps * 60.0) / rotations_per_min;
 
         #[allow(clippy::redundant_field_names)]
@@ -173,6 +184,7 @@ impl RHEEDStreamer {
             rotational_period: fpr,
             rotations_per_min,
             fps_capture_rate: fps,
+            project_id,
         };
 
         let base_endpoint = self.endpoint.clone();
@@ -192,9 +204,24 @@ impl RHEEDStreamer {
                 &data_id,
                 &sample_name,
             );
-            self.rt
+            let sample_id = self
+                .rt
                 .block_on(physical_sample_fut)
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+
+            // If project_id was provided, update the project's tracking_physical_sample_id
+            if let Some(ref proj_id) = settings.project_id {
+                let update_project_fut = update_project_tracking_sample(
+                    &self.client,
+                    &base_endpoint,
+                    &self.api_key,
+                    proj_id,
+                    &sample_id,
+                );
+                self.rt
+                    .block_on(update_project_fut)
+                    .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+            }
         }
 
         self.fps = Some(fps);

tests/_mock_http_server.py

@@ -0,0 +1,122 @@
+"""Subprocess-based mock HTTP server for testing Rust HTTP clients.
+
+This module is designed to be run as a subprocess to provide true process
+isolation when testing Rust HTTP clients (like reqwest) from Python tests.
+
+Usage:
+    python -m tests._mock_http_server <port> <json_response>
+    python -m tests._mock_http_server <port> <json_routes_dict>
+
+The server:
+- Listens on 127.0.0.1:<port>
+- Prints "READY:<port>" to stdout when ready
+- For simple mode: handles one POST request, prints "BODY:<json>" to stdout
+- For routes mode: handles multiple requests based on path matching
+- Responds with the provided JSON response
+"""
+import json
+import sys
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+
+class CaptureHandler(BaseHTTPRequestHandler):
+    """HTTP handler that captures request body and returns configured response."""
+
+    def _handle_request(self, method: str):
+        """Common handler for all HTTP methods."""
+        # Read request body if present
+        length = int(self.headers.get("Content-Length", 0))
+        body = self.rfile.read(length) if length > 0 else b""
+
+        path = self.path
+        routes = getattr(self.server, "routes", None)
+
+        if routes:
+            # Routes mode: find matching route
+            response_data = None
+            for route_path, route_response in routes.items():
+                if path.startswith(route_path):
+                    response_data = route_response
+                    break
+
+            if response_data is None:
+                # Default response for unmatched routes
+                response_data = '""'
+
+            # Print request info for debugging
+            print(f"REQUEST:{method}:{path}:{body.decode() if body else ''}", flush=True)
+        else:
+            # Simple mode: single response for all requests
+            response_data = self.server.response_data
+            print(f"BODY:{body.decode()}", flush=True)
+
+        # Send response
+        self.send_response(200)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", len(response_data))
+        self.end_headers()
+        self.wfile.write(response_data.encode())
+
+    def do_GET(self):
+        self._handle_request("GET")
+
+    def do_POST(self):
+        self._handle_request("POST")
+
+    def do_PUT(self):
+        self._handle_request("PUT")
+
+    def log_message(self, format, *args):
+        """Suppress default logging."""
+        pass
+
+
+class MultiRequestServer(HTTPServer):
+    """HTTP server that can handle multiple requests."""
+
+    def __init__(self, *args, max_requests: int = 1, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_requests = max_requests
+        self.request_count = 0
+
+    def handle_requests(self):
+        """Handle up to max_requests requests."""
+        while self.request_count < self.max_requests:
+            self.handle_request()
+            self.request_count += 1
+
+
+def run_server(port: int, response_data: str) -> None:
+    """Run the mock HTTP server."""
+    # Try to parse as routes dict
+    try:
+        routes = json.loads(response_data)
+        if isinstance(routes, dict) and routes.get("__routes__"):
+            # Routes mode
+            del routes["__routes__"]
+            max_requests = routes.pop("__max_requests__", 10)
+            server = MultiRequestServer(
+                ("127.0.0.1", port), CaptureHandler, max_requests=max_requests
+            )
+            server.routes = routes
+            print(f"READY:{port}", flush=True)
+            server.handle_requests()
+            return
+    except (json.JSONDecodeError, TypeError):
+        pass
+
+    # Simple mode: single request with static response
+    server = HTTPServer(("127.0.0.1", port), CaptureHandler)
+    server.response_data = response_data
+    print(f"READY:{port}", flush=True)
+    server.handle_request()
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 3:
+        print(f"Usage: {sys.argv[0]} <port> <json_response>", file=sys.stderr)
+        sys.exit(1)
+
+    port = int(sys.argv[1])
+    response_data = sys.argv[2]
+    run_server(port, response_data)