gistfile1.txt

from sshtunnel import SSHTunnelForwarder
import pymysql
from pandas import read_sql
from log import logger


class SSHSQLConn:
    def __init__(
        self,
        db_uname="wolf_rendall",
        db_pass="",
        db_host="avails-provisioned-replica.storefront-db-cf.premieredigital.net",
        db_port=3306,
        db_database="avails",
        ssh_pkey="/Users/wolf.rendall/.ssh/id_rsa",
        ssh_pass="",
        bastion_host='13.56.247.79',
        bastion_port=22,
        bastion_uname="ubuntu",
        localhost='0.0.0.0',
        localhost_port=3306
    ):
        self.db_uname = db_uname
        self.db_pass = db_pass
        self.db_host = db_host
        self.db_port = db_port
        self.db_database = db_database
        self.ssh_pkey = ssh_pkey
        self.ssh_pass = ssh_pass
        self.bastion_host = bastion_host
        self.bastion_port = bastion_port
        self.bastion_uname = bastion_uname
        self.localhost = localhost
        self.localhost_port = localhost_port

    def sql_query(self, query="select count(*) from avails.main_itemmaster"):
        with SSHTunnelForwarder(
                (self.bastion_host, self.bastion_port),
                ssh_username=self.bastion_uname,
                ssh_pkey=self.ssh_pkey,
                ssh_private_key_password=self.ssh_pass,
                remote_bind_address=(self.db_host, self.db_port),
                local_bind_address=(self.localhost, self.localhost_port)
        ) as tunnel:
            connection = pymysql.connect(
                host=tunnel.local_bind_host,
                user=self.db_uname,
                passwd=self.db_pass,
                port=tunnel.local_bind_port,
                database=self.db_database
            )

            data = read_sql(query, connection)

            logger.info("got ssh tunnel db conn. Query results shape: {}".format(data.shape))

            connection.close()

            return data