Merge pull request #5 from anecdotes-ai/feature/add-sync

Add the option to run the queries sync

.github/workflows/sanity.yml

@@ -15,7 +15,7 @@ jobs:
         with:
           snowflake_account: ${{ secrets.SNOWFLAKE_ACCOUNT }}
           snowflake_warehouse: ${{ secrets.SNOWFLAKE_WAREHOUSE }}
-          snowflake_username: ${{ secrets.SNOWFLAKE_USER }}
+          snowflake_username: ${{ secrets.SNOWFLAKE_USERNAME }}
           snowflake_password: ${{ secrets.SNOWFLAKE_PASSWORD }}
           queries: 'call system$wait(1);
                     select CURRENT_VERSION();
@@ -45,7 +45,7 @@ jobs:
         with:
           snowflake_account: ${{ secrets.SNOWFLAKE_ACCOUNT }}
           snowflake_warehouse: ${{ secrets.SNOWFLAKE_WAREHOUSE }}
-          snowflake_username: ${{ secrets.SNOWFLAKE_USER }}
+          snowflake_username: ${{ secrets.SNOWFLAKE_USERNAME }}
           snowflake_password: ${{ secrets.SNOWFLAKE_PASSWORD }}
           queries: 'call system$wait(${{matrix.sleep}})'
 

.gitignore

@@ -129,4 +129,7 @@ dmypy.json
 .pyre/
 
 # visual studio code
-.vscode
+.vscode
+
+# PyCharm
+.idea

action.yml

@@ -25,6 +25,9 @@ inputs:
   queries:
     description: 'List of SQL queries, seperated by ;'
     required: true
+  sync:
+    description: 'Whether to run the queries sync (default is async)'
+    required: false
 
 outputs:
   queries_results:

main.py

@@ -9,7 +9,9 @@
 
 def main():
     load_dotenv() # only on local run
+    print(os.environ)
     queries_list = os.environ['INPUT_QUERIES'].split(';')
+    sync = os.environ.get("INPUT_SYNC", False)
     warehouse = os.environ['INPUT_SNOWFLAKE_WAREHOUSE']
     snowflake_account = os.environ['INPUT_SNOWFLAKE_ACCOUNT']
     snowflake_username = os.environ['INPUT_SNOWFLAKE_USERNAME']
@@ -26,16 +28,27 @@ def main():
         con.set_db_warehouse(warehouse)
 
         query_results = []
-        for query in queries_list:
-            query_result = con.query(query)
-            query_results.append(query_result)
-            print("### Running query ###")
-            print(f"[!] Query id - {query_result.query_id}")
-            print(f"[!] Running query ### - {query}")
-
-        json_results = asyncio.run(utils.gather_all_results(query_results))
+        # default, run all queries async
+        if not sync:
+            for query in queries_list:
+                query_result = con.query(query)
+                query_results.append(query_result)
+                print("### Running query ###")
+                print(f"[!] Query id - {query_result.query_id}")
+                print(f"[!] Running query ### - {query}")
+            json_results = asyncio.run(utils.gather_all_results(query_results))
+        # o/w, run them sync
+        else:
+            json_results = {}
+            for query in queries_list:
+                query_result = con.query(query)
+                print("### Running query ###")
+                print(f"[!] Query id - {query_result.query_id}")
+                print(f"[!] Running query ### - {query}")
+                json_results[query_result.query_id] = query_result.fetch_results_sync()
 
     utils.set_github_action_output('queries_results', json.dumps(json_results))
-
+
+
 if __name__ == '__main__':
     main()

snowflake_connector.py

@@ -64,4 +64,8 @@ async def fetch_results(self):
         while self.is_query_running():
             await asyncio.sleep(0.1)
 
-        return self._fetch_results()
+        return self._fetch_results()
+
+    def fetch_results_sync(self):
+        self.cursor.get_results_from_sfqid(self.query_id)
+        return self.cursor.fetchall()