kaqing 2.0.171__py3-none-any.whl → 2.0.203__py3-none-any.whl

adam/sql/sql_completer.py

@@ -2,7 +2,7 @@ from enum import Enum
 from typing import Callable
 
 import sqlparse
-from sqlparse.sql import Statement, Token
+from sqlparse.sql import Token
 
 from adam.sql.term_completer import TermCompleter
 from adam.utils_repl.automata_completer import AutomataCompleter
@@ -26,11 +26,9 @@ class SqlCompleter(AutomataCompleter[Token]):
         tokens = []
 
         stmts = sqlparse.parse(text)
-        if not stmts:
-            tokens = []
-        else:
-            statement: Statement = stmts[0]
-            tokens = statement.tokens
+        if stmts:
+            for stmt in stmts:
+                tokens.extend(stmt.tokens)
 
         return tokens
 

adam/sql/sql_state_machine.py

@@ -2,6 +2,7 @@ from typing import Callable
 from sqlparse.sql import Token
 from sqlparse import tokens as TOKEN
 
+from adam.utils import log_exc
 from adam.utils_repl.state_machine import StateMachine, State
 
 __all__ = [
@@ -77,8 +78,7 @@ SQL_SPEC = [
     'select_from_sq_                 > as               > select_from_x_as                        ^ as',
     'select_from_x_comma_            > name|audit       > select_from_x                           ^ tables',
     'select_from_x_                                                                               ^ as,where,inner join,left outer join,right outer join,full outer join,group by,order by,limit,&',
-    '-                               > &                > select_from_x$',
-    'select_from_x_as_x_             > ,                > select_from_x_comma_                    ^ where,inner join,left outer join,right outer join,full outer join,group by,order by,limit',
+    'select_from_x_as_x_             > ,                > select_from_x_comma_                    ^ where,inner join,left outer join,right outer join,full outer join,group by,order by,limit,&',
     '-                               > as               > select_from_x_as',
     '-                               > where            > select_where',
     '-                               > order            > select_order',
@@ -97,6 +97,7 @@ SQL_SPEC = [
     '-                               > full             > select_from_x_full',
     '-                               > full outer join  > select_join',
     '-                               > ;                > ',
+    '-                               > &                > select_from_x$',
     'select_from_x_as_               > name             > select_from_x_as_x                      ^ x,y,z',
     'select_from_x_as_x              > ,                > select_from_x_as_x_comma_',
     '-                               > ;                > ',
@@ -119,34 +120,38 @@ SQL_SPEC = [
     'select_where_a_not_op           > name|single|num  > select_where_sc                         ^ single',
     'select_where_a_op               > name|single|num  > select_where_sc                         ^ single',
     'select_where_sc                 > ;                > ',
-    'select_where_sc_                > and|or           > select_where                            ^ and,or,order by,group by,limit',
+    'select_where_sc_                > and|or           > select_where                            ^ and,or,order by,group by,limit,&',
     '-                               > group            > select_group',
     '-                               > group by         > select_group_by',
     '-                               > order            > select_order',
     '-                               > order by         > select_order_by',
     '-                               > limit            > select_where_sc_limit',
     '-                               > ;                > ',
+    '-                               > &                > select_from_x$',
     'select_group_                   > by               > select_group_by                         ^ by',
     'select_group_by_                > name             > select_group_by_a                       ^ columns',
     'select_group_by_a               > ,                > select_group_by_a_comma_                ^ columns',
     '-                               > ;                > ',
     'select_group_by_a_comma_        > name             > select_group_by_a                       ^ columns',
-    'select_group_by_a_              > limit            > select_where_sc_limit                   ^ limit,order by',
+    'select_group_by_a_              > limit            > select_where_sc_limit                   ^ limit,order by,&',
     '-                               > order            > select_order',
     '-                               > order by         > select_order_by',
     '-                               > ;                > ',
+    '-                               > &                > select_from_x$',
     'select_order_                   > by               > select_order_by                         ^ by',
     'select_order_by_                > name             > select_order_by_a                       ^ columns',
     'select_order_by_a               > ,                > select_order_by_a_comma_',
     '-                               > ;                > ',
     'select_order_by_a_comma_        > name             > select_order_by_a                       ^ columns',
-    'select_order_by_a_              > desc|asc         > select_order_by_a_desc                  ^ desc,asc,limit',
+    'select_order_by_a_              > desc|asc         > select_order_by_a_desc                  ^ desc,asc,limit,&',
     '-                               > limit            > select_where_sc_limit',
     '-                               > ;                > ',
+    '-                               > &                > select_from_x$',
     'select_order_by_a_desc          > ,                > select_order_by_a_comma_',
     '-                               > ;                > ',
-    'select_order_by_a_desc_         > limit            > select_where_sc_limit                   ^ limit',
+    'select_order_by_a_desc_         > limit            > select_where_sc_limit                   ^ limit,&',
     '-                               > ;                > ',
+    '-                               > &                > select_from_x$',
     'select_where_sc_limit_          > num              > select_where_sc_limit_num               ^ 1',
     'select_where_sc_limit_num       > ;                > ',
     'select_where_sc_limit_num_rp__  > as               > select_from_x_as                        ^ as',
@@ -202,7 +207,9 @@ SQL_SPEC = [
     'insert_values                   > (                > insert_values_lp_',
     'insert_values_lp_               > name|single|num  > insert_values_lp_v                      ^ single',
     'insert_values_lp_v              > ,                > insert_values_lp_v_comma_',
+    '-                               > )                > insert_values_lp_v_rp_',
     'insert_values_lp_v_comma_       > name|single|num  > insert_values_lp_v',
+    'insert_values_lp_v_rp__         > &                > insert_values_lp_v_rp_$                 ^ &',
 
     # <update_statement> ::= UPDATE <table_name>
     #                        SET <set_clause_list>
@@ -223,8 +230,9 @@ SQL_SPEC = [
     'update_set_a_op                 > name|single|num  > update_set_sc                           ^ single',
     'update_set_sc                   > ,                > update_set_sc_comma_',
     'update_set_sc_comma_            > name             > update_set_a                            ^ id',
-    'update_set_sc_                  > ,                > update_set_sc_comma_                    ^ where',
+    'update_set_sc_                  > ,                > update_set_sc_comma_                    ^ where,&',
     '-                               > where            > update_where',
+    '-                               > &                > update_set_sc$                          ^ &',
     'update_where_                   > name             > update_where_a                          ^ id',
     'update_where_a                  > comparison       > update_where_a_op',
     'update_where_a_                 > comparison       > update_where_a_op                       ^ =,<,<=,>,>=,<>,like,not,in',
@@ -241,7 +249,7 @@ SQL_SPEC = [
     'update_where_a_in_lp_a_comma_   > name|single|num  > update_where_a_in_lp_a                  ^ single',
     'update_where_a_not_op           > name|single|num  > update_where_sc                         ^ single',
     'update_where_a_op               > name|single|num  > update_where_sc                         ^ single',
-    'update_where_sc_                > and|or           > update_where                            ^ and,or',
+    'update_where_sc_                > and|or           > update_where                            ^ and,or,&',
 
     # <delete_statement> ::= DELETE FROM <table_name> [ WHERE <search_condition> ]
 
@@ -403,9 +411,15 @@ CQL_SPEC = SQL_SPEC + [
     'export_table_to_                > athena|sqlite|csv > export_table_to$                        ^ athena,sqlite,csv',
 
     '                                > import           > import',
-    'import_                         > session          > import_session                          ^ session',
+    'import_                         > session          > import_session                          ^ session,files',
+    '-                               > files            > import_files',
     'import_session_                 > name             > import_session_s                        ^ export-sessions-incomplete',
     'import_session_s_               > to               > import_session_to                       ^ to',
+    'import_files_                   > name             > import_files_f',
+    'import_files_f                  > ,                > import_files',
+    'import_files_f_                 > as               > import_files_f_as                       ^ as',
+    'import_files_f_as_              > name             > import_files_f_as_a',
+    'import_files_f_as_a_            > to               > import_session_to                       ^ to',
     'import_session_to_              > athena|sqlite    > import_session_to$                      ^ athena,sqlite',
 
     '                                > consistency      > consistency',
@@ -433,7 +447,7 @@ CQL_SPEC = SQL_SPEC + [
 CQL_KEYWORDS = SQL_KEYWORDS + [
     'schema', 'keyspace', 'keyspaces', 'tables', 'export', 'copy', 'consistency',
     'quorum', 'all', 'serial', 'one', 'each_quorum', 'local_quorum', 'any', 'local_one', 'two', 'three', 'local_serial', 'to',
-    'database', 'databases', 'session', 'sessions', 'clean', 'up', 'athena', 'sqlite', 'csv', 'import'
+    'database', 'databases', 'session', 'sessions', 'clean', 'up', 'athena', 'sqlite', 'csv', 'import', 'files'
 ]
 
 CQL_EXPANDABLE_NAMES = EXPANDABLE_NAMES | {
@@ -547,7 +561,7 @@ class SqlStateMachine(StateMachine[Token]):
                 elif token.ttype == TOKEN.Operator.Comparison:
                     it = 'comparison'
 
-                try:
+                with log_exc(False):
                     # print(f'\n{state.to_s} > {it} > ', end='')
                     if comeback_state:
                         state = comeback_state
@@ -558,8 +572,6 @@ class SqlStateMachine(StateMachine[Token]):
 
                     if last_name:
                         state.context['last_name'] = last_name
-                except:
-                    pass
 
         return state
 

adam/sso/authn_ad.py

@@ -1,6 +1,5 @@
 import json
 import re
-import traceback
 import jwt
 import requests
 from urllib.parse import urlparse, parse_qs
@@ -8,7 +7,7 @@ from urllib.parse import urlparse, parse_qs
 from adam.log import Log
 from adam.sso.authenticator import Authenticator
 from adam.sso.id_token import IdToken
-from adam.utils import debug
+from adam.utils import debug, log_exc
 from .idp_login import IdpLogin
 from adam.config import Config
 
@@ -139,7 +138,7 @@ class AdAuthenticator(Authenticator):
 
     def parse_id_token(self, id_token: str) -> IdToken:
         jwks_url = Config().get('idps.ad.jwks-uri', '')
-        try:
+        with log_exc():
             jwks_client = jwt.PyJWKClient(jwks_url, cache_jwk_set=True, lifespan=360)
             signing_key = jwks_client.get_signing_key_from_jwt(id_token)
             data = jwt.decode(
@@ -164,7 +163,5 @@ class AdAuthenticator(Authenticator):
                 nbf=data['nbf'] if 'nbf' in data else 0,
                 exp=data['exp'] if 'exp' in data else 0
             )
-        except:
-            debug(traceback.format_exc())
 
         return None

adam/sso/authn_okta.py

@@ -8,7 +8,7 @@ from adam.sso.id_token import IdToken
 
 from .idp_login import IdpLogin
 from adam.config import Config
-from adam.utils import debug, log2
+from adam.utils import debug, log2, log_exc
 
 class OktaException(Exception):
     pass
@@ -95,7 +95,7 @@ class OktaAuthenticator(Authenticator):
             return None
 
         jwks_url = Config().get('idps.okta.jwks-uri', 'https://c3energy.okta.com/oauth2/v1/keys')
-        try:
+        with log_exc():
             jwks_client = jwt.PyJWKClient(jwks_url, cache_jwk_set=True, lifespan=360)
             signing_key = jwks_client.get_signing_key_from_jwt(id_token)
             data = jwt.decode(
@@ -121,7 +121,5 @@ class OktaAuthenticator(Authenticator):
                 nbf=data['nbf'] if 'nbf' in data else 0,
                 exp=data['exp'] if 'exp' in data else 0
             )
-        except:
-            pass
 
         return None

adam/sso/cred_cache.py

@@ -1,10 +1,9 @@
 import os
 from pathlib import Path
-import traceback
 from dotenv import load_dotenv
 
 from adam.config import Config
-from adam.utils import debug
+from adam.utils import debug, log_exc
 from adam.utils_k8s.kube_context import KubeContext
 
 class CredCache:
@@ -35,10 +34,8 @@ class CredCache:
     def cache(self, username: str, password: str = None):
         if os.path.exists(self.env_f):
             with open(self.env_f, 'w') as file:
-                try:
+                with log_exc():
                     file.truncate()
-                except:
-                    debug(traceback.format_exc())
 
         updated = []
         updated.append(f'IDP_USERNAME={username}')

adam/sso/idp.py

@@ -3,7 +3,6 @@ import getpass
 import os
 import sys
 import termios
-import traceback
 from typing import Callable, TypeVar
 import requests
 from kubernetes import config
@@ -15,7 +14,7 @@ from .cred_cache import CredCache
 from .idp_session import IdpSession
 from .idp_login import IdpLogin
 from adam.config import Config
-from adam.utils import debug, log
+from adam.utils import debug, log, log_exc
 
 T = TypeVar('T')
 
@@ -31,6 +30,8 @@ class Idp:
     def login(app_host: str, username: str = None, idp_uri: str = None, forced = False, use_token_from_env = True, use_cached_creds = True, verify = True) -> IdpLogin:
         session: IdpSession = IdpSession.create(username, app_host, app_host, idp_uri=idp_uri)
 
+        debug(f'Idp.login({username})')
+
         if use_token_from_env:
             if l0 := session.login_from_env_var():
                 return l0
@@ -39,11 +40,9 @@ class Idp:
                 token_server = Config().get('app.login.token-server-url', 'http://localhost:{port}').replace('{port}', port)
                 res: requests.Response = requests.get(token_server)
                 if res.status_code == 200 and res.text:
-                    try:
+                    with log_exc():
                         # may fail if the idp token is not complete
                         return session.login_from_token(res.text)
-                    except:
-                        pass
 
         r: IdpLogin = None
         try:
@@ -59,8 +58,9 @@ class Idp:
                     default_user = CredCache().get_username()
                     debug(f'User read from cache: {default_user}')
 
-                if from_env := os.getenv('USERNAME'):
-                    default_user = from_env
+                # no value in using USERNAME
+                # if from_env := os.getenv('USERNAME') and in_docker():
+                #     default_user = from_env
                 if default_user and default_user != username:
                     session = IdpSession.create(default_user, app_host, app_host)
 
@@ -125,7 +125,7 @@ class Idp:
             termios.tcsetattr(fd, termios.TCSADRAIN, old)
 
     def try_kubeconfig(username: str, kubeconfig: str):
-        try:
+        with log_exc():
             if kubeconfig[0] == '\t':
                 kubeconfig = kubeconfig[1:]
             kubeconfig_string = base64.b64decode(kubeconfig.encode('ascii') + b'==').decode('utf-8')
@@ -136,8 +136,5 @@ class Idp:
                 Secrets.list_secrets(os.getenv('NAMESPACE'))
 
                 return IdpLogin(None, None, None, username)
-        except:
-            debug(traceback.format_exc())
-            pass
 
         return None