comparison lemuriformes/mysql.py @ 15:0d1b8bb1d97b

SQL + data related functionality
author Jeff Hammel <k0scist@gmail.com>
date Sun, 10 Dec 2017 17:16:52 -0800
parents
children
comparison
equal deleted inserted replaced
14:756dbd3e391e 15:0d1b8bb1d97b
1 """
2 MySQL database connection class + helpers
3 """
4
5 import pymysql # https://github.com/PyMySQL/PyMySQL
6 import pymysql.cursors
7 import sys
8 from .cli import ConfigurationParser
9 from .sql import SQLConnection
10
11
12 class MySQLConnection(SQLConnection):
13 """connection to a MySQL database"""
14
15 placeholder = '%s' # VALUE placeholder
16 connect_data_keys = ['host', 'user', 'password', 'db', 'port', 'charset']
17
18 def __init__(self, host, user, password, db, port=3306, charset='utf8mb4'):
19 self.connect_data = {}
20 for key in self.connect_data_keys:
21 self.connect_data[key] = locals()[key]
22
23
24 def connect(self):
25 return pymysql.connect(**self.connect_data)
26
27 def __call__(self, sql, *args):
28
29 with self.connect() as cursor:
30 try:
31 cursor.execute(sql, args)
32 except TypeError:
33 print ((sql, args))
34 raise
35 try:
36 result = cursor.fetchall()
37 except:
38 result = None
39 self.connect().commit()
40 return result
41
42 def tables(self):
43 """return tables"""
44
45 data = self("show tables")
46 return [item[0] for item in data]
47
48 def drop(self, table):
49
50 if table not in self.tables():
51 return
52 self("drop table if exists {table}".format(table=table))
53
54 def create(self, table, *columns):
55 """
56 columns -- each column should be a 2-tuple
57 """
58
59 sql = "CREATE TABLE {table} ({columns}) DEFAULT CHARSET=utf8mb4"
60
61 # format columns
62 _columns = ', '.join(["{0} {1}".format(column, _type)
63 for column, _type in columns])
64
65 # execute query
66 self(sql.format(table=table, columns=_columns))
67
68 def insert(self, table, **row):
69 """insert a `row` into `table`"""
70
71 assert row
72 keys = row.keys()
73 values = [row[key] for key in keys]
74 self(self.insert_sql(table=table,
75 columns=keys),
76 *values)
77
78 def insert_many(self, table, columns, values):
79 """
80 insert many rows into `table`
81 columns -- list of columns to insert
82 """
83
84 # https://stackoverflow.com/questions/13020908/sql-multiple-inserts-with-python
85 # It may be more efficient to flatten the string
86 # instead of using `.executemany`; see
87 # https://stackoverflow.com/questions/14011160/how-to-use-python-mysqldb-to-insert-many-rows-at-once
88
89 with self.connect() as connection:
90 sql = self.insert_sql(table=table, columns=columns)
91 connection.executemany(sql, values)
92 self.connect().commit()
93
94
95 class MySQLParser(ConfigurationParser):
96 """command line parser for MySQL"""
97 # TODO: obsolete!
98
99 def add_arguments(self):
100 self.add_argument('host', help="SQL host")
101 self.add_argument('db', help="database to use")
102 self.add_argument('-u', '--user', dest='user', default='root',
103 help="MySQL user [DEFAULT: %(default)s]")
104 self.add_argument('-p', '--password', dest='password',
105 help="MySQL password [DEFAULT: %(default)s]")
106
107 def connection(self):
108 if self.options is None:
109 raise Exception("parse_args not called successfully!")
110
111 return MySQLConnection(host=self.options.host,
112 user=self.options.user,
113 password=self.options.password,
114 db=self.options.db)