Merge pull request #40 from ksyusha123/bq-upgrade

Bq upgrade

logica.py

@@ -219,8 +219,17 @@ def main(argv):
         parsed_rules, predicates_list).RetrieveTypes(filename)
       return 0
     elif engine == 'bigquery':
+      from google import auth as terminal_auth
+      credentials, project = terminal_auth.default()
+      if not project:
+        from google.colab import auth as colab_auth
+        colab_auth.authenticate_user()
+        print("Please enter project_id to use for BigQuery queries.")
+        project = input()
+        print("project_id is set to %s" % project)
+        print("You can change it with logica.colab_logica.SetProject command.")
       type_retrieval_service.BigQueryTypeRetrievalService(
-        parsed_rules, predicates_list).RetrieveTypes(filename)
+        parsed_rules, predicates_list, credentials, project).RetrieveTypes(filename)
       return 0
     else:
       raise unsupported_engine_exception.UnsupportedEngineException(engine)

type_inference/bad_schema_exception.py

@@ -21,7 +21,7 @@
   from type_inference import type_retrieval_exception
 else:
   from ..common import color
-  import type_retrieval_exception
+  from ..type_inference import type_retrieval_exception
 
 
 class BadSchemaException(type_retrieval_exception.TypeRetrievalException):

type_inference/bigquery_type_retriever.py

@@ -20,8 +20,8 @@
   from type_inference import bigquery_type_parser
   from type_inference import unknown_bigquery_type_exception
 else:
-  import bigquery_type_parser
-  import unknown_bigquery_type_exception
+  from ..type_inference import bigquery_type_parser
+  from ..type_inference import unknown_bigquery_type_exception
 
 
 class BigQueryTypeRetriever:

type_inference/postgresql_type_retriever.py

@@ -19,7 +19,7 @@
 if '.' not in __package__:
   from type_inference import postgresql_type_parser
 else:
-  import postgresql_type_parser
+  from ..type_inference import postgresql_type_parser
 
 
 class PostgresqlTypeRetriever:

type_inference/type_retrieval_service.py

@@ -16,17 +16,11 @@
 
 import json
 from typing import Dict
-import psycopg2
-from google.cloud import bigquery
 
 if '.' not in __package__:
   from type_inference import bad_schema_exception
-  from type_inference import postgresql_type_retriever
-  from type_inference import bigquery_type_retriever
 else:
-  import bad_schema_exception
-  import postgresql_type_retriever
-  import bigquery_type_retriever
+  from ..type_inference import bad_schema_exception
 
 
 def ValidateRuleAndGetTableName(rule: dict, lower_table_name: bool = False) -> str:
@@ -59,6 +53,10 @@ class PostgresqlTypeRetrievalService:
   """The class is an entry point for type retrieval using postgresql."""
   def __init__(self, parsed_rules, predicate_names,
                connection_string='dbname=logica user=logica password=logica host=127.0.0.1'):
+    if '.' not in __package__:
+      from type_inference import postgresql_type_retriever
+    else:
+      from ..type_inference import postgresql_type_retriever
     predicate_names_as_set = set(predicate_names)
     self.parsed_rules = [r for r in parsed_rules if r['head']['predicate_name'] in predicate_names_as_set]
     self.connection_string = connection_string
@@ -72,6 +70,7 @@ def ValidateParsedRulesAndGetTableNames(self) -> Dict[str, str]:
   def RetrieveTypes(self, filename):
     filename = filename.replace('.l', '_schema.l')
 
+    import psycopg2
     with psycopg2.connect(self.connection_string) as conn:
       with conn.cursor() as cursor:
         # for each given table this SQL query returns json object
@@ -101,11 +100,16 @@ def RetrieveTypes(self, filename):
 class BigQueryTypeRetrievalService:
   """The class is an entry point for type retrieval using bigquery."""
   def __init__(self, parsed_rules, predicate_names,
-               project='bigquery-logica'):
+               credentials=None, project='bigquery-logica'):
+    if '.' not in __package__:
+      from type_inference import bigquery_type_retriever
+    else:
+      from ..type_inference import bigquery_type_retriever
     predicate_names_as_set = set(predicate_names)
     self.parsed_rules = [r for r in parsed_rules
                          if r['head']['predicate_name'] in predicate_names_as_set]
     self.project = project
+    self.credentials = credentials
     self.table_names = self.ValidateParsedRulesAndGetTableNames()
     self.type_retriever = bigquery_type_retriever.BigQueryTypeRetriever()
 
@@ -116,7 +120,9 @@ def ValidateParsedRulesAndGetTableNames(self) -> Dict[str, str]:
   def RetrieveTypes(self, filename):
     filename = filename.replace('.l', '_schema.l')
 
-    client = bigquery.Client(project=self.project)
+    from google.cloud import bigquery
+
+    client = bigquery.Client(credentials=self.credentials, project=self.project) # it works for us even if we don't give any credentials
     job_config = bigquery.QueryJobConfig(
       query_parameters=[
         bigquery.ArrayQueryParameter("tables", "STRING", self.table_names),

type_inference/unsupported_engine_exception.py

@@ -21,7 +21,7 @@
   from type_inference import type_retrieval_exception
 else:
   from ..common import color
-  import type_retrieval_exception
+  from ..type_inference import type_retrieval_exception
 
 
 class UnsupportedEngineException(type_retrieval_exception.TypeRetrievalException):