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

quollio_core/__init__.py

@@ -1,4 +1,4 @@
 """Quollio Core"""
 
-__version__ = "0.4.19"
+__version__ = "0.5.0"
 __author__ = "Quollio Technologies, Inc"

quollio_core/dbt_projects/snowflake/models/quollio_stats_profiling_columns.sql

@@ -87,7 +87,7 @@ WITH columns AS (
     , data_type
     , case when data_type in('NUMBER','DECIMAL', 'DEC', 'NUMERIC',
                              'INT', 'INTEGER', 'BIGINT', 'SMALLINT',
-                             'TINYINT', 'BYTEINT')
+                             'TINYINT', 'BYTEINT', 'FLOAT')
                              THEN true
            else false END AS is_calculable
   FROM

quollio_core/helper/core.py

@@ -35,3 +35,10 @@ def setup_dbt_profile(connections_json: Dict[str, str], template_path: str, temp
 
 def trim_prefix(s: str, prefix: str) -> str:
     return s.lstrip(prefix)
+
+
+def is_valid_domain(domain: str, domain_type: str) -> bool:
+    if domain_type == "VPC_ENDPOINT":
+        return domain.endswith("/api")
+    else:
+        return domain.endswith(".com")

quollio_core/profilers/sqllineage.py

@@ -67,15 +67,19 @@ class SQLLineage:
             dest_schema = dest_schema.upper() if dest_schema is not None else None
 
         # MEMO: Complement sql with dialect, source database and source schema info.
-        optimized_stmt: sqlglot.Expression = optimizer.qualify.qualify(
-            statement,
-            dialect=dialect,
-            catalog=src_db,
-            db=src_schema,
-            qualify_columns=False,
-            validate_qualify_columns=False,
-            identify=False,
-        )
+        # MEMO: Skipping qualify because it normalizes the table names.
+        if dialect == "teradata":
+            optimized_stmt = statement
+        else:
+            optimized_stmt: sqlglot.Expression = optimizer.qualify.qualify(
+                statement,
+                dialect=dialect,
+                catalog=src_db,
+                db=src_schema,
+                qualify_columns=False,
+                validate_qualify_columns=False,
+                identify=False,
+            )
 
         orig_dest_table = Table(table="")
         dest_table = Table(table="")

quollio_core/profilers/teradata/lineage.py

@@ -0,0 +1,172 @@
+import os
+from collections import OrderedDict
+from typing import Dict, List, Set, Tuple, Union
+
+from sqlglot import ParseError
+
+from quollio_core.helper.log_utils import error_handling_decorator, logger
+from quollio_core.profilers.sqllineage import SQLLineage, Table
+from quollio_core.repository import qdc
+from quollio_core.repository import teradata as teradata_repo
+
+
+@error_handling_decorator
+def load_lineage(
+    conn_config: teradata_repo.TeradataConfig,
+    endpoint: str = None,
+    tenant_id: str = None,
+    qdc_client: qdc.QDCExternalAPIClient = None,
+    page_size: int = None,
+) -> None:
+    page_size = page_size or int(os.environ.get("TERADATA_PAGE_SIZE", 1000))
+    offset = 0
+    all_lineage_results = []
+
+    with teradata_repo.new_teradata_client(conn_config) as conn:
+        while True:
+            query = f"""
+                SELECT
+                    a.QueryID,
+                    TRIM(a.SqlTextInfo) AS SqlTextInfo,
+                    a.SqlRowNo,
+                    TRIM(d.DatabaseName) AS DefaultDatabase
+                FROM DBC.QryLogSQLV a
+                JOIN DBC.QryLogV b
+                    ON a.QueryID = b.QueryID
+                JOIN DBC.DatabasesV d
+                    ON b.DefaultDatabase = d.DatabaseName
+                WHERE
+                    UPPER(TRIM(SqlTextInfo)) LIKE 'CREATE TABLE%'
+                    OR UPPER(TRIM(SqlTextInfo)) LIKE 'CREATE VIEW%'
+                    OR UPPER(TRIM(SqlTextInfo)) LIKE 'INSERT%'
+                    OR UPPER(TRIM(SqlTextInfo)) LIKE 'MERGE%'
+                    OR UPPER(TRIM(SqlTextInfo)) LIKE 'UPDATE%'
+                QUALIFY ROW_NUMBER() OVER (ORDER BY a.QueryID, a.SqlRowNo) > {offset}
+                    AND ROW_NUMBER() OVER (ORDER BY a.QueryID, a.SqlRowNo) <= {offset + page_size}
+            """
+
+            rows = teradata_repo.execute_query(query, conn)
+            if not rows:
+                break
+
+            logger.info(f"Concatenating split queries for page {offset // page_size + 1}...")
+            concatenated_queries = concatenate_split_queries(rows)
+
+            logger.info("Processing SQL statements and extracting lineage...")
+            lineage_results = process_sql_statements(concatenated_queries)
+            all_lineage_results.extend(lineage_results)
+
+            if len(rows) < page_size:
+                break
+
+            offset += page_size
+
+    logger.info(f"Lineage extraction complete. Found {len(all_lineage_results)} unique entries.")
+    for entry in all_lineage_results:
+        if len(entry) > 1:
+            logger.debug(f"Destination table: {entry[1]}")
+        else:
+            logger.debug("Destination table: Not available (out of bounds)")
+
+        if len(entry) > 0 and isinstance(entry[0], list):
+            logger.debug("Source tables:")
+            for src_table in entry[0]:
+                logger.debug(f"  - {src_table}")
+        else:
+            logger.debug("Source tables: Not available (out of bounds or invalid type)")
+
+        logger.debug("---")
+
+    sql_lineage = SQLLineage()
+    update_table_lineage_inputs = [
+        sql_lineage.gen_lineage_input(
+            tenant_id=tenant_id, endpoint=endpoint, src_tables=src_tables, dest_table=dest_table
+        )
+        for src_tables, dest_table in all_lineage_results
+    ]
+
+    table_req_count = 0
+    logger.info(f"Starting to update lineage information for {len(update_table_lineage_inputs)} tables.")
+    for update_table_lineage_input in update_table_lineage_inputs:
+        logger.info(
+            f"Generating table lineage. downstream: {update_table_lineage_input.downstream_database_name}"
+            f" -> {update_table_lineage_input.downstream_table_name}"
+        )
+        try:
+            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:
+                table_req_count += 1
+            else:
+                logger.error(
+                    f"Failed to update lineage for {update_table_lineage_input.downstream_table_name}.\
+                        Status code: {status_code}"
+                )
+        except Exception as e:
+            logger.error(
+                f"Exception occurred while updating lineage for {update_table_lineage_input.downstream_table_name}: {e}"
+            )
+    logger.info(f"Generating table lineage is finished. {table_req_count} lineages are ingested.")
+
+
+@error_handling_decorator
+def extract_lineage(sql_statement: str, default_database: str = None) -> Tuple[Set[Table], Table]:
+    try:
+        logger.debug(f"Parsing SQL: {sql_statement}")
+        sql_lineage = SQLLineage()
+        source_tables, dest_table = sql_lineage.get_table_level_lineage_source(sql=sql_statement, dialect="teradata")
+
+        source_tables = {Table(db=t.db_schema or default_database, db_schema="", table=t.table) for t in source_tables}
+        dest_table = Table(db=dest_table.db_schema or default_database, db_schema="", table=dest_table.table)
+
+        return source_tables, dest_table
+    except ParseError as e:
+        logger.error(f"Error parsing SQL: {e}")
+        logger.debug(f"Problematic SQL: {sql_statement}")
+    except AttributeError as e:
+        logger.error(f"Attribute error while extracting lineage: {e}")
+        logger.debug(f"Problematic SQL: {sql_statement}")
+    except Exception as e:
+        logger.error(f"Unexpected error while extracting lineage: {e}")
+        logger.debug(f"Problematic SQL: {sql_statement}")
+    return set(), Table(db="", table="")
+
+
+@error_handling_decorator
+def process_sql_statements(queries: List[Union[str, Dict[str, Union[str, int]]]]) -> List[Tuple[Set[Table], Table]]:
+    lineage_dict = OrderedDict()
+    for query in queries:
+        if isinstance(query, str):
+            sql = query
+            default_database = None
+        else:
+            sql = query["SqlTextInfo"]
+            default_database = query.get("DefaultDatabase")
+
+        source_tables, dest_table = extract_lineage(sql, default_database)
+        if dest_table.table and source_tables:
+            if dest_table in lineage_dict:
+                logger.info(f"Merging duplicate entry for {dest_table}")
+                # Merge source tables
+                lineage_dict[dest_table] = lineage_dict[dest_table].union(source_tables)
+            else:
+                lineage_dict[dest_table] = source_tables
+    return [(src_tables, dest_table) for dest_table, src_tables in lineage_dict.items()]
+
+
+def concatenate_split_queries(rows: List[Dict[str, Union[str, int]]]) -> List[Dict[str, Union[str, int]]]:
+    queries = {}
+    for row in rows:
+        query_id = row["QueryID"]
+        sql_text = row["SqlTextInfo"]
+        default_database = row["DefaultDatabase"]
+        if query_id not in queries:
+            queries[query_id] = {"SqlTextInfo": [], "DefaultDatabase": default_database}
+        queries[query_id]["SqlTextInfo"].append(sql_text)
+
+    return [
+        {"SqlTextInfo": "".join(query["SqlTextInfo"]), "DefaultDatabase": query["DefaultDatabase"]}
+        for query in queries.values()
+    ]

quollio_core/profilers/teradata/stats.py

@@ -0,0 +1,218 @@
+from typing import Any, Dict, List, Optional
+
+from quollio_core.helper.log_utils import error_handling_decorator, logger
+from quollio_core.profilers.stats import gen_table_stats_payload
+from quollio_core.repository import qdc
+from quollio_core.repository import teradata as teradata_repo
+
+NUMERIC_TYPES = ["D", "F", "I1", "I2", "I8", "I", "N"]
+
+# I, I1, I2, I8 - INT TYPES  INTEGER, BYTEINT, SMALLINT, BIGINT
+# F - Float
+# D - Decimal
+# N - Number
+
+
+def quote_identifier(identifier: str) -> str:
+    return f'"{identifier}"'
+
+
+@error_handling_decorator
+def load_stats(
+    conn_config: teradata_repo.TeradataConfig,
+    sample_percent: Optional[float] = None,
+    endpoint: Optional[str] = None,
+    tenant_id: Optional[str] = None,
+    qdc_client: Optional[qdc.QDCExternalAPIClient] = None,
+    target_databases: Optional[List[str]] = None,
+    target_databases_method: str = "DENYLIST",
+    stats_items: Optional[List[str]] = None,
+) -> None:
+    stats_list = []
+    numerical_columns = 0
+    non_numerical_columns = 0
+    logger.info(
+        f"Starting statistics collection. " f"Sample percent: {sample_percent if sample_percent is not None else 'N/A'}"
+    )
+
+    with teradata_repo.new_teradata_client(conn_config) as conn:
+        try:
+            tables = teradata_repo.get_table_list(conn, target_databases, target_databases_method)
+            for table in tables:
+                logger.debug(f"Processing table: {table}")
+                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)
+                logger.debug(f"Columns: {columns}")
+
+                for column in columns:
+                    column_name = column["ColumnName"]
+                    column_type = column["ColumnType"]
+                    if column_type is None:
+                        column_type = ""
+                    else:
+                        column_type = column_type.strip()
+
+                    is_numerical = column_type in NUMERIC_TYPES
+                    if is_numerical:
+                        numerical_columns += 1
+                    else:
+                        non_numerical_columns += 1
+
+                    stats_sql = generate_column_statistics_sql(
+                        database_name,
+                        table_name,
+                        column_name,
+                        column_type,
+                        sample_percent if is_numerical else None,
+                        stats_items,
+                    )
+                    logger.debug(f"Generated SQL for column {column_name}: {stats_sql}")
+
+                    try:
+                        result = teradata_repo.execute_query(stats_sql, conn)
+                        logger.debug(f"Query result for column {column_name}: {result}")
+                        if result:
+                            column_stats = parse_column_statistics_result(
+                                result[0], database_name, table_name, column_name, stats_items, is_numerical
+                            )
+                            stats_list.append(column_stats)
+                    except Exception as e:
+                        logger.error(
+                            f"Failed to collect statistics for {database_name}.{table_name}.{column_name}: {e}"
+                        )
+
+        except Exception as e:
+            logger.error(f"Error during statistics collection: {e}")
+
+        logger.info("Statistics collection completed successfully.")
+
+    logger.debug(f"Stats list: {stats_list}")
+    payloads = gen_table_stats_payload(stats=stats_list, tenant_id=tenant_id, endpoint=endpoint)
+    logger.debug(f"Generated payloads: {payloads}")
+
+    req_count = 0
+    for payload in payloads:
+        logger.info(f"Generating table stats. asset: {payload.db} -> {payload.table} -> {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"Loading statistics is finished. {req_count} statistics are ingested. "
+        f"Numerical columns: {numerical_columns}, Non-numerical columns: {non_numerical_columns}"
+    )
+
+
+@error_handling_decorator
+def parse_column_statistics_result(
+    result: Dict[str, Any],
+    database_name: str,
+    table_name: str,
+    column_name: str,
+    stats_items: Optional[List[str]] = None,
+    is_numerical: bool = False,
+) -> Dict[str, Any]:
+    stats_dict = {
+        "DB_NAME": database_name,
+        "SCHEMA_NAME": "",
+        "TABLE_NAME": table_name,
+        "COLUMN_NAME": column_name,
+    }
+
+    if stats_items:
+        for item in stats_items:
+            if item == "cardinality" and "num_uniques" in result:
+                stats_dict["CARDINALITY"] = result["num_uniques"]
+            elif item == "number_of_null" and "num_nulls" in result:
+                stats_dict["NULL_COUNT"] = result["num_nulls"]  # Changed from NUM_NULLS to NULL_COUNT
+
+            if is_numerical:
+                if item == "min" and "min_value" in result:
+                    stats_dict["MIN_VALUE"] = str(result["min_value"])
+                elif item == "max" and "max_value" in result:
+                    stats_dict["MAX_VALUE"] = str(result["max_value"])
+                elif item == "median" and "median_value" in result:
+                    stats_dict["MEDIAN_VALUE"] = str(result["median_value"])
+                elif item == "mean" and "avg_value" in result:
+                    stats_dict["AVG_VALUE"] = str(result["avg_value"])
+                elif item == "stddev" and "stddev_value" in result:
+                    stats_dict["STDDEV_VALUE"] = str(result["stddev_value"])
+                elif item == "mode" and "mode_value" in result and is_numerical:
+                    stats_dict["MODE_VALUE"] = str(result["mode_value"])
+
+    return stats_dict
+
+
+@error_handling_decorator
+def generate_column_statistics_sql(
+    database_name: str,
+    table_name: str,
+    column_name: str,
+    column_type: str,
+    sample_percent: Optional[float] = None,
+    stats_items: Optional[List[str]] = None,
+) -> str:
+    quoted_column = quote_identifier(column_name)
+    quoted_database = quote_identifier(database_name)
+
+    # Handle the case where table_name might include a database
+    if "." in table_name:
+        schema, table = table_name.split(".", 1)
+        quoted_table = f"{quote_identifier(schema)}.{quote_identifier(table)}"
+    else:
+        quoted_table = quote_identifier(table_name)
+
+    stats_clauses = []
+    mode_query = ""
+
+    if stats_items:
+        if "cardinality" in stats_items:
+            stats_clauses.append(f"COUNT(DISTINCT {quoted_column}) AS num_uniques")
+        if "number_of_null" in stats_items:
+            stats_clauses.append(f"SUM(CASE WHEN {quoted_column} IS NULL THEN 1 ELSE 0 END) AS num_nulls")
+
+        if column_type in NUMERIC_TYPES:
+            if "min" in stats_items:
+                stats_clauses.append(f"MIN(CAST({quoted_column} AS FLOAT)) AS min_value")
+            if "max" in stats_items:
+                stats_clauses.append(f"MAX(CAST({quoted_column} AS FLOAT)) AS max_value")
+            if "median" in stats_items:
+                stats_clauses.append(f"MEDIAN(CAST({quoted_column} AS FLOAT)) AS median_value")
+            if "mean" in stats_items:
+                stats_clauses.append(f"AVG(CAST({quoted_column} AS FLOAT)) AS avg_value")
+            if "stddev" in stats_items:
+                stats_clauses.append(f"STDDEV_SAMP(CAST({quoted_column} AS FLOAT)) AS stddev_value")
+            if "mode" in stats_items:
+                mode_query = (
+                    f"WITH MODE_VALUE AS ("
+                    f"    SELECT {quoted_column}, COUNT(*) as freq "
+                    f"    FROM {quoted_database}.{quoted_table} "
+                )
+
+                if sample_percent is not None and 0 < sample_percent <= 99:
+                    sample_fraction = sample_percent / 100
+                    mode_query += f" SAMPLE {sample_fraction} "
+
+                mode_query += (
+                    f"    GROUP BY {quoted_column} " f" QUALIFY ROW_NUMBER() OVER (ORDER BY COUNT(*) DESC) = 1" f") "
+                )
+                stats_clauses.append(f"(SELECT {quoted_column} FROM MODE_VALUE) AS mode_value")
+
+    if not stats_clauses:
+        logger.warning(f"No statistics selected for column {column_name}. Skipping this column.")
+        return ""
+
+    query = f"{mode_query}" f"SELECT {', '.join(stats_clauses)} " f"FROM {quoted_database}.{quoted_table}"
+
+    if sample_percent is not None and 0 < sample_percent <= 99:
+        sample_fraction = sample_percent / 100
+        query += f" SAMPLE {sample_fraction}"
+
+    logger.debug(f"Generated SQL query for {quoted_database}.{quoted_table}.{quoted_column}: {query}")
+    return query

quollio_core/repository/qdc.py

@@ -25,9 +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
         """
-        is_domain_valid = is_valid_domain(domain=self.base_url)
-        if not is_domain_valid:
-            raise ValueError("The format of quollio API URL is invalid. The URL must end with `.com`")
 
         url = f"{self.base_url}/oauth2/token"
         creds = f"{self.client_id}:{self.client_secret}"
@@ -108,7 +105,3 @@ class QDCExternalAPIClient:
 
 def initialize_qdc_client(api_url: str, client_id: str, client_secret: str) -> QDCExternalAPIClient:
     return QDCExternalAPIClient(base_url=api_url, client_id=client_id, client_secret=client_secret)
-
-
-def is_valid_domain(domain: str) -> bool:
-    return domain.endswith(".com")

quollio_core/repository/ssm.py

@@ -0,0 +1,59 @@
+import logging
+import os
+from typing import Tuple
+
+import boto3
+from botocore.exceptions import ClientError
+
+logger = logging.getLogger(__name__)
+
+
+def get_parameter_by_assume_role(key: str, region: str = "ap-northeast-1") -> Tuple[str, Exception]:
+    tenant_id = os.getenv("TENANT_ID")
+    if not _is_str_valid(tenant_id):
+        return ("", Exception("TENANT_ID is not set in get_parameter_by_assume_role."))
+    qdc_account_id = os.getenv("QDC_ACCOUNT_ID")
+    if not _is_valid_aws_account_id(qdc_account_id):
+        return ("", Exception("QDC_ACCOUNT_ID is not set in get_parameter_by_assume_role."))
+    qdc_region = os.getenv("QDC_REGION")
+    if not _is_str_valid(qdc_region):
+        return ("", Exception("QDC_REGION is not set in get_parameter_by_assume_role."))
+
+    sts_assume_role_arn = "arn:aws:iam::{account_id}:role/qdc-{tenant_id}-cross-account-access".format(
+        account_id=qdc_account_id, tenant_id=tenant_id
+    )
+
+    session = boto3.Session(region_name=region)
+    sts = session.client("sts", endpoint_url="https://sts.{region}.amazonaws.com".format(region=qdc_region))
+    assumed_role_object = sts.assume_role(
+        RoleArn=sts_assume_role_arn,
+        RoleSessionName="AssumeRoleSession",
+    )
+    credentials = assumed_role_object["Credentials"]
+
+    try:
+        ssm = session.client(
+            "ssm",
+            endpoint_url="https://ssm.{region}.amazonaws.com".format(region=qdc_region),
+            aws_access_key_id=credentials["AccessKeyId"],
+            aws_secret_access_key=credentials["SecretAccessKey"],
+            aws_session_token=credentials["SessionToken"],
+        )
+        res = ssm.get_parameter(Name=key, WithDecryption=True)
+        return (res["Parameter"]["Value"], None)
+    except ClientError as e:
+        logger.error(
+            "Failed to run ssm.get_parameter().\
+             Please check the value stored in parameter store is correct. error: {err}".format(
+                err=e
+            )
+        )
+        return ("", e)
+
+
+def _is_valid_aws_account_id(s: str) -> bool:
+    return s is not None and len(s) == 12 and s.isdigit()
+
+
+def _is_str_valid(s: str) -> bool:
+    return s is not None and s != ""

quollio_core/repository/teradata.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import teradatasql
+
+from quollio_core.helper.log_utils import error_handling_decorator, logger
+
+
+@dataclass
+class TeradataConfig:
+    host: str
+    port: int
+    username: str
+    password: str
+    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
+    ) -> "TeradataConfig":
+        return cls(
+            host=host,
+            port=int(port),
+            username=credentials["username"],
+            password=credentials["password"],
+            additional_params=additional_params or {},
+        )
+
+    def get_connection_params(self) -> Dict[str, Any]:
+        params = {
+            "host": self.host,
+            "user": self.username,
+            "password": self.password,
+            "database": self.database,
+            "dbs_port": self.port,
+            "encryptdata": str(self.encrypt_data).lower(),
+        }
+        params.update(self.additional_params)
+        return params
+
+
+@error_handling_decorator
+def new_teradata_client(config: TeradataConfig) -> teradatasql.connect:
+    conn = teradatasql.connect(**config.get_connection_params())
+    return conn
+
+
+@error_handling_decorator
+def get_table_list(
+    conn: teradatasql.connect, target_databases: Optional[List[str]] = None, target_databases_method: str = "DENYLIST"
+) -> List[Dict[str, str]]:
+    if target_databases_method == "DENYLIST":
+        operator = "NOT"
+    else:
+        operator = ""
+
+    query_tables = f"""
+    SELECT DatabaseName, TableName
+    FROM DBC.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.")
+    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]]:
+    query_columns = f"""
+    SELECT ColumnName, ColumnType
+    FROM DBC.ColumnsV
+    WHERE DatabaseName = '{database_name}'
+        AND TableName = '{table_name}'
+    """
+    logger.debug(f"Executing query to retrieve columns for {database_name}.{table_name}.")
+    columns = execute_query(query_columns, conn)
+    logger.debug(f"Retrieved columns: {columns}")
+    return columns
+
+
+@error_handling_decorator
+def execute_query(query: str, con: teradatasql.connect) -> List[Dict[str, Any]]:
+    try:
+        with con.cursor() as cur:
+            logger.debug(f"Executing SQL query: {query}")
+            cur.execute(query)
+            logger.debug(f"Column descriptions: {cur.description}")
+            columns = [desc[0] for desc in cur.description]
+            rows = [dict(zip(columns, row)) for row in cur.fetchall()]
+            logger.debug(f"Fetched {len(rows)} rows from Teradata.")
+            return rows
+    except teradatasql.OperationalError as e:
+        logger.error(f"Teradata Operational Error: {e}")
+        raise
+    except teradatasql.ProgrammingError as e:
+        logger.error(f"Teradata Programming Error: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error fetching data from Teradata: {e}")
+        raise

quollio_core/snowflake.py

@@ -3,7 +3,7 @@ import logging
 import os
 import shutil
 
-from quollio_core.helper.core import setup_dbt_profile
+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.snowflake import (
@@ -13,7 +13,7 @@ from quollio_core.profilers.snowflake import (
     snowflake_table_to_table_lineage,
 )
 from quollio_core.profilers.stats import get_column_stats_items
-from quollio_core.repository import dbt, qdc, snowflake
+from quollio_core.repository import dbt, qdc, snowflake, ssm
 
 logger = logging.getLogger(__name__)
 
@@ -298,6 +298,16 @@ if __name__ == "__main__":
         required=False,
         help="Whether to ingest column lineage into QDIC or not. Default value is False",
     )
+    parser.add_argument(
+        "--external_api_access",
+        type=str,
+        choices=["PUBLIC", "VPC_ENDPOINT"],
+        action=env_default("EXTERNAL_API_ACCESS"),
+        default="PUBLIC",
+        required=False,
+        help="Access method to Quollio API. Default 'PUBLIC'. Choose 'VPC_ENDPOINT'\
+             if you use API Gateway VPC Endpoint, DefaultValue is set to PUBLIC.",
+    )
 
     stats_items = get_column_stats_items()
     parser.add_argument(
@@ -336,9 +346,20 @@ if __name__ == "__main__":
             log_level=args.log_level,
             dbt_macro_source=args.dbt_macro_source,
         )
+    api_url = args.api_url
+    if args.external_api_access == "VPC_ENDPOINT":
+        api_url, err = ssm.get_parameter_by_assume_role(args.api_url)
+        if err is not None:
+            logger.error("Fail to ssm.get_parameter_by_assume_role. {err}".format(err=err))
+            raise Exception("Fail to ssm.get_parameter_by_assume_role. {err}".format(err=err))
+
+    is_domain_valid = is_valid_domain(domain=api_url, domain_type=args.external_api_access)
+    if not is_domain_valid:
+        raise ValueError("The format of quollio API URL is invalid. The URL must end with `.com`")
+
     if "load_lineage" in args.commands:
         qdc_client = qdc.QDCExternalAPIClient(
-            base_url=args.api_url,
+            base_url=api_url,
             client_id=args.client_id,
             client_secret=args.client_secret,
         )
@@ -350,7 +371,7 @@ if __name__ == "__main__":
         )
     if "load_stats" in args.commands:
         qdc_client = qdc.QDCExternalAPIClient(
-            base_url=args.api_url,
+            base_url=api_url,
             client_id=args.client_id,
             client_secret=args.client_secret,
         )
@@ -362,7 +383,7 @@ if __name__ == "__main__":
         )
     if "load_sqllineage" in args.commands:
         qdc_client = qdc.QDCExternalAPIClient(
-            base_url=args.api_url,
+            base_url=api_url,
             client_id=args.client_id,
             client_secret=args.client_secret,
         )

quollio_core/teradata.py

@@ -0,0 +1,254 @@
+import argparse
+import json
+
+from quollio_core.helper.env_default import env_default
+from quollio_core.helper.log_utils import configure_logging, error_handling_decorator, logger
+from quollio_core.profilers.stats import get_column_stats_items
+from quollio_core.profilers.teradata.lineage import load_lineage
+from quollio_core.profilers.teradata.stats import load_stats
+from quollio_core.repository import qdc
+from quollio_core.repository import teradata as teradata_repo
+
+DEFAULT_SYSTEM_DATABASES = [
+    "DBC",
+    "GLOBAL_FUNCTIONS",
+    "gs_tables_db",
+    "modelops",
+    "system",
+    "tapidb",
+    "TDaaS_BAR",
+    "TDaaS_DB",
+    "TDaaS_Maint",
+    "TDaaS_Monitor",
+    "TDaaS_Support",
+    "TDaaS_TDBCMgmt1",
+    "TDaaS_TDBCMgmt2",
+    "TDBCMgmt",
+    "Crashdumps",
+    "dbcmngr",
+    "DemoNow_Monitor",
+    "External_AP",
+    "LockLogShredder",
+    "mldb",
+    "SQLJ",
+    "SysAdmin",
+    "SYSBAR",
+    "SYSJDBC",
+    "SYSLIB",
+    "SYSSPATIAL",
+    "SystemFe",
+    "SYSUDTLIB",
+    "SYSUIF",
+    "Sys_Calendar",
+    "TDMaps",
+    "TDPUSER",
+    "TDQCD",
+    "TDStats",
+    "tdwm",
+    "TD_ANALYTICS_DB",
+    "TD_SERVER_DB",
+    "TD_SYSFNLIB",
+    "TD_SYSGPL",
+    "TD_SYSXML",
+    "val",
+]
+
+
+@error_handling_decorator
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="Quollio Intelligence Agent for Teradata",
+        description="Load lineage and stats to Quollio from Teradata",
+        epilog="Copyright (c) 2024 Quollio Technologies, Inc.",
+    )
+    parser.add_argument(
+        "commands",
+        choices=["load_lineage", "load_stats"],
+        type=str,
+        nargs="+",
+        help="""
+        The command to execute.
+        'load_lineage': Load lineage data from Teradata to Quollio,
+        'load_stats': Load stats from Teradata to Quollio
+        """,
+    )
+    parser.add_argument(
+        "--log_level",
+        type=str,
+        choices=["debug", "info", "warn", "error", "none"],
+        action=env_default("LOG_LEVEL"),
+        default="info",
+        required=False,
+        help="The log level for commands. Default value is info",
+    )
+    parser.add_argument(
+        "--tenant_id",
+        type=str,
+        action=env_default("TENANT_ID"),
+        required=False,
+        help="The tenant id (company id) where the lineage and stats are loaded",
+    )
+    parser.add_argument(
+        "--teradata_host",
+        type=str,
+        action=env_default("TERADATA_HOST"),
+        required=True,
+        help="Teradata host",
+    )
+    parser.add_argument(
+        "--teradata_port",
+        type=str,
+        action=env_default("TERADATA_PORT"),
+        required=True,
+        help="Teradata port",
+    )
+    parser.add_argument(
+        "--teradata_user",
+        type=str,
+        action=env_default("TERADATA_USER_NAME"),
+        required=True,
+        help="Teradata username",
+    )
+    parser.add_argument(
+        "--teradata_password",
+        type=str,
+        action=env_default("TERADATA_PASSWORD"),
+        required=True,
+        help="Teradata password",
+    )
+    parser.add_argument(
+        "--teradata_connection_parameters",
+        type=str,
+        action=env_default("TERADATA_CONNECTION_PARAMETERS"),
+        default="{}",
+        help="Additional Teradata connection parameters as a JSON string",
+    )
+    parser.add_argument(
+        "--api_url",
+        type=str,
+        action=env_default("QDC_API_URL"),
+        required=False,
+        help="The base URL of Quollio External API",
+    )
+    parser.add_argument(
+        "--client_id",
+        type=str,
+        action=env_default("QDC_CLIENT_ID"),
+        required=False,
+        help="The client id that is created on Quollio console to let clients access Quollio External API",
+    )
+    parser.add_argument(
+        "--client_secret",
+        type=str,
+        action=env_default("QDC_CLIENT_SECRET"),
+        required=False,
+        help="The client secret that is created on Quollio console to let clients access Quollio External API",
+    )
+    parser.add_argument(
+        "--sample_percent",
+        type=float,
+        action=env_default("SAMPLE_PERCENT"),
+        default=1,
+        required=False,
+        help="Percentage of data to sample when collecting statistics (e.g., 10 for 10%). Default is 1%.",
+    )
+    parser.add_argument(
+        "--teradata_target_databases",
+        type=str,
+        action=env_default("TERADATA_TARGET_DATABASES"),
+        required=False,
+        default=None,
+        help="Comma-separated list of Teradata target databases. If not provided,\
+            DEFAULT_SYSTEM_DATABASES will be used.",
+    )
+    parser.add_argument(
+        "--teradata_target_databases_method",
+        type=str,
+        choices=["ALLOWLIST", "DENYLIST"],
+        action=env_default("TERADATA_TARGET_DATABASE_METHOD"),
+        default="DENYLIST",
+        help="Method to use for teradata_target_databases (allowlist or denylist)",
+    )
+    parser.add_argument(
+        "--teradata_page_size",
+        type=int,
+        action=env_default("TERADATA_PAGE_SIZE"),
+        default=1000,
+        required=False,
+        help="Page size for Teradata queries. Default is 1000.",
+    )
+    parser.add_argument(
+        "--target_stats_items",
+        type=str,
+        nargs="*",
+        choices=get_column_stats_items(),
+        default=get_column_stats_items(),
+        action=env_default("TERADATA_STATS_ITEMS"),
+        required=False,
+        help="The items for statistic values.\
+              You can choose the items to be aggregated for stats.\
+              Default is full stats.",
+    )
+
+    args = parser.parse_args()
+
+    configure_logging(args.log_level)
+
+    logger.info("Starting Quollio Intelligence Agent for Teradata")
+
+    credentials = {
+        "username": args.teradata_user,
+        "password": args.teradata_password,
+    }
+
+    # Parse additional connection parameters
+    try:
+        additional_params = json.loads(args.teradata_connection_parameters)
+    except json.JSONDecodeError:
+        logger.warning("Invalid JSON in TERADATA_CONNECTION_PARAMETERS. Using empty dict.")
+        additional_params = {}
+
+    logger.info("Initializing QDC client")
+    qdc_client = qdc.initialize_qdc_client(args.api_url, args.client_id, args.client_secret)
+
+    logger.info("Initializing Teradata client")
+    config = teradata_repo.TeradataConfig.from_dict(
+        credentials, args.teradata_host, args.teradata_port, additional_params
+    )
+
+    if "load_lineage" in args.commands:
+        logger.info("Starting lineage loading process")
+        load_lineage(
+            conn_config=config,
+            tenant_id=args.tenant_id,
+            endpoint=args.teradata_host,
+            qdc_client=qdc_client,
+            page_size=args.teradata_page_size,
+        )
+        logger.info("Lineage loading process completed")
+
+    if "load_stats" in args.commands:
+        logger.info("Starting statistics loading process")
+        logger.info(f"Selected stats items: {args.target_stats_items}")
+        target_databases = (
+            DEFAULT_SYSTEM_DATABASES
+            if args.teradata_target_databases is None
+            else args.teradata_target_databases.split(",")
+        )
+        load_stats(
+            conn_config=config,
+            sample_percent=args.sample_percent,
+            tenant_id=args.tenant_id,
+            endpoint=args.teradata_host,
+            qdc_client=qdc_client,
+            target_databases=target_databases,
+            target_databases_method=args.teradata_target_databases_method.upper(),
+            stats_items=args.target_stats_items,
+        )
+        logger.info("Statistics loading process completed")
+
+    logger.info("Quollio Intelligence Agent for Teradata completed successfully")
+
+
+if __name__ == "__main__":
+    main()

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: quollio-core
-Version: 0.4.19
+Version: 0.5.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>
@@ -37,6 +37,7 @@ Requires-Dist: google-cloud-bigquery==3.22.0
 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: 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.4.19.dist-info → quollio_core-0.5.0.dist-info}/RECORD

@@ -1,8 +1,9 @@
-quollio_core/__init__.py,sha256=y0SJ_NkmDeB1AzUwy_oG9Ivh80sUSTMebFbwH83kHe8,84
+quollio_core/__init__.py,sha256=JDD0AXVIwawcOMm_hOs62bzrTO4xouhl7kPLgO7yn_c,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=D-d26OwCUIpXIqfxZnkv4Ei1GJ1mdw9z8YA8K0G-bSE,12216
+quollio_core/snowflake.py,sha256=3dHVys5c94s2pkEXkb2cWBpkxwkLQ23QENCoU-pfhuo,13202
+quollio_core/teradata.py,sha256=muOCg40SPMs9Ro6f1h8AjzLhGPNXfX1n8IvTtqnQDg0,7739
 quollio_core/dbt_projects/databricks/.gitignore,sha256=1jJAyXSzJ3YUm0nx3i7wUSE4RjQMX3ad6F8O88UbtzI,29
 quollio_core/dbt_projects/databricks/README.md,sha256=ZpRQyhFAODAiS8dc1Kb_ndkul4cu4o4udN_EMa49CU4,440
 quollio_core/dbt_projects/databricks/dbt_project.yml,sha256=3sH98RNk7TnphvI3yEdXDstb92kW5BNxr-cT0tXhwzk,480
@@ -55,14 +56,14 @@ quollio_core/dbt_projects/snowflake/models/quollio_sqllineage_sources.sql,sha256
 quollio_core/dbt_projects/snowflake/models/quollio_sqllineage_sources.yml,sha256=qgazupx3ca4P8R0loY5F9hyCz2fmAcWqZ6iOySo_NoY,377
 quollio_core/dbt_projects/snowflake/models/quollio_stats_columns.sql,sha256=BzvP9gKMFItmwqEQ4bDgtS-Invxhhe6L73Qe1ucxfHo,284
 quollio_core/dbt_projects/snowflake/models/quollio_stats_columns.yml,sha256=V_BESPk6IqE52ExT26-78As9l9AlWW86-Geb5PIhThU,67
-quollio_core/dbt_projects/snowflake/models/quollio_stats_profiling_columns.sql,sha256=IDvYtl84vcQaE3ImOKhDfgO_8aYWTudmGggCiMF9yi0,2251
+quollio_core/dbt_projects/snowflake/models/quollio_stats_profiling_columns.sql,sha256=1RFE5VP1XbLMBlQAShJSeqMNDg5GFmwLXab-q2mnGcQ,2260
 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/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=wbu4FWI7YiFEttXGSuj3tMyAhtPAFlHOjDpWJGNXOHA,1202
+quollio_core/helper/core.py,sha256=ddV5VLa49Je11eHUjlRN5m3nhtqPMMUSeG4135HuZC8,1390
 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
@@ -72,16 +73,20 @@ quollio_core/profilers/databricks.py,sha256=ik4RiR_GOeU3S7s6C6Y9SGe1D_Y_f98BDWJV
 quollio_core/profilers/lineage.py,sha256=4FyxIuPBrUFihqZryqTQBcfB0Z7634lKl_WwkD82vzE,6865
 quollio_core/profilers/redshift.py,sha256=p6ONDCkhndZAOcKAwEyQ5fsi-jsQrlwHHb7LTI_m1uk,6473
 quollio_core/profilers/snowflake.py,sha256=m9Ivv2LRwnrmgKS36a039AhrO27sR1EaOOdqNF26PhI,11156
-quollio_core/profilers/sqllineage.py,sha256=XkF7hwDWIGNtyEP5cv2wETBgMfdQxeHolv7qPIkntSQ,5066
+quollio_core/profilers/sqllineage.py,sha256=h0FT6CYb0A20zSc68GELZ7Q8bDbaHLQnZQHsXBEXBug,5261
 quollio_core/profilers/stats.py,sha256=OLQrdrh0y64jo9rmzvGlDdxy_c7gMz_GnlXPJzWkBjM,7343
+quollio_core/profilers/teradata/lineage.py,sha256=M1FOrFLZrgJwIp-qEeARugoP3W6YxuXkRHoDTPcj-wA,7169
+quollio_core/profilers/teradata/stats.py,sha256=oIC0pp0vJeAx5VWzgVYPrI7sCbCka86ctF2ksczCOuU,9022
 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=_5ygUD6h-zs02f4rzj6evxXqD1JjSgtWc-oHPEofaig,4902
+quollio_core/repository/qdc.py,sha256=IPGiYafnJYkuD7_kLapVM98-9ZhEwq5S-dGY2bO8fVk,4624
 quollio_core/repository/redshift.py,sha256=p2ouEuYcDCjx1oBhc6H1ekQsvEqHGd3bFu3PW0ngYBc,2880
 quollio_core/repository/snowflake.py,sha256=zL9-xi98AIftdW9MuKI-M3pZ1kQuuH-UiZH8HcJvmk4,1769
-quollio_core-0.4.19.dist-info/LICENSE,sha256=V8j_M8nAz8PvAOZQocyRDX7keai8UJ9skgmnwqETmdY,34520
-quollio_core-0.4.19.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-quollio_core-0.4.19.dist-info/METADATA,sha256=cGPdKKcHhb6b3D9rz46w4FMygML4fUZ-ZQgtfI2DqLY,6887
-quollio_core-0.4.19.dist-info/RECORD,,
+quollio_core/repository/ssm.py,sha256=xpm1FzbBnIsBptuYPUNnPgkKU2AH3XxI-ZL0bEetvW0,2182
+quollio_core/repository/teradata.py,sha256=DkoutYaRspPdwCvDGKxpWX0dU0d12S2sUrzvZJpY-3Q,3420
+quollio_core-0.5.0.dist-info/LICENSE,sha256=V8j_M8nAz8PvAOZQocyRDX7keai8UJ9skgmnwqETmdY,34520
+quollio_core-0.5.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+quollio_core-0.5.0.dist-info/METADATA,sha256=erQi_pK7TboXRC-qBE59sqsrqAINKOlFEGDx0SF1JjA,6924
+quollio_core-0.5.0.dist-info/RECORD,,