Showing posts with label orm. Show all posts
Showing posts with label orm. Show all posts

Friday, August 28, 2020

Postgrex Ecto Types

One thing I found confusing about Postgrex, the excellent PostgreSQL adapter for Elixir, was how to use PostgresSQL-specific data types (cidr, inet, interval, lexeme, range, etc) with Ecto. As far as I can tell, while Postgrex includes structs that can be converted to each type (like Postgrex.INET etc), you still have to write your own Ecto.Type implementation for each type you use in an Ecto schema.

For example, to implement an Ecto schema for a table like the following:

CREATE TABLE hits ( id BIGSERIAL NOT NULL PRIMARY KEY, url TEXT NOT NULL, ip INET NOT NULL, inserted_at TIMESTAMP NOT NULL );

You'd need minimally to create an Ecto type implementation like the following:

# lib/my_app/inet_type.ex defmodule MyApp.InetType do @moduledoc """ `Ecto.Type` implementation for postgres `INET` type. """ use Ecto.Type def type, do: :inet def cast(term), do: {:ok, term} def dump(term), do: {:ok, term} def load(term), do: {:ok, term} end

So that you could then define the Ecto schema with your custom Ecto type:

# lib/my_app/hits/hit.ex defmodule MyApp.Hits.Hit do @moduledoc """ The Hit schema. """ use Ecto.Schema schema "hits" do field :url, :string field :ip, MyApp.InetType timestamps updated_at: false end def changeset(hit, attrs) do hit |> cast(attrs, [:url, :ip]) |> validate_required([:url, :ip]) end end

Note that in your migration code, you use the native database type name (eg inet), not your custom type name:

# priv/repo/migrations/202001010000_create_hits.exs defmodule MyApp.Repo.Migrations.CreateHits do use Ecto.Migration def change do create table(:hits) do add :url, :text, null: false add :ip, :inet, null: false timestamps updated_at: false end end end

A Fancier Type

However, the above basic InetType implementation limits this example Hit schema to working only with Postgrex.INET structs for its ip field — so, while creating a Hit record with the IP address specified via a Postgrex.INET struct works nicely:

iex> ( ...> %MyApp.Hits.Hit{} ...> |> MyApp.Hits.Hit.changeset(%{url: "/", ip: %Postgrex.INET{address: {127, 0, 0, 1}}}) ...> |> MyApp.Repo.insert!() ...> |> Map.get(:ip) ...> ) %Postgrex.INET{address: {127, 0, 0, 1}}

Creating one with the IP address specified as a string (or even a plain tuple like {127, 0, 0, 1}) won't work:

iex> ( iex> %MyApp.Hits.Hit{} ...> |> MyApp.Hits.Hit.changeset(%{url: "/", ip: "127.0.0.1"}}) ...> |> MyApp.Repo.insert!() ...> |> Map.get(:ip) ...> ) ** (Ecto.InvalidChangesetError) could not perform insert because changeset is invalid.

This can be solved by implementing a fancier version of the cast function in the MyApp.InetType module, enabling Ecto to cast strings (and tuples) to the Postgrex.INET type. Here's a version of MyApp.InetType that does that, as well as allows the Postgrex.INET struct to be serialized as a string (including when serialized to JSON with the Jason library, or when rendered as part of a Phoenix HTML template):

# lib/my_app/inet_type.ex defmodule MyApp.InetType do @moduledoc """ `Ecto.Type` implementation for postgres `INET` type. """ use Bitwise use Ecto.Type alias Postgrex.INET def type, do: :inet def cast(nil), do: {:ok, nil} def cast(""), do: {:ok, nil} def cast(%INET{address: nil}), do: {:ok, nil} def cast(%INET{} = term), do: {:ok, term} def cast(term) when is_tuple(term), do: {:ok, %INET{address: term}} def cast(term) when is_binary(term) do [addr | mask] = String.split(term, "/", parts: 2) with {:ok, address} <- parse_address(addr), {:ok, number} <- parse_netmask(mask), {:ok, netmask} <- validate_netmask(number, address) do {:ok, %INET{address: address, netmask: netmask}} else message -> {:error, [message: message]} end end def cast(_), do: :error def dump(term), do: {:ok, term} def load(term), do: {:ok, term} defp parse_address(addr) do case :inet.parse_strict_address(String.to_charlist(addr)) do {:ok, address} -> {:ok, address} _ -> "not a valid IP address" end end defp parse_netmask([]), do: {:ok, nil} defp parse_netmask([mask]) do case Integer.parse(mask) do {number, ""} -> {:ok, number} _ -> "not a CIDR netmask" end end defp validate_netmask(nil, _addr), do: {:ok, nil} defp validate_netmask(mask, _addr) when mask < 0 do "CIDR netmask cannot be negative" end defp validate_netmask(mask, addr) when mask > 32 and tuple_size(addr) == 4 do "CIDR netmask cannot be greater than 32" end defp validate_netmask(mask, _addr) when mask > 128 do "CIDR netmask cannot be greater than 128" end defp validate_netmask(mask, addr) do ipv4 = tuple_size(addr) == 4 max = if ipv4, do: 32, else: 128 subnet = if ipv4, do: 8, else: 16 bits = addr |> Tuple.to_list() |> Enum.reverse() |> Enum.with_index() |> Enum.reduce(0, fn {value, index}, acc -> acc + (value <<< (index * subnet)) end) bitmask = ((1 <<< max) - 1) ^^^ ((1 <<< (max - mask)) - 1) if (bits &&& bitmask) == bits do {:ok, mask} else "masked bits of IP address all must be 0s" end end end defimpl String.Chars, for: Postgrex.INET do def to_string(%{address: address, netmask: netmask}) do "#{address_to_string(address)}#{netmask_to_string(netmask)}" end defp address_to_string(nil), do: "" defp address_to_string(address), do: address |> :inet.ntoa() defp netmask_to_string(nil), do: "" defp netmask_to_string(netmask), do: "/#{netmask}" end defimpl Jason.Encoder, for: Postgrex.INET do def encode(term, opts), do: term |> to_string() |> Jason.Encode.string(opts) end defimpl Phoenix.HTML.Safe, for: Postgrex.INET do def to_iodata(term), do: term |> to_string() end

Alternative Canonical Representation

Note that an alternative way of implementing your Ecto type would be to make the dump and load functions round-trip the Postgrex.INET struct to and from some more convenient canonical representation (like a plain string). For example, a MyApp.InetType like the following would allow you to use plain strings to represent IP address values in your schemas (instead of Postgrex.INET structs). It would dump each such string to a Postgrex.INET struct when Ecto attempts to save the value to the database, and load the value from a Postgrex.INET struct into a string when Ecto attempts to load the value from the database:

# lib/my_app/inet_type.ex defmodule MyApp.InetType do @moduledoc """ `Ecto.Type` implementation for postgres `INET` type. """ use Ecto.Type alias Postgrex.INET def type, do: :inet def cast(nil), do: {:ok, ""} def cast(term) when is_tuple(term), do: {:ok, address_to_string(term)} def cast(term) when is_binary(term), do: {:ok, term} def cast(_), do: :error def dump(nil), do: {:ok, nil} def dump(""), do: {:ok, nil} def dump(term) when is_binary(term) do [addr | mask] = String.split(term, "/", parts: 2) with {:ok, address} <- parse_address(addr), {:ok, number} <- parse_netmask(mask), {:ok, netmask} <- validate_netmask(number, address) do {:ok, %INET{address: address, netmask: netmask}} else message -> {:error, [message: message]} end end def dump(_), do: :error def load(nil), do: {:ok, ""} def load(%INET{address: address, netmask: netmask}) do "#{address_to_string(address)}#{netmask_to_string(netmask)}" end def load(_), do: :error defp parse_address(addr) do case :inet.parse_strict_address(String.to_charlist(addr)) do {:ok, address} -> {:ok, address} _ -> "not a valid IP address" end end defp parse_netmask([]), do: {:ok, nil} defp parse_netmask([mask]) do case Integer.parse(mask) do {number, ""} -> {:ok, number} _ -> "not a CIDR netmask" end end defp validate_netmask(nil, _addr), do: {:ok, nil} defp validate_netmask(mask, _addr) when mask < 0 do "CIDR netmask cannot be negative" end defp validate_netmask(mask, addr) when mask > 32 and tuple_size(addr) == 4 do "CIDR netmask cannot be greater than 32" end defp validate_netmask(mask, _addr) when mask > 128 do "CIDR netmask cannot be greater than 128" end defp validate_netmask(mask, addr) do ipv4 = tuple_size(addr) == 4 max = if ipv4, do: 32, else: 128 subnet = if ipv4, do: 8, else: 16 bits = addr |> Tuple.to_list() |> Enum.reverse() |> Enum.with_index() |> Enum.reduce(0, fn {value, index}, acc -> acc + (value <<< (index * subnet)) end) bitmask = ((1 <<< max) - 1) ^^^ ((1 <<< (max - mask)) - 1) if (bits &&& bitmask) == bits do {:ok, mask} else "masked bits of IP address all must be 0s" end end defp address_to_string(nil), do: "" defp address_to_string(address), do: address |> :inet.ntoa() defp netmask_to_string(nil), do: "" defp netmask_to_string(netmask), do: "/#{netmask}" end

Friday, October 25, 2019

Adapting PostgreSQL Timestamps To Arrow With Psycopg2

I did some digging the other day to try to figure out how to use the excellent Python datetime library Arrow with the workhorse psycopg2 Python-PostgreSQL database adapter (plus the nifty Peewee ORM on top of psycopg2). I was pleasantly surprised how easy and painless it was to implement, with help from a blog post by Omar Rayward, and the psycopg2 docs (and source code) as a guide.

There are 5 core PostgreSQL date/time types that Arrow can handle, which psycopg2 maps to the 3 core Python date/time classes — by default through 4 core psycopg2 datatypes:

PostgreSQL Type Example Output Psycopg2 Type Python Type
timestamp [without time zone] 2001-02-03 04:05:06 PYDATETIME datetime
timestamp with time zone 2001-02-03 04:05:06-07 PYDATETIMETZ datetime
date 2001-02-03 PYDATE date
time [without time zone] 04:05:06 PYTIME time
time with time zone 04:05:06-07 PYTIME time

Arrow can be used to handle each of these 5 types, via its single Arrow class. Here's how you set up the mappings:

import arrow import psycopg2.extensions def adapt_arrow_to_psql(value): """Formats an Arrow object as a quoted string for use in a SQL statement.""" # assume Arrow object is being used for TIME datatype if date is 1900 or earlier if value.year <= 1900: value = value.format("HH:mm:ss.SZ") elif value == arrow.Arrow.max: value = "infinity" elif value == arrow.Arrow.min: value = "-infinity" return psycopg2.extensions.AsIs("'{}'".format(value)) # register adapter to format Arrow objects when passed as parameters to SQL statements psycopg2.extensions.register_adapter(arrow.Arrow, adapt_arrow_to_psql) def cast_psql_date_to_arrow(value, conn): """Parses a SQL timestamp or date string to an Arrow object.""" # handle NULL and special "infinity"/"-infinity" values if not value: return None elif value == "infinity": return arrow.Arrow.max elif value == "-infinity": return arrow.Arrow.min return arrow.get(value) def cast_psql_time_to_arrow(value, conn): """Parses a SQL time string to an Arrow object.""" # handle NULL if not value: return None # handle TIME, TIME with fractional seconds (.S), and TIME WITH TIME ZONE (Z) return arrow.get(value, ["HH:mm:ss", "HH:mm:ss.S", "HH:mm:ssZ", "HH:mm:ss.SZ"]) # override default timestamp/date converters # to convert from SQL timestamp/date results to Arrow objects psycopg2.extensions.register_type(psycopg2.extensions.new_type( ( psycopg2.extensions.PYDATETIME.values + psycopg2.extensions.PYDATETIMETZ.values + psycopg2.extensions.PYDATE.values ), "ARROW", cast_psql_date_to_arrow, )) # override default time converter to convert from SQL time results to Arrow objects psycopg2.extensions.register_type(psycopg2.extensions.new_type( psycopg2.extensions.PYTIME.values, "ARROW_TIME", cast_psql_time_to_arrow ))

The 3 slightly tricky bits are:

  1. Deciding whether to format an Arrow object as a date or a time (in adapt_arrow_to_psql()) — you may want to handle it differently, but since Arrow will parse times without dates as occurring on "0001-01-01", the simplest thing to do is assume a date with an early year (like 1900 or earlier) represents a time instead of a date (which allows round-tripping of times from PostgreSQL to Arrow and back).
  2. Handling PostgreSQL's special "-infinity" and "infinity" values when converting between PostgreSQL and Arrow dates (in adapt_arrow_to_psql() and cast_psql_date_to_arrow()) — Arrow.min and Arrow.max are the closest equivalents.
  3. Handling the 4 different time variants that PostgreSQL emits (in cast_psql_time_to_arrow()):
    • "12:34:56" (no fractional seconds or time zone)
    • "12:34:56.123456" (fractional seconds but no time zone)
    • "12:34:56-07" (no fractional seconds but time zone)
    • "12:34:56.123456-07" (fractional seconds and time zone)

With those mappings in place, you can now use Arrow objects natively with psycopg2:

import arrow import psycopg2 def test_datetimes(): conn = psycopg2.connect(dbname="mydbname", user="myuser") try: cur = conn.cursor() cur.execute(""" CREATE TABLE foo ( id SERIAL PRIMARY KEY, dt TIMESTAMP, dtz TIMESTAMP WITH TIME ZONE, d DATE, t TIME, twtz TIME WITH TIME ZONE ) """) cur.execute( "INSERT INTO foo (dt, dtz, d, t, twtz) VALUES (%s, %s, %s, %s, %s)", ( arrow.get("2001-02-03 04:05:06"), arrow.get("2001-02-03 04:05:06-07"), arrow.get("2001-02-03"), arrow.get("04:05:06", "HH:mm:ss"), arrow.get("04:05:06-07", "HH:mm:ssZ"), ), ) cur.execute("SELECT * FROM foo") result = cur.fetchone() assert result[1] == arrow.get("2001-02-03 04:05:06") assert result[2] == arrow.get("2001-02-03 04:05:06-07") assert result[3] == arrow.get("2001-02-03") assert result[4] == arrow.get("04:05:06", "HH:mm:ss") assert result[5] == arrow.get("04:05:06-07", "HH:mm:ssZ") finally: conn.rollback()

Or with the Peewee ORM, you can use Peewee's built-in date/time fields, and pass and receive Arrow objects to/from those fields:

import arrow import peewee import playhouse.postgres_ext db = playhouse.postgres_ext.PostgresqlExtDatabase("mydbname", user="myuser") class Foo(peewee.Model): dt = peewee.DateTimeField( default=arrow.utcnow, constraints=[peewee.SQL("DEFAULT (CURRENT_TIMESTAMP AT TIME ZONE 'UTC')")], ) dtz = playhouse.postgres_ext.DateTimeTZField( default=arrow.utcnow, constraints=[peewee.SQL("DEFAULT (CURRENT_TIMESTAMP AT TIME ZONE 'UTC')")], ) d = peewee.DateField( default=arrow.utcnow, constraints=[peewee.SQL("DEFAULT (CURRENT_DATE AT TIME ZONE 'UTC')")], ) t = peewee.TimeField( default=lambda: arrow.utcnow().time(), constraints=[peewee.SQL("DEFAULT (CURRENT_TIME AT TIME ZONE 'UTC')")], ) class Meta: database = db def test_datetimes(): with db.transaction() as tx: try: Foo.create_table() result = Foo.get_by_id( Foo.create( dt=arrow.get("2001-02-03 04:05:06"), dtz=arrow.get("2001-02-03 04:05:06-07"), d=arrow.get("2001-02-03"), t=arrow.get("04:05:06", "HH:mm:ss"), ).id ) assert result.dt == arrow.get("2001-02-03 04:05:06") assert result.dtz == arrow.get("2001-02-03 04:05:06-07") assert result.d == arrow.get("2001-02-03") assert result.t == arrow.get("04:05:06", "HH:mm:ss") finally: tx.rollback()