Removed additional add_document function from couchbaseserver.py and modified overwritten changes to non-edge server code

Author: barkha06

client/src/cbltest/api/couchbaseserver.py

@@ -15,15 +15,14 @@
     BucketAlreadyExistsException,
     BucketDoesNotExistException,
     CollectionAlreadyExistsException,
-    DocumentExistsException,
     DocumentNotFoundException,
     QueryIndexAlreadyExistsException,
     ScopeAlreadyExistsException,
 )
 from couchbase.management.buckets import CreateBucketSettings
-from couchbase.management.collections import CollectionSpec
 from couchbase.management.options import CreatePrimaryQueryIndexOptions
 from couchbase.options import ClusterOptions
+from couchbase.subdocument import upsert
 from opentelemetry.trace import get_tracer
 
 from cbltest.api.error import CblTestError
@@ -86,7 +85,7 @@ def create_collections(self, bucket: str, scope: str, names: list[str]) -> None:
                 ):
                     try:
                         if name != "_default":
-                            c.create_collection(CollectionSpec(name, scope))
+                            c.create_collection(scope_name=scope, collection_name=name)
                     except CollectionAlreadyExistsException:
                         pass
 
@@ -258,37 +257,6 @@ def run_query(
 
             return list(dict(result) for result in query_obj.execute())
 
-    def add_document(
-        self, bucket_name: str, document_id: str, document_data: dict
-    ) -> None:
-        """
-        Adds a document to the specified bucket in the Couchbase cluster.
-
-        :param bucket_name: The name of the bucket where the document will be added
-        :param document_id: The unique ID of the document
-        :param document_data: The content of the document as a dictionary
-        :raises DocumentExistsException: If a document with the same ID already exists
-        """
-        with self.__tracer.start_as_current_span("add_document") as span:
-            try:
-                bucket = self.__cluster.bucket(bucket_name)
-                collection = bucket.default_collection()
-
-                collection.insert(document_id, document_data)
-                span.set_attribute("couchbase.bucket", bucket_name)
-                span.set_attribute("couchbase.document_id", document_id)
-
-                print(
-                    f"Document '{document_id}' successfully added to bucket '{bucket_name}'."
-                )
-            except DocumentExistsException:
-                print(
-                    f"Document with ID '{document_id}' already exists in bucket '{bucket_name}'."
-                )
-            except Exception as e:
-                print(f"An error occurred while adding document '{document_id}': {e}")
-                raise
-
     def upsert_document(
         self,
         bucket: str,
@@ -324,6 +292,151 @@ def upsert_document(
                     f"Failed to insert document '{doc_id}' into {bucket}.{scope}.{collection}: {e}"
                 )
 
+    def delete_document(
+        self,
+        bucket: str,
+        doc_id: str,
+        scope: str = "_default",
+        collection: str = "_default",
+    ) -> None:
+        """
+        Deletes a document from the specified bucket.scope.collection.
+        """
+        with self.__tracer.start_as_current_span(
+            "delete_document",
+            attributes={
+                "cbl.bucket.name": bucket,
+                "cbl.scope.name": scope,
+                "cbl.collection.name": collection,
+                "cbl.document.id": doc_id,
+            },
+        ):
+            try:
+                bucket_obj = _try_n_times(10, 1, False, self.__cluster.bucket, bucket)
+                coll = bucket_obj.scope(scope).collection(collection)
+                coll.remove(doc_id)
+            except DocumentNotFoundException:
+                pass
+            except Exception as e:
+                raise CblTestError(
+                    f"Failed to delete document '{doc_id}' from {bucket}.{scope}.{collection}: {e}"
+                )
+
+    def get_document(
+        self,
+        bucket: str,
+        doc_id: str,
+        scope: str = "_default",
+        collection: str = "_default",
+    ) -> dict | None:
+        """
+        Gets a document from the specified bucket.scope.collection.
+
+        :param bucket: The bucket name.
+        :param doc_id: The document ID.
+        :param scope: The scope name.
+        :param collection: The collection name.
+        :return: The document content as a dictionary, or None if not found.
+        """
+        with self.__tracer.start_as_current_span(
+            "get_document",
+            attributes={
+                "cbl.bucket.name": bucket,
+                "cbl.scope.name": scope,
+                "cbl.collection.name": collection,
+                "cbl.document.id": doc_id,
+            },
+        ):
+            try:
+                bucket_obj = _try_n_times(10, 1, False, self.__cluster.bucket, bucket)
+                coll = bucket_obj.scope(scope).collection(collection)
+                result = coll.get(doc_id)
+                return result.content_as[dict] if result else None
+            except DocumentNotFoundException:
+                return None
+            except Exception as e:
+                raise CblTestError(
+                    f"Failed to get document '{doc_id}' from {bucket}.{scope}.{collection}: {e}"
+                )
+
+    def upsert_document_xattr(
+        self,
+        bucket: str,
+        doc_id: str,
+        xattr_key: str,
+        xattr_value: str,
+        scope: str = "_default",
+        collection: str = "_default",
+    ) -> None:
+        """
+        Upserts an xattr on a document using subdocument operations
+
+        :param bucket: The bucket containing the document
+        :param doc_id: The ID of the document to update
+        :param xattr_key: The xattr key to upsert
+        :param xattr_value: The value to set for the xattr
+        :param scope: The scope containing the document (default '_default')
+        :param collection: The collection containing the document (default '_default')
+        """
+        with self.__tracer.start_as_current_span(
+            "upsert_document_xattr",
+            attributes={
+                "cbl.bucket": bucket,
+                "cbl.scope": scope,
+                "cbl.collection": collection,
+                "cbl.document.id": doc_id,
+                "cbl.xattr.key": xattr_key,
+            },
+        ):
+            try:
+                col = self.__cluster.bucket(bucket).scope(scope).collection(collection)
+                col.mutate_in(
+                    doc_id,
+                    [upsert(xattr_key, xattr_value, xattr=True, create_parents=True)],
+                )
+            except Exception as e:
+                raise CblTestError(
+                    f"Failed to upsert xattr '{xattr_key}' on document '{doc_id}' in {bucket}.{scope}.{collection}: {e}"
+                )
+
+    def delete_document_xattr(
+        self,
+        bucket: str,
+        doc_id: str,
+        xattr_key: str,
+        scope: str = "_default",
+        collection: str = "_default",
+    ) -> None:
+        """
+        Deletes an xattr from a document using subdocument operations
+
+        :param bucket: The bucket containing the document
+        :param doc_id: The ID of the document
+        :param xattr_key: The xattr key to delete
+        :param scope: The scope containing the document (default '_default')
+        :param collection: The collection containing the document (default '_default')
+        """
+        with self.__tracer.start_as_current_span(
+            "delete_document_xattr",
+            attributes={
+                "cbl.bucket": bucket,
+                "cbl.scope": scope,
+                "cbl.collection": collection,
+                "cbl.document.id": doc_id,
+                "cbl.xattr.key": xattr_key,
+            },
+        ):
+            try:
+                from couchbase.subdocument import remove
+
+                col = self.__cluster.bucket(bucket).scope(scope).collection(collection)
+                col.mutate_in(
+                    doc_id,
+                    [remove(xattr_key, xattr=True)],
+                )
+            except Exception:
+                pass
+
     def start_xdcr(self, target: "CouchbaseServer", bucket_name: str) -> None:
         """
         Starts an XDCR replication from this cluster to the target cluster
@@ -479,4 +592,4 @@ def stop_xcdr(self, target: "CouchbaseServer", bucket_name: str) -> None:
                     resp = session.delete(
                         f"http://{self.__hostname}:8091/controller/cancelXDCR/{encoded}",
                     )
-                    resp.raise_for_status()
+                    resp.raise_for_status()

client/src/cbltest/api/jsongenerator.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from json import dumps, loads
 from pathlib import Path
-from typing import Any, Optional, cast
+from typing import Any, cast
 from urllib.parse import urljoin
 
 import requests
@@ -390,7 +390,7 @@ def __init__(
         )
 
     def _create_session(
-        self, secure: bool, scheme: str, url: str, port: int, auth: Optional[BasicAuth]
+        self, secure: bool, scheme: str, url: str, port: int, auth: BasicAuth | None
     ) -> ClientSession:
         if secure:
             ssl_context = ssl.create_default_context(cadata=_SGW_CA_CERT)
@@ -408,9 +408,9 @@ async def _send_request(
         self,
         method: str,
         path: str,
-        payload: Optional[JSONSerializable] = None,
-        params: Optional[dict[str, str]] = None,
-        session: Optional[ClientSession] = None,
+        payload: JSONSerializable | None = None,
+        params: dict[str, str] | None = None,
+        session: ClientSession | None = None,
     ) -> Any:
         if session is None:
             session = self.__admin_session
@@ -955,7 +955,7 @@ async def get_document(
         doc_id: str,
         scope: str = "_default",
         collection: str = "_default",
-    ) -> Optional[RemoteDocument]:
+    ) -> RemoteDocument | None:
         """
         Gets a document from Sync Gateway
 

client/src/cbltest/api/syncgateway.py

@@ -34,7 +34,7 @@ async def test_blobs_create_delete(
                 "channels": ["public"],
                 "timestamp": datetime.utcnow().isoformat(),
             }
-            server.add_document(bucket_name, doc["id"], doc)
+            server.upsert_document(bucket_name, doc["id"], doc)
         logger.info("2 documents created in Couchbase Server.")
 
         self.mark_test_step(
@@ -753,7 +753,7 @@ async def test_blob_special_characters(
                 "channels": ["public"],
                 "timestamp": datetime.utcnow().isoformat(),
             }
-            server.add_document(bucket_name, doc["id"], doc)
+            server.upsert_document(bucket_name, doc["id"], doc)
         logger.info("2 documents created in Couchbase Server.")
 
         self.mark_test_step(

tests/QE/edge_server/test_blobs.py

@@ -33,7 +33,7 @@ async def test_changes_feed_longpoll(
                 "channels": ["public"],
                 "timestamp": datetime.utcnow().isoformat(),
             }
-            server.add_document(bucket_name, doc["id"], doc)
+            server.upsert_document(bucket_name, doc["id"], doc)
         logger.info("5 documents created in Couchbase Server.")
 
         self.mark_test_step(

tests/QE/edge_server/test_changes_feed.py

@@ -11,7 +11,7 @@
 from cbltest.api.cloud import CouchbaseCloud
 from cbltest.api.edgeserver import BulkDocOperation
 from cbltest.api.httpclient import ClientFactory, HTTPClient
-from cbltest.api.jsongenerator import JSONGenerator
+from cbltest.api.json_generator import JSONGenerator
 
 
 class TestEdgeServerChaos(CBLTestClass):

tests/QE/edge_server/test_chaos_scenarios.py

@@ -32,7 +32,7 @@ async def test_conflicts(self, cblpytest: CBLPyTest, dataset_path: Path) -> None
                 "channels": ["public"],
                 "timestamp": datetime.utcnow().isoformat(),
             }
-            server.add_document(bucket_name, doc["id"], doc)
+            server.upsert_document(bucket_name, doc["id"], doc)
         logger.info("10 documents created in Couchbase Server.")
 
         self.mark_test_step(