openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.6.0.0rc1__py3-none-any.whl

metadata/ingestion/source/database/snowflake/profiler/system.py

@@ -0,0 +1,436 @@
+"""Snowflake system metrics source"""
+
+import hashlib
+import re
+import traceback
+from typing import List, Optional, Tuple
+
+import sqlalchemy.orm
+from pydantic import TypeAdapter
+
+from metadata.generated.schema.entity.data.table import DmlOperationType, SystemProfile
+from metadata.ingestion.source.database.snowflake.models import (
+    SnowflakeQueryLogEntry,
+    SnowflakeQueryResult,
+)
+from metadata.profiler.metrics.system.dml_operation import DatabaseDMLOperations
+from metadata.profiler.metrics.system.system import (
+    CacheProvider,
+    EmptySystemMetricsSource,
+    SQASessionProvider,
+    SystemMetricsComputer,
+)
+from metadata.utils.collections import CaseInsensitiveString
+from metadata.utils.logger import profiler_logger
+from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
+from metadata.utils.profiler_utils import get_identifiers_from_string
+from metadata.utils.time_utils import datetime_to_timestamp
+
+PUBLIC_SCHEMA = "PUBLIC"
+logger = profiler_logger()
+RESULT_SCAN = """
+    SELECT *
+    FROM TABLE(RESULT_SCAN('{query_id}'));
+    """
+QUERY_PATTERN = r"(?:(INSERT\s*INTO\s*|INSERT\s*OVERWRITE\s*INTO\s*|UPDATE\s*|MERGE\s*INTO\s*|DELETE\s*FROM\s*))([\w._\"\'()]+)(?=[\s*\n])"  # pylint: disable=line-too-long
+IDENTIFIER_PATTERN = r"(IDENTIFIER\(\')([\w._\"]+)(\'\))"
+
+
+def sha256_hash(text: str) -> str:
+    """Return the SHA256 hash of the text"""
+
+    return hashlib.sha256(text.encode()).hexdigest()
+
+
+cache = LRUCache(LRU_CACHE_SIZE)
+
+
+@cache.wrap(key_func=lambda query: sha256_hash(query.strip()))
+def _parse_query(query: str) -> Optional[str]:
+    """Parse snowflake queries to extract the identifiers"""
+    match = re.match(QUERY_PATTERN, query, re.IGNORECASE)
+    try:
+        # This will match results like `DATABASE.SCHEMA.TABLE1` or IDENTIFIER('TABLE1')
+        # If we have `IDENTIFIER` type of queries coming from Stored Procedures, we'll need to further clean it up.
+        identifier = match.group(2)
+
+        match_internal_identifier = re.match(
+            IDENTIFIER_PATTERN, identifier, re.IGNORECASE
+        )
+        internal_identifier = (
+            match_internal_identifier.group(2) if match_internal_identifier else None
+        )
+        if internal_identifier:
+            return internal_identifier
+
+        return identifier
+    except (IndexError, AttributeError):
+        logger.debug("Could not find identifier in query. Skipping row.")
+        return None
+
+
+class SnowflakeTableResovler:
+    """A class the resolves snowflake tables by mimicking snowflake's default resolution logic:
+    https://docs.snowflake.com/en/sql-reference/name-resolution
+
+    This default specification searches in the following order:
+    - The explicitly provided schema
+    - The current schema
+    - The public schema
+
+    This can be altered by changing the SEARCH_PATH session parameter. If the users change
+    this paramter, this resolver will might return wrong values.
+
+    There is no way to extract the SEARCH_PATH from the query after it has been executed. Hence, we can
+    only rely on the default behavior and maybe allow the users to configure the search path
+    at the connection level (TODO).
+    """
+
+    def __init__(self, session: sqlalchemy.orm.Session):
+        self._cache = LRUCache[bool](LRU_CACHE_SIZE)
+        self.session = session
+
+    def show_tables(self, db, schema, table):
+        return self.session.execute(
+            f'SHOW TABLES LIKE \'{table}\' IN SCHEMA "{db}"."{schema}" LIMIT 1;'
+        ).fetchone()
+
+    def table_exists(self, db, schema, table):
+        """Return True if the table exists in Snowflake. Uses cache to store the results.
+
+        Args:
+            db (str): Database name
+            schema (str): Schema name
+            table (str): Table name
+
+        Returns:
+            bool: True if the table exists in Snowflake
+        """
+        if f"{db}.{schema}.{table}" in self._cache:
+            return self._cache.get(f"{db}.{schema}.{table}")
+        table = self.show_tables(db, schema, table)
+        if table:
+            self._cache.put(f"{db}.{schema}.{table}", True)
+            return True
+        return False
+
+    def resolve_implicit_fqn(
+        self,
+        context_database: str,
+        context_schema: Optional[str],
+        table_name: str,
+    ) -> Tuple[str, str, str]:
+        """Resolve the fully qualified name of the table from snowflake based on the following logic:
+        1. If the schema is provided:
+            a. search for the table in the schema
+            b. if not found, go to (2)
+        2. Search for the table in the public schema.
+
+        Args:
+            context_database (str): Database name
+            context_schema (Optional[str]): Schema name. If not provided, we'll search in the public schema.
+            table_name (str): Table name
+        Returns:
+            tuple: Tuple of database, schema and table names
+        Raises:
+            RuntimeError: If the table is not found in the metadata or if there are duplicate results
+            (there shouldn't be)
+
+        """
+        search_paths = []
+        if context_schema and self.table_exists(
+            context_database, context_schema, table_name
+        ):
+            search_paths += ".".join([context_database, context_schema, table_name])
+            return context_database, context_schema, table_name
+        if context_schema != PUBLIC_SCHEMA and self.table_exists(
+            context_database, PUBLIC_SCHEMA, table_name
+        ):
+            search_paths += ".".join([context_database, PUBLIC_SCHEMA, table_name])
+            return context_database, PUBLIC_SCHEMA, table_name
+        raise RuntimeError(
+            "Could not find the table {search_paths}.".format(  # pylint: disable=consider-using-f-string
+                search_paths=" OR ".join(map(lambda x: f"[{x}]", search_paths))
+            )
+        )
+
+    def resolve_snowflake_fqn(
+        self,
+        context_database: str,
+        context_schema: Optional[str],
+        identifier: str,
+    ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+        """Get query identifiers from the query text. If the schema is not provided in the query, we'll look for
+        the table under "PUBLIC" in Snowflake.
+        Database can be retrieved from the query or the query context.
+        If the schema doesnt exist in the query but does in the context, we need to check with Snowflake if table
+        exists in (1) the context schema or ib (2) the public schema in order to imitate the behavior of the query
+        engine. There are edge cases where the table was deleted (and hence not found in the metadata). In such cases,
+        the function will raise an error. It is advised to set the profier window such that there will be minimal
+        drift between the query execution and the profiler run.
+
+        Args:
+            context_database (str): Database name from the query context
+            context_schema (Optional[str]): Schema name from the query context
+            identifier (str): Identifier string extracted from a query (can be
+             'db.schema.table', 'schema.table' or just 'table')
+        Returns:
+            Tuple[Optional[str], Optional[str], Optional[str]]: Tuple of database, schema and table names
+        Raises:
+            RuntimeError: If the table name is not found in the query or if fqn resolution fails
+        """
+        (
+            database_identifier,
+            schema_identifier,
+            table_name,
+        ) = get_identifiers_from_string(identifier)
+        if not table_name:
+            raise RuntimeError("Could not extract the table name.")
+        if not context_database and not database_identifier:
+            logger.debug(
+                f"Could not resolve database name. {identifier=}, {context_database=}"
+            )
+            raise RuntimeError("Could not resolve database name.")
+        if schema_identifier is not None:
+            return (
+                database_identifier or context_database,
+                schema_identifier,
+                table_name,
+            )
+        logger.debug(
+            "Missing schema info from the query. We'll look for it in Snowflake for [%s] or [%s]",
+            (
+                ".".join(
+                    [
+                        database_identifier or context_database,
+                        context_schema,
+                        table_name,
+                    ]
+                )
+                if context_schema
+                else None
+            ),
+            ".".join(
+                [database_identifier or context_database, PUBLIC_SCHEMA, table_name]
+            ),
+        )
+        # If the schema is not explicitly provided in the query, we'll need to resolve it from OpenMetadata
+        # by cascading the search from the context to the public schema.
+        result = self.resolve_implicit_fqn(
+            context_database=context_database,
+            context_schema=context_schema,
+            table_name=table_name,
+        )
+        logger.debug("Resolved table [%s]", ".".join(result))
+        return result
+
+
+def get_snowflake_system_queries(
+    query_log_entry: SnowflakeQueryLogEntry,
+    resolver: SnowflakeTableResovler,
+) -> Optional[SnowflakeQueryResult]:
+    """
+    Run a regex lookup on the query to identify which operation ran against the table.
+
+    If the query does not have the complete set of `database.schema.table` when it runs,
+    we'll use ES to pick up the table, if we find it.
+
+    Args:
+        query_log_entry (dict): row from the snowflake system queries table
+        resolver (SnowflakeTableResolver): resolver to get the table identifiers
+    Returns:
+        QueryResult: namedtuple with the query result
+    """
+
+    try:
+        logger.debug(f"Parsing snowflake query [{query_log_entry.query_id}]")
+        identifier = _parse_query(query_log_entry.query_text)
+        if not identifier:
+            raise RuntimeError("Could not identify the table from the query.")
+
+        database_name, schema_name, table_name = resolver.resolve_snowflake_fqn(
+            identifier=identifier,
+            context_database=query_log_entry.database_name,
+            context_schema=query_log_entry.schema_name,
+        )
+
+        if not all([database_name, schema_name, table_name]):
+            raise RuntimeError(
+                f"Could not extract the identifiers from the query [{query_log_entry.query_id}]."
+            )
+
+        return SnowflakeQueryResult(
+            query_id=query_log_entry.query_id,
+            database_name=database_name.lower(),
+            schema_name=schema_name.lower(),
+            table_name=table_name.lower(),
+            query_text=query_log_entry.query_text,
+            query_type=query_log_entry.query_type,
+            start_time=query_log_entry.start_time,
+            rows_inserted=query_log_entry.rows_inserted,
+            rows_updated=query_log_entry.rows_updated,
+            rows_deleted=query_log_entry.rows_deleted,
+        )
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(
+            f"""Error while processing query with id [{query_log_entry.query_id}]: {exc}\n
+            To investigate the query run:
+            SELECT * FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY WHERE query_id = '{query_log_entry.query_id}'
+            """
+        )
+    return None
+
+
+class SnowflakeSystemMetricsSource(
+    SQASessionProvider, EmptySystemMetricsSource, CacheProvider[SnowflakeQueryLogEntry]
+):
+    """Snowflake system metrics source"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.resolver = SnowflakeTableResovler(
+            session=super().get_session(),
+        )
+
+    def get_kwargs(self, **kwargs):
+        table = kwargs.get("table")
+        return {
+            "table": table.__table__.name,
+            "database": self.get_session().get_bind().url.database,
+            "schema": table.__table__.schema,
+        }
+
+    def get_inserts(self, **kwargs) -> List[SystemProfile]:
+        database, schema, table = (
+            kwargs.get("database"),
+            kwargs.get("schema"),
+            kwargs.get("table"),
+        )
+        return self.get_system_profile(
+            database,
+            schema,
+            table,
+            list(
+                self.get_queries_by_operation(
+                    table,
+                    [
+                        DatabaseDMLOperations.INSERT,
+                        DatabaseDMLOperations.MERGE,
+                    ],
+                )
+            ),
+            "rows_inserted",
+            DmlOperationType.INSERT,
+        )
+
+    def get_updates(self, **kwargs) -> List[SystemProfile]:
+        database, schema, table = (
+            kwargs.get("database"),
+            kwargs.get("schema"),
+            kwargs.get("table"),
+        )
+        return self.get_system_profile(
+            database,
+            schema,
+            table,
+            list(
+                self.get_queries_by_operation(
+                    table,
+                    [
+                        DatabaseDMLOperations.UPDATE,
+                        DatabaseDMLOperations.MERGE,
+                    ],
+                )
+            ),
+            "rows_updated",
+            DmlOperationType.UPDATE,
+        )
+
+    def get_deletes(self, **kwargs) -> List[SystemProfile]:
+        database, schema, table = (
+            kwargs.get("database"),
+            kwargs.get("schema"),
+            kwargs.get("table"),
+        )
+        return self.get_system_profile(
+            database,
+            schema,
+            table,
+            list(
+                self.get_queries_by_operation(
+                    table,
+                    [
+                        DatabaseDMLOperations.DELETE,
+                    ],
+                )
+            ),
+            "rows_deleted",
+            DmlOperationType.DELETE,
+        )
+
+    @staticmethod
+    def get_system_profile(
+        db: str,
+        schema: str,
+        table: str,
+        query_results: List[SnowflakeQueryResult],
+        rows_affected_field: str,
+        operation: DmlOperationType,
+    ) -> List[SystemProfile]:
+        if not SnowflakeQueryResult.model_fields.get(rows_affected_field):
+            raise ValueError(
+                f"rows_affected_field [{rows_affected_field}] is not a valid field in SnowflakeQueryResult."
+            )
+        return TypeAdapter(List[SystemProfile]).validate_python(
+            [
+                {
+                    "timestamp": datetime_to_timestamp(q.start_time, milliseconds=True),
+                    "operation": operation,
+                    "rowsAffected": getattr(q, rows_affected_field),
+                }
+                for q in query_results
+                if getattr(q, rows_affected_field) > 0
+                # snowflake SQL identifiers are case insensitive. All identifiers are stored in upper case.
+                and (
+                    CaseInsensitiveString(db),
+                    CaseInsensitiveString(schema),
+                    CaseInsensitiveString(table),
+                )
+                == (
+                    q.database_name,
+                    q.schema_name,
+                    q.table_name,
+                )
+            ]
+        )
+
+    def get_queries_by_operation(
+        self, table: str, operations: List[DatabaseDMLOperations]
+    ):
+        ops = [op.value for op in operations]
+        yield from (
+            query for query in self.get_queries(table) if query.query_type in ops
+        )
+
+    def get_queries(self, table: str) -> List[SnowflakeQueryResult]:
+        queries = self.get_or_update_cache(
+            table,
+            SnowflakeQueryLogEntry.get_for_table,
+            session=super().get_session(),
+            tablename=table,
+        )
+        results = [
+            get_snowflake_system_queries(
+                query_log_entry=row,
+                resolver=self.resolver,
+            )
+            for row in queries
+        ]
+        return [result for result in results if result is not None]
+
+
+class SnowflakeSystemMetricsComputer(
+    SystemMetricsComputer, SnowflakeSystemMetricsSource
+):
+    pass

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

@@ -245,6 +245,10 @@ SNOWFLAKE_TEST_GET_TABLES = """
 SELECT TABLE_NAME FROM "{database_name}".information_schema.tables LIMIT 1
 """
 
+SNOWFLAKE_TEST_GET_VIEWS = """
+SELECT TABLE_NAME FROM "{database_name}".information_schema.views LIMIT 1
+"""
+
 SNOWFLAKE_GET_DATABASES = "SHOW DATABASES"
 
 
@@ -355,3 +359,26 @@ ORDER BY PROCEDURE_START_TIME DESC
 SNOWFLAKE_GET_TABLE_DDL = """
 SELECT GET_DDL('TABLE','{table_name}') AS \"text\"
 """
+SNOWFLAKE_QUERY_LOG_QUERY = """
+    SELECT
+        QUERY_ID,
+        QUERY_TEXT,
+        QUERY_TYPE,
+        START_TIME,
+        DATABASE_NAME,
+        SCHEMA_NAME,
+        ROWS_INSERTED,
+        ROWS_UPDATED,
+        ROWS_DELETED
+    FROM "SNOWFLAKE"."ACCOUNT_USAGE"."QUERY_HISTORY"
+    WHERE
+    start_time>= DATEADD('DAY', -1, CURRENT_TIMESTAMP)
+    AND QUERY_TEXT ILIKE '%{tablename}%'
+    AND QUERY_TYPE IN (
+        '{insert}',
+        '{update}',
+        '{delete}',
+        '{merge}'
+    )
+    AND EXECUTION_STATUS = 'SUCCESS';
+"""

metadata/ingestion/source/database/snowflake/service_spec.py

@@ -0,0 +1,14 @@
+from metadata.ingestion.source.database.snowflake.lineage import SnowflakeLineageSource
+from metadata.ingestion.source.database.snowflake.metadata import SnowflakeSource
+from metadata.ingestion.source.database.snowflake.profiler.profiler import (
+    SnowflakeProfiler,
+)
+from metadata.ingestion.source.database.snowflake.usage import SnowflakeUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=SnowflakeSource,
+    lineage_source_class=SnowflakeLineageSource,
+    usage_source_class=SnowflakeUsageSource,
+    profiler_class=SnowflakeProfiler,
+)

metadata/ingestion/source/database/sql_column_handler.py

@@ -196,6 +196,26 @@ class SqlColumnHandlerMixin:
             ]
         return Column(**parsed_string)
 
+    @calculate_execution_time()
+    def process_json_type_column_fields(  # pylint: disable=too-many-locals
+        self, schema_name: str, table_name: str, column_name: str, inspector: Inspector
+    ) -> Optional[List[Column]]:
+        """
+        Parse fields column with json data types
+        """
+        try:
+            if hasattr(inspector, "get_json_fields_and_type"):
+                result = inspector.get_json_fields_and_type(
+                    table_name, column_name, schema_name
+                )
+                return result
+
+        except NotImplementedError:
+            logger.debug(
+                "Cannot parse json fields for table column [{schema_name}.{table_name}.{col_name}]: NotImplementedError"
+            )
+        return None
+
     @calculate_execution_time()
     def get_columns_and_constraints(  # pylint: disable=too-many-locals
         self, schema_name: str, table_name: str, db_name: str, inspector: Inspector
@@ -271,6 +291,12 @@ class SqlColumnHandlerMixin:
                         precision,
                     )
                     col_data_length = 1 if col_data_length is None else col_data_length
+
+                    if col_type == "JSON":
+                        children = self.process_json_type_column_fields(
+                            schema_name, table_name, column.get("name"), inspector
+                        )
+
                     om_column = Column(
                         name=ColumnName(
                             root=column["name"]

metadata/ingestion/source/database/sqlite/lineage.py

@@ -0,0 +1,46 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
+"""
+Sqlite lineage module
+"""
+from typing import Optional
+
+from metadata.generated.schema.entity.services.connections.database.sqliteConnection import (
+    SQLiteConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.lineage_source import LineageSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class SqliteLineageSource(LineageSource):
+    """
+    Sqlite lineage source implements view lineage
+    """
+
+    @classmethod
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
+        """Create class instance"""
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: SQLiteConnection = config.serviceConnection.root.config
+        if not isinstance(connection, SQLiteConnection):
+            raise InvalidSourceException(
+                f"Expected SQLiteConnection, but got {connection}"
+            )
+        return cls(config, metadata)

metadata/ingestion/source/database/sqlite/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.database.sqlite.metadata import SqliteSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=SqliteSource)