feat: implement management command support (#51)

* feat: implement msi auth support

* doc: document msi resource id usage support

* refactor: expose `cmd` package for library usage

* feat: implement management command

* test: add test for management command

* doc: update readme

* Update internal/kusto/management.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Initial plan

* Initial plan to fix debug message for management commands

Co-authored-by: bcho <1975118+bcho@users.noreply.github.com>

* Fix debug message for management commands from 'file ingestion settings' to 'management command settings'

Co-authored-by: bcho <1975118+bcho@users.noreply.github.com>

* Remove kusto-ingest binary from version control and add to .gitignore

Co-authored-by: bcho <1975118+bcho@users.noreply.github.com>

* fix: replace deprecreated opts

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: bcho <1975118+bcho@users.noreply.github.com>

Author: 3 people

.gitignore

@@ -8,6 +8,9 @@
 *.so
 *.dylib
 
+# Built binary
+kusto-ingest
+
 # Test binary, built with `go test -c`
 *.test
 

README.md

@@ -2,7 +2,10 @@
 
 Ingest logs data to Kusto with `github.com/Azure/azure-kusto-go`.
 
-## Usage - Ingest file
+
+## Usage
+
+### Ingest file
 
 ```
 $ kusto-ingest file ./testdata/logs.multijson \
@@ -13,32 +16,58 @@ $ kusto-ingest file ./testdata/logs.multijson \
     --kusto-table="TestTable"
 ```
 
-### Authentication - AZCLI
+### Management commands
+
+Run Kusto management commands from a file (e.g., create tables, update policies):
+
+```
+$ kusto-ingest management ./testdata/commands.kql \
+    --auth-azcli \
+    --kusto-endpoint="https://test.kusto.windows.net" \
+    --kusto-database="Test"
+```
 
-`kusto-ingest` supports using azcli for authentication. This option is helpful for OIDC based pipeline usage.
-You can use `--auth-azcli` to enable it.
+Where `./testdata/commands.kql` contains Kusto management commands, e.g.:
+
+```
+.create table TestTable (Timestamp: datetime, Message: string)
+.alter table TestTable policy update @'{"SoftDeletePeriod": "P365D"}'
+```
+
+#### Options for management subcommand
+
+- `--auth-azcli` or other authentication options (see below)
+- `--kusto-endpoint` (required)
+- `--kusto-database` (required)
+
+
+### Authentication
+
+#### AZCLI
+
+Use Azure CLI for authentication (helpful for OIDC-based pipeline usage):
 
 ```
 $ kusto-ingest file ./testdata/logs.multijson \
     # ... other options
     --auth-azcli
 ```
 
-### Authentication - Managed Identity
 
-`kusto-ingest` supports using managed identity for authentication. This option is helpful for Azure services running in Azure.
-You can use `--auth-managed-identity-resource-id=<mi-resource-id>` to enable it.
+#### Managed Identity
+
+Use managed identity for authentication (helpful for Azure services running in Azure):
 
 ```
 $ kusto-ingest file ./testdata/logs.multijson \
     # ... other options
     --auth-managed-identity-resource-id=<mi-resource-id>
 ```
 
-### Authentication - Service Principal ID and Secret (not recommended)
 
-`kusto-ingest` supports using service principal ID and secret for authentication. This is helpful for existing
-pipeline usage with service principal ID and secret. But this is not recommended for new pipeline usage.
+#### Service Principal ID and Secret (not recommended)
+
+Use service principal ID and secret for authentication (not recommended for new pipelines):
 
 ```
 $ kusto-ingest file ./testdata/logs.multijson \

cmd/cli.go

@@ -10,7 +10,8 @@ import (
 var CLI struct {
 	Verbose bool `short:"v" help:"Enable verbose logging."`
 
-	File kusto.FileIngestOptions `cmd:"" help:"Ingest data from local file."`
+	File       kusto.FileIngestOptions `cmd:"" help:"Ingest data from local file."`
+	Management kusto.ManagementOptions `cmd:"" aliases:"mgmt" help:"Run Kusto management commands from a file."`
 }
 
 // Main is the entry point for the CLI application.

internal/kusto/management.go

@@ -0,0 +1,39 @@
+package kusto
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/Azure/azure-kusto-go/kusto/kql"
+	"github.com/Azure/kusto-ingest/internal/cli"
+)
+
+func (m ManagementOptions) Run(cli cli.Provider) error {
+	cli.Logger().Debug(
+		"management command settings",
+		"target.endpoint", m.KustoTarget.Endpoint,
+		"target.database", m.KustoTarget.Database,
+		"auth.tenant", m.Auth.TenantID,
+		"auth.clientID", m.Auth.ClientID,
+	)
+
+	queryer, err := m.createQueryClient(m.KustoTarget, m.Auth)
+	if err != nil {
+		return fmt.Errorf("create Kusto query client: %w", err)
+	}
+	defer func() { _ = queryer.Close() }()
+
+	ctx, cancel := cli.Context()
+	defer cancel()
+
+	cli.Logger().Info("executing management commands")
+	start := time.Now()
+	stmt := kql.New("").AddUnsafe(string(m.Source))
+	_, err = queryer.Mgmt(ctx, m.KustoTarget.Database, stmt)
+	if err != nil {
+		return fmt.Errorf("execute management commands: %w", err)
+	}
+	cli.Logger().Info("management commands executed", "duration", time.Since(start))
+
+	return nil
+}

internal/kusto/management_test.go

@@ -0,0 +1,85 @@
+package kusto
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/Azure/azure-kusto-go/kusto"
+	"github.com/Azure/azure-kusto-go/kusto/ingest"
+	"github.com/Azure/kusto-ingest/internal/cli/testingcli"
+	"github.com/Azure/kusto-ingest/internal/kusto/testingkusto"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_ManagementOptions_Run_Success(t *testing.T) {
+	cli := testingcli.New()
+
+	executed := false
+	q := testingkusto.NewQueryClient(func(qc *testingkusto.QueryClient) {
+		qc.MgmtFn = func(_ context.Context, db string, stmt kusto.Statement, _ ...kusto.QueryOption) (*kusto.RowIterator, error) {
+			executed = true
+			assert.Equal(t, "TestDatabase", db)
+			return nil, nil
+		}
+	})
+
+	opts := ManagementOptions{
+		Source:      []byte(".show tables"),
+		Auth:        newTestAuth(),
+		KustoTarget: newTestKustoTarget(),
+		ingestorBuildSettings: ingestorBuildSettings{
+			CreateQueryClient: func(target KustoTargetOptions, auth AuthOptions) (ingest.QueryClient, error) {
+				return q, nil
+			},
+		},
+	}
+
+	err := opts.Run(cli)
+	assert.NoError(t, err)
+	assert.True(t, executed, "expected management command to execute")
+}
+
+func Test_ManagementOptions_Run_CreateClientError(t *testing.T) {
+	cli := testingcli.New()
+
+	opts := ManagementOptions{
+		Source:      []byte(".show tables"),
+		Auth:        newTestAuth(),
+		KustoTarget: newTestKustoTarget(),
+		ingestorBuildSettings: ingestorBuildSettings{
+			CreateQueryClient: func(target KustoTargetOptions, auth AuthOptions) (ingest.QueryClient, error) {
+				return nil, errors.New("boom")
+			},
+		},
+	}
+
+	err := opts.Run(cli)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "create Kusto query client")
+}
+
+func Test_ManagementOptions_Run_MgmtError(t *testing.T) {
+	cli := testingcli.New()
+
+	q := testingkusto.NewQueryClient(func(qc *testingkusto.QueryClient) {
+		qc.MgmtFn = func(_ context.Context, db string, stmt kusto.Statement, _ ...kusto.QueryOption) (*kusto.RowIterator, error) {
+			return nil, errors.New("mgmt failed")
+		}
+	})
+
+	opts := ManagementOptions{
+		Source:      []byte(".show tables"),
+		Auth:        newTestAuth(),
+		KustoTarget: newTestKustoTarget(),
+		ingestorBuildSettings: ingestorBuildSettings{
+			CreateQueryClient: func(target KustoTargetOptions, auth AuthOptions) (ingest.QueryClient, error) {
+				return q, nil
+			},
+		},
+	}
+
+	err := opts.Run(cli)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "execute management commands")
+}

internal/kusto/options.go

@@ -31,3 +31,14 @@ type FileIngestOptions struct {
 	// for unit test
 	ingestorBuildSettings `kong:"-"`
 }
+
+// ManagementOptions provides the configuration for management commands.
+type ManagementOptions struct {
+	Source []byte `arg:"" type:"filecontent" required:"" help:"The source file to execute."`
+
+	Auth        AuthOptions        `embed:"" prefix:"auth-"`
+	KustoTarget KustoTargetOptions `embed:"" prefix:"kusto-"`
+
+	// for unit test
+	ingestorBuildSettings `kong:"-"`
+}

internal/kusto/testingkusto/query_client.go

@@ -0,0 +1,72 @@
+package testingkusto
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/Azure/azure-kusto-go/kusto"
+	"github.com/Azure/azure-kusto-go/kusto/ingest"
+)
+
+// QueryClient is a fake implementation of ingest.QueryClient for tests.
+type QueryClient struct {
+	MgmtFn  func(ctx context.Context, db string, query kusto.Statement, options ...kusto.QueryOption) (*kusto.RowIterator, error)
+	QueryFn func(ctx context.Context, db string, query kusto.Statement, options ...kusto.QueryOption) (*kusto.RowIterator, error)
+	CloseFn func() error
+
+	// Captured calls for assertions.
+	MgmtCalls  []MgmtCall
+	QueryCalls []QueryCall
+}
+
+// MgmtCall captures a Mgmt invocation.
+type MgmtCall struct {
+	DB    string
+	Query kusto.Statement
+}
+
+// QueryCall captures a Query invocation.
+type QueryCall struct {
+	DB    string
+	Query kusto.Statement
+}
+
+var _ ingest.QueryClient = (*QueryClient)(nil)
+
+func (qc *QueryClient) Close() error {
+	if qc.CloseFn != nil {
+		return qc.CloseFn()
+	}
+	return nil
+}
+
+func (qc *QueryClient) Auth() kusto.Authorization { return kusto.Authorization{} }
+func (qc *QueryClient) Endpoint() string          { return "https://example.kusto.windows.net" }
+
+func (qc *QueryClient) Query(ctx context.Context, db string, query kusto.Statement, options ...kusto.QueryOption) (*kusto.RowIterator, error) {
+	qc.QueryCalls = append(qc.QueryCalls, QueryCall{DB: db, Query: query})
+	if qc.QueryFn != nil {
+		return qc.QueryFn(ctx, db, query, options...)
+	}
+	return nil, nil
+}
+
+func (qc *QueryClient) Mgmt(ctx context.Context, db string, query kusto.Statement, options ...kusto.QueryOption) (*kusto.RowIterator, error) {
+	qc.MgmtCalls = append(qc.MgmtCalls, MgmtCall{DB: db, Query: query})
+	if qc.MgmtFn != nil {
+		return qc.MgmtFn(ctx, db, query, options...)
+	}
+	return nil, nil
+}
+
+func (qc *QueryClient) HttpClient() *http.Client            { return nil }
+func (qc *QueryClient) ClientDetails() *kusto.ClientDetails { return &kusto.ClientDetails{} }
+
+// NewQueryClient creates a QueryClient with optional mutators.
+func NewQueryClient(ms ...func(*QueryClient)) *QueryClient {
+	qc := &QueryClient{}
+	for _, m := range ms {
+		m(qc)
+	}
+	return qc
+}