openmetadata-ingestion 1.5.4.0__py3-none-any.whl → 1.5.5.0__py3-none-any.whl

metadata/generated/schema/type/entityReferenceList.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityReferenceList.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityRelationship.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityRelationship.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityUsage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityUsage.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/filterPattern.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/filterPattern.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/function.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/function.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/include.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/include.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/jdbcConnection.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/jdbcConnection.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/lifeCycle.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/lifeCycle.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/paging.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/paging.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/profile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/profile.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/queryParserData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/queryParserData.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/reaction.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/reaction.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schedule.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schedule.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schema.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schema.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableQuery.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableQuery.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-25T04:53:10+00:00
 
 from __future__ import annotations
 

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -331,19 +331,9 @@ class ESMixin(Generic[T]):
                 else:
                     break
 
-            # Get the data
-            for hit in response.hits.hits:
-                try:
-                    yield self.get_by_name(
-                        entity=entity,
-                        fqn=hit.source["fullyQualifiedName"],
-                        fields=fields,
-                        nullable=False,  # Raise an error if we don't find the Entity
-                    )
-                except Exception as exc:
-                    logger.warning(
-                        f"Error while getting {hit.source['fullyQualifiedName']} - {exc}"
-                    )
+            yield from self._yield_hits_from_api(
+                response=response, entity=entity, fields=fields
+            )
 
             # Get next page
             last_hit = response.hits.hits[-1] if response.hits.hits else None
@@ -362,3 +352,20 @@ class ESMixin(Generic[T]):
             logger.debug(traceback.format_exc())
             logger.warning(f"Error while getting ES response: {exc}")
         return None
+
+    def _yield_hits_from_api(
+        self, response: ESResponse, entity: Type[T], fields: Optional[List[str]]
+    ) -> Iterator[T]:
+        """Get the data from the API based on ES responses"""
+        for hit in response.hits.hits:
+            try:
+                yield self.get_by_name(
+                    entity=entity,
+                    fqn=hit.source["fullyQualifiedName"],
+                    fields=fields,
+                    nullable=False,  # Raise an error if we don't find the Entity
+                )
+            except Exception as exc:
+                logger.warning(
+                    f"Error while getting {hit.source['fullyQualifiedName']} - {exc}"
+                )

metadata/ingestion/ometa/mixins/lineage_mixin.py

@@ -22,13 +22,14 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.type.basic import FullyQualifiedEntityName, Uuid
 from metadata.generated.schema.type.entityLineage import EntitiesEdge
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 from metadata.ingestion.lineage.parser import LINEAGE_PARSING_TIMEOUT
 from metadata.ingestion.models.patch_request import build_patch
 from metadata.ingestion.ometa.client import REST, APIError
-from metadata.ingestion.ometa.utils import get_entity_type, quote
+from metadata.ingestion.ometa.utils import get_entity_type, model_str, quote
 from metadata.utils.logger import ometa_logger
 from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
 
@@ -248,7 +249,7 @@ class OMetaLineageMixin(Generic[T]):
     def get_lineage_by_id(
         self,
         entity: Union[Type[T], str],
-        entity_id: str,
+        entity_id: Union[str, Uuid],
         up_depth: int = 1,
         down_depth: int = 1,
     ) -> Optional[Dict[str, Any]]:
@@ -260,13 +261,16 @@ class OMetaLineageMixin(Generic[T]):
         :param down_depth: Downstream depth of lineage (default=1, min=0, max=3)
         """
         return self._get_lineage(
-            entity=entity, path=entity_id, up_depth=up_depth, down_depth=down_depth
+            entity=entity,
+            path=model_str(entity_id),
+            up_depth=up_depth,
+            down_depth=down_depth,
         )
 
     def get_lineage_by_name(
         self,
         entity: Union[Type[T], str],
-        fqn: str,
+        fqn: Union[str, FullyQualifiedEntityName],
         up_depth: int = 1,
         down_depth: int = 1,
     ) -> Optional[Dict[str, Any]]:
@@ -279,7 +283,7 @@ class OMetaLineageMixin(Generic[T]):
         """
         return self._get_lineage(
             entity=entity,
-            path=f"name/{quote(fqn)}",
+            path=f"name/{quote(model_str(fqn))}",
             up_depth=up_depth,
             down_depth=down_depth,
         )

metadata/ingestion/sink/metadata_rest.py

@@ -565,9 +565,9 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             f"Successfully ingested profile metrics for {record.table.fullyQualifiedName.root}"
         )
 
-        if record.sample_data:
+        if record.sample_data and record.sample_data.store:
             table_data = self.metadata.ingest_table_sample_data(
-                table=record.table, sample_data=record.sample_data
+                table=record.table, sample_data=record.sample_data.data
             )
             if not table_data:
                 self.status.failed(

metadata/ingestion/source/database/common_db_source.py

@@ -11,8 +11,11 @@
 """
 Generic source to build SQL connectors.
 """
+import math
+import time
 import traceback
 from abc import ABC
+from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
 from typing import Any, Iterable, List, Optional, Tuple, Union, cast
 
@@ -60,6 +63,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.connections.session import create_and_bind_thread_safe_session
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
+from metadata.ingestion.models.topology import Queue
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import (
     get_connection,
@@ -73,6 +77,7 @@ from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
 from metadata.utils.db_utils import get_view_lineage
 from metadata.utils.execution_time_tracker import (
+    ExecutionTimeTrackerContextMap,
     calculate_execution_time,
     calculate_execution_time_generator,
 )
@@ -583,9 +588,78 @@ class CommonDbSourceService(
                 )
             )
 
-    @calculate_execution_time_generator()
-    def yield_view_lineage(self) -> Iterable[Either[OMetaLineageRequest]]:
-        logger.info("Processing Lineage for Views")
+    def multithread_process_view_lineage(self) -> Iterable[Either[OMetaLineageRequest]]:
+        """Multithread Processing of a Node"""
+
+        views_list = list(self.context.get().table_views or [])
+        views_length = len(views_list)
+
+        if views_length != 0:
+            chunksize = int(math.ceil(views_length / self.source_config.threads))
+            chunks = [
+                views_list[i : i + chunksize] for i in range(0, views_length, chunksize)
+            ]
+
+            thread_pool = ThreadPoolExecutor(max_workers=self.source_config.threads)
+            queue = Queue()
+
+            futures = [
+                thread_pool.submit(
+                    self._process_view_def_chunk,
+                    chunk,
+                    queue,
+                    self.context.get_current_thread_id(),
+                )
+                for chunk in chunks
+            ]
+
+            while True:
+                if queue.has_tasks():
+                    yield from queue.process()
+
+                else:
+                    if not futures:
+                        break
+
+                    for i, future in enumerate(futures):
+                        if future.done():
+                            future.result()
+                            futures.pop(i)
+
+                time.sleep(0.01)
+
+    def _process_view_def_chunk(
+        self, chunk: List[TableView], queue: Queue, thread_id: int
+    ) -> None:
+        """
+        Process a chunk of view definitions
+        """
+        self.context.copy_from(thread_id)
+        ExecutionTimeTrackerContextMap().copy_from_parent(thread_id)
+        for view in [v for v in chunk if v.view_definition is not None]:
+            for lineage in get_view_lineage(
+                view=view,
+                metadata=self.metadata,
+                service_name=self.context.get().database_service,
+                connection_type=self.service_connection.type.value,
+                timeout_seconds=self.source_config.queryParsingTimeoutLimit,
+            ):
+                if lineage.right is not None:
+                    queue.put(
+                        Either(
+                            right=OMetaLineageRequest(
+                                lineage_request=lineage.right,
+                                override_lineage=self.source_config.overrideViewLineage,
+                            )
+                        )
+                    )
+                else:
+                    queue.put(lineage)
+
+    def _process_view_def_serial(self) -> Iterable[Either[OMetaLineageRequest]]:
+        """
+        Process view definitions serially
+        """
         for view in [
             v for v in self.context.get().table_views if v.view_definition is not None
         ]:
@@ -606,6 +680,14 @@ class CommonDbSourceService(
                 else:
                     yield lineage
 
+    @calculate_execution_time_generator()
+    def yield_view_lineage(self) -> Iterable[Either[OMetaLineageRequest]]:
+        logger.info("Processing Lineage for Views")
+        if self.source_config.threads > 1:
+            yield from self.multithread_process_view_lineage()
+        else:
+            yield from self._process_view_def_serial()
+
     def _get_foreign_constraints(self, foreign_columns) -> List[TableConstraint]:
         """
         Search the referred table for foreign constraints

metadata/ingestion/source/database/salesforce/connection.py

@@ -15,7 +15,6 @@ Source connection handler
 from typing import Optional
 
 from simple_salesforce.api import Salesforce
-from sqlalchemy.engine import Engine
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -27,14 +26,17 @@ from metadata.ingestion.connections.test_connections import test_connection_step
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 
 
-def get_connection(connection: SalesforceConnection) -> Engine:
+def get_connection(connection: SalesforceConnection) -> Salesforce:
     """
     Create connection
     """
     return Salesforce(
         username=connection.username,
         password=connection.password.get_secret_value(),
-        security_token=connection.securityToken.get_secret_value(),
+        security_token=connection.securityToken.get_secret_value()
+        if connection.securityToken
+        else "",
+        organizationId=connection.organizationId if connection.organizationId else "",
         domain=connection.salesforceDomain,
         version=connection.salesforceApiVersion,
         **connection.connectionArguments.root if connection.connectionArguments else {},

metadata/ingestion/source/database/snowflake/metadata.py

@@ -65,9 +65,6 @@ from metadata.ingestion.source.database.external_table_lineage_mixin import (
 from metadata.ingestion.source.database.incremental_metadata_extraction import (
     IncrementalConfig,
 )
-from metadata.ingestion.source.database.life_cycle_query_mixin import (
-    LifeCycleQueryMixin,
-)
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.snowflake.models import (
     STORED_PROC_LANGUAGE_MAP,
@@ -148,7 +145,6 @@ SnowflakeDialect._get_schema_foreign_keys = get_schema_foreign_keys
 
 
 class SnowflakeSource(
-    LifeCycleQueryMixin,
     StoredProcedureMixin,
     ExternalTableLineageMixin,
     CommonDbSourceService,

metadata/pii/ner.py

@@ -51,7 +51,8 @@ class NEREntity(Enum):
     UK_NHS = TagType.SENSITIVE.value
 
     # Spain
-    NIF = TagType.SENSITIVE.value
+    ES_NIF = TagType.SENSITIVE.value
+    ES_NIE = TagType.SENSITIVE.value
 
     # Italy
     IT_FISCAL_CODE = TagType.SENSITIVE.value
@@ -60,8 +61,25 @@ class NEREntity(Enum):
     IT_PASSPORT = TagType.SENSITIVE.value
     IT_IDENTITY_CARD = TagType.SENSITIVE.value
 
+    # Poland
+    PL_PESEL = TagType.SENSITIVE.value
+
+    # Singapore
+    SG_NRIC_FIN = TagType.SENSITIVE.value
+    SG_UEN = TagType.SENSITIVE.value
+
     # Australia
     AU_ABN = TagType.SENSITIVE.value
     AU_ACN = TagType.SENSITIVE.value
     AU_TFN = TagType.SENSITIVE.value
     AU_MEDICARE = TagType.SENSITIVE.value
+
+    # India
+    IN_PAN = TagType.SENSITIVE.value
+    IN_AADHAAR = TagType.SENSITIVE.value
+    IN_VEHICLE_REGISTRATION = TagType.SENSITIVE.value
+    IN_VOTER = TagType.SENSITIVE.value
+    IN_PASSPORT = TagType.SENSITIVE.value
+
+    # Finland
+    FI_PERSONAL_IDENTITY_CODE = TagType.SENSITIVE.value

metadata/pii/processor.py

@@ -171,7 +171,7 @@ class PIIProcessor(Processor):
                 col_tags = self.process_column(
                     idx=idx,
                     column=column,
-                    table_data=record.sample_data,
+                    table_data=record.sample_data.data,
                     confidence_threshold=self.confidence_threshold,
                 )
                 if col_tags:

metadata/pii/scanners/ner_scanner.py

@@ -14,6 +14,7 @@ NER Scanner based on Presidio.
 Supported Entities https://microsoft.github.io/presidio/supported_entities/
 """
 import json
+import logging
 import traceback
 from collections import defaultdict
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -26,9 +27,11 @@ from metadata.pii.models import TagAndConfidence
 from metadata.pii.ner import NEREntity
 from metadata.pii.scanners.base import BaseScanner
 from metadata.utils import fqn
-from metadata.utils.logger import pii_logger
+from metadata.utils.logger import METADATA_LOGGER, pii_logger
 
 logger = pii_logger()
+SUPPORTED_LANG = "en"
+PRESIDIO_LOGGER = "presidio-analyzer"
 
 
 class StringAnalysis(BaseModel):
@@ -40,11 +43,16 @@ class StringAnalysis(BaseModel):
     appearances: int
 
 
+class NLPEngineModel(BaseModel):
+    """Required to pass the nlp_engine as {"lang_code": "en", "model_name": "en_core_web_lg"}"""
+
+    lang_code: str
+    model_name: str
+
+
 # pylint: disable=import-outside-toplevel
 class NERScanner(BaseScanner):
-    """
-    Based on https://microsoft.github.io/presidio/
-    """
+    """Based on https://microsoft.github.io/presidio/"""
 
     def __init__(self):
         import spacy
@@ -60,8 +68,19 @@ class NERScanner(BaseScanner):
             download(SPACY_EN_MODEL)
             spacy.load(SPACY_EN_MODEL)
 
+        nlp_engine_model = NLPEngineModel(
+            lang_code=SUPPORTED_LANG, model_name=SPACY_EN_MODEL
+        )
+
+        # Set the presidio logger to talk less about internal entities unless we are debugging
+        logging.getLogger(PRESIDIO_LOGGER).setLevel(
+            logging.INFO
+            if logging.getLogger(METADATA_LOGGER).level == logging.DEBUG
+            else logging.ERROR
+        )
+
         self.analyzer = AnalyzerEngine(
-            nlp_engine=SpacyNlpEngine(models={"en": SPACY_EN_MODEL})
+            nlp_engine=SpacyNlpEngine(models=[nlp_engine_model.model_dump()])
         )
 
     @staticmethod

metadata/profiler/api/models.py

@@ -17,9 +17,10 @@ multiple profilers per table and columns.
 """
 from typing import List, Optional, Type, Union
 
-from pydantic import ConfigDict
+from pydantic import ConfigDict, Field
 from sqlalchemy import Column
 from sqlalchemy.orm import DeclarativeMeta
+from typing_extensions import Annotated
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.data.createTableProfile import (
@@ -37,6 +38,7 @@ from metadata.generated.schema.entity.services.connections.connectionBasicType i
 )
 from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
+from metadata.ingestion.models.custom_pydantic import BaseModel
 from metadata.ingestion.models.table_metadata import ColumnTag
 from metadata.profiler.metrics.core import Metric, MetricTypes
 from metadata.profiler.processor.models import ProfilerDef
@@ -104,6 +106,15 @@ class ProfilerProcessorConfig(ConfigModel):
     databaseConfig: Optional[List[DatabaseAndSchemaConfig]] = []
 
 
+class SampleData(BaseModel):
+    """TableData wrapper to handle ephemeral SampleData"""
+
+    data: Annotated[TableData, Field(None, description="Table Sample Data")]
+    store: Annotated[
+        bool, Field(False, description="Is the sample data should be stored or not")
+    ]
+
+
 class ProfilerResponse(ConfigModel):
     """
     ORM Profiler processor response.
@@ -114,7 +125,7 @@ class ProfilerResponse(ConfigModel):
 
     table: Table
     profile: CreateTableProfileRequest
-    sample_data: Optional[TableData] = None
+    sample_data: Optional[SampleData] = None
     column_tags: Optional[List[ColumnTag]] = None
 
     def __str__(self):

metadata/profiler/processor/core.py

@@ -33,7 +33,6 @@ from metadata.generated.schema.entity.data.table import (
     ColumnProfile,
     ColumnProfilerConfig,
     SystemProfile,
-    TableData,
     TableProfile,
 )
 from metadata.generated.schema.settings.settings import Settings
@@ -41,7 +40,7 @@ from metadata.generated.schema.tests.customMetric import (
     CustomMetric as CustomMetricEntity,
 )
 from metadata.generated.schema.type.basic import Timestamp
-from metadata.profiler.api.models import ProfilerResponse, ThreadPoolMetrics
+from metadata.profiler.api.models import ProfilerResponse, SampleData, ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import (
     ComposedMetric,
@@ -492,7 +491,12 @@ class Profiler(Generic[TMetric]):
             )
             self.compute_metrics()
 
-        if self.source_config.generateSampleData:
+        # We need the sample data for Sample Data or PII Sensitive processing.
+        # We'll nullify the Sample Data after the PII processing so that it's not stored.
+        if (
+            self.source_config.generateSampleData
+            or self.source_config.processPiiSensitive
+        ):
             sample_data = self.generate_sample_data()
         else:
             sample_data = None
@@ -510,7 +514,7 @@ class Profiler(Generic[TMetric]):
         return table_profile
 
     @calculate_execution_time(store=False)
-    def generate_sample_data(self) -> Optional[TableData]:
+    def generate_sample_data(self) -> Optional[SampleData]:
         """Fetch and ingest sample data
 
         Returns:
@@ -532,7 +536,10 @@ class Profiler(Generic[TMetric]):
                     SAMPLE_DATA_DEFAULT_COUNT, self.profiler_interface.sample_data_count
                 )
             ]
-            return table_data
+            return SampleData(
+                data=table_data, store=self.source_config.generateSampleData
+            )
+
         except Exception as err:
             logger.debug(traceback.format_exc())
             logger.warning(f"Error fetching sample data: {err}")