Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

MINOR: Fix Profiler for SSL Enabled Source #16613

Merged
merged 3 commits into from
Jun 12, 2024
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading