quollio-core 0.4.19__tar.gz → 0.5.3__tar.gz

{quollio_core-0.4.19 → quollio_core-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: quollio-core
-Version: 0.4.19
+Version: 0.5.3
 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 → quollio_core-0.5.3}/pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
     ,"google-cloud-datacatalog==3.19.0"
     ,"google-cloud-datacatalog-lineage==0.3.6"
     ,"google-api-python-client==2.131.0"
+    ,"teradatasql==20.0.0.15"
 ]
 dynamic = ["version", "description"]
 

{quollio_core-0.4.19 → quollio_core-0.5.3}/quollio_core/__init__.py

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

{quollio_core-0.4.19 → quollio_core-0.5.3}/quollio_core/dbt_projects/snowflake/models/quollio_lineage_column_level.sql

@@ -95,6 +95,23 @@ UNION
         {{ source('account_usage', 'TABLES') }}
     WHERE
         DELETED IS NULL
+		AND (
+			{% if var('target_databases_method') == 'ALLOWLIST' %}
+				{% if var('target_databases') %}
+					TABLE_CATALOG LIKE ANY ({{ var('target_databases')|join(",") }})
+				{% else %}
+					1=0  -- If no databases specified in allowlist, deny all
+				{% endif %}
+			{% elif var('target_databases_method') == 'DENYLIST' %}
+				{% if var('target_databases') %}
+					NOT (TABLE_CATALOG LIKE ANY ({{ var('target_databases')|join(",") }}))
+				{% else %}
+					1=1  -- If no databases specified in denylist, include all
+				{% endif %}
+			{% else %}
+				1=1  -- Default case: allow all databases
+			{% endif %}
+		)
 ), exists_upstream_column_lineage AS (
 	SELECT
 	    downstream_table_name

{quollio_core-0.4.19 → quollio_core-0.5.3}/quollio_core/dbt_projects/snowflake/models/quollio_lineage_table_level.sql

@@ -49,6 +49,23 @@ WITH table_lineage_history AS (
         {{ source('account_usage', 'TABLES') }}
     WHERE
         DELETED IS NULL
+        AND (
+            {% if var('target_databases_method') == 'ALLOWLIST' %}
+                {% if var('target_databases') %}
+                    TABLE_CATALOG LIKE ANY ({{ var('target_databases')|join(",") }})
+                {% else %}
+                    1=0  -- If no databases specified in allowlist, deny all
+                {% endif %}
+            {% elif var('target_databases_method') == 'DENYLIST' %}
+                {% if var('target_databases') %}
+                    NOT (TABLE_CATALOG LIKE ANY ({{ var('target_databases')|join(",") }}))
+                {% else %}
+                    1=1  -- If no databases specified in denylist, include all
+                {% endif %}
+            {% else %}
+                1=1  -- Default case: allow all databases
+            {% endif %}
+        )
 ), upstream_exists_table AS (
     SELECT
     	downstream_table_name AS "DOWNSTREAM_TABLE_NAME"

{quollio_core-0.4.19 → quollio_core-0.5.3}/quollio_core/dbt_projects/snowflake/models/quollio_sqllineage_sources.sql

@@ -48,3 +48,21 @@ on
     lst.query_id = qt.query_id
 where
     qt.query_id is not null
+    AND (
+        {% if var('target_databases_method') == 'ALLOWLIST' %}
+            {% if var('target_databases') %}
+                database_name LIKE ANY ({{ var('target_databases')|join(",") }})
+            {% else %}
+                1=0  -- If no databases specified in allowlist, deny all
+            {% endif %}
+        {% elif var('target_databases_method') == 'DENYLIST' %}
+            {% if var('target_databases') %}
+                NOT (database_name LIKE ANY ({{ var('target_databases')|join(",") }}))
+            {% else %}
+                1=1  -- If no databases specified in denylist, include all
+            {% endif %}
+        {% else %}
+            1=1  -- Default case: allow all databases
+        {% endif %}
+    )
+

{quollio_core-0.4.19 → quollio_core-0.5.3}/quollio_core/dbt_projects/snowflake/models/quollio_stats_profiling_columns.sql

@@ -28,7 +28,7 @@ WITH columns AS (
     FROM
       {{ source('account_usage', 'GRANTS_TO_ROLES') }}
     WHERE
-      granted_on in ('TABLE', 'MATERIALIZED VIEW')
+      granted_on in ('TABLE', 'VIEW', 'MATERIALIZED VIEW')
       AND grantee_name = '{{ var("query_role") }}'
       AND privilege in ('SELECT', 'OWNERSHIP')
       AND deleted_on IS NULL
@@ -87,10 +87,26 @@ 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
     implicit_columns_removed
-)
+  WHERE
+      {% if var('target_databases_method') == 'ALLOWLIST' %}
+          {% if var('target_databases') %}
+              TABLE_CATALOG LIKE ANY ({{ var('target_databases')|join(",") }})
+          {% else %}
+              1=0  -- If no databases specified in allowlist, deny all
+          {% endif %}
+      {% elif var('target_databases_method') == 'DENYLIST' %}
+          {% if var('target_databases') %}
+              NOT (TABLE_CATALOG LIKE ANY ({{ var('target_databases')|join(",") }}))
+          {% else %}
+              1=1  -- If no databases specified in denylist, include all
+          {% endif %}
+      {% else %}
+          1=1  -- Default case: allow all databases
+      {% endif %}
+  )
 select * from final

{quollio_core-0.4.19 → quollio_core-0.5.3}/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-0.4.19 → quollio_core-0.5.3}/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-0.5.3/quollio_core/profilers/teradata/lineage.py

@@ -0,0 +1,176 @@
+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,
+    system_database: str = None,
+) -> None:
+    page_size = page_size or int(os.environ.get("TERADATA_PAGE_SIZE", 1000))
+    offset = 0
+    all_lineage_results = []
+
+    # 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:
+        while True:
+            query = f"""
+                SELECT
+                    a.QueryID,
+                    TRIM(a.SqlTextInfo) AS SqlTextInfo,
+                    a.SqlRowNo,
+                    TRIM(d.DatabaseName) AS DefaultDatabase
+                FROM {system_database}.QryLogSQLV a
+                JOIN {system_database}.QryLogV b
+                    ON a.QueryID = b.QueryID
+                JOIN {system_database}.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-0.5.3/quollio_core/profilers/teradata/stats.py

@@ -0,0 +1,224 @@
+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,
+    system_database: Optional[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'}"
+    )
+
+    # 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, system_database)
+            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, system_database=system_database
+                )
+                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-0.4.19 → quollio_core-0.5.3}/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-0.5.3/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 != ""