kaqing 2.0.115__py3-none-any.whl → 2.0.172__py3-none-any.whl

adam/repl_state.py

@@ -20,31 +20,8 @@ class BashSession:
     def pwd(self, state: 'ReplState'):
         command = f'cat /tmp/.qing-{self.session_id}'
 
-        if state.device == ReplState.A and state.app_app:
-            if state.app_pod:
-                rs = [AppPods.exec(state.app_pod, state.namespace, command, show_out=False)]
-            else:
-                pods = AppPods.pod_names(state.namespace, state.app_env, state.app_pod)
-                rs = AppClusters.exec(pods, state.namespace, command, show_out=False)
-        elif state.pod:
-            rs = [CassandraNodes.exec(state.pod, state.namespace, command, show_out=False)]
-        elif state.sts:
-            rs = CassandraClusters.exec(state.sts, state.namespace, command, action='bash', show_out=False)
-
-        dir = None
-        for r in rs:
-            if r.exit_code(): # if fails to read the session file, ignore
-                continue
-
-            dir0 = r.stdout.strip(' \r\n')
-            if dir:
-                if dir != dir0:
-                    log2('Inconsitent working dir found across multiple pods.')
-                    return None
-            else:
-                dir = dir0
-
-        return dir
+        with device(state) as pods:
+            return pods.exec(command, action='bash', show_out=False)
 
 class RequiredState(Enum):
     CLUSTER = 'cluster'
@@ -124,6 +101,8 @@ class ReplState:
             msg = f'{ReplState.L}:'
         elif self.device == ReplState.X:
             msg = f'{ReplState.X}:'
+            if self.export_session:
+                msg += self.export_session
         else:
             msg = f'{ReplState.C}:'
             if self.pod:
@@ -361,13 +340,13 @@ class ReplState:
                 return (False, error)
 
         elif required == RequiredState.EXPORT_DB:
-            if self.device != ReplState.X:
+            if self.device not in [ReplState.C, ReplState.X]:
                 return (False, None)
 
             if not self.export_session:
                 def error():
                     if self.in_repl:
-                        log2('use an export database first.')
+                        log2("Select an export database first with 'use' command.")
                     else:
                         log2('* export database is missing.')
                         log2()
@@ -424,6 +403,56 @@ class ReplState:
             self.app_env = o.app_env
             self.app_app = o.app_app
             self.app_pod = o.app_pod
+            # self.export_session = o.export_session
             self.namespace = o.namespace
 
-            self.original_state = None
+            self.original_state = None
+
+class DevicePodService:
+    def __init__(self, handler: 'DeviceExecHandler'):
+        self.handler = handler
+
+    def exec(self, command: str, action='bash', show_out = True):
+        state = self.handler.state
+
+        rs = None
+        if state.device == ReplState.A and state.app_app:
+            if state.app_pod:
+                rs = [AppPods.exec(state.app_pod, state.namespace, command, show_out=show_out)]
+            else:
+                pods = AppPods.pod_names(state.namespace, state.app_env, state.app_app)
+                rs = AppClusters.exec(pods, state.namespace, command, show_out=show_out)
+        elif state.pod:
+            rs = [CassandraNodes.exec(state.pod, state.namespace, command, show_out=show_out)]
+        elif state.sts:
+            rs = CassandraClusters.exec(state.sts, state.namespace, command, action=action, show_out=show_out)
+        # assume that pg-agent or ops pod is single pod
+
+        dir = None
+        if rs:
+            for r in rs:
+                if r.exit_code(): # if fails to read the session file, ignore
+                    continue
+
+                dir0 = r.stdout.strip(' \r\n')
+                if dir:
+                    if dir != dir0:
+                        log2('Inconsitent working dir found across multiple pods.')
+                        return None
+                else:
+                    dir = dir0
+
+        return dir
+
+class DeviceExecHandler:
+    def __init__(self, state: ReplState):
+        self.state = state
+
+    def __enter__(self):
+        return DevicePodService(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return False
+
+def device(state: ReplState):
+    return DeviceExecHandler(state)

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
@@ -15,38 +16,39 @@ __all__ = [
 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,40 +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':
             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.columns([n]))
+                terms.extend(self._call_expandable(word, [n]))
             else:
-                terms.extend(self.columns([]))
+                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 _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, columns=self.columns, partition_columns=self.partition_columns,
-                                table_props=self.table_props, variant=self.variant)}
+            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__ = [
@@ -76,8 +77,8 @@ 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_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',
@@ -96,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_',
     '-                               > ;                > ',
@@ -118,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',
@@ -201,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>
@@ -222,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',
@@ -240,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> ]
 
@@ -335,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,consistency',
+    '                                > 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>
@@ -357,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',
@@ -365,20 +374,24 @@ 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                            ^ with,tables',
+    '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',
+    '-                               > ,                > export_table_comma_',
     '-                               > .                > export_table                            ^ tables',
-    'export_table_                   > (                > export_table_lp_                        ^ as,(,comma,with consistency',
-    '-                               > ,                > export',
+    '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_',
@@ -387,25 +400,56 @@ CQL_SPEC = SQL_SPEC + [
     '-                               > ,                > export                                  ^ with consistency',
     '-                               > with             > export_with',
     'export_as_                      > name             > export_as_f',
-    'export_as_f                     > ,                > export',
-    'export_as_f_                    > ,                > export                                  ^ with consistency',
+    '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',
+    'import_session_                 > name             > import_session_s                        ^ export-sessions-incomplete',
+    'import_session_s_               > 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              > ;                > '
+    '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', 'export', 'consistency',
-    'quorum', 'all', 'serial', 'one', 'each_quorum', 'local_quorum', 'any', 'local_one', 'two', 'three', 'local_serial', 'to'
+    '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'
 ]
 
+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',
+    '                                > 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',
@@ -421,13 +465,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', '&select'
+    '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)
@@ -498,7 +555,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
@@ -509,8 +566,6 @@ class SqlStateMachine(StateMachine[Token]):
 
                     if last_name:
                         state.context['last_name'] = last_name
-                except:
-                    pass
 
         return state
 
@@ -552,6 +607,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)
@@ -561,3 +619,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