[Multi] Fix references to the SDK models (#7061)

And update CI automation to run on the reduced version of Azure management SDKs.

.travis.yml

@@ -40,7 +40,9 @@ jobs:
       env: TOXENV=py27
       script: ./scripts/ci/unittest.sh
     - stage: verify
-      env: PURPOSE='Automation'
+      env:
+      - PURPOSE='Automation'
+      - REDUCE_SDK='True'
       script: ./scripts/ci/test_automation.sh
       python: 3.6
     - stage: verify

scripts/ci/test_automation.sh

@@ -32,6 +32,17 @@ pip install -qqq $ALL_MODULES
 title 'Installed packages'
 pip freeze
 
+if [ "$REDUCE_SDK" == "True" ]
+then
+    title 'azure.mgmt file counts'
+    (cd $(dirname $(which python)); cd ../lib/*/site-packages/azure/mgmt; find . -name '*.py' | wc)
+
+    python $(cd $(dirname $0); cd ..; pwd)/sdk_process/patch_models.py
+
+    title 'azure.mgmt file counts after reduce'
+    (cd $(dirname $(which python)); cd ../lib/*/site-packages/azure/mgmt; find . -name '*.py' | wc)
+fi
+
 target_profile=${AZURE_CLI_TEST_TARGET_PROFILE:-latest}
 if [ "$target_profile" != "latest" ]; then
     # example: 2017-03-09-profile

scripts/sdk_process/patch_models.py

@@ -0,0 +1,226 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from enum import Enum
+import importlib
+import inspect
+import logging
+from pathlib import Path
+import pkgutil
+import shutil
+import sys
+import tempfile
+
+from msrest.serialization import Model
+from msrest.paging import Paged
+
+_LOGGER = logging.getLogger(__name__)
+
+copyright_header = b"""# coding=utf-8
+# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+#
+# Code generated by Microsoft (R) AutoRest Code Generator.
+# Changes may cause incorrect behavior and will be lost if the code is
+# regenerated.
+# --------------------------------------------------------------------------
+
+"""
+
+header = copyright_header + b"""from msrest.serialization import Model
+from msrest.exceptions import HttpOperationError
+"""
+
+paging_header = copyright_header + b"""from msrest.paging import Paged
+"""
+
+init_file = """
+try:
+    from .{} import *
+except (SyntaxError, ImportError):
+    from .{} import *
+from .{} import *
+"""
+
+
+def parse_input(input_parameter):
+    """From a syntax like package_name#submodule, build a package name
+    and complete module name.
+    """
+    split_package_name = input_parameter.split('#')
+    package_name = split_package_name[0]
+    module_name = package_name.replace("-", ".")
+    if len(split_package_name) >= 2:
+        module_name = ".".join([module_name, split_package_name[1]])
+    return package_name, module_name
+
+
+def solve_mro(models):
+    for models_module in models:
+        models_path = models_module.__path__[0]
+        _LOGGER.info("Working on %s", models_path)
+        if Path(models_path, "models_py3.py").exists():
+            _LOGGER.info("Skipping since already patched")
+            return
+
+        # Build the new files in a temp folder
+        with tempfile.TemporaryDirectory() as temp_folder:
+            final_models_path = Path(temp_folder, "models")
+            final_models_path.mkdir()
+            solve_one_model(models_module, final_models_path)
+
+            # Switch the files
+            shutil.rmtree(models_path)
+            shutil.move(final_models_path, models_path)
+
+
+def solve_one_model(models_module, output_folder):
+    """Will build the compacted models in the output_folder"""
+
+    models_classes = [
+        (len(model_class.__mro__), inspect.getfile(model_class), model_class) for model_name, model_class in
+        vars(models_module).items()
+        if model_name[0].isupper() and Model in model_class.__mro__
+    ]
+    # Only sort based on the first element in the tuple
+    models_classes.sort(key=lambda x: x[0])
+
+    py2_models_classes = [
+        (len_mro, path.replace("_py3.py", ".py"), None)
+        for len_mro, path, _ in models_classes
+    ]
+
+    paged_models_classes = [
+        (inspect.getfile(model_class), model_class) for model_name, model_class in vars(models_module).items()
+        if model_name[0].isupper() and Paged in model_class.__mro__
+    ]
+
+    enum_models_classes = [
+        (inspect.getfile(model_class), model_class) for model_name, model_class in vars(models_module).items()
+        if model_name[0].isupper() and Enum in model_class.__mro__
+    ]
+    if enum_models_classes:
+        enum_file = Path(enum_models_classes[0][0])
+        shutil.copyfile(enum_file, Path(output_folder, enum_file.name))
+        enum_file_module_name = enum_file.with_suffix('').name
+    else:
+        enum_file_module_name = None
+
+    write_model_file(Path(output_folder, "models_py3.py"), models_classes)
+    write_model_file(Path(output_folder, "models.py"), py2_models_classes)
+    write_paging_file(Path(output_folder, "paged_models.py"), paged_models_classes)
+    write_init(
+        Path(output_folder, "__init__.py"),
+        "models_py3",
+        "models",
+        "paged_models",
+        enum_file_module_name
+    )
+
+
+def write_model_file(output_file_path, classes_to_write):
+    with open(output_file_path, "bw") as write_fd:
+        write_fd.write(header)
+
+        for model in classes_to_write:
+            _, model_file_path, _ = model
+
+            with open(model_file_path, "rb") as read_fd:
+                lines = read_fd.readlines()
+                # Skip until it's "class XXXX"
+                while lines:
+                    if lines[0].startswith(b"class "):
+                        break
+                    lines.pop(0)
+                else:
+                    raise ValueError("Never found any class definition!")
+                # Now I keep everything
+                write_fd.write(b'\n')
+                write_fd.write(b'\n')
+                write_fd.writelines(lines)
+
+
+def write_paging_file(output_file_path, classes_to_write):
+    with open(output_file_path, "bw") as write_fd:
+        write_fd.write(paging_header)
+
+        for model in classes_to_write:
+            model_file_path, _ = model
+
+            with open(model_file_path, "rb") as read_fd:
+                # Skip the first 15 lines (based on Autorest deterministic behavior)
+                # If we want this less random, look for the first line starts with "class"
+                lines = read_fd.readlines()[14:]
+                write_fd.write(b'\n')
+                write_fd.write(b'\n')
+                write_fd.writelines(lines)
+
+
+def write_init(output_file_path, model_file_name, model_file_name_py2, paging_file_name, enum_file_name):
+    with open(output_file_path, "bw") as write_fd:
+        write_fd.write(copyright_header)
+
+        write_fd.write(init_file.format(
+            model_file_name,
+            model_file_name_py2,
+            paging_file_name,
+        ).encode('utf8'))
+        if enum_file_name:
+            write_fd.write(
+                "from .{} import *".format(enum_file_name).encode('utf8')
+            )
+
+
+def find_models_to_change(module_name):
+    """Will figure out if the package is a multi-api one,
+    and understand what to generate.
+    """
+    main_module = importlib.import_module(module_name)
+    try:
+        models_module = main_module.models
+        models_module.__path__
+        # It didn't fail, that's a single API package
+        return [models_module]
+    except AttributeError:
+        # This means I loaded the fake module "models"
+        # and it's multi-api, load all models
+        return [
+            importlib.import_module('.' + label + '.models', main_module.__name__)
+            for (_, label, ispkg) in pkgutil.iter_modules(main_module.__path__)
+            if ispkg
+        ]
+
+
+def find_autorest_generated_folder(module_prefix="azure.mgmt"):
+    """Find all Autorest generated code in that module prefix.
+
+    This actually looks for a "models" package only. We could be smarter if necessary.
+    """
+    _LOGGER.info("Looking for Autorest generated package in %s", module_prefix)
+    result = []
+    prefix_module = importlib.import_module(module_prefix)
+    for _, sub_package, ispkg in pkgutil.iter_modules(prefix_module.__path__, module_prefix + "."):
+        try:
+            _LOGGER.debug("Try %s", sub_package)
+            importlib.import_module(".models", sub_package)
+            # If not exception, we found it
+            _LOGGER.info("Found %s", sub_package)
+            result.append(sub_package)
+        except ModuleNotFoundError:
+            # No model, might dig deeper
+            if ispkg:
+                result += find_autorest_generated_folder(sub_package)
+    return result
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+
+    prefix = sys.argv[1] if len(sys.argv) >= 2 else "azure.mgmt"
+    for autorest_package in find_autorest_generated_folder(prefix):
+        models = find_models_to_change(autorest_package)
+        solve_mro(models)

src/command_modules/azure-cli-ams/HISTORY.rst

@@ -3,6 +3,10 @@
 Release History
 ===============
 
+0.2.3
++++++
+* Minor fixes
+
 0.2.2
 +++++
 * Minor changes

src/command_modules/azure-cli-ams/azure/cli/command_modules/ams/_exception_handler.py

@@ -5,7 +5,7 @@
 
 
 def ams_exception_handler(ex):
-    from azure.mgmt.media.models.api_error import ApiErrorException
+    from azure.mgmt.media.models import ApiErrorException
     from knack.util import CLIError
 
     if isinstance(ex, ApiErrorException) \

src/command_modules/azure-cli-ams/setup.py

@@ -14,7 +14,7 @@
     logger.warn("Wheel is not available, disabling bdist_wheel hook")
     cmdclass = {}
 
-VERSION = "0.2.2"
+VERSION = "0.2.3"
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers

src/command_modules/azure-cli-batch/HISTORY.rst

@@ -3,6 +3,10 @@
 Release History
 ===============
 
+3.3.3
++++++
+* Minor fixes
+
 3.3.2
 +++++
 * Update Key Vault SDK dependency

src/command_modules/azure-cli-batch/azure/cli/command_modules/batch/_params.py

@@ -7,8 +7,7 @@
 
 from knack.arguments import CLIArgumentType
 
-from azure.mgmt.batch.models.batch_management_client_enums import \
-    (AccountKeyType)
+from azure.mgmt.batch.models import AccountKeyType
 from azure.batch.models.batch_service_client_enums import \
     (ComputeNodeDeallocationOption)
 

src/command_modules/azure-cli-batch/setup.py

@@ -15,7 +15,7 @@
     logger.warn("Wheel is not available, disabling bdist_wheel hook")
     cmdclass = {}
 
-VERSION = "3.3.2"
+VERSION = "3.3.3"
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers
 CLASSIFIERS = [

src/command_modules/azure-cli-cosmosdb/HISTORY.rst

@@ -3,6 +3,10 @@
 Release History
 ===============
 
+0.2.1
++++++
+* Minor fixes
+
 0.2.0
 +++++
 * BREAKING CHANGE: 'show' commands log error message and fail with exit code of 3 upon a missing resource.

src/command_modules/azure-cli-cosmosdb/azure/cli/command_modules/cosmosdb/_params.py

@@ -17,7 +17,7 @@
 
 def load_arguments(self, _):
 
-    from azure.mgmt.cosmosdb.models.cosmos_db_enums import KeyKind, DefaultConsistencyLevel, DatabaseAccountKind
+    from azure.mgmt.cosmosdb.models import KeyKind, DefaultConsistencyLevel, DatabaseAccountKind
 
     with self.argument_context('cosmosdb') as c:
         c.argument('account_name', arg_type=name_type, help='Name of the Cosmos DB database account', completer=get_resource_name_completion_list('Microsoft.DocumentDb/databaseAccounts'), id_part='name')

src/command_modules/azure-cli-cosmosdb/azure/cli/command_modules/cosmosdb/_validators.py

@@ -6,7 +6,7 @@
 
 def validate_failover_policies(ns):
     """ Extracts multiple space-separated failoverPolicies in regionName=failoverPriority format """
-    from azure.mgmt.cosmosdb.models.failover_policy import FailoverPolicy
+    from azure.mgmt.cosmosdb.models import FailoverPolicy
     fp_dict = []
     for item in ns.failover_policies:
         comps = item.split('=', 1)
@@ -16,7 +16,7 @@ def validate_failover_policies(ns):
 
 def validate_locations(ns):
     """ Extracts multiple space-separated locations in regionName=failoverPriority format """
-    from azure.mgmt.cosmosdb.models.location import Location
+    from azure.mgmt.cosmosdb.models import Location
     if ns.locations is None:
         ns.locations = []
         return
@@ -34,7 +34,7 @@ def validate_ip_range_filter(ns):
 
 def validate_capabilities(ns):
     """ Extracts multiple space-separated capabilities """
-    from azure.mgmt.cosmosdb.models.capability import Capability
+    from azure.mgmt.cosmosdb.models import Capability
     if ns.capabilities is not None:
         capabilties_list = []
         for item in ns.capabilities:
@@ -44,7 +44,7 @@ def validate_capabilities(ns):
 
 def validate_virtual_network_rules(ns):
     """ Extracts multiple space-separated virtual network rules """
-    from azure.mgmt.cosmosdb.models.virtual_network_rule import VirtualNetworkRule
+    from azure.mgmt.cosmosdb.models import VirtualNetworkRule
     if ns.virtual_network_rules is not None:
         virtual_network_rules_list = []
         for item in ns.virtual_network_rules: