255 lines
8.1 KiB
Ruby
255 lines
8.1 KiB
Ruby
# typed: strict
|
|
module Graph
|
|
module CreateGraphTable
|
|
extend T::Sig
|
|
extend T::Helpers
|
|
requires_ancestor { ActiveRecord::Migration }
|
|
|
|
sig do
|
|
params(
|
|
graph_name: T.any(String, Symbol),
|
|
block: T.proc.bind(ChangeGraphContext).void,
|
|
).void
|
|
end
|
|
def change_graph(graph_name, &block)
|
|
context = ChangeGraphContext.new(graph_name, self)
|
|
context.instance_eval(&block)
|
|
end
|
|
|
|
sig do
|
|
params(
|
|
graph_name: T.any(String, Symbol),
|
|
block: T.nilable(T.proc.bind(ChangeGraphContext).void),
|
|
).void
|
|
end
|
|
def create_graph(graph_name, &block)
|
|
create_graph_enums(graph_name)
|
|
create_graph_nodes_table(graph_name)
|
|
create_graph_edges_table(graph_name)
|
|
change_graph(graph_name, &block) if block
|
|
end
|
|
|
|
sig do
|
|
params(
|
|
graph_name: T.any(String, Symbol),
|
|
type_name: T.class_of(Graph::Base),
|
|
index_on_type: T::Boolean,
|
|
).void
|
|
end
|
|
def register_graph_type(graph_name, type_name, index_on_type: true)
|
|
kind_plural, kind_singular =
|
|
if type_name < Graph::Node
|
|
%i[nodes node]
|
|
elsif type_name < Graph::Edge
|
|
%i[edges edge]
|
|
else
|
|
raise "unknown node type #{type_name}"
|
|
end
|
|
|
|
partition_name = type_name.partition
|
|
|
|
reversible do |dir|
|
|
dir.up do
|
|
execute <<-SQL
|
|
ALTER TYPE #{graph_name}_#{kind_singular}_partition ADD VALUE IF NOT EXISTS '#{partition_name}';
|
|
ALTER TYPE #{graph_name}_#{kind_singular}_type ADD VALUE IF NOT EXISTS '#{type_name.name}';
|
|
SQL
|
|
|
|
execute <<-SQL
|
|
CREATE TABLE IF NOT EXISTS #{graph_name}_#{kind_plural}_#{partition_name}
|
|
PARTITION OF #{graph_name}_#{kind_plural}
|
|
FOR VALUES IN ('#{partition_name}');
|
|
SQL
|
|
|
|
execute <<-SQL if index_on_type
|
|
CREATE INDEX IF NOT EXISTS index_#{graph_name}_#{kind_plural}_#{partition_name}_on_type
|
|
ON #{graph_name}_#{kind_plural}_#{partition_name} (type);
|
|
SQL
|
|
end
|
|
|
|
# TODO: remove this once we have a way to drop partitions
|
|
# partition_table_name = "#{graph_name}_#{kind_plural}_#{partition_name}"
|
|
# dir.down { execute <<-SQL }
|
|
# -- ALTER TABLE #{graph_name}_#{kind_plural}
|
|
# -- DETACH PARTITION #{partition_table_name};
|
|
# -- DROP TABLE IF EXISTS #{partition_table_name};
|
|
# SQL
|
|
end
|
|
end
|
|
|
|
sig do
|
|
params(
|
|
graph_name: T.any(String, Symbol),
|
|
node_type: T.class_of(Graph::Base),
|
|
virtual_column:
|
|
T.any(T::Array[T.any(String, Symbol)], T.any(String, Symbol)),
|
|
options: T.untyped,
|
|
).void
|
|
end
|
|
def add_virtual_index(graph_name, node_type, virtual_column, **options)
|
|
kind_singular, kind_plural =
|
|
if node_type < Graph::Node
|
|
%i[node nodes]
|
|
elsif node_type < Graph::Edge
|
|
%i[edge edges]
|
|
else
|
|
raise "unknown node type #{node_type}"
|
|
end
|
|
|
|
# validate that the enum exists for type_name in graph_node_type
|
|
partition = node_type.partition
|
|
up_only do
|
|
unless node_type.attribute_names.include?(virtual_column.to_s)
|
|
raise "virtual column '#{virtual_column}' not found in node type '#{node_type}'"
|
|
end
|
|
|
|
tn = execute <<-SQL
|
|
SELECT *
|
|
FROM unnest(enum_range(NULL::#{graph_name}_#{kind_singular}_type)) as t(name)
|
|
WHERE name::text = '#{node_type}'
|
|
SQL
|
|
raise "graph node type '#{node_type}' not registered" if tn.first.nil?
|
|
|
|
pn = execute <<-SQL
|
|
SELECT *
|
|
FROM unnest(enum_range(NULL::#{graph_name}_#{kind_singular}_partition)) as t(name)
|
|
WHERE name::text = '#{partition}'
|
|
SQL
|
|
if pn.first.nil?
|
|
raise "graph node partition '#{partition}' not registered"
|
|
end
|
|
end
|
|
|
|
index_name =
|
|
"index_#{graph_name}_on_#{node_type.name&.underscore&.gsub("/", "_")}_#{virtual_column}"
|
|
add_index(
|
|
"#{graph_name}_#{kind_plural}_#{partition}",
|
|
"(data::jsonb->>'#{virtual_column}')",
|
|
**options.merge(where: "type = '#{node_type}'", name: index_name),
|
|
)
|
|
end
|
|
|
|
private
|
|
|
|
sig { params(graph_name: T.any(String, Symbol)).void }
|
|
def create_graph_enums(graph_name)
|
|
create_enum :"#{graph_name}_node_partition",
|
|
["#{graph_name}_nodes_default"]
|
|
create_enum :"#{graph_name}_edge_partition",
|
|
["#{graph_name}_edges_default"]
|
|
create_enum :"#{graph_name}_node_type", []
|
|
create_enum :"#{graph_name}_edge_type", []
|
|
end
|
|
|
|
sig { params(graph_name: T.any(String, Symbol)).void }
|
|
def create_graph_nodes_table(graph_name)
|
|
reversible do |dir|
|
|
dir.up do
|
|
execute <<-SQL
|
|
CREATE TABLE IF NOT EXISTS #{graph_name}_nodes
|
|
(
|
|
type #{graph_name}_node_type NOT NULL,
|
|
partition #{graph_name}_node_partition NOT NULL,
|
|
id bigserial NOT NULL,
|
|
data jsonb NOT NULL DEFAULT '{}'::jsonb,
|
|
created_at timestamp(6) without time zone NOT NULL,
|
|
updated_at timestamp(6) without time zone NOT NULL,
|
|
CONSTRAINT #{graph_name}_nodes_pkey PRIMARY KEY (partition, id)
|
|
) PARTITION BY LIST (partition);
|
|
SQL
|
|
|
|
execute <<-SQL
|
|
CREATE TABLE IF NOT EXISTS #{graph_name}_nodes_default PARTITION OF #{graph_name}_nodes DEFAULT;
|
|
CREATE INDEX IF NOT EXISTS index_#{graph_name}_nodes_on_type ON #{graph_name}_nodes_default (type);
|
|
SQL
|
|
end
|
|
|
|
dir.down { execute <<-SQL }
|
|
DROP TABLE IF EXISTS #{graph_name}_nodes;
|
|
SQL
|
|
end
|
|
end
|
|
|
|
sig { params(graph_name: T.any(String, Symbol)).void }
|
|
def create_graph_edges_table(graph_name)
|
|
reversible do |dir|
|
|
dir.up do
|
|
execute <<-SQL
|
|
CREATE TABLE IF NOT EXISTS #{graph_name}_edges
|
|
(
|
|
type #{graph_name}_edge_type NOT NULL,
|
|
partition #{graph_name}_edge_partition NOT NULL,
|
|
id bigserial NOT NULL,
|
|
from_node_partition #{graph_name}_node_partition NOT NULL,
|
|
from_node_id bigint NOT NULL,
|
|
to_node_partition #{graph_name}_node_partition NOT NULL,
|
|
to_node_id bigint NOT NULL,
|
|
data jsonb NOT NULL DEFAULT '{}'::jsonb,
|
|
created_at timestamp(6) without time zone NOT NULL,
|
|
updated_at timestamp(6) without time zone NOT NULL,
|
|
CONSTRAINT #{graph_name}_edges_pkey PRIMARY KEY (partition, id)
|
|
) PARTITION BY LIST (partition);
|
|
SQL
|
|
|
|
execute <<-SQL
|
|
CREATE TABLE IF NOT EXISTS #{graph_name}_edges_default PARTITION OF #{graph_name}_edges DEFAULT;
|
|
CREATE INDEX IF NOT EXISTS index_#{graph_name}_edges_on_type ON #{graph_name}_edges_default (type);
|
|
SQL
|
|
end
|
|
|
|
dir.down { execute <<-SQL }
|
|
DROP TABLE IF EXISTS #{graph_name}_edges;
|
|
SQL
|
|
end
|
|
end
|
|
end
|
|
|
|
class ChangeGraphContext
|
|
extend T::Sig
|
|
extend T::Helpers
|
|
|
|
sig do
|
|
params(
|
|
graph_name: T.any(String, Symbol),
|
|
migration: Graph::CreateGraphTable,
|
|
).void
|
|
end
|
|
def initialize(graph_name, migration)
|
|
@graph_name = graph_name
|
|
@migration = migration
|
|
end
|
|
|
|
sig do
|
|
params(
|
|
type_names: T.class_of(Graph::Base),
|
|
index_on_type: T::Boolean,
|
|
).void
|
|
end
|
|
def register_type(*type_names, index_on_type: true)
|
|
type_names.each do |type_name|
|
|
@migration.register_graph_type(
|
|
@graph_name,
|
|
type_name,
|
|
index_on_type: index_on_type,
|
|
)
|
|
end
|
|
end
|
|
|
|
sig do
|
|
params(
|
|
node_type: T.class_of(Graph::Base),
|
|
column_name: T.any(String, Symbol),
|
|
options: T.untyped,
|
|
).void
|
|
end
|
|
def add_virtual_index(node_type, column_name, **options)
|
|
@migration.add_virtual_index(
|
|
@graph_name,
|
|
node_type,
|
|
column_name,
|
|
**options,
|
|
)
|
|
end
|
|
end
|
|
end
|