support-custom-field-primary-key (piccolo-orm#136)

* draft

* restore PrimaryKey with deprecation

* fix lint and skip test_primary_key

* add support for sqlite versions before 3.35.0

* bump coverages

* modify sqlite coverage

* reverting coverage changes

* check version using sqlite3

* fix minor comments

* change example app migration files

The example app had some new migration files added by mistake. I've removed these, and made sure migrations exist for all tables.

* add failing tests

* add custom primary key definition

* revert insert callback

* Fix lint and some tests

* fix more tests

* fix remaining tests

* fix circular import error

When testing this branch with some Piccolo apps, I sometimes got a circular import error.

* make sure id column is assigned to Table class

* make sure foreign keys reference the correct primary key name

* fix bug with convert_to_sql_value

* black tweaks

Adding an explicit black version to `dev-requirements.txt`, which definitely supports Python 3.9.

* fixing a bug with `remove`

* improve the `remove` unittest

It was pretty terrible before, and let a bug slip through.

* making sure the foreign key type matches the referenced table's primary key type

* merge the primary key tests into a single file

* replace `primary` and `key` kwargs with `primary_key`

Seems to work fine in terms of being backwards compatible with existing migrations.

* draft migration change

* add automatic `id` to `non_default_columns` too

* create table with columns, and avoid adding columns

* add `id_column` to migration_manager_tests

* fix `test_migrations_manager` adding columns after create table

* add new migrations from `example_app` in `tests`

* undo latest `example_app` migration

* add `primary_key` docs

* add automatic `id` column to `default_columns`

* ignore lgtm error

Co-authored-by: Daniel Townsend <[email protected]>

Author: aminalaee

.gitignore

@@ -12,6 +12,7 @@ dist/
 piccolo.sqlite
 _build/
 .coverage
+coverage.xml
 *.sqlite
 htmlcov/
 prof/

docs/src/piccolo/schema/defining.rst

@@ -23,20 +23,40 @@ For a full list of columns, see :ref:`ColumnTypes`.
 
 -------------------------------------------------------------------------------
 
-Default columns
+Primary Key
 ---------------
 
-id
-~~
-
-Each table is automatically given a ``PrimaryKey`` column called ``id``, which
-is an auto incrementing integer.
+You can specify your ``PrimaryKey`` with any column type by passing ``primary_key`` to the ``Column``.
 
 It is used to uniquely identify a row, and is referenced by ``ForeignKey``
 columns on other tables.
 
-If you specify your own ``id`` column, you may get unexpected behaviour, so
-it's not recommended at the moment.
+.. code-block:: python
+
+    # tables.py
+    from piccolo.table import Table
+    from piccolo.columns import UUID, Varchar
+
+
+    class Band(Table):
+        id = UUID(primary_key=True)
+        name = Varchar(length=100)
+
+If you don't specify a ``PrimaryKey``, the table is automatically given a ``PrimaryKey`` column called ``id``, which
+is an auto incrementing integer.
+
+This is equivalent to:
+
+.. code-block:: python
+
+    # tables.py
+    from piccolo.table import Table
+    from piccolo.columns import Serial, Varchar
+
+
+    class Band(Table):
+        id = Serial(primary_key=True)
+        name = Varchar(length=100)
 
 -------------------------------------------------------------------------------
 

piccolo/apps/migrations/auto/migration_manager.py

@@ -569,20 +569,23 @@ async def _run_add_tables(self, backwards=False):
                 ).run()
         else:
             for add_table in self.add_tables:
+                add_columns: t.List[
+                    AddColumnClass
+                ] = self.add_columns.for_table_class_name(add_table.class_name)
                 _Table: t.Type[Table] = create_table_class(
                     class_name=add_table.class_name,
                     class_kwargs={"tablename": add_table.tablename},
                     class_members={
-                        column._meta.name: column
-                        for column in add_table.columns
+                        add_column.column._meta.name: add_column.column
+                        for add_column in add_columns
                     },
                 )
 
                 await _Table.create_table().run()
 
     async def _run_add_columns(self, backwards=False):
         """
-        Add columns, which belong to new and existing tables
+        Add columns, which belong to existing tables
         """
         if backwards:
             for add_column in self.add_columns.add_columns:
@@ -601,6 +604,9 @@ async def _run_add_columns(self, backwards=False):
                 await _Table.alter().drop_column(add_column.column).run()
         else:
             for table_class_name in self.add_columns.table_class_names:
+                if table_class_name in [i.class_name for i in self.add_tables]:
+                    continue  # No need to add columns to new tables
+
                 add_columns: t.List[
                     AddColumnClass
                 ] = self.add_columns.for_table_class_name(table_class_name)

piccolo/apps/user/tables.py

@@ -134,7 +134,7 @@ async def login(cls, username: str, password: str) -> t.Optional[int]:
         """
         query = (
             cls.select()
-            .columns(cls.id, cls.password)
+            .columns(cls._meta.primary_key, cls.password)
             .where((cls.username == username))
             .first()
         )

piccolo/columns/base.py

@@ -121,8 +121,7 @@ class ColumnMeta:
 
     # General attributes:
     null: bool = False
-    primary: bool = False
-    key: bool = False
+    primary_key: bool = False
     unique: bool = False
     index: bool = False
     index_method: IndexMethod = IndexMethod.btree
@@ -253,12 +252,9 @@ class Column(Selectable):
     :param null:
         Whether the column is nullable.
 
-    :param primary:
+    :param primary_key:
         If set, the column is used as a primary key.
 
-    :param key:
-        If set, the column is treated as a key.
-
     :param default:
         The column value to use if not specified by the user.
 
@@ -291,8 +287,7 @@ class Column(Selectable):
     def __init__(
         self,
         null: bool = False,
-        primary: bool = False,
-        key: bool = False,
+        primary_key: bool = False,
         unique: bool = False,
         index: bool = False,
         index_method: IndexMethod = IndexMethod.btree,
@@ -301,6 +296,12 @@ def __init__(
         choices: t.Optional[t.Type[Enum]] = None,
         **kwargs,
     ) -> None:
+        # This is for backwards compatibility - originally there were two
+        # separate arguments `primary` and `key`, but they have now been merged
+        # into `primary_key`.
+        if (kwargs.get("primary") is True) and (kwargs.get("key") is True):
+            primary_key = True
+
         # Used for migrations.
         # We deliberately omit 'required', and 'help_text' as they don't effect
         # the actual schema.
@@ -309,8 +310,7 @@ def __init__(
         kwargs.update(
             {
                 "null": null,
-                "primary": primary,
-                "key": key,
+                "primary_key": primary_key,
                 "unique": unique,
                 "index": index,
                 "index_method": index_method,
@@ -329,8 +329,7 @@ def __init__(
 
         self._meta = ColumnMeta(
             null=null,
-            primary=primary,
-            key=key,
+            primary_key=primary_key,
             unique=unique,
             index=index,
             index_method=index_method,
@@ -575,10 +574,8 @@ def querystring(self) -> QueryString:
         Used when creating tables.
         """
         query = f'"{self._meta.name}" {self.column_type}'
-        if self._meta.primary:
-            query += " PRIMARY"
-        if self._meta.key:
-            query += " KEY"
+        if self._meta.primary_key:
+            query += " PRIMARY KEY"
         if self._meta.unique:
             query += " UNIQUE"
         if not self._meta.null:
@@ -588,16 +585,18 @@ def querystring(self) -> QueryString:
             self, "_foreign_key_meta", None
         )
         if foreign_key_meta:
-            tablename = foreign_key_meta.resolved_references._meta.tablename
+            references = foreign_key_meta.resolved_references
+            tablename = references._meta.tablename
             on_delete = foreign_key_meta.on_delete.value
             on_update = foreign_key_meta.on_update.value
+            primary_key_name = references._meta.primary_key._meta.name
             query += (
-                f" REFERENCES {tablename} (id)"
+                f" REFERENCES {tablename} ({primary_key_name})"
                 f" ON DELETE {on_delete}"
                 f" ON UPDATE {on_update}"
             )
 
-        if not self._meta.primary:
+        if not self._meta.primary_key:
             default = self.get_default_value()
             sql_value = self.get_sql_value(value=default)
             # Escape the value if it contains a pair of curly braces, otherwise

piccolo/columns/column_types.py

@@ -2,6 +2,7 @@
 
 import copy
 import decimal
+import inspect
 import typing as t
 import uuid
 from datetime import date, datetime, time, timedelta
@@ -28,6 +29,7 @@
 from piccolo.columns.reference import LazyTableReference
 from piccolo.querystring import QueryString, Unquoted
 from piccolo.utils.encoding import dump_json
+from piccolo.utils.warnings import colored_warning
 
 if t.TYPE_CHECKING:  # pragma: no cover
     from piccolo.columns.base import ColumnMeta
@@ -493,20 +495,15 @@ def column_type(self):
 ###############################################################################
 
 
+DEFAULT = Unquoted("DEFAULT")
+NULL = Unquoted("null")
+
+
 class Serial(Column):
     """
     An alias to an autoincrementing integer column in Postgres.
     """
 
-    def __init__(self, **kwargs) -> None:
-        super().__init__(**kwargs)
-
-
-DEFAULT = Unquoted("DEFAULT")
-NULL = Unquoted("null")
-
-
-class PrimaryKey(Column):
     @property
     def column_type(self):
         engine_type = self._meta.table._meta.db.engine_type
@@ -524,10 +521,19 @@ def default(self):
             return NULL
         raise Exception("Unrecognized engine type")
 
+
+class PrimaryKey(Serial):
     def __init__(self, **kwargs) -> None:
         # Set the index to False, as a database should automatically create
         # an index for a PrimaryKey column.
-        kwargs.update({"primary": True, "key": True, "index": False})
+        kwargs.update({"primary_key": True, "index": False})
+
+        colored_warning(
+            "`PrimaryKey` is deprecated and "
+            "will be removed in future versions.",
+            category=DeprecationWarning,
+        )
+
         super().__init__(**kwargs)
 
 
@@ -974,9 +980,10 @@ class Float(Real):
 ###############################################################################
 
 
-class ForeignKey(Integer):
+class ForeignKey(Column):  # lgtm [py/missing-equals]
     """
-    Used to reference another table. Uses the ``int`` type for values.
+    Used to reference another table. Uses the same type as the primary key
+    column on the table it references.
 
     **Example**
 
@@ -1139,33 +1146,60 @@ class Band(Table):
 
     """  # noqa: E501
 
-    column_type = "INTEGER"
+    _foreign_key_meta: ForeignKeyMeta
+
+    @property
+    def column_type(self):
+        """
+        A ForeignKey column needs to have the same type as the primary key
+        column of the table being referenced.
+        """
+        referenced_table = self._foreign_key_meta.resolved_references
+        pk_column = referenced_table._meta.primary_key
+        if isinstance(pk_column, Serial):
+            return Integer().column_type
+        else:
+            return pk_column.column_type
 
     def __init__(
         self,
         references: t.Union[t.Type[Table], LazyTableReference, str],
-        default: t.Union[int, Enum, None] = None,
+        default: t.Any = None,
         null: bool = True,
         on_delete: OnDelete = OnDelete.cascade,
         on_update: OnUpdate = OnUpdate.cascade,
         **kwargs,
     ) -> None:
-        self._validate_default(default, (int, None))
+        from piccolo.table import Table
+
+        if inspect.isclass(references):
+            references = t.cast(t.Type, references)
+            if issubclass(references, Table):
+                # Using this to validate the default value - will raise a
+                # ValueError if incorrect.
+                if isinstance(references._meta.primary_key, Serial):
+                    Integer(default=default, null=null)
+                else:
+                    references._meta.primary_key.__class__(
+                        default=default, null=null
+                    )
+
+        self.default = default
 
         kwargs.update(
             {
                 "references": references,
                 "on_delete": on_delete,
                 "on_update": on_update,
+                "null": null,
             }
         )
-        super().__init__(default=default, null=null, **kwargs)
+
+        super().__init__(**kwargs)
 
         # This is here just for type inference - the actual value is set by
         # the Table metaclass. We can't set the actual value here, as
         # only the metaclass has access to the table.
-        from piccolo.table import Table
-
         self._foreign_key_meta = ForeignKeyMeta(
             Table, OnDelete.cascade, OnUpdate.cascade
         )

piccolo/engine/base.py

@@ -3,6 +3,7 @@
 import typing as t
 from abc import ABCMeta, abstractmethod
 
+from piccolo.querystring import QueryString
 from piccolo.utils.sync import run_sync
 from piccolo.utils.warnings import Level, colored_warning
 
@@ -41,6 +42,10 @@ async def prep_database(self):
     async def batch(self, query: Query, batch_size: int = 100) -> Batch:
         pass
 
+    @abstractmethod
+    async def run_querystring(self, querystring: QueryString, in_pool: bool):
+        pass
+
     async def check_version(self):
         """
         Warn if the database version isn't supported.