Skip to content

Commit 257d617

Browse files
author
Erik Cederberg
committed
Add primary key constraint name comparison
Change usage of SQLAlchemy's get_primary_keys to get_pk_constraint which returns a dict with 'constrained_columns' containing the primary key fields and optionally 'name' with the name of the pk constraint. Update get_primary_keys_info to compare the entire dicts and apply ignores on the constraint name rather than the key names if a 'name' field exists. If 'name' does not exist, the primary key columns are compared instead.
1 parent 278a406 commit 257d617

File tree

2 files changed

+72
-14
lines changed

2 files changed

+72
-14
lines changed

sqlalchemydiff/comparer.py

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -278,21 +278,34 @@ def _get_foreign_keys(inspector, table_name):
278278
def _get_primary_keys_info(
279279
left_inspector, right_inspector, table_name, ignores
280280
):
281-
left_pk_list = _get_primary_keys(left_inspector, table_name)
282-
right_pk_list = _get_primary_keys(right_inspector, table_name)
281+
left_pk_constraint = _get_primary_keys(left_inspector, table_name)
282+
right_pk_constraint = _get_primary_keys(right_inspector, table_name)
283283

284-
left_pk_list = _discard_ignores(left_pk_list, ignores)
285-
right_pk_list = _discard_ignores(right_pk_list, ignores)
284+
pk_constraint_has_name = 'name' in left_pk_constraint
286285

287-
# process into dict
288-
left_pk = dict((elem, elem) for elem in left_pk_list)
289-
right_pk = dict((elem, elem) for elem in right_pk_list)
286+
if pk_constraint_has_name:
287+
left_pk = ({left_pk_constraint['name']: left_pk_constraint}
288+
if _discard_ignores_by_name([left_pk_constraint], ignores)
289+
else {})
290+
right_pk = ({right_pk_constraint['name']: right_pk_constraint}
291+
if _discard_ignores_by_name([right_pk_constraint], ignores)
292+
else {})
293+
else:
294+
left_pk_list = left_pk_constraint['constrained_columns']
295+
right_pk_list = right_pk_constraint['constrained_columns']
296+
297+
left_pk_list = _discard_ignores(left_pk_list, ignores)
298+
right_pk_list = _discard_ignores(right_pk_list, ignores)
299+
300+
# process into dict
301+
left_pk = dict((elem, elem) for elem in left_pk_list)
302+
right_pk = dict((elem, elem) for elem in right_pk_list)
290303

291304
return _diff_dicts(left_pk, right_pk)
292305

293306

294307
def _get_primary_keys(inspector, table_name):
295-
return inspector.get_primary_keys(table_name)
308+
return inspector.get_pk_constraint(table_name)
296309

297310

298311
def _get_indexes_info(left_inspector, right_inspector, table_name, ignores):

test/unit/test_comparer.py

Lines changed: 51 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -431,8 +431,8 @@ def test__get_foreign_keys(self):
431431
def test__get_primary_keys_info(
432432
self, _diff_dicts_mock, _get_primary_keys_mock):
433433
_get_primary_keys_mock.side_effect = [
434-
['pk_left_1', 'pk_left_2'],
435-
['pk_right_1']
434+
{'constrained_columns': ['pk_left_1', 'pk_left_2']},
435+
{'constrained_columns': ['pk_right_1']}
436436
]
437437
left_inspector, right_inspector = Mock(), Mock()
438438

@@ -449,8 +449,8 @@ def test__get_primary_keys_info(
449449
def test__get_primary_keys_info_ignores(
450450
self, _diff_dicts_mock, _get_primary_keys_mock):
451451
_get_primary_keys_mock.side_effect = [
452-
['pk_left_1', 'pk_left_2'],
453-
['pk_right_1', 'pk_right_2']
452+
{'constrained_columns': ['pk_left_1', 'pk_left_2']},
453+
{'constrained_columns': ['pk_right_1', 'pk_right_2']},
454454
]
455455
left_inspector, right_inspector = Mock(), Mock()
456456
ignores = ['pk_left_1', 'pk_right_2']
@@ -465,13 +465,58 @@ def test__get_primary_keys_info_ignores(
465465

466466
assert _diff_dicts_mock.return_value == result
467467

468+
def test__get_primary_keys_info_with_pk_constraint_name(
469+
self, _diff_dicts_mock, _get_primary_keys_mock):
470+
_get_primary_keys_mock.side_effect = [
471+
{'name': 'left', 'constrained_columns': ['pk_left_1']},
472+
{'name': 'right', 'constrained_columns': ['pk_right_1']}
473+
]
474+
left_inspector, right_inspector = Mock(), Mock()
475+
476+
result = _get_primary_keys_info(
477+
left_inspector, right_inspector, 'table_A', [])
478+
479+
_diff_dicts_mock.assert_called_once_with(
480+
{
481+
'left': {'name': 'left',
482+
'constrained_columns': ['pk_left_1']}
483+
},
484+
{
485+
'right': {'name': 'right',
486+
'constrained_columns': ['pk_right_1']}
487+
}
488+
)
489+
assert _diff_dicts_mock.return_value == result
490+
491+
def test__get_primary_keys_info_ignores_with_pk_constraint_name(
492+
self, _diff_dicts_mock, _get_primary_keys_mock):
493+
_get_primary_keys_mock.side_effect = [
494+
{'name': 'left_1', 'constrained_columns': ['pk_left_1']},
495+
{'name': 'right_1', 'constrained_columns': ['pk_right_1']},
496+
]
497+
left_inspector, right_inspector = Mock(), Mock()
498+
ignores = ['left_1', 'left_2', 'right_2']
499+
500+
result = _get_primary_keys_info(
501+
left_inspector, right_inspector, 'table_A', ignores)
502+
503+
_diff_dicts_mock.assert_called_once_with(
504+
dict(),
505+
{
506+
'right_1': {'name': 'right_1',
507+
'constrained_columns': ['pk_right_1']},
508+
}
509+
)
510+
511+
assert _diff_dicts_mock.return_value == result
512+
468513
def test__get_primary_keys(self):
469514
inspector = Mock()
470515

471516
result = _get_primary_keys(inspector, 'table_A')
472517

473-
inspector.get_primary_keys.assert_called_once_with('table_A')
474-
assert inspector.get_primary_keys.return_value == result
518+
inspector.get_pk_constraint.assert_called_once_with('table_A')
519+
assert inspector.get_pk_constraint.return_value == result
475520

476521
def test__get_indexes_info(
477522
self, _diff_dicts_mock, _get_indexes_mock):

0 commit comments

Comments
 (0)