Skip to content

Commit 64d514f

Browse files
author
Jacob Huff
committed
merging fixes for query issues with old sql repo
1 parent c0bc842 commit 64d514f

8 files changed

Lines changed: 261 additions & 2 deletions

File tree

CHANGES.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,15 @@
11
# Change Log
22

3+
## 1.1.4
4+
5+
- Added the ability to return rows if there were results to return from procedure call.
6+
Contributed by Bradley Bishop (Encore Technologies)
7+
8+
## 1.1.3
9+
10+
- Added new action to run SQL Procedures
11+
Contributed by Bradley Bishop (Encore Technologies)
12+
313
## 1.1.2
414

515
* Fixed internal working of the query action compatible with sqlalchemy 2.0+

actions/lib/base_action.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,18 @@ def generate_values(self, sql_obj, data_dict):
106106

107107
return sql_obj
108108

109+
def build_connection(self, kwargs_dict):
110+
# Get the connection details from either config or from action params
111+
connection = self.resolve_connection(kwargs_dict)
112+
113+
# Update Driver with a connector
114+
default_driver = DEFAULT_KNOWN_DRIVER_CONNECTORS.get(connection['drivername'], None)
115+
if default_driver:
116+
connection['drivername'] = default_driver
117+
118+
# Format the connection string
119+
return URL(**connection)
120+
109121
@contextmanager
110122
def db_connection(self, kwargs_dict):
111123
"""Connect to the database and instantiate necessary methods to be used

actions/procedure.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
from lib.base_action import BaseAction
2+
import sqlalchemy
3+
from sqlalchemy.orm import sessionmaker
4+
5+
6+
class SQLProcedureAction(BaseAction):
7+
def __init__(self, config):
8+
"""Creates a new BaseAction given a StackStorm config object (kwargs works too)
9+
:param config: StackStorm configuration object for the pack
10+
:returns: a new BaseAction
11+
"""
12+
super(SQLProcedureAction, self).__init__(config)
13+
14+
def format_data(self, proc_data_obj):
15+
proc_data_string = ""
16+
if proc_data_obj:
17+
proc_data_list = []
18+
for name, value in proc_data_obj.items():
19+
proc_data_list.append("@{0}='{1}'".format(name, value))
20+
21+
proc_data_string = ",".join(proc_data_list)
22+
23+
return proc_data_string
24+
25+
def run(self, **kwargs):
26+
"""Main entry point for the StackStorm actions to execute the operation.
27+
:returns: the dns adapters and the number of network adapters to be
28+
on the VM.
29+
"""
30+
kwargs_dict = dict(kwargs)
31+
32+
proc_data = self.get_del_arg('procedure_data', kwargs_dict)
33+
proc_data_string = self.format_data(proc_data)
34+
35+
proc_name = self.get_del_arg('procedure_name', kwargs_dict)
36+
37+
exec_stmt = "EXEC {} {}".format(proc_name, proc_data_string)
38+
39+
database_connection_string = self.build_connection(kwargs_dict)
40+
engine = sqlalchemy.create_engine(database_connection_string)
41+
session = sessionmaker(bind=engine)()
42+
43+
return_result = None
44+
try:
45+
exec_result = session.execute(exec_stmt)
46+
47+
if exec_result.returns_rows:
48+
return_result = []
49+
all_results = exec_result.fetchall()
50+
for row in all_results:
51+
# Rows are returned as tuples with keys.
52+
# Convert that to a dictionary for return
53+
return_result.append(self.row_to_dict(row))
54+
else:
55+
return_result = {'affected_rows': exec_result.rowcount}
56+
57+
session.commit()
58+
except Exception as error:
59+
session.rollback()
60+
61+
# Return error to the user
62+
raise error
63+
finally:
64+
session.close()
65+
66+
return return_result

actions/procedure.yaml

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
---
2+
name: procedure
3+
runner_type: "python-script"
4+
description: "Run a Stored Procedure into a SQL database"
5+
enabled: true
6+
entry_point: procedure.py
7+
parameters:
8+
connection:
9+
type: string
10+
description: "Name of <connection> from this pack's configuration that specifies how to connect to a database server."
11+
required: false
12+
host:
13+
type: string
14+
description: >
15+
Optional override of the database host in <connection> (required if <connection> is not specified). Database server to connect to. If not using a default port add that here. ex. host.domain.tld or host.domain.tld:1234
16+
required: false
17+
username:
18+
type: string
19+
description: "Optional override of the username in <connection> (required if <connection> is not specified). Username for authentication"
20+
required: false
21+
password:
22+
type: string
23+
description: "Optional override of the password in <connection> (required if <connection> is not specified). Password of the specified username"
24+
secret: true
25+
required: false
26+
database:
27+
type: string
28+
description: "Optional override of the database in <connection> (required if <connection> is not specified). Database to connect to, to run querys against."
29+
required: false
30+
port:
31+
description: "Port to connect to database on. If Default leave blank"
32+
type: integer
33+
required: false
34+
drivername:
35+
type: string
36+
description: "Optional override of the database_type in <connection> (required if <connection> is not specified). The type of database that is being connected to."
37+
required: false
38+
procedure_name:
39+
type: string
40+
description: "Database procedure to call"
41+
required: true
42+
procedure_data:
43+
type: object
44+
description: >
45+
Dictionary of data to be passed to the prodecure
46+
{
47+
'column_1': 'data_to_insert_1',
48+
'column_2': 'data_to_insert_2',
49+
'column_3': 'data_to_insert_3',
50+
'column_4': 'data_to_insert_4',
51+
}
52+
default: {}

pack.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ keywords:
66
- Postgres
77
- MySQL
88
- MsSQL
9-
version: 1.1.2
9+
version: 1.1.4
1010
author: Encore Technologies
1111
email: code@encore.tech
1212
python_versions:

tests/fixtures/config_good.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ connections:
1212
password: PassworD!
1313
port: 1234
1414
database: test
15-
drivername: mysql
15+
drivername: mysql+pymysql
1616
sqlite:
1717
database: path/to/test
1818
drivername: sqlite

tests/test_action_lib_base_action.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from lib.base_action import CONFIG_CONNECTION_KEYS
55
from st2common.runners.base_action import Action
66
from insert import SQLInsertAction
7+
from sqlalchemy.engine.url import URL
78
import decimal
89
import datetime
910

@@ -214,6 +215,24 @@ def test_generate_values_multiple(self):
214215
result = action.generate_values(mock_sql_obj, test_dict)
215216
self.assertEqual(result, expected_output)
216217

218+
def test_build_connection(self):
219+
action = self.get_action_instance(self.config_good)
220+
connection_name = 'full'
221+
connection_config = self.config_good['connections'][connection_name]
222+
expected_output = URL(**connection_config)
223+
224+
result = action.build_connection(connection_config)
225+
self.assertEqual(result, expected_output)
226+
227+
def test_build_connection_sqlite(self):
228+
action = self.get_action_instance(self.config_good)
229+
connection_name = 'sqlite'
230+
connection_config = self.config_good['connections'][connection_name]
231+
expected_output = URL(**connection_config)
232+
233+
result = action.build_connection(connection_config)
234+
self.assertEqual(result, expected_output)
235+
217236
@mock.patch('lib.base_action.sqlalchemy')
218237
def test_connect_to_db(self, mock_sqlalchemy):
219238
action = self.get_action_instance(self.config_good)

tests/test_action_procedure.py

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
from sql_base_action_test_case import SqlBaseActionTestCase
2+
3+
from lib.base_action import BaseAction
4+
from st2common.runners.base_action import Action
5+
from procedure import SQLProcedureAction
6+
from sqlalchemy.engine.url import URL
7+
8+
import mock
9+
10+
__all__ = [
11+
'TestActionSQLProcedureAction'
12+
]
13+
14+
15+
class TestActionSQLProcedureAction(SqlBaseActionTestCase):
16+
__test__ = True
17+
action_cls = SQLProcedureAction
18+
19+
def test_init(self):
20+
action = self.get_action_instance({})
21+
self.assertIsInstance(action, SQLProcedureAction)
22+
self.assertIsInstance(action, BaseAction)
23+
self.assertIsInstance(action, Action)
24+
25+
def test_format_data(self):
26+
action = self.get_action_instance({})
27+
test_dict = {"key1": "value1",
28+
"key2": "value2"}
29+
expected_result = "@key1='value1',@key2='value2'"
30+
result = action.format_data(test_dict)
31+
self.assertEqual(result, expected_result)
32+
33+
def test_format_data_none(self):
34+
action = self.get_action_instance({})
35+
test_dict = {}
36+
expected_result = ""
37+
result = action.format_data(test_dict)
38+
self.assertEqual(result, expected_result)
39+
40+
@mock.patch('procedure.sessionmaker')
41+
@mock.patch('procedure.sqlalchemy')
42+
def test_run(self, mock_sqlalchemy, mock_sessionmaker):
43+
action = self.get_action_instance(self.config_good)
44+
connection_name = 'full'
45+
connection_config = self.config_good['connections'][connection_name]
46+
test_dict = {
47+
'procedure_name': 'Test_Procedure',
48+
'procedure_data': {
49+
'column_1': 'value_1',
50+
'column_2': 'value_2'
51+
}
52+
}
53+
test_dict.update(connection_config)
54+
mock_sqlalchemy.create_engine.return_value = 'Mock Engine'
55+
mock_session = mock.Mock()
56+
mock_exec = mock.Mock(rowcount=1, returns_rows=False)
57+
mock_session().execute.return_value = mock_exec
58+
mock_sessionmaker.return_value = mock_session
59+
expected_value = {'affected_rows': 1}
60+
expected_call = "EXEC Test_Procedure @column_1='value_1',@column_2='value_2'"
61+
62+
result = action.run(**test_dict)
63+
self.assertEqual(result, expected_value)
64+
mock_sqlalchemy.create_engine.assert_called_once_with(URL(**connection_config))
65+
mock_sessionmaker.assert_called_once_with(bind='Mock Engine')
66+
mock_session().execute.assert_called_once_with(expected_call)
67+
68+
@mock.patch('procedure.sessionmaker')
69+
@mock.patch('procedure.sqlalchemy')
70+
def test_run_return_rows(self, mock_sqlalchemy, mock_sessionmaker):
71+
action = self.get_action_instance(self.config_good)
72+
connection_name = 'full'
73+
connection_config = self.config_good['connections'][connection_name]
74+
test_dict = {
75+
'procedure_name': 'Test_Procedure',
76+
'procedure_data': {
77+
'column_1': 'value_1',
78+
'column_2': 'value_2'
79+
}
80+
}
81+
test_dict.update(connection_config)
82+
mock_sqlalchemy.create_engine.return_value = 'Mock Engine'
83+
mock_session = mock.Mock()
84+
test_row = mock.Mock(test1='value', test2='value2')
85+
test_row.keys.return_value = ['test1', 'test2']
86+
mock_exec = mock.Mock(rowcount=-1, returns_rows=True)
87+
mock_exec.fetchall.return_value = [test_row]
88+
mock_session().execute.return_value = mock_exec
89+
mock_sessionmaker.return_value = mock_session
90+
expected_result = [{
91+
'test1': 'value',
92+
'test2': 'value2'
93+
}]
94+
expected_call = "EXEC Test_Procedure @column_1='value_1',@column_2='value_2'"
95+
96+
result = action.run(**test_dict)
97+
self.assertEqual(result, expected_result)
98+
mock_sqlalchemy.create_engine.assert_called_once_with(URL(**connection_config))
99+
mock_sessionmaker.assert_called_once_with(bind='Mock Engine')
100+
mock_session().execute.assert_called_once_with(expected_call)

0 commit comments

Comments
 (0)