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

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/profilers/stats.py

@@ -77,7 +77,6 @@ def gen_table_stats_payload(tenant_id: str, endpoint: str, stats: List[Dict[str,
         table_global_id = new_global_id(
             tenant_id=tenant_id, cluster_id=endpoint, data_id=global_id_arg, data_type="column"
         )
-
         stats_request = StatsRequest(
             global_id=table_global_id,
             db=stat["DB_NAME"],

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,
@@ -16,12 +17,10 @@ logger = logging.getLogger(__name__)
 
 def build_view(
     conn: redshift.RedshiftConnectionConfig,
-    skip_heavy: bool,
+    aggregate_all: bool = False,
     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()
@@ -29,9 +28,9 @@ def build_view(
     project_path = f"{current_dir}/dbt_projects/redshift"
     template_path = f"{current_dir}/dbt_projects/redshift/profiles"
     template_name = "profiles_template.yml"
-    options = '{{"query_user": {query_user}, "skip_heavy": {skip_heavy}, "target_database": {database}}}'.format(
+    options = '{{"query_user": {query_user}, "aggregate_all": {aggregate_all}, "target_database": {database}}}'.format(
         query_user=conn.query_user,
-        skip_heavy=skip_heavy,
+        aggregate_all=aggregate_all,
         database=conn.database,
     )
 
@@ -50,11 +49,10 @@ def build_view(
         options=["--no-use-colors", "--log-level", log_level, "--vars", options],
     )
     run_options = ["--no-use-colors", "--log-level", log_level, "--vars", options]
-    target_tables_list = target_tables.split()
-    if target_tables_list is not None:
-        if "quollio_stats_columns" in target_tables_list:
-            target_tables_list.append("quollio_stats_profiling_columns")
-        target_tables_str = " ".join(target_tables_list)
+    if target_tables is not None:
+        if "quollio_stats_columns" in target_tables:
+            target_tables.append("quollio_stats_profiling_columns")
+        target_tables_str = " ".join(target_tables)
         run_options.append("--select")
         run_options.append(target_tables_str)
 
@@ -75,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,
@@ -102,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(
@@ -120,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(
@@ -211,12 +206,12 @@ if __name__ == "__main__":
         help="Target schema name where the views are built by dbt",
     )
     parser.add_argument(
-        "--skip_heavy",
+        "--aggregate_all",
         type=bool,
-        action=env_default("REDSHIFT_SKIP_HEAVY"),
-        default=True,
+        action=env_default("REDSHIFT_AGGREGATE_ALL", store_true=True),
+        default=False,
         required=False,
-        help="Skip heavy queries when building views by dbt",
+        help="Aggregate all stats values. False by default.",
     )
     parser.add_argument(
         "--target_tables",
@@ -267,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,
@@ -284,7 +281,7 @@ if __name__ == "__main__":
     if "build_view" in args.commands:
         build_view(
             conn=conn,
-            skip_heavy=args.skip_heavy,
+            aggregate_all=args.aggregate_all,
             target_tables=args.target_tables,
             log_level=args.log_level,
         )

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/repository/databricks.py

@@ -0,0 +1,62 @@
+import logging
+from dataclasses import asdict, dataclass
+from typing import Dict, List, Optional
+
+from databricks.sdk.core import Config, HeaderFactory, oauth_service_principal
+from databricks.sql.client import Connection, connect
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DatabricksConnectionConfig:
+    host: str
+    http_path: str
+    client_id: str
+    client_secret: str
+    catalog: str
+    schema: str
+
+    def as_dict(self) -> Dict[str, str]:
+        return asdict(self)
+
+
+class DatabricksQueryExecutor:
+    def __init__(self, config: DatabricksConnectionConfig) -> None:
+        self.config = config
+        self.conn = self.__initialize()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.conn.close()
+
+    def __initialize(self) -> Connection:
+        conn = connect(
+            server_hostname=self.config.host,
+            http_path=self.config.http_path,
+            credentials_provider=self.credential_provider,
+        )
+        return conn
+
+    def get_query_results(self, query: str) -> List[Dict[str, str]]:
+        results_asdict: List[Dict[str, str]] = []
+        with self.conn.cursor() as cur:
+            try:
+                cur.execute(query)
+                result: List[Dict[str, str]] = cur.fetchall()
+            except Exception as e:
+                logger.error(query, exc_info=True)
+                logger.error("databricks get_query_results failed. %s", e)
+                raise
+
+            for row in result:
+                results_asdict.append(row.asDict())
+        return results_asdict
+
+    def credential_provider(self) -> Optional[HeaderFactory]:
+        config = Config(
+            host=f"https://{self.config.host}", client_id=self.config.client_id, client_secret=self.config.client_secret
+        )
+        return oauth_service_principal(config)

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)