Repository: incubator-airflow Updated Branches: refs/heads/master ba84b6f4a -> 432ac718b
[AIRFLOW-2525] Fix PostgresHook.copy_expert to work with "COPY FROM" For now PostgresHook.copy_expert supports "COPY TO" but not "COPY FROM", because it opens a file with write mode and doesn't commit operations. This PR fixes it by opening a file with read and write mode and committing operations at last. Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/dabf1b96 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/dabf1b96 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/dabf1b96 Branch: refs/heads/master Commit: dabf1b962dcd4323a6ea723076afcd2a20fcb354 Parents: e4e7b55 Author: Kengo Seki <[email protected]> Authored: Fri May 25 00:04:19 2018 -0400 Committer: Kengo Seki <[email protected]> Committed: Fri May 25 10:52:14 2018 -0400 ---------------------------------------------------------------------- airflow/hooks/postgres_hook.py | 9 +++++---- tests/hooks/test_postgres_hook.py | 11 +++++------ 2 files changed, 10 insertions(+), 10 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/dabf1b96/airflow/hooks/postgres_hook.py ---------------------------------------------------------------------- diff --git a/airflow/hooks/postgres_hook.py b/airflow/hooks/postgres_hook.py index 7e89d93..bbf125b 100644 --- a/airflow/hooks/postgres_hook.py +++ b/airflow/hooks/postgres_hook.py @@ -64,10 +64,11 @@ class PostgresHook(DbApiHook): Executes SQL using psycopg2 copy_expert method Necessary to execute COPY command without access to a superuser """ - f = open(filename, 'w') - with closing(self.get_conn()) as conn: - with closing(conn.cursor()) as cur: - cur.copy_expert(sql, f) + with open(filename, 'w+') as f: + with closing(self.get_conn()) as conn: + with closing(conn.cursor()) as cur: + cur.copy_expert(sql, f) + conn.commit() @staticmethod def _serialize_cell(cell, conn): http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/dabf1b96/tests/hooks/test_postgres_hook.py ---------------------------------------------------------------------- diff --git a/tests/hooks/test_postgres_hook.py b/tests/hooks/test_postgres_hook.py index a740264..f636b5a 100644 --- a/tests/hooks/test_postgres_hook.py +++ b/tests/hooks/test_postgres_hook.py @@ -43,17 +43,16 @@ class TestPostgresHook(unittest.TestCase): def test_copy_expert(self): m = mock.mock_open(read_data='{"some": "json"}') - with mock.patch('airflow.hooks.postgres_hook.open', m, create=True) as m: + with mock.patch('airflow.hooks.postgres_hook.open', m): statement = "SQL" filename = "filename" self.cur.fetchall.return_value = None - f = m(filename, 'w') - def test_open(filename, mode): - return f - self.assertEqual(None, self.db_hook.copy_expert(statement, filename, open=test_open)) + self.assertEqual(None, self.db_hook.copy_expert(statement, filename, open=m)) self.conn.close.assert_called_once() self.cur.close.assert_called_once() - self.cur.copy_expert.assert_called_once_with(statement, f) + self.conn.commit.assert_called_once() + self.cur.copy_expert.assert_called_once_with(statement, m.return_value) + m.assert_called_once_with(filename, "w+")
