diff --git a/main.py b/main.py index 234f61b..a28c41b 100644 --- a/main.py +++ b/main.py @@ -27,19 +27,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}") - + # 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 = utils.gather_all_results_sync(query_results) + 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_all() utils.set_github_action_output('queries_results', json.dumps(json_results)) - + + if __name__ == '__main__': main() \ No newline at end of file diff --git a/snowflake_connector.py b/snowflake_connector.py index e153e1f..202f68a 100644 --- a/snowflake_connector.py +++ b/snowflake_connector.py @@ -64,8 +64,4 @@ async def fetch_results(self): while self.is_query_running(): await asyncio.sleep(0.1) - return self._fetch_results() - - def fetch_results_sync(self): - self.cursor.get_results_from_sfqid(self.query_id) - return self.cursor.fetchall() \ No newline at end of file + return self._fetch_results() \ No newline at end of file diff --git a/utils.py b/utils.py index 79d691e..1acb1c2 100644 --- a/utils.py +++ b/utils.py @@ -35,19 +35,4 @@ async def gather_all_results(query_result_list: List[QueryResult]) -> dict: print(row) json_total_results[done_task.get_name()].append(str(row)) - return json_total_results - - -def gather_all_results_sync(query_results_list: List[QueryResult]) -> dict: - """ - - Args: - query_results_list: - - Returns: - - """ - json_total_results = {} - for query in query_results_list: - json_total_results[query.query_id] = query.fetch_results_sync() return json_total_results \ No newline at end of file