acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/sql/mssql/source.py

@@ -1,7 +1,7 @@
 import logging
 import re
 import urllib.parse
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 import pydantic
 import sqlalchemy.dialects.mssql
@@ -10,9 +10,10 @@ from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine.base import Connection
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.exc import ProgrammingError, ResourceClosedError
+from sqlalchemy.sql import quoted_name
 
 import datahub.metadata.schema_classes as models
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.pattern_utils import UUID_REGEX
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
@@ -27,6 +28,7 @@ from datahub.ingestion.api.decorators import (
 from datahub.ingestion.api.source import StructuredLogLevel
 from datahub.ingestion.api.source_helpers import auto_workunit
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.sql.mssql.job_models import (
     JobStep,
     MSSQLDataFlow,
@@ -40,7 +42,6 @@ from datahub.ingestion.source.sql.mssql.job_models import (
 )
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
-    SqlWorkUnit,
     register_custom_type,
 )
 from datahub.ingestion.source.sql.sql_config import (
@@ -74,7 +75,7 @@ DEFAULT_TEMP_TABLES_PATTERNS = [
 class SQLServerConfig(BasicSQLAlchemyConfig):
     # defaults
     host_port: str = Field(default="localhost:1433", description="MSSQL host URL.")
-    scheme: str = Field(default="mssql+pytds", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="mssql+pytds")
 
     # TODO: rename to include_procedures ?
     include_stored_procedures: bool = Field(
@@ -130,10 +131,18 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
         "match the entire table name in database.schema.table format. Defaults are to set in such a way "
         "to ignore the temporary staging tables created by known ETL tools.",
     )
+    quote_schemas: bool = Field(
+        default=False,
+        description="Represent a schema identifiers combined with quoting preferences. See [sqlalchemy quoted_name docs](https://docs.sqlalchemy.org/en/20/core/sqlelement.html#sqlalchemy.sql.expression.quoted_name).",
+    )
+    is_aws_rds: Optional[bool] = Field(
+        default=None,
+        description="Indicates if the SQL Server instance is running on AWS RDS. When None (default), automatic detection will be attempted using server name analysis.",
+    )
 
     @pydantic.validator("uri_args")
     def passwords_match(cls, v, values, **kwargs):
-        if values["use_odbc"] and "driver" not in v:
+        if values["use_odbc"] and not values["sqlalchemy_uri"] and "driver" not in v:
             raise ValueError("uri_args must contain a 'driver' option")
         elif not values["use_odbc"] and v:
             raise ValueError("uri_args is not supported when ODBC is disabled")
@@ -144,22 +153,36 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
         uri_opts: Optional[Dict[str, Any]] = None,
         current_db: Optional[str] = None,
     ) -> str:
+        current_db = current_db or self.database
+
         if self.use_odbc:
             # Ensure that the import is available.
             import pyodbc  # noqa: F401
 
             self.scheme = "mssql+pyodbc"
 
+            # ODBC requires a database name, otherwise it will interpret host_port
+            # as a pre-defined ODBC connection name.
+            current_db = current_db or "master"
+
         uri: str = self.sqlalchemy_uri or make_sqlalchemy_uri(
             self.scheme,  # type: ignore
             self.username,
             self.password.get_secret_value() if self.password else None,
             self.host_port,  # type: ignore
-            current_db if current_db else self.database,
+            current_db,
             uri_opts=uri_opts,
         )
         if self.use_odbc:
-            uri = f"{uri}?{urllib.parse.urlencode(self.uri_args)}"
+            final_uri_args = self.uri_args.copy()
+            if final_uri_args and current_db:
+                final_uri_args.update({"database": current_db})
+
+            uri = (
+                f"{uri}?{urllib.parse.urlencode(final_uri_args)}"
+                if final_uri_args
+                else uri
+            )
         return uri
 
     @property
@@ -174,7 +197,22 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Enabled by default to get lineage for stored procedures via `include_lineage` and for views via `include_view_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.STORED_PROCEDURE,
+        SourceCapabilityModifier.VIEW,
+    ],
+)
+@capability(
+    SourceCapability.LINEAGE_FINE,
+    "Enabled by default to get lineage for stored procedures via `include_lineage` and for views via `include_view_column_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.STORED_PROCEDURE,
+        SourceCapabilityModifier.VIEW,
+    ],
+)
 class SQLServerSource(SQLAlchemySource):
     """
     This plugin extracts the following:
@@ -323,32 +361,186 @@ class SQLServerSource(SQLAlchemySource):
             try:
                 yield from self.loop_jobs(inspector, self.config)
             except Exception as e:
-                self.report.report_failure(
-                    "jobs",
-                    f"Failed to list jobs due to error {e}",
+                self.report.failure(
+                    message="Failed to list jobs",
+                    title="SQL Server Jobs Extraction",
+                    context="Error occurred during database-level job extraction",
+                    exc=e,
                 )
 
-    def get_schema_level_workunits(
-        self,
-        inspector: Inspector,
-        schema: str,
-        database: str,
-    ) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
-        yield from super().get_schema_level_workunits(
-            inspector=inspector,
-            schema=schema,
-            database=database,
-        )
-        if self.config.include_stored_procedures:
-            try:
-                yield from self.loop_stored_procedures(inspector, schema, self.config)
-            except Exception as e:
-                self.report.report_failure(
-                    "jobs",
-                    f"Failed to list jobs due to error {e}",
+    def _detect_rds_environment(self, conn: Connection) -> bool:
+        """
+        Detect if we're running in an RDS/managed environment vs on-premises.
+        Uses explicit configuration if provided, otherwise attempts automatic detection.
+        Returns True if RDS/managed, False if on-premises.
+        """
+        if self.config.is_aws_rds is not None:
+            logger.info(
+                f"Using explicit is_aws_rds configuration: {self.config.is_aws_rds}"
+            )
+            return self.config.is_aws_rds
+
+        try:
+            result = conn.execute("SELECT @@servername AS server_name")
+            server_name_row = result.fetchone()
+            if server_name_row:
+                server_name = server_name_row["server_name"].lower()
+
+                aws_indicators = ["amazon", "amzn", "amaz", "ec2", "rds.amazonaws.com"]
+                is_rds = any(indicator in server_name for indicator in aws_indicators)
+                if is_rds:
+                    logger.info(f"AWS RDS detected based on server name: {server_name}")
+                else:
+                    logger.info(
+                        f"Non-RDS environment detected based on server name: {server_name}"
+                    )
+
+                return is_rds
+            else:
+                logger.warning(
+                    "Could not retrieve server name, assuming non-RDS environment"
                 )
+                return False
+
+        except Exception as e:
+            logger.warning(
+                f"Failed to detect RDS/managed vs on-prem env, assuming non-RDS environment ({e})"
+            )
+            return False
 
     def _get_jobs(self, conn: Connection, db_name: str) -> Dict[str, Dict[str, Any]]:
+        """
+        Get job information with environment detection to choose optimal method first.
+        """
+        jobs: Dict[str, Dict[str, Any]] = {}
+
+        # Detect environment to choose optimal method first
+        is_rds = self._detect_rds_environment(conn)
+
+        if is_rds:
+            # Managed environment - try stored procedures first
+            try:
+                jobs = self._get_jobs_via_stored_procedures(conn, db_name)
+                logger.info(
+                    "Successfully retrieved jobs using stored procedures (managed environment)"
+                )
+                return jobs
+            except Exception as sp_error:
+                logger.warning(
+                    f"Failed to retrieve jobs via stored procedures in managed environment: {sp_error}"
+                )
+                # Try direct query as fallback (might work in some managed environments)
+                try:
+                    jobs = self._get_jobs_via_direct_query(conn, db_name)
+                    logger.info(
+                        "Successfully retrieved jobs using direct query fallback in managed environment"
+                    )
+                    return jobs
+                except Exception as direct_error:
+                    self.report.failure(
+                        message="Failed to retrieve jobs in managed environment",
+                        title="SQL Server Jobs Extraction",
+                        context="Both stored procedures and direct query methods failed",
+                        exc=direct_error,
+                    )
+        else:
+            # On-premises environment - try direct query first (usually faster)
+            try:
+                jobs = self._get_jobs_via_direct_query(conn, db_name)
+                logger.info(
+                    "Successfully retrieved jobs using direct query (on-premises environment)"
+                )
+                return jobs
+            except Exception as direct_error:
+                logger.warning(
+                    f"Failed to retrieve jobs via direct query in on-premises environment: {direct_error}"
+                )
+                # Try stored procedures as fallback
+                try:
+                    jobs = self._get_jobs_via_stored_procedures(conn, db_name)
+                    logger.info(
+                        "Successfully retrieved jobs using stored procedures fallback in on-premises environment"
+                    )
+                    return jobs
+                except Exception as sp_error:
+                    self.report.failure(
+                        message="Failed to retrieve jobs in on-premises environment",
+                        title="SQL Server Jobs Extraction",
+                        context="Both direct query and stored procedures methods failed",
+                        exc=sp_error,
+                    )
+
+        return jobs
+
+    def _get_jobs_via_stored_procedures(
+        self, conn: Connection, db_name: str
+    ) -> Dict[str, Dict[str, Any]]:
+        jobs: Dict[str, Dict[str, Any]] = {}
+
+        # First, get all jobs
+        jobs_result = conn.execute("EXEC msdb.dbo.sp_help_job")
+        jobs_data = {}
+
+        # SQLAlchemy 1.3 support was dropped in Sept 2023 (PR #8810)
+        # SQLAlchemy 1.4+ returns LegacyRow objects that don't support dictionary-style .get() method
+        # Use .mappings() to get MappingResult with dictionary-like rows that support .get()
+        for row in jobs_result.mappings():
+            job_id = str(row["job_id"])
+            jobs_data[job_id] = {
+                "job_id": job_id,
+                "name": row["name"],
+                "description": row.get("description", ""),
+                "date_created": row.get("date_created"),
+                "date_modified": row.get("date_modified"),
+                "enabled": row.get("enabled", 1),
+            }
+
+        # Now get job steps for each job, filtering by database
+        for job_id, job_info in jobs_data.items():
+            try:
+                # Get steps for this specific job
+                steps_result = conn.execute(
+                    f"EXEC msdb.dbo.sp_help_jobstep @job_id = '{job_id}'"
+                )
+
+                job_steps = {}
+                # Use .mappings() for dictionary-like access (SQLAlchemy 1.4+ compatibility)
+                for step_row in steps_result.mappings():
+                    # Only include steps that run against our target database
+                    step_database = step_row.get("database_name", "")
+                    if step_database.lower() == db_name.lower() or not step_database:
+                        step_data = {
+                            "job_id": job_id,
+                            "job_name": job_info["name"],
+                            "description": job_info["description"],
+                            "date_created": job_info["date_created"],
+                            "date_modified": job_info["date_modified"],
+                            "step_id": step_row["step_id"],
+                            "step_name": step_row["step_name"],
+                            "subsystem": step_row.get("subsystem", ""),
+                            "command": step_row.get("command", ""),
+                            "database_name": step_database,
+                        }
+                        job_steps[step_row["step_id"]] = step_data
+
+                # Only add job if it has relevant steps
+                if job_steps:
+                    jobs[job_info["name"]] = job_steps
+
+            except Exception as step_error:
+                logger.warning(
+                    f"Failed to get steps for job {job_info['name']}: {step_error}"
+                )
+                continue
+
+        return jobs
+
+    def _get_jobs_via_direct_query(
+        self, conn: Connection, db_name: str
+    ) -> Dict[str, Dict[str, Any]]:
+        """
+        Original method using direct table access for on-premises SQL Server.
+        """
         jobs_data = conn.execute(
             f"""
             SELECT
@@ -371,6 +563,7 @@ class SQLServerSource(SQLAlchemySource):
             where database_name = '{db_name}'
             """
         )
+
         jobs: Dict[str, Dict[str, Any]] = {}
         for row in jobs_data:
             step_data = dict(
@@ -383,11 +576,13 @@ class SQLServerSource(SQLAlchemySource):
                 step_name=row["step_name"],
                 subsystem=row["subsystem"],
                 command=row["command"],
+                database_name=row["database_name"],
             )
             if row["name"] in jobs:
                 jobs[row["name"]][row["step_id"]] = step_data
             else:
                 jobs[row["name"]] = {row["step_id"]: step_data}
+
         return jobs
 
     def loop_jobs(
@@ -397,21 +592,59 @@ class SQLServerSource(SQLAlchemySource):
     ) -> Iterable[MetadataWorkUnit]:
         """
         Loop MS SQL jobs as dataFlow-s.
-        :return:
+        Now supports both managed and on-premises SQL Server.
         """
         db_name = self.get_db_name(inspector)
-        with inspector.engine.connect() as conn:
-            jobs = self._get_jobs(conn, db_name)
-            for job_name, job_steps in jobs.items():
-                job = MSSQLJob(
-                    name=job_name,
-                    env=sql_config.env,
-                    db=db_name,
-                    platform_instance=sql_config.platform_instance,
+
+        try:
+            with inspector.engine.connect() as conn:
+                jobs = self._get_jobs(conn, db_name)
+
+                if not jobs:
+                    logger.info(f"No jobs found for database: {db_name}")
+                    return
+
+                logger.info(f"Found {len(jobs)} jobs for database: {db_name}")
+
+                for job_name, job_steps in jobs.items():
+                    try:
+                        job = MSSQLJob(
+                            name=job_name,
+                            env=sql_config.env,
+                            db=db_name,
+                            platform_instance=sql_config.platform_instance,
+                        )
+                        data_flow = MSSQLDataFlow(entity=job)
+                        yield from self.construct_flow_workunits(data_flow=data_flow)
+                        yield from self.loop_job_steps(job, job_steps)
+
+                    except Exception as job_error:
+                        logger.warning(f"Failed to process job {job_name}: {job_error}")
+                        self.report.warning(
+                            message=f"Failed to process job {job_name}",
+                            title="SQL Server Jobs Extraction",
+                            context="Error occurred while processing individual job",
+                            exc=job_error,
+                        )
+                        continue
+
+        except Exception as e:
+            error_message = f"Failed to retrieve jobs for database {db_name}: {e}"
+            logger.error(error_message)
+
+            # Provide specific guidance for permission issues
+            if "permission" in str(e).lower() or "denied" in str(e).lower():
+                permission_guidance = (
+                    "For managed SQL Server services, ensure the following permissions are granted:\n"
+                    "GRANT EXECUTE ON msdb.dbo.sp_help_job TO datahub_read;\n"
+                    "GRANT EXECUTE ON msdb.dbo.sp_help_jobstep TO datahub_read;\n"
+                    "For on-premises SQL Server, you may also need:\n"
+                    "GRANT SELECT ON msdb.dbo.sysjobs TO datahub_read;\n"
+                    "GRANT SELECT ON msdb.dbo.sysjobsteps TO datahub_read;"
                 )
-                data_flow = MSSQLDataFlow(entity=job)
-                yield from self.construct_flow_workunits(data_flow=data_flow)
-                yield from self.loop_job_steps(job, job_steps)
+                logger.info(permission_guidance)
+
+            raise e
 
     def loop_job_steps(
         self, job: MSSQLJob, job_steps: Dict[str, Any]
@@ -431,7 +664,7 @@ class SQLServerSource(SQLAlchemySource):
         self,
         inspector: Inspector,
         schema: str,
-        sql_config: SQLServerConfig,
+        sql_config: SQLServerConfig,  # type: ignore
     ) -> Iterable[MetadataWorkUnit]:
         """
         Loop schema data for get stored procedures as dataJob-s.
@@ -740,25 +973,29 @@ class SQLServerSource(SQLAlchemySource):
         url = self.config.get_sql_alchemy_url()
         logger.debug(f"sql_alchemy_url={url}")
         engine = create_engine(url, **self.config.options)
-        with engine.connect() as conn:
-            if self.config.database and self.config.database != "":
-                inspector = inspect(conn)
-                yield inspector
-            else:
+
+        if (
+            self.config.database
+            and self.config.database != ""
+            or (self.config.sqlalchemy_uri and self.config.sqlalchemy_uri != "")
+        ):
+            inspector = inspect(engine)
+            yield inspector
+        else:
+            with engine.begin() as conn:
                 databases = conn.execute(
                     "SELECT name FROM master.sys.databases WHERE name NOT IN \
                   ('master', 'model', 'msdb', 'tempdb', 'Resource', \
                        'distribution' , 'reportserver', 'reportservertempdb'); "
-                )
-                for db in databases:
-                    if self.config.database_pattern.allowed(db["name"]):
-                        url = self.config.get_sql_alchemy_url(current_db=db["name"])
-                        with create_engine(
-                            url, **self.config.options
-                        ).connect() as conn:
-                            inspector = inspect(conn)
-                            self.current_database = db["name"]
-                            yield inspector
+                ).fetchall()
+
+            for db in databases:
+                if self.config.database_pattern.allowed(db["name"]):
+                    url = self.config.get_sql_alchemy_url(current_db=db["name"])
+                    engine = create_engine(url, **self.config.options)
+                    inspector = inspect(engine)
+                    self.current_database = db["name"]
+                    yield inspector
 
     def get_identifier(
         self, *, schema: str, entity: str, inspector: Inspector, **kwargs: Any
@@ -838,3 +1075,45 @@ class SQLServerSource(SQLAlchemySource):
             if self.config.convert_urns_to_lowercase
             else table_ref_str
         )
+
+    def get_allowed_schemas(self, inspector: Inspector, db_name: str) -> Iterable[str]:
+        for schema in super().get_allowed_schemas(inspector, db_name):
+            if self.config.quote_schemas:
+                yield quoted_name(schema, True)
+            else:
+                yield schema
+
+    def get_db_name(self, inspector: Inspector) -> str:
+        engine = inspector.engine
+
+        try:
+            if (
+                engine
+                and hasattr(engine, "url")
+                and hasattr(engine.url, "database")
+                and engine.url.database
+            ):
+                return str(engine.url.database).strip('"')
+
+            if (
+                engine
+                and hasattr(engine, "url")
+                and hasattr(engine.url, "query")
+                and "odbc_connect" in engine.url.query
+            ):
+                # According to the ODBC connection keywords: https://learn.microsoft.com/en-us/sql/connect/odbc/dsn-connection-string-attribute?view=sql-server-ver17#supported-dsnconnection-string-keywords-and-connection-attributes
+                database = re.search(
+                    r"DATABASE=([^;]*);",
+                    urllib.parse.unquote_plus(str(engine.url.query["odbc_connect"])),
+                    flags=re.IGNORECASE,
+                )
+
+                if database and database.group(1):
+                    return database.group(1)
+
+            return ""
+
+        except Exception as e:
+            raise RuntimeError(
+                "Unable to get database name from Sqlalchemy inspector"
+            ) from e