Coverage for airflow.hooks.presto_hook: 38%

# Coverage for airflow.hooks.presto_hook : 38%

## 47 statements 18 run 29 missing 0 excluded

Hot-keys on this page

r m x p toggle line displays

j k next/prev highlighted chunk

0 (zero) top of page

1 (one) first highlighted chunk

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | from builtins import str from pyhive import presto from pyhive.exc import DatabaseError from airflow.hooks.dbapi_hook import DbApiHook import logging logging.getLogger("pyhive").setLevel(logging.INFO) class PrestoException(Exception): pass class PrestoHook(DbApiHook): """ Interact with Presto through PyHive! >>> ph = PrestoHook() >>> sql = "SELECT count(1) AS num FROM airflow.static_babynames" >>> ph.get_records(sql) [[340698]] """ conn_name_attr = 'presto_conn_id' default_conn_name = 'presto_default' def get_conn(self): """Returns a connection object""" db = self.get_connection(self.presto_conn_id) return presto.connect( host=db.host, port=db.port, username=db.login, catalog=db.extra_dejson.get('catalog', 'hive'), schema=db.schema) @staticmethod def _strip_sql(sql): return sql.strip().rstrip(';') def get_records(self, hql, parameters=None): """ Get a set of records from Presto """ try: return super(PrestoHook, self).get_records( self._strip_sql(hql), parameters) except DatabaseError as e: obj = eval(str(e)) raise PrestoException(obj['message']) def get_first(self, hql, parameters=None): """ Returns only the first row, regardless of how many rows the query returns. """ try: return super(PrestoHook, self).get_first( self._strip_sql(hql), parameters) except DatabaseError as e: obj = eval(str(e)) raise PrestoException(obj['message']) def get_pandas_df(self, hql, parameters=None): """ Get a pandas dataframe from a sql query. """ import pandas cursor = self.get_cursor() cursor.execute(self._strip_sql(hql), parameters) try: data = cursor.fetchall() except DatabaseError as e: obj = eval(str(e)) raise PrestoException(obj['message']) column_descriptions = cursor.description if data: df = pandas.DataFrame(data) df.columns = [c[0] for c in column_descriptions] else: df = pandas.DataFrame() return df def run(self, hql, parameters=None): """ Execute the statement against Presto. Can be used to create views. """ return super(PrestoHook, self).run(self._strip_sql(hql), parameters) def insert_rows(self): raise NotImplemented()