Skip to content

feat: Add distribution graph below 20 distinct count [#DCS-1244] #352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AnshumanX304 merged 6 commits into from
Nov 28, 2025
Merged

feat: Add distribution graph below 20 distinct count [#DCS-1244] #352

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d859779
feat: enhance distinct count handling for JSON columns and add distri…
Kanai2003 Nov 28, 2025
a11a16f
chore: add environment and installation log files and enhance excepti…
Kanai2003 Nov 28, 2025
7350c04
feat: format metrics data structure for improved clarity in distribut…
Kanai2003 Nov 28, 2025
13ffe95
refactor: remove commented-out code and clean up distribution graph l…
Kanai2003 Nov 28, 2025
e2f9cf0
refactor: remove outdated comments and simplify distinct expression h…
Kanai2003 Nov 28, 2025
d8c7045
refactor: remove unnecessary comment regarding JSON grouping in Postg…
Kanai2003 Nov 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 64 additions & 4 deletions dcs_core/integrations/databases/postgres.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,12 +331,16 @@ def build_table_metrics_query(
for col in column_info:
name = col["column_name"]
dtype = col["data_type"].lower()
quoted = self.quote_column(name)

if dtype in ("json", "jsonb"):
distinct_expr = f"{quoted}::text"
else:
distinct_expr = f"{quoted}"

query_parts.append(f'COUNT(DISTINCT {distinct_expr}) AS "{name}_distinct"')
query_parts.append(
f'COUNT(DISTINCT {self.quote_column(name)}) AS "{name}_distinct"'
)
query_parts.append(
f'COUNT(*) - COUNT(DISTINCT {self.quote_column(name)}) AS "{name}_duplicate"'
f'COUNT(*) - COUNT(DISTINCT {distinct_expr}) AS "{name}_duplicate"'
)
query_parts.append(
f'SUM(CASE WHEN {self.quote_column(name)} IS NULL THEN 1 ELSE 0 END) AS "{name}_is_null"'
Expand Down Expand Up @@ -415,6 +419,62 @@ def _normalize_metrics(value):
col_metrics[metric_name] = _normalize_metrics(value)

column_wise.append({"column_name": name, "metrics": col_metrics})

for col_data in column_wise:
metrics = col_data["metrics"]
distinct_count = metrics.get("distinct")
col_name = col_data["column_name"]
dtype = next(
c["data_type"].lower()
for c in column_info
if c["column_name"] == col_name
)
Comment on lines +427 to +431

@coderabbitai coderabbitai Bot Nov 28, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Add a default to prevent StopIteration exception.

The next() call without a default value will raise StopIteration if col_name is not found in column_info. While this scenario should not occur under normal circumstances, defensive coding warrants a fallback.

Apply this diff to add a default:

-            dtype = next(
-                c["data_type"].lower()
-                for c in column_info
-                if c["column_name"] == col_name
-            )
+            dtype = next(
+                (c["data_type"].lower()
+                for c in column_info
+                if c["column_name"] == col_name),
+                None
+            )
+            
+            if dtype is None:
+                continue  # Skip if column info not found
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
dtype = next(
c["data_type"].lower()
for c in column_info
if c["column_name"] == col_name
)
dtype = next(
(c["data_type"].lower()
for c in column_info
if c["column_name"] == col_name),
None
)
if dtype is None:
continue # Skip if column info not found
🤖 Prompt for AI Agents 
In dcs_core/integrations/databases/postgres.py around lines 431 to 435, the
next() call used to find dtype will raise StopIteration if col_name is missing;
change the call to provide a default (e.g., next(..., None) or next(...,
"unknown")) and then handle that default: either raise a clear ValueError with
context if dtype is None/unknown, or set dtype to a safe fallback and proceed.
Ensure the chosen fallback is consistent with downstream expectations and
include a brief log/error message when the default branch is taken.


if isinstance(distinct_count, (int, float)) and distinct_count < 20:
quoted = self.quote_column(col_name)

if dtype in ("json", "jsonb"):
group_expr = f"{quoted}::text"
else:
group_expr = quoted

dist_query = (
f"SELECT {group_expr}, COUNT(*) "
f"FROM {qualified_table} GROUP BY {group_expr} ORDER BY COUNT(*) DESC"
)
Comment on lines +441 to +444

@coderabbitai coderabbitai Bot Nov 28, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# First, let's locate and examine the postgres.py file around lines 446-449
cat -n dcs_core/integrations/databases/postgres.py | sed -n '430,470p'

Repository: datachecks/dcs-core

Length of output: 1865

🏁 Script executed:

#!/bin/bash
# Find the quote_column implementation in postgres.py
rg -n "def quote_column" dcs_core/integrations/databases/postgres.py -A 5

Repository: datachecks/dcs-core

Length of output: 260

🏁 Script executed:

#!/bin/bash
# Check for quote_column implementations across all datasources
rg -n "def quote_column" dcs_core/integrations/databases/ -A 3 --type py

Repository: datachecks/dcs-core

Length of output: 2268

🏁 Script executed:

#!/bin/bash
# Check for logging imports in postgres.py
rg -n "^import|^from" dcs_core/integrations/databases/postgres.py | head -20

Repository: datachecks/dcs-core

Length of output: 502

🏁 Script executed:

#!/bin/bash
# Find the function containing this code to understand column_info origin
rg -n "def.*distribution" dcs_core/integrations/databases/postgres.py -B 5 -A 40 | head -60

Repository: datachecks/dcs-core

Length of output: 45

🏁 Script executed:

#!/bin/bash
# Get full quote_column implementation in postgres.py
sed -n '69,80p' dcs_core/integrations/databases/postgres.py

Repository: datachecks/dcs-core

Length of output: 393

🏁 Script executed:

#!/bin/bash
# Get quote_column implementations from other datasources to compare
echo "=== MySQL quote_column ===" && sed -n '87,95p' dcs_core/integrations/databases/mysql.py
echo ""
echo "=== Oracle quote_column ===" && sed -n '83,91p' dcs_core/integrations/databases/oracle.py

Repository: datachecks/dcs-core

Length of output: 605

🏁 Script executed:

#!/bin/bash
# Search for the function containing the distribution graph code
rg -n "def.*" dcs_core/integrations/databases/postgres.py | rg -B 5 -A 5 "430:"

Repository: datachecks/dcs-core

Length of output: 45

🏁 Script executed:

#!/bin/bash
# Find what function contains line 430
awk '/^[[:space:]]*(def|class) / { func=$0; line=NR } NR==430 { print func " at line " line }' dcs_core/integrations/databases/postgres.py

Repository: datachecks/dcs-core

Length of output: 232

🏁 Script executed:

#!/bin/bash
# Better approach: find function by looking backwards from line 430
sed -n '1,430p' dcs_core/integrations/databases/postgres.py | tail -50 | grep -n "def "</method> to see function name

Repository: datachecks/dcs-core

Length of output: 116

🏁 Script executed:

#!/bin/bash
# Find the function containing line 430 by searching backwards
sed -n '400,430p' dcs_core/integrations/databases/postgres.py | grep -E "^\s*(def|class) " | tail -1

Repository: datachecks/dcs-core

Length of output: 45

🏁 Script executed:

#!/bin/bash
# Alternative: look at more context before line 430
sed -n '350,435p' dcs_core/integrations/databases/postgres.py | head -100

Repository: datachecks/dcs-core

Length of output: 3081

🏁 Script executed:

#!/bin/bash
# Search for where column_info is used and its origin
rg -n "column_info" dcs_core/integrations/databases/postgres.py -B 3 -A 3 | head -40

Repository: datachecks/dcs-core

Length of output: 1403

🏁 Script executed:

#!/bin/bash
# Check if logging is used elsewhere in postgres.py for error handling
rg -n "logging\.|logger\." dcs_core/integrations/databases/postgres.py

Repository: datachecks/dcs-core

Length of output: 45

🏁 Script executed:

#!/bin/bash
# Check for other exception handlers using print() in postgres.py
rg -n "except.*:|print\(" dcs_core/integrations/databases/postgres.py | grep -A 2 "except"

Repository: datachecks/dcs-core

Length of output: 280

🏁 Script executed:

#!/bin/bash
# Check how other datasources handle exceptions (logging vs print)
rg -n "except.*:|print\(|logging\." dcs_core/integrations/databases/mysql.py | head -20

Repository: datachecks/dcs-core

Length of output: 131

🏁 Script executed:

#!/bin/bash
# Check if logging is imported/used in the base SQLDataSource class
rg -n "import logging|logger" dcs_core/core/datasource/sql_datasource.py | head -10

Repository: datachecks/dcs-core

Length of output: 477

🏁 Script executed:

#!/bin/bash
# Check what the exception handlers in mysql.py do
sed -n '72,80p' dcs_core/integrations/databases/mysql.py
sed -n '423,430p' dcs_core/integrations/databases/mysql.py

Repository: datachecks/dcs-core

Length of output: 647

Use loguru logger instead of print() for exception handling; SQL injection risk is low but inherent to quote_column pattern.

The distribution graph query at lines 446-449 has two concerns:

  1. Exception handling (lines 467-470): Uses print() instead of the logging pattern established in the base SQLDataSource class, which imports and uses loguru.logger. Other datasources like mysql.py use logger.error() correctly. Replace with logger.error() for consistency.

  2. SQL injection via unescaped quotes: The quote_column() method returns f'"{column}"' without escaping embedded double-quotes. While col_name originates from database metadata (trusted source), PostgreSQL allows embedded quotes in identifier names. This is a pre-existing pattern across all datasources (mysql uses backticks, oracle uses double quotes—none escape). The risk is low given the trusted data source, but if hardening is desired, escape embedded quotes: f'"{column.replace(chr(34), chr(34)+chr(34))}"'.

🧰 Tools
🪛 Ruff (0.14.6)

447-448: Possible SQL injection vector through string-based query construction

(S608)

🤖 Prompt for AI Agents 
In dcs_core/integrations/databases/postgres.py around lines 446-470, replace the
use of print() in the exception handler with the module's loguru logger
(logger.error(...)) and include the exception info in the log call to match the
SQLDataSource logging pattern; additionally, harden quote_column to escape
embedded double-quotes by doubling them (e.g., replace " with "" before wrapping
in quotes) so identifiers containing quotes are safely quoted.


try:
dist_result = self.connection.execute(text(dist_query)).fetchall()

distribution = []
for r in dist_result:
val = _normalize_metrics(r[0])
distribution.append(
{
"col_val": val,
"count": r[1],
}
)

metrics["distribution_graph"] = distribution

except Exception as e:
print(
f"Failed to generate distribution graph for column {col_name}: {e}"
)

for col_data in column_wise:
metrics = col_data["metrics"]
formatted_metrics_data = {
"general_data": {
key: value
for key, value in metrics.items()
if key != "distribution_graph"
},
"distribution_data": metrics.get("distribution_graph", []),
}
col_data["metrics"] = formatted_metrics_data

return column_wise

def get_table_foreign_key_info(self, table_name: str, schema: str | None = None):
Expand Down
Loading