<commit_before>import select
import traceback
import sys
import time
from sshtunnel import SSHTunnelForwarder
import sshtunnel
import logging
import threading

sshtunnel.DEFAULT_LOGLEVEL = 1
logging.basicConfig(
    format='%(asctime)s| %(levelname)-4.3s|%(threadName)10.9s/%(lineno)04d@%(module)-10.9s| %(message)s', level=1)

SSH_SERVER_ADDRESS = ('127.0.0.1', 2222)
SSH_SERVER_USERNAME = 'linuxserver'
SSH_SERVER_REMOTE_SIDE_ADDRESS_PG = ('10.5.0.5', 5432)
SSH_SERVER_REMOTE_SIDE_ADDRESS_MYSQL = ('10.5.0.6', 3306)
SSH_SERVER_REMOTE_SIDE_ADDRESS_MONGO = ('10.5.0.7', 27017)

PG_DATABASE_NAME = 'main'
PG_USERNAME = 'postgres'
PG_PASSWORD = '<PASSWORD>'
PG_QUERY = 'select version()'
PG_EXPECT = eval(
    """('PostgreSQL 13.0 (Debian 13.0-1.pgdg100+1) on x86_64-pc-linux-gnu, compiled by gcc (Debian 8.3.0-6) 8.3.0, 64-bit',)""")

MYSQL_DATABASE_NAME = 'main'
MYSQL_USERNAME = 'mysql'
MYSQL_PASSWORD = '<PASSWORD>'
MYSQL_QUERY = 'select version()'
MYSQL_EXPECT = (('8.0.22',),)

MONGO_DATABASE_NAME = 'main'
MONGO_USERNAME = 'mongo'
MONGO_PASSWORD = '<PASSWORD>'
MONGO_QUERY = lambda client, db: client.server_info()
MONGO_EXPECT = eval(
    """{'version': '3.6.21', 'gitVersion': '1cd2db51dce4b16f4bc97a75056269df0dc0bddb', 'modules': [], 'allocator': 'tcmalloc', 'javascriptEngine': 'mozjs', 'sysInfo': 'deprecated', 'versionArray': [3, 6, 21, 0], 'openssl': {'running': 'OpenSSL 1.0.2g  1 Mar 2016', 'compiled': 'OpenSSL 1.0.2g  1 Mar 2016'}, 'buildEnvironment': {'distmod': 'ubuntu1604', 'distarch': 'x86_64', 'cc': '/opt/mongodbtoolchain/v2/bin/gcc: gcc (GCC) 5.4.0', 'ccflags': '-fno-omit-frame-pointer -fno-strict-aliasing -ggdb -pthread -Wall -Wsign-compare -Wno-unknown-pragmas -Winvalid-pch -Werror -O2 -Wno-unused-local-typedefs -Wno-unused-function -Wno-deprecated-declarations -Wno-unused-but-set-variable -Wno-missing-braces -fstack-protector-strong -fno-builtin-memcmp', 'cxx': '/opt/mongodbtoolchain/v2/bin/g++: g++ (GCC) 5.4.0', 'cxxflags': '-Woverloaded-virtual -Wno-maybe-uninitialized -std=c++14', 'linkflags': '-pthread -Wl,-z,now -rdynamic -Wl,--fatal-warnings -fstack-protector-strong -fuse-ld=gold -Wl,--build-id -Wl,--hash-style=gnu -Wl,-z,noexecstack -Wl,--warn-execstack -Wl,-z,relro', 'target_arch': 'x86_64', 'target_os': 'linux'}, 'bits': 64, 'debug': False, 'maxBsonObjectSize': 16777216, 'storageEngines': ['devnull', 'ephemeralForTest', 'mmapv1', 'wiredTiger'], 'ok': 1.0}""")


def run_postgres_query(port, query=PG_QUERY):
    import psycopg2

    ASYNC_OK = 1
    ASYNC_READ_TIMEOUT = 2
    ASYNC_WRITE_TIMEOUT = 3
    ASYNC_TIMEOUT = 0.2

    def wait(conn):
        while 1:
            state = conn.poll()
            if state == psycopg2.extensions.POLL_OK:
                break
            elif state == psycopg2.extensions.POLL_WRITE:
                select.select([], [conn.fileno()], [])
            elif state == psycopg2.extensions.POLL_READ:
                select.select([conn.fileno()], [], [])
            else:
                raise psycopg2.OperationalError(
                    "poll() returned %s from _wait function" % state)

    def wait_timeout(conn):
        while 1:
            state = conn.poll()
            if state == psycopg2.extensions.POLL_OK:
                return ASYNC_OK
            elif state == psycopg2.extensions.POLL_WRITE:
                # Wait for the given time and then check the return status
                # If three empty lists are returned then the time-out is
                # reached.
                timeout_status = select.select(
                    [], [conn.fileno()], [], ASYNC_TIMEOUT
                )
                if timeout_status == ([], [], []):
                    return ASYNC_WRITE_TIMEOUT
            elif state == psycopg2.extensions.POLL_READ:
                # Wait for the given time and then check the return status
                # If three empty lists are returned then the time-out is
                # reached.
                timeout_status = select.select(
                    [conn.fileno()], [], [], ASYNC_TIMEOUT
                )
                if timeout_status == ([], [], []):
                    return ASYNC_READ_TIMEOUT
            else:
                raise psycopg2.OperationalError(
                    "poll() returned %s from _wait_timeout function" % state
                )

    pg_conn = psycopg2.connect(
        host='127.0.0.1',
        hostaddr='127.0.0.1',
        port=port,
        database=PG_DATABASE_NAME,
        user=PG_USERNAME,
        password=<PASSWORD>,
        sslmode='disable',
        async_=1
    )
    wait(pg_conn)
    cur = pg_conn.cursor()
    cur.execute(query)
    res = wait_timeout(cur.connection)
    while res != ASYNC_OK:
        res = wait_timeout(cur.connection)
    return cur.fetchone()


def run_mysql_query(port, query=MYSQL_QUERY):
    import pymysql
    conn = pymysql.connect(
        host='127.0.0.1',
        port=port,
        user=MYSQL_USERNAME,
        password=<PASSWORD>,
        database=MYSQL_DATABASE_NAME,
        connect_timeout=5,
        read_timeout=5)
    cursor = conn.cursor()
    cursor.execute(query)
    return cursor.fetchall()


def run_mongo_query(port, query=MONGO_QUERY):
    import pymongo
    client = pymongo.MongoClient('127.0.0.1', port)
    db = client[MONGO_DATABASE_NAME]
    return query(client, db)


def create_tunnel():
    logging.info('Creating SSHTunnelForwarder... (v%s)', sshtunnel.__version__)
    tunnel = SSHTunnelForwarder(
        SSH_SERVER_ADDRESS,
        ssh_username=SSH_SERVER_USERNAME,
        ssh_pkey='./tests/ssh-server-config/ssh_host_rsa_key',
        remote_bind_addresses=[
            SSH_SERVER_REMOTE_SIDE_ADDRESS_PG, SSH_SERVER_REMOTE_SIDE_ADDRESS_MYSQL,
            SSH_SERVER_REMOTE_SIDE_ADDRESS_MONGO,
        ],
    )<commit_msg>run_docker_e2e_tests: extract SSH_PKEY<commit_after>import select
import traceback
import sys
import time
from sshtunnel import SSHTunnelForwarder
import sshtunnel
import logging
import threading

sshtunnel.DEFAULT_LOGLEVEL = 1
logging.basicConfig(
    format='%(asctime)s| %(levelname)-4.3s|%(threadName)10.9s/%(lineno)04d@%(module)-10.9s| %(message)s', level=1)

SSH_SERVER_ADDRESS = ('127.0.0.1', 2222)
SSH_SERVER_USERNAME = 'linuxserver'
SSH_PKEY = './ssh-server-config/ssh_host_rsa_key'
SSH_SERVER_REMOTE_SIDE_ADDRESS_PG = ('10.5.0.5', 5432)
SSH_SERVER_REMOTE_SIDE_ADDRESS_MYSQL = ('10.5.0.6', 3306)
SSH_SERVER_REMOTE_SIDE_ADDRESS_MONGO = ('10.5.0.7', 27017)

PG_DATABASE_NAME = 'main'
PG_USERNAME = 'postgres'
PG_PASSWORD = '<PASSWORD>'
PG_QUERY = 'select version()'
PG_EXPECT = eval(
    """('PostgreSQL 13.0 (Debian 13.0-1.pgdg100+1) on x86_64-pc-linux-gnu, compiled by gcc (Debian 8.3.0-6) 8.3.0, 64-bit',)""")

MYSQL_DATABASE_NAME = 'main'
MYSQL_USERNAME = 'mysql'
MYSQL_PASSWORD = '<PASSWORD>'
MYSQL_QUERY = 'select version()'
MYSQL_EXPECT = (('8.0.22',),)

MONGO_DATABASE_NAME = 'main'
MONGO_USERNAME = 'mongo'
MONGO_PASSWORD = '<PASSWORD>'
MONGO_QUERY = lambda client, db: client.server_info()
MONGO_EXPECT = eval(
    """{'version': '3.6.21', 'gitVersion': '1cd2db51dce4b16f4bc97a75056269df0dc0bddb', 'modules': [], 'allocator': 'tcmalloc', 'javascriptEngine': 'mozjs', 'sysInfo': 'deprecated', 'versionArray': [3, 6, 21, 0], 'openssl': {'running': 'OpenSSL 1.0.2g  1 Mar 2016', 'compiled': 'OpenSSL 1.0.2g  1 Mar 2016'}, 'buildEnvironment': {'distmod': 'ubuntu1604', 'distarch': 'x86_64', 'cc': '/opt/mongodbtoolchain/v2/bin/gcc: gcc (GCC) 5.4.0', 'ccflags': '-fno-omit-frame-pointer -fno-strict-aliasing -ggdb -pthread -Wall -Wsign-compare -Wno-unknown-pragmas -Winvalid-pch -Werror -O2 -Wno-unused-local-typedefs -Wno-unused-function -Wno-deprecated-declarations -Wno-unused-but-set-variable -Wno-missing-braces -fstack-protector-strong -fno-builtin-memcmp', 'cxx': '/opt/mongodbtoolchain/v2/bin/g++: g++ (GCC) 5.4.0', 'cxxflags': '-Woverloaded-virtual -Wno-maybe-uninitialized -std=c++14', 'linkflags': '-pthread -Wl,-z,now -rdynamic -Wl,--fatal-warnings -fstack-protector-strong -fuse-ld=gold -Wl,--build-id -Wl,--hash-style=gnu -Wl,-z,noexecstack -Wl,--warn-execstack -Wl,-z,relro', 'target_arch': 'x86_64', 'target_os': 'linux'}, 'bits': 64, 'debug': False, 'maxBsonObjectSize': 16777216, 'storageEngines': ['devnull', 'ephemeralForTest', 'mmapv1', 'wiredTiger'], 'ok': 1.0}""")


def run_postgres_query(port, query=PG_QUERY):
    import psycopg2

    ASYNC_OK = 1
    ASYNC_READ_TIMEOUT = 2
    ASYNC_WRITE_TIMEOUT = 3
    ASYNC_TIMEOUT = 0.2

    def wait(conn):
        while 1:
            state = conn.poll()
            if state == psycopg2.extensions.POLL_OK:
                break
            elif state == psycopg2.extensions.POLL_WRITE:
                select.select([], [conn.fileno()], [])
            elif state == psycopg2.extensions.POLL_READ:
                select.select([conn.fileno()], [], [])
            else:
                raise psycopg2.OperationalError(
                    "poll() returned %s from _wait function" % state)

    def wait_timeout(conn):
        while 1:
            state = conn.poll()
            if state == psycopg2.extensions.POLL_OK:
                return ASYNC_OK
            elif state == psycopg2.extensions.POLL_WRITE:
                # Wait for the given time and then check the return status
                # If three empty lists are returned then the time-out is
                # reached.
                timeout_status = select.select(
                    [], [conn.fileno()], [], ASYNC_TIMEOUT
                )
                if timeout_status == ([], [], []):
                    return ASYNC_WRITE_TIMEOUT
            elif state == psycopg2.extensions.POLL_READ:
                # Wait for the given time and then check the return status
                # If three empty lists are returned then the time-out is
                # reached.
                timeout_status = select.select(
                    [conn.fileno()], [], [], ASYNC_TIMEOUT
                )
                if timeout_status == ([], [], []):
                    return ASYNC_READ_TIMEOUT
            else:
                raise psycopg2.OperationalError(
                    "poll() returned %s from _wait_timeout function" % state
                )

    pg_conn = psycopg2.connect(
        host='127.0.0.1',
        hostaddr='127.0.0.1',
        port=port,
        database=PG_DATABASE_NAME,
        user=PG_USERNAME,
        password=<PASSWORD>,
        sslmode='disable',
        async_=1
    )
    wait(pg_conn)
    cur = pg_conn.cursor()
    cur.execute(query)
    res = wait_timeout(cur.connection)
    while res != ASYNC_OK:
        res = wait_timeout(cur.connection)
    return cur.fetchone()


def run_mysql_query(port, query=MYSQL_QUERY):
    import pymysql
    conn = pymysql.connect(
        host='127.0.0.1',
        port=port,
        user=MYSQL_USERNAME,
        password=<PASSWORD>,
        database=MYSQL_DATABASE_NAME,
        connect_timeout=5,
        read_timeout=5)
    cursor = conn.cursor()
    cursor.execute(query)
    return cursor.fetchall()


def run_mongo_query(port, query=MONGO_QUERY):
    import pymongo
    client = pymongo.MongoClient('127.0.0.1', port)
    db = client[MONGO_DATABASE_NAME]
    return query(client, db)


def create_tunnel():
    logging.info('Creating SSHTunnelForwarder... (v%s)', sshtunnel.__version__)
    tunnel = SSHTunnelForwarder(
        SSH_SERVER_ADDRESS,
        ssh_username=SSH_SERVER_USERNAME,
        ssh_pkey=SSH_PKEY,
        remote_bind_addresses=[
            SSH_SERVER_REMOTE_SIDE_ADDRESS_PG, SSH_SERVER_REMOTE_SIDE_ADDRESS_MYSQL,
            SSH_SERVER_REMOTE_SIDE_ADDRESS_MONGO,
        ],
    )
{'fraction_non_alphanumeric': 0.13592149436171128, 'fraction_numerical': 0.03503486270121374, 'mean_word_length': 2.7562237310054964, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 16, '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': '3201997', 'n_tokens_mistral': 4332, 'n_tokens_neox': 4054, 'n_words': 766}