Mercurial > hg > Lemuriformes
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) |
