Статьи

Реализация клиентской защиты на уровне строк с помощью jOOQ

Некоторое время назад мы пообещали продолжить нашу статью « Ограничения на просмотры» с продолжением, показывающим, как реализовать защиту на уровне строк на стороне клиента с помощью jOOQ .

Что такое безопасность на уровне строк?

Некоторые базы данных, такие как Oracle или готовящийся к выпуску PostgreSQL 9.5, предоставляют встроенную поддержку безопасности на уровне строк, что потрясающе, но не каждая база данных имеет эту функцию.

Безопасность на уровне строк означает, что данный сеанс базы данных может обращаться только к некоторым строкам в базе данных, но не к другим. Например, вот что Джон видит в своей сессии:

1
2
3
4
5
id  account_owner  account_name             amount
--------------------------------------------------
1   John           savings                  500.00
2   Jane           savings                 1300.00
3   John           secret poker stash  85193065.00

В приведенном выше примере предположим, что Джон и Джейн — супружеская пара, имеющая доступ к банковским счетам друг друга, за исключением того, что Джон хочет скрыть свой секретный покерный тайник от Джейн, потому что он планирует бежать с Джилл на Багамские острова, живя хорошей жизнью , Таким образом, вышеизложенное будет моделировать взгляд Джона на набор данных, тогда как ниже будет моделироваться взгляд Джейн:

1
2
3
4
id  account_owner  account_name             amount
--------------------------------------------------
1   John           savings                  500.00
2   Jane           savings                 1300.00

Приятной особенностью безопасности на уровне строк является тот факт, что данные полностью скрыты от сеанса базы данных, как будто их там даже не было. Это включает в себя агрегацию или фильтрацию, как видно из приведенных ниже примеров представления Джона:

1
2
3
4
5
6
7
id  account_owner  account_name             amount
--------------------------------------------------
1   John           savings                  500.00
2   Jane           savings                 1300.00
3   John           secret poker stash  85193065.00
--------------------------------------------------
Total John+Jane                        85194865.00

против мнения Джейн:

1
2
3
4
5
6
id  account_owner  account_name             amount
--------------------------------------------------
1   John           savings                  500.00
2   Jane           savings                 1300.00
--------------------------------------------------
Total John+Jane                            1800.00

Если ваша база данных не поддерживает безопасность на уровне строк, вам нужно как-то эмулировать ее. Один из способов сделать это — использовать представления и, возможно, временные таблицы или некоторые контекстные переменные:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
CREATE VIEW my_accounts AS
SELECT
  a.id,
  a.account_owner,
  a.account_name,
  a.amount
FROM
  accounts a
JOIN
  account_privileges p
ON
  a.id = p.a_id
WHERE
  p.privilege_owner = SYS_CONTEXT('banking', 'user');

В приведенном выше примере я использую функцию Oracle SYS_CONTEXT , которая позволяет получить доступ к глобальному контексту из текущего сеанса. Этот контекст может быть инициализирован, например, каждый раз, когда соединение JDBC выбирается из пула соединений.

Что если ваша база данных не поддерживает эти вещи?

JOOQ на помощь! Начиная с версии jOOQ 3.2 именно по этой причине был представлен интерфейс поставщика услуг (SPI) VisitListener, который позволяет пользователям jOOQ выполнять преобразования SQL AST во время генерации оператора SQL.

Мы предполагаем:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
CREATE TABLE accounts (
  id BIGINT NOT NULL PRIMARY KEY,
  account_owner VARCHAR(20) NOT NULL,
  account_name VARCHAR(20) NOT NULL,
  amount DECIMAL(18, 2) NOT NULL
);
 
CREATE TABLE transactions (
  id BIGINT NOT NULL PRIMARY KEY,
  account_id BIGINT NOT NULL,
  amount DECIMAL(18, 2) NOT NULL
);
 
INSERT INTO accounts (
  account_owner, account_name, amount
)
VALUES (1, 'John', 'savings', 500.0),
       (2, 'Jane', 'savings', 1300.0),
       (3, 'John', 'secret poker stash', 85193065.00);
 
INSERT INTO transactions (
  id, account_id, amount
)
VALUES (1, 1, 200.0),
       (2, 1, 300.0),
       (3, 2, 300.0),
       (4, 2, 800.0),
       (5, 2, 200.0),
       (6, 3, 85193065.00);

Добавление простого WHERE

Это простейший вариант использования преобразования:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
-- turn this...
SELECT
  accounts.account_name,
  accounts.amount
FROM
  accounts
 
-- ... into this
SELECT
  accounts.account_name,
  accounts.amount
FROM
  accounts
WHERE
  accounts.id IN (?, ?) -- Predicate, for simplicity

Добавление предложения AND к существующему WHERE

Преобразование должно все еще работать, если уже существует предикат

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
-- turn this...
SELECT
  accounts.account_name,
  accounts.amount
FROM
  accounts
WHERE
  accounts.account_owner = 'John'
 
-- ... into this
SELECT
  accounts.account_name,
  accounts.amount
FROM
  accounts
WHERE
  accounts.account_owner = 'John'
AND
  accounts.id IN (?, ?)

Добавление предиката также для псевдонимов таблиц

При самостоятельном присоединении или по любой другой причине вы могли применить псевдоним

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
-- turn this...
SELECT
  a.account_name,
  a.amount
FROM
  accounts a
 
-- ... into this
SELECT
  a.account_name,
  a.amount
FROM
  accounts a
WHERE
  a.id IN (?, ?)

Добавление предиката также в подзапросах / соединениях / полусоединениях

Конечно, мы не должны ограничиваться исправлением SELECT верхнего уровня. Также должно быть применено следующее преобразование:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
-- turn this...
SELECT
  t.amount,
  t.balance
FROM
  transactions t
WHERE
  t.account_id IN (
    SELECT
      a.id
    FROM
      accounts a
    WHERE
      a.account_owner = 'John'
  )
 
-- ... into this
SELECT
  t.amount,
  t.balance
FROM
  transactions t
WHERE
  t.account_id IN (
    SELECT
      a.id
    FROM
      accounts a
    WHERE
      a.account_owner = 'John'
    AND
      a.id IN (?, ?)
  )

Добавление предиката к ссылкам на внешний ключ

Это может быть легко забыто, но на самом деле, мы также хотим добавить дополнительный предикат ко всем ссылкам внешнего ключа accounts.id , а именно transactions.account_id , особенно когда таблица transactions не присоединена к таблице accounts :

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
-- turn this...
SELECT
  t.amount,
  t.balance
FROM
  transactions t
 
-- ... into this
SELECT
  t.amount,
  t.balance
FROM
  transactions t
WHERE
  t.account_id IN (?, ?)

Короче говоря, мы хотели бы найти все запросы, которые ссылаются на таблицу accounts в той или иной форме, и добавить к ней простой предикат, реализующий управление доступом.

ОТКАЗ ОТ ОТВЕТСТВЕННОСТИ : Как всегда с безопасностью, вы должны реализовать безопасность на нескольких уровнях. Преобразование AST в SQL не является тривиальным, и приведенные выше сценарии являются неполными. Кроме того, они работают только для запросов, построенных с использованием jOOQ AST, а не для простых запросов SQL или для запросов, которые выполняются непосредственно через JDBC, через Hibernate или через хранимые процедуры, представления (которые, в свою очередь, ссылаются на таблицу accounts ) или простые синонимы таблицы.

Итак, прочитайте этот пост как учебник, показывающий, как выполнять преобразование AST, а не как полное решение для защиты на уровне строк

Как это сделать с помощью JOOQ?

Теперь начинается интересная часть. Мы собираемся сделать все это с JOOQ. Прежде всего, помните, что для реализации обратно совместимой эволюции SPI мы всегда предоставляем реализацию по умолчанию для наших SPI. В этом случае мы собираемся расширить DefaultVisitListener вместо непосредственной реализации VisitListener .

Основа нашего VisitListener будет следующей:

1
2
3
4
5
6
7
8
public class AccountIDFilter extends DefaultVisitListener {
 
    final Integer[] ids;
 
    public AccountIDFilter(Integer... ids) {
        this.ids = ids;
    }
}

Другими словами, фильтрующий прослушиватель, который фильтрует заданный набор идентификаторов для помещения в предикат IN .

Теперь, во-первых, нам нужно немного вспомогательных методов. Следующие две утилиты помещают или выталкивают некоторые объекты поверх стека:

1
2
3
4
5
6
7
8
9
void push(VisitContext context) {
    conditionStack(context).push(new ArrayList<>());
    whereStack(context).push(false);
}
 
void pop(VisitContext context) {
    whereStack(context).pop();
    conditionStack(context).pop();
}

… и стек можно увидеть здесь:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
Deque<List<Condition>> conditionStack(
        VisitContext context) {
    Deque<List<Condition>> data = (Deque<List<Condition>>)
        context.data("conditions");
 
    if (data == null) {
        data = new ArrayDeque<>();
        context.data("conditions", data);
    }
 
    return data;
}
 
Deque<Boolean> whereStack(VisitContext context) {
    Deque<Boolean> data = (Deque<Boolean>)
        context.data("predicates");
 
    if (data == null) {
        data = new ArrayDeque<>();
        context.data("predicates", data);
    }
 
    return data;
}

В whereStack случае conditionStack поддерживает стек условий для каждого подзапроса, тогда как whereStack поддерживает стек флагов для каждого подзапроса. Условия — это условия, которые должны быть сгенерированы в WHERE данного подзапроса, в то время как флаги указывают, присутствует ли уже WHERE (т.е. должны ли новые условия добавляться с использованием AND , а не WHERE ).

Для удобства добавим также следующие утилиты:

01
02
03
04
05
06
07
08
09
10
11
12
List<Condition> conditions(VisitContext context) {
    return conditionStack(context).peek();
}
 
boolean where(VisitContext context) {
    return whereStack(context).peek();
}
 
void where(VisitContext context, boolean value) {
    whereStack(context).pop();
    whereStack(context).push(value);
}

Эти утилиты позволяют получить доступ как к условиям, так и к флагам в верхней части стека (в текущем подзапросе), а также для замены флага в верхней части стека.

Зачем нам нужен стек?

Это просто. Мы хотим применить предикат только локально для текущего подзапроса, в то время как jOOQ преобразует и генерирует ваш оператор SQL. Помните, когда мы преобразовывали следующее:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
SELECT
  t.amount,
  t.balance
FROM
  transactions t
WHERE
  t.account_id IN (
    SELECT
      a.id
    FROM
      accounts a
    WHERE
      a.account_owner = 'John'
    AND
      a.id IN (?, ?)
  )
AND
  t.account_id IN (?, ?)

… В конце мы хотим получить два дополнительных предиката, сгенерированных в приведенном выше запросе. Один в подзапросе выбирает из accounts , а другой в запросе верхнего уровня выбирает из transactions , но эти два предиката не должны мешать друг другу, т.е. когда jOOQ генерирует подзапрос, мы хотим видеть только объекты, которые имеют отношение к подзапрос (вершина стека).

Итак, давайте посмотрим, как и когда мы помещаем вещи в стек. Прежде всего, нам нужно прослушать начало и конец событий предложений SQL:

начало предложения SQL

Это просто. Каждый раз, когда мы вводим новый оператор SQL, мы хотим поместить новый набор данных (условия, флаги) в стек. В некотором смысле, мы создаем локальную область видимости для подзапроса:

01
02
03
04
05
06
07
08
09
10
11
12
@Override
public void clauseStart(VisitContext context) {
 
    // Enter a new SELECT clause / nested select
    // or DML statement
    if (context.clause() == SELECT ||
        context.clause() == UPDATE ||
        context.clause() == DELETE ||
        context.clause() == INSERT) {
        push(context);
    }
}

Конечно, эта область должна быть очищена в конце того же пункта:

01
02
03
04
05
06
07
08
09
10
11
12
13
@Override
public void clauseEnd(VisitContext context) {
    // [ ... more code will follow ... ]
 
    // Leave a SELECT clause / nested select
    // or DML statement
    if (context.clause() == SELECT ||
        context.clause() == UPDATE ||
        context.clause() == DELETE ||
        context.clause() == INSERT) {
        pop(context);
    }
}

Это была легкая часть. Теперь становится немного интереснее. Когда мы заканчиваем предложение, и это предложение является предложением WHERE SELECT , UPDATE или DELETE , мы хотим отобразить дополнительное ключевое слово и предикат:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
@Override
public void clauseEnd(VisitContext context) {
 
    // Append all collected predicates to the WHERE
    // clause if any
    if (context.clause() == SELECT_WHERE ||
        context.clause() == UPDATE_WHERE ||
        context.clause() == DELETE_WHERE) {
        List<Condition> conditions =
            conditions(context);
 
        if (conditions.size() > 0) {
            context.context()
                   .formatSeparator()
                   .keyword(where(context)
                   ? "and"
                   : "where"
                   )
                   .sql(' ');
 
            context.context().visit(
                DSL.condition(Operator.AND, conditions)
            );
        }
    }
 
    // [ ... code from previous snippet ... ]

Таким образом, единственное, что нам еще нужно сделать, это собрать этот List<Condition> в стеке текущего запроса, чтобы мы могли добавить его к WHERE , а также флаг, который выбирает между "and" и "where" Это можно сделать путем переопределения visitEnd() , который вызывается в конце посещения QueryPart (в отличие от предложения):

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
@Override
public void visitEnd(VisitContext context) {
 
    // We'll see what this means in a bit...
    pushConditions(context, ACCOUNTS,
        ACCOUNTS.ID, ids);
    pushConditions(context, TRANSACTIONS,
        TRANSACTIONS.ACCOUNT_ID, ids);
 
    // Check if we're rendering any condition within
    // the WHERE clause In this case, we can be sure
    // that jOOQ will render a WHERE keyword
    if (context.queryPart() instanceof Condition) {
        List<Clause> clauses = clauses(context);
 
        if (clauses.contains(SELECT_WHERE) ||
            clauses.contains(UPDATE_WHERE) ||
            clauses.contains(DELETE_WHERE)) {
            where(context, true);
        }
    }
}
 
List<Clause> clauses(VisitContext context) {
    List<Clause> result = asList(context.clauses());
    int index = result.lastIndexOf(SELECT);
 
    if (index > 0)
        return result.subList(index, result.size() - 1);
    else
        return result;
}

В конце каждой части запроса…

  • мы пытаемся поместить соответствующие условия в стек, если это применимо
  • мы проверяем наличие WHERE и устанавливаем соответствующий флаг

Итак, наконец, единственная пропавшая основная логика, содержащаяся в pushConditions() , это:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
<E> void pushConditions(
        VisitContext context,
        Table<?> table,
        Field<E> field,
        E... values) {
 
    // Check if we're visiting the given table
    if (context.queryPart() == table) {
        List<Clause> clauses = clauses(context);
 
        // ... and if we're in the context of the current
        //  subselect's FROM clause
        if (clauses.contains(SELECT_FROM) ||
            clauses.contains(UPDATE_UPDATE) ||
            clauses.contains(DELETE_DELETE)) {
 
            // If we're declaring a TABLE_ALIAS...
            // (e.g. "ACCOUNTS" as "a")
            if (clauses.contains(TABLE_ALIAS)) {
                QueryPart[] parts = context.queryParts();
 
                // ... move up the QueryPart visit path to find the
                // defining aliased table, and extract the aliased
                // field from it. (i.e. the "a" reference)
                for (int i = parts.length - 2; i >= 0; i--) {
                    if (parts[i] instanceof Table) {
                        field = ((Table<?>) parts[i]).field(field);
                        break;
                    }
                }
            }
 
            // Push a condition for the field of the
            // (potentially aliased) table
            conditions(context).add(field.in(values));
        }
    }
}

И мы сделали! Ух ты, ну, конечно, это было не так просто. АСТ трансформация никогда не бывает. Но приведенный выше алгоритм является надежным и может использоваться для выполнения всех вышеупомянутых запросов.

Тестирование выше

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
Configuration configuration = create().configuration();
 
// This configuration has full access to all rows
DSLContext fullaccess = DSL.using(configuration);
 
// This configuration has restricted access to IDs 1 and 2
DSLContext restricted = DSL.using(
    configuration.derive(
        DefaultVisitListenerProvider.providers(
            new AccountIDFilter(1, 2)
        )
    )
);
 
// Fetching accounts
assertEquals(3, fullaccess.fetch(ACCOUNTS).size());
assertEquals(2, restricted.fetch(ACCOUNTS).size());

Генерация следующего SQL:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
select
  "PUBLIC"."ACCOUNTS"."ID",
  "PUBLIC"."ACCOUNTS"."ACCOUNT_OWNER",
  "PUBLIC"."ACCOUNTS"."ACCOUNT_NAME",
  "PUBLIC"."ACCOUNTS"."AMOUNT"
from "PUBLIC"."ACCOUNTS"
---------------------------------------
select
  "PUBLIC"."ACCOUNTS"."ID",
  "PUBLIC"."ACCOUNTS"."ACCOUNT_OWNER",
  "PUBLIC"."ACCOUNTS"."ACCOUNT_NAME",
  "PUBLIC"."ACCOUNTS"."AMOUNT"
from "PUBLIC"."ACCOUNTS"
where "PUBLIC"."ACCOUNTS"."ID" in (
  1, 2
)
1
2
3
// Fetching transactions
assertEquals(6, fullaccess.fetch(TRANSACTIONS).size());
assertEquals(5, restricted.fetch(TRANSACTIONS).size());

Генерация следующего SQL:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
select
  "PUBLIC"."TRANSACTIONS"."ID",
  "PUBLIC"."TRANSACTIONS"."ACCOUNT_ID",
  "PUBLIC"."TRANSACTIONS"."AMOUNT"
from "PUBLIC"."TRANSACTIONS"
---------------------------------------
select
  "PUBLIC"."TRANSACTIONS"."ID",
  "PUBLIC"."TRANSACTIONS"."ACCOUNT_ID",
  "PUBLIC"."TRANSACTIONS"."AMOUNT"
from "PUBLIC"."TRANSACTIONS"
where "PUBLIC"."TRANSACTIONS"."ACCOUNT_ID" in (
  1, 2
)
01
02
03
04
05
06
07
08
09
10
11
12
13
// Fetching John's accounts
assertEquals(asList(1, 3), fullaccess.fetchValues(
    select(ACCOUNTS.ID)
    .from(ACCOUNTS)
    .where(ACCOUNTS.ACCOUNT_OWNER.eq("John"))
    .orderBy(1)
));
assertEquals(asList(1   ), restricted.fetchValues(
    select(ACCOUNTS.ID)
    .from(ACCOUNTS)
    .where(ACCOUNTS.ACCOUNT_OWNER.eq("John"))
    .orderBy(1)
));

Генерация следующего SQL:

01
02
03
04
05
06
07
08
09
10
11
12
select "PUBLIC"."ACCOUNTS"."ID"
from "PUBLIC"."ACCOUNTS"
where "PUBLIC"."ACCOUNTS"."ACCOUNT_OWNER" = 'John'
order by 1 asc
---------------------------------------
select "PUBLIC"."ACCOUNTS"."ID"
from "PUBLIC"."ACCOUNTS"
where "PUBLIC"."ACCOUNTS"."ACCOUNT_OWNER" = 'John'
and "PUBLIC"."ACCOUNTS"."ID" in (
  1, 2
)
order by 1 asc
01
02
03
04
05
06
07
08
09
10
11
12
13
14
// Fetching John's accounts via an aliased table
Accounts a = ACCOUNTS.as("a");
assertEquals(asList(1, 3), fullaccess.fetchValues(
    select(a.ID)
    .from(a)
    .where(a.ACCOUNT_OWNER.eq("John"))
    .orderBy(1)
));
assertEquals(asList(1   ), restricted.fetchValues(
    select(a.ID)
    .from(a)
    .where(a.ACCOUNT_OWNER.eq("John"))
    .orderBy(1)
));

Генерация следующего SQL:

01
02
03
04
05
06
07
08
09
10
11
12
select "a"."ID"
from "PUBLIC"."ACCOUNTS" "a"
where "a"."ACCOUNT_OWNER" = 'John'
order by 1 asc
---------------------------------------
select "a"."ID"
from "PUBLIC"."ACCOUNTS" "a"
where "a"."ACCOUNT_OWNER" = 'John'
and "a"."ID" in (
  1, 2
)
order by 1 asc
01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
// Fetching John's transactions
Transactions t = TRANSACTIONS.as("t");
assertEquals(asList(1, 2, 6), fullaccess.fetchValues(
    select(t.ID)
    .from(t)
    .where(t.ACCOUNT_ID.in(
        select(a.ID)
        .from(a)
        .where(a.ACCOUNT_OWNER.eq("John"))
    ))
    .orderBy(1)
));
assertEquals(asList(1, 2   ), restricted.fetchValues(
    select(t.ID)
    .from(t)
    .where(t.ACCOUNT_ID.in(
        select(a.ID)
        .from(a)
        .where(a.ACCOUNT_OWNER.eq("John"))
    ))
    .orderBy(1)
));

Генерация следующего SQL:

01
02
03
04
05
06
07
08
09
10
11
12
13
14
15
16
17
18
19
20
21
22
23
select "t"."ID"
from "PUBLIC"."TRANSACTIONS" "t"
where "t"."ACCOUNT_ID" in (
  select "a"."ID"
  from "PUBLIC"."ACCOUNTS" "a"
  where "a"."ACCOUNT_OWNER" = 'John'
)
order by 1 asc
---------------------------------------
select "t"."ID"
from "PUBLIC"."TRANSACTIONS" "t"
where "t"."ACCOUNT_ID" in (
  select "a"."ID"
  from "PUBLIC"."ACCOUNTS" "a"
  where "a"."ACCOUNT_OWNER" = 'John'
  and "a"."ID" in (
    1, 2
  )
)
and "t"."ACCOUNT_ID" in (
  1, 2
)
order by 1 asc

Вывод

Вышеприведенные примеры показали, как относительно легко реализовать защиту на уровне VisitListener с помощью VisitListener VisitListener , очень мощного интерфейса провайдера услуг для преобразования AST в SQL на стороне клиента .

Приложения на этом не заканчиваются. Легко увидеть, как можно реализовать VisitListener который выдает исключение каждый раз, когда вы запускаете запрос DML, у которого нет WHERE .

где придаточного

Другим приложением может быть замена таблицы на аналогичную таблицу, когда выполняется определенное условие.

И самое лучшее: с jOOQ вам не нужно анализировать SQL для его преобразования (что очень сложно, в зависимости от диалекта SQL). Вы уже вручную создали Абстрактное Синтаксическое Дерево, используя свободный API jOOQ , так что вы получаете все эти функции бесплатно.

Итак: Счастливого преобразования SQL!

jooq-The-лучший способ к записи-SQL-в-Java-маленький

Приложение: полный код

001
002
003
004
005
006
007
008
009
010
011
012
013
014
015
016
017
018
019
020
021
022
023
024
025
026
027
028
029
030
031
032
033
034
035
036
037
038
039
040
041
042
043
044
045
046
047
048
049
050
051
052
053
054
055
056
057
058
059
060
061
062
063
064
065
066
067
068
069
070
071
072
073
074
075
076
077
078
079
080
081
082
083
084
085
086
087
088
089
090
091
092
093
094
095
096
097
098
099
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
import static java.util.Arrays.asList;
import static org.jooq.Clause.DELETE;
import static org.jooq.Clause.DELETE_DELETE;
import static org.jooq.Clause.DELETE_WHERE;
import static org.jooq.Clause.INSERT;
import static org.jooq.Clause.SELECT;
import static org.jooq.Clause.SELECT_FROM;
import static org.jooq.Clause.SELECT_WHERE;
import static org.jooq.Clause.TABLE_ALIAS;
import static org.jooq.Clause.UPDATE;
import static org.jooq.Clause.UPDATE_UPDATE;
import static org.jooq.Clause.UPDATE_WHERE;
import static org.jooq.test.h2.generatedclasses.Tables.ACCOUNTS;
import static org.jooq.test.h2.generatedclasses.Tables.TRANSACTIONS;
 
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.List;
 
import org.jooq.Clause;
import org.jooq.Condition;
import org.jooq.Field;
import org.jooq.Operator;
import org.jooq.QueryPart;
import org.jooq.Table;
import org.jooq.VisitContext;
import org.jooq.impl.DSL;
import org.jooq.impl.DefaultVisitListener;
 
@SuppressWarnings("unchecked")
public class AccountIDFilter extends DefaultVisitListener {
 
    final Integer[] ids;
 
    public AccountIDFilter(Integer... ids) {
        this.ids = ids;
    }
 
    void push(VisitContext context) {
        conditionStack(context).push(new ArrayList<>());
        whereStack(context).push(false);
    }
 
    void pop(VisitContext context) {
        whereStack(context).pop();
        conditionStack(context).pop();
    }
 
    Deque<List<Condition>> conditionStack(VisitContext context) {
        Deque<List<Condition>> data = (Deque<List<Condition>>) context.data("conditions");
 
        if (data == null) {
            data = new ArrayDeque<>();
            context.data("conditions", data);
        }
 
        return data;
    }
 
    Deque<Boolean> whereStack(VisitContext context) {
        Deque<Boolean> data = (Deque<Boolean>) context.data("predicates");
 
        if (data == null) {
            data = new ArrayDeque<>();
            context.data("predicates", data);
        }
 
        return data;
    }
 
    List<Condition> conditions(VisitContext context) {
        return conditionStack(context).peek();
    }
 
    boolean where(VisitContext context) {
        return whereStack(context).peek();
    }
 
    void where(VisitContext context, boolean value) {
        whereStack(context).pop();
        whereStack(context).push(value);
    }
 
    <E> void pushConditions(VisitContext context, Table<?> table, Field<E> field, E... values) {
 
        // Check if we're visiting the given table
        if (context.queryPart() == table) {
            List<Clause> clauses = clauses(context);
 
            // ... and if we're in the context of the current subselect's
            // FROM clause
            if (clauses.contains(SELECT_FROM) ||
                clauses.contains(UPDATE_UPDATE) ||
                clauses.contains(DELETE_DELETE)) {
 
                // If we're declaring a TABLE_ALIAS... (e.g. "T_BOOK" as "b")
                if (clauses.contains(TABLE_ALIAS)) {
                    QueryPart[] parts = context.queryParts();
 
                    // ... move up the QueryPart visit path to find the
                    // defining aliased table, and extract the aliased
                    // field from it. (i.e. the "b" reference)
                    for (int i = parts.length - 2; i >= 0; i--) {
                        if (parts[i] instanceof Table) {
                            field = ((Table<?>) parts[i]).field(field);
                            break;
                        }
                    }
                }
 
                // Push a condition for the field of the (potentially aliased) table
                conditions(context).add(field.in(values));
            }
        }
    }
 
    /**
     * Retrieve all clauses for the current subselect level, starting with
     * the last {@link Clause#SELECT}.
     */
    List<Clause> clauses(VisitContext context) {
        List<Clause> result = asList(context.clauses());
        int index = result.lastIndexOf(SELECT);
 
        if (index > 0)
            return result.subList(index, result.size() - 1);
        else
            return result;
    }
 
    @Override
    public void clauseStart(VisitContext context) {
 
        // Enter a new SELECT clause / nested select, or DML statement
        if (context.clause() == SELECT ||
            context.clause() == UPDATE ||
            context.clause() == DELETE ||
            context.clause() == INSERT) {
            push(context);
        }
    }
 
    @Override
    public void clauseEnd(VisitContext context) {
 
        // Append all collected predicates to the WHERE clause if any
        if (context.clause() == SELECT_WHERE ||
            context.clause() == UPDATE_WHERE ||
            context.clause() == DELETE_WHERE) {
            List<Condition> conditions = conditions(context);
 
            if (conditions.size() > 0) {
                context.context()
                       .formatSeparator()
                       .keyword(where(context) ? "and" : "where")
                       .sql(' ');
 
                context.context().visit(DSL.condition(Operator.AND, conditions));
            }
        }
 
        // Leave a SELECT clause / nested select, or DML statement
        if (context.clause() == SELECT ||
            context.clause() == UPDATE ||
            context.clause() == DELETE ||
            context.clause() == INSERT) {
            pop(context);
        }
    }
 
    @Override
    public void visitEnd(VisitContext context) {
        pushConditions(context, ACCOUNTS, ACCOUNTS.ID, ids);
        pushConditions(context, TRANSACTIONS, TRANSACTIONS.ACCOUNT_ID, ids);
 
        // Check if we're rendering any condition within the WHERE clause
        // In this case, we can be sure that jOOQ will render a WHERE keyword
        if (context.queryPart() instanceof Condition) {
            List<Clause> clauses = clauses(context);
 
            if (clauses.contains(SELECT_WHERE) ||
                clauses.contains(UPDATE_WHERE) ||
                clauses.contains(DELETE_WHERE)) {
                where(context, true);
            }
        }
    }
}