quollio-core 0.4.9__tar.gz → 0.4.10__tar.gz

{quollio_core-0.4.9 → quollio_core-0.4.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: quollio-core
-Version: 0.4.9
+Version: 0.4.10
 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 → quollio_core-0.4.10}/quollio_core/__init__.py

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

{quollio_core-0.4.9 → quollio_core-0.4.10}/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-0.4.9 → quollio_core-0.4.10}/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-0.4.10/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-0.4.9 → quollio_core-0.4.10}/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"""

{quollio_core-0.4.9 → quollio_core-0.4.10}/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-0.4.10/quollio_core/profilers/snowflake.py

@@ -0,0 +1,225 @@
+import logging
+
+from quollio_core.profilers.lineage import (
+    gen_column_lineage_payload,
+    gen_table_lineage_payload,
+    parse_snowflake_results,
+)
+from quollio_core.profilers.sqllineage import SQLLineage
+from quollio_core.profilers.stats import gen_table_stats_payload
+from quollio_core.repository import qdc, snowflake
+
+logger = logging.getLogger(__name__)
+
+
+def snowflake_table_to_table_lineage(
+    conn: snowflake.SnowflakeConnectionConfig,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+) -> None:
+    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,
+            )
+        )
+        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,
+                )
+            )
+            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 snowflake_column_to_column_lineage(
+    conn: snowflake.SnowflakeConnectionConfig,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+) -> None:
+    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,
+            )
+        )
+        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,
+                )
+            )
+            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
+
+
+def snowflake_table_level_sqllineage(
+    conn: snowflake.SnowflakeConnectionConfig,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+) -> None:
+    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,
+            )
+        )
+        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 snowflake_table_stats(
+    conn: snowflake.SnowflakeConnectionConfig,
+    qdc_client: qdc.QDCExternalAPIClient,
+    tenant_id: str,
+) -> None:
+    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 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["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:
+                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
+        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-0.4.9 → quollio_core-0.4.10}/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-0.4.9 → quollio_core-0.4.10}/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-0.4.9 → quollio_core-0.4.10}/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-0.4.9 → quollio_core-0.4.10}/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-0.4.9 → quollio_core-0.4.10}/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-0.4.9 → quollio_core-0.4.10}/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-0.4.9 → quollio_core-0.4.10}/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-0.4.9 → quollio_core-0.4.10}/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/quollio_core/profilers/snowflake.py

@@ -1,256 +0,0 @@
-import logging
-
-from quollio_core.profilers.lineage import (
-    gen_column_lineage_payload,
-    gen_table_lineage_payload,
-    parse_snowflake_results,
-)
-from quollio_core.profilers.sqllineage import SQLLineage
-from quollio_core.profilers.stats import gen_table_stats_payload
-from quollio_core.repository import qdc, snowflake
-
-logger = logging.getLogger(__name__)
-
-
-def snowflake_table_to_table_lineage(
-    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)
-    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,
-            )
-        )
-        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 snowflake_column_to_column_lineage(
-    conn: snowflake.SnowflakeConnectionConfig,
-    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,
-            )
-        )
-        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
-
-
-def snowflake_table_level_sqllineage(
-    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)
-    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,
-            )
-        )
-        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)
-
-    req_count = 0
-    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["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"],
-            )
-            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:
-                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
-    logger.info(f"Generating table stats is finished. {req_count} stats are ingested.")