From a09accc8f88e808fdc78a804b2fe70ec9b6af0eb Mon Sep 17 00:00:00 2001
From: Kent Huang <kent@infuseai.io>
Date: Wed, 12 Jun 2024 15:35:46 +0800
Subject: [PATCH] [Feature] Change to use RecceStateLoader to export file

Signed-off-by: Kent Huang <kent@infuseai.io>
---
 recce/cli.py   | 16 ++++++++++++++--
 recce/run.py   | 11 +++++++++--
 recce/state.py | 24 ++++++++++++++++++------
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/recce/cli.py b/recce/cli.py
index 105cfcbd..843ce11c 100644
--- a/recce/cli.py
+++ b/recce/cli.py
@@ -209,6 +209,7 @@ def server(host, port, state_file=None, **kwargs):
 @add_options(dbt_related_options)
 @add_options(sqlmesh_related_options)
 @add_options(recce_options)
+@add_options(recce_cloud_options)
 def run(output, **kwargs):
     is_github_action, pr_url = check_github_ci_env(**kwargs)
     if is_github_action is True and pr_url is not None:
@@ -217,7 +218,16 @@ def run(output, **kwargs):
     # Initialize Recce Config
     RecceConfig(config_file=kwargs.get('config'))
 
-    return asyncio.run(cli_run(output, **kwargs))
+    cloud_mode = kwargs.get('cloud', False)
+    state_file = kwargs.get('state_file')
+    cloud_options = {
+        'host': kwargs.get('state_file_host'),
+        'secret': kwargs.get('state_file_secret'),
+    } if cloud_mode else None
+    recce_state = RecceStateLoader(review_mode=False, cloud_mode=cloud_mode,
+                                   state_file=state_file, cloud_options=cloud_options)
+
+    return asyncio.run(cli_run(output, recce_state=recce_state, **kwargs))
 
 
 @cli.command(cls=TrackCommand)
@@ -229,9 +239,11 @@ def summary(state_file, **kwargs):
     from rich.console import Console
     from .core import load_context
     console = Console()
+    recce_state = RecceStateLoader(review_mode=False, cloud_mode=False,
+                                   state_file=state_file, cloud_options=None)
     try:
         # Load context in review mode, won't need to check dbt_project.yml file.
-        ctx = load_context(**kwargs, state_file=state_file, review=True)
+        ctx = load_context(**kwargs, recce_state=recce_state, review=True)
     except Exception as e:
         console.print("[[red]Error[/red]] Failed to generate summary:")
         console.print(f"{e}")
diff --git a/recce/run.py b/recce/run.py
index 9c78e2b9..5c0d9008 100644
--- a/recce/run.py
+++ b/recce/run.py
@@ -13,7 +13,7 @@
 from recce.core import RecceContext
 from recce.models.types import RunType
 from recce.pull_request import fetch_pr_metadata_from_event_path
-from recce.state import RecceState, PullRequestInfo
+from recce.state import RecceState, PullRequestInfo, RecceStateLoader
 
 
 def check_github_ci_env(**kwargs):
@@ -253,6 +253,13 @@ async def cli_run(output_state_file: str, **kwargs):
     state: RecceState = ctx.export_state()
     state.pull_request = fetch_pr_metadata(**kwargs)
 
-    state.to_state_file(output_state_file)
+    cloud_mode = kwargs.get('cloud', False)
+    cloud_options = {
+        'host': kwargs.get('state_file_host'),
+        'secret': kwargs.get('state_file_secret'),
+    } if cloud_mode else None
+    recce_state = RecceStateLoader(review_mode=False, cloud_mode=cloud_mode, state_file=output_state_file,
+                                   cloud_options=cloud_options)
+    recce_state.export(state)
     print(f'The state file is stored at [{output_state_file}]')
     return rc
diff --git a/recce/state.py b/recce/state.py
index 7f2b0f2a..3655d955 100644
--- a/recce/state.py
+++ b/recce/state.py
@@ -154,6 +154,7 @@ def __init__(self,
 
         # Load the state
         self.load()
+
         pass
 
     def verify(self) -> bool:
@@ -178,11 +179,14 @@ def update(self, state: RecceState):
         self.state = state
 
     def load(self) -> RecceState:
-        if self.cloud_mode:
-            # TODO: Load the state from cloud storage
+        try:
+            if self.cloud_mode:
+                # TODO: Load the state from cloud storage
+                pass
+            else:
+                self.state = RecceState.from_file(self.state_file)
+        except Exception:
             pass
-        else:
-            self.state = RecceState.from_file(self.state_file)
         return self.state
 
     def export(self, state: RecceState = None):
@@ -194,11 +198,19 @@ def export(self, state: RecceState = None):
         else:
             self._export_state_to_file()
 
-    def _export_state_to_file(self):
+    def _export_state_to_cloud(self):
         # TODO: export the state to remote cloud storage
+        print('TODO: export the state to remote cloud storage')
+        print(f"        remote host: {self.cloud_options.get('host')}")
         pass
 
-    def _export_state_to_cloud(self):
+    def _export_state_to_recce_cloud(self):
+        pass
+
+    def _export_state_to_s3_bucket(self):
+        pass
+
+    def _export_state_to_file(self):
         """
                 Store the state to a file. Store happens when terminating the server or run instance.
                 """