Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion sdv/metadata/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -432,7 +432,14 @@ def add_column_relationship(
super().add_column_relationship(table_name, relationship_type, column_names)

def set_primary_key(self, column_name, table_name=None):
"""Set the primary key of a table."""
"""Set the primary key of a table.

Args:
column_name (str, list[str]):
Name (or list of names) of the primary key column(s).
table_name (str):
Name of the table to set the primary key.
"""
table_name = self._handle_table_name(table_name)
super().set_primary_key(table_name, column_name)

Expand Down
4 changes: 2 additions & 2 deletions sdv/metadata/multi_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -689,8 +689,8 @@ def set_primary_key(self, table_name, column_name):
Args:
table_name (str):
Name of the table to set the primary key.
column_name (str, tulple[str]):
Name (or tuple of names) of the primary key column(s).
column_name (str, list[str]):
Name (or list of names) of the primary key column(s).
"""
self._validate_table_exists(table_name)
self.tables[table_name].set_primary_key(column_name)
Expand Down
4 changes: 2 additions & 2 deletions sdv/metadata/single_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -855,8 +855,8 @@ def set_primary_key(self, column_name):
"""Set the metadata primary key.

Args:
column_name (str):
Name of the primary key column(s).
column_name (str, list[str]):
Name (or list of names) of the primary key column(s).
"""
if isinstance(column_name, list) and len(column_name) == 1:
column_name = column_name[0]
Expand Down
Loading