Test Case: python-instruct-1643

  1import websocket  2import json  3import sqlite3  4​  5def update_database_via_websocket(data, db_path="database.db", websocket_url="ws://localhost:8765"):  6    """  7    Updates a database using a websocket connection.  8​  9    Args: 10        data (dict): A dictionary containing the data to update. 11                     Expected keys: 'problem_id', 'user_id', 'timestamp'. 12        db_path (str): The path to the SQLite database file. 13        websocket_url (str): The URL of the websocket server. 14    """ 15    try: 16        # Establish websocket connection 17        ws = websocket.create_connection(websocket_url) 18        print("Websocket connection established.") 19​ 20        # Connect to the database 21        conn = sqlite3.connect(db_path) 22        cursor = conn.cursor() 23        print("Database connection established.") 24​ 25        problem_id = data.get('problem_id') 26        user_id = data.get('user_id') 27        timestamp = data.get('timestamp') 28​ 29        if not all([problem_id, user_id, timestamp]): 30            print("Error: Missing required data (problem_id, user_id, timestamp).") 31            return 32​ 33        # Update the problem table 34        update_query = """ 35        UPDATE problems 36        SET solved_count = solved_count + 1, 37            solved_user = ?, 38            last_solved_date = ? 39        WHERE id = ? 40        """ 41        cursor.execute(update_query, (user_id, timestamp, problem_id)) 42        print(f"Problem {problem_id} updated in database.") 43​ 44        # Commit changes 45        conn.commit() 46        print("Database changes committed.") 47​ 48        # Send confirmation or relevant data back through websocket 49        response_data = { 50            "status": "success", 51            "message": f"Problem {problem_id} updated successfully.", 52            "problem_id": problem_id, 53            "user_id": user_id, 54            "timestamp": timestamp 55        } 56        ws.send(json.dumps(response_data)) 57        print("Sent confirmation via websocket.") 58​ 59        # Close connections 60        ws.close() 61        conn.close() 62        print("Connections closed.") 63​ 64    except ConnectionRefusedError: 65        print(f"Error: Connection refused. Is the websocket server running at {websocket_url}?") 66    except sqlite3.Error as e: 67        print(f"Database error: {e}") 68        if conn: 69            conn.rollback()  # Rollback changes if an error occurs 70            conn.close() 71    except Exception as e: 72        print(f"An unexpected error occurred: {e}") 73        if ws: 74            ws.close() 75        if conn: 76            conn.close() 77​ 78if __name__ == '__main__': 79    # Example usage: 80    # This part is for demonstration and testing purposes. 81    # You would typically call update_database_via_websocket from your application logic. 82​ 83    # Create a dummy database and table for testing 84    db_file = "test_database.db" 85    conn_test = sqlite3.connect(db_file) 86    cursor_test = conn_test.cursor() 87    cursor_test.execute("DROP TABLE IF EXISTS problems") 88    cursor_test.execute(""" 89        CREATE TABLE problems ( 90            id INTEGER PRIMARY KEY, 91            name TEXT NOT NULL, 92            solved_count INTEGER DEFAULT 0, 93            solved_user TEXT, 94            last_solved_date TEXT 95        ) 96    """) 97    cursor_test.execute("INSERT INTO problems (id, name) VALUES (1, 'Problem A')") 98    cursor_test.execute("INSERT INTO problems (id, name) VALUES (2, 'Problem B')") 99    conn_test.commit()100    conn_test.close()101    print("Dummy database created for testing.")102​103    # Dummy websocket server (for testing purposes)104    # In a real scenario, this would be a separate running websocket server.105    import threading106    import time107​108    def run_dummy_server():109        from websocket_server import WebsocketServer110        server = WebsocketServer(port=8765, host='localhost')111​112        def new_client(client, server):113            print(f"New client connected: {client['id']}")114​115        def message_received(client, server, message):116            print(f"Message received from client {client['id']}: {message}")117            # In a real server, you'd process this message and potentially send a response.118            # For this example, we're not expecting messages from the server to the client119            # that trigger the database update. The update is initiated by the client.120​121        server.set_fn_new_client(new_client)122        server.set_fn_message_received(message_received)123        print("Starting dummy websocket server on ws://localhost:8765...")124        server.run_forever()125​126    server_thread = threading.Thread(target=run_dummy_server)127    server_thread.daemon = True  # Allow the main thread to exit even if this thread is running128    server_thread.start()129    time.sleep(1) # Give the server a moment to start130​131    # Data to update132    update_data = {133        'problem_id': 1,134        'user_id': 'user123',135        'timestamp': '2023-10-27 10:00:00'136    }137​138    print("\nCalling update_database_via_websocket...")139    update_database_via_websocket(update_data, db_path=db_file)140​141    # Verify the update142    conn_verify = sqlite3.connect(db_file)143    cursor_verify = conn_verify.cursor()144    cursor_verify.execute("SELECT solved_count, solved_user, last_solved_date FROM problems WHERE id = 1")145    result = cursor_verify.fetchone()146    print(f"\nVerification: Problem 1 - Solved Count: {result[0]}, Solved User: {result[1]}, Last Solved Date: {result[2]}")147    conn_verify.close()148​149    # Example with missing data150    print("\nCalling update_database_via_websocket with missing data...")151    update_database_via_websocket({'problem_id': 2, 'user_id': 'user456'}, db_path=db_file)152​153    # Example with non-existent problem ID154    print("\nCalling update_database_via_websocket with non-existent problem ID...")155    update_database_via_websocket({'problem_id': 99, 'user_id': 'user789', 'timestamp': '2023-10-27 11:00:00'}, db_path=db_file)156​157    # Keep the main thread alive to see server output if needed158    # In a real application, this would not be necessary.159    # input("Press Enter to exit...")

None

  1import json  2import datetime  3from django.db import connection  4from django.contrib.auth.decorators import login_required  5from channels.generic.websocket import AsyncWebsocketConsumer  6from django.core.exceptions import ObjectDoesNotExist  7​  8# Assume you have a Django model named 'Problem' with fields:  9# id, solved_count, solved_by_user, last_solved_date 10​ 11class DatabaseUpdateConsumer(AsyncWebsocketConsumer): 12    async def connect(self): 13        # Authenticate the user if necessary. For simplicity, we'll assume 14        # authentication is handled by Django's auth system and passed via scope. 15        # In a real-world scenario, you'd likely use token authentication or similar. 16        if self.scope["user"].is_anonymous: 17            await self.close() 18            return 19        self.user = self.scope["user"] 20        await self.accept() 21​ 22    async def disconnect(self, close_code): 23        pass 24​ 25    async def receive(self, text_data): 26        try: 27            data = json.loads(text_data) 28            await self.update_database_and_problem(data) 29        except json.JSONDecodeError: 30            await self.send(text_data=json.dumps({"error": "Invalid JSON received."})) 31        except Exception as e: 32            await self.send(text_data=json.dumps({"error": f"An error occurred: {str(e)}"})) 33​ 34    async def update_database_and_problem(self, update_data): 35        problem_id = update_data.get("problem_id") 36        table_name = update_data.get("table_name") 37        update_values = update_data.get("update_values") 38​ 39        if not all([problem_id, table_name, update_values]): 40            await self.send(text_data=json.dumps({"error": "Missing required fields: problem_id, table_name, update_values."})) 41            return 42​ 43        # --- Security: Validate table_name against a known list of allowed tables --- 44        ALLOWED_TABLES = ['your_problem_table_name'] # Replace with your actual table name 45        if table_name not in ALLOWED_TABLES: 46            await self.send(text_data=json.dumps({"error": f"Invalid table name: {table_name}. Only allowed tables can be updated."})) 47            return 48​ 49        # --- Security: Sanitize and validate update_values to prevent SQL injection --- 50        # This is a basic example. For complex updates, consider a more robust approach. 51        # We'll assume update_values is a dictionary of column: value pairs. 52        # We'll only allow updates to specific, known columns. 53        ALLOWED_COLUMNS_FOR_UPDATE = ['column1', 'column2'] # Replace with your actual columns 54        sanitized_updates = {} 55        for key, value in update_values.items(): 56            if key in ALLOWED_COLUMNS_FOR_UPDATE: 57                # Further validation based on expected data types for each column 58                # For example: 59                # if key == 'column1' and not isinstance(value, int): 60                #     await self.send(text_data=json.dumps({"error": f"Invalid data type for {key}."})) 61                #     return 62                sanitized_updates[key] = value 63            else: 64                await self.send(text_data=json.dumps({"error": f"Invalid column for update: {key}."})) 65                return 66​ 67        if not sanitized_updates: 68            await self.send(text_data=json.dumps({"error": "No valid columns to update."})) 69            return 70​ 71        # --- Security: Construct SQL query using parameterized queries --- 72        # This prevents SQL injection. 73        set_clauses = [] 74        params = [] 75        for key, value in sanitized_updates.items(): 76            set_clauses.append(f"{key} = %s") 77            params.append(value) 78​ 79        # Ensure problem_id is an integer to prevent injection 80        try: 81            problem_id = int(problem_id) 82            params.append(problem_id) 83        except ValueError: 84            await self.send(text_data=json.dumps({"error": "Invalid problem_id format."})) 85            return 86​ 87        sql_update_query = f"UPDATE {table_name} SET {', '.join(set_clauses)} WHERE id = %s" 88​ 89        # --- Security: Increment solved counter, update solved user and date --- 90        # This part assumes you have a 'Problem' model and are updating it. 91        # If your 'Problem' model is not directly the table being updated, 92        # you'll need to adjust this logic. 93        try: 94            from .models import Problem # Adjust import path as needed 95            problem_instance = await Problem.objects.aget(id=problem_id) 96​ 97            problem_instance.solved_count += 1 98            problem_instance.solved_by_user = self.user 99            problem_instance.last_solved_date = datetime.date.today()100            await problem_instance.asave()101​102        except ObjectDoesNotExist:103            await self.send(text_data=json.dumps({"error": f"Problem with ID {problem_id} not found."}))104            return105        except Exception as e:106            await self.send(text_data=json.dumps({"error": f"Error updating Problem model: {str(e)}"}))107            return108​109        # --- Execute the main table update using parameterized query ---110        try:111            with connection.cursor() as cursor:112                cursor.execute(sql_update_query, params)113            # Commit is handled by Django's transaction management by default,114            # but if you're doing manual transactions, you'd commit here.115            # For simplicity and safety, rely on Django's default transaction handling.116​117            await self.send(text_data=json.dumps({"success": True, "message": "Database updated successfully."}))118​119        except Exception as e:120            # Rollback if an error occurs during the main update121            # Again, Django's transaction management usually handles this.122            await self.send(text_data=json.dumps({"error": f"Database update failed: {str(e)}"}))

None