kaqing 2.0.110__py3-none-any.whl → 2.0.214__py3-none-any.whl

adam/sql/sql_completer.py

@@ -1,7 +1,8 @@
+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
@@ -12,41 +13,42 @@ __all__ = [
     "SqlCompleter",
 ]
 
-def default_columns(_: list[str]):
+def default_columns(x: list[str]):
     return 'id,x.,y.,z.'.split(',')
 
+class SqlVariant(Enum):
+    SQL = 'sql'
+    CQL = 'cql'
+    ATHENA = 'athena'
+
 class SqlCompleter(AutomataCompleter[Token]):
     def tokens(self, text: str) -> list[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
 
     def __init__(self,
                  tables: Callable[[], list[str]],
                  dml: str = None,
-                 columns: Callable[[list[str]], list[str]] = default_columns,
-                 partition_columns: Callable[[list[str]], list[str]] = lambda x: [],
-                 table_props: Callable[[], dict[str,list[str]]] = lambda: [],
-                 variant = 'sql',
+                 expandables: dict = {},
+                 variant: SqlVariant = SqlVariant.SQL,
                  debug = False):
         machine = SqlStateMachine(debug=debug)
-        if variant == 'cql':
+        if variant == SqlVariant.CQL:
             machine = CqlStateMachine(debug=debug)
-        elif variant == 'athena':
+        elif variant == SqlVariant.ATHENA:
             machine = AthenaStateMachine(debug=debug)
         super().__init__(machine, dml, debug)
 
         self.tables = tables
-        self.columns = columns
-        self.partition_columns = partition_columns
-        self.table_props = table_props
+        if 'columns' not in expandables:
+            expandables['columns'] = default_columns
+        self.expandables = expandables
         self.variant = variant
         self.debug = debug
 
@@ -63,31 +65,54 @@ class SqlCompleter(AutomataCompleter[Token]):
     def _terms(self, state: State, word: str) -> list[str]:
         terms = []
 
-        if word == 'tables':
+        if word.startswith('`') and word.endswith('`'):
+            terms.append(word.strip('`'))
+        elif word == 'tables':
             terms.extend(self.tables())
-        elif word == '`tables`':
-            terms.append('tables')
         elif word == 'columns':
-            terms.extend(self.columns([]))
+            if 'last_name' in state.context and (n := state.context['last_name']):
+                if 'last_namespace' in state.context and (ns := state.context['last_namespace']):
+                    n = f'{ns}.{n}'
+                terms.extend(self._call_expandable(word, [n]))
+            else:
+                terms.extend(self._call_expandable(word, []))
         elif word == 'partition-columns':
-            terms.extend(self.partition_columns([]))
+            terms.extend(self._call_expandable(word, []))
         elif word == 'table-props':
-            terms.extend(self.table_props().keys())
+            terms.extend(self._call_expandable(word).keys())
         elif word == 'table-prop-values':
             if 'last_name' in state.context and state.context['last_name']:
-                terms.extend(self.table_props()[state.context['last_name']])
+                table_props = self._call_expandable('table-props')
+                terms.extend(table_props[state.context['last_name']])
         elif word == 'single':
             terms.append("'")
         elif word == 'comma':
             terms.append(",")
+        elif word in self.machine.expandable_names():
+            terms.extend(self._call_expandable(word))
         else:
             terms.append(word)
 
         return terms
 
-    def completions_for_nesting(self):
+    def _call_expandable(self, name: str, *args):
+        if name in self.expandables:
+            c = self.expandables[name]
+            if args:
+                return c(args)
+            else:
+                return c()
+
+        return []
+
+    def completions_for_nesting(self, dml: str = None):
+        if dml:
+            return {dml: SqlCompleter(self.tables, dml, expandables=self.expandables, variant=self.variant)}
+
         return {
-            word : SqlCompleter(self.tables, word, columns=self.columns, partition_columns=self.partition_columns,
-                                table_props=self.table_props, variant=self.variant)
+            word: SqlCompleter(self.tables, word, expandables=self.expandables, variant=self.variant)
             for word in self.machine.suggestions[''].strip(' ').split(',')
-        }
+        }
+
+    def __str__(self):
+        return f'{self.variant}, {self.first_term}'

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__ = [
@@ -73,10 +74,11 @@ SQL_SPEC = [
     '-                               < )                > select_from_sq',
     'select_from_lp_                 > select           > select',
     'select_from_x                   > ,                > select_from_x_comma_                    ^ (select,tables',
+    '-                               > ;                > ',
     '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_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,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',
@@ -94,8 +96,11 @@ SQL_SPEC = [
     '-                               > right outer join > select_join',
     '-                               > 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_',
+    '-                               > ;                > ',
     'select_from_x_as_x_comma_       > name|audit       > select_from_x                           ^ tables',
     'select_where_                   > name             > select_where_a                          ^ columns',
     'select_where_a                  > name             > select_where_a                          ^ columns,=,<,<=,>,>=,<>',
@@ -114,28 +119,41 @@ SQL_SPEC = [
     'select_where_a_in_lp_a_comma_   > name|single|num  > select_where_a_in_lp_a                  ^ single',
     '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_                > and|or           > select_where                            ^ and,or,order by,group by,limit',
+    'select_where_sc                 > ;                > ',
+    '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',
     'select_where_x_inner_           > join             > select_join',
     'select_join_                    > name|audit       > select_x_join_y                         ^ tables',
@@ -158,6 +176,7 @@ SQL_SPEC = [
     'select_x_join_y_on_a_           > comparison       > select_x_join_y_on_a_op                 ^ =',
     'select_x_join_y_on_a_op         > name             > select_x_join_y_on_a_op_b               ^ columns',
     'select_x_join_y_on_a_op_b       > _                > select_from_x_as_x_',
+    '-                               > ;                > ',
 
     # <insert_statement> ::= INSERT INTO <table_name> [ ( <column_list> ) ]
     #                        VALUES ( <value_list> )
@@ -188,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>
@@ -209,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',
@@ -227,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> ]
 
@@ -322,10 +344,10 @@ SQL_KEYWORDS = [
     'describe', 'preview'
 ]
 
-EXPANDABLE_NAMES = {'tables', 'columns', 'partition-columns', 'table-props', 'table-props-values'}
+EXPANDABLE_NAMES = {'keyspaces', 'tables', 'columns', 'partition-columns', 'table-props', 'table-props-values'}
 
 CQL_SPEC = SQL_SPEC + [
-    '                                > select           > select                                  ^ select,insert,update,delete,alter,describe,preview',
+    '                                > select           > select                                  ^ select,insert,update,delete,alter,describe,preview,consistency,export,import,drop,clean',
 
     # ALTER TABLE [ <keyspace_name> . ] <table_name>
     #     ( ALTER <column_name> TYPE <cql_type>
@@ -344,7 +366,7 @@ CQL_SPEC = SQL_SPEC + [
     'alter_table_with_p_op_v_        > --include-reaper > alter_table_with_p_op_v_$               ^ --include-reaper',
 
     '                                > describe         > describe',
-    'describe_                       > table            > desc_table                              ^ table,`tables`,keyspace,keyspaces,schema',
+    'describe_                       > table            > desc_table                              ^ table,`tables`,keyspace,`keyspaces`,schema',
     '-                               > tables           > desc_tables',
     '-                               > keyspace         > desc_keyspace',
     '-                               > keyspaces        > desc_keyspaces',
@@ -352,17 +374,88 @@ CQL_SPEC = SQL_SPEC + [
     'desc_table_                     > name             > desc_table_t                            ^ tables',
     'desc_table_t_                   > &                > desc_table_t_$                          ^ &',
     'desc_tables_                    > &                > desc_tables_$                           ^ &',
-    'desc_keyspace_                  > name             > desc_keyspace_k',
+    'desc_keyspace_                  > name             > desc_keyspace_k                         ^ keyspaces',
     'desc_keyspace_k_                > &                > desc_keyspace_k_$                       ^ &',
     'desc_schema_                    > &                > desc_schema_$                           ^ &',
+
+    '                                > export           > export',
+    'export_                         > name             > export_table                            ^ *,tables',
+    '-                               > *                > export_all',
+    'export_all_                     > in               > export_in                               ^ in,with',
+    '-                               > with             > export_with',
+    'export_table                    > (                > export_table_lp_                        ^ (,comma,with,tables',
+    '-                               > ,                > export_table_comma_',
+    '-                               > .                > export_table                            ^ tables',
+    'export_table_                   > (                > export_table_lp_                        ^ as,(,comma,with consistency,to',
+    '-                               > ,                > export_table_comma_',
+    '-                               > as               > export_as',
+    '-                               > with             > export_with',
+    '-                               > to               > export_table_to',
+    'export_table_comma_             > name             > export_table                            ^ tables',
+    'export_table_lp_                > name             > export_table_lp_a                       ^ columns',
+    'export_table_lp_a               > ,                > export_table_lp_a_comma_',
+    '-                               > )                > export_table_lp_a_comma_rp_',
+    'export_table_lp_a_comma_        > name             > export_table_lp_a                       ^ columns',
+    'export_table_lp_a_comma_rp_     > as               > export_as                               ^ as,with consistency',
+    '-                               > ,                > export                                  ^ with consistency',
+    '-                               > with             > export_with',
+    'export_as_                      > name             > export_as_f',
+    'export_as_f                     > ,                > export_table_comma_',
+    'export_as_f_                    > ,                > export_table_comma_                     ^ with consistency,to',
+    '-                               > with             > export_with',
+    '-                               > to               > export_table_to',
+    'export_in_                      > name             > export_in_k                             ^ keyspaces',
+    'export_in_k_                    > with             > export_with                             ^ with consistency',
+    'export_with_                    > consistency      > export_with_consistency                 ^ consistency',
+    'export_with_consistency_        > quorum|all|serial|one|each_quorum|local_quorum|any|local_one|two|three|local_serial > export_with_quorum ^ quorum,all,serial,one,each_quorum,local_quorum,any,local_one,two,three,local_serial',
+    'export_table_to_                > athena|sqlite|csv > export_table_to$                        ^ athena,sqlite,csv',
+
+    '                                > import           > import',
+    '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',
+    'consistency_                    > quorum|all|serial|one|each_quorum|local_quorum|any|local_one|two|three|local_serial > consistency_quorum ^ quorum,all,serial,one,each_quorum,local_quorum,any,local_one,two,three,local_serial',
+    'consistency_quorum              > ;                > ',
+
+    '                                > drop             > drop',
+    'drop_                           > all              > drop_all                                ^ all export databases,export database',
+    '-                               > export           > drop_export',
+    'drop_all_                       > export           > drop_all_export                         ^ export databases',
+    'drop_all_export_                > databases        > drop_all_dbs                            ^ databases',
+    'drop_export_                    > database         > drop_export_db                          ^ database',
+    'drop_export_db_                 > name             > drop_export_db$                         ^ export-dbs',
+
+    '                                > clean            > clean',
+    'clean_                          > up               > clean_up                                ^ up all export sessions,up export session',
+    'clean_up_                       > all              > clean_up_all                            ^ all export sessions,export sessions',
+    '-                               > export           > clean_up_export',
+    'clean_up_all_                   > export           > clean_up_all_export                     ^ export sessions',
+    'clean_up_all_export_            > sessions         > clean_up_all_sessions                   ^ sessions',
+    'clean_up_export_                > sessions         > clean_up_export_sessions                ^ sessions',
+    'clean_up_export_sessions_       > name             > clean_up_export_sessions$               ^ export-sessions',
 ]
 
 CQL_KEYWORDS = SQL_KEYWORDS + [
-    'schema', 'keyspace', 'keyspaces', 'tables'
+    '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', 'files'
 ]
 
+CQL_EXPANDABLE_NAMES = EXPANDABLE_NAMES | {
+    'export-dbs', 'export-sessions', 'export-sessions-incomplete'
+}
+
 ATHENA_SPEC = SQL_SPEC + [
-    '                                > select           > select                                  ^ select,insert,update,delete,alter,describe,preview',
+    '                                > select           > select                                  ^ select,insert,update,delete,alter,describe,preview,drop',
 
     'alter_table_t_                  > add              > alter_table_add                         ^ add partition,drop partition',
     'alter_table_add_                > partition        > alter_partition                         ^ partition',
@@ -378,13 +471,26 @@ ATHENA_SPEC = SQL_SPEC + [
     'describe_                       > name             > desc_t                                  ^ tables',
     'desc_t_                         > name             > desc_t_',
 
-    'repair'
+    'repair',
+
+    '                                > drop             > drop',
+    'drop_                           > all              > drop_all                                ^ all export databases,export database',
+    '-                               > export           > drop_export',
+    'drop_all_                       > export           > drop_all_export                         ^ export databases',
+    'drop_all_export_                > databases        > drop_all_dbs                            ^ databases',
+    'drop_export_                    > database         > drop_export_db                          ^ database',
+    'drop_export_db_                 > name             > drop_export_db$                         ^ export-dbs',
 ]
 
 ATHENA_KEYWORDS = SQL_KEYWORDS + [
-    'partition'
+    'partition',
+    'database', 'databases', 'session', 'sessions', 'clean', 'up', 'all', 'export'
 ]
 
+ATHENA_EXPANDABLE_NAMES = EXPANDABLE_NAMES | {
+    'export-dbs'
+}
+
 class SqlStateMachine(StateMachine[Token]):
     def __init__(self, indent=0, push_level = 0, debug = False):
         super().__init__(indent, push_level, debug)
@@ -449,11 +555,13 @@ class SqlStateMachine(StateMachine[Token]):
                     elif it == ')':
                         self.push_level -= 1
                         comeback_state = handle_pop()
+                    elif it == '.' and 'last_name' in state.context and (ln := state.context['last_name']):
+                        state.context['last_namespace'] = ln
 
                 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
@@ -464,8 +572,6 @@ class SqlStateMachine(StateMachine[Token]):
 
                     if last_name:
                         state.context['last_name'] = last_name
-                except:
-                    pass
 
         return state
 
@@ -507,6 +613,9 @@ class CqlStateMachine(SqlStateMachine):
     def keywords(self):
         return CQL_KEYWORDS
 
+    def expandable_names(self):
+        return CQL_EXPANDABLE_NAMES
+
 class AthenaStateMachine(SqlStateMachine):
     def __init__(self, indent=0, push_level = 0, debug = False):
         super().__init__(indent, push_level, debug)
@@ -516,3 +625,6 @@ class AthenaStateMachine(SqlStateMachine):
 
     def keywords(self):
         return ATHENA_KEYWORDS
+
+    def expandable_names(self):
+        return ATHENA_EXPANDABLE_NAMES

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,6 +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, log_exc
 from .idp_login import IdpLogin
 from adam.config import Config
 
@@ -33,7 +33,7 @@ class AdAuthenticator(Authenticator):
 
         session = requests.Session()
         r = session.get(idp_uri)
-        Config().debug(f'{r.status_code} {idp_uri}')
+        debug(f'{r.status_code} {idp_uri}')
 
         config = self.validate_and_return_config(r)
 
@@ -52,7 +52,7 @@ class AdAuthenticator(Authenticator):
         r = session.post(login_uri, data=body, headers={
             'Content-Type': 'application/x-www-form-urlencoded'
         })
-        Config().debug(f'{r.status_code} {login_uri}')
+        debug(f'{r.status_code} {login_uri}')
 
         config = self.validate_and_return_config(r)
 
@@ -69,7 +69,7 @@ class AdAuthenticator(Authenticator):
         r = session.post(kmsi_uri, data=body, headers={
             'Content-Type': 'application/x-www-form-urlencoded'
         })
-        Config().debug(f'{r.status_code} {kmsi_uri}')
+        debug(f'{r.status_code} {kmsi_uri}')
 
         if (config := self.extract_config_object(r.text)):
             if 'sErrorCode' in config and config['sErrorCode'] == '50058':
@@ -101,7 +101,7 @@ class AdAuthenticator(Authenticator):
 
     def validate_and_return_config(self, r: requests.Response):
         if r.status_code < 200 or r.status_code >= 300:
-            Config().debug(r.text)
+            debug(r.text)
 
             return None
 
@@ -138,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(
@@ -163,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:
-            Config().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 log2
+from adam.utils import debug, log2, log_exc
 
 class OktaException(Exception):
     pass
@@ -49,7 +49,7 @@ class OktaAuthenticator(Authenticator):
 
         session = requests.Session()
         response = session.post(authn_uri, headers=headers, data=json.dumps(payload))
-        Config().debug(f'{response.status_code} {authn_uri}')
+        debug(f'{response.status_code} {authn_uri}')
         auth_response = response.json()
 
         if 'sessionToken' not in auth_response:
@@ -59,7 +59,7 @@ class OktaAuthenticator(Authenticator):
 
         url = f'{idp_uri}&sessionToken={session_token}'
         r = session.get(url)
-        Config().debug(f'{r.status_code} {url}')
+        debug(f'{r.status_code} {url}')
 
         id_token = OktaAuthenticator().extract(r.text, r'.*name=\"id_token\" value=\"(.*?)\".*')
         if not id_token:
@@ -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,9 +1,8 @@
 import os
 from pathlib import Path
-import traceback
 from dotenv import load_dotenv
 
-from adam.config import Config
+from adam.utils import creating_dir, debug, log_exc
 from adam.utils_k8s.kube_context import KubeContext
 
 class CredCache:
@@ -15,7 +14,7 @@ class CredCache:
 
     def __init__(self):
         if not hasattr(self, 'env_f'):
-            self.dir = f'{Path.home()}/.kaqing'
+            self.dir = creating_dir(f'{Path.home()}/.kaqing')
             self.env_f = f'{self.dir}/.credentials'
             # immutable - cannot reload with different file content
             load_dotenv(dotenv_path=self.env_f)
@@ -34,10 +33,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:
-                    Config().debug(traceback.format_exc())
 
         updated = []
         updated.append(f'IDP_USERNAME={username}')
@@ -46,8 +43,6 @@ class CredCache:
             updated.append(f'IDP_PASSWORD={password}')
 
         if updated:
-            if not os.path.exists(self.env_f):
-                os.makedirs(self.dir, exist_ok=True)
             with open(self.env_f, 'w') as file:
                 file.write('\n'.join(updated))
 
@@ -56,4 +51,4 @@ class CredCache:
             if password:
                 self.overrides['IDP_PASSWORD'] = password
 
-            Config().debug(f'Cached username: {username}, password: {password}, try load: {self.get_username()}')
+            debug(f'Cached username: {username}, password: {password}, try load: {self.get_username()}')