openmetadata-ingestion 1.8.7.0__py3-none-any.whl → 1.8.7.2__py3-none-any.whl

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2025-07-24T06:26:23+00:00
+#   timestamp: 2025-08-04T10:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2025-07-24T06:26:23+00:00
+#   timestamp: 2025-08-04T10:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2025-07-24T06:26:23+00:00
+#   timestamp: 2025-08-04T10:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2025-07-24T06:26:23+00:00
+#   timestamp: 2025-08-04T10:21:08+00:00
 
 from __future__ import annotations
 

metadata/ingestion/models/custom_pydantic.py

@@ -199,6 +199,7 @@ class _CustomSecretStr(SecretStr):
             and SecretsManagerFactory().get_secrets_manager()
         ):
             secret_id = self._secret_value.replace(SECRET, "")
+            logger.info(f"Getting secret value for {secret_id}")
             try:
                 return (
                     SecretsManagerFactory()

metadata/ingestion/ometa/ometa_api.py

@@ -59,7 +59,12 @@ from metadata.ingestion.ometa.mixins.user_mixin import OMetaUserMixin
 from metadata.ingestion.ometa.mixins.version_mixin import OMetaVersionMixin
 from metadata.ingestion.ometa.models import EntityList
 from metadata.ingestion.ometa.routes import ROUTES
-from metadata.ingestion.ometa.utils import get_entity_type, model_str, quote
+from metadata.ingestion.ometa.utils import (
+    decode_jwt_token,
+    get_entity_type,
+    model_str,
+    quote,
+)
 from metadata.utils.logger import ometa_logger
 from metadata.utils.secrets.secrets_manager_factory import SecretsManagerFactory
 from metadata.utils.ssl_registry import get_verify_ssl_fn
@@ -156,6 +161,7 @@ class OpenMetadata(
         ).get_secrets_manager()
 
         self._auth_provider = OpenMetadataAuthenticationProvider.create(self.config)
+        self.log_user_name_from_jwt_token()
 
         get_verify_ssl = get_verify_ssl_fn(self.config.verifySSL)
 
@@ -178,6 +184,29 @@ class OpenMetadata(
         if self.config.enableVersionValidation:
             self.validate_versions()
 
+    def log_user_name_from_jwt_token(self) -> None:
+        """
+        Log user name from JWT token.
+        """
+        # Log user name from JWT token if authProvider is openmetadata
+        if (
+            self.config.authProvider
+            and self.config.authProvider.value == "openmetadata"
+        ):
+            try:
+                # Get the JWT token from the auth provider
+                jwt_token, _ = self._auth_provider.get_access_token()
+                if jwt_token:
+                    # Decode the JWT token to extract user information
+                    payload = decode_jwt_token(jwt_token)
+                    if payload:
+                        if payload.get("sub"):
+                            logger.debug(f"Authenticated user: {payload.get('sub')}")
+                        else:
+                            logger.debug("Could not extract user name from JWT token")
+            except Exception as e:
+                logger.debug(f"Error processing JWT token: {e}")
+
     @classmethod
     def from_env(cls) -> "OpenMetadata":
         settings = OpenMetadataSettings()

metadata/ingestion/ometa/utils.py

@@ -12,9 +12,11 @@
 Helper functions to handle OpenMetadata Entities' properties
 """
 
+import base64
+import json
 import re
 import string
-from typing import Any, Type, TypeVar, Union
+from typing import Any, Dict, Optional, Type, TypeVar, Union
 
 from pydantic import BaseModel
 from requests.utils import quote as url_quote
@@ -98,3 +100,32 @@ def build_entity_reference(entity: T) -> EntityReference:
         description=entity.description,
         href=entity.href,
     )
+
+
+def decode_jwt_token(jwt_token: str) -> Optional[Dict[str, Any]]:
+    """
+    Decode JWT token to extract payload without verification.
+    JWT tokens have three parts: header.payload.signature
+    We only decode the payload part to get user information.
+    """
+    from metadata.utils.logger import ometa_logger
+
+    logger = ometa_logger()
+    try:
+        # Split the JWT token into its three parts
+        parts = jwt_token.split(".")
+        if len(parts) != 3:
+            logger.debug("Invalid JWT token format: expected 3 parts")
+            return None
+
+        # Decode the payload (second part)
+        payload = parts[1]
+        # Add padding if needed for base64 decoding
+        payload += "=" * (4 - len(payload) % 4) if len(payload) % 4 else ""
+
+        # Decode base64 and parse JSON
+        decoded_payload = base64.urlsafe_b64decode(payload)
+        return json.loads(decoded_payload.decode("utf-8"))
+    except Exception as e:
+        logger.debug(f"Failed to decode JWT token: {e}")
+        return None

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

@@ -17,6 +17,7 @@ from typing import Optional
 
 from databricks.sdk import WorkspaceClient
 from sqlalchemy.engine import Engine
+from sqlalchemy.exc import DatabaseError
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -34,13 +35,13 @@ from metadata.ingestion.connections.builders import (
 )
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
 from metadata.ingestion.source.database.unitycatalog.models import DatabricksTable
 from metadata.ingestion.source.database.unitycatalog.queries import (
     UNITY_CATALOG_GET_ALL_SCHEMA_TAGS,
     UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS,
     UNITY_CATALOG_GET_ALL_TABLE_TAGS,
     UNITY_CATALOG_GET_CATALOGS_TAGS,
+    UNITY_CATALOG_SQL_STATEMENT_TEST,
 )
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.db_utils import get_host_from_host_port
@@ -93,8 +94,20 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    client = UnityCatalogClient(service_connection)
     table_obj = DatabricksTable()
+    engine = get_sqlalchemy_connection(service_connection)
+
+    def test_database_query(engine: Engine, statement: str):
+        """
+        Method used to execute the given query and fetch a result
+        to test if user has access to the tables specified
+        in the sql statement
+        """
+        try:
+            connection = engine.connect()
+            connection.execute(statement).fetchone()
+        except DatabaseError as soe:
+            logger.debug(f"Failed to fetch catalogs due to: {soe}")
 
     def get_catalogs(connection: WorkspaceClient, table_obj: DatabricksTable):
         for catalog in connection.catalogs.list():
@@ -149,7 +162,11 @@ def test_connection(
         "GetSchemas": partial(get_schemas, connection, table_obj),
         "GetTables": partial(get_tables, connection, table_obj),
         "GetViews": partial(get_tables, connection, table_obj),
-        "GetQueries": client.test_query_api_access,
+        "GetQueries": partial(
+            test_database_query,
+            engine=engine,
+            statement=UNITY_CATALOG_SQL_STATEMENT_TEST,
+        ),
         "GetTags": partial(get_tags, service_connection, table_obj),
     }
 

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

@@ -1,7 +1,19 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  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 for Unity Catalog
+SQL Queries used during ingestion
 """
 
+import textwrap
+
 UNITY_CATALOG_GET_CATALOGS_TAGS = """
 SELECT * FROM `{database}`.information_schema.catalog_tags;
 """
@@ -17,3 +29,27 @@ SELECT * FROM `{database}`.information_schema.table_tags WHERE schema_name = '{s
 UNITY_CATALOG_GET_ALL_TABLE_COLUMNS_TAGS = """
 SELECT * FROM `{database}`.information_schema.column_tags WHERE schema_name = '{schema}';
 """
+
+UNITY_CATALOG_SQL_STATEMENT = textwrap.dedent(
+    """
+    SELECT
+      statement_type AS query_type,
+      statement_text AS query_text,
+      executed_by AS user_name,
+      start_time AS start_time,
+      null AS database_name,
+      null AS schema_name,
+      end_time AS end_time,
+      total_duration_ms/1000 AS duration
+    from system.query.history
+    WHERE statement_text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
+    AND statement_text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
+    AND start_time between to_timestamp('{start_time}') and to_timestamp('{end_time}')
+    {filters}
+    LIMIT {result_limit}
+    """
+)
+
+UNITY_CATALOG_SQL_STATEMENT_TEST = """
+ SELECT statement_text from  system.query.history LIMIT 1
+"""

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

@@ -27,6 +27,9 @@ from metadata.ingestion.source.database.databricks.query_parser import (
 )
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
+from metadata.ingestion.source.database.unitycatalog.connection import (
+    get_sqlalchemy_connection,
+)
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -55,6 +58,7 @@ class UnityCatalogQueryParserSource(
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         self._init_super(config=config, metadata=metadata)
         self.client = UnityCatalogClient(self.service_connection)
+        self.sql_client = get_sqlalchemy_connection(self.service_connection)
 
     @classmethod
     def create(

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

@@ -11,12 +11,10 @@
 """
 unity catalog usage module
 """
-import traceback
-from datetime import datetime
-from typing import Iterable
 
-from metadata.generated.schema.type.basic import DateTime
-from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
+from metadata.ingestion.source.database.unitycatalog.queries import (
+    UNITY_CATALOG_SQL_STATEMENT,
+)
 from metadata.ingestion.source.database.unitycatalog.query_parser import (
     UnityCatalogQueryParserSource,
 )
@@ -35,36 +33,24 @@ class UnitycatalogUsageSource(UnityCatalogQueryParserSource, UsageSource):
     the same API for fetching Usage Queries
     """
 
-    def yield_table_queries(self) -> Iterable[TableQuery]:
+    sql_stmt = UNITY_CATALOG_SQL_STATEMENT
+
+    filters = """
+        AND statement_type NOT IN ('SHOW', 'DESCRIBE', 'USE')
+    """
+
+    def get_engine(self):
+        yield self.sql_client
+
+    def get_sql_statement(self, start_time, end_time):
         """
-        Method to yield TableQueries
+        returns sql statement to fetch query logs.
+
+        Override if we have specific parameters
         """
-        queries = []
-        data = self.client.list_query_history(
-            start_date=self.start,
-            end_date=self.end,
+        return self.sql_stmt.format(
+            start_time=start_time,
+            end_time=end_time,
+            filters=self.get_filters(),
+            result_limit=self.source_config.resultLimit,
         )
-        for row in data or []:
-            try:
-                if self.client.is_query_valid(row):
-                    queries.append(
-                        TableQuery(
-                            dialect=self.dialect.value,
-                            query=row.get("query_text"),
-                            userName=row.get("user_name"),
-                            startTime=str(row.get("query_start_time_ms")),
-                            endTime=str(row.get("execution_end_time_ms")),
-                            analysisDate=DateTime(datetime.now()),
-                            serviceName=self.config.serviceName,
-                            duration=row.get("duration")
-                            if row.get("duration")
-                            else None,
-                        )
-                    )
-            except Exception as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Failed to process query {row.get('query_text')} due to: {err}"
-                )
-
-        yield TableQueries(queries=queries)

metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -187,36 +187,44 @@ class AirflowSource(PipelineServiceSource):
         """
         Return the DagRuns of given dag
         """
-        dag_run_list = (
-            self.session.query(
-                DagRun.dag_id,
-                DagRun.run_id,
-                DagRun.queued_at,
-                DagRun.execution_date,
-                DagRun.start_date,
-                DagRun.state,
+        try:
+            dag_run_list = (
+                self.session.query(
+                    DagRun.dag_id,
+                    DagRun.run_id,
+                    DagRun.queued_at,
+                    DagRun.execution_date,
+                    DagRun.start_date,
+                    DagRun.state,
+                )
+                .filter(DagRun.dag_id == dag_id)
+                .order_by(DagRun.execution_date.desc())
+                .limit(self.config.serviceConnection.root.config.numberOfStatus)
+                .all()
             )
-            .filter(DagRun.dag_id == dag_id)
-            .order_by(DagRun.execution_date.desc())
-            .limit(self.config.serviceConnection.root.config.numberOfStatus)
-            .all()
-        )
 
-        dag_run_dict = [dict(elem) for elem in dag_run_list]
-
-        # Build DagRun manually to not fall into new/old columns from
-        # different Airflow versions
-        return [
-            DagRun(
-                dag_id=elem.get("dag_id"),
-                run_id=elem.get("run_id"),
-                queued_at=elem.get("queued_at"),
-                execution_date=elem.get("execution_date"),
-                start_date=elem.get("start_date"),
-                state=elem.get("state"),
+            dag_run_dict = [dict(elem) for elem in dag_run_list]
+
+            # Build DagRun manually to not fall into new/old columns from
+            # different Airflow versions
+            return [
+                DagRun(
+                    dag_id=elem.get("dag_id"),
+                    run_id=elem.get("run_id"),
+                    queued_at=elem.get("queued_at"),
+                    execution_date=elem.get("execution_date"),
+                    start_date=elem.get("start_date"),
+                    state=elem.get("state"),
+                )
+                for elem in dag_run_dict
+            ]
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Could not get pipeline status for {dag_id}. "
+                f"This might be due to Airflow version incompatibility - {exc}"
             )
-            for elem in dag_run_dict
-        ]
+            return []
 
     def get_task_instances(
         self, dag_id: str, run_id: str, serialized_tasks: List[AirflowTask]
@@ -369,16 +377,27 @@ class AirflowSource(PipelineServiceSource):
                 break
             for serialized_dag in results:
                 try:
-                    dag_model = (
-                        self.session.query(DagModel)
-                        .filter(DagModel.dag_id == serialized_dag[0])
-                        .one_or_none()
-                    )
-                    pipeline_state = (
-                        PipelineState.Active.value
-                        if dag_model and not dag_model.is_paused
-                        else PipelineState.Inactive.value
-                    )
+                    # Query only the is_paused column from DagModel
+                    try:
+                        is_paused_result = (
+                            self.session.query(DagModel.is_paused)
+                            .filter(DagModel.dag_id == serialized_dag[0])
+                            .scalar()
+                        )
+                        pipeline_state = (
+                            PipelineState.Active.value
+                            if not is_paused_result
+                            else PipelineState.Inactive.value
+                        )
+                    except Exception as exc:
+                        logger.debug(traceback.format_exc())
+                        logger.warning(
+                            f"Could not query DagModel.is_paused for {serialized_dag[0]}. "
+                            f"Using default pipeline state - {exc}"
+                        )
+                        # If we can't query is_paused, assume the pipeline is active
+                        pipeline_state = PipelineState.Active.value
+
                     data = serialized_dag[1]["dag"]
                     dag = AirflowDagDetails(
                         dag_id=serialized_dag[0],

metadata/ingestion/source/pipeline/airflow/utils.py

@@ -23,6 +23,7 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
+# pylint: disable=too-many-branches,too-many-return-statements,too-many-nested-blocks
 def get_schedule_interval(pipeline_data: Dict[str, Any]) -> Optional[str]:
     """
     Fetch Schedule Intervals from Airflow Dags
@@ -40,7 +41,49 @@ def get_schedule_interval(pipeline_data: Dict[str, Any]) -> Optional[str]:
 
             expression_class = timetable.get("__type")
             if expression_class:
-                return import_from_module(expression_class)().summary
+                try:
+                    # Try to instantiate the timetable class safely
+                    timetable_class = import_from_module(expression_class)
+
+                    # Handle special cases for classes that require arguments
+                    if "DatasetTriggeredTimetable" in expression_class:
+                        # DatasetTriggeredTimetable requires datasets argument
+                        # For now, return a descriptive string since we can't instantiate it properly
+                        return "Dataset Triggered"
+                    if "CronDataIntervalTimetable" in expression_class:
+                        # Handle cron-based timetables
+                        try:
+                            return timetable_class().summary
+                        except (TypeError, AttributeError):
+                            return "Cron Based"
+                    else:
+                        # Try to instantiate with no arguments
+                        try:
+                            return timetable_class().summary
+                        except (TypeError, AttributeError):
+                            # If summary attribute doesn't exist, try to get a string representation
+                            try:
+                                instance = timetable_class()
+                                return str(instance)
+                            except TypeError:
+                                # If instantiation fails, return the class name
+                                return f"Custom Timetable ({expression_class.split('.')[-1]})"
+                except ImportError as import_error:
+                    logger.debug(
+                        f"Could not import timetable class {expression_class}: {import_error}"
+                    )
+                    return f"Custom Timetable ({expression_class.split('.')[-1]})"
+                except TypeError as type_error:
+                    # If instantiation fails due to missing arguments, log and continue
+                    logger.debug(
+                        f"Could not instantiate timetable class {expression_class}: {type_error}"
+                    )
+                    return f"Custom Timetable ({expression_class.split('.')[-1]})"
+                except Exception as inst_error:
+                    logger.debug(
+                        f"Error instantiating timetable class {expression_class}: {inst_error}"
+                    )
+                    return f"Custom Timetable ({expression_class.split('.')[-1]})"
 
         if schedule:
             if isinstance(schedule, str):
@@ -57,7 +100,6 @@ def get_schedule_interval(pipeline_data: Dict[str, Any]) -> Optional[str]:
 
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        logger.warning(
-            f"Couldn't fetch schedule interval for dag {pipeline_data.get('_dag_id'): [{exc}]}"
-        )
+        dag_id = pipeline_data.get("_dag_id", "unknown")
+        logger.warning(f"Couldn't fetch schedule interval for dag {dag_id}: {exc}")
     return None