Fix undefined method `postgresql?` during migrationrs-gitlab-database-module

2 files changed, 34 insertions, 6 deletions

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 6ebb8027553..6f9da69983a 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -17,7 +17,7 @@ module Gitlab
     end
 
     def true_value
-      if self.class.postgresql?
+      if Gitlab::Database.postgresql?
         "'t'"
       else
         1
@@ -25,7 +25,7 @@ module Gitlab
     end
 
     def false_value
-      if self.class.postgresql?
+      if Gitlab::Database.postgresql?
         "'f'"
       else
         0
@@ -47,9 +47,5 @@ module Gitlab
         row.first
       end
     end
-
-    def connection
-      self.class.connection
-    end
   end
 end
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index bd8688fefa1..d0a447753b7 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -1,5 +1,9 @@
 require 'spec_helper'
 
+class MigrationTest
+  include Gitlab::Database
+end
+
 describe Gitlab::Database, lib: true do
   # These are just simple smoke tests to check if the methods work (regardless
   # of what they may return).
@@ -34,4 +38,32 @@ describe Gitlab::Database, lib: true do
       end
     end
   end
+
+  describe '#true_value' do
+    it 'returns correct value for PostgreSQL' do
+      expect(described_class).to receive(:postgresql?).and_return(true)
+
+      expect(MigrationTest.new.true_value).to eq "'t'"
+    end
+
+    it 'returns correct value for MySQL' do
+      expect(described_class).to receive(:postgresql?).and_return(false)
+
+      expect(MigrationTest.new.true_value).to eq 1
+    end
+  end
+
+  describe '#false_value' do
+    it 'returns correct value for PostgreSQL' do
+      expect(described_class).to receive(:postgresql?).and_return(true)
+
+      expect(MigrationTest.new.false_value).to eq "'f'"
+    end
+
+    it 'returns correct value for MySQL' do
+      expect(described_class).to receive(:postgresql?).and_return(false)
+
+      expect(MigrationTest.new.false_value).to eq 0
+    end
+  end
 end