Airflow: How to push xcom value from PostgreOperator?

If i'm correct, airflow automatically pushes to xcom when a query returns a value. However, when you look at the code of the postgresoperator you see that it has an execute method that calls the run method of the PostgresHook (extension of dbapi_hook). Both methods do not return anything, as such it pushes nothing to xcom. What we did to fix this is create a CustomPostgresSelectOperator, a copy of the PostgresOperator, but instead of 'hook.run(..)' do 'return hook.get_records(..)'.

Hope that helps you.


Finally, I created a new Sensor ExecuteSqlOperator in the plugin manager under $AIRFLOW_HOME/plugins.

I used CheckOperator as an example and I modified the returned value: the basic running of this operator was exactly the reverse of what I needed.

Here's the of the default ExecuteSqlOperator: CheckOperator

and here is my customized SqlSensor: ReverseSqlSensor

class SqlExecuteOperator(BaseOperator):
    """
    Performs checks against a db. The ``CheckOperator`` expects
    a sql query that will return a single row.

    Note that this is an abstract class and get_db_hook
    needs to be defined. Whereas a get_db_hook is hook that gets a
    single record from an external source.
    :param sql: the sql to be executed
    :type sql: string
    """

    template_fields = ('sql',)
    template_ext = ('.hql', '.sql',)
    ui_color = '#fff7e6'

    @apply_defaults
    def __init__(
            self, sql,
            conn_id=None,
            *args, **kwargs):
        super(SqlExecuteOperator, self).__init__(*args, **kwargs)
        self.conn_id = conn_id
        self.sql = sql

    def execute(self, context=None):
        logging.info('Executing SQL statement: ' + self.sql)
        records = self.get_db_hook().get_first(self.sql)
        logging.info("Record: " + str(records))
        records_int = int(records[0])
        print (records_int)
        return records_int

    def get_db_hook(self):
        return BaseHook.get_hook(conn_id=self.conn_id)