search

Canner / wren-engine

Wren Engine is the backbone of the Wren AI project - The semantic engine for LLMs, bringing business context to AI agents.
https://getwren.ai
Apache License 2.0
89 stars 21 forks source link

Handle Postgres data type Enum in the sqlglot and in the ibis #617

Open grieve54706 opened 1 week ago

grieve54706 commented 1 week ago

If you select a table with a data type that is enum in the Postgres, ibis will create a view and get the schema of the view. Ibis get data_type from pg_attribute is enum_operation. Ibis will put this text into sqlglot to parse to DataType. But the enum_operation is not in the sqlglot.expressions.DataType.

   File "/app/model/connector.py", line 29, in query
     return self.connection.sql(rewritten_sql, dialect="trino").to_pandas()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/app/.venv/lib/python3.11/site-packages/ibis/backends/sql/__init__.py", line 200, in sql
     schema = self._get_schema_using_query(query)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/app/.venv/lib/python3.11/site-packages/ibis/backends/postgres/__init__.py", line 610, in _get_schema_using_query
     return self.get_schema(name)
            ^^^^^^^^^^^^^^^^^^^^^
   File "/app/.venv/lib/python3.11/site-packages/ibis/backends/postgres/__init__.py", line 588, in get_schema
     {
   File "/app/.venv/lib/python3.11/site-packages/ibis/backends/postgres/__init__.py", line 589, in <dictcomp>
     col: type_mapper.from_string(typestr, nullable=nullable)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/app/.venv/lib/python3.11/site-packages/ibis/backends/sql/datatypes.py", line 426, in from_string
     sgtype = sg.parse_one(text, into=sge.DataType, read=cls.dialect)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/app/.venv/lib/python3.11/site-packages/sqlglot/__init__.py", line 137, in parse_one
     result = dialect.parse_into(into, sql, **opts)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/app/.venv/lib/python3.11/site-packages/sqlglot/dialects/dialect.py", line 511, in parse_into
     return self.parser(**opts).parse_into(expression_type, self.tokenize(sql), sql)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   File "/app/.venv/lib/python3.11/site-packages/sqlglot/parser.py", line 1210, in parse_into
     raise ParseError(
 sqlglot.errors.ParseError: Failed to parse 'enum_operation' into <class 'sqlglot.expressions.DataType'>

May we fix the enum_operation to be enum in the ibis.

grieve54706 commented 3 days ago

Follow https://github.com/ibis-project/ibis/issues/9295