Update YdbDatabaseModelFactory.cs (#301)

Author: KirillKurdyukov

src/EfCore.Ydb/src/Scaffolding/Internal/YdbDatabaseModelFactory.cs

@@ -21,85 +21,97 @@ public override DatabaseModel Create(string connectionString, DatabaseModelFacto
     public override DatabaseModel Create(DbConnection connection, DatabaseModelFactoryOptions options)
     {
         var ydbConnection = (YdbConnection)connection;
-        if (connection.State != ConnectionState.Open)
-        {
-            connection.Open();
-        }
-
-        var tableNames = new List<string>();
-        tableNames.AddRange(options.Tables);
+        var connectionStartedOpen = ydbConnection.State == ConnectionState.Open;
 
-        if (tableNames.Count == 0)
+        if (!connectionStartedOpen)
         {
-            tableNames.AddRange(
-                from ydbObject in YdbSchema.SchemaObjects(ydbConnection).GetAwaiter().GetResult()
-                where ydbObject.Type is SchemeType.Table or SchemeType.ColumnTable or SchemeType.ExternalTable &&
-                      !ydbObject.IsSystem
-                select ydbObject.Name
-            );
+            ydbConnection.Open();
         }
 
-        var databaseModel = new DatabaseModel
+        try
         {
-            DatabaseName = connection.Database
-        };
+            var tableNames = new List<string>();
+            tableNames.AddRange(options.Tables);
 
-        foreach (var ydbTable in tableNames.Select(tableName =>
-                     YdbSchema.DescribeTable(ydbConnection, tableName).GetAwaiter().GetResult()))
-        {
-            var databaseTable = new DatabaseTable
+            if (tableNames.Count == 0)
             {
-                Name = ydbTable.Name,
-                Database = databaseModel
-            };
+                tableNames.AddRange(
+                    from ydbObject in YdbSchema.SchemaObjects(ydbConnection).GetAwaiter().GetResult()
+                    where ydbObject.Type is SchemeType.Table or SchemeType.ColumnTable or SchemeType.ExternalTable &&
+                          !ydbObject.IsSystem
+                    select ydbObject.Name
+                );
+            }
 
-            var columnNameToDatabaseColumn = new Dictionary<string, DatabaseColumn>();
+            var databaseModel = new DatabaseModel
+            {
+                DatabaseName = connection.Database
+            };
 
-            foreach (var column in ydbTable.Columns)
+            foreach (var ydbTable in tableNames.Select(tableName =>
+                         YdbSchema.DescribeTable(ydbConnection, tableName).GetAwaiter().GetResult()))
             {
-                var databaseColumn = new DatabaseColumn
+                var databaseTable = new DatabaseTable
                 {
-                    Name = column.Name,
-                    Table = databaseTable,
-                    StoreType = column.StorageType,
-                    IsNullable = column.IsNullable
+                    Name = ydbTable.Name,
+                    Database = databaseModel
                 };
 
-                databaseTable.Columns.Add(databaseColumn);
-                columnNameToDatabaseColumn[column.Name] = databaseColumn;
-            }
+                var columnNameToDatabaseColumn = new Dictionary<string, DatabaseColumn>();
 
-            foreach (var index in ydbTable.Indexes)
-            {
-                var databaseIndex = new DatabaseIndex
+                foreach (var column in ydbTable.Columns)
                 {
-                    Name = index.Name,
-                    Table = databaseTable,
-                    IsUnique = index.Type == YdbTableIndex.IndexType.GlobalUniqueIndex
+                    var databaseColumn = new DatabaseColumn
+                    {
+                        Name = column.Name,
+                        Table = databaseTable,
+                        StoreType = column.StorageType,
+                        IsNullable = column.IsNullable
+                    };
+
+                    databaseTable.Columns.Add(databaseColumn);
+                    columnNameToDatabaseColumn[column.Name] = databaseColumn;
+                }
+
+                foreach (var index in ydbTable.Indexes)
+                {
+                    var databaseIndex = new DatabaseIndex
+                    {
+                        Name = index.Name,
+                        Table = databaseTable,
+                        IsUnique = index.Type == YdbTableIndex.IndexType.GlobalUniqueIndex
+                    };
+
+                    foreach (var columnName in index.IndexColumns)
+                    {
+                        databaseIndex.Columns.Add(columnNameToDatabaseColumn[columnName]);
+                        databaseIndex.IsDescending.Add(false);
+                    }
+
+                    databaseTable.Indexes.Add(databaseIndex);
+                }
+
+                databaseTable.PrimaryKey = new DatabasePrimaryKey
+                {
+                    Name = null // YDB does not have a primary key named
                 };
 
-                foreach (var columnName in index.IndexColumns)
+                foreach (var columnName in ydbTable.PrimaryKey)
                 {
-                    databaseIndex.Columns.Add(columnNameToDatabaseColumn[columnName]);
-                    databaseIndex.IsDescending.Add(false);
+                    databaseTable.PrimaryKey.Columns.Add(columnNameToDatabaseColumn[columnName]);
                 }
 
-                databaseTable.Indexes.Add(databaseIndex);
+                databaseModel.Tables.Add(databaseTable);
             }
 
-            databaseTable.PrimaryKey = new DatabasePrimaryKey
-            {
-                Name = null // YDB does not have a primary key named
-            };
-
-            foreach (var columnName in ydbTable.PrimaryKey)
+            return databaseModel;
+        }
+        finally
+        {
+            if (!connectionStartedOpen)
             {
-                databaseTable.PrimaryKey.Columns.Add(columnNameToDatabaseColumn[columnName]);
+                ydbConnection.Close();
             }
-
-            databaseModel.Tables.Add(databaseTable);
         }
-
-        return databaseModel;
     }
 }