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

metadata/ingestion/source/database/dbt/constants.py

@@ -22,6 +22,14 @@ REQUIRED_MANIFEST_KEYS = ["name", "schema", "resource_type"]
 # Based on https://schemas.getdbt.com/dbt/catalog/v1.json
 REQUIRED_CATALOG_KEYS = ["name", "type", "index"]
 
+REQUIRED_CONSTRAINT_KEYS = [
+    "type",
+    "name",
+    "expression",
+    "warn_unenforced",
+    "warn_unsupported",
+]
+
 REQUIRED_RESULTS_KEYS = {
     "status",
     "timing",
@@ -82,6 +90,7 @@ NONE_KEYWORDS_LIST = ["none", "null"]
 DBT_CATALOG_FILE_NAME = "catalog.json"
 DBT_MANIFEST_FILE_NAME = "manifest.json"
 DBT_RUN_RESULTS_FILE_NAME = "run_results"
+DBT_SOURCES_FILE_NAME = "sources.json"
 
 
 class SkipResourceTypeEnum(Enum):
@@ -91,6 +100,7 @@ class SkipResourceTypeEnum(Enum):
 
     ANALYSIS = "analysis"
     TEST = "test"
+    SOURCE = "source"
 
 
 class CompiledQueriesEnum(Enum):
@@ -127,6 +137,7 @@ class DbtTestFailureEnum(Enum):
 
     FAILURE = "failure"
     FAIL = "fail"
+    ERROR = "error"
 
 
 class DbtCommonEnum(Enum):
@@ -137,6 +148,7 @@ class DbtCommonEnum(Enum):
     OWNER = "owner"
     NODES = "nodes"
     SOURCES = "sources"
+    SOURCES_FILE = "sources_file"
     SOURCE = "source"
     RESOURCETYPE = "resource_type"
     MANIFEST_NODE = "manifest_node"

metadata/ingestion/source/database/dbt/dbt_config.py

@@ -43,6 +43,7 @@ from metadata.ingestion.source.database.dbt.constants import (
     DBT_CATALOG_FILE_NAME,
     DBT_MANIFEST_FILE_NAME,
     DBT_RUN_RESULTS_FILE_NAME,
+    DBT_SOURCES_FILE_NAME,
 )
 from metadata.ingestion.source.database.dbt.models import DbtFiles
 from metadata.readers.file.config_source_factory import get_reader
@@ -85,6 +86,7 @@ def _(config: DbtLocalConfig):
             config.dbtManifestFilePath,
             config.dbtCatalogFilePath,
             config.dbtRunResultsFilePath,
+            config.dbtSourcesFilePath,
         ]
         yield from download_dbt_files(
             blob_grouped_by_directory=blob_grouped_by_directory,
@@ -123,12 +125,22 @@ def _(config: DbtHttpConfig):
             dbt_catalog = requests.get(  # pylint: disable=missing-timeout
                 config.dbtCatalogHttpPath
             )
+
+        dbt_sources = None
+        if config.dbtSourcesHttpPath:
+            logger.debug(
+                f"Requesting [dbtSourcesHttpPath] to: {config.dbtSourcesHttpPath}"
+            )
+            dbt_sources = requests.get(  # pylint: disable=missing-timeout
+                config.dbtSourcesHttpPath
+            )
         if not dbt_manifest:
             raise DBTConfigException("Manifest file not found in file server")
         yield DbtFiles(
             dbt_catalog=dbt_catalog.json() if dbt_catalog else None,
             dbt_manifest=dbt_manifest.json(),
             dbt_run_results=[dbt_run_results.json()] if dbt_run_results else None,
+            dbt_sources=dbt_sources.json() if dbt_sources else None,
         )
     except DBTConfigException as exc:
         raise exc
@@ -164,7 +176,7 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
         logger.debug(
             "Requesting [dbt_catalog], [dbt_manifest] and [dbt_run_results] data"
         )
-        params_data = {"order_by": "-finished_at", "limit": "1", "status": "10"}
+        params_data = {"order_by": "-finished_at", "limit": "1"}
         if project_id:
             params_data["project_id"] = project_id
 
@@ -243,6 +255,7 @@ def get_blobs_grouped_by_dir(blobs: List[str]) -> Dict[str, List[str]]:
     return blob_grouped_by_directory
 
 
+# pylint: disable=too-many-locals, too-many-branches
 def download_dbt_files(
     blob_grouped_by_directory: Dict, config, client, bucket_name: Optional[str]
 ) -> Iterable[DbtFiles]:
@@ -255,6 +268,7 @@ def download_dbt_files(
     ) in blob_grouped_by_directory.items():
         dbt_catalog = None
         dbt_manifest = None
+        dbt_sources = None
         dbt_run_results = []
         kwargs = {}
         if bucket_name:
@@ -285,12 +299,16 @@ def download_dbt_files(
                             logger.warning(
                                 f"{DBT_RUN_RESULTS_FILE_NAME} not found in {key}: {exc}"
                             )
+                    if DBT_SOURCES_FILE_NAME == blob_file_name.lower():
+                        logger.debug(f"{DBT_SOURCES_FILE_NAME} found in {key}")
+                        dbt_sources = reader.read(path=blob, **kwargs)
             if not dbt_manifest:
                 raise DBTConfigException(f"Manifest file not found at: {key}")
             yield DbtFiles(
                 dbt_catalog=json.loads(dbt_catalog) if dbt_catalog else None,
                 dbt_manifest=json.loads(dbt_manifest),
                 dbt_run_results=dbt_run_results if dbt_run_results else None,
+                dbt_sources=json.loads(dbt_sources) if dbt_sources else None,
             )
         except DBTConfigException as exc:
             logger.warning(exc)

metadata/ingestion/source/database/dbt/dbt_service.py

@@ -15,7 +15,12 @@ DBT service Topology.
 from abc import ABC, abstractmethod
 from typing import Iterable, List
 
-from dbt_artifacts_parser.parser import parse_catalog, parse_manifest, parse_run_results
+from dbt_artifacts_parser.parser import (
+    parse_catalog,
+    parse_manifest,
+    parse_run_results,
+    parse_sources,
+)
 from pydantic import Field
 from typing_extensions import Annotated
 
@@ -38,6 +43,7 @@ from metadata.ingestion.models.topology import (
 )
 from metadata.ingestion.source.database.database_service import DataModelLink
 from metadata.ingestion.source.database.dbt.constants import (
+    REQUIRED_CONSTRAINT_KEYS,
     REQUIRED_NODE_KEYS,
     REQUIRED_RESULTS_KEYS,
 )
@@ -182,6 +188,20 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
                 ]
                 for key in keys_to_delete:
                     del value[key]
+                if value.get("columns"):
+                    for col_name, value in value[
+                        "columns"
+                    ].items():  # pylint: disable=unused-variable
+                        if value.get("constraints"):
+                            keys_to_delete = [
+                                key
+                                for key in value
+                                if key.lower() not in REQUIRED_CONSTRAINT_KEYS
+                            ]
+                            for key in keys_to_delete:
+                                del value[key]
+                        else:
+                            value["constraints"] = None
 
     def remove_run_result_non_required_keys(self, run_results: List[dict]):
         """
@@ -209,11 +229,15 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
             self.remove_run_result_non_required_keys(
                 run_results=self.context.get().dbt_file.dbt_run_results
             )
+
         dbt_objects = DbtObjects(
             dbt_catalog=parse_catalog(self.context.get().dbt_file.dbt_catalog)
             if self.context.get().dbt_file.dbt_catalog
             else None,
             dbt_manifest=parse_manifest(self.context.get().dbt_file.dbt_manifest),
+            dbt_sources=parse_sources(self.context.get().dbt_file.dbt_sources)
+            if self.context.get().dbt_file.dbt_sources
+            else None,
             dbt_run_results=[
                 parse_run_results(run_result_file)
                 for run_result_file in self.context.get().dbt_file.dbt_run_results

metadata/ingestion/source/database/dbt/dbt_utils.py

@@ -44,6 +44,20 @@ def create_test_case_parameter_definitions(dbt_test):
                 }
             ]
             return test_case_param_definition
+        if hasattr(dbt_test, "freshness"):
+            test_case_param_definition = [
+                {
+                    "name": "warn_after",
+                    "displayName": "warn_after",
+                    "required": False,
+                },
+                {
+                    "name": "error_after",
+                    "displayName": "error_after",
+                    "required": False,
+                },
+            ]
+            return test_case_param_definition
     except Exception as err:  # pylint: disable=broad-except
         logger.debug(traceback.format_exc())
         logger.error(
@@ -67,6 +81,21 @@ def create_test_case_parameter_values(dbt_test):
                 {"name": manifest_node.test_metadata.name, "value": dbt_test_values}
             ]
             return test_case_param_values
+        if hasattr(manifest_node, "freshness"):
+            warn_after = manifest_node.freshness.warn_after
+            error_after = manifest_node.freshness.error_after
+
+            test_case_param_values = [
+                {
+                    "name": "error_after",
+                    "value": f"{error_after.count} {error_after.period.value}",
+                },
+                {
+                    "name": "warn_after",
+                    "value": f"{warn_after.count} {warn_after.period.value}",
+                },
+            ]
+            return test_case_param_values
     except Exception as err:  # pylint: disable=broad-except
         logger.debug(traceback.format_exc())
         logger.error(

metadata/ingestion/source/database/dbt/metadata.py

@@ -13,6 +13,7 @@
 DBT source methods.
 """
 import traceback
+from copy import deepcopy
 from datetime import datetime
 from typing import Any, Iterable, List, Optional, Union
 
@@ -324,7 +325,41 @@ class DbtSource(DbtServiceSource):
             None,
         )
 
-    # pylint: disable=too-many-locals, too-many-branches
+    def _add_dbt_freshness_test_from_sources(
+        self, key: str, manifest_node, manifest_entities, dbt_objects: DbtObjects
+    ):
+        # in dbt manifest sources node name is table/view name (not test name like with test nodes)
+        # so in order for the test creation to be named precisely I am amending manifest node name within it's deepcopy
+        manifest_node_new = deepcopy(manifest_node)
+        manifest_node_new.name = manifest_node_new.name + "_freshness"
+
+        freshness_test_result = next(
+            (item for item in dbt_objects.dbt_sources.results if item.unique_id == key),
+            None,
+        )
+
+        if freshness_test_result:
+            self.context.get().dbt_tests[key + "_freshness"] = {
+                DbtCommonEnum.MANIFEST_NODE.value: manifest_node_new
+            }
+            self.context.get().dbt_tests[key + "_freshness"][
+                DbtCommonEnum.UPSTREAM.value
+            ] = self.parse_upstream_nodes(manifest_entities, manifest_node)
+            self.context.get().dbt_tests[key + "_freshness"][
+                DbtCommonEnum.RESULTS.value
+            ] = freshness_test_result
+
+    def add_dbt_sources(
+        self, key: str, manifest_node, manifest_entities, dbt_objects: DbtObjects
+    ) -> None:
+        """
+        Method to append dbt test cases based on sources file for later processing
+        """
+        self._add_dbt_freshness_test_from_sources(
+            key, manifest_node, manifest_entities, dbt_objects
+        )
+
+    # pylint: disable=too-many-locals, too-many-branches, too-many-statements
     def yield_data_models(
         self, dbt_objects: DbtObjects
     ) -> Iterable[Either[DataModelLink]]:
@@ -376,6 +411,17 @@ class DbtSource(DbtServiceSource):
                         )
                         continue
 
+                    if (
+                        dbt_objects.dbt_sources
+                        and resource_type == SkipResourceTypeEnum.SOURCE.value
+                    ):
+                        self.add_dbt_sources(
+                            key,
+                            manifest_node=manifest_node,
+                            manifest_entities=manifest_entities,
+                            dbt_objects=dbt_objects,
+                        )
+
                     # Skip the ephemeral nodes since it is not materialized
                     if check_ephemeral_node(manifest_node):
                         logger.debug(f"Skipping ephemeral DBT node: {key}.")
@@ -549,6 +595,29 @@ class DbtSource(DbtServiceSource):
                         f"Failed to parse the DBT node {node} to get upstream nodes: {exc}"
                     )
                     continue
+
+        if dbt_node.resource_type == SkipResourceTypeEnum.SOURCE.value:
+            parent_fqn = fqn.build(
+                self.metadata,
+                entity_type=Table,
+                service_name="*",
+                database_name=get_corrected_name(dbt_node.database),
+                schema_name=get_corrected_name(dbt_node.schema_),
+                table_name=dbt_node.name,
+            )
+
+            # check if the parent table exists in OM before adding it to the upstream list
+            parent_table_entity: Optional[
+                Union[Table, List[Table]]
+            ] = get_entity_from_es_result(
+                entity_list=self.metadata.es_search_from_fqn(
+                    entity_type=Table, fqn_search_string=parent_fqn
+                ),
+                fetch_multiple_entities=False,
+            )
+            if parent_table_entity:
+                upstream_nodes.append(parent_fqn)
+
         return upstream_nodes
 
     def parse_data_model_columns(

metadata/ingestion/source/database/dbt/models.py

@@ -20,12 +20,14 @@ from pydantic import BaseModel
 class DbtFiles(BaseModel):
     dbt_catalog: Optional[dict] = None
     dbt_manifest: dict
+    dbt_sources: Optional[dict] = None
     dbt_run_results: Optional[List[dict]] = None
 
 
 class DbtObjects(BaseModel):
     dbt_catalog: Optional[Any] = None
     dbt_manifest: Any
+    dbt_sources: Optional[Any] = None
     dbt_run_results: Optional[List[Any]] = None
 
 

metadata/ingestion/source/database/dynamodb/service_spec.py

@@ -1,7 +1,10 @@
 from metadata.ingestion.source.database.dynamodb.metadata import DynamodbSource
 from metadata.profiler.interface.nosql.profiler_interface import NoSQLProfilerInterface
+from metadata.sampler.nosql.sampler import NoSQLSampler
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=DynamodbSource, profiler_class=NoSQLProfilerInterface
+    metadata_source_class=DynamodbSource,
+    profiler_class=NoSQLProfilerInterface,
+    sampler_class=NoSQLSampler,
 )

metadata/ingestion/source/database/mongodb/service_spec.py

@@ -1,7 +1,10 @@
 from metadata.ingestion.source.database.mongodb.metadata import MongodbSource
 from metadata.profiler.interface.nosql.profiler_interface import NoSQLProfilerInterface
+from metadata.sampler.nosql.sampler import NoSQLSampler
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=MongodbSource, profiler_class=NoSQLProfilerInterface
+    metadata_source_class=MongodbSource,
+    profiler_class=NoSQLProfilerInterface,
+    sampler_class=NoSQLSampler,
 )

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

@@ -38,6 +38,7 @@ from metadata.ingestion.connections.test_connections import (
     test_connection_db_schema_sources,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.mysql.queries import MYSQL_TEST_GET_QUERIES
 from metadata.utils.constants import THREE_MIN
 
 
@@ -73,10 +74,14 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
+    queries = {
+        "GetQueries": MYSQL_TEST_GET_QUERIES,
+    }
     return test_connection_db_schema_sources(
         metadata=metadata,
         engine=engine,
         service_connection=service_connection,
         automation_workflow=automation_workflow,
         timeout_seconds=timeout_seconds,
+        queries=queries,
     )

metadata/ingestion/source/database/mysql/lineage.py

@@ -9,38 +9,20 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 """
-Mysql lineage module
+MYSQL lineage module
 """
-from typing import Optional
-
-from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
-    MysqlConnection,
-)
-from metadata.generated.schema.metadataIngestion.workflow import (
-    Source as WorkflowSource,
-)
-from metadata.ingestion.api.steps import InvalidSourceException
-from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.lineage_source import LineageSource
-from metadata.utils.logger import ingestion_logger
-
-logger = ingestion_logger()
+from metadata.ingestion.source.database.mysql.queries import MYSQL_SQL_STATEMENT
+from metadata.ingestion.source.database.mysql.query_parser import MysqlQueryParserSource
 
 
-class MysqlLineageSource(LineageSource):
+class MysqlLineageSource(MysqlQueryParserSource, LineageSource):
+    sql_stmt = MYSQL_SQL_STATEMENT
+    filters = """
+        AND (
+            lower(argument) LIKE '%%create%%table%%select%%'
+            OR lower(argument) LIKE '%%insert%%into%%select%%'
+            OR lower(argument) LIKE '%%update%%'
+            OR lower(argument) LIKE '%%merge%%'
+        )
     """
-    Mysql lineage source implements view lineage
-    """
-
-    @classmethod
-    def create(
-        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
-    ):
-        """Create class instance"""
-        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
-        connection: MysqlConnection = config.serviceConnection.root.config
-        if not isinstance(connection, MysqlConnection):
-            raise InvalidSourceException(
-                f"Expected MysqlConnection, but got {connection}"
-            )
-        return cls(config, metadata)

metadata/ingestion/source/database/mysql/queries.py

@@ -0,0 +1,43 @@
+#  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.
+"""
+SQL Queries used during ingestion
+"""
+import textwrap
+
+MYSQL_SQL_STATEMENT = textwrap.dedent(
+    """
+SELECT 
+	NULL `database_name`,
+	argument `query_text`,
+	event_time `start_time`,
+    NULL `end_time`,
+	NULL `duration`,
+	NULL `schema_name`,
+    NULL `query_type`,
+    NULL `user_name`,
+    NULL `aborted`
+FROM mysql.general_log
+WHERE command_type = 'Query' 
+    AND event_time between '{start_time}' and '{end_time}'
+    AND argument NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
+    AND argument NOT LIKE '/* {{"app": "dbt", %%}} */%%'
+    {filters}
+ORDER BY event_time desc
+LIMIT {result_limit};
+"""
+)
+
+MYSQL_TEST_GET_QUERIES = textwrap.dedent(
+    """
+SELECT `argument` from mysql.general_log limit 1;
+"""
+)

metadata/ingestion/source/database/mysql/query_parser.py

@@ -0,0 +1,46 @@
+#  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.
+"""
+Mysql query parser module
+"""
+from abc import ABC
+from typing import Optional
+
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.query_parser_source import QueryParserSource
+
+
+class MysqlQueryParserSource(QueryParserSource, ABC):
+    """
+    Mysql base for Usage and Lineage
+    """
+
+    filters: str
+
+    @classmethod
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
+        """Create class instance"""
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: MysqlConnection = config.serviceConnection.root.config
+        if not isinstance(connection, MysqlConnection):
+            raise InvalidSourceException(
+                f"Expected MysqlConnection, but got {connection}"
+            )
+        return cls(config, metadata)

metadata/ingestion/source/database/mysql/service_spec.py

@@ -1,4 +1,10 @@
+from metadata.ingestion.source.database.mysql.lineage import MysqlLineageSource
 from metadata.ingestion.source.database.mysql.metadata import MysqlSource
+from metadata.ingestion.source.database.mysql.usage import MysqlUsageSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=MysqlSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=MysqlSource,
+    lineage_source_class=MysqlLineageSource,
+    usage_source_class=MysqlUsageSource,
+)

metadata/{data_quality/builders/pandas_validator_builder.py → ingestion/source/database/mysql/usage.py}

@@ -8,22 +8,17 @@
 #  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.
-
 """
-Builder defining the structure of builders for validators for Pandas sources
+MYSQL usage module
 """
-
-from typing import TYPE_CHECKING
-
-from metadata.data_quality.builders.i_validator_builder import IValidatorBuilder
-
-if TYPE_CHECKING:
-    pass
+from metadata.ingestion.source.database.mysql.queries import MYSQL_SQL_STATEMENT
+from metadata.ingestion.source.database.mysql.query_parser import MysqlQueryParserSource
+from metadata.ingestion.source.database.usage_source import UsageSource
 
 
-class PandasValidatorBuilder(IValidatorBuilder):
-    """Builder for Pandas validators"""
+class MysqlUsageSource(MysqlQueryParserSource, UsageSource):
+    sql_stmt = MYSQL_SQL_STATEMENT
+    filters = ""
 
-    def _get_source_type(self) -> str:
-        """Return the test case"""
-        return "pandas"
+    def format_query(self, query: bytes) -> str:
+        return query.decode(errors="ignore").replace("\\n", "\n")

metadata/ingestion/source/database/redshift/profiler/system.py

@@ -18,6 +18,7 @@ from metadata.profiler.metrics.system.system import (
     SQASessionProvider,
     SystemMetricsComputer,
 )
+from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.logger import profiler_logger
 from metadata.utils.profiler_utils import QueryResult
 from metadata.utils.time_utils import datetime_to_timestamp
@@ -45,11 +46,11 @@ class RedshiftSystemMetricsSource(
         return get_metric_result(queries, table)
 
     def get_kwargs(self, **kwargs):
-        table = kwargs.get("table")
+        runner: QueryRunner = kwargs.get("runner")
         return {
-            "table": table.__table__.name,
-            "database": self.get_session().get_bind().url.database,
-            "schema": table.__table__.schema,
+            "table": runner.table_name,
+            "database": runner.session.get_bind().url.database,
+            "schema": runner.schema_name,
         }
 
     def get_deletes(self, **kwargs) -> List[SystemProfile]:

metadata/ingestion/source/database/sample_data.py

@@ -136,7 +136,7 @@ from metadata.parsers.schema_parsers import (
     InvalidSchemaTypeException,
     schema_parser_config_registry,
 )
-from metadata.profiler.api.models import ProfilerResponse, SampleData
+from metadata.sampler.models import SampleData, SamplerResponse
 from metadata.utils import entity_link, fqn
 from metadata.utils.constants import UTF_8
 from metadata.utils.fqn import FQN_SEPARATOR
@@ -899,22 +899,14 @@ class SampleDataSource(
 
                 self.metadata.ingest_table_sample_data(
                     table_entity,
-                    ProfilerResponse(
+                    sample_data=SamplerResponse(
                         table=table_entity,
-                        profile=CreateTableProfileRequest(
-                            tableProfile=TableProfile(
-                                timestamp=Timestamp(
-                                    int(datetime.now().timestamp() * 1000)
-                                ),
-                                columnCount=1.0,
-                                rowCount=3.0,
-                            )
-                        ),
                         sample_data=SampleData(
                             data=TableData(
                                 rows=table["sampleData"]["rows"],
                                 columns=table["sampleData"]["columns"],
-                            )
+                            ),
+                            store=True,
                         ),
                     ).sample_data.data,
                 )

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

@@ -121,7 +121,7 @@ def get_connection(connection: SnowflakeConnection) -> Engine:
             )
         p_key = serialization.load_pem_private_key(
             bytes(connection.privateKey.get_secret_value(), "utf-8"),
-            password=snowflake_private_key_passphrase.encode(),
+            password=snowflake_private_key_passphrase.encode() or None,
             backend=default_backend(),
         )
         pkb = p_key.private_bytes(
@@ -137,11 +137,16 @@ def get_connection(connection: SnowflakeConnection) -> Engine:
             "client_session_keep_alive"
         ] = connection.clientSessionKeepAlive
 
-    return create_generic_db_connection(
+    engine = create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args_common,
     )
+    if connection.connectionArguments.root and connection.connectionArguments.root.get(
+        "private_key"
+    ):
+        del connection.connectionArguments.root["private_key"]
+    return engine
 
 
 def test_connection(
@@ -225,7 +230,7 @@ def execute_inspector_func(engine_wrapper: SnowflakeEngineWrapper, func_name: st
     the function with name `func_name` and executes it
     """
     _init_database(engine_wrapper)
-    engine_wrapper.engine.execute(f"USE DATABASE {engine_wrapper.database_name}")
+    engine_wrapper.engine.execute(f'USE DATABASE "{engine_wrapper.database_name}"')
     inspector = inspect(engine_wrapper.engine)
     inspector_fn = getattr(inspector, func_name)
     inspector_fn()