Forem: Arash Outadi

Ensure auto-generated code is always up-to-date with compile-time assertions in Go

Arash Outadi — Thu, 30 Jul 2020 17:58:34 +0000

Introduction

Do you get annoyed when others (or you) forget to run the code generation command and now nothing works/code is out-dated?

Compile Time Assertions to the rescue!

Scenario

Suppose you have an interface:

type MyInterface interface {
    WithMethod(method string) MyInterface
}

This interface has an associated implementation:

type MyInterfaceImpl struct {
    privateField string
}
func (my *MyInterfaceImpl) MyInterface {
    // ... implementation code
}

Which you've have been auto-generating mock objects for using the awesome mockery package.

Why Though

You are doing this because you want users of your library to be able to mock your library.

Thus allowing them to isolate their own code when writing unit tests.
Meaning we don't actually use this mock in OUR tests, it is used by OTHERS

So the interface definition above is augmented with:

// Generate mocks by running "go generate ./..."
// go:generate mockery --name MyInteface
type MyInterface interface {
    WithMethod(method string) MyInterface
}

Which produces when go generate ./... is ran:
mocks/MyInterface.go

// MyInterface is an autogenerated mock type for the MyInterface type
type MyInterface struct {
    mock.Mock
}

// WithMethod provides a mock function with given fields: method
func (_m *MyInterface) WithMethod(method string) (pkg.MyInterface) {
// ... some mock implementation code

Problem

Now what happens when Greg who isn't aware of the auto-generated code, adds another method to MyInterface and MyInterfaceImpl:

type MyInterface interface {
    WithMethod(method string) MyInterface
    WithAnotherMethod(method string) MyInterface // New
}

Since the MyInterface has changed but Greg didn't auto-generate the mocks using go generate ./..., the mock object will not fulfill the interface and cannot be used for testing by downstream users (Thanks Greg)!

Solution

Add a compile-assertion above the interface definition after you have initially generated the mock object (Also add --inpackage flag to go generate command to avoid import cycles).

// Generate mocks by running "go generate ./..."
// go:generate mockery --name MyInteface --inpackage
var _ MyInterface = &MockMyInterface{} // COMPILE-TIME Assertion
type MyInterface interface {
    WithMethod(method string) MyInterface
}

Now if Greg adds another method but forgets to run go generate

var _ MyInterface = &MockMyInterface{}
type MyInterface interface {
    WithMethod(method string) MyInterface
    WithAnotherMethod(method string) MyInterface // New
}

He will get hit with:

cannot use &MockMyInterface literal (type *MockMyInterface) as type MyInterface in assignment:
    *MockMyInterface does not implement MyInterface (missing WithAnotherMethod method)

At which point, he might realize that he needs to add the method to the mock object using go generate.

Caveat: You will need to comment out the compile-time assertion var _ MyInterface = &MockMyInterface{} when running go generate/mockery tool

Adding Contexts via Go AST (Code Instrumentation)

Arash Outadi — Wed, 25 Dec 2019 06:54:42 +0000

Problem

You changed one function to require a ctx context.Context and now you have to change the function signatures of the
all the upstream functions in your codebase 😠!

Do you have to make all these changes manually or can you automate the process somehow?

We can automate it using Go's AST.

Demo

Check out this playground link to see how to source code can be programmatically changed to include the proper arguments/parameters: PLAYGROUND
We'll describe how it all works throughout the tutorial

Illustration

We want to change this function to what is commented out in the TODO:

And imagine the function sits inside the fictitious file below, with many other functions calling it.

Again the TODO:s indicate what needs to change to make this file compile.

NOTE: The other functions are not important, only there to illustrate we have to change all the functions that call our changed function

Fictitious File Notice all the TODOs

Manual Solution

Fixing this manually, usually involves a trial and error solution process:

Adjust the function body to correctly call your new function with ctx arguments
Adjust the function definition/signature to include a ctx context.Context parameter
Iterate until the compiler no longer complains

Programmatic Solution (Pseudo-Code)

The algorithm for the manual solution is very simple, so on a high level to automate this process we can:

Parse the Go code
Generate an AST(Abstract Syntax Tree) [3]
Programmatically determine where we need "inject" code, specifically:
- ctx argument to a function call
- ctx context.Context function parameter to function declaration
Edit the AST
Iterate until there are no more places where an "injection" is required
Convert the AST back into the text representation -> Our new Golang source code

Pre-Requisites

In this tutorial I assume that:

You are proficient with Go
Somewhat familiar with what an AST is
- If you want to learn more about Go's AST, I recommend this post from Eno Compton[7]

Tutorial Conventions

Playground Links

After a code example, I will provide a full working example via the Go Playground see you can run the code for yourself.
Look out for PLAYGROUND

Brevity

I will often only include the minimal amount of code to demonstrate a new concept in the code examples and will generally cut out any boiler-plate code.
Look for the Playground links for full working examples.

Setup

Libraries

Note in this tutorial we will be using these libraries:

github.com/dave/dst (Alternative to go/ast)

You can download them with the go get command:

go get github.com/dave/dst

github.com/dave/dst is a fork of the official go/ast package that is meant specifically for instrumenting go code.

In contrast go/ast was primarily meant for code generation.

This is an important difference because in code instrumentation, we only want to change a very specific region of code and leave the rest of the AST exactly as it was. go/ast has a difficulties achieving this, especially with comments [4]

Code

Step 0 - Visualize the AST

Using this go-ast visualizer, we can get an idea of what the Go AST looks like and what we need to look for when injecting new code [5].

So go to http://goast.yuroyoro.net/ and play around with these above gists

Our primary focus should be on the on the FuncDecl and CallExpr Nodes since our injection points will be either when we are:

Defining a function <-- Might need to add a ctx context.Context parameter
Calling a function <-- Might need to add a ctx argument

Step 1 - Use "dst" to parse Go code

Before we start jumping into the logic of the program, let's just see a quick demo of how we parse Go code using the "dave/dst" package and print out the AST representation of the FuncDecl nodes.

PLAYGROUND

I've left comments in the above code snippet, so be sure to read those before moving on.

Step 2 - Helper Functions

Remember that we either need to

Add a ctx context.Context parameter to a function declaration (FuncDecl Node)
Add a ctx argument to a function call (CallExpr) The Go AST structs for these two actions can be defined as follows:

Step 3 - Editing the AST

To demonstrate how to add arguments and parameters using our helper functions, observe the following "naive" applyFunc that adds an additional ctx argument and ctx context.Context parameter to every function call and function definition respectively.

Let's use our applyFunc and print out the source code! Notice I added another utility function for converting the AST representation into actual Go code.

PLAYGROUND

Notice that this applyFunc doesn't check if it is actually necessary to inject additional code. It just does it.

To see why this applyFunc is not adequate try running changing the srcCodeString to:

package test

import (
    "fmt"
    "context"
)

func alreadyHasContext(ctx context.Context) {
    fmt.Println("Do some important work...")
    makeDownstreamRequest(ctx, "Some important data!") 
}

Or if you prefer go to playground link below.

PLAYGROUND

Step 4 - Being Selective

Instead of indiscriminately adding ctx everywhere like in the naive implementation, this time we will examine the nodes in the AST to determine where we need to inject a ctx argument or ctx content.Context function parameter.

FuncDecl Node

Let's look at the FuncDecl, node for the following code.

Function Signature

Specifically let's start the with FuncType, which essentially describes the function signature

The feature that we care about most is SelectorExpr that contains the Ident for the Context parameter. With this in mind, we can construct a function to check if the FuncDecl contains a context.Context as a parameter in the function signature

If it already has a context.Context then we don't need to do anything and can move on.

Function Body

To determine if we need to add a context.Context to the function signature, we need to examine the function body to check if there are any calls to functions which require a context.Context parameter.

At this point, you might be asking how we can know if a function requires a Context in the first place?

Below are two possible ways of determining that:

If there is a function call to a function we have already determined needs a Context (Infinite Recursion 😆)
- Obviously, we would need a mechanism for recording which functions have Context parameters
- Also we need a "seed" function which already has Context parameter, so we have at-least one function which we KNOW requires Context
We can naively look for function calls that have a argument named ctx

In this tutorial, we'll go over method 2 as it slightly simpler and doesn't require an initial scan.

Functions To Examine The FuncDecl & CallExpr

Note: the global map needsContextsFuncs (However, we use the map as a Set)

Step 5 - Selective ApplyFunc

To make the our ApplyFunc function more selective we will have to examine the FuncDecl and CallExpr nodes

Selective Apply Function

Step 6 - Is this it?

So now we have a Selective ApplyFunc will it be able to add all the Contexts now?

PLAYGROUND

Spoilers... It doesn't

Step 7 - Iterative Solution

Our new ApplyFunc only manages to change the immediate ancestors, but if want the changes to propagate up further we need an iterative solution.
Remember Step 5 of our pseudo code algorithm

Infinite For Loop?

To make sure that all the ancestors are updated we could simply run the Apply with our ApplyFunc over and over again inside a for {} loop, but when should we stop?
A simple idea that will work is to stop when the previous code is the same as the current code generated.

That is when the ApplyFunc deems that there are no new areas to add ctx arguments or ctx context.Context parameters.

Code

PLAYGROUND

Optional - Add comments describing iteration

Lastly, as a illustrative exercise, below is a playground link for adding comments that describe in which iteration the ctx/ctx context.Context was added

PLAYGROUND

Output:

package test

import (
    "context"
    "fmt"
)

// Added on 'ctx context.Context' parameter on iteration 0
func changedFn(ctx context.Context) {
    fmt.Println("Do some important work...")
    // Now also make a downstream call
    makeDownstreamRequest(ctx, "Some important data!")
}

// Added on 'ctx context.Context' parameter on iteration 1
func needsctx1(ctx context.Context, n int) {
    if true {
        changedFn(ctx) // Added on ctx arg on iteration 0
    }
}

// Added on 'ctx context.Context' parameter on iteration 2
func needsctx2(ctx context.Context) bool {
    for index := 0; index < 3; index++ {
        needsctx1(ctx, 1) // Added on ctx arg on iteration 1
    }
    return true
}

// Added on 'ctx context.Context' parameter on iteration 1
func needsctx3(ctx context.Context) {
    if needsctx2(ctx) { // Added on ctx arg on iteration 2

        changedFn(ctx) // Added on ctx arg on iteration 0
    }
}

type SS struct{}

// Added on 'ctx context.Context' parameter on iteration 2
func (rec *SS) save(ctx context.Context, s string, n int) {
    needsctx1(ctx, 2) // Added on ctx arg on iteration 1
}

Future Work

It's not to hard to extend this to changing the AST of a file via ParseFile and entire directories with ParseDir

If there is any interest, I'm happy to make a second tutorial covering these topics and additionally how to handle imported packages.

Thanks

Thanks for reading this tutorial about using Go's AST I hope you found it useful.

Let me know if you liked it or hated it!

Also huge thanks to Dave for making the awesome dst package!

As well as all the other sources I used to make this tutorial

Sources

https://faiface.github.io/post/context-should-go-away-go2/
https://medium.com/@cep21/how-to-correctly-use-context-context-in-go-1-7-8f2c0fafdf39
https://en.wikipedia.org/wiki/Abstract_syntax_tree
https://github.com/golang/go/issues/20744 "Free-floating comments are single-biggest issue when manipulating the AST"
http://goast.yuroyoro.net/ Golang AST visualizer
https://godoc.org/github.com/dave/dst GoDoc for github.com/dave/dst package
https://commandercoriander.net/blog/2016/12/30/reading-go-ast/

Wrapper functions for dangerous commands

Arash Outadi — Fri, 20 Dec 2019 04:37:25 +0000

Have you ever had a brain fart at work and accidentally deleted the PRODUCTION database because you thought you were in the development environment? Yeah, me too...

Here's a simple bash script that might save your life (multiple times).

Let me draw the scenario:

Your company has different environments (prod, dev, etc...)
It's possible to accidentally run a bash command meant for development environment in a production environment, which will have fatal consequences

Below is a fictitious wrapper function for a dangerous_cmd only meant to be used in a development environment.

function dangerous_cmd(){
    echo "Running dangerous_cmd wrapper"
    # Assume that 'dangerous_cmd environment' returns the environment you are currently in
    env=$(dangerous_cmd environment) 
    if [[ $env == "development" ]]; then
        # 'command' is a bash built-in that avoids shell function lookup
        # We need this because otherwise we would be stuck in a infinite loop,
        # calling 'dangeroud_cmd` over and over again
        # '$@' passes all the arguments passed to the function to this call
        command dangerous_cmd $@; 
        return
    fi
    echo "Not in development environment! Not running dangerous_cmd"
}

The main workhorse of of this function is the bash built-ins command and $@

command(StackOverflow Description) runs our command without doing a shell function look-up, which we need to avoid infinitely calling dangerous_cmd
$@ passes all the parameters to this point. Useful illustration

The good thing about using a bash function wrapper with the same name is that:

Autocomplete still works as usual
You can add more complex checking logic as you see fit
You don't have to remember any obscure commands

I put bash functions like these in my .bashrc or bash_profile (Mac)

Here is an example of a wrapper function that only allows running something like skaffold delete (A command that cleans up your development environment by deleting all pods and services related to your application) in a minikube (Basically development) Kubernetes environment.
Trust me you don't want to run skaffold delete in production.

function skaffold(){
    echo "Running skaffold wrapper"
    context=$(kubectl config current-context)
    if [[ $context == "minikube" ]]; then
        command skaffold $@;
        return
    fi
    echo "Not in minikube context! Not running skaffold command"
}

Forem: Arash Outadi

Ensure auto-generated code is always up-to-date with compile-time assertions in Go

Introduction

Scenario

Why Though

Problem

Solution

Adding Contexts via Go AST (Code Instrumentation)

Problem

Demo

Illustration

Manual Solution

Programmatic Solution (Pseudo-Code)

Pre-Requisites

Tutorial Conventions

Playground Links

Brevity

Setup

Libraries

Code

Step 0 - Visualize the AST

Step 1 - Use "dst" to parse Go code

Step 2 - Helper Functions

Step 3 - Editing the AST

Step 4 - Being Selective

FuncDecl Node

Function Signature

Function Body

Step 5 - Selective ApplyFunc

Step 6 - Is this it?

Step 7 - Iterative Solution

Infinite For Loop?

Code

Optional - Add comments describing iteration

Future Work

Thanks

Related Reading

Sources

Wrapper functions for dangerous commands