quollio-core 0.4.6__tar.gz → 0.4.7__tar.gz

{quollio_core-0.4.6 → quollio_core-0.4.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: quollio-core
-Version: 0.4.6
+Version: 0.4.7
 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>
@@ -31,6 +31,11 @@ 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: google-cloud==0.34.0
+Requires-Dist: google-cloud-bigquery==3.22.0
+Requires-Dist: google-cloud-datacatalog==3.19.0
+Requires-Dist: google-cloud-datacatalog-lineage==0.3.6
+Requires-Dist: google-api-python-client==2.131.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"

{quollio_core-0.4.6 → quollio_core-0.4.7}/pyproject.toml

@@ -43,6 +43,11 @@ dependencies = [
     ,"databricks-sdk==0.17.0"
     ,"databricks-sql-connector==2.9.5"
     ,"sqlglot==20.8.0"
+    ,"google-cloud==0.34.0"
+    ,"google-cloud-bigquery==3.22.0"
+    ,"google-cloud-datacatalog==3.19.0"
+    ,"google-cloud-datacatalog-lineage==0.3.6"
+    ,"google-api-python-client==2.131.0"
 ]
 dynamic = ["version", "description"]
 

{quollio_core-0.4.6 → quollio_core-0.4.7}/quollio_core/__init__.py

@@ -1,4 +1,4 @@
 """Quollio Core"""
 
-__version__ = "0.4.6"
+__version__ = "0.4.7"
 __author__ = "Quollio Technologies, Inc"

quollio_core-0.4.7/quollio_core/bigquery.py

@@ -0,0 +1,114 @@
+import argparse
+import json
+import logging
+
+from quollio_core.helper.env_default import env_default
+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="+",
+    )
+
+    args = parser.parse_args()
+
+    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-0.4.6 → quollio_core-0.4.7}/quollio_core/bricks.py

@@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
 
 def build_view(
     conn: db.DatabricksConnectionConfig,
-    target_tables: str,
+    target_tables: str = "",
     log_level: str = "info",
 ) -> None:
     logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
@@ -41,7 +41,13 @@ def build_view(
         options=["--no-use-colors", "--log-level", log_level],
     )
 
-    run_options = ["--no-use-colors", "--log-level", log_level, "--select", target_tables]
+    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,
@@ -106,7 +112,6 @@ if __name__ == "__main__":
         '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(
@@ -193,8 +198,8 @@ if __name__ == "__main__":
     parser.add_argument(
         "--target_tables",
         type=str,
-        nargs="*",
-        choices=["quollio_lineage_table_level", "quollio_lineage_view_level"],
+        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. \
@@ -203,6 +208,16 @@ if __name__ == "__main__":
               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",
+    )
+
     args = parser.parse_args()
 
     conn = db.DatabricksConnectionConfig(
@@ -234,4 +249,9 @@ if __name__ == "__main__":
         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)
+        databricks_column_stats(
+            conn=conn,
+            qdc_client=qdc_client,
+            tenant_id=args.tenant_id,
+            monitoring_table_suffix=args.monitoring_table_suffix,
+        )

{quollio_core-0.4.6 → quollio_core-0.4.7}/quollio_core/dbt_projects/snowflake/macros/materialization/divided_view.sql

@@ -24,6 +24,7 @@ SELECT * FROM {{  ref('quollio_stats_profiling_columns')  }} WHERE NOT startswit
   {%- set build_sql %}
   {%- for record in records[i: i+chunk] -%}
     {%- if not loop.first %}UNION{% endif %}
+
     SELECT
       DISTINCT
       '{{record[0]}}' as db_name
@@ -38,7 +39,7 @@ SELECT * FROM {{  ref('quollio_stats_profiling_columns')  }} WHERE NOT startswit
       , {% if record[5] == true %}median("{{record[3]}}"){% else %}null{% endif %} AS median_value
       , {% if record[5] == true %}approx_top_k("{{record[3]}}")[0][0]{% else %}null{% endif %} AS mode_value
       , {% if record[5] == true %}stddev("{{record[3]}}"){% else %}null{% endif %} AS stddev_value
-    FROM {{ record[0] }}.{{ record[1] }}.{{ record[2] }} {{ var("sample_method") }}
+    FROM "{{record[0]}}"."{{record[1]}}"."{{record[2]}}" {{ var("sample_method") }}
   {% endfor -%}
   {%- endset %}
   -- create a view with a index as suffix

quollio_core-0.4.7/quollio_core/dbt_projects/snowflake/models/quollio_stats_profiling_columns.sql

@@ -0,0 +1,96 @@
+WITH columns AS (
+    SELECT
+      table_catalog
+      , table_schema
+      , table_name
+      , column_name
+      , data_type
+    FROM
+      {{ source('account_usage', 'COLUMNS') }}
+    WHERE
+      deleted is null
+      AND table_name NOT LIKE 'QUOLLIO_%%'
+    GROUP BY
+      table_catalog
+      , table_schema
+      , table_name
+      , column_name
+      , data_type
+    ORDER BY
+      table_catalog
+      , table_schema
+      , table_name
+), accessible_tables AS (
+    SELECT
+      table_catalog
+      , table_schema
+      , name
+    FROM
+      {{ source('account_usage', 'GRANTS_TO_ROLES') }}
+    WHERE
+      granted_on in ('TABLE', 'MATERIALIZED VIEW')
+      AND grantee_name = '{{ var("query_role") }}'
+      AND privilege in ('SELECT', 'OWNERSHIP', 'REFERENCES')
+      AND deleted_on IS NULL
+    GROUP BY
+      table_catalog
+      , table_schema
+      , name  
+), m_view_sys_columns AS (
+  SELECT
+    cols.table_catalog
+    , cols.table_schema
+    , cols.table_name
+    , cols.column_name
+    , cols.data_type
+  FROM
+    {{ source('account_usage', 'COLUMNS') }} cols
+  LEFT OUTER JOIN
+    {{ source('account_usage', 'TABLES') }} tbls
+  ON
+    cols.table_catalog = tbls.table_catalog
+    AND cols.table_schema = tbls.table_schema
+    AND cols.table_name = tbls.table_name
+  WHERE
+    tbls.table_type = 'MATERIALIZED VIEW'
+    AND cols.column_name = 'SYS_MV_SOURCE_PARTITION'
+), implicit_columns_removed AS (
+  SELECT
+    c.table_catalog
+    , c.table_schema
+    , c.table_name
+    , c.column_name
+    , c.data_type
+  FROM
+    columns c
+  INNER JOIN
+    accessible_tables a
+  ON
+    c.table_catalog = a.table_catalog
+    AND c.table_schema = a.table_schema
+    AND c.table_name = a.name
+  MINUS
+  SELECT
+    table_catalog
+    , table_schema
+    , table_name
+    , column_name
+    , data_type
+  FROM
+    m_view_sys_columns
+), final AS (
+  SELECT
+    table_catalog
+    , table_schema
+    , table_name
+    , column_name
+    , data_type
+    , case when data_type in('NUMBER','DECIMAL', 'DEC', 'NUMERIC',
+                             'INT', 'INTEGER', 'BIGINT', 'SMALLINT',
+                             'TINYINT', 'BYTEINT')
+                             THEN true
+           else false END AS is_calculable
+  FROM
+    implicit_columns_removed
+)
+select * from final

quollio_core-0.4.7/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-0.4.6 → quollio_core-0.4.7}/quollio_core/profilers/databricks.py

@@ -99,7 +99,7 @@ def databricks_column_level_lineage(
 
 
 def _get_monitoring_tables(
-    conn: databricks.DatabricksConnectionConfig, monitoring_table_id: str = "_profile_metrics"
+    conn: databricks.DatabricksConnectionConfig, monitoring_table_suffix: str = "_profile_metrics"
 ) -> List[Dict[str, str]]:
     tables = []
     query = f"""
@@ -110,7 +110,7 @@ def _get_monitoring_tables(
             CONCAT(table_catalog, '.', table_schema, '.', table_name) AS table_fqdn
         FROM
             system.information_schema.tables
-        WHERE table_name LIKE "%{monitoring_table_id}"
+        WHERE table_name LIKE "%{monitoring_table_suffix}"
         """
     with databricks.DatabricksQueryExecutor(config=conn) as databricks_executor:
         tables = databricks_executor.get_query_results(query)
@@ -123,9 +123,9 @@ def _get_monitoring_tables(
 
 
 def _get_column_stats(
-    conn: databricks.DatabricksConnectionConfig, monitoring_table_id: str = "_profile_metrics"
+    conn: databricks.DatabricksConnectionConfig, monitoring_table_suffix: str = "_profile_metrics"
 ) -> List[Dict[str, str]]:
-    tables = _get_monitoring_tables(conn, monitoring_table_id)
+    tables = _get_monitoring_tables(conn, monitoring_table_suffix)
     if not tables:
         return []
     stats = []
@@ -178,9 +178,9 @@ def databricks_column_stats(
     conn: databricks.DatabricksConnectionConfig,
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
-    monitoring_table_id: str = "_profile_metrics",
+    monitoring_table_suffix: str = "_profile_metrics",
 ) -> None:
-    table_stats = _get_column_stats(conn, monitoring_table_id)
+    table_stats = _get_column_stats(conn, monitoring_table_suffix)
     for table in table_stats:
         stats = gen_table_stats_payload(tenant_id, conn.host, table)
         for stat in stats:

{quollio_core-0.4.6 → quollio_core-0.4.7}/quollio_core/profilers/lineage.py

@@ -153,3 +153,17 @@ def parse_databricks_table_lineage(results: List) -> List[Dict[str, Dict]]:
         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-0.4.7/quollio_core/repository/bigquery.py

@@ -0,0 +1,61 @@
+import logging
+
+from google.cloud.bigquery import Client
+from google.cloud.datacatalog_lineage_v1 import EntityReference, LineageClient, SearchLinksRequest
+from google.oauth2.service_account import Credentials
+from googleapiclient.discovery import build
+
+logger = logging.getLogger(__name__)
+
+
+class BigQueryClient:
+    def __init__(self, credentials: Credentials) -> None:
+        self.client = self.__initialze(credentials=credentials)
+
+    def __initialze(self, credentials: Credentials) -> Client:
+        client = Client(credentials=credentials)
+        return client
+
+    def list_datasets(self, project_id) -> list:
+        datasets = list(self.client.list_datasets(project_id))
+        logger.debug("Found %s datasets in project %s", len(datasets), project_id)
+        return datasets
+
+    def list_tables(self, dataset_id) -> list:
+        tables = list(self.client.list_tables(dataset_id))
+        logger.debug("Found %s tables in dataset %s", len(tables), dataset_id)
+        return list(self.client.list_tables(dataset_id))
+
+
+class GCPLineageClient:
+    def __init__(self, credentials: Credentials) -> None:
+        self.client = self.__initialze(credentials=credentials)
+
+    def __initialze(self, credentials: Credentials) -> LineageClient:
+        client = LineageClient(credentials=credentials)
+        return client
+
+    def get_links(self, request: SearchLinksRequest) -> list:
+        response = self.client.search_links(request)
+        return response.links
+
+
+def get_entitiy_reference() -> EntityReference:
+    return EntityReference()
+
+
+def get_search_request(downstream_table: EntityReference, project_id: str, region: str) -> SearchLinksRequest:
+    return SearchLinksRequest(target=downstream_table, parent=f"projects/{project_id}/locations/{region.lower()}")
+
+
+def get_credentials(credentials_json: dict) -> Credentials:
+    return Credentials.from_service_account_info(credentials_json)
+
+
+def get_org_id(credentials_json: dict) -> str:
+    credentials = get_credentials(credentials_json)
+    crm_service = build("cloudresourcemanager", "v1", credentials=credentials)
+    project_id = credentials_json["project_id"]
+    project = crm_service.projects().get(projectId=project_id).execute()
+    org_id = project["parent"]["id"]
+    return org_id

quollio_core-0.4.6/quollio_core/dbt_projects/snowflake/models/quollio_stats_profiling_columns.sql

@@ -1,59 +0,0 @@
-WITH columns AS (
-    SELECT
-      table_catalog
-      , table_schema
-      , table_name
-      , column_name
-      , data_type
-    FROM
-      {{ source('account_usage', 'COLUMNS') }}
-    WHERE
-      deleted is null
-      AND table_name NOT LIKE 'QUOLLIO_%%'
-    GROUP BY
-      table_catalog
-      , table_schema
-      , table_name
-      , column_name
-      , data_type
-    ORDER BY
-      table_catalog
-      , table_schema
-      , table_name
-), accessible_tables AS (
-    SELECT
-      table_catalog
-      , table_schema
-      , name
-    FROM
-      {{ source('account_usage', 'GRANTS_TO_ROLES') }}
-    WHERE
-      granted_on in ('TABLE', 'MATERIALIZED VIEW')
-      AND grantee_name = '{{ var("query_role") }}'
-      AND privilege in ('SELECT', 'OWNERSHIP', 'REFERENCES')
-      AND deleted_on IS NULL
-    GROUP BY
-      table_catalog
-      , table_schema
-      , name  
-)
-
-SELECT
-  c.table_catalog
-  , c.table_schema
-  , c.table_name
-  , c.column_name
-  , c.data_type
-  , case when c.data_type in('NUMBER','DECIMAL', 'DEC', 'NUMERIC',
-                           'INT', 'INTEGER', 'BIGINT', 'SMALLINT',
-                           'TINYINT', 'BYTEINT')
-                           THEN true
-         else false END AS is_calculable
-FROM
-  columns c
-INNER JOIN
-  accessible_tables a
-ON
-  c.table_catalog = a.table_catalog
-  AND c.table_schema = a.table_schema
-  AND c.table_name = a.name