quollio-core 0.4.4__py3-none-any.whl → 0.4.10__py3-none-any.whl

quollio_core/profilers/bigquery.py

@@ -0,0 +1,81 @@
+import logging
+from typing import Any, Dict, List
+
+from quollio_core.profilers.lineage import gen_table_lineage_payload, parse_bigquery_table_lineage
+from quollio_core.repository import qdc
+from quollio_core.repository.bigquery import BigQueryClient, GCPLineageClient, get_entitiy_reference, get_search_request
+
+logger = logging.getLogger(__name__)
+
+
+def bigquery_table_lineage(
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+    project_id: str,
+    regions: list,
+    org_id: str,
+    credentials: Any,
+):
+    lineage_client = GCPLineageClient(credentials)
+    bq_client = BigQueryClient(credentials)
+
+    datasets = bq_client.list_datasets(project_id)
+    all_tables = generate_table_list(datasets, bq_client)
+    lineage_links = generate_lineage_links(all_tables, lineage_client, project_id, regions)
+    lineage_links = parse_bigquery_table_lineage(lineage_links)
+
+    update_table_lineage_inputs = gen_table_lineage_payload(tenant_id=tenant_id, endpoint=org_id, tables=lineage_links)
+
+    req_count = 0
+    for update_table_lineage_input in update_table_lineage_inputs:
+        logger.info(
+            "Generating table lineage. downstream: %s -> %s-> %s",
+            update_table_lineage_input.downstream_database_name,
+            update_table_lineage_input.downstream_schema_name,
+            update_table_lineage_input.downstream_table_name,
+        )
+        status_code = qdc_client.update_lineage_by_id(
+            global_id=update_table_lineage_input.downstream_global_id,
+            payload=update_table_lineage_input.upstreams.as_dict(),
+        )
+        if status_code == 200:
+            req_count += 1
+    logger.info("Generating table lineage is finished. %s lineages are ingested.", req_count)
+
+
+def generate_table_list(datasets: List[str], bq_client: BigQueryClient) -> List[str]:
+    all_tables = []
+    for dataset in datasets:
+        all_tables.extend(
+            [
+                table
+                for table in bq_client.list_tables(dataset.dataset_id)
+                if table.table_type in ["TABLE", "VIEW", "MATERIALIZED_VIEW"]
+            ]
+        )
+
+    all_table_names = []
+    for table in all_tables:
+        all_table_names.append(f"{table.project}.{table.dataset_id}.{table.table_id}")
+
+    return all_table_names
+
+
+def generate_lineage_links(
+    all_tables: List[str], lineage_client: GCPLineageClient, project_id: str, regions: List[str]
+) -> Dict[str, List[str]]:
+    lineage_links = {}
+    for table in all_tables:
+        downstream = get_entitiy_reference()
+        downstream.fully_qualified_name = f"bigquery:{table}"
+
+        for region in regions:
+            request = get_search_request(downstream_table=downstream, project_id=project_id, region=region)
+            response = lineage_client.get_links(request=request)
+            for lineage in response:
+                target_table = str(lineage.target.fully_qualified_name).replace("bigquery:", "")
+                if target_table not in lineage_links:
+                    lineage_links[target_table] = []
+                lineage_links[target_table].append(str(lineage.source.fully_qualified_name).replace("bigquery:", ""))
+
+    return lineage_links

quollio_core/profilers/databricks.py

@@ -0,0 +1,198 @@
+import logging
+from typing import Dict, List
+
+from quollio_core.profilers.lineage import (
+    gen_column_lineage_payload,
+    gen_table_lineage_payload,
+    parse_databricks_table_lineage,
+)
+from quollio_core.profilers.stats import gen_table_stats_payload
+from quollio_core.repository import databricks, qdc
+
+logger = logging.getLogger(__name__)
+
+
+def databricks_table_level_lineage(
+    conn: databricks.DatabricksConnectionConfig,
+    endpoint: str,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+    dbt_table_name: str = "quollio_lineage_table_level",
+) -> None:
+    with databricks.DatabricksQueryExecutor(config=conn) as databricks_executor:
+        results = databricks_executor.get_query_results(
+            query=f"""
+            SELECT
+                DOWNSTREAM_TABLE_NAME,
+                UPSTREAM_TABLES
+            FROM {conn.catalog}.{conn.schema}.{dbt_table_name}
+            """
+        )
+        tables = parse_databricks_table_lineage(results)
+        update_table_lineage_inputs = gen_table_lineage_payload(
+            tenant_id=tenant_id,
+            endpoint=endpoint,
+            tables=tables,
+        )
+
+        req_count = 0
+        for update_table_lineage_input in update_table_lineage_inputs:
+            logger.info(
+                "Generating table lineage. downstream: %s -> %s-> %s",
+                update_table_lineage_input.downstream_database_name,
+                update_table_lineage_input.downstream_schema_name,
+                update_table_lineage_input.downstream_table_name,
+            )
+            status_code = qdc_client.update_lineage_by_id(
+                global_id=update_table_lineage_input.downstream_global_id,
+                payload=update_table_lineage_input.upstreams.as_dict(),
+            )
+            if status_code == 200:
+                req_count += 1
+        logger.info("Generating table lineage is finished. %s lineages are ingested.", req_count)
+    return
+
+
+def databricks_column_level_lineage(
+    conn: databricks.DatabricksConnectionConfig,
+    endpoint: str,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+    dbt_table_name: str = "quollio_lineage_column_level",
+) -> None:
+    with databricks.DatabricksQueryExecutor(config=conn) as databricks_executor:
+        results = databricks_executor.get_query_results(
+            query=f"""
+            SELECT
+                *
+            FROM
+                {conn.catalog}.{conn.schema}.{dbt_table_name}
+            """
+        )
+
+    update_column_lineage_inputs = gen_column_lineage_payload(
+        tenant_id=tenant_id,
+        endpoint=endpoint,
+        columns=results,
+    )
+
+    req_count = 0
+    for update_column_lineage_input in update_column_lineage_inputs:
+        logger.info(
+            "Generating column lineage. downstream: %s -> %s -> %s -> %s",
+            update_column_lineage_input.downstream_database_name,
+            update_column_lineage_input.downstream_schema_name,
+            update_column_lineage_input.downstream_table_name,
+            update_column_lineage_input.downstream_column_name,
+        )
+        status_code = qdc_client.update_lineage_by_id(
+            global_id=update_column_lineage_input.downstream_global_id,
+            payload=update_column_lineage_input.upstreams.as_dict(),
+        )
+        if status_code == 200:
+            req_count += 1
+    logger.info(
+        "Generating column lineage is finished. %s lineages are ingested.",
+        req_count,
+    )
+    return
+
+
+def _get_monitoring_tables(
+    conn: databricks.DatabricksConnectionConfig, monitoring_table_suffix: str = "_profile_metrics"
+) -> List[Dict[str, str]]:
+    tables = []
+    query = f"""
+        SELECT
+            table_catalog,
+            table_schema,
+            table_name,
+            CONCAT(table_catalog, '.', table_schema, '.', table_name) AS table_fqdn
+        FROM
+            system.information_schema.tables
+        WHERE
+            table_name LIKE "%{monitoring_table_suffix}"
+            AND table_name NOT LIKE ('quollio_%')
+        """
+    with databricks.DatabricksQueryExecutor(config=conn) as databricks_executor:
+        tables = databricks_executor.get_query_results(query)
+    if len(tables) > 0:
+        logger.info("Found %s monitoring tables.", len(tables))
+        return tables
+    else:
+        logger.info("No monitoring tables found.")
+        return []
+
+
+def _get_column_stats(
+    conn: databricks.DatabricksConnectionConfig, monitoring_table_suffix: str = "_profile_metrics"
+) -> List[Dict[str, str]]:
+    tables = _get_monitoring_tables(conn, monitoring_table_suffix)
+    if not tables:
+        return []
+    stats = []
+    for table in tables:
+        monitored_table = table["table_fqdn"].removesuffix("_profile_metrics")
+        monitored_table = monitored_table.split(".")
+        if len(monitored_table) != 3:
+            raise ValueError(f"Invalid table name: {table['table_fqdn']}")
+        with databricks.DatabricksQueryExecutor(config=conn) as databricks_executor:
+            query = """
+                SELECT
+                    "{monitored_table_catalog}" as DB_NAME,
+                    "{monitored_table_schema}" as SCHEMA_NAME,
+                    "{monitored_table_name}" as TABLE_NAME,
+                    t.COLUMN_NAME,
+                    t.DATA_TYPE,
+                    t.distinct_count as CARDINALITY,
+                    t.MAX as MAX_VALUE,
+                    t.MIN as MIN_VALUE,
+                    t.AVG as AVG_VALUE,
+                    t.MEDIAN as MEDIAN_VALUE,
+                    t.STDDEV as STDDEV_VALUE,
+                    t.NUM_NULLS as NULL_COUNT,
+                    t.frequent_items[0].item AS MODE_VALUE,
+                    MAX(t.window) AS LATEST
+                FROM
+                    {monitoring_table} t
+                WHERE
+                    t.column_name not in (':table')
+                GROUP BY
+                    t.COLUMN_NAME,
+                    t.DATA_TYPE,
+                    t.distinct_count,
+                    t.MAX,
+                    t.MIN,
+                    t.AVG,
+                    t.MEDIAN,
+                    t.STDDEV,
+                    t.NUM_NULLS,
+                    t.frequent_items
+                """.format(
+                monitoring_table=table["table_fqdn"],
+                monitored_table_catalog=monitored_table[0],
+                monitored_table_schema=monitored_table[1],
+                monitored_table_name=monitored_table[2],
+            )
+            stats.append(databricks_executor.get_query_results(query))
+    return stats
+
+
+def databricks_column_stats(
+    conn: databricks.DatabricksConnectionConfig,
+    endpoint: str,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+    monitoring_table_suffix: str = "_profile_metrics",
+) -> None:
+    table_stats = _get_column_stats(conn, monitoring_table_suffix)
+    for table in table_stats:
+        stats = gen_table_stats_payload(tenant_id=tenant_id, endpoint=endpoint, stats=table)
+        for stat in stats:
+            status_code = qdc_client.update_stats_by_id(
+                global_id=stat.global_id,
+                payload=stat.body.as_dict(),
+            )
+            if status_code == 200:
+                logger.info("Stats for %s is successfully ingested.", stat.global_id)
+    return

quollio_core/profilers/lineage.py

@@ -141,3 +141,29 @@ def parse_snowflake_results(results: List[Dict[str, str]]):
         payload["UPSTREAM_TABLES"] = json.loads(result["UPSTREAM_TABLES"])
         payloads.append(payload)
     return payloads
+
+
+def parse_databricks_table_lineage(results: List) -> List[Dict[str, Dict]]:
+    # Parses results from Quollio Databricks lineage table
+    # Returns tuple of downstream_table_name (0) and upstream_tables (1)
+    payloads = list()
+    for result in results:
+        payload = dict()
+        payload["DOWNSTREAM_TABLE_NAME"] = result["DOWNSTREAM_TABLE_NAME"]
+        payload["UPSTREAM_TABLES"] = json.loads(result["UPSTREAM_TABLES"])
+        payloads.append(payload)
+    return payloads
+
+
+def parse_bigquery_table_lineage(tables: Dict) -> List[Dict[str, Dict]]:
+    payloads = list()
+    for downstream, upstream in tables.items():
+        payload = {
+            "DOWNSTREAM_TABLE_NAME": "",
+            "UPSTREAM_TABLES": [],
+        }
+        payload["DOWNSTREAM_TABLE_NAME"] = downstream
+        for upstream_table in upstream:
+            payload["UPSTREAM_TABLES"].append({"upstream_object_name": upstream_table})
+        payloads.append(payload)
+    return payloads

quollio_core/profilers/redshift.py

@@ -14,7 +14,6 @@ def redshift_table_level_lineage(
     tenant_id: str,
     dbt_table_name: str,
 ) -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
     with redshift.RedshiftQueryExecutor(config=conn) as redshift_executor:
         results = redshift_executor.get_query_results(
             query="""
@@ -55,22 +54,7 @@ def redshift_table_level_lineage(
     return
 
 
-def _get_target_tables_query(db: str, schema: str) -> str:
-    query = """
-        SELECT
-            DISTINCT
-            database_name
-            , schema_name
-            , table_name
-        FROM
-            {db}.{schema}.quollio_stats_profiling_columns
-        """.format(
-        db=db, schema=schema
-    )
-    return query
-
-
-def _get_stats_tables_query(db: str, schema: str) -> str:
+def _gen_get_stats_views_query(db: str, schema: str) -> str:
     query = """
         SELECT
             DISTINCT
@@ -93,70 +77,54 @@ def redshift_table_stats(
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
 ) -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
 
     with redshift.RedshiftQueryExecutor(config=conn) as redshift_executor:
-        req_count = 0
-        target_query = _get_target_tables_query(
+        stats_query = _gen_get_stats_views_query(
             db=conn.database,
             schema=conn.schema,
         )
-        target_assets = redshift_executor.get_query_results(query=target_query)
+        stats_views = redshift_executor.get_query_results(query=stats_query)
 
-        stats_query = _get_stats_tables_query(
-            db=conn.database,
-            schema=conn.schema,
-        )
-        stats_columns = redshift_executor.get_query_results(query=stats_query)
-        for target_asset in target_assets:
-            for stats_column in stats_columns:
-                stats_query = """
-                SELECT
-                    db_name
-                    , schema_name
-                    , table_name
-                    , column_name
-                    , max_value
-                    , min_value
-                    , null_count
-                    , cardinality
-                    , avg_value
-                    , median_value
-                    , mode_value
-                    , stddev_value
-                FROM
-                    {db}.{schema}.{table}
-                WHERE
-                    db_name = '{target_db}'
-                    and schema_name = '{target_schema}'
-                    and table_name = '{target_table}'
-                """.format(
-                    db=stats_column[0],
-                    schema=stats_column[1],
-                    table=stats_column[2],
-                    target_db=target_asset[0],
-                    target_schema=target_asset[1],
-                    target_table=target_asset[2],
+        req_count = 0
+        for stats_view in stats_views:
+            stats_query = """
+            SELECT
+                db_name
+                , schema_name
+                , table_name
+                , column_name
+                , max_value
+                , min_value
+                , null_count
+                , cardinality
+                , avg_value
+                , median_value
+                , mode_value
+                , stddev_value
+            FROM
+                {db}.{schema}.{table}
+            """.format(
+                db=stats_view[0],
+                schema=stats_view[1],
+                table=stats_view[2],
+            )
+            stats_result = redshift_executor.get_query_results(query=stats_query)
+            payloads = gen_table_stats_payload_from_tuple(tenant_id=tenant_id, endpoint=conn.host, stats=stats_result)
+            for payload in payloads:
+                logger.info(
+                    "Generating table stats. asset: {db} -> {schema} -> {table} -> {column}".format(
+                        db=payload.db,
+                        schema=payload.schema,
+                        table=payload.table,
+                        column=payload.column,
+                    )
                 )
-                stats_result = redshift_executor.get_query_results(query=stats_query)
-                payloads = gen_table_stats_payload_from_tuple(
-                    tenant_id=tenant_id, endpoint=conn.host, stats=stats_result
+                status_code = qdc_client.update_stats_by_id(
+                    global_id=payload.global_id,
+                    payload=payload.body.get_column_stats(),
                 )
-                for payload in payloads:
-                    logger.info(
-                        "Generating table stats. asset: {db} -> {schema} -> {table} -> {column}".format(
-                            db=payload.db,
-                            schema=payload.schema,
-                            table=payload.table,
-                            column=payload.column,
-                        )
-                    )
-                    status_code = qdc_client.update_stats_by_id(
-                        global_id=payload.global_id,
-                        payload=payload.body.get_column_stats(),
-                    )
-                    if status_code == 200:
-                        req_count += 1
+                if status_code == 200:
+                    req_count += 1
     logger.info(f"Generating table stats is finished. {req_count} stats are ingested.")
     return
 
@@ -166,7 +134,6 @@ def redshift_table_level_sqllineage(
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
 ) -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
     redshift_connector = redshift.RedshiftQueryExecutor(conn)
     results = redshift_connector.get_query_results(
         query="""