quollio-core 0.5.3__py3-none-any.whl → 0.6.1__py3-none-any.whl

quollio_core/__init__.py

@@ -1,4 +1,4 @@
 """Quollio Core"""
 
-__version__ = "0.5.3"
+__version__ = "0.6.1"
 __author__ = "Quollio Technologies, Inc"

quollio_core/dbt_projects/snowflake/profiles/profiles_template.yml

@@ -8,6 +8,11 @@ quollio_intelligence_snowflake:
       schema: {{ account_schema }}
       type: snowflake
       user: {{ account_user }}
-      password: {{ account_password }}
       warehouse: {{ account_warehouse }}
       threads: {{ threads }}
+      {% if private_key is defined %}
+      private_key: |
+        {{ private_key | indent(8) }}
+      {% else %}
+      password: {{ account_password }}
+      {% endif %}

quollio_core/helper/core.py

@@ -8,6 +8,7 @@ from jinja2 import Environment, FileSystemLoader
 def new_global_id(tenant_id: str, cluster_id: str, data_id: str, data_type: str) -> str:
     prefix = ""
     data_types = {
+        "data_source": "dsrc-",
         "schema": "schm-",
         "table": "tbl-",
         "column": "clmn-",

quollio_core/models/avroasset.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+from typing import List, Optional
+
+from dataclasses_avroschema import AvroModel
+
+
+@dataclass
+class AvroAsset(AvroModel):
+    "AvroAsset"
+
+    id: str
+    object_type: str
+    parents: List[str]
+    name: str
+    stats_max: Optional[str] = None
+    stats_min: Optional[str] = None
+    stats_mean: Optional[str] = None
+    stats_median: Optional[str] = None
+    stats_mode: Optional[str] = None
+    stats_stddev: Optional[str] = None
+    stats_number_of_null: Optional[str] = None
+    stats_number_of_unique: Optional[str] = None
+    upstream: Optional[List[str]] = None

quollio_core/models/qdc.py

@@ -0,0 +1,36 @@
+from dataclasses import asdict, dataclass
+from typing import Dict
+
+
+@dataclass
+class GetImportURLRequest:
+    service_name: str
+    source_name: str
+    file_name: str
+    override_logical_name: str
+    update_mode: str
+
+    def as_dict(self) -> Dict[str, str]:
+        return asdict(self)
+
+
+@dataclass
+class DataSourceMetadataResponseBody:
+    user_id: str
+    job_key: str
+    service_name: str
+    source_name: str
+    source_type: str
+    override_logical_name: str
+
+    def as_dict(self) -> Dict[str, str]:
+        return asdict(self)
+
+
+@dataclass
+class GetImportURLResponse:
+    location: str
+    datasource_metadata_response_body: DataSourceMetadataResponseBody
+
+    def as_dict(self) -> Dict[str, str]:
+        return asdict(self)

quollio_core/profilers/lineage.py

@@ -3,6 +3,7 @@ from dataclasses import asdict, dataclass
 from typing import Dict, List, Tuple, Union
 
 from quollio_core.helper.core import new_global_id
+from quollio_core.models.avroasset import AvroAsset
 
 
 @dataclass
@@ -23,6 +24,101 @@ class LineageInputs:
     upstreams: LineageInput
 
 
+def gen_table_avro_lineage_payload(
+    tenant_id: str,
+    endpoint: str,
+    tables: List[Dict[str, Union[Dict[str, str], str]]],
+    existing_global_ids: Dict[str, bool],
+) -> List[Dict[str, str]]:
+    payload = list()
+    for table in tables:
+        downstream_table_fqn = table["DOWNSTREAM_TABLE_NAME"].split(".")
+        if len(downstream_table_fqn) != 3:
+            continue
+        else:
+            global_id_arg = "{db}{schema}{table}".format(
+                db=downstream_table_fqn[0], schema=downstream_table_fqn[1], table=downstream_table_fqn[2]
+            )
+            downstream_table_global_id = new_global_id(
+                tenant_id=tenant_id, cluster_id=endpoint, data_id=global_id_arg, data_type="table"
+            )
+            if existing_global_ids.get(downstream_table_global_id) is not True:
+                continue
+            upstreams = list()
+            for upstream_table in table["UPSTREAM_TABLES"]:
+                upstream_table_fqn = upstream_table["upstream_object_name"].split(".")
+                if len(upstream_table_fqn) != 3:
+                    continue
+                else:
+                    upstream_global_id_arg = "{db}{schema}{table}".format(
+                        db=upstream_table_fqn[0], schema=upstream_table_fqn[1], table=upstream_table_fqn[2]
+                    )
+                    upstream_table_global_id = new_global_id(
+                        tenant_id=tenant_id, cluster_id=endpoint, data_id=upstream_global_id_arg, data_type="table"
+                    )
+                    upstreams.append(upstream_table_global_id)
+
+            avro_assets = AvroAsset(
+                id=downstream_table_global_id,
+                object_type="table",
+                parents=[downstream_table_fqn[0], downstream_table_fqn[1]],
+                name=downstream_table_fqn[2],
+                upstream=upstreams,
+            )
+            payload.append(avro_assets.to_dict())
+    return payload
+
+
+def gen_column_avro_lineage_payload(
+    tenant_id: str, endpoint: str, columns: List[Dict[str, str]], existing_global_ids: Dict[str, bool]
+) -> List[Dict[str, str]]:
+    payload = list()
+    for column in columns:
+        downstream_table_fqn = column["DOWNSTREAM_TABLE_NAME"].split(".")
+        if len(downstream_table_fqn) != 3:
+            continue
+        else:
+            global_id_arg = "{db}{schema}{table}{column}".format(
+                db=downstream_table_fqn[0],
+                schema=downstream_table_fqn[1],
+                table=downstream_table_fqn[2],
+                column=column["DOWNSTREAM_COLUMN_NAME"],
+            )
+            downstream_column_global_id = new_global_id(
+                tenant_id=tenant_id, cluster_id=endpoint, data_id=global_id_arg, data_type="column"
+            )
+            if existing_global_ids.get(downstream_column_global_id) is not True:
+                continue
+            upstream_columns: List[Dict[str, str]] = json.loads(column["UPSTREAM_COLUMNS"])
+            upstreams = list()
+            for upstream_column in upstream_columns:
+                upstream_table_fqn = upstream_column["upstream_table_name"].split(".")
+                if len(upstream_table_fqn) != 3:
+                    continue
+                elif not upstream_column.get("upstream_column_name"):
+                    continue
+                else:
+                    upstream_global_id_arg = "{db}{schema}{table}{column}".format(
+                        db=upstream_table_fqn[0],
+                        schema=upstream_table_fqn[1],
+                        table=upstream_table_fqn[2],
+                        column=upstream_column["upstream_column_name"],
+                    )
+                    upstream_column_global_id = new_global_id(
+                        tenant_id=tenant_id, cluster_id=endpoint, data_id=upstream_global_id_arg, data_type="column"
+                    )
+                    upstreams.append(upstream_column_global_id)
+            avro_assets = AvroAsset(
+                id=downstream_column_global_id,
+                object_type="column",
+                parents=[downstream_table_fqn[0], downstream_table_fqn[1], downstream_table_fqn[2]],
+                name=column["DOWNSTREAM_COLUMN_NAME"],
+                upstream=upstreams,
+            )
+            payload.append(avro_assets.to_dict())
+    return payload
+
+
 def gen_table_lineage_payload(
     tenant_id: str, endpoint: str, tables: List[Dict[str, Union[Dict[str, str], str]]]
 ) -> List[LineageInputs]:

quollio_core/profilers/qdc.py

@@ -0,0 +1,27 @@
+import logging
+from io import BytesIO
+from typing import Dict
+
+import fastavro
+
+from quollio_core.helper.core import new_global_id
+from quollio_core.models.avroasset import AvroAsset
+from quollio_core.repository import qdc
+
+logger = logging.getLogger(__name__)
+
+
+def gen_existing_global_id_dict(avro_content: bytes) -> Dict[str, bool]:
+    byte_io = BytesIO(avro_content)
+    avro_schema = AvroAsset.avro_schema_to_python()
+    reader = fastavro.reader(byte_io, avro_schema)
+    records = {record["id"]: True for record in reader}
+    return records
+
+
+def get_avro_file_content(tenant_id: str, account_id: str, qdc_client: qdc.QDCExternalAPIClient) -> bytes:
+    datasource_id = new_global_id(tenant_id=tenant_id, cluster_id=account_id, data_id="", data_type="data_source")
+    logger.debug("Datasource id: {dsrc_id}".format(dsrc_id=datasource_id))
+    res = qdc_client.get_export_url(datasource_id=datasource_id)
+    file_content = qdc_client.download_file(res).content
+    return file_content

quollio_core/profilers/snowflake.py

@@ -1,13 +1,20 @@
+import io
 import logging
-from typing import List
+import os
+from typing import Dict, List
 
+from fastavro import writer
+
+from quollio_core.helper.core import new_global_id
+from quollio_core.models.avroasset import AvroAsset
+from quollio_core.models.qdc import GetImportURLRequest
 from quollio_core.profilers.lineage import (
-    gen_column_lineage_payload,
-    gen_table_lineage_payload,
+    gen_column_avro_lineage_payload,
+    gen_table_avro_lineage_payload,
     parse_snowflake_results,
 )
 from quollio_core.profilers.sqllineage import SQLLineage
-from quollio_core.profilers.stats import gen_table_stats_payload, get_is_target_stats_items, render_sql_for_stats
+from quollio_core.profilers.stats import gen_table_stats_avro_payload, get_is_target_stats_items, render_sql_for_stats
 from quollio_core.repository import qdc, snowflake
 
 logger = logging.getLogger(__name__)
@@ -17,6 +24,7 @@ def snowflake_table_to_table_lineage(
     conn: snowflake.SnowflakeConnectionConfig,
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
+    existing_global_ids: Dict[str, bool],
 ) -> None:
     with snowflake.SnowflakeQueryExecutor(conn) as sf_executor:
         results, err = sf_executor.get_query_results(
@@ -38,28 +46,41 @@ def snowflake_table_to_table_lineage(
             )
             return
         parsed_results = parse_snowflake_results(results=results)
-        update_table_lineage_inputs = gen_table_lineage_payload(
+        update_table_lineage_inputs = gen_table_avro_lineage_payload(
             tenant_id=tenant_id,
             endpoint=conn.account_id,
             tables=parsed_results,
+            existing_global_ids=existing_global_ids,
+        )
+        stack_name = os.getenv("CF_STACK")
+        import_req = GetImportURLRequest(
+            service_name="snowflake",
+            source_name=stack_name,
+            file_name="{name}.avro".format(name=stack_name),
+            override_logical_name="false",
+            update_mode="partial",
+        )
+        datasource_id = new_global_id(
+            tenant_id=tenant_id, cluster_id=conn.account_id, data_id="", data_type="data_source"
         )
+        logger.debug("Datasource id: {dsrc_id}".format(dsrc_id=datasource_id))
+        import_res = qdc_client.get_import_url(datasource_id=datasource_id, payload=import_req)
+        if import_res is None:
+            logger.error("get_import_url failed. Please retry `load_lineage` again")
+            return
+        logger.debug("ImportResponse: {res}".format(res=import_res))
 
-        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.")
+        avro_schema = AvroAsset.avro_schema_to_python()
+
+        buffer = io.BytesIO()
+        writer(buffer, avro_schema, update_table_lineage_inputs)
+        res = qdc_client.upload_file(
+            url=import_res.location,
+            metadata=import_res.datasource_metadata_response_body,
+            buffer=buffer.getbuffer().tobytes(),
+        )
+        if res == 200:
+            logger.info("Upload table lineage is finished.")
     return
 
 
@@ -67,6 +88,7 @@ def snowflake_column_to_column_lineage(
     conn: snowflake.SnowflakeConnectionConfig,
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
+    existing_global_ids: Dict[str, bool],
 ) -> None:
     with snowflake.SnowflakeQueryExecutor(conn) as sf_executor:
         results, err = sf_executor.get_query_results(
@@ -87,29 +109,39 @@ def snowflake_column_to_column_lineage(
                 "No lineage data in ACCOUNT_USAGE.SNOWFLAKE. Please check the data in `QUOLLIO_LINEAGE_COLUMN_LEVEL`."
             )
             return
-        update_column_lineage_inputs = gen_column_lineage_payload(
-            tenant_id=tenant_id,
-            endpoint=conn.account_id,
-            columns=results,
+        update_column_lineage_inputs = gen_column_avro_lineage_payload(
+            tenant_id=tenant_id, endpoint=conn.account_id, columns=results, existing_global_ids=existing_global_ids
         )
 
-        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.")
+        stack_name = os.getenv("CF_STACK")
+        import_req = GetImportURLRequest(
+            service_name="snowflake",
+            source_name=stack_name,
+            file_name="{name}.avro".format(name=stack_name),
+            override_logical_name="false",
+            update_mode="partial",
+        )
+        datasource_id = new_global_id(
+            tenant_id=tenant_id, cluster_id=conn.account_id, data_id="", data_type="data_source"
+        )
+        logger.debug("Datasource id: {dsrc_id}".format(dsrc_id=datasource_id))
+        import_res = qdc_client.get_import_url(datasource_id=datasource_id, payload=import_req)
+        if import_res is None:
+            logger.error("get_import_url failed. Please retry load_lineage again")
+            return
+        logger.debug("ImportResponse: {res}".format(res=import_res))
+
+        avro_schema = AvroAsset.avro_schema_to_python()
+
+        buffer = io.BytesIO()
+        writer(buffer, avro_schema, update_column_lineage_inputs)
+        res = qdc_client.upload_file(
+            url=import_res.location,
+            metadata=import_res.datasource_metadata_response_body,
+            buffer=buffer.getbuffer().tobytes(),
+        )
+        if res == 200:
+            logger.info("Upload column lineage is finished.")
     return
 
 
@@ -177,6 +209,7 @@ def snowflake_table_stats(
     qdc_client: qdc.QDCExternalAPIClient,
     tenant_id: str,
     stats_items: List[str],
+    existing_global_ids: Dict[str, bool],
 ) -> None:
     with snowflake.SnowflakeQueryExecutor(conn) as sf_executor:
         get_stats_view_query = _gen_get_stats_views_query(
@@ -193,8 +226,8 @@ and fix it or grant usage permission to both `{conn.account_database}` and `{con
 and select permissions to views begins with `QUOLLIO_STATS_COLUMNS_`."
             )
             return
-        req_count = 0
         is_aggregate_items = get_is_target_stats_items(stats_items=stats_items)
+        update_stats_inputs = list()
         for stats_view in stats_views:
             table_fqn = '"{catalog}"."{schema}"."{table}"'.format(
                 catalog=stats_view["TABLE_CATALOG"], schema=stats_view["TABLE_SCHEMA"], table=stats_view["TABLE_NAME"]
@@ -210,23 +243,43 @@ and select permissions to views begins with `QUOLLIO_STATS_COLUMNS_`."
 or user has select permission to it."
                 )
                 continue
-            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.")
+            payloads = gen_table_stats_avro_payload(
+                tenant_id=tenant_id,
+                endpoint=conn.account_id,
+                stats=stats_result,
+                existing_global_ids=existing_global_ids,
+            )
+            update_stats_inputs += payloads
+
+        stack_name = os.getenv("CF_STACK")
+        import_req = GetImportURLRequest(
+            service_name="snowflake",
+            source_name=stack_name,
+            file_name="{name}.avro".format(name=stack_name),
+            override_logical_name="false",
+            update_mode="partial",
+        )
+        datasource_id = new_global_id(
+            tenant_id=tenant_id, cluster_id=conn.account_id, data_id="", data_type="data_source"
+        )
+        logger.debug("Datasource id: {dsrc_id}".format(dsrc_id=datasource_id))
+        import_res = qdc_client.get_import_url(datasource_id=datasource_id, payload=import_req)
+        if import_res is None:
+            logger.error("get_import_url failed. Please retry load_stats again")
+            return
+        logger.debug("ImportResponse: {res}".format(res=import_res))
+
+        avro_schema = AvroAsset.avro_schema_to_python()
+
+        buffer = io.BytesIO()
+        writer(buffer, avro_schema, update_stats_inputs)
+        res = qdc_client.upload_file(
+            url=import_res.location,
+            metadata=import_res.datasource_metadata_response_body,
+            buffer=buffer.getbuffer().tobytes(),
+        )
+        if res == 200:
+            logger.info("Generating table stats is finished.")
     return
 
 

quollio_core/profilers/stats.py

@@ -6,6 +6,7 @@ from typing import Dict, List, Tuple, Union
 from jinja2 import Template
 
 from quollio_core.helper.core import new_global_id
+from quollio_core.models.avroasset import AvroAsset
 
 logger = logging.getLogger(__name__)
 
@@ -70,6 +71,41 @@ def convert_value_type(obj, cast_str: bool = False):
     return obj
 
 
+def gen_table_stats_avro_payload(
+    tenant_id: str, endpoint: str, stats: List[Dict[str, str]], existing_global_ids: Dict[str, bool]
+) -> List[Dict[str, str]]:
+    payloads = list()
+    for stat in stats:
+        db_name = stat.get("DB_NAME", stat.get("db_name"))
+        schema_name = stat.get("SCHEMA_NAME", stat.get("schema_name"))
+        table_name = stat.get("TABLE_NAME", stat.get("table_name"))
+        column_name = stat.get("COLUMN_NAME", stat.get("column_name"))
+        global_id_arg = "{db}{schema}{table}{column}".format(
+            db=db_name, schema=schema_name, table=table_name, column=column_name
+        )
+        column_global_id = new_global_id(
+            tenant_id=tenant_id, cluster_id=endpoint, data_id=global_id_arg, data_type="column"
+        )
+        if existing_global_ids.get(column_global_id) is not True:
+            continue
+        avro_assets = AvroAsset(
+            id=column_global_id,
+            object_type="column",
+            parents=[db_name, schema_name, table_name],
+            name=column_name,
+            stats_max=convert_value_type(stat.get("MAX_VALUE", stat.get("max_value")), True),
+            stats_min=convert_value_type(stat.get("MIN_VALUE", stat.get("min_value")), True),
+            stats_mean=convert_value_type(stat.get("AVG_VALUE", stat.get("avg_value")), True),
+            stats_median=convert_value_type(stat.get("MEDIAN_VALUE", stat.get("median_value")), True),
+            stats_mode=convert_value_type(stat.get("MODE_VALUE", stat.get("mode_value")), True),
+            stats_stddev=convert_value_type(stat.get("STDDEV_VALUE", stat.get("stddev_value")), True),
+            stats_number_of_null=convert_value_type(stat.get("NULL_COUNT", stat.get("null_count")), True),
+            stats_number_of_unique=convert_value_type(stat.get("CARDINALITY", stat.get("cardinality")), True),
+        )
+        payloads.append(avro_assets.to_dict())
+    return payloads
+
+
 def gen_table_stats_payload(tenant_id: str, endpoint: str, stats: List[Dict[str, str]]) -> List[StatsRequest]:
     payloads = list()
     for stat in stats:

quollio_core/repository/qdc.py

@@ -8,6 +8,8 @@ import jwt
 import requests  # type: ignore
 from requests.exceptions import ConnectionError, HTTPError, RequestException, Timeout
 
+from quollio_core.models.qdc import DataSourceMetadataResponseBody, GetImportURLRequest, GetImportURLResponse
+
 logger = logging.getLogger(__name__)
 
 
@@ -64,6 +66,95 @@ class QDCExternalAPIClient:
         session.mount("https://", requests.adapters.HTTPAdapter(max_retries=retry))
         return session
 
+    def get_export_url(self, datasource_id: str) -> GetImportURLResponse:
+        self._refresh_token_if_expired()
+        headers = {"content-type": "application/json", "authorization": f"Bearer {self.auth_token}"}
+        endpoint = "{base_url}/v2/sources/{dsrc_id}/export-avro".format(base_url=self.base_url, dsrc_id=datasource_id)
+        try:
+            res = self.session.post(endpoint, headers=headers, data={})
+            res.raise_for_status()
+        except ConnectionError as ce:
+            logger.error(f"Connection Error: {ce} global_id: {datasource_id}.")
+        except HTTPError as he:
+            logger.error(f"HTTP Error: {he} global_id: {datasource_id}.")
+        except Timeout as te:
+            logger.error(f"Timeout Error: {te} global_id: {datasource_id}.")
+        except RequestException as re:
+            logger.error(f"RequestException Error: {re} global_id: {datasource_id}.")
+        else:
+            res = json.loads(res.text)
+            location = res.get("data").get("location")
+            return location
+
+    def download_file(self, url: str) -> requests.Response:
+        self._refresh_token_if_expired()
+
+        try:
+            res = self.session.get(url)
+            res.raise_for_status()
+        except ConnectionError as ce:
+            logger.error(f"Connection Error: {ce}.")
+        except HTTPError as he:
+            logger.error(f"HTTP Error: {he}.")
+        except Timeout as te:
+            logger.error(f"Timeout Error: {te}")
+        except RequestException as re:
+            logger.error(f"RequestException Error: {re}")
+        else:
+            return res
+
+    def get_import_url(self, datasource_id: str, payload: GetImportURLRequest) -> GetImportURLResponse:
+        self._refresh_token_if_expired()
+        headers = {"content-type": "application/json", "authorization": f"Bearer {self.auth_token}"}
+        endpoint = "{base_url}/v2/sources/{dsrc_id}/import".format(base_url=self.base_url, dsrc_id=datasource_id)
+        try:
+            payload_dict = payload.as_dict()
+            res = self.session.post(endpoint, headers=headers, json=payload_dict)
+            logger.debug(f"Got the result of import_url request: {res.text}")
+            res.raise_for_status()
+        except ConnectionError as ce:
+            logger.error(f"Connection Error: {ce} global_id: {datasource_id}.")
+        except HTTPError as he:
+            logger.error(f"HTTP Error: {he} global_id: {datasource_id}.")
+        except Timeout as te:
+            logger.error(f"Timeout Error: {te} global_id: {datasource_id}.")
+        except RequestException as re:
+            logger.error(f"RequestException Error: {re} global_id: {datasource_id}.")
+        else:
+            res = json.loads(res.text)
+            datasource_metadata_response = DataSourceMetadataResponseBody(**res.get("data").get("metadata"))
+            location = res.get("data").get("location")
+            response = GetImportURLResponse(
+                location=location, datasource_metadata_response_body=datasource_metadata_response
+            )
+            return response
+
+    def upload_file(self, url: str, metadata: DataSourceMetadataResponseBody, buffer: bytes):
+        self._refresh_token_if_expired()
+        headers = {
+            "Content-Type": "application/octet-stream",
+            "x-amz-meta-user_id": metadata.user_id,
+            "x-amz-meta-job_key": metadata.job_key,
+            "x-amz-meta-service_name": metadata.service_name,
+            "x-amz-meta-source_name": metadata.source_name,
+            "x-amz-meta-source_type": metadata.source_type,
+            "x-amz-meta-override_logical_name": metadata.override_logical_name,
+            "Content-Length": str(len(buffer)),
+        }
+        try:
+            res = self.session.put(url, headers=headers, data=buffer)
+            res.raise_for_status()
+        except ConnectionError as ce:
+            logger.error(f"Connection Error: {ce}.")
+        except HTTPError as he:
+            logger.error(f"HTTP Error: {he}.")
+        except Timeout as te:
+            logger.error(f"Timeout Error: {te}")
+        except RequestException as re:
+            logger.error(f"RequestException Error: {re}")
+        else:
+            return res.status_code
+
     def update_stats_by_id(self, global_id: str, payload: Dict[str, List[str]]) -> int:
         self._refresh_token_if_expired()
         headers = {"content-type": "application/json", "authorization": f"Bearer {self.auth_token}"}

quollio_core/repository/snowflake.py

@@ -1,7 +1,9 @@
 import logging
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass
 from typing import Dict, List, Tuple
 
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
 from snowflake.connector import DictCursor, connect, errors
 from snowflake.connector.connection import SnowflakeConnection
 
@@ -12,16 +14,67 @@ logger = logging.getLogger(__name__)
 class SnowflakeConnectionConfig:
     account_id: str
     account_user: str
-    account_password: str
     account_build_role: str
     account_query_role: str
     account_warehouse: str
     account_database: str
     account_schema: str
+    account_password: str = None
+    private_key: str = None
     threads: int = 3
 
     def as_dict(self) -> Dict[str, str]:
-        return asdict(self)
+        """Convert config to dictionary, handling both auth methods for DBT."""
+        base_params = {
+            "account_id": self.account_id,
+            "account_user": self.account_user,
+            "account_build_role": self.account_build_role,
+            "account_query_role": self.account_query_role,
+            "account_warehouse": self.account_warehouse,
+            "account_database": self.account_database,
+            "account_schema": self.account_schema,
+            "threads": self.threads,
+        }
+
+        # Add auth parameters based on method
+        if self.private_key:
+            # Keep private key as is, template will handle formatting
+            base_params["private_key"] = self.private_key
+        elif self.account_password:
+            base_params["account_password"] = self.account_password
+
+        return {k: v for k, v in base_params.items() if v is not None}
+
+    def get_connection_params(self) -> Dict[str, str]:
+        """Get the appropriate connection parameters based on authentication method."""
+        params = {
+            "user": self.account_user,
+            "account": self.account_id,
+            "warehouse": self.account_warehouse,
+            "database": self.account_database,
+            "schema": self.account_schema,
+            "role": self.account_query_role,
+        }
+
+        # Add authentication parameters based on method
+        if self.private_key:
+            try:
+                # Parse private key content into RSA key object
+                pkey = serialization.load_pem_private_key(
+                    self.private_key.encode("utf-8"),
+                    password=None,
+                    backend=default_backend(),
+                )
+                params["private_key"] = pkey
+            except Exception as e:
+                logger.error(f"Failed to parse private key: {str(e)}")
+                raise
+        elif self.account_password:
+            params["password"] = self.account_password
+        else:
+            raise ValueError("Either password or private key authentication must be configured")
+
+        return params
 
 
 class SnowflakeQueryExecutor:
@@ -35,16 +88,12 @@ class SnowflakeQueryExecutor:
         self.conn.close()
 
     def __initialize(self, config: SnowflakeConnectionConfig) -> SnowflakeConnection:
-        conn: SnowflakeConnection = connect(
-            user=config.account_user,
-            password=config.account_password,
-            role=config.account_query_role,
-            account=config.account_id,
-            warehouse=config.account_warehouse,
-            database=config.account_database,
-            schema=config.account_schema,
-        )
-        return conn
+        try:
+            conn: SnowflakeConnection = connect(**config.get_connection_params())
+            return conn
+        except Exception as e:
+            logger.error(f"Failed to initialize Snowflake connection: {str(e)}")
+            raise
 
     def get_query_results(self, query: str) -> Tuple[List[Dict[str, str]], Exception]:
         with self.conn.cursor(DictCursor) as cur:

quollio_core/snowflake.py

@@ -6,6 +6,7 @@ import shutil
 from quollio_core.helper.core import is_valid_domain, 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.qdc import gen_existing_global_id_dict, get_avro_file_content
 from quollio_core.profilers.snowflake import (
     snowflake_column_to_column_lineage,
     snowflake_table_level_sqllineage,
@@ -93,10 +94,15 @@ def load_lineage(
 ) -> None:
     logger.info("Generate Snowflake table to table lineage.")
 
-    snowflake_table_to_table_lineage(
-        conn=conn,
-        qdc_client=qdc_client,
+    file_content = get_avro_file_content(
         tenant_id=tenant_id,
+        account_id=conn.account_id,
+        qdc_client=qdc_client,
+    )
+    existing_global_ids = gen_existing_global_id_dict(avro_content=file_content)
+
+    snowflake_table_to_table_lineage(
+        conn=conn, qdc_client=qdc_client, tenant_id=tenant_id, existing_global_ids=existing_global_ids
     )
 
     if enable_column_lineage:
@@ -104,9 +110,7 @@ def load_lineage(
             f"enable_column_lineage is set to {enable_column_lineage}.Generate Snowflake column to column lineage."
         )
         snowflake_column_to_column_lineage(
-            conn=conn,
-            qdc_client=qdc_client,
-            tenant_id=tenant_id,
+            conn=conn, qdc_client=qdc_client, tenant_id=tenant_id, existing_global_ids=existing_global_ids
         )
     else:
         logger.info("Skip column lineage ingestion. Set enable_column_lineage to True if you ingest column lineage.")
@@ -124,6 +128,13 @@ def load_stats(
 ) -> None:
     logger.info("Generate Snowflake stats.")
 
+    file_content = get_avro_file_content(
+        tenant_id=tenant_id,
+        account_id=conn.account_id,
+        qdc_client=qdc_client,
+    )
+    existing_global_ids = gen_existing_global_id_dict(avro_content=file_content)
+
     if stats_items is None:
         raise ValueError("No stats items are not selected. Please specify any value to `stats_items` param.")
 
@@ -133,6 +144,7 @@ def load_stats(
         qdc_client=qdc_client,
         tenant_id=tenant_id,
         stats_items=stats_items,
+        existing_global_ids=existing_global_ids,
     )
 
     logger.info("Stats data is successfully finished.")
@@ -237,7 +249,11 @@ if __name__ == "__main__":
         "--target_tables",
         type=str,
         nargs="*",
-        choices=["quollio_lineage_column_level", "quollio_lineage_table_level", "quollio_stats_columns"],
+        choices=[
+            "quollio_lineage_column_level",
+            "quollio_lineage_table_level",
+            "quollio_stats_columns",
+        ],
         action=env_default("SNOWFLAKE_TARGET_TABLES"),
         required=False,
         help="Target table name if you want to create only specific tables. \
@@ -335,6 +351,22 @@ if __name__ == "__main__":
         help="Access method to Quollio API. Default 'PUBLIC'. Choose 'VPC_ENDPOINT'\
              if you use API Gateway VPC Endpoint, DefaultValue is set to PUBLIC.",
     )
+    parser.add_argument(
+        "--auth_type",
+        type=str,
+        choices=["PASSWORD", "KEYPAIR"],
+        action=env_default("SNOWFLAKE_AUTH_TYPE"),
+        default="PASSWORD",
+        required=False,
+        help="Authentication method to use (PASSWORD or KEYPAIR)",
+    )
+    parser.add_argument(
+        "--private_key",
+        type=str,
+        action=env_default("SNOWFLAKE_PRIVATE_KEY"),
+        required=False,
+        help="Private key content for keypair authentication",
+    )
 
     stats_items = get_column_stats_items()
     parser.add_argument(
@@ -351,16 +383,31 @@ if __name__ == "__main__":
     args = parser.parse_args()
     set_log_level(level=args.log_level)
 
-    conn = snowflake.SnowflakeConnectionConfig(
-        account_id=args.account_id,
-        account_user=args.user,
-        account_password=args.password,
-        account_build_role=args.build_role,
-        account_query_role=args.query_role,
-        account_warehouse=args.warehouse,
-        account_database=args.database,
-        account_schema=args.schema,
-    )
+    # Update authentication handling
+    auth_params = {
+        "account_id": args.account_id,
+        "account_user": args.user,
+        "account_build_role": args.build_role,
+        "account_query_role": args.query_role,
+        "account_warehouse": args.warehouse,
+        "account_database": args.database,
+        "account_schema": args.schema,
+    }
+
+    # Add authentication specific parameters based on method
+    if args.auth_type == "KEYPAIR":
+        if not args.private_key:
+            raise ValueError("private_key is required when using keypair authentication")
+        auth_params["private_key"] = args.private_key
+        logger.info("Using keypair authentication")
+    else:
+        if not args.password:
+            raise ValueError("password is required when using password authentication")
+        auth_params["account_password"] = args.password
+        logger.info("Using password authentication")
+        logger.warning("Password authentication is being deprecated. Please consider using keypair authentication.")
+
+    conn = snowflake.SnowflakeConnectionConfig(**auth_params)
 
     if len(args.commands) == 0:
         raise ValueError("No command is provided")

{quollio_core-0.5.3.dist-info → quollio_core-0.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: quollio-core
-Version: 0.5.3
+Version: 0.6.1
 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>
@@ -17,6 +17,7 @@ Classifier: Operating System :: OS Independent
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
+License-File: LICENSE
 Requires-Dist: blake3==0.3.3
 Requires-Dist: dbt-core==1.7.10
 Requires-Dist: dbt-snowflake==1.7.0
@@ -38,6 +39,8 @@ Requires-Dist: google-cloud-datacatalog==3.19.0
 Requires-Dist: google-cloud-datacatalog-lineage==0.3.6
 Requires-Dist: google-api-python-client==2.131.0
 Requires-Dist: teradatasql==20.0.0.15
+Requires-Dist: dataclasses_avroschema==0.63.9
+Requires-Dist: fastavro==1.9.7
 Requires-Dist: black>=22.3.0 ; extra == "test"
 Requires-Dist: coverage>=7.3.2 ; extra == "test"
 Requires-Dist: isort>=5.10.1 ; extra == "test"

{quollio_core-0.5.3.dist-info → quollio_core-0.6.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
-quollio_core/__init__.py,sha256=YrefYO__k4ewIsYGUmshPK7ySvzQxHKIdRv7kTYGaVk,83
+quollio_core/__init__.py,sha256=UKsQb_ngmTusubotU4aVGiWhbSJfs4FBWUasBuPnlg0,83
 quollio_core/bigquery.py,sha256=6Oq4DVGpa3X21Es_nbrsb8pK3vaxwb9Egnvq3huo95k,5894
 quollio_core/bricks.py,sha256=8h3kbI2b6lGH2s-56jE_Q5-R5-nIsQYMfvtRrkFOzoU,10784
 quollio_core/redshift.py,sha256=KcdljY95xYf9JYrsaMOBoP_XxQQ8wFVE5ue_XEMVSFc,11504
-quollio_core/snowflake.py,sha256=wKBfyoqdoLN-kYsceT0ctqWeay_Sn8OpKMSYedlBjD4,14558
+quollio_core/snowflake.py,sha256=YM9rnlk7iKY7rg2BxEPQgKtCuHf17Tb5HTr4a7YV8oo,16467
 quollio_core/teradata.py,sha256=H2VUcJvr8W-M2wvm3710Gf1ENb-BSscrDRKNm8gdHJE,8227
 quollio_core/dbt_projects/databricks/.gitignore,sha256=1jJAyXSzJ3YUm0nx3i7wUSE4RjQMX3ad6F8O88UbtzI,29
 quollio_core/dbt_projects/databricks/README.md,sha256=ZpRQyhFAODAiS8dc1Kb_ndkul4cu4o4udN_EMa49CU4,440
@@ -59,34 +59,37 @@ quollio_core/dbt_projects/snowflake/models/quollio_stats_columns.yml,sha256=V_BE
 quollio_core/dbt_projects/snowflake/models/quollio_stats_profiling_columns.sql,sha256=ubMEzZNHq55zCncg7HbzdMKMSdqHnwbJmVKYpet8Otc,2968
 quollio_core/dbt_projects/snowflake/models/quollio_stats_profiling_columns.yml,sha256=W39VAmFnnX6RBoW7B_4CConC1lm0Jm9o50Jsz9bYZzY,538
 quollio_core/dbt_projects/snowflake/models/sources.yml,sha256=vGSV33cNj4UUyPUcYS-JFgc3r8KvSLfiA7qhbDCUU9s,10975
-quollio_core/dbt_projects/snowflake/profiles/profiles_template.yml,sha256=gcZsgdGP461QuUM9jLbBKdadT8cHTXgNarq_azOOMhk,379
+quollio_core/dbt_projects/snowflake/profiles/profiles_template.yml,sha256=6yIlaFSKxWxqnX-LboO1xVC-7DtPYcAKrRjn3rPGNfk,511
 quollio_core/dbt_projects/snowflake/seeds/.gitkeep,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 quollio_core/dbt_projects/snowflake/snapshots/.gitkeep,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 quollio_core/helper/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-quollio_core/helper/core.py,sha256=ddV5VLa49Je11eHUjlRN5m3nhtqPMMUSeG4135HuZC8,1390
+quollio_core/helper/core.py,sha256=PMpc80fWVf2tdf3J1qtK9WghXgU2l7i72ufamDeHwcw,1422
 quollio_core/helper/env_default.py,sha256=H6gbSGUPrEDZr4YDrL49hbOpw6RntI4U82kX1q6vUnI,2148
 quollio_core/helper/log.py,sha256=flxyZZ44G79l1TaUp3OT58uCHcnE5z_pCduwoeI6IUs,645
 quollio_core/helper/log_utils.py,sha256=QontLKETHjSAbQniJ7YqS0RY2AYvFHSjrlPiGr317LE,1391
+quollio_core/models/avroasset.py,sha256=YZHzOS62N0_sidneXI3IZ2MA8Bz1vFVgF6F9_UilC3s,603
+quollio_core/models/qdc.py,sha256=UObaUpvAQ4vOhI6jfwvNFrJ3--6AX2v9yl9_d3Juy7M,739
 quollio_core/profilers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 quollio_core/profilers/bigquery.py,sha256=mEr7CFQNgBFqWR8XfCOk8WTm5k5qZhLF8ODVWfskPRI,5797
 quollio_core/profilers/databricks.py,sha256=ik4RiR_GOeU3S7s6C6Y9SGe1D_Y_f98BDWJVlEJXL4U,7868
-quollio_core/profilers/lineage.py,sha256=4FyxIuPBrUFihqZryqTQBcfB0Z7634lKl_WwkD82vzE,6865
+quollio_core/profilers/lineage.py,sha256=GMWue6lgiz7wFYnNpaHVFivprA-iqhbCHf63IsVB8Vk,11260
+quollio_core/profilers/qdc.py,sha256=P0STRfe5G4d-UI7RdVbYmAfP_tAn1HbFUjeirxWipz4,995
 quollio_core/profilers/redshift.py,sha256=p6ONDCkhndZAOcKAwEyQ5fsi-jsQrlwHHb7LTI_m1uk,6473
-quollio_core/profilers/snowflake.py,sha256=m9Ivv2LRwnrmgKS36a039AhrO27sR1EaOOdqNF26PhI,11156
+quollio_core/profilers/snowflake.py,sha256=ewvULWIlcq2h0jOyRzUpedW0NS8QlkSgICS-dZDYl18,13027
 quollio_core/profilers/sqllineage.py,sha256=h0FT6CYb0A20zSc68GELZ7Q8bDbaHLQnZQHsXBEXBug,5261
-quollio_core/profilers/stats.py,sha256=OLQrdrh0y64jo9rmzvGlDdxy_c7gMz_GnlXPJzWkBjM,7343
+quollio_core/profilers/stats.py,sha256=PwMNyr7JwdrVbSSfVd-XsgCZV7sKoRjihaRhhZ88uyM,9328
 quollio_core/profilers/teradata/lineage.py,sha256=2wNksBQD8vC6UTQwCglPsF53YMEVIkAb2CWTmpiTHDU,7368
 quollio_core/profilers/teradata/stats.py,sha256=OagvkTRFiWVbiLABwZwR3wQ7y36edwOViDetHsYiyxI,9277
 quollio_core/repository/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 quollio_core/repository/bigquery.py,sha256=3AyGcJNYGnUyMweyc6lGm4quwrOzd-ZBS2zNnFwafII,3990
 quollio_core/repository/databricks.py,sha256=9Cgdv8qBnVaHqu3RA-IUBieAqb69moQ-KAAMVSf5Ds4,1877
 quollio_core/repository/dbt.py,sha256=cnLwJPywLi8VowVW7zfIBa9jxVwDWO7xzzNRn1vWiuw,659
-quollio_core/repository/qdc.py,sha256=IPGiYafnJYkuD7_kLapVM98-9ZhEwq5S-dGY2bO8fVk,4624
+quollio_core/repository/qdc.py,sha256=Ni0rk9CX8ienqM_HYLuWLBKTwycvTENC7x7wGWhzjXs,8978
 quollio_core/repository/redshift.py,sha256=p2ouEuYcDCjx1oBhc6H1ekQsvEqHGd3bFu3PW0ngYBc,2880
-quollio_core/repository/snowflake.py,sha256=zL9-xi98AIftdW9MuKI-M3pZ1kQuuH-UiZH8HcJvmk4,1769
+quollio_core/repository/snowflake.py,sha256=yCYXrYf4I5GL_ITNTXoggj0xNbQsdwxPSmsVvZYwUVU,3869
 quollio_core/repository/ssm.py,sha256=xpm1FzbBnIsBptuYPUNnPgkKU2AH3XxI-ZL0bEetvW0,2182
 quollio_core/repository/teradata.py,sha256=1AExxRBTswpSyF4OVyAUkoiZ0yVRfqt4T99FdllkTEI,3763
-quollio_core-0.5.3.dist-info/LICENSE,sha256=V8j_M8nAz8PvAOZQocyRDX7keai8UJ9skgmnwqETmdY,34520
-quollio_core-0.5.3.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
-quollio_core-0.5.3.dist-info/METADATA,sha256=aQkusHnQ6mw4E6KUM5Yo4LrzFFFHqDbIgIyzjHZFJc4,6924
-quollio_core-0.5.3.dist-info/RECORD,,
+quollio_core-0.6.1.dist-info/licenses/LICENSE,sha256=V8j_M8nAz8PvAOZQocyRDX7keai8UJ9skgmnwqETmdY,34520
+quollio_core-0.6.1.dist-info/WHEEL,sha256=_2ozNFCLWc93bK4WKHCO-eDUENDlo-dgc9cU3qokYO4,82
+quollio_core-0.6.1.dist-info/METADATA,sha256=ZdX8sgnt97nINrc021GvEp-o-i2Ph01vRLtjc_3J8fg,7023
+quollio_core-0.6.1.dist-info/RECORD,,

{quollio_core-0.5.3.dist-info → quollio_core-0.6.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: flit 3.10.1
+Generator: flit 3.11.0
 Root-Is-Purelib: true
 Tag: py3-none-any