Support postgresql enums

## Support postgresql enums

Adds support for postgresql enums, addressing
#415.

This uses the
[pg_type](https://www.postgresql.org/docs/current/catalog-pg-type.html)
and
[pg_enum](https://www.postgresql.org/docs/current/catalog-pg-enum.html)
tables to discover the valid enum values for any columns using

## Limitations

There are other [postgresql custom
types](https://www.postgresql.org/docs/current/sql-createtype.html) that
I am not attempting to support in this PR:

* composite type
* a range type
* a base type
* a shell type


## Validation

Using the enum example in the [postgres
docs](https://www.postgresql.org/docs/current/datatype-enum.html), the
integration test creates a new custom enum type, uses it in the test
table, and expects the recap `EnumType` to be found.

I also added a unit test for basically the same thing.

recap/clients/postgresql.py

@@ -90,8 +90,11 @@ def schema(self, catalog: str, schema: str, table: str) -> StructType:
  f"""
  SELECT
  information_schema.columns.*,
- pg_attribute.attndims
+ pg_attribute.attndims,
+ enums.enum_values
  FROM information_schema.columns
+
+ -- Join to get the array dimensions
  JOIN pg_catalog.pg_namespace
  ON pg_catalog.pg_namespace.nspname = information_schema.columns.table_schema
  JOIN pg_catalog.pg_class
@@ -100,6 +103,18 @@ def schema(self, catalog: str, schema: str, table: str) -> StructType:
  JOIN pg_catalog.pg_attribute
  ON pg_catalog.pg_attribute.attrelid = pg_catalog.pg_class.oid
  AND pg_catalog.pg_attribute.attname = information_schema.columns.column_name
+
+ -- Join to get the enum values
+ LEFT JOIN pg_catalog.pg_type
+ ON pg_catalog.pg_type.oid = pg_catalog.pg_attribute.atttypid
+ AND pg_catalog.pg_type.typtype = 'e' -- Ensuring it's an enum type
+ LEFT JOIN (
+ SELECT
+ enumtypid,
+ array_agg(enumlabel) AS enum_values
+ FROM pg_catalog.pg_enum
+ GROUP BY enumtypid
+ ) enums ON enums.enumtypid = pg_catalog.pg_type.oid
  WHERE table_name = {self.param_style}
  AND table_schema = {self.param_style}
  AND table_catalog = {self.param_style}

recap/converters/postgresql.py

@@ -5,6 +5,7 @@
 from recap.types import (
  BoolType,
  BytesType,
+ EnumType,
  FloatType,
  IntType,
  ListType,
@@ -138,6 +139,10 @@ def _parse_type(self, column_props: dict[str, Any]) -> RecapType:
  ),
  )
  self.registry.register_alias(base_type)
+ elif data_type == "user-defined" and column_props["ENUM_VALUES"]:
+ base_type = EnumType(
+ symbols=column_props["ENUM_VALUES"],
+ )
  else:
  raise ValueError(f"Unknown data type: {data_type}")
 

tests/integration/clients/test_postgresql.py

@@ -6,6 +6,7 @@
 from recap.types import (
  BoolType,
  BytesType,
+ EnumType,
  FloatType,
  IntType,
  ListType,
@@ -30,8 +31,19 @@ def setup_class(cls):
  dbname="testdb",
  )
 
- # Create tables
+ # Create custom types
  cursor = cls.connection.cursor()
+ cursor.execute(
+ """
+ CREATE TYPE test_enum_type_mood AS ENUM (
+ 'sad',
+ 'ok',
+ 'happy'
+ );
+ """
+ )
+
+ # Create tables
  cursor.execute(
  """
  CREATE TABLE IF NOT EXISTS test_types (
@@ -55,7 +67,8 @@ def setup_class(cls):
  test_bit_array BIT(8)[],
  test_not_null_array INTEGER[] NOT NULL,
  test_int_array_2d INTEGER[][],
- test_text_array_3d TEXT[][][]
+ test_text_array_3d TEXT[][][],
+ test_enum_mood test_enum_type_mood
  );
  """
  )
@@ -66,6 +79,7 @@ def teardown_class(cls):
  # Delete tables
  cursor = cls.connection.cursor()
  cursor.execute("DROP TABLE IF EXISTS test_types;")
+ cursor.execute("DROP TYPE IF EXISTS test_enum_type_mood;")
  cls.connection.commit()
 
  # Close the connection
@@ -266,6 +280,14 @@ def test_struct_method_arrays_no_enforce_dimensions(self):
  ),
  ],
  ),
+ UnionType(
+ default=None,
+ name="test_enum_mood",
+ types=[
+ NullType(),
+ EnumType(symbols=["sad", "ok", "happy"]),
+ ],
+ ),
  ]
  validate_results(test_types_struct, expected_fields)
 
@@ -448,6 +470,14 @@ def test_struct_method_arrays_enforce_dimensions(self):
  ),
  ],
  ),
+ UnionType(
+ default=None,
+ name="test_enum_mood",
+ types=[
+ NullType(),
+ EnumType(symbols=["sad", "ok", "happy"]),
+ ],
+ ),
  ]
  validate_results(test_types_struct, expected_fields)
 

tests/unit/converters/test_postgresql.py

@@ -4,6 +4,7 @@
 from recap.types import (
  BoolType,
  BytesType,
+ EnumType,
  FloatType,
  IntType,
  ListType,
@@ -239,6 +240,20 @@
  scale=0,
  ),
  ),
+ (
+ {
+ "COLUMN_NAME": "test_enum",
+ "DATA_TYPE": "USER-DEFINED",
+ "CHARACTER_MAXIMUM_LENGTH": None,
+ "CHARACTER_OCTET_LENGTH": None,
+ "ENUM_VALUES": ["sad", "ok", "happy"],
+ "UDT_NAME": None,
+ "ATTNDIMS": 0,
+ },
+ EnumType(
+ symbols=["sad", "ok", "happy"],
+ ),
+ ),
  ],
 )
 def test_postgresql_converter(column_props, expected):