<commit_before>import subprocess
import StringIO

from airflow import settings
from airflow.models import DatabaseConnection
from airflow.hooks.base_hook import BaseHook
from airflow.hooks.presto.presto_client import PrestoClient

class PrestoException(Exception):
    pass

class PrestoHook(BaseHook):
    """
    Interact with Presto!
    """
    def __init__(self, presto_dbid=settings.PRESTO_DEFAULT_DBID):
        session = settings.Session()
        db = session.query(
            DatabaseConnection).filter(
                DatabaseConnection.db_id == presto_dbid)
        if db.count() == 0:
            raise Exception("The presto_dbid you provided isn't defined")
        else:
            db = db.all()[0]
        self.host = db.host
        self.db = db.schema
        self.port = db.port
        session.commit()
        session.close()

        self.client = PrestoClient(
            self.host, in_port=self.port, in_catalog=self.db,
            in_user='airflow')

    def get_records(self, hql, schema="default"):
        if self.client.runquery(hql, schema):
            return self.client.getdata()
        else:
            raise PrestoException(self.client.getlasterrormessage())

    def get_pandas_df(self, hql, schema="default"):
        import pandas
        client = self.client
        if client.runquery(hql, schema):
            data = client.getdata()
            df = pandas.DataFrame(data)
            df.columns = [c['name'] for c in client.getcolumns()]
            return df
        else:
            raise PrestoException(self.client.getlasterrormessage())

<commit_msg>Implementing run method in presto_hook (for view creation)<commit_after>import subprocess
import StringIO

from airflow import settings
from airflow.models import DatabaseConnection
from airflow.hooks.base_hook import BaseHook
from airflow.hooks.presto.presto_client import PrestoClient

class PrestoException(Exception):
    pass

class PrestoHook(BaseHook):
    """
    Interact with Presto!
    """
    def __init__(self, presto_dbid=settings.PRESTO_DEFAULT_DBID):
        session = settings.Session()
        db = session.query(
            DatabaseConnection).filter(
                DatabaseConnection.db_id == presto_dbid)
        if db.count() == 0:
            raise Exception("The presto_dbid you provided isn't defined")
        else:
            db = db.all()[0]
        self.host = db.host
        self.db = db.schema
        self.port = db.port
        session.commit()
        session.close()

        self.client = PrestoClient(
            self.host, in_port=self.port, in_catalog=self.db,
            in_user='airflow')

    def get_records(self, hql, schema="default"):
        if self.client.runquery(hql, schema):
            return self.client.getdata()
        else:
            raise PrestoException(self.client.getlasterrormessage())

    def get_pandas_df(self, hql, schema="default"):
        import pandas
        client = self.client
        if client.runquery(hql, schema):
            data = client.getdata()
            df = pandas.DataFrame(data)
            df.columns = [c['name'] for c in client.getcolumns()]
            return df
        else:
            raise PrestoException(self.client.getlasterrormessage())

    def run(self, hql, schema="default"):
        if not self.client.runquery(hql, schema):
            raise PrestoException(self.client.getlasterrormessage())
{'fraction_non_alphanumeric': 0.09768938286048552, 'fraction_numerical': 0.00116993272886809, 'mean_word_length': 2.4029850746268657, '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': '1539183', 'n_tokens_mistral': 987, 'n_tokens_neox': 986, 'n_words': 223}