quollio-core 0.4.9__py3-none-any.whl → 0.4.11__py3-none-any.whl

quollio_core/__init__.py

@@ -1,4 +1,4 @@
 """Quollio Core"""
 
-__version__ = "0.4.9"
+__version__ = "0.4.11"
 __author__ = "Quollio Technologies, Inc"

quollio_core/bigquery.py

@@ -3,6 +3,7 @@ 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
@@ -88,14 +89,22 @@ if __name__ == "__main__":
         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
         )

quollio_core/bricks.py

@@ -4,6 +4,7 @@ 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,
@@ -20,7 +21,6 @@ def build_view(
     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
@@ -64,7 +64,6 @@ def load_lineage(
     tenant_id: str,
     enable_column_lineage: bool = False,
 ) -> 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(
@@ -99,7 +98,6 @@ def load_column_stats(
     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(
@@ -243,6 +241,7 @@ if __name__ == "__main__":
     )
 
     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.

quollio_core/helper/log.py

@@ -0,0 +1,17 @@
+import logging
+
+
+def set_log_level(level: str = "info") -> None:
+    fmt = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
+    if level == "info":
+        logging.basicConfig(level=logging.INFO, format=fmt)
+    elif level == "debug":
+        logging.basicConfig(level=logging.DEBUG, format=fmt)
+    elif level == "warn":
+        logging.basicConfig(level=logging.WARNING, format=fmt)
+    elif level == "error":
+        logging.basicConfig(level=logging.ERROR, format=fmt)
+    elif level == "critical":
+        logging.basicConfig(level=logging.CRITICAL, format=fmt)
+    else:
+        logging.basicConfig(level=logging.NOTSET, format=fmt)

quollio_core/profilers/databricks.py

@@ -19,7 +19,6 @@ def databricks_table_level_lineage(
     tenant_id: str,
     dbt_table_name: str = "quollio_lineage_table_level",
 ) -> None:
-    logging.basicConfig(level=logging.info, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
     with databricks.DatabricksQueryExecutor(config=conn) as databricks_executor:
         results = databricks_executor.get_query_results(
             query=f"""
@@ -61,7 +60,6 @@ def databricks_column_level_lineage(
     tenant_id: str,
     dbt_table_name: str = "quollio_lineage_column_level",
 ) -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
     with databricks.DatabricksQueryExecutor(config=conn) as databricks_executor:
         results = databricks_executor.get_query_results(
             query=f"""
@@ -140,42 +138,47 @@ def _get_column_stats(
             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
+                    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
+                            , 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(
                 monitoring_table=table["table_fqdn"],
                 monitored_table_catalog=monitored_table[0],
                 monitored_table_schema=monitored_table[1],
                 monitored_table_name=monitored_table[2],
             )
+            logger.debug(f"The following sql will be fetched to retrieve stats values. {query}")
             stats.append(databricks_executor.get_query_results(query))
     return stats
 

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="""

quollio_core/profilers/snowflake.py

@@ -17,42 +17,41 @@ def snowflake_table_to_table_lineage(
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
 ) -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
-    sf_executor = snowflake.SnowflakeQueryExecutor(conn)
-    results = sf_executor.get_query_results(
-        query="""
-        SELECT
-            *
-        FROM
-           {db}.{schema}.QUOLLIO_LINEAGE_TABLE_LEVEL
-        """.format(
-            db=conn.account_database,
-            schema=conn.account_schema,
-        )
-    )
-    parsed_results = parse_snowflake_results(results=results)
-    update_table_lineage_inputs = gen_table_lineage_payload(
-        tenant_id=tenant_id,
-        endpoint=conn.account_id,
-        tables=parsed_results,
-    )
-
-    req_count = 0
-    for update_table_lineage_input in update_table_lineage_inputs:
-        logger.info(
-            "Generating table lineage. downstream: {db} -> {schema} -> {table}".format(
-                db=update_table_lineage_input.downstream_database_name,
-                schema=update_table_lineage_input.downstream_schema_name,
-                table=update_table_lineage_input.downstream_table_name,
+    with snowflake.SnowflakeQueryExecutor(conn) as sf_executor:
+        results = sf_executor.get_query_results(
+            query="""
+            SELECT
+                *
+            FROM
+               {db}.{schema}.QUOLLIO_LINEAGE_TABLE_LEVEL
+            """.format(
+                db=conn.account_database,
+                schema=conn.account_schema,
             )
         )
-        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(),
+        parsed_results = parse_snowflake_results(results=results)
+        update_table_lineage_inputs = gen_table_lineage_payload(
+            tenant_id=tenant_id,
+            endpoint=conn.account_id,
+            tables=parsed_results,
         )
-        if status_code == 200:
-            req_count += 1
-    logger.info(f"Generating table lineage is finished. {req_count} lineages are ingested.")
+
+        req_count = 0
+        for update_table_lineage_input in update_table_lineage_inputs:
+            logger.info(
+                "Generating table lineage. downstream: {db} -> {schema} -> {table}".format(
+                    db=update_table_lineage_input.downstream_database_name,
+                    schema=update_table_lineage_input.downstream_schema_name,
+                    table=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(f"Generating table lineage is finished. {req_count} lineages are ingested.")
     return
 
 
@@ -61,41 +60,41 @@ def snowflake_column_to_column_lineage(
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
 ) -> None:
-    sf_executor = snowflake.SnowflakeQueryExecutor(conn)
-    results = sf_executor.get_query_results(
-        query="""
-        SELECT
-            *
-        FROM
-            {db}.{schema}.QUOLLIO_LINEAGE_COLUMN_LEVEL
-        """.format(
-            db=conn.account_database,
-            schema=conn.account_schema,
-        )
-    )
-    update_column_lineage_inputs = gen_column_lineage_payload(
-        tenant_id=tenant_id,
-        endpoint=conn.account_id,
-        columns=results,
-    )
-
-    req_count = 0
-    for update_column_lineage_input in update_column_lineage_inputs:
-        logger.info(
-            "Generating column lineage. downstream: {db} -> {schema} -> {table} -> {column}".format(
-                db=update_column_lineage_input.downstream_database_name,
-                schema=update_column_lineage_input.downstream_schema_name,
-                table=update_column_lineage_input.downstream_table_name,
-                column=update_column_lineage_input.downstream_column_name,
+    with snowflake.SnowflakeQueryExecutor(conn) as sf_executor:
+        results = sf_executor.get_query_results(
+            query="""
+            SELECT
+                *
+            FROM
+                {db}.{schema}.QUOLLIO_LINEAGE_COLUMN_LEVEL
+            """.format(
+                db=conn.account_database,
+                schema=conn.account_schema,
             )
         )
-        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(),
+        update_column_lineage_inputs = gen_column_lineage_payload(
+            tenant_id=tenant_id,
+            endpoint=conn.account_id,
+            columns=results,
         )
-        if status_code == 200:
-            req_count += 1
-    logger.info(f"Generating column lineage is finished. {req_count} lineages are ingested.")
+
+        req_count = 0
+        for update_column_lineage_input in update_column_lineage_inputs:
+            logger.info(
+                "Generating column lineage. downstream: {db} -> {schema} -> {table} -> {column}".format(
+                    db=update_column_lineage_input.downstream_database_name,
+                    schema=update_column_lineage_input.downstream_schema_name,
+                    table=update_column_lineage_input.downstream_table_name,
+                    column=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(f"Generating column lineage is finished. {req_count} lineages are ingested.")
     return
 
 
@@ -104,110 +103,67 @@ def snowflake_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")
-    sf_executor = snowflake.SnowflakeQueryExecutor(conn)
-    results = sf_executor.get_query_results(
-        query="""
-        SELECT
-            database_name
-            , schema_name
-            , query_text
-        FROM
-           {db}.{schema}.QUOLLIO_SQLLINEAGE_SOURCES
-        """.format(
-            db=conn.account_database,
-            schema=conn.account_schema,
-        )
-    )
-    update_table_lineage_inputs_list = list()
-    sql_lineage = SQLLineage()
-    for result in results:
-        src_tables, dest_table = sql_lineage.get_table_level_lineage_source(
-            sql=result["QUERY_TEXT"],
-            dialect="snowflake",
-            dest_db=result["DATABASE_NAME"],
-            dest_schema=result["SCHEMA_NAME"],
-        )
-        update_table_lineage_inputs = sql_lineage.gen_lineage_input(
-            tenant_id=tenant_id, endpoint=conn.account_id, src_tables=src_tables, dest_table=dest_table
-        )
-        update_table_lineage_inputs_list.append(update_table_lineage_inputs)
-
-    req_count = 0
-    for update_table_lineage_input in update_table_lineage_inputs_list:
-        logger.info(
-            "Generating table lineage. downstream: {db} -> {schema} -> {table}".format(
-                db=update_table_lineage_input.downstream_database_name,
-                schema=update_table_lineage_input.downstream_schema_name,
-                table=update_table_lineage_input.downstream_table_name,
+    with snowflake.SnowflakeQueryExecutor(conn) as sf_executor:
+        results = sf_executor.get_query_results(
+            query="""
+            SELECT
+                database_name
+                , schema_name
+                , query_text
+            FROM
+               {db}.{schema}.QUOLLIO_SQLLINEAGE_SOURCES
+            """.format(
+                db=conn.account_database,
+                schema=conn.account_schema,
             )
         )
-        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(f"Generating table lineage is finished. {req_count} lineages are ingested.")
+        update_table_lineage_inputs_list = list()
+        sql_lineage = SQLLineage()
+        for result in results:
+            src_tables, dest_table = sql_lineage.get_table_level_lineage_source(
+                sql=result["QUERY_TEXT"],
+                dialect="snowflake",
+                dest_db=result["DATABASE_NAME"],
+                dest_schema=result["SCHEMA_NAME"],
+            )
+            update_table_lineage_inputs = sql_lineage.gen_lineage_input(
+                tenant_id=tenant_id, endpoint=conn.account_id, src_tables=src_tables, dest_table=dest_table
+            )
+            update_table_lineage_inputs_list.append(update_table_lineage_inputs)
+
+        req_count = 0
+        for update_table_lineage_input in update_table_lineage_inputs_list:
+            logger.info(
+                "Generating table lineage. downstream: {db} -> {schema} -> {table}".format(
+                    db=update_table_lineage_input.downstream_database_name,
+                    schema=update_table_lineage_input.downstream_schema_name,
+                    table=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(f"Generating table lineage is finished. {req_count} lineages are ingested.")
     return
 
 
-def _get_target_tables_query(db: str, schema: str) -> str:
-    query = """
-        SELECT
-            DISTINCT
-            TABLE_CATALOG
-            , TABLE_SCHEMA
-            , 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:
-    query = """
-        SELECT
-            DISTINCT
-            TABLE_CATALOG
-            , TABLE_SCHEMA
-            , TABLE_NAME
-        FROM
-            {db}.INFORMATION_SCHEMA.TABLES
-        WHERE
-            startswith(TABLE_NAME, 'QUOLLIO_STATS_COLUMNS_')
-            AND TABLE_SCHEMA = UPPER('{schema}')
-        """.format(
-        db=db, schema=schema
-    )
-    return query
-
-
 def snowflake_table_stats(
     conn: snowflake.SnowflakeConnectionConfig,
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
 ) -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
-    sf_executor = snowflake.SnowflakeQueryExecutor(conn)
-
-    target_query = _get_target_tables_query(
-        db=conn.account_database,
-        schema=conn.account_schema,
-    )
-    target_assets = sf_executor.get_query_results(query=target_query)
-
-    stats_query = _get_stats_tables_query(
-        db=conn.account_database,
-        schema=conn.account_schema,
-    )
-    stats_columns = sf_executor.get_query_results(query=stats_query)
+    with snowflake.SnowflakeQueryExecutor(conn) as sf_executor:
+        stats_query = _gen_get_stats_views_query(
+            db=conn.account_database,
+            schema=conn.account_schema,
+        )
+        stats_views = sf_executor.get_query_results(query=stats_query)
 
-    req_count = 0
-    for target_asset in target_assets:
-        for stats_column in stats_columns:
+        req_count = 0
+        for stats_view in stats_views:
             stats_query = """
             SELECT
                 db_name
@@ -224,18 +180,12 @@ def snowflake_table_stats(
                 , 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["TABLE_CATALOG"],
-                schema=stats_column["TABLE_SCHEMA"],
-                table=stats_column["TABLE_NAME"],
-                target_db=target_asset["TABLE_CATALOG"],
-                target_schema=target_asset["TABLE_SCHEMA"],
-                target_table=target_asset["TABLE_NAME"],
+                db=stats_view["TABLE_CATALOG"],
+                schema=stats_view["TABLE_SCHEMA"],
+                table=stats_view["TABLE_NAME"],
             )
+            logger.debug(f"The following sql will be fetched to retrieve stats values. {stats_query}")
             stats_result = sf_executor.get_query_results(query=stats_query)
             payloads = gen_table_stats_payload(tenant_id=tenant_id, endpoint=conn.account_id, stats=stats_result)
             for payload in payloads:
@@ -253,4 +203,23 @@ def snowflake_table_stats(
                 )
                 if status_code == 200:
                     req_count += 1
-    logger.info(f"Generating table stats is finished. {req_count} stats are ingested.")
+        logger.info(f"Generating table stats is finished. {req_count} stats are ingested.")
+    return
+
+
+def _gen_get_stats_views_query(db: str, schema: str) -> str:
+    query = """
+        SELECT
+            DISTINCT
+            TABLE_CATALOG
+            , TABLE_SCHEMA
+            , TABLE_NAME
+        FROM
+            {db}.INFORMATION_SCHEMA.TABLES
+        WHERE
+            startswith(TABLE_NAME, 'QUOLLIO_STATS_COLUMNS_')
+            AND TABLE_SCHEMA = UPPER('{schema}')
+        """.format(
+        db=db, schema=schema
+    )
+    return query

quollio_core/profilers/sqllineage.py

@@ -54,7 +54,6 @@ class SQLLineage:
         dest_db: str = None,
         dest_schema: str = None,
     ) -> Tuple[Set[Table], Table]:
-        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
         try:
             statement: sqlglot.Expression = sqlglot.parse_one(sql=sql, error_level=sqlglot.ErrorLevel.RAISE)
         except ParseError as e:

quollio_core/redshift.py

@@ -4,6 +4,7 @@ import os
 
 from quollio_core.helper.core import setup_dbt_profile
 from quollio_core.helper.env_default import env_default
+from quollio_core.helper.log import set_log_level
 from quollio_core.profilers.redshift import (
     redshift_table_level_lineage,
     redshift_table_level_sqllineage,
@@ -20,8 +21,6 @@ def build_view(
     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()
@@ -74,7 +73,6 @@ def load_lineage(
     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 redshift table to table lineage.")
     redshift_table_level_lineage(
         conn=conn,
@@ -101,7 +99,6 @@ def load_stats(
     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 redshift stats.")
     redshift_table_stats(
@@ -119,7 +116,6 @@ def load_sqllineage(
     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 Redshift sqllineage.")
     redshift_table_level_sqllineage(
@@ -266,6 +262,8 @@ if __name__ == "__main__":
         help="The client secrete that is created on Quollio console to let clients access Quollio External API",
     )
     args = parser.parse_args()
+    set_log_level(level=args.log_level)
+
     conn = redshift.RedshiftConnectionConfig(
         host=args.host,
         build_user=args.build_user,

quollio_core/repository/databricks.py

@@ -5,7 +5,7 @@ from typing import Dict, List, Optional
 from databricks.sdk.core import Config, HeaderFactory, oauth_service_principal
 from databricks.sql.client import Connection, connect
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -47,8 +47,8 @@ class DatabricksQueryExecutor:
                 cur.execute(query)
                 result: List[Dict[str, str]] = cur.fetchall()
             except Exception as e:
-                logging.error(query, exc_info=True)
-                logging.error("databricks get_query_results failed. %s", e)
+                logger.error(query, exc_info=True)
+                logger.error("databricks get_query_results failed. %s", e)
                 raise
 
             for row in result:

quollio_core/repository/dbt.py

@@ -11,7 +11,6 @@ class DBTClient:
         self.dbt = dbtRunner()
 
     def invoke(self, cmd: str, project_dir: str, profile_dir: str, options: List[str] = None) -> dbtRunnerResult:
-        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
         req = [cmd, "--project-dir", project_dir, "--profiles-dir", profile_dir]
         if options is not None:
             req.extend(options)

quollio_core/repository/qdc.py

@@ -25,7 +25,6 @@ class QDCExternalAPIClient:
         Tried to find a package for oauth0 client credentials flow,
         but any of them contains bugs or lacks of features to handle the token refresh when it's expired
         """
-        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
         url = f"{self.base_url}/oauth2/token"
         creds = f"{self.client_id}:{self.client_secret}"
         encoded_creds = base64.b64encode(creds.encode()).decode()
@@ -65,7 +64,6 @@ class QDCExternalAPIClient:
         return session
 
     def update_stats_by_id(self, global_id: str, payload: Dict[str, List[str]]) -> int:
-        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
         self._refresh_token_if_expired()
         headers = {"content-type": "application/json", "authorization": f"Bearer {self.auth_token}"}
         endpoint = f"{self.base_url}/v2/assets/{global_id}/stats"
@@ -85,7 +83,6 @@ class QDCExternalAPIClient:
             return res.status_code
 
     def update_lineage_by_id(self, global_id: str, payload: Dict[str, List[str]]) -> int:
-        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
         self._refresh_token_if_expired()
         headers = {"content-type": "application/json", "authorization": f"Bearer {self.auth_token}"}
         endpoint = f"{self.base_url}/v2/lineage/{global_id}"

quollio_core/repository/redshift.py

@@ -67,7 +67,6 @@ class RedshiftQueryExecutor:
         return conn
 
     def get_query_results(self, query: str) -> Tuple[List[str]]:
-        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
         with self.conn.cursor() as cur:
             try:
                 cur.execute(query)

quollio_core/repository/snowflake.py

@@ -28,6 +28,12 @@ class SnowflakeQueryExecutor:
     def __init__(self, config: SnowflakeConnectionConfig) -> None:
         self.conn = self.__initialize(config)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.conn.close()
+
     def __initialize(self, config: SnowflakeConnectionConfig) -> SnowflakeConnection:
         conn: SnowflakeConnection = connect(
             user=config.account_user,
@@ -41,7 +47,6 @@ class SnowflakeQueryExecutor:
         return conn
 
     def get_query_results(self, query: str) -> List[Dict[str, str]]:
-        logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
         with self.conn.cursor(DictCursor) as cur:
             try:
                 cur.execute(query)

quollio_core/snowflake.py

@@ -4,6 +4,7 @@ import os
 
 from quollio_core.helper.core import setup_dbt_profile
 from quollio_core.helper.env_default import env_default
+from quollio_core.helper.log import set_log_level
 from quollio_core.profilers.snowflake import (
     snowflake_column_to_column_lineage,
     snowflake_table_level_sqllineage,
@@ -21,7 +22,6 @@ def build_view(
     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
@@ -74,9 +74,8 @@ def load_lineage(
     tenant_id: str,
     enable_column_lineage: bool = False,
 ) -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")
-
     logger.info("Generate Snowflake table to table lineage.")
+
     snowflake_table_to_table_lineage(
         conn=conn,
         qdc_client=qdc_client,
@@ -105,7 +104,6 @@ def load_stats(
     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 Snowflake stats.")
     snowflake_table_stats(
@@ -124,7 +122,6 @@ def load_sqllineage(
     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 Snowflake sqllineage.")
     snowflake_table_level_sqllineage(
@@ -279,6 +276,8 @@ if __name__ == "__main__":
         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 = snowflake.SnowflakeConnectionConfig(
         account_id=args.account_id,
         account_user=args.user,

{quollio_core-0.4.9.dist-info → quollio_core-0.4.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: quollio-core
-Version: 0.4.9
+Version: 0.4.11
 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>

{quollio_core-0.4.9.dist-info → quollio_core-0.4.11.dist-info}/RECORD

@@ -1,8 +1,8 @@
-quollio_core/__init__.py,sha256=143doeugqWpXanEUQkIVJmQ5p5zRqMXEa3Ab81BNjFw,83
-quollio_core/bigquery.py,sha256=2DrUMo4evcH4BHiUtnY48IjmsdAsQMoPGtNx8SRoyzQ,3528
-quollio_core/bricks.py,sha256=BVwh9clJMXe_YXbd78ku6Y9470cYYE3CflcuCRzrY3I,9689
-quollio_core/redshift.py,sha256=wVqtNnv1fmZ5QwYq4A2JVLHH3VrCYPEZBx39NoCqRIM,10220
-quollio_core/snowflake.py,sha256=WRtnXxc01AaV4oTa382MTppYvmOn4mOY7jFDNDnzxDA,10922
+quollio_core/__init__.py,sha256=_jIDStu07dkGu5ouTAUr-ImjBx7pmvosgMOUjUT25Pc,84
+quollio_core/bigquery.py,sha256=RguUznaY5YjROzJtXimoS8yCNH9jgGphpzd5v_JgSQM,3884
+quollio_core/bricks.py,sha256=Lehv-qsBSMNNE9BGVvidGOXJsxLSSsbNtmiEZH4lSUg,9458
+quollio_core/redshift.py,sha256=1d-mHnalB1jtiGPgzsGd3lRwLHCxaBJlUMEV2dh4f60,9882
+quollio_core/snowflake.py,sha256=G3tImWbZgMlycYuw1b5WnNBp3zWo3hyrbOX5ARLIs7A,10585
 quollio_core/dbt_projects/databricks/.gitignore,sha256=1jJAyXSzJ3YUm0nx3i7wUSE4RjQMX3ad6F8O88UbtzI,29
 quollio_core/dbt_projects/databricks/README.md,sha256=ZpRQyhFAODAiS8dc1Kb_ndkul4cu4o4udN_EMa49CU4,440
 quollio_core/dbt_projects/databricks/dbt_project.yml,sha256=3sH98RNk7TnphvI3yEdXDstb92kW5BNxr-cT0tXhwzk,480
@@ -64,22 +64,23 @@ quollio_core/dbt_projects/snowflake/snapshots/.gitkeep,sha256=47DEQpj8HBSa-_TImW
 quollio_core/helper/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 quollio_core/helper/core.py,sha256=wbu4FWI7YiFEttXGSuj3tMyAhtPAFlHOjDpWJGNXOHA,1202
 quollio_core/helper/env_default.py,sha256=H6gbSGUPrEDZr4YDrL49hbOpw6RntI4U82kX1q6vUnI,2148
+quollio_core/helper/log.py,sha256=flxyZZ44G79l1TaUp3OT58uCHcnE5z_pCduwoeI6IUs,645
 quollio_core/profilers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 quollio_core/profilers/bigquery.py,sha256=e1Y8cZR-LxI9mSsYb0DurQyy0eCjM_kAKLfvl4IuQLE,3262
-quollio_core/profilers/databricks.py,sha256=aNUIk48GuwSDEnmNzhZ6yv6TqAc1dBJt9e81NiTGBjo,7496
+quollio_core/profilers/databricks.py,sha256=pFWV6ecTefOS8evbnnXTcqR1jmbLsqn1V89ISyM2uqU,7743
 quollio_core/profilers/lineage.py,sha256=4FyxIuPBrUFihqZryqTQBcfB0Z7634lKl_WwkD82vzE,6865
-quollio_core/profilers/redshift.py,sha256=obdHVIsOM1bwHGdvYKalsJcTXwLK02kAKQMSBzSvsDo,7862
-quollio_core/profilers/snowflake.py,sha256=C1LC19ZaUMwNoXjsbnez0xANydJYs8oNRt6tixWKDq8,9090
-quollio_core/profilers/sqllineage.py,sha256=oCyl4tpXL5bkfguXAzTHSB9kZBL3tQK_rfcJ4XQMrLo,5177
+quollio_core/profilers/redshift.py,sha256=6_4amsBL4QW0ZajWhS-TW3f_cjKKa6TpXClMgBC-fZo,6440
+quollio_core/profilers/snowflake.py,sha256=nitlP5pmDm2RhLGO4f_WTzkw41EmOTY2uWN1HZkCHbI,8465
+quollio_core/profilers/sqllineage.py,sha256=XkF7hwDWIGNtyEP5cv2wETBgMfdQxeHolv7qPIkntSQ,5066
 quollio_core/profilers/stats.py,sha256=PG1NbbUSpc1JuEYvBzD66rd24tp0C13_Y5Y7vRjYG1c,4720
 quollio_core/repository/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 quollio_core/repository/bigquery.py,sha256=KMJTeF4OUxtaJt0ymoJ4tkrMKq8yLyMYaMxNvU5yd_Y,2271
-quollio_core/repository/databricks.py,sha256=m68tja5N-QxH3VqEq-mOJKBeR2qldSgj_L9iIxvWwm0,1945
-quollio_core/repository/dbt.py,sha256=HXqW_xa4xYPh9CnKkg4L1gwG3SGjj2BAYoWgzWMFU4U,770
-quollio_core/repository/qdc.py,sha256=VCmzAUvjLemw1os5TaPtfBFkMCOMuPeftjZmUPhFj2Y,4702
-quollio_core/repository/redshift.py,sha256=UVHIpYzDQ2AbBTAGa8DgmEenG0NZsHfYroR1MmEPQGA,2991
-quollio_core/repository/snowflake.py,sha256=1YVMDfb9euJKvikv1pk_IxVF6SVsiemSvZ-WMTSbY7E,1874
-quollio_core-0.4.9.dist-info/LICENSE,sha256=V8j_M8nAz8PvAOZQocyRDX7keai8UJ9skgmnwqETmdY,34520
-quollio_core-0.4.9.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-quollio_core-0.4.9.dist-info/METADATA,sha256=GW_XzITMPnzDHVXIMzgmAkYuzVkW2ZD3uPF89fCxrxA,6803
-quollio_core-0.4.9.dist-info/RECORD,,
+quollio_core/repository/databricks.py,sha256=9Cgdv8qBnVaHqu3RA-IUBieAqb69moQ-KAAMVSf5Ds4,1877
+quollio_core/repository/dbt.py,sha256=cnLwJPywLi8VowVW7zfIBa9jxVwDWO7xzzNRn1vWiuw,659
+quollio_core/repository/qdc.py,sha256=qEpMF6rKdic23dPJoDYmbIcyCKDuSFqbDF2_jqmqoZw,4369
+quollio_core/repository/redshift.py,sha256=p2ouEuYcDCjx1oBhc6H1ekQsvEqHGd3bFu3PW0ngYBc,2880
+quollio_core/repository/snowflake.py,sha256=J9rHshfWdOSnjQWxwGEYPpAU2lY7Tu5UFB_BNakkAX0,1892
+quollio_core-0.4.11.dist-info/LICENSE,sha256=V8j_M8nAz8PvAOZQocyRDX7keai8UJ9skgmnwqETmdY,34520
+quollio_core-0.4.11.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+quollio_core-0.4.11.dist-info/METADATA,sha256=0tv9AKZYWKt9yPa9iaziHZyTzhj708v-Aq22MFEbYAs,6804
+quollio_core-0.4.11.dist-info/RECORD,,