0004
0004

Reputation: 1260

Custom Data Type via Alembic for H3 (Hexagonal hierarchical geospatial indexing system)

I'm currently developing a project that requires storing H3Index (https://h3geo.org/docs/core-library/h3Indexing/) values in my database using SQLAlchemy and Alembic for migrations. However, my search for existing libraries supporting H3Index has been unsuccessful. I'm considering alternatives to implement this custom datatype. Given that I've successfully utilized libraries like geoalchemy2 for similar custom types, I'm contemplating whether a similar approach could be used for H3Index. Here's my proposed approach and the issues I've encountered with it:

Snippet from custom types that have worked -- but there is a library

from geoalchemy2 import Geometry

# Example of using geoalchemy2 for a custom type
city_center_point = Column(
    Geometry('POINT', 4326, from_text='ST_GeomFromEWKT', name='geometry'), nullable=False)

# Example of using a custom type with server_default
is_disabled = db.Column(db.Boolean, nullable=False,
                        server_default=db.FetchedValue())

However, since there's no readily available library for H3Index (unless someone knows of an h3 library for alembic?), I haven't been able to implement it as easily as the example above. My proposed solution was to create a custom datatype for H3Index using SQLAlchemy's TypeDecorator. Here's the code for the proposed custom datatype and Alembic migration :

from sqlalchemy import types
from sqlalchemy.dialects.postgresql import UUID

class H3Index(types.TypeDecorator):
    impl = UUID

    def process_bind_param(self, value, dialect):
        if value is not None:
            # Convert H3Index to string representation
            return str(value)
        return None

    def process_result_value(self, value, dialect):
        if value is not None:
            # Convert string representation to H3Index
            return H3Index(value)
        return None

    """Add H3Index column
    
    Revision ID: abc123
    Revises: xyz789
    Create Date: 2024-03-21 10:00:00
    """
    
    from alembic import op
    import sqlalchemy as sa
    from sqlalchemy import Column
    
    # Import the custom datatype class
    from yourmodule import H3Index
    
    # revision identifiers, used by Alembic.
    revision = 'abc123'
    down_revision = 'xyz789'
    branch_labels = None
    depends_on = None
    
    def upgrade():
        # Add the column with H3Index datatype
        op.add_column('your_table', 
                      Column('h3_index', H3Index, nullable=True))
    
    def downgrade():
        # Drop the column
        op.drop_column('your_table', 'h3_index')

Despite writing this solution, I know it won't work I was just trying to get a ball rolling. Could someone provide guidance on whether this approach is feasible for implementing a custom datatype for H3Index in SQLAlchemy with Alembic migrations? Additionally, are there alternative approaches I could consider? Any insights or suggestions would be greatly appreciated. Thank you!

Upvotes: 0

Views: 112

Answers (1)

Jonas Frei
Jonas Frei

Reputation: 382

The following solution isn't perfect, but it works. If anyone has a better approach, please feel free to share it! :)

  1. Create a custom H3Index SQLAlchemy data type:
import sqlalchemy.types as types


class H3Index(types.UserDefinedType):
    def __init__(self, *args):
        self._args = args

    def get_col_spec(self):
        return 'h3index'

    def convert_bind_param(self, value, engine):
        return value

    def convert_result_value(self, value, engine):
        return value
  1. Define a model using this custom data type:
from sqlalchemy import Column, String

from app.database.base_class import Base
from app.data_types import H3Index


class Test(Base):
    __tablename__ = 'test'
    __table_args__ = {'schema': 'schema_name'}

    id = Column(String, primary_key=True)
    h3_index = Column(H3Index)
  1. Generate an Alembic migration:
alembic revision --autogenerate -m "command"
  1. Modify the generated migration file. In the newly created migration file, ensure the H3Index import is included:
"""003

Revision ID: 1346ea689d44
Revises: 160b09d1a4e4
Create Date: 2024-08-27 07:55:19.981987

"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

# ----> THIS LINE NEEDS TO BE ADDED <----
from app.data_types import H3Index

# revision identifiers, used by Alembic.
revision: str = '1346ea689d44'
down_revision: Union[str, None] = '160b09d1a4e4'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        'test',
        sa.Column('id', sa.String(), nullable=False),
        sa.Column('h3_index', H3Index(), nullable=True),
        sa.PrimaryKeyConstraint('id'),
        schema='schema_name',
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('test', schema='schema_name')
    # ### end Alembic commands ###

The main disadvantage of this solution is that you need to manually edit the Alembic migration file (step 4). Additionally, the migration file references your existing code, meaning that if you later change the location or name of your custom H3Index SQLAlchemy data type, you will need to update the old Alembic migration file as well.

Upvotes: 0

Related Questions