updates

TonicAI · Jul 24, 2019 · 98e11f5 · 98e11f5
1 parent e2cec47
commit 98e11f5
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/mysql_database_helper.py b/mysql_database_helper.py
@@ -74,6 +74,19 @@ def copy_to_temp_table(conn, query, target_table, pk_columns = None):
     finally:
         cur.close()
 
+def clean_temp_table_cells(fk_table, fk_columns, target_table, target_columns, conn):
+    fk_alias = 'tonic_subset_398dhjr23_fk'
+    target_alias = 'tonic_subset_398dhjr23_target'
+
+    fk_table = fully_qualified_table(source_db_temp_table(fk_table))
+    target_table = fully_qualified_table(source_db_temp_table(target_table))
+    assignment_list = ','.join(['{}.{} = NULL'.format(fk_alias, quoter(c)) for c in fk_columns])
+    column_matching = ' AND '.join(['{}.{} = {}.{}'.format(fk_alias, quoter(fc), target_alias, quoter(tc)) for fc, tc in zip(fk_columns, target_columns)])
+    target_columns_null = ' AND '.join(['{}.{} IS NULL'.format(target_alias, quoter(tc)) for tc in target_columns]
+        + ['{}.{} IS NOT NULL'.format(fk_alias, quoter(c)) for c in fk_columns])
+    q = 'UPDATE {} {} LEFT JOIN {} {} ON {} SET {} WHERE {}'.format(fk_table, fk_alias, target_table, target_alias, column_matching, assignment_list, target_columns_null)
+    run_query(q, conn)
+
 def source_db_temp_table(target_table):
     return temp_db + '.' + schema_name(target_table) + '_' + table_name(target_table)
 

diff --git a/psql_database_helper.py b/psql_database_helper.py
@@ -2,7 +2,7 @@
 import config_reader
 from pathlib import Path
 from psycopg2.extras import execute_values, register_default_json, register_default_jsonb
-from subset_utils import columns_joined, columns_tupled, schema_name, table_name, fully_qualified_table, redact_relationships
+from subset_utils import columns_joined, columns_tupled, schema_name, table_name, fully_qualified_table, redact_relationships, quoter
 
 register_default_json(loads=lambda x: str(x))
 register_default_jsonb(loads=lambda x: str(x))
@@ -75,6 +75,17 @@ def copy_to_temp_table(conn, query, target_table, pk_columns = None):
         cur.execute('INSERT INTO ' + temp_table + ' ' + query)
         conn.commit()
 
+def clean_temp_table_cells(fk_table, fk_columns, target_table, target_columns, conn):
+    fk_alias = 'tonic_subset_398dhjr23_fk'
+    target_alias = 'tonic_subset_398dhjr23_target'
+
+    fk_table = fully_qualified_table(source_db_temp_table(fk_table))
+    target_table = fully_qualified_table(source_db_temp_table(target_table))
+    assignment_list = ','.join(['{} = NULL'.format(quoter(c)) for c in fk_columns])
+    column_matching = ' AND '.join(['{}.{} = {}.{}'.format(fk_alias, quoter(fc), target_alias, quoter(tc)) for fc, tc in zip(fk_columns, target_columns)])
+    q = 'UPDATE {} {} SET {} WHERE NOT EXISTS (SELECT 1 FROM {} {} WHERE {})'.format(fk_table, fk_alias, assignment_list, target_table, target_alias, column_matching)
+    run_query(q, conn)
+
 def get_redacted_table_references(table_name, tables, conn):
     relationships = get_unredacted_fk_relationships(tables, conn)
     redacted = redact_relationships(relationships)