chore(tool): add manifest extract for query_local_run (#1191)

Author: goldmedal

ibis-server/tools/query_local_run.py

@@ -22,6 +22,7 @@
 from wren_core import SessionContext
 from app.model.data_source import BigQueryConnectionInfo
 from app.model.data_source import DataSourceExtension
+import wren_core
 
 if sys.stdin.isatty():
     print("please provide the SQL query via stdin, e.g. `python query_local_run.py < test.sql`", file=sys.stderr)
@@ -57,6 +58,14 @@
 with open(connection_info_path) as file:
     connection_info = json.load(file)
 
+# Extract the requried tables from the SQL query
+extractor = wren_core.ManifestExtractor(encoded_str)
+tables = extractor.resolve_used_table_names(sql)
+print("# Tables used in the SQL query:", tables)
+# Extract the manifest for the required tables
+manifest = extractor.extract_by(tables)
+encoded_str = wren_core.to_json_base64(manifest)
+
 print("### Starting the session context ###")
 print("#")
 session_context = SessionContext(encoded_str, function_list_path + f"/{data_source}.csv")