Fixes for BigQuery, REPL

Author: erezsh

sqeleton/databases/base.py

@@ -487,7 +487,7 @@ def _normalize_table_path(self, path: DbPath) -> DbPath:
     def parse_table_name(self, name: str) -> DbPath:
         return parse_table_name(name)
 
-    def _query_cursor(self, c, sql_code: str):
+    def _query_cursor(self, c, sql_code: str) -> QueryResult:
         assert isinstance(sql_code, str), sql_code
         try:
             c.execute(sql_code)
@@ -499,7 +499,7 @@ def _query_cursor(self, c, sql_code: str):
             # logger.error(f'Caused by SQL: {sql_code}')
             raise
 
-    def _query_conn(self, conn, sql_code: Union[str, ThreadLocalInterpreter]) -> list:
+    def _query_conn(self, conn, sql_code: Union[str, ThreadLocalInterpreter]) -> QueryResult:
         c = conn.cursor()
         callback = partial(self._query_cursor, c)
         return apply_query(callback, sql_code)
@@ -542,7 +542,7 @@ def set_conn(self):
         except Exception as e:
             self._init_error = e
 
-    def _query(self, sql_code: Union[str, ThreadLocalInterpreter]):
+    def _query(self, sql_code: Union[str, ThreadLocalInterpreter]) -> QueryResult:
         r = self._queue.submit(self._query_in_worker, sql_code)
         return r.result()
 

sqeleton/databases/bigquery.py

@@ -19,7 +19,7 @@
 )
 from ..abcs import Compilable
 from ..queries import this, table, SKIP, code
-from .base import BaseDialect, Database, import_helper, parse_table_name, ConnectError, apply_query
+from .base import BaseDialect, Database, import_helper, parse_table_name, ConnectError, apply_query, QueryResult
 from .base import TIMESTAMP_PRECISION_POS, ThreadLocalInterpreter, Mixin_RandomSample
 
 
@@ -161,16 +161,18 @@ def _query_atom(self, sql_code: str):
         from google.cloud import bigquery
 
         try:
-            res = list(self._client.query(sql_code))
+            result = self._client.query(sql_code).result()
+            columns = [c.name for c in result.schema]
+            rows = list(result)
         except Exception as e:
             msg = "Exception when trying to execute SQL code:\n    %s\n\nGot error: %s"
             raise ConnectError(msg % (sql_code, e))
 
-        if res and isinstance(res[0], bigquery.table.Row):
-            res = [tuple(self._normalize_returned_value(v) for v in row.values()) for row in res]
-        return res
+        if rows and isinstance(rows[0], bigquery.table.Row):
+            rows = [tuple(self._normalize_returned_value(v) for v in row.values()) for row in rows]
+        return QueryResult(rows, columns)
 
-    def _query(self, sql_code: Union[str, ThreadLocalInterpreter]):
+    def _query(self, sql_code: Union[str, ThreadLocalInterpreter]) -> QueryResult:
         return apply_query(self._query_atom, sql_code)
 
     def close(self):

sqeleton/repl.py

@@ -47,7 +47,9 @@ def repl(uri):
                 help()
                 continue
             try:
-                schema = db.query_table_schema((table_name,))
+                path = db.parse_table_name(table_name)
+                print('->', path)
+                schema = db.query_table_schema(path)
             except Exception as e:
                 logging.error(e)
             else: