22# -*- coding: utf-8 -*-
33from __future__ import unicode_literals
44import json
5- try :
6- import mysql .connector # MySQL Connector
7- from mysql .connector import errorcode
8- except :
9- raise ValueError (
10- 'Please, install mysql-connector module before using plugin.'
11- )
125
136
147class DBConnect :
15- """
16- Light database connection object
17- """
8+ """Light database connection object."""
9+ settings = {}
1810 def _check_settings (self ):
1911 """
2012 Check configuration file
2113 :return: True if all settings are correct
2214 """
23- keys = ['host' , 'user' , 'password' , 'database' ]
15+ keys = ['host' , 'user' , 'password' ]
2416 if not all (key in self .settings .keys () for key in keys ):
2517 raise ValueError (
2618 'Please check credentials file for correct keys: host, user, '
2719 'password, database'
2820 )
21+ if self .engine == "mysql" and 'database' not in self .settings .keys ():
22+ raise ValueError (
23+ 'database parameter is missing in credentials'
24+ )
25+ # @NOTE PostgreSQL uses dbname and is automatically set to username
2926
3027 def connect (self ):
3128 """
3229 Creates connection to database, sets connection and cursor
3330 Connection to database can be loosed,
3431 if that happens you can use this function to reconnect to database
3532 """
36- try :
37- self . connection = mysql . connector . connect ( ** self . settings )
38- except mysql .connector . Error as err :
39- if err . errno == errorcode . ER_ACCESS_DENIED_ERROR :
40- raise ValueError ( "Wrong credentials, ACCESS DENIED" )
41- elif err . errno == errorcode . ER_BAD_DB_ERROR :
33+ if self . engine == "mysql" :
34+ try :
35+ import mysql .connector # MySQL Connector
36+ from mysql . connector import errorcode
37+ self . connection = mysql . connector . connect ( ** self . settings )
38+ except ImportError :
4239 raise ValueError (
43- "Database %s does not exists" % ( self . settings [ 'database' ])
40+ 'Please, install mysql-connector module before using plugin.'
4441 )
45- else :
46- raise ValueError (err )
42+ except mysql .connector .Error as err :
43+ if err .errno == errorcode .ER_ACCESS_DENIED_ERROR :
44+ raise ValueError ("Wrong credentials, ACCESS DENIED" )
45+ elif err .errno == errorcode .ER_BAD_DB_ERROR :
46+ raise ValueError (
47+ "Database %s does not exists" % (self .settings ['database' ])
48+ )
49+ else :
50+ raise ValueError (err )
51+ elif self .engine == "postgres" :
52+ try :
53+ import psycopg2
54+ except ImportError :
55+ raise ValueError (
56+ 'Please, install psycopg2 module before using plugin.'
57+ )
58+ self .connection = psycopg2 .connect (** self .settings )
59+ else :
60+ raise NotImplementedError (
61+ "Database engine %s not implemented!" % self .engine
62+ )
63+
4764 self .cursor = self .connection .cursor ()
4865
4966 def __init__ (self , credentials_file = None , charset = 'utf8' ,
50- port = 3306 , ** kwargs ):
67+ port = 3306 , engine = "mysql" , ** kwargs ):
5168 """
5269 Initialise object with credentials file provided
5370 You can choose between providing file or connection details
@@ -63,6 +80,10 @@ def __init__(self, credentials_file=None, charset='utf8',
6380 self .settings ['charset' ] = charset
6481 # Merge with kwargs
6582 self .settings .update (** kwargs )
83+ self .engine = self .settings .pop ('engine' , engine )
84+ # @NOTE Charset parameter not supported in PostgreSQL
85+ if self .engine == 'postgres' :
86+ self .settings .pop ('charset' , None )
6687 self ._check_settings ()
6788 self .connection = None
6889 self .cursor = None
@@ -197,14 +218,30 @@ def insert(self, data, table, commit=True, update=None):
197218 query = query_insert + query_value
198219 if update and bool (update ):
199220 # bool(dict) checks if dict is not empty
200- query += ' ON DUPLICATE KEY UPDATE '
201- for key in update :
202- query += key + ' = '
203- if isinstance (update [key ], int ):
204- query += update [key ] + ', '
205- else :
206- query += '"' + update [key ] + '", '
207- query = query .rstrip (', ' )
221+ if self .engine == "mysql" :
222+ query += ' ON DUPLICATE KEY UPDATE '
223+ for key in update :
224+ query += key + ' = '
225+ if isinstance (update [key ], int ):
226+ query += update [key ] + ', '
227+ else :
228+ query += '"' + update [key ] + '", '
229+ query = query .rstrip (', ' )
230+ elif self .engine == "postgres" :
231+ query += ' ON CONFLICT ON CONSTRAINT '
232+ query += table + '_pkey'
233+ query += ' DO UPDATE SET '
234+ for key in update :
235+ query = key + ' = '
236+ if isinstance (update [key ], int ):
237+ query += update [key ] + ', '
238+ else :
239+ query += '"' + update [key ] + '", '
240+ query = query .rstrip (', ' )
241+ else :
242+ raise NotImplementedError (
243+ "Update on insert not implemented for choosen engine"
244+ )
208245 # Format, execute and send to database:
209246 self .cursor .execute (query , data )
210247 if commit :
@@ -283,22 +320,24 @@ def delete(self, table, filters=None, case='AND', commit=True):
283320 if commit :
284321 self .commit ()
285322
286- def increment (self , table , columns , steps = 1 , filters = None ,
323+ def increment (self , table , fields , steps = 1 , filters = None ,
287324 case = "AND" , commit = True ):
288325 """
289326 Increment column in table
290327 :param table: str table name
291- :param columns : list column names to increment
328+ :param fields : list column names to increment
292329 :param steps: int steps to increment, default is 1
293330 :param filters: dict filters for rows to use
294331 :param case: Search case, Should be 'AND' or 'OR'
295332 :param commit: Commit at the end or add to pool
296333 :note: If you use safe update mode, filters should be provided
297334 """
298- if not columns :
299- raise ValueError ("You must provide which columns to update" )
335+ if not fields :
336+ raise ValueError (
337+ "You must provide which columns (fields) to update"
338+ )
300339 query = "UPDATE %s SET " % str (table )
301- for column in columns :
340+ for column in fields :
302341 query += "{column} = {column} + {steps}, " .format (
303342 column = column , steps = steps )
304343 query = query .rstrip (', ' )
@@ -317,6 +356,41 @@ def increment(self, table, columns, steps=1, filters=None,
317356 self .commit ()
318357 return {'status' : True , 'message' : "Columns incremented" }
319358
359+ def value_sum (self , table , fields , filters = None , case = 'AND' ):
360+ """
361+ Get total sum of a numeric column(s)
362+ :param table: name of the table
363+ :type table: str
364+ :param fields: fields to get sum of
365+ :type fields: list
366+ :param filters: filters to get custom results (where)
367+ :type filters: dict
368+ :param case: [AND, OR] for filter type
369+ :type case: str
370+ :return: dict with column name and value as Decimal
371+ """
372+ query = 'SELECT '
373+ for field in fields :
374+ query += 'SUM(' + field + '), '
375+ query = query .rstrip (', ' ) + ' FROM ' + str (table )
376+ data = None
377+ if filters :
378+ data = {}
379+ query += ' WHERE '
380+ update_query , where_data = self ._where_builder (filters , case )
381+ query += update_query
382+ for key in where_data :
383+ data ['where_' + key ] = where_data [key ]
384+ if data :
385+ self .cursor .execute (query , data )
386+ else :
387+ self .cursor .execute (query )
388+ row = self .cursor .fetchone ()
389+ result = {}
390+ for i in range (len (row )):
391+ result [fields [i ]] = row [i ]
392+ return result
393+
320394 def commit (self ):
321395 """
322396 Commit collected data for making changes to database
0 commit comments