<commit_before>require 'bricolage/exception'
require 'pg'

module Bricolage

  class PostgreSQLException < SQLException; end

  class PostgresConnection
    def initialize(connection, ds, logger)
      @connection = connection
      @ds = ds
      @logger = logger
    end

    def source
      @connection
    end

    def execute_update(query)
      @logger.info "[#{@ds.name}] #{query}"
      log_elapsed_time {
        rs = @connection.exec(query)
        result = rs.to_a
        rs.clear
        result
      }
    rescue PG::Error => ex
      raise PostgreSQLException.wrap(ex)
    end

    alias execute execute_update
    alias update execute_update

    def select(table, &block)
      execute_query("select * from #{table}", &block)
    end

    def execute_query(query, &block)
      @logger.info "[#{@ds.name}] #{query}"
      rs = log_elapsed_time {
        @connection.exec(query)
      }
      result = yield rs
      rs.clear
      result
    rescue PG::Error => ex
      raise PostgreSQLException.wrap(ex)
    end

    alias query execute_query

    def query_batch(query, batch_size = 5000, &block)
      open_cursor(query) {|cur|
        cur.each_result_set(batch_size, &block)
      }
    end

    def streaming_execute_query(query, &block)
      @logger.info "[#{@ds.name}] #{query}"
      log_elapsed_time {
        @connection.send_query(query)
      }
      @connection.set_single_row_mode
      while rs = @connection.get_result
        begin
          rs.check
          yield rs
        ensure
          rs.clear
        end
      end
    rescue PG::Error => ex
      raise PostgreSQLException.wrap(ex)
    end

    def in_transaction?
      @connection.transaction_status == PG::Constants::PQTRANS_INTRANS
    end

    def transaction
      execute 'begin transaction'
      yield
      execute 'commit'
    end

    def open_cursor(query, name = nil, &block)
      unless in_transaction?
        transaction {
          return open_cursor(query, &block)
        }
      end
      name ||= make_unique_cursor_name
      execute "declare #{name} cursor for #{query}"
      yield Cursor.new(name, self, @logger)
    end

    Thread.current['bricolage_cursor_seq'] = 0

    def make_unique_cursor_name
      seq = (Thread.current['bricolage_cursor_seq'] += 1)
      "cur_bric_#{$$}_#{'%X' % Thread.current.object_id}_#{seq}"
    end

    class Cursor
      def initialize(name, conn, logger)
        @name = name
        @conn = conn
        @logger = logger
      end

      attr_reader :name

      def each_result_set(fetch_size = 5000)
        while true
          @conn.execute_query("fetch #{fetch_size} in #{@name}") {|rs|
            return if rs.values.empty?
            yield rs
          }
        end
      end
    end

    def drop_table(name)
      execute "drop table #{name} cascade;"
    end

    def drop_table_force(name)
      drop_table name
    rescue PostgreSQLException => err
      @logger.error err.message<commit_msg>Use DROP TABLE IF EXISTS for force drop<commit_after>require 'bricolage/exception'
require 'pg'

module Bricolage

  class PostgreSQLException < SQLException; end

  class PostgresConnection
    def initialize(connection, ds, logger)
      @connection = connection
      @ds = ds
      @logger = logger
    end

    def source
      @connection
    end

    def execute_update(query)
      @logger.info "[#{@ds.name}] #{query}"
      log_elapsed_time {
        rs = @connection.exec(query)
        result = rs.to_a
        rs.clear
        result
      }
    rescue PG::Error => ex
      raise PostgreSQLException.wrap(ex)
    end

    alias execute execute_update
    alias update execute_update

    def select(table, &block)
      execute_query("select * from #{table}", &block)
    end

    def execute_query(query, &block)
      @logger.info "[#{@ds.name}] #{query}"
      rs = log_elapsed_time {
        @connection.exec(query)
      }
      result = yield rs
      rs.clear
      result
    rescue PG::Error => ex
      raise PostgreSQLException.wrap(ex)
    end

    alias query execute_query

    def query_batch(query, batch_size = 5000, &block)
      open_cursor(query) {|cur|
        cur.each_result_set(batch_size, &block)
      }
    end

    def streaming_execute_query(query, &block)
      @logger.info "[#{@ds.name}] #{query}"
      log_elapsed_time {
        @connection.send_query(query)
      }
      @connection.set_single_row_mode
      while rs = @connection.get_result
        begin
          rs.check
          yield rs
        ensure
          rs.clear
        end
      end
    rescue PG::Error => ex
      raise PostgreSQLException.wrap(ex)
    end

    def in_transaction?
      @connection.transaction_status == PG::Constants::PQTRANS_INTRANS
    end

    def transaction
      execute 'begin transaction'
      yield
      execute 'commit'
    end

    def open_cursor(query, name = nil, &block)
      unless in_transaction?
        transaction {
          return open_cursor(query, &block)
        }
      end
      name ||= make_unique_cursor_name
      execute "declare #{name} cursor for #{query}"
      yield Cursor.new(name, self, @logger)
    end

    Thread.current['bricolage_cursor_seq'] = 0

    def make_unique_cursor_name
      seq = (Thread.current['bricolage_cursor_seq'] += 1)
      "cur_bric_#{$$}_#{'%X' % Thread.current.object_id}_#{seq}"
    end

    class Cursor
      def initialize(name, conn, logger)
        @name = name
        @conn = conn
        @logger = logger
      end

      attr_reader :name

      def each_result_set(fetch_size = 5000)
        while true
          @conn.execute_query("fetch #{fetch_size} in #{@name}") {|rs|
            return if rs.values.empty?
            yield rs
          }
        end
      end
    end

    def drop_table(name)
      execute "drop table #{name} cascade;"
    end

    def drop_table_force(name)
      execute "drop table if exists #{name} cascade;"
{'fraction_non_alphanumeric': 0.09595278246205734, 'fraction_numerical': 0.003372681281618887, 'mean_word_length': 2.1869962385814077, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4920340', 'n_tokens_mistral': 1828, 'n_tokens_neox': 1783, 'n_words': 440}