<commit_before>require 'bricolage/datasource'
require 'bricolage/s3datasource'
require 'bricolage/sqlstatement'
require 'bricolage/commandutils'
require 'bricolage/postgresconnection'
require 'bricolage/vacuumlock'
require 'bricolage/exception'
require 'pathname'

module Bricolage

  class PSQLDataSource < DataSource
    declare_type 'psql'

    include CommandUtils

    def initialize(
        host: 'localhost',
        port: 5439,
        database: 'dev',
        username: ENV['LOGNAME'],
        password: <PASSWORD>,
        pgpass: nil,
        encoding: nil,
        psql: 'psql',
        tmpdir: Dir.tmpdir)
      @host = host
      @port = port
      @database = database
      @user = username
      @password = password
      @pgpass = pgpass
      @encoding = encoding
      @psql = psql
      @tmpdir = tmpdir
      raise ParameterError, "missing psql host" unless @host
      raise ParameterError, "missing psql port" unless @port
      raise ParameterError, "missing psql database" unless @database
      raise ParameterError, "missing psql username" unless @user
      unless @pgpass or @password
        raise ParameterError, "missing psql password"
      end
    end

    attr_reader :host
    attr_reader :port
    attr_reader :database
    attr_reader :user

    def new_task
      PSQLTask.new(self)
    end

    def open_for_batch
      # do not call #open
      yield
    end

    def execute(source, options = [])
      make_tmpfile(source, tmpdir: @tmpdir) {|path|
        st = command @psql, "--no-psqlrc", "--host=#{@host}", "--port=#{@port}",
            "--username=#{@user}", @database,
            '--echo-all',
            '-v', 'ON_ERROR_STOP=true',
            '-f', path,
            '--no-password',
            *options,
            env: get_psql_env
        msg = retrieve_last_match_from_stderr(/^psql:.*?:\d+: ERROR: (.*)/, 1) unless st.success?
        JobResult.for_process_status(st, msg)
      }
    end

    def get_psql_env
      env = {}
      if @pgpass
        env["PGPASSFILE"] = @pgpass
      elsif @password
        env["PGPASSWORD"] = <PASSWORD>
      end
      env
    end

    #
    # Ruby Library Interface
    #

    def password
      # FIXME: same user must not exist
      @password ||= read_password_from_pgpass(@pgpass, @user)
    end

    def read_password_from_pgpass(path, user)
      File.read(path).slice(/:#{user}:([^:\r\n]+)$/, 1) or
          raise ParameterError, "could not read password: #{path}, #{user}"
    end

    def open(&block)
      conn = PG::Connection.open(host: @host, port: @port, dbname: @database, user: @user, password: password)
      yield PostgresConnection.new(conn, self, logger)
    ensure
      conn.close if conn
    end

    def query_batch(query, batch_size = 5000, &block)
      open {|conn| conn.query_batch(query, batch_size, &block) }
    end

    def drop_table(name)
      open {|conn| conn.drop_table(name) }
    end

    def drop_table_force(name)
      open {|conn| conn.drop_table_force(name) }
    end

    def select(table, &block)
      open {|conn| conn.select(table, &block) }
    end

    include VacuumLock

    def vacuum(table)
      serialize_vacuum {
        open {|conn| conn.vacuum(table) }
      }
    end

    def vacuum_sort_only(table)
      serialize_vacuum {
        open {|conn| conn.vacuum_sort_only(table) }
      }
    end

    def analyze(table)
      open {|conn| conn.analyze(table) }
    end
  end

  class DataSource   # reopen
    def redshift_loader_source?
      false
    end
  end

  class S3DataSource   # reopen
    def redshift_loader_source?
      true
    end
  end

  # We don't support dynamodb still now

  class PSQLTask < DataSourceTask
    def exec(stmt)
      add SQLAction.new(stmt)
    end

    def each_statement
      each_action do |action|
        yield action.statement
      end
    end

    class SQLAction < Action
      def initialize(stmt)
        @statement = stmt
      end

      attr_reader :statement

      def_delegator '@statement', :bind
    end

    # override
    def source
      buf = StringIO.new
      buf.puts '\timing on'
      each_statement do |stmt|
        buf.puts
        buf.puts "-- #{stmt.location}" if stmt.location
        buf.puts stmt.stripped_source
      end
      buf.string
    end

    # override
    def run
      VacuumLock.using {
        @ds.execute source
      }
    end

    def run_explain
      @ds.execute explain_source
    end

    def explain_source
      buf = StringIO.new
      each_statement do |stmt|
        buf.puts
        buf.puts "-- #{stmt.location}" if stmt.location<commit_msg>Insert location as a block comment<commit_after>require 'bricolage/datasource'
require 'bricolage/s3datasource'
require 'bricolage/sqlstatement'
require 'bricolage/commandutils'
require 'bricolage/postgresconnection'
require 'bricolage/vacuumlock'
require 'bricolage/exception'
require 'pathname'

module Bricolage

  class PSQLDataSource < DataSource
    declare_type 'psql'

    include CommandUtils

    def initialize(
        host: 'localhost',
        port: 5439,
        database: 'dev',
        username: ENV['LOGNAME'],
        password: <PASSWORD>,
        pgpass: <PASSWORD>,
        encoding: nil,
        psql: 'psql',
        tmpdir: Dir.tmpdir)
      @host = host
      @port = port
      @database = database
      @user = username
      @password = password
      @pgpass = <PASSWORD>
      @encoding = encoding
      @psql = psql
      @tmpdir = tmpdir
      raise ParameterError, "missing psql host" unless @host
      raise ParameterError, "missing psql port" unless @port
      raise ParameterError, "missing psql database" unless @database
      raise ParameterError, "missing psql username" unless @user
      unless @pgpass or @password
        raise ParameterError, "missing psql password"
      end
    end

    attr_reader :host
    attr_reader :port
    attr_reader :database
    attr_reader :user

    def new_task
      PSQLTask.new(self)
    end

    def open_for_batch
      # do not call #open
      yield
    end

    def execute(source, options = [])
      make_tmpfile(source, tmpdir: @tmpdir) {|path|
        st = command @psql, "--no-psqlrc", "--host=#{@host}", "--port=#{@port}",
            "--username=#{@user}", @database,
            '--echo-all',
            '-v', 'ON_ERROR_STOP=true',
            '-f', path,
            '--no-password',
            *options,
            env: get_psql_env
        msg = retrieve_last_match_from_stderr(/^psql:.*?:\d+: ERROR: (.*)/, 1) unless st.success?
        JobResult.for_process_status(st, msg)
      }
    end

    def get_psql_env
      env = {}
      if @pgpass
        env["PGPASSFILE"] = @pgpass
      elsif @password
        env["PGPASSWORD"] = @password
      end
      env
    end

    #
    # Ruby Library Interface
    #

    def password
      # FIXME: same user must not exist
      @password ||= read_password_from_pgpass(@pgpass, @<PASSWORD>)
    end

    def read_password_from_pgpass(path, user)
      File.read(path).slice(/:#{user}:([^:\r\n]+)$/, 1) or
          raise ParameterError, "could not read password: #{path}, #{user}"
    end

    def open(&block)
      conn = PG::Connection.open(host: @host, port: @port, dbname: @database, user: @user, password: password)
      yield PostgresConnection.new(conn, self, logger)
    ensure
      conn.close if conn
    end

    def query_batch(query, batch_size = 5000, &block)
      open {|conn| conn.query_batch(query, batch_size, &block) }
    end

    def drop_table(name)
      open {|conn| conn.drop_table(name) }
    end

    def drop_table_force(name)
      open {|conn| conn.drop_table_force(name) }
    end

    def select(table, &block)
      open {|conn| conn.select(table, &block) }
    end

    include VacuumLock

    def vacuum(table)
      serialize_vacuum {
        open {|conn| conn.vacuum(table) }
      }
    end

    def vacuum_sort_only(table)
      serialize_vacuum {
        open {|conn| conn.vacuum_sort_only(table) }
      }
    end

    def analyze(table)
      open {|conn| conn.analyze(table) }
    end
  end

  class DataSource   # reopen
    def redshift_loader_source?
      false
    end
  end

  class S3DataSource   # reopen
    def redshift_loader_source?
      true
    end
  end

  # We don't support dynamodb still now

  class PSQLTask < DataSourceTask
    def exec(stmt)
      add SQLAction.new(stmt)
    end

    def each_statement
      each_action do |action|
        yield action.statement
      end
    end

    class SQLAction < Action
      def initialize(stmt)
        @statement = stmt
      end

      attr_reader :statement

      def_delegator '@statement', :bind
    end

    # override
    def source
      buf = StringIO.new
      buf.puts '\timing on'
      each_statement do |stmt|
        buf.puts
        buf.puts "/* #{stmt.location} */" if stmt.location
        buf.puts stmt.stripped_source
      end
      buf.string
    end

    # override
    def run
      VacuumLock.using {
        @ds.execute source
      }
    end

    def run_explain
      @ds.execute explain_source
    end

    def explain_source
      buf = StringIO.new
      each_statement do |stmt|
        buf.puts
        buf.puts "-- #{stmt.location}" if stmt.location
{'fraction_non_alphanumeric': 0.10164182328796717, 'fraction_numerical': 0.002592352559948153, 'mean_word_length': 2.209358752166378, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 9, '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': '2245818', 'n_tokens_mistral': 2977, 'n_tokens_neox': 2886, 'n_words': 750}