Skip to content
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

fix(ibis): added primary key management on postgres and updated mssql and mysql to manage composite keys #1106

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open

fix(ibis): added primary key management on postgres and updated mssql and mysql to manage composite keys #1106

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e159f5a
fix(ibis): added primary key management on postgres and updated mssql…
yalexbadan Mar 25, 2025
0b5eb0a
fix(ibis): runned just for code formatting
yalexbadan Mar 26, 2025
bfcd4ed
fix(ibis): fixed tests
yalexbadan Mar 28, 2025
8076da8
fix(ibis): fmt code
goldmedal Mar 31, 2025
7bcce3b
fix(ibis): fixed clickhouse and trino metadata classes
yalexbadan Mar 31, 2025
dbff7d8
fix(ibis): fixed test models
yalexbadan Mar 31, 2025
592ed28
fix(ibis): updated java servers with new models
yalexbadan Mar 31, 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
2 changes: 1 addition & 1 deletion ibis-server/app/model/metadata/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ def get_table(row) -> Table:
catalog=row["table_catalog"],
table=row["table_name"],
),
primaryKey="",
primaryKey=[],
)

def is_root_column(row) -> bool:
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/app/model/metadata/canner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ def _build_table(cls, data: dict) -> Table:
name=data["sqlName"],
description=data["dataMetadata"]["metadata"]["description"],
columns=cls._build_columns(data["columns"]),
primaryKey="",
primaryKey=[],
properties=(
TableProperties(
catalog="canner",
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/app/model/metadata/clickhouse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ def get_table_list(self) -> list[Table]:
schema=row["table_schema"],
table=row["table_name"],
),
primaryKey="",
primaryKey=[],
)

# table exists, and add column to the table
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/app/model/metadata/dto.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ class Table(BaseModel):
columns: list[Column]
description: str | None = None
properties: TableProperties = None
primaryKey: str | None = None
primaryKey: list[str] | None = None


class ConstraintType(Enum):
Expand Down
5 changes: 3 additions & 2 deletions ibis-server/app/model/metadata/mssql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ def get_table_list(self) -> list[Table]:
catalog=row["catalog"],
table=row["table_name"],
),
primaryKey="",
primaryKey=[],
)

# table exists, and add column to the table
Expand All @@ -96,7 +96,8 @@ def get_table_list(self) -> list[Table]:
)
# if column is primary key
if row["is_pk"] == "YES":
unique_tables[schema_table].primaryKey = row["column_name"]
unique_tables[schema_table].primaryKey.append(row["column_name"])

return list(unique_tables.values())

def get_constraints(self) -> list[Constraint]:
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/app/model/metadata/mysql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ def get_table_list(self) -> list[Table]:
catalog="",
table=row["table_name"],
),
primaryKey="",
primaryKey=[],
)

# table exists, and add column to the table
Expand All @@ -71,7 +71,7 @@ def get_table_list(self) -> list[Table]:
)
# if column is primary key
if row["column_key"] == "PRI":
unique_tables[schema_table].primaryKey = row["column_name"]
unique_tables[schema_table].primaryKey.append(row["column_name"])
return list(unique_tables.values())

def get_constraints(self) -> list[Constraint]:
Expand Down
4 changes: 3 additions & 1 deletion ibis-server/app/model/metadata/oracle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def get_table_list(self) -> list[Table]:
catalog="", # Oracle doesn't use catalogs.
table=row["TABLE_NAME"],
),
primaryKey="",
primaryKey=[],
)

unique_tables[schema_table].columns.append(
Expand All @@ -101,6 +101,8 @@ def get_table_list(self) -> list[Table]:
)
)

# TODO: manage primary key

return list(unique_tables.values())

def get_constraints(self) -> list[Constraint]:
Expand Down
13 changes: 12 additions & 1 deletion ibis-server/app/model/metadata/postgres.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ def get_table_list(self) -> list[Table]:
c.data_type,
c.is_nullable,
c.ordinal_position,
tc.constraint_type,
obj_description(cls.oid) AS table_comment,
col_description(cls.oid, a.attnum) AS column_comment
FROM
Expand All @@ -44,6 +45,12 @@ def get_table_list(self) -> list[Table]:
pg_attribute a
ON a.attrelid = cls.oid
AND a.attname = c.column_name
LEFT JOIN information_schema.key_column_usage kcu
ON c.table_name = kcu.table_name
AND c.column_name = kcu.column_name
AND c.table_schema = kcu.table_schema
LEFT JOIN information_schema.table_constraints tc
ON kcu.constraint_name = tc.constraint_name
WHERE
t.table_type IN ('BASE TABLE', 'VIEW')
AND t.table_schema NOT IN ('information_schema', 'pg_catalog');
Expand All @@ -67,7 +74,7 @@ def get_table_list(self) -> list[Table]:
catalog=row["table_catalog"],
table=row["table_name"],
),
primaryKey="",
primaryKey=[],
)

# table exists, and add column to the table
Expand All @@ -80,6 +87,10 @@ def get_table_list(self) -> list[Table]:
properties=None,
)
)

if row["constraint_type"] == "PRIMARY KEY":
unique_tables[schema_table].primaryKey.append(row["column_name"])

return list(unique_tables.values())

def get_constraints(self) -> list[Constraint]:
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/app/model/metadata/snowflake.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ def get_table_list(self) -> list[Table]:
catalog=row["TABLE_CATALOG"],
table=row["TABLE_NAME"],
),
primaryKey="",
primaryKey=[],
)

# table exists, and add column to the table
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/app/model/metadata/trino.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ def get_table_list(self) -> list[Table]:
catalog=row["table_catalog"],
table=row["table_name"],
),
primaryKey="",
primaryKey=[],
Copy link

Choose a reason for hiding this comment

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

💡 Verification agent

🧩 Analysis chain

Initialization changed but no logic to populate primary keys

The change from initializing primaryKey as a string to a list is consistent with the modifications in other database metadata classes to support composite keys. However, unlike the MSSQL, MySQL, and PostgreSQL implementations mentioned in the PR objectives, there's no logic in this file to populate the primaryKey list with actual primary key column names.

This change may be incomplete without corresponding logic to populate the primary key list. Please verify if Trino supports retrieving primary key information and, if so, update either the SQL query in get_table_list or implement the get_constraints method to populate the primaryKey list. This might be related to the test failures mentioned in the PR comments for test_trino.py.

🏁 Script executed:

#!/bin/bash
# Check how primary keys are handled in other metadata modules

echo "Looking at how other metadata modules handle primary keys..."
echo "==============================================="

# Look for primaryKey initialization and population in other metadata modules
echo "Checking primaryKey initialization in other metadata modules:"
rg -A 1 "primaryKey\s*=\s*\[\]" ibis-server/app/model/metadata/

echo "==============================================="

# Look for code that populates primary keys in other modules 
echo "Checking how primary keys are populated in other modules:"
rg -A 10 "\.primaryKey\.append" ibis-server/app/model/metadata/

echo "==============================================="

# Check the Table class definition in dto.py to verify the type change
echo "Examining Table class definition in dto.py:"
rg -A 20 "class Table" ibis-server/app/model/metadata/dto.py

echo "==============================================="

# Check the failing test to understand requirements
echo "Looking for trino test file mentioned in PR comments:"
fd "test_trino.py" ibis-server/ --exec cat {}

Length of output: 20170

Action Required: Address Missing Primary Key Population for TRINO

The TRINO metadata module now initializes primaryKey as an empty list—which aligns with the recent changes in other drivers—but, unlike the MSSQL, MySQL, and PostgreSQL implementations, there is no logic to populate this list with actual primary key column names. The search confirms that while other modules append to primaryKey (e.g., via their respective get_constraints methods), the TRINO module lacks any similar population strategy. This may be contributing to the test failures in test_trino.py.

  • Verify if Trino supports retrieving primary key details.
  • If supported, update the SQL query in get_table_list or implement a get_constraints method to append the primary key column data.
  • If not supported, please document this behavior to clarify the intentional absence of key population.

)

# table exists, and add column to the table
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/tests/routers/v2/connector/test_canner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@
"type": "timestamp",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
],
}
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v2/connector/test_clickhouse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@
"expression": "customer.name",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
{
"name": "Customer",
Expand All @@ -91,7 +91,7 @@
"expression": "sum(orders.totalprice)",
},
],
"primaryKey": "custkey",
"primaryKey": ["custkey"],
},
],
"relationships": [
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v2/connector/test_gcs_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
"type": "varchar",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
Copy link

Choose a reason for hiding this comment

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

💡 Verification agent

🧩 Analysis chain

Verify all test assertions are compatible with the new primary key format.

While the primary key format changes look good, it's important to ensure that any assertions or validations in the test suite that might reference these primary keys are also updated to expect lists instead of strings.

Also applies to: 69-69

🏁 Script executed:

#!/bin/bash
# Scan for assertions or validations that might expect primaryKey to be a string
rg -A 5 -B 5 "primaryKey" --glob "*.py" | grep -E "assert|test" | grep -v "\[\".*\"\]"

Length of output: 30078

Attention: Update Test Assertions for Primary Key Format Change

The primary key update in ibis-server/tests/routers/v2/connector/test_gcs_file.py (lines 44 and 69) now uses a list format ("primaryKey": ["orderkey"]). However, our verification indicates that several tests—such as in wren-core-py/tests/test_modeling_core.py—still reference the primary key as a string (e.g., "primaryKey": "c_custkey", "primaryKey": "o_orderkey", etc.).

  • ibis-server/tests/routers/v2/connector/test_gcs_file.py (lines 44 and 69): Confirm that any assertions or validations reflect the new list format.
  • Other test files (e.g., wren-core-py/tests/test_modeling_core.py): Review and update any test assertions or validations that still expect the primary key as a string to ensure consistency with the updated format.

Please update the tests accordingly to prevent potential mismatches and unexpected behavior during test execution.

},
{
"name": "Customer",
Expand All @@ -66,7 +66,7 @@
"expression": "sum(orders.totalprice)",
},
],
"primaryKey": "custkey",
"primaryKey": ["custkey"],
},
],
"relationships": [
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v2/connector/test_local_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
"type": "varchar",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
{
"name": "Customer",
Expand All @@ -61,7 +61,7 @@
"expression": "sum(orders.totalprice)",
},
],
"primaryKey": "custkey",
"primaryKey": ["custkey"],
},
],
"relationships": [
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v2/connector/test_minio_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@
"type": "varchar",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
{
"name": "Customer",
Expand All @@ -64,7 +64,7 @@
"expression": "sum(orders.totalprice)",
},
],
"primaryKey": "custkey",
"primaryKey": ["custkey"],
},
],
"relationships": [
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/tests/routers/v2/connector/test_mssql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@
"type": "bytea",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
],
}
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v2/connector/test_mysql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@
"type": "bytea",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
{
"name": "Customer",
Expand All @@ -73,7 +73,7 @@
{"name": "custkey", "expression": "c_custkey", "type": "integer"},
{"name": "name", "expression": "c_name", "type": "varchar"},
],
"primaryKey": "custkey",
"primaryKey": ["custkey"],
},
],
}
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/tests/routers/v2/connector/test_oracle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@
"type": "blob",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
}
],
}
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/tests/routers/v2/connector/test_postgres.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@
"type": "bytea",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
{
"name": "Customer",
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v2/connector/test_s3_file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
"type": "varchar",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
{
"name": "Customer",
Expand All @@ -66,7 +66,7 @@
"expression": "sum(orders.totalprice)",
},
],
"primaryKey": "custkey",
"primaryKey": ["custkey"],
},
],
"relationships": [
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/tests/routers/v2/connector/test_snowflake.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@
"type": "timestamp",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
],
}
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/tests/routers/v2/connector/test_trino.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@
"type": "bytea",
},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
],
}
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v2/test_analysis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
{"name": "mktsegment", "expression": "mktsegment", "type": "varchar"},
{"name": "comment", "expression": "comment", "type": "varchar"},
],
"primaryKey": "custkey",
"primaryKey": ["custkey"],
},
{
"name": "orders",
Expand All @@ -44,7 +44,7 @@
},
{"name": "comment", "expression": "comment", "type": "varchar"},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
{
"name": "lineitem",
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/tests/routers/v2/test_relationship_valid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
{"name": "id", "type": "integer"},
{"name": "many_col", "type": "integer"},
],
"primaryKey": "id",
"primaryKey": ["id"],
},
{
"name": "t2",
Expand Down
2 changes: 1 addition & 1 deletion ibis-server/tests/routers/v3/connector/bigquery/test_query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@
"type": "timestamptz",
},
],
"primaryKey": "o_orderkey",
"primaryKey": ["o_orderkey"],
},
],
}
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v3/connector/local_file/test_query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
},
{"name": "orderdate", "expression": "o_orderdate", "type": "date"},
],
"primaryKey": "orderkey",
"primaryKey": ["orderkey"],
},
{
"name": "Customer",
Expand All @@ -54,7 +54,7 @@
"expression": 'sum("Orders".totalprice)',
},
],
"primaryKey": "custkey",
"primaryKey": ["custkey"],
},
],
"relationships": [
Expand Down
4 changes: 2 additions & 2 deletions ibis-server/tests/routers/v3/connector/postgres/test_query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@
"type": "timestamptz",
},
],
"primaryKey": "o_orderkey",
"primaryKey": ["o_orderkey"],
},
{
"name": "customer",
Expand All @@ -78,7 +78,7 @@
"expression": "sum(orders.o_totalprice_double)",
},
],
"primaryKey": "c_custkey",
"primaryKey": ["c_custkey"],
},
],
"relationships": [
Expand Down
Loading
Loading