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

quollio_core/__init__.py

@@ -1,4 +1,4 @@
 """Quollio Core"""
 
-__version__ = "0.4.4"
+__version__ = "0.4.10"
 __author__ = "Quollio Technologies, Inc"

quollio_core/bigquery.py

@@ -0,0 +1,123 @@
+import argparse
+import json
+import logging
+
+from quollio_core.helper.env_default import env_default
+from quollio_core.helper.log import set_log_level
+from quollio_core.profilers.bigquery import bigquery_table_lineage
+from quollio_core.repository import qdc
+from quollio_core.repository.bigquery import get_credentials, get_org_id
+
+logger = logging.getLogger(__name__)
+
+
+def load_lineage(
+    qdc_client: qdc.QDCExternalAPIClient, project_id: str, regions: list, tenant_id: str, credentials: dict, org_id: str
+):
+    bigquery_table_lineage(
+        qdc_client=qdc_client,
+        tenant_id=tenant_id,
+        project_id=project_id,
+        regions=regions,
+        credentials=credentials,
+        org_id=org_id,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="Quollio Intelligence Agent for Google BigQuery",
+        description="Collect lineage and stats from Google BigQuery and load to Quollio Data Catalog",
+        epilog="Copyright (c) 2024 Quollio Technologies, Inc.",
+    )
+    parser.add_argument(
+        "commands",
+        choices=["load_lineage"],
+        type=str,
+        nargs="+",
+        help="""
+        The command to execute.
+        'load_lineage': Load lineage data from Google Data Catalog to Quollio,
+        """,
+    )
+    parser.add_argument(
+        "--credentials",
+        type=str,
+        action=env_default("GOOGLE_APPLICATION_CREDENTIALS"),
+        help="Crendentials for Google Cloud Platform",
+    )
+    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(
+        "--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(
+        "--project_id",
+        type=str,
+        action=env_default("GCP_PROJECT_ID"),
+        required=False,
+        help="GCP Project ID",
+    )
+    parser.add_argument(
+        "--regions",
+        type=str,
+        action=env_default("GCP_REGIONS"),
+        required=False,
+        help="GCP regions where the data is located. Multiple regions can be provided separated by space.",
+        nargs="+",
+    )
+    parser.add_argument(
+        "--log_level",
+        type=str,
+        choices=["debug", "info", "warn", "error", "none"],
+        action=env_default("LOG_LEVEL"),
+        required=False,
+        help="The log level for dbt commands. Default value is info",
+    )
+
+    args = parser.parse_args()
+    set_log_level(level=args.log_level)
+
+    if len(args.commands) == 0:
+        raise ValueError("No command is provided")
+
+    if "load_lineage" in args.commands:
+        qdc_client = qdc.QDCExternalAPIClient(
+            base_url=args.api_url, client_id=args.client_id, client_secret=args.client_secret
+        )
+
+        credentials_json = json.loads(args.credentials)
+        credentials = get_credentials(credentials_json=credentials_json)
+        org_id = get_org_id(credentials_json=credentials_json)
+
+        load_lineage(
+            qdc_client=qdc_client,
+            project_id=args.project_id,
+            regions=args.regions,
+            tenant_id=args.tenant_id,
+            credentials=credentials,
+            org_id=org_id,
+        )

quollio_core/bricks.py

@@ -0,0 +1,288 @@
+import argparse
+import logging
+import os
+
+from quollio_core.helper.core import setup_dbt_profile, trim_prefix
+from quollio_core.helper.env_default import env_default
+from quollio_core.helper.log import set_log_level
+from quollio_core.profilers.databricks import (
+    databricks_column_level_lineage,
+    databricks_column_stats,
+    databricks_table_level_lineage,
+)
+from quollio_core.repository import databricks as db
+from quollio_core.repository import dbt, qdc
+
+logger = logging.getLogger(__name__)
+
+
+def build_view(
+    conn: db.DatabricksConnectionConfig,
+    target_tables: str = "",
+    log_level: str = "info",
+) -> None:
+
+    logger.info("Build profiler views using dbt")
+    # set parameters
+    dbt_client = dbt.DBTClient()
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    project_path = f"{current_dir}/dbt_projects/databricks"
+    template_path = f"{current_dir}/dbt_projects/databricks/profiles"
+    template_name = "profiles_template.yml"
+
+    # build views using dbt
+    setup_dbt_profile(connections_json=conn.as_dict(), template_path=template_path, template_name=template_name)
+    # FIXME: when executing some of the commands, directory changes due to the library bug.
+    # https://github.com/dbt-labs/dbt-core/issues/8997
+    dbt_client.invoke(
+        cmd="deps",
+        project_dir=project_path,
+        profile_dir=template_path,
+        options=["--no-use-colors", "--log-level", log_level],
+    )
+
+    run_options = ["--no-use-colors", "--log-level", log_level]
+
+    if target_tables is not None:
+        target_tables_str = " ".join(target_tables)
+        run_options.append("--select")
+        run_options.append(target_tables_str)
+
+    dbt_client.invoke(
+        cmd="run",
+        project_dir=project_path,
+        profile_dir=template_path,
+        options=run_options,
+    )
+    return
+
+
+def load_lineage(
+    conn: db.DatabricksConnectionConfig,
+    endpoint: str,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+    enable_column_lineage: bool = False,
+) -> None:
+
+    logger.info("Generate Databricks table to table lineage.")
+    databricks_table_level_lineage(
+        conn=conn,
+        endpoint=endpoint,
+        qdc_client=qdc_client,
+        tenant_id=tenant_id,
+        dbt_table_name="quollio_lineage_table_level",
+    )
+
+    if enable_column_lineage:
+        logger.info(
+            f"enable_column_lineage is set to {enable_column_lineage}.Generate Databricks column to column lineage."
+        )
+        databricks_column_level_lineage(
+            conn=conn,
+            endpoint=endpoint,
+            qdc_client=qdc_client,
+            tenant_id=tenant_id,
+            dbt_table_name="quollio_lineage_column_level",
+        )
+    else:
+        logger.info("Skip column lineage ingestion. Set enable_column_lineage to True if you ingest column lineage.")
+
+    logger.info("Lineage data is successfully loaded.")
+    return
+
+
+def load_column_stats(
+    conn: db.DatabricksConnectionConfig,
+    endpoint: str,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+) -> None:
+
+    logger.info("Generate Databricks column stats.")
+    databricks_column_stats(
+        conn=conn,
+        endpoint=endpoint,
+        qdc_client=qdc_client,
+        tenant_id=tenant_id,
+    )
+
+    logger.info("Column stats are successfully loaded.")
+    return
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="Quollio Intelligence Agent for Databricks",
+        description="Build views and load lineage and stats to Quollio from Databricks using dbt.",
+        epilog="Copyright (c) 2024 Quollio Technologies, Inc.",
+    )
+    parser.add_argument(
+        "commands",
+        choices=["build_view", "load_lineage", "load_stats"],
+        type=str,
+        nargs="+",
+        help="""
+        The command to execute.
+        'build_view': Build views using dbt,
+        'load_lineage': Load lineage data from created views to Quollio,
+        'load_stats': Load stats from created views to Quollio,
+        """,
+    )
+    parser.add_argument(
+        "--host", type=str, action=env_default("DATABRICKS_HOST"), required=False, help="Host for Databricks workspace"
+    )
+    parser.add_argument(
+        "--http_path",
+        type=str,
+        action=env_default("DATABRICKS_HTTP_PATH"),
+        required=False,
+        help="HTTP path for a Databricks compute resource (i.e warehouse)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        action=env_default("DATABRICKS_PORT"),
+        required=False,
+        help="Port for Databricks compute resource",
+    )
+    parser.add_argument(
+        "--databricks_client_secret",
+        type=str,
+        action=env_default("DATABRICKS_CLIENT_SECRET"),
+        required=False,
+        help="Secret for the service principal",
+    )
+    parser.add_argument(
+        "--databricks_client_id",
+        type=str,
+        action=env_default("DATABRICKS_CLIENT_ID"),
+        required=False,
+        help="Client id for the service principal",
+    )
+    parser.add_argument(
+        "--catalog",
+        type=str,
+        required=False,
+        action=env_default("DATABRICKS_TARGET_CATALOG"),
+        help="Target database name where the views are built by dbt",
+    )
+    parser.add_argument(
+        "--schema",
+        type=str,
+        action=env_default("DATABRICKS_TARGET_SCHEMA"),
+        required=False,
+        help="Target schema name where the views are built by dbt",
+    )
+    parser.add_argument(
+        "--log_level",
+        type=str,
+        choices=["debug", "info", "warn", "error", "none"],
+        action=env_default("LOG_LEVEL"),
+        required=False,
+        help="The log level for dbt commands. Default value is info",
+    )
+    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 secrete that is created on Quollio console to let clients access Quollio External API",
+    )
+    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(
+        "--target_tables",
+        type=str,
+        nargs="+",
+        choices=["quollio_lineage_table_level", "quollio_lineage_column_level"],
+        action=env_default("DATABRICKS_TARGET_TABLES"),
+        required=False,
+        help="Target tables you want to create with dbt module. \
+              You need to specify this parameter if you want to specify tables, not all ones. \
+              Please specify table name with blank delimiter like tableA tableB \
+              if you want to create two or more tables",
+    )
+    parser.add_argument(
+        "--monitoring_table_suffix",
+        type=str,
+        action=env_default("DATABRICKS_MONITORING_TABLE_SUFFIX"),
+        required=False,
+        help="Sets the monitoring tables suffix for databricks. \
+              This is used to identify the monitoring tables created by the databricks monitoring tool. \
+              Default value is _profile_metrics",
+    )
+    parser.add_argument(
+        "--enable_column_lineage",
+        type=bool,
+        action=env_default("ENABLE_COLUMN_LINEAGE", store_true=True),
+        default=False,
+        required=False,
+        help="Whether to ingest column lineage into QDIC or not. Default value is False",
+    )
+
+    args = parser.parse_args()
+    set_log_level(level=args.log_level)
+
+    conn = db.DatabricksConnectionConfig(
+        # MEMO: Metadata agent allows the string 'https://' as a host name but is not allowed by intelligence agent.
+        host=trim_prefix(args.host, "https://"),
+        http_path=args.http_path,
+        client_id=args.databricks_client_id,
+        client_secret=args.databricks_client_secret,
+        catalog=args.catalog,
+        schema=args.schema,
+    )
+
+    if len(args.commands) == 0:
+        raise ValueError("No command is provided")
+
+    if "build_view" in args.commands:
+        build_view(
+            conn=conn,
+            target_tables=args.target_tables,
+            log_level=args.log_level,
+        )
+
+    if "load_lineage" in args.commands:
+        qdc_client = qdc.QDCExternalAPIClient(
+            base_url=args.api_url, client_id=args.client_id, client_secret=args.client_secret
+        )
+        load_lineage(
+            conn=conn,
+            endpoint=args.host,
+            qdc_client=qdc_client,
+            tenant_id=args.tenant_id,
+            enable_column_lineage=args.enable_column_lineage,
+        )
+
+    if "load_stats" in args.commands:
+        qdc_client = qdc.QDCExternalAPIClient(
+            base_url=args.api_url, client_id=args.client_id, client_secret=args.client_secret
+        )
+        databricks_column_stats(
+            conn=conn,
+            endpoint=args.host,
+            qdc_client=qdc_client,
+            tenant_id=args.tenant_id,
+            monitoring_table_suffix=args.monitoring_table_suffix,
+        )

quollio_core/dbt_projects/databricks/.gitignore

@@ -0,0 +1,4 @@
+
+target/
+dbt_packages/
+logs/

quollio_core/dbt_projects/databricks/README.md

@@ -0,0 +1,5 @@
+###  Quollio Intelligence Agent Support For Databricks
+Notable Files:
+1. [quollio_lineage_table_level.sql](models/quollio_lineage_table_level.sql) - Generates table lineage data from Databricks system tables.
+2. [quollio_lineage_column_level.sql](models/quollio_lineage_table_level.sql) - Generates column lineage data from Databricks system tables.
+3. [sources.yml](models/sources.yml) - Refrences sources in the Databricks system catalog.

quollio_core/dbt_projects/databricks/dbt_project.yml

@@ -0,0 +1,21 @@
+name: 'quollio_intelligence_databricks'
+version: '1.0.0'
+config-version: 2
+
+profile: 'quollio_intelligence_databricks'
+
+model-paths: ["models"]
+analysis-paths: ["analyses"]
+test-paths: ["tests"]
+seed-paths: ["seeds"]
+macro-paths: ["macros"]
+snapshot-paths: ["snapshots"]
+
+clean-targets:
+  - "target"
+  - "dbt_packages"
+
+models:
+  +dbt-osmosis: "{model}.yml"
+  # Databricks automatically enables grants on SQL endpoints
+  # https://docs.getdbt.com/reference/resource-configs/grants

quollio_core/dbt_projects/databricks/models/quollio_lineage_column_level.sql

@@ -0,0 +1,73 @@
+-- Gets  full table lineage from Databricks
+WITH columns_lineage_history AS (
+  SELECT
+    -- The databricks columns table does not have a full table name, create with CONCAT()
+    source_table_full_name AS upstream_table,
+    target_table_full_name as downstream_table,
+    source_column_name as upstream_column,
+    target_column_name as downstream_column,
+    event_time,
+    RANK() OVER (
+      PARTITION BY target_table_full_name
+      ORDER BY
+        event_time DESC
+    ) AS rank
+  FROM
+    {{ source('access','column_lineage')  }}
+  WHERE
+    source_table_full_name IS NOT NULL
+    AND target_table_full_name IS NOT NULL
+    AND source_table_full_name NOT LIKE "%quollio%"
+    AND target_table_full_name NOT LIKE "%quollio%"
+),
+-- Gets list of existing columns in catalogs
+existing_columns (
+  SELECT
+    CONCAT(table_catalog, '.', table_schema, '.', table_name) AS table_full_name,
+    column_name
+  FROM
+   {{ source('inf_sch','columns')  }}
+),
+
+-- Checks if the downstream tables exists and group operations.
+downstream_column_exists (
+  SELECT
+    upstream_table AS UPSTREAM_TABLE_NAME,
+    upstream_column AS UPSTREAM_COLUMN_NAME,
+    downstream_table AS DOWNSTREAM_TABLE_NAME,
+    downstream_column AS DOWNSTREAM_COLUMN_NAME,
+    event_time
+  FROM
+    columns_lineage_history clh
+    INNER JOIN existing_columns ec ON clh.downstream_table = ec.table_full_name
+    AND clh.downstream_column = ec.column_name
+  WHERE
+    rank = 1
+  GROUP BY UPSTREAM_TABLE, UPSTREAM_COLUMN, DOWNSTREAM_TABLE, DOWNSTREAM_COLUMN, EVENT_TIME
+),
+
+-- Aggregates the column lineage
+aggregated_column_lineage AS (
+  SELECT
+      downstream_table_name,
+      downstream_column_name,
+      collect_set(
+          named_struct(
+              'upstream_table_name', upstream_table_name,
+              'upstream_column_name', upstream_column_name
+          )
+      ) AS upstream_columns
+  FROM
+    downstream_column_exists
+  GROUP BY
+      downstream_table_name,
+      downstream_column_name
+)
+
+SELECT
+  downstream_table_name AS DOWNSTREAM_TABLE_NAME,
+  downstream_column_name AS DOWNSTREAM_COLUMN_NAME,
+  to_json(upstream_columns) AS UPSTREAM_COLUMNS
+FROM
+  aggregated_column_lineage
+  

quollio_core/dbt_projects/databricks/models/quollio_lineage_column_level.yml

@@ -0,0 +1,14 @@
+version: 2
+
+model:
+  - name: quollio_lineage_column_level
+    columns: 
+      - name: UPSTREAM_COLUMNS
+        description: 'String column with all upstream columns in JSON format'
+        type: string
+      - name: DOWNSTREAM_TABLE_NAME
+        description: 'Full downstream table name in <catalog>.<schema>.<table> format'
+        type: string
+      - name: DOWNSTREAM_COLUMN_NAME
+        description: 'Downstream column name'
+        type: string

quollio_core/dbt_projects/databricks/models/quollio_lineage_table_level.sql

@@ -0,0 +1,63 @@
+-- Gets  full table lineage from Databricks
+WITH table_lineage_history AS (
+  SELECT
+    source_table_full_name as upstream_table,
+    target_table_full_name as downstream_table,
+    target_type,
+    event_time,
+    RANK() OVER (
+      PARTITION BY target_table_full_name
+      ORDER BY
+        event_time DESC
+    ) AS rank
+  FROM
+    {{ source('access','table_lineage')  }}
+  WHERE
+    source_table_full_name IS NOT NULL
+    AND target_table_full_name IS NOT NULL
+    AND source_table_full_name NOT LIKE "%quollio%"
+    AND target_table_full_name NOT LIKE "%quollio%"
+),
+-- Gets list of existing tables in catalogs
+existing_tables (
+  SELECT
+    CONCAT(table_catalog, '.', table_schema, '.', table_name) AS table_full_name
+  FROM
+   {{ source('inf_sch','tables')  }}
+),
+
+-- Checks if the downstream tables exists and group operations.
+downstream_table_exists (
+  SELECT
+    upstream_table,
+    downstream_table,
+    target_type,
+    event_time
+  FROM
+    table_lineage_history tlh
+    INNER JOIN existing_tables et ON tlh.downstream_table = et.table_full_name
+  WHERE
+    rank = 1
+  GROUP BY upstream_table, downstream_table, target_type, event_time
+),
+
+aggregated_table_lineage AS (
+  SELECT
+      downstream_table,
+      collect_set(
+          named_struct(
+              'upstream_object_name', upstream_table
+          )
+      ) AS upstream_tables
+  FROM
+    downstream_table_exists
+  GROUP BY
+      downstream_table
+)
+SELECT
+  downstream_table as DOWNSTREAM_TABLE_NAME,
+  to_json(upstream_tables) as UPSTREAM_TABLES
+
+FROM
+  aggregated_table_lineage
+  

quollio_core/dbt_projects/databricks/models/quollio_lineage_table_level.yml

@@ -0,0 +1,11 @@
+version: 2
+
+model:
+  - name: quollio_lineage_table_level
+    columns: 
+      - name: UPSTREAM_TABLES
+        description: 'String column with all upstream tables in JSON format'
+        type: string
+      - name: DOWNSTREAM_TABLE_NAME
+        description: 'Full downstream table name in <catalog>.<schema>.<table> format'
+        type: string