Merge pull request #573 from dimitri-yatsenko/alter2

Altering table definition (Fix #110)

Author: guzman-raphael

datajoint/connection.py

@@ -137,9 +137,9 @@ def is_connected(self):
         """
         try:
             self.ping()
-            return True
         except:
             return False
+        return True
 
     def query(self, query, args=(), as_dict=False, suppress_warnings=True, reconnect=None):
         """

datajoint/declare.py

@@ -5,9 +5,10 @@
 import re
 import pyparsing as pp
 import logging
-
 from .errors import DataJointError
 
+from .utils import OrderedDict
+
 UUID_DATA_TYPE = 'binary(16)'
 MAX_TABLE_NAME_LENGTH = 64
 CONSTANT_LITERALS = {'CURRENT_TIMESTAMP'}  # SQL literals to be used without quotes (case insensitive)
@@ -218,20 +219,7 @@ def compile_foreign_key(line, context, attributes, primary_key, attr_sql, foreig
         index_sql.append('UNIQUE INDEX ({attrs})'.format(attrs='`,`'.join(ref.primary_key)))
 
 
-def declare(full_table_name, definition, context):
-    """
-    Parse declaration and create new SQL table accordingly.
-
-    :param full_table_name: full name of the table
-    :param definition: DataJoint table definition
-    :param context: dictionary of objects that might be referred to in the table.
-    """
-    table_name = full_table_name.strip('`').split('.')[1]
-    if len(table_name) > MAX_TABLE_NAME_LENGTH:
-        raise DataJointError(
-            'Table name `{name}` exceeds the max length of {max_length}'.format(
-                name=table_name,
-                max_length=MAX_TABLE_NAME_LENGTH))
+def prepare_declare(definition, context):
     # split definition into lines
     definition = re.split(r'\s*\n\s*', definition.strip())
     # check for optional table comment
@@ -266,7 +254,28 @@ def declare(full_table_name, definition, context):
             if name not in attributes:
                 attributes.append(name)
                 attribute_sql.append(sql)
-    # compile SQL
+
+    return table_comment, primary_key, attribute_sql, foreign_key_sql, index_sql, external_stores
+
+
+def declare(full_table_name, definition, context):
+    """
+    Parse declaration and generate the SQL CREATE TABLE code
+    :param full_table_name: full name of the table
+    :param definition: DataJoint table definition
+    :param context: dictionary of objects that might be referred to in the table
+    :return: SQL CREATE TABLE statement, list of external stores used
+    """
+    table_name = full_table_name.strip('`').split('.')[1]
+    if len(table_name) > MAX_TABLE_NAME_LENGTH:
+        raise DataJointError(
+            'Table name `{name}` exceeds the max length of {max_length}'.format(
+                name=table_name,
+                max_length=MAX_TABLE_NAME_LENGTH))
+
+    table_comment, primary_key, attribute_sql, foreign_key_sql, index_sql, external_stores = prepare_declare(
+        definition, context)
+
     if not primary_key:
         raise DataJointError('Table must have a primary key')
 
@@ -276,6 +285,94 @@ def declare(full_table_name, definition, context):
         '\n) ENGINE=InnoDB, COMMENT "%s"' % table_comment), external_stores
 
 
+def _make_attribute_alter(new, old, primary_key):
+    """
+    :param new: new attribute declarations
+    :param old: old attribute declarations
+    :param primary_key: primary key attributes
+    :return: list of SQL ALTER commands
+    """
+
+    # parse attribute names
+    name_regexp = re.compile(r"^`(?P<name>\w+)`")
+    original_regexp = re.compile(r'COMMENT "\{\s*(?P<name>\w+)\s*\}')
+    matched = ((name_regexp.match(d), original_regexp.search(d)) for d in new)
+    new_names = OrderedDict((d.group('name'), n and n.group('name')) for d, n in matched)
+    old_names = [name_regexp.search(d).group('name') for d in old]
+
+    # verify that original names are only used once
+    renamed = set()
+    for v in new_names.values():
+        if v:
+            if v in renamed:
+                raise DataJointError('Alter attempted to rename attribute {%s} twice.' % v)
+            renamed.add(v)
+
+    # verify that all renamed attributes existed in the old definition
+    try:
+        raise DataJointError(
+            "Attribute {} does not exist in the original definition".format(
+                next(attr for attr in renamed if attr not in old_names)))
+    except StopIteration:
+        pass
+
+    # dropping attributes
+    to_drop = [n for n in old_names if n not in renamed and n not in new_names]
+    sql = ['DROP `%s`' % n for n in to_drop]
+    old_names = [name for name in old_names if name not in to_drop]
+
+    # add or change attributes in order
+    prev = None
+    for new_def, (new_name, old_name) in zip(new, new_names.items()):
+        if new_name not in primary_key:
+            after = None  # if None, then must include the AFTER clause
+            if prev:
+                try:
+                    idx = old_names.index(old_name or new_name)
+                except ValueError:
+                    after = prev[0]
+                else:
+                    if idx >= 1 and old_names[idx - 1] != (prev[1] or prev[0]):
+                        after = prev[0]
+            if new_def not in old or after:
+                sql.append('{command} {new_def} {after}'.format(
+                    command=("ADD" if (old_name or new_name) not in old_names else
+                             "MODIFY" if not old_name else
+                             "CHANGE `%s`" % old_name),
+                    new_def=new_def,
+                    after="" if after is None else "AFTER `%s`" % after))
+        prev = new_name, old_name
+
+    return sql
+
+
+def alter(definition, old_definition, context):
+    """
+    :param definition: new table definition
+    :param old_definition: current table definition
+    :param context: the context in which to evaluate foreign key definitions
+    :return: string SQL ALTER command, list of new stores used for external storage
+    """
+    table_comment, primary_key, attribute_sql, foreign_key_sql, index_sql, external_stores = prepare_declare(
+        definition, context)
+    table_comment_, primary_key_, attribute_sql_, foreign_key_sql_, index_sql_, external_stores_ = prepare_declare(
+        old_definition, context)
+
+    # analyze differences between declarations
+    sql = list()
+    if primary_key != primary_key_:
+        raise NotImplementedError('table.alter cannot alter the primary key (yet).')
+    if foreign_key_sql != foreign_key_sql_:
+        raise NotImplementedError('table.alter cannot alter foreign keys (yet).')
+    if index_sql != index_sql_:
+        raise NotImplementedError('table.alter cannot alter indexes (yet)')
+    if attribute_sql != attribute_sql_:
+        sql.extend(_make_attribute_alter(attribute_sql, attribute_sql_, primary_key))
+    if table_comment != table_comment_:
+        sql.append('COMMENT="%s"' % table_comment)
+    return sql, [e for e in external_stores if e not in external_stores_]
+
+
 def compile_index(line, index_sql):
     match = index_parser.parseString(line)
     index_sql.append('{unique} index ({attrs})'.format(

datajoint/table.py

@@ -9,7 +9,7 @@
 import uuid
 from pymysql import OperationalError, InternalError, IntegrityError
 from .settings import config
-from .declare import declare
+from .declare import declare, alter
 from .expression import QueryExpression
 from . import attach, blob
 from .utils import user_choice
@@ -56,14 +56,15 @@ def heading(self):
 
     def declare(self, context=None):
         """
-        Use self.definition to declare the table in the schema.
+        Declare the table in the schema based on self.definition.
+        :param context: the context for foreign key resolution. If None, foreign keys are not allowed.
         """
         if self.connection.in_transaction:
             raise DataJointError('Cannot declare new tables inside a transaction, '
                                  'e.g. from inside a populate/make call')
+        sql, external_stores = declare(self.full_table_name, self.definition, context)
+        sql = sql.format(database=self.database)
         try:
-            sql, external_stores = declare(self.full_table_name, self.definition, context)
-            sql = sql.format(database=self.database)
             # declare all external tables before declaring main table
             for store in external_stores:
                 self.connection.schemas[self.database].external[store]
@@ -77,6 +78,41 @@ def declare(self, context=None):
         else:
             self._log('Declared ' + self.full_table_name)
 
+    def alter(self, prompt=True, context=None):
+        """
+        Alter the table definition from self.definition
+        """
+        if self.connection.in_transaction:
+            raise DataJointError('Cannot update table declaration inside a transaction, '
+                                 'e.g. from inside a populate/make call')
+        if context is None:
+            frame = inspect.currentframe().f_back
+            context = dict(frame.f_globals, **frame.f_locals)
+            del frame
+        old_definition = self.describe(context=context, printout=False)
+        sql, external_stores = alter(self.definition, old_definition, context)
+        if not sql:
+            if prompt:
+                print('Nothing to alter.')
+        else:
+            sql = "ALTER TABLE {tab}\n\t".format(tab=self.full_table_name) + ",\n\t".join(sql)
+            if not prompt or user_choice(sql + '\n\nExecute?') == 'yes':
+                try:
+                    # declare all external tables before declaring main table
+                    for store in external_stores:
+                        self.connection.schemas[self.database].external[store]
+                    self.connection.query(sql)
+                except pymysql.OperationalError as error:
+                    # skip if no create privilege
+                    if error.args[0] == server_error_codes['command denied']:
+                        logger.warning(error.args[1])
+                    else:
+                        raise
+                else:
+                    if prompt:
+                        print('Table altered')
+                    self._log('Altered ' + self.full_table_name)
+
     @property
     def from_clause(self):
         """
@@ -544,7 +580,7 @@ def _update(self, attrname, value=None):
 
             Example
 
-            >>> (v2p.Mice() & key).update('mouse_dob',   '2011-01-01')
+            >>> (v2p.Mice() & key).update('mouse_dob', '2011-01-01')
             >>> (v2p.Mice() & key).update( 'lens')   # set the value to NULL
 
         """

datajoint/user_tables.py

@@ -12,7 +12,7 @@
 
 # attributes that trigger instantiation of user classes
 supported_class_attrs = {
-    'key_source', 'describe', 'heading', 'populate', 'progress', 'primary_key', 'proj', 'aggr',
+    'key_source', 'describe', 'alter', 'heading', 'populate', 'progress', 'primary_key', 'proj', 'aggr',
     'fetch', 'fetch1','head', 'tail',
     'insert', 'insert1', 'drop', 'drop_quick', 'delete', 'delete_quick'}
 

tests/schema_advanced.py

@@ -50,7 +50,7 @@ def fill(self):
         def make_parent(pid, parent):
             return dict(person_id=pid,
                         parent=parent,
-                        parent_sex=(Person() & dict(person_id=parent)).fetch('sex')[0])
+                        parent_sex=(Person & {'person_id': parent}).fetch1('sex'))
 
         self.insert(make_parent(*r) for r in (
             (0, 2), (0, 3), (1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 7), (4, 7),

tests/test_alter.py

@@ -0,0 +1,42 @@
+from nose.tools import assert_equal, assert_not_equal
+
+from .schema import *
+
+
+@schema
+class Experiment(dj.Imported):
+
+    original_definition = """  # information about experiments
+    -> Subject
+    experiment_id  :smallint  # experiment number for this subject
+    ---
+    experiment_date  :date   # date when experiment was started
+    -> [nullable] User
+    data_path=""     :varchar(255)  # file path to recorded data
+    notes=""         :varchar(2048) # e.g. purpose of experiment
+    entry_time=CURRENT_TIMESTAMP :timestamp   # automatic timestamp
+    """
+
+    definition1 = """  # Experiment
+    -> Subject
+    experiment_id  :smallint  # experiment number for this subject
+    ---
+    data_path     : int  # some number
+    extra=null : longblob  # just testing
+    -> [nullable] User
+    subject_notes=null         :varchar(2048) # {notes} e.g. purpose of experiment
+    entry_time=CURRENT_TIMESTAMP :timestamp   # automatic timestamp
+    """
+
+
+def test_alter():
+
+    original = schema.connection.query("SHOW CREATE TABLE " + Experiment.full_table_name).fetchone()[1]
+    Experiment.definition = Experiment.definition1
+    Experiment.alter(prompt=False)
+    altered = schema.connection.query("SHOW CREATE TABLE " + Experiment.full_table_name).fetchone()[1]
+    Experiment.definition = Experiment.original_definition
+    Experiment().alter(prompt=False)
+    restored = schema.connection.query("SHOW CREATE TABLE " + Experiment.full_table_name).fetchone()[1]
+    assert_equal(original, restored)
+    assert_not_equal(original, altered)