1 files changed, 48 insertions, 16 deletions

diff --git a/app/models/concerns/sha_attribute.rb b/app/models/concerns/sha_attribute.rb
index e49f4d03bda..701d2fda5c5 100644
--- a/app/models/concerns/sha_attribute.rb
+++ b/app/models/concerns/sha_attribute.rb
@@ -3,39 +3,71 @@
 module ShaAttribute
   extend ActiveSupport::Concern
 
-  # Needed for the database method
-  include DatabaseReflection
+  class ShaAttributeTypeMismatchError < StandardError
+    def initialize(column_name, column_type)
+      @column_name = column_name
+      @column_type = column_type
+    end
+
+    def message
+      "sha_attribute :#{@column_name} should be a :binary column but it is :#{@column_type}"
+    end
+  end
+
+  class Sha256AttributeTypeMismatchError < ShaAttributeTypeMismatchError
+    def message
+      "sha256_attribute :#{@column_name} should be a :binary column but it is :#{@column_type}"
+    end
+  end
 
   class_methods do
     def sha_attribute(name)
       return if ENV['STATIC_VERIFICATION']
 
-      validate_binary_column_exists!(name) if Rails.env.development? || Rails.env.test?
+      sha_attribute_fields << name
 
       attribute(name, Gitlab::Database::ShaAttribute.new)
     end
 
+    def sha_attribute_fields
+      @sha_attribute_fields ||= []
+    end
+
+    def sha256_attribute(name)
+      return if ENV['STATIC_VERIFICATION']
+
+      sha256_attribute_fields << name
+
+      attribute(name, Gitlab::Database::Sha256Attribute.new)
+    end
+
+    def sha256_attribute_fields
+      @sha256_attribute_fields ||= []
+    end
+
     # This only gets executed in non-production environments as an additional check to ensure
     # the column is the correct type.  In production it should behave like any other attribute.
     # See https://gitlab.com/gitlab-org/gitlab/merge_requests/5502 for more discussion
-    def validate_binary_column_exists!(name)
-      return unless database_exists?
-      return unless table_exists?
+    def load_schema!
+      super
 
-      column = columns.find { |c| c.name == name.to_s }
+      return if Rails.env.production?
 
-      return unless column
+      sha_attribute_fields.each do |field|
+        column = columns_hash[field.to_s]
 
-      unless column.type == :binary
-        raise ArgumentError, "sha_attribute #{name.inspect} is invalid since the column type is not :binary"
+        if column && column.type != :binary
+          raise ShaAttributeTypeMismatchError.new(column.name, column.type)
+        end
       end
-    rescue StandardError => error
-      Gitlab::AppLogger.error "ShaAttribute initialization: #{error.message}"
-      raise
-    end
 
-    def database_exists?
-      database.exists?
+      sha256_attribute_fields.each do |field|
+        column = columns_hash[field.to_s]
+
+        if column && column.type != :binary
+          raise Sha256AttributeTypeMismatchError.new(column.name, column.type)
+        end
+      end
     end
   end
 end