Added result limits to usage queries (#6748)

* Added result limits to queries

* Fixed formatting

Co-authored-by: Onkar Ravgan <onkarravgan@Onkars-MacBook-Pro.local>

catalog-rest-service/src/main/resources/json/schema/metadataIngestion/databaseServiceQueryLineagePipeline.json

@@ -29,7 +29,7 @@
     "resultLimit": {
       "description": "Configuration to set the limit for query logs",
       "type": "integer",
-      "default": "100"
+      "default": "1000"
     },
     "schemaFilterPattern": {
       "description": "Regex to only fetch tables or databases that matches the pattern.",

catalog-rest-service/src/main/resources/json/schema/metadataIngestion/databaseServiceQueryUsagePipeline.json

@@ -30,7 +30,7 @@
     "resultLimit": {
       "description": "Configuration to set the limit for query logs",
       "type": "integer",
-      "default": "100"
+      "default": "1000"
     },
     "queryLogFilePath": {
       "description": "Configuration to set the file path for query logs",

ingestion/examples/workflows/mssql_lineage.yaml

@@ -0,0 +1,15 @@
+source:
+  type: mssql-lineage
+  serviceName: local_mssql
+  sourceConfig:
+    config:
+      type: DatabaseLineage
+      queryLogDuration: 1
+      resultLimit: 10000
+sink:
+  type: metadata-rest
+  config: {}
+workflowConfig:
+  openMetadataServerConfig:
+    hostPort: http://localhost:8585/api
+    authProvider: no-auth

ingestion/src/metadata/ingestion/source/database/bigquery_query_parser.py

@@ -59,6 +59,7 @@ class BigqueryQueryParserSource(QueryParserSource, ABC):
             end_time=end_time,
             region=self.connection.usageLocation,
             filters=self.filters,
+            result_limit=self.source_config.resultLimit,
         )
 
     @staticmethod

ingestion/src/metadata/ingestion/source/database/clickhouse_query_parser.py

@@ -62,4 +62,5 @@ class ClickhouseQueryParserSource(QueryParserSource, ABC):
             start_time=start_time,
             end_time=end_time,
             filters=self.filters,
+            result_limit=self.source_config.resultLimit,
         )

ingestion/src/metadata/ingestion/source/database/mssql_lineage.py

@@ -0,0 +1,23 @@
+#  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 lineage module
+"""
+from metadata.ingestion.source.database.lineage_source import LineageSource
+from metadata.ingestion.source.database.mssql_query_parser import MssqlQueryParserSource
+from metadata.utils.sql_queries import MSSQL_SQL_STATEMENT
+
+
+class MssqlLineageSource(MssqlQueryParserSource, LineageSource):
+
+    sql_stmt = MSSQL_SQL_STATEMENT
+
+    filters = ""  # No filtering in the queries

ingestion/src/metadata/ingestion/source/database/query_parser_source.py

@@ -95,7 +95,10 @@ class QueryParserSource(Source[Union[TableQuery, AddLineageRequest]], ABC):
         Override if we have specific parameters
         """
         return self.sql_stmt.format(
-            start_time=start_time, end_time=end_time, filters=self.filters
+            start_time=start_time,
+            end_time=end_time,
+            filters=self.filters,
+            result_limit=self.source_config.resultLimit,
         )
 
     def get_report(self):

ingestion/src/metadata/ingestion/source/database/redshift_query_parser.py

@@ -55,4 +55,5 @@ class RedshiftQueryParserSource(QueryParserSource, ABC):
             start_time=start_time,
             end_time=end_time,
             filters=self.filters,
+            result_limit=self.source_config.resultLimit,
         )

ingestion/src/metadata/utils/sql_queries.py

@@ -29,6 +29,7 @@ REDSHIFT_SQL_STATEMENT = textwrap.dedent(
           AND aborted = 0
           AND starttime >= '{start_time}'
           AND starttime < '{end_time}'
+          LIMIT {result_limit}
           
   ),
   full_queries AS (
@@ -352,7 +353,7 @@ VERTICA_VIEW_DEFINITION = textwrap.dedent(
 
 MSSQL_SQL_STATEMENT = textwrap.dedent(
     """
-      SELECT
+      SELECT TOP {result_limit}
         db.NAME database_name,
         t.text query_text,
         s.last_execution_time start_time,
@@ -396,6 +397,7 @@ CLICKHOUSE_SQL_STATEMENT = textwrap.dedent(
         and query NOT LIKE '/* {{"app": "dbt", %%}} */%%'
         {filters}
         and (`type`='QueryFinish' or `type`='QueryStart')
+        LIMIT {result_limit}
 """
 )
 
@@ -445,6 +447,7 @@ WHERE creation_time BETWEEN "{start_time}" AND "{end_time}"
   AND IFNULL(statement_type, "NO") not in ("NO", "DROP_TABLE", "CREATE_TABLE")
   AND query NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
   AND query NOT LIKE '/* {{"app": "dbt", %%}} */%%'
+  LIMIT {result_limit}
 """
 )