1 # -*- coding: utf-8 -*-

     2 #

     3 # PgDiff - capture differences of database metadata

     4 #

     5 # Depends on PgBrowser

     6 #

     7 # Copyright (c) 2011  Radek Brich <radek.brich@devl.cz>

     8 #

     9 # Permission is hereby granted, free of charge, to any person obtaining a copy

    10 # of this software and associated documentation files (the "Software"), to deal

    11 # in the Software without restriction, including without limitation the rights

    12 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

    13 # copies of the Software, and to permit persons to whom the Software is

    14 # furnished to do so, subject to the following conditions:

    15 #

    16 # The above copyright notice and this permission notice shall be included in

    17 # all copies or substantial portions of the Software.

    18 #

    19 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

    20 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

    21 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

    22 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

    23 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

    24 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

    25 # THE SOFTWARE.

    26 

    27 

    28 from pycolib.ansicolor import *

    29 

    30 import re

    31 import difflib

    32 

    33 

    34 class PgDiffError(Exception):

    35     pass

    36 

    37 

    38 class DiffBase:

    39     COLORS = {

    40         '+' : BOLD | GREEN,

    41         '-' : BOLD | RED,

    42         '*' : BOLD | YELLOW,

    43     }

    44 

    45     COMMANDS = {

    46         '+' : 'CREATE',

    47         '-' : 'DROP',

    48         '*' : 'ALTER',

    49     }

    50 

    51     def __init__(self):

    52         self.changes = None

    53 

    54     def format(self):

    55         out = ['  ' * self.level]

    56 

    57         out.append(highlight(1, self.COLORS[self.change]))

    58         out.append(self.change)

    59 

    60         out += [' ', self.type, ' ', self.name, highlight(0)]

    61 

    62         if self.changes:

    63             out += [highlight(1, WHITE), ' (', self._formatchanges(), ')', highlight(0)]

    64 

    65         return ''.join(out)

    66 

    67     def _formatnotnull(self, notnull):

    68         if notnull:

    69             return 'NOT NULL'

    70         else:

    71             return None

    72 

    73     def _formatchanges(self):

    74         res = []

    75         for type, a, b in self.changes:

    76             if type == 'notnull':

    77                 type = ''

    78                 a = self._formatnotnull(a)

    79                 b = self._formatnotnull(b)

    80 

    81             if a and b:

    82                 s = ''.join(['Changed ', type, ' from ',

    83                     highlight(1,15), a, highlight(0), ' to ',

    84                     highlight(1,15), b, highlight(0), '.'])

    85             elif a and not b:

    86                 l = ['Removed ']

    87                 if type:

    88                     l += [type, ' ']

    89                 l += [highlight(1,15), a, highlight(0), '.']

    90                 s = ''.join(l)

    91             elif b and not a:

    92                 l = ['Added ']

    93                 if type:

    94                     l += [type, ' ']

    95                 l += [highlight(1,15), b, highlight(0), '.']

    96                 s = ''.join(l)

    97             res.append(s)

    98         return ' '.join(res)

    99 

   100     def format_patch(self):

   101         if self.change == '*' and self.type in ('schema', 'table'):

   102             return None

   103         return ['%s %s %s;' % (self.COMMANDS[self.change], self.type.upper(), self.name)]

   104 

   105 

   106 class DiffSchema(DiffBase):

   107     def __init__(self, change, schema):

   108         DiffBase.__init__(self)

   109         self.level = 0

   110         self.type = 'schema'

   111         self.change = change

   112         self.schema = schema

   113         self.name = schema

   114 

   115 

   116 class DiffTable(DiffBase):

   117     def __init__(self, change, schema, table):

   118         DiffBase.__init__(self)

   119         self.level = 1

   120         self.type = 'table'

   121         self.change = change

   122         self.schema = schema

   123         self.table = table

   124         self.name = table

   125 

   126 

   127 class DiffArgument(DiffBase):

   128     def __init__(self, change, schema, function, argument):

   129         DiffBase.__init__(self)

   130         self.level = 2

   131         self.type = 'argument'

   132         self.change = change

   133         self.schema = schema

   134         self.function = function

   135         self.argument = argument

   136         self.name = argument

   137 

   138 

   139 class DiffFunction(DiffBase):

   140     def __init__(self, change, schema, function, definition, show_body_diff=False):

   141         DiffBase.__init__(self)

   142         self.level = 1

   143         self.type = 'function'

   144         self.change = change

   145         self.schema = schema

   146         self.function = function

   147         #: New function definition

   148         self.definition = definition

   149         self.name = function

   150         self.show_body_diff = show_body_diff

   151 

   152     def _formatchanges(self):

   153         res = []

   154         for x in self.changes:

   155             type, a, b = x

   156             if type == 'source':

   157                 if self.show_body_diff:

   158                     lines = ['Source differs:\n']

   159                     for line in difflib.unified_diff(a, b, lineterm=''):

   160                         if line[:3] in ('---', '+++'):

   161                             continue

   162                         color = {' ': WHITE, '-': YELLOW, '+': GREEN, '@': WHITE|BOLD}[line[0]]

   163                         lines.append(highlight(1, color) + line + highlight(0) + '\n')

   164                     res.append(''.join(lines))

   165                 else:

   166                     res.append('Source differs.')

   167             else:

   168                 res.append(''.join(['Changed ', type, ' from ',

   169                     highlight(1,15), a, highlight(0), ' to ',

   170                     highlight(1,15), b, highlight(0), '.']))

   171         return ' '.join(res)

   172 

   173     def format_patch(self):

   174         return [self.definition]

   175 

   176 

   177 class DiffColumn(DiffBase):

   178     ALTER_COMMANDS = {

   179         '+' : 'ADD',

   180         '-' : 'DROP',

   181         '*' : 'ALTER',

   182     }

   183 

   184     def __init__(self, change, schema, table, column, columntype, columndefault, columnnotnull, changes=None):

   185         DiffBase.__init__(self)

   186         self.level = 2

   187         self.type = 'column'

   188         self.change = change

   189         self.schema = schema

   190         self.table = table

   191         self.column = column

   192         self.columntype = columntype

   193         self.columndefault = columndefault

   194         self.columnnotnull = columnnotnull

   195         self.name = column

   196         self.changes = changes

   197 

   198     def format_patch(self):

   199         alter_table = 'ALTER TABLE %s.%s %s COLUMN %s' % (

   200             self.schema,

   201             self.table,

   202             self.ALTER_COMMANDS[self.change],

   203             self.name,

   204         )

   205         out = []

   206         if self.change == '-':

   207             out.append('%s;' % alter_table);

   208         if self.change == '+':

   209             notnull = ''

   210             if self.columnnotnull:

   211                 notnull = ' NOT NULL'

   212             default = ''

   213             if self.columndefault:

   214                 default = ' DEFAULT %s' % self.columndefault

   215             out.append('%s %s%s%s;'

   216                 % (alter_table, self.columntype, notnull, default));

   217         if self.change == '*':

   218             for type, a, b in self.changes:

   219                 if type == 'type':

   220                     out.append('%s TYPE %s;' % (alter_table, b))

   221                 if type == 'notnull':

   222                     if a and not b:

   223                         out.append('%s DROP NOT NULL;' % alter_table)

   224                     if not a and b:

   225                         out.append('%s SET NOT NULL;' % alter_table)

   226                 if type == 'default':

   227                     if b:

   228                         out.append('%s SET DEFAULT %s;' % (alter_table, b))

   229                     else:

   230                         out.append('%s DROP DEFAULT;' % alter_table)

   231         return out

   232 

   233 

   234 class DiffConstraint(DiffBase):

   235     def __init__(self, change, schema, table, constraint, definition, changes=None):

   236         DiffBase.__init__(self)

   237         self.level = 2

   238         self.type = 'constraint'

   239         self.change = change

   240         self.schema = schema

   241         self.table = table

   242         self.constraint = constraint

   243         self.name = constraint

   244         self.definition = definition

   245         self.changes = changes

   246 

   247     def format_patch(self):

   248         q_alter = 'ALTER TABLE %s.%s' % (self.schema, self.table)

   249         q_drop = '%s DROP CONSTRAINT %s;' % (q_alter, self.constraint)

   250         q_add = '%s ADD CONSTRAINT %s %s;' % (q_alter, self.constraint, self.definition)

   251         if self.change == '*':

   252             out = [q_drop, q_add]

   253         if self.change == '+':

   254             out = [q_add]

   255         if self.change == '-':

   256             out = [q_drop]

   257         return out

   258 

   259 

   260 class DiffIndex(DiffBase):

   261     def __init__(self, change, schema, table, index, definition, changes=None):

   262         DiffBase.__init__(self)

   263         self.level = 2

   264         self.type = 'index'

   265         self.change = change

   266         self.schema = schema

   267         self.table = table

   268         self.index = index

   269         self.name = index

   270         self.definition = definition

   271         self.changes = changes

   272 

   273     def format_patch(self):

   274         q_drop = 'DROP INDEX %s;' % (self.index,)

   275         q_add = '%s;' % (self.definition,)

   276         if self.change == '*':

   277             out = [q_drop, q_add]

   278         if self.change == '+':

   279             out = [q_add]

   280         if self.change == '-':

   281             out = [q_drop]

   282         return out

   283 

   284 

   285 class DiffType(DiffBase):

   286     def __init__(self, change, schema, name):

   287         DiffBase.__init__(self)

   288         self.level = 1

   289         self.type = 'type'

   290         self.change = change

   291         self.schema = schema

   292         self.name = name

   293 

   294 

   295 class PgDiff:

   296     def __init__(self, srcbrowser=None, dstbrowser=None):

   297         self.allowcolor = False

   298         self.src = srcbrowser

   299         self.dst = dstbrowser

   300         self.include_schemas = set()  # if not empty, consider only these schemas for diff

   301         self.exclude_schemas = set()  # exclude these schemas from diff

   302         self.include_tables = set()

   303         self.exclude_tables = set()

   304         self.function_regex = re.compile(r"")

   305         self.function_body_diff = False

   306 

   307     def _test_schema(self, schema):

   308         if self.include_schemas and schema not in self.include_schemas:

   309             return False

   310         if schema in self.exclude_schemas:

   311             return False

   312         return True

   313 

   314     def _test_table(self, table):

   315         if self.include_tables and table not in self.include_tables:

   316             return False

   317         if table in self.exclude_tables:

   318             return False

   319         return True

   320 

   321     def _test_function(self, function):

   322         return bool(self.function_regex.match(function))

   323 

   324     def _diff_names(self, src, dst):

   325         for x in src:

   326             if x in dst:

   327                 yield ('*', x)

   328             else:

   329                 yield ('-', x)

   330         for x in dst:

   331             if x not in src:

   332                 yield ('+', x)

   333 

   334     def _compare_columns(self, a, b):

   335         diff = []

   336         if a.type != b.type:

   337             diff.append(('type', a.type, b.type))

   338         if a.notnull != b.notnull:

   339             diff.append(('notnull', a.notnull, b.notnull))

   340         if a.default != b.default:

   341             diff.append(('default', a.default, b.default))

   342         return diff

   343 

   344     def _compare_constraints(self, a, b):

   345         diff = []

   346         if a.type != b.type:

   347             diff.append(('type', a.type, b.type))

   348         if a.definition != b.definition:

   349             diff.append(('definition', a.definition, b.definition))

   350         return diff

   351 

   352     def _compare_indexes(self, a, b):

   353         diff = []

   354         if a.definition != b.definition:

   355             diff.append(('definition', a.definition, b.definition))

   356         return diff

   357 

   358     def _compare_functions(self, a, b):

   359         diff = []

   360         if a.result != b.result:

   361             diff.append(('result', a.result, b.result))

   362         # function source may differ in newlines (\n vs \r\n)

   363         # split lines before comparison, so that these differencies are ignored

   364         a_source = a.source.splitlines()

   365         b_source = b.source.splitlines()

   366         if a_source != b_source:

   367             diff.append(('source', a_source, b_source))

   368         return diff

   369 

   370     def _compare_arguments(self, a, b):

   371         diff = []

   372         if a.type != b.type:

   373             diff.append(('type', a.type, b.type))

   374         if a.mode != b.mode:

   375             diff.append(('mode', a.mode, b.mode))

   376         if a.default != b.default:

   377             diff.append(('default', a.default, b.default))

   378         return diff

   379 

   380     def _compare_types(self, a, b):

   381         diff = []

   382         if a.type != b.type:

   383             diff.append(('type', a.type, b.type))

   384         if a.elements != b.elements:

   385             diff.append(('elements', repr(a.elements), repr(b.elements)))

   386         return diff

   387 

   388     def _diff_columns(self, schema, table, src_columns, dst_columns):

   389         for nd in self._diff_names(src_columns, dst_columns):

   390             if nd[1] in dst_columns:

   391                 dst_type = dst_columns[nd[1]].type

   392                 dst_default = dst_columns[nd[1]].default

   393                 dst_notnull = dst_columns[nd[1]].notnull

   394             else:

   395                 dst_type = None

   396                 dst_default = None

   397                 dst_notnull = None

   398             cdo = DiffColumn(change=nd[0], schema=schema, table=table, column=nd[1],

   399                 columntype=dst_type, columndefault=dst_default, columnnotnull=dst_notnull)

   400             if nd[0] == '*':

   401                 a = src_columns[nd[1]]

   402                 b = dst_columns[nd[1]]

   403                 cdo.changes = self._compare_columns(a, b)

   404                 if cdo.changes:

   405                     yield cdo

   406             else:

   407                 yield cdo

   408 

   409     def _diff_constraints(self, schema, table, src_constraints, dst_constraints):

   410         for nd in self._diff_names(src_constraints, dst_constraints):

   411             if nd[1] in dst_constraints:

   412                 dst_definition = dst_constraints[nd[1]].definition

   413             else:

   414                 dst_definition = None

   415             cdo = DiffConstraint(change=nd[0], schema=schema, table=table, constraint=nd[1],

   416                 definition=dst_definition)

   417             if nd[0] == '*':

   418                 a = src_constraints[nd[1]]

   419                 b = dst_constraints[nd[1]]

   420                 cdo.changes = self._compare_constraints(a, b)

   421                 if cdo.changes:

   422                     yield cdo

   423             else:

   424                 yield cdo

   425 

   426     def _diff_indexes(self, schema, table, src_indexes, dst_indexes):

   427         for nd in self._diff_names(src_indexes, dst_indexes):

   428             if nd[1] in dst_indexes:

   429                 dst_definition = dst_indexes[nd[1]].definition

   430             else:

   431                 dst_definition = None

   432             ido = DiffIndex(change=nd[0], schema=schema, table=table, index=nd[1],

   433                 definition=dst_definition)

   434             if nd[0] == '*':

   435                 a = src_indexes[nd[1]]

   436                 b = dst_indexes[nd[1]]

   437                 ido.changes = self._compare_indexes(a, b)

   438                 if ido.changes:

   439                     yield ido

   440             else:

   441                 yield ido

   442 

   443     def _diff_tables(self, schema, src_tables, dst_tables):

   444         for nd in self._diff_names(src_tables, dst_tables):

   445             if not self._test_table(nd[1]):

   446                 continue

   447             tdo = DiffTable(change=nd[0], schema=schema, table=nd[1])

   448             if nd[0] == '*':

   449                 # columns

   450                 src_columns = src_tables[nd[1]].columns

   451                 dst_columns = dst_tables[nd[1]].columns

   452                 for cdo in self._diff_columns(schema, nd[1], src_columns, dst_columns):

   453                     if tdo:

   454                         yield tdo

   455                         tdo = None

   456                     yield cdo

   457                 # constraints

   458                 src_constraints = src_tables[nd[1]].constraints

   459                 dst_constraints = dst_tables[nd[1]].constraints

   460                 for cdo in self._diff_constraints(schema, nd[1], src_constraints, dst_constraints):

   461                     if tdo:

   462                         yield tdo

   463                         tdo = None

   464                     yield cdo

   465                 # indexes

   466                 src_indexes = src_tables[nd[1]].indexes

   467                 dst_indexes = dst_tables[nd[1]].indexes

   468                 for ido in self._diff_indexes(schema, nd[1], src_indexes, dst_indexes):

   469                     if tdo:

   470                         yield tdo

   471                         tdo = None

   472                     yield ido

   473             else:

   474                 yield tdo

   475 

   476     def _diff_arguments(self, schema, function, src_args, dst_args):

   477         for nd in self._diff_names(src_args, dst_args):

   478             ado = DiffArgument(change=nd[0], schema=schema, function=function, argument=nd[1])

   479             if nd[0] == '*':

   480                 a = src_args[nd[1]]

   481                 b = dst_args[nd[1]]

   482                 ado.changes = self._compare_arguments(a, b)

   483                 if ado.changes:

   484                     yield ado

   485             else:

   486                 yield ado

   487 

   488     def _diff_functions(self, schema, src_functions, dst_functions):

   489         for nd in self._diff_names(src_functions, dst_functions):

   490             if not self._test_function(nd[1]):

   491                 continue

   492             if nd[1] in dst_functions:

   493                 dst_definition = dst_functions[nd[1]].definition

   494             else:

   495                 dst_definition = None

   496             fdo = DiffFunction(change=nd[0], schema=schema, function=nd[1],

   497                                definition=dst_definition,

   498                                show_body_diff=self.function_body_diff)

   499             if nd[0] == '*':

   500                 # compare function body and result

   501                 a = src_functions[nd[1]]

   502                 b = dst_functions[nd[1]]

   503                 fdo.changes = self._compare_functions(a, b)

   504                 if fdo.changes:

   505                     yield fdo

   506                     fdo = None

   507                 # arguments

   508                 src_args = src_functions[nd[1]].arguments

   509                 dst_args = dst_functions[nd[1]].arguments

   510                 for ado in self._diff_arguments(schema, nd[1], src_args, dst_args):

   511                     if fdo:

   512                         yield fdo

   513                         fdo = None

   514                     yield ado

   515             else:

   516                 yield fdo

   517 

   518     def _diff_types(self, schema, src_types, dst_types):

   519         for nd in self._diff_names(src_types, dst_types):

   520             tdo = DiffType(change=nd[0], schema=schema, name=nd[1])

   521             if nd[0] == '*':

   522                 a = src_types[nd[1]]

   523                 b = dst_types[nd[1]]

   524                 tdo.changes = self._compare_types(a, b)

   525                 if tdo.changes:

   526                     yield tdo

   527             else:

   528                 yield tdo

   529 

   530     def iter_diff(self):

   531         '''Return diff between src and dst database schema.

   532 

   533         Yields one line at the time. Each line is in form of object

   534         iherited from DiffBase. This object contains all information

   535         about changes. See format() method.

   536 

   537         '''

   538         src_schemas = self.src.schemas

   539         dst_schemas = self.dst.schemas

   540         src = [x.name for x in src_schemas.values() if not x.system and self._test_schema(x.name)]

   541         dst = [x.name for x in dst_schemas.values() if not x.system and self._test_schema(x.name)]

   542         for nd in self._diff_names(src, dst):

   543             sdo = DiffSchema(change=nd[0], schema=nd[1])

   544             if nd[0] == '*':

   545                 # tables

   546                 src_tables = src_schemas[nd[1]].tables

   547                 dst_tables = dst_schemas[nd[1]].tables

   548                 for tdo in self._diff_tables(nd[1], src_tables, dst_tables):

   549                     if sdo:

   550                         yield sdo

   551                         sdo = None

   552                     yield tdo

   553                 # functions

   554                 src_functions = src_schemas[nd[1]].functions

   555                 dst_functions = dst_schemas[nd[1]].functions

   556                 for fdo in self._diff_functions(nd[1], src_functions, dst_functions):

   557                     if sdo:

   558                         yield sdo

   559                         sdo = None

   560                     yield fdo

   561                 # types

   562                 src_types = src_schemas[nd[1]].types

   563                 dst_types = dst_schemas[nd[1]].types

   564                 for tdo in self._diff_types(nd[1], src_types, dst_types):

   565                     if sdo:

   566                         yield sdo

   567                         sdo = None

   568                     yield tdo

   569             else:

   570                 yield sdo

   571 

   572     def print_diff(self):

   573         '''Print diff between src and dst database schema.

   574 

   575         The output is in human readable form.

   576 

   577         Set allowcolor=True of PgDiff instance to get colored output.

   578 

   579         '''

   580         for ln in self.iter_diff():

   581             print(ln.format())

   582 

   583     def print_patch(self):

   584         '''Print patch for updating from src schema to dst schema.

   585 

   586         Supports table drop, add, column drop, add and following

   587         changes of columns:

   588           - type

   589           - set/remove not null

   590           - default value

   591 

   592         This is experimental, not tested very much.

   593         Do not use without checking the commands.

   594         Even if it works as intended, it can cause table lock ups

   595         and/or loss of data. You have been warned.

   596 

   597         '''

   598         for ln in self.iter_diff():

   599             patch = ln.format_patch()

   600             if patch:

   601                 print('\n'.join(patch))

   602 

   603     def filter_schemas(self, include=[], exclude=[]):

   604         '''Modify list of schemas which are used for computing diff.

   605 

   606         include (list) -- if not empty, consider only these schemas for diff

   607         exclude (list) -- exclude these schemas from diff

   608 

   609         Order: include, exclude

   610         include=[] means include everything

   611 

   612         Raises:

   613             PgDiffError: when schema from include list is not found in src db

   614 

   615         '''

   616         for schema in include:

   617             self._check_schema_exist(schema)

   618         self.include_schemas.clear()

   619         self.include_schemas.update(include)

   620         self.exclude_schemas.clear()

   621         self.exclude_schemas.update(exclude)

   622 

   623     def filter_tables(self, include=[], exclude=[]):

   624         self.include_tables.clear()

   625         self.include_tables.update(include)

   626         self.exclude_tables.clear()

   627         self.exclude_tables.update(exclude)

   628 

   629     def filter_functions(self, regex=''):

   630         self.function_regex = re.compile(regex)

   631 

   632     def _check_schema_exist(self, schema):

   633         if not schema in self.src.schemas:

   634             raise PgDiffError('Schema "%s" not found in source database.' % schema)

   635