Fix MODEL_TYPE fragment classification in CreateExternalModelStatement

SqlScriptDom/Parser/TSql/Ast.xml

@@ -1,4 +1,4 @@
-<?xml version="1.0"?>
+<?xml version="1.0"?>
 <!--
  <copyright file="Ast.xml" company="Microsoft">
          Copyright (c) Microsoft Corporation.  All rights reserved.
@@ -1364,7 +1364,7 @@
     <Member Name="Name" Type="Identifier" Summary="The external model name."/>
     <Member Name="Location" Type="Literal" Summary="The external model location name."/>
     <Member Name="ApiFormat" Type="Literal" Summary="The external model api format name."/>
-    <Member Name="ModelType" Type="ExternalModelTypeOption?" GenerateUpdatePositionInfoCall="false" Summary="The external model type name."/>
+    <Member Name="ModelType" Type="ExternalModelTypeOption" Summary="The external model type name."/>
     <Member Name="ModelName" Type="Literal" Summary="The external model name to be used to generate embeddings."/>
     <Member Name="Credential" Type="Identifier" Summary="The external model credentials name."/>
     <Member Name="Parameters" Type="Literal" Summary="The external model parameters as key-value pairs."/>
@@ -1380,6 +1380,9 @@
   <Class Name="AlterExternalModelStatement" Base="ExternalModelStatement" Summary="Represents a ALTER EXTERNAL MODEL statement.">
     <InheritedClass Name="ExternalModelStatement"/>
   </Class>
+  <Class Name="ExternalModelTypeOption" Base="TSqlFragment" Summary="Represents the MODEL_TYPE option in CREATE/ALTER EXTERNAL MODEL.">
+    <Member Name="OptionKind" Type="ExternalModelTypeOptionKind" GenerateUpdatePositionInfoCall="false" Summary="The option kind."/>
+  </Class>
 
   <Class Name="ExternalFileFormatStatement" Abstract="true" Base="TSqlStatement" Summary="Base class for all external file format statement objects.">
     <Member Name="Name" Type="Identifier" Summary="The external file format name."/>

SqlScriptDom/Parser/TSql/ExternalModelTypeOption.cs

@@ -1,4 +1,4 @@
-//------------------------------------------------------------------------------
+//------------------------------------------------------------------------------
 // <copyright file="ExternalModelTypeOption.cs" company="Microsoft">
 //         Copyright (c) Microsoft Corporation.  All rights reserved.
 // </copyright>
@@ -11,14 +11,14 @@ namespace Microsoft.SqlServer.TransactSql.ScriptDom
 
     /// <summary>
     /// The enumeration specifies the external model type
-    /// Currently, we support EMBEDDINGS only.
+    /// Currently, we support Embeddings only.
     /// </summary>
-    public enum ExternalModelTypeOption
+    public enum ExternalModelTypeOptionKind
     {
         /// <summary>
         /// MODEL_TYPE = EMBEDDINGS
         /// </summary>
-        EMBEDDINGS = 0,
+        Embeddings = 0,
 
     }
 

SqlScriptDom/Parser/TSql/TSql170.g

@@ -24170,20 +24170,26 @@ StringLiteral vApiFormat;
 ;
 
 externalModelModelType[ExternalModelStatement vParent]
+{
+    ExternalModelTypeOption vModelTypeOption = null;
+}
   :
         tModelType:Identifier
         {
             Match(tModelType, CodeGenerationSupporter.ModelType);
-            UpdateTokenInfo(vParent, tModelType);
+            vModelTypeOption = this.FragmentFactory.CreateFragment<ExternalModelTypeOption>();
+            UpdateTokenInfo(vModelTypeOption, tModelType);
         }
         EqualsSign
         (
             tEmbeddings:Identifier
             {
                 if (TryMatch(tEmbeddings, CodeGenerationSupporter.Embeddings))
                 {
-                    vParent.ModelType = ExternalModelTypeOption.EMBEDDINGS;
-                    UpdateTokenInfo(vParent, tEmbeddings);
+                    vModelTypeOption.OptionKind = ExternalModelTypeOptionKind.Embeddings;
+                    UpdateTokenInfo(vModelTypeOption, tEmbeddings);
+                    vParent.ModelType = vModelTypeOption;
+                    vParent.UpdateTokenInfo(vModelTypeOption);
                 }
                 else
                 {

SqlScriptDom/Parser/TSql/TSql180.g

@@ -24170,20 +24170,26 @@ StringLiteral vApiFormat;
 ;
 
 externalModelModelType[ExternalModelStatement vParent]
+{
+    ExternalModelTypeOption vModelTypeOption = null;
+}
   :
         tModelType:Identifier
         {
             Match(tModelType, CodeGenerationSupporter.ModelType);
-            UpdateTokenInfo(vParent, tModelType);
+            vModelTypeOption = this.FragmentFactory.CreateFragment<ExternalModelTypeOption>();
+            UpdateTokenInfo(vModelTypeOption, tModelType);
         }
         EqualsSign
         (
             tEmbeddings:Identifier
             {
                 if (TryMatch(tEmbeddings, CodeGenerationSupporter.Embeddings))
                 {
-                    vParent.ModelType = ExternalModelTypeOption.EMBEDDINGS;
-                    UpdateTokenInfo(vParent, tEmbeddings);
+                    vModelTypeOption.OptionKind = ExternalModelTypeOptionKind.Embeddings;
+                    UpdateTokenInfo(vModelTypeOption, tEmbeddings);
+                    vParent.ModelType = vModelTypeOption;
+                    vParent.UpdateTokenInfo(vModelTypeOption);
                 }
                 else
                 {

SqlScriptDom/ScriptDom/SqlServer/ScriptGenerator/SqlScriptGenerator.AlterExternalModelStatement.cs

@@ -1,4 +1,4 @@
-//------------------------------------------------------------------------------
+//------------------------------------------------------------------------------
 // <copyright file="SqlScriptGeneratorVisitor.AlterExternalModelStatement.cs" company="Microsoft">
 //         Copyright (c) Microsoft Corporation.  All rights reserved.
 // </copyright>
@@ -53,17 +53,15 @@ protected void GenerateAlterExternalModelStatementBody(AlterExternalModelStateme
             }
 
             // external model Model Type options
-            if (node.ModelType == ExternalModelTypeOption.EMBEDDINGS)
+            if (node.ModelType != null)
             {
                 if (!ifFirst)
                 {
                     GenerateSymbol(TSqlTokenType.Comma);
                 }
                 ifFirst = false;
-                ExternalModelTypeOption typeOption = ExternalModelTypeOption.EMBEDDINGS;
-                string externalModelTypeOption = GetValueForEnumKey(_externalModelTypeOption, typeOption);
                 NewLine();
-                GenerateNameEqualsValue(CodeGenerationSupporter.ModelType, externalModelTypeOption);
+                GenerateFragmentIfNotNull(node.ModelType);
             }
 
             // external model name options

SqlScriptDom/ScriptDom/SqlServer/ScriptGenerator/SqlScriptGeneratorVisitor.CreateExternalModelStatement.cs

@@ -1,4 +1,4 @@
-//------------------------------------------------------------------------------
+//------------------------------------------------------------------------------
 // <copyright file="SqlScriptGeneratorVisitor.CreateExternalModelStatement.cs" company="Microsoft">
 //         Copyright (c) Microsoft Corporation.  All rights reserved.
 // </copyright>
@@ -17,9 +17,9 @@ public override void ExplicitVisit(CreateExternalModelStatement node)
             GenerateSpaceAndIdentifier(CodeGenerationSupporter.Model);
             GenerateCreateExternalModelStatementBody(node);
         }
-        protected static Dictionary<ExternalModelTypeOption, string> _externalModelTypeOption = new Dictionary<ExternalModelTypeOption, string>()
+        protected static Dictionary<ExternalModelTypeOptionKind, string> _externalModelTypeOptionKind = new Dictionary<ExternalModelTypeOptionKind, string>()
         {
-            {ExternalModelTypeOption.EMBEDDINGS, CodeGenerationSupporter.Embeddings}
+            {ExternalModelTypeOptionKind.Embeddings, CodeGenerationSupporter.Embeddings}
         };
 
         protected void GenerateCreateExternalModelStatementBody(CreateExternalModelStatement node)
@@ -54,17 +54,15 @@ protected void GenerateCreateExternalModelStatementBody(CreateExternalModelState
             }
 
             // external model Model Type options
-            if (node.ModelType == ExternalModelTypeOption.EMBEDDINGS)
+            if (node.ModelType != null)
             {
                 if (!ifFirst)
                 {
                     GenerateSymbol(TSqlTokenType.Comma);
                 }
                 ifFirst = false;
-                ExternalModelTypeOption typeOption = ExternalModelTypeOption.EMBEDDINGS;
-                string externalModelTypeOption = GetValueForEnumKey(_externalModelTypeOption, typeOption);
                 NewLine();
-                GenerateNameEqualsValue(CodeGenerationSupporter.ModelType, externalModelTypeOption);
+                GenerateFragmentIfNotNull(node.ModelType);
             }
 
             // external model name options
@@ -118,5 +116,11 @@ protected void GenerateCreateExternalModelStatementBody(CreateExternalModelState
             NewLine();
             GenerateKeyword(TSqlTokenType.RightParenthesis);
         }
+
+        public override void ExplicitVisit(ExternalModelTypeOption node)
+        {
+            string optionKindString = GetValueForEnumKey(_externalModelTypeOptionKind, node.OptionKind);
+            GenerateNameEqualsValue(CodeGenerationSupporter.ModelType, optionKindString);
+        }
     }
 }

Test/SqlDom/Only170SyntaxTests.cs

@@ -392,5 +392,37 @@ public void TSql100SyntaxIn170ParserTest()
                 ParserTest.ParseAndVerify(parser, scriptGen, ti._scriptFilename, ti._result100);
             }
         }
+
+        [TestMethod]
+        [Priority(0)]
+        [SqlStudioTestCategory(Category.UnitTest)]
+        public void TestExternalModelModelTypeVisitor()
+        {
+            string sql = "CREATE EXTERNAL MODEL simple_model WITH (LOCATION = '/models/simple', API_FORMAT = 'OpenAI', MODEL_TYPE = EMBEDDINGS, MODEL = 'gpt-3.5-turbo');";
+            TSql170Parser parser = new TSql170Parser(true);
+            IList<ParseError> errors;
+            TSqlFragment fragment = parser.Parse(new StringReader(sql), out errors);
+
+            Assert.AreEqual(0, errors.Count);
+            Assert.IsNotNull(fragment);
+
+            // Collect all visited fragments using a custom visitor
+            var visitor = new ExternalModelTypeOptionVisitor();
+            fragment.Accept(visitor);
+
+            Assert.AreEqual(1, visitor.VisitedOptions.Count);
+            Assert.AreEqual(ExternalModelTypeOptionKind.Embeddings, visitor.VisitedOptions[0].OptionKind);
+        }
+
+        private class ExternalModelTypeOptionVisitor : TSqlFragmentVisitor
+        {
+            public List<ExternalModelTypeOption> VisitedOptions { get; } = new List<ExternalModelTypeOption>();
+
+            public override void Visit(ExternalModelTypeOption node)
+            {
+                VisitedOptions.Add(node);
+                base.Visit(node);
+            }
+        }
     }
 }