def testCheckValidSumAverageArgument(self): stack = [ util.FieldToken('Year'), util.FieldToken('Year'), util.OperatorToken('*', 2), util.HomomorphicIntToken('Invoice_Price'), util.OperatorToken('+', 2) ] expected_stack = [[['Year', 'Year', '*'], [util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price']], True, True] self.assertEqual(interpreter.CheckValidSumAverageArgument(stack), expected_stack) stack = [ 2, util.FieldToken('Year'), util.FieldToken('Year'), util.OperatorToken('*', 2), util.HomomorphicIntToken('Invoice_Price'), util.OperatorToken('+', 2), util.OperatorToken('*', 2) ] expected_stack = [[ [2, 'Year', 'Year', '*', '*'], [2, util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price', '*'] ], True, True] self.assertEqual(interpreter.CheckValidSumAverageArgument(stack), expected_stack) stack = [util.ProbabilisticToken('Price')] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, interpreter.CheckValidSumAverageArgument, stack) stack = [ util.HomomorphicIntToken('Invoice_Price'), util.HomomorphicFloatToken('Holdback_Percentage'), util.OperatorToken('*', 2) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, interpreter.CheckValidSumAverageArgument, stack) stack = [ util.HomomorphicIntToken('Invoice_Price'), util.FieldToken('Year'), util.OperatorToken('*', 2) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, interpreter.CheckValidSumAverageArgument, stack) stack = [ util.HomomorphicIntToken('Invoice_Price'), util.FieldToken('Year'), util.OperatorToken('/', 2) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, interpreter.CheckValidSumAverageArgument, stack)
def RewriteField(field): """Rewrite fields for real query with server.""" if not isinstance(field, util.FieldToken): return field row = util.GetEntryFromSchema(field, schema) if not row: return field if row['encrypt'].startswith('probabilistic'): return util.ProbabilisticToken(str(field)) elif row['encrypt'] == 'pseudonym': return util.PseudonymToken(str(field)) elif row['encrypt'] == 'homomorphic' and row['type'] == 'integer': return util.HomomorphicIntToken(str(field)) elif row['encrypt'] == 'homomorphic' and row['type'] == 'float': return util.HomomorphicFloatToken(str(field)) elif row['encrypt'] == 'searchwords': return util.SearchwordsToken(str(field)) return field
def testRewriteAggregations(self): stack = [ util.CountStarToken(), util.AggregationFunctionToken('COUNT', 1) ] rewritten_stack = ['COUNT(*)'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.ProbabilisticToken('Price'), util.AggregationFunctionToken('COUNT', 1) ] rewritten_stack = ['COUNT(' + util.PROBABILISTIC_PREFIX + 'Price)'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.ProbabilisticToken('Price'), 4, util.AggregationFunctionToken('COUNT', 2) ] rewritten_stack = ['COUNT(' + util.PROBABILISTIC_PREFIX + 'Price, 4)'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.FieldToken('Year'), 5, util.AggregationFunctionToken('DISTINCTCOUNT', 2), util.FieldToken('Year'), util.AggregationFunctionToken('COUNT', 1), util.OperatorToken('+', 2) ] rewritten_stack = ['COUNT(DISTINCT Year, 5)', 'COUNT(Year)', '+'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ 0, util.BuiltInFunctionToken('cos'), util.AggregationFunctionToken('COUNT', 1) ] rewritten_stack = ['COUNT(1.0)'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.StringLiteralToken('"Hello"'), 2, util.BuiltInFunctionToken('left'), util.StringLiteralToken('"y"'), util.BuiltInFunctionToken('concat'), util.AggregationFunctionToken('GROUP_CONCAT', 1) ] rewritten_stack = ['GROUP_CONCAT("Hey")'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.FieldToken('Year'), util.FieldToken('Year'), util.OperatorToken('*', 2), util.AggregationFunctionToken('SUM', 1) ] rewritten_stack = ['SUM((Year * Year))'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.HomomorphicIntToken('Invoice_Price'), util.AggregationFunctionToken('SUM', 1) ] rewritten_stack = [ 0.0, 'COUNT(' + util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price)', '*', 1.0, 'TO_BASE64(BYTES(PAILLIER_SUM(FROM_BASE64(' + util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price), \'0\')))', '*', '+' ] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.HomomorphicFloatToken('Holdback_Percentage'), util.AggregationFunctionToken('AVG', 1) ] rewritten_stack = [ 0.0, 'COUNT(' + util.HOMOMORPHIC_FLOAT_PREFIX + 'Holdback_Percentage)', '*', 1.0, 'TO_BASE64(BYTES(PAILLIER_SUM(FROM_BASE64(' + util.HOMOMORPHIC_FLOAT_PREFIX + 'Holdback_Percentage), \'0\')))', '*', '+', 'COUNT(' + util.HOMOMORPHIC_FLOAT_PREFIX + 'Holdback_Percentage)', '/' ] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.HomomorphicIntToken('Invoice_Price'), 2, util.OperatorToken('+', 2), 5, util.OperatorToken('*', 2), util.AggregationFunctionToken('SUM', 1) ] rewritten_stack = [ 0.0, 'COUNT(' + util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price)', '*', 5.0, 'TO_BASE64(BYTES(PAILLIER_SUM(FROM_BASE64(' + util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price), \'0\')))', '*', '+', 0.0, 'COUNT(' + util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price)', '*', 1.0, 'SUM((2 * 5))', '*', '+', '+' ] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.PseudonymToken('Make'), 2, util.AggregationFunctionToken('DISTINCTCOUNT', 2) ] rewritten_stack = [ 'COUNT(DISTINCT ' + util.PSEUDONYM_PREFIX + 'Make, 2)' ] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.FieldToken('Year'), util.AggregationFunctionToken('TOP', 1) ] rewritten_stack = ['TOP(Year)'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.PseudonymToken('Make'), 5, 1, util.AggregationFunctionToken('TOP', 3) ] rewritten_stack = ['TOP(' + util.PSEUDONYM_PREFIX + 'Make, 5, 1)'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.FieldToken('Year'), util.BuiltInFunctionToken('cos'), util.HomomorphicIntToken('Invoice_Price'), util.OperatorToken('+', 2), util.AggregationFunctionToken('SUM', 1) ] rewritten_stack = [ 0.0, 'COUNT(' + util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price)', '*', 1.0, 'SUM(cos(Year))', '*', '+', 0.0, 'COUNT(' + util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price)', '*', 1.0, 'TO_BASE64(BYTES(PAILLIER_SUM(FROM_BASE64(' + util.HOMOMORPHIC_INT_PREFIX + 'Invoice_Price),' ' \'0\')))', '*', '+', '+' ] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.ProbabilisticToken('Model'), util.AggregationFunctionToken('DISTINCTCOUNT', 1) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, query_lib._RewriteAggregations, [stack], _TEST_NSQUARE) stack = [ util.ProbabilisticToken('Price'), util.AggregationFunctionToken('SUM', 1) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, query_lib._RewriteAggregations, [stack], _TEST_NSQUARE) stack = [ util.HomomorphicIntToken('Invoice_Price'), util.HomomorphicFloatToken('Holdback_Percentage'), util.OperatorToken('*', 2), util.AggregationFunctionToken('SUM', 1) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, query_lib._RewriteAggregations, [stack], _TEST_NSQUARE) stack = [ util.HomomorphicFloatToken('Holdback_Percentage'), util.BuiltInFunctionToken('cos'), util.AggregationFunctionToken('SUM', 1) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, query_lib._RewriteAggregations, [stack], _TEST_NSQUARE) stack = [ util.HomomorphicIntToken('Invoice_Price'), util.AggregationFunctionToken('TOP', 1) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, query_lib._RewriteAggregations, [stack], _TEST_NSQUARE) stack = [ util.FieldToken('Year'), util.AggregationFunctionToken('SUM', 1), util.AggregationFunctionToken('SUM', 1) ] rewritten_stack = ['SUM(SUM(Year))'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.HomomorphicIntToken('Invoice_Price'), util.AggregationFunctionToken('SUM', 1), util.AggregationFunctionToken('SUM', 1) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, query_lib._RewriteAggregations, [stack], _TEST_NSQUARE) stack = [ util.FieldToken('Year'), util.AggregationFunctionToken('GROUP_CONCAT', 1), util.AggregationFunctionToken('GROUP_CONCAT', 1) ] rewritten_stack = ['GROUP_CONCAT(GROUP_CONCAT(Year))'] self.assertEqual( query_lib._RewriteAggregations([stack], _TEST_NSQUARE), [rewritten_stack]) stack = [ util.PseudonymToken('Make'), util.AggregationFunctionToken('GROUP_CONCAT', 1), util.AggregationFunctionToken('GROUP_CONCAT', 1) ] self.assertRaises(bigquery_client.BigqueryInvalidQueryError, query_lib._RewriteAggregations, [stack], _TEST_NSQUARE)