def execute(self, context: Dict[str, str]):
mssql = MsSqlHook(mssql_conn_id=self.mssql_conn_id)
self.log.info(
"Dumping Microsoft SQL Server query results to local file")
with mssql.get_conn() as conn:
with conn.cursor() as cursor:
cursor.execute(self.sql)
with NamedTemporaryFile("w") as tmp_file:
csv_writer = csv.writer(tmp_file,
delimiter=self.delimiter,
encoding='utf-8')
field_dict = OrderedDict()
col_count = 0
for field in cursor.description:
col_count += 1
col_position = "Column{position}".format(
position=col_count)
field_dict[col_position if field[0] ==
'' else field[0]] = self.type_map(field[1])
csv_writer.writerows(cursor)
tmp_file.flush()
hive = HiveCliHook(hive_cli_conn_id=self.hive_cli_conn_id)
self.log.info("Loading file into Hive")
hive.load_file(tmp_file.name,
self.hive_table,
field_dict=field_dict,
create=self.create,
partition=self.partition,
delimiter=self.delimiter,
recreate=self.recreate,
tblproperties=self.tblproperties)
def test_get_conn_should_return_connection(self, get_connection, mssql_get_conn):
get_connection.return_value = PYMSSQL_CONN
mssql_get_conn.return_value = mock.Mock()
hook = MsSqlHook()
conn = hook.get_conn()
self.assertEqual(mssql_get_conn.return_value, conn)
mssql_get_conn.assert_called_once()
def test_get_autocommit_should_return_autocommit_state(self, get_connection, mssql_get_conn):
get_connection.return_value = PYMSSQL_CONN
mssql_get_conn.return_value = mock.Mock()
mssql_get_conn.return_value.autocommit_state = 'autocommit_state'
hook = MsSqlHook()
conn = hook.get_conn()
mssql_get_conn.assert_called_once()
self.assertEqual(hook.get_autocommit(conn), 'autocommit_state')
def query(self):
"""
Queries MSSQL and returns a cursor of results.
:return: mssql cursor
"""
mssql = MsSqlHook(mssql_conn_id=self.mssql_conn_id)
conn = mssql.get_conn()
cursor = conn.cursor()
cursor.execute(self.sql)
return cursor
def test_set_autocommit_should_invoke_autocommit(self, get_connection, mssql_get_conn):
get_connection.return_value = PYMSSQL_CONN
mssql_get_conn.return_value = mock.Mock()
autocommit_value = mock.Mock()
hook = MsSqlHook()
conn = hook.get_conn()
hook.set_autocommit(conn, autocommit_value)
mssql_get_conn.assert_called_once()
mssql_get_conn.return_value.autocommit.assert_called_once_with(autocommit_value)