# Copyright 2018 Braxton Mckee

#

# Licensed 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.

from typed_python import Alternative, TupleOf, OneOf, ConstDict

from typed_python.SerializationContext import SerializationContext

from object_database.schema import Indexed, Index, Schema

from object_database.core_schema import core_schema

from object_database.view import RevisionConflictException, DisconnectedException, ObjectDoesntExistException

from object_database.database_connection import TransactionListener, DatabaseConnection, SetWithEdits

from object_database.tcp_server import TcpServer

from object_database.inmem_server import InMemServer

from object_database.persistence import InMemoryPersistence, RedisPersistence

from object_database.util import configureLogging, genToken

from object_database.test_util import currentMemUsageMb

import object_database.messages as messages

import queue

import unittest

import tempfile

import numpy

import redis

import subprocess

import os

import threading

import random

import time

import ssl

class BlockingCallback:

def __init__(self):

self.callbackArgs = queue.Queue()

self.is_released = queue.Queue()

def callback(self, arg=None):

self.callbackArgs.put(arg)

self.is_released.get(timeout=1.0)

def waitForCallback(self, timeout):

return self.callbackArgs.get(timeout=timeout)

def releaseCallback(self):

self.is_released.put(True)

expr = Alternative(

"Expr",

Constant={'value': int},

# Add = {'l': expr, 'r': expr},

# Sub = {'l': expr, 'r': expr},

# Mul = {'l': expr, 'r': expr}

)

schema = Schema("test_schema")

schema.expr = expr

@schema.define

class Root:

obj = OneOf(None, schema.Object)

k = int

@schema.define

class Object:

k = Indexed(expr)

other = OneOf(None, schema.Object)

@property

def otherK(self):

if self.other is not None:

return self.other.k

@schema.define

class ThingWithDicts:

x = ConstDict(str, bytes)

@schema.define

class Counter:

k = Indexed(int)

x = int

def f(self):

return self.k + 1

def __str__(self):

return "Counter(k=%s)" % self.k

@schema.define

class StringIndexed:

name = Indexed(str)

class ObjectDatabaseTests:

@classmethod

def setUpClass(cls):

configureLogging("database_test")

cls.PERFORMANCE_FACTOR = 1.0 if os.environ.get('TRAVIS_CI', None) is None else 2.0

def test_assigning_dicts(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

z = ThingWithDicts()

z.x = {'a': b'b'}

with db.transaction():

z2 = ThingWithDicts()

z2.x = z.x

def test_subscribe_excluding(self):

db = self.createNewDb()

db.subscribeToSchema(schema, excluding=[ThingWithDicts])

with db.view():

with self.assertRaises(Exception):

ThingWithDicts.lookupAll()

Counter.lookupAll()

db.subscribeToSchema(schema)

with db.view():

ThingWithDicts.lookupAll()

def test_subscribe_to_objects(self):

db1 = self.createNewDb()

db1.subscribeToSchema(schema)

db2 = self.createNewDb()

with db1.transaction():

someThings = [Counter(k=i) for i in range(10)]

db2.subscribeToObjects(someThings[::2])

with db2.view():

for i in range(10):

if i % 2 == 0:

self.assertTrue(someThings[i].exists())

else:

self.assertFalse(someThings[i].exists())

def test_serialization_contexts(self):

db = self.createNewDb()

class ArbitraryBaseClass:

def __init__(self, x):

self.x = x

class ArbitrarySubclass(ArbitraryBaseClass):

def __init__(self, x, y):

super().__init__(x)

self.y = y

db.setSerializationContext(SerializationContext({'ABC': ArbitraryBaseClass, 'SUB': ArbitrarySubclass}))

schema = Schema("test_schema")

@schema.define

class HoldsArbitrary:

holding = ArbitraryBaseClass

@schema.define

class HoldsObject:

holding = object

db.subscribeToSchema(schema)

with db.transaction():

x = HoldsArbitrary(holding=ArbitraryBaseClass(10))

self.assertEqual(x.holding.x, 10)

with db.transaction():

self.assertEqual(x.holding.x, 10)

self.assertIsInstance(x.holding, ArbitraryBaseClass)

x.holding = ArbitrarySubclass(10, 20)

with db.transaction():

self.assertEqual(x.holding.x, 10)

self.assertEqual(x.holding.y, 20)

self.assertIsInstance(x.holding, ArbitrarySubclass)

with self.assertRaises(Exception):

with db.transaction():

x.holding = "hi"

with db.transaction():

x = HoldsObject(holding="hi")

x.holding = 10

with db.transaction():

self.assertEqual(x.holding, 10)

def test_disconnecting(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

db.disconnect()

with self.assertRaises(DisconnectedException):

with db.view():

pass

usage = currentMemUsageMb(residentOnly=False)

for i in range(100):

db = self.createNewDb()

db.subscribeToSchema(schema)

db.flush()

db.disconnect()

usage = currentMemUsageMb(residentOnly=False)

for i in range(500):

db = self.createNewDb()

db.subscribeToSchema(schema)

db.flush()

db.disconnect()

assert currentMemUsageMb(residentOnly=False) < usage + 100

def test_disconnecting_is_immediate(self):

db1 = self.createNewDb()

db2 = self.createNewDb()

db1.subscribeToSchema(core_schema)

db2.subscribeToSchema(core_schema)

with db2.view():

assert db1.connectionObject.exists()

with db1.view():

assert db2.connectionObject.exists()

db2Connection = db2.connectionObject

db2.disconnect()

self.assertTrue(

db1.waitForCondition(

lambda: not db2Connection.exists(),

timeout=2.0*self.PERFORMANCE_FACTOR)

)

def test_lazy_subscriptions(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

loadedIDs = queue.Queue()

self.server._lazyLoadCallback = loadedIDs.put

with db.transaction():

c = Counter(k=2, x=3)

db2 = self.createNewDb()

db2.subscribeToSchema(schema, lazySubscription=True)

with db2.view():

# lookup in the index doesn't dirty the object because we have to load

# the index values when we first subscribe

self.assertEqual(Counter.lookupAll(k=2), (c,))

with db2.view():

self.assertEqual(c.k, 2)

self.assertEqual(loadedIDs.get_nowait(), c._identity)

# at this point, the value is loaded

with db2.view():

self.assertEqual(c.x, 3)

with self.assertRaises(queue.Empty):

loadedIDs.get_nowait()

def test_methods(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

counter = Counter()

counter.k = 2

self.assertEqual(counter.f(), 3)

self.assertEqual(str(counter), "Counter(k=2)")

def test_property_object(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

counter = Object(k=expr.Constant(value=10))

counter2 = Object(other=counter, k=expr.Constant(value=0))

self.assertEqual(counter2.otherK, counter.k)

def test_identity_transfer(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

root = Root()

root2 = Root.fromIdentity(root._identity)

root.obj = Object(k=expr.Constant(value=23))

self.assertEqual(root2.obj.k.value, 23)

def test_adding_fields_to_type(self):

schema = Schema("schema")

@schema.define

class Test:

i = int

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

t = Test(i=1)

schema2 = Schema("schema")

@schema2.define

class Test:

i = int

k = int

db2 = self.createNewDb()

db2.subscribeToSchema(schema2)

t = Test.fromIdentity(t._identity)

with db2.view():

self.assertEqual(t.i, 1)

self.assertEqual(t.k, 0)

def test_subclassing(self):

schema = Schema("schema")

@schema.define

class Test:

i = int

def f(self):

return 1

def g(self):

return 2

@schema.define

class SubclassTesting(Test):

y = int

def g(self):

return 3

def h(self):

return 4

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

t = Test(i=1)

t2 = SubclassTesting(i=2, y=3)

self.assertEqual(t.f(), 1)

self.assertEqual(t.g(), 2)

self.assertEqual(t2.f(), 1)

self.assertEqual(t2.g(), 3)

self.assertEqual(t2.h(), 4)

self.assertEqual(t2.y, 3)

def test_many_subscriptions(self):

OK = []

FINISHED = []

count = 10

threadCount = 10

def worker(index):

db = self.createNewDb()

indices = list(range(count))

numpy.random.shuffle(indices)

for i in indices:

db.subscribeToIndex(Counter, k=i)

with db.transaction():

Counter(k=i, x=index)

FINISHED.append(True)

db.waitForCondition(

lambda: len(FINISHED) == threadCount,

10.0*self.PERFORMANCE_FACTOR

)

db.flush()

with db.view():

actuallyVisible = len(Counter.lookupAll())

if actuallyVisible != count * threadCount:

print("TOTAL is ", actuallyVisible, " != ", count*threadCount)

else:

OK.append(True)

threads = [threading.Thread(target=worker, args=(i,)) for i in range(threadCount)]

for t in threads:

t.daemon = True

t.start()

for t in threads:

t.join()

db1 = self.createNewDb()

db1.subscribeToSchema(schema)

with db1.view():

self.assertEqual(len(Counter.lookupAll()), count*threadCount)

db2 = self.createNewDb()

for i in range(count):

db2.subscribeToIndex(Counter, k=i)

db2.flush()

with db2.view():

self.assertEqual(len(Counter.lookupAll()), count*threadCount)

self.assertEqual(len(OK), 10)

def test_transaction_handlers(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

didOne = threading.Event()

def handler(changed):

didOne.set()

with TransactionListener(db, handler):

with db.transaction():

Root()

didOne.wait()

assert didOne.isSet()

def test_basic(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

root = Root()

self.assertTrue(root.obj is None, root.obj)

root.obj = Object(k=expr.Constant(value=23))

db2 = self.createNewDb()

db2.subscribeToSchema(schema)

with db2.view():

self.assertEqual(root.obj.k.value, 23)

def test_throughput(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

root = Root()

root.obj = Object(k=expr.Constant(value=0))

t0 = time.time()

while time.time() < t0 + 1.0:

with db.transaction():

root.obj.k = expr.Constant(value=root.obj.k.value + 1)

with db.view():

self.assertTrue(root.obj.k.value > 500, root.obj.k.value)

print(root.obj.k.value, "transactions per second")

def test_delayed_transactions(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

confirmed = queue.Queue()

with db.transaction():

root = Root()

root.obj = Object(k=expr.Constant(value=0))

for i in range(1000):

with db.transaction().onConfirmed(confirmed.put):

root.obj.k = expr.Constant(value=root.obj.k.value + 1)

self.assertTrue(confirmed.qsize() < 1000)

good = 0

for i in range(1000):

if confirmed.get().matches.Success:

good += 1

self.assertGreater(good, 0)

self.assertLess(good, 1000)

def test_exists(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

root = Root()

self.assertTrue(root.exists())

self.assertEqual(root.k, 0)

root.delete()

self.assertFalse(root.exists())

with self.assertRaises(ObjectDoesntExistException):

root.k

with db.view():

self.assertFalse(root.exists())

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.view():

self.assertFalse(root.exists())

def test_read_performance(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

objects = {}

with db.transaction():

for i in range(100):

root = Root()

e = expr.Constant(value=i)

root.obj = Object(k=e)

objects[i] = root

db = self.createNewDb()

db.subscribeToSchema(schema)

t0 = time.time()

count = 0

steps = 0

while time.time() < t0 + 1.0:

with db.transaction():

for i in range(100):

count += objects[i].obj.k.value

steps += 1

def test_transactions(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

root = Root()

views = [db.view()]

for i in [1, 2, 3]:

with db.transaction():

root.obj = Object(k=expr.Constant(value=i))

views.append(db.view())

vals = []

for v in views:

with v:

if root.obj is None:

vals.append(None)

else:

vals.append(root.obj.k.value)

self.assertEqual(vals, [None, 1, 2, 3])

def test_conflicts(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

root = Root()

root.obj = Object(k=expr.Constant(value=0))

for ordering in [0, 1]:

t1 = db.transaction()

t2 = db.transaction()

if ordering:

t1, t2 = t2, t1

with t1:

root.obj.k = expr.Constant(value=root.obj.k.value + 1)

with self.assertRaises(RevisionConflictException):

with t2:

root.obj.k = expr.Constant(value=root.obj.k.value + 1)

def test_conflicts_dont_cause_view_leaks(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

root = Root()

root.obj = Object(k=expr.Constant(value=0))

t1 = db.transaction()

t2 = db.transaction()

with t1:

root.obj.k = expr.Constant(value=root.obj.k.value + 1)

try:

with t2:

root.obj.k = expr.Constant(value=root.obj.k.value + 1)

except RevisionConflictException:

pass

for i in range(100):

with db.transaction():

root.obj.k = expr.Constant(value=root.obj.k.value + 1)

self.assertTrue(db._noViewsOutstanding())

def test_object_versions_robust(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

counters = []

counter_vals_by_tn = {}

views_by_tn = {}

random.seed(123)

# expect nothing initially

views_by_tn[db._cur_transaction_num] = db.view()

counter_vals_by_tn[db._cur_transaction_num] = {}

# seed the initial state

with db.transaction():

for i in range(20):

counter = Counter(_identity="C_%s" % i)

counter.k = int(random.random() * 100)

counters.append(counter)

counter_vals_by_tn[db._cur_transaction_num + 1] = {c: c.k for c in counters}

total_writes = 0

for passIx in range(1000):

with db.transaction():

didOne = False

for subix in range(int(random.random() * 5 + 1)):

counter = counters[int(random.random() * len(counters))]

if counter.exists():

if random.random() < .001:

counter.delete()

else:

counter.k = int(random.random() * 100)

total_writes += 1

didOne = True

if didOne:

counter_vals_by_tn[db._cur_transaction_num + 1] = {c: c.k for c in counters if c.exists()}

if didOne:

views_by_tn[db._cur_transaction_num] = db.view()

while views_by_tn and random.random() < .5 or len(views_by_tn) > 10:

# pick a random view and check that it's consistent

all_tids = list(views_by_tn)

tid = all_tids[int(random.random() * len(all_tids))]

with views_by_tn[tid]:

for c in counters:

if not c.exists():

assert c not in counter_vals_by_tn[tid], tid

else:

self.assertEqual(c.k, counter_vals_by_tn[tid][c])

del views_by_tn[tid]

if random.random() < .05 and views_by_tn:

with db.view():

curCounterVals = {c: c.k for c in counters if c.exists()}

# reset the database

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.view():

newCounterVals = {c: c.k for c in counters if c.exists()}

self.assertEqual(curCounterVals, newCounterVals)

views_by_tn = {}

counter_vals_by_tn = {}

# we may have one or two for connection objects, and we have two values for every indexed thing

self.assertLess(self.mem_store.storedStringCount(), 203)

self.assertTrue(total_writes > 500)

def test_flush_db_works(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

counters = []

with db.transaction():

for _ in range(10):

counters.append(Counter(k=1))

self.assertTrue(self.mem_store.values)

view = db.view()

with db.transaction():

for c in counters:

c.delete()

# database doesn't have this

t0 = time.time()

while time.time() - t0 < 1.0 and self.mem_store.storedStringCount() >= 2:

time.sleep(.01)

self.assertLess(self.mem_store.storedStringCount(), 4)

# but the view does!

with view:

for c in counters:

self.assertTrue(c.exists())

def test_read_write_conflict(self):

db = self.createNewDb()

schema = Schema("test_schema")

@schema.define

class Counter:

k = int

db.subscribeToSchema(schema)

with db.transaction():

o1 = Counter()

o2 = Counter()

for consistency in [True, False]:

if consistency:

t1 = db.transaction().consistency(reads=True)

t2 = db.transaction().consistency(reads=True)

else:

t1 = db.transaction().consistency(none=True)

t2 = db.transaction().consistency(none=True)

with t1.nocommit():

o1.k = o2.k + 1

with t2.nocommit():

o2.k = o1.k + 1

t1.commit()

if consistency:

with self.assertRaises(RevisionConflictException):

t2.commit()

else:

t2.commit()

def test_indices(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.view():

self.assertEqual(Counter.lookupAll(k=20), ())

self.assertEqual(Counter.lookupAll(k=30), ())

with db.transaction():

o1 = Counter(k=20)

with db.view():

self.assertEqual(Counter.lookupAll(k=20), (o1,))

self.assertEqual(Counter.lookupAll(k=30), ())

with db.transaction():

o1.k = 30

with db.view():

self.assertEqual(Counter.lookupAll(k=20), ())

self.assertEqual(Counter.lookupAll(k=30), (o1,))

with db.transaction():

o1.delete()

with db.view():

self.assertEqual(Counter.lookupAll(k=20), ())

self.assertEqual(Counter.lookupAll(k=30), ())

def test_indices_multiple_values(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

k1 = Counter(k=20)

Counter(k=20)

self.assertEqual(len(Counter.lookupAll(k=20)), 2)

k1.k = 30

self.assertEqual(len(Counter.lookupAll(k=20)), 1)

k1.k = 20

self.assertEqual(len(Counter.lookupAll(k=20)), 2)

with db.transaction():

self.assertEqual(len(Counter.lookupAll(k=20)), 2)

k1.k = 30

self.assertEqual(len(Counter.lookupAll(k=20)), 1)

k1.k = 20

self.assertEqual(len(Counter.lookupAll(k=20)), 2)

def test_indices_across_invocations(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

o = Counter(k=1)

o.x = 10

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

o = Counter.lookupOne(k=1)

self.assertEqual(o.x, 10)

o.k = 2

o.x = 11

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

o = Counter.lookupOne(k=2)

o.k = 3

self.assertEqual(o.x, 11)

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

self.assertFalse(Counter.lookupAny(k=2))

o = Counter.lookupOne(k=3)

o.k = 3

self.assertEqual(o.x, 11)

def test_index_consistency(self):

db = self.createNewDb()

schema = Schema("test_schema")

@schema.define

class Object:

x = int

y = int

pair = Index('x', 'y')

db.subscribeToSchema(schema)

with db.transaction():

o = Object(x=0, y=0)

t1 = db.transaction()

t2 = db.transaction()

with t1.nocommit():

o.x = 1

with t2.nocommit():

o.y = 1

t1.commit()

with self.assertRaises(RevisionConflictException):

t2.commit()

with self.assertRaises(Exception):

with db.transaction().consistency(writes=True):

o.y = 2

def test_indices_of_algebraics(self):

db = self.createNewDb()

db.subscribeToSchema(schema)

with db.transaction():

o1 = Object(k=expr.Constant(value=123))

with db.view():

self.assertEqual(Object.lookupAll(k=expr.Constant(value=123)), (o1,))

def test_frozen_schema(self):

schema = Schema("test_schema")

@schema.define

class Object:

x = int

y = int

Object.fromIdentity("hi")

with self.assertRaises(AttributeError):

schema.SomeOtherObject

def test_freezing_schema_with_undefined_fails(self):

schema = Schema("test_schema")

@schema.define

class Object:

x = schema.Object2

y = int

with self.assertRaises(Exception):

schema.freeze()

@schema.define

class Object2:

x = int

schema.freeze()

def test_index_functions(self):

db = self.createNewDb()

schema = Schema("test_schema")

@schema.define

class Object:

k = Indexed(int)

pair_index = Index('k', 'k')

db.subscribeToSchema(schema)

with db.transaction():

o1 = Object(k=10)

with db.view():

self.assertEqual(Object.lookupAll(k=10), (o1,))

self.assertEqual(Object.lookupAll(k=20), ())

self.assertEqual(Object.lookupAll(pair_index=(10, 10)), (o1,))

self.assertEqual(Object.lookupAll(pair_index=(10, 11)), ())

with self.assertRaises(Exception):

self.assertEqual(Object.lookupAll(pair_index=(10, "hi")), (o1,))

def test_indices_update_during_transactions(self):

db = self.createNewDb()

schema = Schema("test_schema")

@schema.define

class Object:

k = Indexed(int)

db.subscribeToSchema(schema)

with db.transaction():

self.assertEqual(Object.lookupAll(k=10), ())

o1 = Object(k=10)

self.assertEqual(Object.lookupAll(k=10), (o1,))

o1.k = 20

self.assertEqual(Object.lookupAll(k=10), ())

self.assertEqual(Object.lookupAll(k=20), (o1,))