acryl-datahub 1.1.1rc3__py3-none-any.whl → 1.2.0__py3-none-any.whl

datahub/ingestion/source/sql/clickhouse.py

@@ -379,7 +379,9 @@ clickhouse_datetime_format = "%Y-%m-%d %H:%M:%S"
 @platform_name("ClickHouse")
 @config_class(ClickHouseConfig)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 class ClickHouseSource(TwoTierSQLAlchemySource):
     """

datahub/ingestion/source/sql/cockroachdb.py

@@ -26,7 +26,6 @@ class CockroachDBConfig(PostgresConfig):
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 class CockroachDBSource(PostgresSource):
     config: CockroachDBConfig
 

datahub/ingestion/source/sql/hana.py

@@ -27,7 +27,9 @@ class HanaConfig(BasicSQLAlchemyConfig):
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 class HanaSource(SQLAlchemySource):
     def __init__(self, config: HanaConfig, ctx: PipelineContext):
         super().__init__(config, ctx, "hana")

datahub/ingestion/source/sql/hive_metastore.py

@@ -52,7 +52,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import Dataset
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaField
 from datahub.metadata.schema_classes import (
-    ChangeTypeClass,
     DatasetPropertiesClass,
     SubTypesClass,
     ViewPropertiesClass,
@@ -161,7 +160,9 @@ class HiveMetastore(BasicSQLAlchemyConfig):
 @platform_name("Hive Metastore")
 @config_class(HiveMetastore)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 @capability(SourceCapability.DATA_PROFILING, "Not Supported", False)
 @capability(SourceCapability.CLASSIFICATION, "Not Supported", False)
 @capability(
@@ -599,10 +600,7 @@ class HiveMetastoreSource(SQLAlchemySource):
                 yield dpi_aspect
 
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="subTypes",
                 aspect=SubTypesClass(typeNames=[self.table_subtype]),
             ).as_workunit()
 
@@ -808,10 +806,7 @@ class HiveMetastoreSource(SQLAlchemySource):
 
             # Add views subtype
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="subTypes",
                 aspect=SubTypesClass(typeNames=[self.view_subtype]),
             ).as_workunit()
 
@@ -822,10 +817,7 @@ class HiveMetastoreSource(SQLAlchemySource):
                 viewLogic=dataset.view_definition if dataset.view_definition else "",
             )
             yield MetadataChangeProposalWrapper(
-                entityType="dataset",
-                changeType=ChangeTypeClass.UPSERT,
                 entityUrn=dataset_urn,
-                aspectName="viewProperties",
                 aspect=view_properties_aspect,
             ).as_workunit()
 

datahub/ingestion/source/sql/mariadb.py

@@ -15,7 +15,6 @@ from datahub.ingestion.source.sql.mysql import MySQLConfig, MySQLSource
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 class MariaDBSource(MySQLSource):
     def get_platform(self):
         return "mariadb"

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

@@ -27,6 +27,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,
@@ -174,7 +175,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,9 +339,11 @@ 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(
@@ -343,12 +361,158 @@ class SQLServerSource(SQLAlchemySource):
             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}",
+                self.report.failure(
+                    message="Failed to list stored procedures",
+                    title="SQL Server Stored Procedures Extraction",
+                    context="Error occurred during schema-level stored procedure extraction",
+                    exc=e,
                 )
 
+    def _detect_rds_environment(self, conn: Connection) -> bool:
+        """
+        Detect if we're running in an RDS/managed environment vs on-premises.
+        Returns True if RDS/managed, False if on-premises.
+        """
+        try:
+            # Try to access system tables directly - this typically fails in RDS
+            conn.execute("SELECT TOP 1 * FROM msdb.dbo.sysjobs")
+            logger.debug(
+                "Direct table access successful - likely on-premises environment"
+            )
+            return False
+        except Exception:
+            logger.debug("Direct table access failed - likely RDS/managed environment")
+            return True
+
     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 = {}
+
+        for row in jobs_result:
+            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 = {}
+                for step_row in steps_result:
+                    # 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 +535,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 +548,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 +564,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]
@@ -740,25 +945,25 @@ 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 != "":
+            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

datahub/ingestion/source/sql/mysql.py

@@ -65,7 +65,6 @@ class MySQLConfig(MySQLConnectionConfig, TwoTierSQLAlchemyConfig):
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 class MySQLSource(TwoTierSQLAlchemySource):
     """
     This plugin extracts the following:

datahub/ingestion/source/sql/oracle.py

@@ -441,7 +441,7 @@ class OracleInspectorObjectWrapper:
             "\nac.constraint_name,"
             "\nac.constraint_type,"
             "\nacc.column_name AS local_column,"
-            "\nac.r_table_name AS remote_table,"
+            "\nac.table_name AS remote_table,"
             "\nrcc.column_name AS remote_column,"
             "\nac.r_owner AS remote_owner,"
             "\nacc.position AS loc_pos,"

datahub/ingestion/source/sql/postgres.py

@@ -131,7 +131,6 @@ class PostgresConfig(BasePostgresConfig):
 @capability(SourceCapability.DOMAINS, "Enabled by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
-@capability(SourceCapability.LINEAGE_COARSE, "Optionally enabled via configuration")
 class PostgresSource(SQLAlchemySource):
     """
     This plugin extracts the following: