Merge pull request #378 from DataRecce/feature/drc-531-support-to-sen…

…d-webhook-to-github-app-to-update-the-check

[Feature] DRC-531 --state-file-host also support GitHub App

recce/cli.py

@@ -381,7 +381,7 @@ def purge_cloud_state(**kwargs):
                 host = cloud_options.get('host')
                 token = cloud_options.get('token')
                 pr_info = fetch_pr_metadata(github_token=token)
-                rc, err_msg = RecceStateLoader.purge_cloud_state(host=host, pr_info=pr_info, token=token)
+                rc, err_msg = RecceStateLoader.purge_cloud_state(token=token, pr_info=pr_info, host=host)
                 if rc is True:
                     console.rule('Purged Successfully')
                 else:

recce/state.py

@@ -18,6 +18,7 @@
 from recce.models.types import Run, Check
 from recce.pull_request import fetch_pr_metadata, PullRequestInfo
 from recce.util.pydantic_model import pydantic_model_json_dump, pydantic_model_dump
+from recce.util.recce_cloud import RecceCloud, PresignedUrlMethod, RecceCloudException
 
 logger = logging.getLogger('uvicorn')
 
@@ -249,17 +250,12 @@ def info(self):
 
     def purge(self) -> bool:
         if self.cloud_mode is True:
-            if self.cloud_options.get('host', '').startswith('s3://'):
-                rc, err_msg = RecceStateLoader._purge_state_from_s3_bucket(self.cloud_options.get('host', ''),
-                                                                           self.pr_info)
-                if err_msg:
-                    self.error_message = err_msg
-                return rc
-            else:
-                rc, err_msg = RecceStateLoader.purge_cloud_state(self.cloud_options.get('token'), self.pr_info)
-                if err_msg:
-                    self.error_message = err_msg
-                return rc
+            host = self.cloud_options.get('host')
+            token = self.cloud_options.get('token')
+            rc, err_msg = RecceStateLoader.purge_cloud_state(token, self.pr_info, host)
+            if err_msg:
+                self.error_message = err_msg
+            return rc
         else:
             if self.state_file is not None:
                 try:
@@ -272,31 +268,15 @@ def purge(self) -> bool:
                 return False
 
     @staticmethod
-    def purge_cloud_state(host: str, pr_info: PullRequestInfo, token: str = None) -> (bool, str):
+    def purge_cloud_state(token: str, pr_info: PullRequestInfo, host: str) -> (bool, str):
         if host.startswith('s3://'):
-            return RecceStateLoader._purge_state_from_s3_bucket(host, pr_info)
+            return RecceStateLoader._purge_state_from_s3_bucket(token, pr_info, host)
         else:
-            return RecceStateLoader._purge_state_from_cloud(token, pr_info)
-
-    def _get_presigned_url(self, pr_info: PullRequestInfo, artifact_name: str, method: str = 'upload',
-                           metadata: dict = None) -> str:
-        import requests
-        # Step 1: Get the token
-        token = self.cloud_options.get('token')
-        if token is None:
-            raise Exception('No token is provided to access Recce Cloud.')
-
-        # Step 2: Call Recce Cloud API to get presigned URL
-        api_url = f'{RECCE_CLOUD_API_HOST}/api/v1/{pr_info.repository}/pulls/{pr_info.id}/artifacts/{method}?artifact_name={artifact_name}'
-        headers = {
-            'Authorization': f'Bearer {token}'
-        }
-        response = requests.post(api_url, headers=headers, json=metadata)
-        if response.status_code != 200:
-            self.error_message = response.text
-            raise Exception('Failed to get presigned URL from Recce Cloud.')
-        presigned_url = response.json().get('presigned_url')
-        return presigned_url
+            try:
+                RecceCloud(token=token).purge_artifacts(pr_info)
+            except RecceCloudException as e:
+                return False, str(e)
+            return True, None
 
     def _load_state_from_file(self, file_path: Optional[str] = None) -> RecceState:
         file_path = file_path or self.state_file
@@ -317,7 +297,8 @@ def _load_state_from_recce_cloud(self) -> Union[RecceState, None]:
         import tempfile
         import requests
 
-        presigned_url = self._get_presigned_url(self.pr_info, RECCE_STATE_COMPRESSED_FILE, method='download')
+        presigned_url = RecceCloud(token=self.cloud_options.get('token')).get_presigned_url(
+            method=PresignedUrlMethod.DOWNLOAD, pr_info=self.pr_info, artifact_name=RECCE_STATE_COMPRESSED_FILE)
 
         compress_passwd = self.cloud_options.get('password')
         with tempfile.NamedTemporaryFile() as tmp:
@@ -377,8 +358,10 @@ def _export_state_to_recce_cloud(self, metadata: dict = None) -> Union[str, None
         import tempfile
         import requests
 
-        presigned_url = self._get_presigned_url(self.pr_info, RECCE_STATE_COMPRESSED_FILE, method='upload',
-                                                metadata=metadata)
+        presigned_url = RecceCloud(token=self.cloud_options.get('token')).get_presigned_url(
+            method=PresignedUrlMethod.UPLOAD, pr_info=self.pr_info, artifact_name=RECCE_STATE_COMPRESSED_FILE,
+            metadata=metadata)
+
         compress_passwd = self.cloud_options.get('password')
         with tempfile.NamedTemporaryFile() as tmp:
             self._export_state_to_file(tmp.name, compress=True, compress_passwd=compress_passwd)
@@ -406,23 +389,9 @@ def _export_state_to_s3_bucket(self, metadata: dict = None) -> Union[str, None]:
             s3_client.upload_file(tmp.name, s3_bucket_name, s3_bucket_key,
                                   # Casting all the values under metadata to string
                                   ExtraArgs={'Metadata': {k: str(v) for k, v in metadata.items()}})
+        RecceCloud(token=self.cloud_options.get('token')).update_github_pull_request_check(self.pr_info, metadata)
         return f'The state file is uploaded to \' s3://{s3_bucket_name}/{s3_bucket_key}\''
 
-    def _get_artifact_metadata_from_recce_cloud(self, artifact_name: str) -> dict:
-        import requests
-        token = self.cloud_options.get('token')
-        if token is None:
-            raise Exception('No token is provided to access Recce Cloud.')
-        api_url = f'{RECCE_CLOUD_API_HOST}/api/v1/{self.pr_info.repository}/pulls/{self.pr_info.id}/artifacts/{artifact_name}/metadata'
-        headers = {
-            'Authorization': f'Bearer {token}'
-        }
-        response = requests.get(api_url, headers=headers)
-        if response.status_code != 200:
-            self.error_message = response.text
-            raise Exception('Failed to get artifact metadata from Recce Cloud.')
-        return response.json()
-
     def _get_artifact_metadata_from_s3_bucket(self, artifact_name: str) -> Union[dict, None]:
         import boto3
         s3_client = boto3.client('s3')
@@ -462,28 +431,14 @@ def _export_state_to_file(self, file_path: Optional[str] = None, compress: bool
         return f'The state file is stored at \'{file_path}\''
 
     @staticmethod
-    def _purge_state_from_cloud(token: str, pr_info: PullRequestInfo) -> (bool, str):
-        import requests
-        logger.debug('Purging the state from Recce Cloud...')
-        token = token
-        api_url = f'{RECCE_CLOUD_API_HOST}/api/v1/{pr_info.repository}/pulls/{pr_info.id}/artifacts'
-        headers = {
-            'Authorization': f'Bearer {token}'
-        }
-        response = requests.delete(api_url, headers=headers)
-        if response.status_code != 204:
-            return False, response.text
-        return True, None
-
-    @staticmethod
-    def _purge_state_from_s3_bucket(host: str, pr_info: PullRequestInfo) -> (bool, str):
+    def _purge_state_from_s3_bucket(token: str, pr_info: PullRequestInfo, s3_bucket: str) -> (bool, str):
         import boto3
         from rich.console import Console
         console = Console()
         delete_objects = []
         logger.debug('Purging the state from AWS S3 bucket...')
         s3_client = boto3.client('s3')
-        s3_bucket_name = host.replace('s3://', '')
+        s3_bucket_name = s3_bucket.replace('s3://', '')
         s3_key_prefix = f'github/{pr_info.repository}/pulls/{pr_info.id}/'
         list_response = s3_client.list_objects_v2(Bucket=s3_bucket_name, Prefix=s3_key_prefix)
         if 'Contents' in list_response:
@@ -497,4 +452,5 @@ def _purge_state_from_s3_bucket(host: str, pr_info: PullRequestInfo) -> (bool, s
         delete_response = s3_client.delete_objects(Bucket=s3_bucket_name, Delete={'Objects': delete_objects})
         if 'Deleted' not in delete_response:
             return False, 'Failed to delete the state file from the S3 bucket.'
+        RecceCloud(token=token).update_github_pull_request_check(pr_info)
         return True, None

recce/util/recce_cloud.py

@@ -0,0 +1,82 @@
+import logging
+import os
+
+import requests
+
+from recce.pull_request import PullRequestInfo
+
+RECCE_CLOUD_API_HOST = os.environ.get('RECCE_CLOUD_API_HOST', 'https://staging.cloud.datarecce.io')
+
+logger = logging.getLogger('uvicorn')
+
+
+class PresignedUrlMethod:
+    UPLOAD = 'upload'
+    DOWNLOAD = 'download'
+
+
+class RecceCloudException(Exception):
+    def __init__(self, message: str, reason: str, status_code: int):
+        super().__init__(message)
+        self.reason = reason
+        self.status_code = status_code
+
+
+class RecceCloud:
+    def __init__(self, token: str):
+        self.token = token
+        self.base_url = f'{RECCE_CLOUD_API_HOST}/api/v1'
+
+    def _request(self, method, url, data=None):
+        headers = {
+            'Authorization': f'Bearer {self.token}'
+        }
+        return requests.request(method, url, headers=headers, json=data)
+
+    def get_presigned_url(self,
+                          method: PresignedUrlMethod,
+                          pr_info: PullRequestInfo,
+                          artifact_name: str,
+                          metadata: dict = None) -> str:
+        api_url = f'{self.base_url}/{pr_info.repository}/pulls/{pr_info.id}/artifacts/{method}?artifact_name={artifact_name}'
+        response = self._request('POST', api_url, data=metadata)
+        if response.status_code != 200:
+            raise RecceCloudException(
+                message='Failed to {method} artifact {preposition} Recce Cloud.'.format(
+                    method=method,
+                    preposition='from' if method == PresignedUrlMethod.DOWNLOAD else 'to'
+                ),
+                reason=response.text,
+                status_code=response.status_code
+            )
+        presigned_url = response.json().get('presigned_url')
+        return presigned_url
+
+    def get_artifact_metadata(self, pr_info: PullRequestInfo, artifact_name: str) -> dict:
+        api_url = f'{self.base_url}/{pr_info.repository}/pulls/{pr_info.id}/artifacts/{artifact_name}/metadata'
+        response = self._request('GET', api_url)
+        if response.status_code != 200:
+            raise RecceCloudException(
+                message='Failed to get artifact metadata from Recce Cloud.',
+                reason=response.text,
+                status_code=response.status_code
+            )
+        return response.json()
+
+    def purge_artifacts(self, pr_info: PullRequestInfo):
+        api_url = f'{self.base_url}/{pr_info.repository}/pulls/{pr_info.id}/artifacts'
+        response = self._request('DELETE', api_url)
+        if response.status_code != 204:
+            raise RecceCloudException(
+                message='Failed to purge artifacts from Recce Cloud.',
+                reason=response.text,
+                status_code=response.status_code
+            )
+
+    def update_github_pull_request_check(self, pr_info: PullRequestInfo, metadata: dict = None):
+        api_url = f'{self.base_url}/{pr_info.repository}/pulls/{pr_info.id}/github/checks'
+        try:
+            self._request('POST', api_url, data=metadata)
+        except Exception as e:
+            # We don't care the response of this request, so we don't need to raise any exception.
+            logger.debug(f'Failed to update the GitHub PR check. Reason: {str(e)}')