ccsbase2/db.py at main · abant07/ccsbase2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import sqlite3

import pandas as pd


class Database:
    def __init__(self, db_filename):
        self.db_filename = db_filename

    def connect(self):
        return sqlite3.connect(self.db_filename)

    def read(self, sql, params=None):
        conn = self.connect()
        try:
            cur = conn.cursor()
            cur.execute(sql, params or [])
            return cur.fetchall()
        finally:
            conn.close()

    def read_df(self, sql, params=None):
        conn = self.connect()
        try:
            return pd.read_sql_query(sql, conn, params=params)
        finally:
            conn.close()

    def write(self, sql, params=None):
        conn = self.connect()
        try:
            cur = conn.cursor()
            cur.execute(sql, params or [])
            conn.commit()
        finally:
            conn.close()

    def write_many(self, sql, params_list):
        conn = self.connect()
        try:
            cur = conn.cursor()
            cur.executemany(sql, params_list)
            conn.commit()
        finally:
            conn.close()

    def write_df(self, df, table_name, if_exists="append"):
        conn = self.connect()
        try:
            df.to_sql(table_name, conn, if_exists=if_exists, index=False)
            conn.commit()
        finally:
            conn.close()