quollio-core 0.4.12__tar.gz → 0.4.13__tar.gz

{quollio_core-0.4.12 → quollio_core-0.4.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: quollio-core
-Version: 0.4.12
+Version: 0.4.13
 Summary: Quollio Core
 Author-email: quollio-dev <qt.dev@quollio.com>
 Maintainer-email: RyoAriyama <ryo.arym@gmail.com>, tharuta <35373297+TakumiHaruta@users.noreply.github.com>
@@ -22,6 +22,7 @@ Requires-Dist: dbt-core==1.7.10
 Requires-Dist: dbt-snowflake==1.7.0
 Requires-Dist: dbt-redshift==1.7.1
 Requires-Dist: dbt-databricks==1.7.1
+Requires-Dist: db-dtypes==1.2.0
 Requires-Dist: jinja2==3.1.3
 Requires-Dist: PyYAML==6.0.1
 Requires-Dist: requests==2.31.0

{quollio_core-0.4.12 → quollio_core-0.4.13}/pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
     ,"dbt-snowflake==1.7.0"
     ,"dbt-redshift==1.7.1"
     ,"dbt-databricks==1.7.1"
+    ,"db-dtypes==1.2.0"
     ,"jinja2==3.1.3"
     ,"PyYAML==6.0.1"
     ,"requests==2.31.0"

{quollio_core-0.4.12 → quollio_core-0.4.13}/quollio_core/__init__.py

@@ -1,4 +1,4 @@
 """Quollio Core"""
 
-__version__ = "0.4.12"
+__version__ = "0.4.13"
 __author__ = "Quollio Technologies, Inc"

quollio_core-0.4.13/quollio_core/bigquery.py

@@ -0,0 +1,183 @@
+import argparse
+import json
+
+from google.auth.credentials import Credentials
+
+from quollio_core.helper.env_default import env_default
+from quollio_core.helper.log_utils import configure_logging, error_handling_decorator, logger
+from quollio_core.profilers.bigquery import bigquery_table_lineage, bigquery_table_stats
+from quollio_core.repository import qdc
+from quollio_core.repository.bigquery import BigQueryClient, get_credentials, get_org_id
+
+
+def initialize_credentials(credentials_json: str) -> Credentials:
+    return get_credentials(json.loads(credentials_json))
+
+
+def initialize_org_id(credentials_json: str) -> str:
+    return get_org_id(json.loads(credentials_json))
+
+
+def initialize_bq_client(credentials: Credentials, project_id: str) -> BigQueryClient:
+    return BigQueryClient(credentials=credentials, project_id=project_id)
+
+
+@error_handling_decorator
+def load_lineage(
+    tenant_id: str,
+    project_id: str,
+    regions: list,
+    org_id: str,
+    credentials: Credentials,
+    qdc_client: qdc.QDCExternalAPIClient,
+) -> None:
+    logger.info("Loading lineage data.")
+    bigquery_table_lineage(
+        qdc_client=qdc_client,
+        tenant_id=tenant_id,
+        project_id=project_id,
+        regions=regions,
+        credentials=credentials,
+        org_id=org_id,
+    )
+    logger.info("Lineage data loaded successfully.")
+
+
+@error_handling_decorator
+def load_stats(
+    conn: BigQueryClient,
+    tenant_id: str,
+    org_id: str,
+    qdc_client: qdc.QDCExternalAPIClient,
+    dataplex_stats_tables: list,
+) -> None:
+    logger.info("Loading statistics data.")
+    bigquery_table_stats(
+        bq_client=conn,
+        qdc_client=qdc_client,
+        tenant_id=tenant_id,
+        org_id=org_id,
+        dataplex_stats_tables=dataplex_stats_tables,
+    )
+    logger.info("Statistics data loaded successfully.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="Quollio Intelligence Agent for BigQuery",
+        description="Load lineage and stats to Quollio from BigQuery using Dataplex and BigQuery APIs",
+        epilog="Copyright (c) 2024 Quollio Technologies, Inc.",
+    )
+    parser.add_argument(
+        "commands",
+        choices=["load_lineage", "load_stats"],
+        type=str,
+        nargs="+",
+        help="""
+        The command to execute.
+        'load_lineage': Load lineage data from created views to Quollio,
+        'load_stats': Load stats from created views to Quollio,
+        """,
+    )
+    parser.add_argument(
+        "--log_level",
+        type=str,
+        choices=["debug", "info", "warn", "error", "none"],
+        action=env_default("LOG_LEVEL"),
+        default="info",
+        required=False,
+        help="The log level for dbt commands. Default value is info",
+    )
+    parser.add_argument(
+        "--tenant_id",
+        type=str,
+        action=env_default("TENANT_ID"),
+        required=False,
+        help="The tenant id (company id) where the lineage and stats are loaded",
+    )
+    parser.add_argument(
+        "--project_id",
+        type=str,
+        default=None,
+        required=False,
+        help="Project ID of the BigQuery project to load lineage and stats from (default is loaded from credentials)",
+    )
+    parser.add_argument(
+        "--regions",
+        type=str,
+        action=env_default("GCP_REGIONS"),
+        required=True,
+        help="Comma-separated list of regions BigQuery data is in",
+    )
+    parser.add_argument(
+        "--credentials_json",
+        type=str,
+        action=env_default("GOOGLE_APPLICATION_CREDENTIALS"),
+        required=True,
+        help="Credentials JSON",
+    )
+    parser.add_argument(
+        "--api_url",
+        type=str,
+        action=env_default("QDC_API_URL"),
+        required=False,
+        help="The base URL of Quollio External API",
+    )
+    parser.add_argument(
+        "--client_id",
+        type=str,
+        action=env_default("QDC_CLIENT_ID"),
+        required=False,
+        help="The client id that is created on Quollio console to let clients access Quollio External API",
+    )
+    parser.add_argument(
+        "--client_secret",
+        type=str,
+        action=env_default("QDC_CLIENT_SECRET"),
+        required=False,
+        help="The client secret that is created on Quollio console to let clients access Quollio External API",
+    )
+
+    parser.add_argument(
+        "--dataplex_stats_tables",
+        type=str,
+        action=env_default("DATAPLEX_STATS_TABLES"),
+        required=False,
+        help="Comma-separated list of dataplex stats tables - <project_id>.<dataset_id>.<table_id>",
+    )
+
+    args = parser.parse_args()
+
+    # Validate that dataplex_stats_tables is provided if load_stats is in commands
+    if "load_stats" in args.commands and not args.dataplex_stats_tables:
+        parser.error("--dataplex_stats_tables is required when 'load_stats' command is used")
+
+    configure_logging(args.log_level)
+
+    credentials = initialize_credentials(args.credentials_json)
+    org_id = initialize_org_id(args.credentials_json)
+    qdc_client = qdc.initialize_qdc_client(args.api_url, args.client_id, args.client_secret)
+    bq_client = initialize_bq_client(credentials, args.project_id)
+    if args.project_id is None:
+        args.project_id = json.loads(args.credentials_json)["project_id"]
+    regions = args.regions.split(",")
+
+    if "load_lineage" in args.commands:
+        load_lineage(
+            tenant_id=args.tenant_id,
+            project_id=args.project_id,
+            regions=regions,
+            org_id=org_id,
+            credentials=credentials,
+            qdc_client=qdc_client,
+        )
+
+    if "load_stats" in args.commands:
+        tables = args.dataplex_stats_tables.split(",")
+        load_stats(
+            conn=bq_client,
+            tenant_id=args.tenant_id,
+            org_id=org_id,
+            qdc_client=qdc_client,
+            dataplex_stats_tables=tables,
+        )

{quollio_core-0.4.12 → quollio_core-0.4.13}/quollio_core/bricks.py

@@ -10,6 +10,7 @@ from quollio_core.profilers.databricks import (
     databricks_column_stats,
     databricks_table_level_lineage,
 )
+from quollio_core.profilers.stats import get_column_stats_items
 from quollio_core.repository import databricks as db
 from quollio_core.repository import dbt, qdc
 
@@ -21,7 +22,6 @@ def build_view(
     target_tables: str = "",
     log_level: str = "info",
 ) -> None:
-
     logger.info("Build profiler views using dbt")
     # set parameters
     dbt_client = dbt.DBTClient()
@@ -64,7 +64,6 @@ def load_lineage(
     tenant_id: str,
     enable_column_lineage: bool = False,
 ) -> None:
-
     logger.info("Generate Databricks table to table lineage.")
     databricks_table_level_lineage(
         conn=conn,
@@ -98,7 +97,6 @@ def load_column_stats(
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
 ) -> None:
-
     logger.info("Generate Databricks column stats.")
     databricks_column_stats(
         conn=conn,
@@ -240,6 +238,19 @@ if __name__ == "__main__":
         help="Whether to ingest column lineage into QDIC or not. Default value is False",
     )
 
+    stats_items = get_column_stats_items()
+    parser.add_argument(
+        "--target_stats_items",
+        type=str,
+        nargs="*",
+        choices=stats_items,
+        default=stats_items,
+        action=env_default("DATABRICKS_STATS_ITEMS"),
+        required=False,
+        help="The items for statistic values.\
+              You can choose the items to be aggregated for stats. All items are selected by default.",
+    )
+
     args = parser.parse_args()
     set_log_level(level=args.log_level)
 
@@ -284,5 +295,6 @@ if __name__ == "__main__":
             endpoint=args.host,
             qdc_client=qdc_client,
             tenant_id=args.tenant_id,
+            stats_items=args.target_stats_items,
             monitoring_table_suffix=args.monitoring_table_suffix,
         )

quollio_core-0.4.13/quollio_core/helper/log_utils.py

@@ -0,0 +1,48 @@
+import inspect
+import logging
+
+LOG_LEVELS = {
+    "critical": logging.CRITICAL,
+    "error": logging.ERROR,
+    "warning": logging.WARNING,
+    "info": logging.INFO,
+    "debug": logging.DEBUG,
+    "notset": logging.NOTSET,
+}
+
+logger = logging.getLogger(__name__)
+
+
+def configure_logging(level: str = "INFO"):
+    """Configure logging settings."""
+    log_level = LOG_LEVELS.get(level.lower())
+    if log_level is None:
+        raise ValueError(f"Unknown log level: {level}")
+
+    logging.basicConfig(
+        level=log_level,
+        format="%(asctime)s - %(levelname)s - %(funcName)s:%(lineno)d - %(message)s",
+    )
+    logger.setLevel(log_level)
+    logger.info(f"Logging is configured to {level} level.")
+
+
+def error_handling_decorator(func):
+    """Decorator for consistent error handling in CLI commands."""
+
+    def wrapper(*args, **kwargs):
+        func_name = func.__name__
+        try:
+            logger.debug(f"Starting {func_name}")
+            result = func(*args, **kwargs)
+            logger.debug(f"Completed {func_name} successfully")
+            return result
+        except Exception as e:
+
+            current_frame = inspect.currentframe()
+            error_frame = current_frame.f_back
+            line_number = error_frame.f_lineno
+            logger.error(f"Error in {func_name} at line {line_number}: {str(e)}", exc_info=True)
+            raise
+
+    return wrapper

quollio_core-0.4.13/quollio_core/profilers/bigquery.py

@@ -0,0 +1,145 @@
+from typing import Dict, List
+
+from google.auth.credentials import Credentials
+
+from quollio_core.helper.log_utils import error_handling_decorator, logger
+from quollio_core.profilers.lineage import gen_table_lineage_payload, parse_bigquery_table_lineage
+from quollio_core.profilers.stats import gen_table_stats_payload
+from quollio_core.repository import qdc
+from quollio_core.repository.bigquery import BigQueryClient, GCPLineageClient, get_entitiy_reference, get_search_request
+
+
+@error_handling_decorator
+def bigquery_table_lineage(
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+    project_id: str,
+    regions: list,
+    org_id: str,
+    credentials: Credentials,
+) -> None:
+    lineage_client = GCPLineageClient(credentials)
+    bq_client = BigQueryClient(credentials, project_id)
+
+    datasets = bq_client.list_dataset_ids()
+    all_tables = generate_table_list(bq_client, datasets)
+    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)
+
+
+@error_handling_decorator
+def bigquery_table_stats(
+    qdc_client: qdc.QDCExternalAPIClient,
+    bq_client: BigQueryClient,
+    tenant_id: str,
+    org_id: str,
+    dataplex_stats_tables: list,
+) -> None:
+    profiling_results = []
+    for table in dataplex_stats_tables:
+        logger.info("Profiling columns using Dataplex stats table: %s", table)
+        profiling_results.extend(column_stats_from_dataplex(bq_client, table))
+
+    stats = gen_table_stats_payload(tenant_id, org_id, profiling_results)
+
+    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 column %s -> %s -> %s -> %s is successfully ingested.",
+                stat.db,
+                stat.schema,
+                stat.table,
+                stat.column,
+            )
+            logger.debug("Stats for column id %s is successfully ingested.", stat.global_id)
+
+
+def generate_table_list(bq_client: BigQueryClient, datasets: List[str]) -> List[str]:
+    all_tables = []
+    for dataset in datasets:
+        all_tables.extend(
+            [
+                table
+                for table in bq_client.list_tables(dataset)
+                if table["table_type"] in ["TABLE", "VIEW", "MATERIALIZED_VIEW"]
+            ],
+        )
+
+    all_table_names = []
+    for table in all_tables:
+        all_table_names.append(f"{bq_client.client.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:
+        if "quollio" in table.lower():
+            continue
+        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:", "")
+                source_table = str(lineage.source.fully_qualified_name).replace("bigquery:", "")
+                if target_table not in lineage_links:
+                    lineage_links[target_table] = []
+                if source_table not in lineage_links[target_table]:
+                    lineage_links[target_table].append(source_table)
+
+    return lineage_links
+
+
+def column_stats_from_dataplex(bq_client: BigQueryClient, profiling_table: str) -> List[Dict]:
+    query = f"""
+    SELECT
+        data_source.table_project_id AS DB_NAME,
+        data_source.dataset_id AS SCHEMA_NAME,
+        data_source.table_id AS TABLE_NAME,
+        column_name AS COLUMN_NAME,
+        min_value AS MIN_VALUE,
+        max_value AS MAX_VALUE,
+        average_value AS AVG_VALUE,
+        quartile_median AS MEDIAN_VALUE,
+        standard_deviation AS STDDEV_VALUE,
+        top_n[0][0] AS MODE_VALUE,
+        CAST((percent_null / 100) * job_rows_scanned AS INT) as NULL_COUNT,
+        CAST((percent_unique / 100) * job_rows_scanned AS INT) as CARDINALITY
+    FROM `{profiling_table}`
+    """
+    logger.debug(f"Executing Query: {query}")
+    results = bq_client.client.query(query).result()
+
+    # Convert RowIterator to a list of dictionaries
+    return [dict(row) for row in results]

{quollio_core-0.4.12 → quollio_core-0.4.13}/quollio_core/profilers/databricks.py

@@ -6,7 +6,7 @@ from quollio_core.profilers.lineage import (
     gen_table_lineage_payload,
     parse_databricks_table_lineage,
 )
-from quollio_core.profilers.stats import gen_table_stats_payload
+from quollio_core.profilers.stats import gen_table_stats_payload, get_is_target_stats_items, render_sql_for_stats
 from quollio_core.repository import databricks, qdc
 
 logger = logging.getLogger(__name__)
@@ -125,59 +125,63 @@ def _get_monitoring_tables(
 
 
 def _get_column_stats(
-    conn: databricks.DatabricksConnectionConfig, monitoring_table_suffix: str = "_profile_metrics"
+    conn: databricks.DatabricksConnectionConfig,
+    stats_items: List[str],
+    monitoring_table_suffix: str = "_profile_metrics",
 ) -> List[Dict[str, str]]:
     tables = _get_monitoring_tables(conn, monitoring_table_suffix)
     if not tables:
         return []
     stats = []
+    is_aggregate_items = get_is_target_stats_items(stats_items=stats_items)
     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 = """
-                    WITH profile_record_history AS (
-                        SELECT
-                            COLUMN_NAME
-                            , distinct_count as CARDINALITY
-                            , MAX as MAX_VALUE
-                            , MIN as MIN_VALUE
-                            , AVG as AVG_VALUE
-                            , MEDIAN as MEDIAN_VALUE
-                            , STDDEV as STDDEV_VALUE
-                            , NUM_NULLS as NULL_COUNT
-                            , get(frequent_items, 0).item AS MODE_VALUE
-                            , row_number() over(partition by column_name order by window desc) rownum
-                        FROM
-                            {monitoring_table}
-                        WHERE
-                            column_name not in (':table')
-                    )
-                    SELECT
-                        "{monitored_table_catalog}" as DB_NAME
-                        , "{monitored_table_schema}" as SCHEMA_NAME
-                        , "{monitored_table_name}" as TABLE_NAME
-                        , COLUMN_NAME
-                        , CARDINALITY
-                        , MAX_VALUE
-                        , MIN_VALUE
-                        , AVG_VALUE
-                        , MEDIAN_VALUE
-                        , STDDEV_VALUE
-                        , NULL_COUNT
-                        , MODE_VALUE
-                    FROM
-                        profile_record_history
-                    WHERE
-                        rownum = 1
-                """.format(
+            cte = """
+            WITH profile_record_history AS (
+                SELECT
+                    COLUMN_NAME
+                    , distinct_count as cardinality
+                    , MAX as max_value
+                    , MIN as min_value
+                    , AVG as avg_value
+                    , MEDIAN as median_value
+                    , STDDEV as stddev_value
+                    , NUM_NULLS as null_count
+                    , get(frequent_items, 0).item AS mode_value
+                    , row_number() over(partition by column_name order by window desc) rownum
+                FROM
+                    {monitoring_table}
+                WHERE
+                    column_name not in (':table')
+            ), profile_record AS (
+            SELECT
+                "{monitored_table_catalog}" as db_name
+                , "{monitored_table_schema}" as schema_name
+                , "{monitored_table_name}" as table_name
+                , column_name
+                , max_value
+                , min_value
+                , null_count
+                , cardinality
+                , avg_value
+                , median_value
+                , mode_value
+                , stddev_value
+            FROM
+                profile_record_history
+            WHERE
+                rownum = 1
+            )""".format(
                 monitoring_table=table["table_fqdn"],
                 monitored_table_catalog=monitored_table[0],
                 monitored_table_schema=monitored_table[1],
                 monitored_table_name=monitored_table[2],
             )
+            query = render_sql_for_stats(is_aggregate_items=is_aggregate_items, table_fqn="profile_record", cte=cte)
             logger.debug(f"The following sql will be fetched to retrieve stats values. {query}")
             stats.append(databricks_executor.get_query_results(query))
     return stats
@@ -188,9 +192,10 @@ def databricks_column_stats(
     endpoint: str,
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
+    stats_items: List[str],
     monitoring_table_suffix: str = "_profile_metrics",
 ) -> None:
-    table_stats = _get_column_stats(conn, monitoring_table_suffix)
+    table_stats = _get_column_stats(conn, stats_items, monitoring_table_suffix)
     for table in table_stats:
         logger.debug("Table %s will be aggregated.", table)
         stats = gen_table_stats_payload(tenant_id=tenant_id, endpoint=endpoint, stats=table)

{quollio_core-0.4.12 → quollio_core-0.4.13}/quollio_core/profilers/redshift.py

@@ -1,8 +1,13 @@
 import logging
+from typing import List
 
 from quollio_core.profilers.lineage import gen_table_lineage_payload, gen_table_lineage_payload_inputs
 from quollio_core.profilers.sqllineage import SQLLineage
-from quollio_core.profilers.stats import gen_table_stats_payload_from_tuple
+from quollio_core.profilers.stats import (
+    gen_table_stats_payload_from_tuple,
+    get_is_target_stats_items,
+    render_sql_for_stats,
+)
 from quollio_core.repository import qdc, redshift
 
 logger = logging.getLogger(__name__)
@@ -76,38 +81,24 @@ def redshift_table_stats(
     conn: redshift.RedshiftConnectionConfig,
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
+    stats_items: List[str],
 ) -> None:
-
+    is_aggregate_items = get_is_target_stats_items(stats_items=stats_items)
     with redshift.RedshiftQueryExecutor(config=conn) as redshift_executor:
         stats_query = _gen_get_stats_views_query(
             db=conn.database,
             schema=conn.schema,
         )
         stats_views = redshift_executor.get_query_results(query=stats_query)
+        logger.info("Found %s for table statistics.", len(stats_views))
 
         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],
+            table_fqn = "{catalog}.{schema}.{table}".format(
+                catalog=stats_view[0], schema=stats_view[1], table=stats_view[2]
             )
+            stats_query = render_sql_for_stats(is_aggregate_items=is_aggregate_items, table_fqn=table_fqn)
+            logger.debug(f"The following sql will be fetched to retrieve stats values. {stats_query}")
             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: