openmetadata-ingestion 1.4.6.0__py3-none-any.whl → 1.5.0.0rc1__py3-none-any.whl

metadata/ingestion/ometa/mixins/pipeline_mixin.py

@@ -43,7 +43,7 @@ class OMetaPipelineMixin:
         """
         resp = self.client.put(
             f"{self.get_suffix(Pipeline)}/{fqn}/status",
-            data=status.json(),
+            data=status.model_dump_json(),
         )
 
         return Pipeline(**resp)
@@ -86,7 +86,7 @@ class OMetaPipelineMixin:
             startDate=pipeline.startDate,
             service=pipeline.service.fullyQualifiedName,
             tasks=all_tasks,
-            owner=pipeline.owner,
+            owners=pipeline.owners,
             tags=pipeline.tags,
         )
 
@@ -113,7 +113,7 @@ class OMetaPipelineMixin:
             startDate=pipeline.startDate,
             service=pipeline.service.fullyQualifiedName,
             tasks=[task for task in pipeline.tasks if task.name in task_ids],
-            owner=pipeline.owner,
+            owners=pipeline.owners,
             tags=pipeline.tags,
         )
 

metadata/ingestion/ometa/mixins/query_mixin.py

@@ -41,10 +41,10 @@ class OMetaQueryMixin:
         return str(result.hexdigest())
 
     def _get_or_create_query(self, query: CreateQueryRequest) -> Optional[Query]:
-        query_hash = self._get_query_hash(query=query.query.__root__)
+        query_hash = self._get_query_hash(query=query.query.root)
         query_entity = self.get_by_name(entity=Query, fqn=query_hash)
         if query_entity is None:
-            resp = self.client.put(self.get_suffix(Query), data=query.json())
+            resp = self.client.put(self.get_suffix(Query), data=query.model_dump_json())
             if resp and resp.get("id"):
                 query_entity = Query(**resp)
         return query_entity
@@ -63,9 +63,9 @@ class OMetaQueryMixin:
                 query = self._get_or_create_query(create_query)
                 if query:
                     # Add Query Usage
-                    table_ref = EntityReference(id=entity.id.__root__, type="table")
+                    table_ref = EntityReference(id=entity.id.root, type="table")
                     # convert object to json array string
-                    table_ref_json = "[" + table_ref.json() + "]"
+                    table_ref_json = "[" + table_ref.model_dump_json() + "]"
                     self.client.put(
                         f"{self.get_suffix(Query)}/{model_str(query.id)}/usage",
                         data=table_ref_json,

metadata/ingestion/ometa/mixins/role_policy_mixin.py

@@ -125,7 +125,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
                     if previous is None
                     else PatchOperation.REPLACE,
                     PatchField.PATH: path.format(rule_index=rule_index),
-                    PatchField.VALUE: str(current.__root__),
+                    PatchField.VALUE: str(current.root),
                 }
             ]
         return data
@@ -158,10 +158,10 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
         if not instance:
             return None
 
-        policy_index: int = len(instance.policies.__root__) - 1
+        policy_index: int = len(instance.policies.root) - 1
         data: List
         if operation is PatchOperation.REMOVE:
-            if len(instance.policies.__root__) == 1:
+            if len(instance.policies.root) == 1:
                 logger.error(
                     f"The Role with id [{model_str(entity_id)}] has only one (1)"
                     f" policy. Unable to remove."
@@ -177,7 +177,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
 
             index: int = 0
             is_policy_found: bool = False
-            for policy in instance.policies.__root__:
+            for policy in instance.policies.root:
                 if model_str(policy.id) == model_str(policy_id):
                     is_policy_found = True
                     continue
@@ -187,7 +187,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
                         PatchField.PATH: PatchPath.POLICIES_DESCRIPTION.format(
                             index=index
                         ),
-                        PatchField.VALUE: model_str(policy.description.__root__),
+                        PatchField.VALUE: model_str(policy.description.root),
                     }
                 )
                 data.append(
@@ -294,7 +294,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
         if not instance:
             return None
 
-        rule_index: int = len(instance.rules.__root__) - 1
+        rule_index: int = len(instance.rules.root) - 1
         data: List[Dict]
         if operation == PatchOperation.ADD:
             data = [
@@ -303,7 +303,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
                     PatchField.PATH: PatchPath.RULES.format(rule_index=rule_index + 1),
                     PatchField.VALUE: {
                         PatchValue.NAME: rule.name,
-                        PatchValue.CONDITION: rule.condition.__root__,
+                        PatchValue.CONDITION: rule.condition.root,
                         PatchValue.EFFECT: rule.effect.name,
                         PatchValue.OPERATIONS: [
                             operation.name for operation in rule.operations
@@ -314,12 +314,12 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
             ]
             if rule.description is not None:
                 data[0][PatchField.VALUE][PatchValue.DESCRIPTION] = str(
-                    rule.description.__root__
+                    rule.description.root
                 )
 
             if rule.fullyQualifiedName is not None:
                 data[0][PatchField.VALUE][PatchValue.FQN] = str(
-                    rule.fullyQualifiedName.__root__
+                    rule.fullyQualifiedName.root
                 )
 
         else:
@@ -334,8 +334,8 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
                 }
             ]
 
-            for rule_index in range(len(instance.rules.__root__) - 1, -1, -1):
-                current_rule: Rule = instance.rules.__root__[rule_index]
+            for rule_index in range(len(instance.rules.root) - 1, -1, -1):
+                current_rule: Rule = instance.rules.root[rule_index]
                 if current_rule.name == rule.name:
                     break
 
@@ -345,7 +345,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
                     )
                     return None
 
-                previous_rule: Rule = instance.rules.__root__[rule_index - 1]
+                previous_rule: Rule = instance.rules.root[rule_index - 1]
                 # Condition
                 data.append(
                     {
@@ -353,7 +353,7 @@ class OMetaRolePolicyMixin(OMetaPatchMixinBase):
                         PatchField.PATH: PatchPath.RULES_CONDITION.format(
                             rule_index=rule_index - 1
                         ),
-                        PatchField.VALUE: current_rule.condition.__root__,
+                        PatchField.VALUE: current_rule.condition.root,
                     }
                 )
                 # Description - Optional

metadata/ingestion/ometa/mixins/search_index_mixin.py

@@ -47,13 +47,13 @@ class OMetaSearchIndexMixin:
         resp = None
         try:
             resp = self.client.put(
-                f"{self.get_suffix(SearchIndex)}/{search_index.id.__root__}/sampleData",
-                data=sample_data.json(),
+                f"{self.get_suffix(SearchIndex)}/{search_index.id.root}/sampleData",
+                data=sample_data.model_dump_json(),
             )
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
-                f"Error trying to PUT sample data for {search_index.fullyQualifiedName.__root__}: {exc}"
+                f"Error trying to PUT sample data for {search_index.fullyQualifiedName.root}: {exc}"
             )
 
         if resp:
@@ -63,13 +63,13 @@ class OMetaSearchIndexMixin:
                 logger.debug(traceback.format_exc())
                 logger.warning(
                     "Unicode Error parsing the sample data response "
-                    f"from {search_index.fullyQualifiedName.__root__}: {err}"
+                    f"from {search_index.fullyQualifiedName.root}: {err}"
                 )
             except Exception as exc:
                 logger.debug(traceback.format_exc())
                 logger.warning(
                     "Error trying to parse sample data results"
-                    f"from {search_index.fullyQualifiedName.__root__}: {exc}"
+                    f"from {search_index.fullyQualifiedName.root}: {exc}"
                 )
 
         return None

metadata/ingestion/ometa/mixins/server_mixin.py

@@ -21,7 +21,6 @@ from metadata.__version__ import (
     match_versions,
 )
 from metadata.generated.schema.settings.settings import Settings, SettingType
-from metadata.ingestion.models.encoders import show_secrets_encoder
 from metadata.ingestion.ometa.client import REST
 from metadata.ingestion.ometa.routes import ROUTES
 from metadata.utils.logger import ometa_logger
@@ -91,9 +90,9 @@ class OMetaServerMixin:
         Returns:
             Settings
         """
-        data = settings.json(encoder=show_secrets_encoder)
+        data = settings.model_dump_json()
         response = self.client.put(ROUTES.get(Settings.__name__), data)
-        return Settings.parse_obj(response)
+        return Settings.model_validate(response)
 
     def get_settings_by_name(self, setting_type: SettingType) -> Optional[Settings]:
         """Get setting by name
@@ -106,7 +105,7 @@ class OMetaServerMixin:
         )
         if not response:
             return None
-        return Settings.parse_obj(response)
+        return Settings.model_validate(response)
 
     def get_profiler_config_settings(self) -> Optional[Settings]:
         """Get profiler config setting
@@ -117,4 +116,4 @@ class OMetaServerMixin:
         response = self.client.get("/system/settings/profilerConfiguration")
         if not response:
             return None
-        return Settings.parse_obj(response)
+        return Settings.model_validate(response)

metadata/ingestion/ometa/mixins/service_mixin.py

@@ -55,8 +55,8 @@ class OMetaServiceMixin:
         create_entity_class = self.get_create_entity_type(entity=entity)
         return create_entity_class(
             name=config.serviceName,
-            serviceType=config.serviceConnection.__root__.config.type.value,
-            connection=config.serviceConnection.__root__
+            serviceType=config.serviceConnection.root.config.type.value,
+            connection=config.serviceConnection.root
             if self.config.storeServiceConnection
             else None,
         )

metadata/ingestion/ometa/mixins/suggestions_mixin.py

@@ -34,8 +34,8 @@ class OMetaSuggestionsMixin:
         Update an existing Suggestion with new fields
         """
         resp = self.client.put(
-            f"{self.get_suffix(Suggestion)}/{str(suggestion.id.__root__)}",
-            data=suggestion.json(),
+            f"{self.get_suffix(Suggestion)}/{str(suggestion.root.id.root)}",
+            data=suggestion.model_dump_json(),
         )
 
         return Suggestion(**resp)

metadata/ingestion/ometa/mixins/table_mixin.py

@@ -69,13 +69,13 @@ class OMetaTableMixin:
         resp = None
         try:
             resp = self.client.put(
-                f"{self.get_suffix(Table)}/{table.id.__root__}/sampleData",
-                data=sample_data.json(),
+                f"{self.get_suffix(Table)}/{table.id.root}/sampleData",
+                data=sample_data.model_dump_json(),
             )
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
-                f"Error trying to PUT sample data for {table.fullyQualifiedName.__root__}: {exc}"
+                f"Error trying to PUT sample data for {table.fullyQualifiedName.root}: {exc}"
             )
 
         if resp:
@@ -84,12 +84,12 @@ class OMetaTableMixin:
             except UnicodeError as err:
                 logger.debug(traceback.format_exc())
                 logger.warning(
-                    f"Unicode Error parsing the sample data response from {table.fullyQualifiedName.__root__}: {err}"
+                    f"Unicode Error parsing the sample data response from {table.fullyQualifiedName.root}: {err}"
                 )
             except Exception as exc:
                 logger.debug(traceback.format_exc())
                 logger.warning(
-                    f"Error trying to parse sample data results from {table.fullyQualifiedName.__root__}: {exc}"
+                    f"Error trying to parse sample data results from {table.fullyQualifiedName.root}: {exc}"
                 )
 
         return None
@@ -103,12 +103,12 @@ class OMetaTableMixin:
         resp = None
         try:
             resp = self.client.get(
-                f"{self.get_suffix(Table)}/{table.id.__root__}/sampleData",
+                f"{self.get_suffix(Table)}/{table.id.root}/sampleData",
             )
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
-                f"Error trying to GET sample data for {table.fullyQualifiedName.__root__}: {exc}"
+                f"Error trying to GET sample data for {table.fullyQualifiedName.root}: {exc}"
             )
 
         if resp:
@@ -117,12 +117,12 @@ class OMetaTableMixin:
             except UnicodeError as err:
                 logger.debug(traceback.format_exc())
                 logger.warning(
-                    f"Unicode Error parsing the sample data response from {table.fullyQualifiedName.__root__}: {err}"
+                    f"Unicode Error parsing the sample data response from {table.fullyQualifiedName.root}: {err}"
                 )
             except Exception as exc:
                 logger.debug(traceback.format_exc())
                 logger.warning(
-                    f"Error trying to parse sample data results from {table.fullyQualifiedName.__root__}: {exc}"
+                    f"Error trying to parse sample data results from {table.fullyQualifiedName.root}: {exc}"
                 )
 
         return None
@@ -137,8 +137,8 @@ class OMetaTableMixin:
         :param table_profile: Profile data to add
         """
         resp = self.client.put(
-            f"{self.get_suffix(Table)}/{table.id.__root__}/tableProfile",
-            data=profile_request.json(),
+            f"{self.get_suffix(Table)}/{table.id.root}/tableProfile",
+            data=profile_request.model_dump_json(),
         )
         return Table(**resp)
 
@@ -150,8 +150,8 @@ class OMetaTableMixin:
         :param data_model: Model to add
         """
         resp = self.client.put(
-            f"{self.get_suffix(Table)}/{table.id.__root__}/dataModel",
-            data=data_model.json(),
+            f"{self.get_suffix(Table)}/{table.id.root}/dataModel",
+            data=data_model.model_dump_json(),
         )
         return Table(**resp)
 
@@ -165,7 +165,7 @@ class OMetaTableMixin:
         :param table_usage_request: Usage data to add
         """
         resp = self.client.post(
-            f"/usage/table/{table.id.__root__}", data=table_usage_request.json()
+            f"/usage/table/{table.id.root}", data=table_usage_request.model_dump_json()
         )
         logger.debug("published table usage %s", resp)
 
@@ -179,10 +179,10 @@ class OMetaTableMixin:
         :param table_join_request: Join data to add
         """
 
-        logger.info("table join request %s", table_join_request.json())
+        logger.info("table join request %s", table_join_request.model_dump_json())
         resp = self.client.put(
-            f"{self.get_suffix(Table)}/{table.id.__root__}/joins",
-            data=table_join_request.json(),
+            f"{self.get_suffix(Table)}/{table.id.root}/joins",
+            data=table_join_request.model_dump_json(),
         )
         logger.debug("published frequently joined with %s", resp)
 
@@ -203,7 +203,7 @@ class OMetaTableMixin:
         """
         resp = self.client.put(
             f"{self.get_suffix(Table)}/{model_str(table_id)}/tableProfilerConfig",
-            data=table_profiler_config.json(),
+            data=table_profiler_config.model_dump_json(),
         )
         return Table(**resp)
 
@@ -290,7 +290,7 @@ class OMetaTableMixin:
         Returns:
             Optional[Table]: OM table object
         """
-        return self._get(Table, f"{quote(model_str(fqn))}/tableProfile/latest")
+        return self._get(Table, f"{quote(model_str(fqn), safe='')}/tableProfile/latest")
 
     def create_or_update_custom_metric(
         self, custom_metric: CreateCustomMetricRequest, table_id: str
@@ -303,6 +303,6 @@ class OMetaTableMixin:
         """
         resp = self.client.put(
             f"{self.get_suffix(Table)}/{table_id}/customMetric",
-            data=custom_metric.json(),
+            data=custom_metric.model_dump_json(),
         )
         return Table(**resp)

metadata/ingestion/ometa/mixins/tests_mixin.py

@@ -15,7 +15,7 @@ To be used by OpenMetadata class
 """
 
 import traceback
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import List, Optional, Type, Union
 from urllib.parse import quote
 from uuid import UUID
@@ -45,7 +45,6 @@ from metadata.generated.schema.tests.testDefinition import (
 )
 from metadata.generated.schema.tests.testSuite import TestSuite
 from metadata.generated.schema.type.entityReference import EntityReference
-from metadata.ingestion.models.encoders import show_secrets_encoder
 from metadata.ingestion.ometa.client import REST
 from metadata.ingestion.ometa.utils import model_str
 from metadata.utils.logger import ometa_logger
@@ -78,7 +77,7 @@ class OMetaTestsMixin:
         """
         resp = self.client.put(
             f"{self.get_suffix(TestCase)}/{quote(test_case_fqn,safe='')}/testCaseResult",
-            test_results.json(),
+            test_results.model_dump_json(),
         )
 
         return resp
@@ -88,7 +87,7 @@ class OMetaTestsMixin:
         test_suite_name: str,
         test_suite_description: Optional[
             str
-        ] = f"Test Suite created on {datetime.now(timezone.utc).strftime('%Y-%m-%d')}",
+        ] = f"Test Suite created on {datetime.now().strftime('%Y-%m-%d')}",
     ) -> TestSuite:
         """Get or create a TestSuite
 
@@ -229,8 +228,8 @@ class OMetaTestsMixin:
             return table_entity.testSuite
 
         create_test_suite = CreateTestSuiteRequest(
-            name=f"{table_entity.fullyQualifiedName.__root__}.TestSuite",
-            executableEntityReference=table_entity.fullyQualifiedName.__root__,
+            name=f"{table_entity.fullyQualifiedName.root}.TestSuite",
+            executableEntityReference=table_entity.fullyQualifiedName.root,
         )  # type: ignore
         test_suite = self.create_or_update_executable_test_suite(create_test_suite)
         return test_suite
@@ -260,7 +259,7 @@ class OMetaTestsMixin:
         )
 
         if resp:
-            return [TestCaseResult.parse_obj(entity) for entity in resp["data"]]
+            return [TestCaseResult.model_validate(entity) for entity in resp["data"]]
         return None
 
     def create_or_update_executable_test_suite(
@@ -277,9 +276,9 @@ class OMetaTestsMixin:
         entity = data.__class__
         entity_class = self.get_entity_from_create(entity)
         path = self.get_suffix(entity) + "/executable"
-        resp = self.client.put(path, data=data.json(encoder=show_secrets_encoder))
+        resp = self.client.put(path, data=data.model_dump_json())
 
-        return entity_class.parse_obj(resp)
+        return entity_class.model_validate(resp)
 
     def delete_executable_test_suite(
         self,
@@ -307,7 +306,7 @@ class OMetaTestsMixin:
             data (CreateLogicalTestCases): logical test cases
         """
         path = self.get_suffix(TestCase) + "/logicalTestCases"
-        self.client.put(path, data=data.json(encoder=show_secrets_encoder))
+        self.client.put(path, data=data.model_dump_json())
 
     def create_test_case_resolution(
         self, data: CreateTestCaseResolutionStatus
@@ -321,12 +320,15 @@ class OMetaTestsMixin:
             TestCaseResolutionStatus
         """
         path = self.get_suffix(TestCase) + "/testCaseIncidentStatus"
-        response = self.client.post(path, data=data.json(encoder=show_secrets_encoder))
+        response = self.client.post(path, data=data.model_dump_json())
 
         return TestCaseResolutionStatus(**response)
 
     def ingest_failed_rows_sample(
-        self, test_case: TestCase, failed_rows: TableData
+        self,
+        test_case: TestCase,
+        failed_rows: TableData,
+        validate=True,
     ) -> Optional[TableData]:
         """
         PUT sample failed data for a test case.
@@ -336,14 +338,15 @@ class OMetaTestsMixin:
         """
         resp = None
         try:
+            params = "" if validate else "validate=false"
             resp = self.client.put(
-                f"{self.get_suffix(TestCase)}/{test_case.id.__root__}/failedRowsSample",
-                data=failed_rows.json(),
+                f"{self.get_suffix(TestCase)}/{test_case.id.root}/failedRowsSample?{params}",
+                data=failed_rows.model_dump_json(),
             )
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
-                f"Error trying to PUT sample data for {test_case.fullyQualifiedName.__root__}: {exc}"
+                f"Error trying to PUT sample data for {test_case.fullyQualifiedName.root}: {exc}"
             )
 
         if resp:
@@ -352,13 +355,13 @@ class OMetaTestsMixin:
             except UnicodeError as err:
                 logger.debug(traceback.format_exc())
                 logger.warning(
-                    f"Unicode Error parsing the sample data response from {test_case.fullyQualifiedName.__root__}: "
+                    f"Unicode Error parsing the sample data response from {test_case.fullyQualifiedName.root}: "
                     f"{err}"
                 )
             except Exception as exc:
                 logger.debug(traceback.format_exc())
                 logger.warning(
-                    f"Error trying to parse sample data results from {test_case.fullyQualifiedName.__root__}: {exc}"
+                    f"Error trying to parse sample data results from {test_case.fullyQualifiedName.root}: {exc}"
                 )
 
         return None
@@ -373,7 +376,7 @@ class OMetaTestsMixin:
         :param inspection_query: SQL query to inspect the failed rows
         """
         resp = self.client.put(
-            f"{self.get_suffix(TestCase)}/{test_case.id.__root__}/inspectionQuery",
+            f"{self.get_suffix(TestCase)}/{test_case.id.root}/inspectionQuery",
             data=inspection_query,
         )
         return TestCase(**resp)

metadata/ingestion/ometa/mixins/topic_mixin.py

@@ -40,7 +40,7 @@ class OMetaTopicMixin:
         :param sample_data: Data to add
         """
         resp = self.client.put(
-            f"{self.get_suffix(Topic)}/{topic.id.__root__}/sampleData",
-            data=sample_data.json(),
+            f"{self.get_suffix(Topic)}/{topic.id.root}/sampleData",
+            data=sample_data.model_dump_json(),
         )
         return TopicSampleData(**resp["sampleData"])

metadata/ingestion/ometa/mixins/user_mixin.py

@@ -16,9 +16,10 @@ To be used by OpenMetadata class
 from functools import lru_cache
 from typing import Optional, Type
 
-from metadata.generated.schema.entity.teams.team import Team
+from metadata.generated.schema.entity.teams.team import Team, TeamType
 from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.ingestion.api.common import T
 from metadata.ingestion.ometa.client import REST
 from metadata.utils.constants import ENTITY_REFERENCE_TYPE_MAP
@@ -118,7 +119,7 @@ class OMetaUserMixin:
         from_count: int = 0,
         size: int = 1,
         fields: Optional[list] = None,
-    ) -> Optional[EntityReference]:
+    ) -> Optional[EntityReferenceList]:
         """
         Get a User or Team Entity Reference by searching by its mail
         """
@@ -126,22 +127,30 @@ class OMetaUserMixin:
             entity=User, email=email, from_count=from_count, size=size, fields=fields
         )
         if maybe_user:
-            return EntityReference(
-                id=maybe_user.id.__root__,
-                type=ENTITY_REFERENCE_TYPE_MAP[User.__name__],
-                name=maybe_user.name.__root__,
-                displayName=maybe_user.displayName,
+            return EntityReferenceList(
+                root=[
+                    EntityReference(
+                        id=maybe_user.id.root,
+                        type=ENTITY_REFERENCE_TYPE_MAP[User.__name__],
+                        name=maybe_user.name.root,
+                        displayName=maybe_user.displayName,
+                    )
+                ]
             )
 
         maybe_team = self._search_by_email(
             entity=Team, email=email, from_count=from_count, size=size, fields=fields
         )
         if maybe_team:
-            return EntityReference(
-                id=maybe_team.id.__root__,
-                type=ENTITY_REFERENCE_TYPE_MAP[Team.__name__],
-                name=maybe_team.name.__root__,
-                displayName=maybe_team.displayName,
+            return EntityReferenceList(
+                root=[
+                    EntityReference(
+                        id=maybe_team.id.root,
+                        type=ENTITY_REFERENCE_TYPE_MAP[Team.__name__],
+                        name=maybe_team.name.root,
+                        displayName=maybe_team.displayName,
+                    )
+                ]
             )
 
         return None
@@ -153,7 +162,8 @@ class OMetaUserMixin:
         from_count: int = 0,
         size: int = 1,
         fields: Optional[list] = None,
-    ) -> Optional[EntityReference]:
+        is_owner: bool = False,
+    ) -> Optional[EntityReferenceList]:
         """
         Get a User or Team Entity Reference by searching by its name
         """
@@ -161,22 +171,33 @@ class OMetaUserMixin:
             entity=User, name=name, from_count=from_count, size=size, fields=fields
         )
         if maybe_user:
-            return EntityReference(
-                id=maybe_user.id.__root__,
-                type=ENTITY_REFERENCE_TYPE_MAP[User.__name__],
-                name=maybe_user.name.__root__,
-                displayName=maybe_user.displayName,
+            return EntityReferenceList(
+                root=[
+                    EntityReference(
+                        id=maybe_user.id.root,
+                        type=ENTITY_REFERENCE_TYPE_MAP[User.__name__],
+                        name=maybe_user.name.root,
+                        displayName=maybe_user.displayName,
+                    )
+                ]
             )
 
         maybe_team = self._search_by_name(
             entity=Team, name=name, from_count=from_count, size=size, fields=fields
         )
         if maybe_team:
-            return EntityReference(
-                id=maybe_team.id.__root__,
-                type=ENTITY_REFERENCE_TYPE_MAP[Team.__name__],
-                name=maybe_team.name.__root__,
-                displayName=maybe_team.displayName,
+            # if is_owner is True, we only want to return the team if it is a group
+            if is_owner and maybe_team.teamType != TeamType.Group:
+                return None
+            return EntityReferenceList(
+                root=[
+                    EntityReference(
+                        id=maybe_team.id.root,
+                        type=ENTITY_REFERENCE_TYPE_MAP[Team.__name__],
+                        name=maybe_team.name.root,
+                        displayName=maybe_team.displayName,
+                    )
+                ]
             )
 
         return None

metadata/ingestion/ometa/models.py

@@ -17,7 +17,7 @@ from pydantic import BaseModel
 T = TypeVar("T", bound=BaseModel)
 
 
-class EntityList(Generic[T], BaseModel):
+class EntityList(BaseModel, Generic[T]):
     """
     Pydantic Entity list model