From d3671d396629c904f894f4d8082be99d7f08d9e4 Mon Sep 17 00:00:00 2001
From: Thomas Bleijendaal <t.bleijendaal@wearetriple.com>
Date: Tue, 22 Aug 2023 19:30:57 +0200
Subject: [PATCH] Allow custom schemas using IDataSourceSchemaProvider

---
 .../ResourceProvider/ResourceRegistryTest.cs  | 36 ++++++++++++++++---
 .../DataSourceRegistration.cs                 |  5 +++
 .../IDataSourceSchemaProvider.cs              |  6 ++++
 .../ResourceProvider/ResourceRegistry.cs      | 17 ++++++---
 ...erviceCollectionResourceRegistryContext.cs |  4 +--
 5 files changed, 58 insertions(+), 10 deletions(-)
 create mode 100644 TerraformPluginDotNet/ResourceProvider/DataSourceRegistration.cs
 create mode 100644 TerraformPluginDotNet/ResourceProvider/IDataSourceSchemaProvider.cs

diff --git a/TerraformPluginDotNet.Test/ResourceProvider/ResourceRegistryTest.cs b/TerraformPluginDotNet.Test/ResourceProvider/ResourceRegistryTest.cs
index 7da487a..072943c 100644
--- a/TerraformPluginDotNet.Test/ResourceProvider/ResourceRegistryTest.cs
+++ b/TerraformPluginDotNet.Test/ResourceProvider/ResourceRegistryTest.cs
@@ -1,4 +1,5 @@
-using Microsoft.Extensions.Logging.Abstractions;
+using System.Collections.Generic;
+using Microsoft.Extensions.Logging.Abstractions;
 using NUnit.Framework;
 using TerraformPluginDotNet.ResourceProvider;
 using TerraformPluginDotNet.Schemas;
@@ -27,18 +28,45 @@ public void TestRegisterDataSource()
     {
         var registry = BuildRegistry();
 
-        Assert.That(registry.DataSchemas.Keys, Is.EquivalentTo(new[] { "test_data_source" }));
+        Assert.That(registry.DataSchemas.Keys, Is.EquivalentTo(new[] { "test_data_source", "test_provided_data_source" }));
         Assert.That(registry.DataSchemas["test_data_source"], Is.InstanceOf<Schema>());
+        Assert.That(registry.DataSchemas["test_provided_data_source"], Is.InstanceOf<Schema>());
 
-        Assert.That(registry.DataTypes.Keys, Is.EquivalentTo(new[] { "test_data_source" }));
+        Assert.That(registry.DataTypes.Keys, Is.EquivalentTo(new[] { "test_data_source", "test_provided_data_source" }));
         Assert.That(registry.DataTypes["test_data_source"], Is.EqualTo(typeof(TestResource)));
+        Assert.That(registry.DataTypes["test_provided_data_source"], Is.EqualTo(typeof(TestResource)));
     }
 
     private static ResourceRegistry BuildRegistry()
     {
+        var schemaBuilder = new SchemaBuilder(NullLogger<SchemaBuilder>.Instance, new TerraformTypeBuilder());
         var registration = new ResourceRegistryRegistration("test_resource", typeof(TestResource));
         var dataSourceRegistration = new DataSourceRegistryRegistration("test_data_source", typeof(TestResource));
-        var registry = new ResourceRegistry(new SchemaBuilder(NullLogger<SchemaBuilder>.Instance, new TerraformTypeBuilder()), new[] { registration }, new[] { dataSourceRegistration });
+        var dataSourceSchemaProvider = new SchemaProvider(schemaBuilder);
+        var registry = new ResourceRegistry(
+            schemaBuilder,
+            new[] { registration },
+            new[] { dataSourceRegistration },
+            new[] { dataSourceSchemaProvider });
         return registry;
     }
+
+    private class SchemaProvider : IDataSourceSchemaProvider
+    {
+        private readonly SchemaBuilder _schemaBuilder;
+
+        public SchemaProvider(
+            SchemaBuilder schemaBuilder)
+        {
+            _schemaBuilder = schemaBuilder;
+        }
+
+        public IEnumerable<DataSourceRegistration> GetSchemas()
+        {
+            yield return new DataSourceRegistration(
+                "test_provided_data_source",
+                typeof(TestResource),
+                _schemaBuilder.BuildSchema(typeof(TestResource)));
+        }
+    }
 }
diff --git a/TerraformPluginDotNet/ResourceProvider/DataSourceRegistration.cs b/TerraformPluginDotNet/ResourceProvider/DataSourceRegistration.cs
new file mode 100644
index 0000000..763728b
--- /dev/null
+++ b/TerraformPluginDotNet/ResourceProvider/DataSourceRegistration.cs
@@ -0,0 +1,5 @@
+using Tfplugin5;
+
+namespace TerraformPluginDotNet.ResourceProvider;
+
+public record DataSourceRegistration(string ResourceName, Type Type, Schema Schema);
diff --git a/TerraformPluginDotNet/ResourceProvider/IDataSourceSchemaProvider.cs b/TerraformPluginDotNet/ResourceProvider/IDataSourceSchemaProvider.cs
new file mode 100644
index 0000000..f82100d
--- /dev/null
+++ b/TerraformPluginDotNet/ResourceProvider/IDataSourceSchemaProvider.cs
@@ -0,0 +1,6 @@
+namespace TerraformPluginDotNet.ResourceProvider;
+
+public interface IDataSourceSchemaProvider
+{
+    IEnumerable<DataSourceRegistration> GetSchemas();
+}
diff --git a/TerraformPluginDotNet/ResourceProvider/ResourceRegistry.cs b/TerraformPluginDotNet/ResourceProvider/ResourceRegistry.cs
index 75f5ce5..55fbe4d 100644
--- a/TerraformPluginDotNet/ResourceProvider/ResourceRegistry.cs
+++ b/TerraformPluginDotNet/ResourceProvider/ResourceRegistry.cs
@@ -6,20 +6,29 @@ namespace TerraformPluginDotNet.ResourceProvider;
 class ResourceRegistry
 {
     public ResourceRegistry(
-        ISchemaBuilder schemaBulder,
+        ISchemaBuilder schemaBuilder,
         IEnumerable<ResourceRegistryRegistration> resourceRegistrations,
-        IEnumerable<DataSourceRegistryRegistration> dataSourceRegistrations)
+        IEnumerable<DataSourceRegistryRegistration> dataSourceRegistrations,
+        IEnumerable<IDataSourceSchemaProvider> dataSourceSchemaProviders)
     {
         foreach (var registration in resourceRegistrations)
         {
-            Schemas.Add(registration.ResourceName, schemaBulder.BuildSchema(registration.Type));
+            Schemas.Add(registration.ResourceName, schemaBuilder.BuildSchema(registration.Type));
             Types.Add(registration.ResourceName, registration.Type);
         }
         foreach (var registration in dataSourceRegistrations)
         {
-            DataSchemas.Add(registration.ResourceName, schemaBulder.BuildSchema(registration.Type));
+            DataSchemas.Add(registration.ResourceName, schemaBuilder.BuildSchema(registration.Type));
             DataTypes.Add(registration.ResourceName, registration.Type);
         }
+        foreach (var provider in dataSourceSchemaProviders)
+        {
+            foreach (var registration in provider.GetSchemas())
+            {
+                DataSchemas.Add(registration.ResourceName, registration.Schema);
+                DataTypes.Add(registration.ResourceName, registration.Type);
+            }
+        }
     }
 
     public Dictionary<string, Schema> Schemas { get; } = new Dictionary<string, Schema>();
diff --git a/TerraformPluginDotNet/ResourceProvider/ServiceCollectionResourceRegistryContext.cs b/TerraformPluginDotNet/ResourceProvider/ServiceCollectionResourceRegistryContext.cs
index dae3bf2..f35f60e 100644
--- a/TerraformPluginDotNet/ResourceProvider/ServiceCollectionResourceRegistryContext.cs
+++ b/TerraformPluginDotNet/ResourceProvider/ServiceCollectionResourceRegistryContext.cs
@@ -18,11 +18,11 @@ public void RegisterResource<T>(string resourceName)
         _services.AddSingleton(new ResourceRegistryRegistration(resourceName, typeof(T)));
     }
 
-    public void RegisterDataSource<T>(string resourceName)
+    public void RegisterDataSource<T>(string dataSourceName)
     {
         EnsureValidType<T>();
 
-        _services.AddSingleton(new DataSourceRegistryRegistration(resourceName, typeof(T)));
+        _services.AddSingleton(new DataSourceRegistryRegistration(dataSourceName, typeof(T)));
     }
 
     private static void EnsureValidType<T>()