schematic.store.database.synapse_database

  1"""SynapseDatabase"""
  2
  3import pandas as pd
  4import synapseclient as sc  # type: ignore
  5
  6from schematic.store.database.synapse_database_wrapper import Synapse
  7from schematic.store.synapse_tracker import SynapseEntityTracker
  8
  9
 10class SynapseDatabaseMissingTableAnnotationsError(Exception):
 11    """Raised when a table is missing expected annotations"""
 12
 13    def __init__(self, message: str, table_name: str) -> None:
 14        self.message = message
 15        self.table_name = table_name
 16        super().__init__(self.message)
 17
 18    def __str__(self) -> str:
 19        return f"{self.message}; " f"name: {self.table_name};"
 20
 21
 22class InputDataframeMissingColumn(Exception):
 23    """Raised when an input dataframe is missing a needed column(s)"""
 24
 25    def __init__(
 26        self, message: str, table_columns: list[str], missing_columns: list[str]
 27    ) -> None:
 28        self.message = message
 29        self.table_columns = table_columns
 30        self.missing_columns = missing_columns
 31        super().__init__(self.message)
 32
 33    def __str__(self) -> str:
 34        return (
 35            f"{self.message}; "
 36            f"table_columns: {self.table_columns}; "
 37            f"missing_columns: {self.missing_columns}"
 38        )
 39
 40
 41class SynapseDatabase:
 42    """Represents a database stored as Synapse tables"""
 43
 44    def __init__(
 45        self,
 46        auth_token: str,
 47        project_id: str,
 48        synapse_entity_tracker: SynapseEntityTracker = None,
 49        syn: sc.Synapse = None,
 50    ) -> None:
 51        """Init
 52
 53        Args:
 54            auth_token (str): A Synapse auth_token
 55            project_id (str): A Synapse id for a project
 56            synapse_entity_tracker: Tracker for a pull-through cache of Synapse entities
 57        """
 58        self.synapse = Synapse(
 59            auth_token=auth_token,
 60            project_id=project_id,
 61            synapse_entity_tracker=synapse_entity_tracker,
 62            syn=syn,
 63        )
 64
 65    def upsert_table_rows(self, table_name: str, data: pd.DataFrame) -> None:
 66        """Upserts rows into the given table
 67
 68        Args:
 69            table_name (str): The name of the table to be upserted into.
 70            data (pd.DataFrame): The table the rows will come from
 71
 72        Raises:
 73            SynapseDatabaseMissingTableAnnotationsError: Raised when the table has no
 74             primary key annotation.
 75        """
 76        table_id = self.synapse.get_synapse_id_from_table_name(table_name)
 77        annotations = self.synapse.get_entity_annotations(table_id)
 78        if "primary_key" not in annotations:
 79            raise SynapseDatabaseMissingTableAnnotationsError(
 80                "Table has no primary_key annotation", table_name
 81            )
 82        primary_key = annotations["primary_key"][0]
 83        self._upsert_table_rows(table_id, data, primary_key)
 84
 85    def _upsert_table_rows(
 86        self, table_id: str, data: pd.DataFrame, primary_key: str
 87    ) -> None:
 88        """Upserts rows into the given table
 89
 90        Args:
 91            table_id (str): The Synapse id of the table to be upserted into.
 92            data (pd.DataFrame): The table the rows will come from
 93            primary_key (str): The primary key of the table used to identify
 94              which rows to update
 95
 96        Raises:
 97            InputDataframeMissingColumn: Raised when the input dataframe has
 98              no column that matches the primary key argument.
 99        """
100        if primary_key not in list(data.columns):
101            raise InputDataframeMissingColumn(
102                "Input dataframe missing primary key column.",
103                list(data.columns),
104                [primary_key],
105            )
106
107        table = self._create_primary_key_table(table_id, primary_key)
108        merged_table = pd.merge(
109            data, table, how="left", on=primary_key, validate="one_to_one"
110        )
111        self.synapse.upsert_table_rows(table_id, merged_table)
112
113    def _create_primary_key_table(
114        self, table_id: str, primary_key: str
115    ) -> pd.DataFrame:
116        """Creates a dataframe with just the primary key of the table
117
118        Args:
119            table_id (str): The id of the table to query
120            primary_key (str): The name of the primary key
121
122        Returns:
123            pd.DataFrame: The table in pandas.DataFrame form with the primary key, ROW_ID, and
124             ROW_VERSION columns
125
126        Raises:
127            InputDataframeMissingColumn: Raised when the synapse table has no column that
128              matches the primary key argument.
129        """
130        table = self.synapse.query_table(table_id, include_row_data=True)
131        if primary_key not in list(table.columns):
132            raise InputDataframeMissingColumn(
133                "Synapse table missing primary key column",
134                list(table.columns),
135                [primary_key],
136            )
137        table = table[["ROW_ID", "ROW_VERSION", primary_key]]
138        return table