stock-operation/user_events/data/storage/DataStorage.py

# -*- coding: utf-8 -*-
"""
@author: zsc
@time: 2024/11/18
@desc:  数据处理
"""

import sqlite3


class DataStorage:
    """
    数据存储模块：将数据存储到数据库中。
    """

    def __init__(self, db_path):
        """
        初始化DataStorage实例。

        :param db_path: 数据库文件路径。
        """
        self.conn = sqlite3.connect(db_path)
        self.create_table()

    def create_table(self):
        """
        创建用于存储用户行为的数据库表。
        """
        create_table_query = '''
            CREATE TABLE IF NOT EXISTS user_behavior (
                user_id INTEGER,
                action TEXT,
                timestamp TEXT
            )
        '''
        self.conn.execute(create_table_query)

    def save_data(self, df):
        """
        将DataFrame中的数据保存到数据库表中。

        :param df: 包含用户行为数据的DataFrame。
        """
        if df is not None:
            # 将DataFrame追加到数据库表中
            df.to_sql('user_behavior', self.conn, if_exists='append', index=False)