Don't assume default client search path

lib/multi_schema/behaviors.rb

@@ -1,6 +1,8 @@
 module MultiSchema
   module Behaviors
     @@disable_message = false
+    @@default_search_path = '"$user",public'
+    @@service_schema = 'public'
 
     def disable_message=(val)
       @@disable_message = val
@@ -10,6 +12,18 @@ def disable_message
       @@disable_message
     end
 
+    def default_search_path=(val)
+      @@default_search_path = val
+    end
+
+    def default_search_path
+      @@default_search_path
+    end
+
+    def service_schema
+      @@service_schema
+    end
+
     def with_in_schemas(options = nil)
       options = unify_type(options, Hash) { |items| {:only => items} }
       options[:only] = unify_type(options[:only], Array) { |item| item.nil? ? all_schemas : [item] }.map { |item| item.to_s }
@@ -29,7 +43,7 @@ def with_in_schemas(options = nil)
 
     def all_schemas
       ActiveRecord::Base.connection.select_values <<-END
-      SELECT *
+      SELECT nspname
       FROM pg_namespace
       WHERE
         nspname NOT IN ('information_schema') AND
@@ -42,8 +56,11 @@ def current_schema
     end
 
     def reset_schema_path
-      puts "--- Restore Schema to public" unless disable_message
-      ActiveRecord::Base.connection.schema_search_path = 'public'
+      puts '--- Restore Schema to "$user", public' unless disable_message
+      ::ActiveRecord::Base.connection.schema_search_path = '"$user", public'
+      @@service_schema = 'public'
+      clear_cache
+      nil
     end
 
     def schema_path
@@ -52,7 +69,14 @@ def schema_path
 
     def set_schema_path(schema)
       puts "--- Select Schema: #{schema} " unless disable_message
+      @@service_schema = schema.split(/\s*,\s*/).reject { |x| x == '"$user"' }.first
       ActiveRecord::Base.connection.schema_search_path = schema
+      clear_cache
+      nil
+    end
+
+    def set_schema_and_public(schema)
+      set_schema_path "#{schema}, public"
     end
 
     def push_schema_to_path(schema)
@@ -65,6 +89,34 @@ def pop_schema_from_path
       set_schema_path(new_path)
     end
 
+    def ensure_schema_reset
+      old = ::ActiveRecord::Base.connection.schema_search_path
+
+      begin
+        yield
+      ensure
+        e = $ERROR_INFO
+        e = e.cause while e.respond_to?(:cause)
+        set_schema_path old unless e.is_a?(PG::Error)
+      end
+    end
+
+    def with_schema_and_public(schema)
+      ensure_schema_reset do
+        set_schema_and_public schema
+        yield
+      end
+    end
+
+    def each_schema_and_public
+      ensure_schema_reset do
+        (all_schemas - ['public']).each do |schema|
+          set_schema_and_public schema
+          yield schema
+        end
+      end
+    end
+
     private
 
     def unify_type(input, type)
@@ -80,5 +132,18 @@ def unify_array_item_type(input, type, &block)
         unify_type item, type, &block
       end
     end
+
+    def clear_cache
+      ::ActiveRecord::Base.connection.schema_cache.clear!
+
+      ::ActiveRecord::Base.descendants.map(&:base_class).uniq.each do |klass|
+        # reset_sequence_name would needlessly fetch them all now.
+        unless klass.instance_variable_get(:@explicit_sequence_name)
+          klass.instance_variable_set(:@sequence_name, nil)
+        end
+      end
+
+      nil
+    end
   end
 end