Add argument to specify schema name

dmarkey · Feb 9, 2021 · 43fbf16 · 43fbf16
1 parent 225de1f
commit 43fbf16
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/pgdatadiff/main.py b/pgdatadiff/main.py
@@ -1,13 +1,14 @@
 """
 Usage:
-  pgdatadiff --firstdb=<firstconnectionstring> --seconddb=<secondconnectionstring> [--only-data|--only-sequences] [--count-only] [--chunk-size=<size>] [--exclude-tables=<table1,table2>]
+  pgdatadiff --firstdb=<firstconnectionstring> --seconddb=<secondconnectionstring> [--schema=<schema>] [--only-data|--only-sequences] [--count-only] [--chunk-size=<size>] [--exclude-tables=<table1,table2>]
   pgdatadiff --version
 
 Options:
   -h --help          Show this screen.
   --version          Show version.
   --firstdb=postgres://postgres:password@localhost/firstdb        The connection string of the first DB
   --seconddb=postgres://postgres:password@localhost/seconddb         The connection string of the second DB
+  --schema="public"         The schema of tables in comparison
   --only-data        Only compare data, exclude sequences
   --only-sequences   Only compare seqences, exclude data
   --exclude-tables=""   Exclude tables from data comparison         Must be a comma separated string [default: empty string]
@@ -35,7 +36,8 @@ def main():
     differ = DBDiff(first_db_connection_string, second_db_connection_string,
                     chunk_size=arguments['--chunk-size'],
                     count_only=arguments['--count-only'],
-                    exclude_tables=arguments['--exclude-tables'])
+                    exclude_tables=arguments['--exclude-tables'],
+                    schema=arguments['--schema'])
 
     if not arguments['--only-sequences']:
         if differ.diff_all_table_data():

diff --git a/pgdatadiff/pgdatadiff.py b/pgdatadiff/pgdatadiff.py
@@ -19,7 +19,7 @@ def make_session(connection_string):
 
 class DBDiff(object):
 
-    def __init__(self, firstdb, seconddb, chunk_size=10000, count_only=False, exclude_tables=""):
+    def __init__(self, firstdb, seconddb, schema, chunk_size=10000, count_only=False, exclude_tables=""):
         firstsession, firstengine = make_session(firstdb)
         secondsession, secondengine = make_session(seconddb)
         self.firstsession = firstsession
@@ -33,6 +33,7 @@ def __init__(self, firstdb, seconddb, chunk_size=10000, count_only=False, exclud
         self.chunk_size = int(chunk_size)
         self.count_only = count_only
         self.exclude_tables = exclude_tables.split(',')
+        self.schema = schema or 'public'
 
     def diff_table_data(self, tablename):
         try:
@@ -62,7 +63,7 @@ def diff_table_data(self, tablename):
         SELECT md5(array_agg(md5((t.*)::varchar))::varchar)
         FROM (
                 SELECT *
-                FROM {tablename}
+                FROM {self.schema}.{tablename}
                 ORDER BY {pk} limit :row_limit offset :row_offset
             ) AS t;
                         """
@@ -91,7 +92,7 @@ def get_all_sequences(self):
                 self.firstsession.execute(GET_SEQUENCES_SQL).fetchall()]
 
     def diff_sequence(self, seq_name):
-        GET_SEQUENCES_VALUE_SQL = f"SELECT last_value FROM {seq_name};"
+        GET_SEQUENCES_VALUE_SQL = f"SELECT last_value FROM {self.schema}.{seq_name};"
 
         try:
             firstvalue = \
@@ -141,7 +142,7 @@ def diff_all_table_data(self):
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", category=sa_exc.SAWarning)
             tables = sorted(
-                self.firstinspector.get_table_names(schema="public"))
+                self.firstinspector.get_table_names(schema=self.schema))
             for table in tables:
                 if table in self.exclude_tables:
                     print(bold(yellow(f"Ignoring table {table}")))