feat: support psycopg (#425)

Author: waketzheng

.github/workflows/ci.yml

@@ -89,14 +89,21 @@ jobs:
           POSTGRES_PORT: 5432
         run: poetry run make _testall
       - name: Verify aiomysql support
+        # Only check the latest version of tortoise
         if: matrix.tortoise-orm == 'tortoise024'
         run: |
             poetry run pip uninstall -y asyncmy
-            poetry run pip install aiomysql
             poetry run make test_mysql
-            poetry run pip uninstall -y aiomysql
             poetry run pip install asyncmy
         env:
           MYSQL_PASS: root
           MYSQL_HOST: 127.0.0.1
           MYSQL_PORT: 3306
+      - name: Verify psycopg support
+        # Only check the latest version of tortoise
+        if: matrix.tortoise-orm == 'tortoise024'
+        run: poetry run make test_psycopg
+        env:
+          POSTGRES_PASS: 123456
+          POSTGRES_HOST: 127.0.0.1
+          POSTGRES_PORT: 5432

CHANGELOG.md

@@ -5,10 +5,11 @@
 ### [0.8.2]**(Unreleased)**
 
 #### Added
+- feat: support psycopg. ([#425])
+- Support run `poetry add aerich` in project that inited by poetry v2. ([#424])
 - feat: support command `python -m aerich`. ([#417])
 - feat: add --fake to upgrade/downgrade. ([#398])
 - Support ignore table by settings `managed=False` in `Meta` class. ([#397])
-- Support run `poetry add aerich` in project that inited by poetry v2. ([#424])
 
 #### Fixed
 - fix: aerich migrate raises tortoise.exceptions.FieldError when `index.INDEX_TYPE` is not empty. ([#415])
@@ -27,6 +28,7 @@
 [#415]: https://github.com/tortoise/aerich/pull/415
 [#417]: https://github.com/tortoise/aerich/pull/417
 [#424]: https://github.com/tortoise/aerich/pull/424
+[#425]: https://github.com/tortoise/aerich/pull/425
 
 ### [0.8.1](../../releases/tag/v0.8.1) - 2024-12-27
 

Makefile

@@ -37,6 +37,9 @@ test_mysql:
 test_postgres:
 	$(py_warn) TEST_DB="postgres://postgres:$(POSTGRES_PASS)@$(POSTGRES_HOST):$(POSTGRES_PORT)/test_\{\}" pytest -vv -s
 
+test_psycopg:
+	$(py_warn) TEST_DB="psycopg://postgres:$(POSTGRES_PASS)@$(POSTGRES_HOST):$(POSTGRES_PORT)/test_\{\}" pytest -vv -s
+
 _testall: test_sqlite test_postgres test_mysql
 testall: deps _testall
 

aerich/coder.py

@@ -22,22 +22,25 @@ def default(self, obj) -> Any:
             return super().default(obj)
 
 
+def object_hook(obj) -> Any:
+    if (type_ := obj.get("type")) and type_ == "index" and (val := obj.get("val")):
+        return pickle.loads(base64.b64decode(val))  # nosec: B301
+    return obj
+
+
 def load_index(obj: dict) -> Index:
     """Convert a dict that generated by `Index.decribe()` to a Index instance"""
-    index = Index(fields=obj["fields"] or obj["expressions"], name=obj.get("name"))
+    try:
+        index = Index(fields=obj["fields"] or obj["expressions"], name=obj.get("name"))
+    except KeyError:
+        return object_hook(obj)
     if extra := obj.get("extra"):
         index.extra = extra
     if idx_type := obj.get("type"):
         index.INDEX_TYPE = idx_type
     return index
 
 
-def object_hook(obj) -> Any:
-    if (type_ := obj.get("type")) and type_ == "index" and (val := obj.get("val")):
-        return pickle.loads(base64.b64decode(val))  # nosec: B301
-    return obj
-
-
 def encoder(obj: dict) -> str:
     return json.dumps(obj, cls=JsonEncoder)
 

aerich/ddl/postgres/__init__.py

@@ -3,14 +3,14 @@
 from typing import cast
 
 from tortoise import Model
-from tortoise.backends.asyncpg.schema_generator import AsyncpgSchemaGenerator
+from tortoise.backends.base_postgres.schema_generator import BasePostgresSchemaGenerator
 
 from aerich.ddl import BaseDDL
 
 
 class PostgresDDL(BaseDDL):
-    schema_generator_cls = AsyncpgSchemaGenerator
-    DIALECT = AsyncpgSchemaGenerator.DIALECT
+    schema_generator_cls = BasePostgresSchemaGenerator
+    DIALECT = BasePostgresSchemaGenerator.DIALECT
     _ADD_INDEX_TEMPLATE = 'CREATE {unique}INDEX IF NOT EXISTS "{index_name}" ON "{table_name}" {index_type}({column_names}){extra}'
     _DROP_INDEX_TEMPLATE = 'DROP INDEX IF EXISTS "{index_name}"'
     _ALTER_NULL_TEMPLATE = 'ALTER TABLE "{table_name}" ALTER COLUMN "{column}" {set_drop} NOT NULL'

aerich/inspectdb/postgres.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from typing import TYPE_CHECKING
 
 from aerich.inspectdb import Column, FieldMapDict, Inspect
@@ -60,6 +61,8 @@ async def get_columns(self, table: str) -> list[Column]:
 where c.table_catalog = $1
   and c.table_name = $2
   and c.table_schema = $3"""  # nosec:B608
+        if "psycopg" in str(type(self.conn)).lower():
+            sql = re.sub(r"\$[123]", "%s", sql)
         ret = await self.conn.execute_query_dict(sql, [self.database, table, self.schema])
         for row in ret:
             columns.append(

conftest.py

@@ -8,7 +8,7 @@
 
 import pytest
 from tortoise import Tortoise, expand_db_url
-from tortoise.backends.asyncpg.schema_generator import AsyncpgSchemaGenerator
+from tortoise.backends.base_postgres.schema_generator import BasePostgresSchemaGenerator
 from tortoise.backends.mysql.schema_generator import MySQLSchemaGenerator
 from tortoise.backends.sqlite.schema_generator import SqliteSchemaGenerator
 from tortoise.contrib.test import MEMORY_SQLITE
@@ -64,7 +64,7 @@ async def initialize_tests(event_loop, request) -> None:
         Migrate.ddl = MysqlDDL(client)
     elif client.schema_generator is SqliteSchemaGenerator:
         Migrate.ddl = SqliteDDL(client)
-    elif client.schema_generator is AsyncpgSchemaGenerator:
+    elif issubclass(client.schema_generator, BasePostgresSchemaGenerator):
         Migrate.ddl = PostgresDDL(client)
     Migrate.dialect = Migrate.ddl.DIALECT
     request.addfinalizer(lambda: event_loop.run_until_complete(Tortoise._drop_databases()))