-
Notifications
You must be signed in to change notification settings - Fork 198
feat(ibis): postgreSQL, PostGIS support for IBIS Server. #1188
Changes from 1 commit
1da65b7
5c8957e
4dc8002
d71e3e7
dea0da1
ec18377
a79b58f
3ae6bfb
5d5c1d6
c3335ff
f819515
86bb245
7e6bb0f
e20d7e8
36ea433
dfec903
2d97880
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -11,6 +11,92 @@ | |||||||||||||
| from app.model.metadata.metadata import Metadata | ||||||||||||||
|
|
||||||||||||||
|
|
||||||||||||||
| class ExtensionHandler: | ||||||||||||||
| def __init__(self, connection): | ||||||||||||||
| self.connection = connection | ||||||||||||||
|
|
||||||||||||||
| self.handlers = { | ||||||||||||||
| "postgis": self.postgis_handler, | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| def augment(self, tables: list[Table]) -> list[Table]: | ||||||||||||||
| # Get the list of extensions from the database | ||||||||||||||
| extensions = self.get_extensions() | ||||||||||||||
|
|
||||||||||||||
| # Iterate through the extensions and call the appropriate handler | ||||||||||||||
| for ext in extensions: | ||||||||||||||
| ext_name = ext["extension_name"] | ||||||||||||||
| schema_name = ext["schema_name"] | ||||||||||||||
|
|
||||||||||||||
| if ext_name in self.handlers: | ||||||||||||||
| handler = self.handlers[ext_name] | ||||||||||||||
| tables = handler(tables, schema_name) | ||||||||||||||
|
|
||||||||||||||
| return tables | ||||||||||||||
|
|
||||||||||||||
| def get_extensions(self) -> list[str]: | ||||||||||||||
| sql = """ | ||||||||||||||
| SELECT | ||||||||||||||
| e.extname AS extension_name, | ||||||||||||||
| n.nspname AS schema_name | ||||||||||||||
| FROM | ||||||||||||||
| pg_extension e | ||||||||||||||
| JOIN | ||||||||||||||
| pg_namespace n ON n.oid = e.extnamespace; | ||||||||||||||
| """ | ||||||||||||||
| df = self.connection.sql(sql).to_pandas() | ||||||||||||||
| if df.empty: | ||||||||||||||
| return [] | ||||||||||||||
| response = df.to_dict(orient="records") | ||||||||||||||
| return response | ||||||||||||||
|
|
||||||||||||||
| def postgis_handler(self, tables: list[Table], schema_name: str) -> list[Table]: | ||||||||||||||
| # Get the list of geometry and geography columns | ||||||||||||||
| sql = f""" | ||||||||||||||
| SELECT | ||||||||||||||
| f_table_schema, | ||||||||||||||
| f_table_name, | ||||||||||||||
| f_geometry_column AS column_name, | ||||||||||||||
| 'geometry' AS column_type | ||||||||||||||
| FROM | ||||||||||||||
| {schema_name}.geometry_columns | ||||||||||||||
| UNION ALL | ||||||||||||||
| SELECT | ||||||||||||||
| f_table_schema, | ||||||||||||||
| f_table_name, | ||||||||||||||
| f_geography_column AS column_name, | ||||||||||||||
| 'geography' AS column_type | ||||||||||||||
| FROM | ||||||||||||||
| {schema_name}.geography_columns; | ||||||||||||||
| """ | ||||||||||||||
| response = self.connection.sql(sql).to_pandas().to_dict(orient="records") | ||||||||||||||
|
|
||||||||||||||
| # Update tables | ||||||||||||||
| for row in response: | ||||||||||||||
| table_name = f"{row['f_table_schema']}.{row['f_table_name']}" # ? Might want to use a global `_format_postgres_compact_table_name` function. | ||||||||||||||
| table = tables[table_name] | ||||||||||||||
| for column in table.columns: | ||||||||||||||
| if column.name == row["column_name"]: | ||||||||||||||
| column.type = self._transform_postgres_column_type( | ||||||||||||||
| row["column_type"] | ||||||||||||||
| ) | ||||||||||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
I got the following message when testing: I think it's better to cast to |
||||||||||||||
| break | ||||||||||||||
|
|
||||||||||||||
| return tables | ||||||||||||||
|
|
||||||||||||||
| def _transform_postgres_column_type(self, data_type): | ||||||||||||||
| # lower case the data_type | ||||||||||||||
| data_type = data_type.lower() | ||||||||||||||
|
|
||||||||||||||
| # Extension types | ||||||||||||||
| switcher = { | ||||||||||||||
| "geometry": RustWrenEngineColumnType.GEOMETRY, | ||||||||||||||
| "geography": RustWrenEngineColumnType.GEOGRAPHY, | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| return switcher.get(data_type, RustWrenEngineColumnType.UNKNOWN) | ||||||||||||||
|
|
||||||||||||||
|
|
||||||||||||||
| class PostgresMetadata(Metadata): | ||||||||||||||
| def __init__(self, connection_info: PostgresConnectionInfo): | ||||||||||||||
| super().__init__(connection_info) | ||||||||||||||
|
|
@@ -80,6 +166,8 @@ def get_table_list(self) -> list[Table]: | |||||||||||||
| properties=None, | ||||||||||||||
| ) | ||||||||||||||
| ) | ||||||||||||||
| extension_handler = ExtensionHandler(self.connection) | ||||||||||||||
| unique_tables = extension_handler.augment(unique_tables) | ||||||||||||||
| return list(unique_tables.values()) | ||||||||||||||
|
Comment on lines
+169
to
171
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 💡 Verification agent 🧩 Analysis chainVerify extension handler augmentation behavior The 🏁 Script executed: #!/bin/bash
# Check if the augment method modifies the tables in place
echo "Examining ExtensionHandler.augment method:"
rg -A 5 "def augment" --type py
echo "------------------------------------------------"
echo "Checking reference to unique_tables before and after augment call:"
rg -A 5 -B 2 "extension_handler.augment" --type pyLength of output: 1566 🏁 Script executed: #!/bin/bash
# Show full implementation of the augment method including its return behavior
rg -n "def augment" -A 50 ibis-server/app/model/metadata/postgres.py
# Show context around unique_tables before and after its definition and usage
rg -n "unique_tables" -C 5 ibis-server/app/model/metadata/postgres.pyLength of output: 3870 🏁 Script executed: #!/bin/bash
# Show complete ExtensionHandler class, including handler registrations
rg -n "class ExtensionHandler" -A 200 ibis-server/app/model/metadata/postgres.py
echo "---- handlers mapping ----"
rg -n "self.handlers" -A 20 ibis-server/app/model/metadata/postgres.py
echo "---- postgis_handler implementation ----"
rg -n "def postgis_handler" -A 200 ibis-server/app/model/metadata/postgres.pyLength of output: 19056 Fix type annotations for ExtensionHandler.augment and handlers The • In ibis-server/app/model/metadata/postgres.py, adjust: - class ExtensionHandler:
- def augment(self, tables: list[Table]) -> list[Table]:
+ class ExtensionHandler:
+ def augment(self, tables: dict[str, Table]) -> dict[str, Table]:
# …
- def postgis_handler(self, tables: list[Table], schema_name: str) -> list[Table]:
+ def postgis_handler(self, tables: dict[str, Table], schema_name: str) -> dict[str, Table]:
# …• Optionally remove the assignment extension_handler.augment(unique_tables)
return list(unique_tables.values())without re-assigning These changes will make the code’s behavior and its annotations consistent. |
||||||||||||||
|
|
||||||||||||||
| def get_constraints(self) -> list[Constraint]: | ||||||||||||||
|
|
||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The comment is like a to-do enhancement. It's better to add the TODO word for it.