<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals
import logging
import os

import jpype
from future.utils import iteritems

from pyathenajdbc import (ATHENA_DRIVER_CLASS_NAME,
                          ATHENA_CONNECTION_STRING,
                          ATHENA_JAR)
from pyathenajdbc.converter import JDBCTypeConverter
from pyathenajdbc.cursor import Cursor
from pyathenajdbc.error import ProgrammingError, NotSupportedError
from pyathenajdbc.formatter import ParameterFormatter


_logger = logging.getLogger(__name__)


class Connection(object):

    _ENV_S3_STAGING_DIR = 'AWS_ATHENA_S3_STAGING_DIR'

    def __init__(self, s3_staging_dir=None, access_key=None, secret_key=None,
                 region_name=None, profile_name=None, credential_file=None,
                 jvm_options=None, converter=None, formatter=None, jvm_path=None,
                 **driver_args):
        if s3_staging_dir:
            self.s3_staging_dir = s3_staging_dir
        else:
            self.s3_staging_dir = os.getenv(self._ENV_S3_STAGING_DIR, None)
        assert self.s3_staging_dir, 'Required argument `s3_staging_dir` not found.'

        if credential_file:
            self.access_key = None
            self.secret_key = None
            self.token = None
            self.credential_file = credential_file
            assert self.credential_file, 'Required argument `credential_file` not found.'
            self.region_name = region_name
            assert self.region_name, 'Required argument `region_name` not found.'
        else:
            import botocore.session
            session = botocore.session.get_session()
            if access_key and secret_key:
                session.set_credentials(access_key, secret_key)
            if profile_name:
                session.set_config_variable('profile', profile_name)
            if region_name:
                session.set_config_variable('region', region_name)
            credentials = session.get_credentials()
            self.access_key = credentials.access_key
            assert self.access_key, 'Required argument `access_key` not found.'
            self.secret_key = credentials.secret_key
            assert self.secret_key, 'Required argument `secret_key` not found.'
            self.token = credentials.token
            self.credential_file = None
            self.region_name = session.get_config_variable('region')
            assert self.region_name, 'Required argument `region_name` not found.'

        self._start_jvm(jvm_options, jvm_path)

        props = self._build_driver_args(**driver_args)
        jpype.JClass(ATHENA_DRIVER_CLASS_NAME)
        self._jdbc_conn = jpype.java.sql.DriverManager.getConnection(
            ATHENA_CONNECTION_STRING.format(region=self.region_name), props)

        self._converter = converter if converter else JDBCTypeConverter()
        self._formatter = formatter if formatter else ParameterFormatter()

    @classmethod
    def _start_jvm(cls, options, jvm_path):
        if jvm_path is None:
            jvm_path = jpype.get_default_jvm_path()
        if not jpype.isJVMStarted():
            _logger.debug('JVM path: %s', jvm_path)
            args = ['-server', '-Djava.class.path={0}'.format(
                os.path.join(os.path.dirname(__file__), ATHENA_JAR))]
            if options:
                args.extend(options)
            _logger.debug('JVM args: %s', args)
            jpype.startJVM(jvm_path, *args)
        if not jpype.isThreadAttachedToJVM():
            jpype.attachThreadToJVM()

    def _build_driver_args(self, **kwargs):
        props = jpype.java.util.Properties()
        if self.credential_file:
            props.setProperty('aws_credentials_provider_class',
                              'com.amazonaws.auth.PropertiesFileCredentialsProvider')
            props.setProperty('aws_credentials_provider_arguments',
                              self.credential_file)
        if self.token:
            props.setProperty('aws_credentials_provider_class',
                              'com.amazonaws.athena.jdbc.CustomSessionCredentialsProvider')
            props.setProperty('aws_credentials_provider_arguments',
                              ','.join([self.access_key, self.secret_key, self.token]))
        else:
            props.setProperty('user', self.access_key)
            props.setProperty('password', self.secret_key)
        props.setProperty('s3_staging_dir', self.s3_staging_dir)
        for k, v in iteritems(kwargs):<commit_msg>Fix credentials provider configuration<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals
import logging
import os

import jpype
from future.utils import iteritems

from pyathenajdbc import (ATHENA_DRIVER_CLASS_NAME,
                          ATHENA_CONNECTION_STRING,
                          ATHENA_JAR)
from pyathenajdbc.converter import JDBCTypeConverter
from pyathenajdbc.cursor import Cursor
from pyathenajdbc.error import ProgrammingError, NotSupportedError
from pyathenajdbc.formatter import ParameterFormatter


_logger = logging.getLogger(__name__)


class Connection(object):

    _ENV_S3_STAGING_DIR = 'AWS_ATHENA_S3_STAGING_DIR'

    def __init__(self, s3_staging_dir=None, access_key=None, secret_key=None,
                 region_name=None, profile_name=None, credential_file=None,
                 jvm_options=None, converter=None, formatter=None, jvm_path=None,
                 **driver_args):
        if s3_staging_dir:
            self.s3_staging_dir = s3_staging_dir
        else:
            self.s3_staging_dir = os.getenv(self._ENV_S3_STAGING_DIR, None)
        assert self.s3_staging_dir, 'Required argument `s3_staging_dir` not found.'

        if credential_file:
            self.access_key = None
            self.secret_key = None
            self.token = None
            self.credential_file = credential_file
            assert self.credential_file, 'Required argument `credential_file` not found.'
            self.region_name = region_name
            assert self.region_name, 'Required argument `region_name` not found.'
        else:
            import botocore.session
            session = botocore.session.get_session()
            if access_key and secret_key:
                session.set_credentials(access_key, secret_key)
            if profile_name:
                session.set_config_variable('profile', profile_name)
            if region_name:
                session.set_config_variable('region', region_name)
            credentials = session.get_credentials()
            self.access_key = credentials.access_key
            assert self.access_key, 'Required argument `access_key` not found.'
            self.secret_key = credentials.secret_key
            assert self.secret_key, 'Required argument `secret_key` not found.'
            self.token = credentials.token
            self.credential_file = None
            self.region_name = session.get_config_variable('region')
            assert self.region_name, 'Required argument `region_name` not found.'

        self._start_jvm(jvm_options, jvm_path)

        props = self._build_driver_args(**driver_args)
        jpype.JClass(ATHENA_DRIVER_CLASS_NAME)
        self._jdbc_conn = jpype.java.sql.DriverManager.getConnection(
            ATHENA_CONNECTION_STRING.format(region=self.region_name), props)

        self._converter = converter if converter else JDBCTypeConverter()
        self._formatter = formatter if formatter else ParameterFormatter()

    @classmethod
    def _start_jvm(cls, options, jvm_path):
        if jvm_path is None:
            jvm_path = jpype.get_default_jvm_path()
        if not jpype.isJVMStarted():
            _logger.debug('JVM path: %s', jvm_path)
            args = ['-server', '-Djava.class.path={0}'.format(
                os.path.join(os.path.dirname(__file__), ATHENA_JAR))]
            if options:
                args.extend(options)
            _logger.debug('JVM args: %s', args)
            jpype.startJVM(jvm_path, *args)
        if not jpype.isThreadAttachedToJVM():
            jpype.attachThreadToJVM()

    def _build_driver_args(self, **kwargs):
        props = jpype.java.util.Properties()
        if self.credential_file:
            props.setProperty('aws_credentials_provider_class',
                              'com.amazonaws.athena.jdbc.shaded.' +
                              'com.amazonaws.auth.PropertiesFileCredentialsProvider')
            props.setProperty('aws_credentials_provider_arguments',
                              self.credential_file)
        elif self.token:
            props.setProperty('aws_credentials_provider_class',
                              'com.amazonaws.athena.jdbc.shaded.' +
                              'com.amazonaws.auth.InstanceProfileCredentialsProvider')
        else:
            props.setProperty('user', self.access_key)
            props.setProperty('password', self.secret_key)
        props.setProperty('s3_staging_dir', self.s3_staging_dir)
        for k, v in iteritems(kwargs):
{'fraction_non_alphanumeric': 0.08066637439719421, 'fraction_numerical': 0.003068829460762823, 'mean_word_length': 2.44079939668175, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '784576', 'n_tokens_mistral': 2659, 'n_tokens_neox': 2604, 'n_words': 488}