diff --git a/ingestion/src/metadata/ingestion/source/database/redshift/connection.py b/ingestion/src/metadata/ingestion/source/database/redshift/connection.py index 16c845459918..828683cd69e9 100644 --- a/ingestion/src/metadata/ingestion/source/database/redshift/connection.py +++ b/ingestion/src/metadata/ingestion/source/database/redshift/connection.py @@ -16,6 +16,7 @@ from typing import Optional from sqlalchemy.engine import Engine +from sqlalchemy.exc import ProgrammingError from sqlalchemy.sql import text from metadata.generated.schema.entity.automations.workflow import ( @@ -41,13 +42,17 @@ ) from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.ingestion.source.connections_utils import kill_active_connections +from metadata.ingestion.source.database.redshift.models import RedshiftInstanceType from metadata.ingestion.source.database.redshift.queries import ( REDSHIFT_GET_ALL_RELATIONS, REDSHIFT_GET_DATABASE_NAMES, - REDSHIFT_TEST_GET_QUERIES, + REDSHIFT_TEST_GET_QUERIES_MAP, REDSHIFT_TEST_PARTITION_DETAILS, ) from metadata.utils.constants import THREE_MIN +from metadata.utils.logger import ingestion_logger + +logger = ingestion_logger() def get_connection(connection: RedshiftConnection) -> Engine: @@ -61,6 +66,42 @@ def get_connection(connection: RedshiftConnection) -> Engine: ) +def get_redshift_instance_type(engine: Engine) -> RedshiftInstanceType: + """ + Detect whether the connected Amazon Redshift deployment is Provisioned + or Serverless by probing for STL system table availability. + + Serverless deployments do not have access to STL_* system tables due to + their architecture. Use SYS_* views instead for Serverless compatibility. + + Reference: https://docs.aws.amazon.com/redshift/latest/dg/cm_chap_system-tables.html#sys_view_migration-use_cases + + Args: + engine (Engine): SQLAlchemy engine connected to a Redshift endpoint. + + Returns: + RedshiftInstanceType: PROVISIONED if STL tables are accessible, + SERVERLESS otherwise. + """ + probe_query = text("SELECT 1 FROM pg_catalog.stl_query LIMIT 1") + + try: + with engine.connect() as conn: + conn.execute(probe_query) + + logger.info( + "Redshift instance type detected: PROVISIONED (STL tables accessible)" + ) + return RedshiftInstanceType.PROVISIONED + + except ProgrammingError: + logger.info( + "Redshift instance type detected: SERVERLESS " + "(STL tables not accessible, will use SYS_* views)" + ) + return RedshiftInstanceType.SERVERLESS + + def test_connection( metadata: OpenMetadata, engine: Engine, @@ -80,12 +121,27 @@ def test_connection( def test_get_queries_permissions(engine_: Engine): """Check if we have the right permissions to list queries""" + redshift_instance_type = get_redshift_instance_type(engine_) + with engine_.connect() as conn: - res = conn.execute(REDSHIFT_TEST_GET_QUERIES).fetchone() - if not all(res): - raise SourceConnectionException( - f"We don't have the right permissions to list queries - {res}" - ) + if redshift_instance_type == RedshiftInstanceType.PROVISIONED: + res = conn.execute( + REDSHIFT_TEST_GET_QUERIES_MAP[RedshiftInstanceType.PROVISIONED] + ).fetchone() + if not all(res): + raise SourceConnectionException( + "We don't have the right permissions to list queries from stl views (Redshift Provisioned)" + f" - {res}" + ) + else: + res = conn.execute( + REDSHIFT_TEST_GET_QUERIES_MAP[RedshiftInstanceType.SERVERLESS] + ).fetchone() + if not all(res): + raise SourceConnectionException( + "We don't have the right permissions to list queries from sys views (Redshift Serverless)" + f" - {res}" + ) test_fn = { "CheckAccess": partial(test_connection_engine_step, engine), diff --git a/ingestion/src/metadata/ingestion/source/database/redshift/lineage.py b/ingestion/src/metadata/ingestion/source/database/redshift/lineage.py index 9a3acddb24b1..cc9a6803b55b 100644 --- a/ingestion/src/metadata/ingestion/source/database/redshift/lineage.py +++ b/ingestion/src/metadata/ingestion/source/database/redshift/lineage.py @@ -35,12 +35,18 @@ from metadata.generated.schema.type.tableQuery import TableQuery from metadata.ingestion.source.database.lineage_source import LineageSource +from metadata.ingestion.source.database.redshift.connection import ( + get_redshift_instance_type, +) +from metadata.ingestion.source.database.redshift.models import RedshiftInstanceType from metadata.ingestion.source.database.redshift.queries import ( - REDSHIFT_GET_STORED_PROCEDURE_QUERIES, - REDSHIFT_SQL_STATEMENT, + REDSHIFT_GET_STORED_PROCEDURE_QUERIES_MAP, + REDSHIFT_SQL_STATEMENT_MAP, ) from metadata.ingestion.source.database.redshift.query_parser import ( + OpenMetadata, RedshiftQueryParserSource, + WorkflowSource, ) from metadata.ingestion.source.database.stored_procedures_mixin import ( StoredProcedureLineageMixin, @@ -54,7 +60,7 @@ class RedshiftLineageSource( RedshiftQueryParserSource, StoredProcedureLineageMixin, LineageSource ): - filters = """ + provisioned_filters = """ AND ( querytxt ILIKE '%%create%%table%%as%%select%%' OR querytxt ILIKE '%%insert%%into%%select%%' @@ -63,7 +69,30 @@ class RedshiftLineageSource( ) """ - sql_stmt = REDSHIFT_SQL_STATEMENT + serverless_filters = """ + AND ( + (query_text ILIKE '%%create%%table%%as%%select%%' AND query_type = 'CTAS') + OR (query_text ILIKE '%%insert%%into%%select%%' AND query_type = 'INSERT') + OR (query_text ILIKE '%%update%%' AND query_type = 'UPDATE') + OR (query_text ILIKE '%%merge%%' AND query_type = 'MERGE') + ) + """ + + def __init__(self, config: WorkflowSource, metadata: OpenMetadata): + super().__init__(config, metadata) + + self.redshift_instance_type = get_redshift_instance_type(self.engine) + + if self.redshift_instance_type == RedshiftInstanceType.PROVISIONED: + self.sql_stmt = REDSHIFT_SQL_STATEMENT_MAP[RedshiftInstanceType.PROVISIONED] + self.filters = self.provisioned_filters + logger.info( + "Using STL views for lineage processing of Redshift Provisioned" + ) + else: + self.sql_stmt = REDSHIFT_SQL_STATEMENT_MAP[RedshiftInstanceType.SERVERLESS] + self.filters = self.serverless_filters + logger.info("Using SYS views for lineage processing of Redshift Serverless") def yield_table_query(self) -> Iterator[TableQuery]: """ @@ -101,6 +130,8 @@ def get_stored_procedure_sql_statement(self) -> str: Return the SQL statement to get the stored procedure queries """ start, _ = get_start_and_end(self.source_config.queryLogDuration) - query = REDSHIFT_GET_STORED_PROCEDURE_QUERIES.format(start_date=start) + query = REDSHIFT_GET_STORED_PROCEDURE_QUERIES_MAP[ + self.redshift_instance_type + ].format(start_date=start) return query diff --git a/ingestion/src/metadata/ingestion/source/database/redshift/metadata.py b/ingestion/src/metadata/ingestion/source/database/redshift/metadata.py index 1d4acada9335..a8f8b76a4288 100644 --- a/ingestion/src/metadata/ingestion/source/database/redshift/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/redshift/metadata.py @@ -216,9 +216,11 @@ def query_table_names_and_types( result = self.connection.execute( sql.text( REDSHIFT_GET_ALL_RELATION_INFO.format( - view_filter="OR c.relkind IN ('v', 'm')" - if self.source_config.includeViews - else "AND c.relkind NOT IN ('v', 'm')" + view_filter=( + "OR c.relkind IN ('v', 'm')" + if self.source_config.includeViews + else "AND c.relkind NOT IN ('v', 'm')" + ) ) ), {"schema": schema_name}, @@ -263,7 +265,7 @@ def get_database_names_raw(self) -> Iterable[str]: yield from self._execute_database_query(REDSHIFT_GET_DATABASE_NAMES) def _set_incremental_table_processor(self, database: str): - """Prepares the needed data for doing incremental metadata extration for a given database. + """Prepares the needed data for doing incremental metadata extraction for a given database. 1. Queries Redshift to get the changes done after the `self.incremental.start_datetime_utc` 2. Sets the table map with the changes within the RedshiftIncrementalTableProcessor diff --git a/ingestion/src/metadata/ingestion/source/database/redshift/models.py b/ingestion/src/metadata/ingestion/source/database/redshift/models.py index 3423b8afa3a6..2ea274dc261a 100644 --- a/ingestion/src/metadata/ingestion/source/database/redshift/models.py +++ b/ingestion/src/metadata/ingestion/source/database/redshift/models.py @@ -12,6 +12,7 @@ Redshift models """ import re +from enum import Enum from typing import Dict, List, Optional, Tuple from pydantic import BaseModel @@ -20,6 +21,13 @@ SchemaName = str +class RedshiftInstanceType(Enum): + """Redshift Instance Types""" + + PROVISIONED = "PROVISIONED" + SERVERLESS = "SERVERLESS" + + class RedshiftStoredProcedure(BaseModel): """Redshift stored procedure list query results""" diff --git a/ingestion/src/metadata/ingestion/source/database/redshift/queries.py b/ingestion/src/metadata/ingestion/source/database/redshift/queries.py index fbcf6eaf5536..b183a5beaacb 100644 --- a/ingestion/src/metadata/ingestion/source/database/redshift/queries.py +++ b/ingestion/src/metadata/ingestion/source/database/redshift/queries.py @@ -13,10 +13,8 @@ """ import textwrap -from typing import List -from metadata.utils.profiler_utils import QueryResult -from metadata.utils.time_utils import datetime_to_timestamp +from metadata.ingestion.source.database.redshift.models import RedshiftInstanceType # Not able to use SYS_QUERY_HISTORY here. Few users not getting any results REDSHIFT_SQL_STATEMENT = textwrap.dedent( @@ -50,7 +48,7 @@ LISTAGG(CASE WHEN LEN(RTRIM(text)) = 0 THEN text ELSE RTRIM(text) END, '') WITHIN GROUP (ORDER BY sequence) AS query_text FROM deduped_querytext - WHERE sequence < 327 -- each chunk contains up to 200, RS has a maximum str length of 65535. + WHERE sequence < 327 -- each chunk contains up to 200, RS has a maximum str length of 65535. GROUP BY query ), raw_scans AS ( @@ -59,11 +57,11 @@ FROM pg_catalog.stl_scan ), scans AS ( - SELECT DISTINCT - query, - sti.database AS database_name, + SELECT DISTINCT + query, + sti.database AS database_name, sti.schema AS schema_name - FROM raw_scans AS s + FROM raw_scans AS s INNER JOIN pg_catalog.svv_table_info AS sti ON (s.tbl)::oid = sti.table_id ) @@ -90,6 +88,75 @@ ) +REDSHIFT_SERVERLESS_SQL_STATEMENT = textwrap.dedent( + """ +WITH queries AS ( + SELECT * + FROM SYS_QUERY_HISTORY + WHERE user_id > 1 + {filters} + -- Filter out all automated & cursor queries + AND LOWER(query_label) NOT IN ('maintenance', 'metrics', 'health') + AND query_text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%' + AND query_text NOT LIKE '/* {{"app": "dbt", %%}} */%%' + AND LOWER(status) = 'success' + AND start_time >= '{start_time}' + AND start_time < '{end_time}' +), +deduped_querytext AS ( + -- Sometimes rows are duplicated, causing LISTAGG to fail in the full_queries CTE. + SELECT DISTINCT + qt.* + FROM SYS_QUERY_TEXT AS qt + INNER JOIN queries AS q + ON qt.query_id = q.query_id +), +full_queries AS ( + SELECT + query_id, + LISTAGG(CASE WHEN LEN(RTRIM(text)) = 0 THEN text ELSE RTRIM(text) END, '') + WITHIN GROUP (ORDER BY sequence) AS query_text + FROM deduped_querytext + WHERE sequence < 327 -- each chunk contains up to 200, RS has a maximum str length of 65535. + GROUP BY query_id +), +query_detail AS ( + SELECT DISTINCT + query_id, + COALESCE(SPLIT_PART(table_name, '.', 1), 'unknown') AS database_name, + COALESCE(SPLIT_PART(table_name, '.', 2), 'unknown') AS schema_name, + COALESCE(SPLIT_PART(table_name, '.', 3), table_name) AS table_name + FROM SYS_QUERY_DETAIL + WHERE LOWER(step_name) IN ('insert', 'merge', 'delete') + AND table_name <> '' +) +SELECT DISTINCT + q.user_id, + q.query_id, + RTRIM(u.usename) AS user_name, + fq.query_text, + qd.database_name, + qd.schema_name, + q.start_time AS start_time, + q.end_time AS end_time, + datediff(millisecond, q.start_time, q.end_time) AS duration, + q.status, + q.query_type +FROM queries AS q + -- instances where query_detail has no entries of step 'insert', 'merge' or 'delete' + -- means no table was affected so perform inner join to filter those out + INNER JOIN query_detail AS qd + ON q.query_id = qd.query_id + INNER JOIN full_queries AS fq + ON q.query_id = fq.query_id + INNER JOIN pg_catalog.pg_user AS u + ON q.user_id = u.usesysid +ORDER BY q.start_time DESC +LIMIT {result_limit} +""" +) + + REDSHIFT_GET_ALL_RELATION_INFO = textwrap.dedent( """ SELECT @@ -231,17 +298,24 @@ """ REDSHIFT_TEST_GET_QUERIES = """ -SELECT +SELECT has_table_privilege('svv_table_info', 'SELECT') as can_access_svv_table_info, has_table_privilege('stl_querytext', 'SELECT') as can_access_stl_querytext, has_table_privilege('stl_query', 'SELECT') as can_access_stl_query; """ +REDSHIFT_SERVERLESS_TEST_GET_QUERIES = """ +SELECT + has_table_privilege('svv_table_info', 'SELECT') as can_access_svv_table_info, + has_table_privilege('sys_query_text', 'SELECT') as can_access_sys_query_text, + has_table_privilege('sys_query_history', 'SELECT') as can_access_sys_query_history; +""" + REDSHIFT_TEST_PARTITION_DETAILS = "select * from SVV_TABLE_INFO limit 1" REDSHIFT_GET_ALL_CONSTRAINTS = """ -select +select n.nspname as "schema", c.relname as "table_name", t.contype as "constraint_type", @@ -389,6 +463,66 @@ """ ) +REDSHIFT_SERVERLESS_GET_STORED_PROCEDURE_QUERIES = textwrap.dedent( + """ +WITH SP_HISTORY AS ( + SELECT + spcall.query_text AS procedure_text, + spcall.start_time AS procedure_start_time, + spcall.end_time AS procedure_end_time, + spcall.query_id AS procedure_query_id, + qh.session_id AS procedure_session_id + FROM SYS_PROCEDURE_CALL spcall + LEFT JOIN SYS_QUERY_HISTORY qh + ON spcall.query_id = qh.query_id + WHERE LOWER(spcall.status) = 'success' + AND spcall.start_time >= '{start_date}' +), +Q_HISTORY AS ( + SELECT + query_text, + CASE + WHEN query_text ILIKE '%%MERGE%%' THEN 'MERGE' + WHEN query_text ILIKE '%%UPDATE%%' THEN 'UPDATE' + WHEN query_text ILIKE '%%CREATE%%AS%%' THEN 'CREATE_TABLE_AS_SELECT' + WHEN query_text ILIKE '%%INSERT%%' THEN 'INSERT' + ELSE 'UNKNOWN' END query_type, + database_name AS query_database_name, + session_id as query_session_id, + start_time AS query_start_time, + end_time AS query_end_time, + CAST(b.usename AS varchar) AS query_user_name + FROM SYS_QUERY_HISTORY q + JOIN pg_catalog.pg_user b + ON b.usesysid = q.user_id + WHERE LOWER(q.query_label) NOT IN ('maintenance', 'metrics', 'health') + AND q.query_text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%' + AND q.query_text NOT LIKE '/* {{"app": "dbt", %%}} */%%' + AND LOWER(q.status) = 'success' + AND q.start_time >= '{start_date}' + AND q.user_id <> 1 +) +SELECT + TRIM(sp.procedure_text) procedure_text, + sp.procedure_start_time, + sp.procedure_end_time, + TRIM(q.query_text) query_text, + q.query_type, + TRIM(q.query_database_name) query_database_name, + null AS query_schema_name, + q.query_start_time, + q.query_end_time, + TRIM(q.query_user_name) query_user_name +FROM SP_HISTORY sp + JOIN Q_HISTORY q + ON sp.procedure_session_id = q.query_session_id + AND q.query_start_time BETWEEN sp.procedure_start_time AND sp.procedure_end_time + AND q.query_end_time BETWEEN sp.procedure_start_time AND sp.procedure_end_time +ORDER BY procedure_start_time DESC + """ +) + + REDSHIFT_LIFE_CYCLE_QUERY = textwrap.dedent( """ select "table" as table_name, @@ -415,51 +549,106 @@ """ -STL_QUERY = """ - with data as ( - select - {alias}.* - from - pg_catalog.stl_insert si - {join_type} join pg_catalog.stl_delete sd on si.query = sd.query - where - {condition} - ) - SELECT - SUM(data."rows") AS "rows", - sti."database", - sti."schema", - sti."table", - DATE_TRUNC('second', data.starttime) AS starttime - FROM - data - INNER JOIN pg_catalog.svv_table_info sti ON data.tbl = sti.table_id +REDSHIFT_SYSTEM_METRICS_QUERY = """ +with data as ( + select + {alias}.* + from + pg_catalog.stl_insert si + {join_type} join pg_catalog.stl_delete sd on si.query = sd.query where - sti."database" = '{database}' AND - sti."schema" = '{schema}' AND - "rows" != 0 AND - DATE(data.starttime) >= CURRENT_DATE - 1 - GROUP BY 2,3,4,5 - ORDER BY 5 DESC + {condition} +) +SELECT + SUM(data."rows") AS "rows", + sti."database", + sti."schema", + sti."table", + DATE_TRUNC('second', data.starttime) AS starttime +FROM + data + INNER JOIN pg_catalog.svv_table_info sti ON data.tbl = sti.table_id +where + sti."database" = '{database}' AND + sti."schema" = '{schema}' AND + "rows" != 0 AND + DATE(data.starttime) >= CURRENT_DATE - 1 +GROUP BY 2,3,4,5 +ORDER BY 5 DESC """ +# output_rows from SYS_QUERY_DETAIL should match rows from stl_insert/stl_delete +# It’s often wrong (usually too high). I noticed that sometimes the 'scan' step +# with plan_parent_id > 0 and plan_node_id > 0 gives the correct count, taking the +# min id if there are multiple scans. It worked in all the cases I tried, but it’s +# not really reliable for general use. +# For now, we just use the number of queries as a placeholder until we figure out +# a proper fix. +REDSHIFT_SERVERLESS_SYSTEM_METRICS_QUERY = """ +WITH data AS ( + SELECT + {alias}.output_rows AS rows, + COALESCE(SPLIT_PART({alias}.table_name, '.', 1), 'unknown') AS database, + COALESCE(SPLIT_PART({alias}.table_name, '.', 2), 'unknown') AS schema, + COALESCE(SPLIT_PART({alias}.table_name, '.', 3), {alias}.table_name) AS table, + DATE_TRUNC('second', {alias}.start_time) AS starttime + FROM ( + SELECT * + FROM SYS_QUERY_DETAIL + WHERE lower(step_name) = 'insert' + ) si + {join_type} JOIN ( + SELECT * + FROM SYS_QUERY_DETAIL + WHERE lower(step_name) = 'delete' + ) sd + ON si.query_id = sd.query_id + WHERE + {condition} + AND {alias}.table_name <> '' + AND DATE({alias}.start_time) >= CURRENT_DATE - 1 + AND {alias}.output_rows <> 0 +) +SELECT + COUNT(data.rows) as rows, + data.database, + data.schema, + data.table, + data.starttime +FROM data +WHERE + lower(data.database) = lower('{database}') + AND lower(data.schema) = lower('{schema}') +GROUP BY + data.database, + data.schema, + data.table, + data.starttime +ORDER BY data.starttime DESC; +""" -def get_metric_result(ddls: List[QueryResult], table_name: str) -> List: - """Given query results, retur the metric result +# Ideally, all serverless specific queries defined here should work with +# both Redshift Serverless and Provisioned since sys views are available +# in both instances. However, it still needs to be tested in Provisioned +# clusters. +# Ref: https://github.com/open-metadata/OpenMetadata/pull/6568/files#diff-65e5e8591345679be6a347ea29c4d283d5ca9aa723ef788c9a2524344de49ff3R17 - Args: - ddls (List[QueryResult]): list of query results - table_name (str): table name +REDSHIFT_TEST_GET_QUERIES_MAP = { + RedshiftInstanceType.PROVISIONED: REDSHIFT_TEST_GET_QUERIES, + RedshiftInstanceType.SERVERLESS: REDSHIFT_SERVERLESS_TEST_GET_QUERIES, +} - Returns: - List: - """ - return [ - { - "timestamp": datetime_to_timestamp(ddl.start_time, milliseconds=True), - "operation": ddl.query_type, - "rowsAffected": ddl.rows, - } - for ddl in ddls - if ddl.table_name == table_name - ] +REDSHIFT_SQL_STATEMENT_MAP = { + RedshiftInstanceType.PROVISIONED: REDSHIFT_SQL_STATEMENT, + RedshiftInstanceType.SERVERLESS: REDSHIFT_SERVERLESS_SQL_STATEMENT, +} + +REDSHIFT_GET_STORED_PROCEDURE_QUERIES_MAP = { + RedshiftInstanceType.PROVISIONED: REDSHIFT_GET_STORED_PROCEDURE_QUERIES, + RedshiftInstanceType.SERVERLESS: REDSHIFT_SERVERLESS_GET_STORED_PROCEDURE_QUERIES, +} + +REDSHIFT_SYSTEM_METRICS_QUERY_MAP = { + RedshiftInstanceType.PROVISIONED: REDSHIFT_SYSTEM_METRICS_QUERY, + RedshiftInstanceType.SERVERLESS: REDSHIFT_SERVERLESS_SYSTEM_METRICS_QUERY, +} diff --git a/ingestion/src/metadata/ingestion/source/database/redshift/usage.py b/ingestion/src/metadata/ingestion/source/database/redshift/usage.py index ab679ce50912..1b85dadf9d35 100644 --- a/ingestion/src/metadata/ingestion/source/database/redshift/usage.py +++ b/ingestion/src/metadata/ingestion/source/database/redshift/usage.py @@ -11,18 +11,48 @@ """ Redshift usage module """ -from metadata.ingestion.source.database.redshift.queries import REDSHIFT_SQL_STATEMENT +from metadata.ingestion.source.database.redshift.connection import ( + get_redshift_instance_type, +) +from metadata.ingestion.source.database.redshift.models import RedshiftInstanceType +from metadata.ingestion.source.database.redshift.queries import ( + REDSHIFT_SQL_STATEMENT_MAP, +) from metadata.ingestion.source.database.redshift.query_parser import ( RedshiftQueryParserSource, ) from metadata.ingestion.source.database.usage_source import UsageSource +from metadata.utils.logger import ingestion_logger + +logger = ingestion_logger() class RedshiftUsageSource(RedshiftQueryParserSource, UsageSource): - filters = """ + """Redshift Usage Source with support for both Provisioned and Serverless deployments.""" + + provisioned_filters = """ AND querytxt NOT ILIKE 'fetch%%' AND querytxt NOT ILIKE 'padb_fetch_sample:%%' AND querytxt NOT ILIKE 'Undoing%%transactions%%on%%table%%with%%current%%xid%%' """ - sql_stmt = REDSHIFT_SQL_STATEMENT + # Serverless uses SYS views instead of STL views and have query_text column instead of querytxt + serverless_filters = """ + AND query_text NOT ILIKE 'fetch%%' + AND query_text NOT ILIKE 'padb_fetch_sample:%%' + AND query_text NOT ILIKE 'Undoing%%transactions%%on%%table%%with%%current%%xid%%' + """ + + def __init__(self, config, metadata_config): + super().__init__(config, metadata_config) + + self.redshift_instance_type = get_redshift_instance_type(self.engine) + + if self.redshift_instance_type == RedshiftInstanceType.PROVISIONED: + self.sql_stmt = REDSHIFT_SQL_STATEMENT_MAP[RedshiftInstanceType.PROVISIONED] + self.filters = self.provisioned_filters + logger.info("Using STL views for usage processing of Redshift Provisioned") + else: + self.sql_stmt = REDSHIFT_SQL_STATEMENT_MAP[RedshiftInstanceType.SERVERLESS] + self.filters = self.serverless_filters + logger.info("Using SYS views for usage processing of Redshift Serverless") diff --git a/ingestion/src/metadata/ingestion/source/database/redshift/utils.py b/ingestion/src/metadata/ingestion/source/database/redshift/utils.py index 0eeb1dff12bc..5531a3f32158 100644 --- a/ingestion/src/metadata/ingestion/source/database/redshift/utils.py +++ b/ingestion/src/metadata/ingestion/source/database/redshift/utils.py @@ -109,7 +109,7 @@ def _get_column_info(self, *args, **kwargs): )._get_column_info(*args, **kwdrs) # raw_data_type is not included in column_info as - # redhift doesn't support complex data types directly + # redshift doesn't support complex data types directly # https://docs.aws.amazon.com/redshift/latest/dg/c_Supported_data_types.html if "info" not in column_info: diff --git a/ingestion/src/metadata/profiler/metrics/system/redshift/system.py b/ingestion/src/metadata/profiler/metrics/system/redshift/system.py index 376d28ae49ad..c6506dc54980 100644 --- a/ingestion/src/metadata/profiler/metrics/system/redshift/system.py +++ b/ingestion/src/metadata/profiler/metrics/system/redshift/system.py @@ -8,7 +8,13 @@ from sqlalchemy.orm import Session from metadata.generated.schema.entity.data.table import SystemProfile -from metadata.ingestion.source.database.redshift.queries import STL_QUERY +from metadata.ingestion.source.database.redshift.connection import ( + get_redshift_instance_type, +) +from metadata.ingestion.source.database.redshift.models import RedshiftInstanceType +from metadata.ingestion.source.database.redshift.queries import ( + REDSHIFT_SYSTEM_METRICS_QUERY_MAP, +) from metadata.profiler.metrics.system.dml_operation import DatabaseDMLOperations from metadata.profiler.metrics.system.system import ( CacheProvider, @@ -37,6 +43,9 @@ def __init__( self.table = runner.table_name self.database = runner.session.get_bind().url.database self.schema = runner.schema_name + self.engine = runner.session.get_bind() + + self.redshift_instance_type = get_redshift_instance_type(self.engine) def get_inserts(self) -> List[SystemProfile]: queries = self.get_or_update_cache( @@ -66,13 +75,26 @@ def get_updates(self) -> List[SystemProfile]: return get_metric_result(queries, self.table) def _get_insert_queries(self, database: str, schema: str) -> List[QueryResult]: - insert_query = STL_QUERY.format( - alias="si", - join_type="LEFT", - condition="sd.query is null", - database=database, - schema=schema, - ) + if self.redshift_instance_type == RedshiftInstanceType.PROVISIONED: + insert_query = REDSHIFT_SYSTEM_METRICS_QUERY_MAP[ + RedshiftInstanceType.PROVISIONED + ].format( + alias="si", + join_type="LEFT", + condition="sd.query is null", + database=database, + schema=schema, + ) + else: + insert_query = REDSHIFT_SYSTEM_METRICS_QUERY_MAP[ + RedshiftInstanceType.SERVERLESS + ].format( + alias="si", + join_type="LEFT", + condition="sd.query_id is null", + database=database, + schema=schema, + ) return self._get_query_results( self.session, insert_query, @@ -80,13 +102,26 @@ def _get_insert_queries(self, database: str, schema: str) -> List[QueryResult]: ) def _get_delete_queries(self, database: str, schema: str) -> List[QueryResult]: - delete_query = STL_QUERY.format( - alias="sd", - join_type="RIGHT", - condition="si.query is null", - database=database, - schema=schema, - ) + if self.redshift_instance_type == RedshiftInstanceType.PROVISIONED: + delete_query = REDSHIFT_SYSTEM_METRICS_QUERY_MAP[ + RedshiftInstanceType.PROVISIONED + ].format( + alias="sd", + join_type="RIGHT", + condition="si.query is null", + database=database, + schema=schema, + ) + else: + delete_query = REDSHIFT_SYSTEM_METRICS_QUERY_MAP[ + RedshiftInstanceType.SERVERLESS + ].format( + alias="sd", + join_type="RIGHT", + condition="si.query_id is null", + database=database, + schema=schema, + ) return self._get_query_results( self.session, delete_query, @@ -94,13 +129,26 @@ def _get_delete_queries(self, database: str, schema: str) -> List[QueryResult]: ) def _get_update_queries(self, database: str, schema: str) -> List[QueryResult]: - update_query = STL_QUERY.format( - alias="si", - join_type="INNER", - condition="sd.query is not null", - database=database, - schema=schema, - ) + if self.redshift_instance_type == RedshiftInstanceType.PROVISIONED: + update_query = REDSHIFT_SYSTEM_METRICS_QUERY_MAP[ + RedshiftInstanceType.PROVISIONED + ].format( + alias="si", + join_type="INNER", + condition="sd.query is not null", + database=database, + schema=schema, + ) + else: + update_query = REDSHIFT_SYSTEM_METRICS_QUERY_MAP[ + RedshiftInstanceType.SERVERLESS + ].format( + alias="si", + join_type="INNER", + condition="sd.query_id is not null", + database=database, + schema=schema, + ) return self._get_query_results( self.session, update_query, @@ -109,7 +157,7 @@ def _get_update_queries(self, database: str, schema: str) -> List[QueryResult]: def get_metric_result(ddls: List[QueryResult], table_name: str) -> List[SystemProfile]: - """Given query results, retur the metric result + """Given query results, return the metric result Args: ddls (List[QueryResult]): list of query results diff --git a/ingestion/tests/unit/topology/database/test_redshift.py b/ingestion/tests/unit/topology/database/test_redshift.py index bc2d5128518b..babdfaaa108f 100644 --- a/ingestion/tests/unit/topology/database/test_redshift.py +++ b/ingestion/tests/unit/topology/database/test_redshift.py @@ -14,12 +14,19 @@ """ from unittest import TestCase -from unittest.mock import patch +from unittest.mock import Mock, patch from metadata.generated.schema.metadataIngestion.workflow import ( OpenMetadataWorkflowConfig, ) +from metadata.ingestion.source.database.redshift.connection import ( + get_redshift_instance_type, +) from metadata.ingestion.source.database.redshift.metadata import RedshiftSource +from metadata.ingestion.source.database.redshift.models import RedshiftInstanceType +from metadata.ingestion.source.database.redshift.queries import ( + REDSHIFT_SQL_STATEMENT_MAP, +) mock_redshift_config = { "source": { @@ -84,3 +91,26 @@ def test_partition_parse_columns(self): def test_close_connection(self, engine, connection): connection.return_value = True self.redshift_source.close() + + def test_detect_provisioned_when_stl_accessible(self): + """Verify Provisioned is detected when STL tables are accessible""" + mock_engine = Mock() + mock_conn = Mock() + mock_context = Mock() + mock_context.__enter__ = Mock(return_value=mock_conn) + mock_context.__exit__ = Mock(return_value=False) + mock_engine.connect.return_value = mock_context + mock_conn.execute.return_value = Mock() # STL query succeeds + + result = get_redshift_instance_type(mock_engine) + + self.assertEqual(result, RedshiftInstanceType.PROVISIONED) + mock_conn.execute.assert_called_once() + + def test_provisioned_uses_stl_queries(self): + """Verify Provisioned uses STL-based queries""" + sql = REDSHIFT_SQL_STATEMENT_MAP[RedshiftInstanceType.PROVISIONED] + + # Must use STL tables + self.assertIn("stl_query", sql) + self.assertIn("stl_querytext", sql) diff --git a/ingestion/tests/unit/topology/database/test_redshift_serverless.py b/ingestion/tests/unit/topology/database/test_redshift_serverless.py new file mode 100644 index 000000000000..94e6f6d1863a --- /dev/null +++ b/ingestion/tests/unit/topology/database/test_redshift_serverless.py @@ -0,0 +1,73 @@ +# Copyright 2025 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# 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. + +""" +Test Redshift Serverless functionality +""" + +from unittest import TestCase +from unittest.mock import Mock + +from sqlalchemy.exc import ProgrammingError + +from metadata.ingestion.source.database.redshift.connection import ( + get_redshift_instance_type, +) +from metadata.ingestion.source.database.redshift.models import RedshiftInstanceType +from metadata.ingestion.source.database.redshift.queries import ( + REDSHIFT_GET_STORED_PROCEDURE_QUERIES_MAP, + REDSHIFT_SQL_STATEMENT_MAP, + REDSHIFT_SYSTEM_METRICS_QUERY_MAP, +) + + +class TestRedshiftServerlessInstance(TestCase): + """Test Redshift Serverless instance type detection and query selection""" + + def test_detect_serverless_when_stl_not_accessible(self): + """Verify Serverless is detected when STL tables fail""" + mock_engine = Mock() + mock_conn = Mock() + mock_context = Mock() + mock_context.__enter__ = Mock(return_value=mock_conn) + mock_context.__exit__ = Mock(return_value=False) + mock_engine.connect.return_value = mock_context + mock_conn.execute.side_effect = ProgrammingError( + "relation does not exist", params=None, orig=None + ) + + result = get_redshift_instance_type(mock_engine) + + self.assertEqual(result, RedshiftInstanceType.SERVERLESS) + + def test_serverless_uses_sys_queries(self): + """Verify Serverless uses SYS views not STL tables""" + sql = REDSHIFT_SQL_STATEMENT_MAP[RedshiftInstanceType.SERVERLESS] + + self.assertIn("SYS_QUERY_HISTORY", sql) + self.assertIn("SYS_QUERY_DETAIL", sql) + self.assertNotIn("stl_query", sql.lower()) + + def test_serverless_stored_procedures_use_sys(self): + """Verify Serverless stored procedures use SYS views""" + query = REDSHIFT_GET_STORED_PROCEDURE_QUERIES_MAP[ + RedshiftInstanceType.SERVERLESS + ] + + self.assertIn("SYS_PROCEDURE_CALL", query) + self.assertIn("SYS_QUERY_HISTORY", query) + + def test_serverless_system_metrics_use_sys(self): + """Verify Serverless system metrics use SYS views""" + query = REDSHIFT_SYSTEM_METRICS_QUERY_MAP[RedshiftInstanceType.SERVERLESS] + + self.assertIn("SYS_QUERY_DETAIL", query) + self.assertNotIn("stl_insert", query.lower())