-
Notifications
You must be signed in to change notification settings - Fork 0
/
sql_db.py
63 lines (52 loc) · 1.76 KB
/
sql_db.py
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
54
55
56
57
58
59
60
61
62
63
import sqlite3
from sqlite3 import Error
import random
from datetime import date, timedelta
import pandas as pd
def create_connection(db_file):
"""Create or connect to an SQLite database."""
conn = None
try:
conn = sqlite3.connect(db_file)
except Error as e:
print(e)
return conn
def create_table(conn, create_table_sql):
"""Create a table with the specified SQL command."""
try:
c = conn.cursor()
c.execute(create_table_sql)
except Error as e:
print(e)
def insert_data(conn, table_name, data_dict):
"""Insert new data into a table."""
columns = ', '.join(data_dict.keys())
placeholders = ', '.join('?' * len(data_dict))
sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
cur = conn.cursor()
cur.execute(sql, list(data_dict.values()))
conn.commit()
return cur.lastrowid
def query_database(query, db_file):
"""Run an SQL query and return results in a DataFrame."""
conn = create_connection(db_file)
df = pd.read_sql_query(query, conn)
conn.close()
return df
def get_schema_representation(db_file):
"""Get the database schema in a JSON-like format."""
conn = create_connection(db_file)
cursor = conn.cursor()
# Get all table names
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
tables = cursor.fetchall()
db_schema = {}
for table in tables:
table_name = table[0]
# Get column details for each table
cursor.execute(f"PRAGMA table_info({table_name});")
columns = cursor.fetchall()
column_details = {column[1]: column[2] for column in columns}
db_schema[table_name] = column_details
conn.close()
return db_schema