diff --git a/mostlyai/sdk/_data/base.py b/mostlyai/sdk/_data/base.py index 9c29af8a..d5b9715e 100644 --- a/mostlyai/sdk/_data/base.py +++ b/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)] diff --git a/mostlyai/sdk/_data/context.py b/mostlyai/sdk/_data/context.py index ad40a56b..d3d319fd 100644 --- a/mostlyai/sdk/_data/context.py +++ b/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(), ) ) diff --git a/mostlyai/sdk/_data/pull_utils.py b/mostlyai/sdk/_data/pull_utils.py index a9d6be64..4c3512a6 100644 --- a/mostlyai/sdk/_data/pull_utils.py +++ b/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 diff --git a/tests/_data/unit/test_pull.py b/tests/_data/unit/test_pull.py index aedb4847..448df580 100644 --- a/tests/_data/unit/test_pull.py +++ b/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"}