fix regression caused by check for existing replica schema

the4thdoctor · Jan 21, 2025 · 6b4a2dd · 6b4a2dd
1 parent bc8ff99
commit 6b4a2dd
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/pg_chameleon/lib/global_lib.py b/pg_chameleon/lib/global_lib.py
@@ -134,9 +134,11 @@ def __init__(self, args):
         self.mysql_source.sources = self.config["sources"]
         self.mysql_source.type_override = self.config["type_override"]
         self.mysql_source.notifier = self.notifier
-        self.mysql_source.net_read_timeout = int(self.mysql_source.sources[self.mysql_source.source].get('net_read_timeout', '600'))
-
-        #pgsql_source instance initialisation
+        try:
+            self.mysql_source.net_read_timeout = int(self.mysql_source.sources[self.mysql_source.source].get('net_read_timeout', '600'))
+        except:
+            self.mysql_source.net_read_timeout = 600
+            #pgsql_source instance initialisation
         self.pgsql_source = pgsql_source()
         self.pgsql_source.source = self.args.source
         self.pgsql_source.tables = self.args.tables
@@ -162,10 +164,16 @@ def __init__(self, args):
 
         if self.args.source != '*' and self.args.command != 'add_source':
             self.pg_engine.connect_db()
-            source_count = self.pg_engine.check_source()
-            self.pg_engine.disconnect_db()
+            try:
+                source_count = self.pg_engine.check_source()
+            except Exception as e:
+                if type(e).__name__ == "UndefinedTable" and self.count_replica_schema() ==  0:
+                    print("ERROR - Could not find the replica schema. Did you run the command create_replica_schema?")
+                self.pg_engine.disconnect_db()
+                sys.exit()
             if source_count == 0:
                 print("FATAL, The source %s is not registered. Please add it with the command add_source" % (self.args.source))
+                self.pg_engine.disconnect_db()
                 sys.exit()
 
 

diff --git a/scripts/chameleon.py b/scripts/chameleon.py
@@ -62,11 +62,3 @@
     except AttributeError:
         print("ERROR - Invalid command" )
         print(command_help)
-    except Exception as exception:
-        if type(exception).__name__ == "UndefinedTable":
-            if replica.count_replica_schema() == 0:
-                print("ERROR - Could not find the replica schema. Did you run the command create_replica_schema?")
-            else:
-                raise exception
-        else:
-            raise exception