[AIRFLOW-3767] Correct bulk insert function (#4773)

* [AIRFLOW-3767] Correct bulk insert function

Fix Oracle hook bulk_insert bug when
param target_fields is None or rows
is empty iterable

* change without overwriting variables as Fokko said
This commit is contained in:
zhongjiajie 2019-03-04 22:20:45 +08:00 коммит произвёл Fokko Driesprong
Родитель 717ddaf737
Коммит d665ac19f3
2 изменённых файлов: 17 добавлений и 4 удалений

Просмотреть файл

@ -199,13 +199,15 @@ class OracleHook(DbApiHook):
Default 5000. Set greater than 0. Set 1 to insert each row in each transaction
:type commit_every: int
"""
if not rows:
raise ValueError("parameter rows could not be None or empty iterable")
conn = self.get_conn()
cursor = conn.cursor()
values = ', '.join(':%s' % i for i in range(1, len(target_fields) + 1))
prepared_stm = 'insert into {tablename} ({columns}) values ({values})'.format(
values_base = target_fields if target_fields else rows[0]
prepared_stm = 'insert into {tablename} {columns} values ({values})'.format(
tablename=table,
columns=', '.join(target_fields),
values=values,
columns='({})'.format(', '.join(target_fields)) if target_fields else '',
values=', '.join(':%s' % i for i in range(1, len(values_base) + 1)),
)
row_count = 0
# Chunk the rows

Просмотреть файл

@ -240,3 +240,14 @@ class TestOracleHook(unittest.TestCase):
self.cur.prepare.assert_called_with(
"insert into table (col1, col2, col3) values (:1, :2, :3)")
self.cur.executemany.assert_called_with(None, rows[2:])
def test_bulk_insert_rows_without_fields(self):
rows = [(1, 2, 3), (4, 5, 6), (7, 8, 9)]
self.db_hook.bulk_insert_rows('table', rows)
self.cur.prepare.assert_called_once_with(
"insert into table values (:1, :2, :3)")
self.cur.executemany.assert_called_once_with(None, rows)
def test_bulk_insert_rows_no_rows(self):
rows = []
self.assertRaises(ValueError, self.db_hook.bulk_insert_rows, 'table', rows)