From 4f2b97781bebc34716e2e39c29c8af7cbffebff2 Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Tue, 17 Jun 2025 13:15:56 +0200 Subject: [PATCH 1/7] fix: excluded PK on DB table --- mostlyai/sdk/_data/base.py | 16 ++++++++++------ mostlyai/sdk/_data/context.py | 2 +- mostlyai/sdk/_data/pull_utils.py | 2 +- 3 files changed, 12 insertions(+), 8 deletions(-) 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..334b52bb 100644 --- a/mostlyai/sdk/_data/context.py +++ b/mostlyai/sdk/_data/context.py @@ -259,7 +259,7 @@ def get_table_chain_to_tgt(schema: Schema, tables: list[str], tgt: str) -> tuple nodes, path = sub_schema.path_to(tgt) path.append( ContextRelation( - parent=DataIdentifier(table=tgt, column=schema.tables[tgt].primary_key), + parent=DataIdentifier(table=tgt, column=schema.get_primary_key(tgt)), child=DataIdentifier(), ) ) diff --git a/mostlyai/sdk/_data/pull_utils.py b/mostlyai/sdk/_data/pull_utils.py index a9d6be64..23ab934a 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 = schema.get_primary_key(tgt) 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 From 8fe331a99c754e0c371250a18006232cc5aac635 Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Tue, 17 Jun 2025 18:07:26 +0200 Subject: [PATCH 2/7] fix: excluded PK on DB table --- mostlyai/sdk/_data/context.py | 2 +- mostlyai/sdk/_data/pull_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mostlyai/sdk/_data/context.py b/mostlyai/sdk/_data/context.py index 334b52bb..6d514867 100644 --- a/mostlyai/sdk/_data/context.py +++ b/mostlyai/sdk/_data/context.py @@ -259,7 +259,7 @@ def get_table_chain_to_tgt(schema: Schema, tables: list[str], tgt: str) -> tuple nodes, path = sub_schema.path_to(tgt) path.append( ContextRelation( - parent=DataIdentifier(table=tgt, column=schema.get_primary_key(tgt)), + parent=DataIdentifier(table=tgt, column=schema.get_primary_key(tgt).column), child=DataIdentifier(), ) ) diff --git a/mostlyai/sdk/_data/pull_utils.py b/mostlyai/sdk/_data/pull_utils.py index 23ab934a..d076ba17 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 = schema.get_primary_key(tgt) + tgt_pk = schema.get_primary_key(tgt).column 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 From 162c8fc9f1801c72ff91b6a969ed40daef375a6d Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Tue, 17 Jun 2025 18:13:51 +0200 Subject: [PATCH 3/7] fix: excluded PK on DB table --- mostlyai/sdk/_data/context.py | 3 ++- mostlyai/sdk/_data/pull_utils.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/mostlyai/sdk/_data/context.py b/mostlyai/sdk/_data/context.py index 6d514867..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.get_primary_key(tgt).column), + 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 d076ba17..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 = schema.get_primary_key(tgt).column + 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 From d2ec985954ed8952e3a1beefb8052642c5795b5f Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Tue, 17 Jun 2025 18:16:17 +0200 Subject: [PATCH 4/7] fix: excluded PK on DB table --- mostlyai/sdk/_data/context.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/mostlyai/sdk/_data/context.py b/mostlyai/sdk/_data/context.py index d3d319fd..286008e0 100644 --- a/mostlyai/sdk/_data/context.py +++ b/mostlyai/sdk/_data/context.py @@ -257,11 +257,5 @@ 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=tgt_pk), - child=DataIdentifier(), - ) - ) + path.append(ContextRelation(parent=schema.get_primary_key(tgt), child=DataIdentifier())) return nodes, path From 38bf228cc6c44eb4e6d7fcd61e35407b798b4894 Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Tue, 17 Jun 2025 18:17:40 +0200 Subject: [PATCH 5/7] fix: excluded PK on DB table --- mostlyai/sdk/_data/pull_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mostlyai/sdk/_data/pull_utils.py b/mostlyai/sdk/_data/pull_utils.py index 4c3512a6..4d5c66bd 100644 --- a/mostlyai/sdk/_data/pull_utils.py +++ b/mostlyai/sdk/_data/pull_utils.py @@ -287,13 +287,13 @@ 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 = pk.column if (pk := schema.get_primary_key(tgt)) else None + tgt_pk = schema.get_primary_key(tgt) for rel in schema.subset(relations_to=tgt).relations: - if rel.child == DataIdentifier(tgt, tgt_pk): + if rel.child == tgt_pk: # having the same field as PK and FK is not supported # FK is prioritised, thus removing PK role from it tgt_pk = None - if tgt_pk and tgt_pk in tgt_table.columns: + if tgt_pk and tgt_pk.column in tgt_table.columns: tgt_keys = {"primary_key": tgt_pk} else: tgt_keys = {} From 058e9f3bc1f5e9edece93bd6c9cf0a1f559934ad Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Tue, 17 Jun 2025 18:31:10 +0200 Subject: [PATCH 6/7] fix: excluded PK on DB table --- mostlyai/sdk/_data/context.py | 8 +++++++- mostlyai/sdk/_data/pull_utils.py | 6 +++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/mostlyai/sdk/_data/context.py b/mostlyai/sdk/_data/context.py index 286008e0..d3d319fd 100644 --- a/mostlyai/sdk/_data/context.py +++ b/mostlyai/sdk/_data/context.py @@ -257,5 +257,11 @@ 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) - path.append(ContextRelation(parent=schema.get_primary_key(tgt), child=DataIdentifier())) + tgt_pk = pk.column if (pk := schema.get_primary_key(tgt)) else None + path.append( + ContextRelation( + parent=DataIdentifier(table=tgt, column=tgt_pk), + child=DataIdentifier(), + ) + ) return nodes, path diff --git a/mostlyai/sdk/_data/pull_utils.py b/mostlyai/sdk/_data/pull_utils.py index 4d5c66bd..4c3512a6 100644 --- a/mostlyai/sdk/_data/pull_utils.py +++ b/mostlyai/sdk/_data/pull_utils.py @@ -287,13 +287,13 @@ 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 = schema.get_primary_key(tgt) + 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 == tgt_pk: + if rel.child == DataIdentifier(tgt, tgt_pk): # having the same field as PK and FK is not supported # FK is prioritised, thus removing PK role from it tgt_pk = None - if tgt_pk and tgt_pk.column in tgt_table.columns: + if tgt_pk and tgt_pk in tgt_table.columns: tgt_keys = {"primary_key": tgt_pk} else: tgt_keys = {} From 52ca17dea848f952eae8ed90aefd71ca5fbe0243 Mon Sep 17 00:00:00 2001 From: Lukasz Kolodziejczyk Date: Tue, 17 Jun 2025 18:46:49 +0200 Subject: [PATCH 7/7] fix: excluded PK on DB table --- tests/_data/unit/test_pull.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) 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"}