Skip to content

Commit

Permalink
MINOR: Fix Profiler for SSL Enabled Source (#16613)
Browse files Browse the repository at this point in the history
  • Loading branch information
ulixius9 authored Jun 12, 2024
1 parent 4f50e0b commit 7359d62
Show file tree
Hide file tree
Showing 8 changed files with 28 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,6 @@
from metadata.generated.schema.tests.customMetric import CustomMetric
from metadata.ingestion.api.status import Status
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.ingestion.source.connections import get_connection
from metadata.profiler.api.models import (
DatabaseAndSchemaConfig,
ProfilerProcessorConfig,
Expand All @@ -63,6 +62,7 @@
from metadata.profiler.processor.runner import QueryRunner
from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
from metadata.utils.partition import get_partition_details
from metadata.utils.ssl_manager import get_ssl_connection


class ProfilerProcessorStatus(Status):
Expand Down Expand Up @@ -110,7 +110,7 @@ def __init__(
self.ometa_client = ometa_client
self.source_config = source_config
self.service_connection_config = service_connection_config
self.connection = get_connection(self.service_connection_config)
self.connection = get_ssl_connection(self.service_connection_config)
self.status = ProfilerProcessorStatus()
try:
fqn = self.table_entity.fullyQualifiedName
Expand Down
4 changes: 2 additions & 2 deletions ingestion/src/metadata/profiler/source/metadata_ext.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,6 @@
)
from metadata.ingestion.api.models import Either
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.ingestion.source.connections import get_connection
from metadata.profiler.source.metadata import (
OpenMetadataSource,
ProfilerSourceAndEntity,
Expand All @@ -52,6 +51,7 @@
from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
from metadata.utils.importer import import_source_class
from metadata.utils.logger import profiler_logger
from metadata.utils.ssl_manager import get_ssl_connection

logger = profiler_logger()

Expand Down Expand Up @@ -112,7 +112,7 @@ def set_inspector(self, database_name: str = None) -> None:
if database_name:
logger.info(f"Ingesting from database: {database_name}")
new_service_connection.database = database_name
self.engine = get_connection(new_service_connection)
self.engine = get_ssl_connection(new_service_connection)
self.inspector = inspect(self.engine)
self._connection = None # Lazy init as well

Expand Down
16 changes: 16 additions & 0 deletions ingestion/src/metadata/utils/ssl_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
"""
import os
import tempfile
import traceback
from functools import singledispatch, singledispatchmethod
from typing import Optional, Union, cast

Expand Down Expand Up @@ -43,6 +44,10 @@
)
from metadata.generated.schema.security.ssl import verifySSLConfig
from metadata.ingestion.connections.builders import init_empty_connection_arguments
from metadata.ingestion.source.connections import get_connection
from metadata.utils.logger import utils_logger

logger = utils_logger()


class SSLManager:
Expand Down Expand Up @@ -173,3 +178,14 @@ def _(connection):
ca=connection.sslConfig.root.caCertificate if connection.sslConfig else None
)
return None


def get_ssl_connection(service_config):
try:
ssl_manager: SSLManager = check_ssl_and_init(service_config)
if ssl_manager:
service_config = ssl_manager.setup_ssl(service_config)
except Exception:
logger.debug("Failed to setup SSL for the connection")
logger.debug(traceback.format_exc())
return get_connection(service_config)
5 changes: 3 additions & 2 deletions ingestion/src/metadata/workflow/profiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,14 @@
OpenMetadataWorkflowConfig,
)
from metadata.ingestion.api.steps import Processor, Sink
from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
from metadata.ingestion.source.connections import get_test_connection_fn
from metadata.pii.processor import PIIProcessor
from metadata.profiler.processor.processor import ProfilerProcessor
from metadata.profiler.source.metadata import OpenMetadataSource
from metadata.profiler.source.metadata_ext import OpenMetadataSourceExt
from metadata.utils.importer import import_sink_class
from metadata.utils.logger import profiler_logger
from metadata.utils.ssl_manager import get_ssl_connection
from metadata.workflow.ingestion import IngestionWorkflow

logger = profiler_logger()
Expand Down Expand Up @@ -62,7 +63,7 @@ def set_steps(self):

def test_connection(self):
service_config = self.config.source.serviceConnection.root.config
conn = get_connection(service_config)
conn = get_ssl_connection(service_config)

test_connection_fn = get_test_connection_fn(service_config)
test_connection_fn(self.metadata, conn, service_config)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ class DatalakeMetricsTest(TestCase):

@classmethod
@mock.patch(
"metadata.profiler.interface.profiler_interface.get_connection",
"metadata.profiler.interface.profiler_interface.get_ssl_connection",
return_value=FakeConnection,
)
@mock.patch(
Expand Down
2 changes: 1 addition & 1 deletion ingestion/tests/unit/profiler/pandas/test_profiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ class ProfilerTest(TestCase):

@classmethod
@mock.patch(
"metadata.profiler.interface.profiler_interface.get_connection",
"metadata.profiler.interface.profiler_interface.get_ssl_connection",
return_value=FakeConnection,
)
@mock.patch(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,7 @@ class PandasInterfaceTest(TestCase):

@classmethod
@mock.patch(
"metadata.profiler.interface.profiler_interface.get_connection",
"metadata.profiler.interface.profiler_interface.get_ssl_connection",
return_value=FakeConnection,
)
@mock.patch(
Expand Down
4 changes: 2 additions & 2 deletions ingestion/tests/unit/profiler/pandas/test_sample.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ class DatalakeSampleTest(TestCase):

@classmethod
@mock.patch(
"metadata.profiler.interface.profiler_interface.get_connection",
"metadata.profiler.interface.profiler_interface.get_ssl_connection",
return_value=FakeConnection,
)
@mock.patch(
Expand Down Expand Up @@ -169,7 +169,7 @@ def test_random_sampler(self):
assert res < 5

@mock.patch(
"metadata.profiler.interface.profiler_interface.get_connection",
"metadata.profiler.interface.profiler_interface.get_ssl_connection",
return_value=FakeConnection,
)
@mock.patch(
Expand Down

0 comments on commit 7359d62

Please sign in to comment.