[DS-285]: Add id_name field to Couchbase connection_config (#252)

* Add id_name field to couchbase downloader config

* Couchbase downloader config id_name field; changelog and version update

* Include id_name in couchbase src e2e test

* Fix id_name for e2e test

* Fix spellcheck

* Use id instead of cbmid as id_name for couchbase src test

* parse string

* Rename id_name to collection_id

* Fix formatting

* Update description of Couchbase collection_id field

Author: mpolomdeepsense

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.3.4-dev1
+
+### Features
+
+* **Add `collection_id` field to Couchbase `downloader_config`**
+
 ## 0.3.3
 
 ### Enhancements

test_e2e/env_setup/couchbase/common/constants.env

@@ -4,4 +4,5 @@ CB_PASSWORD=password
 CB_BUCKET=unstructured
 CB_SCOPE=_default
 CB_COLLECTION=_default
-CB_INDEX_NAME=unstructured_test_search
+CB_INDEX_NAME=unstructured_test_search
+CB_COLLECTION_ID=id

test_e2e/env_setup/couchbase/source_connector/ingest_source_setup_cluster.py

@@ -17,6 +17,7 @@ class ClusterConfig:
     bucket_name: str
     scope_name: str
     collection_name: str
+    collection_id: str
 
 
 def get_client(cluster_config: ClusterConfig) -> Cluster:
@@ -33,6 +34,7 @@ def setup_cluster(cluster_config: ClusterConfig, source_file: str):
     bucket = cluster.bucket(cluster_config.bucket_name)
     scope = bucket.scope(cluster_config.scope_name)
     collection = scope.collection(cluster_config.collection_name)
+    collection_id = cluster_config.collection_id
 
     cluster.query(
         f"Create primary index on "
@@ -45,8 +47,7 @@ def setup_cluster(cluster_config: ClusterConfig, source_file: str):
         for line in file:
             try:
                 doc = json.loads(line)
-                # Assuming 'cbmid' is the document ID. Adjust as necessary.
-                doc_id = doc.get("cbmid", uuid.uuid4())
+                doc_id = str(doc.get(collection_id, uuid.uuid4()))
                 if doc_id:
                     collection.upsert(doc_id, doc)
                 else:
@@ -67,6 +68,7 @@ def setup_cluster(cluster_config: ClusterConfig, source_file: str):
     parser.add_argument("--bucket_name", required=True, help="Couchbase bucket name")
     parser.add_argument("--scope_name", required=True, help="Couchbase scope name")
     parser.add_argument("--collection_name", required=True, help="Couchbase collection name")
+    parser.add_argument("--collection_id", required=True, help="Couchbase collection id key")
     parser.add_argument("--source_file", required=True, help="Source file to ingest")
 
     args = parser.parse_args()
@@ -78,6 +80,7 @@ def setup_cluster(cluster_config: ClusterConfig, source_file: str):
         bucket_name=args.bucket_name,
         scope_name=args.scope_name,
         collection_name=args.collection_name,
+        collection_id=args.collection_id,
     )
 
     setup_cluster(config, args.source_file)

test_e2e/src/couchbase.sh

@@ -19,8 +19,8 @@ source "$SCRIPT_DIR"/cleanup.sh
 source "$SCRIPT_DIR"/env_setup/couchbase/common/constants.env
 
 # Check if all necessary environment variables are set
-if [ -z "$CB_USERNAME" ] || [ -z "$CB_PASSWORD" ] || [ -z "$CB_CONN_STR" ] || [ -z "$CB_BUCKET" ] || [ -z "$CB_SCOPE" ] || [ -z "$CB_COLLECTION" ]; then
-  echo "Error: One or more environment variables are not set. Please set CB_CONN_STR, CB_USERNAME, CB_PASSWORD, CB_BUCKET, CB_SCOPE, and CB_COLLECTION."
+if [ -z "$CB_USERNAME" ] || [ -z "$CB_PASSWORD" ] || [ -z "$CB_CONN_STR" ] || [ -z "$CB_BUCKET" ] || [ -z "$CB_SCOPE" ] || [ -z "$CB_COLLECTION" ] || [ -z "$CB_COLLECTION_ID" ]; then
+  echo "Error: One or more environment variables are not set. Please set CB_CONN_STR, CB_USERNAME, CB_PASSWORD, CB_BUCKET, CB_SCOPE, CB_COLLECTION and CB_COLLECTION_ID."
   exit 1
 fi
 
@@ -51,6 +51,7 @@ python "$SCRIPT_DIR"/env_setup/couchbase/source_connector/ingest_source_setup_cl
   --bucket_name "$CB_BUCKET" \
   --scope_name "$CB_SCOPE" \
   --collection_name "$CB_COLLECTION" \
+  --collection_id "$CB_COLLECTION_ID" \
   --source_file "$SCRIPT_DIR"/env_setup/couchbase/source_connector/airline_sample.jsonlines
 wait
 
@@ -69,6 +70,7 @@ PYTHONPATH=. ./unstructured_ingest/main.py \
   --password "$CB_PASSWORD" \
   --scope "$CB_SCOPE" \
   --collection "$CB_COLLECTION" \
+  --collection-id "$CB_COLLECTION_ID" \
   --work-dir "$WORK_DIR" \
   --preserve-downloads \
   --reprocess \

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.3"  # pragma: no cover
+__version__ = "0.3.4-dev1"  # pragma: no cover

unstructured_ingest/v2/processes/connectors/couchbase.py

@@ -219,6 +219,9 @@ def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
 
 
 class CouchbaseDownloaderConfig(DownloaderConfig):
+    collection_id: str = Field(
+        default="id", description="The unique key of the id field in the collection"
+    )
     fields: list[str] = field(default_factory=list)
 
 
@@ -250,7 +253,7 @@ def map_cb_results(self, cb_results: dict) -> str:
     def generate_download_response(
         self, result: dict, bucket: str, file_data: FileData
     ) -> DownloadResponse:
-        record_id = result["id"]
+        record_id = result[self.download_config.collection_id]
         filename_id = self.get_identifier(bucket=bucket, record_id=record_id)
         filename = f"{filename_id}.txt"
         download_path = self.download_dir / Path(filename)