File Coverage

lib/URI/Shortener.pm

Criterion	Covered	Total	%
statement	114	145	78.6
branch	28	52	53.8
condition	7	40	17.5
subroutine	18	21	85.7
pod	5	6	83.3
total	172	264	65.1

package URI::Shortener 1.006;

#ABSTRACT: Shorten URIs so that you don't have to rely on external services

use strict;
use warnings;

use v5.012;

use Capture::Tiny qw{capture_merged};
use Carp::Always;
use POSIX qw{floor};
use DBI;
use File::Touch;
use Crypt::PRNG;


my $SCHEMA_NAMES = {
    uri_tablename    => 'uris',
    prefix_tablename => 'prefix',
    uri_idxname      => 'uri_idx',
    prefix_idxname   => 'prefix_idx',
    cipher_idxname   => 'cipher_idx',
    created_idxname  => 'created_idx',
};

our $SCHEMA_SQLITE = qq{
CREATE TABLE IF NOT EXISTS prefix_tablename (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    prefix TEXT NOT NULL UNIQUE
);

CREATE TABLE IF NOT EXISTS uri_tablename (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    prefix_id INTEGER NOT NULL REFERENCES prefix_tablename(id) ON DELETE CASCADE,
    uri TEXT NOT NULL UNIQUE,
    cipher TEXT DEFAULT NULL UNIQUE,
    created INTEGER
);

CREATE INDEX IF NOT EXISTS uri_idxname     ON uri_tablename(uri);
CREATE INDEX IF NOT EXISTS prefix_idxname  ON prefix_tablename(prefix);
CREATE INDEX IF NOT EXISTS cipher_idxname  ON uri_tablename(cipher);
CREATE INDEX IF NOT EXISTS created_idxname ON uri_tablename(created);
};

our $SCHEMA_PG = qq{
CREATE TABLE IF NOT EXISTS prefix_tablename (
    id SERIAL PRIMARY KEY,
    prefix TEXT NOT NULL UNIQUE
);

CREATE TABLE IF NOT EXISTS uri_tablename (
    id BIGSERIAL PRIMARY KEY,
    prefix_id INTEGER NOT NULL REFERENCES prefix_tablename(id) ON DELETE CASCADE,
    uri TEXT NOT NULL UNIQUE,
    cipher TEXT DEFAULT NULL UNIQUE,
    created INTEGER
);

CREATE INDEX IF NOT EXISTS uri_idxname     ON uri_tablename(uri);
CREATE INDEX IF NOT EXISTS prefix_idxname  ON prefix_tablename(prefix);
CREATE INDEX IF NOT EXISTS cipher_idxname  ON uri_tablename(cipher);
CREATE INDEX IF NOT EXISTS created_idxname ON uri_tablename(created);
};

our $SCHEMA_MYSQL = qq{
CREATE TABLE IF NOT EXISTS prefix_tablename (
    id INTEGER AUTO_INCREMENT,
    prefix TEXT NOT NULL,
    PRIMARY KEY(id)
);

CREATE TABLE IF NOT EXISTS uri_tablename (
    id BIGINT AUTO_INCREMENT,
    prefix_id INTEGER NOT NULL REFERENCES prefix_tablename(id) ON DELETE CASCADE,
    uri TEXT NOT NULL,
    cipher VARCHAR(180) DEFAULT NULL UNIQUE,
    created INTEGER,
    PRIMARY KEY(id)
);

};


sub new {
    my $class = shift;
    my %options = (
        %$SCHEMA_NAMES,
        @_
    );
    $options{dbtype} //= 'sqlite';

    $options{domain} ||= join('',('a'..'z','A'..'Z'));

    foreach my $required (qw{domain prefix dbname seed}) {
        die "$required required" unless $options{$required};
    }
    $options{length} = 12 if !$options{length} || $options{length} < 0;

    # Strip trailing slash from prefix
    $options{prefix} =~ s|/+$||;

    $options{sqlite_schema} = $SCHEMA_SQLITE;
    $options{mysql_schema}  = $SCHEMA_MYSQL;
    $options{pg_schema}     = $SCHEMA_PG;

    # Mongle the schema appropriately
    foreach my $sql_obj (keys(%$SCHEMA_NAMES)) {
        $options{sqlite_schema} =~ s/\Q$sql_obj\E/$options{$sql_obj}/gmx;
        $options{mysql_schema}  =~ s/\Q$sql_obj\E/$options{$sql_obj}/gmx;
        $options{pg_schema}     =~ s/\Q$sql_obj\E/$options{$sql_obj}/gmx;
    }

    $options{dbh} = {};

    return bless( \%options, $class );
}


sub cipher {
    my ( $self, $id ) = @_;
    my $rr = Crypt::PRNG->new('Fortuna', $self->{seed} + $id);
    return $rr->string_from($self->{domain}, $self->{length});
}


my $smash=0;
sub shorten {
    my ( $self, $uri ) = @_;

    my $query = "SELECT id, cipher FROM $self->{uri_tablename} WHERE uri=?";

    my $rows = $self->_dbh()->selectall_arrayref( $query, { Slice => {} }, $uri );
    $rows //= [];
    if (@$rows) {
        return $self->{prefix}."/".$rows->[0]{cipher} if $rows->[0]{cipher};
        my $ciphered = $self->cipher( $rows->[0]{id} );
        my $worked = $self->_dbh()->do( "UPDATE $self->{uri_tablename} SET cipher=? WHERE id=?", undef, $ciphered, $rows->[0]{id} );
        # In the (incredibly rare) event of a collision, just burn the row and move on.
        if (!$worked) {
            warn "DANGER: cipher collision detected.";
            $self->_dbh()->do( "UPDATE $self->{uri_tablename} SET uri=? WHERE id=?", undef, "$uri-BURNED$smash", $rows->[0]{id} ) or die "Could not burn row";
            $smash++;
            die "Too many failures to avoid name collisions encountered, prune your DB!" if $smash > 64;
            goto \&shorten;
        }
        return $self->{prefix} . "/" . $ciphered;
    }

    # Otherwise we need to store the URI and retrieve the ID.
    my $pis        = "SELECT id FROM $self->{prefix_tablename} WHERE prefix=?";
    my $has_prefix = $self->_dbh->selectall_arrayref( $pis, { Slice => {} }, $self->{prefix} );
    unless (@$has_prefix) {
        $self->_dbh()->do( "INSERT INTO $self->{prefix_tablename} (prefix) VALUES (?)", undef, $self->{prefix} ) or die $self->_dbh()->errstr;
    }

    my $qq = "INSERT INTO $self->{uri_tablename} (uri,created,prefix_id) VALUES (?,?,(SELECT id FROM $self->{prefix_tablename} WHERE prefix=?))";
    $self->_dbh()->do( $qq, undef, $uri, time(), $self->{prefix} ) or die $self->_dbh()->errstr;
    goto \&shorten;
}


sub lengthen {
    my ( $self, $uri ) = @_;
    my ($cipher) = $uri =~ m|^\Q$self->{prefix}\E/(.*)$|;

    my $query = "SELECT uri FROM $self->{uri_tablename} WHERE cipher=? AND prefix_id IN (SELECT id FROM $self->{prefix_tablename} WHERE prefix=?)";

    my $rows = $self->_dbh()->selectall_arrayref( $query, { Slice => {} }, $cipher, $self->{prefix} );
    $rows //= [];
    return unless @$rows;
    return $rows->[0]{uri};
}


sub prune_before {
    my ( $self, $when ) = @_;
    $self->_dbh()->do( "DELETE FROM $self->{uri_tablename} WHERE created < ?", undef, $when ) or die $self->_dbh()->errstr;
    return 1;
}

my %db_dispatch = (
    sqlite => \&_sqlite_dbh,
    pg     => \&_pg_dbh,
    mysql  => \&_my_dbh,
);

sub _dbh {
    my ($self) = @_;
    return $db_dispatch{$self->{dbtype}}->(@_);
}

sub _sqlite_dbh {
    my ($self) = @_;
    my $dbname = $self->{dbname};
    return $self->{dbh}->{$dbname} if exists $self->{dbh}->{$dbname};

    # Some systems splash down without this.  YMMV.
    File::Touch::touch($dbname) if $dbname ne ':memory:' && !-f $dbname;

    my $db = DBI->connect( "dbi:SQLite:dbname=$dbname", "", "" );
    $db->{sqlite_allow_multiple_statements} = 1;
    $db->do($self->{sqlite_schema}) or die "Could not ensure database consistency: " . $db->errstr;
    $db->{sqlite_allow_multiple_statements} = 0;
    $self->{dbh}->{$dbname} = $db;

    # Turn on fkeys
    $db->do("PRAGMA foreign_keys = ON") or die "Could not enable foreign keys";

    # Turn on WALmode, performance
    $db->do("PRAGMA journal_mode = WAL") or die "Could not enable WAL mode";

    return $db;
}

sub _pg_dbh {
    my ($self) = @_;
    my $dbname = $self->{dbname};
    return $self->{dbh}->{$dbname} if exists $self->{dbh}->{$dbname};

    my $host = $self->{dbhost} // $ENV{PGHOST} || 'localhost';
    my $port = $self->{dbport} // $ENV{PGPORT} || 5432;
    my $user = $self->{dbuser} // $ENV{PGUSER};
    my $pass = $self->{dbpass} // $ENV{PGPASSWORD};

    my $db = DBI->connect("dbi:Pg:dbname=$dbname;host=$host;port=$port", $user, $pass);

    #XXX pg is noisy even when you say 'IF NOT EXISTS'
    my $result;
    capture_merged { $result = $db->do($self->{pg_schema}) };
    die "Could not ensure database consistency: " . $db->errstr unless $result;

    $self->{dbh}->{$dbname} = $db;
    return $db;
}

sub _my_dbh {
    my ($self) = @_;
    my $dbname = $self->{dbname};
    return $self->{dbh}->{$dbname} if exists $self->{dbh}->{$dbname};

    my $host = $self->{dbhost} // $ENV{MYSQL_HOST} || 'localhost';
    my $port = $self->{dbport} // $ENV{MYSQL_TCP_PORT} || 3306;
    my $user = $self->{dbuser} // $ENV{DBI_USER};
    my $pass = $self->{dbpass} // $ENV{MYSQL_PWD};

    # Handle the mysql defaults file
    my $defaults_file = $self->{mysql_read_default_file} // "$ENV{HOME}/.my.cnf";
    my $defaults_group = $self->{mysql_read_default_group} // 'client';
    my $df = "";
    $df .= "mysql_read_default_file=$defaults_file;"   if -f $defaults_file;
    $df .= "mysql_read_default_group=$defaults_group;" if $defaults_group;

    my $dsn = "dbi:mysql:mysql_multi_statements=1;database=$dbname;".$df."host=$host;port=$port";

    my $db = DBI->connect($dsn, $user, $pass);
    $db->do($self->{mysql_schema}) or die "Could not ensure database consistency: " . $db->errstr;

    $self->{dbh}->{$dbname} = $db;
    return $db;
}


sub migrate {
    my ($self, $new) = @_;
    my $from_dbh = $self->_dbh();
    my $to_dbh   = $new->_dbh();

    my $prefixes = $from_dbh->selectall_arrayref(qq|SELECT * FROM $self->{prefix_tablename}|);
    _batch_insert($to_dbh, $new->{prefix_tablename}, @$prefixes) if @$prefixes;

    my $uris = $from_dbh->selectall_arrayref(qq|SELECT * FROM $self->{uri_tablename}|);
    _batch_insert($to_dbh, $new->{uri_tablename}, @$uris) if @$uris;
}

# Gotta batch stuff cuz sqlite has a 10k param limit
sub _batch_insert {
    my ($dbh, $tbl, @data) = @_;
    my $ncols = @{$data[0]};
    die "No columns in table $tbl" unless $ncols;

    while (my @batch = splice(@data, 0,int(10_000 / $ncols)) ) {
        my $param = join(',', map { '?' } @{$batch[0]});
        my $bind = join(',', map { "($param)" } @batch);

        my $query = "INSERT INTO $tbl VALUES $bind";

        print "Migrating ".scalar(@batch)." rows (".scalar(@data)." left) to $tbl...";
        $dbh->do($query, undef, map { @$_ } @batch);
        print "Done.\n"
    }
}

1;

__END__