recce-nightly 1.16.0.20250818__py3-none-any.whl → 1.16.0.20250819.post1__py3-none-any.whl

recce/state/cloud.py

@@ -51,6 +51,7 @@ class CloudStateLoader(RecceStateLoader):
             cloud_options=cloud_options,
             initial_state=initial_state,
         )
+        self.recce_cloud = RecceCloud(token=self.token)
 
     def verify(self) -> bool:
         if self.catalog == "github":
@@ -74,157 +75,393 @@ class CloudStateLoader(RecceStateLoader):
                     'Please provide a share URL in the command argument with option "--share-url <share-url>"'
                 )
                 return False
+        elif self.catalog == "snapshot":
+            if self.cloud_options.get("api_token") is None:
+                self.error_message = RECCE_API_TOKEN_MISSING.error_message
+                self.hint_message = RECCE_API_TOKEN_MISSING.hint_message
+                return False
+            if self.cloud_options.get("snapshot_id") is None:
+                self.error_message = "No snapshot ID is provided for the snapshot catalog."
+                self.hint_message = (
+                    'Please provide a snapshot ID in the command argument with option "--snapshot-id <snapshot-id>"'
+                )
+                return False
         return True
 
-    def _load_state(self) -> Tuple[RecceState, str]:
-        return self._load_state_from_cloud()
-
-    def _export_state(self, state: RecceState = None) -> Union[str, None]:
-        return self._export_state_to_cloud()
-
     def purge(self) -> bool:
         rc, err_msg = RecceCloudStateManager(self.cloud_options).purge_cloud_state()
         if err_msg:
             self.error_message = err_msg
         return rc
 
-    def _load_state_from_cloud(self) -> Tuple[RecceState, str]:
+    def _load_state(self) -> Tuple[RecceState, str]:
         """
-        Load the state from Recce Cloud.
+        Load the state from Recce Cloud based on catalog type.
 
         Returns:
             RecceState: The state object.
-            str: The etag of the state file.
+            str: The etag of the state file (only used for GitHub).
         """
         if self.catalog == "github":
-            if (self.pr_info is None) or (self.pr_info.id is None) or (self.pr_info.repository is None):
-                raise RecceException("Cannot get the pull request information from GitHub.")
+            return self._load_state_from_github()
         elif self.catalog == "preview":
-            if self.share_id is None:
-                raise RecceException("Cannot load the share state from Recce Cloud. No share ID is provided.")
+            return self._load_state_from_preview()
+        elif self.catalog == "snapshot":
+            return self._load_state_from_snapshot(), None
+        else:
+            raise RecceException(f"Unsupported catalog type: {self.catalog}")
+
+    def _load_state_from_github(self) -> Tuple[RecceState, str]:
+        """Load state from GitHub PR with etag checking."""
+        if (self.pr_info is None) or (self.pr_info.id is None) or (self.pr_info.repository is None):
+            raise RecceException("Cannot get the pull request information from GitHub.")
 
-        logger.debug("Fetching state from Recce Cloud...")
+        logger.debug("Fetching GitHub state from Recce Cloud...")
+
+        # Check metadata and etag for GitHub only
         metadata = self._get_metadata_from_recce_cloud()
-        if metadata:
-            state_etag = metadata.get("etag")
-        else:
-            state_etag = None
+        state_etag = metadata.get("etag") if metadata else None
+
+        # Return cached state if etag matches
         if self.state_etag and state_etag == self.state_etag:
             return self.state, self.state_etag
 
-        return self._load_state_from_recce_cloud(), state_etag
+        # Download state from GitHub
+        presigned_url = self.recce_cloud.get_presigned_url_by_github_repo(
+            method=PresignedUrlMethod.DOWNLOAD,
+            pr_id=self.pr_info.id,
+            repository=self.pr_info.repository,
+            artifact_name=RECCE_STATE_COMPRESSED_FILE,
+        )
 
-    def _get_metadata_from_recce_cloud(self) -> Union[dict, None]:
-        recce_cloud = RecceCloud(token=self.token)
-        return recce_cloud.get_artifact_metadata(pr_info=self.pr_info) if self.pr_info else None
+        password = self.cloud_options.get("password")
+        if password is None:
+            raise RecceException(RECCE_CLOUD_PASSWORD_MISSING.error_message)
 
-    def _load_state_from_recce_cloud(self) -> Union[RecceState, None]:
-        import tempfile
+        headers = s3_sse_c_headers(password)
+        loaded_state = self._download_state_from_url(presigned_url, SupportedFileTypes.GZIP, headers)
 
-        import requests
+        # Handle the case where download returns None (404 error)
+        if loaded_state is None:
+            return None, state_etag
 
-        recce_cloud = RecceCloud(token=self.token)
-        password = None
+        return loaded_state, state_etag
 
-        if self.catalog == "github":
-            presigned_url = recce_cloud.get_presigned_url_by_github_repo(
-                method=PresignedUrlMethod.DOWNLOAD,
-                pr_id=self.pr_info.id,
-                repository=self.pr_info.repository,
-                artifact_name=RECCE_STATE_COMPRESSED_FILE,
-            )
+    def _load_state_from_preview(self) -> Tuple[RecceState, None]:
+        """Load state from preview share (no etag checking needed)."""
+        if self.share_id is None:
+            raise RecceException("Cannot load the share state from Recce Cloud. No share ID is provided.")
 
-            password = self.cloud_options.get("password")
-            if password is None:
-                raise RecceException(RECCE_CLOUD_PASSWORD_MISSING.error_message)
-        elif self.catalog == "preview":
-            share_id = self.cloud_options.get("share_id")
-            presigned_url = recce_cloud.get_presigned_url_by_share_id(
-                method=PresignedUrlMethod.DOWNLOAD, share_id=share_id
-            )
+        logger.debug("Fetching preview state from Recce Cloud...")
 
-        with tempfile.NamedTemporaryFile() as tmp:
-            from .cloud import s3_sse_c_headers
+        # Download state from preview share
+        presigned_url = self.recce_cloud.get_presigned_url_by_share_id(
+            method=PresignedUrlMethod.DOWNLOAD, share_id=self.share_id
+        )
+
+        loaded_state = self._download_state_from_url(presigned_url, SupportedFileTypes.FILE)
+
+        # Handle the case where download returns None (404 error)
+        if loaded_state is None:
+            return None, None
+
+        return loaded_state, None
+
+    def _get_metadata_from_recce_cloud(self) -> Union[dict, None]:
+        return self.recce_cloud.get_artifact_metadata(pr_info=self.pr_info) if self.pr_info else None
 
-            headers = s3_sse_c_headers(password) if password else None
+    def _download_state_from_url(
+        self, presigned_url: str, file_type: SupportedFileTypes, headers: dict = None
+    ) -> RecceState:
+        """Download state file from presigned URL and convert to RecceState."""
+        import tempfile
+
+        import requests
+
+        with tempfile.NamedTemporaryFile() as tmp:
             response = requests.get(presigned_url, headers=headers)
+
             if response.status_code == 404:
                 self.error_message = "The state file is not found in Recce Cloud."
                 return None
             elif response.status_code != 200:
                 self.error_message = response.text
-                raise RecceException(
-                    f"{response.status_code} Failed to download the state file from Recce Cloud. The password could be wrong."
-                )
+                error_msg = f"{response.status_code} Failed to download the state file from Recce Cloud."
+                if headers:  # GitHub case with password
+                    error_msg += " The password could be wrong."
+                raise RecceException(error_msg)
+
             with open(tmp.name, "wb") as f:
                 f.write(response.content)
 
-            file_type = SupportedFileTypes.GZIP if self.catalog == "github" else SupportedFileTypes.FILE
             return RecceState.from_file(tmp.name, file_type=file_type)
 
-    def _export_state_to_cloud(self) -> Tuple[Union[str, None], str]:
+    def _load_state_from_snapshot(self) -> RecceState:
+        """
+        Load state from snapshot by:
+        1. Get snapshot info
+        2. Download artifacts for both base and current snapshots
+        3. Download recce_state if available, otherwise create empty state with artifacts
+        """
+        if self.snapshot_id is None:
+            raise RecceException("Cannot load the snapshot state from Recce Cloud. No snapshot ID is provided.")
+
+        # 1. Get snapshot information
+        logger.debug(f"Getting snapshot {self.snapshot_id}")
+        snapshot = self.recce_cloud.get_snapshot(self.snapshot_id)
+
+        org_id = snapshot.get("org_id")
+        project_id = snapshot.get("project_id")
+
+        if not org_id or not project_id:
+            raise RecceException(f"Snapshot {self.snapshot_id} does not belong to a valid organization or project.")
+
+        # 2. Download manifests and catalogs for both snapshots
+        logger.debug(f"Downloading current snapshot artifacts for {self.snapshot_id}")
+        current_artifacts = self._download_snapshot_artifacts(self.recce_cloud, org_id, project_id, self.snapshot_id)
+
+        logger.debug(f"Downloading base snapshot artifacts for project {project_id}")
+        base_artifacts = self._download_base_snapshot_artifacts(self.recce_cloud, org_id, project_id)
+
+        # 3. Try to download existing recce_state, otherwise create new state
+        try:
+            logger.debug(f"Downloading recce_state for snapshot {self.snapshot_id}")
+            state = self._download_snapshot_recce_state(self.recce_cloud, org_id, project_id, self.snapshot_id)
+        except Exception as e:
+            logger.debug(f"No existing recce_state found, creating new state: {e}")
+            state = RecceState()
+
+        # Set artifacts regardless of whether we loaded existing state
+        state.artifacts.base = base_artifacts
+        state.artifacts.current = current_artifacts
+
+        return state
+
+    def _download_snapshot_artifacts(self, recce_cloud, org_id: str, project_id: str, snapshot_id: str) -> dict:
+        """Download manifest and catalog for a snapshot, return JSON data directly."""
+        import requests
+
+        # Get download URLs
+        presigned_urls = recce_cloud.get_download_urls_by_snapshot_id(org_id, project_id, snapshot_id)
+
+        artifacts = {}
+
+        # Download manifest
+        response = requests.get(presigned_urls["manifest_url"])
+        if response.status_code == 200:
+            artifacts["manifest"] = response.json()
+        else:
+            raise RecceException(f"Failed to download manifest for snapshot {snapshot_id}")
+
+        # Download catalog
+        response = requests.get(presigned_urls["catalog_url"])
+        if response.status_code == 200:
+            artifacts["catalog"] = response.json()
+        else:
+            raise RecceException(f"Failed to download catalog for snapshot {snapshot_id}")
+
+        return artifacts
+
+    def _download_snapshot_recce_state(self, recce_cloud, org_id: str, project_id: str, snapshot_id: str) -> RecceState:
+        """Download recce_state for a snapshot."""
+        # Get download URLs (now includes recce_state_url)
+        presigned_urls = recce_cloud.get_download_urls_by_snapshot_id(org_id, project_id, snapshot_id)
+        recce_state_url = presigned_urls.get("recce_state_url")
+
+        if not recce_state_url:
+            raise RecceException(f"No recce_state_url found for snapshot {snapshot_id}")
+
+        # Reuse the existing download method
+        state = self._download_state_from_url(recce_state_url, SupportedFileTypes.FILE)
+
+        if state is None:
+            raise RecceException(f"Failed to download recce_state for snapshot {snapshot_id}")
+
+        return state
+
+    def _download_base_snapshot_artifacts(self, recce_cloud, org_id: str, project_id: str) -> dict:
+        """Download manifest and catalog for the base snapshot, return JSON data directly."""
+        import requests
+
+        # Get download URLs for base snapshot
+        presigned_urls = recce_cloud.get_base_snapshot_download_urls(org_id, project_id)
+
+        artifacts = {}
+
+        # Download manifest
+        response = requests.get(presigned_urls["manifest_url"])
+        if response.status_code == 200:
+            artifacts["manifest"] = response.json()
+        else:
+            raise RecceException(f"Failed to download base snapshot manifest for project {project_id}")
+
+        # Download catalog
+        response = requests.get(presigned_urls["catalog_url"])
+        if response.status_code == 200:
+            artifacts["catalog"] = response.json()
+        else:
+            raise RecceException(f"Failed to download base snapshot catalog for project {project_id}")
+
+        return artifacts
+
+    def _export_state(self) -> Tuple[Union[str, None], str]:
+        """
+        Export state to Recce Cloud based on catalog type.
+
+        Returns:
+            str: A message indicating the result of the export operation.
+            str: The etag of the exported state file (only used for GitHub).
+        """
+        logger.info("Store recce state to Recce Cloud")
+
         if self.catalog == "github":
-            if (self.pr_info is None) or (self.pr_info.id is None) or (self.pr_info.repository is None):
-                raise RecceException("Cannot get the pull request information from GitHub.")
+            return self._export_state_to_github()
         elif self.catalog == "preview":
-            pass
+            return self._export_state_to_preview()
+        elif self.catalog == "snapshot":
+            return self._export_state_to_snapshot()
+        else:
+            raise RecceException(f"Unsupported catalog type: {self.catalog}")
 
+    def _export_state_to_github(self) -> Tuple[Union[str, None], str]:
+        """Export state to GitHub PR with metadata and etag."""
+        if (self.pr_info is None) or (self.pr_info.id is None) or (self.pr_info.repository is None):
+            raise RecceException("Cannot get the pull request information from GitHub.")
+
+        # Generate metadata for GitHub only
         check_status = CheckDAO().status()
         metadata = {
             "total_checks": check_status.get("total", 0),
             "approved_checks": check_status.get("approved", 0),
         }
 
-        logger.info("Store recce state to Recce Cloud")
-        message = self._export_state_to_recce_cloud(metadata=metadata)
-        metadata = self._get_metadata_from_recce_cloud()
-        state_etag = metadata.get("etag") if metadata else None
+        # Upload to Cloud
+        presigned_url = self.recce_cloud.get_presigned_url_by_github_repo(
+            method=PresignedUrlMethod.UPLOAD,
+            repository=self.pr_info.repository,
+            artifact_name=RECCE_STATE_COMPRESSED_FILE,
+            pr_id=self.pr_info.id,
+            metadata=metadata,
+        )
+        message = self._upload_state_to_url(
+            presigned_url=presigned_url,
+            file_type=SupportedFileTypes.GZIP,
+            password=self.cloud_options.get("password"),
+            metadata=metadata,
+        )
+
+        # Get updated etag after upload
+        metadata_response = self._get_metadata_from_recce_cloud()
+        state_etag = metadata_response.get("etag") if metadata_response else None
+
         if message:
             logger.warning(message)
         return message, state_etag
 
-    def _export_state_to_recce_cloud(self, metadata: dict = None) -> Union[str, None]:
+    def _export_state_to_preview(self) -> Tuple[Union[str, None], None]:
+        """Export state to preview share (no metadata or etag needed)."""
+        share_id = self.cloud_options.get("share_id")
+        presigned_url = self.recce_cloud.get_presigned_url_by_share_id(
+            method=PresignedUrlMethod.UPLOAD,
+            share_id=share_id,
+            metadata=None,
+        )
+        message = self._upload_state_to_url(
+            presigned_url=presigned_url, file_type=SupportedFileTypes.FILE, password=None, metadata=None
+        )
+
+        if message:
+            logger.warning(message)
+        return message, None
+
+    def _export_state_to_snapshot(self) -> Tuple[Union[str, None], None]:
+        """Export state to snapshot (upload recce_state with empty artifacts)."""
+        if self.snapshot_id is None:
+            raise RecceException("Cannot export state to snapshot. No snapshot ID is provided.")
+
+        # Get snapshot information
+        snapshot = self.recce_cloud.get_snapshot(self.snapshot_id)
+        org_id = snapshot.get("org_id")
+        project_id = snapshot.get("project_id")
+
+        if not org_id or not project_id:
+            raise RecceException(f"Snapshot {self.snapshot_id} does not belong to a valid organization or project.")
+
+        # Get upload URLs (now includes recce_state_url)
+        presigned_urls = self.recce_cloud.get_upload_urls_by_snapshot_id(org_id, project_id, self.snapshot_id)
+        recce_state_url = presigned_urls.get("recce_state_url")
+
+        if not recce_state_url:
+            raise RecceException(f"No recce_state_url found for snapshot {self.snapshot_id}")
+
+        # Create a copy of the state with empty artifacts for upload
+        upload_state = RecceState()
+        upload_state.runs = self.state.runs.copy() if self.state.runs else []
+        upload_state.checks = self.state.checks.copy() if self.state.checks else []
+        # Keep artifacts empty (don't copy self.state.artifacts)
+
+        # Upload the state with empty artifacts
+        message = self._upload_state_to_url(
+            presigned_url=recce_state_url,
+            file_type=SupportedFileTypes.FILE,
+            password=None,
+            metadata=None,
+            state=upload_state,
+        )
+
+        if message:
+            logger.warning(message)
+        return message, None
+
+    def _upload_state_to_url(
+        self,
+        presigned_url: str,
+        file_type: SupportedFileTypes,
+        password: str = None,
+        metadata: dict = None,
+        state: RecceState = None,
+    ) -> Union[str, None]:
+        """Upload state file to presigned URL."""
         import tempfile
 
         import requests
 
-        if self.catalog == "github":
-            presigned_url = RecceCloud(token=self.token).get_presigned_url_by_github_repo(
-                method=PresignedUrlMethod.UPLOAD,
-                repository=self.pr_info.repository,
-                artifact_name=RECCE_STATE_COMPRESSED_FILE,
-                pr_id=self.pr_info.id,
-                metadata=metadata,
-            )
-        elif self.catalog == "preview":
-            share_id = self.cloud_options.get("share_id")
-            presigned_url = RecceCloud(token=self.token).get_presigned_url_by_share_id(
-                method=PresignedUrlMethod.UPLOAD,
-                share_id=share_id,
-                metadata=metadata,
-            )
-        compress_passwd = self.cloud_options.get("password")
-        if compress_passwd:
-            headers = s3_sse_c_headers(compress_passwd)
-        else:
-            headers = {}
+        # Use provided state or default to self.state
+        upload_state = state or self.state
 
+        # Prepare headers
+        headers = {}
+        if password:
+            headers.update(s3_sse_c_headers(password))
         if metadata:
             headers["x-amz-tagging"] = urlencode(metadata)
+
         with tempfile.NamedTemporaryFile() as tmp:
-            if self.catalog == "github":
-                file_type = SupportedFileTypes.GZIP
-            elif self.catalog == "preview":
-                file_type = SupportedFileTypes.FILE
-            self._export_state_to_file(tmp.name, file_type=file_type)
+            # Use the specified state to export to file
+            json_data = upload_state.to_json()
+            io = file_io_factory(file_type)
+            io.write(tmp.name, json_data)
 
             with open(tmp.name, "rb") as fd:
                 response = requests.put(presigned_url, data=fd.read(), headers=headers)
+
             if response.status_code not in [200, 204]:
                 self.error_message = response.text
                 return "Failed to upload the state file to Recce Cloud. Reason: " + response.text
+
         return None
 
+    def check_conflict(self) -> bool:
+        if self.catalog != "github":
+            return False
+
+        metadata = self._get_metadata_from_recce_cloud()
+        if not metadata:
+            return False
+
+        state_etag = metadata.get("etag")
+        return state_etag != self.state_etag
+
 
 class RecceCloudStateManager:
     error_message: str

recce/state/state_loader.py

@@ -35,8 +35,9 @@ class RecceStateLoader(ABC):
         self.state_lock = threading.Lock()
         self.state_etag = None
         self.pr_info = None
-        self.catalog: Literal["github", "preview"] = "github"
+        self.catalog: Literal["github", "preview", "snapshot"] = "github"
         self.share_id = None
+        self.snapshot_id = None
 
         if self.cloud_mode:
             if self.cloud_options.get("github_token"):
@@ -47,11 +48,19 @@ class RecceStateLoader(ABC):
                 if self.pr_info.id is None:
                     raise RecceException("Cannot get the pull request information from GitHub.")
             elif self.cloud_options.get("api_token"):
-                self.catalog = "preview"
-                self.share_id = self.cloud_options.get("share_id")
+                if self.cloud_options.get("snapshot_id"):
+                    self.catalog = "snapshot"
+                    self.snapshot_id = self.cloud_options.get("snapshot_id")
+                else:
+                    self.catalog = "preview"
+                    self.share_id = self.cloud_options.get("share_id")
             else:
                 raise RecceException(RECCE_CLOUD_TOKEN_MISSING.error_message)
 
+    @property
+    def token(self):
+        return self.cloud_options.get("github_token") or self.cloud_options.get("api_token")
+
     @abstractmethod
     def verify(self) -> bool:
         """
@@ -61,10 +70,6 @@ class RecceStateLoader(ABC):
         """
         raise NotImplementedError("Subclasses must implement this method.")
 
-    @property
-    def token(self):
-        return self.cloud_options.get("github_token") or self.cloud_options.get("api_token")
-
     @property
     def error_and_hint(self) -> (Union[str, None], Union[str, None]):
         return self.error_message, self.hint_message
@@ -118,7 +123,7 @@ class RecceStateLoader(ABC):
         return message
 
     @abstractmethod
-    def _export_state(self, state: RecceState = None) -> Tuple[Union[str, None], str]:
+    def _export_state(self) -> Tuple[Union[str, None], str]:
         """
         Export the current Recce state to a file or cloud storage.
         Returns:
@@ -143,17 +148,9 @@ class RecceStateLoader(ABC):
         return new_state
 
     def check_conflict(self) -> bool:
-        if not self.cloud_mode:
-            return False
-
-        metadata = self._get_metadata_from_recce_cloud()
-        if not metadata:
-            return False
-
-        state_etag = metadata.get("etag")
-        return state_etag != self.state_etag
+        return False
 
-    def info(self):
+    def info(self) -> dict:
         if self.state is None:
             self.error_message = "No state is loaded."
             return None

recce/util/lineage.py

@@ -7,19 +7,17 @@ def find_upstream(node_ids: Iterable, parent_map):
     visited = set()
     upstream = set()
 
-    def dfs(current):
+    stack = list(node_ids)
+    while stack:
+        current = stack.pop()
         if current in visited:
-            return
+            continue
         visited.add(current)
-
         parents = parent_map.get(current, [])
         for parent in parents:
-            upstream.add(parent)
-            dfs(parent)
-
-    for node_id in node_ids:
-        dfs(node_id)
-
+            if parent not in upstream:
+                upstream.add(parent)
+                stack.append(parent)
     return upstream
 
 
@@ -27,19 +25,17 @@ def find_downstream(node_ids: Iterable, child_map):
     visited = set()
     downstream = set()
 
-    def dfs(current):
+    stack = list(node_ids)
+    while stack:
+        current = stack.pop()
         if current in visited:
-            return
+            continue
         visited.add(current)
-
         children = child_map.get(current, [])
         for child in children:
-            downstream.add(child)
-            dfs(child)
-
-    for node_id in node_ids:
-        dfs(node_id)
-
+            if child not in downstream:
+                downstream.add(child)
+                stack.append(child)
     return downstream