Merge pull request #6 from schemacrawler/command

Use command pattern for AI functions

Author: sualeh

.github/workflows/quick_build.yml

@@ -25,6 +25,28 @@ jobs:
       uses: sualeh/prepare-maven-build@v1.4.0
       with:
         java-version: 21
+    - id: install-graphviz
+      name: Install Graphviz
+      uses: sualeh/install-graphviz@v1.0.3
+
+# BUILD DEPENDENCIES
+    - id: checkout-schemacrawler
+      name: Checkout SchemaCrawler
+      uses: actions/checkout@v4
+      with:
+        repository: schemacrawler/SchemaCrawler
+        path: SchemaCrawler
+    - id: build-schemacrawler
+      name: Build SchemaCrawler for local Maven repository
+      shell: bash
+      run: |
+        # Build SchemaCrawler
+        cd SchemaCrawler
+        mvn \
+          --no-transfer-progress \
+          --batch-mode \
+          -Dmaven.test.skip=true \
+          clean install
 
 # BUILD AND TEST
     - id: build-test

src/main/java/schemacrawler/tools/command/aichat/AiChatCommand.java

@@ -53,8 +53,7 @@ public void checkAvailability() throws RuntimeException {
 
   @Override
   public void execute() {
-    try (AiChatConsole aiChatConsole =
-        new AiChatConsole(commandOptions, catalog, connection); ) {
+    try (AiChatConsole aiChatConsole = new AiChatConsole(commandOptions, catalog, connection); ) {
       aiChatConsole.console();
     }
   }

src/main/java/schemacrawler/tools/command/aichat/AiChatCommandProvider.java

@@ -55,8 +55,7 @@ public PluginCommand getCommandLineCommand() {
         .addOption("api-key", String.class, "OpenAI API key")
         .addOption(
             "api-key:env", String.class, "OpenAI API key, from an environmental variable value")
-        .addOption(
-            "model", String.class, "AI chat model", "Optional, defaults to 'gpt-4o-mini'")
+        .addOption("model", String.class, "AI chat model", "Optional, defaults to 'gpt-4o-mini'")
         .addOption(
             "timeout",
             Integer.class,
@@ -86,7 +85,7 @@ public AiChatCommand newSchemaCrawlerCommand(final String command, final Config
           AiChatCommandOptionsBuilder.builder().fromConfig(config).toOptions();
 
       final AiChatCommand scCommand = new AiChatCommand();
-      scCommand.setCommandOptions(options);
+      scCommand.configure(options);
       return scCommand;
     } catch (final Exception e) {
       throw new ExecutionRuntimeException(e);

src/main/java/schemacrawler/tools/command/aichat/FunctionExecutor.java

@@ -28,18 +28,12 @@
 
 package schemacrawler.tools.command.aichat;
 
-import java.sql.Connection;
 import java.util.UUID;
-import java.util.concurrent.Callable;
-import schemacrawler.schema.Catalog;
+import schemacrawler.tools.executable.Command;
 
-public interface FunctionExecutor<P extends FunctionParameters> extends Callable<FunctionReturn> {
+public interface FunctionExecutor<P extends FunctionParameters> extends Command<P, FunctionReturn> {
 
   String getDescription();
 
   UUID getExecutorInstanceId();
-
-  String getName();
-
-  void initialize(P args, Catalog catalog, Connection connection);
 }

src/main/java/schemacrawler/tools/command/aichat/functions/AbstractFunctionExecutor.java

@@ -54,6 +54,11 @@ protected AbstractFunctionExecutor(final PropertyName functionName) {
     executorInstanceId = UUID.randomUUID();
   }
 
+  @Override
+  public void configure(final P args) {
+    this.args = requireNonNull(args, "No arguments provided");
+  }
+
   @Override
   public boolean equals(final Object obj) {
     if (this == obj) {
@@ -66,6 +71,16 @@ public boolean equals(final Object obj) {
     return Objects.equals(args, other.args);
   }
 
+  @Override
+  public Catalog getCatalog() {
+    return catalog;
+  }
+
+  @Override
+  public Connection getConnection() {
+    return connection;
+  }
+
   @Override
   public final String getDescription() {
     return functionName.getDescription();
@@ -86,6 +101,24 @@ public int hashCode() {
     return Objects.hash(args);
   }
 
+  @Override
+  public void initialize() {
+    // No-op
+  }
+
+  @Override
+  public void setCatalog(final Catalog catalog) {
+    this.catalog = requireNonNull(catalog, "Catalog is not provided");
+  }
+
+  @Override
+  public void setConnection(final Connection connection) {
+    if (!usesConnection()) {
+      throw new ExecutionRuntimeException("Function does not use a connection");
+    }
+    this.connection = requireNonNull(connection, "Connection is not provided");
+  }
+
   @Override
   public String toString() {
     return String.format(
@@ -94,15 +127,4 @@ public String toString() {
         new KebabCaseStrategy().translate(args.getClass().getSimpleName()),
         getDescription());
   }
-
-  @Override
-  public void initialize(final P args, final Catalog catalog, final Connection connection) {
-    this.args = requireNonNull(args, "No arguments provided");
-    if (catalog == null) {
-      throw new ExecutionRuntimeException("Catalog is not provided");
-    }
-    this.catalog = catalog;
-    // Connection can be null
-    this.connection = connection;
-  }
 }

src/main/java/schemacrawler/tools/command/aichat/functions/LintFunctionExecutor.java

@@ -75,4 +75,9 @@ protected String getCommand() {
   protected Function<Catalog, Boolean> getResultsChecker() {
     return catalog -> !catalog.getTables().isEmpty();
   }
+
+  @Override
+  public boolean usesConnection() {
+    return true;
+  }
 }

src/main/java/schemacrawler/tools/command/aichat/utility/AiChatUtility.java

@@ -77,7 +77,12 @@ public static <P extends FunctionParameters> String execute(
     try {
       final schemacrawler.tools.command.aichat.FunctionExecutor<P> functionExecutor =
           functionDefinitionToCall.newExecutor();
-      functionExecutor.initialize(parameters, catalog, connection);
+      functionExecutor.configure(parameters);
+      functionExecutor.initialize();
+      functionExecutor.setCatalog(catalog);
+      if (functionExecutor.usesConnection()) {
+        functionExecutor.setConnection(connection);
+      }
       functionReturn = functionExecutor.call();
       return functionReturn.get();
     } catch (final Exception e) {
@@ -108,7 +113,7 @@ public static boolean isExitCondition(final List<ChatMessage> completions) {
   public static FunctionExecutor newFunctionExecutor() {
 
     final List<FunctionDef> chatFunctions = new ArrayList<>();
-    for (final FunctionDefinition functionDefinition :
+    for (final FunctionDefinition<?> functionDefinition :
         FunctionDefinitionRegistry.getFunctionDefinitionRegistry().getFunctionDefinitions()) {
       if (functionDefinition.getFunctionType() != FunctionType.USER) {
         continue;

src/test/java/schemacrawler/tools/command/aichat/function/test/DatabaseObjectDescriptionFunctionTest.java

@@ -159,7 +159,8 @@ private void describeDatabaseObject(
     try (final TestWriter out = testout) {
       final FunctionExecutor<DatabaseObjectDescriptionFunctionParameters> executor =
           functionDefinition.newExecutor();
-      executor.initialize(args, catalog, null);
+      executor.configure(args);
+      executor.setCatalog(catalog);
       final FunctionReturn functionReturn = executor.call();
       out.write(functionReturn.get());
     }

src/test/java/schemacrawler/tools/command/aichat/function/test/DatabaseObjectListFunctionTest.java

@@ -142,7 +142,8 @@ private void databaseObjects(
     try (final TestWriter out = testout) {
       final FunctionExecutor<DatabaseObjectListFunctionParameters> executor =
           functionDefinition.newExecutor();
-      executor.initialize(args, catalog, null);
+      executor.configure(args);
+      executor.setCatalog(catalog);
       final FunctionReturn functionReturn = executor.call();
       out.write(functionReturn.get());
     }

src/test/java/schemacrawler/tools/command/aichat/function/test/ExitFunctionTest.java

@@ -59,7 +59,8 @@ public void exit(final TestContext testContext) throws Exception {
     final TestWriter testout = new TestWriter();
     try (final TestWriter out = testout) {
       final FunctionExecutor<NoParameters> executor = functionDefinition.newExecutor();
-      executor.initialize(args, mock(Catalog.class), null);
+      executor.configure(args);
+      executor.setCatalog(mock(Catalog.class));
       final FunctionReturn functionReturn = executor.call();
       out.write(functionReturn.get());
     }