openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/ingestion/source/pipeline/dbtcloud/client.py

@@ -35,6 +35,7 @@ from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
+API_VERSION = "api/v2"
 
 
 class DBTCloudClient:
@@ -44,9 +45,13 @@ class DBTCloudClient:
 
     def __init__(self, config: DBTCloudConnection):
         self.config = config
+
+        self.job_ids = self.config.jobIds
+        self.project_ids = self.config.projectIds
+
         client_config: ClientConfig = ClientConfig(
             base_url=clean_uri(self.config.host),
-            api_version="api/v2",
+            api_version=API_VERSION,
             auth_header=AUTHORIZATION_HEADER,
             auth_token=lambda: (self.config.token.get_secret_value(), 0),
             allow_redirects=True,
@@ -63,27 +68,43 @@ class DBTCloudClient:
         self.client = REST(client_config)
         self.graphql_client = REST(graphql_client_config)
 
-    def test_get_jobs(self) -> Optional[List[DBTJob]]:
+    def _get_jobs(
+        self, job_id: str = None, project_id: str = None
+    ) -> Optional[List[DBTJob]]:
         """
-        test fetch jobs for an account in dbt cloud
+        fetch jobs for an account in dbt cloud
         """
-        job_path = f"{self.config.jobId}/" if self.config.jobId else ""
-        result = self.client.get(f"/accounts/{self.config.accountId}/jobs/{job_path}")
-        job_list = (
-            [DBTJob.model_validate(result["data"])]
-            if self.config.jobId
-            else DBTJobList.model_validate(result).Jobs
-        )
+        job_list = []
+        try:
+            job_path = f"{job_id}/" if job_id else ""
+            project_path = f"?project_id={project_id}" if project_id else ""
+            result = self.client.get(
+                f"/accounts/{self.config.accountId}/jobs/{job_path}{project_path}"
+            )
+            job_list = (
+                [DBTJob.model_validate(result["data"])]
+                if job_id
+                else DBTJobList.model_validate(result).Jobs
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(
+                f"Failed to get job info for project_id: `{project_id}` or job_id: `{job_id}` : {exc}"
+            )
         return job_list
 
-    def test_get_runs(self, job_id: int) -> Optional[List[DBTRun]]:
+    def test_get_jobs(self) -> List[DBTJob]:
+        """
+        test fetch jobs for an account in dbt cloud
+        """
+        job_list = self.client.get(f"/accounts/{self.config.accountId}/jobs/")
+        return DBTJobList.model_validate(job_list).Jobs
+
+    def test_get_runs(self) -> List[DBTRun]:
         """
         test fetch runs for a job in dbt cloud
         """
-        result = self.client.get(
-            f"/accounts/{self.config.accountId}/runs/",
-            data={"job_definition_id": job_id},
-        )
+        result = self.client.get(f"/accounts/{self.config.accountId}/runs/")
         run_list = DBTRunList.model_validate(result).Runs
         return run_list
 
@@ -92,17 +113,29 @@ class DBTCloudClient:
         list jobs for an account in dbt cloud
         """
         try:
-            job_path = f"{self.config.jobId}/" if self.config.jobId else ""
-            result = self.client.get(
-                f"/accounts/{self.config.accountId}/jobs/{job_path}"
-            )
-            if result:
-                job_list = (
-                    [DBTJob.model_validate(result.get("data"))]
-                    if self.config.jobId
-                    else DBTJobList.model_validate(result).Jobs
-                )
-                return job_list
+            jobs = []
+            # case when job_ids are specified and project_ids are not
+            if self.job_ids and not self.project_ids:
+                for job_id in self.job_ids:
+                    jobs.extend(self._get_jobs(job_id=job_id))
+            # case when project_ids are specified or both are specified
+            elif self.project_ids:
+                for project_id in self.project_ids:
+                    results = self._get_jobs(project_id=project_id)
+                    if self.job_ids:
+                        jobs.extend(
+                            [
+                                result
+                                for result in results
+                                if str(result.id) in self.job_ids
+                            ]
+                        )
+                    else:
+                        jobs.extend(results)
+            else:
+                results = self._get_jobs()
+                jobs.extend(results)
+            return jobs
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(f"Unable to get job info :{exc}")
@@ -141,6 +174,9 @@ class DBTCloudClient:
 
             if result.get("data") and result["data"].get("job"):
                 model_list = DBTModelList.model_validate(result["data"]["job"]).models
+                logger.debug(
+                    f"Successfully fetched models from dbt for job_id:{job_id} run_id:{run_id}: {model_list}"
+                )
                 return model_list
 
         except Exception as exc:
@@ -150,7 +186,7 @@ class DBTCloudClient:
 
     def get_models_and_seeds_details(self, job_id: int, run_id: int):
         """
-        get model details for a job in dbt cloud for lineage
+        get parent model details for a job in dbt cloud for lineage
         """
         try:
             query_params = {
@@ -163,9 +199,12 @@ class DBTCloudClient:
             if result.get("data") and result["data"].get("job"):
                 result = DBTModelList.model_validate(result["data"]["job"])
                 parents_list = result.models + result.seeds
+                logger.debug(
+                    f"Successfully fetched parent models from dbt for job_id:{job_id} run_id:{run_id}: {parents_list}"
+                )
                 return parents_list
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(f"Unable to get model info :{exc}")
+            logger.warning(f"Unable to get parents model info :{exc}")
         return None

metadata/ingestion/source/pipeline/dbtcloud/connection.py

@@ -50,11 +50,9 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    job_id = int(service_connection.jobId) if service_connection.jobId else 0
-
     test_fn = {
         "GetJobs": client.test_get_jobs,
-        "GetRuns": partial(client.test_get_runs, job_id=job_id),
+        "GetRuns": partial(client.test_get_runs),
     }
 
     return test_connection_steps(

metadata/ingestion/source/pipeline/dbtcloud/models.py

@@ -35,7 +35,7 @@ class DBTJob(BaseModel):
 
 
 class DBTJobList(BaseModel):
-    Jobs: Optional[List[DBTJob]] = Field([], alias="data")
+    Jobs: List[DBTJob] = Field(alias="data")
 
 
 class DBTRun(BaseModel):

metadata/ingestion/source/pipeline/kafkaconnect/client.py

@@ -54,7 +54,7 @@ class KafkaConnectClient:
         auth = None
         ssl_verify = config.verifySSL
         if config.KafkaConnectConfig:
-            auth = f"{config.KafkaConnectConfig.username}:{config.KafkaConnectConfig.password}"
+            auth = f"{config.KafkaConnectConfig.username}:{config.KafkaConnectConfig.password.get_secret_value()}"
         self.client = KafkaConnect(url=url, auth=auth, ssl_verify=ssl_verify)
 
     def get_cluster_info(self) -> Optional[dict]:

metadata/ingestion/source/search/elasticsearch/metadata.py

@@ -12,6 +12,7 @@
 Elasticsearch source to extract metadata
 """
 import shutil
+import traceback
 from pathlib import Path
 from typing import Any, Iterable, Optional
 
@@ -21,6 +22,7 @@ from metadata.generated.schema.api.data.createSearchIndex import (
     CreateSearchIndexRequest,
 )
 from metadata.generated.schema.entity.data.searchIndex import (
+    IndexType,
     SearchIndex,
     SearchIndexSampleData,
 )
@@ -103,6 +105,7 @@ class ElasticsearchSource(SearchServiceSource):
                 fields=parse_es_index_mapping(
                     search_index_details.get(index_name, {}).get("mappings")
                 ),
+                indexType=IndexType.Index,
             )
             yield Either(right=search_index_request)
             self.register_record(search_index_request=search_index_request)
@@ -143,6 +146,56 @@ class ElasticsearchSource(SearchServiceSource):
                 )
             )
 
+    def get_search_index_template_list(self) -> Iterable[dict]:
+        """
+        Get List of all search index template
+        """
+        yield from self.client.indices.get_index_template().get("index_templates", [])
+
+    def get_search_index_template_name(
+        self, search_index_template_details: dict
+    ) -> Optional[str]:
+        """
+        Get Search Index Template Name
+        """
+        return search_index_template_details and search_index_template_details["name"]
+
+    def yield_search_index_template(
+        self, search_index_template_details: Any
+    ) -> Iterable[Either[CreateSearchIndexRequest]]:
+        """
+        Method to Get Search Index Template Entity
+        """
+        try:
+            if self.source_config.includeIndexTemplate:
+                index_name = self.get_search_index_template_name(
+                    search_index_template_details
+                )
+                index_template = search_index_template_details["index_template"]
+                if index_name:
+                    search_index_template_request = CreateSearchIndexRequest(
+                        name=EntityName(index_name),
+                        displayName=index_name,
+                        searchIndexSettings=index_template.get("template", {}).get(
+                            "settings", {}
+                        ),
+                        service=FullyQualifiedEntityName(
+                            self.context.get().search_service
+                        ),
+                        fields=parse_es_index_mapping(
+                            index_template.get("template", {}).get("mappings")
+                        ),
+                        indexType=IndexType.IndexTemplate,
+                        description=index_template.get("_meta", {}).get("description"),
+                    )
+                    yield Either(right=search_index_template_request)
+                    self.register_record(
+                        search_index_request=search_index_template_request
+                    )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(f"Could not include index templates due to {exc}")
+
     def close(self):
         try:
             if Path(self.service_connection.sslConfig.certificates.stagingDir).exists():

metadata/ingestion/source/search/search_service.py

@@ -83,7 +83,7 @@ class SearchServiceTopology(ServiceTopology):
                 cache_entities=True,
             ),
         ],
-        children=["search_index"],
+        children=["search_index", "search_index_template"],
         post_process=["mark_search_indexes_as_deleted"],
     )
     search_index: Annotated[
@@ -107,6 +107,21 @@ class SearchServiceTopology(ServiceTopology):
         ],
     )
 
+    search_index_template: Annotated[
+        TopologyNode, Field(description="Search Index Template Processing Node")
+    ] = TopologyNode(
+        producer="get_search_index_template",
+        stages=[
+            NodeStage(
+                type_=SearchIndex,
+                context="search_index_template",
+                processor="yield_search_index_template",
+                consumer=["search_service"],
+                use_cache=True,
+            )
+        ],
+    )
+
 
 class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
     """
@@ -178,6 +193,34 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
                 continue
             yield index_details
 
+    def yield_search_index_template(
+        self, search_index_template_details: Any
+    ) -> Iterable[Either[CreateSearchIndexRequest]]:
+        """Method to Get Search Index Templates"""
+
+    def get_search_index_template_list(self) -> Optional[List[Any]]:
+        """Get list of all search index templates"""
+
+    def get_search_index_template_name(self, search_index_template_details: Any) -> str:
+        """Get Search Index Template Name"""
+
+    def get_search_index_template(self) -> Any:
+        if self.source_config.includeIndexTemplate:
+            for index_template_details in self.get_search_index_template_list():
+                if search_index_template_name := self.get_search_index_template_name(
+                    index_template_details
+                ):
+                    if filter_by_search_index(
+                        self.source_config.searchIndexFilterPattern,
+                        search_index_template_name,
+                    ):
+                        self.status.filter(
+                            search_index_template_name,
+                            "Search Index Template Filtered Out",
+                        )
+                        continue
+                    yield index_template_details
+
     def yield_create_request_search_service(
         self, config: WorkflowSource
     ) -> Iterable[Either[CreateSearchServiceRequest]]:

metadata/mixins/pandas/pandas_mixin.py

@@ -13,7 +13,6 @@
 Interfaces with database for all database engine
 supporting sqlalchemy abstraction layer
 """
-import math
 import random
 from typing import cast
 
@@ -23,7 +22,6 @@ from metadata.data_quality.validations.table.pandas.tableRowInsertedCountToBeBet
 from metadata.generated.schema.entity.data.table import (
     PartitionIntervalTypes,
     PartitionProfilerConfig,
-    ProfileSampleType,
 )
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.utils.datalake.datalake_utils import fetch_dataframe
@@ -81,9 +79,7 @@ class PandasInterfaceMixin:
             for df in dfs
         ]
 
-    def return_ometa_dataframes_sampled(
-        self, service_connection_config, client, table, profile_sample_config
-    ):
+    def get_dataframes(self, service_connection_config, client, table):
         """
         returns sampled ometa dataframes
         """
@@ -94,35 +90,10 @@ class PandasInterfaceMixin:
                 key=table.name.root,
                 bucket_name=table.databaseSchema.name,
                 file_extension=table.fileFormat,
+                separator=None,
             ),
         )
         if data:
             random.shuffle(data)
-            # sampling data based on profiler config (if any)
-            if hasattr(profile_sample_config, "profile_sample"):
-                if (
-                    profile_sample_config.profile_sample_type
-                    == ProfileSampleType.PERCENTAGE
-                ):
-                    return [
-                        df.sample(
-                            frac=profile_sample_config.profile_sample / 100,
-                            random_state=random.randint(0, 100),
-                            replace=True,
-                        )
-                        for df in data
-                    ]
-                if profile_sample_config.profile_sample_type == ProfileSampleType.ROWS:
-                    sample_rows_per_chunk: int = math.floor(
-                        profile_sample_config.profile_sample / len(data)
-                    )
-                    return [
-                        df.sample(
-                            n=sample_rows_per_chunk,
-                            random_state=random.randint(0, 100),
-                            replace=True,
-                        )
-                        for df in data
-                    ]
             return data
         raise TypeError(f"Couldn't fetch {table.name.root}")

metadata/mixins/sqalchemy/sqa_mixin.py

@@ -20,6 +20,7 @@ from typing import List, Optional
 from sqlalchemy import Column, MetaData, inspect
 from sqlalchemy.orm import DeclarativeMeta
 
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
     DatabricksConnection,
 )
@@ -29,16 +30,19 @@ from metadata.generated.schema.entity.services.connections.database.snowflakeCon
 from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
     UnityCatalogConnection,
 )
+from metadata.ingestion.models.custom_pydantic import BaseModel
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_SESSION_TAG_QUERY,
 )
 from metadata.profiler.orm.converter.base import ometa_to_sqa_orm
 from metadata.utils.collaborative_super import Root
+from metadata.utils.constants import NON_SQA_DATABASE_CONNECTIONS
 
 
 class SQAInterfaceMixin(Root):
-    """SQLAlchemy inteface mixin grouping shared methods between sequential and threaded executor"""
+    """SQLAlchemy interface mixin grouping shared methods between sequential and threaded executor"""
 
     def _get_engine(self):
         """Get engine for database
@@ -48,26 +52,13 @@ class SQAInterfaceMixin(Root):
         Returns:
             sqlalchemy engine
         """
-        engine = get_connection(self.service_connection_config)
+        engine = get_connection(super().service_connection_config)
 
         return engine
 
-    def _convert_table_to_orm_object(
-        self,
-        sqa_metadata_obj: Optional[MetaData] = None,
-    ) -> DeclarativeMeta:
-        """Given a table entity return a SQA ORM object
-
-        Args:
-            sqa_metadata_obj: sqa metadata registry
-        Returns:
-            DeclarativeMeta
-        """
-        return ometa_to_sqa_orm(self.table_entity, self.ometa_client, sqa_metadata_obj)
-
     def get_columns(self) -> Column:
         """get columns from an orm object"""
-        return inspect(self.table).c
+        return inspect(super().table).c
 
     def set_session_tag(self, session) -> None:
         """
@@ -115,3 +106,12 @@ class SQAInterfaceMixin(Root):
             for column in self.table.__table__.columns
             if column.name in {col.name.root for col in self.table_entity.columns}
         ]
+
+    def build_table_orm(
+        self, table: Table, service_conn_config: BaseModel, ometa_client: OpenMetadata
+    ) -> Optional[DeclarativeMeta]:
+        """Build the ORM table if needed for the sampler and profiler interfaces"""
+        if service_conn_config.type.value not in NON_SQA_DATABASE_CONNECTIONS:
+            orm_obj = ometa_to_sqa_orm(table, ometa_client, MetaData())
+            return orm_obj
+        return None

metadata/pii/processor.py

@@ -19,8 +19,8 @@ from metadata.generated.schema.entity.data.table import Column, TableData
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
-from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
-    DatabaseServiceProfilerPipeline,
+from metadata.generated.schema.metadataIngestion.databaseServiceAutoClassificationPipeline import (
+    DatabaseServiceAutoClassificationPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
@@ -40,7 +40,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.pii.constants import PII
 from metadata.pii.scanners.column_name_scanner import ColumnNameScanner
 from metadata.pii.scanners.ner_scanner import NERScanner
-from metadata.profiler.api.models import ProfilerResponse
+from metadata.sampler.models import SamplerResponse
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
@@ -61,8 +61,9 @@ class PIIProcessor(Processor):
         self.metadata = metadata
 
         # Init and type the source config
-        self.source_config: DatabaseServiceProfilerPipeline = cast(
-            DatabaseServiceProfilerPipeline, self.config.source.sourceConfig.config
+        self.source_config: DatabaseServiceAutoClassificationPipeline = cast(
+            DatabaseServiceAutoClassificationPipeline,
+            self.config.source.sourceConfig.config,
         )  # Used to satisfy type checked
 
         self._ner_scanner = None
@@ -71,7 +72,7 @@ class PIIProcessor(Processor):
 
     @property
     def name(self) -> str:
-        return "PII Processor"
+        return "Auto Classification Processor"
 
     @property
     def ner_scanner(self) -> NERScanner:
@@ -152,8 +153,8 @@ class PIIProcessor(Processor):
 
     def _run(
         self,
-        record: ProfilerResponse,
-    ) -> Either[ProfilerResponse]:
+        record: SamplerResponse,
+    ) -> Either[SamplerResponse]:
         """
         Main entrypoint for the scanner.
 
@@ -162,7 +163,7 @@ class PIIProcessor(Processor):
         """
 
         # We don't always need to process
-        if not self.source_config.processPiiSensitive:
+        if not self.source_config.enableAutoClassification:
             return Either(right=record)
 
         column_tags = []

metadata/profiler/api/models.py

@@ -18,87 +18,22 @@ multiple profilers per table and columns.
 
 from typing import List, Optional, Type, Union
 
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict
 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 (
     CreateTableProfileRequest,
 )
-from metadata.generated.schema.entity.data.table import (
-    ColumnProfilerConfig,
-    PartitionProfilerConfig,
-    ProfileSampleType,
-    SamplingMethodType,
-    Table,
-    TableData,
-)
-from metadata.generated.schema.entity.services.connections.connectionBasicType import (
-    SampleDataStorageConfig,
-)
+from metadata.generated.schema.entity.data.table import Table
 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
+from metadata.sampler.models import DatabaseAndSchemaConfig, TableConfig
 from metadata.utils.sqa_like_column import SQALikeColumn
 
 
-class ColumnConfig(ConfigModel):
-    """Column config for profiler"""
-
-    excludeColumns: Optional[List[str]] = None
-    includeColumns: Optional[List[ColumnProfilerConfig]] = None
-
-
-class BaseProfileConfig(ConfigModel):
-    """base profile config"""
-
-    fullyQualifiedName: FullyQualifiedEntityName
-    profileSample: Optional[Union[float, int]] = None
-    profileSampleType: Optional[ProfileSampleType] = None
-    samplingMethodType: Optional[SamplingMethodType] = None
-    sampleDataCount: Optional[int] = 100
-
-
-class TableConfig(BaseProfileConfig):
-    """table profile config"""
-
-    profileQuery: Optional[str] = None
-    partitionConfig: Optional[PartitionProfilerConfig] = None
-    columnConfig: Optional[ColumnConfig] = None
-
-    @classmethod
-    def from_database_and_schema_config(
-        cls, config: "DatabaseAndSchemaConfig", table_fqn: str
-    ):
-        table_config = TableConfig(
-            fullyQualifiedName=table_fqn,
-            profileSample=config.profileSample,
-            profileSampleType=config.profileSampleType,
-            sampleDataCount=config.sampleDataCount,
-            samplingMethodType=config.samplingMethodType,
-        )
-        return table_config
-
-
-class DatabaseAndSchemaConfig(BaseProfileConfig):
-    """schema profile config"""
-
-    sampleDataStorageConfig: Optional[SampleDataStorageConfig] = None
-
-
-class ProfileSampleConfig(ConfigModel):
-    """Profile Sample Config"""
-
-    profile_sample: Optional[Union[float, int]] = None
-    profile_sample_type: Optional[ProfileSampleType] = ProfileSampleType.PERCENTAGE
-    sampling_method_type: Optional[SamplingMethodType] = None
-
-
 class ProfilerProcessorConfig(ConfigModel):
     """
     Defines how we read the processor information
@@ -111,15 +46,6 @@ 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.
@@ -130,8 +56,6 @@ class ProfilerResponse(ConfigModel):
 
     table: Table
     profile: CreateTableProfileRequest
-    sample_data: Optional[SampleData] = None
-    column_tags: Optional[List[ColumnTag]] = None
 
     def __str__(self):
         """Return the table name being processed"""

metadata/profiler/config.py

@@ -0,0 +1,39 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Profiler configuration helpers
+"""
+from typing import Optional
+
+from metadata.generated.schema.entity.data.database import (
+    Database,
+    DatabaseProfilerConfig,
+)
+from metadata.generated.schema.entity.data.databaseSchema import (
+    DatabaseSchema,
+    DatabaseSchemaProfilerConfig,
+)
+
+
+def get_database_profiler_config(
+    database_entity: Optional[Database],
+) -> Optional[DatabaseProfilerConfig]:
+    if database_entity and database_entity.databaseProfilerConfig:
+        return database_entity.databaseProfilerConfig
+    return None
+
+
+def get_schema_profiler_config(
+    schema_entity: Optional[DatabaseSchema],
+) -> Optional[DatabaseSchemaProfilerConfig]:
+    if schema_entity and schema_entity.databaseSchemaProfilerConfig:
+        return schema_entity.databaseSchemaProfilerConfig
+    return None