I'm writing a small Python script which needs to fetch live data from a database (I'm using SQLite 3 for now). I may reuse this function for a lot of scripts. Can this code be made better?
def fetch_database(database, table, filter_dict, case=None):
"""
Takes filter parameter as a dict
dict = {'column_name':['filter', 'filter2'], 'column_name2'=[], ..}
case = wrap around the dictionary key(UPPER, LOWER, etc..), None
database (with pathname ex: database.db
table
column = dict.keys()
"""
conn = sqlite3.connect(database)
c = conn.cursor()
filter_dict = dict_format(filter_dict) # A separate function deletes empty keys from dictionary
keys_list = filter_dict.keys()
statement = 'SELECT * FROM ' + table
if len(keys_list) > 0:
statement += ' WHERE '
for keys in keys_list:
if case == None:
key = keys
else:
key = case + '(' + keys + ')'
temp_data = filter_dict[keys]
temp_size = len(temp_data)
if keys_list.index(keys) != 0:
statement += ' AND '
if temp_size > 0:
for data in temp_data:
if temp_data.index(data) == 0:
statement += key + '="' + data + '"'
else:
statement += ' OR ' + key + '="' + data + '"'
dataset = c.execute(statement)
dataset = dataset.fetchall()
return dataset
2 Answers 2
When it comes to SQLite, I like to first wrap everything related to database into its own class:
class SqliteDB:
def __init__(self):
self.parent_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
self.db_path = os.path.join(self.parent_directory, 'db.sqlite3')
self.db_connection = sqlite3.connect(self.db_path)
self.db_cursor = self.db_connection.cursor()
def execute(self, query, args=''):
return self.db_cursor.execute(query, args)
def commit(self):
self.db_connection.commit()
def close(self):
self.db_connection.close()
That way, you can easily use the class into whatever project you need. Now, given the above, you can modify your function to take advantage of the new class:
def fetch_database(database, table, filter_dict, case=None):
filter_dict = dict_format(filter_dict)
keys_list = filter_dict.keys()
statement = 'SELECT * FROM {}'.format(table)
if len(keys_list) > 0:
statement += ' WHERE '
for keys in keys_list:
if case is None:
key = keys
else:
key = '{}({})'.format(case, keys)
temp_data = filter_dict[keys]
temp_size = len(temp_data)
if keys_list.index(keys) != 0:
statement += ' AND '
if temp_size > 0:
for data in temp_data:
if temp_data.index(data) == 0:
statement += '{}="{}"'.format(key, data)
else:
statement += ' OR {}="{}"'.format(key, data)
return database.execute(statement).fetchall()
Then you can use it like this:
if __name__ == '__main__':
database = SqliteDB()
fetch_database(database, 'some_table', {'some': 'dict'}, case=None)
Other changes that I did to what you had:
- I've used the
.format()
method instead of the deprecated+
method - I've modified
if case == None
toif case is None
. - I added
if __name__ == '__main__'
guard
-
\$\begingroup\$ seems like you should have mentioned how this is prone to sql injection attacks \$\endgroup\$james– james2017年04月27日 15:30:23 +00:00Commented Apr 27, 2017 at 15:30
-
\$\begingroup\$ (Its it is - first sentence.) \$\endgroup\$Peter Mortensen– Peter Mortensen2017年04月27日 18:26:26 +00:00Commented Apr 27, 2017 at 18:26
-
1\$\begingroup\$ @PeterMortensen thanks. I know the difference. It was just a typo \$\endgroup\$Grajdeanu Alex– Grajdeanu Alex2017年04月27日 19:46:50 +00:00Commented Apr 27, 2017 at 19:46
I'm surprised nobody mentioned the code being vulnerable to SQL injection attacks :
Her daughter is named Help I'm trapped in a driver's licence factory.
Using string formatting and string concatenation to construct SQL queries is, first of all, unsafe and also, fragile - you will have to handle the argument types, balancing and escaping quotes yourself manually.
Instead, you should parameterize your queries - this way, you are not only protected from SQL injections, but also let the database driver worry about the python-to-database type conversions, sample:
query = """
SELECT column1
FROM table
WHERE column2 = ?"""
cursor.execute(query, (data, ))
Note how parameters are passed in a separate argument to execute()
.
On the other hand, it feels like you may be reinventing the wheel in a certain way - please check if switching to an ORM like SQLAlchemy
or peewee
would help to avoid re-writing boilerplate code.
-
1\$\begingroup\$ I wanted to say this, because I remembered I've even given an answer to a related question here but I don't know why I didn't do it ^_^. Nicely spot! \$\endgroup\$Grajdeanu Alex– Grajdeanu Alex2017年04月27日 09:09:38 +00:00Commented Apr 27, 2017 at 9:09
-
1\$\begingroup\$ Thanks for mentioning SQLAlchemy, I will check that. The user input is spell corrected (using personnel dictionary), classified & hence a dictionary format for filtering. It can't contain anything that is not listed. \$\endgroup\$Abhishek Kulkarni– Abhishek Kulkarni2017年04月27日 09:10:10 +00:00Commented Apr 27, 2017 at 9:10
-
2\$\begingroup\$ @MrGrj but you've still provided a very nice SQLite db Python wrapper! I've decided to add Bobby Tables's story here too :) Thanks. \$\endgroup\$alecxe– alecxe2017年04月27日 09:45:04 +00:00Commented Apr 27, 2017 at 9:45