lib/cassanity/schema.rb in cassanity-0.1.0 vs lib/cassanity/schema.rb in cassanity-0.2.0

@@ -1,11 +1,14 @@
 module Cassanity
   class Schema
     # Internal
-    attr_reader :primary_key
+    attr_reader :primary_keys
 
     # Internal
+    alias_method :primary_key, :primary_keys
+
+    # Internal
     attr_reader :columns
 
     # Internal
     attr_reader :with
 
@@ -15,12 +18,46 @@
     #        :primary_key - The String or Symbol key or Array of String/Symbol
     #                       keys to use as primary key.
     #        :columns - The Hash of columns where the name is the column name
     #                   and the value is the column type.
     #        :with - The Hash of options for the WITH clause.
+    # Raises KeyError if missing required argument key.
+    # Raises ArgumentError if primary key is not included in the columns.
     def initialize(args = {})
-      @primary_key = args.fetch(:primary_key)
+      @primary_keys = Array(args.fetch(:primary_key))
       @columns = args.fetch(:columns)
+
+      ensure_primary_keys_are_columns
+
       @with = args[:with] || {}
+      @composite_primary_key = @primary_keys.size > 1
+    end
+
+    # Public
+    def composite_primary_key?
+      @composite_primary_key == true
+    end
+
+    # Public: Returns an array of the column names
+    def column_names
+      @column_names ||= @columns.keys
+    end
+
+    # Public: Returns an array of the column types
+    def column_types
+      @column_types ||= @columns.values
+    end
+
+    # Private
+    def ensure_primary_keys_are_columns
+      unless primary_keys_are_defined_as_columns?
+        raise ArgumentError, "Not all primary keys (#{primary_keys.inspect}) were defined as columns (#{column_names.inspect})"
+      end
+    end
+
+    # Private
+    def primary_keys_are_defined_as_columns?
+      shared_columns = column_names & @primary_keys
+      shared_columns == @primary_keys
     end
   end
 end