openmetadata-ingestion 1.2.5.2__py3-none-any.whl → 1.3.0.0__py3-none-any.whl

metadata/ingestion/source/database/mssql/metadata.py

@@ -10,29 +10,57 @@
 #  limitations under the License.
 """MSSQL source module"""
 import traceback
-from typing import Iterable, Optional
+from typing import Dict, Iterable, List, Optional
 
 from sqlalchemy.dialects.mssql.base import MSDialect, ischema_names
 
+from metadata.generated.schema.api.data.createStoredProcedure import (
+    CreateStoredProcedureRequest,
+)
 from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
 from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
     MssqlConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
+from metadata.generated.schema.type.basic import EntityName
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.common_db_source import CommonDbSourceService
-from metadata.ingestion.source.database.mssql.queries import MSSQL_GET_DATABASE
+from metadata.ingestion.source.database.mssql.models import (
+    STORED_PROC_LANGUAGE_MAP,
+    MssqlStoredProcedure,
+)
+from metadata.ingestion.source.database.mssql.queries import (
+    MSSQL_GET_DATABASE,
+    MSSQL_GET_STORED_PROCEDURE_QUERIES,
+    MSSQL_GET_STORED_PROCEDURES,
+)
 from metadata.ingestion.source.database.mssql.utils import (
     get_columns,
+    get_foreign_keys,
+    get_pk_constraint,
     get_table_comment,
+    get_table_names,
+    get_unique_constraints,
     get_view_definition,
+    get_view_names,
 )
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
+from metadata.ingestion.source.database.stored_procedures_mixin import (
+    QueryByProcedure,
+    StoredProcedureMixin,
+)
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database
+from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqa_utils import update_mssql_ischema_names
 from metadata.utils.sqlalchemy_utils import (
@@ -53,9 +81,14 @@ MSDialect.get_view_definition = get_view_definition
 MSDialect.get_all_view_definitions = get_all_view_definitions
 MSDialect.get_all_table_comments = get_all_table_comments
 MSDialect.get_columns = get_columns
+MSDialect.get_pk_constraint = get_pk_constraint
+MSDialect.get_unique_constraints = get_unique_constraints
+MSDialect.get_foreign_keys = get_foreign_keys
+MSDialect.get_table_names = get_table_names
+MSDialect.get_view_names = get_view_names
 
 
-class MssqlSource(CommonDbSourceService, MultiDBSource):
+class MssqlSource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource):
     """
     Implements the necessary methods to extract
     Database metadata from MSSQL Source
@@ -81,7 +114,6 @@ class MssqlSource(CommonDbSourceService, MultiDBSource):
         yield from self._execute_database_query(MSSQL_GET_DATABASE)
 
     def get_database_names(self) -> Iterable[str]:
-
         if not self.config.serviceConnection.__root__.config.ingestAllDatabases:
             configured_db = self.config.serviceConnection.__root__.config.database
             self.set_inspector(database_name=configured_db)
@@ -112,3 +144,75 @@ class MssqlSource(CommonDbSourceService, MultiDBSource):
                     logger.error(
                         f"Error trying to connect to database {new_database}: {exc}"
                     )
+
+    def get_stored_procedures(self) -> Iterable[MssqlStoredProcedure]:
+        """List Snowflake stored procedures"""
+        if self.source_config.includeStoredProcedures:
+            results = self.engine.execute(
+                MSSQL_GET_STORED_PROCEDURES.format(
+                    database_name=self.context.database,
+                    schema_name=self.context.database_schema,
+                )
+            ).all()
+            for row in results:
+                try:
+                    stored_procedure = MssqlStoredProcedure.parse_obj(dict(row))
+                    yield stored_procedure
+                except Exception as exc:
+                    logger.error()
+                    self.status.failed(
+                        error=StackTraceError(
+                            name=dict(row).get("name", "UNKNOWN"),
+                            error=f"Error parsing Stored Procedure payload: {exc}",
+                            stackTrace=traceback.format_exc(),
+                        )
+                    )
+
+    def yield_stored_procedure(
+        self, stored_procedure: MssqlStoredProcedure
+    ) -> Iterable[Either[CreateStoredProcedureRequest]]:
+        """Prepare the stored procedure payload"""
+
+        try:
+            stored_procedure_request = CreateStoredProcedureRequest(
+                name=EntityName(__root__=stored_procedure.name),
+                description=None,
+                storedProcedureCode=StoredProcedureCode(
+                    language=STORED_PROC_LANGUAGE_MAP.get(stored_procedure.language),
+                    code=stored_procedure.definition,
+                ),
+                databaseSchema=fqn.build(
+                    metadata=self.metadata,
+                    entity_type=DatabaseSchema,
+                    service_name=self.context.database_service,
+                    database_name=self.context.database,
+                    schema_name=self.context.database_schema,
+                ),
+            )
+            yield Either(right=stored_procedure_request)
+            self.register_record_stored_proc_request(stored_procedure_request)
+
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name=stored_procedure.name,
+                    error=f"Error yielding Stored Procedure [{stored_procedure.name}] due to [{exc}]",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
+    def get_stored_procedure_queries_dict(self) -> Dict[str, List[QueryByProcedure]]:
+        """
+        Return the dictionary associating stored procedures to the
+        queries they triggered
+        """
+        start, _ = get_start_and_end(self.source_config.queryLogDuration)
+        query = MSSQL_GET_STORED_PROCEDURE_QUERIES.format(
+            start_date=start,
+        )
+
+        queries_dict = self.procedure_queries_dict(
+            query=query,
+        )
+
+        return queries_dict

metadata/ingestion/source/database/mssql/models.py

@@ -0,0 +1,30 @@
+#  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.
+"""MSSQL models"""
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from metadata.generated.schema.entity.data.storedProcedure import Language
+
+STORED_PROC_LANGUAGE_MAP = {
+    "SQL": Language.SQL,
+    "EXTERNAL": Language.External,
+}
+
+
+class MssqlStoredProcedure(BaseModel):
+    """MSSQL stored procedure list query results"""
+
+    name: str = Field(...)
+    owner: Optional[str] = Field(None)
+    language: str = Field(Language.SQL)
+    definition: str = Field(None)

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

@@ -81,8 +81,186 @@ MSSQL_TEST_GET_QUERIES = textwrap.dedent(
       FROM sys.dm_exec_cached_plans AS p
       INNER JOIN sys.dm_exec_query_stats AS s
         ON p.plan_handle = s.plan_handle
-      CROSS APPLY sys.Dm_exec_sql_text(p.plan_handle) AS t
+      CROSS APPLY sys.dm_exec_sql_text(p.plan_handle) AS t
       INNER JOIN sys.databases db
         ON db.database_id = t.dbid
 """
 )
+
+MSSQL_GET_FOREIGN_KEY = """\
+WITH fk_info AS (
+    SELECT
+        ischema_ref_con.constraint_schema,
+        ischema_ref_con.constraint_name,
+        ischema_key_col.ordinal_position,
+        ischema_key_col.[table_schema],
+        ischema_key_col.table_name,
+        ischema_ref_con.unique_constraint_schema,
+        ischema_ref_con.unique_constraint_name,
+        ischema_ref_con.match_option,
+        ischema_ref_con.update_rule,
+        ischema_ref_con.delete_rule,
+        ischema_key_col.column_name AS constrained_column
+    FROM
+        INFORMATION_SCHEMA.REFERENTIAL_CONSTRAINTS ischema_ref_con
+        INNER JOIN
+        INFORMATION_SCHEMA.KEY_COLUMN_USAGE ischema_key_col ON
+            ischema_key_col.[table_schema] = ischema_ref_con.constraint_schema
+            AND ischema_key_col.constraint_name =
+            ischema_ref_con.constraint_name
+    WHERE ischema_key_col.table_name = :tablename
+        AND ischema_key_col.[table_schema] = :owner
+),
+constraint_info AS (
+    SELECT
+        ischema_key_col.constraint_schema,
+        ischema_key_col.constraint_name,
+        ischema_key_col.ordinal_position,
+        ischema_key_col.[table_schema],
+        ischema_key_col.table_name,
+        ischema_key_col.column_name
+    FROM
+        INFORMATION_SCHEMA.KEY_COLUMN_USAGE ischema_key_col
+),
+index_info AS (
+    SELECT
+        sys.schemas.name AS index_schema,
+        sys.indexes.name AS index_name,
+        sys.index_columns.key_ordinal AS ordinal_position,
+        sys.schemas.name AS [table_schema],
+        sys.objects.name AS table_name,
+        sys.columns.name AS column_name
+    FROM
+        sys.indexes
+        INNER JOIN
+        sys.objects ON
+            sys.objects.object_id = sys.indexes.object_id
+        INNER JOIN
+        sys.schemas ON
+            sys.schemas.schema_id = sys.objects.schema_id
+        INNER JOIN
+        sys.index_columns ON
+            sys.index_columns.object_id = sys.objects.object_id
+            AND sys.index_columns.index_id = sys.indexes.index_id
+        INNER JOIN
+        sys.columns ON
+            sys.columns.object_id = sys.indexes.object_id
+            AND sys.columns.column_id = sys.index_columns.column_id
+)
+    SELECT
+        fk_info.constraint_schema,
+        fk_info.constraint_name,
+        fk_info.ordinal_position,
+        fk_info.constrained_column,
+        constraint_info.[table_schema] AS referred_table_schema,
+        constraint_info.table_name AS referred_table_name,
+        constraint_info.column_name AS referred_column,
+        fk_info.match_option,
+        fk_info.update_rule,
+        fk_info.delete_rule
+    FROM
+        fk_info INNER JOIN constraint_info ON
+            constraint_info.constraint_schema =
+                fk_info.unique_constraint_schema
+            AND constraint_info.constraint_name =
+                fk_info.unique_constraint_name
+            AND constraint_info.ordinal_position = fk_info.ordinal_position
+    UNION
+    SELECT
+        fk_info.constraint_schema,
+        fk_info.constraint_name,
+        fk_info.ordinal_position,
+        fk_info.constrained_column,
+        index_info.[table_schema] AS referred_table_schema,
+        index_info.table_name AS referred_table_name,
+        index_info.column_name AS referred_column,
+        fk_info.match_option,
+        fk_info.update_rule,
+        fk_info.delete_rule
+    FROM
+        fk_info INNER JOIN index_info ON
+            index_info.index_schema = fk_info.unique_constraint_schema
+            AND index_info.index_name = fk_info.unique_constraint_name
+            AND index_info.ordinal_position = fk_info.ordinal_position
+
+    ORDER BY fk_info.constraint_schema, fk_info.constraint_name,
+        fk_info.ordinal_position
+"""
+
+MSSQL_GET_STORED_PROCEDURES = textwrap.dedent(
+    """
+SELECT
+  ROUTINE_NAME AS name,
+  NULL AS owner,
+  ROUTINE_BODY AS language,
+  ROUTINE_DEFINITION AS definition
+FROM INFORMATION_SCHEMA.ROUTINES
+ WHERE ROUTINE_TYPE = 'PROCEDURE'
+   AND ROUTINE_CATALOG = '{database_name}'
+   AND ROUTINE_SCHEMA = '{schema_name}' 
+   AND LEFT(ROUTINE_NAME, 3) NOT IN ('sp_', 'xp_', 'ms_')
+    """
+)
+
+MSSQL_GET_STORED_PROCEDURE_QUERIES = textwrap.dedent(
+    """
+WITH SP_HISTORY (start_time, end_time, procedure_name, query_text) AS (
+  select 
+    s.last_execution_time start_time,
+    DATEADD(s, s.total_elapsed_time/1000, s.last_execution_time) end_time,
+    OBJECT_NAME(object_id, database_id) as procedure_name,
+    text as query_text
+  from sys.dm_exec_procedure_stats s
+  CROSS APPLY sys.dm_exec_sql_text(s.plan_handle)
+  WHERE OBJECT_NAME(object_id, database_id) IS NOT NULL
+    AND s.last_execution_time > '{start_date}'
+),
+Q_HISTORY (database_name, query_text, start_time, end_time, duration,query_type, schema_name, user_name) AS (
+  select    
+    db.NAME database_name,
+    t.text query_text,
+    s.last_execution_time start_time,
+    DATEADD(s, s.total_elapsed_time/1000, s.last_execution_time) end_time,
+    s.total_elapsed_time/1000 duration,
+    case
+        when t.text LIKE '%%MERGE%%' then 'MERGE'
+        when t.text LIKE '%%UPDATE%%' then 'UPDATE'
+        when t.text LIKE '%%SELECT%%INTO%%' then 'CREATE_TABLE_AS_SELECT'
+        when t.text LIKE '%%INSERT%%' then 'INSERT'
+    else 'UNKNOWN' end query_type,
+    NULL schema_name,
+    NULL user_name
+  FROM sys.dm_exec_cached_plans AS p
+  INNER JOIN sys.dm_exec_query_stats AS s
+    ON p.plan_handle = s.plan_handle
+  CROSS APPLY sys.dm_exec_sql_text(p.plan_handle) AS t
+  INNER JOIN sys.databases db
+    ON db.database_id = t.dbid
+  WHERE s.last_execution_time between '2024-01-13' and '2024-01-20'
+    AND t.text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
+    AND t.text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
+    AND p.objtype NOT IN ('Prepared', 'Proc')
+    AND s.last_execution_time > '{start_date}'
+)
+select 
+  Q.query_type AS QUERY_TYPE,
+  Q.database_name  AS QUERY_DATABASE_NAME,
+  Q.schema_name AS QUERY_SCHEMA_NAME,
+  Q.query_text AS QUERY_TEXT,
+  Q.user_name AS QUERY_USER_NAME,
+  Q.start_time AS QUERY_START_TIME,
+  Q.duration AS QUERY_DURATION,
+  SP.procedure_name AS PROCEDURE_NAME,
+  SP.query_text AS PROCEDURE_TEXT,
+  SP.start_time AS PROCEDURE_START_TIME,
+  SP.end_time AS PROCEDURE_END_TIME
+from SP_HISTORY SP
+JOIN Q_HISTORY Q
+  ON (
+    Q.start_time BETWEEN SP.start_time and SP.end_time
+    OR Q.end_time BETWEEN SP.start_time and SP.end_time
+    )
+order by PROCEDURE_START_TIME desc
+;
+    """
+)

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

@@ -12,7 +12,7 @@
 MSSQL SQLAlchemy Helper Methods
 """
 
-from sqlalchemy import Column, Integer, MetaData, String, Table, alias, sql
+from sqlalchemy import Column, Integer, MetaData, String, Table, alias, sql, text
 from sqlalchemy import types as sqltypes
 from sqlalchemy import util
 from sqlalchemy.dialects.mssql import information_schema as ischema
@@ -25,7 +25,9 @@ from sqlalchemy.dialects.mssql.base import (
     MSString,
     MSText,
     MSVarBinary,
-    _db_plus_owner,
+    _owner_plus_db,
+    _switch_db,
+    update_wrapper,
 )
 from sqlalchemy.engine import reflection
 from sqlalchemy.sql import func
@@ -34,6 +36,7 @@ from sqlalchemy.util import compat
 
 from metadata.ingestion.source.database.mssql.queries import (
     MSSQL_ALL_VIEW_DEFINITIONS,
+    MSSQL_GET_FOREIGN_KEY,
     MSSQL_GET_TABLE_COMMENTS,
 )
 from metadata.utils.logger import ingestion_logger
@@ -59,8 +62,39 @@ def get_table_comment(
     )
 
 
+def db_plus_owner_listing(fn):
+    def wrap(dialect, connection, schema=None, **kw):
+        schema = f"[{schema}]" if schema and "." in schema else schema
+        dbname, owner = _owner_plus_db(dialect, schema)
+        return _switch_db(
+            dbname, connection, fn, dialect, connection, dbname, owner, schema, **kw
+        )
+
+    return update_wrapper(wrap, fn)
+
+
+def db_plus_owner(fn):
+    def wrap(dialect, connection, tablename, schema=None, **kw):
+        schema = f"[{schema}]" if schema and "." in schema else schema
+        dbname, owner = _owner_plus_db(dialect, schema)
+        return _switch_db(
+            dbname,
+            connection,
+            fn,
+            dialect,
+            connection,
+            tablename,
+            dbname,
+            owner,
+            schema,
+            **kw,
+        )
+
+    return update_wrapper(wrap, fn)
+
+
 @reflection.cache
-@_db_plus_owner
+@db_plus_owner
 def get_columns(
     self, connection, tablename, dbname, owner, schema, **kw
 ):  # pylint: disable=unused-argument, too-many-locals, disable=too-many-branches, too-many-statements
@@ -271,7 +305,7 @@ def get_columns(
 
 
 @reflection.cache
-@_db_plus_owner
+@db_plus_owner
 def get_view_definition(
     self, connection, viewname, dbname, owner, schema, **kw
 ):  # pylint: disable=unused-argument
@@ -282,3 +316,172 @@ def get_view_definition(
         schema=owner,
         query=MSSQL_ALL_VIEW_DEFINITIONS,
     )
+
+
+@reflection.cache
+@db_plus_owner
+def get_pk_constraint(
+    self, connection, tablename, dbname, owner=None, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    This function overrides to get pk constraint
+    """
+    pkeys = []
+    tc_ = ischema.constraints
+    c_key_constaint = ischema.key_constraints.alias("C")
+
+    # Primary key constraints
+    query_ = (
+        sql.select(
+            c_key_constaint.c.column_name,
+            tc_.c.constraint_type,
+            c_key_constaint.c.constraint_name,
+        )
+        .where(
+            sql.and_(
+                tc_.c.constraint_name == c_key_constaint.c.constraint_name,
+                tc_.c.table_schema == c_key_constaint.c.table_schema,
+                c_key_constaint.c.table_name == tablename,
+                c_key_constaint.c.table_schema == owner,
+            ),
+        )
+        .order_by(tc_.c.constraint_name, c_key_constaint.c.ordinal_position)
+    )
+    cursor = connection.execution_options(future_result=True).execute(query_)
+    constraint_name = None
+    for row in cursor.mappings():
+        if "PRIMARY" in row[tc_.c.constraint_type.name]:
+            pkeys.append(row["COLUMN_NAME"])
+            if constraint_name is None:
+                constraint_name = row[c_key_constaint.c.constraint_name.name]
+    return {"constrained_columns": pkeys, "name": constraint_name}
+
+
+@reflection.cache
+def get_unique_constraints(self, connection, table_name, schema=None, **kw):
+    raise NotImplementedError()
+
+
+@reflection.cache
+@db_plus_owner
+def get_foreign_keys(
+    self, connection, tablename, dbname, owner=None, schema=None, **kw
+):  # pylint: disable=unused-argument, too-many-locals
+    """
+    This function overrides to get foreign key constraint
+    """
+    query_ = (
+        text(MSSQL_GET_FOREIGN_KEY)
+        .bindparams(
+            sql.bindparam("tablename", tablename, ischema.CoerceUnicode()),
+            sql.bindparam("owner", owner, ischema.CoerceUnicode()),
+        )
+        .columns(
+            constraint_schema=sqltypes.Unicode(),
+            constraint_name=sqltypes.Unicode(),
+            table_schema=sqltypes.Unicode(),
+            table_name=sqltypes.Unicode(),
+            constrained_column=sqltypes.Unicode(),
+            referred_table_schema=sqltypes.Unicode(),
+            referred_table_name=sqltypes.Unicode(),
+            referred_column=sqltypes.Unicode(),
+        )
+    )
+
+    # group rows by constraint ID, to handle multi-column FKs
+    fkeys = []
+
+    def fkey_rec():
+        return {
+            "name": None,
+            "constrained_columns": [],
+            "referred_schema": None,
+            "referred_table": None,
+            "referred_columns": [],
+            "options": {},
+        }
+
+    fkeys = util.defaultdict(fkey_rec)
+
+    for row_ in connection.execute(query_).fetchall():
+        (
+            _,  # constraint schema
+            rfknm,
+            _,  # ordinal position
+            scol,
+            rschema,
+            rtbl,
+            rcol,
+            # TODO: we support match=<keyword> for foreign keys so
+            # we can support this also, PG has match=FULL for example
+            # but this seems to not be a valid value for SQL Server
+            _,  # match rule
+            fkuprule,
+            fkdelrule,
+        ) = row_
+
+        rec = fkeys[rfknm]
+        rec["name"] = rfknm
+
+        if fkuprule != "NO ACTION":
+            rec["options"]["onupdate"] = fkuprule
+
+        if fkdelrule != "NO ACTION":
+            rec["options"]["ondelete"] = fkdelrule
+
+        if not rec["referred_table"]:
+            rec["referred_table"] = rtbl
+            if schema is not None or owner != rschema:
+                if dbname:
+                    rschema = dbname + "." + rschema
+                rec["referred_schema"] = rschema
+
+        local_cols, remote_cols = (
+            rec["constrained_columns"],
+            rec["referred_columns"],
+        )
+
+        local_cols.append(scol)
+        remote_cols.append(rcol)
+
+    return list(fkeys.values())
+
+
+@reflection.cache
+@db_plus_owner_listing
+def get_table_names(
+    self, connection, dbname, owner, schema, **kw
+):  # pylint: disable=unused-argument
+    tables = ischema.tables
+    query_ = (
+        sql.select(tables.c.table_name)
+        .where(
+            sql.and_(
+                tables.c.table_schema == owner,
+                tables.c.table_type == "BASE TABLE",
+            )
+        )
+        .order_by(tables.c.table_name)
+    )
+    table_names = [r[0] for r in connection.execute(query_)]
+    return table_names
+
+
+@reflection.cache
+@db_plus_owner_listing
+def get_view_names(
+    self, connection, dbname, owner, schema, **kw
+):  # pylint: disable=unused-argument
+    tables = ischema.tables
+    query_ = (
+        sql.select(tables.c.table_name)
+        .where(
+            sql.and_(
+                tables.c.table_schema == owner,
+                tables.c.table_type == "VIEW",
+            )
+        )
+        .order_by(tables.c.table_name)
+    )
+    view_names = [r[0] for r in connection.execute(query_)]
+    return view_names

metadata/ingestion/source/database/mysql/metadata.py

@@ -42,8 +42,6 @@ class MysqlSource(CommonDbSourceService):
     @classmethod
     def create(cls, config_dict, metadata: OpenMetadata):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        if config.serviceConnection is None:
-            raise InvalidSourceException("Missing service connection")
         connection = cast(MysqlConnection, config.serviceConnection.__root__.config)
         if not isinstance(connection, MysqlConnection):
             raise InvalidSourceException(