Add support for multiple schemas

Author: Carlos Silva

lib/torque/postgresql.rb

@@ -20,12 +20,13 @@
 require 'torque/postgresql/attributes'
 require 'torque/postgresql/autosave_association'
 require 'torque/postgresql/auxiliary_statement'
-require 'torque/postgresql/base'
 require 'torque/postgresql/inheritance'
+require 'torque/postgresql/base' # Needs to be after inheritance
 require 'torque/postgresql/insert_all'
 require 'torque/postgresql/migration'
 require 'torque/postgresql/relation'
 require 'torque/postgresql/reflection'
 require 'torque/postgresql/schema_cache'
+require 'torque/postgresql/table_name'
 
 require 'torque/postgresql/railtie' if defined?(Rails)

lib/torque/postgresql/adapter.rb

@@ -31,9 +31,9 @@ def version
         )
       end
 
-      # Add `inherits` to the list of extracted table options
+      # Add `inherits` and `schema` to the list of extracted table options
       def extract_table_options!(options)
-        super.merge(options.extract!(:inherits))
+        super.merge(options.extract!(:inherits, :schema))
       end
 
       # Allow filtered bulk insert by adding the where clause. This method is

lib/torque/postgresql/adapter/database_statements.rb

@@ -12,6 +12,26 @@ def dump_mode!
           @_dump_mode = !!!@_dump_mode
         end
 
+        # List of schemas blocked by the application in the current connection
+        def schemas_blacklist
+          @schemas_blacklist ||= Torque::PostgreSQL.config.schemas.blacklist +
+            (@config.dig(:schemas, 'blacklist') || [])
+        end
+
+        # List of schemas used by the application in the current connection
+        def schemas_whitelist
+          @schemas_whitelist ||= Torque::PostgreSQL.config.schemas.whitelist +
+            (@config.dig(:schemas, 'whitelist') || [])
+        end
+
+        # A list of schemas on the search path sanitized
+        def schemas_search_path_sanitized
+          @schemas_search_path_sanitized ||= begin
+            db_user = @config[:username] || ENV['USER'] || ENV['USERNAME']
+            schema_search_path.split(',').map { |item| item.strip.sub('"$user"', db_user) }
+          end
+        end
+
         # Check if a given type is valid.
         def valid_type?(type)
           super || extended_types.include?(type)
@@ -22,6 +42,17 @@ def extended_types
           EXTENDED_DATABASE_TYPES
         end
 
+        # Checks if a given schema exists in the database. If +filtered+ is
+        # given as false, then it will check regardless of whitelist and
+        # blacklist
+        def schema_exists?(name, filtered: true)
+          return user_defined_schemas.include?(name.to_s) if filtered
+
+          query_value(<<-SQL) == 1
+            SELECT 1 FROM pg_catalog.pg_namespace WHERE nspname = '#{name}'
+          SQL
+        end
+
         # Returns true if type exists.
         def type_exists?(name)
           user_defined_types.key? name.to_s
@@ -124,18 +155,41 @@ def user_defined_types(*categories)
         # Get the list of inherited tables associated with their parent tables
         def inherited_tables
           tables = query(<<-SQL, 'SCHEMA')
-            SELECT child.relname             AS table_name,
-                   array_agg(parent.relname) AS inheritances
+            SELECT inhrelid::regclass  AS table_name,
+                   inhparent::regclass AS inheritances
             FROM pg_inherits
             JOIN pg_class parent ON pg_inherits.inhparent = parent.oid
             JOIN pg_class child  ON pg_inherits.inhrelid  = child.oid
-            GROUP BY child.relname, pg_inherits.inhrelid
-            ORDER BY pg_inherits.inhrelid
+            ORDER BY inhrelid
           SQL
 
-          tables.map do |(table, refs)|
-            [table, PG::TextDecoder::Array.new.decode(refs)]
-          end.to_h
+          tables.each_with_object({}) do |(child, parent), result|
+            (result[child] ||= []) << parent
+          end
+        end
+
+        # Get the list of schemas that were created by the user
+        def user_defined_schemas
+          query_values(user_defined_schemas_sql, 'SCHEMA')
+        end
+
+        # Build the query for allowed schemas
+        def user_defined_schemas_sql
+          conditions = []
+          conditions << <<-SQL if schemas_blacklist.any?
+            nspname NOT LIKE ANY (ARRAY['#{schemas_blacklist.join("', '")}'])
+          SQL
+
+          conditions << <<-SQL if schemas_whitelist.any?
+            nspname LIKE ANY (ARRAY['#{schemas_whitelist.join("', '")}'])
+          SQL
+
+          <<-SQL.squish
+            SELECT nspname
+            FROM pg_catalog.pg_namespace
+            WHERE 1=1 AND #{conditions.join(' AND ')}
+            ORDER BY oid
+          SQL
         end
 
         # Get the list of columns, and their definition, but only from the

lib/torque/postgresql/adapter/schema_dumper.rb

@@ -12,6 +12,11 @@ def dump(stream) # :nodoc:
           stream
         end
 
+        def extensions(stream) # :nodoc:
+          super
+          user_defined_schemas(stream)
+        end
+
         # Translate +:enum_set+ into +:enum+
         def schema_type(column)
           column.type == :enum_set ? :enum : super
@@ -34,7 +39,9 @@ def around_tables(stream)
 
           def dump_tables(stream)
             inherited_tables = @connection.inherited_tables
-            sorted_tables = @connection.tables.sort - @connection.views
+            sorted_tables = (@connection.tables - @connection.views).sort_by do |table_name|
+              table_name.split(/(?:public)?\./).reverse
+            end
 
             stream.puts "  # These are the common tables"
             (sorted_tables - inherited_tables.keys).each do |table_name|
@@ -51,7 +58,7 @@ def dump_tables(stream)
 
                 # Add the inherits setting
                 sub_stream.rewind
-                inherits.map!(&:to_sym)
+                inherits.map! { |parent| parent.to_s.sub(/\Apublic\./, '') }
                 inherits = inherits.first if inherits.size === 1
                 inherits = ", inherits: #{inherits.inspect} do |t|"
                 table_dump = sub_stream.read.gsub(/ do \|t\|$/, inherits)
@@ -70,6 +77,20 @@ def dump_tables(stream)
             end
           end
 
+          # Make sure to remove the schema from the table name
+          def remove_prefix_and_suffix(table)
+            super(table.sub(/\A[a-z0-9_]*\./, ''))
+          end
+
+          # Dump user defined schemas
+          def user_defined_schemas(stream)
+            return if (list = (@connection.user_defined_schemas - ['public'])).empty?
+
+            stream.puts "  # Custom schemas defined in this database."
+            list.each { |name| stream.puts "  create_schema \"#{name}\", force: :cascade" }
+            stream.puts
+          end
+
           def fx_functions_position
             return unless defined?(::Fx::SchemaDumper::Function)
             Fx.configuration.dump_functions_at_beginning_of_schema ? :beginning : :end

lib/torque/postgresql/adapter/schema_statements.rb

@@ -7,6 +7,21 @@ module SchemaStatements
 
         TableDefinition = ActiveRecord::ConnectionAdapters::PostgreSQL::TableDefinition
 
+        # Create a new schema
+        def create_schema(name, options = {})
+          drop_schema(name, options) if options[:force]
+
+          check = 'IF NOT EXISTS' if options.fetch(:check, true)
+          execute("CREATE SCHEMA #{check} #{quote_schema_name(name.to_s)}")
+        end
+
+        # Drop an existing schema
+        def drop_schema(name, options = {})
+          force = options.fetch(:force, '').upcase
+          check = 'IF EXISTS' if options.fetch(:check, true)
+          execute("DROP SCHEMA #{check} #{quote_schema_name(name.to_s)} #{force}")
+        end
+
         # Drops a type.
         def drop_type(name, options = {})
           force = options.fetch(:force, '').upcase
@@ -64,12 +79,37 @@ def enum_values(name)
 
         # Rewrite the method that creates tables to easily accept extra options
         def create_table(table_name, **options, &block)
+          table_name = "#{options[:schema]}.#{table_name}" if options[:schema].present?
+
           options[:id] = false if options[:inherits].present? &&
             options[:primary_key].blank? && options[:id].blank?
 
           super table_name, **options, &block
         end
 
+        # Add the schema option when extracting table options
+        def table_options(table_name)
+          parts = table_name.split('.').reverse
+          return super unless parts.size == 2 && parts[1] != 'public'
+
+          (super || {}).merge(schema: parts[1])
+        end
+
+        # When dumping the schema we need to add all schemas, not only those
+        # active for the current +schema_search_path+
+        def quoted_scope(name = nil, type: nil)
+          return super unless name.nil?
+
+          super.merge(schema: "ANY ('{#{user_defined_schemas.join(',')}}')")
+        end
+
+        # Fix the query to include the schema on tables names when dumping
+        def data_source_sql(name = nil, type: nil)
+          return super unless name.nil?
+
+          super.sub('SELECT c.relname FROM', "SELECT n.nspname || '.' || c.relname FROM")
+        end
+
         private
 
           def quote_enum_values(name, values, options)

lib/torque/postgresql/base.rb

@@ -5,15 +5,27 @@ module PostgreSQL
     module Base
       extend ActiveSupport::Concern
 
+      ##
+      # :singleton-method: schema
+      # :call-seq: schema
+      #
+      # The schema to which the table belongs to.
+
       included do
         mattr_accessor :belongs_to_many_required_by_default, instance_accessor: false
+        class_attribute :schema, instance_writer: false
       end
 
       module ClassMethods
         delegate :distinct_on, :with, :itself_only, :cast_records, to: :all
 
-        # Wenever it's inherited, add a new list of auxiliary statements
-        # It also adds an auxiliary statement to load inherited records' relname
+        # Make sure that table name is an instance of TableName class
+        def reset_table_name
+          self.table_name = TableName.new(self, super)
+        end
+
+        # Whenever the base model is inherited, add a list of auxiliary
+        # statements like the one that loads inherited records' relname
         def inherited(subclass)
           super
 
@@ -22,32 +34,14 @@ def inherited(subclass)
 
           record_class = ActiveRecord::Relation._record_class_attribute
 
-          # Define helper methods to return the class of the given records
-          subclass.auxiliary_statement record_class do |cte|
-            pg_class = ::Arel::Table.new('pg_class')
-            arel_query = ::Arel::SelectManager.new(pg_class)
-            arel_query.project(pg_class['oid'], pg_class['relname'].as(record_class.to_s))
-
-            cte.query 'pg_class', arel_query.to_sql
-            cte.attributes col(record_class) => record_class
-            cte.join tableoid: :oid
-          end
-
           # Define the dynamic attribute that returns the same information as
           # the one provided by the auxiliary statement
           subclass.dynamic_attribute(record_class) do
-            next self.class.table_name unless self.class.physically_inheritances?
-
-            pg_class = ::Arel::Table.new('pg_class')
-            source = ::Arel::Table.new(subclass.table_name, as: 'source')
-            quoted_id = ::Arel::Nodes::Quoted.new(id)
-
-            query = ::Arel::SelectManager.new(pg_class)
-            query.join(source).on(pg_class['oid'].eq(source['tableoid']))
-            query.where(source[subclass.primary_key].eq(quoted_id))
-            query.project(pg_class['relname'])
+            klass = self.class
+            next klass.table_name unless klass.physically_inheritances?
 
-            self.class.connection.select_value(query)
+            query = klass.unscoped.where(subclass.primary_key => id)
+            query.pluck(klass.arel_table['tableoid'].cast('regclass')).first
           end
         end
 

lib/torque/postgresql/config.rb

@@ -39,6 +39,19 @@ def config.irregular_models=(hash)
 
     end
 
+    # Configure multiple schemas
+    config.nested(:schemas) do |schemas|
+
+      # Defines a list of LIKE-based schemas to not consider for a multiple
+      # schema database
+      schemas.blacklist = %w[information_schema pg_%]
+
+      # Defines a list of LIKE-based schemas to consider for a multiple schema
+      # database
+      schemas.whitelist = %w[public]
+
+    end
+
     # Configure auxiliary statement features
     config.nested(:auxiliary_statement) do |cte|
 

lib/torque/postgresql/migration/command_recorder.rb

@@ -5,24 +5,24 @@ module PostgreSQL
     module Migration
       module CommandRecorder
 
-        # Records the rename operation for types.
+        # Records the rename operation for types
         def rename_type(*args, &block)
           record(:rename_type, args, &block)
         end
 
-        # Inverts the type name.
+        # Inverts the type rename operation
         def invert_rename_type(args)
           [:rename_type, args.reverse]
         end
 
-        # Records the creation of the enum to be reverted.
-        def create_enum(*args, &block)
-          record(:create_enum, args, &block)
+        # Records the creation of a schema
+        def create_schema(*args, &block)
+          record(:create_schema, args, &block)
         end
 
-        # Inverts the creation of the enum.
-        def invert_create_enum(args)
-          [:drop_type, [args.first]]
+        # Inverts the creation of a schema
+        def invert_create_schema(args)
+          [:drop_schema, [args.first]]
         end
 
       end

lib/torque/postgresql/schema_cache.rb

@@ -109,7 +109,6 @@ def associations(table_name)
 
       # Try to find a model based on a given table
       def lookup_model(table_name, scoped_class = '')
-        # byebug if table_name == 'activities'
         scoped_class = scoped_class.name if scoped_class.is_a?(Class)
         return @data_sources_model_names[table_name] \
           if @data_sources_model_names.key?(table_name)
@@ -118,6 +117,12 @@ def lookup_model(table_name, scoped_class = '')
         scopes = scoped_class.scan(/(?:::)?[A-Z][a-z]+/)
         scopes.unshift('Object::')
 
+        # Check if the table name comes with a schema
+        if table_name.include?('.')
+          schema, table_name = table_name.split('.')
+          scopes.insert(1, schema.camelize) if schema != 'public'
+        end
+
         # Consider the maximum namespaced possible model name
         max_name = table_name.tr('_', '/').camelize.split(/(::)/)
         max_name[-1] = max_name[-1].singularize