Skip to content

Commit 2f61cbd

Browse files
fix: PostgreSQL compatibility for diagrams and date functions
- Fix Diagram node discovery to handle PostgreSQL double-quote format - Fix indexes dict to filter out None column names - Add null check for heading.indexes in describe() - Add TIMESTAMPDIFF translation (YEAR, MONTH, DAY units) - Add CURDATE() → CURRENT_DATE translation - Add NOW() → CURRENT_TIMESTAMP translation Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent fd31b22 commit 2f61cbd

File tree

4 files changed

+35
-3
lines changed

4 files changed

+35
-3
lines changed

src/datajoint/adapters/postgres.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1046,6 +1046,36 @@ def replace_group_concat(match):
10461046

10471047
expr = re.sub(r"GROUP_CONCAT\s*\((.+?)\)", replace_group_concat, expr, flags=re.IGNORECASE)
10481048

1049+
# TIMESTAMPDIFF(YEAR, d1, d2) → EXTRACT(YEAR FROM AGE(d2, d1))::int
1050+
# TIMESTAMPDIFF(MONTH, d1, d2) → year*12 + month from AGE
1051+
# TIMESTAMPDIFF(DAY, d1, d2) → (d2::date - d1::date)
1052+
def replace_timestampdiff(match):
1053+
unit = match.group(1).upper()
1054+
date1 = match.group(2).strip()
1055+
date2 = match.group(3).strip()
1056+
if unit == "YEAR":
1057+
return f"EXTRACT(YEAR FROM AGE({date2}, {date1}))::int"
1058+
elif unit == "MONTH":
1059+
return f"(EXTRACT(YEAR FROM AGE({date2}, {date1})) * 12 + EXTRACT(MONTH FROM AGE({date2}, {date1})))::int"
1060+
elif unit == "DAY":
1061+
return f"({date2}::date - {date1}::date)"
1062+
else:
1063+
# For other units, fall back to extracting from interval
1064+
return f"EXTRACT({unit} FROM AGE({date2}, {date1}))::int"
1065+
1066+
expr = re.sub(
1067+
r"TIMESTAMPDIFF\s*\(\s*(\w+)\s*,\s*(.+?)\s*,\s*(.+?)\s*\)",
1068+
replace_timestampdiff,
1069+
expr,
1070+
flags=re.IGNORECASE,
1071+
)
1072+
1073+
# CURDATE() → CURRENT_DATE
1074+
expr = re.sub(r"CURDATE\s*\(\s*\)", "CURRENT_DATE", expr, flags=re.IGNORECASE)
1075+
1076+
# NOW() → CURRENT_TIMESTAMP (already works but ensure compatibility)
1077+
expr = re.sub(r"\bNOW\s*\(\s*\)", "CURRENT_TIMESTAMP", expr, flags=re.IGNORECASE)
1078+
10491079
return expr
10501080

10511081
# =========================================================================

src/datajoint/diagram.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,8 @@ def __init__(self, source, context=None) -> None:
134134
except AttributeError:
135135
raise DataJointError("Cannot plot Diagram for %s" % repr(source))
136136
for node in self:
137-
if node.startswith("`%s`" % database):
137+
# Handle both MySQL backticks and PostgreSQL double quotes
138+
if node.startswith("`%s`" % database) or node.startswith('"%s"' % database):
138139
self.nodes_to_show.add(node)
139140

140141
@classmethod

src/datajoint/heading.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -572,11 +572,12 @@ def _init_from_database(self) -> None:
572572
nullable=nullable,
573573
)
574574
self.indexes = {
575-
tuple(item[k]["column"] for k in sorted(item.keys())): dict(
575+
tuple(item[k]["column"] for k in sorted(item.keys()) if item[k]["column"] is not None): dict(
576576
unique=item[1]["unique"],
577577
nullable=any(v["nullable"] for v in item.values()),
578578
)
579579
for item in keys.values()
580+
if any(item[k]["column"] is not None for k in item.keys())
580581
}
581582

582583
def select(self, select_list, rename_map=None, compute_map=None):

src/datajoint/table.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1143,7 +1143,7 @@ def describe(self, context=None, printout=False):
11431143
definition = "# " + self.heading.table_status["comment"] + "\n" if self.heading.table_status["comment"] else ""
11441144
attributes_thus_far = set()
11451145
attributes_declared = set()
1146-
indexes = self.heading.indexes.copy()
1146+
indexes = self.heading.indexes.copy() if self.heading.indexes else {}
11471147
for attr in self.heading.attributes.values():
11481148
if in_key and not attr.in_key:
11491149
definition += "---\n"

0 commit comments

Comments
 (0)