diff --git a/ibis-server/app/model/metadata/bigquery.py b/ibis-server/app/model/metadata/bigquery.py index dbcf92a18..be6665c19 100644 --- a/ibis-server/app/model/metadata/bigquery.py +++ b/ibis-server/app/model/metadata/bigquery.py @@ -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: diff --git a/ibis-server/app/model/metadata/canner.py b/ibis-server/app/model/metadata/canner.py index a7699d59d..a598d8ff5 100644 --- a/ibis-server/app/model/metadata/canner.py +++ b/ibis-server/app/model/metadata/canner.py @@ -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", diff --git a/ibis-server/app/model/metadata/clickhouse.py b/ibis-server/app/model/metadata/clickhouse.py index 07dc335a9..1a6c3f2ae 100644 --- a/ibis-server/app/model/metadata/clickhouse.py +++ b/ibis-server/app/model/metadata/clickhouse.py @@ -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 diff --git a/ibis-server/app/model/metadata/dto.py b/ibis-server/app/model/metadata/dto.py index 85ec3f8d9..6687fa045 100644 --- a/ibis-server/app/model/metadata/dto.py +++ b/ibis-server/app/model/metadata/dto.py @@ -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): diff --git a/ibis-server/app/model/metadata/mssql.py b/ibis-server/app/model/metadata/mssql.py index c3a135a84..6b9e2e434 100644 --- a/ibis-server/app/model/metadata/mssql.py +++ b/ibis-server/app/model/metadata/mssql.py @@ -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 @@ -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]: diff --git a/ibis-server/app/model/metadata/mysql.py b/ibis-server/app/model/metadata/mysql.py index 13f34d357..6c2318210 100644 --- a/ibis-server/app/model/metadata/mysql.py +++ b/ibis-server/app/model/metadata/mysql.py @@ -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 @@ -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]: diff --git a/ibis-server/app/model/metadata/oracle.py b/ibis-server/app/model/metadata/oracle.py index 253dc3769..56d5ad6f2 100644 --- a/ibis-server/app/model/metadata/oracle.py +++ b/ibis-server/app/model/metadata/oracle.py @@ -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( @@ -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]: diff --git a/ibis-server/app/model/metadata/postgres.py b/ibis-server/app/model/metadata/postgres.py index 6a5b5720d..5589fa307 100644 --- a/ibis-server/app/model/metadata/postgres.py +++ b/ibis-server/app/model/metadata/postgres.py @@ -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 @@ -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'); @@ -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 @@ -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]: diff --git a/ibis-server/app/model/metadata/snowflake.py b/ibis-server/app/model/metadata/snowflake.py index 7998b9f09..098102f43 100644 --- a/ibis-server/app/model/metadata/snowflake.py +++ b/ibis-server/app/model/metadata/snowflake.py @@ -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 diff --git a/ibis-server/app/model/metadata/trino.py b/ibis-server/app/model/metadata/trino.py index 0f063d91b..aa7cf927c 100644 --- a/ibis-server/app/model/metadata/trino.py +++ b/ibis-server/app/model/metadata/trino.py @@ -62,7 +62,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 diff --git a/ibis-server/tests/routers/v2/connector/test_canner.py b/ibis-server/tests/routers/v2/connector/test_canner.py index f43d88320..37ccdc373 100644 --- a/ibis-server/tests/routers/v2/connector/test_canner.py +++ b/ibis-server/tests/routers/v2/connector/test_canner.py @@ -66,7 +66,7 @@ "type": "timestamp", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, ], } diff --git a/ibis-server/tests/routers/v2/connector/test_clickhouse.py b/ibis-server/tests/routers/v2/connector/test_clickhouse.py index 2abae3d45..1fd3fb9cc 100644 --- a/ibis-server/tests/routers/v2/connector/test_clickhouse.py +++ b/ibis-server/tests/routers/v2/connector/test_clickhouse.py @@ -71,7 +71,7 @@ "expression": "customer.name", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "Customer", @@ -91,7 +91,7 @@ "expression": "sum(orders.totalprice)", }, ], - "primaryKey": "custkey", + "primaryKey": ["custkey"], }, ], "relationships": [ diff --git a/ibis-server/tests/routers/v2/connector/test_gcs_file.py b/ibis-server/tests/routers/v2/connector/test_gcs_file.py index 47d53ce9a..aeccef113 100644 --- a/ibis-server/tests/routers/v2/connector/test_gcs_file.py +++ b/ibis-server/tests/routers/v2/connector/test_gcs_file.py @@ -41,7 +41,7 @@ "type": "varchar", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "Customer", @@ -66,7 +66,7 @@ "expression": "sum(orders.totalprice)", }, ], - "primaryKey": "custkey", + "primaryKey": ["custkey"], }, ], "relationships": [ diff --git a/ibis-server/tests/routers/v2/connector/test_local_file.py b/ibis-server/tests/routers/v2/connector/test_local_file.py index 0f0b6c99c..3be569d4c 100644 --- a/ibis-server/tests/routers/v2/connector/test_local_file.py +++ b/ibis-server/tests/routers/v2/connector/test_local_file.py @@ -36,7 +36,7 @@ "type": "varchar", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "Customer", @@ -61,7 +61,7 @@ "expression": "sum(orders.totalprice)", }, ], - "primaryKey": "custkey", + "primaryKey": ["custkey"], }, ], "relationships": [ diff --git a/ibis-server/tests/routers/v2/connector/test_minio_file.py b/ibis-server/tests/routers/v2/connector/test_minio_file.py index a0441e76a..4dcf94982 100644 --- a/ibis-server/tests/routers/v2/connector/test_minio_file.py +++ b/ibis-server/tests/routers/v2/connector/test_minio_file.py @@ -39,7 +39,7 @@ "type": "varchar", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "Customer", @@ -64,7 +64,7 @@ "expression": "sum(orders.totalprice)", }, ], - "primaryKey": "custkey", + "primaryKey": ["custkey"], }, ], "relationships": [ diff --git a/ibis-server/tests/routers/v2/connector/test_mssql.py b/ibis-server/tests/routers/v2/connector/test_mssql.py index 414bdd842..d1a5274ae 100644 --- a/ibis-server/tests/routers/v2/connector/test_mssql.py +++ b/ibis-server/tests/routers/v2/connector/test_mssql.py @@ -62,7 +62,7 @@ "type": "bytea", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, ], } diff --git a/ibis-server/tests/routers/v2/connector/test_mysql.py b/ibis-server/tests/routers/v2/connector/test_mysql.py index 7f1fb7753..e83a60d54 100644 --- a/ibis-server/tests/routers/v2/connector/test_mysql.py +++ b/ibis-server/tests/routers/v2/connector/test_mysql.py @@ -64,7 +64,7 @@ "type": "bytea", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "Customer", @@ -73,7 +73,7 @@ {"name": "custkey", "expression": "c_custkey", "type": "integer"}, {"name": "name", "expression": "c_name", "type": "varchar"}, ], - "primaryKey": "custkey", + "primaryKey": ["custkey"], }, ], } diff --git a/ibis-server/tests/routers/v2/connector/test_oracle.py b/ibis-server/tests/routers/v2/connector/test_oracle.py index 1d808d600..ed621f4c7 100644 --- a/ibis-server/tests/routers/v2/connector/test_oracle.py +++ b/ibis-server/tests/routers/v2/connector/test_oracle.py @@ -71,7 +71,7 @@ "type": "blob", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], } ], } diff --git a/ibis-server/tests/routers/v2/connector/test_postgres.py b/ibis-server/tests/routers/v2/connector/test_postgres.py index 4c6a11a6b..f98f30141 100644 --- a/ibis-server/tests/routers/v2/connector/test_postgres.py +++ b/ibis-server/tests/routers/v2/connector/test_postgres.py @@ -66,7 +66,7 @@ "type": "bytea", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "Customer", diff --git a/ibis-server/tests/routers/v2/connector/test_s3_file.py b/ibis-server/tests/routers/v2/connector/test_s3_file.py index 8e855fc9b..af1bc6772 100644 --- a/ibis-server/tests/routers/v2/connector/test_s3_file.py +++ b/ibis-server/tests/routers/v2/connector/test_s3_file.py @@ -41,7 +41,7 @@ "type": "varchar", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "Customer", @@ -66,7 +66,7 @@ "expression": "sum(orders.totalprice)", }, ], - "primaryKey": "custkey", + "primaryKey": ["custkey"], }, ], "relationships": [ diff --git a/ibis-server/tests/routers/v2/connector/test_snowflake.py b/ibis-server/tests/routers/v2/connector/test_snowflake.py index ee374d7bb..4d62df50a 100644 --- a/ibis-server/tests/routers/v2/connector/test_snowflake.py +++ b/ibis-server/tests/routers/v2/connector/test_snowflake.py @@ -61,7 +61,7 @@ "type": "timestamp", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, ], } diff --git a/ibis-server/tests/routers/v2/connector/test_trino.py b/ibis-server/tests/routers/v2/connector/test_trino.py index 2ff1a1f5c..19623cd62 100644 --- a/ibis-server/tests/routers/v2/connector/test_trino.py +++ b/ibis-server/tests/routers/v2/connector/test_trino.py @@ -59,7 +59,7 @@ "type": "bytea", }, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, ], } diff --git a/ibis-server/tests/routers/v2/test_analysis.py b/ibis-server/tests/routers/v2/test_analysis.py index 622de6bbd..31f1a372f 100644 --- a/ibis-server/tests/routers/v2/test_analysis.py +++ b/ibis-server/tests/routers/v2/test_analysis.py @@ -20,7 +20,7 @@ {"name": "mktsegment", "expression": "mktsegment", "type": "varchar"}, {"name": "comment", "expression": "comment", "type": "varchar"}, ], - "primaryKey": "custkey", + "primaryKey": ["custkey"], }, { "name": "orders", @@ -44,7 +44,7 @@ }, {"name": "comment", "expression": "comment", "type": "varchar"}, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "lineitem", diff --git a/ibis-server/tests/routers/v2/test_relationship_valid.py b/ibis-server/tests/routers/v2/test_relationship_valid.py index 50837bdfc..286c8248d 100644 --- a/ibis-server/tests/routers/v2/test_relationship_valid.py +++ b/ibis-server/tests/routers/v2/test_relationship_valid.py @@ -17,7 +17,7 @@ {"name": "id", "type": "integer"}, {"name": "many_col", "type": "integer"}, ], - "primaryKey": "id", + "primaryKey": ["id"], }, { "name": "t2", diff --git a/ibis-server/tests/routers/v3/connector/bigquery/test_query.py b/ibis-server/tests/routers/v3/connector/bigquery/test_query.py index ebdd0fac8..278dacede 100644 --- a/ibis-server/tests/routers/v3/connector/bigquery/test_query.py +++ b/ibis-server/tests/routers/v3/connector/bigquery/test_query.py @@ -53,7 +53,7 @@ "type": "timestamptz", }, ], - "primaryKey": "o_orderkey", + "primaryKey": ["o_orderkey"], }, ], } diff --git a/ibis-server/tests/routers/v3/connector/local_file/test_query.py b/ibis-server/tests/routers/v3/connector/local_file/test_query.py index 19c924086..7a9f8c2e3 100644 --- a/ibis-server/tests/routers/v3/connector/local_file/test_query.py +++ b/ibis-server/tests/routers/v3/connector/local_file/test_query.py @@ -29,7 +29,7 @@ }, {"name": "orderdate", "expression": "o_orderdate", "type": "date"}, ], - "primaryKey": "orderkey", + "primaryKey": ["orderkey"], }, { "name": "Customer", @@ -54,7 +54,7 @@ "expression": 'sum("Orders".totalprice)', }, ], - "primaryKey": "custkey", + "primaryKey": ["custkey"], }, ], "relationships": [ diff --git a/ibis-server/tests/routers/v3/connector/postgres/test_query.py b/ibis-server/tests/routers/v3/connector/postgres/test_query.py index 039f70405..93cf17d9f 100644 --- a/ibis-server/tests/routers/v3/connector/postgres/test_query.py +++ b/ibis-server/tests/routers/v3/connector/postgres/test_query.py @@ -59,7 +59,7 @@ "type": "timestamptz", }, ], - "primaryKey": "o_orderkey", + "primaryKey": ["o_orderkey"], }, { "name": "customer", @@ -78,7 +78,7 @@ "expression": "sum(orders.o_totalprice_double)", }, ], - "primaryKey": "c_custkey", + "primaryKey": ["c_custkey"], }, ], "relationships": [ diff --git a/wren-base/src/main/java/io/wren/base/dto/Model.java b/wren-base/src/main/java/io/wren/base/dto/Model.java index 23dab92c4..ec66916cc 100644 --- a/wren-base/src/main/java/io/wren/base/dto/Model.java +++ b/wren-base/src/main/java/io/wren/base/dto/Model.java @@ -34,7 +34,7 @@ public class Model private final String baseObject; private final TableReference tableReference; private final List columns; - private final String primaryKey; + private final List primaryKey; private final boolean cached; private final Duration refreshTime; @@ -48,17 +48,17 @@ public static Model model(String name, String refSql, List columns, bool return new Model(name, refSql, null, null, columns, null, cached, null); } - public static Model model(String name, String refSql, List columns, String primaryKey) + public static Model model(String name, String refSql, List columns, List primaryKey) { return new Model(name, refSql, null, null, columns, primaryKey, false, null); } - public static Model onBaseObject(String name, String baseObject, List columns, String primaryKey) + public static Model onBaseObject(String name, String baseObject, List columns, List primaryKey) { return new Model(name, null, baseObject, null, columns, primaryKey, false, null); } - public static Model onTableReference(String name, TableReference tableReference, List columns, String primaryKey) + public static Model onTableReference(String name, TableReference tableReference, List columns, List primaryKey) { return new Model(name, null, null, tableReference, columns, primaryKey, false, null); } @@ -70,7 +70,7 @@ public Model( @JsonProperty("baseObject") String baseObject, @JsonProperty("tableReference") TableReference tableReference, @JsonProperty("columns") List columns, - @JsonProperty("primaryKey") String primaryKey, + @JsonProperty("primaryKey") List primaryKey, @JsonProperty("cached") boolean cached, @JsonProperty("refreshTime") Duration refreshTime) { @@ -130,7 +130,7 @@ public List getColumns() } @JsonProperty - public String getPrimaryKey() + public List getPrimaryKey() { return primaryKey; } diff --git a/wren-base/src/test/java/io/wren/base/dto/TestMacro.java b/wren-base/src/test/java/io/wren/base/dto/TestMacro.java index 1a29a3655..8cb49efd1 100644 --- a/wren-base/src/test/java/io/wren/base/dto/TestMacro.java +++ b/wren-base/src/test/java/io/wren/base/dto/TestMacro.java @@ -80,7 +80,7 @@ public void testOneParameterCall() Column.column("custkey_pass2Macro", WrenTypes.INTEGER, null, true, "{{pass2Macro('custkey', addOne, addTwo)}}"), Column.column("custkey_sum_addOne", WrenTypes.INTEGER, null, true, "{{addOne('sum(custkey)')}}"), Column.column("name", WrenTypes.VARCHAR, null, true)), - "pk"))) + List.of("pk")))) .setMacros(List.of( Macro.macro("addOne", "(text: Expression) => {{ text }} + 1"), Macro.macro("addTwo", "(text: Expression) => {{ text }} + 2"), @@ -115,7 +115,7 @@ public void testTwoParameterCall() Column.column("custkey_concat_name", WrenTypes.INTEGER, null, true, "{{concat('custkey', 'name')}}"), Column.column("custkey_callAddOne", WrenTypes.INTEGER, null, true, "{{addPrefixOne('custkey')}}"), Column.column("custkey_pass1Macro", WrenTypes.INTEGER, null, true, "{{pass1Macro('custkey', 'name', concat)}}")), - "pk"))) + List.of("pk")))) .setMacros(List.of( Macro.macro("concat", "(text: Expression, text2: Expression) => {{ text }} || {{ text2 }}"), Macro.macro("addPrefixOne", "(text: Expression) => {{concat(\"'1'\", text)}}"), @@ -145,7 +145,7 @@ public void testZeroParameterCall() Column.column("standardTime", WrenTypes.INTEGER, null, true, "{{standardTime()}}"), Column.column("callStandardTime", WrenTypes.INTEGER, null, true, "{{callStandardTime()}}"), Column.column("passStandardTime", WrenTypes.INTEGER, null, true, "{{passStandardTime(standardTime)}}")), - "pk"))) + List.of("pk")))) .setMacros(List.of( Macro.macro("standardTime", "() => standardTime"), Macro.macro("callStandardTime", "() => {{callStandardTime()}}"), diff --git a/wren-base/src/test/java/io/wren/base/dto/TestManifestSerDe.java b/wren-base/src/test/java/io/wren/base/dto/TestManifestSerDe.java index 05cd8a185..f7d0e5f99 100644 --- a/wren-base/src/test/java/io/wren/base/dto/TestManifestSerDe.java +++ b/wren-base/src/test/java/io/wren/base/dto/TestManifestSerDe.java @@ -59,7 +59,7 @@ private static Manifest createManifest() new Column("shippriority", "integer", null, false, true, null), new Column("comment", "string", null, false, true, null), new Column("customer", "CustomerModel", "OrdersCustomer", false, true, null)), - "orderkey", + List.of("orderkey"), false, null), new Model("LineitemModel", @@ -89,7 +89,7 @@ private static Manifest createManifest() new Column("orders", "OrdersModel", "OrdersCustomer", false, true, null), // calculated field new Column("orders_totalprice", WrenTypes.VARCHAR, null, true, false, "SUM(orders.totalprice)")), - "custkey", + List.of("custkey"), false, null))) .setRelationships(List.of( diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/AbstractTestModel.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/AbstractTestModel.java index 9aacb64cf..8dd8dbd81 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/AbstractTestModel.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/AbstractTestModel.java @@ -259,7 +259,7 @@ public void testModelOnModel() ImmutableList.of( Column.column("mom_custkey", "VARCHAR", null, true, "custkey"), Column.column("mom_totalprice", "VARCHAR", null, true, "totalprice")), - "mom_custkey"); + List.of("mom_custkey")); Manifest manifest = withDefaultCatalogSchema() .setModels(List.of(newCustomer, onCustomer, orders)) .setRelationships(List.of(ordersCustomer)) diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestAllRulesRewrite.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestAllRulesRewrite.java index f9d44b350..1b5e9c7ab 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestAllRulesRewrite.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestAllRulesRewrite.java @@ -60,7 +60,7 @@ public TestAllRulesRewrite() Column.column("status", "Inventory", null, true), Column.column("statusA", "InventoryA", null, true), Column.relationshipColumn("orders", "Order", "AlbumOrder")), - "id"), + List.of("id")), Model.model("Band", "select * from (values (1, 'ZUTOMAYO'), " + "(2, 'Yorushika')) " + @@ -69,12 +69,12 @@ public TestAllRulesRewrite() Column.column("id", WrenTypes.INTEGER, null, true), Column.column("name", WrenTypes.VARCHAR, null, true), Column.relationshipColumn("albums", "Album", "AlbumBand")), - "id"), + List.of("id")), Model.model("Order", "select * from (values (1, 1), (2, 1), (3, 2), (4, 3)) Orders(orderkey, albumId)", List.of( Column.column("orderkey", WrenTypes.INTEGER, null, true), Column.column("albumId", WrenTypes.INTEGER, null, true)), - "orderkey"))) + List.of("orderkey")))) .setRelationships(List.of( Relationship.relationship("AlbumBand", List.of("Album", "Band"), JoinType.MANY_TO_ONE, "Album.bandId = Band.id"), Relationship.relationship("AlbumOrder", List.of("Album", "Order"), JoinType.ONE_TO_MANY, diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestCumulativeMetric.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestCumulativeMetric.java index dd1060f07..6f677f2f5 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestCumulativeMetric.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestCumulativeMetric.java @@ -115,7 +115,7 @@ public void testModelOnCumulativeMetric() List.of( Column.column("totalprice", WrenTypes.INTEGER, null, false), Column.column("orderdate", "DATE", null, false)), - "orderdate")) + List.of("orderdate"))) .build(); WrenMDL mdl = WrenMDL.fromManifest( copyOf(manifest) diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestEnumRewrite.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestEnumRewrite.java index 08b390945..1349c20ae 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestEnumRewrite.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestEnumRewrite.java @@ -50,7 +50,7 @@ public TestEnumRewrite() Column.column("name", WrenTypes.VARCHAR, null, true), Column.column("sex", "Sex", null, true), Column.column("country", "Country", null, true)), - "userId"))) + List.of("userId")))) .setEnumDefinitions(List.of( EnumDefinition.enumDefinition("Sex", List.of(EnumValue.enumValue("MALE"), EnumValue.enumValue("FEMALE"))), EnumDefinition.enumDefinition("Country", List.of(EnumValue.enumValue("TAIWAN", "tw"), EnumValue.enumValue("JAPAN", "jp"), EnumValue.enumValue("USA", "us"))))) diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestMetric.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestMetric.java index b8304f8e1..0ad1bab2e 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestMetric.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestMetric.java @@ -64,7 +64,7 @@ public TestMetric() Column.calculatedColumn("cumstomer_address", WrenTypes.VARCHAR, "customer.address"), Column.column("customer", "Customer", "OrdersCustomer", true), Column.column("lineitem", "Lineitem", "OrdersLineitem", true)), - "orderkey"), + List.of("orderkey")), Model.model("Customer", "select * from main.customer", List.of( @@ -77,7 +77,7 @@ public TestMetric() Column.column("mktsegment", WrenTypes.VARCHAR, null, true), Column.column("comment", WrenTypes.VARCHAR, null, true), Column.column("orders", "Orders", "OrdersCustomer", true)), - "custkey"), + List.of("custkey")), Model.model("Lineitem", "select * from main.lineitem", List.of( @@ -99,7 +99,7 @@ public TestMetric() Column.column("comment", WrenTypes.VARCHAR, null, true), Column.column("orderkey_linenumber", WrenTypes.VARCHAR, null, true, "concat(orderkey, '-', linenumber)"), Column.column("order_record", "Orders", "OrdersLineitem", true)), - "orderkey_linenumber"))) + List.of("orderkey_linenumber")))) .setRelationships(List.of( Relationship.relationship("OrdersCustomer", List.of("Orders", "Customer"), JoinType.MANY_TO_ONE, "Orders.custkey = Customer.custkey"), Relationship.relationship("OrdersLineitem", List.of("Orders", "Lineitem"), JoinType.ONE_TO_MANY, "Orders.orderkey = Lineitem.orderkey"))) @@ -279,7 +279,7 @@ public void testModelOnMetric() List.of( Column.column("name", WrenTypes.VARCHAR, null, true), Column.column("revenue", WrenTypes.INTEGER, null, true, "totalprice")), - "name")) + List.of("name"))) .build(); WrenMDL mdl = WrenMDL.fromManifest( copyOf(manifest) diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestModelRefSql.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestModelRefSql.java index 84873309b..14f4ef8d0 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestModelRefSql.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestModelRefSql.java @@ -14,6 +14,8 @@ package io.wren.base.sqlrewrite; +import java.util.List; + import static io.wren.base.dto.Model.model; public class TestModelRefSql @@ -22,8 +24,8 @@ public class TestModelRefSql public TestModelRefSql() { super(); - orders = model("Orders", "select * from main.orders", ordersColumns, "orderkey"); - lineitem = model("Lineitem", "select * from main.lineitem", lineitemColumns, "orderkey_linenumber"); - customer = model("Customer", "select * from main.customer", customerColumns, "custkey"); + orders = model("Orders", "select * from main.orders", ordersColumns, List.of("orderkey")); + lineitem = model("Lineitem", "select * from main.lineitem", lineitemColumns, List.of("orderkey_linenumber")); + customer = model("Customer", "select * from main.customer", customerColumns, List.of("custkey")); } } diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestModelTableReference.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestModelTableReference.java index da52fd03c..f7950c668 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestModelTableReference.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestModelTableReference.java @@ -17,14 +17,16 @@ import io.wren.base.dto.Model; import io.wren.base.dto.TableReference; +import java.util.List; + public class TestModelTableReference extends AbstractTestModel { public TestModelTableReference() { super(); - orders = Model.onTableReference("Orders", TableReference.tableReference("memory", "main", "orders"), ordersColumns, "orderkey"); - lineitem = Model.onTableReference("Lineitem", TableReference.tableReference("memory", "main", "lineitem"), lineitemColumns, "orderkey_linenumber"); - customer = Model.onTableReference("Customer", TableReference.tableReference("memory", "main", "customer"), customerColumns, "custkey"); + orders = Model.onTableReference("Orders", TableReference.tableReference("memory", "main", "orders"), ordersColumns, List.of("orderkey")); + lineitem = Model.onTableReference("Lineitem", TableReference.tableReference("memory", "main", "lineitem"), lineitemColumns, List.of("orderkey_linenumber")); + customer = Model.onTableReference("Customer", TableReference.tableReference("memory", "main", "customer"), customerColumns, List.of("custkey")); } } diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestView.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestView.java index fc1f51a73..eeeeb376e 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestView.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestView.java @@ -52,7 +52,7 @@ public TestView() Column.column("shippriority", WrenTypes.INTEGER, null, true), Column.column("comment", WrenTypes.VARCHAR, null, true), Column.column("lineitem", "Lineitem", "OrdersLineitem", true)), - "orderkey"); + List.of("orderkey")); } @Override diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestWrenDataLineage.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestWrenDataLineage.java index 67a8d7d68..65e7aa8fb 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/TestWrenDataLineage.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/TestWrenDataLineage.java @@ -62,7 +62,7 @@ public TestWrenDataLineage() Column.column("acctbal", WrenTypes.INTEGER, null, true), Column.column("mktsegment", WrenTypes.VARCHAR, null, true), Column.column("comment", WrenTypes.VARCHAR, null, true)), - "custkey"); + List.of("custkey")); orders = Model.model("Orders", "select * from main.orders", List.of( @@ -76,7 +76,7 @@ public TestWrenDataLineage() Column.column("shippriority", WrenTypes.INTEGER, null, true), Column.column("comment", WrenTypes.VARCHAR, null, true), Column.column("lineitem", "Lineitem", "OrdersLineitem", true)), - "orderkey"); + List.of("orderkey")); lineitem = Model.model("Lineitem", "select * from main.lineitem", List.of( @@ -97,7 +97,7 @@ public TestWrenDataLineage() Column.column("shipmode", WrenTypes.VARCHAR, null, true), Column.column("comment", WrenTypes.VARCHAR, null, true), Column.column("orderkey_linenumber", WrenTypes.VARCHAR, null, true, "concat(orderkey, '-', linenumber)")), - "orderkey_linenumber"); + List.of("orderkey_linenumber")); ordersCustomer = Relationship.relationship("OrdersCustomer", List.of("Orders", "Customer"), JoinType.MANY_TO_ONE, "Orders.custkey = Customer.custkey"); ordersLineitem = Relationship.relationship("OrdersLineitem", List.of("Orders", "Lineitem"), JoinType.ONE_TO_MANY, "Orders.orderkey = Lineitem.orderkey"); } @@ -211,7 +211,7 @@ public void testAnalyzeModelOnModel() Column.column("mom_name", "VARCHAR", null, true, "name"), Column.column("mom_custkey", "VARCHAR", null, true, "custkey"), Column.column("mom_totalprice", "VARCHAR", null, true, "total_price")), - "mom_custkey"); + List.of("mom_custkey")); Model newOrders = addColumnsToModel( orders, Column.column("on_customer", "OnCustomer", "OrdersOnCustomer", true), @@ -447,7 +447,7 @@ public void testAnalyzeModelOnCumulativeMetric() ImmutableList.of( Column.column("c_totalprice", WrenTypes.INTEGER, null, true, "totalprice"), Column.column("c_orderdate", WrenTypes.DATE, null, true, "orderdate")), - "orderdate"); + List.of("orderdate")); Manifest manifest = withDefaultCatalogSchema() .setModels(List.of(onDailyRevenue, orders)) diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java index 8aabd66e6..f41a94404 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestDecisionPointAnalyzer.java @@ -98,8 +98,8 @@ public TestDecisionPointAnalyzer() mdl = WrenMDL.fromManifest(Manifest.builder() .setCatalog(DEFAULT_SESSION_CONTEXT.getCatalog().orElseThrow()) .setSchema(DEFAULT_SESSION_CONTEXT.getSchema().orElseThrow()) - .setModels(List.of(onTableReference("customer", tableReference(null, "main", "customer"), customerColumns, "custkey"), - onTableReference("orders", tableReference(null, "main", "orders"), ordersColumns, "orderkey"), + .setModels(List.of(onTableReference("customer", tableReference(null, "main", "customer"), customerColumns, List.of("custkey")), + onTableReference("orders", tableReference(null, "main", "orders"), ordersColumns, List.of("orderkey")), onTableReference("lineitem", tableReference(null, "main", "lineitem"), lineitemColumns, null))) .setRelationships(List.of(Relationship.relationship("CustomerOrders", List.of("customer", "orders"), JoinType.ONE_TO_MANY, "customer.custkey = orders.custkey"))) .build()); diff --git a/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestStatementAnalyzer.java b/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestStatementAnalyzer.java index c71b323f8..d4321ed5e 100644 --- a/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestStatementAnalyzer.java +++ b/wren-base/src/test/java/io/wren/base/sqlrewrite/analyzer/TestStatementAnalyzer.java @@ -250,7 +250,7 @@ public void testTargetDotAll() column("custkey", "integer", null, false, "o_custkey"), column("customer", "Customer", "CustomerOrders", false), calculatedColumn("customer_name", "varchar", "customer.name")), - "orderkey"), + List.of("orderkey")), model("Customer", "SELECT * FROM tpch.customer", List.of(column("custkey", "integer", null, false, "c_custkey"), column("name", "varchar", null, false, "c_name"))))) diff --git a/wren-base/src/test/resources/tpch_mdl.json b/wren-base/src/test/resources/tpch_mdl.json index a2a4c925e..a1c53ba46 100644 --- a/wren-base/src/test/resources/tpch_mdl.json +++ b/wren-base/src/test/resources/tpch_mdl.json @@ -42,7 +42,7 @@ "relationship": "OrdersLineitem" } ], - "primaryKey": "orderkey" + "primaryKey": ["orderkey"] }, { "name": "Customer", @@ -74,7 +74,7 @@ "relationship": "CustomerNation" } ], - "primaryKey": "custkey" + "primaryKey": ["custkey"] }, { "name": "Lineitem", @@ -126,7 +126,7 @@ "expression": "concat(l_orderkey, l_linenumber)" } ], - "primaryKey": "orderkey_linenumber" + "primaryKey": ["orderkey_linenumber"] }, { "name": "Part", @@ -143,7 +143,7 @@ "type": "varchar" } ], - "primaryKey": "partkey" + "primaryKey": ["partkey"] }, { "name": "Nation", @@ -185,7 +185,7 @@ "relationship": "NationSupplier" } ], - "primaryKey": "nationkey" + "primaryKey": ["nationkey"] }, { "name": "Region", @@ -212,7 +212,7 @@ "relationship": "NationRegion" } ], - "primaryKey": "regionkey" + "primaryKey": ["regionkey"] } ], "relationships": [ diff --git a/wren-core-base/tests/data/mdl.json b/wren-core-base/tests/data/mdl.json index e588b4045..08a01d391 100644 --- a/wren-core-base/tests/data/mdl.json +++ b/wren-core-base/tests/data/mdl.json @@ -34,7 +34,7 @@ } } ], - "primaryKey": "c_custkey", + "primaryKey": ["c_custkey"], "properties": { "description": "This is a customer table", "maintainer": "test" @@ -70,7 +70,7 @@ "expression": "sum(customer.orders.o_totalprice)" } ], - "primaryKey": "p_custkey" + "primaryKey": ["p_custkey"] }, { "name": "orders", @@ -135,7 +135,7 @@ } } ], - "primaryKey": "o_orderkey" + "primaryKey": ["o_orderkey"] } ], "relationships": [ diff --git a/wren-core-py/tests/test_modeling_core.py b/wren-core-py/tests/test_modeling_core.py index 81d54a692..d9612ba64 100644 --- a/wren-core-py/tests/test_modeling_core.py +++ b/wren-core-py/tests/test_modeling_core.py @@ -25,7 +25,7 @@ {"name": "c_name", "type": "varchar"}, {"name": "orders", "type": "orders", "relationship": "orders_customer"}, ], - "primaryKey": "c_custkey", + "primaryKey": ["c_custkey"], }, { "name": "orders", @@ -43,7 +43,7 @@ "relationship": "orders_lineitem", }, ], - "primaryKey": "o_orderkey", + "primaryKey": ["o_orderkey"], }, { "name": "lineitem", @@ -56,7 +56,7 @@ {"name": "l_quantity", "type": "decimal"}, {"name": "l_extendedprice", "type": "decimal"}, ], - "primaryKey": "l_orderkey", + "primaryKey": ["l_orderkey"], }, ], "relationships": [ diff --git a/wren-core/core/tests/data/mdl.json b/wren-core/core/tests/data/mdl.json index f7da5b57f..5a2823ad3 100644 --- a/wren-core/core/tests/data/mdl.json +++ b/wren-core/core/tests/data/mdl.json @@ -34,7 +34,7 @@ } } ], - "primaryKey": "c_custkey", + "primaryKey": ["c_custkey"], "properties": { "description": "This is a customer table", "maintainer": "test" @@ -70,7 +70,7 @@ "expression": "sum(customer.orders.o_totalprice)" } ], - "primaryKey": "p_custkey" + "primaryKey": ["p_custkey"] }, { "name": "orders", @@ -116,7 +116,7 @@ "isCalculated": true } ], - "primaryKey": "o_orderkey" + "primaryKey": ["o_orderkey"] } ], "relationships": [ diff --git a/wren-tests/src/test/java/io/wren/testing/TestAnalysisResource.java b/wren-tests/src/test/java/io/wren/testing/TestAnalysisResource.java index c6c880a85..1a90a0c58 100644 --- a/wren-tests/src/test/java/io/wren/testing/TestAnalysisResource.java +++ b/wren-tests/src/test/java/io/wren/testing/TestAnalysisResource.java @@ -111,8 +111,8 @@ private void initData() manifest = Manifest.builder() .setCatalog(DEFAULT_SESSION_CONTEXT.getCatalog().orElseThrow()) .setSchema(DEFAULT_SESSION_CONTEXT.getSchema().orElseThrow()) - .setModels(List.of(onTableReference("customer", tableReference(null, "main", "customer"), customerColumns, "custkey"), - onTableReference("orders", tableReference(null, "main", "orders"), ordersColumns, "orderkey"), + .setModels(List.of(onTableReference("customer", tableReference(null, "main", "customer"), customerColumns, List.of("custkey")), + onTableReference("orders", tableReference(null, "main", "orders"), ordersColumns, List.of("orderkey")), onTableReference("lineitem", tableReference(null, "main", "lineitem"), lineitemColumns, null))) .build(); } diff --git a/wren-tests/src/test/java/io/wren/testing/TestMDLResource.java b/wren-tests/src/test/java/io/wren/testing/TestMDLResource.java index 71c02fb7d..523dc7cda 100644 --- a/wren-tests/src/test/java/io/wren/testing/TestMDLResource.java +++ b/wren-tests/src/test/java/io/wren/testing/TestMDLResource.java @@ -115,7 +115,7 @@ public void testDryRunAndDryPlan() column("custkey", "integer", null, false, "o_custkey"), column("customer", "Customer", "CustomerOrders", false), calculatedColumn("customer_name", "varchar", "customer.name")), - "orderkey"))) + List.of("orderkey")))) .setRelationships(List.of(relationship("CustomerOrders", List.of("Customer", "Orders"), JoinType.ONE_TO_MANY, "Customer.custkey = Orders.custkey"))) .build(); diff --git a/wren-tests/src/test/java/io/wren/testing/TestMDLResourceV2.java b/wren-tests/src/test/java/io/wren/testing/TestMDLResourceV2.java index 5378ab94b..1c17390e3 100644 --- a/wren-tests/src/test/java/io/wren/testing/TestMDLResourceV2.java +++ b/wren-tests/src/test/java/io/wren/testing/TestMDLResourceV2.java @@ -77,7 +77,7 @@ public void testDryPlan() column("custkey", "integer", null, false, "o_custkey"), column("customer", "Customer", "CustomerOrders", false), calculatedColumn("customer_name", "varchar", "customer.name")), - "orderkey"))) + List.of("orderkey")))) .setRelationships(List.of(relationship("CustomerOrders", List.of("Customer", "Orders"), JoinType.ONE_TO_MANY, "Customer.custkey = Orders.custkey"))) .build(); @@ -207,7 +207,7 @@ public void testSetManyToMany() column("custkey", "integer", null, false, "o_custkey"), column("customer", "Customer", "CustomerOrders", false), calculatedColumn("customer_name", "varchar", "customer.name")), - "orderkey"))) + List.of("orderkey")))) .setRelationships(List.of(relationship("CustomerOrders", List.of("Customer", "Orders"), JoinType.MANY_TO_MANY, "Customer.custkey = Orders.custkey"))) .build(); @@ -262,7 +262,7 @@ public void testPlanCountWithClause() column("custkey", "integer", null, false, "o_custkey"), column("customer", "Customer", "CustomerOrders", false), calculatedColumn("customer_name", "varchar", "customer.name")), - "orderkey"))) + List.of("orderkey")))) .setRelationships(List.of(relationship("CustomerOrders", List.of("Customer", "Orders"), JoinType.MANY_TO_MANY, "Customer.custkey = Orders.custkey"))) .build(); String manifestStr = base64Encode(toJson(manifest));