ERD model/figures refactoring.

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/editor/ERDEditorEmbedded.java

@@ -179,9 +179,9 @@ public class ERDEditorEmbedded extends ERDEditorPart implements IDatabaseEditor,
         }
         EntityDiagram diagram;
         if (!dbObject.isPersisted()) {
-            diagram = new EntityDiagram(dbObject, "New Object");
+            diagram = new EntityDiagram(createDecorator(), dbObject, "New Object");
         } else {
-            diagram = new EntityDiagram(dbObject, dbObject.getName());
+            diagram = new EntityDiagram(createDecorator(), dbObject, dbObject.getName());
 
             diagram.fillTables(
                 monitor,

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/editor/ERDEditorPart.java

@@ -73,6 +73,8 @@ import org.jkiss.dbeaver.ext.erd.editor.tools.ChangeZOrderAction;
 import org.jkiss.dbeaver.ext.erd.editor.tools.SetPartColorAction;
 import org.jkiss.dbeaver.ext.erd.export.ERDExportFormatHandler;
 import org.jkiss.dbeaver.ext.erd.export.ERDExportFormatRegistry;
+import org.jkiss.dbeaver.ext.erd.model.ERDDecorator;
+import org.jkiss.dbeaver.ext.erd.model.ERDDecoratorDefault;
 import org.jkiss.dbeaver.ext.erd.model.ERDNote;
 import org.jkiss.dbeaver.ext.erd.model.EntityDiagram;
 import org.jkiss.dbeaver.ext.erd.part.DiagramPart;
@@ -141,6 +143,10 @@ public abstract class ERDEditorPart extends GraphicalEditorWithFlyoutPalette
     {
     }
 
+    protected ERDDecorator createDecorator() {
+        return new ERDDecoratorDefault();
+    }
+
     @Override
     protected ERDGraphicalViewer getGraphicalViewer() {
         return (ERDGraphicalViewer) super.getGraphicalViewer();
@@ -353,7 +359,7 @@ public abstract class ERDEditorPart extends GraphicalEditorWithFlyoutPalette
         initializeGraphicalViewer();
 
         // Set initial (empty) contents
-        viewer.setContents(new EntityDiagram(null, "empty"));
+        viewer.setContents(new EntityDiagram(createDecorator(), null, "empty"));
 
         // Set context menu
         ContextMenuProvider provider = new ERDEditorContextMenuProvider(this);

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/editor/ERDEditorStandalone.java

@@ -160,7 +160,7 @@ public class ERDEditorStandalone extends ERDEditorPart implements DBPContextProv
         final IFile file = getEditorFile();
 
         final DiagramPart diagramPart = getDiagramPart();
-        EntityDiagram entityDiagram = new EntityDiagram(null, file.getName());
+        EntityDiagram entityDiagram = new EntityDiagram(createDecorator(), null, file.getName());
         entityDiagram.clear();
         entityDiagram.setLayoutManualAllowed(true);
         entityDiagram.setLayoutManualDesired(true);

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/figures/AttributeItemFigure.java

@@ -19,30 +19,63 @@
  */
 package org.jkiss.dbeaver.ext.erd.figures;
 
+import org.eclipse.draw2d.*;
+import org.jkiss.code.NotNull;
+import org.jkiss.code.Nullable;
+import org.jkiss.dbeaver.ext.erd.model.ERDDecorator;
 import org.jkiss.dbeaver.ext.erd.model.ERDEntityAttribute;
+import org.jkiss.dbeaver.ext.erd.part.AttributePart;
 import org.jkiss.dbeaver.model.DBPImage;
 import org.jkiss.dbeaver.ui.DBeaverIcons;
 
+import java.util.List;
+
 /**
  * Figure used to hold the column labels
  * @author Serge Rider
  */
-public class AttributeItemFigure extends EditableLabel
+public class AttributeItemFigure extends Figure
 {
-    private final ERDEntityAttribute attribute;
+    private final AttributePart part;
 
-	public AttributeItemFigure(ERDEntityAttribute attribute)
+	public AttributeItemFigure(AttributePart part)
 	{
-        super(attribute.getLabelText());
-        this.attribute = attribute;
-        DBPImage labelImage = attribute.getLabelImage();
+        super();
+        this.part = part;
+
+        ToolbarLayout layout = new ToolbarLayout(true);
+
+        setLayoutManager(layout);
+
+        boolean showCheckboxes = part.getDiagramPart().getDiagram().getDecorator().showCheckboxes();
+        if (showCheckboxes) {
+            CustomCheckBoxFigure attrCheckbox = new CustomCheckBoxFigure();
+            add(attrCheckbox);
+        }
+        EditableLabel attrNameLabel = new EditableLabel(part.getAttribute().getName());
+        DBPImage labelImage = part.getAttribute().getLabelImage();
         if (labelImage != null) {
-            setIcon(DBeaverIcons.getImage(labelImage));
+            attrNameLabel.setIcon(DBeaverIcons.getImage(labelImage));
         }
-
+        add(attrNameLabel);
 	}
 
     public ERDEntityAttribute getAttribute() {
-        return attribute;
+        return part.getAttribute();
     }
-}
\ No newline at end of file
+
+    @Nullable
+    public CustomCheckBoxFigure getCheckBox() {
+	    if (getChildren().size() < 2) {
+	        return null;
+        }
+        return (CustomCheckBoxFigure) getChildren().get(0);
+    }
+
+    @NotNull
+    public EditableLabel getLabel() {
+        List children = getChildren();
+        return (EditableLabel) children.get(children.size() == 1 ? 0 : 1);
+    }
+
+}

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/figures/AttributeListFigure.java

@@ -38,17 +38,18 @@ public class AttributeListFigure extends Figure
 
 	public AttributeListFigure(ERDEntity entity, boolean key)
 	{
-		FlowLayout layout = new FlowLayout();
+/*
+		GridLayout layout = new GridLayout(1, false);
+		layout.verticalSpacing = 0;
+		layout.marginHeight = 0;
+		layout.marginWidth = 0;
+*/
+		FlowLayout layout = new FlowLayout(false);
 		layout.setMinorAlignment(FlowLayout.ALIGN_TOPLEFT);
-		layout.setStretchMinorAxis(false);
-		layout.setHorizontal(false);
+		layout.setStretchMinorAxis(true);
+		layout.setMinorSpacing(2);
 		setLayoutManager(layout);
 		setBorder(new ColumnFigureBorder());
-        if (entity.isPrimary()) {
-            //setBackgroundColor(EntityFigure.primaryTableColor);
-        } else {
-		    //setBackgroundColor(ColorConstants.tooltipBackground);
-        }
 		ColorRegistry colorRegistry = UIUtils.getColorRegistry();
         setBackgroundColor(colorRegistry.get(ERDConstants.COLOR_ERD_ATTR_BACKGROUND));
 		setForegroundColor(colorRegistry.get(ERDConstants.COLOR_ERD_ATTR_FOREGROUND));

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/figures/CustomCheckBoxFigure.java

+//
+// Source code recreated from a .class file by IntelliJ IDEA
+// (powered by Fernflower decompiler)
+//
+
+package org.jkiss.dbeaver.ext.erd.figures;
+
+import org.eclipse.draw2d.ChangeEvent;
+import org.eclipse.draw2d.ChangeListener;
+import org.eclipse.draw2d.Label;
+import org.eclipse.draw2d.Toggle;
+import org.eclipse.swt.graphics.Image;
+import org.jkiss.dbeaver.ui.DBeaverIcons;
+import org.jkiss.dbeaver.ui.UIIcon;
+
+public final class CustomCheckBoxFigure extends Toggle {
+    private Label label;
+    private static final Image UNCHECKED = DBeaverIcons.getImage(UIIcon.CHECK_OFF);
+    private static final Image CHECKED = DBeaverIcons.getImage(UIIcon.CHECK_ON);
+
+    public CustomCheckBoxFigure() {
+        this("");
+    }
+
+    public CustomCheckBoxFigure(String text) {
+        this.label = null;
+        this.setContents(this.label = new Label(text, UNCHECKED));
+    }
+
+    protected void handleSelectionChanged() {
+        if (this.isSelected()) {
+            this.label.setIcon(CHECKED);
+        } else {
+            this.label.setIcon(UNCHECKED);
+        }
+
+    }
+
+    protected void init() {
+        super.init();
+        this.addChangeListener(new ChangeListener() {
+            public void handleStateChanged(ChangeEvent changeEvent) {
+                if (changeEvent.getPropertyName().equals("selected")) {
+                    CustomCheckBoxFigure.this.handleSelectionChanged();
+                }
+
+            }
+        });
+    }
+}

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/figures/EditableLabel.java

@@ -21,6 +21,7 @@ package org.jkiss.dbeaver.ext.erd.figures;
 
 import org.eclipse.draw2d.Graphics;
 import org.eclipse.draw2d.Label;
+import org.eclipse.draw2d.geometry.Dimension;
 import org.eclipse.draw2d.geometry.Insets;
 import org.eclipse.draw2d.geometry.Rectangle;
 import org.jkiss.dbeaver.ext.erd.ERDConstants;
@@ -79,6 +80,11 @@ public class EditableLabel extends Label
 		super.paintFigure(graphics);
 	}
 
+	@Override
+	public Dimension getPreferredSize(int wHint, int hHint) {
+		return super.getPreferredSize(wHint, hHint);
+	}
+
 	/**
 	 * Sets the selection state of this SimpleActivityLabel
 	 * 

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/model/ERDDecorator.java

+/*
+ * DBeaver - Universal Database Manager
+ * Copyright (C) 2010-2017 Serge Rider (serge@jkiss.org)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jkiss.dbeaver.ext.erd.model;
+
+/**
+ * ERD object adapter
+ */
+public interface ERDDecorator {
+
+    boolean showCheckboxes();
+
+}

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/model/ERDDecoratorDefault.java

+/*
+ * DBeaver - Universal Database Manager
+ * Copyright (C) 2010-2017 Serge Rider (serge@jkiss.org)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jkiss.dbeaver.ext.erd.model;
+
+/**
+ * ERD object adapter
+ */
+public class ERDDecoratorDefault implements ERDDecorator {
+
+    @Override
+    public boolean showCheckboxes() {
+        return false;
+    }
+}

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/model/EntityDiagram.java

@@ -48,6 +48,7 @@ public class EntityDiagram extends ERDObject<DBSObject>
         public int zOrder = 0;
     }
 
+    private ERDDecorator decorator;
 	private String name;
 	private List<ERDEntity> entities = new ArrayList<>();
 	private boolean layoutManualDesired = true;
@@ -61,9 +62,10 @@ public class EntityDiagram extends ERDObject<DBSObject>
 
     private List<String> errorMessages = new ArrayList<>();
 
-    public EntityDiagram(DBSObject container, String name)
+    public EntityDiagram(ERDDecorator decorator, DBSObject container, String name)
 	{
 		super(container);
+		this.decorator = decorator;
 		if (name == null)
 			throw new IllegalArgumentException("Name cannot be null");
 		this.name = name;
@@ -72,6 +74,10 @@ public class EntityDiagram extends ERDObject<DBSObject>
         this.attributeStyles = ERDViewStyle.getDefaultStyles(store);
 	}
 
+    public ERDDecorator getDecorator() {
+        return decorator;
+    }
+
     public boolean hasAttributeStyle(ERDViewStyle style) {
         return ArrayUtils.contains(attributeStyles, style);
     }
@@ -235,7 +241,7 @@ public class EntityDiagram extends ERDObject<DBSObject>
 
     public EntityDiagram copy()
     {
-        EntityDiagram copy = new EntityDiagram(getObject(), getName());
+        EntityDiagram copy = new EntityDiagram(decorator, getObject(), getName());
         copy.entities.addAll(this.entities);
         copy.tableMap.putAll(this.tableMap);
         copy.layoutManualDesired = this.layoutManualDesired;

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/navigator/DiagramCreateWizard.java

@@ -30,6 +30,7 @@ import org.eclipse.ui.IWorkbench;
 import org.jkiss.dbeaver.core.DBeaverCore;
 import org.jkiss.dbeaver.ext.erd.ERDMessages;
 import org.jkiss.dbeaver.ext.erd.model.DiagramObjectCollector;
+import org.jkiss.dbeaver.ext.erd.model.ERDDecoratorDefault;
 import org.jkiss.dbeaver.ext.erd.model.EntityDiagram;
 import org.jkiss.dbeaver.model.navigator.DBNDatabaseNode;
 import org.jkiss.dbeaver.model.navigator.DBNNode;
@@ -49,7 +50,7 @@ import java.util.List;
 public class DiagramCreateWizard extends Wizard implements INewWizard {
 
     private IFolder folder;
-    private EntityDiagram diagram = new EntityDiagram(null, "");
+    private EntityDiagram diagram = new EntityDiagram(new ERDDecoratorDefault(), null, "");
     private DiagramCreateWizardPage pageContent;
 	private String errorMessage;
     private IStructuredSelection entitySelection;

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/navigator/ERDResourceHandler.java

@@ -28,6 +28,7 @@ import org.jkiss.dbeaver.core.DBeaverCore;
 import org.jkiss.dbeaver.ext.erd.editor.ERDEditorInput;
 import org.jkiss.dbeaver.ext.erd.editor.ERDEditorStandalone;
 import org.jkiss.dbeaver.ext.erd.model.DiagramLoader;
+import org.jkiss.dbeaver.ext.erd.model.ERDDecoratorDefault;
 import org.jkiss.dbeaver.ext.erd.model.EntityDiagram;
 import org.jkiss.dbeaver.model.DBPDataSourceContainer;
 import org.jkiss.dbeaver.model.navigator.DBNNode;
@@ -146,7 +147,9 @@ public class ERDResourceHandler extends AbstractResourceHandler {
                 public void run(DBRProgressMonitor monitor) throws InvocationTargetException, InterruptedException
                 {
                     try {
-                        EntityDiagram newDiagram = copyFrom == null ? new EntityDiagram(null, "<Diagram>") : copyFrom.copy();
+                        EntityDiagram newDiagram = copyFrom == null ?
+                                new EntityDiagram(new ERDDecoratorDefault(), null, "<Diagram>") :
+                                copyFrom.copy();
                         newDiagram.setName(title);
                         newDiagram.setLayoutManualAllowed(true);
                         newDiagram.setLayoutManualDesired(true);

plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/part/AttributePart.java

@@ -44,12 +44,12 @@ import java.beans.PropertyChangeEvent;
 
 /**
  * Represents an editable Column object in the model
+ *
  * @author Serge Rider
  */
-public class AttributePart extends PropertyAwarePart
-{
+public class AttributePart extends PropertyAwarePart {
 
-	protected DirectEditManager manager;
+    protected DirectEditManager manager;
 
     @Override
     public boolean isSelectable() {
@@ -57,23 +57,24 @@ public class AttributePart extends PropertyAwarePart
     }
 
     /**
-	 * @return the ColumnLabel representing the Column
-	 */
-	@Override
-    protected AttributeItemFigure createFigure()
-	{
-		ERDEntityAttribute column = (ERDEntityAttribute) getModel();
-        AttributeItemFigure attributeFigure = new AttributeItemFigure(column);
-
-        DiagramPart diagramPart = (DiagramPart) getParent().getParent();
-		boolean showNullability = diagramPart.getDiagram().hasAttributeStyle(ERDViewStyle.NULLABILITY);
+     * @return the ColumnLabel representing the Column
+     */
+    @Override
+    protected AttributeItemFigure createFigure() {
+        ERDEntityAttribute column = (ERDEntityAttribute) getModel();
+        AttributeItemFigure attributeFigure = new AttributeItemFigure(this);
+
+        DiagramPart diagramPart = getDiagramPart();
+        boolean showNullability = diagramPart.getDiagram().hasAttributeStyle(ERDViewStyle.NULLABILITY);
         Font columnFont = diagramPart.getNormalFont();
         Color columnColor = diagramPart.getContentPane().getForegroundColor();
         if (column.isInPrimaryKey()) {
             columnFont = diagramPart.getBoldFont();
+/*
             if (showNullability && !column.getObject().isRequired()) {
                 columnFont = diagramPart.getBoldItalicFont();
             }
+*/
 /*
             if (!column.isInForeignKey()) {
                 columnFont = diagramPart.getBoldFont();
@@ -82,9 +83,11 @@ public class AttributePart extends PropertyAwarePart
             }
 */
         } else {
+/*
             if (showNullability && !column.getObject().isRequired()) {
                 columnFont = diagramPart.getItalicFont();
             }
+*/
         }
         if (column.isInForeignKey()) {
             //columnColor = Display.getDefault().getSystemColor(SWT.COLOR_DARK_BLUE);
@@ -94,27 +97,28 @@ public class AttributePart extends PropertyAwarePart
         return attributeFigure;
     }
 
-	@Override
-	public AttributeItemFigure getFigure() {
-		return (AttributeItemFigure)super.getFigure();
-	}
-
-	/**
-	 * Create EditPolicies for the column label
-	 */
-	@Override
-    protected void createEditPolicies()
-	{
-		//installEditPolicy(EditPolicy.COMPONENT_ROLE, new AttributeEditPolicy());
-		//installEditPolicy(EditPolicy.DIRECT_EDIT_ROLE, new ColumnDirectEditPolicy());
-		//installEditPolicy(EditPolicy.LAYOUT_ROLE, null);
-	}
-
-	@Override
-    public void performRequest(Request request)
-	{
-		if (request.getType() == RequestConstants.REQ_DIRECT_EDIT)
-		{
+    public DiagramPart getDiagramPart() {
+        return (DiagramPart) getParent().getParent();
+    }
+
+    @Override
+    public AttributeItemFigure getFigure() {
+        return (AttributeItemFigure) super.getFigure();
+    }
+
+    /**
+     * Create EditPolicies for the column label
+     */
+    @Override
+    protected void createEditPolicies() {
+        //installEditPolicy(EditPolicy.COMPONENT_ROLE, new AttributeEditPolicy());
+        //installEditPolicy(EditPolicy.DIRECT_EDIT_ROLE, new ColumnDirectEditPolicy());
+        //installEditPolicy(EditPolicy.LAYOUT_ROLE, null);
+    }
+
+    @Override
+    public void performRequest(Request request) {
+        if (request.getType() == RequestConstants.REQ_DIRECT_EDIT) {
 /*
 			if (request instanceof DirectEditRequest
 					&& !directEditHitTest(((DirectEditRequest) request).getLocation().getCopy()))
@@ -124,98 +128,88 @@ public class AttributePart extends PropertyAwarePart
         } else if (request.getType() == RequestConstants.REQ_OPEN) {
             getAttribute().openEditor();
         }
-	}
+    }
 
-	private boolean directEditHitTest(Point requestLoc)
-	{
-		IFigure figure = getFigure();
-		figure.translateToRelative(requestLoc);
+    private boolean directEditHitTest(Point requestLoc) {
+        IFigure figure = getFigure();
+        figure.translateToRelative(requestLoc);
         return figure.containsPoint(requestLoc);
     }
 
-	protected void performDirectEdit()
-	{
-		if (manager == null)
-		{
-			ERDGraphicalViewer viewer = (ERDGraphicalViewer) getViewer();
-			ValidationMessageHandler handler = viewer.getValidationHandler();
-
-			Label l = getFigure();
-			ColumnNameTypeCellEditorValidator columnNameTypeCellEditorValidator = new ColumnNameTypeCellEditorValidator(
-					handler);
-
-			manager = new ExtendedDirectEditManager(this, TextCellEditor.class, new LabelCellEditorLocator(l), l,
-					columnNameTypeCellEditorValidator);
-		}
-		manager.show();
-	}
-
-	/**
-	 * Sets the width of the line when selected
-	 */
-	@Override
-    public void setSelected(int value)
-	{
-		super.setSelected(value);
-		EditableLabel columnLabel = getFigure();
-		if (value != EditPart.SELECTED_NONE)
-			columnLabel.setSelected(true);
-		else
-			columnLabel.setSelected(false);
-		columnLabel.repaint();
-	}
-
-	public void handleNameChange(String textValue)
-	{
-		EditableLabel label = getFigure();
-		label.setVisible(false);
-		setSelected(EditPart.SELECTED_NONE);
-		label.revalidate();
-	}
-
-	/**
-	 * Handles when successfully applying direct edit
-	 */
-	@Override
-    protected void commitNameChange(PropertyChangeEvent evt)
-	{
-		AttributeItemFigure label = getFigure();
-		label.setText(getAttribute().getLabelText());
-		setSelected(EditPart.SELECTED_PRIMARY);
-		label.revalidate();
-	}
-
-
-	/**
-	 * Reverts state back to prior edit state
-	 */
-	public void revertNameChange(String oldValue)
-	{
-		AttributeItemFigure label = getFigure();
-		label.setVisible(true);
-		setSelected(EditPart.SELECTED_PRIMARY);
-		label.revalidate();
-	}
-
-	/**
-	 * We don't need to explicitly handle refresh visuals because the times when
-	 * this needs to be done it is handled by the table e.g. handleNameChange()
-	 */
-	@Override
-    protected void refreshVisuals()
-	{
-		ERDEntityAttribute column = (ERDEntityAttribute) getModel();
-		getFigure().setText(column.getLabelText());
-	}
-	
-	public ERDEntityAttribute getAttribute()
-	{
-		return (ERDEntityAttribute) getModel();
-	}
-	
+    protected void performDirectEdit() {
+        if (manager == null) {
+            ERDGraphicalViewer viewer = (ERDGraphicalViewer) getViewer();
+            ValidationMessageHandler handler = viewer.getValidationHandler();
+
+            Label l = getFigure().getLabel();
+            ColumnNameTypeCellEditorValidator columnNameTypeCellEditorValidator = new ColumnNameTypeCellEditorValidator(
+                    handler);
+
+            manager = new ExtendedDirectEditManager(this, TextCellEditor.class, new LabelCellEditorLocator(l), l,
+                    columnNameTypeCellEditorValidator);
+        }
+        manager.show();
+    }
+
+    /**
+     * Sets the width of the line when selected
+     */
+    @Override
+    public void setSelected(int value) {
+        super.setSelected(value);
+        EditableLabel columnLabel = getFigure().getLabel();
+        if (value != EditPart.SELECTED_NONE)
+            columnLabel.setSelected(true);
+        else
+            columnLabel.setSelected(false);
+        columnLabel.repaint();
+    }
+
+    public void handleNameChange(String textValue) {
+        EditableLabel label = getFigure().getLabel();
+        label.setVisible(false);
+        setSelected(EditPart.SELECTED_NONE);
+        label.revalidate();
+    }
+
+    /**
+     * Handles when successfully applying direct edit
+     */
+    @Override
+    protected void commitNameChange(PropertyChangeEvent evt) {
+        AttributeItemFigure label = getFigure();
+        label.getLabel().setText(getAttribute().getLabelText());
+        setSelected(EditPart.SELECTED_PRIMARY);
+        label.revalidate();
+    }
+
+
+    /**
+     * Reverts state back to prior edit state
+     */
+    public void revertNameChange(String oldValue) {
+        AttributeItemFigure label = getFigure();
+        label.setVisible(true);
+        setSelected(EditPart.SELECTED_PRIMARY);
+        label.revalidate();
+    }
+
+    /**
+     * We don't need to explicitly handle refresh visuals because the times when
+     * this needs to be done it is handled by the table e.g. handleNameChange()
+     */
+    @Override
+    protected void refreshVisuals() {
+        ERDEntityAttribute column = (ERDEntityAttribute) getModel();
+        getFigure().getLabel().setText(column.getLabelText());
+    }
+
+    public ERDEntityAttribute getAttribute() {
+        return (ERDEntityAttribute) getModel();
+    }
+
     @Override
-    public String toString()
-    {
+    public String toString() {
         return ERDMessages.column_ + getAttribute().getLabelText();
     }