Skip to content

fix: Return the link to the connection pool after the node execution is completed #3957

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shaohuzhang1 merged 1 commit into from
Aug 28, 2025
Merged

fix: Return the link to the connection pool after the node execution is completed #3957

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions apps/application/flow/step_node/tool_lib_node/i_tool_lib_node.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
"""
from typing import Type

from django.db import connection
from django.db.models import QuerySet
from django.utils.translation import gettext_lazy as _
from rest_framework import serializers
Expand All @@ -31,6 +32,8 @@ class FunctionLibNodeParamsSerializer(serializers.Serializer):
def is_valid(self, *, raise_exception=False):
super().is_valid(raise_exception=True)
f_lib = QuerySet(Tool).filter(id=self.data.get('tool_lib_id')).first()
# 归还链接到连接池
connection.close()
if f_lib is None:
raise Exception(_('The function has been deleted'))

Copy link
Contributor Author

@shaohuzhang1 shaohuzhang1 Aug 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The code contains a few recommendations:

  1. Connection Management: It's important to manage database connections properly, especially when dealing with Django applications. Closing the connection at all times can lead to performance issues and resource leaks. Instead of closing it here, you might want to ensure that any other operations involving databases do not leave open unnecessary connections.

  2. Validation Logic: There's redundancy in handling exceptions for the missing tool library. You should only catch this exception once and raise the appropriate error message without duplicating logic within is_valid.

Here's an improved version of the code:

from typing import Type

+import functools

from django.db import connection
from django.db.models import QuerySet
from django.utils.translation import gettext_lazy as _
from rest_framework import serializers


class FunctionLibNodeParamsSerializer(serializers.Serializer):

    def __init__(self, *args, **kwargs):
        kwargs['context'] = {**kwargs.get('context', {}), 'connection': connection}
        super().__init__(*args, **kwargs)

    @functools.lru_cache(maxsize=None)
    def is_valid(self, *, raise_exception=False):
        errors = self._validate_data()

        if errors:
            if not raise_exception:
                return False
            else:
                for field, messages in errors.items():
                    self.add_error(field=field, errors=messages)
                return False
        return True

    def _validate_data(self):        
        super().is_valid(raise_exception=True)
        try:
            query_set = Tool.objects.filter(id=self.validated_data["tool_lib_id"])
            if not query_set.exists():  # Check existence before fetching first item
                raise Exception(_("The function has been deleted"))
            f_lib = query_set.first()          
        except KeyError as e:
            errors = {'tool_lib_id': [_("Invalid tool lib ID"), str(e)]}     
            return errors

Key Improvements:

  • Context Handling: The serializer now passes the database connection context into the serializer via its initialization method.
  • Exception Handling: Only one occurrence of checking for a null f_lib is made using exception handling.
  • Redundancy Elimination: Removed redundant logic for catching key errors during data validation.

These changes help optimize the serializer's logic while ensuring proper database management, which is crucial for maintaining application performance.

Expand Down
4 changes: 3 additions & 1 deletion apps/application/flow/workflow_manage.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from functools import reduce
from typing import List, Dict

from django.db import close_old_connections
from django.db import close_old_connections, connection
from django.utils import translation
from django.utils.translation import get_language
from langchain_core.prompts import PromptTemplate
Expand Down Expand Up @@ -433,6 +433,8 @@ def hand_event_node_result(self, current_node, node_result_future):
return None
finally:
current_node.node_chunk.end()
# 归还链接到连接池
connection.close()

def run_node_async(self, node):
future = executor.submit(self.run_node, node)
Copy link
Contributor Author

@shaohuzhang1 shaohuzhang1 Aug 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No issues found and no optimizations required for this code. It's structured well, clear, and follows best practices. The connection.close() call after releasing node chunks is appropriate for resource management in Django applications.

Expand Down
2 changes: 2 additions & 0 deletions apps/models_provider/tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,8 @@ def is_valid_credential(provider, model_type, model_name, model_credential: Dict

def get_model_by_id(_id, workspace_id):
model = QuerySet(Model).filter(id=_id).first()
# 归还链接到连接池
connection.close()
get_authorized_model = DatabaseModelManage.get_model("get_authorized_model")
if model and model.workspace_id != workspace_id and get_authorized_model is not None:
model = get_authorized_model(QuerySet(Model).filter(id=_id), workspace_id).first()
Copy link
Contributor Author

@shaohuzhang1 shaohuzhang1 Aug 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is an issue with this code that needs to be addressed:

The function get_model_by_id does not manage database connections correctly. Specifically, it attempts to close the connection after filtering a queryset but before returning the model result. This can lead to data inconsistencies because the query may have started on one connection and the subsequent filter operation might happen on another.

To fix this, you should ensure that the same database connection is used throughout the entire process of querying and filtering. Here's how you can modify the function:

from django.db import connection

def get_model_by_id(_id, workspace_id):
    try:
        with connection.cursor() as cursor:
            # Assuming Model and DatabaseModelManage are models from Django
            query = f"SELECT * FROM {Model._meta.db_table} WHERE id=%s AND workspace_id=%s"
            cursor.execute(query, (_id, workspace_id))
            
            model_data = cursor.fetchone()
            if model_data:
                model = Model.objects.create(**model_data)
                
                if model and model.workspace_id != workspace_id:
                    authorized_model_query = "SELECT * FROM auth_authorized_models WHERE model_id = %s AND workspace_id = %s;"
                    authorized_cursor = connection.cursor()
                    authorized_cursor.execute(authorized_model_query, (model.id, workspace_id,))
                    
                    authorized_model_data = authorized_cursor.fetchone()
                    if authorized_model_data:
                        authorized_model = DatabaseModelManage.queryset_class().create(**authorized_model_data)
                        
                        return authorized_model
                        
                return model
                
    except Exception as e:
        print(f"An error occurred: {e}")
        return None
    
    finally:
        # Always use commit=False when using .fetchone() without creating a new object
        connection.commit()

Key Changes:

  1. Use of Context Manager: Encapsulate database operations within a context manager (with connection.cursor()), which automatically manages opening and closing the connection.
  2. Database Operations: Perform all necessary database interactions through the same connection to avoid cross-correlation issues.
  3. Commit Operation: Ensure that connection.commit() is called after fetching data to ensure changes are saved.
  4. Safe Data Handling: Use parameterized queries (%s) instead of string interpolation for security purposes.

By making these adjustments, the code will handle database connections properly and mitigate potential data consistency issues related to different connection contexts.

Expand Down
Loading