resolve static type checking errors

Author: jsingh-yelp

ingestion/src/metadata/ingestion/source/pipeline/openlineage/metadata.py

@@ -817,7 +817,7 @@ def yield_pipeline(self, pipeline_details: OpenLineageEvent) -> Iterable[Either[
             )
             description = f"""```json
             {json.dumps(pipeline_details.run_facet, indent=4).strip()}```"""
-            request = CreatePipelineRequest(
+            request = CreatePipelineRequest(  # pyright: ignore[reportCallIssue]
                 name=pipeline_name,
                 service=self._current_pipeline_service,
                 description=description,

ingestion/src/metadata/ingestion/source/pipeline/openlineage/ownership_resolver.py

@@ -42,19 +42,20 @@ class OpenLineageOwnerResolver:
     def __init__(
         self,
         metadata: OpenMetadata,
-        include_owners: bool = True,
-        ownership_update_mode: OwnershipUpdateMode = OwnershipUpdateMode.replace,
+        include_owners: bool | None,
+        ownership_update_mode: OwnershipUpdateMode | str | None,
     ):
         self.metadata = metadata
-        self.include_owners = include_owners
-        self.ownership_update_mode = (
-            ownership_update_mode
+        self.include_owners = bool(include_owners)
+        self.ownership_update_mode = OwnershipUpdateMode(
+            ownership_update_mode.value
             if isinstance(ownership_update_mode, OwnershipUpdateMode)
-            else OwnershipUpdateMode(ownership_update_mode)
+            else ownership_update_mode or OwnershipUpdateMode.replace.value
         )
-        self._team_owner_ref_by_name: Optional[Dict[str, EntityReference]] = None  # noqa: UP006, UP045
-        self._user_owner_ref_by_name: Optional[Dict[str, EntityReference]] = None  # noqa: UP006, UP045
-        self._owner_refs_by_pipeline_fqn: Optional[Dict[str, List[EntityReference]]] = None  # noqa: UP006, UP045
+        self._owner_cache_loaded = False
+        self._team_owner_ref_by_name: Dict[str, EntityReference] = {}  # noqa: UP006
+        self._user_owner_ref_by_name: Dict[str, EntityReference] = {}  # noqa: UP006
+        self._owner_refs_by_pipeline_fqn: Dict[str, List[EntityReference]] = {}  # noqa: UP006
 
     def get_pipeline_job_owners(
         self,
@@ -154,11 +155,7 @@ def _ensure_pipeline_owner_cache(self) -> None:
         """
         Load OpenMetadata Group teams, users, and pipeline ownership references.
         """
-        if (
-            self._team_owner_ref_by_name is not None
-            and self._user_owner_ref_by_name is not None
-            and self._owner_refs_by_pipeline_fqn is not None
-        ):
+        if self._owner_cache_loaded:
             return
 
         team_owner_ref_by_name: Dict[str, EntityReference] = {}  # noqa: UP006
@@ -179,7 +176,7 @@ def _ensure_pipeline_owner_cache(self) -> None:
                 if team_name:
                     team_ref = team_owner_ref_by_name.setdefault(
                         team_name,
-                        EntityReference(
+                        EntityReference(  # pyright: ignore[reportCallIssue]
                             id=model_str(team.id),
                             type=OWNER_TEAM_ENTITY,
                             name=model_str(team.name),
@@ -214,7 +211,7 @@ def _ensure_pipeline_owner_cache(self) -> None:
                 if user_name:
                     user_ref = user_owner_ref_by_name.setdefault(
                         user_name,
-                        EntityReference(
+                        EntityReference(  # pyright: ignore[reportCallIssue]
                             id=model_str(user.id),
                             type=OWNER_USER_ENTITY,
                             name=model_str(user.name),
@@ -237,3 +234,4 @@ def _ensure_pipeline_owner_cache(self) -> None:
         self._team_owner_ref_by_name = team_owner_ref_by_name
         self._user_owner_ref_by_name = user_owner_ref_by_name
         self._owner_refs_by_pipeline_fqn = owner_refs_by_pipeline_fqn
+        self._owner_cache_loaded = True

ingestion/tests/unit/topology/pipeline/test_openlineage.py

@@ -36,8 +36,6 @@
 from metadata.generated.schema.type.entityLineage import ColumnLineage
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
-from metadata.ingestion.api.models import Either
-from metadata.ingestion.source.pipeline.openlineage.metadata import OpenlineageSource
 from metadata.ingestion.source.pipeline.openlineage.metadata import (
     RESOLUTION_CACHE_MAXSIZE,
     OpenlineageSource,

ingestion/tests/unit/topology/pipeline/test_openlineage_ownership.py

@@ -136,6 +136,19 @@ def test_does_not_build_cache_when_include_owners_is_disabled():
     metadata.list_all_entities.assert_not_called()
 
 
+def test_does_not_build_cache_when_include_owners_is_none():
+    resolver, metadata = build_resolver(
+        users=[build_user("jdoe")],
+        teams=[build_team("data-platform")],
+        include_owners=None,
+    )
+
+    owners = resolver.get_pipeline_job_owners(build_job(["user:jdoe", "team:data-platform"]))
+
+    assert owners is None
+    metadata.list_all_entities.assert_not_called()
+
+
 def test_builds_pipeline_owner_cache_with_filtered_owns_api():
     pipeline_ref = EntityReference(
         id=uuid4(),
@@ -189,6 +202,30 @@ def test_replace_mode_does_not_carry_existing_pipeline_owners():
     assert [(owner.type, owner.name) for owner in owners.root] == [("user", "jdoe")]
 
 
+def test_none_ownership_update_mode_defaults_to_replace():
+    pipeline_ref = EntityReference(
+        id=uuid4(),
+        type="pipeline",
+        name="daily_orders",
+        fullyQualifiedName="airflow.daily_orders",
+    )
+    existing_team = build_team("data-platform", owns=[pipeline_ref])
+    new_user = build_user("jdoe")
+    resolver, _ = build_resolver(
+        users=[new_user],
+        teams=[existing_team],
+        ownership_update_mode=None,
+    )
+
+    owners = resolver.get_pipeline_job_owners(
+        build_job(["user:jdoe"]),
+        pipeline_fqn="airflow.daily_orders",
+    )
+
+    assert owners is not None
+    assert [(owner.type, owner.name) for owner in owners.root] == [("user", "jdoe")]
+
+
 def test_append_mode_merges_existing_pipeline_owners_and_updates_cache():
     pipeline_ref = EntityReference(
         id=uuid4(),

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java

@@ -2227,7 +2227,7 @@ public ResultList<T> listAfter(
       String afterId = cursorMap.get("id");
       List<String> jsons = dao.listAfter(filter, limitParam + 1, afterName, afterId);
 
-      entities = listInternal(jsons, fields, uriInfo);
+      entities = listInternal(jsons, fields, uriInfo, filter);
 
       String beforeCursor;
       String afterCursor = null;
@@ -2249,18 +2249,19 @@ public ResultList<T> listAfterWithOffset(
     int total = ListCountCache.getOrCompute(entityType, filter, () -> dao.listCount(filter));
     List<String> jsons = dao.listAfter(filter, limit, offset);
 
-    List<T> entities = listInternal(jsons, fields, uriInfo);
+    List<T> entities = listInternal(jsons, fields, uriInfo, filter);
 
     return new ResultList<>(entities, offset, limit, total);
   }
 
-  private List<T> listInternal(List<String> jsons, Fields fields, UriInfo uriInfo) {
+  private List<T> listInternal(
+      List<String> jsons, Fields fields, UriInfo uriInfo, ListFilter filter) {
     List<T> entities;
     try (var ignored = phase("jsonDeserialize")) {
       entities = JsonUtils.readObjects(jsons, entityClass);
     }
     try (var ignored = phase("setFieldsBulk")) {
-      setFieldsInBulk(fields, entities);
+      setFieldsInBulk(fields, entities, filter);
     }
     entities.forEach(entity -> withHref(uriInfo, entity));
     return entities;
@@ -2284,7 +2285,8 @@ public ResultList<T> listAfterKeyset(
       boolean hasMoreData = jsons.size() > limitParam;
       List<String> jsonsToProcess = hasMoreData ? jsons.subList(0, limitParam) : jsons;
 
-      Iterator<Either<T, EntityError>> iterator = serializeJsons(jsonsToProcess, fields, null);
+      Iterator<Either<T, EntityError>> iterator =
+          serializeJsons(jsonsToProcess, fields, null, filter);
       while (iterator.hasNext()) {
         Either<T, EntityError> either = iterator.next();
         if (either.right().isPresent()) {
@@ -2343,7 +2345,7 @@ public ResultList<T> listBefore(
     List<String> jsons = dao.listBefore(filter, limitParam + 1, beforeName, beforeId);
 
     List<T> entities = JsonUtils.readObjects(jsons, entityClass);
-    setFieldsInBulk(fields, entities);
+    setFieldsInBulk(fields, entities, filter);
     entities.forEach(entity -> withHref(uriInfo, entity));
 
     String beforeCursor = null;
@@ -2443,7 +2445,7 @@ public final ResultList<T> listWithOffset(
     String beforeOffset = getBeforeOffset(offsetInt, limitParam);
     if (limitParam > 0) {
       List<String> jsons = callable.apply(filter, limitParam, offsetInt);
-      Iterator<Either<T, EntityError>> iterator = serializeJsons(jsons, fields, uriInfo);
+      Iterator<Either<T, EntityError>> iterator = serializeJsons(jsons, fields, uriInfo, filter);
       while (iterator.hasNext()) {
         Either<T, EntityError> either = iterator.next();
         if (either.right().isPresent()) {
@@ -10926,7 +10928,7 @@ List<String> entityListToStrings(List<T> entities) {
   }
 
   private Iterator<Either<T, EntityError>> serializeJsons(
-      List<String> jsons, Fields fields, UriInfo uriInfo) {
+      List<String> jsons, Fields fields, UriInfo uriInfo, ListFilter filter) {
     List<Either<T, EntityError>> results = new ArrayList<>();
     List<T> entities = new ArrayList<>();
 
@@ -10945,7 +10947,7 @@ private Iterator<Either<T, EntityError>> serializeJsons(
 
     if (!entities.isEmpty()) {
       try {
-        setFieldsInBulk(fields, entities);
+        setFieldsInBulk(fields, entities, filter);
         if (!nullOrEmpty(uriInfo)) {
           entities.forEach(entity -> withHref(uriInfo, entity));
         }