activerecord7-redshift-adapter-pennylane 1.0.2 → 1.0.3

data/lib/active_record/connection_adapters/redshift_7_1_adapter.rb

@@ -23,28 +23,14 @@ require 'ipaddr'
 ActiveRecord::Tasks::DatabaseTasks.register_task(/redshift/, 'ActiveRecord::Tasks::PostgreSQLDatabaseTasks')
 module ActiveRecord
   module ConnectionHandling # :nodoc:
-    RS_VALID_CONN_PARAMS = %i[host hostaddr port dbname user password connect_timeout
-                              client_encoding options application_name fallback_application_name
-                              keepalives keepalives_idle keepalives_interval keepalives_count
-                              tty sslmode requiressl sslcompression sslcert sslkey
-                              sslrootcert sslcrl requirepeer krbsrvname gsslib service].freeze
+
+    def redshift_adapter_class
+      ConnectionAdapters::RedshiftAdapter
+    end
 
     # Establishes a connection to the database that's used by all Active Record objects
     def redshift_connection(config)
-      conn_params = config.symbolize_keys
-
-      conn_params.delete_if { |_, v| v.nil? }
-
-      # Map ActiveRecords param names to PGs.
-      conn_params[:user] = conn_params.delete(:username) if conn_params[:username]
-      conn_params[:dbname] = conn_params.delete(:database) if conn_params[:database]
-
-      # Forward only valid config params to PG::Connection.connect.
-      conn_params.keep_if { |k, _| RS_VALID_CONN_PARAMS.include?(k) }
-
-      # The postgres drivers don't allow the creation of an unconnected PG::Connection object,
-      # so just pass a nil connection object for the time being.
-      ConnectionAdapters::RedshiftAdapter.new(nil, logger, conn_params, config)
+      redshift_adapter_class.new(config)
     end
   end
 
@@ -78,6 +64,12 @@ module ActiveRecord
     class RedshiftAdapter < AbstractAdapter
       ADAPTER_NAME = 'Redshift'
 
+      RS_VALID_CONN_PARAMS = %i[host hostaddr port dbname user password connect_timeout
+                              client_encoding options application_name fallback_application_name
+                              keepalives keepalives_idle keepalives_interval keepalives_count
+                              tty sslmode requiressl sslcompression sslcert sslkey
+                              sslrootcert sslcrl requirepeer krbsrvname gsslib service].freeze
+
       NATIVE_DATABASE_TYPES = {
         primary_key: 'integer identity primary key',
         string: { name: 'varchar' },
@@ -143,54 +135,63 @@ module ActiveRecord
         end
 
         def next_key
-          "a#{@counter + 1}"
-        end
-
-        def []=(sql, key)
-          super.tap { @counter += 1 }
+          "a#{@counter += 1}"
         end
 
         private
 
         def dealloc(key)
-          @raw_connection.query "DEALLOCATE #{key}" if connection_active?
+          # This is ugly, but safe: the statement pool is only
+          # accessed while holding the connection's lock. (And we
+          # don't need the complication of with_raw_connection because
+          # a reconnect would invalidate the entire statement pool.)
+          if conn = @connection.instance_variable_get(:@raw_connection)
+            conn.query "DEALLOCATE #{key}" if conn.status == PG::CONNECTION_OK
+          end
         rescue PG::Error
         end
+      end
 
-        def connection_active?
-          @raw_connection.status == PG::CONNECTION_OK
-        rescue PG::Error
-          false
+      class << self
+        def new_client(conn_params)
+          PG.connect(**conn_params)
+        rescue ::PG::Error => error
+          if conn_params && conn_params[:dbname] == "postgres"
+            raise ActiveRecord::ConnectionNotEstablished, error.message
+          elsif conn_params && conn_params[:dbname] && error.message.include?(conn_params[:dbname])
+            raise ActiveRecord::NoDatabaseError.db_error(conn_params[:dbname])
+          elsif conn_params && conn_params[:user] && error.message.include?(conn_params[:user])
+            raise ActiveRecord::DatabaseConnectionError.username_error(conn_params[:user])
+          elsif conn_params && conn_params[:host] && error.message.include?(conn_params[:host])
+            raise ActiveRecord::DatabaseConnectionError.hostname_error(conn_params[:host])
+          else
+            raise ActiveRecord::ConnectionNotEstablished, error.message
+          end
         end
       end
 
       # Initializes and connects a PostgreSQL adapter.
-      def initialize(connection, logger, connection_parameters, config)
-        super(connection, logger, config)
+      def initialize(...)
+        super
 
-        @visitor = Arel::Visitors::PostgreSQL.new self
-        @visitor.extend(ConnectionAdapters::DetermineIfPreparableVisitor) if defined?(ConnectionAdapters::DetermineIfPreparableVisitor)
-        @prepared_statements = false
+        conn_params = @config.compact
 
-        @raw_connection_parameters = connection_parameters
+        # Map ActiveRecords param names to PGs.
+        conn_params[:user] = conn_params.delete(:username) if conn_params[:username]
+        conn_params[:dbname] = conn_params.delete(:database) if conn_params[:database]
 
-        # @local_tz is initialized as nil to avoid warnings when connect tries to use it
-        @local_tz = nil
-        @table_alias_length = nil
+        # Forward only valid config params to PG::Connection.connect.
+        conn_params.slice!(*RS_VALID_CONN_PARAMS)
 
-        connect
-        @statements = StatementPool.new @raw_connection,
-                                        self.class.type_cast_config_to_integer(config[:statement_limit])
+        @connection_parameters = conn_params
 
-        @type_map = Type::HashLookupTypeMap.new
-        initialize_type_map(type_map)
-        @local_tz = execute('SHOW TIME ZONE', 'SCHEMA').first['TimeZone']
-        @use_insert_returning = @config.key?(:insert_returning) ? self.class.type_cast_config_to_boolean(@config[:insert_returning]) : false
-      end
+        @max_identifier_length = nil
+        @type_map = nil
+        @raw_connection = nil
+        @notice_receiver_sql_warnings = []
+        @table_alias_length = nil
 
-      # Clears the prepared statements cache.
-      def clear_cache!(new_connection: false)
-        @statements.clear
+        @use_insert_returning = @config.key?(:insert_returning) ? self.class.type_cast_config_to_boolean(@config[:insert_returning]) : true
       end
 
       def truncate(table_name, name = nil)
@@ -199,44 +200,55 @@ module ActiveRecord
 
       # Is this connection alive and ready for queries?
       def active?
-        @raw_connection.query 'SELECT 1'
+        @lock.synchronize do
+          return false unless @raw_connection
+          @raw_connection.query ";"
+        end
         true
       rescue PG::Error
         false
       end
 
-      def reload_type_map
-        type_map.clear
-        initialize_type_map
-      end
+      def reload_type_map # :nodoc:
+        @lock.synchronize do
+          if @type_map
+            type_map.clear
+          else
+            @type_map = Type::HashLookupTypeMap.new
+          end
 
-      # Close then reopen the connection.
-      def reconnect!
-        super
-        @raw_connection.reset
-        configure_connection
-        reload_type_map
+          initialize_type_map
+        end
       end
 
       def reset!
-        clear_cache!
-        reset_transaction
-        @raw_connection.query 'ROLLBACK' unless @raw_connection.transaction_status == ::PG::PQTRANS_IDLE
-        @raw_connection.query 'DISCARD ALL'
-        configure_connection
+        @lock.synchronize do
+          return connect! unless @raw_connection
+
+          unless @raw_connection.transaction_status == ::PG::PQTRANS_IDLE
+            @raw_connection.query "ROLLBACK"
+          end
+
+          super
+        end
       end
 
       # Disconnects from the database if already connected. Otherwise, this
       # method does nothing.
       def disconnect!
-        super
-        begin
-          @raw_connection.close
-        rescue StandardError
-          nil
+        @lock.synchronize do
+          super
+          @raw_connection&.close rescue nil
+          @raw_connection = nil
         end
       end
 
+      def discard! # :nodoc:
+        super
+        @raw_connection&.socket_io&.reopen(IO::NULL) rescue nil
+        @raw_connection = nil
+      end
+
       def native_database_types # :nodoc:
         NATIVE_DATABASE_TYPES
       end
@@ -275,14 +287,22 @@ module ActiveRecord
         true
       end
 
-      def enable_extension(name); end
+      def enable_extension(name)
+        ;
+      end
 
-      def disable_extension(name); end
+      def disable_extension(name)
+        ;
+      end
 
       def extension_enabled?(_name)
         false
       end
 
+      def supports_common_table_expressions?
+        true
+      end
+
       # Returns the configured supported identifier length supported by PostgreSQL
       def table_alias_length
         @table_alias_length ||= query('SHOW max_identifier_length', 'SCHEMA')[0][0].to_i
@@ -302,23 +322,26 @@ module ActiveRecord
         !native_database_types[type].nil?
       end
 
-      def update_table_definition(table_name, base) # :nodoc:
+      def update_table_definition(table_name, base)
+        # :nodoc:
         Redshift::Table.new(table_name, base)
       end
 
-      def lookup_cast_type(sql_type) # :nodoc:
+      def lookup_cast_type(sql_type)
+        # :nodoc:
         oid = execute("SELECT #{quote(sql_type)}::regtype::oid", 'SCHEMA').first['oid'].to_i
         super(oid)
       end
 
-      def column_name_for_operation(operation, _node) # :nodoc:
+      def column_name_for_operation(operation, _node)
+        # :nodoc:
         OPERATION_ALIASES.fetch(operation) { operation.downcase }
       end
 
       OPERATION_ALIASES = { # :nodoc:
-        'maximum' => 'max',
-        'minimum' => 'min',
-        'average' => 'avg'
+                            'maximum' => 'max',
+                            'minimum' => 'min',
+                            'average' => 'avg'
       }.freeze
 
       protected
@@ -342,7 +365,8 @@ module ActiveRecord
       end
 
       class << self
-        def initialize_type_map(m) # :nodoc:
+        def initialize_type_map(m)
+          # :nodoc:
           m.register_type 'int2', Type::Integer.new(limit: 2)
           m.register_type 'int4', Type::Integer.new(limit: 4)
           m.register_type 'int8', Type::Integer.new(limit: 8)
@@ -388,7 +412,8 @@ module ActiveRecord
 
       private
 
-      def get_oid_type(oid, fmod, column_name, sql_type = '') # :nodoc:
+      def get_oid_type(oid, fmod, column_name, sql_type = '')
+        # :nodoc:
         load_additional_types(type_map, [oid]) unless type_map.key?(oid)
 
         type_map.fetch(oid, fmod, sql_type) do
@@ -408,7 +433,8 @@ module ActiveRecord
         load_additional_types(m)
       end
 
-      def extract_limit(sql_type) # :nodoc:
+      def extract_limit(sql_type)
+        # :nodoc:
         case sql_type
         when /^bigint/i, /^int8/i
           8
@@ -420,7 +446,8 @@ module ActiveRecord
       end
 
       # Extracts the value from a PostgreSQL column default definition.
-      def extract_value_from_default(default) # :nodoc:
+      def extract_value_from_default(default)
+        # :nodoc:
         case default
           # Quoted types
         when /\A[(B]?'(.*)'::/m
@@ -441,19 +468,22 @@ module ActiveRecord
         end
       end
 
-      def extract_default_function(default_value, default) # :nodoc:
+      def extract_default_function(default_value, default)
+        # :nodoc:
         default if has_default_function?(default_value, default)
       end
 
-      def has_default_function?(default_value, default) # :nodoc:
+      def has_default_function?(default_value, default)
+        # :nodoc:
         !default_value && (/\w+\(.*\)/ === default)
       end
 
-      def load_additional_types(type_map, oids = nil) # :nodoc:
+      def load_additional_types(type_map, oids = nil)
+        # :nodoc:
         initializer = OID::TypeMapInitializer.new(type_map)
 
         load_types_queries(initializer, oids) do |query|
-          execute_and_clear(query, 'SCHEMA', []) do |records|
+          execute_and_clear(query, 'SCHEMA', [], allow_retry: true, materialize_transactions: false) do |records|
             initializer.run(records)
           end
         end
@@ -483,58 +513,69 @@ module ActiveRecord
 
       FEATURE_NOT_SUPPORTED = '0A000' # :nodoc:
 
-      def execute_and_clear(sql, name, binds, prepare: false, async: false)
-        result =
-          if without_prepared_statement?(binds)
-            exec_no_cache(sql, name, [])
-          elsif !prepare
-            exec_no_cache(sql, name, binds)
-          else
-            exec_cache(sql, name, binds)
-          end
+      def execute_and_clear(sql, name, binds, prepare: false, async: false, allow_retry: false, materialize_transactions: true)
+        sql = transform_query(sql)
+        check_if_write_query(sql)
 
-        ret = yield result
-        result.clear
+        if !prepare || without_prepared_statement?(binds)
+          result = exec_no_cache(sql, name, binds, async: async, allow_retry: allow_retry, materialize_transactions: materialize_transactions)
+        else
+          result = exec_cache(sql, name, binds, async: async, allow_retry: allow_retry, materialize_transactions: materialize_transactions)
+        end
+        begin
+          ret = yield result
+        ensure
+          result.clear
+        end
         ret
       end
 
-      def exec_no_cache(sql, name, binds)
-        materialize_transactions
+      def exec_no_cache(sql, name, binds, async:, allow_retry:, materialize_transactions:)
+        mark_transaction_written_if_write(sql)
 
         # make sure we carry over any changes to ActiveRecord.default_timezone that have been
         # made since we established the connection
         update_typemap_for_default_timezone
 
         type_casted_binds = type_casted_binds(binds)
-        log(sql, name, binds, type_casted_binds) do
-          ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
-            @raw_connection.exec_params(sql, type_casted_binds)
+        log(sql, name, binds, type_casted_binds, async: async) do
+          with_raw_connection do |conn|
+            result = conn.exec_params(sql, type_casted_binds)
+            verified!
+            result
           end
         end
       end
 
-      def exec_cache(sql, name, binds)
-        materialize_transactions
+      def exec_cache(sql, name, binds, async:, allow_retry:, materialize_transactions:)
+        mark_transaction_written_if_write(sql)
+
         update_typemap_for_default_timezone
 
-        stmt_key = prepare_statement(sql, binds)
-        type_casted_binds = type_casted_binds(binds)
+        with_raw_connection do |conn|
+          stmt_key = prepare_statement(sql, binds, conn)
+          type_casted_binds = type_casted_binds(binds)
 
-        log(sql, name, binds, type_casted_binds, stmt_key) do
-          ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
-            @raw_connection.exec_prepared(stmt_key, type_casted_binds)
+          log(sql, name, binds, type_casted_binds, stmt_key, async: async) do
+            result = conn.exec_prepared(stmt_key, type_casted_binds)
+            verified!
+            result
           end
         end
       rescue ActiveRecord::StatementInvalid => e
         raise unless is_cached_plan_failure?(e)
-        raise ActiveRecord::PreparedStatementCacheExpired, e.cause.message if in_transaction?
 
-        @lock.synchronize do
-          # outside of transactions we can simply flush this query and retry
-          @statements.delete sql_key(sql)
+        # Nothing we can do if we are in a transaction because all commands
+        # will raise InFailedSQLTransaction
+        if in_transaction?
+          raise ActiveRecord::PreparedStatementCacheExpired.new(e.cause.message)
+        else
+          @lock.synchronize do
+            # outside of transactions we can simply flush this query and retry
+            @statements.delete sql_key(sql)
+          end
+          retry
         end
-
-        retry
       end
 
       # Annoyingly, the code for prepared statements whose return value may
@@ -547,6 +588,7 @@ module ActiveRecord
       # Check here for more details:
       # https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/cache/plancache.c#l573
       CACHED_PLAN_HEURISTIC = 'cached plan must not change result type'
+
       def is_cached_plan_failure?(e)
         pgerror = e.cause
         code = pgerror.result.result_error_field(PG::PG_DIAG_SQLSTATE)
@@ -563,51 +605,60 @@ module ActiveRecord
 
       # Prepare the statement if it hasn't been prepared, return
       # the statement key.
-      def prepare_statement(sql, binds)
-        @lock.synchronize do
-          sql_key = sql_key(sql)
-          unless @statements.key? sql_key
-            nextkey = @statements.next_key
-            begin
-              @raw_connection.prepare nextkey, sql
-            rescue StandardError => e
-              raise translate_exception_class(e, sql, binds)
-            end
-            # Clear the queue
-            @raw_connection.get_last_result
-            @statements[sql_key] = nextkey
+      def prepare_statement(sql, binds, conn)
+        sql_key = sql_key(sql)
+        unless @statements.key? sql_key
+          nextkey = @statements.next_key
+          begin
+            conn.prepare nextkey, sql
+          rescue => e
+            raise translate_exception_class(e, sql, binds)
           end
-          @statements[sql_key]
+          # Clear the queue
+          conn.get_last_result
+          @statements[sql_key] = nextkey
         end
+        @statements[sql_key]
       end
 
       # Connects to a PostgreSQL server and sets up the adapter depending on the
       # connected server's characteristics.
       def connect
-        @raw_connection = PG.connect(@raw_connection_parameters)
-        configure_connection
-        add_pg_encoders
-        add_pg_decoders
+        @raw_connection = self.class.new_client(@connection_parameters)
+      rescue ConnectionNotEstablished => ex
+        raise ex.set_pool(@pool)
+      end
+
+      def reconnect
+        begin
+          @raw_connection&.reset
+        rescue PG::ConnectionBad
+          @raw_connection = nil
+        end
+
+        connect unless @raw_connection
+      end
+
+      def reconnect
+        begin
+          @raw_connection&.reset
+        rescue PG::ConnectionBad
+          @raw_connection = nil
+        end
+
+        connect unless @raw_connection
       end
 
       # Configures the encoding, verbosity, schema search path, and time zone of the connection.
       # This is called by #connect and should not be called manually.
       def configure_connection
-        @raw_connection.set_client_encoding(@config[:encoding]) if @config[:encoding]
+        if @config[:encoding]
+          @raw_connection.set_client_encoding(@config[:encoding])
+        end
         self.schema_search_path = @config[:schema_search_path] || @config[:schema_order]
 
         variables = @config.fetch(:variables, {}).stringify_keys
 
-        # If using Active Record's time zone support configure the connection to return
-        # TIMESTAMP WITH ZONE types in UTC.
-        unless variables['timezone']
-          if ActiveRecord.default_timezone == :utc
-            variables['timezone'] = 'UTC'
-          elsif @local_tz
-            variables['timezone'] = @local_tz
-          end
-        end
-
         # SET statements from :variables config hash
         # https://www.postgresql.org/docs/current/static/sql-set.html
         variables.map do |k, v|
@@ -618,9 +669,32 @@ module ActiveRecord
             execute("SET #{k} TO #{quote(v)}", 'SCHEMA')
           end
         end
+
+        add_pg_encoders
+        add_pg_decoders
+
+        reload_type_map
+      end
+
+      def reconfigure_connection_timezone
+        variables = @config.fetch(:variables, {}).stringify_keys
+
+        # If it's been directly configured as a connection variable, we don't
+        # need to do anything here; it will be set up by configure_connection
+        # and then never changed.
+        return if variables["timezone"]
+
+        # If using Active Record's time zone support configure the connection
+        # to return TIMESTAMP WITH ZONE types in UTC.
+        if default_timezone == :utc
+          internal_execute("SET timezone TO 'UTC'")
+        else
+          internal_execute("SET timezone TO DEFAULT")
+        end
       end
 
-      def last_insert_id_result(sequence_name) # :nodoc:
+      def last_insert_id_result(sequence_name)
+        # :nodoc:
         exec_query("SELECT currval('#{sequence_name}')", 'SQL')
       end
 
@@ -642,7 +716,8 @@ module ActiveRecord
       # Query implementation notes:
       #  - format_type includes the column size constraint, e.g. varchar(50)
       #  - ::regclass is a function that gives the id for a table name
-      def column_definitions(table_name) # :nodoc:
+      def column_definitions(table_name)
+        # :nodoc:
         query(<<-END_SQL, 'SCHEMA')
               SELECT a.attname, format_type(a.atttypid, a.atttypmod),
                      pg_get_expr(d.adbin, d.adrelid), a.attnotnull, a.atttypid, a.atttypmod
@@ -670,23 +745,23 @@ module ActiveRecord
       def can_perform_case_insensitive_comparison_for?(column)
         @case_insensitive_cache ||= {}
         @case_insensitive_cache[column.sql_type] ||= begin
-          sql = <<~SQL
-            SELECT exists(
-              SELECT * FROM pg_proc
-              WHERE proname = 'lower'
-                AND proargtypes = ARRAY[#{quote column.sql_type}::regtype]::oidvector
-            ) OR exists(
-              SELECT * FROM pg_proc
-              INNER JOIN pg_cast
-                ON ARRAY[casttarget]::oidvector = proargtypes
-              WHERE proname = 'lower'
-                AND castsource = #{quote column.sql_type}::regtype
-            )
-          SQL
-          execute_and_clear(sql, 'SCHEMA', []) do |result|
-            result.getvalue(0, 0)
-          end
-        end
+                                                       sql = <<~SQL
+                                                         SELECT exists(
+                                                           SELECT * FROM pg_proc
+                                                           WHERE proname = 'lower'
+                                                             AND proargtypes = ARRAY[#{quote column.sql_type}::regtype]::oidvector
+                                                         ) OR exists(
+                                                           SELECT * FROM pg_proc
+                                                           INNER JOIN pg_cast
+                                                             ON ARRAY[casttarget]::oidvector = proargtypes
+                                                           WHERE proname = 'lower'
+                                                             AND castsource = #{quote column.sql_type}::regtype
+                                                         )
+                                                       SQL
+                                                       execute_and_clear(sql, 'SCHEMA', []) do |result|
+                                                         result.getvalue(0, 0)
+                                                       end
+                                                     end
       end
 
       def add_pg_encoders
@@ -698,22 +773,22 @@ module ActiveRecord
       end
 
       def update_typemap_for_default_timezone
-        return if @default_timezone == ActiveRecord.default_timezone || !@timestamp_decoder
+        if @raw_connection && @mapped_default_timezone != default_timezone && @timestamp_decoder
+          decoder_class = default_timezone == :utc ?
+                            PG::TextDecoder::TimestampUtc :
+                            PG::TextDecoder::TimestampWithoutTimeZone
 
-        decoder_class =
-          if ActiveRecord.default_timezone == :utc
-            PG::TextDecoder::TimestampUtc
-          else
-            PG::TextDecoder::TimestampWithoutTimeZone
-          end
+          @timestamp_decoder = decoder_class.new(**@timestamp_decoder.to_h)
+          @raw_connection.type_map_for_results.add_coder(@timestamp_decoder)
 
-        @timestamp_decoder = decoder_class.new(@timestamp_decoder.to_h)
-        @raw_connection.type_map_for_results.add_coder(@timestamp_decoder)
-        @default_timezone = ActiveRecord.default_timezone
+          @mapped_default_timezone = default_timezone
 
-        # if default timezone has changed, we need to reconfigure the connection
-        # (specifically, the session time zone)
-        configure_connection
+          # if default timezone has changed, we need to reconfigure the connection
+          # (specifically, the session time zone)
+          reconfigure_connection_timezone
+
+          true
+        end
       end
 
       def add_pg_decoders
@@ -742,7 +817,7 @@ module ActiveRecord
           FROM pg_type as t
           WHERE t.typname IN (%s)
         SQL
-        coders = execute_and_clear(query, 'SCHEMA', []) do |result|
+        coders = execute_and_clear(query, 'SCHEMA', [], allow_retry: true, materialize_transactions: false) do |result|
           result.filter_map { |row| construct_coder(row, coders_by_name[row['typname']]) }
         end
 
@@ -750,6 +825,9 @@ module ActiveRecord
         coders.each { |coder| map.add_coder(coder) }
         @raw_connection.type_map_for_results = map
 
+        @type_map_for_results = PG::TypeMapByOid.new
+        @type_map_for_results.default_type_map = map
+
         # extract timestamp decoder for use in update_typemap_for_default_timezone
         @timestamp_decoder = coders.find { |coder| coder.name == 'timestamp' }
         update_typemap_for_default_timezone
@@ -761,7 +839,8 @@ module ActiveRecord
         coder_class.new(oid: row['oid'].to_i, name: row['typname'])
       end
 
-      def create_table_definition(*args, **options) # :nodoc:
+      def create_table_definition(*args, **options)
+        # :nodoc:
         Redshift::TableDefinition.new(self, *args, **options)
       end
     end