Repository: spotify/luigi

Issue:
Update SQLA module
@erikbern @Tarrasch 
I have been testing the sqla module to incorporate sqlalchemy based task. I have identified some improvements that bring enhancements and fixes, namely:
- Insertions in marker table is NULL in MySQL
- Enhance sqla.CopyToTable to allow table updates apart from insertions

This is a placeholder to bring in these improvements.



Discussion:

Resolving patch:
diff --git a/luigi/contrib/sqla.py b/luigi/contrib/sqla.py
--- a/luigi/contrib/sqla.py
+++ b/luigi/contrib/sqla.py
@@ -127,6 +127,10 @@ def requires(self):
 a transaction at a time. Depending on the size of the inserts, this value can be tuned
 for performance.
 
+See here for a `tutorial on building task pipelines using luigi
+<http://gouthamanbalaraman.com/blog/building-luigi-task-pipeline.html>`_ and
+using `SQLAlchemy in workflow pipelines <http://gouthamanbalaraman.com/blog/sqlalchemy-luigi-workflow-pipeline.html>`_.
+
 Author: Gouthaman Balaraman
 Date: 01/02/2015
 """
@@ -136,11 +140,9 @@ def requires(self):
 import datetime
 import itertools
 import logging
-
 import luigi
 import sqlalchemy
-
-logger = logging.getLogger('luigi-interface')
+import os
 
 
 class SQLAlchemyTarget(luigi.Target):
@@ -153,22 +155,37 @@ class SQLAlchemyTarget(luigi.Target):
     to create a task to write to the database.
     """
     marker_table = None
+    _engine = None  # sqlalchemy engine
+    _pid = None  # the pid of the sqlalchemy engine object
 
     def __init__(self, connection_string, target_table, update_id, echo=False):
         """
         Constructor for the SQLAlchemyTarget.
 
-        :param connection_string: (str) SQLAlchemy connection string
-        :param target_table: (str) The table name for the data
-        :param update_id: (str) An identifier for this data set
-        :param echo: (bool) Flag to setup SQLAlchemy logging
+        :param connection_string: SQLAlchemy connection string
+        :type connection_string: str
+        :param target_table: The table name for the data
+        :type target_table: str
+        :param update_id: An identifier for this data set
+        :type update_id: str
+        :param echo: Flag to setup SQLAlchemy logging
+        :type echo: bool
         :return:
         """
         self.target_table = target_table
         self.update_id = update_id
-        self.engine = sqlalchemy.create_engine(connection_string, echo=echo)
+        self.connection_string = connection_string
+        self.echo = echo
         self.marker_table_bound = None
 
+    @property
+    def engine(self):
+        pid = os.getpid()
+        if (SQLAlchemyTarget._engine is None) or (SQLAlchemyTarget._pid != pid):
+            SQLAlchemyTarget._engine = sqlalchemy.create_engine(self.connection_string, echo=self.echo)
+            SQLAlchemyTarget._pid = pid
+        return SQLAlchemyTarget._engine
+
     def touch(self):
         """
         Mark this update as complete.
@@ -177,13 +194,16 @@ def touch(self):
             self.create_marker_table()
 
         table = self.marker_table_bound
+        id_exists = self.exists()
         with self.engine.begin() as conn:
-            id_exists = self.exists()
             if not id_exists:
-                ins = table.insert().values(update_id=self.update_id, target_table=self.target_table)
-            else:
-                ins = table.update().values(update_id=self.update_id, target_table=self.target_table,
+                ins = table.insert().values(update_id=self.update_id, target_table=self.target_table,
                                             inserted=datetime.datetime.now())
+            else:
+                ins = table.update().where(sqlalchemy.and_(table.c.update_id == self.update_id,
+                                                           table.c.target_table == self.target_table)).\
+                    values(update_id=self.update_id, target_table=self.target_table,
+                           inserted=datetime.datetime.now())
             conn.execute(ins)
         assert self.exists()
 
@@ -193,7 +213,8 @@ def exists(self):
             self.create_marker_table()
         with self.engine.begin() as conn:
             table = self.marker_table_bound
-            s = sqlalchemy.select([table]).where(table.c.update_id == self.update_id).limit(1)
+            s = sqlalchemy.select([table]).where(sqlalchemy.and_(table.c.update_id == self.update_id,
+                                                                 table.c.target_table == self.target_table)).limit(1)
             row = conn.execute(s).fetchone()
         return row is not None
 
@@ -206,7 +227,9 @@ def create_marker_table(self):
         if self.marker_table is None:
             self.marker_table = luigi.configuration.get_config().get('sqlalchemy', 'marker-table', 'table_updates')
 
-        with self.engine.begin() as con:
+        engine = self.engine
+
+        with engine.begin() as con:
             metadata = sqlalchemy.MetaData()
             if not con.dialect.has_table(con, self.marker_table):
                 self.marker_table_bound = sqlalchemy.Table(
@@ -214,9 +237,9 @@ def create_marker_table(self):
                     sqlalchemy.Column("update_id", sqlalchemy.String(128), primary_key=True),
                     sqlalchemy.Column("target_table", sqlalchemy.String(128)),
                     sqlalchemy.Column("inserted", sqlalchemy.DateTime, default=datetime.datetime.now()))
-                metadata.create_all(self.engine)
+                metadata.create_all(engine)
             else:
-                metadata.reflect(bind=self.engine)
+                metadata.reflect(bind=engine)
                 self.marker_table_bound = metadata.tables[self.marker_table]
 
     def open(self, mode):
@@ -231,6 +254,8 @@ class CopyToTable(luigi.Task):
 
     * subclass and override the required `connection_string`, `table` and `columns` attributes.
     """
+    _logger = logging.getLogger('luigi-interface')
+
     echo = False
 
     @abc.abstractmethod
@@ -268,6 +293,8 @@ def create_table(self, engine):
 
         If overridden, use the provided connection object for setting up the table in order to
         create the table and insert data using the same transaction.
+        :param engine: The sqlalchemy engine instance
+        :type engine: object
         """
         def construct_sqla_columns(columns):
             retval = [sqlalchemy.Column(*c[0], **c[1]) for c in columns]
@@ -290,7 +317,7 @@ def construct_sqla_columns(columns):
                         metadata.reflect(bind=engine)
                         self.table_bound = metadata.tables[self.table]
                 except Exception as e:
-                    logger.exception(self.table + str(e))
+                    self._logger.exception(self.table + str(e))
 
     def update_id(self):
         """
@@ -303,8 +330,7 @@ def output(self):
             connection_string=self.connection_string,
             target_table=self.table,
             update_id=self.update_id(),
-            echo=self.echo
-        )
+            echo=self.echo)
 
     def rows(self):
         """
@@ -317,18 +343,33 @@ def rows(self):
                 yield line.strip("\n").split(self.column_separator)
 
     def run(self):
-        logger.info('Running task copy to table for update id %s for table %s', self.update_id(), self.table)
+        self._logger.info("Running task copy to table for update id %s for table %s" % (self.update_id(), self.table))
         output = self.output()
-        self.create_table(output.engine)
-        with output.engine.begin() as conn:
+        engine = output.engine
+        self.create_table(engine)
+        with engine.begin() as conn:
             rows = iter(self.rows())
-            ins_rows = [dict(zip((c.key for c in self.table_bound.c), row))
+            ins_rows = [dict(zip(("_" + c.key for c in self.table_bound.c), row))
                         for row in itertools.islice(rows, self.chunk_size)]
             while ins_rows:
-                ins = self.table_bound.insert()
-                conn.execute(ins, ins_rows)
-                ins_rows = [dict(zip((c.key for c in self.table_bound.c), row))
+                self.copy(conn, ins_rows, self.table_bound)
+                ins_rows = [dict(zip(("_" + c.key for c in self.table_bound.c), row))
                             for row in itertools.islice(rows, self.chunk_size)]
-                logger.info('Finished inserting %d rows into SQLAlchemy target', len(ins_rows))
+                self._logger.info("Finished inserting %d rows into SQLAlchemy target" % len(ins_rows))
         output.touch()
-        logger.info("Finished inserting rows into SQLAlchemy target")
+        self._logger.info("Finished inserting rows into SQLAlchemy target")
+
+    def copy(self, conn, ins_rows, table_bound):
+        """
+        This method does the actual insertion of the rows of data given by ins_rows into the
+        database. A task that needs row updates instead of insertions should overload this method.
+        :param conn: The sqlalchemy connection object
+        :param ins_rows: The dictionary of rows with the keys in the format _<column_name>. For example
+        if you have a table with a column name "property", then the key in the dictionary
+        would be "_property". This format is consistent with the bindparam usage in sqlalchemy.
+        :param table_bound: The object referring to the table
+        :return:
+        """
+        bound_cols = dict((c, sqlalchemy.bindparam("_" + c.key)) for c in table_bound.columns)
+        ins = table_bound.insert().values(bound_cols)
+        conn.execute(ins, ins_rows)


Test patch:
diff --git a/test/contrib/sqla_test.py b/test/contrib/sqla_test.py
--- a/test/contrib/sqla_test.py
+++ b/test/contrib/sqla_test.py
@@ -57,12 +57,14 @@ class SQLATask(sqla.CopyToTable):
     ]
     connection_string = CONNECTION_STRING
     table = "item_property"
+    chunk_size = 1
 
     def requires(self):
         return BaseTask()
 
 
 class TestSQLA(unittest.TestCase):
+    NUM_WORKERS = 1
 
     def _clear_tables(self):
         meta = sqlalchemy.MetaData()
@@ -93,6 +95,7 @@ class TestSQLData(sqla.CopyToTable):
                 (["name", sqlalchemy.String(64)], {}),
                 (["value", sqlalchemy.String(64)], {})
             ]
+            chunk_size = 1
 
             def output(self):
                 pass
@@ -114,6 +117,7 @@ class TestSQLData(sqla.CopyToTable):
             connection_string = CONNECTION_STRING
             table = "test_table"
             columns = []
+            chunk_size = 1
 
         def output(self):
             pass
@@ -140,7 +144,7 @@ def _check_entries(self, engine):
 
     def test_rows(self):
         task, task0 = SQLATask(), BaseTask()
-        luigi.build([task, task0], local_scheduler=True)
+        luigi.build([task, task0], local_scheduler=True, workers=self.NUM_WORKERS)
 
         for i, row in enumerate(task.rows()):
             given = TASK_LIST[i].strip("\n").split("\t")
@@ -158,7 +162,7 @@ def test_run(self):
         self._check_entries(self.engine)
 
         # rerun and the num entries should be the same
-        luigi.build([task0, task], local_scheduler=True)
+        luigi.build([task0, task], local_scheduler=True, workers=self.NUM_WORKERS)
         self._check_entries(self.engine)
 
     def test_run_with_chunk_size(self):
@@ -169,7 +173,7 @@ def test_run_with_chunk_size(self):
         task, task0 = SQLATask(), BaseTask()
         self.engine = sqlalchemy.create_engine(task.connection_string)
         task.chunk_size = 2  # change chunk size and check it runs ok
-        luigi.build([task, task0], local_scheduler=True)
+        luigi.build([task, task0], local_scheduler=True, workers=self.NUM_WORKERS)
         self._check_entries(self.engine)
 
     def test_reflect(self):
@@ -178,17 +182,27 @@ def test_reflect(self):
         completely skip the columns part. It is not even required at that point.
         :return:
         """
-
         class AnotherSQLATask(sqla.CopyToTable):
             connection_string = CONNECTION_STRING
             table = "item_property"
             reflect = True
+            chunk_size = 1
 
             def requires(self):
                 return SQLATask()
 
+            def copy(self, conn, ins_rows, table_bound):
+                ins = table_bound.update().\
+                    where(table_bound.c.property == sqlalchemy.bindparam("_property")).\
+                    values({table_bound.c.item: sqlalchemy.bindparam("_item")})
+                conn.execute(ins, ins_rows)
+
+            def rows(self):
+                for line in TASK_LIST:
+                    yield line.strip("\n").split("\t")
+
         task0, task1, task2 = AnotherSQLATask(), SQLATask(), BaseTask()
-        luigi.build([task0, task1, task2], local_scheduler=True)
+        luigi.build([task0, task1, task2], local_scheduler=True, workers=self.NUM_WORKERS)
         self._check_entries(self.engine)
 
     def test_create_marker_table(self):
@@ -221,6 +235,7 @@ class SQLARowOverloadTest(sqla.CopyToTable):
             ]
             connection_string = CONNECTION_STRING
             table = "item_property"
+            chunk_size = 1
 
             def rows(self):
                 tasks = [("item0", "property0"), ("item1", "property1"), ("item2", "property2"), ("item3", "property3"),
@@ -230,11 +245,14 @@ def rows(self):
                     yield row
 
         task = SQLARowOverloadTest()
-        luigi.build([task], local_scheduler=True)
+        luigi.build([task], local_scheduler=True, workers=self.NUM_WORKERS)
         self._check_entries(self.engine)
 
     def test_column_row_separator(self):
-
+        """
+        Test alternate column row separator works
+        :return:
+        """
         class ModBaseTask(luigi.Task):
 
             def output(self):
@@ -255,10 +273,98 @@ class ModSQLATask(sqla.CopyToTable):
             connection_string = CONNECTION_STRING
             table = "item_property"
             column_separator = ","
+            chunk_size = 1
 
             def requires(self):
                 return ModBaseTask()
 
         task1, task2 = ModBaseTask(), ModSQLATask()
-        luigi.build([task1, task2], local_scheduler=True)
+        luigi.build([task1, task2], local_scheduler=True, workers=self.NUM_WORKERS)
+        self._check_entries(self.engine)
+
+    def test_update_rows_test(self):
+        """
+        Overload the copy() method and implement an update action.
+        :return:
+        """
+        class ModBaseTask(luigi.Task):
+
+            def output(self):
+                return MockFile("BaseTask", mirror_on_stderr=True)
+
+            def run(self):
+                out = self.output().open("w")
+                for task in TASK_LIST:
+                    out.write("dummy_" + task)
+                out.close()
+
+        class ModSQLATask(sqla.CopyToTable):
+            connection_string = CONNECTION_STRING
+            table = "item_property"
+            columns = [
+                (["item", sqlalchemy.String(64)], {}),
+                (["property", sqlalchemy.String(64)], {})
+            ]
+            chunk_size = 1
+
+            def requires(self):
+                return ModBaseTask()
+
+        class UpdateSQLATask(sqla.CopyToTable):
+            connection_string = CONNECTION_STRING
+            table = "item_property"
+            reflect = True
+            chunk_size = 1
+
+            def requires(self):
+                return ModSQLATask()
+
+            def copy(self, conn, ins_rows, table_bound):
+                ins = table_bound.update().\
+                    where(table_bound.c.property == sqlalchemy.bindparam("_property")).\
+                    values({table_bound.c.item: sqlalchemy.bindparam("_item")})
+                conn.execute(ins, ins_rows)
+
+            def rows(self):
+                for task in TASK_LIST:
+                    yield task.strip("\n").split("\t")
+
+        # Running only task1, and task2 should fail
+        task1, task2, task3 = ModBaseTask(), ModSQLATask(), UpdateSQLATask()
+        luigi.build([task1, task2, task3], local_scheduler=True, workers=self.NUM_WORKERS)
+        self._check_entries(self.engine)
+
+    def test_multiple_tasks(self):
+        """
+        Test a case where there are multiple tasks
+        :return:
+        """
+        class SmallSQLATask(sqla.CopyToTable):
+            item = luigi.Parameter()
+            property = luigi.Parameter()
+            columns = [
+                (["item", sqlalchemy.String(64)], {}),
+                (["property", sqlalchemy.String(64)], {})
+            ]
+            connection_string = CONNECTION_STRING
+            table = "item_property"
+            chunk_size = 1
+
+            def rows(self):
+                yield (self.item, self.property)
+
+        class ManyBaseTask(luigi.Task):
+            def requires(self):
+                for t in TASK_LIST:
+                    item, property = t.strip().split("\t")
+                    yield SmallSQLATask(item=item, property=property)
+
+        task2 = ManyBaseTask()
+        luigi.build([task2], local_scheduler=True, workers=self.NUM_WORKERS)
         self._check_entries(self.engine)
+
+
+class TestSQLA2(TestSQLA):
+    """ 2 workers version
+    """
+    NUM_WORKERS = 2