From 804fa563f526de6dee8cbd70efebe0f721b7b535 Mon Sep 17 00:00:00 2001
From: xPythonic <gr33nmyth@aol.com>
Date: Fri, 26 Sep 2014 00:50:41 -0400
Subject: [PATCH] AstronDatabaseInterface: Add support for querying specific
 fields of an object

---
 .../distributed/AstronDatabaseInterface.py    | 53 +++++++++++++++----
 .../distributed/AstronInternalRepository.py   |  2 +
 2 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/direct/src/distributed/AstronDatabaseInterface.py b/direct/src/distributed/AstronDatabaseInterface.py
index 4b425d05e6..01d670be2f 100644
--- a/direct/src/distributed/AstronDatabaseInterface.py
+++ b/direct/src/distributed/AstronDatabaseInterface.py
@@ -20,6 +20,7 @@ class AstronDatabaseInterface:
         self.air = air
 
         self._callbacks = {}
+        self._dclasses = {}
 
     def createObject(self, databaseId, dclass, fields={}, callback=None):
         """
@@ -73,7 +74,7 @@ class AstronDatabaseInterface:
 
         del self._callbacks[ctx]
 
-    def queryObject(self, databaseId, doId, callback):
+    def queryObject(self, databaseId, doId, callback, dclass=None, fieldNames=()):
         """
         Query object `doId` out of the database.
 
@@ -85,21 +86,44 @@ class AstronDatabaseInterface:
         # Save the callback:
         ctx = self.air.getContext()
         self._callbacks[ctx] = callback
+        self._dclasses[ctx] = dclass
 
         # Generate and send the datagram:
         dg = PyDatagram()
-        dg.addServerHeader(databaseId, self.air.ourChannel, DBSERVER_OBJECT_GET_ALL)
+
+        if not fieldNames:
+            dg.addServerHeader(databaseId, self.air.ourChannel,
+                               DBSERVER_OBJECT_GET_ALL)
+        else:
+            # We need a dclass in order to convert the field names into field IDs:
+            assert dclass is not None
+
+            if len(fieldNames) > 1:
+                dg.addServerHeader(databaseId, self.air.ourChannel,
+                                   DBSERVER_OBJECT_GET_FIELDS)
+            else:
+                dg.addServerHeader(databaseId, self.air.ourChannel,
+                                   DBSERVER_OBJECT_GET_FIELD)
+
         dg.addUint32(ctx)
         dg.addUint32(doId)
+        if len(fieldNames) > 1:
+            dg.addUint16(len(fieldNames))
+        for fieldName in fieldNames:
+            field = dclass.getFieldByName(fieldName)
+            if field is None:
+                self.notify.error('Bad field named %s in query for'
+                                  ' %s object' % (fieldName, dclass.getName()))
+            dg.addUint16(field.getNumber())
         self.air.send(dg)
 
-    def handleQueryObjectResp(self, di):
+    def handleQueryObjectResp(self, msgType, di):
         ctx = di.getUint32()
         success = di.getUint8()
 
         if ctx not in self._callbacks:
-            self.notify.warning('Received unexpected DBSERVER_OBJECT_GET_ALL_RESP'
-                                ' (ctx %d)' % (ctx))
+            self.notify.warning('Received unexpected %s'
+                                ' (ctx %d)' % (MsgId2Names[msgType], ctx))
             return
 
         try:
@@ -108,14 +132,20 @@ class AstronDatabaseInterface:
                     self._callbacks[ctx](None, None)
                 return
 
-            dclassId = di.getUint16()
-            dclass = self.air.dclassesByNumber.get(dclassId)
+            if msgType == DBSERVER_OBJECT_GET_ALL_RESP:
+                dclassId = di.getUint16()
+                dclass = self.air.dclassesByNumber.get(dclassId)
+            else:
+                dclass = self._dclasses[ctx]
 
             if not dclass:
                 self.notify.error('Received bad dclass %d in'
                                   ' DBSERVER_OBJECT_GET_ALL_RESP' % (dclassId))
 
-            fieldCount = di.getUint16()
+            if msgType == DBSERVER_OBJECT_GET_FIELD:
+                fieldCount = 1
+            else:
+                fieldCount = di.getUint16()
             unpacker = DCPacker()
             unpacker.setUnpackData(di.getRemainingBytes())
             fields = {}
@@ -127,7 +157,6 @@ class AstronDatabaseInterface:
                     self.notify.error('Received bad field %d in query for'
                                       ' %s object' % (fieldId, dclass.getName()))
 
-
                 unpacker.beginUnpack(field)
                 fields[field.getName()] = field.unpackArgs(unpacker)
                 unpacker.endUnpack()
@@ -267,8 +296,10 @@ class AstronDatabaseInterface:
     def handleDatagram(self, msgType, di):
         if msgType == DBSERVER_CREATE_OBJECT_RESP:
             self.handleCreateObjectResp(di)
-        elif msgType == DBSERVER_OBJECT_GET_ALL_RESP:
-            self.handleQueryObjectResp(di)
+        elif msgType in (DBSERVER_OBJECT_GET_ALL_RESP,
+                         DBSERVER_OBJECT_GET_FIELDS_RESP,
+                         DBSERVER_OBJECT_GET_FIELD_RESP):
+            self.handleQueryObjectResp(msgType, di)
         elif msgType == DBSERVER_OBJECT_SET_FIELD_IF_EQUALS_RESP:
             self.handleUpdateObjectResp(di, False)
         elif msgType == DBSERVER_OBJECT_SET_FIELDS_IF_EQUALS_RESP:
diff --git a/direct/src/distributed/AstronInternalRepository.py b/direct/src/distributed/AstronInternalRepository.py
index f56168084e..69e4da3312 100644
--- a/direct/src/distributed/AstronInternalRepository.py
+++ b/direct/src/distributed/AstronInternalRepository.py
@@ -233,6 +233,8 @@ class AstronInternalRepository(ConnectionRepository):
             self.handleObjLocation(di)
         elif msgType in (DBSERVER_CREATE_OBJECT_RESP,
                          DBSERVER_OBJECT_GET_ALL_RESP,
+                         DBSERVER_OBJECT_GET_FIELDS_RESP,
+                         DBSERVER_OBJECT_GET_FIELD_RESP,
                          DBSERVER_OBJECT_SET_FIELD_IF_EQUALS_RESP,
                          DBSERVER_OBJECT_SET_FIELDS_IF_EQUALS_RESP):
             self.dbInterface.handleDatagram(msgType, di)