quollio-core 0.4.4__tar.gz → 0.4.5__tar.gz

{quollio_core-0.4.4 → quollio_core-0.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: quollio-core
-Version: 0.4.4
+Version: 0.4.5
 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>
@@ -21,14 +21,18 @@ Requires-Dist: blake3==0.3.3
 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: jinja2==3.1.3
 Requires-Dist: PyYAML==6.0.1
 Requires-Dist: requests==2.31.0
 Requires-Dist: pyjwt==2.8.0
 Requires-Dist: redshift-connector==2.0.915
 Requires-Dist: snowflake-connector-python==3.5.0
+Requires-Dist: databricks-sdk==0.17.0
+Requires-Dist: databricks-sql-connector==2.9.5
 Requires-Dist: sqlglot==20.8.0
 Requires-Dist: black>=22.3.0 ; extra == "test"
+Requires-Dist: coverage>=7.3.2 ; extra == "test"
 Requires-Dist: isort>=5.10.1 ; extra == "test"
 Requires-Dist: pyproject-flake8>=0.0.1-alpha.2 ; extra == "test"
 Requires-Dist: pytest>=5.2 ; extra == "test"

{quollio_core-0.4.4 → quollio_core-0.4.5}/pyproject.toml

@@ -33,12 +33,15 @@ dependencies = [
     ,"dbt-core==1.7.10"
     ,"dbt-snowflake==1.7.0"
     ,"dbt-redshift==1.7.1"
+    ,"dbt-databricks==1.7.1"
     ,"jinja2==3.1.3"
     ,"PyYAML==6.0.1"
     ,"requests==2.31.0"
     ,"pyjwt==2.8.0"
     ,"redshift-connector==2.0.915"
     ,"snowflake-connector-python==3.5.0"
+    ,"databricks-sdk==0.17.0"
+    ,"databricks-sql-connector==2.9.5"
     ,"sqlglot==20.8.0"
 ]
 dynamic = ["version", "description"]
@@ -50,6 +53,7 @@ Home = "https://quollio.com"
 [project.optional-dependencies]
 test = [
   "black>=22.3.0"
+  ,"coverage>=7.3.2"
   ,"isort>=5.10.1"
   ,"pyproject-flake8>=0.0.1-alpha.2"
   ,"pytest>=5.2"

{quollio_core-0.4.4 → quollio_core-0.4.5}/quollio_core/__init__.py

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

quollio_core-0.4.5/quollio_core/bricks.py

@@ -0,0 +1,237 @@
+import argparse
+import logging
+import os
+
+from quollio_core.helper.core import setup_dbt_profile
+from quollio_core.helper.env_default import env_default
+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:
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
+
+    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, "--select", target_tables]
+    dbt_client.invoke(
+        cmd="run",
+        project_dir=project_path,
+        profile_dir=template_path,
+        options=run_options,
+    )
+    return
+
+
+def load_lineage(
+    conn: db.DatabricksConnectionConfig,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+) -> None:
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
+
+    logger.info("Generate Databricks table to table lineage.")
+    databricks_table_level_lineage(
+        conn=conn, qdc_client=qdc_client, tenant_id=tenant_id, dbt_table_name="quollio_lineage_table_level"
+    )
+
+    logger.info("Generate Databricks column to column lineage.")
+    databricks_column_level_lineage(
+        conn=conn, qdc_client=qdc_client, tenant_id=tenant_id, dbt_table_name="quollio_lineage_column_level"
+    )
+
+    logger.info("Lineage data is successfully loaded.")
+    return
+
+
+def load_column_stats(
+    conn: db.DatabricksConnectionConfig,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+) -> None:
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
+
+    logger.info("Generate Databricks column stats.")
+    databricks_column_stats(
+        conn=conn,
+        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,
+        'load_sqllineage': Load lineage data from sql parse result(alpha),
+        """,
+    )
+    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_view_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",
+    )
+
+    args = parser.parse_args()
+
+    conn = db.DatabricksConnectionConfig(
+        host=args.host,
+        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, qdc_client=qdc_client, tenant_id=args.tenant_id)
+
+    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, qdc_client=qdc_client, tenant_id=args.tenant_id)

quollio_core-0.4.5/quollio_core/dbt_projects/databricks/.gitignore

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

quollio_core-0.4.5/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-0.4.5/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-0.4.5/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-0.4.5/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-0.4.5/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-0.4.5/quollio_core/dbt_projects/databricks/models/quollio_lineage_table_level.yml

@@ -0,0 +1,11 @@
+version: 2
+
+model:
+  - name: quollio_lineage_column_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

quollio_core-0.4.5/quollio_core/dbt_projects/databricks/models/sources.yml

@@ -0,0 +1,84 @@
+version: 2
+
+sources: 
+  - name: access
+    database: system
+    schema: access
+    tables: 
+      - name: table_lineage
+        description: Describes table level lineage
+        columns:
+          - name: source_table_full_name
+            description: ''
+            type: string
+          - name: target_table_full_name
+            description: ''
+            type: string
+          - name: target_type
+            description: ''
+            type: string
+          - name: event_time
+            description: ''
+            type: timestamp
+
+      - name: column_lineage
+        description: Describes column level lineage
+        columns:
+          - name: source_table_full_name
+            description: ''
+            type: string
+          - name: target_table_full_name
+            description: ''
+            type: string
+          - name: event_time
+            description: ''
+            type: timestamp
+          - name: source_column_name
+            description: ''
+            type: string
+          - name: target_column_name
+            description: ''
+            type: string
+  
+  - name: inf_sch
+    database: system
+    schema: information_schema
+    tables: 
+      - name: tables
+        description: Lists existing tables (i.e., not deleted).
+        columns:
+          - name: table_catalog
+            description: ''
+            type: string
+          - name: table_schema
+            description: ''
+            type: string
+          - name: table_name
+            description: ''
+            type: string
+
+      - name: views
+        description: Lists existing views (i.e., not deleted). Views are treated as tables.
+        columns:
+          - name: table_catalog
+            description: ''
+            type: string
+          - name: table_schema
+            description: ''
+            type: string
+          - name: table_name
+            description: ''
+            type: string
+
+      - name: columns
+        description: ''
+        columns:
+          - name: table_catalog
+            description: ''
+            type: string
+          - name: table_schema
+            description: ''
+            type: string
+          - name: table_name
+            description: ''
+            type: string

quollio_core-0.4.5/quollio_core/dbt_projects/databricks/package-lock.yml

@@ -0,0 +1,14 @@
+packages:
+- package: dbt-labs/dbt_utils
+  version: 1.1.1
+- package: dbt-labs/spark_utils
+  version: 0.3.0
+- package: dbt-labs/codegen
+  version: 0.12.1
+- package: dbt-labs/dbt_external_tables
+  version: 0.8.7
+- package: dbt-labs/dbt_project_evaluator
+  version: 0.8.1
+- package: brooklyn-data/dbt_artifacts
+  version: 2.6.2
+sha1_hash: cbb324267dbf6c6fb7de11b162e4fbafd1e32a9c

quollio_core-0.4.5/quollio_core/dbt_projects/databricks/packages.yml

@@ -0,0 +1,13 @@
+packages:
+  - package: dbt-labs/dbt_utils
+    version: [">=0.0.0", "<2.0.0"]
+  - package: dbt-labs/spark_utils
+    version: [">=0.0.0", "<1.0.0"]
+  - package: dbt-labs/codegen
+    version: [">=0.0.0", "<1.0.0"]
+  - package: dbt-labs/dbt_external_tables
+    version: [">=0.0.0", "<1.0.0"]
+  - package: dbt-labs/dbt_project_evaluator
+    version: [">=0.0.0", "<1.0.0"]
+  - package: brooklyn-data/dbt_artifacts
+    version: [">=2.0.0", "<3.0.0"]

quollio_core-0.4.5/quollio_core/dbt_projects/databricks/profiles/profiles_template.yml

@@ -0,0 +1,14 @@
+quollio_intelligence_databricks:
+  target: project
+  outputs:
+    project:
+      type: databricks
+      host: {{ host }}
+      http_path: {{ http_path }}
+      catalog: {{ catalog }}
+      schema: {{ schema }}
+      auth_type: oauth
+      client_id: {{ client_id }}
+      client_secret: {{ client_secret }}
+      databricks_port: {{ databricks_port }}
+