quollio-core 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

quollio_core/profilers/teradata/stats.py

@@ -27,6 +27,7 @@ def load_stats(
     target_databases: Optional[List[str]] = None,
     target_databases_method: str = "DENYLIST",
     stats_items: Optional[List[str]] = None,
+    system_database: Optional[str] = None,
 ) -> None:
     stats_list = []
     numerical_columns = 0
@@ -35,16 +36,21 @@ def load_stats(
         f"Starting statistics collection. " f"Sample percent: {sample_percent if sample_percent is not None else 'N/A'}"
     )
 
+    # Use system_database from config if not provided
+    system_database = system_database or conn_config.system_database
+
     with teradata_repo.new_teradata_client(conn_config) as conn:
         try:
-            tables = teradata_repo.get_table_list(conn, target_databases, target_databases_method)
+            tables = teradata_repo.get_table_list(conn, target_databases, target_databases_method, system_database)
             for table in tables:
                 logger.debug(f"Processing table: {table}")
-                database_name = table["DataBaseName"]
+                database_name = table["DatabaseName"]
                 table_name = table["TableName"]
 
                 logger.info(f"Processing table {database_name}.{table_name}")
-                columns = teradata_repo.get_column_list(conn, database_name=database_name, table_name=table_name)
+                columns = teradata_repo.get_column_list(
+                    conn, database_name=database_name, table_name=table_name, system_database=system_database
+                )
                 logger.debug(f"Columns: {columns}")
 
                 for column in columns:

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

@@ -13,18 +13,25 @@ class TeradataConfig:
     username: str
     password: str
     database: str = "DBC"
+    system_database: str = "DBC"
     encrypt_data: bool = True
     additional_params: Dict[str, Any] = field(default_factory=dict)
 
     @classmethod
     def from_dict(
-        cls, credentials: Dict[str, str], host: str, port: str, additional_params: Dict[str, Any] = None
+        cls,
+        credentials: Dict[str, str],
+        host: str,
+        port: str,
+        additional_params: Dict[str, Any] = None,
+        system_database: str = "DBC",
     ) -> "TeradataConfig":
         return cls(
             host=host,
             port=int(port),
             username=credentials["username"],
             password=credentials["password"],
+            system_database=system_database,
             additional_params=additional_params or {},
         )
 
@@ -49,7 +56,10 @@ def new_teradata_client(config: TeradataConfig) -> teradatasql.connect:
 
 @error_handling_decorator
 def get_table_list(
-    conn: teradatasql.connect, target_databases: Optional[List[str]] = None, target_databases_method: str = "DENYLIST"
+    conn: teradatasql.connect,
+    target_databases: Optional[List[str]] = None,
+    target_databases_method: str = "DENYLIST",
+    system_database: str = "DBC",
 ) -> List[Dict[str, str]]:
     if target_databases_method == "DENYLIST":
         operator = "NOT"
@@ -58,24 +68,28 @@ def get_table_list(
 
     query_tables = f"""
     SELECT DatabaseName, TableName
-    FROM DBC.TablesV
+    FROM {system_database}.TablesV
     WHERE TableKind IN ('T', 'O', 'Q')
         AND DatabaseName {operator} IN ({','.join("'" + db + "'" for db in target_databases)})
     """
     logger.debug("Executing query to retrieve table names.")
+    logger.debug(f"Query: {query_tables}")
     tables = execute_query(query_tables, conn)
     return tables
 
 
 @error_handling_decorator
-def get_column_list(conn: teradatasql.connect, database_name: str, table_name: str) -> List[Dict[str, str]]:
+def get_column_list(
+    conn: teradatasql.connect, database_name: str, table_name: str, system_database: str = "DBC"
+) -> List[Dict[str, str]]:
     query_columns = f"""
     SELECT ColumnName, ColumnType
-    FROM DBC.ColumnsV
+    FROM {system_database}.ColumnsV
     WHERE DatabaseName = '{database_name}'
         AND TableName = '{table_name}'
     """
     logger.debug(f"Executing query to retrieve columns for {database_name}.{table_name}.")
+    logger.debug(f"Query: {query_columns}")
     columns = execute_query(query_columns, conn)
     logger.debug(f"Retrieved columns: {columns}")
     return columns

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,
@@ -24,6 +25,8 @@ def build_view(
     target_tables: str = "",
     log_level: str = "info",
     dbt_macro_source: str = "hub",
+    target_databases_method: str = "DENYLIST",
+    target_databases: list[str] = [],
 ) -> None:
     logger.info("Build profiler views using dbt")
     # set parameters
@@ -32,10 +35,16 @@ def build_view(
     project_path = f"{current_dir}/dbt_projects/snowflake"
     template_path = f"{current_dir}/dbt_projects/snowflake/profiles"
     template_name = "profiles_template.yml"
-    options = '{{"query_role": {query_role}, "sample_method": {sample_method}}}'.format(
+
+    options = '{{"query_role": "{query_role}", "sample_method": "{sample_method}",\
+          "target_databases_method": "{target_databases_method}",\
+              "target_databases": {target_databases}}}'.format(
         query_role=conn.account_query_role,
         sample_method=stats_sample_method,
+        target_databases_method=target_databases_method,
+        target_databases=target_databases,
     )
+
     new_package_file = f"{project_path}/packages.yml"
     if dbt_macro_source == "local":
         shutil.copyfile(f"{project_path}/packages_local.yml", new_package_file)
@@ -85,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:
@@ -96,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.")
@@ -116,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.")
 
@@ -125,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.")
@@ -229,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. \
@@ -237,6 +261,25 @@ if __name__ == "__main__":
               Please specify table name with blank delimiter like tableA tableB \
               if you want to create two or more tables.",
     )
+    parser.add_argument(
+        "--target_databases_method",
+        type=str,
+        choices=["ALLOWLIST", "DENYLIST"],
+        action=env_default("SNOWFLAKE_TARGET_DATABASE_METHOD"),
+        required=False,
+        help="Method to filter databases. 'ALLOWLIST' to only include listed databases,\
+              'DENNYLIST' to exclude listed databases",
+    )
+    parser.add_argument(
+        "--target_databases",
+        type=str,
+        nargs="*",
+        action=env_default("SNOWFLAKE_TARGET_DATABASES"),
+        required=False,
+        help='List of databases to allow or deny based on target_database_method\
+            please specify database names with blank space as delimiter\
+            wildcards (%) are supported "DATABASE%" ',
+    )
     parser.add_argument(
         "--sample_method",
         type=str,
@@ -308,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_method",
+        type=str,
+        choices=["PASSWORD", "KEYPAIR"],
+        action=env_default("SNOWFLAKE_AUTH_METHOD"),
+        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(
@@ -324,27 +383,50 @@ 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_method == "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")
 
     if "build_view" in args.commands:
+
+        if args.target_databases:
+            target_databases = ["'" + db + "'" for db in args.target_databases[0].split(",")]
+        else:
+            target_databases = []
+
         build_view(
             conn=conn,
             stats_sample_method=args.sample_method,
             target_tables=args.target_tables,
             log_level=args.log_level,
             dbt_macro_source=args.dbt_macro_source,
+            target_databases_method=args.target_databases_method,
+            target_databases=target_databases,
         )
     api_url = args.api_url
     if args.external_api_access == "VPC_ENDPOINT":

quollio_core/teradata.py

@@ -189,6 +189,14 @@ def main() -> None:
               You can choose the items to be aggregated for stats.\
               Default is full stats.",
     )
+    parser.add_argument(
+        "--teradata_system_database",
+        type=str,
+        action=env_default("TERADATA_SYSTEM_DATABASE"),
+        default="DBC",
+        help="Name of the Teradata system database.\
+              Default is DBC",
+    )
 
     args = parser.parse_args()
 
@@ -213,7 +221,11 @@ def main() -> None:
 
     logger.info("Initializing Teradata client")
     config = teradata_repo.TeradataConfig.from_dict(
-        credentials, args.teradata_host, args.teradata_port, additional_params
+        credentials=credentials,
+        host=args.teradata_host,
+        port=args.teradata_port,
+        additional_params=additional_params,
+        system_database=args.teradata_system_database,
     )
 
     if "load_lineage" in args.commands:
@@ -224,6 +236,7 @@ def main() -> None:
             endpoint=args.teradata_host,
             qdc_client=qdc_client,
             page_size=args.teradata_page_size,
+            system_database=args.teradata_system_database,
         )
         logger.info("Lineage loading process completed")
 
@@ -244,6 +257,7 @@ def main() -> None:
             target_databases=target_databases,
             target_databases_method=args.teradata_target_databases_method.upper(),
             stats_items=args.target_stats_items,
+            system_database=args.teradata_system_database,
         )
         logger.info("Statistics loading process completed")
 

{quollio_core-0.5.0.dist-info → quollio_core-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: quollio-core
-Version: 0.5.0
+Version: 0.6.0
 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"