[baskets] Edit basket support (includes tests)

modelbaker/dbconnector/db_connector.py

@@ -349,6 +349,12 @@ def create_basket(self, dataset_tid, topic, tilitid_value=None):
         """
         return False, None
 
+    def edit_basket(self, basket_config: dict) -> tuple[bool, str]:
+        """
+        Returns the state and the errormessage
+        """
+        return False, None
+
     def get_tid_handling(self):
         """
         Returns `True` if a tid handling is enabled according to the settings table (when the database has been created with `--createTidCol`).

modelbaker/dbconnector/gpkg_connector.py

@@ -55,6 +55,7 @@ def __init__(self, uri, schema):
         self.tid = "T_Id"
         self.tilitid = "T_Ili_Tid"
         self.dispName = "dispName"
+        self.attachmentKey = "attachmentKey"
         self.basket_table_name = GPKG_BASKET_TABLE
         self.dataset_table_name = GPKG_DATASET_TABLE
 
@@ -976,6 +977,47 @@ def create_basket(self, dataset_tid, topic, tilitid_value=None):
                 ).format(topic, error_message)
         return False, self.tr('Could not create basket for topic "{}".').format(topic)
 
+    def edit_basket(self, basket_config: dict) -> tuple[bool, str]:
+        if self._table_exists(GPKG_BASKET_TABLE):
+            cursor = self.conn.cursor()
+            try:
+                cursor.execute(
+                    """
+                    UPDATE {basket_table}
+                    SET dataset = ?,
+                        {t_ili_tid} = ?,
+                        {attachment_key} = ?
+                    WHERE {tid_name} = ?
+                    """.format(
+                        basket_table=GPKG_BASKET_TABLE,
+                        t_ili_tid=self.tilitid,
+                        attachment_key=self.attachmentKey,
+                        tid_name=self.tid,
+                    ),
+                    (
+                        basket_config["dataset_t_id"],
+                        basket_config["bid_value"],
+                        basket_config["attachmentkey"],
+                        basket_config["basket_t_id"],
+                    ),
+                )
+                self.conn.commit()
+                cursor.close()
+                return True, self.tr(
+                    'Successfully edited basket for topic "{}" and dataset "{}".'
+                ).format(basket_config["topic"], basket_config["datasetname"])
+            except sqlite3.Error as e:
+                cursor.close()
+                error_message = " ".join(e.args)
+                return False, self.tr(
+                    'Could not edit basket for topic "{}" and dataset "{}": {}'
+                ).format(
+                    basket_config["topic"], basket_config["datasetname"], error_message
+                )
+        return False, self.tr(
+            'Could not edit basket for topic "{}" and dataset "{}"'
+        ).format(basket_config["topic"], basket_config["datasetname"])
+
     def get_tid_handling(self):
         if self._table_exists(GPKG_SETTINGS_TABLE):
             cursor = self.conn.cursor()

modelbaker/dbconnector/mssql_connector.py

@@ -51,6 +51,7 @@ def __init__(self, uri, schema):
         self.iliCodeName = "iliCode"
         self.tid = "T_Id"
         self.tilitid = "T_Ili_Tid"
+        self.attachmentKey = "attachmentkey"
         self.dispName = "dispName"
         self.basket_table_name = BASKET_TABLE
         self.dataset_table_name = DATASET_TABLE
@@ -1103,6 +1104,46 @@ def create_basket(self, dataset_tid, topic, tilitid_value=None):
                 ).format(topic, error_message)
         return False, self.tr('Could not create basket for topic "{}".').format(topic)
 
+    def edit_basket(self, basket_config: dict) -> tuple[bool, str]:
+        if self.schema and self._table_exists(BASKET_TABLE):
+            cur = self.conn.cursor()
+            try:
+                cur.execute(
+                    """
+                    UPDATE {schema}.{basket_table}
+                    SET dataset = ?,
+                        {t_ili_tid} = ?,
+                        {attachment_key} = ?
+                    WHERE {tid_name} = ?
+                    """.format(
+                        schema=self.schema,
+                        basket_table=BASKET_TABLE,
+                        t_ili_tid=self.tilitid,
+                        attachment_key=self.attachmentKey,
+                        tid_name=self.tid,
+                    ),
+                    (
+                        basket_config["dataset_t_id"],
+                        basket_config["bid_value"],
+                        basket_config["attachmentkey"],
+                        basket_config["basket_t_id"],
+                    ),
+                )
+                self.conn.commit()
+                return True, self.tr(
+                    'Successfully edited basket for topic "{}" and dataset "{}".'
+                ).format(basket_config["topic"], basket_config["datasetname"])
+            except pyodbc.errors.Error as e:
+                error_message = " ".join(e.args)
+                return False, self.tr(
+                    'Could not edit basket for topic "{}" and dataset "{}": {}'
+                ).format(
+                    basket_config["topic"], basket_config["datasetname"], error_message
+                )
+        return False, self.tr(
+            'Could not edit basket for topic "{}" and dataset "{}"'
+        ).format(basket_config["topic"], basket_config["datasetname"])
+
     def get_tid_handling(self):
         if self.schema and self._table_exists(SETTINGS_TABLE):
             cur = self.conn.cursor()

modelbaker/dbconnector/pg_connector.py

@@ -49,6 +49,7 @@ def __init__(self, uri, schema):
         self.iliCodeName = "ilicode"
         self.tid = "t_id"
         self.tilitid = "t_ili_tid"
+        self.attachmentKey = "attachmentkey"
         self.dispName = "dispname"
         self.basket_table_name = PG_BASKET_TABLE
         self.dataset_table_name = PG_DATASET_TABLE
@@ -1154,6 +1155,48 @@ def create_basket(self, dataset_tid, topic, tilitid_value=None):
                 ).format(topic, error_message)
         return False, self.tr('Could not create basket for topic "{}".').format(topic)
 
+    def edit_basket(self, basket_config: dict) -> tuple[bool, str]:
+        if self.schema and self._table_exists(PG_BASKET_TABLE):
+            cur = self.conn.cursor()
+            try:
+                cur.execute(
+                    sql.SQL(
+                        """
+                        UPDATE {schema}.{basket_table}
+                        SET dataset = %s,
+                            {t_ili_tid} = %s,
+                            {attachment_key} = %s
+                        WHERE {tid_name} = %s
+                        """
+                    ).format(
+                        schema=sql.Identifier(self.schema),
+                        basket_table=sql.Identifier(PG_BASKET_TABLE),
+                        t_ili_tid=sql.Identifier(self.tilitid),
+                        attachment_key=sql.Identifier(self.attachmentKey),
+                        tid_name=sql.Identifier(self.tid),
+                    ),
+                    (
+                        basket_config["dataset_t_id"],
+                        basket_config["bid_value"],
+                        basket_config["attachmentkey"],
+                        basket_config["basket_t_id"],
+                    ),
+                )
+                self.conn.commit()
+                return True, self.tr(
+                    'Successfully edited basket for topic "{}" and dataset "{}".'
+                ).format(basket_config["topic"], basket_config["datasetname"])
+            except psycopg2.errors.Error as e:
+                error_message = " ".join(e.args)
+                return False, self.tr(
+                    'Could not edit basket for topic "{}" and dataset "{}": {}'
+                ).format(
+                    basket_config["topic"], basket_config["datasetname"], error_message
+                )
+        return False, self.tr(
+            'Could not edit basket for topic "{}" and dataset "{}"'
+        ).format(basket_config["topic"], basket_config["datasetname"])
+
     def get_tid_handling(self):
         if self.schema and self._table_exists(PG_SETTINGS_TABLE):
             cur = self.conn.cursor()

tests/test_dataset_handling.py

@@ -296,6 +296,46 @@ def check_dataset_mutations(self, db_connector):
             record["datasetname"] for record in db_connector.get_datasets_info()
         } == {"Winti", "Seuzach", "Glarus West"}
 
+        # Edit basket for topic PipeBasketTest.Infrastructure and dataset Glarus West
+        baskets_info = db_connector.get_baskets_info()
+        for record in baskets_info:
+            if (
+                record["topic"] == "PipeBasketTest.Infrastructure"
+                and record["datasetname"] == "Glarus West"
+            ):
+                basket_info = record
+                break
+        basket_t_id = basket_info["basket_t_id"]
+        dataset_t_id = basket_info["dataset_t_id"]
+
+        # Info to be set to existing basket
+        basket_config = {
+            "topic": "PipeBasketTest.Infrastructure",
+            "basket_t_id": basket_t_id,
+            "bid_value": "3aa70ca6-13c6-482f-a415-a59694cfd658",
+            "attachmentkey": "my own attachment key",
+            "dataset_t_id": dataset_t_id,
+            "datasetname": "Glarus West",
+        }
+        res, msg = db_connector.edit_basket(basket_config)
+        assert res, msg
+
+        baskets_info = db_connector.get_baskets_info()
+        assert len(baskets_info) == 6
+        for record in baskets_info:
+            if record["basket_t_id"] == basket_t_id:
+                edited_basket_info = record
+                break
+        assert edited_basket_info["basket_t_id"] == basket_t_id
+        assert edited_basket_info["dataset_t_id"] == dataset_t_id
+        assert (
+            edited_basket_info["basket_t_ili_tid"]
+            == "3aa70ca6-13c6-482f-a415-a59694cfd658"
+        )
+        assert edited_basket_info["attachmentkey"] == "my own attachment key"
+        assert edited_basket_info["datasetname"] == "Glarus West"
+        assert edited_basket_info["topic"] == "PipeBasketTest.Infrastructure"
+
     def print_info(self, text):
         logging.info(text)