Small fixes.

py/vtdb/topology.py

@@ -180,3 +180,9 @@ def get_host_port_by_name(topo_client, db_key, encrypted=False):
     return encrypted_host_port_list
   random.shuffle(host_port_list)
   return host_port_list
+
+
+def is_sharded_keyspace(keyspace_name, db_type):
+  ks = get_keyspace(keyspace_name)
+  shard_count = ks.get_shard_count(db_type)
+  return shard_count > 1

py/vtdb/vtgate_cursor.py

@@ -1,17 +1,23 @@
 # Copyright 2012, Google Inc. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can
 # be found in the LICENSE file.
+
 import itertools
 
 from vtdb import cursor
 from vtdb import dbexceptions
+from vtdb import keyrange_constants
+from vtdb import topology
 
 
 write_sql_pattern = re.compile('\s*(insert|update|delete)', re.IGNORECASE)
+
+
 class __EmptyBindVariables(frozenset):
   pass
-EmptyBindVariables = __EmptyBindVariables()
 
+
+EmptyBindVariables = __EmptyBindVariables()
 MAX_QUERY_SIZE = 32 * 1024
 EXCESSIVE_BIND_VAR_COUNT = 1000
 MAX_BIND_VAR_COUNT = 4000
@@ -73,12 +79,17 @@ class VTGateCursor(object):
     if write_query:
       if not self.is_writable():
         raise dbexceptions.DatabaseError('DML on a non-writable cursor', sql)
-      if keyspace_ids is None or len(keyspace_ids) != 1:
-        raise dbexceptions.ProgrammingError('DML on zero or multiple keyspace ids is not allowed')
+
+      # FIXME(shrutip): these checks maybe better on vtgate server.
+      if topology.is_sharded_keyspace(self.keyspace, self.tablet_type):
+        if keyspace_ids is None or len(keyspace_ids) != 1:
+          raise dbexceptions.ProgrammingError('DML on zero or multiple keyspace ids is not allowed')
       else:
-        # FIXME(shrutip): this could potentially be moved to vtgate server.
-        sql += self._binlog_hint(keyspace_ids[0])
-    
+        if keyrange is None or keyrange != keyrange_constants.NON_PARTIAL_KEYRANGE:
+          raise dbexceptions.ProgrammingError('Keyrange not correct for non-sharded keyspace')
+
+      # FIXME(shrutip): this could potentially be done on vtgate server.
+      sql += self._binlog_hint(keyspace_ids[0])
 
     self.results, self.rowcount, self.lastrowid, self.description = self.connection._execute(sql,
                                                                                              bind_variables,

py/vtdb/vtgatev2.py

@@ -12,10 +12,6 @@ from vtdb import cursor
 from vtdb import dbexceptions
 from vtdb import field_types
 
-# This is the shard name for when the keyrange covers the entire space
-# for unsharded database.
-SHARD_ZERO = "0"
-
 
 _errno_pattern = re.compile('\(errno (\d+)\)')
 # Map specific errors to specific classes.
@@ -23,8 +19,10 @@ _errno_map = {
     1062: dbexceptions.IntegrityError,
 }
 
+
 __error_counter_callback = None
 
+
 def register_error_counter_callback(func):
   global __error_counter_callback
   __error_counter_callback = func
@@ -164,7 +162,7 @@ class VtgateConnection(object):
     else:
       raise dbexceptions.ProgrammingError('_execute called without specifying keyspace_ids or keyranges')
 
- 
+
     self._add_session(req)
 
     fields = []