Merge pull request #1579 from datajoint/fix/fk-core-type-marker

fix: a column copied by a foreign key keeps the parent's core type

Author: dimitri-yatsenko

src/datajoint/declare.py

@@ -203,6 +203,7 @@ def compile_foreign_key(
     adapter,
     fk_attribute_map: dict[str, tuple[str, str]] | None = None,
     fk_index_candidates: list[list[str]] | None = None,
+    column_comments: dict[str, str] | None = None,
 ) -> None:
     """
     Parse a foreign key line and update declaration components.
@@ -233,6 +234,11 @@ def compile_foreign_key(
         candidate supporting index; the redundancy/coverage decision is made
         post-parse (once the full primary key and index list are known).
         Updated in place.
+    column_comments : dict, optional
+        Mapping of ``column -> comment`` for backends that set column comments
+        out of line (PostgreSQL's ``COMMENT ON COLUMN``). Each copied column's
+        comment, with its core-type marker (e.g. ``:int32:``), is added so the
+        child keeps the parent's declared type. Updated in place.
 
     Raises
     ------
@@ -288,14 +294,17 @@ def compile_foreign_key(
             # Enum type names start with "enum_" (generated hash-based names)
             if sql_type.startswith("enum_") and adapter.backend == "postgresql":
                 sql_type = f"{adapter.quote_identifier(ref.database)}.{adapter.quote_identifier(sql_type)}"
+            comment = parent_attr.sql_comment
             col_def = adapter.format_column_definition(
                 name=attr,
                 sql_type=sql_type,
                 nullable=is_nullable,
                 default=None,
-                comment=parent_attr.sql_comment,
+                comment=comment,
             )
             attr_sql.append(col_def)
+            if comment and column_comments is not None:
+                column_comments[attr] = comment
 
         # Track FK attribute mapping for lineage: child_attr -> (parent_table, parent_attr)
         if fk_attribute_map is not None:
@@ -459,6 +468,7 @@ def add_attribute(line: str, in_key: bool) -> None:
                 adapter,
                 fk_attribute_map,
                 fk_index_candidates,
+                column_comments,
             )
         elif re.match(r"^(unique\s+)?index\s*\(.*\)\s*(#.*)?$", line, re.I):  # index
             compile_index(re.sub(r"\s*#.*$", "", line), index_sql, adapter)

src/datajoint/heading.py

@@ -127,13 +127,23 @@ def sql_comment(self) -> str:
         """
         Return the full SQL comment including type markers.
 
+        A core type (``int32``, ``uuid``, ``float64``, ...) is recorded in the column
+        comment as ``:type:`` so it can be reconstructed from the database; heading
+        parsing strips the marker into ``original_type``. Restoring it here keeps the
+        core type on columns copied by a foreign key.
+
         Returns
         -------
         str
-            Comment with optional ``:uuid:`` prefix.
+            Comment with the core-type prefix, e.g. ``:int32:``, when there is one.
         """
-        # UUID info is stored in the comment for reconstruction
-        return (":uuid:" if self.uuid else "") + (self.comment or "")
+        if self.original_type:
+            marker = f":{self.original_type}:"
+        elif self.uuid:
+            marker = ":uuid:"
+        else:
+            marker = ""
+        return marker + (self.comment or "")
 
     @property
     def sql(self) -> str:

tests/integration/test_fk_core_types.py

@@ -0,0 +1,130 @@
+"""
+A column copied by a foreign key keeps the parent's core type.
+
+Core types (``int32``, ``int16``, ``uuid``, ...) are recorded in the column comment as
+``:type:``. A foreign key copies the parent's primary-key columns into the child; the copy
+must carry the same marker, or the child's heading reports the native SQL type and every
+table derived from it (a jobs table, a further child) inherits the loss.
+"""
+
+import logging
+
+import pytest
+
+import datajoint as dj
+
+
+class _Records(logging.Handler):
+    def __init__(self):
+        super().__init__(level=logging.WARNING)
+        self.messages = []
+
+    def emit(self, record):
+        self.messages.append(record.getMessage())
+
+
+@pytest.mark.backend_agnostic
+def test_fk_columns_keep_core_types(connection_by_backend, backend, prefix):
+    schema = dj.Schema(f"{prefix}_fk_core_types_{backend}", connection=connection_by_backend)
+    try:
+
+        @schema
+        class Subject(dj.Manual):
+            definition = """
+            subject_id : int32
+            """
+
+        @schema
+        class Session(dj.Manual):
+            definition = """
+            -> Subject
+            session_idx : int16
+            """
+
+        @schema
+        class Summary(dj.Manual):
+            definition = """
+            -> Session
+            ---
+            n : int32
+            """
+
+        assert Subject.heading["subject_id"].original_type == "int32"
+        # inherited one level down
+        assert Session.heading["subject_id"].original_type == "int32"
+        assert Session.heading["session_idx"].original_type == "int16"
+        # inherited two levels down
+        assert Summary.heading["subject_id"].original_type == "int32"
+        assert Summary.heading["session_idx"].original_type == "int16"
+        # the comment text itself is unchanged
+        assert Summary.heading["subject_id"].comment == Subject.heading["subject_id"].comment
+    finally:
+        schema.drop(prompt=False)
+
+
+@pytest.mark.backend_agnostic
+def test_fk_columns_keep_uuid(connection_by_backend, backend, prefix):
+    schema = dj.Schema(f"{prefix}_fk_core_types_uuid_{backend}", connection=connection_by_backend)
+    try:
+
+        @schema
+        class Item(dj.Manual):
+            definition = """
+            item_id : uuid       # identifier of the item
+            """
+
+        @schema
+        class Tag(dj.Manual):
+            definition = """
+            -> Item
+            tag : varchar(16)
+            """
+
+        assert Tag.heading["item_id"].uuid
+        assert Tag.heading["item_id"].original_type == "uuid"
+        assert Tag.heading["item_id"].comment == "identifier of the item"
+    finally:
+        schema.drop(prompt=False)
+
+
+@pytest.mark.backend_agnostic
+def test_jobs_table_keeps_inherited_core_types(connection_by_backend, backend, prefix):
+    """A jobs table built from inherited keys declares core types and warns about nothing."""
+    schema = dj.Schema(f"{prefix}_fk_core_types_jobs_{backend}", connection=connection_by_backend)
+    records = _Records()
+    logger = logging.getLogger("datajoint")
+    try:
+
+        @schema
+        class Subject(dj.Manual):
+            definition = """
+            subject_id : int32
+            """
+
+        @schema
+        class Session(dj.Manual):
+            definition = """
+            -> Subject
+            session_idx : int16
+            """
+
+        @schema
+        class Summary(dj.Computed):
+            definition = """
+            -> Session
+            ---
+            n : int32
+            """
+
+            def make(self, key):
+                self.insert1(dict(key, n=1))
+
+        logger.addHandler(records)
+        jobs = Summary.jobs
+        assert jobs.heading["subject_id"].original_type == "int32"
+        assert jobs.heading["session_idx"].original_type == "int16"
+        native = [m for m in records.messages if m.startswith("Native type")]
+        assert native == [], native
+    finally:
+        logger.removeHandler(records)
+        schema.drop(prompt=False)