Skip to content

Instantly share code, notes, and snippets.

@DeaconDesperado
Last active November 11, 2021 18:33
Show Gist options
  • Save DeaconDesperado/416f11bf91d0ca60b3c30b22626f4178 to your computer and use it in GitHub Desktop.
Save DeaconDesperado/416f11bf91d0ca60b3c30b22626f4178 to your computer and use it in GitHub Desktop.
Attempting calcite field extraction
package com.test.parsersandbox;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.calcite.runtime.CalciteContextException;
import org.apache.calcite.sql.SqlCall;
import org.apache.calcite.sql.SqlIdentifier;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlSelect;
import org.apache.calcite.sql.SqlWith;
import org.apache.calcite.sql.SqlWithItem;
import org.apache.calcite.sql.validate.SqlQualified;
import org.apache.calcite.sql.validate.SqlScopedShuttle;
import org.apache.calcite.sql.validate.SqlValidatorImpl;
import org.apache.calcite.sql.validate.SqlValidatorNamespace;
import org.apache.calcite.sql.validate.SqlValidatorScope;
import org.checkerframework.checker.nullness.qual.Nullable;
public class AccumulatingFieldsShuttle extends SqlScopedShuttle {
private List<List<String>> accumulatedFields = new ArrayList<>();
public AccumulatingFieldsShuttle(SqlValidatorScope initialScope) {
super(initialScope);
}
public static AccumulatingFieldsShuttle of(SqlValidatorImpl validator, SqlNode sql) {
if (sql instanceof SqlSelect) {
SqlSelect select = (SqlSelect) sql;
final SqlValidatorScope selectScope = validator.getSelectScope(select);
return new AccumulatingFieldsShuttle(validator, selectScope);
} else if (sql instanceof SqlWith) {
SqlWith withSql = (SqlWith) sql;
SqlWithItem withItem = (SqlWithItem) withSql.withList.get(0);
final SqlValidatorScope withScope = validator.getWithScope(withItem);
return new AccumulatingFieldsShuttle(validator, withScope);
}
throw new UnsupportedOperationException(sql.getKind().toString());
}
@Override
public @Nullable SqlNode visit(SqlIdentifier id) {
try {
SqlQualified qualified = getScope().fullyQualify(id);
SqlValidatorNamespace namespace = qualified.namespace;
if (namespace != null && namespace.getTable() != null) {
accumulatedFields.add(
Stream.concat(
namespace.getTable().getQualifiedName().stream(), qualified.suffix().stream())
.collect(Collectors.toList()));
}
} catch (CalciteContextException exc) {
// Ignore identifiers not found in the catalog for now
}
return super.visit(id);
}
public List<List<String>> getAccumulatedFields() {
return accumulatedFields;
}
@Override
protected @Nullable SqlNode visitScoped(SqlCall call) {
return call.getOperator().acceptCall(this, call);
}
}
display_name: Simple join
statement: |-
SELECT
a.gid as foobar,
b.gid as artistGid
FROM
`metadata-entities.artist.artist_YYYYMMDD` as a
JOIN
`foobar-entities.artist.foo` as b
ON
a.h = b.h
WHERE
b.artistId IS NOT NULL
assertions:
- ['metadata-entities','artist','artist_YYYYMMDD','gid']
- ['metadata-entities','artist','artist_YYYYMMDD','h']
- ['foobar-entities','artist','foo','gid']
- ['foobar-entities','artist','foo','h']
- ['foobar-entities','artist','foo','artistId']
# PASS
display_name: UNNEST
statement: |-
SELECT
a.invoice_id,
item.sku,
item.name
FROM
`orders.invoices.invoices_nov` as a,
UNNEST(line_items) item
assertions:
- ['orders', 'invoices', 'invoices_nov', 'invoice_id']
- ['orders', 'invoices', 'invoices_nov', 'items']
- ['orders', 'invoices', 'invoices_nov', 'items', 'sku']
- ['orders', 'invoices', 'invoices_nov', 'items', 'name']
# java.lang.AssertionError:
# Expected: iterable with items [<[orders, invoices, invoices_nov, invoice_id]>, <[orders, invoices, invoices_nov, line_items]>, <[orders, invoices, invoices_nov, line_items, name]>, <[orders, invoices, invoices_nov, line_items, sku]>] in any order
# but: not matched: <[orders, invoices, invoices_nov, name]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment