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

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

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

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

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

{quollio_core-0.5.0 → 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.5.0 → 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.5.0 → 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.5.0 → 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
@@ -92,5 +92,21 @@ WITH columns AS (
            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.5.0 → quollio_core-0.5.3}/quollio_core/profilers/teradata/lineage.py

@@ -17,11 +17,15 @@ def load_lineage(
     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"""
@@ -30,10 +34,10 @@ def load_lineage(
                     TRIM(a.SqlTextInfo) AS SqlTextInfo,
                     a.SqlRowNo,
                     TRIM(d.DatabaseName) AS DefaultDatabase
-                FROM DBC.QryLogSQLV a
-                JOIN DBC.QryLogV b
+                FROM {system_database}.QryLogSQLV a
+                JOIN {system_database}.QryLogV b
                     ON a.QueryID = b.QueryID
-                JOIN DBC.DatabasesV d
+                JOIN {system_database}.DatabasesV d
                     ON b.DefaultDatabase = d.DatabaseName
                 WHERE
                     UPPER(TRIM(SqlTextInfo)) LIKE 'CREATE TABLE%'

{quollio_core-0.5.0 → quollio_core-0.5.3}/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-0.5.0 → quollio_core-0.5.3}/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-0.5.0 → quollio_core-0.5.3}/quollio_core/snowflake.py

@@ -24,6 +24,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 +34,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)
@@ -237,6 +245,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,
@@ -339,12 +366,20 @@ if __name__ == "__main__":
         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-0.5.0 → quollio_core-0.5.3}/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")