Tùy chỉnh trình sinh mã SQLAlchemy với sqlacodegen: Hướng dẫn phát triển bộ tạo mở rộng

sqlacodegen là công cụ sinh mã mô hình SQLAlchemy từ cấu trúc cơ sở dữ liệu, hỗ trợ tự động hóa việc xây dựng lớp ORM. Bài viết này tập trung vào cách xây dựng và tích hợp các bộ tạo tùy chỉnh — không chỉ mở rộng chức năng sẵn có mà còn điều chỉnh hành vi sinh mã theo chuẩn kiến trúc dự án cụ thể.

Cấu trúc kiến trúc bộ tạo

Hệ thống dựa trên lớp trừu tượng CodeGenerator, nơi mọi bộ tạo đều kế thừa và triển khai phương thức generate(). Thiết kế hướng giao diện giúp dễ dàng thay thế hoặc mở rộng từng thành phần xử lý: bảng, cột, ràng buộc, hoặc toàn bộ lớp mô hình.

from abc import ABCMeta, abstractmethod
from sqlalchemy import MetaData
from sqlalchemy.engine import Connection, Engine

class CodeGenerator(metaclass=ABCMeta):
    valid_options: set[str] = set()

    def __init__(self, metadata: MetaData, bind: Connection | Engine, options: list[str]):
        self.metadata = metadata
        self.bind = bind
        self.options = set(options)

    @abstractmethod
    def generate(self) -> str:
        ...

Các bộ tạo được cung cấp sẵn bao gồm:

  • DeclarativeGenerator: Tạo lớp Python dùng Base__tablename__
  • DataclassGenerator: Dùng decorator @dataclassField cho kiểu dữ liệu rõ ràng
  • SQLModelGenerator: Tương thích với thư viện SQLModel của pydantic

Xây dựng bộ tạo tùy chỉnh

Để thêm logic sinh mã riêng, bạn tạo lớp con kế thừa từ một trong các bộ tạo hiện hữu — ví dụ DeclarativeGenerator — rồi ghi đè các phương thức cần thiết.

Bước 1: Khai báo lớp và tùy chọn dòng lệnh

Mở rộng tập hợp valid_options để kích hoạt xử lý tùy chọn mới khi chạy CLI:

from sqlacodegen.generators import DeclarativeGenerator

class AuditEnabledGenerator(DeclarativeGenerator):
    valid_options = DeclarativeGenerator.valid_options | {"audit_fields", "soft_delete"}

Bước 2: Ghi đè phương thức sinh lớp

Phương thức generate_class() là nơi kiểm soát toàn bộ nội dung lớp — tên, thuộc tính, phương thức, và chú thích:

from datetime import datetime
from sqlalchemy import Column, DateTime, Boolean, String

def generate_class(self, table: Table, columns: list[Column], constraints: list[Constraint]) -> str:
    base_code = super().generate_class(table, columns, constraints)
    
    if "audit_fields" in self.options:
        base_code += (
            "\n    created_at = Column(DateTime, default=datetime.utcnow)\n"
            "    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)\n"
        )
    
    if "soft_delete" in self.options:
        base_code += "\n    is_active = Column(Boolean, default=True)\n"
    
    return base_code

Bước 3: Đăng ký vào trình điều khiển CLI

Thêm đối số dòng lệnh mới bằng cách mở rộng parser trong tệp cli.py:

def configure_custom_parser(parser):
    parser.add_argument(
        "--audit",
        action="store_const",
        dest="generator",
        const=AuditEnabledGenerator,
        help="Sinh lớp với trường audit (created_at/updated_at)"
    )
    parser.add_argument(
        "--soft-delete",
        action="append_const",
        dest="options",
        const="soft_delete",
        help="Thêm cột is_active cho xóa mềm"
    )

configure_custom_parser(generate_parser())

Tối ưu hóa ánh xạ kiểu và quản lý import

Thay vì phụ thuộc hoàn toàn vào mặc định, bạn có thể kiểm soát chi tiết từng bước:

  • Ánh xạ kiểu cột: Ghi đè get_column_type() để chuyển đổi kiểu SQL sang annotation Python chính xác hơn:
def get_column_type(self, column: Column) -> str:
    if isinstance(column.type, JSON):
        return "dict | list"
    elif isinstance(column.type, Enum):
        return f"{column.type.enum_class.__name__}"
    return super().get_column_type(column)
  • Quản lý import: Thêm import động dựa trên tùy chọn:
def generate_imports(self) -> str:
    imports = super().generate_imports()
    if "audit_fields" in self.options:
        imports += "from datetime import datetime\n"
    if "soft_delete" in self.options:
        imports += "from sqlalchemy import Boolean\n"
    return imports

Xử lý trường hợp đặc biệt

Hỗ trợ view: Kích hoạt bằng thuộc tính views_supported và triển khai generate_view():

@property
def views_supported(self) -> bool:
    return True

def generate_view(self, view: Table) -> str:
    return f"class {self.class_name(view)}(Base):\n    __table__ = view\n"

Dùng lớp cơ sở tùy chỉnh: Ghi đè generate_class_definition() để thay đổi phần khai báo lớp:

def generate_class_definition(self, table: Table) -> str:
    base_class = "AuditBase" if "audit_fields" in self.options else "Base"
    return f"class {self.class_name(table)}({base_class}):"

Tích hợp vào quy trình CI/CD

Đặt lệnh sinh mã vào script tiền-commit hoặc pipeline để đảm bảo mô hình luôn đồng bộ với schema:

# .pre-commit-config.yaml
- repo: local
  hooks:
    - id: sync-models
      name: Cập nhật models.py từ DB
      entry: >
        sqlacodegen --audit --soft-delete postgresql://$DB_URL > models.py
      language: system
      pass_filenames: false
      files: ^$

Thẻ: sqlacodegen sqlalchemy python orm code-generation

Đăng vào ngày 14 tháng 6 lúc 05:49