openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

metadata/data_quality/validations/table/sqlalchemy/tableCustomSQLQuery.py

@@ -13,21 +13,278 @@
 Validator for table custom SQL Query test case
 """
 
+from typing import Optional, Tuple, cast
+
+import sqlparse
 from sqlalchemy import text
+from sqlalchemy.sql import func, select
+from sqlparse.sql import Statement, Token, Where
+from sqlparse.tokens import Keyword
 
 from metadata.data_quality.validations.mixins.sqa_validator_mixin import (
     SQAValidatorMixin,
 )
+from metadata.data_quality.validations.models import (
+    TableCustomSQLQueryRuntimeParameters,
+)
 from metadata.data_quality.validations.table.base.tableCustomSQLQuery import (
     BaseTableCustomSQLQueryValidator,
     Strategy,
 )
+from metadata.generated.schema.tests.basic import TestCaseResult
+from metadata.profiler.metrics.registry import Metrics
+from metadata.profiler.orm.functions.table_metric_computer import TableMetricComputer
+from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.helpers import is_safe_sql_query
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class TableCustomSQLQueryValidator(BaseTableCustomSQLQueryValidator, SQAValidatorMixin):
     """Validator for table custom SQL Query test case"""
 
+    def _replace_where_clause(
+        self, sql_query: str, partition_expression: str
+    ) -> Optional[str]:
+        """Replace or add WHERE clause in SQL query using sqlparse.
+
+        This method properly handles:
+        - Queries with existing WHERE clause (replaces it)
+        - Queries without WHERE clause (adds it)
+        - Complex queries with joins, subqueries, CTEs
+        - Preserves GROUP BY, ORDER BY, LIMIT, etc.
+
+        Args:
+            sql_query: Original SQL query
+            partition_expression: New WHERE condition (without WHERE keyword)
+
+        Returns:
+            Modified SQL query with partition_expression as WHERE clause
+        """
+        parsed = sqlparse.parse(sql_query)
+        if not parsed or len(parsed) == 0:
+            return None
+
+        statement: Statement = parsed[0]
+        tokens = list(statement.tokens)
+
+        where_idx, where_end_idx, insert_before_idx = self._find_clause_positions(
+            tokens
+        )
+        new_tokens = self._build_new_tokens(
+            tokens, where_idx, where_end_idx, insert_before_idx, partition_expression
+        )
+
+        return "".join(str(token) for token in new_tokens)
+
+    def _find_clause_positions(
+        self, tokens: list
+    ) -> Tuple[Optional[int], Optional[int], Optional[int]]:
+        """Find positions of WHERE clause and insertion points in token list.
+
+        Args:
+            tokens: List of parsed SQL tokens
+
+        Returns:
+            Tuple of (where_idx, where_end_idx, insert_before_idx)
+        """
+        where_idx = None
+        where_end_idx = None
+        insert_before_idx = None
+        paren_depth = 0
+
+        for i, token in enumerate(tokens):
+            paren_depth = self._update_parentheses_depth(token, paren_depth)
+
+            if isinstance(token, Where) and paren_depth == 0:
+                where_idx = i
+                where_end_idx = i + 1
+                break
+
+            if self._should_insert_before_token(token, insert_before_idx, paren_depth):
+                insert_before_idx = i
+
+        return where_idx, where_end_idx, insert_before_idx
+
+    def _update_parentheses_depth(self, token: Token, current_depth: int) -> int:
+        """Update parentheses depth based on token content.
+
+        Args:
+            token: SQL token to analyze
+            current_depth: Current parentheses depth
+
+        Returns:
+            Updated parentheses depth
+        """
+        if token.ttype is None and hasattr(token, "tokens"):
+            paren_count = str(token).count("(") - str(token).count(")")
+            return current_depth + paren_count
+        elif token.value == "(":
+            return current_depth + 1
+        elif token.value == ")":
+            return current_depth - 1
+        return current_depth
+
+    def _should_insert_before_token(
+        self, token: Token, insert_before_idx: Optional[int], paren_depth: int
+    ) -> bool:
+        """Check if WHERE clause should be inserted before this token.
+
+        Args:
+            token: SQL token to check
+            insert_before_idx: Current insertion index (None if not set)
+            paren_depth: Current parentheses depth
+
+        Returns:
+            True if WHERE should be inserted before this token
+        """
+        if insert_before_idx is not None or paren_depth != 0:
+            return False
+
+        if token.ttype is not Keyword:
+            return False
+
+        clause_keywords = {
+            "GROUP BY",
+            "ORDER BY",
+            "HAVING",
+            "LIMIT",
+            "OFFSET",
+            "UNION",
+            "EXCEPT",
+            "INTERSECT",
+        }
+
+        return any(keyword in token.value.upper() for keyword in clause_keywords)
+
+    def _build_new_tokens(
+        self,
+        tokens: list,
+        where_idx: Optional[int],
+        where_end_idx: Optional[int],
+        insert_before_idx: Optional[int],
+        partition_expression: str,
+    ) -> list:
+        """Build new token list with WHERE clause inserted or replaced.
+
+        Args:
+            tokens: Original token list
+            where_idx: Index of existing WHERE clause (None if not found)
+            where_end_idx: End index of existing WHERE clause
+            insert_before_idx: Index to insert WHERE before (None if append)
+            partition_expression: WHERE condition expression
+
+        Returns:
+            New list of tokens with WHERE clause
+        """
+        if where_idx is not None:
+            return self._replace_existing_where(
+                tokens, where_idx, where_end_idx, partition_expression
+            )
+        elif insert_before_idx is not None:
+            return self._insert_where_before_clause(
+                tokens, insert_before_idx, partition_expression
+            )
+        else:
+            return self._append_where_clause(tokens, partition_expression)
+
+    def _replace_existing_where(
+        self,
+        tokens: list,
+        where_idx: int,
+        where_end_idx: int,
+        partition_expression: str,
+    ) -> list:
+        """Replace existing WHERE clause with new expression.
+
+        Args:
+            tokens: Original token list
+            where_idx: Index of WHERE clause to replace
+            where_end_idx: End index of WHERE clause
+            partition_expression: New WHERE condition
+
+        Returns:
+            Token list with replaced WHERE clause
+        """
+        original_where = str(tokens[where_idx])
+        trailing_whitespace = self._extract_trailing_whitespace(original_where)
+
+        return (
+            tokens[:where_idx]
+            + [
+                Token(Keyword, "WHERE"),
+                Token(None, f" {partition_expression}{trailing_whitespace}"),
+            ]
+            + tokens[where_end_idx:]
+        )
+
+    def _insert_where_before_clause(
+        self, tokens: list, insert_before_idx: int, partition_expression: str
+    ) -> list:
+        """Insert WHERE clause before specified token index.
+
+        Args:
+            tokens: Original token list
+            insert_before_idx: Index to insert WHERE clause before
+            partition_expression: WHERE condition expression
+
+        Returns:
+            Token list with WHERE clause inserted
+        """
+        return (
+            tokens[:insert_before_idx]
+            + [Token(Keyword, "WHERE"), Token(None, f" {partition_expression} ")]
+            + tokens[insert_before_idx:]
+        )
+
+    def _append_where_clause(self, tokens: list, partition_expression: str) -> list:
+        """Append WHERE clause to end of token list.
+
+        Args:
+            tokens: Original token list
+            partition_expression: WHERE condition expression
+
+        Returns:
+            Token list with WHERE clause appended
+        """
+        return tokens + [
+            Token(None, " "),
+            Token(Keyword, "WHERE"),
+            Token(None, f" {partition_expression}"),
+        ]
+
+    def _extract_trailing_whitespace(self, where_clause: str) -> str:
+        """Extract trailing whitespace from WHERE clause string.
+
+        Args:
+            where_clause: Original WHERE clause string
+
+        Returns:
+            Trailing whitespace string
+        """
+        if not where_clause.split():
+            return ""
+
+        last_word = where_clause.split()[-1]
+        where_content_end = where_clause.rfind(last_word) + len(last_word)
+
+        if where_content_end < len(where_clause):
+            return where_clause[where_content_end:]
+
+        return ""
+
+    def run_validation(self) -> TestCaseResult:
+        """Run validation for the given test case
+
+        Returns:
+            TestCaseResult:
+        """
+        self.runtime_params = self.get_runtime_parameters(
+            TableCustomSQLQueryRuntimeParameters
+        )
+        return super().run_validation()
+
     def _run_results(self, sql_expression: str, strategy: Strategy = Strategy.ROWS):
         """compute result of the test case"""
         if not is_safe_sql_query(sql_expression):
@@ -48,3 +305,66 @@ class TableCustomSQLQueryValidator(BaseTableCustomSQLQueryValidator, SQAValidato
         except Exception as exc:
             self.runner._session.rollback()  # pylint: disable=protected-access
             raise exc
+
+    def compute_row_count(self) -> Optional[int]:
+        """Compute row count for the given column
+
+        Raises:
+            NotImplementedError:
+        """
+        partition_expression = next(
+            (
+                param.value
+                for param in self.test_case.parameterValues
+                if param.name == "partitionExpression"
+            ),
+            None,
+        )
+        if partition_expression:
+            custom_sql = self.get_test_case_param_value(
+                self.test_case.parameterValues,  # type: ignore
+                "sqlExpression",
+                str,
+            )
+
+            if custom_sql:
+                modified_query = self._replace_where_clause(
+                    custom_sql, partition_expression
+                )
+                if modified_query is None:
+                    return None
+                count_query = f"SELECT COUNT(*) FROM ({modified_query}) AS test_results"
+
+                try:
+                    result = self.runner.session.execute(text(count_query)).scalar()
+                    return result
+                except Exception as exc:
+                    logger.error(
+                        "Failed to execute custom SQL with partition expression. "
+                        f"Query: {count_query}\n"
+                        f"Error: {exc}\n",
+                        exc_info=True,
+                    )
+                    self.runner.session.rollback()
+                    raise exc
+            else:
+                stmt = (
+                    select(func.count())
+                    .select_from(self.runner.table)
+                    .filter(text(partition_expression))
+                )
+                return self.runner.session.execute(stmt).scalar()
+
+        self.runner = cast(QueryRunner, self.runner)
+        dialect = self.runner._session.get_bind().dialect.name
+        table_metric_computer: TableMetricComputer = TableMetricComputer(
+            dialect,
+            runner=self.runner,
+            metrics=[Metrics.ROW_COUNT],
+            conn_config=self.runtime_params.conn_config,
+            entity=self.runtime_params.entity,
+        )
+        row = table_metric_computer.compute()
+        if row:
+            return dict(row).get(Metrics.ROW_COUNT.value.name())
+        return None

metadata/data_quality/validations/table/sqlalchemy/tableDiff.py

@@ -70,6 +70,8 @@ SUPPORTED_DIALECTS = [
     Dialects.Oracle,
     Dialects.Trino,
     SapHanaScheme.hana.value,
+    Dialects.Databricks,
+    Dialects.UnityCatalog,
 ]
 
 
@@ -192,6 +194,7 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             )
             return result
         except UnsupportedDialectError as e:
+            logger.warning(f"[Data Diff]: Unsupported dialect: {e}")
             result = TestCaseResult(
                 timestamp=self.execution_date,  # type: ignore
                 testCaseStatus=TestCaseStatus.Aborted,
@@ -259,6 +262,7 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
         Returns:
             List[str]: A list of column names that have incomparable types
         """
+
         table1 = data_diff.connect_to_table(
             self.runtime_params.table1.serviceUrl,
             self.runtime_params.table1.path,

metadata/examples/workflows/airbyte.yaml

@@ -5,8 +5,14 @@ source:
     config:
       type: Airbyte
       hostPort: http://localhost:8000
-      username: username
-      password: password
+      auth:
+      # OAuth 2.0 Authentication (for Airbyte Cloud)
+        clientId: <client id>
+        clientSecret: client secret
+      # Basic Authentication (for self-hosted Airbyte)
+      # username: your_username
+      # password: your_password
+      
   sourceConfig:
     config:
       type: PipelineMetadata

metadata/examples/workflows/bigquery.yaml

@@ -5,6 +5,7 @@ source:
     config:
       type: BigQuery
       taxonomyProjectID: [ project-id-where-policy-tags-exist ]
+      # includePolicyTags: false
       credentials:
         gcpConfig:
           type: service_account

metadata/examples/workflows/confluent_cdc.yaml

@@ -0,0 +1,70 @@
+source:
+  type: kafkaconnect
+  serviceName: confluent_cdc_mysql_postgres
+  serviceConnection:
+    config:
+      type: KafkaConnect
+      hostPort: http://localhost:8083
+      # For Kafka Connect, choose one of the following authentication methods:
+
+      # Option 1: Username/Password Authentication (for self-hosted Kafka Connect)
+      KafkaConnectConfig:
+        username: admin
+        password: admin_password
+
+      # Option 2: API Key Authentication (for Confluent Cloud)
+      # KafkaConnectConfig:
+      #   username: YOUR_CONFLUENT_CLOUD_API_KEY
+      #   password: YOUR_CONFLUENT_CLOUD_API_SECRET
+
+      # Option 3: No Authentication
+      # KafkaConnectConfig: null
+
+      verifySSL: true
+      messagingServiceName: KafkaProd
+
+      # Optional: Filter CDC connectors/pipelines using regex patterns
+      # pipelineFilterPattern:
+      #   includes:
+      #     - "mysql-cdc-.*"
+      #     - "postgres-sink-.*"
+      #   excludes:
+      #     - ".*test.*"
+      #     - ".*dev.*"
+
+  sourceConfig:
+    config:
+      type: PipelineMetadata
+      lineageInformation:
+        dbServiceNames:
+          - "MysqlProd"
+          - "PostgresProd"
+        storageServiceNames: []
+
+sink:
+  type: metadata-rest
+  config: {}
+
+workflowConfig:
+  loggerLevel: INFO  # DEBUG, INFO, WARN or ERROR
+  openMetadataServerConfig:
+    hostPort: http://localhost:8585/api
+    authProvider: openmetadata
+    securityConfig:
+      jwtToken: <your-jwt-token-here>
+
+# To run this workflow:
+# metadata ingest -c /path/to/confluent_cdc.yaml
+
+# Prerequisites:
+# 1. OpenMetadata server running at http://localhost:8585
+# 2. MySQL database service already ingested into OpenMetadata
+# 3. Postgres database service already ingested into OpenMetadata
+# 4. Kafka messaging service already ingested into OpenMetadata
+# 5. Kafka Connect connectors configured and running
+#
+# The connector will create:
+# - Pipeline entities for each CDC connector
+# - Table → Topic lineage with column-level mapping
+# - Topic → Table lineage with column-level mapping
+# - Complete end-to-end lineage: MySQL columns → Kafka topic fields → Postgres columns

metadata/examples/workflows/confluent_cdc_cloud.yaml

@@ -0,0 +1,86 @@
+# Example: Confluent Cloud CDC Configuration
+# This example shows how to connect to Confluent Cloud CDC connectors
+# that replicate data from MySQL to PostgreSQL via Kafka
+
+source:
+  type: kafkaconnect
+  serviceName: confluent_cloud_cdc_production
+  serviceConnection:
+    config:
+      type: KafkaConnect
+      # Confluent Cloud Kafka Connect REST API endpoint
+      hostPort: https://pkc-xxxxx.us-east-1.aws.confluent.cloud:443
+
+      # Confluent Cloud API Key authentication
+      KafkaConnectConfig:
+        username: YOUR_CONFLUENT_CLOUD_API_KEY
+        password: YOUR_CONFLUENT_CLOUD_API_SECRET
+
+      verifySSL: true
+      messagingServiceName: KafkaCloudProd
+
+      # Filter to include only production CDC connectors
+      pipelineFilterPattern:
+        includes:
+          - "prod-mysql-cdc-.*"
+          - "prod-postgres-sink-.*"
+        excludes:
+          - ".*staging.*"
+          - ".*test.*"
+
+  sourceConfig:
+    config:
+      type: PipelineMetadata
+      lineageInformation:
+        dbServiceNames:
+          - "MysqlProd"
+          - "PostgresProd"
+        storageServiceNames: []
+
+sink:
+  type: metadata-rest
+  config: {}
+
+workflowConfig:
+  loggerLevel: INFO
+  openMetadataServerConfig:
+    hostPort: https://your-openmetadata-server.com/api
+    authProvider: openmetadata
+    securityConfig:
+      jwtToken: <your-jwt-token>
+
+# Example CDC Connector Configuration that this workflow will discover:
+#
+# MySQL CDC Source Connector:
+# {
+#   "name": "prod-mysql-cdc-customers",
+#   "config": {
+#     "connector.class": "io.debezium.connector.mysql.MySqlConnector",
+#     "database.hostname": "mysql.example.com",
+#     "database.name": "ecommerce",
+#     "table.include.list": "ecommerce.customers,ecommerce.orders",
+#     "database.server.name": "prod-mysql",
+#     "topics": "prod-mysql.ecommerce.customers"
+#   }
+# }
+#
+# Postgres Sink Connector:
+# {
+#   "name": "prod-postgres-sink-customers",
+#   "config": {
+#     "connector.class": "io.confluent.connect.jdbc.JdbcSinkConnector",
+#     "connection.url": "jdbc:postgresql://postgres.example.com:5432/analytics",
+#     "topics": "prod-mysql.ecommerce.customers",
+#     "table.name.format": "customers",
+#     "schema": "public"
+#   }
+# }
+#
+# This will create lineage:
+# MySQL: ecommerce.customers (id, name, email, created_at)
+#   ↓
+# Kafka: prod-mysql.ecommerce.customers (id, name, email, created_at) [topic schema fields]
+#   ↓
+# Postgres: analytics.public.customers (id, name, email, created_at)
+#
+# With column-level lineage showing exact field mappings at each hop!

metadata/examples/workflows/confluent_cdc_local.yaml

@@ -0,0 +1,120 @@
+# Example: Local Confluent CDC Setup for Development
+# This example shows a typical local development setup with:
+# - Local Kafka Connect cluster
+# - MySQL source database
+# - PostgreSQL target database
+# - Local Kafka broker
+
+source:
+  type: kafkaconnect
+  serviceName: local_cdc_dev
+  serviceConnection:
+    config:
+      type: KafkaConnect
+      # Local Kafka Connect REST API
+      hostPort: http://localhost:8083
+
+      # No authentication for local development
+      # KafkaConnectConfig: null
+
+      # Or use basic auth if configured
+      KafkaConnectConfig:
+        username: connect-user
+        password: connect-password
+
+      verifySSL: false  # Typically disabled for local development
+      messagingServiceName: KafkaLocal
+
+  sourceConfig:
+    config:
+      type: PipelineMetadata
+      lineageInformation:
+        dbServiceNames:
+          - "MysqlLocal"
+          - "PostgresLocal"
+        storageServiceNames: []
+
+sink:
+  type: metadata-rest
+  config: {}
+
+workflowConfig:
+  loggerLevel: DEBUG  # Use DEBUG for development to see detailed logs
+  openMetadataServerConfig:
+    hostPort: http://localhost:8585/api
+    authProvider: openmetadata
+    securityConfig:
+      jwtToken: eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg
+
+# ============================================================================
+# Local Development Setup Guide
+# ============================================================================
+#
+# 1. Start local infrastructure:
+#    docker-compose up -d  # Start MySQL, Postgres, Kafka, Kafka Connect
+#
+# 2. Ingest database metadata into OpenMetadata:
+#    # MySQL
+#    metadata ingest -c mysql_local.yaml
+#
+#    # PostgreSQL
+#    metadata ingest -c postgres_local.yaml
+#
+#    # Kafka
+#    metadata ingest -c kafka_local.yaml
+#
+# 3. Configure CDC connectors in Kafka Connect:
+#
+#    # MySQL CDC Source Connector (Debezium)
+#    curl -X POST http://localhost:8083/connectors \
+#      -H "Content-Type: application/json" \
+#      -d '{
+#        "name": "mysql-source-customers",
+#        "config": {
+#          "connector.class": "io.debezium.connector.mysql.MySqlConnector",
+#          "database.hostname": "mysql",
+#          "database.port": "3306",
+#          "database.user": "debezium",
+#          "database.password": "dbz",
+#          "database.server.id": "184054",
+#          "database.server.name": "local-mysql",
+#          "database.name": "testdb",
+#          "table.include.list": "testdb.customers,testdb.orders",
+#          "database.history.kafka.bootstrap.servers": "kafka:9092",
+#          "database.history.kafka.topic": "schema-changes.testdb"
+#        }
+#      }'
+#
+#    # PostgreSQL Sink Connector
+#    curl -X POST http://localhost:8083/connectors \
+#      -H "Content-Type: application/json" \
+#      -d '{
+#        "name": "postgres-sink-customers",
+#        "config": {
+#          "connector.class": "io.confluent.connect.jdbc.JdbcSinkConnector",
+#          "connection.url": "jdbc:postgresql://postgres:5432/warehouse",
+#          "connection.user": "postgres",
+#          "connection.password": "postgres",
+#          "topics": "local-mysql.testdb.customers",
+#          "table.name.format": "customers",
+#          "auto.create": "true",
+#          "auto.evolve": "true",
+#          "insert.mode": "upsert",
+#          "pk.mode": "record_key",
+#          "delete.enabled": "true"
+#        }
+#      }'
+#
+# 4. Run this CDC workflow:
+#    metadata ingest -c confluent_cdc_local.yaml
+#
+# 5. View lineage in OpenMetadata UI:
+#    http://localhost:8585
+#    Navigate to: Pipeline Services → local_cdc_dev → Pipelines
+#    Click on a pipeline to see column-level lineage
+#
+# Expected Lineage:
+# testdb.customers.id       → local-mysql.testdb.customers.id       → warehouse.public.customers.id
+# testdb.customers.name     → local-mysql.testdb.customers.name     → warehouse.public.customers.name
+# testdb.customers.email    → local-mysql.testdb.customers.email    → warehouse.public.customers.email
+# testdb.customers.created  → local-mysql.testdb.customers.created  → warehouse.public.customers.created

metadata/examples/workflows/databricks.yaml

@@ -5,7 +5,15 @@ source:
     config:
       catalog: hive_metastore
       databaseSchema: default
-      token: <databricks token>
+      
+      authType:
+        token: <databricks token>
+        # clientId: databricks service principal client id
+        # clientSecret: databricks service principal client secret
+        # azureClientSecret: azure client secret
+        # azureClientId: azure client id
+        # azureTenantId: azure tenant id
+
       hostPort: localhost:443
       connectionTimeout: 120
       connectionArguments: