Improving PrestoHook to use REST, DatabaseConnection model

.gitignore

@@ -1,8 +1,9 @@
-airflow.db
-env
-logs
-build
-dist
-MANIFEST
 *.pyc
 docs/_*
+airflow.db
+build
+dist
+env
+initdb.py
+logs
+MANIFEST

TODO.md

@@ -27,6 +27,7 @@ TODO
 #### Backend
 * CeleryExecutor
 * Clear should kill running jobs
+* Mysql port should carry through
 
 #### Misc
 * Write an hypervisor, looks for dead jobs without a heartbeat and kills

airflow/bin/airflow

@@ -195,11 +195,29 @@ def initdb(args):
         # Creating the local_mysql DB connection
         session = settings.Session()
         session.query(models.DatabaseConnection).delete()
-        mysqldb = models.DatabaseConnection(
-            db_id='local_mysql', db_type='mysql',
-            host='localhost', login='airflow', password='airflow',
-            schema='airflow')
-        session.add(mysqldb)
+        session.add(
+            models.DatabaseConnection(
+                db_id='local_mysql', db_type='mysql',
+                host='localhost', login='airflow', password='airflow',
+                schema='airflow'))
+        session.commit()
+        session.add(
+            models.DatabaseConnection(
+                db_id='mysql_default', db_type='mysql',
+                host='localhost', login='airflow', password='airflow',
+                schema='airflow'))
+        session.commit()
+        session.add(
+            models.DatabaseConnection(
+                db_id='presto_default', db_type='presto',
+                host='localhost', 
+                schema='hive', port=10001))
+        session.commit()
+        session.add(
+            models.DatabaseConnection(
+                db_id='hive_default', db_type='hive',
+                host='localhost', 
+                schema='default', port=10000))
         session.commit()
 
 

airflow/hooks/__init__.py

@@ -1,3 +1,3 @@
 from mysql_hook import MySqlHook
 from hive_hook import HiveHook
-from presto_hook import PrestoHook
+from airflow.hooks.presto.presto_hook import PrestoHook

airflow/hooks/presto/presto_hook.py

@@ -0,0 +1,52 @@
+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())
+
+

airflow/hooks/presto_hook.py

@@ -1,40 +0,0 @@
-import csv
-import subprocess
-import StringIO
-
-from airflow.hooks.base_hook import BaseHook
-
-PRESTO_BIN = "java -jar /usr/bin/presto-cli.jar --server http://i-5966e3b2:3400 --catalog hive --debug"
-
-
-class PrestoHook(BaseHook):
-    """
-    Interact with Presto! This is just a thin wrapper on the Presto CLI
-    """
-    def __init__(self):
-        pass
-
-    def query(hql, schema="default", headers=False):
-        csv_format = 'CSV_HEADER' if headers else 'CSV'
-
-        cmd = (PRESTO_BIN +
-            ' --execute "{hql}" --schema {schema} ' +
-            ' --output-format {csv_format}').format(**locals())
-        sp = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
-        out, err = sp.communicate()
-        return out
-
-    def get_records(self, hql, schema="default"):
-        results = self.query(hql, schema)
-        return [row for row in csv.reader(StringIO.StringIO(results))]
-
-    def get_pandas_df(self, hql, schema="default"):
-        import pandas
-        results = self.query(hql, schema, headers=True)
-        return pandas.DataFrame.from_csv(
-            StringIO.StringIO(results),
-            header=True,
-        )
-
-    def run(self, sql):
-        raise NotImplemented()

airflow/models.py

@@ -108,24 +108,25 @@ class DatabaseConnection(Base):
     """
     __tablename__ = "db_connection"
 
-    id = Column(Integer, primary_key=True)
-    db_id = Column(String(ID_LEN), unique=True)
+    db_id = Column(String(ID_LEN), primary_key=True)
     db_type = Column(String(500))
     host = Column(String(500))
     schema = Column(String(500))
     login = Column(String(500))
     password = Column(String(500))
+    port = Column(Integer())
 
     def __init__(
             self, db_id=None, db_type=None,
             host=None, login=None, password=None,
-            schema=None):
+            schema=None, port=None):
         self.db_id = db_id
         self.db_type = db_type
         self.host = host
         self.login = login
         self.password = password
         self.schema = schema
+        self.port = port
 
 
 class DagPickle(Base):
@@ -1015,7 +1016,7 @@ class DAG(Base):
         session.commit()
 
     def run(
-            self, start_date=None, end_date=None, mark_success=False, 
+            self, start_date=None, end_date=None, mark_success=False,
             executor=DEFAULT_EXECUTOR):
         session = settings.Session()
         job = BackfillJob(executor=executor)

airflow/settings.py

@@ -4,6 +4,14 @@ from sqlalchemy.orm import sessionmaker
 from sqlalchemy import create_engine
 
 
+HEADER = """\
+       .__         _____.__
+_____  |__|_______/ ____\  |   ______  _  __
+\__  \ |  \_  __ \   __\|  |  /  _ \ \/ \/ /
+ / __ \|  ||  | \/|  |  |  |_(  <_> )     /
+(____  /__||__|   |__|  |____/\____/ \/\_/
+     \/"""
+
 if 'AIRFLOW_HOME' not in os.environ:
     os.environ['AIRFLOW_HOME'] = os.path.join(os.path.dirname(__file__), "..")
 AIRFLOW_HOME = os.environ['AIRFLOW_HOME']
@@ -20,14 +28,10 @@ ID_LEN = 250  # Used for dag_id and task_id VARCHAR length
 LOG_FORMAT = \
     '[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s'
 
+PRESTO_DEFAULT_DBID = "presto_default"
+HIVE_DEFAULT_DBID = "hive_default"
+
 Session = sessionmaker()
 #engine = create_engine('mysql://airflow:airflow@localhost/airflow')
 engine = create_engine('sqlite:///' + BASE_FOLDER + '/airflow.db' )
 Session.configure(bind=engine)
-HEADER = """\
-       .__         _____.__                 
-_____  |__|_______/ ____\  |   ______  _  __
-\__  \ |  \_  __ \   __\|  |  /  _ \ \/ \/ /
- / __ \|  ||  | \/|  |  |  |_(  <_> )     / 
-(____  /__||__|   |__|  |____/\____/ \/\_/  
-     \/"""

airflow/www/app.py

@@ -452,9 +452,11 @@ admin.add_view(mv)
 
 
 class DatabaseConnectionModelView(ModelView):
-    column_exclude_list = ('login', 'password',)
+    column_list = ('db_id', 'db_type', 'host', 'port')
     form_choices = {
         'db_type': [
+            ('hive', 'Hive',),
+            ('presto', 'Presto',),
             ('mysql', 'MySQL',),
             ('oracle', 'Oracle',),
         ]