openmetadata-ingestion 1.8.2.0__py3-none-any.whl → 1.8.4.0__py3-none-any.whl

metadata/ingestion/source/database/domodatabase/metadata.py

@@ -109,12 +109,13 @@ class DomodatabaseSource(DatabaseServiceSource):
     def yield_database(
         self, database_name: str
     ) -> Iterable[Either[CreateDatabaseRequest]]:
-        yield Either(
-            right=CreateDatabaseRequest(
-                name=EntityName(database_name),
-                service=self.context.get().database_service,
-            )
+
+        database_request = CreateDatabaseRequest(
+            name=EntityName(database_name),
+            service=self.context.get().database_service,
         )
+        yield Either(right=database_request)
+        self.register_record_database_request(database_request=database_request)
 
     def get_database_schema_names(self) -> Iterable[str]:
         scheme_name = "default"
@@ -123,19 +124,19 @@ class DomodatabaseSource(DatabaseServiceSource):
     def yield_database_schema(
         self, schema_name: str
     ) -> Iterable[Either[CreateDatabaseSchemaRequest]]:
-        yield Either(
-            right=CreateDatabaseSchemaRequest(
-                name=EntityName(schema_name),
-                database=FullyQualifiedEntityName(
-                    fqn.build(
-                        metadata=self.metadata,
-                        entity_type=Database,
-                        service_name=self.context.get().database_service,
-                        database_name=self.context.get().database,
-                    )
-                ),
-            )
+        schema_request = CreateDatabaseSchemaRequest(
+            name=EntityName(schema_name),
+            database=FullyQualifiedEntityName(
+                fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                )
+            ),
         )
+        yield Either(right=schema_request)
+        self.register_record_schema_request(schema_request=schema_request)
 
     def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, str]]]:
         schema_name = self.context.get().database_schema

metadata/ingestion/source/database/glue/metadata.py

@@ -174,12 +174,12 @@ class GlueSource(ExternalTableLineageMixin, DatabaseServiceSource):
         From topology.
         Prepare a database request and pass it to the sink
         """
-        yield Either(
-            right=CreateDatabaseRequest(
-                name=database_name,
-                service=self.context.get().database_service,
-            )
+        database_request = CreateDatabaseRequest(
+            name=database_name,
+            service=self.context.get().database_service,
         )
+        yield Either(right=database_request)
+        self.register_record_database_request(database_request=database_request)
 
     def get_database_schema_names(self) -> Iterable[str]:
         """
@@ -224,24 +224,24 @@ class GlueSource(ExternalTableLineageMixin, DatabaseServiceSource):
         From topology.
         Prepare a database schema request and pass it to the sink
         """
-        yield Either(
-            right=CreateDatabaseSchemaRequest(
-                name=EntityName(schema_name),
-                description=self.schema_description_map.get(schema_name),
-                database=FullyQualifiedEntityName(
-                    fqn.build(
-                        metadata=self.metadata,
-                        entity_type=Database,
-                        service_name=self.context.get().database_service,
-                        database_name=self.context.get().database,
-                    )
-                ),
-                sourceUrl=self.get_source_url(
+        schema_request = CreateDatabaseSchemaRequest(
+            name=EntityName(schema_name),
+            description=self.schema_description_map.get(schema_name),
+            database=FullyQualifiedEntityName(
+                fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Database,
+                    service_name=self.context.get().database_service,
                     database_name=self.context.get().database,
-                    schema_name=schema_name,
-                ),
-            )
+                )
+            ),
+            sourceUrl=self.get_source_url(
+                database_name=self.context.get().database,
+                schema_name=schema_name,
+            ),
         )
+        yield Either(right=schema_request)
+        self.register_record_schema_request(schema_request=schema_request)
 
     def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, str]]]:
         """

metadata/ingestion/source/database/iceberg/metadata.py

@@ -108,12 +108,12 @@ class IcebergSource(DatabaseServiceSource):
 
         Also, update the self.inspector value to the current db.
         """
-        yield Either(
-            right=CreateDatabaseRequest(
-                name=database_name,
-                service=self.context.get().database_service,
-            )
+        database_request = CreateDatabaseRequest(
+            name=database_name,
+            service=self.context.get().database_service,
         )
+        yield Either(right=database_request)
+        self.register_record_database_request(database_request=database_request)
 
     def get_database_schema_names(self) -> Iterable[str]:
         """
@@ -155,19 +155,19 @@ class IcebergSource(DatabaseServiceSource):
         From topology.
         Prepare a database request and pass it to the sink.
         """
-        yield Either(
-            right=CreateDatabaseSchemaRequest(
-                name=EntityName(schema_name),
-                database=FullyQualifiedEntityName(
-                    fqn.build(
-                        metadata=self.metadata,
-                        entity_type=Database,
-                        service_name=self.context.get().database_service,
-                        database_name=self.context.get().database,
-                    )
-                ),
-            )
+        schema_request = CreateDatabaseSchemaRequest(
+            name=EntityName(schema_name),
+            database=FullyQualifiedEntityName(
+                fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                )
+            ),
         )
+        yield Either(right=schema_request)
+        self.register_record_schema_request(schema_request=schema_request)
 
     def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, str]]]:
         """

metadata/ingestion/source/database/postgres/utils.py

@@ -509,13 +509,32 @@ def get_postgres_time_column_name(engine) -> str:
     """
     Return the correct column name for the time column based on postgres version
     """
-    time_column_name = "total_exec_time"
-    postgres_version = get_postgres_version(engine)
-    if postgres_version and version.parse(postgres_version) < version.parse(
-        OLD_POSTGRES_VERSION
-    ):
-        time_column_name = "total_time"
-    return time_column_name
+    # Try to check the column in pg_stat_statements, fallback to version check if fails
+    try:
+        with engine.connect() as conn:
+            result = conn.execute(
+                "SELECT column_name FROM information_schema.columns WHERE table_name = 'pg_stat_statements'"
+            )
+            columns = {row[0] for row in result}
+            if "total_exec_time" in columns:
+                return "total_exec_time"
+            elif "total_time" in columns:
+                return "total_time"
+            else:
+                logger.warning(
+                    "Neither 'total_exec_time' nor 'total_time' found in pg_stat_statements. Defaulting to 'total_exec_time'."
+                )
+                return "total_exec_time"
+    except Exception as ex:
+        logger.debug(f"Failed to check columns in pg_stat_statements: {ex}")
+        # Fallback to version check
+        time_column_name = "total_exec_time"
+        postgres_version = get_postgres_version(engine)
+        if postgres_version and version.parse(postgres_version) < version.parse(
+            OLD_POSTGRES_VERSION
+        ):
+            time_column_name = "total_time"
+        return time_column_name
 
 
 @reflection.cache

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

@@ -119,12 +119,12 @@ class SalesforceSource(DatabaseServiceSource):
         From topology.
         Prepare a database request and pass it to the sink
         """
-        yield Either(
-            right=CreateDatabaseRequest(
-                name=database_name,
-                service=self.context.get().database_service,
-            )
+        database_request = CreateDatabaseRequest(
+            name=database_name,
+            service=self.context.get().database_service,
         )
+        yield Either(right=database_request)
+        self.register_record_database_request(database_request=database_request)
 
     def get_database_schema_names(self) -> Iterable[str]:
         """
@@ -139,19 +139,19 @@ class SalesforceSource(DatabaseServiceSource):
         From topology.
         Prepare a database schema request and pass it to the sink
         """
-        yield Either(
-            right=CreateDatabaseSchemaRequest(
-                name=EntityName(schema_name),
-                database=FullyQualifiedEntityName(
-                    fqn.build(
-                        metadata=self.metadata,
-                        entity_type=Database,
-                        service_name=self.context.get().database_service,
-                        database_name=self.context.get().database,
-                    )
-                ),
-            )
+        schema_request = CreateDatabaseSchemaRequest(
+            name=EntityName(schema_name),
+            database=FullyQualifiedEntityName(
+                fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                )
+            ),
         )
+        yield Either(right=schema_request)
+        self.register_record_schema_request(schema_request=schema_request)
 
     def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, str]]]:
         """

metadata/ingestion/source/database/sas/metadata.py

@@ -845,12 +845,12 @@ class SasSource(
     def yield_database(
         self, database_name: str
     ) -> Iterable[Either[CreateDatabaseRequest]]:
-        yield Either(
-            right=CreateDatabaseRequest(
-                name=EntityName(database_name),
-                service=self.context.get().database_service,
-            )
+        database_request = CreateDatabaseRequest(
+            name=EntityName(database_name),
+            service=self.context.get().database_service,
         )
+        yield Either(right=database_request)
+        self.register_record_database_request(database_request=database_request)
 
     def get_database_schema_names(self) -> Iterable[Tuple[str, str]]:
         for database, database_schemas in self.database_schemas.items():
@@ -860,18 +860,20 @@ class SasSource(
     def yield_database_schema(
         self, schema_name: Tuple[str, str]
     ) -> Iterable[Either[CreateDatabaseSchemaRequest]]:
-        yield Either(
-            right=CreateDatabaseSchemaRequest(
-                name=schema_name[1],
-                database=fqn.build(
-                    metadata=self.metadata,
-                    entity_type=Database,
-                    service_name=self.context.get().database_service,
-                    database_name=schema_name[0],
-                ),
-            )
+
+        schema_request = CreateDatabaseSchemaRequest(
+            name=schema_name[1],
+            database=fqn.build(
+                metadata=self.metadata,
+                entity_type=Database,
+                service_name=self.context.get().database_service,
+                database_name=schema_name[0],
+            ),
         )
 
+        yield Either(right=schema_request)
+        self.register_record_schema_request(schema_request=schema_request)
+
     def yield_tag(
         self, schema_name: str
     ) -> Iterable[Either[OMetaTagAndClassification]]:

metadata/ingestion/source/database/unitycatalog/metadata.py

@@ -193,15 +193,15 @@ class UnitycatalogSource(
         Prepare a database request and pass it to the sink
         """
         catalog = self.client.catalogs.get(database_name)
-        yield Either(
-            right=CreateDatabaseRequest(
-                name=database_name,
-                service=self.context.get().database_service,
-                owners=self.get_owner_ref(catalog.owner),
-                description=catalog.comment,
-                tags=self.get_database_tag_labels(database_name),
-            )
+        database_request = CreateDatabaseRequest(
+            name=database_name,
+            service=self.context.get().database_service,
+            owners=self.get_owner_ref(catalog.owner),
+            description=catalog.comment,
+            tags=self.get_database_tag_labels(database_name),
         )
+        yield Either(right=database_request)
+        self.register_record_database_request(database_request=database_request)
 
     def get_database_schema_names(self) -> Iterable[str]:
         """
@@ -247,22 +247,22 @@ class UnitycatalogSource(
         schema = self.client.schemas.get(
             full_name=f"{self.context.get().database}.{schema_name}"
         )
-        yield Either(
-            right=CreateDatabaseSchemaRequest(
-                name=EntityName(schema_name),
-                database=FullyQualifiedEntityName(
-                    fqn.build(
-                        metadata=self.metadata,
-                        entity_type=Database,
-                        service_name=self.context.get().database_service,
-                        database_name=self.context.get().database,
-                    )
-                ),
-                description=schema.comment,
-                owners=self.get_owner_ref(schema.owner),
-                tags=self.get_schema_tag_labels(schema_name),
-            )
+        schema_request = CreateDatabaseSchemaRequest(
+            name=EntityName(schema_name),
+            database=FullyQualifiedEntityName(
+                fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                )
+            ),
+            description=schema.comment,
+            owners=self.get_owner_ref(schema.owner),
+            tags=self.get_schema_tag_labels(schema_name),
         )
+        yield Either(right=schema_request)
+        self.register_record_schema_request(schema_request=schema_request)
 
     def get_tables_name_and_type(self) -> Iterable[Tuple[str, str]]:
         """

metadata/pii/algorithms/classifiers.py

@@ -73,9 +73,6 @@ class ColumnClassifier(ABC, Generic[T]):
         """
 
 
-# Implementations
-
-
 @final
 class HeuristicPIIClassifier(ColumnClassifier[PIITag]):
     """

metadata/pii/algorithms/presidio_utils.py

@@ -13,7 +13,7 @@ Utilities for working with the Presidio Library.
 """
 import inspect
 import logging
-from typing import Iterable, Optional, Type
+from typing import Iterable, Type, Union
 
 import spacy
 from presidio_analyzer import (
@@ -26,7 +26,7 @@ from presidio_analyzer.nlp_engine import SpacyNlpEngine
 from spacy.cli.download import download  # pyright: ignore[reportUnknownVariableType]
 
 from metadata.pii.constants import PRESIDIO_LOGGER, SPACY_EN_MODEL, SUPPORTED_LANG
-from metadata.utils.logger import METADATA_LOGGER, pii_logger
+from metadata.utils.logger import pii_logger
 
 logger = pii_logger()
 
@@ -64,17 +64,10 @@ def build_analyzer_engine(
     return analyzer_engine
 
 
-def set_presidio_logger_level(log_level: Optional[int] = None) -> None:
+def set_presidio_logger_level(log_level: Union[int, str] = logging.ERROR) -> None:
     """
     Set the presidio logger to talk less about internal entities unless we are debugging.
     """
-    if log_level is None:
-        log_level = (
-            logging.INFO
-            if logging.getLogger(METADATA_LOGGER).level == logging.DEBUG
-            else logging.ERROR
-        )
-
     logging.getLogger(PRESIDIO_LOGGER).setLevel(log_level)
 
 
@@ -87,7 +80,6 @@ def _load_spacy_model(model_name: str) -> None:
     try:
         _ = spacy.load(model_name)
     except OSError:
-
         logger.warning(f"Downloading {model_name} language model for the spaCy")
         download(model_name)
         _ = spacy.load(model_name)

metadata/pii/scanners/ner_scanner.py

@@ -22,6 +22,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 from pydantic import BaseModel, ConfigDict
 
 from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.pii.algorithms.presidio_utils import _load_spacy_model
 from metadata.pii.constants import PII, SPACY_EN_MODEL
 from metadata.pii.models import TagAndConfidence
 from metadata.pii.ner import NEREntity
@@ -56,18 +57,10 @@ class NERScanner(BaseScanner):
     """Based on https://microsoft.github.io/presidio/"""
 
     def __init__(self):
-        import spacy
         from presidio_analyzer import AnalyzerEngine
         from presidio_analyzer.nlp_engine.spacy_nlp_engine import SpacyNlpEngine
 
-        try:
-            spacy.load(SPACY_EN_MODEL)
-        except OSError:
-            logger.warning("Downloading en_core_web_md language model for the spaCy")
-            from spacy.cli import download
-
-            download(SPACY_EN_MODEL)
-            spacy.load(SPACY_EN_MODEL)
+        _load_spacy_model(SPACY_EN_MODEL)
 
         nlp_engine_model = NLPEngineModel(
             lang_code=SUPPORTED_LANG, model_name=SPACY_EN_MODEL