Implement Connection.predefine_table_schema (#422)

pynamodb/connection/base.py

@@ -345,6 +345,51 @@ def dispatch(self, operation_name, operation_kwargs):
             log.debug("%s %s consumed %s units",  data.get(TABLE_NAME, ''), operation_name, capacity)
         return data
 
+    def predefine_table_schema(self,
+                               table_name,
+                               attribute_definitions,
+                               key_schema,
+                               global_secondary_indexes,
+                               local_secondary_indexes):
+        """
+        Saves attribute definitions and a key schema for the given table
+        to avoid additional DescribeTable requests for cases when key schema is known,
+        for instance from Madel._get_schema
+        :param table_name:
+        :param attribute_definitions: dictionary with pythonic keys
+        :param key_schema: dictionary with pythonic keys
+        :param global_secondary_indexes: list of dictionaries with pythonic keys
+        :param local_secondary_indexes: list of dictionaries with pythonic keys
+        :return: None
+        """
+        self._tables[table_name] = MetaTable({
+            ATTR_DEFINITIONS: [
+                {ATTR_NAME: attr.get(pythonic(ATTR_NAME)), ATTR_TYPE: attr.get(pythonic(ATTR_TYPE))}
+                for attr in attribute_definitions
+            ],
+            KEY_SCHEMA: [
+                {ATTR_NAME: attr.get(pythonic(ATTR_NAME)), KEY_TYPE: attr.get(pythonic(KEY_TYPE))}
+                for attr in key_schema
+            ],
+            GLOBAL_SECONDARY_INDEXES: [
+                {
+                    INDEX_NAME: index.get(pythonic(INDEX_NAME)),
+                    KEY_SCHEMA: sorted(index.get(pythonic(KEY_SCHEMA)), key=lambda x: x.get(KEY_TYPE)),
+                    PROJECTION: index.get(pythonic(PROJECTION)),
+                    PROVISIONED_THROUGHPUT: index.get(pythonic(PROVISIONED_THROUGHPUT))
+                }
+                for index in global_secondary_indexes
+            ],
+            LOCAL_SECONDARY_INDEXES: [
+                {
+                    INDEX_NAME: index.get(pythonic(INDEX_NAME)),
+                    KEY_SCHEMA: sorted(index.get(pythonic(KEY_SCHEMA)), key=lambda x: x.get(KEY_TYPE)),
+                    PROJECTION: index.get(pythonic(PROJECTION))
+                }
+                for index in local_secondary_indexes
+            ]
+        })
+
     def send_post_boto_callback(self, operation_name, req_uuid, table_name):
         try:
             post_dynamodb_send.send(self, operation_name=operation_name, table_name=table_name, req_uuid=req_uuid)
@@ -776,7 +821,9 @@ def get_attribute_type(self, table_name, attribute_name, value=None):
 
     def get_identifier_map(self, table_name, hash_key, range_key=None, key=KEY):
         """
-        Builds the identifier map that is common to several operations
+        Builds the identifier map that is common to several operations,
+        if the connection doesn't have predefined table schema,
+        then describe_table request will be sent
         """
         tbl = self.get_meta_table(table_name)
         if tbl is None:

pynamodb/connection/base.pyi

@@ -43,6 +43,7 @@ class Connection:
     def session(self) -> botocore.session.Session: ...
     @property
     def client(self): ...
+    def predefine_table_schema(self, table_name: Text, attribute_definitions: Dict, key_schema: Dict, global_secondary_indexes: List[Dict], local_secondary_indexes: List[Dict]): ...
     def get_meta_table(self, table_name: Text, refresh: bool = ...): ...
     def create_table(self, table_name: Text, attribute_definitions: Optional[Any] = ..., key_schema: Optional[Any] = ..., read_capacity_units: Optional[Any] = ..., write_capacity_units: Optional[Any] = ..., global_secondary_indexes: Optional[Any] = ..., local_secondary_indexes: Optional[Any] = ..., stream_specification: Optional[Any] = ...): ...
     def update_time_to_live(self, table_name: Text, ttl_attribute_name: Text): ...

pynamodb/connection/table.py

@@ -23,7 +23,8 @@ def __init__(self,
                  extra_headers=None,
                  aws_access_key_id=None,
                  aws_secret_access_key=None,
-                 aws_session_token=None):
+                 aws_session_token=None,
+                 predefined_schema=None):
         self._hash_keyname = None
         self._range_keyname = None
         self.table_name = table_name
@@ -35,6 +36,8 @@ def __init__(self,
                                      base_backoff_ms=base_backoff_ms,
                                      max_pool_connections=max_pool_connections,
                                      extra_headers=extra_headers)
+        if predefined_schema:
+            self.connection.predefine_table_schema(table_name, **predefined_schema)
 
         if aws_access_key_id and aws_secret_access_key:
             self.connection.session.set_credentials(aws_access_key_id,

pynamodb/connection/table.pyi

@@ -21,6 +21,7 @@ class TableConnection:
         aws_access_key_id: Optional[str] = ...,
         aws_secret_access_key: Optional[str] = ...,
         aws_access_token: Optional[str] = ...,
+        predefined_schema: Optional[Dict] = ...
     ) -> None: ...
 
     def get_operation_kwargs(

pynamodb/models.py

@@ -735,16 +735,6 @@ def create_table(
                 schema[pythonic(WRITE_CAPACITY_UNITS)] = write_capacity_units
             if billing_mode is not None:
                 schema[pythonic(BILLING_MODE)] = billing_mode
-            index_data = cls._get_indexes()
-            schema[pythonic(GLOBAL_SECONDARY_INDEXES)] = index_data.get(pythonic(GLOBAL_SECONDARY_INDEXES))
-            schema[pythonic(LOCAL_SECONDARY_INDEXES)] = index_data.get(pythonic(LOCAL_SECONDARY_INDEXES))
-            index_attrs = index_data.get(pythonic(ATTR_DEFINITIONS))
-            attr_keys = [attr.get(pythonic(ATTR_NAME)) for attr in schema.get(pythonic(ATTR_DEFINITIONS))]
-            for attr in index_attrs:
-                attr_name = attr.get(pythonic(ATTR_NAME))
-                if attr_name not in attr_keys:
-                    schema[pythonic(ATTR_DEFINITIONS)].append(attr)
-                    attr_keys.append(attr_name)
             cls._get_connection().create_table(
                 **schema
             )
@@ -859,6 +849,16 @@ def _get_schema(cls):
                     pythonic(KEY_TYPE): RANGE,
                     pythonic(ATTR_NAME): attr_cls.attr_name
                 })
+        index_data = cls._get_indexes()
+        schema[pythonic(GLOBAL_SECONDARY_INDEXES)] = index_data.get(pythonic(GLOBAL_SECONDARY_INDEXES))
+        schema[pythonic(LOCAL_SECONDARY_INDEXES)] = index_data.get(pythonic(LOCAL_SECONDARY_INDEXES))
+        index_attrs = index_data.get(pythonic(ATTR_DEFINITIONS))
+        attr_keys = [attr.get(pythonic(ATTR_NAME)) for attr in schema.get(pythonic(ATTR_DEFINITIONS))]
+        for attr in index_attrs:
+            attr_name = attr.get(pythonic(ATTR_NAME))
+            if attr_name not in attr_keys:
+                schema[pythonic(ATTR_DEFINITIONS)].append(attr)
+                attr_keys.append(attr_name)
         return schema
 
     @classmethod
@@ -1060,7 +1060,8 @@ def _get_connection(cls):
                                               extra_headers=cls.Meta.extra_headers,
                                               aws_access_key_id=cls.Meta.aws_access_key_id,
                                               aws_secret_access_key=cls.Meta.aws_secret_access_key,
-                                              aws_session_token=cls.Meta.aws_session_token)
+                                              aws_session_token=cls.Meta.aws_session_token,
+                                              predefined_schema=cls._get_schema())
         return cls._connection
 
     def _deserialize(self, attrs):

tests/test_base_connection.py

@@ -17,6 +17,7 @@
 from pynamodb.connection.base import MetaTable
 from pynamodb.exceptions import (
     TableError, DeleteError, PutError, ScanError, GetError, UpdateError, TableDoesNotExist)
+from pynamodb.types import HASH, RANGE
 from pynamodb.constants import (
     DEFAULT_REGION, UNPROCESSED_ITEMS, STRING_SHORT, BINARY_SHORT, DEFAULT_ENCODING, TABLE_KEY,
     PROVISIONED_BILLING_MODE, PAY_PER_REQUEST_BILLING_MODE)
@@ -1697,3 +1698,43 @@ def test_update_time_to_live_fail(self):
         with patch(PATCH_METHOD) as req:
             req.side_effect = BotoCoreError
             self.assertRaises(TableError, conn.update_time_to_live, 'test table', 'my_ttl')
+
+    def test_get_identifier_map(self):
+        schema = {'attribute_definitions': [{'attribute_name': 'UserName',
+                                             'attribute_type': 'S'},
+                                            {'attribute_name': 'UserRange',
+                                             'attribute_type': 'S'}],
+                  'key_schema': [{'attribute_name': 'UserName',
+                                  'key_type': HASH},
+                                 {'attribute_name': 'UserRange',
+                                  'key_type': RANGE}],
+                  'global_secondary_indexes': [],
+                  'local_secondary_indexes': []}
+        table_name_with_schema = 'table_with_schema'
+        table_name_without_schema = 'table_without_schema'
+        conn = Connection()
+        conn.predefine_table_schema(table_name_with_schema, **schema)
+
+        with patch(PATCH_METHOD) as req:
+            req.return_value = DESCRIBE_TABLE_DATA
+            conn.describe_table(table_name_without_schema)
+
+        with patch(PATCH_METHOD) as mocked_api_call:
+            assert ({'Key': {'ForumName': {'S': 'hash-key'}}}
+                    == conn.get_identifier_map(table_name_without_schema, 'hash-key'))
+            mocked_api_call.assert_not_called()
+
+        with patch(PATCH_METHOD) as mocked_api_call:
+            assert ({'Key': {'ForumName': {'S': 'hash-key'}, 'Subject': {'S': 'range-key'}}}
+                    == conn.get_identifier_map(table_name_without_schema, 'hash-key', 'range-key'))
+            mocked_api_call.assert_not_called()
+
+        with patch(PATCH_METHOD) as mocked_api_call:
+            assert ({'Key': {'UserName': {'S': 'hash-key'}}}
+                    == conn.get_identifier_map(table_name_with_schema, 'hash-key'))
+            mocked_api_call.assert_not_called()
+
+        with patch(PATCH_METHOD) as mocked_api_call:
+            assert ({'Key': {'UserName': {'S': 'hash-key'}, 'UserRange': {'S': 'range-key'}}}
+                    == conn.get_identifier_map(table_name_with_schema, 'hash-key', 'range-key'))
+            mocked_api_call.assert_not_called()