From f5d7d56fa5d90d9634aac20dc650c484b23ed43f Mon Sep 17 00:00:00 2001
From: Waket Zheng <waketzheng@gmail.com>
Date: Fri, 27 Dec 2024 23:49:53 +0800
Subject: [PATCH] fix: inspectdb raise KeyError 'int2' for smallint (#401)

* fix: inspectdb raise KeyError 'int2' for smallint

* fix ci error

* no ask confirm for ci

* docs: update changelog
---
 .github/workflows/ci.yml     |  4 +++-
 CHANGELOG.md                 | 24 +++++++++++++++---------
 aerich/inspectdb/__init__.py |  6 +++---
 aerich/inspectdb/postgres.py |  5 +++--
 4 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 21c6008..6c495da 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -56,7 +56,9 @@ jobs:
         run: poetry run pip install --upgrade "tortoise-orm>=0.23,<0.24"
       - name: Install TortoiseORM develop branch
         if: matrix.tortoise-orm == 'tortoisedev'
-        run: poetry run pip install --upgrade "git+https://github.com/tortoise/tortoise-orm"
+        run: |
+          poetry run pip uninstall -y tortoise-orm
+          poetry run pip install --upgrade "git+https://github.com/tortoise/tortoise-orm"
       - name: CI
         env:
           MYSQL_PASS: root
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1a6bb90..504e1e9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,13 @@
 
 ## 0.8
 
+### [0.8.2]**(Unreleased)**
+
+#### Fixed
+- fix: inspectdb raise KeyError 'int2' for smallint. ([#401])
+
+[#401]: https://github.com/tortoise/aerich/pull/401
+
 ### [0.8.1](../../releases/tag/v0.8.1) - 2024-12-27
 
 #### Fixed
@@ -29,19 +36,18 @@
 [#395]: https://github.com/tortoise/aerich/pull/395
 [#394]: https://github.com/tortoise/aerich/pull/394
 [#393]: https://github.com/tortoise/aerich/pull/393
-[#376]: https://github.com/tortoise/aerich/pull/376
+[#392]: https://github.com/tortoise/aerich/pull/392
+[#388]: https://github.com/tortoise/aerich/pull/388
 [#386]: https://github.com/tortoise/aerich/pull/386
-[#272]: https://github.com/tortoise/aerich/pull/272
-[#334]: https://github.com/tortoise/aerich/pull/334
-[#284]: https://github.com/tortoise/aerich/pull/284
-[#286]: https://github.com/tortoise/aerich/pull/286
-[#302]: https://github.com/tortoise/aerich/pull/302
 [#378]: https://github.com/tortoise/aerich/pull/378
 [#377]: https://github.com/tortoise/aerich/pull/377
-[#271]: https://github.com/tortoise/aerich/pull/271
+[#376]: https://github.com/tortoise/aerich/pull/376
+[#334]: https://github.com/tortoise/aerich/pull/334
+[#302]: https://github.com/tortoise/aerich/pull/302
 [#286]: https://github.com/tortoise/aerich/pull/286
-[#388]: https://github.com/tortoise/aerich/pull/388
-[#392]: https://github.com/tortoise/aerich/pull/392
+[#284]: https://github.com/tortoise/aerich/pull/284
+[#272]: https://github.com/tortoise/aerich/pull/272
+[#271]: https://github.com/tortoise/aerich/pull/271
 
 ### [0.8.0](../../releases/tag/v0.8.0) - 2024-12-04
 
diff --git a/aerich/inspectdb/__init__.py b/aerich/inspectdb/__init__.py
index 7cdd86f..5b5e960 100644
--- a/aerich/inspectdb/__init__.py
+++ b/aerich/inspectdb/__init__.py
@@ -56,7 +56,7 @@ def translate(self) -> ColumnInfoDict:
                 length = ", ".join(length_parts) + ", "
         if self.null:
             null = "null=True, "
-        if self.default is not None:
+        if self.default is not None and not self.pk:
             if self.data_type in ("tinyint", "INT"):
                 default = f"default={'True' if self.default == '1' else 'False'}, "
             elif self.data_type == "bool":
@@ -158,11 +158,11 @@ def char_field(cls, **kwargs) -> str:
 
     @classmethod
     def int_field(cls, **kwargs) -> str:
-        return "{name} = fields.IntField({pk}{index}{comment})".format(**kwargs)
+        return "{name} = fields.IntField({pk}{index}{default}{comment})".format(**kwargs)
 
     @classmethod
     def smallint_field(cls, **kwargs) -> str:
-        return "{name} = fields.SmallIntField({pk}{index}{comment})".format(**kwargs)
+        return "{name} = fields.SmallIntField({pk}{index}{default}{comment})".format(**kwargs)
 
     @classmethod
     def bigint_field(cls, **kwargs) -> str:
diff --git a/aerich/inspectdb/postgres.py b/aerich/inspectdb/postgres.py
index 19fa031..f061878 100644
--- a/aerich/inspectdb/postgres.py
+++ b/aerich/inspectdb/postgres.py
@@ -16,12 +16,13 @@ def __init__(self, conn: "BasePostgresClient", tables: list[str] | None = None)
     @property
     def field_map(self) -> FieldMapDict:
         return {
+            "int2": self.smallint_field,
             "int4": self.int_field,
-            "int8": self.int_field,
+            "int8": self.bigint_field,
             "smallint": self.smallint_field,
+            "bigint": self.bigint_field,
             "varchar": self.char_field,
             "text": self.text_field,
-            "bigint": self.bigint_field,
             "timestamptz": self.datetime_field,
             "float4": self.float_field,
             "float8": self.float_field,