fix: allow user to unset auto-detected primary key in client mode

mostlyai/sdk/_data/base.py

@@ -189,7 +189,7 @@ def get_context_key(self, table_name: str) -> DataIdentifier | None:
 
     def get_primary_key(self, table_name: str) -> DataIdentifier | None:
         primary_key = None
-        if self.tables[table_name].primary_key is not None:
+        if self.tables[table_name].primary_key:
             # 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):
@@ -260,7 +260,7 @@ def get_scp_relations(self, table: str) -> list[ContextRelation]:
 
     def _update_key_encoding_types(self) -> None:
         for tbl_name, tbl_table in self.tables.items():
-            if tbl_table.primary_key is not None:
+            if tbl_table.primary_key:
                 if tbl_table.primary_key in tbl_table.encoding_types:
                     del tbl_table.encoding_types[tbl_table.primary_key]
             for rel in self.get_relations_to_table(tbl_name):

mostlyai/sdk/_data/db/base.py

@@ -876,6 +876,7 @@ class SqlAlchemyTable(DBTable, abc.ABC):
 
     def __init__(self, *args, **kwargs):
         self.is_view = kwargs.get("is_view", False)
+        self.lazy_fetch_primary_key = kwargs.get("lazy_fetch_primary_key", True)
         super().__init__(*args, **kwargs)
 
     def __repr__(self):
@@ -1065,7 +1066,7 @@ def _lazy_fetch(self, item: str) -> None:
         if item == "columns":
             self.columns = [c.name for c in self._sa_table.columns]
         elif item == "primary_key":
-            self.primary_key = self._get_primary_key()
+            self.primary_key = self._get_primary_key() if self.lazy_fetch_primary_key else None
         elif item == "dtypes":
             self.dtypes = self._get_dtypes()
         else:

mostlyai/sdk/_data/file/utils.py

@@ -88,7 +88,9 @@ def _fetch_file_data_table(
     return data_table
 
 
-def make_data_table_from_container(container: DataContainer, is_output=False) -> DataTable:
+def make_data_table_from_container(
+    container: DataContainer, is_output: bool = False, lazy_fetch_primary_key: bool = True
+) -> DataTable:
     if isinstance(container, SqlAlchemyContainer):
         # handle DB containers
         data_table_class = container.table_class()
@@ -100,4 +102,4 @@ def make_data_table_from_container(container: DataContainer, is_output=False) ->
             data_table_class = read_data_table_from_path(container, return_class=True)
     else:
         raise RuntimeError(f"Unknown container type: {type(container)}")
-    return data_table_class(container=container, is_output=is_output)
+    return data_table_class(container=container, is_output=is_output, lazy_fetch_primary_key=lazy_fetch_primary_key)

mostlyai/sdk/_data/pull_utils.py

@@ -448,7 +448,7 @@ def fetch_table_data(
     )
     n_fetched_rows = 0
     for idx, chunk_df in enumerate(iterator):
-        if deduplicate_pks and primary_key is not None and primary_key.column in chunk_df.columns:
+        if deduplicate_pks and primary_key and primary_key.column in chunk_df.columns:
             # consider only the first occurrence of each primary key
             if not chunk_df[primary_key.column].is_unique:
                 drop_idx = chunk_df[primary_key.column].duplicated()
@@ -568,7 +568,7 @@ def fetch_target_table(
             key_fraction_df[FRACTION] = keys[MAX_TGT_ROWS_PER_CTX_KEY] / schema.tables[tgt].row_count
             key_fraction_df = key_fraction_df.rename(columns={tgt_context_key.ref_name(): tgt_context_key.column})
             key_fraction_df = key_fraction_df.drop(columns=[MAX_TGT_ROWS_PER_CTX_KEY])
-    elif tgt_primary_key is not None:
+    elif tgt_primary_key:
         # flat setup with primary key
         # sample target table by target primary key
         # use provided keys to filter and order the target table (max_sample_size=None, do_shuffle=False)

mostlyai/sdk/_local/execution/step_pull_training_data.py

@@ -67,7 +67,7 @@ def _create_training_schema(generator: Generator, connectors: list[Connector]) -
         container = create_container_from_connector(connector)
         container.set_location(table.location)
         # create DataTable
-        data_table = make_data_table_from_container(container)
+        data_table = make_data_table_from_container(container, lazy_fetch_primary_key=False)
         data_table.name = table.name
         data_table.primary_key = table.primary_key
         if table.columns:

mostlyai/sdk/domain.py

@@ -2911,7 +2911,7 @@ def validate_keys_exists_in_columns(cls, values):
         if values.columns:
             column_names = {col.name for col in values.columns}
             pk = values.primary_key
-            if pk is not None and pk not in column_names:
+            if pk and pk not in column_names:
                 raise ValueError(f"Primary key column '{pk}' does not exist in the table's columns.")
             for fk in values.foreign_keys or []:
                 if fk.column not in column_names:
@@ -2923,7 +2923,7 @@ def validate_keys_exists_in_columns(cls, values):
     def validate_pk_and_fks_are_not_overlapping(cls, values):
         primary_key = values.primary_key
         foreign_keys = [fk.column for fk in values.foreign_keys or []]
-        if primary_key is not None and primary_key in foreign_keys:
+        if primary_key and primary_key in foreign_keys:
             raise ValueError(f"Column '{primary_key}' is both a primary key and a foreign key.")
         return values
 

tests/_data/unit/test_pull.py

@@ -2120,7 +2120,15 @@ def test_pull_empty_sequences(self, tmp_path):
 
 
 class TestPullWithDB:
-    def test_excluded_primary_key(self, tmp_path):
+    @pytest.mark.parametrize(
+        "columns,lazy_fetch_primary_key,expected_keys",
+        [
+            (["id", "username", "email"], True, {"primary_key": "id"}),
+            (["id", "username", "email"], False, {}),
+            (["username", "email"], True, {}),
+        ],
+    )
+    def test_excluded_primary_key(self, tmp_path, columns, lazy_fetch_primary_key, expected_keys):
         db_path = tmp_path / "database.db"
         with sqlite3.connect(str(db_path)) as conn:
             cursor = conn.cursor()
@@ -2137,14 +2145,18 @@ def test_excluded_primary_key(self, tmp_path):
             "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"],
+                primary_key=None,
+                columns=columns,
+                lazy_fetch_primary_key=lazy_fetch_primary_key,
             ),
         }
         schema = Schema(tables=tables)
 
         pull(tgt="users", schema=schema, workspace_dir=tmp_path)
 
+        keys = read_json(tmp_path / "OriginalData" / "tgt-meta" / "keys.json")
+        assert keys == expected_keys
+
         df = pd.read_parquet(tmp_path / "OriginalData" / "tgt-data")
         assert len(df) == 2
-        assert set(df.columns) == {"username", "email"}
+        assert set(df.columns) == set(columns)

tools/custom_template/pydantic_v2/BaseModel.jinja2

@@ -786,7 +786,7 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %}  # {{ comme
         if values.columns:
             column_names = {col.name for col in values.columns}
             pk = values.primary_key
-            if pk is not None and pk not in column_names:
+            if pk and pk not in column_names:
                 raise ValueError(f"Primary key column '{pk}' does not exist in the table's columns.")
             for fk in values.foreign_keys or []:
                 if fk.column not in column_names:
@@ -798,7 +798,7 @@ class {{ class_name }}({{ base_class }}):{% if comment is defined %}  # {{ comme
     def validate_pk_and_fks_are_not_overlapping(cls, values):
         primary_key = values.primary_key
         foreign_keys = [fk.column for fk in values.foreign_keys or []]
-        if primary_key is not None and primary_key in foreign_keys:
+        if primary_key and primary_key in foreign_keys:
             raise ValueError(f"Column '{primary_key}' is both a primary key and a foreign key.")
         return values
 {%- endif %}{%- if class_name == "SyntheticTableConfiguration" %}

tools/model.py

@@ -602,7 +602,7 @@ def validate_keys_exists_in_columns(cls, values):
         if values.columns:
             column_names = {col.name for col in values.columns}
             pk = values.primary_key
-            if pk is not None and pk not in column_names:
+            if pk and pk not in column_names:
                 raise ValueError(f"Primary key column '{pk}' does not exist in the table's columns.")
             for fk in values.foreign_keys or []:
                 if fk.column not in column_names:
@@ -614,7 +614,7 @@ def validate_keys_exists_in_columns(cls, values):
     def validate_pk_and_fks_are_not_overlapping(cls, values):
         primary_key = values.primary_key
         foreign_keys = [fk.column for fk in values.foreign_keys or []]
-        if primary_key is not None and primary_key in foreign_keys:
+        if primary_key and primary_key in foreign_keys:
             raise ValueError(f"Column '{primary_key}' is both a primary key and a foreign key.")
         return values