fix: excluded PK on DB table

mostlyai/sdk/_data/base.py

@@ -188,13 +188,17 @@ def get_context_key(self, table_name: str) -> DataIdentifier | None:
             return ctx_rel.child
 
     def get_primary_key(self, table_name: str) -> DataIdentifier | None:
-        context_relations = self.get_child_context_relations(table_name)
-        # first, check primary key on table
+        primary_key = None
         if self.tables[table_name].primary_key is not None:
-            return DataIdentifier(table_name, self.tables[table_name].primary_key)
-        # fall back to checking primary key on schema
-        if context_relations:
-            return context_relations[0].parent
+            # first, check primary key on table
+            primary_key = DataIdentifier(table_name, self.tables[table_name].primary_key)
+        elif context_relations := self.get_child_context_relations(table_name):
+            # fall back to checking primary key on schema
+            primary_key = context_relations[0].parent
+        # lastly, check if the primary key is in the included columns
+        if primary_key and primary_key.column not in self.tables[table_name].columns:
+            primary_key = None
+        return primary_key
 
     def get_child_context_relations(self, parent_table: str) -> list[DataRelation]:
         return [rel for rel in self.relations if rel.parent.table == parent_table and isinstance(rel, ContextRelation)]

mostlyai/sdk/_data/context.py

@@ -257,9 +257,10 @@ def get_table_chain_to_tgt(schema: Schema, tables: list[str], tgt: str) -> tuple
     """
     sub_schema = schema.subset(relation_type=ContextRelation, tables=tables + [tgt])
     nodes, path = sub_schema.path_to(tgt)
+    tgt_pk = pk.column if (pk := schema.get_primary_key(tgt)) else None
     path.append(
         ContextRelation(
-            parent=DataIdentifier(table=tgt, column=schema.tables[tgt].primary_key),
+            parent=DataIdentifier(table=tgt, column=tgt_pk),
             child=DataIdentifier(),
         )
     )

mostlyai/sdk/_data/pull_utils.py

@@ -287,7 +287,7 @@ def write_json(data: dict, fn: Path) -> None:
     tgt_metadata_path = workspace_dir / "OriginalData" / "tgt-meta"
     tgt_metadata_path.mkdir(parents=True, exist_ok=True)
     # tgt-meta / keys.json
-    tgt_pk = tgt_table.primary_key
+    tgt_pk = pk.column if (pk := schema.get_primary_key(tgt)) else None
     for rel in schema.subset(relations_to=tgt).relations:
         if rel.child == DataIdentifier(tgt, tgt_pk):
             # having the same field as PK and FK is not supported

tests/_data/unit/test_pull.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import sqlite3
 from pathlib import Path
 from unittest import mock
 from unittest.mock import patch
@@ -27,6 +28,7 @@
     ForeignKey,
     Schema,
 )
+from mostlyai.sdk._data.db.sqlite import SqliteContainer, SqliteTable
 from mostlyai.sdk._data.dtype import (
     STRING,
     is_float_dtype,
@@ -2115,3 +2117,34 @@ def test_pull_empty_sequences(self, tmp_path):
         ctx_files = sorted([f.name for f in (tmp_path / "OriginalData" / "ctx-data").glob("*.parquet")])
         tgt_files = sorted([f.name for f in (tmp_path / "OriginalData" / "tgt-data").glob("*.parquet")])
         assert ctx_files == tgt_files
+
+
+class TestPullWithDB:
+    def test_excluded_primary_key(self, tmp_path):
+        db_path = tmp_path / "database.db"
+        with sqlite3.connect(str(db_path)) as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+            CREATE TABLE IF NOT EXISTS users (
+                id INTEGER PRIMARY KEY,
+                username TEXT NOT NULL,
+                email TEXT
+            )
+            """)
+            cursor.execute("INSERT INTO users (username, email) VALUES (?, ?)", ("alice", "alice@example.com"))
+            cursor.execute("INSERT INTO users (username, email) VALUES (?, ?)", ("bob", "bob@example.com"))
+        tables = {
+            "users": SqliteTable(
+                name="users",
+                container=SqliteContainer(dbname=str(db_path)),
+                primary_key=None,  # we specifically exclude the existing primary key in this test
+                columns=["username", "email"],
+            ),
+        }
+        schema = Schema(tables=tables)
+
+        pull(tgt="users", schema=schema, workspace_dir=tmp_path)
+
+        df = pd.read_parquet(tmp_path / "OriginalData" / "tgt-data")
+        assert len(df) == 2
+        assert set(df.columns) == {"username", "email"}