1 files changed, 18 insertions, 20 deletions

diff --git a/src/diagramdocument.cpp b/src/diagramdocument.cpp
index 6edd54e..2aa486d 100644
--- a/src/diagramdocument.cpp
+++ b/src/diagramdocument.cpp
@@ -15,6 +15,8 @@
 // 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
 #include "diagramdocument.h"
+#include "diagramobject.h"
+#include "diagramconnection.h"
 #include "databasetable.h"
 #include "databaserelationship.h"
 #include "column.h"
@@ -28,7 +30,6 @@ DiagramDocument::DiagramDocument(QObject *parent)
 	: QGraphicsScene(parent), m_mode(DiagramDocument::Select), m_line(NULL)
 {
 	m_undoStack = new QUndoStack(this);
-	//connect(this, SIGNAL(tableMoved(DatabaseTable *)), SLOT(updatePositions(DatabaseTable *)));
 }
 
 DiagramDocument::Mode
@@ -45,11 +46,10 @@ DiagramDocument::setMode(Mode mode)
 }
 
 void
-DiagramDocument::updatePositions(DatabaseTable *table)
+DiagramDocument::updatePositions(DiagramObject *object)
 {
-	foreach (DatabaseRelationship *relation, findTableRelations(table)) {
-		relation->updatePositions();
-	}
+	foreach (DiagramConnection *connection, findConnections(object))
+		connection->updatePositions();
 }
 
 template <class T> QList<T *>
@@ -57,9 +57,13 @@ DiagramDocument::itemsByType()
 {
 	QList<T *> result;
 	foreach(QGraphicsItem *item, items()) {
-		if (item->type() == T::Type) {
-			result.append(static_cast<T *>(item));
+		T *typedItem = dynamic_cast<T *>(item);
+		if (typedItem) {
+			result.append(typedItem);
 		}
+/*		if (item->type() == T::Type) {
+			result.append(static_cast<T *>(item));
+		}*/
 	}
 	return result;
 }
@@ -165,24 +169,18 @@ DiagramDocument::deleteSelectedItems()
 	foreach (QGraphicsItem *item, selectedItems()) {
 		DatabaseTable *table = qgraphicsitem_cast<DatabaseTable *>(item);
 		if (table) {
-			// FIXME
-			foreach (DatabaseRelationship *relation, findTableRelations(table)) {
-				removeItem(relation);
-			}
-			undoStack()->push(new RemoveItemCommand(this, table));
+			undoStack()->push(new RemoveObjectCommand(this, table));
 		}
 	}
 }
 
-QList<DatabaseRelationship *>
-DiagramDocument::findTableRelations(DatabaseTable *table)
+QList<DiagramConnection *>
+DiagramDocument::findConnections(DiagramObject *object)
 {
-	QList<DatabaseRelationship *> result;
-	foreach (DatabaseRelationship *relation, itemsByType<DatabaseRelationship>()) {
-		if (relation->source() == table || relation->target() == table) {
-			result << relation;
-		}
-	}
+	QList<DiagramConnection *> result;
+	foreach (DiagramConnection *connection, itemsByType<DiagramConnection>())
+		if (connection->source() == object || connection->target() == object)
+			result.append(connection);
 	return result;
 }