diff --git a/src/Tasks/AddIdentifierField.php b/src/Tasks/AddIdentifierField.php
index c85fe77..d6fda4a 100644
--- a/src/Tasks/AddIdentifierField.php
+++ b/src/Tasks/AddIdentifierField.php
@@ -2,13 +2,15 @@
 
 namespace Naoray\BlueprintNovaAddon\Tasks;
 
-use Blueprint\Models\Column;
-use Blueprint\Models\Model;
 use Closure;
+use Blueprint\Models\Model;
 use Illuminate\Support\Arr;
+use Blueprint\Models\Column;
 
 class AddIdentifierField
 {
+    use InteractWithRelationships;
+
     const INDENT = '            ';
 
     /** @var Model */
@@ -21,32 +23,23 @@ public function __construct(Model $model)
 
     public function handle($data, Closure $next): array
     {
-        $column = $this->getIdentifierColumn(
-            $this->model->columns(),
-            $this->model->relationships()
-        );
+        $column = $this->identifierColumn();
 
-        $identifierName = $column->name() === 'id' ? '' : "'".$column->name()."'";
-        $data['fields'] .= 'ID::make('.$identifierName.')->sortable(),'.PHP_EOL.PHP_EOL;
+        $identifierName = $column->name() === 'id' ? '' : "'" . $column->name() . "'";
+        $data['fields'] .= 'ID::make(' . $identifierName . ')->sortable(),' . PHP_EOL . PHP_EOL;
         $data['imports'][] = 'ID';
 
         return $next($data);
     }
 
-    private function getIdentifierColumn(array $columns, array $relations): Column
+    private function identifierColumn(): Column
     {
-        $name = collect($columns)
-            ->filter(function (Column $column) {
-                return $column->dataType() === 'id';
-            })
-            ->map(function (Column $column) {
-                return empty($column->attributes())
-                    ? $column->name()
-                    : implode(':', $column->attributes()).":{$column->name()}";
-            })->values()
-            ->diff(Arr::get($relations, 'belongsTo', []))
+        $name = $this->relationshipIdentifiers($this->model->columns())
+            ->values()
+            // filter out all relationships
+            ->diff(Arr::get($this->model->relationships(), 'belongsTo', []))
             ->first();
 
-        return $columns[$name];
+        return $this->model->columns()[$name];
     }
 }
diff --git a/src/Tasks/InteractWithRelationships.php b/src/Tasks/InteractWithRelationships.php
new file mode 100644
index 0000000..065de9b
--- /dev/null
+++ b/src/Tasks/InteractWithRelationships.php
@@ -0,0 +1,23 @@
+<?php
+
+namespace Naoray\BlueprintNovaAddon\Tasks;
+
+use Blueprint\Models\Column;
+use Illuminate\Support\Collection;
+
+trait InteractWithRelationships
+{
+    private function relationshipIdentifiers(array $columns): Collection
+    {
+        return collect($columns)
+            ->filter(function (Column $column) {
+                return $column->dataType() === 'id';
+            })
+            // map id columns to match related relationships
+            ->map(function (Column $column) {
+                return empty($column->attributes())
+                    ? $column->name()
+                    : implode(':', $column->attributes()) . ":{$column->name()}";
+            });
+    }
+}