feat[langchain-db2]: Added get_pks method, added clean_table function, update patch version (#84)

* feat: Added get_pks method, added clean_table function, update patch version

* Use TRUNCATE command instead of DELETE

Author: MateuszOssGit

libs/langchain-db2/langchain_db2/db2vs.py

@@ -150,6 +150,33 @@ def drop_table(client: Connection, table_name: str) -> None:
     return
 
 
+@_handle_exceptions
+def clear_table(client: Connection, table_name: str) -> None:
+    """Remove all records from the table using TRUNCATE.
+
+    Args:
+        client: The ibm_db_dbi connection object.
+        table_name: The name of the table to clear.
+    """
+    if not _table_exists(client, table_name):
+        logger.info(f"Table {table_name} not found…")
+        return
+
+    cursor = client.cursor()
+    ddl = f"TRUNCATE TABLE {table_name} IMMEDIATE"
+    try:
+        client.commit()
+        cursor.execute(ddl)
+        client.commit()
+        logger.info(f"Table {table_name} cleared successfully.")
+    except Exception:
+        client.rollback()
+        logger.exception(f"Failed to clear table {table_name}. Rolled back.")
+        raise
+    finally:
+        cursor.close()
+
+
 class DB2VS(VectorStore):
     """`DB2VS` vector store.
 
@@ -711,3 +738,28 @@ def from_texts(
         )
         vss.add_texts(texts=list(texts), metadatas=metadatas)
         return vss
+
+    @_handle_exceptions
+    def get_pks(self, expr: Optional[str] = None) -> List[str]:
+        """Get primary keys, optionally filtered by expr.
+
+        Args:
+            expr: SQL boolean expression to filter rows, e.g.:
+                  "id IN ('ABC123','DEF456')" or "title LIKE 'Abc%'".
+                  If None, returns all rows.
+        Returns:
+            List[str]: List of matching primary-key values.
+        """
+        sql = f"SELECT id FROM {self.table_name}"
+
+        if expr:
+            sql += f" WHERE {expr}"
+
+        cursor = self.client.cursor()
+        try:
+            cursor.execute(sql)
+            rows = cursor.fetchall()
+        finally:
+            cursor.close()
+
+        return [row[0] for row in rows]