test/unit/org/apache/cassandra/cql3/functions/masking/ColumnMaskWithUDFTest.java - cassandra - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 package org.apache.cassandra.cql3.functions.masking;

 import org.junit.Test;

 import com.datastax.driver.core.BoundStatement;
 import com.datastax.driver.core.Session;
 import org.apache.cassandra.exceptions.InvalidRequestException;
 import org.assertj.core.api.Assertions;

 import static java.lang.String.format;
 import static org.junit.Assert.assertNull;

 /**
  * {@link ColumnMaskTester} for column masks using a UDF.
  */
 public class ColumnMaskWithUDFTest extends ColumnMaskTester
 {
     @Test
     @SuppressWarnings("resource")
     public void testUDF() throws Throwable
     {
         // create a table masked with a UDF and with a materialized view
         String function = createAddFunction();
         createTable(format("CREATE TABLE %%s (k int, c int, v int MASKED WITH %s(7), PRIMARY KEY (k, c))", function));
         String view = createView("CREATE MATERIALIZED VIEW %s AS SELECT * FROM %s " +
                                  "WHERE k IS NOT NULL AND c IS NOT NULL AND v IS NOT NULL " +
                                  "PRIMARY KEY (v, k, c)");

         // add some data
         execute("INSERT INTO %s (k, c, v) VALUES (0, 0, 10)");
         execute("INSERT INTO %s (k, c, v) VALUES (0, 1, 20)");
         Object[][] clearRows = rows(row(0, 0, 10), row(0, 1, 20));
         Object[][] maskedRows = rows(row(0, 0, 17), row(0, 1, 27));
         waitForViewMutations();

         Session session = sessionNet();

         // query the table and verify that the column mask is applied
         String tableQuery = "SELECT k, c, v FROM %s";
         assertRowsNet(executeNet(tableQuery), maskedRows);
         BoundStatement tablePrepared = session.prepare(formatQuery(tableQuery)).bind();
         assertRowsNet(session.execute(tablePrepared), maskedRows);

         // query the view and verify that the column mask is applied
         String viewQuery = format("SELECT k, c, v FROM %s.%s", KEYSPACE, view);
         BoundStatement viewPrepared = session.prepare(formatQuery(tableQuery)).bind();
         assertRowsNet(executeNet(viewQuery), maskedRows);
         assertRowsNet(session.execute(viewPrepared), maskedRows);

         // drop the column mask and verify that the column is not masked anymore
         alterTable("ALTER TABLE %s ALTER v DROP MASKED");
         assertRowsNet(executeNet(tableQuery), clearRows);
         assertRowsNet(executeNet(viewQuery), clearRows);
         assertRowsNet(session.execute(tablePrepared), clearRows);
         assertRowsNet(session.execute(viewPrepared), clearRows);

         // mask the column again, but this time with a different argument
         alterTable(format("ALTER TABLE %%s ALTER v MASKED WITH %s(8)", function));
         maskedRows = rows(row(0, 0, 18), row(0, 1, 28));
         assertRowsNet(executeNet(tableQuery), maskedRows);
         assertRowsNet(executeNet(viewQuery), maskedRows);
         assertRowsNet(session.execute(tablePrepared), maskedRows);
         assertRowsNet(session.execute(viewPrepared), maskedRows);
     }

     @Test
     public void testUDFWithReversed() throws Throwable
     {
         // create a table masked with a UDF and with a materialized view
         String function = createAddFunction();
         createTable(format("CREATE TABLE %%s (k int, c int MASKED WITH %s(100), v int, PRIMARY KEY (k, c)) " +
                            "WITH CLUSTERING ORDER BY (c DESC)", function));
         String view = createView("CREATE MATERIALIZED VIEW %s AS SELECT * FROM %s " +
                                  "WHERE k IS NOT NULL AND c IS NOT NULL AND v IS NOT NULL " +
                                  "PRIMARY KEY (v, c, K) " +
                                  "WITH CLUSTERING ORDER BY (c DESC, k ASC)");

         // add some data
         execute("INSERT INTO %s (k, c, v) VALUES (0, 1, 0)");
         execute("INSERT INTO %s (k, c, v) VALUES (0, 2, 0)");
         Object[][] clearRows = rows(row(0, 2, 0), row(0, 1, 0));
         Object[][] maskedRows = rows(row(0, 102, 0), row(0, 101, 0));
         waitForViewMutations();

         // query the table and verify that the column mask is applied
         String tableQuery = "SELECT k, c, v FROM %s";
         assertRowsNet(executeNet(tableQuery), maskedRows);

         // query the view and verify that the column mask is applied
         String viewQuery = format("SELECT k, c, v FROM %s.%s", KEYSPACE, view);
         assertRowsNet(executeNet(viewQuery), maskedRows);

         // drop the column mask and verify that the column is not masked anymore
         alterTable("ALTER TABLE %s ALTER c DROP MASKED");
         assertRowsNet(executeNet(tableQuery), clearRows);
         assertRowsNet(executeNet(viewQuery), clearRows);

         // mask the column again, but this time with a different argument
         alterTable(format("ALTER TABLE %%s ALTER c MASKED WITH %s(1000)", function));
         maskedRows = rows(row(0, 1002, 0), row(0, 1001, 0));
         assertRowsNet(executeNet(tableQuery), maskedRows);
         assertRowsNet(executeNet(viewQuery), maskedRows);
     }

     /**
      * Verifies that queries dropping a UDF that is used for masking columns are rejected.
      */
     @Test
     public void testDropUDFWithDependentMasks() throws Throwable
     {
         String function = createAddFunction();
         String mask = format("MASKED WITH %s(7)", function);
         String table1 = createTable(format("CREATE TABLE %%s (k int %s PRIMARY KEY, v int %<s)", mask));
         String table2 = createTable(format("CREATE TABLE %%s (k int %s, c int %<s, v int %<s, s int static %<s, PRIMARY KEY (k, c))", mask));

         String dropFunctionQuery = format("DROP FUNCTION %s", function);

         String message = format("Function '%s' is still referenced by column masks in tables %s, %s", function, table1, table2);
         assertInvalidThrowMessage(message, InvalidRequestException.class, dropFunctionQuery);

         alterTable(format("ALTER TABLE %s.%s ALTER k DROP MASKED", KEYSPACE, table1));
         alterTable(format("ALTER TABLE %s.%s ALTER k DROP MASKED", KEYSPACE, table2));
         assertInvalidThrowMessage(message, InvalidRequestException.class, dropFunctionQuery);

         alterTable(format("ALTER TABLE %s.%s ALTER v DROP MASKED", KEYSPACE, table1));
         alterTable(format("ALTER TABLE %s.%s ALTER v DROP MASKED", KEYSPACE, table2));
         message = format("Function '%s' is still referenced by column masks in tables %s", function, table2);
         assertInvalidThrowMessage(message, InvalidRequestException.class, dropFunctionQuery);

         alterTable(format("ALTER TABLE %s.%s ALTER c DROP MASKED", KEYSPACE, table2));
         assertInvalidThrowMessage(message, InvalidRequestException.class, dropFunctionQuery);

         alterTable(format("ALTER TABLE %s.%s ALTER s DROP MASKED", KEYSPACE, table2));
         schemaChange(dropFunctionQuery);
     }

     @Test
     public void testMissingUDF() throws Throwable
     {
         assertMaskingFails("Unable to find masking function for v, no declared function matches the signature %s()",
                            "missing_udf");
     }

     @Test
     public void testUDFWithNoArguments() throws Throwable
     {
         assertMaskingFails("Invalid number of arguments in call to function %s: 0 required but 1 provided",
                            createFunction(KEYSPACE,
                                           "",
                                           "CREATE FUNCTION %s () " +
                                           "CALLED ON NULL INPUT " +
                                           "RETURNS int " +
                                           "LANGUAGE java " +
                                           "AS 'return 42;'"));
     }

     @Test
     public void testUDFWithWrongArgumentCount() throws Throwable
     {
         assertMaskingFails("Invalid number of arguments in call to function %s: 2 required but 3 provided",
                            createAddFunction(),
                            "(1, 3)");
     }

     @Test
     public void testUDFWithWrongArgumentType() throws Throwable
     {
         assertMaskingFails("Type error: org.apache.cassandra.db.marshal.Int32Type " +
                            "cannot be passed as argument 0 of function %s of type text",
                            createFunction(KEYSPACE,
                                           "text",
                                           "CREATE FUNCTION %s (a text) " +
                                           "CALLED ON NULL INPUT " +
                                           "RETURNS text " +
                                           "LANGUAGE java " +
                                           "AS 'return \"redacted\";'"));
     }

     @Test
     public void testUDFWithWrongReturnType() throws Throwable
     {
         assertMaskingFails("Masking function %s() return type is text",
                            createFunction(KEYSPACE,
                                           "int",
                                           "CREATE FUNCTION %s (a int) " +
                                           "CALLED ON NULL INPUT " +
                                           "RETURNS text " +
                                           "LANGUAGE java " +
                                           "AS 'return \"redacted\";'"));
     }

     @Test
     public void testUDFInOtherKeyspace() throws Throwable
     {
         assertMaskingFails("Masking function %s() doesn't belong to the same keyspace as the table",
                            createFunction(KEYSPACE_PER_TEST,
                                           "int",
                                           "CREATE FUNCTION %s (input int) " +
                                           "CALLED ON NULL INPUT " +
                                           "RETURNS int " +
                                           "LANGUAGE java " +
                                           "AS 'return Integer.valueOf(input);'"));
     }

     @Test
     public void testUDA() throws Throwable
     {
         String fState = createAddFunction();
         String fFinal = createIdentityFunction();
         assertMaskingFails("Aggregate function %s() cannot be used for masking table columns",
                            createAggregate(KEYSPACE,
                                            "int",
                                            "CREATE AGGREGATE %s(int) " +
                                            "SFUNC " + shortFunctionName(fState) + " " +
                                            "STYPE int " +
                                            "FINALFUNC " + shortFunctionName(fFinal) + " " +
                                            "INITCOND 42"));
     }

     private void assertMaskingFails(String message, String function) throws Throwable
     {
         assertMaskingFails(message, function, "()");
     }

     private void assertMaskingFails(String message, String function, String arguments) throws Throwable
     {
         message = format(message, function);

         // create table should fail
         Assertions.assertThatThrownBy(() -> execute(format("CREATE TABLE %s.%s (k int PRIMARY KEY, v int MASKED WITH %s%s)",
                                                            KEYSPACE, createTableName(), function, arguments)))
                   .isInstanceOf(InvalidRequestException.class)
                   .hasMessageContaining(message);
         assertNull(currentTableMetadata());

         // alter table should fail
         String table = createTable("CREATE TABLE %s (k int PRIMARY KEY, v int)");
         Assertions.assertThatThrownBy(() -> execute(format("ALTER TABLE %s.%s ALTER v MASKED WITH %s%s",
                                                            KEYSPACE, table, function, arguments)))
                   .isInstanceOf(InvalidRequestException.class)
                   .hasMessageContaining(message);
         assertColumnIsNotMasked(table, "v");
     }

     private String createAddFunction() throws Throwable
     {
         return createFunction(KEYSPACE,
                               "int, int",
                               "CREATE FUNCTION %s (a int, b int) " +
                               "CALLED ON NULL INPUT " +
                               "RETURNS int " +
                               "LANGUAGE java " +
                               "AS 'return Integer.valueOf(a) + Integer.valueOf(b);'");
     }

     private String createIdentityFunction() throws Throwable
     {
         return createFunction(KEYSPACE,
                               "int, int",
                               "CREATE FUNCTION %s (a int) " +
                               "CALLED ON NULL INPUT " +
                               "RETURNS int " +
                               "LANGUAGE java " +
                               "AS 'return a;'");
     }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	package org.apache.cassandra.cql3.functions.masking;

	import org.junit.Test;

	import com.datastax.driver.core.BoundStatement;
	import com.datastax.driver.core.Session;
	import org.apache.cassandra.exceptions.InvalidRequestException;
	import org.assertj.core.api.Assertions;

	import static java.lang.String.format;
	import static org.junit.Assert.assertNull;

	/**
	* {@link ColumnMaskTester} for column masks using a UDF.
	*/
	public class ColumnMaskWithUDFTest extends ColumnMaskTester
	{
	@Test
	@SuppressWarnings("resource")
	public void testUDF() throws Throwable
	{
	// create a table masked with a UDF and with a materialized view
	String function = createAddFunction();
	createTable(format("CREATE TABLE %%s (k int, c int, v int MASKED WITH %s(7), PRIMARY KEY (k, c))", function));
	String view = createView("CREATE MATERIALIZED VIEW %s AS SELECT * FROM %s " +
	"WHERE k IS NOT NULL AND c IS NOT NULL AND v IS NOT NULL " +
	"PRIMARY KEY (v, k, c)");

	// add some data
	execute("INSERT INTO %s (k, c, v) VALUES (0, 0, 10)");
	execute("INSERT INTO %s (k, c, v) VALUES (0, 1, 20)");
	Object[][] clearRows = rows(row(0, 0, 10), row(0, 1, 20));
	Object[][] maskedRows = rows(row(0, 0, 17), row(0, 1, 27));
	waitForViewMutations();

	Session session = sessionNet();

	// query the table and verify that the column mask is applied
	String tableQuery = "SELECT k, c, v FROM %s";
	assertRowsNet(executeNet(tableQuery), maskedRows);
	BoundStatement tablePrepared = session.prepare(formatQuery(tableQuery)).bind();
	assertRowsNet(session.execute(tablePrepared), maskedRows);

	// query the view and verify that the column mask is applied
	String viewQuery = format("SELECT k, c, v FROM %s.%s", KEYSPACE, view);
	BoundStatement viewPrepared = session.prepare(formatQuery(tableQuery)).bind();
	assertRowsNet(executeNet(viewQuery), maskedRows);
	assertRowsNet(session.execute(viewPrepared), maskedRows);

	// drop the column mask and verify that the column is not masked anymore
	alterTable("ALTER TABLE %s ALTER v DROP MASKED");
	assertRowsNet(executeNet(tableQuery), clearRows);
	assertRowsNet(executeNet(viewQuery), clearRows);
	assertRowsNet(session.execute(tablePrepared), clearRows);
	assertRowsNet(session.execute(viewPrepared), clearRows);

	// mask the column again, but this time with a different argument
	alterTable(format("ALTER TABLE %%s ALTER v MASKED WITH %s(8)", function));
	maskedRows = rows(row(0, 0, 18), row(0, 1, 28));
	assertRowsNet(executeNet(tableQuery), maskedRows);
	assertRowsNet(executeNet(viewQuery), maskedRows);
	assertRowsNet(session.execute(tablePrepared), maskedRows);
	assertRowsNet(session.execute(viewPrepared), maskedRows);
	}

	@Test
	public void testUDFWithReversed() throws Throwable
	{
	// create a table masked with a UDF and with a materialized view
	String function = createAddFunction();
	createTable(format("CREATE TABLE %%s (k int, c int MASKED WITH %s(100), v int, PRIMARY KEY (k, c)) " +
	"WITH CLUSTERING ORDER BY (c DESC)", function));
	String view = createView("CREATE MATERIALIZED VIEW %s AS SELECT * FROM %s " +
	"WHERE k IS NOT NULL AND c IS NOT NULL AND v IS NOT NULL " +
	"PRIMARY KEY (v, c, K) " +
	"WITH CLUSTERING ORDER BY (c DESC, k ASC)");

	// add some data
	execute("INSERT INTO %s (k, c, v) VALUES (0, 1, 0)");
	execute("INSERT INTO %s (k, c, v) VALUES (0, 2, 0)");
	Object[][] clearRows = rows(row(0, 2, 0), row(0, 1, 0));
	Object[][] maskedRows = rows(row(0, 102, 0), row(0, 101, 0));
	waitForViewMutations();

	// query the table and verify that the column mask is applied
	String tableQuery = "SELECT k, c, v FROM %s";
	assertRowsNet(executeNet(tableQuery), maskedRows);

	// query the view and verify that the column mask is applied
	String viewQuery = format("SELECT k, c, v FROM %s.%s", KEYSPACE, view);
	assertRowsNet(executeNet(viewQuery), maskedRows);

	// drop the column mask and verify that the column is not masked anymore
	alterTable("ALTER TABLE %s ALTER c DROP MASKED");
	assertRowsNet(executeNet(tableQuery), clearRows);
	assertRowsNet(executeNet(viewQuery), clearRows);

	// mask the column again, but this time with a different argument
	alterTable(format("ALTER TABLE %%s ALTER c MASKED WITH %s(1000)", function));
	maskedRows = rows(row(0, 1002, 0), row(0, 1001, 0));
	assertRowsNet(executeNet(tableQuery), maskedRows);
	assertRowsNet(executeNet(viewQuery), maskedRows);
	}

	/**
	* Verifies that queries dropping a UDF that is used for masking columns are rejected.
	*/
	@Test
	public void testDropUDFWithDependentMasks() throws Throwable
	{
	String function = createAddFunction();
	String mask = format("MASKED WITH %s(7)", function);
	String table1 = createTable(format("CREATE TABLE %%s (k int %s PRIMARY KEY, v int %<s)", mask));
	String table2 = createTable(format("CREATE TABLE %%s (k int %s, c int %<s, v int %<s, s int static %<s, PRIMARY KEY (k, c))", mask));

	String dropFunctionQuery = format("DROP FUNCTION %s", function);

	String message = format("Function '%s' is still referenced by column masks in tables %s, %s", function, table1, table2);
	assertInvalidThrowMessage(message, InvalidRequestException.class, dropFunctionQuery);

	alterTable(format("ALTER TABLE %s.%s ALTER k DROP MASKED", KEYSPACE, table1));
	alterTable(format("ALTER TABLE %s.%s ALTER k DROP MASKED", KEYSPACE, table2));
	assertInvalidThrowMessage(message, InvalidRequestException.class, dropFunctionQuery);

	alterTable(format("ALTER TABLE %s.%s ALTER v DROP MASKED", KEYSPACE, table1));
	alterTable(format("ALTER TABLE %s.%s ALTER v DROP MASKED", KEYSPACE, table2));
	message = format("Function '%s' is still referenced by column masks in tables %s", function, table2);
	assertInvalidThrowMessage(message, InvalidRequestException.class, dropFunctionQuery);

	alterTable(format("ALTER TABLE %s.%s ALTER c DROP MASKED", KEYSPACE, table2));
	assertInvalidThrowMessage(message, InvalidRequestException.class, dropFunctionQuery);

	alterTable(format("ALTER TABLE %s.%s ALTER s DROP MASKED", KEYSPACE, table2));
	schemaChange(dropFunctionQuery);
	}

	@Test
	public void testMissingUDF() throws Throwable
	{
	assertMaskingFails("Unable to find masking function for v, no declared function matches the signature %s()",
	"missing_udf");
	}

	@Test
	public void testUDFWithNoArguments() throws Throwable
	{
	assertMaskingFails("Invalid number of arguments in call to function %s: 0 required but 1 provided",
	createFunction(KEYSPACE,
	"",
	"CREATE FUNCTION %s () " +
	"CALLED ON NULL INPUT " +
	"RETURNS int " +
	"LANGUAGE java " +
	"AS 'return 42;'"));
	}

	@Test
	public void testUDFWithWrongArgumentCount() throws Throwable
	{
	assertMaskingFails("Invalid number of arguments in call to function %s: 2 required but 3 provided",
	createAddFunction(),
	"(1, 3)");
	}

	@Test
	public void testUDFWithWrongArgumentType() throws Throwable
	{
	assertMaskingFails("Type error: org.apache.cassandra.db.marshal.Int32Type " +
	"cannot be passed as argument 0 of function %s of type text",
	createFunction(KEYSPACE,
	"text",
	"CREATE FUNCTION %s (a text) " +
	"CALLED ON NULL INPUT " +
	"RETURNS text " +
	"LANGUAGE java " +
	"AS 'return \"redacted\";'"));
	}

	@Test
	public void testUDFWithWrongReturnType() throws Throwable
	{
	assertMaskingFails("Masking function %s() return type is text",
	createFunction(KEYSPACE,
	"int",
	"CREATE FUNCTION %s (a int) " +
	"CALLED ON NULL INPUT " +
	"RETURNS text " +
	"LANGUAGE java " +
	"AS 'return \"redacted\";'"));
	}

	@Test
	public void testUDFInOtherKeyspace() throws Throwable
	{
	assertMaskingFails("Masking function %s() doesn't belong to the same keyspace as the table",
	createFunction(KEYSPACE_PER_TEST,
	"int",
	"CREATE FUNCTION %s (input int) " +
	"CALLED ON NULL INPUT " +
	"RETURNS int " +
	"LANGUAGE java " +
	"AS 'return Integer.valueOf(input);'"));
	}

	@Test
	public void testUDA() throws Throwable
	{
	String fState = createAddFunction();
	String fFinal = createIdentityFunction();
	assertMaskingFails("Aggregate function %s() cannot be used for masking table columns",
	createAggregate(KEYSPACE,
	"int",
	"CREATE AGGREGATE %s(int) " +
	"SFUNC " + shortFunctionName(fState) + " " +
	"STYPE int " +
	"FINALFUNC " + shortFunctionName(fFinal) + " " +
	"INITCOND 42"));
	}

	private void assertMaskingFails(String message, String function) throws Throwable
	{
	assertMaskingFails(message, function, "()");
	}

	private void assertMaskingFails(String message, String function, String arguments) throws Throwable
	{
	message = format(message, function);

	// create table should fail
	Assertions.assertThatThrownBy(() -> execute(format("CREATE TABLE %s.%s (k int PRIMARY KEY, v int MASKED WITH %s%s)",
	KEYSPACE, createTableName(), function, arguments)))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining(message);
	assertNull(currentTableMetadata());

	// alter table should fail
	String table = createTable("CREATE TABLE %s (k int PRIMARY KEY, v int)");
	Assertions.assertThatThrownBy(() -> execute(format("ALTER TABLE %s.%s ALTER v MASKED WITH %s%s",
	KEYSPACE, table, function, arguments)))
	.isInstanceOf(InvalidRequestException.class)
	.hasMessageContaining(message);
	assertColumnIsNotMasked(table, "v");
	}

	private String createAddFunction() throws Throwable
	{
	return createFunction(KEYSPACE,
	"int, int",
	"CREATE FUNCTION %s (a int, b int) " +
	"CALLED ON NULL INPUT " +
	"RETURNS int " +
	"LANGUAGE java " +
	"AS 'return Integer.valueOf(a) + Integer.valueOf(b);'");
	}

	private String createIdentityFunction() throws Throwable
	{
	return createFunction(KEYSPACE,
	"int, int",
	"CREATE FUNCTION %s (a int) " +
	"CALLED ON NULL INPUT " +
	"RETURNS int " +
	"LANGUAGE java " +
	"AS 'return a;'");
	}
	}