[PR #23224] Migrate SQLAlchemy from 1.x to 2.0 with automated and manual adjustments #30203

Closed
opened 2026-02-21 20:47:03 -05:00 by yindo · 0 comments
Owner

Original Pull Request: https://github.com/langgenius/dify/pull/23224

State: closed
Merged: Yes


This pull request upgrades our codebase from SQLAlchemy 1.x to 2.0 in the api/core directory only.
Related issue: #22668
Related PR: #22801

The migration was done in two steps:

  1. The initial commit was automatically generated using a Python script powered by ChatGPT. It applied systematic changes to align with SQLAlchemy 2.0’s new syntax and API patterns.
  2. The second commit includes manual revisions after reviewing the automated changes.

Currently, changes are scoped to the core directory only. If the approach proves reliable, I plan to extend the migration to the entire api directory in a future update.

The Python script automatically generated using ChatGPT is as follows:

python script part-1
import os
import libcst as cst
from libcst.metadata import MetadataWrapper, PositionProvider


class QueryRewriteTransformer(cst.CSTTransformer):
    METADATA_DEPENDENCIES = (PositionProvider,)

    def leave_SimpleStatementLine(
        self, original_node: cst.SimpleStatementLine, updated_node: cst.SimpleStatementLine
    ) -> cst.FlattenSentinel[cst.BaseStatement]:

        # 匹配:account = db.session.query(Account).where(...).one_or_none()
        if len(original_node.body) != 1:
            return updated_node
        stmt = original_node.body[0]
        if not isinstance(stmt, cst.Assign):
            return updated_node

        target = stmt.targets[0].target
        value = stmt.value

        # 判断是否是:db.session.query(Account).where(...).one_or_none()
        if not isinstance(value, cst.Call):
            return updated_node
        if not isinstance(value.func, cst.Attribute):
            return updated_node
        method_name = value.func.attr.value
        if method_name not in ("first", "one_or_none", "all"):
            return updated_node

        # value.func.value 应该是 where() 调用
        where_call = value.func.value
        if not isinstance(where_call, cst.Call):
            return updated_node
        if not isinstance(where_call.func, cst.Attribute):
            return updated_node
        if where_call.func.attr.value != "where":
            return updated_node
        where_args = where_call.args

        # where_call.func.value 应该是 query() 调用
        query_call = where_call.func.value
        if not isinstance(query_call, cst.Call):
            return updated_node
        if not isinstance(query_call.func, cst.Attribute):
            return updated_node
        if query_call.func.attr.value != "query":
            return updated_node

        model_expr = query_call.args[0].value  # Account

        # ✅ 构造新的两行代码
        assign_stmt_line = cst.SimpleStatementLine(
            body=[
                cst.Assign(
                    targets=[cst.AssignTarget(cst.Name("stmt"))],
                    value=cst.Call(
                        func=cst.Attribute(
                            value=cst.Call(
                                func=cst.Name("select"),
                                args=[cst.Arg(model_expr)]
                            ),
                            attr=cst.Name("where")
                        ),
                        args=where_args
                    )
                )
            ]
        )

        result_stmt_line = cst.SimpleStatementLine(
            body=[
                cst.Assign(
                    targets=[cst.AssignTarget(target)],
                    value=cst.parse_expression(
                        f"db.session.execute(stmt).scalars().{method_name}()"
                    )
                )
            ]
        )

        return cst.FlattenSentinel([assign_stmt_line, result_stmt_line])


def rewrite_file(filepath: str):
    with open(filepath, "r", encoding="utf-8") as f:
        src = f.read()

    module = cst.parse_module(src)
    wrapper = MetadataWrapper(module)
    modified = wrapper.module.visit(QueryRewriteTransformer())

    if modified.code != src:
        with open(filepath, "w", encoding="utf-8") as f:
            f.write(modified.code)
        print(f"✅ Rewritten: {filepath}")
    else:
        print(f"✅ Skipped (no change): {filepath}")


def scan_directory(path: str):
    for root, _, files in os.walk(path):
        for file in files:
            if file.endswith(".py"):
                rewrite_file(os.path.join(root, file))


if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        print("Usage: python query_to_select_rewriter.py <your_project_path>")
        exit(1)
    scan_directory(sys.argv[1])

The script to resolve review comments is as follows:

python script part-2
import argparse
import difflib
import os
from pathlib import Path
import libcst as cst


class _SessionMatcherMixin:
    @staticmethod
    def _is_session_obj(expr: cst.BaseExpression) -> bool:
        # Match Name("session") or Attribute(..., "session")
        if isinstance(expr, cst.Name) and expr.value == "session":
            return True
        if isinstance(expr, cst.Attribute) and isinstance(expr.attr, cst.Name) and expr.attr.value == "session":
            return True
        return False

    @staticmethod
    def _is_attr_call_with_name(node: cst.CSTNode, name: str) -> bool:
        return (
            isinstance(node, cst.Call)
            and isinstance(node.func, cst.Attribute)
            and isinstance(node.func.attr, cst.Name)
            and node.func.attr.value == name
        )


class ExecScalarsFirstToScalar(_SessionMatcherMixin, cst.CSTTransformer):
    """*.session.execute(...).scalars().first() → *.session.scalar(...)"""

    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression:
        # Must be .first() with no args
        if not isinstance(updated_node.func, cst.Attribute):
            return updated_node
        if not isinstance(updated_node.func.attr, cst.Name) or updated_node.func.attr.value != "first":
            return updated_node
        if updated_node.args:
            return updated_node

        # Receiver must be scalars() with no args
        scalars_call = updated_node.func.value
        if not self._is_attr_call_with_name(scalars_call, "scalars") or scalars_call.args:
            return updated_node

        # Whose receiver must be execute(...)
        execute_call = scalars_call.func.value
        if not self._is_attr_call_with_name(execute_call, "execute"):
            return updated_node

        # execute(...) must be on a session-like object
        session_obj = execute_call.func.value
        if not self._is_session_obj(session_obj):
            return updated_node

        # Build: <session_obj>.scalar(<execute args...>)
        new_func = cst.Attribute(value=session_obj, attr=cst.Name("scalar"))
        new_call = cst.Call(func=new_func, args=execute_call.args)

        # Keep parentheses from the outer call if present
        if original_node.lpar or original_node.rpar:
            new_call = new_call.with_changes(lpar=original_node.lpar, rpar=original_node.rpar)
        return new_call


class ExecScalarsOneOrNoneToScalars(_SessionMatcherMixin, cst.CSTTransformer):
    """*.session.execute(...).scalars().one_or_none() → *.session.scalars(...).one_or_none()"""

    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression:
        # Must be .one_or_none() with no args
        if not isinstance(updated_node.func, cst.Attribute):
            return updated_node
        if not isinstance(updated_node.func.attr, cst.Name) or updated_node.func.attr.value != "one_or_none":
            return updated_node
        if updated_node.args:
            return updated_node

        # Receiver must be scalars() with no args
        scalars_call = updated_node.func.value
        if not self._is_attr_call_with_name(scalars_call, "scalars") or scalars_call.args:
            return updated_node

        # Whose receiver must be execute(...)
        execute_call = scalars_call.func.value
        if not self._is_attr_call_with_name(execute_call, "execute"):
            return updated_node

        # execute(...) must be on a session-like object
        session_obj = execute_call.func.value
        if not self._is_session_obj(session_obj):
            return updated_node

        # Build: <session_obj>.scalars(<execute args...>).one_or_none()
        new_scalars_func = cst.Attribute(value=session_obj, attr=cst.Name("scalars"))
        new_scalars_call = cst.Call(func=new_scalars_func, args=execute_call.args)

        new_postfix_attr = cst.Attribute(value=new_scalars_call, attr=cst.Name("one_or_none"))
        new_postfix_call = cst.Call(func=new_postfix_attr, args=[])

        if original_node.lpar or original_node.rpar:
            new_postfix_call = new_postfix_call.with_changes(lpar=original_node.lpar, rpar=original_node.rpar)
        return new_postfix_call


class ExecScalarsAllToScalars(_SessionMatcherMixin, cst.CSTTransformer):
    """*.session.execute(...).scalars().all() → *.session.scalars(...).all()"""

    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression:
        # Must be .all() with no args
        if not isinstance(updated_node.func, cst.Attribute):
            return updated_node
        if not isinstance(updated_node.func.attr, cst.Name) or updated_node.func.attr.value != "all":
            return updated_node
        if updated_node.args:
            return updated_node

        # Receiver must be scalars() with no args
        scalars_call = updated_node.func.value
        if not self._is_attr_call_with_name(scalars_call, "scalars") or scalars_call.args:
            return updated_node

        # Whose receiver must be execute(...)
        execute_call = scalars_call.func.value
        if not self._is_attr_call_with_name(execute_call, "execute"):
            return updated_node

        # execute(...) must be on a session-like object
        session_obj = execute_call.func.value
        if not self._is_session_obj(session_obj):
            return updated_node

        # Build: <session_obj>.scalars(<execute args...>).all()
        new_scalars_func = cst.Attribute(value=session_obj, attr=cst.Name("scalars"))
        new_scalars_call = cst.Call(func=new_scalars_func, args=execute_call.args)

        new_postfix_attr = cst.Attribute(value=new_scalars_call, attr=cst.Name("all"))
        new_postfix_call = cst.Call(func=new_postfix_attr, args=[])

        if original_node.lpar or original_node.rpar:
            new_postfix_call = new_postfix_call.with_changes(lpar=original_node.lpar, rpar=original_node.rpar)
        return new_postfix_call


def rewrite_source(code: str) -> str:
    mod = cst.parse_module(code)
    # 顺序:先处理 first,再处理 one_or_none,最后 all
    mod = mod.visit(ExecScalarsFirstToScalar())
    mod = mod.visit(ExecScalarsOneOrNoneToScalars())
    mod = mod.visit(ExecScalarsAllToScalars())
    return mod.code


def color_diff(diff_lines):
    for line in diff_lines:
        if line.startswith("-"):
            yield f"\033[31m{line}\033[0m"  # 红色
        elif line.startswith("+"):
            yield f"\033[32m{line}\033[0m"  # 绿色
        else:
            yield line


def process_file(fp: Path, write: bool) -> bool:
    try:
        src = fp.read_text(encoding="utf-8")
    except Exception:
        return False

    try:
        new_src = rewrite_source(src)
    except Exception:
        # skip files that fail to parse
        return False

    if src == new_src:
        return False

    if write:
        fp.write_text(new_src, encoding="utf-8")
        print(f"Updated: {fp}")
    else:
        print(f"DRY-RUN (would update): {fp}")
        diff = difflib.unified_diff(
            src.splitlines(keepends=True),
            new_src.splitlines(keepends=True),
            fromfile=str(fp),
            tofile=f"{fp} (rewritten)",
            n=3,
        )
        for i, line in enumerate(color_diff(diff)):
            if i > 200:
                print("... (diff truncated)")
                break
            print(line, end="")
    return True


def iter_py_files(root: Path):
    if root.is_file() and root.suffix == ".py":
        yield root
        return
    EXCLUDES = {".git", ".hg", ".svn", ".idea", ".vscode", "node_modules", "dist", "build", "__pycache__", ".venv", "venv"}
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if d not in EXCLUDES]
        for name in filenames:
            if name.endswith(".py"):
                yield Path(dirpath) / name


def main():
    ap = argparse.ArgumentParser(
        description="Rewrite session.execute(...).scalars().{first,one_or_none,all}() using LibCST"
    )
    ap.add_argument("path", help="Project root directory or a single .py file")
    ap.add_argument("--write", action="store_true", help="Apply changes (default: dry-run preview)")
    args = ap.parse_args()

    root = Path(args.path).resolve()
    if not root.exists():
        raise SystemExit(f"Path not found: {root}")

    changed = 0
    for fp in iter_py_files(root):
        if process_file(fp, write=args.write):
            changed += 1

    if args.write:
        print(f"Done. Files updated: {changed}")
    else:
        print(f"Dry-run complete. Files that would change: {changed}")


if __name__ == "__main__":
    main()

Checklist

  • This change requires a documentation update, included: Dify Document
  • I understand that this PR may be closed in case there was no previous discussion or issues. (This doesn't apply to typos!)
  • I've added a test for each change that was introduced, and I tried as much as possible to make a single atomic change.
  • I've updated the documentation accordingly.
  • I ran dev/reformat(backend) and cd web && npx lint-staged(frontend) to appease the lint gods
**Original Pull Request:** https://github.com/langgenius/dify/pull/23224 **State:** closed **Merged:** Yes --- This pull request upgrades our codebase from SQLAlchemy 1.x to 2.0 in the `api/core` directory only. Related issue: #22668 Related PR: #22801 The migration was done in two steps: 1. The initial commit was automatically generated using a Python script powered by ChatGPT. It applied systematic changes to align with SQLAlchemy 2.0’s new syntax and API patterns. 2. The second commit includes manual revisions after reviewing the automated changes. Currently, changes are scoped to the `core` directory only. If the approach proves reliable, I plan to extend the migration to the entire `api` directory in a future update. The Python script automatically generated using ChatGPT is as follows: <details> <summary>python script part-1</summary> ``` python import os import libcst as cst from libcst.metadata import MetadataWrapper, PositionProvider class QueryRewriteTransformer(cst.CSTTransformer): METADATA_DEPENDENCIES = (PositionProvider,) def leave_SimpleStatementLine( self, original_node: cst.SimpleStatementLine, updated_node: cst.SimpleStatementLine ) -> cst.FlattenSentinel[cst.BaseStatement]: # 匹配:account = db.session.query(Account).where(...).one_or_none() if len(original_node.body) != 1: return updated_node stmt = original_node.body[0] if not isinstance(stmt, cst.Assign): return updated_node target = stmt.targets[0].target value = stmt.value # 判断是否是:db.session.query(Account).where(...).one_or_none() if not isinstance(value, cst.Call): return updated_node if not isinstance(value.func, cst.Attribute): return updated_node method_name = value.func.attr.value if method_name not in ("first", "one_or_none", "all"): return updated_node # value.func.value 应该是 where() 调用 where_call = value.func.value if not isinstance(where_call, cst.Call): return updated_node if not isinstance(where_call.func, cst.Attribute): return updated_node if where_call.func.attr.value != "where": return updated_node where_args = where_call.args # where_call.func.value 应该是 query() 调用 query_call = where_call.func.value if not isinstance(query_call, cst.Call): return updated_node if not isinstance(query_call.func, cst.Attribute): return updated_node if query_call.func.attr.value != "query": return updated_node model_expr = query_call.args[0].value # Account # ✅ 构造新的两行代码 assign_stmt_line = cst.SimpleStatementLine( body=[ cst.Assign( targets=[cst.AssignTarget(cst.Name("stmt"))], value=cst.Call( func=cst.Attribute( value=cst.Call( func=cst.Name("select"), args=[cst.Arg(model_expr)] ), attr=cst.Name("where") ), args=where_args ) ) ] ) result_stmt_line = cst.SimpleStatementLine( body=[ cst.Assign( targets=[cst.AssignTarget(target)], value=cst.parse_expression( f"db.session.execute(stmt).scalars().{method_name}()" ) ) ] ) return cst.FlattenSentinel([assign_stmt_line, result_stmt_line]) def rewrite_file(filepath: str): with open(filepath, "r", encoding="utf-8") as f: src = f.read() module = cst.parse_module(src) wrapper = MetadataWrapper(module) modified = wrapper.module.visit(QueryRewriteTransformer()) if modified.code != src: with open(filepath, "w", encoding="utf-8") as f: f.write(modified.code) print(f"✅ Rewritten: {filepath}") else: print(f"✅ Skipped (no change): {filepath}") def scan_directory(path: str): for root, _, files in os.walk(path): for file in files: if file.endswith(".py"): rewrite_file(os.path.join(root, file)) if __name__ == "__main__": import sys if len(sys.argv) != 2: print("Usage: python query_to_select_rewriter.py <your_project_path>") exit(1) scan_directory(sys.argv[1]) ``` </details> The script to resolve review comments is as follows: <details> <summary>python script part-2</summary> ``` python import argparse import difflib import os from pathlib import Path import libcst as cst class _SessionMatcherMixin: @staticmethod def _is_session_obj(expr: cst.BaseExpression) -> bool: # Match Name("session") or Attribute(..., "session") if isinstance(expr, cst.Name) and expr.value == "session": return True if isinstance(expr, cst.Attribute) and isinstance(expr.attr, cst.Name) and expr.attr.value == "session": return True return False @staticmethod def _is_attr_call_with_name(node: cst.CSTNode, name: str) -> bool: return ( isinstance(node, cst.Call) and isinstance(node.func, cst.Attribute) and isinstance(node.func.attr, cst.Name) and node.func.attr.value == name ) class ExecScalarsFirstToScalar(_SessionMatcherMixin, cst.CSTTransformer): """*.session.execute(...).scalars().first() → *.session.scalar(...)""" def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression: # Must be .first() with no args if not isinstance(updated_node.func, cst.Attribute): return updated_node if not isinstance(updated_node.func.attr, cst.Name) or updated_node.func.attr.value != "first": return updated_node if updated_node.args: return updated_node # Receiver must be scalars() with no args scalars_call = updated_node.func.value if not self._is_attr_call_with_name(scalars_call, "scalars") or scalars_call.args: return updated_node # Whose receiver must be execute(...) execute_call = scalars_call.func.value if not self._is_attr_call_with_name(execute_call, "execute"): return updated_node # execute(...) must be on a session-like object session_obj = execute_call.func.value if not self._is_session_obj(session_obj): return updated_node # Build: <session_obj>.scalar(<execute args...>) new_func = cst.Attribute(value=session_obj, attr=cst.Name("scalar")) new_call = cst.Call(func=new_func, args=execute_call.args) # Keep parentheses from the outer call if present if original_node.lpar or original_node.rpar: new_call = new_call.with_changes(lpar=original_node.lpar, rpar=original_node.rpar) return new_call class ExecScalarsOneOrNoneToScalars(_SessionMatcherMixin, cst.CSTTransformer): """*.session.execute(...).scalars().one_or_none() → *.session.scalars(...).one_or_none()""" def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression: # Must be .one_or_none() with no args if not isinstance(updated_node.func, cst.Attribute): return updated_node if not isinstance(updated_node.func.attr, cst.Name) or updated_node.func.attr.value != "one_or_none": return updated_node if updated_node.args: return updated_node # Receiver must be scalars() with no args scalars_call = updated_node.func.value if not self._is_attr_call_with_name(scalars_call, "scalars") or scalars_call.args: return updated_node # Whose receiver must be execute(...) execute_call = scalars_call.func.value if not self._is_attr_call_with_name(execute_call, "execute"): return updated_node # execute(...) must be on a session-like object session_obj = execute_call.func.value if not self._is_session_obj(session_obj): return updated_node # Build: <session_obj>.scalars(<execute args...>).one_or_none() new_scalars_func = cst.Attribute(value=session_obj, attr=cst.Name("scalars")) new_scalars_call = cst.Call(func=new_scalars_func, args=execute_call.args) new_postfix_attr = cst.Attribute(value=new_scalars_call, attr=cst.Name("one_or_none")) new_postfix_call = cst.Call(func=new_postfix_attr, args=[]) if original_node.lpar or original_node.rpar: new_postfix_call = new_postfix_call.with_changes(lpar=original_node.lpar, rpar=original_node.rpar) return new_postfix_call class ExecScalarsAllToScalars(_SessionMatcherMixin, cst.CSTTransformer): """*.session.execute(...).scalars().all() → *.session.scalars(...).all()""" def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression: # Must be .all() with no args if not isinstance(updated_node.func, cst.Attribute): return updated_node if not isinstance(updated_node.func.attr, cst.Name) or updated_node.func.attr.value != "all": return updated_node if updated_node.args: return updated_node # Receiver must be scalars() with no args scalars_call = updated_node.func.value if not self._is_attr_call_with_name(scalars_call, "scalars") or scalars_call.args: return updated_node # Whose receiver must be execute(...) execute_call = scalars_call.func.value if not self._is_attr_call_with_name(execute_call, "execute"): return updated_node # execute(...) must be on a session-like object session_obj = execute_call.func.value if not self._is_session_obj(session_obj): return updated_node # Build: <session_obj>.scalars(<execute args...>).all() new_scalars_func = cst.Attribute(value=session_obj, attr=cst.Name("scalars")) new_scalars_call = cst.Call(func=new_scalars_func, args=execute_call.args) new_postfix_attr = cst.Attribute(value=new_scalars_call, attr=cst.Name("all")) new_postfix_call = cst.Call(func=new_postfix_attr, args=[]) if original_node.lpar or original_node.rpar: new_postfix_call = new_postfix_call.with_changes(lpar=original_node.lpar, rpar=original_node.rpar) return new_postfix_call def rewrite_source(code: str) -> str: mod = cst.parse_module(code) # 顺序:先处理 first,再处理 one_or_none,最后 all mod = mod.visit(ExecScalarsFirstToScalar()) mod = mod.visit(ExecScalarsOneOrNoneToScalars()) mod = mod.visit(ExecScalarsAllToScalars()) return mod.code def color_diff(diff_lines): for line in diff_lines: if line.startswith("-"): yield f"\033[31m{line}\033[0m" # 红色 elif line.startswith("+"): yield f"\033[32m{line}\033[0m" # 绿色 else: yield line def process_file(fp: Path, write: bool) -> bool: try: src = fp.read_text(encoding="utf-8") except Exception: return False try: new_src = rewrite_source(src) except Exception: # skip files that fail to parse return False if src == new_src: return False if write: fp.write_text(new_src, encoding="utf-8") print(f"Updated: {fp}") else: print(f"DRY-RUN (would update): {fp}") diff = difflib.unified_diff( src.splitlines(keepends=True), new_src.splitlines(keepends=True), fromfile=str(fp), tofile=f"{fp} (rewritten)", n=3, ) for i, line in enumerate(color_diff(diff)): if i > 200: print("... (diff truncated)") break print(line, end="") return True def iter_py_files(root: Path): if root.is_file() and root.suffix == ".py": yield root return EXCLUDES = {".git", ".hg", ".svn", ".idea", ".vscode", "node_modules", "dist", "build", "__pycache__", ".venv", "venv"} for dirpath, dirnames, filenames in os.walk(root): dirnames[:] = [d for d in dirnames if d not in EXCLUDES] for name in filenames: if name.endswith(".py"): yield Path(dirpath) / name def main(): ap = argparse.ArgumentParser( description="Rewrite session.execute(...).scalars().{first,one_or_none,all}() using LibCST" ) ap.add_argument("path", help="Project root directory or a single .py file") ap.add_argument("--write", action="store_true", help="Apply changes (default: dry-run preview)") args = ap.parse_args() root = Path(args.path).resolve() if not root.exists(): raise SystemExit(f"Path not found: {root}") changed = 0 for fp in iter_py_files(root): if process_file(fp, write=args.write): changed += 1 if args.write: print(f"Done. Files updated: {changed}") else: print(f"Dry-run complete. Files that would change: {changed}") if __name__ == "__main__": main() ``` </details> ## Checklist - [ ] This change requires a documentation update, included: [Dify Document](https://github.com/langgenius/dify-docs) - [x] I understand that this PR may be closed in case there was no previous discussion or issues. (This doesn't apply to typos!) - [x] I've added a test for each change that was introduced, and I tried as much as possible to make a single atomic change. - [x] I've updated the documentation accordingly. - [x] I ran `dev/reformat`(backend) and `cd web && npx lint-staged`(frontend) to appease the lint gods
yindo added the pull-request label 2026-02-21 20:47:03 -05:00
yindo closed this issue 2026-02-21 20:47:03 -05:00
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: langgenius/dify#30203