9 files changed, 106 insertions, 9 deletions
diff --git a/app/models/concerns/time_trackable.rb b/app/models/concerns/time_trackable.rb
index 040e3a2884e..9cf83440784 100644
--- a/app/models/concerns/time_trackable.rb
+++ b/app/models/concerns/time_trackable.rb
@@ -18,7 +18,7 @@ module TimeTrackable
     validates :time_estimate, numericality: { message: 'has an invalid format' }, allow_nil: false
     validate  :check_negative_time_spent
 
-    has_many :timelogs, as: :trackable, dependent: :destroy
+    has_many :timelogs, dependent: :destroy
   end
 
   def spend_time(options)
diff --git a/app/models/timelog.rb b/app/models/timelog.rb
index f768c4e3da5..e166cf69703 100644
--- a/app/models/timelog.rb
+++ b/app/models/timelog.rb
@@ -1,6 +1,22 @@
 class Timelog < ActiveRecord::Base
   validates :time_spent, :user, presence: true
+  validate :issuable_id_is_present
 
-  belongs_to :trackable, polymorphic: true
+  belongs_to :issue
+  belongs_to :merge_request
   belongs_to :user
+
+  def issuable
+    issue || merge_request
+  end
+
+  private
+
+  def issuable_id_is_present
+    if issue_id && merge_request_id
+      errors.add(:base, 'Only Issue ID or Merge Request ID is required')
+    elsif issuable.nil?
+      errors.add(:base, 'Issue or Merge Request ID is required')
+    end
+  end
 end
diff --git a/changelogs/unreleased/26908-make-timelogs-use-foreign-keys b/changelogs/unreleased/26908-make-timelogs-use-foreign-keys
new file mode 100644
index 00000000000..0e8f7093b34
--- /dev/null
+++ b/changelogs/unreleased/26908-make-timelogs-use-foreign-keys
@@ -0,0 +1,4 @@
+---
+title: Refactor Timelogs structure to use foreign keys.
+merge_request: 8769
+author:
diff --git a/db/migrate/20170124174637_add_foreign_keys_to_timelogs.rb b/db/migrate/20170124174637_add_foreign_keys_to_timelogs.rb
new file mode 100644
index 00000000000..7956463e12c
--- /dev/null
+++ b/db/migrate/20170124174637_add_foreign_keys_to_timelogs.rb
@@ -0,0 +1,45 @@
+# See http://doc.gitlab.com/ce/development/migration_style_guide.html
+# for more information on how to write migrations for GitLab.
+
+class AddForeignKeysToTimelogs < ActiveRecord::Migration
+  include Gitlab::Database::MigrationHelpers
+
+  # Set this constant to true if this migration requires downtime.
+  DOWNTIME = true
+  # When a migration requires downtime you **must** uncomment the following
+  # constant and define a short and easy to understand explanation as to why the
+  # migration requires downtime.
+  DOWNTIME_REASON = 'Adding foreign keys'
+
+  # When using the methods "add_concurrent_index" or "add_column_with_default"
+  # you must disable the use of transactions as these methods can not run in an
+  # existing transaction. When using "add_concurrent_index" make sure that this
+  # method is the _only_ method called in the migration, any other changes
+  # should go in a separate migration. This ensures that upon failure _only_ the
+  # index creation fails and can be retried or reverted easily.
+  #
+  # To disable transactions uncomment the following line and remove these
+  # comments:
+  # disable_ddl_transaction!
+
+  def up
+    change_table :timelogs do |t|
+      t.references :issue, index: true, foreign_key: { on_delete: :cascade }
+      t.references :merge_request, index: true, foreign_key: { on_delete: :cascade }
+    end
+
+    Timelog.where(trackable_type: 'Issue').update_all("issue_id = trackable_id")
+    Timelog.where(trackable_type: 'MergeRequest').update_all("merge_request_id = trackable_id")
+
+    remove_columns :timelogs, :trackable_id, :trackable_type
+  end
+
+  def down
+    add_reference :timelogs, :trackable, polymorphic: true, index: true
+
+    Timelog.where('issue_id IS NOT NULL').update_all("trackable_id = issue_id, trackable_type = 'Issue'")
+    Timelog.where('merge_request_id IS NOT NULL').update_all("trackable_id = merge_request_id, trackable_type = 'MergeRequest'")
+
+    remove_columns :timelogs, :issue_id, :merge_request_id
+  end
+end
diff --git a/db/schema.rb b/db/schema.rb
index aeb0d8210f0..cef887f65f1 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -1151,14 +1151,15 @@ ActiveRecord::Schema.define(version: 20170206071414) do
 
   create_table "timelogs", force: :cascade do |t|
     t.integer "time_spent", null: false
-    t.integer "trackable_id"
-    t.string "trackable_type"
     t.integer "user_id"
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
+    t.integer "issue_id"
+    t.integer "merge_request_id"
   end
 
-  add_index "timelogs", ["trackable_type", "trackable_id"], name: "index_timelogs_on_trackable_type_and_trackable_id", using: :btree
+  add_index "timelogs", ["issue_id"], name: "index_timelogs_on_issue_id", using: :btree
+  add_index "timelogs", ["merge_request_id"], name: "index_timelogs_on_merge_request_id", using: :btree
   add_index "timelogs", ["user_id"], name: "index_timelogs_on_user_id", using: :btree
 
   create_table "todos", force: :cascade do |t|
@@ -1340,6 +1341,8 @@ ActiveRecord::Schema.define(version: 20170206071414) do
   add_foreign_key "protected_branch_merge_access_levels", "protected_branches"
   add_foreign_key "protected_branch_push_access_levels", "protected_branches"
   add_foreign_key "subscriptions", "projects", on_delete: :cascade
+  add_foreign_key "timelogs", "issues", on_delete: :cascade
+  add_foreign_key "timelogs", "merge_requests", on_delete: :cascade
   add_foreign_key "trending_projects", "projects", on_delete: :cascade
   add_foreign_key "u2f_registrations", "users"
 end
diff --git a/spec/factories/timelogs.rb b/spec/factories/timelogs.rb
index 12fc4ec4486..6f1545418eb 100644
--- a/spec/factories/timelogs.rb
+++ b/spec/factories/timelogs.rb
@@ -4,6 +4,6 @@ FactoryGirl.define do
   factory :timelog do
     time_spent 3600
     user
-    association :trackable, factory: :issue
+    issue
   end
 end
diff --git a/spec/lib/gitlab/import_export/all_models.yml b/spec/lib/gitlab/import_export/all_models.yml
index 5231ab0ba3f..06617f3b007 100644
--- a/spec/lib/gitlab/import_export/all_models.yml
+++ b/spec/lib/gitlab/import_export/all_models.yml
@@ -203,5 +203,6 @@ award_emoji:
 priorities:
 - label
 timelogs:
-- trackable
+- issue
+- merge_request
 - user
diff --git a/spec/lib/gitlab/import_export/safe_model_attributes.yml b/spec/lib/gitlab/import_export/safe_model_attributes.yml
index 95b230e4f5c..c5ac702d831 100644
--- a/spec/lib/gitlab/import_export/safe_model_attributes.yml
+++ b/spec/lib/gitlab/import_export/safe_model_attributes.yml
@@ -350,8 +350,8 @@ LabelPriority:
 Timelog:
 - id
 - time_spent
-- trackable_id
-- trackable_type
+- merge_request_id
+- issue_id
 - user_id
 - created_at
 - updated_at
diff --git a/spec/models/timelog_spec.rb b/spec/models/timelog_spec.rb
index f08935b6425..ebc694213b6 100644
--- a/spec/models/timelog_spec.rb
+++ b/spec/models/timelog_spec.rb
@@ -2,9 +2,37 @@ require 'rails_helper'
 
 RSpec.describe Timelog, type: :model do
   subject { build(:timelog) }
+  let(:issue) { create(:issue) }
+  let(:merge_request) { create(:merge_request) }
 
   it { is_expected.to be_valid }
 
   it { is_expected.to validate_presence_of(:time_spent) }
   it { is_expected.to validate_presence_of(:user) }
+
+  describe 'Issuable validation' do
+    it 'is invalid if issue_id and merge_request_id are missing' do
+      subject.attributes = { issue: nil, merge_request: nil }
+
+      expect(subject).to be_invalid
+    end
+
+    it 'is invalid if issue_id and merge_request_id are set' do
+      subject.attributes = { issue: issue, merge_request: merge_request }
+
+      expect(subject).to be_invalid
+    end
+
+    it 'is valid if only issue_id is set' do
+      subject.attributes = { issue: issue, merge_request: nil }
+
+      expect(subject).to be_valid
+    end
+
+    it 'is valid if only merge_request_id is set' do
+      subject.attributes = { merge_request: merge_request, issue: nil }
+
+      expect(subject).to be_valid
+    end
+  end
 end