<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)
      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',
            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"] = <PASSWORD>pass
      elsif @password
        env["PGPASSWORD"] = <PASSWORD>
      end
      env
    end

    #
    # Ruby Library Interface
    #
<commit_msg>Pass options to psql command<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 = 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
    #

{'fraction_non_alphanumeric': 0.1093284437313775, 'fraction_numerical': 0.0027504011001604402, 'mean_word_length': 2.2787377911344855, 'pattern_counts': {'":': 0, '<': 11, '<?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': '5193274', 'n_tokens_mistral': 1363, 'n_tokens_neox': 1342, 'n_words': 346}