sortkey, distkey options for columns

lib/active_record/connection_adapters/redshift/column.rb

@@ -1,9 +1,13 @@
 module ActiveRecord
   module ConnectionAdapters
     class RedshiftColumn < Column #:nodoc:
-      def initialize(name, default, cast_type, sql_type = nil, null = true, default_function = nil)
+      attr_reader :sortkey, :distkey
+
+      def initialize(name, default, cast_type, sql_type = nil, null = true, default_function = nil, sortkey = false, distkey = false)
         super name, default, cast_type, sql_type, null
         @default_function = default_function
+        @sortkey = sortkey
+        @distkey = distkey
       end
     end
   end

lib/active_record/connection_adapters/redshift/schema_definitions.rb

@@ -11,7 +11,10 @@ def jsonb(name, options = {})
         end
       end
 
-      class ColumnDefinition < ActiveRecord::ConnectionAdapters::ColumnDefinition
+      class ColumnDefinition < Struct.new(*ActiveRecord::ConnectionAdapters::ColumnDefinition.members, :sortkey, :distkey)
+        def primary_key?
+          primary_key || type.to_sym == :primary_key
+        end
       end
 
       class TableDefinition < ActiveRecord::ConnectionAdapters::TableDefinition
@@ -52,6 +55,13 @@ def primary_key(name, type = :primary_key, options = {})
           column name, type, options
         end
 
+        def new_column_definition(name, type, options) # :nodoc:
+          column = super
+          column.sortkey = options[:sortkey]
+          column.distkey = options[:distkey]
+          column
+        end
+
         private
 
           def create_column_definition(name, type)

lib/active_record/connection_adapters/redshift/schema_statements.rb

@@ -14,20 +14,22 @@ def visit_ColumnDefinition(o)
         end
 
         def add_column_options!(sql, options)
-          column = options.fetch(:column) { return super }
-          if column.type == :uuid && options[:default] =~ /\(\)/
-            sql << " DEFAULT #{options[:default]}"
-          else
-            super
+          if options[:sortkey]
+            sql << " SORTKEY"
           end
-        end
 
-        def type_for_column(column)
-          if column.array
-            @conn.lookup_cast_type("#{column.sql_type}[]")
-          else
-            super
+          if options[:distkey]
+            sql << " DISTKEY"
           end
+
+          super
+        end
+
+        def column_options(o)
+          column_options = super
+          column_options[:sortkey] = o.sortkey
+          column_options[:distkey] = o.distkey
+          column_options
         end
       end
 
@@ -121,16 +123,16 @@ def indexes(table_name, name = nil)
         # Returns the list of all column definitions for a table.
         def columns(table_name)
           # Limit, precision, and scale are all handled by the superclass.
-          column_definitions(table_name).map do |column_name, type, default, notnull, oid, fmod|
+          column_definitions(table_name).map do |column_name, type, default, notnull, oid, fmod, sortkey, distkey|
             oid = get_oid_type(oid.to_i, fmod.to_i, column_name, type)
             default_value = extract_value_from_default(oid, default)
             default_function = extract_default_function(default_value, default)
-            new_column(column_name, default_value, oid, type, notnull == 'f', default_function)
+            new_column(column_name, default_value, oid, type, notnull == 'f', default_function, sortkey == '1', distkey)
           end
         end
 
-        def new_column(name, default, cast_type, sql_type = nil, null = true, default_function = nil) # :nodoc:
-          RedshiftColumn.new(name, default, cast_type, sql_type, null, default_function)
+        def new_column(name, default, cast_type, sql_type = nil, null = true, default_function = nil, sortkey = false, distkey = false) # :nodoc:
+          RedshiftColumn.new(name, default, cast_type, sql_type, null, default_function, sortkey, distkey)
         end
 
         # Returns the current database name.

lib/active_record/connection_adapters/redshift_adapter.rb

@@ -631,7 +631,7 @@ def last_insert_id_result(sequence_name) #:nodoc:
         def column_definitions(table_name) # :nodoc:
           exec_query(<<-end_sql, 'SCHEMA').rows
               SELECT a.attname, format_type(a.atttypid, a.atttypmod),
-                     pg_get_expr(d.adbin, d.adrelid), a.attnotnull, a.atttypid, a.atttypmod
+                     pg_get_expr(d.adbin, d.adrelid), a.attnotnull, a.atttypid, a.atttypmod, a.attsortkeyord, a.attisdistkey
                 FROM pg_attribute a LEFT JOIN pg_attrdef d
                   ON a.attrelid = d.adrelid AND a.attnum = d.adnum
                WHERE a.attrelid = '#{quote_table_name(table_name)}'::regclass