D5sBHkbn

1diff --git a/examples/orm.rb b/examples/orm.rb
2index bfb475d..1fb9717 100755
3--- a/examples/orm.rb
4+++ b/examples/orm.rb
5@@ -1,5 +1,6 @@
6 #!/usr/bin/env ruby
7 require_relative '../lib/swift-orm'
8+require_relative '../lib/swift-orm/types'
9 require 'etc'
10 require 'pp'
11 
12@@ -9,10 +10,11 @@ class User < Swift.resource do
13     property :name,     String,  index: 'info'
14     property :email,    String,  index: 'info'
15     property :optional, String,  default: 'woot'
16+    property :category, Swift::Types::Enum["admin", "user"], default: 'user'
17   end
18 end # User
19 
20-Swift.setup user: Etc.getlogin, driver: 'postgresql', db: 'swift'
21+Swift.setup user: Etc.getlogin, db: 'swift', driver: ARGV[0] || 'postgresql'
22 Swift.trace true
23 
24 puts '-- auto_migrate! --'
25diff --git a/lib/swift-orm/migrations.rb b/lib/swift-orm/migrations.rb
26index 4e66b89..1dbd54e 100644
27--- a/lib/swift-orm/migrations.rb
28+++ b/lib/swift-orm/migrations.rb
29@@ -9,7 +9,7 @@ module Swift
30     end
31 
32     def migrate_up! resource
33-      fields = resource.properties.map{|p| column_definition(p)}.join(', ')
34+      fields = resource.properties.map{|p| column_definition(p, resource)}.join(', ')
35       keys   = resource.properties.keys
36       fields += ', primary key (%s)' % keys.join(', ') unless keys.empty? or resource.properties.serial?
37       @adapter.execute "create table #{resource.store} (#{fields})"
38@@ -19,8 +19,15 @@ module Swift
39       end
40     end
41 
42-    def column_definition property
43-      "#{property.field} #{property.serial ? typemap[:serial] : typemap[property.type] || 'text'}"
44+    def column_definition property, resource
45+      type, *commands = if property.type.respond_to?(:storage_type)
46+        property.type.storage_type(property, resource, @adapter.driver)
47+      else
48+        "#{property.serial ? typemap[:serial] : typemap[property.type] || 'text'}"
49+      end
50+
51+      commands.each {|command| @adapter.execute(command)} unless commands.empty?
52+      "#{property.field} #{type}"
53     end
54 
55     def migrate_down! resource
56diff --git a/lib/swift-orm/types.rb b/lib/swift-orm/types.rb
57new file mode 100644
58index 0000000..45ba416
59--- /dev/null
60+++ b/lib/swift-orm/types.rb
61@@ -0,0 +1,31 @@
62+module Swift
63+  module Types
64+    class Enum
65+      def initialize value
66+        if !self.class.values.include?(value)
67+          raise ArgumentError, "#{value} not found in (%s)" % self.class.values.join(',')
68+        end
69+        @value = value
70+      end
71+
72+      class << self
73+        attr_reader :values
74+        def [] *values
75+          Class.new(self) do
76+            @values = values.flatten
77+          end
78+        end
79+        def storage_type property, resource, driver
80+          defn = 'enum(%s)' % values.map{|v| "'#{v}'"}.join(',')
81+          case driver.to_s
82+            when 'postgresql'
83+              name = "#{resource.store}_#{property.field}_type"
84+              [ name, "drop type if exists #{name}", "create type #{name} as #{defn}" ]
85+            else
86+              defn
87+          end
88+        end
89+      end
90+    end # Enum
91+  end # Types
92+end # Swift