feat(ingest): Couchbase source extra error handling

metadata-ingestion/src/datahub/ingestion/source/couchbase/couchbase_aggregate.py

@@ -4,9 +4,11 @@
 
 from acouchbase.collection import AsyncCollection
 from acouchbase.scope import AsyncScope
+from couchbase.exceptions import DocumentNotFoundException
 from couchbase.result import GetResult
 
 from datahub.ingestion.source.couchbase.couchbase_connect import CouchbaseConnect
+from datahub.ingestion.source.couchbase.retry import retry
 
 logger = logging.getLogger(__name__)
 
@@ -36,9 +38,16 @@ async def init(self):
             self.keyspace
         )
 
-    async def collection_get(self, key: str) -> GetResult:
-        return await self.collection.get(key)
+    @retry()
+    async def collection_get(self, key: str) -> dict:
+        try:
+            result: GetResult = await self.collection.get(key)
+            return result.content_as[dict]
+        except DocumentNotFoundException:
+            logger.warning(f"Document ID {key} not found")
+            return {}
 
+    @retry()
     async def get_keys(self):
         query = f"select meta().id from {self.connector.collection_name}"
         if self.max_sample_size > 0:
@@ -75,8 +84,9 @@ async def get_documents(self) -> AsyncGenerator[List[dict], None]:
                     if isinstance(result, Exception):
                         logger.error(result)
                         errors += 1
-                    elif isinstance(result, GetResult):
-                        batch.append(result.content_as[dict])
+                    elif isinstance(result, dict):
+                        if result:
+                            batch.append(result)
                 yield batch
 
             if errors > 0:

metadata-ingestion/src/datahub/ingestion/source/couchbase/couchbase_connect.py

@@ -180,6 +180,7 @@ def bucket_info(self, bucket_name: str) -> BucketSettings:
         settings: BucketSettings = self.bucket_manager.get_bucket(bucket_name)
         return settings
 
+    @retry()
     def collection_infer(
         self, sample_size: int, sample_values: int, keyspace: str
     ) -> dict:

metadata-ingestion/tests/integration/couchbase/test_couchbase.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import shutil
+import time
 from pathlib import Path
 
 import pytest
@@ -62,6 +63,8 @@ def test_couchbase_ingest(docker_compose_runner, pytestconfig, tmp_path, mock_ti
 
         assert response.status_code == 200
 
+        time.sleep(2)
+
         # Run the metadata ingestion pipeline.
         pipeline = Pipeline.create(
             {

metadata-ingestion/tests/unit/couchbase/test_couchbase_source.py

@@ -1,3 +1,4 @@
+import asyncio
 import base64
 import logging
 
@@ -56,6 +57,8 @@ async def test_couchbase_driver(
 
         assert response.status_code == 200
 
+        await asyncio.sleep(2)
+
         # Run the driver test.
         couchbase_connect = CouchbaseConnect(
             "couchbases://127.0.0.1", "Administrator", "password", 5, 60