Closes #939: Persisted, existing table metadata should be updated.

redash/tasks/queries.py

@@ -303,14 +303,21 @@ def insert_or_update_table_metadata(data_source, existing_tables_set, table_data
         TableMetadata.name.in_(existing_tables_set),
         TableMetadata.data_source_id == data_source.id,
     )
-    persisted_tables.update({"exists": True}, synchronize_session='fetch')
+    persisted_table_data = []
+    for persisted_table in persisted_tables:
+        # Add IDs to persisted table data so it can be used for updates.
+        table_data[persisted_table.name]['id'] = persisted_table.id
+        persisted_table_data.append(table_data[persisted_table.name])
 
-    # Find the tables that need to be created by subtracting the sets:
-    persisted_table_set = set([
-        persisted_table.name for persisted_table in persisted_tables.all()
-    ])
+    models.db.session.bulk_update_mappings(
+        TableMetadata,
+        persisted_table_data
+    )
 
+    # Find the tables that need to be created by subtracting the sets:
+    persisted_table_set = set([col_data['name'] for col_data in persisted_table_data])
     tables_to_create = existing_tables_set.difference(persisted_table_set)
+
     table_metadata = [table_data[table_name] for table_name in tables_to_create]
 
     models.db.session.bulk_insert_mappings(
@@ -335,6 +342,8 @@ def insert_or_update_column_metadata(table, existing_columns_set, column_data):
         ColumnMetadata,
         persisted_column_data
     )
+
+    # Find the columns that need to be created by subtracting the sets:
     persisted_column_set = set([col_data['name'] for col_data in persisted_column_data])
     columns_to_create = existing_columns_set.difference(persisted_column_set)
 
@@ -372,16 +381,17 @@ def refresh_schema(data_source_id):
             existing_tables_set.add(table_name)
 
             table_data[table_name] = {
-                "org_id": ds.org_id,
-                "name": table_name,
-                "data_source_id": ds.id,
-                "column_metadata": "metadata" in table
+                'org_id': ds.org_id,
+                'name': table_name,
+                'data_source_id': ds.id,
+                'column_metadata': "metadata" in table,
+                'exists': True
             }
             new_column_names[table_name] = table['columns']
             new_column_metadata[table_name] = table.get('metadata', None)
 
         insert_or_update_table_metadata(ds, existing_tables_set, table_data)
-        models.db.session.flush()
+        models.db.session.commit()
 
         all_existing_persisted_tables = TableMetadata.query.filter(
             TableMetadata.exists.is_(True),

tests/tasks/test_refresh_schemas.py

@@ -115,6 +115,35 @@ def test_refresh_schema_deleted_table_marked(self):
         table_metadata = TableMetadata.query.all()
         self.assertTrue(table_metadata[0].to_dict()['exists'])
 
+    def test_refresh_schema_table_with_new_metadata_updated(self):
+        refresh_schema(self.factory.data_source.id)
+        table_metadata = TableMetadata.query.all()
+        column_metadata = ColumnMetadata.query.all()
+
+        self.assertEqual(len(table_metadata), 1)
+        self.assertEqual(len(column_metadata), 1)
+        self.assertTrue(table_metadata[0].to_dict()['column_metadata'])
+
+        # Table has no metdata field, `column_metadata` should be False.
+        self.patched_get_schema.return_value = [{
+            'name': 'table',
+            'columns': [self.COLUMN_NAME],
+        }]
+
+        refresh_schema(self.factory.data_source.id)
+        table_metadata = TableMetadata.query.all()
+        column_metadata = ColumnMetadata.query.all()
+
+        self.assertEqual(len(table_metadata), 1)
+        self.assertEqual(len(column_metadata), 1)
+        self.assertFalse(table_metadata[0].to_dict()['column_metadata'])
+
+        # Table metadata field is back, `column_metadata` should be True again.
+        self.patched_get_schema.return_value = self.default_schema_return_value
+        refresh_schema(self.factory.data_source.id)
+        table_metadata = TableMetadata.query.all()
+        self.assertTrue(table_metadata[0].to_dict()['column_metadata'])
+
     def test_refresh_schema_delete_column(self):
         NEW_COLUMN_NAME = 'new_column'
         refresh_schema(self.factory.data_source.id)